Merge "ultrahdr: compress image reads outside bounds for non aligned widths" into udc-dev

3 files changed, 69 insertions, 22 deletions

diff --git a/libs/ultrahdr/jpegencoderhelper.cpp b/libs/ultrahdr/jpegencoderhelper.cpp
index 10a763035f..ab2f8c7b5a 100644
--- a/libs/ultrahdr/jpegencoderhelper.cpp
+++ b/libs/ultrahdr/jpegencoderhelper.cpp
@@ -22,6 +22,8 @@
 
 namespace android::ultrahdr {
 
+#define ALIGNM(x, m)  ((((x) + ((m) - 1)) / (m)) * (m))
+
 // The destination manager that can access |mResultBuffer| in JpegEncoderHelper.
 struct destination_mgr {
 public:
@@ -175,6 +177,37 @@ bool JpegEncoderHelper::compressYuv(jpeg_compress_struct* cinfo, const uint8_t*
     std::unique_ptr<uint8_t[]> empty(new uint8_t[cinfo->image_width]);
     memset(empty.get(), 0, cinfo->image_width);
 
+    const int aligned_width = ALIGNM(cinfo->image_width, kCompressBatchSize);
+    const bool is_width_aligned = (aligned_width == cinfo->image_width);
+    std::unique_ptr<uint8_t[]> buffer_intrm = nullptr;
+    uint8_t* y_plane_intrm = nullptr;
+    uint8_t* u_plane_intrm = nullptr;
+    uint8_t* v_plane_intrm = nullptr;
+    JSAMPROW y_intrm[kCompressBatchSize];
+    JSAMPROW cb_intrm[kCompressBatchSize / 2];
+    JSAMPROW cr_intrm[kCompressBatchSize / 2];
+    JSAMPARRAY planes_intrm[3]{y_intrm, cb_intrm, cr_intrm};
+    if (!is_width_aligned) {
+        size_t mcu_row_size = aligned_width * kCompressBatchSize * 3 / 2;
+        buffer_intrm = std::make_unique<uint8_t[]>(mcu_row_size);
+        y_plane_intrm = buffer_intrm.get();
+        u_plane_intrm = y_plane_intrm + (aligned_width * kCompressBatchSize);
+        v_plane_intrm = u_plane_intrm + (aligned_width * kCompressBatchSize) / 4;
+        for (int i = 0; i < kCompressBatchSize; ++i) {
+            y_intrm[i] = y_plane_intrm + i * aligned_width;
+            memset(y_intrm[i] + cinfo->image_width, 0, aligned_width - cinfo->image_width);
+        }
+        for (int i = 0; i < kCompressBatchSize / 2; ++i) {
+            int offset_intrm = i * (aligned_width / 2);
+            cb_intrm[i] = u_plane_intrm + offset_intrm;
+            cr_intrm[i] = v_plane_intrm + offset_intrm;
+            memset(cb_intrm[i] + cinfo->image_width / 2, 0,
+                   (aligned_width - cinfo->image_width) / 2);
+            memset(cr_intrm[i] + cinfo->image_width / 2, 0,
+                   (aligned_width - cinfo->image_width) / 2);
+        }
+    }
+
     while (cinfo->next_scanline < cinfo->image_height) {
         for (int i = 0; i < kCompressBatchSize; ++i) {
             size_t scanline = cinfo->next_scanline + i;
@@ -183,6 +216,9 @@ bool JpegEncoderHelper::compressYuv(jpeg_compress_struct* cinfo, const uint8_t*
             } else {
                 y[i] = empty.get();
             }
+            if (!is_width_aligned) {
+                memcpy(y_intrm[i], y[i], cinfo->image_width);
+            }
         }
         // cb, cr only have half scanlines
         for (int i = 0; i < kCompressBatchSize / 2; ++i) {
@@ -194,9 +230,13 @@ bool JpegEncoderHelper::compressYuv(jpeg_compress_struct* cinfo, const uint8_t*
             } else {
                 cb[i] = cr[i] = empty.get();
             }
+            if (!is_width_aligned) {
+                memcpy(cb_intrm[i], cb[i], cinfo->image_width / 2);
+                memcpy(cr_intrm[i], cr[i], cinfo->image_width / 2);
+            }
         }
-
-        int processed = jpeg_write_raw_data(cinfo, planes, kCompressBatchSize);
+        int processed = jpeg_write_raw_data(cinfo, is_width_aligned ? planes : planes_intrm,
+                                            kCompressBatchSize);
         if (processed != kCompressBatchSize) {
             ALOGE("Number of processed lines does not equal input lines.");
             return false;
@@ -213,6 +253,23 @@ bool JpegEncoderHelper::compressSingleChannel(jpeg_compress_struct* cinfo, const
     std::unique_ptr<uint8_t[]> empty(new uint8_t[cinfo->image_width]);
     memset(empty.get(), 0, cinfo->image_width);
 
+    const int aligned_width = ALIGNM(cinfo->image_width, kCompressBatchSize);
+    bool is_width_aligned = (aligned_width == cinfo->image_width);
+    std::unique_ptr<uint8_t[]> buffer_intrm = nullptr;
+    uint8_t* y_plane_intrm = nullptr;
+    uint8_t* u_plane_intrm = nullptr;
+    JSAMPROW y_intrm[kCompressBatchSize];
+    JSAMPARRAY planes_intrm[]{y_intrm};
+    if (!is_width_aligned) {
+        size_t mcu_row_size = aligned_width * kCompressBatchSize;
+        buffer_intrm = std::make_unique<uint8_t[]>(mcu_row_size);
+        y_plane_intrm = buffer_intrm.get();
+        for (int i = 0; i < kCompressBatchSize; ++i) {
+            y_intrm[i] = y_plane_intrm + i * aligned_width;
+            memset(y_intrm[i] + cinfo->image_width, 0, aligned_width - cinfo->image_width);
+        }
+    }
+
     while (cinfo->next_scanline < cinfo->image_height) {
         for (int i = 0; i < kCompressBatchSize; ++i) {
             size_t scanline = cinfo->next_scanline + i;
@@ -221,8 +278,12 @@ bool JpegEncoderHelper::compressSingleChannel(jpeg_compress_struct* cinfo, const
             } else {
                 y[i] = empty.get();
             }
+            if (!is_width_aligned) {
+                memcpy(y_intrm[i], y[i], cinfo->image_width);
+            }
         }
-        int processed = jpeg_write_raw_data(cinfo, planes, kCompressBatchSize);
+        int processed = jpeg_write_raw_data(cinfo, is_width_aligned ? planes : planes_intrm,
+                                            kCompressBatchSize);
         if (processed != kCompressBatchSize / 2) {
             ALOGE("Number of processed lines does not equal input lines.");
             return false;
diff --git a/libs/ultrahdr/jpegr.cpp b/libs/ultrahdr/jpegr.cpp
index c250aa02fd..9aadb74ca6 100644
--- a/libs/ultrahdr/jpegr.cpp
+++ b/libs/ultrahdr/jpegr.cpp
@@ -76,9 +76,9 @@ static const int kMinHeight = 2 * kMapDimensionScaleFactor;
 // JPEG encoding / decoding will require block based DCT transform 16 x 16 for luma,
 // and 8 x 8 for chroma.
 // Width must be 16 dividable for luma, and 8 dividable for chroma.
-// If this criteria is not ficilitated, we will pad zeros based on the required block size.
+// If this criteria is not facilitated, we will pad zeros based to each line on the
+// required block size.
 static const size_t kJpegBlock = JpegEncoderHelper::kCompressBatchSize;
-static const size_t kJpegBlockSquare = kJpegBlock * kJpegBlock;
 // JPEG compress quality (0 ~ 100) for gain map
 static const int kMapCompressQuality = 85;
 
@@ -228,13 +228,8 @@ status_t JpegR::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
   metadata.version = kJpegrVersion;
 
   jpegr_uncompressed_struct uncompressed_yuv_420_image;
-  size_t gain_map_length = uncompressed_p010_image->width * uncompressed_p010_image->height * 3 / 2;
-  // Pad a pseudo chroma block (kJpegBlock / 2) x (kJpegBlock / 2)
-  // if width is not kJpegBlock aligned.
-  if (uncompressed_p010_image->width % kJpegBlock != 0) {
-    gain_map_length += kJpegBlockSquare / 4;
-  }
-  unique_ptr<uint8_t[]> uncompressed_yuv_420_image_data = make_unique<uint8_t[]>(gain_map_length);
+  unique_ptr<uint8_t[]> uncompressed_yuv_420_image_data = make_unique<uint8_t[]>(
+      uncompressed_p010_image->width * uncompressed_p010_image->height * 3 / 2);
   uncompressed_yuv_420_image.data = uncompressed_yuv_420_image_data.get();
   JPEGR_CHECK(toneMap(uncompressed_p010_image, &uncompressed_yuv_420_image));
 
diff --git a/libs/ultrahdr/tests/jpegencoderhelper_test.cpp b/libs/ultrahdr/tests/jpegencoderhelper_test.cpp
index 8f18ac0004..f0e1fa4968 100644
--- a/libs/ultrahdr/tests/jpegencoderhelper_test.cpp
+++ b/libs/ultrahdr/tests/jpegencoderhelper_test.cpp
@@ -108,18 +108,9 @@ TEST_F(JpegEncoderHelperTest, encodeAlignedImage) {
     ASSERT_GT(encoder.getCompressedImageSize(), static_cast<uint32_t>(0));
 }
 
-// The width of the "unaligned" image is not 16-aligned, and will fail if encoded directly.
-// Should pass with the padding zero method.
 TEST_F(JpegEncoderHelperTest, encodeUnalignedImage) {
     JpegEncoderHelper encoder;
-    const size_t paddingZeroLength = JpegEncoderHelper::kCompressBatchSize
-            * JpegEncoderHelper::kCompressBatchSize / 4;
-    std::unique_ptr<uint8_t[]> imageWithPaddingZeros(
-            new uint8_t[UNALIGNED_IMAGE_WIDTH * UNALIGNED_IMAGE_HEIGHT * 3 / 2
-            + paddingZeroLength]);
-    memcpy(imageWithPaddingZeros.get(), mUnalignedImage.buffer.get(),
-            UNALIGNED_IMAGE_WIDTH * UNALIGNED_IMAGE_HEIGHT * 3 / 2);
-    EXPECT_TRUE(encoder.compressImage(imageWithPaddingZeros.get(), mUnalignedImage.width,
+    EXPECT_TRUE(encoder.compressImage(mUnalignedImage.buffer.get(), mUnalignedImage.width,
                                       mUnalignedImage.height, JPEG_QUALITY, NULL, 0));
     ASSERT_GT(encoder.getCompressedImageSize(), static_cast<uint32_t>(0));
 }