Use even more spans in fx_crypt_aes.cpp Avoids some UNSAFE_TODO() usage. Bug: 42271176 Change-Id: Ifbcc7339145c9fee8347cf9f9c158f981115f896 Reviewed-on: https://pdfium-review.googlesource.com/c/pdfium/+/121893 Reviewed-by: Tom Sepez <tsepez@google.com> Reviewed-by: Lei Zhang <thestig@chromium.org> Commit-Queue: Tom Sepez <tsepez@chromium.org>

diff --git a/core/fdrm/fx_crypt_aes.cpp b/core/fdrm/fx_crypt_aes.cpp
index 1f373d9..95acb21 100644
--- a/core/fdrm/fx_crypt_aes.cpp
+++ b/core/fdrm/fx_crypt_aes.cpp

@@ -527,18 +527,17 @@
 
 }  // namespace
 
-void CRYPT_AESSetKey(CRYPT_aes_context* ctx,
-                     const uint8_t* key,
-                     uint32_t keylen) {
-  DCHECK(keylen == 16 || keylen == 24 || keylen == 32);
-  auto keyspan = UNSAFE_TODO(pdfium::make_span(key, keylen));
-  int Nk = keylen / 4;
+void CRYPT_AESSetKey(CRYPT_aes_context* ctx, pdfium::span<const uint8_t> key) {
+  size_t keylen = key.size();
+  CHECK(keylen == 16 || keylen == 24 || keylen == 32);
+
+  int Nk = static_cast<int>(keylen / 4);
   ctx->Nb = 4;
   ctx->Nr = 6 + (ctx->Nb > Nk ? ctx->Nb : Nk);
   int rconst = 1;
   for (int i = 0; i < (ctx->Nr + 1) * ctx->Nb; i++) {
     if (i < Nk) {
-      ctx->keysched[i] = fxcrt::GetUInt32MSBFirst(keyspan.subspan(4 * i));
+      ctx->keysched[i] = fxcrt::GetUInt32MSBFirst(key.subspan(4 * i));
     } else {
       uint32_t temp = ctx->keysched[i - 1];
       if (i % Nk == 0) {
@@ -583,11 +582,9 @@
   }
 }
 
-void CRYPT_AESSetIV(CRYPT_aes_context* ctx, const uint8_t* iv) {
+void CRYPT_AESSetIV(CRYPT_aes_context* ctx, pdfium::span<const uint8_t> iv) {
   for (int i = 0; i < ctx->Nb; i++) {
-    // TODO(tsepez): Pass actual span.
-    ctx->iv[i] = fxcrt::GetUInt32MSBFirst(
-        UNSAFE_TODO(pdfium::make_span(iv + 4 * i, 4u)));
+    ctx->iv[i] = fxcrt::GetUInt32MSBFirst(iv.subspan(4u * i, 4u));
   }
 }
 

diff --git a/core/fdrm/fx_crypt_aes.h b/core/fdrm/fx_crypt_aes.h
index ab673d2..cae3790 100644
--- a/core/fdrm/fx_crypt_aes.h
+++ b/core/fdrm/fx_crypt_aes.h

@@ -25,10 +25,8 @@
   std::array<uint32_t, kMaxNb> iv;
 };
 
-void CRYPT_AESSetKey(CRYPT_aes_context* ctx,
-                     const uint8_t* key,
-                     uint32_t keylen);
-void CRYPT_AESSetIV(CRYPT_aes_context* ctx, const uint8_t* iv);
+void CRYPT_AESSetKey(CRYPT_aes_context* ctx, pdfium::span<const uint8_t> key);
+void CRYPT_AESSetIV(CRYPT_aes_context* ctx, pdfium::span<const uint8_t> iv);
 void CRYPT_AESDecrypt(CRYPT_aes_context* ctx,
                       uint8_t* dest,
                       const uint8_t* src,

diff --git a/core/fpdfapi/parser/cpdf_crypto_handler.cpp b/core/fpdfapi/parser/cpdf_crypto_handler.cpp
index 86f2ec7..91f58e6 100644
--- a/core/fpdfapi/parser/cpdf_crypto_handler.cpp
+++ b/core/fpdfapi/parser/cpdf_crypto_handler.cpp

@@ -69,7 +69,9 @@
   }
   if (m_Cipher == Cipher::kAES) {
     CRYPT_AESSetKey(m_pAESContext.get(),
-                    m_KeyLen == 32 ? m_EncryptKey.data() : realkey, m_KeyLen);
+                    m_KeyLen == 32
+                        ? pdfium::span<const uint8_t>(m_EncryptKey)
+                        : pdfium::span<const uint8_t>(realkey).first(m_KeyLen));
 
     constexpr size_t kIVSize = 16;
     constexpr size_t kPaddingSize = 16;
@@ -85,8 +87,7 @@
     for (auto& v : dest_iv_span) {
       v = static_cast<uint8_t>(rand());
     }
-    CRYPT_AESSetIV(m_pAESContext.get(), dest_iv_span.data());
-
+    CRYPT_AESSetIV(m_pAESContext.get(), dest_iv_span);
     CRYPT_AESEncrypt(m_pAESContext.get(), dest_data_span,
                      source.first(source_data_size));
 
@@ -117,7 +118,7 @@
     AESCryptContext* pContext = FX_Alloc(AESCryptContext, 1);
     pContext->m_bIV = true;
     pContext->m_BlockOffset = 0;
-    CRYPT_AESSetKey(&pContext->m_Context, m_EncryptKey.data(), 32);
+    CRYPT_AESSetKey(&pContext->m_Context, m_EncryptKey);
     return pContext;
   }
   uint8_t key1[48];
@@ -136,7 +137,7 @@
     AESCryptContext* pContext = FX_Alloc(AESCryptContext, 1);
     pContext->m_bIV = true;
     pContext->m_BlockOffset = 0;
-    CRYPT_AESSetKey(&pContext->m_Context, realkey, 16);
+    CRYPT_AESSetKey(&pContext->m_Context, realkey);
     return pContext;
   }
   CRYPT_rc4_context* pContext = FX_Alloc(CRYPT_rc4_context, 1);

diff --git a/core/fpdfapi/parser/cpdf_security_handler.cpp b/core/fpdfapi/parser/cpdf_security_handler.cpp
index c03d5db..ec3694a 100644
--- a/core/fpdfapi/parser/cpdf_security_handler.cpp
+++ b/core/fpdfapi/parser/cpdf_security_handler.cpp

@@ -150,12 +150,8 @@
       }
     }
     CHECK_EQ(content.size(), encrypted_output.size());
-    {
-      pdfium::span<uint8_t> key = input.first<16u>();
-      pdfium::span<uint8_t> iv = input.subspan<16u>();
-      CRYPT_AESSetKey(&aes, key.data(), 16);
-      CRYPT_AESSetIV(&aes, iv.data());
-    }
+    CRYPT_AESSetKey(&aes, input.first<16u>());
+    CRYPT_AESSetIV(&aes, input.subspan<16u>());
     CRYPT_AESEncrypt(&aes, encrypted_output_span, content);
 
     input = pdfium::span<uint8_t>();  // Dangling after assignments below.
@@ -366,11 +362,11 @@
     return false;
 
   CRYPT_aes_context aes = {};
-  CRYPT_AESSetKey(&aes, digest, sizeof(digest));
+  CRYPT_AESSetKey(&aes, digest);
   uint8_t iv[16] = {};
   CRYPT_AESSetIV(&aes, iv);
   CRYPT_AESDecrypt(&aes, m_EncryptKey.data(), ekey.unsigned_str(), 32);
-  CRYPT_AESSetKey(&aes, m_EncryptKey.data(), m_EncryptKey.size());
+  CRYPT_AESSetKey(&aes, m_EncryptKey);
   CRYPT_AESSetIV(&aes, iv);
   ByteString perms = m_pEncryptDict->GetByteStringFor("Perms");
   if (perms.IsEmpty())
@@ -648,7 +644,7 @@
     CRYPT_SHA256Finish(&sha2, pdfium::make_span(digest1).first<32>());
   }
   CRYPT_aes_context aes = {};
-  CRYPT_AESSetKey(&aes, digest1, 32);
+  CRYPT_AESSetKey(&aes, pdfium::make_span(digest1).first<32u>());
   uint8_t iv[16] = {};
   CRYPT_AESSetIV(&aes, iv);
   CRYPT_AESEncrypt(&aes, digest1, m_EncryptKey);
@@ -679,7 +675,7 @@
               pdfium::make_span(buf).subspan<12, 4>());
 
   CRYPT_aes_context aes = {};
-  CRYPT_AESSetKey(&aes, m_EncryptKey.data(), m_EncryptKey.size());
+  CRYPT_AESSetKey(&aes, m_EncryptKey);
 
   uint8_t iv[16] = {};
   CRYPT_AESSetIV(&aes, iv);