feat: add ws or http response selection logic

snyk · Mar 1, 2024 · 642bfaa · 642bfaa
1 parent 7664861
commit 642bfaa
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 3 deletions.
diff --git a/lib/common/relay/forwardHttpRequest.ts b/lib/common/relay/forwardHttpRequest.ts
@@ -31,6 +31,9 @@ export const forwardHttpRequest = (
     // If this is the server, we should receive a Snyk-Request-Id header from upstream
     // If this is the client, we will have to generate one
     req.headers['snyk-request-id'] ||= uuid();
+    const responseWantedOverWs = req.headers['x-broker-ws-response']
+      ? true
+      : false;
     const logContext: ExtendedLogContext = {
       url: req.url,
       requestMethod: req.method,
@@ -212,7 +215,10 @@ export const forwardHttpRequest = (
       return res.status(401).send({ message: 'blocked', reason, url: req.url });
     } else {
       incrementHttpRequestsTotal(false, 'inbound-request');
-      if (res?.locals?.capabilities?.includes('post-streams')) {
+      if (
+        res?.locals?.capabilities?.includes('post-streams') &&
+        !responseWantedOverWs
+      ) {
         makeWebsocketRequestWithStreamingResponse(filterResponse);
       } else {
         makeWebsocketRequestWithWebsocketResponse(filterResponse);

diff --git a/lib/common/relay/forwardWebsocketRequest.ts b/lib/common/relay/forwardWebsocketRequest.ts
@@ -43,6 +43,9 @@ export const forwardWebSocketRequest = (
       hashedToken: hashToken(connectionIdentifier),
       transport:
         websocketConnectionHandler?.socket?.transport?.name ?? 'unknown',
+      responseMedium: payload.headers['x-broker-ws-response']
+        ? 'websocket'
+        : 'http',
     };
 
     if (!requestId) {
@@ -97,6 +100,13 @@ export const forwardWebSocketRequest = (
       responseData,
       isResponseFromRequestModule = false,
     ) => {
+      if (responseData) {
+        responseData['headers'] = responseData['headers'] ?? {};
+        responseData.headers['snyk-request-id'] = requestId;
+        responseData.headers['x-broker-ws-response'] =
+          responseData.headers['x-broker-ws-response'] ?? 'true';
+      }
+
       // Traffic over websockets
       if (payload.streamingID) {
         if (isResponseFromRequestModule) {
@@ -135,7 +145,10 @@ export const forwardWebSocketRequest = (
     };
 
     if (
-      websocketConnectionHandler?.capabilities?.includes('receive-post-streams')
+      websocketConnectionHandler?.capabilities?.includes(
+        'receive-post-streams',
+      ) &&
+      !emit
     ) {
       // Traffic over HTTP Post
       emit = postOverrideEmit;

diff --git a/lib/common/types/log.ts b/lib/common/types/log.ts
@@ -24,4 +24,5 @@ export interface ExtendedLogContext extends LogContext {
   responseBodyType?: string;
   ioErrorType?: string;
   ioOriginalBodySize?: string;
+  responseMedium?: string;
 }
diff --git a/test/functional/server-client-universal.test.ts b/test/functional/server-client-universal.test.ts
@@ -85,7 +85,12 @@ describe('proxy requests originating from behind the broker server', () => {
       `http://localhost:${bs.port}/broker/${process.env.BROKER_TOKEN_4}/echo-auth-header-with-bearer-auth/xyz`,
     );
 
-    // const response5 = await axiosClient.get(
+    const response5 = await axiosClient.get(
+      `http://localhost:${bs.port}/broker/${process.env.BROKER_TOKEN_4}/echo-auth-header-with-bearer-auth/xyz`,
+      { headers: { 'x-broker-ws-response': 'whatever' } },
+    );
+
+    // const response6 = await axiosClient.get(
     //   `http://localhost:${bs.port}/broker/${process.env.BROKER_TOKEN_3}/echo-auth-header-with-token-auth/xyz`,
     // );
 
@@ -101,6 +106,10 @@ describe('proxy requests originating from behind the broker server', () => {
     );
     expect(response4.status).toEqual(200);
     expect(response4.data).toEqual(`Bearer ${process.env.JIRA_PAT}`);
+
+    expect(response5.status).toEqual(200);
+    expect(response5.data).toEqual(`Bearer ${process.env.JIRA_PAT}`);
+    expect(response.headers['x-broker-ws-response']).not.toBeNull();
   });
 
   it('successfully warn logs requests without x-snyk-broker-type header', async () => {

diff --git a/test/functional/server-client.test.ts b/test/functional/server-client.test.ts
@@ -93,11 +93,43 @@ describe('proxy requests originating from behind the broker server', () => {
     expect(Buffer.from(response.data)).toEqual(body);
   });
 
+  it('successfully broker exact bytes of POST body with WS response', async () => {
+    // stringify the JSON unusually to ensure an unusual exact body
+    const body = Buffer.from(
+      JSON.stringify({ some: { example: 'json' } }, null, 5),
+    );
+    const response = await axiosClient.post(
+      `http://localhost:${bs.port}/broker/${brokerToken}/echo-body`,
+      body,
+      {
+        headers: {
+          'content-type': 'application/json',
+          'x-broker-ws-response': 'whatever',
+        },
+
+        transformResponse: (r) => r,
+      },
+    );
+    expect(response.status).toEqual(200);
+    expect(Buffer.from(response.data)).toEqual(body);
+    expect(response.headers['x-broker-ws-response']).not.toBeNull();
+  });
   it('successfully broker GET', async () => {
     const response = await axiosClient.get(
       `http://localhost:${bs.port}/broker/${brokerToken}/echo-param/xyz`,
     );
 
+    expect(response.headers['x-broker-ws-response']).not.toBeNull();
+    expect(response.status).toEqual(200);
+    expect(response.data).toEqual('xyz');
+  });
+
+  it('successfully broker GET with WS response', async () => {
+    const response = await axiosClient.get(
+      `http://localhost:${bs.port}/broker/${brokerToken}/echo-param/xyz`,
+      { headers: { 'x-broker-ws-response': 'whatever' } },
+    );
+    expect(response.headers['x-broker-ws-response']).not.toBeNull();
     expect(response.status).toEqual(200);
     expect(response.data).toEqual('xyz');
   });