From e39e4419050eb637461e8330310f22347a937ddd Mon Sep 17 00:00:00 2001
From: Andrea Orru <[email protected]>
Date: Fri, 11 Sep 2026 16:03:37 -0700
Subject: [PATCH] [Extensions] Scope cookies.getPartitionKey to calling profile

CookiesGetPartitionKeyFunction resolves frameId and documentId through
ExtensionApiFrameIdMap, which is a browser-global singleton. Unlike the
tabId branch (which already goes through ExtensionTabUtil::GetTabById
with the caller's BrowserContext), these branches did not verify that
the resolved frame belongs to a context the calling extension can
access.

Verify ExtensionTabUtil::IsWebContentsInContext() after the lookup and
return generic 'Invalid frameId.' / 'Invalid documentId.' errors when it
fails. Also fix a crash when frameId < 1 is passed to
GetRenderFrameHostByFrameId, check frame lifecycle inside the lookup
loop, and fix a crash when documentId and tabId are passed without
frameId.

Add a browser test verifying that an extension without incognito access
cannot resolve an incognito frame via frameId or documentId.

Bug: 497062057
Change-Id: If7ea9d42152e356a65d73d7fafb3de77e095c84f
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/8356347
Commit-Queue: Andrea Orru <[email protected]>
Reviewed-by: Kelvin Jiang <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1696694}
---
diff --git a/chrome/browser/extensions/api/cookies/cookies_api.cc b/chrome/browser/extensions/api/cookies/cookies_api.cc
index df671cb..0def08cd 100644
--- a/chrome/browser/extensions/api/cookies/cookies_api.cc
+++ b/chrome/browser/extensions/api/cookies/cookies_api.cc
@@ -25,6 +25,7 @@
 #include "content/public/browser/web_contents.h"
 #include "extensions/browser/event_router.h"
 #include "extensions/browser/extension_api_frame_id_map.h"
+#include "extensions/browser/extension_util.h"
 #include "extensions/browser/extensions_browser_client.h"
 #include "extensions/browser/safe_browsing_delegate.h"
 #include "extensions/common/constants.h"
@@ -767,16 +768,35 @@
   content::RenderFrameHost* render_frame_host = nullptr;
   content::WebContents* web_contents = nullptr;
   std::optional<int> frame_id = parsed_args_->details.frame_id;
-  std::optional<ExtensionApiFrameIdMap::DocumentId> document_id;
   std::optional<int> tab_id = parsed_args_->details.tab_id;
 
   if (parsed_args_->details.document_id.has_value()) {
-    document_id = ExtensionApiFrameIdMap::DocumentIdFromString(
-        *parsed_args_->details.document_id);
+    ExtensionApiFrameIdMap::DocumentId document_id =
+        ExtensionApiFrameIdMap::DocumentIdFromString(
+            *parsed_args_->details.document_id);
+    if (!document_id) {
+      return RespondNow(Error("Invalid `documentId`."));
+    }
     render_frame_host =
         ExtensionApiFrameIdMap::Get()->GetRenderFrameHostByDocumentId(
-            document_id.value());
+            document_id);
+    if (!render_frame_host) {
+      return RespondNow(Error("Invalid `documentId`."));
+    }
     web_contents = content::WebContents::FromRenderFrameHost(render_frame_host);
+    if (!web_contents ||
+        !util::IsWebContentsInContext(*web_contents, *browser_context(),
+                                      include_incognito_information())) {
+      return RespondNow(Error("Invalid `documentId`."));
+    }
+
+    if ((tab_id.has_value() &&
+         ExtensionTabUtil::GetTabId(web_contents) != tab_id.value()) ||
+        (frame_id.has_value() && ExtensionApiFrameIdMap::GetFrameId(
+                                     render_frame_host) != frame_id.value())) {
+      return RespondNow(
+          Error("Provided `tabId` and `frameId` do not match the frame."));
+    }
   } else if (tab_id.has_value()) {
     if (!frame_id.has_value()) {
       // Default to main frame if no frame is provided.
@@ -791,34 +811,36 @@
     }
     render_frame_host = ExtensionApiFrameIdMap::GetRenderFrameHostById(
         web_contents, frame_id.value());
+    if (!render_frame_host) {
+      return RespondNow(Error("Invalid `frameId`."));
+    }
   } else if (frame_id.has_value()) {
     if (frame_id.value() == 0) {
       return RespondNow(
           Error("`frameId` may not be 0 if no `tabId` is present."));
     }
+    if (frame_id.value() < 0) {
+      return RespondNow(Error("Invalid `frameId`."));
+    }
 
     render_frame_host =
         ExtensionApiFrameIdMap::Get()->GetRenderFrameHostByFrameId(
             frame_id.value());
+    if (!render_frame_host) {
+      return RespondNow(Error("Invalid `frameId`."));
+    }
+    web_contents = content::WebContents::FromRenderFrameHost(render_frame_host);
+    if (!web_contents ||
+        !util::IsWebContentsInContext(*web_contents, *browser_context(),
+                                      include_incognito_information())) {
+      return RespondNow(Error("Invalid `frameId`."));
+    }
   } else {
     return RespondNow(
         Error("Either `documentId` or `tabId` must be specified."));
   }
 
-  if (!render_frame_host) {
-    return RespondNow(document_id.has_value() ? Error("Invalid `documentId`.")
-                                              : Error("Invalid `frameId`."));
-  }
-
-  // If both document_id and tab_id are provided, make sure they match.
-  if (document_id.has_value() && tab_id.has_value()) {
-    if (ExtensionTabUtil::GetTabId(web_contents) != tab_id.value() ||
-        ExtensionApiFrameIdMap::GetFrameId(render_frame_host) !=
-            frame_id.value()) {
-      return RespondNow(
-          Error("Provided `tabId` and `frameId` do not match the frame."));
-    }
-  }
+  CHECK(render_frame_host);
 
   base::expected<net::CookiePartitionKey::SerializedCookiePartitionKey,
                  std::string>
diff --git a/extensions/browser/extension_api_frame_id_map.cc b/extensions/browser/extension_api_frame_id_map.cc
index 3f3daf5..7a408635 100644
--- a/extensions/browser/extension_api_frame_id_map.cc
+++ b/extensions/browser/extension_api_frame_id_map.cc
@@ -162,26 +162,22 @@
     int frame_id) {
   // Frame_id values of 0 are not guaranteed to be unique. Values less than 0
   // are invalid.
-  CHECK_GE(frame_id, 1);
+  if (frame_id < 1) {
+    return nullptr;
+  }
 
-  content::RenderFrameHost* render_frame_host = nullptr;
   for (const auto& iter : document_id_map_) {
-    if (frame_id ==
-        ExtensionApiFrameIdMap::GetFrameId(&iter.second->render_frame_host())) {
-      render_frame_host = &iter.second->render_frame_host();
-      break;
+    content::RenderFrameHost* render_frame_host =
+        &iter.second->render_frame_host();
+    if (frame_id == ExtensionApiFrameIdMap::GetFrameId(render_frame_host) &&
+        (render_frame_host->IsActive() ||
+         render_frame_host->IsInLifecycleState(
+             content::RenderFrameHost::LifecycleState::kPrerendering))) {
+      return render_frame_host;
     }
   }
 
-  // Fail if the frame is not active or in prerendering (e.g. in the
-  // back/forward cache).
-  if (!render_frame_host ||
-      (!render_frame_host->IsActive() &&
-       !render_frame_host->IsInLifecycleState(
-           content::RenderFrameHost::LifecycleState::kPrerendering))) {
-    return nullptr;
-  }
-  return render_frame_host;
+  return nullptr;
 }
 
 ExtensionApiFrameIdMap::DocumentId ExtensionApiFrameIdMap::DocumentIdFromString(
diff --git a/chrome/browser/extensions/api/cookies/cookies_apitest.cc b/chrome/browser/extensions/api/cookies/cookies_apitest.cc
index 7eb85be..894b8560c 100644
--- a/chrome/browser/extensions/api/cookies/cookies_apitest.cc
+++ b/chrome/browser/extensions/api/cookies/cookies_apitest.cc
@@ -5,6 +5,7 @@
 #include "base/command_line.h"
 #include "base/functional/callback_helpers.h"
 #include "base/memory/raw_ptr.h"
+#include "base/strings/stringprintf.h"
 #include "base/test/bind.h"
 #include "build/build_config.h"
 #include "chrome/browser/extensions/api/cookies/cookies_api.h"
@@ -18,6 +19,8 @@
 #include "content/public/browser/web_contents.h"
 #include "content/public/test/browser_test.h"
 #include "content/public/test/browser_test_utils.h"
+#include "extensions/browser/api_test_utils.h"
+#include "extensions/browser/extension_api_frame_id_map.h"
 #include "extensions/browser/extension_registry_observer.h"
 #include "extensions/browser/process_manager.h"
 #include "extensions/buildflags/buildflags.h"
@@ -223,6 +226,116 @@
   ASSERT_TRUE(RunTest("cookies/get_partition_key")) << message_;
 }
 
+IN_PROC_BROWSER_TEST_P(CookiesApiTest, TestGetPartitionKeyContextIsolation) {
+  // Set up an off-the-record window with a cross-site subframe.
+  const std::string default_response = "/defaultresponse";
+  content::WebContents* incognito_contents = PlatformOpenURLOffTheRecord(
+      profile(), embedded_test_server()->GetURL("a.com", default_response));
+  ASSERT_TRUE(incognito_contents);
+
+  const GURL cross_site_url =
+      embedded_test_server()->GetURL("b.com", default_response);
+  std::string script =
+      "var f = document.createElement('iframe');\n"
+      "f.src = '" +
+      cross_site_url.spec() +
+      "';\n"
+      "document.body.appendChild(f);\n";
+  EXPECT_TRUE(ExecJs(incognito_contents, script));
+  EXPECT_TRUE(WaitForLoadStop(incognito_contents));
+
+  content::RenderFrameHost* incognito_subframe =
+      content::ChildFrameAt(incognito_contents->GetPrimaryMainFrame(), 0);
+  ASSERT_TRUE(incognito_subframe);
+  int incognito_frame_id =
+      ExtensionApiFrameIdMap::GetFrameId(incognito_subframe);
+  std::string incognito_document_id =
+      ExtensionApiFrameIdMap::GetDocumentId(incognito_subframe).ToString();
+
+  // Load an extension in the regular profile without off-the-record access.
+  static constexpr char kManifest[] = R"({
+    "name": "Cookies Partition Key Isolation Test",
+    "version": "1.0",
+    "manifest_version": 3,
+    "permissions": ["cookies"],
+    "host_permissions": ["*://*/*"],
+    "background": {"service_worker": "background.js"}
+  })";
+
+  TestExtensionDir test_dir;
+  test_dir.WriteManifest(kManifest);
+  test_dir.WriteFile(FILE_PATH_LITERAL("background.js"), "// background");
+
+  const Extension* extension =
+      LoadExtension(test_dir.UnpackedPath(), {.allow_in_incognito = false});
+  ASSERT_TRUE(extension);
+
+  // Attempting to query the partition key for the off-the-record frame by
+  // frameId from the regular profile extension must be rejected.
+  {
+    auto function = base::MakeRefCounted<CookiesGetPartitionKeyFunction>();
+    function->set_extension(extension);
+    std::string error = api_test_utils::RunFunctionAndReturnError(
+        function.get(),
+        base::StringPrintf(R"([{"frameId": %d}])", incognito_frame_id),
+        profile(), api_test_utils::FunctionMode::kNone);
+    EXPECT_EQ("Invalid `frameId`.", error);
+  }
+
+  // Attempting to query the partition key for the off-the-record frame by
+  // documentId from the regular profile extension must be rejected.
+  {
+    auto function = base::MakeRefCounted<CookiesGetPartitionKeyFunction>();
+    function->set_extension(extension);
+    std::string error = api_test_utils::RunFunctionAndReturnError(
+        function.get(),
+        base::StringPrintf(R"([{"documentId": "%s"}])",
+                           incognito_document_id.c_str()),
+        profile(), api_test_utils::FunctionMode::kNone);
+    EXPECT_EQ("Invalid `documentId`.", error);
+  }
+
+  // Also verify that an extension without host permissions receives the same
+  // rejection error without any URL details.
+  static constexpr char kRestrictedManifest[] = R"({
+    "name": "Cookies Partition Key Restricted Test",
+    "version": "1.0",
+    "manifest_version": 3,
+    "permissions": ["cookies"],
+    "background": {"service_worker": "background.js"}
+  })";
+
+  TestExtensionDir restricted_test_dir;
+  restricted_test_dir.WriteManifest(kRestrictedManifest);
+  restricted_test_dir.WriteFile(FILE_PATH_LITERAL("background.js"),
+                                "// background");
+
+  const Extension* restricted_extension = LoadExtension(
+      restricted_test_dir.UnpackedPath(), {.allow_in_incognito = false});
+  ASSERT_TRUE(restricted_extension);
+
+  {
+    auto function = base::MakeRefCounted<CookiesGetPartitionKeyFunction>();
+    function->set_extension(restricted_extension);
+    std::string error = api_test_utils::RunFunctionAndReturnError(
+        function.get(),
+        base::StringPrintf(R"([{"frameId": %d}])", incognito_frame_id),
+        profile(), api_test_utils::FunctionMode::kNone);
+    EXPECT_EQ("Invalid `frameId`.", error);
+  }
+
+  {
+    auto function = base::MakeRefCounted<CookiesGetPartitionKeyFunction>();
+    function->set_extension(restricted_extension);
+    std::string error = api_test_utils::RunFunctionAndReturnError(
+        function.get(),
+        base::StringPrintf(R"([{"documentId": "%s"}])",
+                           incognito_document_id.c_str()),
+        profile(), api_test_utils::FunctionMode::kNone);
+    EXPECT_EQ("Invalid `documentId`.", error);
+  }
+}
+
 IN_PROC_BROWSER_TEST_F(ExtensionApiTest, OTRReceiverMojoConnectionError) {
   // Test that simulates a mojo connection error with an existing OTR profile.
   // This test verifies that the fix for crbug.com/472076020 works correctly:
diff --git a/chrome/test/data/extensions/api_test/cookies/get_partition_key/background.js b/chrome/test/data/extensions/api_test/cookies/get_partition_key/background.js
index 1b35c18..4ae63c2 100644
--- a/chrome/test/data/extensions/api_test/cookies/get_partition_key/background.js
+++ b/chrome/test/data/extensions/api_test/cookies/get_partition_key/background.js
@@ -74,10 +74,41 @@
         'Error: Invalid `frameId`.');
 
     await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey({tabId: tab.id, frameId: 999999}),
+        'Error: Invalid `frameId`.');
+
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey({frameId: -1}),
+        'Error: Invalid `frameId`.');
+
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey({frameId: 999999}),
+        'Error: Invalid `frameId`.');
+
+    await chrome.test.assertPromiseRejects(
         chrome.cookies.getPartitionKey(
             {tabId: tab.id, frameId: 0, documentId: ''}),
         'Error: Invalid `documentId`.');
 
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey(
+            {documentId: '0123456789abcdef0123456789abcdef'}),
+        'Error: Invalid `documentId`.');
+
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey(
+            {documentId: topLevelFrame.documentId, tabId: (tab.id - 1)}),
+        'Error: Provided `tabId` and `frameId` do not match the frame.');
+
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey(
+            {documentId: crossSiteFrame.documentId, frameId: 999999}),
+        'Error: Provided `tabId` and `frameId` do not match the frame.');
+
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey({}),
+        'Error: Either `documentId` or `tabId` must be specified.');
+
     chrome.test.succeed();
   },
   async function testHostPermissionsRequired() {
@@ -127,6 +158,20 @@
     });
     chrome.test.assertEq(expectedTopLevelKey, actualPartitionKey);
 
+    // Providing documentId and matching tabId without frameId.
+    actualPartitionKey = await chrome.cookies.getPartitionKey({
+      tabId: tab.id,
+      documentId: topLevelDocId,
+    });
+    chrome.test.assertEq(expectedTopLevelKey, actualPartitionKey);
+
+    // Providing documentId and matching frameId without tabId.
+    actualPartitionKey = await chrome.cookies.getPartitionKey({
+      frameId: topLevelFrame.frameId,
+      documentId: topLevelDocId,
+    });
+    chrome.test.assertEq(expectedTopLevelKey, actualPartitionKey);
+
     // Providing no frameId defaults to topLevel frameId (0).
     actualPartitionKey = await chrome.cookies.getPartitionKey({
       tabId: tab.id,
@@ -157,6 +202,20 @@
     });
     chrome.test.assertEq(expectedCrossSiteKey, actualPartitionKey);
 
+    // Providing documentId and matching tabId without frameId.
+    actualPartitionKey = await chrome.cookies.getPartitionKey({
+      tabId: tab.id,
+      documentId: crossSiteDocId,
+    });
+    chrome.test.assertEq(expectedCrossSiteKey, actualPartitionKey);
+
+    // Providing documentId and matching frameId without tabId.
+    actualPartitionKey = await chrome.cookies.getPartitionKey({
+      frameId: crossSiteFrame.frameId,
+      documentId: crossSiteDocId,
+    });
+    chrome.test.assertEq(expectedCrossSiteKey, actualPartitionKey);
+
     chrome.test.succeed();
   },
 ]);

diff --git a/chrome/browser/extensions/api/cookies/cookies_apitest.cc b/chrome/browser/extensions/api/cookies/cookies_apitest.cc
index 7eb85be..894b8560c 100644
--- a/chrome/browser/extensions/api/cookies/cookies_apitest.cc
+++ b/chrome/browser/extensions/api/cookies/cookies_apitest.cc
@@ -5,6 +5,7 @@
 #include "base/command_line.h"
 #include "base/functional/callback_helpers.h"
 #include "base/memory/raw_ptr.h"
+#include "base/strings/stringprintf.h"
 #include "base/test/bind.h"
 #include "build/build_config.h"
 #include "chrome/browser/extensions/api/cookies/cookies_api.h"
@@ -18,6 +19,8 @@
 #include "content/public/browser/web_contents.h"
 #include "content/public/test/browser_test.h"
 #include "content/public/test/browser_test_utils.h"
+#include "extensions/browser/api_test_utils.h"
+#include "extensions/browser/extension_api_frame_id_map.h"
 #include "extensions/browser/extension_registry_observer.h"
 #include "extensions/browser/process_manager.h"
 #include "extensions/buildflags/buildflags.h"
@@ -223,6 +226,116 @@
   ASSERT_TRUE(RunTest("cookies/get_partition_key")) << message_;
 }
 
+IN_PROC_BROWSER_TEST_P(CookiesApiTest, TestGetPartitionKeyContextIsolation) {
+  // Set up an off-the-record window with a cross-site subframe.
+  const std::string default_response = "/defaultresponse";
+  content::WebContents* incognito_contents = PlatformOpenURLOffTheRecord(
+      profile(), embedded_test_server()->GetURL("a.com", default_response));
+  ASSERT_TRUE(incognito_contents);
+
+  const GURL cross_site_url =
+      embedded_test_server()->GetURL("b.com", default_response);
+  std::string script =
+      "var f = document.createElement('iframe');\n"
+      "f.src = '" +
+      cross_site_url.spec() +
+      "';\n"
+      "document.body.appendChild(f);\n";
+  EXPECT_TRUE(ExecJs(incognito_contents, script));
+  EXPECT_TRUE(WaitForLoadStop(incognito_contents));
+
+  content::RenderFrameHost* incognito_subframe =
+      content::ChildFrameAt(incognito_contents->GetPrimaryMainFrame(), 0);
+  ASSERT_TRUE(incognito_subframe);
+  int incognito_frame_id =
+      ExtensionApiFrameIdMap::GetFrameId(incognito_subframe);
+  std::string incognito_document_id =
+      ExtensionApiFrameIdMap::GetDocumentId(incognito_subframe).ToString();
+
+  // Load an extension in the regular profile without off-the-record access.
+  static constexpr char kManifest[] = R"({
+    "name": "Cookies Partition Key Isolation Test",
+    "version": "1.0",
+    "manifest_version": 3,
+    "permissions": ["cookies"],
+    "host_permissions": ["*://*/*"],
+    "background": {"service_worker": "background.js"}
+  })";
+
+  TestExtensionDir test_dir;
+  test_dir.WriteManifest(kManifest);
+  test_dir.WriteFile(FILE_PATH_LITERAL("background.js"), "// background");
+
+  const Extension* extension =
+      LoadExtension(test_dir.UnpackedPath(), {.allow_in_incognito = false});
+  ASSERT_TRUE(extension);
+
+  // Attempting to query the partition key for the off-the-record frame by
+  // frameId from the regular profile extension must be rejected.
+  {
+    auto function = base::MakeRefCounted<CookiesGetPartitionKeyFunction>();
+    function->set_extension(extension);
+    std::string error = api_test_utils::RunFunctionAndReturnError(
+        function.get(),
+        base::StringPrintf(R"([{"frameId": %d}])", incognito_frame_id),
+        profile(), api_test_utils::FunctionMode::kNone);
+    EXPECT_EQ("Invalid `frameId`.", error);
+  }
+
+  // Attempting to query the partition key for the off-the-record frame by
+  // documentId from the regular profile extension must be rejected.
+  {
+    auto function = base::MakeRefCounted<CookiesGetPartitionKeyFunction>();
+    function->set_extension(extension);
+    std::string error = api_test_utils::RunFunctionAndReturnError(
+        function.get(),
+        base::StringPrintf(R"([{"documentId": "%s"}])",
+                           incognito_document_id.c_str()),
+        profile(), api_test_utils::FunctionMode::kNone);
+    EXPECT_EQ("Invalid `documentId`.", error);
+  }
+
+  // Also verify that an extension without host permissions receives the same
+  // rejection error without any URL details.
+  static constexpr char kRestrictedManifest[] = R"({
+    "name": "Cookies Partition Key Restricted Test",
+    "version": "1.0",
+    "manifest_version": 3,
+    "permissions": ["cookies"],
+    "background": {"service_worker": "background.js"}
+  })";
+
+  TestExtensionDir restricted_test_dir;
+  restricted_test_dir.WriteManifest(kRestrictedManifest);
+  restricted_test_dir.WriteFile(FILE_PATH_LITERAL("background.js"),
+                                "// background");
+
+  const Extension* restricted_extension = LoadExtension(
+      restricted_test_dir.UnpackedPath(), {.allow_in_incognito = false});
+  ASSERT_TRUE(restricted_extension);
+
+  {
+    auto function = base::MakeRefCounted<CookiesGetPartitionKeyFunction>();
+    function->set_extension(restricted_extension);
+    std::string error = api_test_utils::RunFunctionAndReturnError(
+        function.get(),
+        base::StringPrintf(R"([{"frameId": %d}])", incognito_frame_id),
+        profile(), api_test_utils::FunctionMode::kNone);
+    EXPECT_EQ("Invalid `frameId`.", error);
+  }
+
+  {
+    auto function = base::MakeRefCounted<CookiesGetPartitionKeyFunction>();
+    function->set_extension(restricted_extension);
+    std::string error = api_test_utils::RunFunctionAndReturnError(
+        function.get(),
+        base::StringPrintf(R"([{"documentId": "%s"}])",
+                           incognito_document_id.c_str()),
+        profile(), api_test_utils::FunctionMode::kNone);
+    EXPECT_EQ("Invalid `documentId`.", error);
+  }
+}
+
 IN_PROC_BROWSER_TEST_F(ExtensionApiTest, OTRReceiverMojoConnectionError) {
   // Test that simulates a mojo connection error with an existing OTR profile.
   // This test verifies that the fix for crbug.com/472076020 works correctly:
diff --git a/chrome/test/data/extensions/api_test/cookies/get_partition_key/background.js b/chrome/test/data/extensions/api_test/cookies/get_partition_key/background.js
index 1b35c18..4ae63c2 100644
--- a/chrome/test/data/extensions/api_test/cookies/get_partition_key/background.js
+++ b/chrome/test/data/extensions/api_test/cookies/get_partition_key/background.js
@@ -74,10 +74,41 @@
         'Error: Invalid `frameId`.');
 
     await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey({tabId: tab.id, frameId: 999999}),
+        'Error: Invalid `frameId`.');
+
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey({frameId: -1}),
+        'Error: Invalid `frameId`.');
+
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey({frameId: 999999}),
+        'Error: Invalid `frameId`.');
+
+    await chrome.test.assertPromiseRejects(
         chrome.cookies.getPartitionKey(
             {tabId: tab.id, frameId: 0, documentId: ''}),
         'Error: Invalid `documentId`.');
 
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey(
+            {documentId: '0123456789abcdef0123456789abcdef'}),
+        'Error: Invalid `documentId`.');
+
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey(
+            {documentId: topLevelFrame.documentId, tabId: (tab.id - 1)}),
+        'Error: Provided `tabId` and `frameId` do not match the frame.');
+
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey(
+            {documentId: crossSiteFrame.documentId, frameId: 999999}),
+        'Error: Provided `tabId` and `frameId` do not match the frame.');
+
+    await chrome.test.assertPromiseRejects(
+        chrome.cookies.getPartitionKey({}),
+        'Error: Either `documentId` or `tabId` must be specified.');
+
     chrome.test.succeed();
   },
   async function testHostPermissionsRequired() {
@@ -127,6 +158,20 @@
     });
     chrome.test.assertEq(expectedTopLevelKey, actualPartitionKey);
 
+    // Providing documentId and matching tabId without frameId.
+    actualPartitionKey = await chrome.cookies.getPartitionKey({
+      tabId: tab.id,
+      documentId: topLevelDocId,
+    });
+    chrome.test.assertEq(expectedTopLevelKey, actualPartitionKey);
+
+    // Providing documentId and matching frameId without tabId.
+    actualPartitionKey = await chrome.cookies.getPartitionKey({
+      frameId: topLevelFrame.frameId,
+      documentId: topLevelDocId,
+    });
+    chrome.test.assertEq(expectedTopLevelKey, actualPartitionKey);
+
     // Providing no frameId defaults to topLevel frameId (0).
     actualPartitionKey = await chrome.cookies.getPartitionKey({
       tabId: tab.id,
@@ -157,6 +202,20 @@
     });
     chrome.test.assertEq(expectedCrossSiteKey, actualPartitionKey);
 
+    // Providing documentId and matching tabId without frameId.
+    actualPartitionKey = await chrome.cookies.getPartitionKey({
+      tabId: tab.id,
+      documentId: crossSiteDocId,
+    });
+    chrome.test.assertEq(expectedCrossSiteKey, actualPartitionKey);
+
+    // Providing documentId and matching frameId without tabId.
+    actualPartitionKey = await chrome.cookies.getPartitionKey({
+      frameId: crossSiteFrame.frameId,
+      documentId: crossSiteDocId,
+    });
+    chrome.test.assertEq(expectedCrossSiteKey, actualPartitionKey);
+
     chrome.test.succeed();
   },
 ]);