Clean up CRYPT_ArcFourCrypt().

Add a unit test too.

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

2 files changed, 228 insertions, 28 deletions

diff --git a/core/fdrm/crypto/fx_crypt.cpp b/core/fdrm/crypto/fx_crypt.cpp
index f5ebdf09d2..d53fa6ba7a 100644
--- a/core/fdrm/crypto/fx_crypt.cpp
+++ b/core/fdrm/crypto/fx_crypt.cpp
@@ -153,23 +153,13 @@ void CRYPT_ArcFourSetup(CRYPT_rc4_context* s,
   }
 }
 
-void CRYPT_ArcFourCrypt(CRYPT_rc4_context* s,
-                        unsigned char* data,
-                        uint32_t length) {
-  int i, x, y, *m, a, b;
-  x = s->x;
-  y = s->y;
-  m = s->m;
-  for (i = 0; i < (int)length; i++) {
-    x = (x + 1) & 0xFF;
-    a = m[x];
-    y = (y + a) & 0xFF;
-    m[x] = b = m[y];
-    m[y] = a;
-    data[i] ^= m[(a + b) & 0xFF];
+void CRYPT_ArcFourCrypt(CRYPT_rc4_context* s, uint8_t* data, uint32_t length) {
+  for (uint32_t i = 0; i < length; ++i) {
+    s->x = (s->x + 1) & 0xFF;
+    s->y = (s->y + s->m[s->x]) & 0xFF;
+    std::swap(s->m[s->x], s->m[s->y]);
+    data[i] ^= s->m[(s->m[s->x] + s->m[s->y]) & 0xFF];
   }
-  s->x = x;
-  s->y = y;
 }
 
 void CRYPT_ArcFourCryptBlock(uint8_t* pData,
diff --git a/core/fdrm/crypto/fx_crypt_unittest.cpp b/core/fdrm/crypto/fx_crypt_unittest.cpp
index 247da07233..9795d97cfe 100644
--- a/core/fdrm/crypto/fx_crypt_unittest.cpp
+++ b/core/fdrm/crypto/fx_crypt_unittest.cpp
@@ -19,6 +19,16 @@ std::string CRYPT_MD5String(const char* str) {
   return GenerateMD5Base16(reinterpret_cast<const uint8_t*>(str), strlen(str));
 }
 
+void CheckArcFourContext(const CRYPT_rc4_context& context,
+                         int32_t expected_x,
+                         int32_t expected_y,
+                         const uint8_t* expected_permutation) {
+  EXPECT_EQ(expected_x, context.x);
+  EXPECT_EQ(expected_y, context.y);
+  for (int32_t i = 0; i < kRC4ContextPermutationLength; ++i)
+    EXPECT_EQ(expected_permutation[i], context.m[i]) << i;
+}
+
 }  // namespace
 
 TEST(FXCRYPT, CryptToBase16) {
@@ -237,12 +247,9 @@ TEST(FXCRYPT, CRYPT_ArcFourSetup) {
         179, 87,  74,  147, 80,  54,  212, 16,  167, 222, 136, 213, 55,  182,
         139, 24,  209, 251, 208, 28,  111, 89,  158, 155, 243, 107, 233, 169,
         117, 184, 31,  39};
-    CRYPT_rc4_context s;
-    CRYPT_ArcFourSetup(&s, nullptr, 0);
-    EXPECT_EQ(0, s.x);
-    EXPECT_EQ(0, s.y);
-    for (int32_t i = 0; i < kRC4ContextPermutationLength; ++i)
-      EXPECT_EQ(kNullPermutation[i], s.m[i]) << i;
+    CRYPT_rc4_context context;
+    CRYPT_ArcFourSetup(&context, nullptr, 0);
+    CheckArcFourContext(context, 0, 0, kNullPermutation);
   }
   {
     const uint8_t kFoobarPermutation[kRC4ContextPermutationLength] = {
@@ -265,12 +272,215 @@ TEST(FXCRYPT, CRYPT_ArcFourSetup) {
         159, 67,  211, 108, 178, 146, 202, 11,  164, 226, 184, 50,  190, 174,
         71,  233, 235, 198, 95,  51,  110, 255, 253, 72,  115, 0,   47,  94,
         29,  45,  14,  111};
-    CRYPT_rc4_context s;
+    CRYPT_rc4_context context;
+    const uint8_t kFooBar[] = "foobar";
+    CRYPT_ArcFourSetup(&context, kFooBar, FX_ArraySize(kFooBar) - 1);
+    CheckArcFourContext(context, 0, 0, kFoobarPermutation);
+  }
+}
+
+TEST(FXCRYPT, CRYPT_ArcFourCrypt) {
+  const uint8_t kDataShort[] = "The Quick Fox Jumped Over The Lazy Brown Dog.";
+  const uint8_t kDataLong[] =
+      "The Quick Fox Jumped Over The Lazy Brown Dog.\n"
+      "1234567890123456789012345678901234567890123456789012345678901234567890\n"
+      "1234567890123456789012345678901234567890123456789012345678901234567890\n"
+      "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ\n"
+      "!@#$%^&*()[]{};':\",.<>/?\\|\r\t\n";
+  {
+    CRYPT_rc4_context context;
+    CRYPT_ArcFourSetup(&context, nullptr, 0);
+
+    uint8_t data_short[FX_ArraySize(kDataShort)];
+    memcpy(data_short, kDataShort, FX_ArraySize(kDataShort));
+    const uint8_t kExpectedEncryptedDataShort[] = {
+        138, 112, 236, 97,  242, 66,  52,  89,  225, 38,  88,  8,
+        47,  78,  216, 24,  170, 106, 26,  199, 208, 131, 157, 242,
+        55,  11,  25,  90,  66,  182, 19,  255, 210, 181, 85,  69,
+        31,  240, 206, 171, 97,  62,  202, 172, 30,  252};
+    static_assert(
+        FX_ArraySize(kExpectedEncryptedDataShort) == FX_ArraySize(data_short),
+        "data_short mismatch");
+    CRYPT_ArcFourCrypt(&context, data_short, FX_ArraySize(data_short));
+    for (size_t i = 0; i < FX_ArraySize(data_short); ++i)
+      EXPECT_EQ(kExpectedEncryptedDataShort[i], data_short[i]) << i;
+
+    const uint8_t kPermutation[kRC4ContextPermutationLength] = {
+        0,   198, 10,  37,  253, 192, 171, 183, 99,  8,   144, 103, 208, 191,
+        149, 9,   228, 243, 94,  150, 169, 151, 210, 206, 221, 235, 32,  186,
+        212, 122, 72,  200, 236, 138, 244, 217, 158, 213, 139, 242, 17,  143,
+        50,  132, 12,  160, 145, 250, 214, 76,  123, 35,  27,  249, 203, 127,
+        64,  62,  33,  60,  248, 85,  177, 6,   142, 83,  110, 140, 41,  135,
+        196, 238, 156, 91,  141, 67,  5,   185, 131, 63,  137, 43,  172, 121,
+        70,  134, 237, 130, 25,  44,  153, 166, 78,  201, 42,  119, 215, 7,
+        126, 114, 97,  11,  53,  4,   254, 45,  102, 133, 230, 88,  193, 129,
+        18,  124, 84,  108, 239, 189, 152, 120, 115, 207, 234, 176, 86,  157,
+        164, 187, 71,  1,   15,  58,  29,  21,  46,  23,  247, 162, 95,  229,
+        13,  226, 159, 175, 56,  100, 96,  202, 101, 178, 154, 47,  205, 106,
+        148, 104, 93,  112, 26,  165, 128, 246, 146, 218, 66,  211, 65,  90,
+        252, 19,  40,  49,  223, 174, 255, 51,  77,  227, 48,  220, 168, 118,
+        224, 98,  75,  105, 125, 199, 73,  82,  57,  181, 81,  173, 68,  52,
+        232, 22,  2,   216, 113, 30,  109, 163, 92,  61,  14,  36,  38,  225,
+        79,  231, 170, 240, 20,  219, 204, 161, 180, 188, 116, 190, 241, 197,
+        179, 87,  74,  147, 80,  54,  69,  16,  167, 222, 136, 245, 55,  182,
+        3,   24,  209, 251, 59,  28,  111, 89,  195, 155, 194, 107, 233, 34,
+        117, 184, 31,  39};
+    CheckArcFourContext(context, 46, 135, kPermutation);
+  }
+  {
+    CRYPT_rc4_context context;
+    CRYPT_ArcFourSetup(&context, nullptr, 0);
+
+    uint8_t data_long[FX_ArraySize(kDataLong)];
+    memcpy(data_long, kDataLong, FX_ArraySize(kDataLong));
+    const uint8_t kExpectedEncryptedDataLong[] = {
+        138, 112, 236, 97,  242, 66,  52,  89,  225, 38,  88,  8,   47,  78,
+        216, 24,  170, 106, 26,  199, 208, 131, 157, 242, 55,  11,  25,  90,
+        66,  182, 19,  255, 210, 181, 85,  69,  31,  240, 206, 171, 97,  62,
+        202, 172, 30,  246, 19,  43,  184, 0,   173, 27,  140, 90,  167, 240,
+        122, 125, 184, 49,  149, 71,  63,  104, 171, 144, 242, 106, 121, 124,
+        209, 149, 61,  1,   66,  186, 252, 47,  51,  170, 253, 75,  95,  41,
+        203, 28,  197, 174, 144, 209, 166, 98,  142, 125, 44,  5,   147, 42,
+        73,  178, 119, 90,  253, 69,  103, 178, 15,  136, 51,  112, 39,  81,
+        37,  111, 129, 232, 106, 159, 126, 142, 120, 124, 48,  140, 253, 12,
+        223, 208, 106, 76,  60,  238, 5,   162, 100, 226, 251, 156, 169, 35,
+        193, 10,  242, 210, 20,  96,  37,  84,  99,  183, 179, 203, 62,  122,
+        54,  6,   51,  239, 142, 250, 238, 41,  223, 58,  48,  101, 29,  187,
+        43,  235, 3,   5,   176, 33,  14,  171, 36,  26,  234, 207, 105, 79,
+        69,  126, 82,  183, 105, 228, 31,  173, 8,   240, 99,  5,   147, 206,
+        215, 140, 48,  190, 165, 50,  41,  232, 29,  105, 156, 64,  229, 165,
+        12,  64,  163, 255, 146, 108, 212, 125, 142, 101, 13,  99,  174, 10,
+        160, 68,  196, 120, 110, 201, 254, 158, 97,  215, 0,   207, 90,  23,
+        208, 161, 105, 226, 164, 114, 80,  137, 58,  107, 109, 42,  110, 100,
+        202, 170, 224, 89,  28,  5,   138, 19,  253, 105, 220, 105, 24,  187,
+        109, 89,  205, 89,  202};
+    static_assert(
+        FX_ArraySize(kExpectedEncryptedDataLong) == FX_ArraySize(data_long),
+        "data_long mismatch");
+    static_assert(FX_ArraySize(data_long) > 256, "too short");
+    CRYPT_ArcFourCrypt(&context, data_long, FX_ArraySize(data_long));
+    for (size_t i = 0; i < FX_ArraySize(data_long); ++i)
+      EXPECT_EQ(kExpectedEncryptedDataLong[i], data_long[i]) << i;
+
+    const uint8_t kPermutation[kRC4ContextPermutationLength] = {
+        172, 59,  196, 72,  101, 21,  215, 210, 212, 52,  243, 73,  47,  213,
+        211, 50,  228, 144, 66,  93,  169, 31,  237, 206, 221, 235, 222, 250,
+        97,  87,  174, 164, 190, 111, 27,  217, 173, 189, 65,  11,  115, 171,
+        104, 132, 12,  170, 205, 114, 7,   105, 37,  83,  78,  134, 236, 70,
+        197, 122, 177, 202, 39,  195, 30,  3,   86,  127, 74,  106, 68,  91,
+        110, 121, 208, 25,  56,  6,   28,  225, 163, 193, 166, 244, 119, 34,
+        23,  88,  108, 123, 162, 159, 242, 61,  230, 227, 254, 14,  4,   156,
+        161, 44,  58,  153, 33,  143, 129, 232, 182, 152, 76,  168, 238, 239,
+        185, 219, 233, 16,  188, 45,  40,  35,  103, 99,  89,  157, 241, 245,
+        192, 180, 248, 8,   85,  231, 146, 154, 252, 181, 107, 126, 98,  80,
+        102, 165, 199, 94,  49,  255, 18,  204, 216, 77,  20,  187, 145, 125,
+        1,   247, 79,  26,  207, 81,  117, 179, 186, 38,  175, 19,  139, 138,
+        149, 54,  64,  109, 249, 135, 142, 118, 17,  13,  201, 184, 55,  224,
+        209, 155, 113, 218, 82,  131, 178, 253, 140, 226, 43,  42,  24,  29,
+        229, 200, 137, 240, 203, 167, 95,  148, 15,  176, 60,  75,  53,  41,
+        150, 112, 160, 96,  22,  10,  234, 116, 130, 158, 214, 36,  9,   67,
+        198, 194, 191, 100, 124, 147, 32,  183, 120, 246, 51,  141, 46,  251,
+        92,  223, 133, 63,  0,   71,  48,  128, 220, 90,  62,  136, 2,   5,
+        69,  57,  151, 84};
+    CheckArcFourContext(context, 15, 222, kPermutation);
+  }
+  {
+    CRYPT_rc4_context context;
+    const uint8_t kFooBar[] = "foobar";
+    CRYPT_ArcFourSetup(&context, kFooBar, FX_ArraySize(kFooBar) - 1);
+
+    uint8_t data_short[FX_ArraySize(kDataShort)];
+    memcpy(data_short, kDataShort, FX_ArraySize(kDataShort));
+    const uint8_t kExpectedEncryptedDataShort[] = {
+        59,  193, 117, 206, 167, 54,  218, 7,   229, 214, 188, 55,
+        90,  205, 196, 25,  36,  114, 199, 218, 161, 107, 122, 119,
+        106, 167, 44,  175, 240, 123, 192, 102, 174, 167, 105, 187,
+        202, 70,  121, 81,  17,  30,  5,   138, 116, 166};
+    static_assert(
+        FX_ArraySize(kExpectedEncryptedDataShort) == FX_ArraySize(data_short),
+        "data_short mismatch");
+    CRYPT_ArcFourCrypt(&context, data_short, FX_ArraySize(data_short));
+    for (size_t i = 0; i < FX_ArraySize(data_short); ++i)
+      EXPECT_EQ(kExpectedEncryptedDataShort[i], data_short[i]) << i;
+
+    const uint8_t kPermutation[kRC4ContextPermutationLength] = {
+        102, 41,  45,  82,  124, 141, 237, 38,  6,   64,  90,  140, 254, 96,
+        220, 109, 99,  49,  27,  227, 205, 75,  191, 37,  17,  54,  83,  196,
+        108, 79,  31,  190, 180, 0,   125, 194, 243, 156, 224, 246, 253, 193,
+        42,  81,  117, 56,  181, 252, 113, 210, 217, 62,  218, 129, 61,  33,
+        128, 9,   153, 59,  43,  13,  206, 48,  131, 18,  213, 118, 173, 122,
+        80,  172, 177, 105, 148, 207, 186, 5,   85,  32,  68,  215, 19,  84,
+        169, 209, 150, 7,   133, 63,  147, 93,  26,  130, 60,  145, 250, 57,
+        24,  247, 200, 127, 136, 66,  112, 107, 212, 154, 70,  170, 185, 138,
+        248, 236, 88,  86,  44,  216, 241, 35,  100, 151, 78,  74,  119, 55,
+        245, 46,  199, 208, 229, 16,  249, 149, 53,  157, 201, 234, 58,  28,
+        142, 238, 182, 163, 179, 144, 12,  114, 176, 10,  183, 239, 104, 40,
+        73,  101, 137, 69,  221, 134, 165, 188, 25,  87,  1,   91,  2,   171,
+        232, 34,  162, 166, 160, 126, 225, 167, 123, 197, 223, 195, 22,  203,
+        189, 244, 103, 139, 222, 175, 23,  143, 152, 192, 21,  231, 228, 132,
+        30,  204, 158, 240, 120, 98,  89,  121, 135, 251, 168, 4,   161, 3,
+        8,   230, 52,  219, 214, 242, 36,  97,  15,  155, 65,  187, 116, 76,
+        159, 67,  211, 20,  178, 146, 202, 11,  164, 226, 184, 50,  77,  174,
+        71,  233, 235, 198, 95,  51,  110, 255, 92,  72,  115, 106, 47,  94,
+        29,  39,  14,  111};
+    CheckArcFourContext(context, 46, 39, kPermutation);
+  }
+  {
+    CRYPT_rc4_context context;
     const uint8_t kFooBar[] = "foobar";
-    CRYPT_ArcFourSetup(&s, kFooBar, 6);
-    EXPECT_EQ(0, s.x);
-    EXPECT_EQ(0, s.y);
-    for (int32_t i = 0; i < kRC4ContextPermutationLength; ++i)
-      EXPECT_EQ(kFoobarPermutation[i], s.m[i]) << i;
+    CRYPT_ArcFourSetup(&context, kFooBar, FX_ArraySize(kFooBar) - 1);
+
+    uint8_t data_long[FX_ArraySize(kDataLong)];
+    memcpy(data_long, kDataLong, FX_ArraySize(kDataLong));
+    const uint8_t kExpectedEncryptedDataLong[] = {
+        59,  193, 117, 206, 167, 54,  218, 7,   229, 214, 188, 55,  90,  205,
+        196, 25,  36,  114, 199, 218, 161, 107, 122, 119, 106, 167, 44,  175,
+        240, 123, 192, 102, 174, 167, 105, 187, 202, 70,  121, 81,  17,  30,
+        5,   138, 116, 172, 169, 50,  160, 116, 237, 117, 108, 241, 127, 61,
+        83,  45,  77,  176, 0,   106, 191, 221, 132, 143, 219, 94,  2,   235,
+        204, 166, 201, 139, 140, 163, 104, 115, 48,  37,  18,  114, 168, 49,
+        235, 163, 179, 131, 182, 218, 120, 200, 9,   90,  60,  47,  55,  235,
+        135, 37,  21,  170, 48,  112, 185, 169, 43,  233, 88,  134, 117, 126,
+        248, 40,  176, 248, 30,  131, 108, 43,  139, 68,  232, 219, 7,   39,
+        223, 45,  199, 243, 54,  171, 31,  37,  161, 24,  38,  251, 13,  144,
+        106, 215, 179, 203, 5,   253, 25,  32,  25,  146, 109, 193, 143, 141,
+        177, 226, 134, 222, 95,  79,  156, 202, 240, 34,  153, 145, 169, 150,
+        231, 63,  113, 242, 156, 39,  136, 249, 108, 50,  181, 22,  22,  180,
+        57,  76,  69,  62,  254, 47,  141, 249, 235, 90,  25,  34,  40,  194,
+        66,  86,  110, 192, 235, 191, 205, 133, 91,  32,  104, 65,  43,  36,
+        140, 36,  228, 156, 105, 251, 169, 168, 203, 189, 238, 221, 64,  200,
+        68,  137, 153, 9,   183, 84,  153, 140, 239, 0,   15,  50,  126, 145,
+        22,  110, 43,  56,  94,  127, 48,  96,  47,  172, 3,   31,  130, 249,
+        243, 73,  206, 89,  9,   93,  156, 167, 205, 166, 75,  227, 36,  34,
+        81,  124, 195, 246, 152};
+    static_assert(
+        FX_ArraySize(kExpectedEncryptedDataLong) == FX_ArraySize(data_long),
+        "data_long mismatch");
+    static_assert(FX_ArraySize(data_long) > 256, "too short");
+    CRYPT_ArcFourCrypt(&context, data_long, FX_ArraySize(data_long));
+    for (size_t i = 0; i < FX_ArraySize(data_long); ++i)
+      EXPECT_EQ(kExpectedEncryptedDataLong[i], data_long[i]) << i;
+
+    const uint8_t kPermutation[kRC4ContextPermutationLength] = {
+        188, 12,  81,  130, 228, 58,  124, 218, 72,  210, 50,  70,  166, 38,
+        110, 111, 73,  49,  27,  227, 249, 21,  1,   226, 17,  54,  53,  16,
+        108, 51,  31,  123, 221, 23,  125, 148, 5,   200, 208, 246, 253, 193,
+        42,  45,  236, 56,  230, 194, 178, 213, 120, 116, 7,   164, 33,  107,
+        189, 20,  133, 114, 173, 161, 59,  128, 3,   238, 65,  69,  144, 179,
+        44,  35,  8,   163, 252, 195, 160, 197, 204, 28,  34,  129, 67,  89,
+        22,  149, 199, 131, 182, 46,  250, 222, 155, 104, 10,  32,  139, 245,
+        90,  41,  132, 224, 83,  242, 135, 75,  74,  61,  62,  141, 43,  127,
+        255, 91,  170, 78,  157, 101, 243, 216, 254, 156, 229, 118, 174, 147,
+        103, 76,  196, 145, 134, 94,  205, 146, 202, 98,  100, 106, 232, 177,
+        187, 13,  80,  137, 151, 11,  82,  40,  167, 175, 25,  219, 168, 240,
+        99,  55,  4,   19,  180, 2,   203, 18,  171, 154, 113, 117, 6,   185,
+        172, 186, 237, 223, 233, 244, 217, 191, 190, 198, 97,  165, 220, 9,
+        214, 150, 184, 143, 206, 24,  209, 207, 36,  142, 87,  15,  159, 71,
+        84,  162, 169, 86,  48,  47,  140, 215, 241, 235, 158, 14,  26,  248,
+        138, 119, 212, 39,  88,  121, 96,  109, 29,  66,  136, 102, 225, 92,
+        201, 126, 122, 192, 60,  0,   64,  239, 183, 37,  57,  63,  234, 181,
+        153, 52,  176, 112, 93,  79,  77,  115, 231, 30,  95,  251, 211, 68,
+        105, 85,  247, 152};
+    CheckArcFourContext(context, 15, 68, kPermutation);
   }
 }