Use more helper macros/methods in JBig2_Image.cpp.

Bundling the test with the accessor is a safer pattern than
performing the check externally.

Add test for CopyLine().

Change-Id: I7056bf33bdca40cb84a89e4928567a389d88ff1c
Reviewed-on: https://pdfium-review.googlesource.com/39431
Reviewed-by: Lei Zhang <thestig@chromium.org>
Commit-Queue: Tom Sepez <tsepez@chromium.org>

3 files changed, 71 insertions, 23 deletions

diff --git a/core/fxcodec/jbig2/JBig2_Image.cpp b/core/fxcodec/jbig2/JBig2_Image.cpp
index f66780aa16..b64a4fdbe7 100644
--- a/core/fxcodec/jbig2/JBig2_Image.cpp
+++ b/core/fxcodec/jbig2/JBig2_Image.cpp
@@ -30,6 +30,7 @@
    (buf)[2] = static_cast<uint8_t>((val) >> 8),  \
    (buf)[3] = static_cast<uint8_t>((val) >> 0))
 
+#define BIT_INDEX_TO_BYTE(x) ((x) >> 3)
 #define BIT_INDEX_TO_ALIGNED_BYTE(x) (((x) >> 5) << 2)
 
 namespace {
@@ -101,43 +102,50 @@ int CJBig2_Image::GetPixel(int32_t x, int32_t y) const {
   if (x < 0 || x >= m_nWidth)
     return 0;
 
-  if (y < 0 || y >= m_nHeight)
+  const uint8_t* pLine = GetLine(y);
+  if (!pLine)
     return 0;
 
-  int32_t m = y * m_nStride + (x >> 3);
+  int32_t m = BIT_INDEX_TO_BYTE(x);
   int32_t n = x & 7;
-  return ((data()[m] >> (7 - n)) & 1);
+  return ((pLine[m] >> (7 - n)) & 1);
 }
 
 void CJBig2_Image::SetPixel(int32_t x, int32_t y, int v) {
   if (!m_pData)
     return;
 
-  if (x < 0 || x >= m_nWidth || y < 0 || y >= m_nHeight)
+  if (x < 0 || x >= m_nWidth)
+    return;
+
+  uint8_t* pLine = GetLine(y);
+  if (!pLine)
     return;
 
-  int32_t m = y * m_nStride + (x >> 3);
+  int32_t m = BIT_INDEX_TO_BYTE(x);
   int32_t n = 1 << (7 - (x & 7));
   if (v)
-    data()[m] |= n;
+    pLine[m] |= n;
   else
-    data()[m] &= ~n;
-}
-
-uint8_t* CJBig2_Image::GetLine(int32_t y) const {
-  return (y >= 0 && y < m_nHeight) ? GetLineUnsafe(y) : nullptr;
+    pLine[m] &= ~n;
 }
 
 void CJBig2_Image::CopyLine(int32_t hTo, int32_t hFrom) {
   if (!m_pData)
     return;
 
-  if (hFrom < 0 || hFrom >= m_nHeight) {
-    memset(data() + hTo * m_nStride, 0, m_nStride);
-  } else {
-    memcpy(data() + hTo * m_nStride, data() + hFrom * m_nStride, m_nStride);
+  uint8_t* pDst = GetLine(hTo);
+  if (!pDst)
+    return;
+
+  const uint8_t* pSrc = GetLine(hFrom);
+  if (!pSrc) {
+    memset(pDst, 0, m_nStride);
+    return;
   }
+  memcpy(pDst, pSrc, m_nStride);
 }
+
 void CJBig2_Image::Fill(bool v) {
   if (!m_pData)
     return;
@@ -280,9 +288,9 @@ bool CJBig2_Image::ComposeToOpt2(CJBig2_Image* pDst,
   uint32_t maskL = 0xffffffff >> d1;
   uint32_t maskR = 0xffffffff << ((32 - (xd1 & 31)) % 32);
   uint32_t maskM = maskL & maskR;
-  uint8_t* lineSrc = data() + ys0 * m_nStride + ((xs0 >> 5) << 2);
-  int32_t lineLeft = m_nStride - ((xs0 >> 5) << 2);
-  uint8_t* lineDst = pDst->data() + yd0 * pDst->m_nStride + ((xd0 >> 5) << 2);
+  uint8_t* lineSrc = GetLineUnsafe(ys0) + BIT_INDEX_TO_ALIGNED_BYTE(xs0);
+  int32_t lineLeft = m_nStride - BIT_INDEX_TO_ALIGNED_BYTE(xs0);
+  uint8_t* lineDst = pDst->GetLineUnsafe(yd0) + BIT_INDEX_TO_ALIGNED_BYTE(xd0);
   if ((xd0 & ~31) == ((xd1 - 1) & ~31)) {
     if ((xs0 & ~31) == ((xs1 - 1) & ~31)) {
       if (s1 > d1) {
@@ -665,11 +673,11 @@ bool CJBig2_Image::ComposeToOpt2WithRect(CJBig2_Image* pDst,
   int32_t maskL = 0xffffffff >> d1;
   int32_t maskR = 0xffffffff << ((32 - (xd1 & 31)) % 32);
   int32_t maskM = maskL & maskR;
-  const uint8_t* lineSrc =
-      data() + (rtSrc.top + ys0) * m_nStride + (((xs0 + rtSrc.left) >> 5) << 2);
+  const uint8_t* lineSrc = GetLineUnsafe(rtSrc.top + ys0) +
+                           BIT_INDEX_TO_ALIGNED_BYTE(xs0 + rtSrc.left);
   const uint8_t* lineSrcEnd = data() + m_nHeight * m_nStride;
-  int32_t lineLeft = m_nStride - ((xs0 >> 5) << 2);
-  uint8_t* lineDst = pDst->data() + yd0 * pDst->m_nStride + ((xd0 >> 5) << 2);
+  int32_t lineLeft = m_nStride - BIT_INDEX_TO_ALIGNED_BYTE(xs0);
+  uint8_t* lineDst = pDst->GetLineUnsafe(yd0) + BIT_INDEX_TO_ALIGNED_BYTE(xd0);
   if ((xd0 & ~31) == ((xd1 - 1) & ~31)) {
     if ((xs0 & ~31) == ((xs1 - 1) & ~31)) {
       if (s1 > d1) {
diff --git a/core/fxcodec/jbig2/JBig2_Image.h b/core/fxcodec/jbig2/JBig2_Image.h
index b1b1489506..b61ce84e0f 100644
--- a/core/fxcodec/jbig2/JBig2_Image.h
+++ b/core/fxcodec/jbig2/JBig2_Image.h
@@ -40,8 +40,11 @@ class CJBig2_Image {
   int GetPixel(int32_t x, int32_t y) const;
   void SetPixel(int32_t x, int32_t y, int bVal);
 
-  uint8_t* GetLine(int32_t y) const;
   uint8_t* GetLineUnsafe(int32_t y) const { return data() + y * m_nStride; }
+  uint8_t* GetLine(int32_t y) const {
+    return (y >= 0 && y < m_nHeight) ? GetLineUnsafe(y) : nullptr;
+  }
+
   void CopyLine(int32_t hTo, int32_t hFrom);
   void Fill(bool v);
 
diff --git a/core/fxcodec/jbig2/JBig2_Image_unittest.cpp b/core/fxcodec/jbig2/JBig2_Image_unittest.cpp
index 2b4d897cea..4b078324e1 100644
--- a/core/fxcodec/jbig2/JBig2_Image_unittest.cpp
+++ b/core/fxcodec/jbig2/JBig2_Image_unittest.cpp
@@ -297,3 +297,40 @@ TEST(fxcodec, JBig2SubImage) {
   EXPECT_EQ(32, sub->width());
   EXPECT_EQ(40, sub->height());
 }
+
+TEST(fxcodec, JBig2CopyLine) {
+  // Horizontal line in image.
+  uint8_t pattern[3][8] = {
+      {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+      {0x00, 0x01, 0xff, 0x80, 0x00, 0x00, 0x00, 0x00},
+      {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+  };
+
+  uint8_t expected_pattern[3][8] = {
+      {0x00, 0x01, 0xff, 0x80, 0x00, 0x00, 0x00, 0x00},
+      {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+      {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+  };
+
+  auto img = pdfium::MakeUnique<CJBig2_Image>(
+      37, 3, 8, reinterpret_cast<uint8_t*>(pattern));
+
+  auto expected = pdfium::MakeUnique<CJBig2_Image>(
+      37, 3, 8, reinterpret_cast<uint8_t*>(expected_pattern));
+
+  // Shuffle.
+  img->CopyLine(2, 1);
+  img->CopyLine(1, 0);
+  img->CopyLine(0, 2);
+
+  // Clear top line via invalid |from| offset.
+  img->CopyLine(2, 3);
+
+  // Copies with invalid |to|s don't mess with things.
+  img->CopyLine(-1, 0);
+  img->CopyLine(4, 0);
+  img->CopyLine(-1, -1);
+  img->CopyLine(4, 4);
+
+  CheckImageEq(expected.get(), img.get(), __LINE__);
+}