1 files changed, 394 insertions, 0 deletions
diff --git a/toolkit/components/places/PageIconProtocolHandler.cpp b/toolkit/components/places/PageIconProtocolHandler.cpp
new file mode 100644
index 0000000000..a294cb2a9c
--- /dev/null
+++ b/toolkit/components/places/PageIconProtocolHandler.cpp
@@ -0,0 +1,394 @@
+/* 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 "PageIconProtocolHandler.h"
+
+#include "mozilla/NullPrincipal.h"
+#include "mozilla/ScopeExit.h"
+#include "mozilla/Components.h"
+#include "nsFaviconService.h"
+#include "nsIChannel.h"
+#include "nsIFaviconService.h"
+#include "nsIIOService.h"
+#include "nsILoadInfo.h"
+#include "nsIOutputStream.h"
+#include "nsIPipe.h"
+#include "nsIRequestObserver.h"
+#include "nsIURIMutator.h"
+#include "nsNetUtil.h"
+#include "SimpleChannel.h"
+
+#define PAGE_ICON_SCHEME "page-icon"
+
+using mozilla::net::IsNeckoChild;
+using mozilla::net::NeckoChild;
+using mozilla::net::RemoteStreamGetter;
+using mozilla::net::RemoteStreamInfo;
+
+namespace mozilla::places {
+
+struct FaviconMetadata {
+  nsCOMPtr<nsIInputStream> mStream;
+  nsCString mContentType;
+  int64_t mContentLength = 0;
+};
+
+StaticRefPtr<PageIconProtocolHandler> PageIconProtocolHandler::sSingleton;
+
+namespace {
+
+class DefaultFaviconObserver final : public nsIRequestObserver {
+ public:
+  explicit DefaultFaviconObserver(nsIOutputStream* aOutputStream)
+      : mOutputStream(aOutputStream) {
+    MOZ_ASSERT(aOutputStream);
+  }
+
+  NS_DECL_ISUPPORTS
+  NS_DECL_NSIREQUESTOBSERVER
+ private:
+  ~DefaultFaviconObserver() = default;
+
+  nsCOMPtr<nsIOutputStream> mOutputStream;
+};
+
+NS_IMPL_ISUPPORTS(DefaultFaviconObserver, nsIRequestObserver);
+
+NS_IMETHODIMP DefaultFaviconObserver::OnStartRequest(nsIRequest*) {
+  return NS_OK;
+}
+
+NS_IMETHODIMP DefaultFaviconObserver::OnStopRequest(nsIRequest*, nsresult) {
+  // We must close the outputStream regardless.
+  mOutputStream->Close();
+  return NS_OK;
+}
+
+}  // namespace
+
+static nsresult MakeDefaultFaviconChannel(nsIURI* aURI, nsILoadInfo* aLoadInfo,
+                                          nsIChannel** aOutChannel) {
+  nsCOMPtr<nsIIOService> ios = mozilla::components::IO::Service();
+  nsCOMPtr<nsIChannel> chan;
+  nsCOMPtr<nsIURI> defaultFaviconURI;
+
+  auto* faviconService = nsFaviconService::GetFaviconService();
+  if (MOZ_UNLIKELY(!faviconService)) {
+    return NS_ERROR_UNEXPECTED;
+  }
+
+  nsresult rv =
+      faviconService->GetDefaultFavicon(getter_AddRefs(defaultFaviconURI));
+  NS_ENSURE_SUCCESS(rv, rv);
+  rv = ios->NewChannelFromURIWithLoadInfo(defaultFaviconURI, aLoadInfo,
+                                          getter_AddRefs(chan));
+  NS_ENSURE_SUCCESS(rv, rv);
+  chan->SetOriginalURI(aURI);
+  chan->SetContentType(nsLiteralCString(FAVICON_DEFAULT_MIMETYPE));
+  chan.forget(aOutChannel);
+  return NS_OK;
+}
+
+static nsresult StreamDefaultFavicon(nsIURI* aURI, nsILoadInfo* aLoadInfo,
+                                     nsIOutputStream* aOutputStream) {
+  auto closeStreamOnError =
+      mozilla::MakeScopeExit([&] { aOutputStream->Close(); });
+
+  auto observer = MakeRefPtr<DefaultFaviconObserver>(aOutputStream);
+
+  nsCOMPtr<nsIStreamListener> listener;
+  nsresult rv = NS_NewSimpleStreamListener(getter_AddRefs(listener),
+                                           aOutputStream, observer);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  nsCOMPtr<nsIChannel> defaultIconChannel;
+  rv = MakeDefaultFaviconChannel(aURI, aLoadInfo,
+                                 getter_AddRefs(defaultIconChannel));
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  rv = defaultIconChannel->AsyncOpen(listener);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  closeStreamOnError.release();
+  return NS_OK;
+}
+
+namespace {
+
+class FaviconDataCallback final : public nsIFaviconDataCallback {
+ public:
+  FaviconDataCallback(nsIURI* aURI, nsILoadInfo* aLoadInfo)
+      : mURI(aURI), mLoadInfo(aLoadInfo) {
+    MOZ_ASSERT(aURI);
+    MOZ_ASSERT(aLoadInfo);
+  }
+
+  NS_DECL_ISUPPORTS
+  NS_DECL_NSIFAVICONDATACALLBACK
+
+  RefPtr<FaviconMetadataPromise> Promise() {
+    return mPromiseHolder.Ensure(__func__);
+  }
+
+ private:
+  ~FaviconDataCallback();
+  nsCOMPtr<nsIURI> mURI;
+  MozPromiseHolder<FaviconMetadataPromise> mPromiseHolder;
+  nsCOMPtr<nsILoadInfo> mLoadInfo;
+};
+
+NS_IMPL_ISUPPORTS(FaviconDataCallback, nsIFaviconDataCallback);
+
+FaviconDataCallback::~FaviconDataCallback() {
+  mPromiseHolder.RejectIfExists(NS_ERROR_FAILURE, __func__);
+}
+
+NS_IMETHODIMP FaviconDataCallback::OnComplete(nsIURI* aURI, uint32_t aDataLen,
+                                              const uint8_t* aData,
+                                              const nsACString& aMimeType,
+                                              uint16_t aWidth) {
+  if (!aDataLen) {
+    mPromiseHolder.Reject(NS_ERROR_NOT_AVAILABLE, __func__);
+    return NS_OK;
+  }
+
+  nsCOMPtr<nsIInputStream> inputStream;
+  nsresult rv =
+      NS_NewByteInputStream(getter_AddRefs(inputStream),
+                            AsChars(Span{aData, aDataLen}), NS_ASSIGNMENT_COPY);
+  if (NS_FAILED(rv)) {
+    mPromiseHolder.Reject(rv, __func__);
+    return rv;
+  }
+
+  FaviconMetadata metadata;
+  metadata.mStream = inputStream;
+  metadata.mContentType = aMimeType;
+  metadata.mContentLength = aDataLen;
+  mPromiseHolder.Resolve(std::move(metadata), __func__);
+
+  return NS_OK;
+}
+
+}  // namespace
+
+NS_IMPL_ISUPPORTS(PageIconProtocolHandler, nsIProtocolHandler,
+                  nsISupportsWeakReference);
+
+NS_IMETHODIMP PageIconProtocolHandler::GetScheme(nsACString& aScheme) {
+  aScheme.AssignLiteral(PAGE_ICON_SCHEME);
+  return NS_OK;
+}
+
+NS_IMETHODIMP PageIconProtocolHandler::AllowPort(int32_t, const char*,
+                                                 bool* aAllow) {
+  *aAllow = false;
+  return NS_OK;
+}
+
+NS_IMETHODIMP PageIconProtocolHandler::NewChannel(nsIURI* aURI,
+                                                  nsILoadInfo* aLoadInfo,
+                                                  nsIChannel** aOutChannel) {
+  // Load the URI remotely if accessed from a child.
+  if (IsNeckoChild()) {
+    MOZ_TRY(SubstituteRemoteChannel(aURI, aLoadInfo, aOutChannel));
+    return NS_OK;
+  }
+
+  nsresult rv = NewChannelInternal(aURI, aLoadInfo, aOutChannel);
+  if (NS_SUCCEEDED(rv)) {
+    return rv;
+  }
+  return MakeDefaultFaviconChannel(aURI, aLoadInfo, aOutChannel);
+}
+
+Result<Ok, nsresult> PageIconProtocolHandler::SubstituteRemoteChannel(
+    nsIURI* aURI, nsILoadInfo* aLoadInfo, nsIChannel** aRetVal) {
+  MOZ_ASSERT(IsNeckoChild());
+  MOZ_TRY(aURI ? NS_OK : NS_ERROR_INVALID_ARG);
+  MOZ_TRY(aLoadInfo ? NS_OK : NS_ERROR_INVALID_ARG);
+
+  RefPtr<RemoteStreamGetter> streamGetter =
+      new RemoteStreamGetter(aURI, aLoadInfo);
+
+  NewSimpleChannel(aURI, aLoadInfo, streamGetter, aRetVal);
+  return Ok();
+}
+
+nsresult PageIconProtocolHandler::NewChannelInternal(nsIURI* aURI,
+                                                     nsILoadInfo* aLoadInfo,
+                                                     nsIChannel** aOutChannel) {
+  // Create a pipe that will give us an output stream that we can use once
+  // we got all the favicon data.
+  nsCOMPtr<nsIAsyncInputStream> pipeIn;
+  nsCOMPtr<nsIAsyncOutputStream> pipeOut;
+  GetStreams(getter_AddRefs(pipeIn), getter_AddRefs(pipeOut));
+
+  // Create our channel.
+  nsCOMPtr<nsIChannel> channel;
+  {
+    // We override the channel's loadinfo below anyway, so using a null
+    // principal here is alright.
+    nsCOMPtr<nsIPrincipal> loadingPrincipal =
+        NullPrincipal::CreateWithoutOriginAttributes();
+    nsresult rv = NS_NewInputStreamChannel(
+        getter_AddRefs(channel), aURI, pipeIn.forget(), loadingPrincipal,
+        nsILoadInfo::SEC_REQUIRE_SAME_ORIGIN_DATA_IS_BLOCKED,
+        nsIContentPolicy::TYPE_INTERNAL_IMAGE);
+    NS_ENSURE_SUCCESS(rv, rv);
+  }
+
+  nsresult rv = channel->SetLoadInfo(aLoadInfo);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  GetFaviconData(aURI, aLoadInfo)
+      ->Then(
+          GetMainThreadSerialEventTarget(), __func__,
+          [pipeOut, channel](const FaviconMetadata& aMetadata) {
+            channel->SetContentType(aMetadata.mContentType);
+            channel->SetContentLength(aMetadata.mContentLength);
+
+            nsresult rv;
+            const nsCOMPtr<nsIEventTarget> target =
+                do_GetService(NS_STREAMTRANSPORTSERVICE_CONTRACTID, &rv);
+
+            if (NS_WARN_IF(NS_FAILED(rv))) {
+              channel->CancelWithReason(NS_BINDING_ABORTED,
+                                        "GetFaviconData failed"_ns);
+              return;
+            }
+
+            NS_AsyncCopy(aMetadata.mStream, pipeOut, target);
+          },
+          [uri = nsCOMPtr{aURI}, loadInfo = nsCOMPtr{aLoadInfo}, pipeOut,
+           channel](nsresult aRv) {
+            // There are a few reasons why this might fail. For example, one
+            // reason is that the URI might not actually be properly parsable.
+            // In that case, we'll try one last time to stream the default
+            // favicon before giving up.
+            channel->SetContentType(nsLiteralCString(FAVICON_DEFAULT_MIMETYPE));
+            channel->SetContentLength(-1);
+            Unused << StreamDefaultFavicon(uri, loadInfo, pipeOut);
+          });
+  channel.forget(aOutChannel);
+  return NS_OK;
+}
+
+RefPtr<FaviconMetadataPromise> PageIconProtocolHandler::GetFaviconData(
+    nsIURI* aPageIconURI, nsILoadInfo* aLoadInfo) {
+  auto* faviconService = nsFaviconService::GetFaviconService();
+  if (MOZ_UNLIKELY(!faviconService)) {
+    return FaviconMetadataPromise::CreateAndReject(NS_ERROR_UNEXPECTED,
+                                                   __func__);
+  }
+
+  uint16_t preferredSize = 0;
+  faviconService->PreferredSizeFromURI(aPageIconURI, &preferredSize);
+
+  nsCOMPtr<nsIURI> pageURI;
+  nsresult rv;
+  {
+    // NOTE: We don't need to strip #size= fragments because
+    // GetFaviconDataForPage strips them when doing the database lookup.
+    nsAutoCString pageQuery;
+    aPageIconURI->GetPathQueryRef(pageQuery);
+    rv = NS_NewURI(getter_AddRefs(pageURI), pageQuery);
+    if (NS_FAILED(rv)) {
+      return FaviconMetadataPromise::CreateAndReject(rv, __func__);
+    }
+  }
+
+  auto faviconCallback =
+      MakeRefPtr<FaviconDataCallback>(aPageIconURI, aLoadInfo);
+  rv = faviconService->GetFaviconDataForPage(pageURI, faviconCallback,
+                                             preferredSize);
+  if (NS_FAILED(rv)) {
+    return FaviconMetadataPromise::CreateAndReject(rv, __func__);
+  }
+
+  return faviconCallback->Promise();
+}
+
+RefPtr<RemoteStreamPromise> PageIconProtocolHandler::NewStream(
+    nsIURI* aChildURI, nsILoadInfo* aLoadInfo, bool* aTerminateSender) {
+  MOZ_ASSERT(!IsNeckoChild());
+  MOZ_ASSERT(NS_IsMainThread());
+
+  if (!aChildURI || !aLoadInfo || !aTerminateSender) {
+    return RemoteStreamPromise::CreateAndReject(NS_ERROR_INVALID_ARG, __func__);
+  }
+
+  *aTerminateSender = true;
+
+  // We should never receive a URI that isn't for a page-icon because
+  // these requests ordinarily come from the child's PageIconProtocolHandler.
+  // Ensure this request is for a page-icon URI. A compromised child process
+  // could send us any URI.
+  bool isPageIconScheme = false;
+  if (NS_FAILED(aChildURI->SchemeIs(PAGE_ICON_SCHEME, &isPageIconScheme)) ||
+      !isPageIconScheme) {
+    return RemoteStreamPromise::CreateAndReject(NS_ERROR_UNKNOWN_PROTOCOL,
+                                                __func__);
+  }
+
+  // For errors after this point, we want to propagate the error to
+  // the child, but we don't force the child process to be terminated.
+  *aTerminateSender = false;
+
+  RefPtr<RemoteStreamPromise::Private> outerPromise =
+      new RemoteStreamPromise::Private(__func__);
+  nsCOMPtr<nsIURI> uri(aChildURI);
+  nsCOMPtr<nsILoadInfo> loadInfo(aLoadInfo);
+  RefPtr<PageIconProtocolHandler> self = this;
+
+  GetFaviconData(uri, loadInfo)
+      ->Then(
+          GetMainThreadSerialEventTarget(), __func__,
+          [outerPromise](const FaviconMetadata& aMetadata) {
+            RemoteStreamInfo info(aMetadata.mStream, aMetadata.mContentType,
+                                  aMetadata.mContentLength);
+            outerPromise->Resolve(std::move(info), __func__);
+          },
+          [self, uri, loadInfo, outerPromise](nsresult aRv) {
+            nsCOMPtr<nsIAsyncInputStream> pipeIn;
+            nsCOMPtr<nsIAsyncOutputStream> pipeOut;
+            self->GetStreams(getter_AddRefs(pipeIn), getter_AddRefs(pipeOut));
+
+            RemoteStreamInfo info(
+                pipeIn, nsLiteralCString(FAVICON_DEFAULT_MIMETYPE), -1);
+            Unused << StreamDefaultFavicon(uri, loadInfo, pipeOut);
+            outerPromise->Resolve(std::move(info), __func__);
+          });
+  return outerPromise;
+}
+
+void PageIconProtocolHandler::GetStreams(nsIAsyncInputStream** inStream,
+                                         nsIAsyncOutputStream** outStream) {
+  static constexpr size_t kSegmentSize = 4096;
+  nsCOMPtr<nsIAsyncInputStream> pipeIn;
+  nsCOMPtr<nsIAsyncOutputStream> pipeOut;
+  NS_NewPipe2(getter_AddRefs(pipeIn), getter_AddRefs(pipeOut), true, true,
+              kSegmentSize,
+              nsIFaviconService::MAX_FAVICON_BUFFER_SIZE / kSegmentSize);
+
+  pipeIn.forget(inStream);
+  pipeOut.forget(outStream);
+}
+
+// static
+void PageIconProtocolHandler::NewSimpleChannel(
+    nsIURI* aURI, nsILoadInfo* aLoadInfo, RemoteStreamGetter* aStreamGetter,
+    nsIChannel** aRetVal) {
+  nsCOMPtr<nsIChannel> channel = NS_NewSimpleChannel(
+      aURI, aLoadInfo, aStreamGetter,
+      [](nsIStreamListener* listener, nsIChannel* simpleChannel,
+         RemoteStreamGetter* getter) -> RequestOrReason {
+        return getter->GetAsync(listener, simpleChannel,
+                                &NeckoChild::SendGetPageIconStream);
+      });
+
+  channel.swap(*aRetVal);
+}
+
+}  // namespace mozilla::places