From cb0a0302395faf5ddfcce4d41e3e1f38f768a651 Mon Sep 17 00:00:00 2001
From: Antoine Arlaud <antoine@snyk.io>
Date: Mon, 11 Nov 2024 12:53:18 +0100
Subject: [PATCH] chore: add server side broker header

---
 lib/common/relay/forwardWebsocketRequest.ts |  5 +++++
 test/functional/client-server.test.ts       | 13 +++++++++++++
 test/functional/webhook.test.ts             | 14 ++++++++++++++
 test/setup/test-web-server.ts               |  8 ++++++++
 4 files changed, 40 insertions(+)

diff --git a/lib/common/relay/forwardWebsocketRequest.ts b/lib/common/relay/forwardWebsocketRequest.ts
index 6d9208ca5..548ec3087 100644
--- a/lib/common/relay/forwardWebsocketRequest.ts
+++ b/lib/common/relay/forwardWebsocketRequest.ts
@@ -229,6 +229,11 @@ export const forwardWebSocketRequest = (
         );
         payload.headers[contentLengthHeader] = computeContentLength(payload);
       }
+      if (options.config.brokerType !== 'client') {
+        preparedRequest.req.headers['x-snyk-broker'] = `${maskToken(
+          connectionIdentifier,
+        )}`;
+      }
 
       incrementHttpRequestsTotal(false, 'outbound-request');
       payload.streamingID
diff --git a/test/functional/client-server.test.ts b/test/functional/client-server.test.ts
index 97aeb21b5..15b2efd6d 100644
--- a/test/functional/client-server.test.ts
+++ b/test/functional/client-server.test.ts
@@ -14,6 +14,7 @@ import {
   waitForBrokerClientConnections,
 } from '../setup/broker-server';
 import { TestWebServer, createTestWebServer } from '../setup/test-web-server';
+import { maskToken } from '../../lib/common/utils/token';
 
 const fixtures = path.resolve(__dirname, '..', 'fixtures');
 const serverAccept = path.join(fixtures, 'server', 'filters.json');
@@ -74,6 +75,18 @@ describe('proxy requests originating from behind the broker client', () => {
     });
   });
 
+  it('successfully broker POST with x-broker-server header', async () => {
+    const response = await axiosClient.post(
+      `http://localhost:${bc.port}/echo-headers`,
+      { some: { example: 'json' } },
+    );
+
+    expect(response.status).toEqual(200);
+    expect(response.data['x-snyk-broker']).toStrictEqual(
+      maskToken(brokerToken),
+    );
+  });
+
   it('successfully broker exact bytes of POST body', async () => {
     // stringify the JSON unusually to ensure an unusual exact body
     const body = Buffer.from(
diff --git a/test/functional/webhook.test.ts b/test/functional/webhook.test.ts
index 97a9e463e..be6cf32c7 100644
--- a/test/functional/webhook.test.ts
+++ b/test/functional/webhook.test.ts
@@ -16,6 +16,7 @@ import {
 } from '../setup/broker-server';
 import { TestWebServer, createTestWebServer } from '../setup/test-web-server';
 import { DEFAULT_TEST_WEB_SERVER_PORT } from '../setup/constants';
+import { maskToken } from '../../lib/common/utils/token';
 
 const fixtures = path.resolve(__dirname, '..', 'fixtures');
 const serverAccept = path.join(fixtures, 'server', 'filters-webhook.json');
@@ -80,4 +81,17 @@ describe('proxy requests originating from behind the broker client', () => {
     expect(response.status).toEqual(200);
     expect(response.data).toStrictEqual('Received webhook via API');
   });
+  it('successfully broker injects x-snyk-broker header to Webhook calls', async () => {
+    await closeBrokerServer(bs);
+
+    const response = await axiosClient.post(
+      `http://localhost:${bc.port}/webhook/github/return-req-headers`,
+      { some: { example: 'json' } },
+    );
+
+    expect(response.status).toEqual(200);
+    expect(response.data['x-snyk-broker']).toStrictEqual(
+      maskToken('broker-token-12345'),
+    );
+  });
 });
diff --git a/test/setup/test-web-server.ts b/test/setup/test-web-server.ts
index d3f1af265..24c5e69a8 100644
--- a/test/setup/test-web-server.ts
+++ b/test/setup/test-web-server.ts
@@ -125,6 +125,14 @@ const applyEchoRoutes = (app: Express) => {
     },
   );
 
+  echoRouter.post(
+    '/webhook/github/return-req-headers',
+    (req: express.Request, resp: express.Response) => {
+      resp.status(200);
+      resp.send(req.headers);
+    },
+  );
+
   echoRouter.get('/test', (_: express.Request, resp: express.Response) => {
     resp.status(200);
     resp.send('All good');