From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Cheng Zhao <zcbenz@gmail.com>
Date: Thu, 14 Dec 2023 21:16:53 +0900
Subject: Enable V8 code cache for custom schemes

Add a new category in ContentClient::AddAdditionalSchemes which allows
embedders to make custom schemes allow V8 code cache.

Chromium CL: https://chromium-review.googlesource.com/c/chromium/src/+/5019665

diff --git a/content/browser/code_cache/generated_code_cache.cc b/content/browser/code_cache/generated_code_cache.cc
index a9e025b1d611c3becb983a77d20fc03da2418293..d04e399638116fef5f2b15b6510157d3a878723a 100644
--- a/content/browser/code_cache/generated_code_cache.cc
+++ b/content/browser/code_cache/generated_code_cache.cc
@@ -7,6 +7,7 @@
 #include <iostream>
 #include <string_view>
 
+#include "base/containers/contains.h"
 #include "base/feature_list.h"
 #include "base/functional/bind.h"
 #include "base/functional/callback_helpers.h"
@@ -27,6 +28,7 @@
 #include "net/http/http_cache.h"
 #include "third_party/blink/public/common/scheme_registry.h"
 #include "url/gurl.h"
+#include "url/url_util.h"
 
 using storage::BigIOBuffer;
 
@@ -39,7 +41,7 @@ constexpr char kSeparator[] = " \n";
 
 // We always expect to receive valid URLs that can be used as keys to the code
 // cache. The relevant checks (for ex: resource_url is valid, origin_lock is
-// not opque etc.,) must be done prior to requesting the code cache.
+// not opaque etc.,) must be done prior to requesting the code cache.
 //
 // This function doesn't enforce anything in the production code. It is here
 // to make the assumptions explicit and to catch any errors when DCHECKs are
@@ -49,33 +51,55 @@ void CheckValidKeys(const GURL& resource_url,
                     GeneratedCodeCache::CodeCacheType cache_type) {
   // If the resource url is invalid don't cache the code.
   DCHECK(resource_url.is_valid());
-  bool resource_url_is_chrome_or_chrome_untrusted =
+
+  // There are 3 kind of URL scheme compatible for the `resource_url`.
+  // 1. http: and https: URLs.
+  // 2. chrome: and chrome-untrusted: URLs.
+  // 3. URLs whose scheme are allowed by the content/ embedder.
+  const bool resource_url_http = resource_url.SchemeIsHTTPOrHTTPS();
+  const bool resource_url_webui =
       resource_url.SchemeIs(content::kChromeUIScheme) ||
       resource_url.SchemeIs(content::kChromeUIUntrustedScheme);
-  DCHECK(resource_url.SchemeIsHTTPOrHTTPS() ||
-         resource_url_is_chrome_or_chrome_untrusted ||
-         blink::CommonSchemeRegistry::IsExtensionScheme(resource_url.scheme()));
-
-  // |origin_lock| should be either empty or should have
-  // Http/Https/chrome/chrome-untrusted schemes and it should not be a URL with
-  // opaque origin. Empty origin_locks are allowed when the renderer is not
-  // locked to an origin.
-  bool origin_lock_is_chrome_or_chrome_untrusted =
+
+  const bool resource_url_embedder =
+      base::Contains(url::GetCodeCacheSchemes(), resource_url.scheme());
+  DCHECK(resource_url_http || resource_url_webui || resource_url_embedder);
+
+  // |origin_lock| should be either empty or should have code cache allowed
+  // schemes (http/https/chrome/chrome-untrusted or other custom schemes added
+  // by url::AddCodeCacheScheme), and it should not be a URL with opaque
+  // origin. Empty origin_locks are allowed when the renderer is not locked to
+  // an origin.
+  const bool origin_lock_empty = origin_lock.is_empty();
+  const bool origin_lock_for_http = origin_lock.SchemeIsHTTPOrHTTPS();
+  const bool origin_lock_for_webui =
       origin_lock.SchemeIs(content::kChromeUIScheme) ||
       origin_lock.SchemeIs(content::kChromeUIUntrustedScheme);
-  DCHECK(
-      origin_lock.is_empty() ||
-      ((origin_lock.SchemeIsHTTPOrHTTPS() ||
-        origin_lock_is_chrome_or_chrome_untrusted ||
-        blink::CommonSchemeRegistry::IsExtensionScheme(origin_lock.scheme())) &&
-       !url::Origin::Create(origin_lock).opaque()));
-
-  // The chrome and chrome-untrusted schemes are only used with the WebUI
-  // code cache type.
-  DCHECK_EQ(origin_lock_is_chrome_or_chrome_untrusted,
-            cache_type == GeneratedCodeCache::kWebUIJavaScript);
-  DCHECK_EQ(resource_url_is_chrome_or_chrome_untrusted,
-            cache_type == GeneratedCodeCache::kWebUIJavaScript);
+  const bool origin_lock_for_embedder =
+      base::Contains(url::GetCodeCacheSchemes(), origin_lock.scheme());
+
+  DCHECK(origin_lock_empty || ((origin_lock_for_http || origin_lock_for_webui ||
+                                origin_lock_for_embedder) &&
+                               !url::Origin::Create(origin_lock).opaque()));
+
+  // The webui schemes are only used with their dedicated code cache type.
+  switch (cache_type) {
+    case GeneratedCodeCache::kJavaScript:
+    case GeneratedCodeCache::kWebAssembly:
+      DCHECK(!origin_lock_for_webui);
+      DCHECK(!resource_url_webui);
+      break;
+    case GeneratedCodeCache::kWebUIJavaScript:
+      DCHECK(origin_lock_for_webui);
+      DCHECK(resource_url_webui);
+      break;
+  }
+
+  // The custom schemes share the cache type with http(s).
+  if (origin_lock_for_embedder || resource_url_embedder) {
+    DCHECK(cache_type == GeneratedCodeCache::kJavaScript ||
+           cache_type == GeneratedCodeCache::kWebAssembly);
+  }
 }
 
 // Generates the cache key for the given |resource_url|, |origin_lock| and
diff --git a/content/browser/code_cache/generated_code_cache.h b/content/browser/code_cache/generated_code_cache.h
index f5c5ff2c89489257003dfe3284ee9de9f517c99b..fdd2e2483171c4d43963590200817dac27d22cf9 100644
--- a/content/browser/code_cache/generated_code_cache.h
+++ b/content/browser/code_cache/generated_code_cache.h
@@ -52,12 +52,14 @@ class CONTENT_EXPORT GeneratedCodeCache {
   // Cache type. Used for collecting statistics for JS and Wasm in separate
   // buckets.
   enum CodeCacheType {
-    // JavaScript from http(s) pages.
+    // JavaScript from pages of http(s) schemes or custom schemes registered by
+    // url::AddCodeCacheScheme.
     kJavaScript,
 
-    // WebAssembly from http(s) pages. This cache allows more total size and
-    // more size per item than the JavaScript cache, since some
-    // WebAssembly programs are very large.
+    // WebAssembly from pages of http(s) schemes or custom schemes registered by
+    // url::AddCodeCacheScheme. This cache allows more total size and more size
+    // per item than the JavaScript cache, since some WebAssembly programs are
+    // very large.
     kWebAssembly,
 
     // JavaScript from chrome and chrome-untrusted pages. The resource URLs are
diff --git a/content/browser/code_cache/generated_code_cache_browsertest.cc b/content/browser/code_cache/generated_code_cache_browsertest.cc
index b6dd7405a5c9275ab699d4b347759427b30ef594..253918a2e54c98ce0075bce4e1a52134032ce367 100644
--- a/content/browser/code_cache/generated_code_cache_browsertest.cc
+++ b/content/browser/code_cache/generated_code_cache_browsertest.cc
@@ -8,13 +8,18 @@
 #include "content/browser/code_cache/generated_code_cache.h"
 #include "content/browser/code_cache/generated_code_cache_context.h"
 #include "content/browser/renderer_host/code_cache_host_impl.h"
+#include "content/browser/storage_partition_impl.h"
+#include "content/common/url_schemes.h"
 #include "content/public/browser/browser_context.h"
+#include "content/public/browser/browser_thread.h"
 #include "content/public/browser/storage_partition.h"
 #include "content/public/test/browser_test.h"
 #include "content/public/test/browser_test_utils.h"
 #include "content/public/test/content_browser_test.h"
 #include "content/public/test/content_browser_test_utils.h"
+#include "content/public/test/test_browser_context.h"
 #include "content/shell/browser/shell.h"
+#include "content/test/test_content_client.h"
 #include "net/dns/mock_host_resolver.h"
 #include "third_party/blink/public/common/features.h"
 #include "third_party/blink/public/common/page/v8_compile_hints_histograms.h"
@@ -23,6 +28,8 @@ namespace content {
 
 namespace {
 
+const std::string kCodeCacheScheme = "test-code-cache";
+
 bool SupportsSharedWorker() {
 #if BUILDFLAG(IS_ANDROID)
   // SharedWorkers are not enabled on Android. https://crbug.com/154571
@@ -714,4 +721,82 @@ IN_PROC_BROWSER_TEST_F(LocalCompileHintsBrowserTest, LocalCompileHints) {
   }
 }
 
+class CodeCacheInCustomSchemeBrowserTest : public ContentBrowserTest,
+                                           public TestContentClient {
+ public:
+  CodeCacheInCustomSchemeBrowserTest() {
+    SetContentClient(this);
+    ReRegisterContentSchemesForTests();
+  }
+
+  ~CodeCacheInCustomSchemeBrowserTest() override { SetContentClient(nullptr); }
+
+ private:
+  void AddAdditionalSchemes(Schemes* schemes) override {
+    schemes->standard_schemes.push_back(kCodeCacheScheme);
+    schemes->code_cache_schemes.push_back(kCodeCacheScheme);
+  }
+
+  url::ScopedSchemeRegistryForTests scheme_registry_;
+};
+
+IN_PROC_BROWSER_TEST_F(CodeCacheInCustomSchemeBrowserTest,
+                       AllowedCustomSchemeCanGenerateCodeCache) {
+  StoragePartitionImpl* partition =
+      static_cast<StoragePartitionImpl*>(shell()
+                                             ->web_contents()
+                                             ->GetBrowserContext()
+                                             ->GetDefaultStoragePartition());
+  scoped_refptr<GeneratedCodeCacheContext> context =
+      partition->GetGeneratedCodeCacheContext();
+  EXPECT_NE(context, nullptr);
+
+  GURL url(kCodeCacheScheme + "://host4/script.js");
+  GURL origin(kCodeCacheScheme + "://host1:1/");
+  ASSERT_TRUE(url.is_valid());
+  ASSERT_TRUE(origin.is_valid());
+  std::string data("SomeData");
+
+  // Add a code cache entry for the custom scheme.
+  base::test::TestFuture<void> add_entry_future;
+  GeneratedCodeCacheContext::RunOrPostTask(
+      context.get(), FROM_HERE,
+      base::BindOnce(
+          [](scoped_refptr<GeneratedCodeCacheContext> context, const GURL& url,
+             const GURL& origin, const std::string& data,
+             base::OnceClosure callback) {
+            context->generated_js_code_cache()->WriteEntry(
+                url, origin, net::NetworkIsolationKey(), base::Time::Now(),
+                std::vector<uint8_t>(data.begin(), data.end()));
+            GetUIThreadTaskRunner({})->PostTask(FROM_HERE, std::move(callback));
+          },
+          context, url, origin, data, add_entry_future.GetCallback()));
+  ASSERT_TRUE(add_entry_future.Wait());
+
+  // Get the code cache entry.
+  base::test::TestFuture<std::string> get_entry_future;
+  GeneratedCodeCacheContext::RunOrPostTask(
+      context.get(), FROM_HERE,
+      base::BindOnce(
+          [](scoped_refptr<GeneratedCodeCacheContext> context, const GURL& url,
+             const GURL& origin,
+             base::OnceCallback<void(std::string)> callback) {
+            context->generated_js_code_cache()->FetchEntry(
+                url, origin, net::NetworkIsolationKey(),
+                base::BindOnce(
+                    [](base::OnceCallback<void(std::string)> callback,
+                       const base::Time& response_time,
+                       mojo_base::BigBuffer buffer) {
+                      std::string data(buffer.data(),
+                                       buffer.data() + buffer.size());
+                      GetUIThreadTaskRunner({})->PostTask(
+                          FROM_HERE, base::BindOnce(std::move(callback), data));
+                    },
+                    std::move(callback)));
+          },
+          context, url, origin, get_entry_future.GetCallback()));
+  ASSERT_TRUE(get_entry_future.Wait());
+  ASSERT_EQ(data, get_entry_future.Get<0>());
+}
+
 }  // namespace content
diff --git a/content/browser/renderer_host/code_cache_host_impl.cc b/content/browser/renderer_host/code_cache_host_impl.cc
index b083cd89c46e676f26a7c28eda091cedbf9a5a97..84736ba6b3ae6111f1cebad862989daadb838234 100644
--- a/content/browser/renderer_host/code_cache_host_impl.cc
+++ b/content/browser/renderer_host/code_cache_host_impl.cc
@@ -6,6 +6,7 @@
 
 #include <utility>
 
+#include "base/containers/contains.h"
 #include "base/functional/bind.h"
 #include "base/functional/callback_helpers.h"
 #include "base/metrics/histogram_functions.h"
@@ -29,6 +30,7 @@
 #include "third_party/blink/public/common/scheme_registry.h"
 #include "url/gurl.h"
 #include "url/origin.h"
+#include "url/url_util.h"
 
 using blink::mojom::CacheStorageError;
 
@@ -36,6 +38,11 @@ namespace content {
 
 namespace {
 
+bool ProcessLockURLIsCodeCacheScheme(const ProcessLock& process_lock) {
+  return base::Contains(url::GetCodeCacheSchemes(),
+                        process_lock.lock_url().scheme());
+}
+
 bool CheckSecurityForAccessingCodeCacheData(
     const GURL& resource_url,
     int render_process_id,
@@ -46,39 +53,57 @@ bool CheckSecurityForAccessingCodeCacheData(
 
   // Code caching is only allowed for http(s) and chrome/chrome-untrusted
   // scripts. Furthermore, there is no way for http(s) pages to load chrome or
+  // Code caching is only allowed for scripts from:
+  // 1. http: and https: schemes.
+  // 2. chrome: and chrome-untrusted: schemes.
+  // 3. Schemes registered by content/ embedder via url::AddCodeCacheScheme.
+  //
+  // Furthermore, we know there are no way for http(s) pages to load chrome or
   // chrome-untrusted scripts, so any http(s) page attempting to store data
   // about a chrome or chrome-untrusted script would be an indication of
   // suspicious activity.
-  if (resource_url.SchemeIs(content::kChromeUIScheme) ||
-      resource_url.SchemeIs(content::kChromeUIUntrustedScheme)) {
-    if (!process_lock.is_locked_to_site()) {
-      // We can't tell for certain whether this renderer is doing something
-      // malicious, but we don't trust it enough to store data.
-      return false;
-    }
+  if (resource_url.SchemeIsHTTPOrHTTPS()) {
     if (process_lock.matches_scheme(url::kHttpScheme) ||
         process_lock.matches_scheme(url::kHttpsScheme)) {
-      if (operation == CodeCacheHostImpl::Operation::kWrite) {
+      return true;
+    }
+    // Pages in custom schemes like isolated-app: are allowed to load http(s)
+    // resources.
+    if (ProcessLockURLIsCodeCacheScheme(process_lock)) {
+      return true;
+    }
+    // It is possible for WebUI pages to include open-web content, but such
+    // usage is rare and we've decided that reasoning about security is easier
+    // if the WebUI code cache includes only WebUI scripts.
+    return false;
+  }
+
+  if (resource_url.SchemeIs(kChromeUIScheme) ||
+      resource_url.SchemeIs(kChromeUIUntrustedScheme)) {
+    if (process_lock.matches_scheme(kChromeUIScheme) ||
+        process_lock.matches_scheme(kChromeUIUntrustedScheme)) {
+      return true;
+    }
+    if (operation == CodeCacheHostImpl::Operation::kWrite) {
+      if (process_lock.matches_scheme(url::kHttpScheme) ||
+          process_lock.matches_scheme(url::kHttpsScheme)) {
         mojo::ReportBadMessage("HTTP(S) pages cannot cache WebUI code");
       }
+      if (ProcessLockURLIsCodeCacheScheme(process_lock)) {
+        mojo::ReportBadMessage(
+            "Page whose scheme are allowed by content/ embedders cannot cache "
+            "WebUI code. Did the embedder misconfigured content/?");
+      }
       return false;
     }
     // Other schemes which might successfully load chrome or chrome-untrusted
     // scripts, such as the PDF viewer, are unsupported but not considered
-    // dangerous.
-    return process_lock.matches_scheme(content::kChromeUIScheme) ||
-           process_lock.matches_scheme(content::kChromeUIUntrustedScheme);
+    // dangerous.  Similarly, the process might not be locked to a site.
+    return false;
   }
-  if (resource_url.SchemeIsHTTPOrHTTPS() ||
-      blink::CommonSchemeRegistry::IsExtensionScheme(resource_url.scheme())) {
-    if (process_lock.matches_scheme(content::kChromeUIScheme) ||
-        process_lock.matches_scheme(content::kChromeUIUntrustedScheme)) {
-      // It is possible for WebUI pages to include open-web content, but such
-      // usage is rare and we've decided that reasoning about security is easier
-      // if the WebUI code cache includes only WebUI scripts.
-      return false;
-    }
-    return true;
+
+  if (base::Contains(url::GetCodeCacheSchemes(), resource_url.scheme())) {
+    return ProcessLockURLIsCodeCacheScheme(process_lock);
   }
 
   if (operation == CodeCacheHostImpl::Operation::kWrite) {
@@ -425,6 +450,7 @@ std::optional<GURL> CodeCacheHostImpl::GetSecondaryKeyForCodeCache(
       process_lock.matches_scheme(url::kHttpsScheme) ||
       process_lock.matches_scheme(content::kChromeUIScheme) ||
       process_lock.matches_scheme(content::kChromeUIUntrustedScheme) ||
+      ProcessLockURLIsCodeCacheScheme(process_lock) ||
       blink::CommonSchemeRegistry::IsExtensionScheme(
           process_lock.lock_url().scheme())) {
     return process_lock.lock_url();
diff --git a/content/common/url_schemes.cc b/content/common/url_schemes.cc
index ce9644d33fe83379127b01bf9a2b1c4badc3bc7c..fd486d4637ae4766ed78571dee7f9cebbd809f38 100644
--- a/content/common/url_schemes.cc
+++ b/content/common/url_schemes.cc
@@ -98,6 +98,14 @@ void RegisterContentSchemes(bool should_lock_registry) {
   for (auto& scheme : schemes.empty_document_schemes)
     url::AddEmptyDocumentScheme(scheme.c_str());
 
+  for (auto& scheme : schemes.code_cache_schemes) {
+    CHECK_NE(scheme, kChromeUIScheme);
+    CHECK_NE(scheme, kChromeUIUntrustedScheme);
+    CHECK_NE(scheme, url::kHttpScheme);
+    CHECK_NE(scheme, url::kHttpsScheme);
+    url::AddCodeCacheScheme(scheme.c_str());
+  }
+
 #if BUILDFLAG(IS_ANDROID)
   if (schemes.allow_non_standard_schemes_in_origins)
     url::EnableNonStandardSchemesForAndroidWebView();
diff --git a/content/public/common/content_client.h b/content/public/common/content_client.h
index 9dc2d5a33858da7c31fd87bbbabe3899301fa52d..ebf0bb23b9aedb7bf9eb8af52b4756dba452183e 100644
--- a/content/public/common/content_client.h
+++ b/content/public/common/content_client.h
@@ -142,6 +142,9 @@ class CONTENT_EXPORT ContentClient {
     // Registers a URL scheme as strictly empty documents, allowing them to
     // commit synchronously.
     std::vector<std::string> empty_document_schemes;
+    // Registers a URL scheme whose js and wasm scripts have V8 code cache
+    // enabled.
+    std::vector<std::string> code_cache_schemes;
     // Registers a URL scheme as extension scheme.
     std::vector<std::string> extension_schemes;
     // Registers a URL scheme with a predefined default custom handler.
diff --git a/url/url_util.cc b/url/url_util.cc
index 301824a7f83e9e7057279723eccd6c9c9ad44618..25e7e087e2ddf0e8fab1cd6f20a56790f7e406b2 100644
--- a/url/url_util.cc
+++ b/url/url_util.cc
@@ -133,6 +133,9 @@ struct SchemeRegistry {
       kMaterializedViewScheme,
   };
 
+  // Embedder schemes that have V8 code cache enabled in js and wasm scripts.
+  std::vector<std::string> code_cache_schemes = {};
+
   // Schemes with a predefined default custom handler.
   std::vector<SchemeWithHandler> predefined_handler_schemes;
 
@@ -714,6 +717,15 @@ const std::vector<std::string>& GetEmptyDocumentSchemes() {
   return GetSchemeRegistry().empty_document_schemes;
 }
 
+void AddCodeCacheScheme(const char* new_scheme) {
+  DoAddScheme(new_scheme,
+              &GetSchemeRegistryWithoutLocking()->code_cache_schemes);
+}
+
+const std::vector<std::string>& GetCodeCacheSchemes() {
+  return GetSchemeRegistry().code_cache_schemes;
+}
+
 void AddPredefinedHandlerScheme(const char* new_scheme, const char* handler) {
   DoAddSchemeWithHandler(
       new_scheme, handler,
diff --git a/url/url_util.h b/url/url_util.h
index e39a44057cebce7cdf90bcb02a7463e88bd271b4..e80b81d2ddeb2ec201b143d86bec6ee54ca49afc 100644
--- a/url/url_util.h
+++ b/url/url_util.h
@@ -115,6 +115,15 @@ COMPONENT_EXPORT(URL) const std::vector<std::string>& GetCSPBypassingSchemes();
 COMPONENT_EXPORT(URL) void AddEmptyDocumentScheme(const char* new_scheme);
 COMPONENT_EXPORT(URL) const std::vector<std::string>& GetEmptyDocumentSchemes();
 
+// Adds an application-defined scheme to the list of schemes that have V8 code
+// cache enabled for the js and wasm scripts.
+// The WebUI schemes (chrome/chrome-untrusted) do not belong to this list, as
+// they are treated as a separate cache type for security purpose.
+// The http(s) schemes do not belong to this list neither, they always have V8
+// code cache enabled.
+COMPONENT_EXPORT(URL) void AddCodeCacheScheme(const char* new_scheme);
+COMPONENT_EXPORT(URL) const std::vector<std::string>& GetCodeCacheSchemes();
+
 // Adds a scheme with a predefined default handler.
 //
 // This pair of strings must be normalized protocol handler parameters as
