From e76e820fecaed6ef6d96e52bc7b09daff7021f63 Mon Sep 17 00:00:00 2001
From: Antoine Arlaud <antoine@snyk.io>
Date: Fri, 6 Dec 2024 15:27:34 +0100
Subject: [PATCH] fix: add auth headers on all flows

---
 lib/client/config/configHelpers.ts                      | 7 +++++++
 lib/client/dispatcher/client/api.ts                     | 9 ++++++---
 lib/client/socket.ts                                    | 6 ++++--
 lib/hybrid-sdk/http/downstream-post-stream-to-server.ts | 6 ++++++
 4 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/lib/client/config/configHelpers.ts b/lib/client/config/configHelpers.ts
index b36103cd6..88291346c 100644
--- a/lib/client/config/configHelpers.ts
+++ b/lib/client/config/configHelpers.ts
@@ -5,6 +5,12 @@ import { getConfig, loadBrokerConfig } from '../../common/config/config';
 import version from '../../common/utils/version';
 import { CONFIGURATION } from '../../common/types/options';
 
+let globalClientOpts: Record<string, any> = {};
+
+export const getClientOpts = () => {
+  return globalClientOpts;
+};
+
 export const reloadConfig = async (clientOpts) => {
   // Reload config with connection
   await loadBrokerConfig();
@@ -14,6 +20,7 @@ export const reloadConfig = async (clientOpts) => {
     clientOpts.config,
     globalConfig.config,
   ) as Record<string, any> as CONFIGURATION;
+  globalClientOpts = clientOpts;
 };
 
 export const getClientConfigMetadata = (
diff --git a/lib/client/dispatcher/client/api.ts b/lib/client/dispatcher/client/api.ts
index 37784bd09..8288eca75 100644
--- a/lib/client/dispatcher/client/api.ts
+++ b/lib/client/dispatcher/client/api.ts
@@ -7,6 +7,7 @@ import {
   DispatcherServiceClient,
   ServerId,
 } from '../dispatcher-service';
+import { getClientOpts } from '../../config/configHelpers';
 
 export class HttpDispatcherServiceClient implements DispatcherServiceClient {
   private readonly version = '2022-12-01~experimental';
@@ -24,12 +25,14 @@ export class HttpDispatcherServiceClient implements DispatcherServiceClient {
       const path = `/hidden/broker/${params.hashedBrokerToken}/connections/${params.brokerClientId}`;
       const url = new URL(path, this.baseUrl);
       url.searchParams.append('version', this.version);
+      const headers = { 'Content-type': 'application/vnd.api+json' };
+      if (getClientOpts().accessToken) {
+        headers['Authorization'] = getClientOpts().accessToken?.authHeader;
+      }
       const req: PostFilterPreparedRequest = {
         url: url.toString(),
         method: 'POST',
-        headers: {
-          'Content-type': 'application/vnd.api+json',
-        },
+        headers,
         body: JSON.stringify({
           data: {
             attributes: {
diff --git a/lib/client/socket.ts b/lib/client/socket.ts
index db23c7132..cd8f5c942 100644
--- a/lib/client/socket.ts
+++ b/lib/client/socket.ts
@@ -138,6 +138,8 @@ export const createWebSocket = (
     socketSettings['transport'] = {
       extraHeaders: {
         Authorization: clientOpts.accessToken?.authHeader,
+        'x-snyk-broker-client-id': identifyingMetadata.clientId,
+        'x-snyk-broker-client-role': identifyingMetadata.role,
       },
     };
   }
@@ -171,8 +173,8 @@ export const createWebSocket = (
         clientOpts.config.brokerClientConfiguration.common.oauth!.clientSecret,
       );
 
-      // websocket.transport.extraHeaders['Authorization'] =
-      //   clientOpts.accessToken!.authHeader;
+      websocket.transport.extraHeaders['Authorization'] =
+        clientOpts.accessToken!.authHeader;
       // websocket.end();
       // websocket.open();
       timeoutHandlerId = setTimeout(
diff --git a/lib/hybrid-sdk/http/downstream-post-stream-to-server.ts b/lib/hybrid-sdk/http/downstream-post-stream-to-server.ts
index 780c87eb6..8384e6775 100644
--- a/lib/hybrid-sdk/http/downstream-post-stream-to-server.ts
+++ b/lib/hybrid-sdk/http/downstream-post-stream-to-server.ts
@@ -9,6 +9,7 @@ import { bootstrap } from 'global-agent';
 import https from 'https';
 import http from 'http';
 import { getConfig } from '../../common/config/config';
+import { getClientOpts } from '../../client/config/configHelpers';
 
 const BROKER_CONTENT_TYPE = 'application/vnd.broker.stream+octet-stream';
 
@@ -68,6 +69,7 @@ class BrokerServerPostResponseHandler {
           this.#streamingId
         }`,
       );
+
       if (this.#serverId) {
         url.searchParams.append('server_id', this.#serverId);
       }
@@ -94,6 +96,10 @@ class BrokerServerPostResponseHandler {
           ? parseInt(this.#config.brokerClientPostTimeout)
           : 1200000,
       };
+      if (getClientOpts().accessToken) {
+        options.headers['authorization'] =
+          getClientOpts().accessToken.authHeader;
+      }
 
       this.#brokerSrvPostRequestHandler = client.request(
         brokerServerPostRequestUrl,