From 642bfaa6e71901ab98b5395d9d659e8def38c250 Mon Sep 17 00:00:00 2001 From: Antoine Arlaud Date: Fri, 1 Mar 2024 21:35:30 +0100 Subject: [PATCH] feat: add ws or http response selection logic --- lib/common/relay/forwardHttpRequest.ts | 8 ++++- lib/common/relay/forwardWebsocketRequest.ts | 15 ++++++++- lib/common/types/log.ts | 1 + .../server-client-universal.test.ts | 11 ++++++- test/functional/server-client.test.ts | 32 +++++++++++++++++++ 5 files changed, 64 insertions(+), 3 deletions(-) diff --git a/lib/common/relay/forwardHttpRequest.ts b/lib/common/relay/forwardHttpRequest.ts index 12555b6bf..8872c84db 100644 --- 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 index 5dd6f4e09..96e70ee1e 100644 --- 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 index f6d074bc4..1cfdd7253 100644 --- 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 index 5aad7651f..4b063f0b6 100644 --- 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 index f65999658..ad7309226 100644 --- 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'); });