From 470b1050f1e3b088b38909ab3631b600d35d0b3b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20M=C3=BCller?= <daniel.mueller@surnet.ch>
Date: Tue, 19 Nov 2024 09:04:09 +0100
Subject: [PATCH 1/2] Add optional function to modify headers

---
 packages/sandbox/src/EmbeddedSandbox.ts           |  8 +++++++-
 .../src/helpers/postMessageRelayHelpers.ts        |  4 ++++
 packages/sandbox/src/setupSandboxEmbedRelay.ts    | 15 +++++++++++++--
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/packages/sandbox/src/EmbeddedSandbox.ts b/packages/sandbox/src/EmbeddedSandbox.ts
index 229fa57c..34e42daf 100644
--- a/packages/sandbox/src/EmbeddedSandbox.ts
+++ b/packages/sandbox/src/EmbeddedSandbox.ts
@@ -1,6 +1,6 @@
 import { EMBEDDABLE_SANDBOX_URL, IFRAME_DOM_ID } from './helpers/constants';
 import { defaultHandleRequest } from './helpers/defaultHandleRequest';
-import type { HandleRequest } from './helpers/postMessageRelayHelpers';
+import type { HandleRequest, ModifyHeaders } from './helpers/postMessageRelayHelpers';
 import { setupSandboxEmbedRelay } from './setupSandboxEmbedRelay';
 import packageJSON from '../package.json';
 import type { JSONObject } from './helpers/types';
@@ -53,6 +53,11 @@ export interface EmbeddableSandboxOptions {
    */
   handleRequest?: HandleRequest;
 
+   /**
+    * optional. Function that accepts the original headers and returns the modified headers.
+    */
+  modifyHeaders?: ModifyHeaders;
+
   /**
    * optional. If this is passed, its value will take precedence over your sandbox connection settings `includeCookies` value.
    * If you pass `handleRequest`, that will override this value and its behavior.
@@ -122,6 +127,7 @@ export class EmbeddedSandbox {
     this.disposable = setupSandboxEmbedRelay({
       embeddedSandboxIFrameElement: this.embeddedSandboxIFrameElement,
       handleRequest: this.handleRequest,
+      modifyHeaders: this.options.modifyHeaders,
       __testLocal__: !!this.__testLocal__,
     });
   }
diff --git a/packages/sandbox/src/helpers/postMessageRelayHelpers.ts b/packages/sandbox/src/helpers/postMessageRelayHelpers.ts
index 3cac55ba..473aed4f 100644
--- a/packages/sandbox/src/helpers/postMessageRelayHelpers.ts
+++ b/packages/sandbox/src/helpers/postMessageRelayHelpers.ts
@@ -41,6 +41,10 @@ export type HandleRequest = (
   endpointUrl: string,
   options: Omit<RequestInit, 'headers'> & { headers: Record<string, string> }
 ) => Promise<Response>;
+export type ModifyHeaders = (
+  endpointUrl: string,
+  headers: Record<string, string> | undefined
+) => Promise<Record<string, string> | undefined>;
 
 export type SocketStatus = 'disconnected' | 'connecting' | 'connected';
 
diff --git a/packages/sandbox/src/setupSandboxEmbedRelay.ts b/packages/sandbox/src/setupSandboxEmbedRelay.ts
index e0ecd6f4..b894ec6f 100644
--- a/packages/sandbox/src/setupSandboxEmbedRelay.ts
+++ b/packages/sandbox/src/setupSandboxEmbedRelay.ts
@@ -12,22 +12,25 @@ import {
   handleAuthenticationPostMessage,
   HandleRequest,
   IncomingEmbedMessage,
+  ModifyHeaders,
   sendPostMessageToEmbed,
 } from './helpers/postMessageRelayHelpers';
 import { executeSubscription } from './helpers/subscriptionPostMessageRelayHelpers';
 
 export function setupSandboxEmbedRelay({
   handleRequest,
+  modifyHeaders,
   embeddedSandboxIFrameElement,
   __testLocal__,
 }: {
   handleRequest: HandleRequest;
+  modifyHeaders?: ModifyHeaders;
   embeddedSandboxIFrameElement: HTMLIFrameElement;
   __testLocal__: boolean;
 }) {
   const embedUrl = EMBEDDABLE_SANDBOX_URL(__testLocal__);
   // Callback definition
-  const onPostMessageReceived = (event: IncomingEmbedMessage) => {
+  const onPostMessageReceived = async (event: IncomingEmbedMessage) => {
     handleAuthenticationPostMessage({
       event,
       embedUrl,
@@ -85,7 +88,7 @@ export function setupSandboxEmbedRelay({
         data.operationId
       ) {
         // Extract the operation details from the event.data object
-        const { operation, variables, operationName, operationId, headers } =
+        const { operation, variables, operationName, operationId, headers: originalHeaders } =
           data;
 
         if (isQueryOrMutation) {
@@ -95,6 +98,10 @@ export function setupSandboxEmbedRelay({
               'Something went wrong, we should not have gotten here. The sandbox endpoint url was not sent.'
             );
           }
+          // If the user has provided a function to modify headers, call it
+          const headers = modifyHeaders
+            ? await modifyHeaders(endpointUrl, originalHeaders)
+            : originalHeaders;
           executeOperation({
             endpointUrl,
             handleRequest,
@@ -112,6 +119,10 @@ export function setupSandboxEmbedRelay({
           });
         } else if (isSubscription) {
           const { httpMultipartParams } = data;
+          // If the user has provided a function to modify headers, call it
+          const headers = modifyHeaders
+            ? await modifyHeaders(data.subscriptionUrl, originalHeaders)
+            : originalHeaders;
           executeSubscription({
             operation,
             operationName,

From 67b7f6072707c4b0e986be1bafa2546e1047b8c7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20M=C3=BCller?= <daniel.mueller@surnet.ch>
Date: Tue, 19 Nov 2024 09:09:02 +0100
Subject: [PATCH 2/2] Add changeset with `npx changeset`

---
 .changeset/heavy-apples-hug.md | 5 +++++
 1 file changed, 5 insertions(+)
 create mode 100644 .changeset/heavy-apples-hug.md

diff --git a/.changeset/heavy-apples-hug.md b/.changeset/heavy-apples-hug.md
new file mode 100644
index 00000000..d03268f9
--- /dev/null
+++ b/.changeset/heavy-apples-hug.md
@@ -0,0 +1,5 @@
+---
+'@apollo/sandbox': minor
+---
+
+Adding support for optional `modifyHeaders` function which can modify the headers before sending the request