1 files changed, 229 insertions, 4 deletions

diff --git a/netwerk/streamconv/converters/nsHTTPCompressConv.cpp b/netwerk/streamconv/converters/nsHTTPCompressConv.cpp
index e06c72222b..7c7404f110 100644
--- a/netwerk/streamconv/converters/nsHTTPCompressConv.cpp
+++ b/netwerk/streamconv/converters/nsHTTPCompressConv.cpp
@@ -5,6 +5,7 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
 #include "nsHTTPCompressConv.h"
+#include "ErrorList.h"
 #include "nsCOMPtr.h"
 #include "nsCRT.h"
 #include "nsError.h"
@@ -19,6 +20,8 @@
 #include "nsIForcePendingChannel.h"
 #include "nsIRequest.h"
 #include "mozilla/UniquePtrExtensions.h"
+#include "nsIThreadRetargetableRequest.h"
+#include "nsIChannel.h"
 
 // brotli headers
 #undef assert
@@ -26,6 +29,8 @@
 #include "state.h"
 #include "brotli/decode.h"
 
+#include "zstd/zstd.h"
+
 namespace mozilla {
 namespace net {
 
@@ -50,6 +55,26 @@ class BrotliWrapper {
   uint64_t mSourceOffset{0};
 };
 
+class ZstdWrapper {
+ public:
+  ZstdWrapper() {
+    mDStream = ZSTD_createDStream();
+    ZSTD_DCtx_setParameter(mDStream, ZSTD_d_windowLogMax, 23 /*8*1024*1024*/);
+  }
+  ~ZstdWrapper() {
+    if (mDStream) {
+      ZSTD_freeDStream(mDStream);
+    }
+  }
+
+  UniquePtr<uint8_t[]> mOutBuffer;
+  nsresult mStatus = NS_OK;
+  nsIRequest* mRequest{nullptr};
+  nsISupports* mContext{nullptr};
+  uint64_t mSourceOffset{0};
+  ZSTD_DStream* mDStream{nullptr};
+};
+
 // nsISupports implementation
 NS_IMPL_ISUPPORTS(nsHTTPCompressConv, nsIStreamConverter, nsIStreamListener,
                   nsIRequestObserver, nsICompressConvStats,
@@ -109,6 +134,12 @@ nsHTTPCompressConv::AsyncConvertData(const char* aFromType, const char* aToType,
   } else if (!nsCRT::strncasecmp(aFromType, HTTP_BROTLI_TYPE,
                                  sizeof(HTTP_BROTLI_TYPE) - 1)) {
     mMode = HTTP_COMPRESS_BROTLI;
+  } else if (!nsCRT::strncasecmp(aFromType, HTTP_ZSTD_TYPE,
+                                 sizeof(HTTP_ZSTD_TYPE) - 1)) {
+    mMode = HTTP_COMPRESS_ZSTD;
+  } else if (!nsCRT::strncasecmp(aFromType, HTTP_ZST_TYPE,
+                                 sizeof(HTTP_ZST_TYPE) - 1)) {
+    mMode = HTTP_COMPRESS_ZSTD;
   }
   LOG(("nsHttpCompresssConv %p AsyncConvertData %s %s mode %d\n", this,
        aFromType, aToType, (CompressMode)mMode));
@@ -128,6 +159,54 @@ nsHTTPCompressConv::GetConvertedType(const nsACString& aFromType,
 }
 
 NS_IMETHODIMP
+nsHTTPCompressConv::MaybeRetarget(nsIRequest* request) {
+  MOZ_ASSERT(NS_IsMainThread());
+  nsresult rv;
+  nsCOMPtr<nsIThreadRetargetableRequest> req = do_QueryInterface(request);
+  if (!req) {
+    return NS_ERROR_NO_INTERFACE;
+  }
+  if (!StaticPrefs::network_decompression_off_mainthread()) {
+    return NS_OK;
+  }
+  nsCOMPtr<nsISerialEventTarget> target;
+  rv = req->GetDeliveryTarget(getter_AddRefs(target));
+  if (NS_FAILED(rv) || !target || target->IsOnCurrentThread()) {
+    nsCOMPtr<nsIChannel> channel(do_QueryInterface(request));
+    int64_t length = -1;
+    if (channel) {
+      channel->GetContentLength(&length);
+      // If this fails we'll retarget
+    }
+    if (length <= 0 ||
+        length >=
+            StaticPrefs::network_decompression_off_mainthread_min_size()) {
+      LOG(("MaybeRetarget: Retargeting to background thread: Length %" PRId64,
+           length));
+      // No retargetting was performed.  Decompress off MainThread,
+      // and dispatch results back to MainThread.
+      // Don't do this if the input is small, if we know the length.
+      // If the length is 0 (unknown), always use OMT.
+      nsCOMPtr<nsISerialEventTarget> backgroundThread;
+      rv = NS_CreateBackgroundTaskQueue("nsHTTPCompressConv",
+                                        getter_AddRefs(backgroundThread));
+      NS_ENSURE_SUCCESS(rv, rv);
+      rv = req->RetargetDeliveryTo(backgroundThread);
+      NS_ENSURE_SUCCESS(rv, rv);
+      if (NS_SUCCEEDED(rv)) {
+        mDispatchToMainThread = true;
+      }
+    } else {
+      LOG(("MaybeRetarget: Not retargeting: Length %" PRId64, length));
+    }
+  } else {
+    LOG(("MaybeRetarget: Don't need to retarget"));
+  }
+
+  return NS_OK;
+}
+
+NS_IMETHODIMP
 nsHTTPCompressConv::OnStartRequest(nsIRequest* request) {
   LOG(("nsHttpCompresssConv %p onstart\n", this));
   nsCOMPtr<nsIStreamListener> listener;
@@ -135,14 +214,33 @@ nsHTTPCompressConv::OnStartRequest(nsIRequest* request) {
     MutexAutoLock lock(mMutex);
     listener = mListener;
   }
-  return listener->OnStartRequest(request);
+  nsresult rv = listener->OnStartRequest(request);
+  if (NS_SUCCEEDED(rv)) {
+    if (XRE_IsContentProcess()) {
+      nsCOMPtr<nsIThreadRetargetableStreamListener> retargetlistener =
+          do_QueryInterface(listener);
+      // |nsHTTPCompressConv| should *always* be dispatched off of the main
+      // thread from a content process, even if its listeners don't support it.
+      //
+      // If its listener chain does not support being retargeted off of the
+      // main thread, it will be dispatched back to the main thread in
+      // |do_OnDataAvailable| and |OnStopRequest|.
+      if (!retargetlistener ||
+          NS_FAILED(retargetlistener->CheckListenerChain())) {
+        mDispatchToMainThread = true;
+      }
+    }
+  }
+  return rv;
 }
 
 NS_IMETHODIMP
 nsHTTPCompressConv::OnStopRequest(nsIRequest* request, nsresult aStatus) {
   nsresult status = aStatus;
-  LOG(("nsHttpCompresssConv %p onstop %" PRIx32 "\n", this,
-       static_cast<uint32_t>(aStatus)));
+  // Bug 1886237 : TRRServiceChannel calls OnStopRequest OMT
+  // MOZ_ASSERT(NS_IsMainThread());
+  LOG(("nsHttpCompresssConv %p onstop %" PRIx32 " mDispatchToMainThread %d\n",
+       this, static_cast<uint32_t>(aStatus), mDispatchToMainThread));
 
   // Framing integrity is enforced for content-encoding: gzip, but not for
   // content-encoding: deflate. Note that gzip vs deflate is NOT determined
@@ -181,6 +279,7 @@ nsHTTPCompressConv::OnStopRequest(nsIRequest* request, nsresult aStatus) {
     MutexAutoLock lock(mMutex);
     listener = mListener;
   }
+
   return listener->OnStopRequest(request, status);
 }
 
@@ -292,6 +391,71 @@ nsresult nsHTTPCompressConv::BrotliHandler(nsIInputStream* stream,
   return self->mBrotli->mStatus;
 }
 
+/* static */
+nsresult nsHTTPCompressConv::ZstdHandler(nsIInputStream* stream, void* closure,
+                                         const char* dataIn, uint32_t,
+                                         uint32_t aAvail, uint32_t* countRead) {
+  MOZ_ASSERT(stream);
+  nsHTTPCompressConv* self = static_cast<nsHTTPCompressConv*>(closure);
+  *countRead = 0;
+
+  const size_t kOutSize = ZSTD_DStreamOutSize();  // normally 128K
+  uint8_t* outPtr;
+  size_t avail = aAvail;
+
+  // Stop decompressing after an error
+  if (self->mZstd->mStatus != NS_OK) {
+    *countRead = aAvail;
+    return NS_OK;
+  }
+
+  if (!self->mZstd->mOutBuffer) {
+    self->mZstd->mOutBuffer = MakeUniqueFallible<uint8_t[]>(kOutSize);
+    if (!self->mZstd->mOutBuffer) {
+      self->mZstd->mStatus = NS_ERROR_OUT_OF_MEMORY;
+      return self->mZstd->mStatus;
+    }
+  }
+  ZSTD_inBuffer inBuffer = {.src = dataIn, .size = aAvail, .pos = 0};
+  uint32_t last_pos = 0;
+  while (inBuffer.pos < inBuffer.size) {
+    outPtr = self->mZstd->mOutBuffer.get();
+
+    LOG(("nsHttpCompresssConv %p zstdhandler decompress %zu\n", self, avail));
+    // Use ZSTD_(de)compressStream to (de)compress the input buffer into the
+    // output buffer, and fill aReadCount with the number of bytes consumed.
+    ZSTD_outBuffer outBuffer{.dst = outPtr, .size = kOutSize};
+    size_t result;
+    bool output_full;
+    do {
+      outBuffer.pos = 0;
+      result =
+          ZSTD_decompressStream(self->mZstd->mDStream, &outBuffer, &inBuffer);
+
+      // If we errored when writing, flag this and abort writing.
+      if (ZSTD_isError(result)) {
+        self->mZstd->mStatus = NS_ERROR_INVALID_CONTENT_ENCODING;
+        return self->mZstd->mStatus;
+      }
+
+      nsresult rv = self->do_OnDataAvailable(
+          self->mZstd->mRequest, self->mZstd->mSourceOffset,
+          reinterpret_cast<const char*>(outPtr), outBuffer.pos);
+      if (NS_FAILED(rv)) {
+        self->mZstd->mStatus = rv;
+        return rv;
+      }
+      self->mZstd->mSourceOffset += inBuffer.pos - last_pos;
+      last_pos = inBuffer.pos;
+      output_full = outBuffer.pos == outBuffer.size;
+      // in the unlikely case that the output buffer was full, loop to
+      // drain it before processing more input
+    } while (output_full);
+  }
+  *countRead = inBuffer.pos;
+  return NS_OK;
+}
+
 NS_IMETHODIMP
 nsHTTPCompressConv::OnDataAvailable(nsIRequest* request, nsIInputStream* iStr,
                                     uint64_t aSourceOffset, uint32_t aCount) {
@@ -525,6 +689,25 @@ nsHTTPCompressConv::OnDataAvailable(nsIRequest* request, nsIInputStream* iStr,
       }
     } break;
 
+    case HTTP_COMPRESS_ZSTD: {
+      if (!mZstd) {
+        mZstd = MakeUnique<ZstdWrapper>();
+      }
+
+      mZstd->mRequest = request;
+      mZstd->mContext = nullptr;
+      mZstd->mSourceOffset = aSourceOffset;
+
+      uint32_t countRead;
+      rv = iStr->ReadSegments(ZstdHandler, this, streamLen, &countRead);
+      if (NS_SUCCEEDED(rv)) {
+        rv = mZstd->mStatus;
+      }
+      if (NS_FAILED(rv)) {
+        return rv;
+      }
+    } break;
+
     default:
       nsCOMPtr<nsIStreamListener> listener;
       {
@@ -553,6 +736,35 @@ nsresult nsHTTPCompressConv::do_OnDataAvailable(nsIRequest* request,
                                                 uint64_t offset,
                                                 const char* buffer,
                                                 uint32_t count) {
+  LOG(("nsHttpCompressConv %p do_OnDataAvailable mDispatchToMainThread %d",
+       this, mDispatchToMainThread));
+  if (mDispatchToMainThread && !NS_IsMainThread()) {
+    nsCOMPtr<nsIInputStream> stream;
+    MOZ_TRY(NS_NewByteInputStream(getter_AddRefs(stream), Span(buffer, count),
+                                  nsAssignmentType::NS_ASSIGNMENT_COPY));
+
+    nsCOMPtr<nsIStreamListener> listener;
+    {
+      MutexAutoLock lock(mMutex);
+      listener = mListener;
+    }
+
+    // This is safe and will always run before OnStopRequest, because
+    // ChanneleventQueue means that we can't enqueue OnStopRequest until after
+    // the OMT OnDataAvailable call has completed.  So Dispatching here will
+    // ensure it's in the MainThread event queue before OnStopRequest
+    nsCOMPtr<nsIRunnable> handler = NS_NewRunnableFunction(
+        "nsHTTPCompressConv::do_OnDataAvailable",
+        [request{RefPtr<nsIRequest>(request)}, stream{std::move(stream)},
+         listener{std::move(listener)}, offset, count]() {
+          LOG(("nsHttpCompressConv Calling OnDataAvailable on Mainthread"));
+          Unused << listener->OnDataAvailable(request, stream, offset, count);
+        });
+
+    mDecodedDataLength += count;
+    return NS_DispatchToMainThread(handler);
+  }
+
   if (!mStream) {
     mStream = do_CreateInstance(NS_STRINGINPUTSTREAM_CONTRACTID);
     NS_ENSURE_STATE(mStream);
@@ -725,12 +937,16 @@ uint32_t nsHTTPCompressConv::check_header(nsIInputStream* iStr,
 
 NS_IMETHODIMP
 nsHTTPCompressConv::CheckListenerChain() {
+  if (XRE_IsContentProcess()) {
+    // handle decompression OMT always.  If the chain needs to be MT,
+    // we'll determine that in OnStartRequest and dispatch to MT
+    return NS_OK;
+  }
   nsCOMPtr<nsIThreadRetargetableStreamListener> listener;
   {
     MutexAutoLock lock(mMutex);
     listener = do_QueryInterface(mListener);
   }
-
   if (!listener) {
     return NS_ERROR_NO_INTERFACE;
   }
@@ -748,6 +964,15 @@ nsHTTPCompressConv::OnDataFinished(nsresult aStatus) {
   }
 
   if (listener) {
+    if (mDispatchToMainThread && !NS_IsMainThread()) {
+      nsCOMPtr<nsIRunnable> handler = NS_NewRunnableFunction(
+          "dispatch", [listener{std::move(listener)}, aStatus]() {
+            Unused << listener->OnDataFinished(aStatus);
+          });
+
+      return NS_DispatchToMainThread(handler);
+    }
+
     return listener->OnDataFinished(aStatus);
   }