Change jpeg decoding to record the number of rows processed, instead of the

fraction processed. Also avoid one case of int overflow.
Change: 117774225

1 files changed, 15 insertions, 12 deletions

diff --git a/tensorflow/core/lib/jpeg/jpeg_mem.cc b/tensorflow/core/lib/jpeg/jpeg_mem.cc
index e67779060d..2ea5e04055 100644
--- a/tensorflow/core/lib/jpeg/jpeg_mem.cc
+++ b/tensorflow/core/lib/jpeg/jpeg_mem.cc
@@ -53,7 +53,7 @@ class FewerArgsForCompiler {
         flags_(flags),
         pnwarn_(nwarn),
         allocate_output_(allocate_output),
-        fraction_read_(0.),
+        height_read_(0),
         height_(0),
         stride_(0) {
     if (pnwarn_ != nullptr) *pnwarn_ = 0;
@@ -63,7 +63,7 @@ class FewerArgsForCompiler {
   const UncompressFlags flags_;
   int64* const pnwarn_;
   std::function<uint8*(int, int, int)> allocate_output_;
-  float fraction_read_;  // fraction of scanline lines successfully read
+  int height_read_;  // number of scanline lines successfully read
   int height_;
   int stride_;
 };
@@ -184,7 +184,7 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
 
   // If there is an error reading a line, this aborts the reading.
   // Save the fraction of the image that has been read.
-  argball->fraction_read_ = 1.0;
+  argball->height_read_ = cinfo.output_height;
   while (cinfo.output_scanline < cinfo.output_height) {
     int num_lines_read = 0;
     if (cinfo.out_color_space == JCS_CMYK) {
@@ -217,8 +217,7 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
       LOG(ERROR) << "Premature end of JPEG data. Stopped at line "
                  << cinfo.output_scanline << "/" << cinfo.output_height;
       if (!flags.try_recover_truncated_jpeg) {
-        argball->fraction_read_ =
-            static_cast<float>(cinfo.output_scanline) / cinfo.output_height;
+        argball->height_read_ = cinfo.output_scanline;
         error = JPEGERRORS_UNEXPECTED_END_OF_DATA;
       } else {
         for (size_t line = cinfo.output_scanline; line < cinfo.output_height;
@@ -232,7 +231,8 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
           }
           output_line += stride;
         }
-        argball->fraction_read_ = 1.0;  // consider all lines as read
+        argball->height_read_ =
+            cinfo.output_height;  // consider all lines as read
         // prevent error-on-exit in libjpeg:
         cinfo.output_scanline = cinfo.output_height;
       }
@@ -249,8 +249,8 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
   // RGBRGBRGB... --> RGBARGBARGBA...
   if (components == 4) {
     // Start on the last line.
-    JSAMPLE* scanlineptr =
-        static_cast<JSAMPLE*>(dstdata + (cinfo.output_height - 1) * stride);
+    JSAMPLE* scanlineptr = static_cast<JSAMPLE*>(
+        dstdata + static_cast<int64>(cinfo.output_height - 1) * stride);
     const JSAMPLE kOpaque = -1;  // All ones appropriate for JSAMPLE.
     const int right_rgb = (cinfo.output_width - 1) * 3;
     const int right_rgba = (cinfo.output_width - 1) * 4;
@@ -335,7 +335,11 @@ uint8* Uncompress(const void* srcdata, int datasize,
                   std::function<uint8*(int, int, int)> allocate_output) {
   FewerArgsForCompiler argball(datasize, flags, nwarn, allocate_output);
   uint8* const dstdata = UncompressLow(srcdata, &argball);
-  const float fraction_read = argball.fraction_read_;
+
+  const float fraction_read =
+      argball.height_ == 0
+          ? 1.0
+          : (static_cast<float>(argball.height_read_) / argball.height_);
   if (dstdata == NULL ||
       fraction_read < std::min(1.0f, flags.min_acceptable_fraction)) {
     // Major failure, none or too-partial read returned; get out
@@ -344,9 +348,8 @@ uint8* Uncompress(const void* srcdata, int datasize,
 
   // If there was an error in reading the jpeg data,
   // set the unread pixels to black
-  if (fraction_read < 1.0) {
-    const int first_bad_line =
-        static_cast<int>(fraction_read * argball.height_);
+  if (argball.height_read_ != argball.height_) {
+    const int first_bad_line = argball.height_read_;
     uint8* start = dstdata + first_bad_line * argball.stride_;
     const int nbytes = (argball.height_ - first_bad_line) * argball.stride_;
     memset(static_cast<void*>(start), 0, nbytes);