Bug 1757833 - Compress Stencil bytecode before writing to cache r=nbp

The SRI hash at the beginning of ScriptLoadRequest::mScriptBytecode is left
uncompressed because ScriptLoader::OnIncrementalData() tries to decode it
as soon as enough data is read (instead of waiting until OnStreamComplete()).

ScriptLoader writes the length of the uncompressed bytecode to the buffer
to make it easy for ScriptLoadHandler to allocate an buffer of the right size
to decompress the bytecode.

These changes are based on the bytecode compression implemented for WASM in
dom/fetch/FetchUtil.cpp.

Differential Revision: https://phabricator.services.mozilla.com/D141524

dom/script/ScriptLoadHandler.cpp

@@ -38,6 +38,15 @@
 #include "nsMimeTypes.h"
 #include "nsString.h"
 #include "nsTArray.h"
+#include "zlib.h"
+
+namespace {
+// A LengthPrefixType is stored at the start of the compressed optimized
+// encoding, allowing the decompressed buffer to be allocated to exactly
+// the right size.
+using LengthPrefixType = uint32_t;
+const unsigned PREFIX_BYTES = sizeof(LengthPrefixType);
+}  // namespace
 
 namespace mozilla::dom {
 
@@ -400,6 +409,47 @@ ScriptLoadHandler::OnStreamComplete(nsIIncrementalStreamLoader* aLoader,
       }
 
       mRequest->mBytecodeOffset = JS::AlignTranscodingBytecodeOffset(sriLength);
+
+      {
+        Vector<uint8_t>
+            compressedBytecode;  // starts with SRI hash, followed by length
+                                 // prefix, then compressed bytecode
+        compressedBytecode.swap(mRequest->mScriptBytecode);
+
+        LengthPrefixType uncompressedLength;
+        memcpy(&uncompressedLength,
+               compressedBytecode.begin() + mRequest->mBytecodeOffset,
+               PREFIX_BYTES);
+        if (!mRequest->mScriptBytecode.resizeUninitialized(
+                mRequest->mBytecodeOffset + uncompressedLength)) {
+          return NS_ERROR_OUT_OF_MEMORY;
+        }
+        memcpy(mRequest->mScriptBytecode.begin(), compressedBytecode.begin(),
+               mRequest->mBytecodeOffset);  // SRI hash
+
+        z_stream zstream{.next_in = compressedBytecode.begin() +
+                                    mRequest->mBytecodeOffset + PREFIX_BYTES,
+                         .avail_in = static_cast<uint32_t>(
+                             compressedBytecode.length() -
+                             mRequest->mBytecodeOffset - PREFIX_BYTES),
+                         .next_out = mRequest->mScriptBytecode.begin() +
+                                     mRequest->mBytecodeOffset,
+                         .avail_out = uncompressedLength};
+        if (inflateInit(&zstream) != Z_OK) {
+          LOG(("ScriptLoadRequest (%p): inflateInit FAILED (%s)",
+               mRequest.get(), zstream.msg));
+          return nsresult::NS_ERROR_UNEXPECTED;
+        }
+        auto autoDestroy = MakeScopeExit([&]() { inflateEnd(&zstream); });
+
+        int ret = inflate(&zstream, Z_NO_FLUSH);
+        bool ok = (ret == Z_OK || ret == Z_STREAM_END) && zstream.avail_in == 0;
+        if (!ok) {
+          LOG(("ScriptLoadRequest (%p): inflate FAILED (%s)", mRequest.get(),
+               zstream.msg));
+          return NS_ERROR_UNEXPECTED;
+        }
+      }
     }
   }
 

dom/script/ScriptLoader.cpp

@@ -9,6 +9,8 @@
 #include "ScriptTrace.h"
 #include "ModuleLoader.h"
 
+#include "zlib.h"
+
 #include "prsystem.h"
 #include "jsapi.h"
 #include "jsfriendapi.h"
@@ -98,6 +100,14 @@ using namespace JS::loader;
 
 using mozilla::Telemetry::LABELS_DOM_SCRIPT_PRELOAD_RESULT;
 
+namespace {  // TODO shared code with ScriptLoadHandler.cpp
+// A LengthPrefixType is stored at the start of the compressed optimized
+// encoding, allowing the decompressed buffer to be allocated to exactly
+// the right size.
+using LengthPrefixType = uint32_t;
+const unsigned PREFIX_BYTES = sizeof(LengthPrefixType);
+}  // namespace
+
 namespace mozilla::dom {
 
 LazyLogModule ScriptLoader::gCspPRLog("CSP");
@@ -2552,7 +2562,47 @@ void ScriptLoader::EncodeRequestBytecode(JSContext* aCx,
     return;
   }
 
-  if (aRequest->mScriptBytecode.length() >= UINT32_MAX) {
+  Vector<uint8_t> compressedBytecode;
+  {
+    // TODO probably need to move this to a helper thread
+    LengthPrefixType uncompressedLength =
+        aRequest->mScriptBytecode.length() - aRequest->mBytecodeOffset;
+    z_stream zstream{.next_in = aRequest->mScriptBytecode.begin() +
+                                aRequest->mBytecodeOffset,
+                     .avail_in = uncompressedLength};
+    auto compressedLength = deflateBound(&zstream, uncompressedLength);
+    if (!compressedBytecode.resizeUninitialized(
+            compressedLength + aRequest->mBytecodeOffset + PREFIX_BYTES)) {
+      return;
+    }
+    memcpy(compressedBytecode.begin(), aRequest->mScriptBytecode.begin(),
+           aRequest->mBytecodeOffset);
+    memcpy(compressedBytecode.begin() + aRequest->mBytecodeOffset,
+           &uncompressedLength, PREFIX_BYTES);
+    zstream.next_out =
+        compressedBytecode.begin() + aRequest->mBytecodeOffset + PREFIX_BYTES;
+    zstream.avail_out = compressedLength;
+
+    const int COMPRESSION = 2;  // TODO find appropriate compression level
+    if (deflateInit(&zstream, COMPRESSION) != Z_OK) {
+      LOG(
+          ("ScriptLoadRequest (%p): Unable to initialize bytecode cache "
+           "compression.",
+           aRequest));
+      return;
+    }
+    auto autoDestroy = MakeScopeExit([&]() { deflateEnd(&zstream); });
+
+    int ret = deflate(&zstream, Z_FINISH);
+    if (ret == Z_MEM_ERROR) {
+      return;
+    }
+    MOZ_RELEASE_ASSERT(ret == Z_STREAM_END);
+
+    compressedBytecode.shrinkTo(zstream.next_out - compressedBytecode.begin());
+  }
+
+  if (compressedBytecode.length() >= UINT32_MAX) {
     LOG(
         ("ScriptLoadRequest (%p): Bytecode cache is too large to be decoded "
          "correctly.",
@@ -2565,7 +2615,8 @@ void ScriptLoader::EncodeRequestBytecode(JSContext* aCx,
   // case, we just ignore the current one.
   nsCOMPtr<nsIAsyncOutputStream> output;
   rv = aRequest->mCacheInfo->OpenAlternativeOutputStream(
-      BytecodeMimeTypeFor(aRequest), aRequest->mScriptBytecode.length(),
+      BytecodeMimeTypeFor(aRequest),
+      static_cast<int64_t>(compressedBytecode.length()),
       getter_AddRefs(output));
   if (NS_FAILED(rv)) {
     LOG(
@@ -2582,17 +2633,17 @@ void ScriptLoader::EncodeRequestBytecode(JSContext* aCx,
   });
 
   uint32_t n;
-  rv = output->Write(reinterpret_cast<char*>(aRequest->mScriptBytecode.begin()),
-                     aRequest->mScriptBytecode.length(), &n);
-  LOG((
-      "ScriptLoadRequest (%p): Write bytecode cache (rv = %X, length = %u, "
+  rv = output->Write(reinterpret_cast<char*>(compressedBytecode.begin()),
+                     compressedBytecode.length(), &n);
+  LOG(
+      ("ScriptLoadRequest (%p): Write bytecode cache (rv = %X, length = %u, "
        "written = %u)",
-      aRequest, unsigned(rv), unsigned(aRequest->mScriptBytecode.length()), n));
+       aRequest, unsigned(rv), unsigned(compressedBytecode.length()), n));
   if (NS_FAILED(rv)) {
     return;
   }
 
-  MOZ_RELEASE_ASSERT(aRequest->mScriptBytecode.length() == n);
+  MOZ_RELEASE_ASSERT(compressedBytecode.length() == n);
 
   bytecodeFailed.release();
   TRACE_FOR_TEST_NONE(aRequest->GetScriptLoadContext()->GetScriptElement(),