Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent memory leak in decoding PNG images.
PiperOrigin-RevId: 409300653
Change-Id: I6182124c545989cef80cefd439b659095920763b
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Nov 12, 2021
1 parent fb5ce99 commit ab51e5b
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions tensorflow/core/kernels/image/decode_image_op.cc
Expand Up @@ -18,6 +18,8 @@ limitations under the License.
#include <cstdint>
#include <memory>

#include "tensorflow/core/lib/gtl/cleanup.h"

#define EIGEN_USE_THREADS

#include "absl/strings/escaping.h"
Expand Down Expand Up @@ -326,6 +328,16 @@ class DecodeImageV2Op : public OpKernel {
context, png::CommonInitDecode(input, channels_, channel_bits, &decode),
errors::InvalidArgument("Invalid PNG. Failed to initialize decoder."));

// If we reach this point, then there is data in `decode` which must be
// freed by the time we end execution in this function. We cannot call
// `png::CommonFreeDecode()` before an `OP_REQUIRES` because if
// `OP_REQUIRES` constraint is satisfied then the data would be freed
// prematurely. Instead, let's use a `Cleanup` object.
auto cleanup = gtl::MakeCleanup([&decode]() {
std::cerr << "Cleanup called...\n";
png::CommonFreeDecode(&decode);
});

// Verify that width and height are not too large:
// - verify width and height don't overflow int.
// - width can later be multiplied by channels_ and sizeof(uint16), so
Expand Down

0 comments on commit ab51e5b

Please sign in to comment.