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 7a775f9faac4296806ec1ab421da6dc7e9463a7a..251fdb8fb65e1b68056cb9c97a88a2da5450e380 100644
--- a/content/browser/code_cache/generated_code_cache.cc
+++ b/content/browser/code_cache/generated_code_cache.cc
@@ -4,6 +4,7 @@
 
 #include "content/browser/code_cache/generated_code_cache.h"
 
+#include <algorithm>
 #include <iostream>
 #include <string_view>
 
@@ -33,6 +34,7 @@
 #include "third_party/blink/public/common/loader/code_cache_util.h"
 #include "third_party/blink/public/common/scheme_registry.h"
 #include "url/gurl.h"
+#include "url/url_util.h"
 
 using storage::BigIOBuffer;
 
@@ -54,13 +56,17 @@ void CheckValidResource(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.GetScheme()));
+  const bool resource_url_embedder = std::ranges::contains(
+      url::GetCodeCacheSchemes(), resource_url.GetScheme());
+  DCHECK(resource_url_http || resource_url_webui || resource_url_embedder);
 
   if (!blink::features::IsPersistentCacheForCodeCacheEnabled()) {
     // The chrome and chrome-untrusted schemes are only used with the WebUI code
@@ -68,35 +74,51 @@ void CheckValidResource(const GURL& resource_url,
     // segments WebUI from non-WebUI in multiple ways to prevent privilege
     // escalation, using both `GetCacheId` and
     // `CheckSecurityForAccessingCodeCacheData`.
-    DCHECK_EQ(resource_url_is_chrome_or_chrome_untrusted,
+    DCHECK_EQ(resource_url_webui,
               cache_type == GeneratedCodeCache::kWebUIJavaScript);
   }
 }
 
 void CheckValidContext(const GURL& origin_lock,
                        GeneratedCodeCache::CodeCacheType cache_type) {
-  // |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 =
+  // |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.GetScheme())) &&
-          !url::Origin::Create(origin_lock).opaque()));
+  const bool origin_lock_for_embedder = std::ranges::contains(
+      url::GetCodeCacheSchemes(), origin_lock.GetScheme());
+
+  DCHECK(origin_lock_empty || ((origin_lock_for_http || origin_lock_for_webui ||
+                                origin_lock_for_embedder) &&
+                               !url::Origin::Create(origin_lock).opaque()));
 
   if (!blink::features::IsPersistentCacheForCodeCacheEnabled()) {
-    // The chrome and chrome-untrusted schemes are only used with the WebUI code
-    // cache type when PersistentCache is not used. Otherwise, PersistentCache
+    // The chrome and chrome-untrusted schemes are only used with their dedicated
+    // code cache type when PersistentCache is not used. Otherwise, PersistentCache
     // segments WebUI from non-WebUI in multiple ways to prevent privilege
     // escalation, using both `GetCacheId` and
     // `CheckSecurityForAccessingCodeCacheData`.
-    DCHECK_EQ(origin_lock_is_chrome_or_chrome_untrusted,
-              cache_type == GeneratedCodeCache::kWebUIJavaScript);
+    switch (cache_type) {
+      case GeneratedCodeCache::kJavaScript:
+      case GeneratedCodeCache::kWebAssembly:
+        DCHECK(!origin_lock_for_webui);
+        break;
+      case GeneratedCodeCache::kWebUIJavaScript:
+        DCHECK(origin_lock_for_webui);
+        break;
+    }
+
+    // The custom schemes share the cache type with http(s).
+    if (origin_lock_for_embedder) {
+      DCHECK(cache_type == GeneratedCodeCache::kJavaScript ||
+             cache_type == GeneratedCodeCache::kWebAssembly);
+    }
   }
 }
 
diff --git a/content/browser/code_cache/generated_code_cache.h b/content/browser/code_cache/generated_code_cache.h
index a01f0d96ef33ce9460a851b072b7ceed5227dee3..f7e39b28cc0ba2251123925c01083a7935f46f56 100644
--- a/content/browser/code_cache/generated_code_cache.h
+++ b/content/browser/code_cache/generated_code_cache.h
@@ -51,12 +51,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 68a3095a49caf472c83b93b5cef66e5549a2d7cc..aa371ba5576f9fbaf5558e39704f7eb87ec7511e 100644
--- a/content/browser/code_cache/generated_code_cache_browsertest.cc
+++ b/content/browser/code_cache/generated_code_cache_browsertest.cc
@@ -16,17 +16,22 @@
 #include "base/time/time.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/browser/web_contents/web_contents_impl.h"
 #include "content/common/renderer.mojom.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/render_process_host.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/content_browser_test_utils_internal.h"
+#include "content/test/test_content_client.h"
 #include "net/base/features.h"
 #include "net/dns/mock_host_resolver.h"
 #include "third_party/blink/public/common/features.h"
@@ -38,6 +43,8 @@ namespace content {
 
 namespace {
 
+const std::string kCodeCacheScheme = "test-code-cache";
+
 bool SupportsSharedWorker() {
   return base::FeatureList::IsEnabled(blink::features::kSharedWorker);
 }
@@ -1063,4 +1070,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 ffc684ad1e2dfbf148bb01b4d185a44619a2dceb..39b15dd921539c17e1038fabc64bf20b0974e20a 100644
--- a/content/browser/renderer_host/code_cache_host_impl.cc
+++ b/content/browser/renderer_host/code_cache_host_impl.cc
@@ -4,6 +4,7 @@
 
 #include "content/browser/renderer_host/code_cache_host_impl.h"
 
+#include <algorithm>
 #include <optional>
 #include <string_view>
 #include <utility>
@@ -42,6 +43,7 @@
 #include "third_party/perfetto/include/perfetto/tracing/track_event_args.h"
 #include "url/gurl.h"
 #include "url/origin.h"
+#include "url/url_util.h"
 
 using blink::mojom::CacheStorageError;
 
@@ -56,6 +58,10 @@ enum class Operation {
   kWrite,
 };
 
+bool ProcessLockURLIsCodeCacheScheme(const ProcessLock& process_lock) {
+  return std::ranges::contains(url::GetCodeCacheSchemes(), process_lock.GetProcessLockURL().scheme());
+}
+
 bool CheckSecurityForAccessingCodeCacheData(const GURL& resource_url,
                                             int render_process_id,
                                             Operation operation) {
@@ -67,42 +73,56 @@ bool CheckSecurityForAccessingCodeCacheData(const GURL& resource_url,
       ChildProcessSecurityPolicyImpl::GetInstance()->GetProcessLock(
           render_process_id);
 
-  // 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.SchemeIsHTTPOrHTTPS()) {
+    if (process_lock.MatchesScheme(url::kHttpScheme) ||
+        process_lock.MatchesScheme(url::kHttpsScheme)) {
+      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(content::kChromeUIScheme) ||
       resource_url.SchemeIs(content::kChromeUIUntrustedScheme)) {
-    if (!process_lock.IsLockedToSite()) {
-      // 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 (process_lock.MatchesScheme(content::kChromeUIScheme) ||
+        process_lock.MatchesScheme(content::kChromeUIUntrustedScheme)) {
+      return true;
     }
-    if (process_lock.MatchesScheme(url::kHttpScheme) ||
-        process_lock.MatchesScheme(url::kHttpsScheme)) {
-      if (operation == Operation::kWrite) {
+    if (operation == Operation::kWrite) {
+      if (process_lock.MatchesScheme(url::kHttpScheme) ||
+          process_lock.MatchesScheme(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.MatchesScheme(content::kChromeUIScheme) ||
-           process_lock.MatchesScheme(content::kChromeUIUntrustedScheme);
+    // dangerous.  Similarly, the process might not be locked to a site.
+    return false;
   }
-  if (resource_url.SchemeIsHTTPOrHTTPS() ||
-      blink::CommonSchemeRegistry::IsExtensionScheme(
-          resource_url.GetScheme())) {
-    if (process_lock.MatchesScheme(content::kChromeUIScheme) ||
-        process_lock.MatchesScheme(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 (std::ranges::contains(url::GetCodeCacheSchemes(), resource_url.GetScheme())) {
+    return ProcessLockURLIsCodeCacheScheme(process_lock);
   }
 
   if (operation == Operation::kWrite) {
@@ -180,6 +200,7 @@ std::optional<GURL> GetOriginLock(int render_process_id) {
       process_lock.MatchesScheme(url::kHttpsScheme) ||
       process_lock.MatchesScheme(content::kChromeUIScheme) ||
       process_lock.MatchesScheme(content::kChromeUIUntrustedScheme) ||
+      ProcessLockURLIsCodeCacheScheme(process_lock) ||
       blink::CommonSchemeRegistry::IsExtensionScheme(
           process_lock.GetProcessLockURL().GetScheme())) {
     return process_lock.GetProcessLockURL();
diff --git a/content/common/url_schemes.cc b/content/common/url_schemes.cc
index 225e017909b8869231b870eaaf161a0b5e93e2a0..846a5251429630b8528a84a3d67ed56cb28df5a1 100644
--- a/content/common/url_schemes.cc
+++ b/content/common/url_schemes.cc
@@ -102,6 +102,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 33e2ff42e4d9da442d522b959a4a21c2f7032b6b..a0d81212327fc17e1f4704e78803c1d7d82b2016 100644
--- a/content/public/common/content_client.h
+++ b/content/public/common/content_client.h
@@ -139,6 +139,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 3d88fa831fdf8f684031fe6cacea09c9446b47fb..21cfd314fd52727e4b07de880bd80a0bec7407d5 100644
--- a/url/url_util.cc
+++ b/url/url_util.cc
@@ -131,6 +131,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;
 
@@ -666,6 +669,15 @@ const std::vector<std::string>& GetEmptyDocumentSchemes() {
   return GetSchemeRegistry().empty_document_schemes;
 }
 
+void AddCodeCacheScheme(const std::string_view new_scheme) {
+  DoAddScheme(new_scheme,
+              &GetSchemeRegistryWithoutLocking()->code_cache_schemes);
+}
+
+const std::vector<std::string>& GetCodeCacheSchemes() {
+  return GetSchemeRegistry().code_cache_schemes;
+}
+
 void AddPredefinedHandlerScheme(std::string_view new_scheme,
                                 std::string_view handler) {
   DoAddSchemeWithHandler(
diff --git a/url/url_util.h b/url/url_util.h
index 1f9304c7ea5488a1c50891e3dabfa3a61af97b7b..f965c1dbd47781748d3091209d140a128ca7192f 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(std::string_view 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(std::string_view 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
