From f62380d3732bff26e3766a58b0e261b6a69ef51e Mon Sep 17 00:00:00 2001
From: jrandolf <101637635+jrandolf@users.noreply.github.com>
Date: Tue, 6 Feb 2024 17:34:41 +0100
Subject: [PATCH] chore: add storage.getCookies and storage.setCookie to
 bidi/core (#11847)

Co-authored-by: Alex Rudenko <OrKoN@users.noreply.github.com>
---
 .../src/bidi/core/BrowsingContext.ts          | 41 +++++++++++++++++++
 .../src/bidi/core/UserContext.ts              | 40 ++++++++++++++++++
 2 files changed, 81 insertions(+)

diff --git a/packages/puppeteer-core/src/bidi/core/BrowsingContext.ts b/packages/puppeteer-core/src/bidi/core/BrowsingContext.ts
index 9bec2a506c2..dbf4eee94ec 100644
--- a/packages/puppeteer-core/src/bidi/core/BrowsingContext.ts
+++ b/packages/puppeteer-core/src/bidi/core/BrowsingContext.ts
@@ -57,6 +57,14 @@ export type SetViewportOptions = Omit<
   'context'
 >;
 
+/**
+ * @internal
+ */
+export type GetCookiesOptions = Omit<
+  Bidi.Storage.GetCookiesParameters,
+  'partition'
+>;
+
 /**
  * @internal
  */
@@ -464,6 +472,39 @@ export class BrowsingContext extends EventEmitter<{
     await this.userContext.browser.removePreloadScript(script);
   }
 
+  @throwIfDisposed<BrowsingContext>(context => {
+    // SAFETY: Disposal implies this exists.
+    return context.#reason!;
+  })
+  async getCookies(
+    options: GetCookiesOptions = {}
+  ): Promise<Bidi.Network.Cookie[]> {
+    const {
+      result: {cookies},
+    } = await this.#session.send('storage.getCookies', {
+      ...options,
+      partition: {
+        type: 'context',
+        context: this.id,
+      },
+    });
+    return cookies;
+  }
+
+  @throwIfDisposed<BrowsingContext>(context => {
+    // SAFETY: Disposal implies this exists.
+    return context.#reason!;
+  })
+  async setCookie(cookie: Bidi.Storage.PartialCookie): Promise<void> {
+    await this.#session.send('storage.setCookie', {
+      cookie,
+      partition: {
+        type: 'context',
+        context: this.id,
+      },
+    });
+  }
+
   [disposeSymbol](): void {
     this.#reason ??=
       'Browsing context already closed, probably because the user context closed.';
diff --git a/packages/puppeteer-core/src/bidi/core/UserContext.ts b/packages/puppeteer-core/src/bidi/core/UserContext.ts
index 1bc2b372cf3..b1c7e6eceb6 100644
--- a/packages/puppeteer-core/src/bidi/core/UserContext.ts
+++ b/packages/puppeteer-core/src/bidi/core/UserContext.ts
@@ -12,6 +12,7 @@ import {inertIfDisposed, throwIfDisposed} from '../../util/decorators.js';
 import {DisposableStack, disposeSymbol} from '../../util/disposable.js';
 
 import type {Browser} from './Browser.js';
+import type {GetCookiesOptions} from './BrowsingContext.js';
 import {BrowsingContext} from './BrowsingContext.js';
 
 /**
@@ -174,6 +175,45 @@ export class UserContext extends EventEmitter<{
     }
   }
 
+  @throwIfDisposed<UserContext>(context => {
+    // SAFETY: Disposal implies this exists.
+    return context.#reason!;
+  })
+  async getCookies(
+    options: GetCookiesOptions = {},
+    sourceOrigin: string | undefined = undefined
+  ): Promise<Bidi.Network.Cookie[]> {
+    const {
+      result: {cookies},
+    } = await this.#session.send('storage.getCookies', {
+      ...options,
+      partition: {
+        type: 'storageKey',
+        userContext: this.#id,
+        sourceOrigin,
+      },
+    });
+    return cookies;
+  }
+
+  @throwIfDisposed<UserContext>(context => {
+    // SAFETY: Disposal implies this exists.
+    return context.#reason!;
+  })
+  async setCookie(
+    cookie: Bidi.Storage.PartialCookie,
+    sourceOrigin?: string
+  ): Promise<void> {
+    await this.#session.send('storage.setCookie', {
+      cookie,
+      partition: {
+        type: 'storageKey',
+        sourceOrigin,
+        userContext: this.id,
+      },
+    });
+  }
+
   [disposeSymbol](): void {
     this.#reason ??=
       'User context already closed, probably because the browser disconnected/closed.';