Adjust CJBig2_Context::HuffmanAssignCode().

It looks a lot like CJBig2_HuffmanTable::InitCodes(). Port over the
UBSAN error fix from commit 76c9a1b1.

BUG=chromium:709781

Change-Id: I5d2f8fb013c09099c82b0565627b77e4fb0f8a98
Reviewed-on: https://pdfium-review.googlesource.com/31536
Commit-Queue: Lei Zhang <thestig@chromium.org>
Reviewed-by: Henrique Nakashima <hnakashima@chromium.org>

3 files changed, 28 insertions, 18 deletions

diff --git a/core/fxcodec/jbig2/JBig2_Context.cpp b/core/fxcodec/jbig2/JBig2_Context.cpp
index ef5f05ba5a..6cb94872c6 100644
--- a/core/fxcodec/jbig2/JBig2_Context.cpp
+++ b/core/fxcodec/jbig2/JBig2_Context.cpp
@@ -1257,7 +1257,8 @@ std::vector<JBig2HuffmanCode> CJBig2_Context::DecodeSymbolIDHuffmanTable(
     if (m_pStream->readNBits(4, &huffman_codes[i].codelen) != 0)
       return std::vector<JBig2HuffmanCode>();
   }
-  HuffmanAssignCode(huffman_codes, kRunCodesSize);
+  if (!HuffmanAssignCode(huffman_codes, kRunCodesSize))
+    return std::vector<JBig2HuffmanCode>();
 
   std::vector<JBig2HuffmanCode> SBSYMCODES(SBNUMSYMS);
   int32_t run = 0;
@@ -1313,29 +1314,38 @@ std::vector<JBig2HuffmanCode> CJBig2_Context::DecodeSymbolIDHuffmanTable(
       ++i;
     }
   }
-  HuffmanAssignCode(SBSYMCODES.data(), SBNUMSYMS);
+  if (!HuffmanAssignCode(SBSYMCODES.data(), SBNUMSYMS))
+    return std::vector<JBig2HuffmanCode>();
   return SBSYMCODES;
 }
 
-void CJBig2_Context::HuffmanAssignCode(JBig2HuffmanCode* SBSYMCODES,
-                                       int NTEMP) {
-  // TODO(thestig) CJBig2_HuffmanTable::ParseFromCodedBuffer() has similar code.
+bool CJBig2_Context::HuffmanAssignCode(JBig2HuffmanCode* SBSYMCODES,
+                                       uint32_t NTEMP) {
+  // TODO(thestig): CJBig2_HuffmanTable::InitCodes() has similar code.
   int LENMAX = 0;
-  for (int i = 0; i < NTEMP; ++i)
-    LENMAX = std::max(LENMAX, SBSYMCODES[i].codelen);
+  for (uint32_t i = 0; i < NTEMP; ++i)
+    LENMAX = std::max(SBSYMCODES[i].codelen, LENMAX);
+
   std::vector<int> LENCOUNT(LENMAX + 1);
   std::vector<int> FIRSTCODE(LENMAX + 1);
-  for (int i = 0; i < NTEMP; ++i)
+  for (uint32_t i = 0; i < NTEMP; ++i)
     ++LENCOUNT[SBSYMCODES[i].codelen];
+
   LENCOUNT[0] = 0;
-  for (int CURLEN = 1; CURLEN <= LENMAX; ++CURLEN) {
-    FIRSTCODE[CURLEN] = (FIRSTCODE[CURLEN - 1] + LENCOUNT[CURLEN - 1]) << 1;
-    int CURCODE = FIRSTCODE[CURLEN];
-    for (int CURTEMP = 0; CURTEMP < NTEMP; ++CURTEMP) {
-      if (SBSYMCODES[CURTEMP].codelen == CURLEN) {
-        SBSYMCODES[CURTEMP].code = CURCODE;
-        CURCODE = CURCODE + 1;
+  for (int i = 1; i <= LENMAX; ++i) {
+    pdfium::base::CheckedNumeric<int> shifted = FIRSTCODE[i - 1];
+    shifted += LENCOUNT[i - 1];
+    shifted <<= 1;
+    if (!shifted.IsValid())
+      return false;
+
+    FIRSTCODE[i] = shifted.ValueOrDie();
+    int CURCODE = FIRSTCODE[i];
+    for (uint32_t j = 0; j < NTEMP; ++j) {
+      if (SBSYMCODES[j].codelen == i) {
+        SBSYMCODES[j].code = CURCODE++;
       }
     }
   }
+  return true;
 }
diff --git a/core/fxcodec/jbig2/JBig2_Context.h b/core/fxcodec/jbig2/JBig2_Context.h
index 3c0e5b6e2b..3f3304d8b7 100644
--- a/core/fxcodec/jbig2/JBig2_Context.h
+++ b/core/fxcodec/jbig2/JBig2_Context.h
@@ -82,7 +82,7 @@ class CJBig2_Context {
 
   std::vector<JBig2HuffmanCode> DecodeSymbolIDHuffmanTable(uint32_t SBNUMSYMS);
 
-  void HuffmanAssignCode(JBig2HuffmanCode* SBSYMCODES, int NTEMP);
+  bool HuffmanAssignCode(JBig2HuffmanCode* SBSYMCODES, uint32_t NTEMP);
 
   std::unique_ptr<CJBig2_Context> m_pGlobalContext;
   std::unique_ptr<CJBig2_BitStream> m_pStream;
diff --git a/core/fxcodec/jbig2/JBig2_HuffmanTable.cpp b/core/fxcodec/jbig2/JBig2_HuffmanTable.cpp
index 83f9fed010..5a49e1e4be 100644
--- a/core/fxcodec/jbig2/JBig2_HuffmanTable.cpp
+++ b/core/fxcodec/jbig2/JBig2_HuffmanTable.cpp
@@ -122,8 +122,8 @@ bool CJBig2_HuffmanTable::InitCodes() {
   FIRSTCODE[0] = 0;
   LENCOUNT[0] = 0;
   for (int i = 1; i <= lenmax; ++i) {
-    pdfium::base::CheckedNumeric<int> shifted;
-    shifted = FIRSTCODE[i - 1] + LENCOUNT[i - 1];
+    pdfium::base::CheckedNumeric<int> shifted = FIRSTCODE[i - 1];
+    shifted += LENCOUNT[i - 1];
     shifted <<= 1;
     if (!shifted.IsValid())
       return false;