Fix AES block size as 128 bits (it is always that).

Regardless of key length, which only affects key scheduling.

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

core/fdrm/fx_crypt_aes.cpp

diff --git a/core/fdrm/fx_crypt_aes.cpp b/core/fdrm/fx_crypt_aes.cpp
index 149e460..f793917 100644
--- a/core/fdrm/fx_crypt_aes.cpp
+++ b/core/fdrm/fx_crypt_aes.cpp
@@ -528,14 +528,14 @@
 
 void CRYPT_AESSetKey(CRYPT_aes_context* ctx, pdfium::span<const uint8_t> key) {
   CHECK(key.size() == 16 || key.size() == 24 || key.size() == 32);
-  size_t Nk = key.size() / 4;
-  ctx->Nb = 4;
-  ctx->Nr = 6 + (ctx->Nb > Nk ? ctx->Nb : Nk);
+  const size_t Nk = key.size() / 4;
+  ctx->Nr = 6 + Nk;
+  const size_t sched_size = CRYPT_aes_context::kBlockSize * (ctx->Nr + 1);
   int rconst = 1;
-  for (size_t i = 0; i < (ctx->Nr + 1) * ctx->Nb; i++) {
+  for (size_t i = 0; i < sched_size; i++) {
     if (i < Nk) {
-      ctx->keysched[i] =
-          fxcrt::GetUInt32MSBFirst(key.subspan(4u * i).first<4u>());
+      ctx->keysched[i] = fxcrt::GetUInt32MSBFirst(key.first<4u>());
+      key = key.subspan<4u>();
     } else {
       uint32_t temp = ctx->keysched[i - 1];
       if (i % Nk == 0) {
@@ -562,8 +562,9 @@
     }
   }
   for (size_t i = 0; i <= ctx->Nr; i++) {
-    for (size_t j = 0; j < ctx->Nb; j++) {
-      uint32_t temp = ctx->keysched[(ctx->Nr - i) * ctx->Nb + j];
+    for (size_t j = 0; j < CRYPT_aes_context::kBlockSize; j++) {
+      uint32_t temp =
+          ctx->keysched[(ctx->Nr - i) * CRYPT_aes_context::kBlockSize + j];
       if (i != 0 && i != ctx->Nr) {
         int a = (temp >> 24) & 0xFF;
         int b = (temp >> 16) & 0xFF;
@@ -574,14 +575,14 @@
         temp ^= D2[Sbox[c]];
         temp ^= D3[Sbox[d]];
       }
-      ctx->invkeysched[i * ctx->Nb + j] = temp;
+      ctx->invkeysched[i * CRYPT_aes_context::kBlockSize + j] = temp;
     }
   }
 }
 
 void CRYPT_AESSetIV(CRYPT_aes_context* ctx,
                     pdfium::span<const uint8_t, 16> iv) {
-  for (size_t i = 0; i < ctx->Nb; i++) {
+  for (size_t i = 0; i < CRYPT_aes_context::kBlockSize; i++) {
     ctx->iv[i] = fxcrt::GetUInt32MSBFirst(iv.subspan(4u * i).first<4u>());
   }
 }

core/fdrm/fx_crypt_aes.h

diff --git a/core/fdrm/fx_crypt_aes.h b/core/fdrm/fx_crypt_aes.h
index 5f1e4ca..e78c057 100644
--- a/core/fdrm/fx_crypt_aes.h
+++ b/core/fdrm/fx_crypt_aes.h
@@ -14,15 +14,15 @@
 #include "core/fxcrt/span.h"
 
 struct CRYPT_aes_context {
-  static constexpr int kMaxNb = 8;
-  static constexpr int kMaxNr = 14;
-  static constexpr int kSchedSize = (kMaxNr + 1) * kMaxNb;
+  static constexpr size_t kBlockSize = 4;       // Words, not bytes.
+  static constexpr size_t kMaxNb = 8;           // Max words in largest key.
+  static constexpr size_t kMaxNr = 6 + kMaxNb;  // Max rounds for largest key.
+  static constexpr size_t kMaxSchedSize = (kMaxNr + 1) * kBlockSize;
 
-  size_t Nb;
-  size_t Nr;
-  std::array<uint32_t, kSchedSize> keysched;
-  std::array<uint32_t, kSchedSize> invkeysched;
-  std::array<uint32_t, kMaxNb> iv;
+  size_t Nr;  // Actual number of rounds based on keysize.
+  std::array<uint32_t, kMaxSchedSize> keysched;
+  std::array<uint32_t, kMaxSchedSize> invkeysched;
+  std::array<uint32_t, kBlockSize> iv;
 };
 
 void CRYPT_AESSetKey(CRYPT_aes_context* ctx, pdfium::span<const uint8_t> key);