Bug 1757833 - Extract methods ScriptBytecodeCompress and ScriptBytecodeDecompress r=nbp

This leaves the code in ScriptLoader and ScriptLoadHandler a lot more readable.

ScriptBytecodeCompressedDataLayout and ScriptBytecodeDataLayout simplify
locating data in the ScriptLoadRequest bytecode buffer when compressing and
decompressing it.

The interface is still error-prone. For example, these classes don't check
that the returned pointers are within the bounds of the buffer.

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

dom/script/ScriptCompression.cpp

@@ -0,0 +1,160 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "zlib.h"
+#include "ScriptLoadRequest.h"
+#include "ScriptLoader.h"
+#include "mozilla/Vector.h"
+#include "mozilla/ScopeExit.h"
+
+using namespace mozilla;
+
+namespace JS::loader {
+
+#undef LOG
+#define LOG(args)                                       \
+  MOZ_LOG(mozilla::dom::ScriptLoader::gScriptLoaderLog, \
+          mozilla::LogLevel::Debug, args)
+
+/*
+ * ScriptBytecodeDataLayout
+ *
+ * ScriptBytecodeDataLayout provides accessors to maintain the correct data
+ * layout of the ScriptLoadRequest's script bytecode buffer.
+ */
+class ScriptBytecodeDataLayout {
+ public:
+  explicit ScriptBytecodeDataLayout(mozilla::Vector<uint8_t>& aBytecode,
+                                    size_t aBytecodeOffset)
+      : mBytecode(aBytecode), mBytecodeOffset(aBytecodeOffset) {}
+
+  uint8_t* prelude() const { return mBytecode.begin(); }
+  size_t preludeLength() const { return mBytecodeOffset; }
+
+  uint8_t* bytecode() const { return prelude() + mBytecodeOffset; }
+  size_t bytecodeLength() const { return mBytecode.length() - preludeLength(); }
+
+  mozilla::Vector<uint8_t>& mBytecode;
+  size_t mBytecodeOffset;
+};
+
+/*
+ * ScriptBytecodeCompressedDataLayout
+ *
+ * ScriptBytecodeCompressedDataLayout provides accessors to maintain the correct
+ * data layout of a compressed script bytecode buffer.
+ */
+class ScriptBytecodeCompressedDataLayout {
+ public:
+  using UncompressedLengthType = uint32_t;
+
+  explicit ScriptBytecodeCompressedDataLayout(
+      mozilla::Vector<uint8_t>& aBytecode, size_t aBytecodeOffset)
+      : mBytecode(aBytecode), mBytecodeOffset(aBytecodeOffset) {}
+
+  uint8_t* prelude() const { return mBytecode.begin(); }
+  size_t preludeLength() const { return mBytecodeOffset; }
+
+  uint8_t* uncompressedLength() const { return prelude() + mBytecodeOffset; }
+  size_t uncompressedLengthLength() const {
+    return sizeof(UncompressedLengthType);
+  }
+
+  uint8_t* bytecode() const {
+    return uncompressedLength() + uncompressedLengthLength();
+  }
+  size_t bytecodeLength() const {
+    return mBytecode.length() - uncompressedLengthLength() - preludeLength();
+  }
+
+  mozilla::Vector<uint8_t>& mBytecode;
+  size_t mBytecodeOffset;
+};
+
+bool ScriptBytecodeCompress(Vector<uint8_t>& aBytecodeBuf,
+                            size_t aBytecodeOffset,
+                            Vector<uint8_t>& aCompressedBytecodeBufOut) {
+  // TODO probably need to move this to a helper thread
+  ScriptBytecodeDataLayout uncompressedLayout(aBytecodeBuf, aBytecodeOffset);
+  ScriptBytecodeCompressedDataLayout compressedLayout(
+      aCompressedBytecodeBufOut, uncompressedLayout.preludeLength());
+  ScriptBytecodeCompressedDataLayout::UncompressedLengthType
+      uncompressedLength = uncompressedLayout.bytecodeLength();
+  z_stream zstream{.next_in = uncompressedLayout.bytecode(),
+                   .avail_in = uncompressedLength};
+  auto compressedLength = deflateBound(&zstream, uncompressedLength);
+  if (!aCompressedBytecodeBufOut.resizeUninitialized(
+          compressedLength + compressedLayout.preludeLength() +
+          compressedLayout.uncompressedLengthLength())) {
+    return false;
+  }
+  memcpy(compressedLayout.prelude(), uncompressedLayout.prelude(),
+         uncompressedLayout.preludeLength());
+  memcpy(compressedLayout.uncompressedLength(), &uncompressedLength,
+         sizeof(uncompressedLength));
+  zstream.next_out = compressedLayout.bytecode();
+  zstream.avail_out = compressedLength;
+
+  const int COMPRESSION = 2;  // TODO find appropriate compression level
+  if (deflateInit(&zstream, COMPRESSION) != Z_OK) {
+    LOG(
+        ("ScriptLoadRequest: Unable to initialize bytecode cache "
+         "compression."));
+    return false;
+  }
+  auto autoDestroy = MakeScopeExit([&]() { deflateEnd(&zstream); });
+
+  int ret = deflate(&zstream, Z_FINISH);
+  if (ret == Z_MEM_ERROR) {
+    return false;
+  }
+  MOZ_RELEASE_ASSERT(ret == Z_STREAM_END);
+
+  aCompressedBytecodeBufOut.shrinkTo(zstream.next_out -
+                                     aCompressedBytecodeBufOut.begin());
+  return true;
+}
+
+bool ScriptBytecodeDecompress(Vector<uint8_t>& aCompressedBytecodeBuf,
+                              size_t aBytecodeOffset,
+                              Vector<uint8_t>& aBytecodeBufOut) {
+  ScriptBytecodeDataLayout uncompressedLayout(aBytecodeBufOut, aBytecodeOffset);
+  ScriptBytecodeCompressedDataLayout compressedLayout(
+      aCompressedBytecodeBuf, uncompressedLayout.preludeLength());
+  ScriptBytecodeCompressedDataLayout::UncompressedLengthType uncompressedLength;
+  memcpy(&uncompressedLength, compressedLayout.uncompressedLength(),
+         compressedLayout.uncompressedLengthLength());
+  if (!aBytecodeBufOut.resizeUninitialized(uncompressedLayout.preludeLength() +
+                                           uncompressedLength)) {
+    return false;
+  }
+  memcpy(uncompressedLayout.prelude(), compressedLayout.prelude(),
+         compressedLayout.preludeLength());
+
+  z_stream zstream{nullptr};
+  zstream.next_in = compressedLayout.bytecode();
+  zstream.avail_in = static_cast<uint32_t>(compressedLayout.bytecodeLength());
+  zstream.next_out = uncompressedLayout.bytecode();
+  zstream.avail_out = uncompressedLength;
+  if (inflateInit(&zstream) != Z_OK) {
+    LOG(("ScriptLoadRequest: inflateInit FAILED (%s)", zstream.msg));
+    return false;
+  }
+  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(("ScriptLoadReques: inflate FAILED (%s)", zstream.msg));
+    return false;
+  }
+
+  return true;
+}
+
+#undef LOG
+
+}  // namespace JS::loader

dom/script/ScriptCompression.h

@@ -0,0 +1,43 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#ifndef js_loader_ScriptCompression_h
+#define js_loader_ScriptCompression_h
+
+#include "ErrorList.h"
+#include "mozilla/Vector.h"
+
+namespace JS::loader {
+
+class ScriptLoadRequest;
+
+/**
+ * Compress the bytecode stored in a buffer. All data before the bytecode is
+ * copied into the output buffer without modification.
+ *
+ * @param aBytecodeBuf buffer containing the uncompressed bytecode
+ * @param aBytecodeOffset offset of the bytecode in the buffer
+ * @param aCompressedBytecodeBufOut buffer to store the compressed bytecode in
+ */
+bool ScriptBytecodeCompress(
+    mozilla::Vector<uint8_t>& aBytecodeBuf, size_t aBytecodeOffset,
+    mozilla::Vector<uint8_t>& aCompressedBytecodeBufOut);
+
+/**
+ * Uncompress the bytecode stored in a buffer. All data before the bytecode is
+ * copied into the output buffer without modification.
+ *
+ * @param aCompressedBytecodeBuf buffer containing the compressed bytecode
+ * @param aBytecodeOffset offset of the bytecode in the buffer
+ * @param aBytecodeBufOut buffer to store the uncompressed bytecode in
+ */
+bool ScriptBytecodeDecompress(mozilla::Vector<uint8_t>& aCompressedBytecodeBuf,
+                              size_t aBytecodeOffset,
+                              mozilla::Vector<uint8_t>& aBytecodeBufOut);
+
+}  // namespace JS::loader
+
+#endif  // js_loader_ScriptCompression_h

dom/script/ScriptLoadHandler.cpp

@@ -8,6 +8,7 @@
 
 #include <stdlib.h>
 #include <utility>
+#include "ScriptCompression.h"
 #include "ScriptLoader.h"
 #include "ScriptTrace.h"
 #include "js/Transcoding.h"
@@ -40,14 +41,6 @@
 #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 {
 
 #undef LOG
@@ -410,45 +403,14 @@ 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;
-        }
+      Vector<uint8_t> compressedBytecode;
+      // mRequest has the compressed bytecode, but will be filled with the
+      // uncompressed bytecode
+      compressedBytecode.swap(mRequest->mScriptBytecode);
+      if (!JS::loader::ScriptBytecodeDecompress(compressedBytecode,
+                                                mRequest->mBytecodeOffset,
+                                                mRequest->mScriptBytecode)) {
+        return NS_ERROR_UNEXPECTED;
       }
     }
   }

dom/script/ScriptLoader.cpp

@@ -19,6 +19,7 @@
 #include "js/ContextOptions.h"        // JS::ContextOptionsRef
 #include "js/friend/ErrorMessages.h"  // js::GetErrorMessage, JSMSG_*
 #include "js/loader/ScriptLoadRequest.h"
+#include "ScriptCompression.h"
 #include "js/loader/LoadedScript.h"
 #include "js/loader/ModuleLoadRequest.h"
 #include "js/MemoryFunctions.h"
@@ -100,14 +101,6 @@ 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");
@@ -2563,43 +2556,10 @@ void ScriptLoader::EncodeRequestBytecode(JSContext* aCx,
   }
 
   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());
+  // TODO probably need to move this to a helper thread
+  if (!ScriptBytecodeCompress(aRequest->mScriptBytecode,
+                              aRequest->mBytecodeOffset, compressedBytecode)) {
+    return;
   }
 
   if (compressedBytecode.length() >= UINT32_MAX) {

dom/script/moz.build

@@ -20,6 +20,7 @@ EXPORTS += [
 EXPORTS.mozilla.dom += [
     "AutoEntryScript.h",
     "ModuleLoader.h",
+    "ScriptCompression.h",
     "ScriptDecoding.h",
     "ScriptElement.h",
     "ScriptLoadContext.h",
@@ -32,6 +33,7 @@ UNIFIED_SOURCES += [
     "AutoEntryScript.cpp",
     "ModuleLoader.cpp",
     "nsIScriptElement.cpp",
+    "ScriptCompression.cpp",
     "ScriptElement.cpp",
     "ScriptLoadContext.cpp",
     "ScriptLoader.cpp",

js/loader/ScriptLoadRequest.h

@@ -325,7 +325,8 @@ class ScriptLoadRequest
   size_t mScriptTextLength;
 
   // Holds the SRI serialized hash and the script bytecode for non-inline
-  // scripts.
+  // scripts. The data is laid out according to ScriptBytecodeDataLayout
+  // or, if compression is enabled, ScriptBytecodeCompressedDataLayout.
   mozilla::Vector<uint8_t> mScriptBytecode;
   uint32_t mBytecodeOffset;  // Offset of the bytecode in mScriptBytecode