这是indexloc提供的服务，不要输入任何密码
Skip to content

Fix decode_bmp crash by adding length check before reading the data in buffer #14967

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Nov 30, 2017
Merged

Fix decode_bmp crash by adding length check before reading the data in buffer #14967

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a046700
Fix decode_bmp crash by adding length check before reading the data i…
yongtang Nov 29, 2017
d3c891a
Sanitize decode_bmp_op.cc with clang-format -i --style=Google
yongtang Nov 29, 2017
d599946
Add additional check to make sure header is safe to access in bmp
yongtang Nov 29, 2017
60a7220
Not require padding (as paddings are not accessed)
yongtang Nov 29, 2017
15156e0
Remove duplicated row_size calculation.
yongtang Nov 29, 2017
36238d3
Add test cases for decoding incomplete bmp
yongtang Nov 29, 2017
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 36 additions & 15 deletions tensorflow/core/kernels/decode_bmp_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,10 +33,11 @@ class DecodeBmpOp : public OpKernel {
public:
explicit DecodeBmpOp(OpKernelConstruction* context) : OpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("channels", &channels_));
OP_REQUIRES(context, channels_ == 0 || channels_ == 1 || channels_ == 3 ||
channels_ == 4,
errors::InvalidArgument("channels must be 0, 1, 3 or 4, got ",
channels_));
OP_REQUIRES(
context,
channels_ == 0 || channels_ == 1 || channels_ == 3 || channels_ == 4,
errors::InvalidArgument("channels must be 0, 1, 3 or 4, got ",
channels_));
}

void Compute(OpKernelContext* context) override {
Expand All @@ -48,6 +49,12 @@ class DecodeBmpOp : public OpKernel {
// Start decoding image to get shape details
const StringPiece input = contents.scalar<string>()();

OP_REQUIRES(context, (32 <= input.size()),
errors::InvalidArgument("Incomplete bmp content, requires at "
"least 32 bytes to find the header "
"size, width, height, and bpp, got ",
input.size(), " bytes"));

const uint8* img_bytes = reinterpret_cast<const uint8*>(input.data());
const int32 header_size = internal::SubtleMustCopy(
*(reinterpret_cast<const int32*>(img_bytes + 10)));
Expand All @@ -73,6 +80,22 @@ class DecodeBmpOp : public OpKernel {
errors::InvalidArgument(
"Number of channels must be 1, 3 or 4, was ", channels_));

// there may be padding bytes when the width is not a multiple of 4 bytes
// 8 * channels == bits per pixel
const int row_size = (8 * channels_ * width + 31) / 32 * 4;

const int last_pixel_offset =
header_size + (abs(height) - 1) * row_size + (width - 1) * channels_;

// [expected file size] = [last pixel offset] + [last pixel size=channels]
const int expected_file_size = last_pixel_offset + channels_;

OP_REQUIRES(
context, (expected_file_size <= input.size()),
errors::InvalidArgument("Incomplete bmp content, requires at least ",
expected_file_size, " bytes, got ",
input.size(), " bytes"));

// if height is negative, data layout is top down
// otherwise, it's bottom up
bool top_down = (height < 0);
Expand All @@ -85,25 +108,23 @@ class DecodeBmpOp : public OpKernel {

const uint8* bmp_pixels = &img_bytes[header_size];

Decode(bmp_pixels, output->flat<uint8>().data(), width, abs(height),
channels_, top_down);
Decode(bmp_pixels, row_size, output->flat<uint8>().data(), width,
abs(height), channels_, top_down);
}

uint8* Decode(const uint8* input, uint8* const output, const int width,
const int height, const int channles, bool top_down);
uint8* Decode(const uint8* input, const int row_size, uint8* const output,
const int width, const int height, const int channles,
bool top_down);

private:
int channels_;
};
REGISTER_KERNEL_BUILDER(Name("DecodeBmp").Device(DEVICE_CPU), DecodeBmpOp);

uint8* DecodeBmpOp::Decode(const uint8* input, uint8* const output,
const int width, const int height,
const int channels, bool top_down) {
// there may be padding bytes when the width is not a multiple of 4 bytes
// 8 * channels == bits per pixel
int row_size = (8 * channels * width + 31) / 32 * 4;

uint8* DecodeBmpOp::Decode(const uint8* input, const int row_size,
uint8* const output, const int width,
const int height, const int channels,
bool top_down) {
for (int i = 0; i < height; i++) {
int src_pos;
int dst_pos;
Expand Down
50 changes: 50 additions & 0 deletions tensorflow/python/kernel_tests/decode_bmp_op_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import image_ops
from tensorflow.python.platform import test
Expand Down Expand Up @@ -99,5 +100,54 @@ def testGrayscale(self):
decoded = decode.eval()
self.assertAllEqual(decoded, img_bytes)

def testIncompleteHeader(self):
# Encoded BMP bytes from Wikipedia
encoded_bytes = [
0x42, 0x40,
0x46, 0, 0, 0,
]

byte_string = bytes(bytearray(encoded_bytes))
img_in = constant_op.constant(byte_string, dtype=dtypes.string)
decode = array_ops.squeeze(image_ops.decode_bmp(img_in))

with self.test_session():
with self.assertRaisesRegexp(errors_impl.InvalidArgumentError,
"requires at least 32 bytes to find the header"):
decoded = decode.eval()

def testIncompleteBody(self):
# Encoded BMP bytes from Wikipedia
encoded_bytes = [
0x42, 0x40,
0x46, 0, 0, 0,
0, 0,
0, 0,
0x36, 0, 0, 0,
0x28, 0, 0, 0,
0x2, 0, 0, 0,
0x2, 0, 0, 0,
0x1, 0,
0x18, 0,
0, 0, 0, 0,
0x10, 0, 0, 0,
0x13, 0xb, 0, 0,
0x13, 0xb, 0, 0,
0, 0, 0, 0,
0, 0, 0, 0,
0, 0, 0xff,
0xff, 0xff, 0xff,
0, 0,
]

byte_string = bytes(bytearray(encoded_bytes))
img_in = constant_op.constant(byte_string, dtype=dtypes.string)
decode = array_ops.squeeze(image_ops.decode_bmp(img_in))

with self.test_session():
with self.assertRaisesRegexp(errors_impl.InvalidArgumentError,
"requires at least 68 bytes, got 62 bytes"):
decoded = decode.eval()

if __name__ == "__main__":
test.main()