From a623dc751952817f4676e9f6836792a1e11de234 Mon Sep 17 00:00:00 2001
From: Alex Rudenko <OrKoN@users.noreply.github.com>
Date: Thu, 22 Feb 2024 11:01:25 +0100
Subject: [PATCH] feat: support user context in permissions (#1805)

This PR adds `goog:` prefixed user context to allow use cases required
by Puppeteer. Once permissions spec issue is resolved
https://github.com/w3c/permissions/pull/438 we can remove the prefix.
---
 .../permissions/PermissionsProcessor.ts       |  7 +++
 src/protocol-parser/protocol-parser.ts        | 12 +++--
 tests/permissions/__init__.py                 |  9 +++-
 tests/permissions/test_set_permission.py      | 52 ++++++++++++++++++-
 4 files changed, 73 insertions(+), 7 deletions(-)

diff --git a/src/bidiMapper/domains/permissions/PermissionsProcessor.ts b/src/bidiMapper/domains/permissions/PermissionsProcessor.ts
index c629c9538..434adf8cd 100644
--- a/src/bidiMapper/domains/permissions/PermissionsProcessor.ts
+++ b/src/bidiMapper/domains/permissions/PermissionsProcessor.ts
@@ -33,8 +33,15 @@ export class PermissionsProcessor {
     params: Permissions.SetPermissionParameters
   ): Promise<EmptyResult> {
     try {
+      const userContextId = (params as {'goog:userContext'?: string})[
+        'goog:userContext'
+      ];
       await this.#browserCdpClient.sendCommand('Browser.setPermission', {
         origin: params.origin,
+        browserContextId:
+          userContextId && userContextId !== 'default'
+            ? userContextId
+            : undefined,
         permission: {
           name: params.descriptor.name,
         },
diff --git a/src/protocol-parser/protocol-parser.ts b/src/protocol-parser/protocol-parser.ts
index 58e29a197..c9c4e753e 100644
--- a/src/protocol-parser/protocol-parser.ts
+++ b/src/protocol-parser/protocol-parser.ts
@@ -350,9 +350,13 @@ export namespace Permissions {
   export function parseSetPermissionsParams(
     params: unknown
   ): Protocol.Permissions.SetPermissionParameters {
-    return parseObject(
-      params,
-      WebDriverBidiPermissions.Permissions.SetPermissionParametersSchema
-    ) as Protocol.Permissions.SetPermissionParameters;
+    return {
+      // TODO: remove once "goog:" attributes are not needed.
+      ...(params as object),
+      ...(parseObject(
+        params,
+        WebDriverBidiPermissions.Permissions.SetPermissionParametersSchema
+      ) as Protocol.Permissions.SetPermissionParameters),
+    };
   }
 }
diff --git a/tests/permissions/__init__.py b/tests/permissions/__init__.py
index 6465f20c4..3794bbc33 100644
--- a/tests/permissions/__init__.py
+++ b/tests/permissions/__init__.py
@@ -45,7 +45,11 @@ async def query_permission(websocket, context_id, name):
     return result['result']['value']
 
 
-async def set_permission(websocket, origin, descriptor, state):
+async def set_permission(websocket,
+                         origin,
+                         descriptor,
+                         state,
+                         user_context=None):
     """ Set a permission via the permissions.setPermission command."""
     return await execute_command(
         websocket, {
@@ -53,6 +57,7 @@ async def set_permission(websocket, origin, descriptor, state):
             'params': {
                 'origin': origin,
                 'descriptor': descriptor,
-                'state': state
+                'state': state,
+                'goog:userContext': user_context
             }
         })
diff --git a/tests/permissions/test_set_permission.py b/tests/permissions/test_set_permission.py
index de6239a23..dc4f531ec 100644
--- a/tests/permissions/test_set_permission.py
+++ b/tests/permissions/test_set_permission.py
@@ -15,7 +15,7 @@
 
 import pytest
 from permissions import get_origin, query_permission, set_permission
-from test_helpers import goto_url
+from test_helpers import execute_command, goto_url
 
 
 @pytest.mark.asyncio
@@ -34,3 +34,53 @@ async def test_permissions_set_permission(websocket, context_id, example_url):
     assert resp == {}
     assert await query_permission(websocket, context_id,
                                   'geolocation') == 'prompt'
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip(reason="See chromium-bidi/issues#1610")
+async def test_permissions_set_permission_in_user_context(
+        websocket, context_id, example_url):
+    await goto_url(websocket, context_id, example_url)
+
+    user_context = await execute_command(websocket, {
+        "method": "browser.createUserContext",
+        "params": {}
+    })
+
+    browsing_context = await execute_command(
+        websocket, {
+            "method": "browsingContext.create",
+            "params": {
+                "type": "tab",
+                "userContext": user_context["userContext"]
+            }
+        })
+
+    origin = get_origin(example_url)
+
+    await goto_url(websocket, browsing_context['context'], example_url)
+
+    assert await query_permission(websocket, context_id,
+                                  'geolocation') == 'prompt'
+    assert await query_permission(websocket, browsing_context['context'],
+                                  'geolocation') == 'prompt'
+
+    resp = await set_permission(websocket,
+                                origin, {'name': 'geolocation'},
+                                'granted',
+                                user_context=user_context["userContext"])
+    assert resp == {}
+    assert await query_permission(websocket, context_id,
+                                  'geolocation') == 'prompt'
+    assert await query_permission(websocket, browsing_context['context'],
+                                  'geolocation') == 'granted'
+
+    resp = await set_permission(websocket,
+                                origin, {'name': 'geolocation'},
+                                'prompt',
+                                user_context=user_context["userContext"])
+    assert resp == {}
+    assert await query_permission(websocket, context_id,
+                                  'geolocation') == 'prompt'
+    assert await query_permission(websocket, browsing_context['context'],
+                                  'geolocation') == 'prompt'