Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix integer overflow in BMP decoder by making the checks in DecodeBmp
more stringent.  Add fuzzer to improve the robustness of the decoder
in the future.

PiperOrigin-RevId: 185780111
  • Loading branch information
David G. Andersen authored and tensorflower-gardener committed Feb 15, 2018
1 parent 109d7af commit 49f73c5
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 5 deletions.
27 changes: 22 additions & 5 deletions tensorflow/core/kernels/decode_bmp_op.cc
Expand Up @@ -91,15 +91,32 @@ class DecodeBmpOp : public OpKernel {
errors::InvalidArgument(
"Number of channels must be 1, 3 or 4, was ", channels_));

OP_REQUIRES(context, width > 0 && header_size >= 0,
errors::InvalidArgument("Width must be positive"));
OP_REQUIRES(context, header_size >= 0,
errors::InvalidArgument("header size must be nonnegative"));

// The real requirement is < 2^31 minus some headers and channel data,
// so rounding down to something that's still ridiculously big.
OP_REQUIRES(
context,
(static_cast<int64>(width) * std::abs(static_cast<int64>(height))) <
static_cast<int64>(std::numeric_limits<int32_t>::max() / 8),
errors::InvalidArgument(
"Total possible pixel bytes must be less than 2^30"));

const int32 abs_height = abs(height);

// 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_;
const int64 last_pixel_offset = static_cast<int64>(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_;
const int64 expected_file_size = last_pixel_offset + channels_;

OP_REQUIRES(
context, (expected_file_size <= input.size()),
Expand All @@ -115,12 +132,12 @@ class DecodeBmpOp : public OpKernel {
Tensor* output = nullptr;
OP_REQUIRES_OK(
context, context->allocate_output(
0, TensorShape({abs(height), width, channels_}), &output));
0, TensorShape({abs_height, width, channels_}), &output));

const uint8* bmp_pixels = &img_bytes[header_size];

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

uint8* Decode(const uint8* input, const int row_size, uint8* const output,
Expand Down
2 changes: 2 additions & 0 deletions tensorflow/core/kernels/fuzzing/BUILD
Expand Up @@ -43,6 +43,8 @@ tf_ops_fuzz_target_lib("decode_base64")

tf_ops_fuzz_target_lib("encode_jpeg")

tf_ops_fuzz_target_lib("decode_bmp")

tf_ops_fuzz_target_lib("decode_png")

tf_ops_fuzz_target_lib("decode_jpeg")
Expand Down
29 changes: 29 additions & 0 deletions tensorflow/core/kernels/fuzzing/decode_bmp_fuzz.cc
@@ -0,0 +1,29 @@
/* Copyright 2018 Google Inc. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/kernels/fuzzing/fuzz_session.h"

namespace tensorflow {
namespace fuzzing {

class FuzzDecodeBmp : public FuzzStringInputOp {
SINGLE_INPUT_OP_BUILDER(DT_STRING, DecodeBmp);
};

STANDARD_TF_FUZZ_FUNCTION(FuzzDecodeBmp);

} // end namespace fuzzing
} // end namespace tensorflow

0 comments on commit 49f73c5

Please sign in to comment.