From c1c9b8ae4b1a0d37acd45daa727deb2e5bbc8b7a Mon Sep 17 00:00:00 2001
From: Jan Varga <jvarga@mozilla.com>
Date: Wed, 19 Jun 2024 08:49:22 +0000
Subject: [PATCH] Bug 1901527 - Query and store decrypted stream size in
 DecryptingInputStream only once; r=hsingh

Differential Revision: https://phabricator.services.mozilla.com/D213125
---
 dom/quota/DecryptingInputStream.h      |   5 ++
 dom/quota/DecryptingInputStream_impl.h | 113 ++++++++++++++-----------
 2 files changed, 69 insertions(+), 49 deletions(-)

diff --git a/dom/quota/DecryptingInputStream.h b/dom/quota/DecryptingInputStream.h
index 73c7eeb60785..3207cfa1f830 100644
--- a/dom/quota/DecryptingInputStream.h
+++ b/dom/quota/DecryptingInputStream.h
@@ -143,6 +143,9 @@ class DecryptingInputStream final : public DecryptingInputStreamBase {
 
   bool EnsureBuffers();
 
+  // This method may change the current position in the stream.
+  nsresult EnsureDecryptedStreamSize();
+
   CipherStrategy mCipherStrategy;
   LazyInitializedOnce<const typename CipherStrategy::KeyType> mKey;
 
@@ -154,6 +157,8 @@ class DecryptingInputStream final : public DecryptingInputStreamBase {
 
   // Buffer storing the resulting plain data.
   nsTArray<uint8_t> mPlainBuffer;
+
+  LazyInitializedOnce<const int64_t> mDecryptedStreamSize;
 };
 
 }  // namespace mozilla::dom::quota
diff --git a/dom/quota/DecryptingInputStream_impl.h b/dom/quota/DecryptingInputStream_impl.h
index 669a648f97f5..aae50e1a89da 100644
--- a/dom/quota/DecryptingInputStream_impl.h
+++ b/dom/quota/DecryptingInputStream_impl.h
@@ -264,6 +264,63 @@ bool DecryptingInputStream<CipherStrategy>::EnsureBuffers() {
   return true;
 }
 
+template <typename CipherStrategy>
+nsresult DecryptingInputStream<CipherStrategy>::EnsureDecryptedStreamSize() {
+  if (mDecryptedStreamSize) {
+    return NS_OK;
+  }
+
+  auto decryptedStreamSizeOrErr = [this]() -> Result<int64_t, nsresult> {
+    nsresult rv = (*mBaseSeekableStream)->Seek(NS_SEEK_SET, 0);
+    if (NS_WARN_IF(NS_FAILED(rv))) {
+      return Err(rv);
+    }
+
+    uint64_t baseStreamSize;
+    rv = (*mBaseStream)->Available(&baseStreamSize);
+    if (NS_WARN_IF(NS_FAILED(rv))) {
+      return Err(rv);
+    }
+
+    if (!baseStreamSize) {
+      return 0;
+    }
+
+    rv = (*mBaseSeekableStream)
+             ->Seek(NS_SEEK_END, -static_cast<int64_t>(*mBlockSize));
+    if (NS_WARN_IF(NS_FAILED(rv))) {
+      return Err(rv);
+    }
+
+    uint32_t bytesRead;
+    rv = ParseNextChunk(&bytesRead);
+    if (NS_WARN_IF(NS_FAILED(rv))) {
+      return Err(rv);
+    }
+    MOZ_ASSERT(bytesRead);
+
+    mPlainBytes = bytesRead;
+
+    mNextByte = bytesRead;
+
+    int64_t current;
+    rv = Tell(&current);
+    if (NS_WARN_IF(NS_FAILED(rv))) {
+      return Err(rv);
+    }
+
+    return current;
+  }();
+
+  if (decryptedStreamSizeOrErr.isErr()) {
+    return decryptedStreamSizeOrErr.unwrapErr();
+  }
+
+  mDecryptedStreamSize.init(decryptedStreamSizeOrErr.inspect());
+
+  return NS_OK;
+}
+
 template <typename CipherStrategy>
 NS_IMETHODIMP DecryptingInputStream<CipherStrategy>::Tell(
     int64_t* const aRetval) {
@@ -314,8 +371,8 @@ NS_IMETHODIMP DecryptingInputStream<CipherStrategy>::Seek(const int32_t aWhence,
     return Err(rv);
   }
 
-  // Can't call this just in NS_SEEK_CUR case, because getting the decrypted
-  // size below changes the current position.
+  // Can't call this just in NS_SEEK_CUR case, because ensuring the decrypted
+  // size below may change the current position.
   int64_t current;
   rv = Tell(&current);
   if (NS_WARN_IF(NS_FAILED(rv))) {
@@ -334,51 +391,9 @@ NS_IMETHODIMP DecryptingInputStream<CipherStrategy>::Seek(const int32_t aWhence,
         nextByte = savedNextByte;
       });
 
-  // XXX The size of the stream could also be queried and stored once only.
-  auto decryptedStreamSizeOrErr = [this]() -> Result<int64_t, nsresult> {
-    nsresult rv = (*mBaseSeekableStream)->Seek(NS_SEEK_SET, 0);
-    if (NS_WARN_IF(NS_FAILED(rv))) {
-      return Err(rv);
-    }
-
-    uint64_t baseStreamSize;
-    rv = (*mBaseStream)->Available(&baseStreamSize);
-    if (NS_WARN_IF(NS_FAILED(rv))) {
-      return Err(rv);
-    }
-
-    if (!baseStreamSize) {
-      return 0;
-    }
-
-    rv = (*mBaseSeekableStream)
-             ->Seek(NS_SEEK_END, -static_cast<int64_t>(*mBlockSize));
-    if (NS_WARN_IF(NS_FAILED(rv))) {
-      return Err(rv);
-    }
-
-    uint32_t bytesRead;
-    rv = ParseNextChunk(&bytesRead);
-    if (NS_WARN_IF(NS_FAILED(rv))) {
-      return Err(rv);
-    }
-    MOZ_ASSERT(bytesRead);
-
-    mPlainBytes = bytesRead;
-
-    mNextByte = bytesRead;
-
-    int64_t current;
-    rv = Tell(&current);
-    if (NS_WARN_IF(NS_FAILED(rv))) {
-      return Err(rv);
-    }
-
-    return current;
-  }();
-
-  if (decryptedStreamSizeOrErr.isErr()) {
-    return decryptedStreamSizeOrErr.unwrapErr();
+  rv = EnsureDecryptedStreamSize();
+  if (NS_WARN_IF(NS_FAILED(rv))) {
+    return rv;
   }
 
   int64_t baseBlocksOffset;
@@ -394,14 +409,14 @@ NS_IMETHODIMP DecryptingInputStream<CipherStrategy>::Seek(const int32_t aWhence,
 
     case NS_SEEK_END:
       // XXX Simplify this without using Seek/Tell.
-      aOffset += decryptedStreamSizeOrErr.inspect();
+      aOffset += *mDecryptedStreamSize;
       break;
 
     default:
       return NS_ERROR_ILLEGAL_VALUE;
   }
 
-  if (aOffset < 0 || aOffset > decryptedStreamSizeOrErr.inspect()) {
+  if (aOffset < 0 || aOffset > *mDecryptedStreamSize) {
     return NS_ERROR_ILLEGAL_VALUE;
   }