[INS-1833] Include Auth Header in Headers mapping for WebSocket Connection (#5120)

* add auth to the header

* remove console log

* remove unneeded async

* add success redirect logic to websocket server

* add unexpected-response handler

Co-authored-by: jackkav <[email protected]>

Author: 2 people

packages/insomnia-smoke-test/server/websocket.ts

@@ -1,13 +1,20 @@
 import { IncomingMessage, Server } from 'http';
+import { Socket } from 'net';
 import { WebSocket, WebSocketServer } from 'ws';
 
 /**
  * Starts an echo WebSocket server that receives messages from a client and echoes them back.
  */
 export function startWebSocketServer(server: Server, httpsServer: Server) {
-  const wsServer = new WebSocketServer({ server });
-  const wssServer = new WebSocketServer({ server: httpsServer });
+  const wsServer = new WebSocketServer({ noServer: true });
+  const wssServer = new WebSocketServer({ noServer: true });
 
+  server.on('upgrade', (request, socket, head) => {
+    upgrade(wsServer, request, socket, head);
+  });
+  httpsServer.on('upgrade', (request, socket, head) => {
+    upgrade(wssServer, request, socket, head);
+  });
   wsServer.on('connection', handleConnection);
   wssServer.on('connection', handleConnection);
 }
@@ -31,3 +38,33 @@ const handleConnection = (ws: WebSocket, req: IncomingMessage) => {
     console.log('WebSocket connection was closed');
   });
 };
+const redirectOnSuccess = (socket: Socket) => {
+  socket.end(`HTTP/1.1 302 Found
+Location: ws://localhost:4010
+
+`);
+  return;
+};
+const upgrade = (wss: WebSocketServer, request: IncomingMessage, socket: Socket, head: Buffer) => {
+  if (request.url === '/redirect') {
+    return redirectOnSuccess(socket);
+  }
+  if (request.url === '/bearer') {
+    if (request.headers.authorization !== 'Bearer insomnia-cool-token-!!!1112113243111') {
+      socket.end('HTTP/1.1 401 Unauthorized\n\n');
+      return;
+    }
+    return redirectOnSuccess(socket);
+  }
+  if (request.url === '/basic-auth') {
+    // login with user:password
+    if (request.headers.authorization !== 'Basic dXNlcjpwYXNzd29yZA==') {
+      socket.end('HTTP/1.1 401 Unauthorized\n\n');
+      return;
+    }
+    return redirectOnSuccess(socket);
+  }
+  wss.handleUpgrade(request, socket, head, ws => {
+    wss.emit('connection', ws, request);
+  });
+};

packages/insomnia/src/main/network/websocket.ts

@@ -1,5 +1,6 @@
 import electron, { ipcMain } from 'electron';
 import fs from 'fs';
+import { IncomingMessage } from 'http';
 import { setDefaultProtocol } from 'insomnia-url';
 import mkdirp from 'mkdirp';
 import path from 'path';
@@ -13,11 +14,15 @@ import {
   WebSocket,
 } from 'ws';
 
+import { AUTH_BASIC, AUTH_BEARER } from '../../common/constants';
 import { generateId } from '../../common/misc';
 import { websocketRequest } from '../../models';
 import * as models from '../../models';
+import { RequestAuthentication, RequestHeader } from '../../models/request';
 import type { Response } from '../../models/response';
 import { BaseWebSocketRequest } from '../../models/websocket-request';
+import { getBasicAuthHeader } from '../../network/basic-auth/get-header';
+import { getBearerAuthHeader } from '../../network/bearer-auth/get-header';
 import { urlMatchesCertHost } from '../../network/url-matches-cert-host';
 
 export interface WebSocketConnection extends WebSocket {
@@ -96,6 +101,23 @@ function dispatchWebSocketEvent(target: Electron.WebContents, eventChannel: stri
   }
 }
 
+const parseResponseAndBuildTimeline = (url: string, incomingMessage: IncomingMessage, clientRequestHeaders: string) => {
+  const statusMessage = incomingMessage.statusMessage || '';
+  const statusCode = incomingMessage.statusCode || 0;
+  const httpVersion = incomingMessage.httpVersion;
+  const responseHeaders = Object.entries(incomingMessage.headers).map(([name, value]) => ({ name, value: value?.toString() || '' }));
+  const headersIn = responseHeaders.map(({ name, value }) => `${name}: ${value}`).join('\n');
+  const timeline = [
+    { value: `Preparing request to ${url}`, name: 'Text', timestamp: Date.now() },
+    { value: `Current time is ${new Date().toISOString()}`, name: 'Text', timestamp: Date.now() },
+    { value: 'Using HTTP 1.1', name: 'Text', timestamp: Date.now() },
+    { value: clientRequestHeaders, name: 'HeaderOut', timestamp: Date.now() },
+    { value: `HTTP/${httpVersion} ${statusCode} ${statusMessage}`, name: 'HeaderIn', timestamp: Date.now() },
+    { value: headersIn, name: 'HeaderIn', timestamp: Date.now() },
+  ];
+  return { timeline, responseHeaders, statusCode, statusMessage, httpVersion };
+};
+
 async function createWebSocketConnection(
   event: Electron.IpcMainInvokeEvent,
   options: { requestId: string; workspaceId: string }
@@ -122,11 +144,24 @@ async function createWebSocketConnection(
   try {
     const eventChannel = `webSocketRequest.connection.${responseId}.event`;
     const readyStateChannel = `webSocketRequest.connection.${request._id}.readyState`;
-
     // @TODO: Render nunjucks tags in these headers
     const reduceArrayToLowerCaseKeyedDictionary = (acc: { [key: string]: string }, { name, value }: BaseWebSocketRequest['headers'][0]) =>
       ({ ...acc, [name.toLowerCase() || '']: value || '' });
-    const headers = request.headers.filter(({ value, disabled }) => !!value && !disabled)
+    const headers = request.headers;
+    if (request.authentication.disabled === false) {
+      if (request.authentication.type === AUTH_BASIC) {
+        const { username, password, useISO88591 } = request.authentication;
+        const encoding = useISO88591 ? 'latin1' : 'utf8';
+        headers.push(getBasicAuthHeader(username, password, encoding));
+      }
+      if (request.authentication.type === AUTH_BEARER) {
+        const { token, prefix } = request.authentication;
+        headers.push(getBearerAuthHeader(token, prefix));
+      }
+    }
+
+    const lowerCasedEnabledHeaders = headers
+      .filter(({ value, disabled }) => !!value && !disabled)
       .reduce(reduceArrayToLowerCaseKeyedDictionary, {});
 
     const settings = await models.settings.getOrCreate();
@@ -158,34 +193,43 @@ async function createWebSocketConnection(
     });
 
     const ws = new WebSocket(request.url, {
-      headers,
+      headers: lowerCasedEnabledHeaders,
       cert: pemCertificates,
       key: pemCertificateKeys,
       pfx: pfxCertificates,
       rejectUnauthorized: settings.validateSSL,
+      followRedirects: true,
     });
     WebSocketConnections.set(options.requestId, ws);
 
-    ws.on('upgrade', async incoming => {
+    ws.on('upgrade', async incomingMessage => {
       // @ts-expect-error -- private property
-      const internalRequest = ws._req;
-      // response
-      const statusMessage = incoming.statusMessage || '';
-      const statusCode = incoming.statusCode || 0;
-      const httpVersion = incoming.httpVersion;
-      const responseHeaders = Object.entries(incoming.headers).map(([name, value]) => ({ name, value: value?.toString() || '' }));
-      const headersIn = responseHeaders.map(({ name, value }) => `${name}: ${value}`).join('\n');
-
-      // @TODO: We may want to add set-cookie handling here.
-      [
-        { value: `Preparing request to ${request.url}`, name: 'Text', timestamp: Date.now() },
-        { value: `Current time is ${new Date().toISOString()}`, name: 'Text', timestamp: Date.now() },
-        { value: 'Using HTTP 1.1', name: 'Text', timestamp: Date.now() },
-        { value: internalRequest._header, name: 'HeaderOut', timestamp: Date.now() },
-        { value: `HTTP/${httpVersion} ${statusCode} ${statusMessage}`, name: 'HeaderIn', timestamp: Date.now() },
-        { value: headersIn, name: 'HeaderIn', timestamp: Date.now() },
-      ].map(t => timelineFileStreams.get(options.requestId)?.write(JSON.stringify(t) + '\n'));
-
+      const internalRequestHeader = ws._req._header;
+      const { timeline, responseHeaders, statusCode, statusMessage, httpVersion } = parseResponseAndBuildTimeline(request.url, incomingMessage, internalRequestHeader);
+      timeline.map(t => timelineFileStreams.get(options.requestId)?.write(JSON.stringify(t) + '\n'));
+      const responsePatch: Partial<Response> = {
+        _id: responseId,
+        parentId: request._id,
+        headers: responseHeaders,
+        url: request.url,
+        statusCode,
+        statusMessage,
+        httpVersion,
+        elapsedTime: performance.now() - start,
+        timelinePath,
+        bodyPath: responseBodyPath,
+        // NOTE: required for legacy zip workaround
+        bodyCompression: null,
+      };
+      const settings = await models.settings.getOrCreate();
+      models.response.create(responsePatch, settings.maxHistoryResponses);
+      models.requestMeta.updateOrCreateByParentId(request._id, { activeResponseId: null });
+    });
+    ws.on('unexpected-response', async (clientRequest, incomingMessage) => {
+      // @ts-expect-error -- private property
+      const internalRequestHeader = clientRequest._header;
+      const { timeline, responseHeaders, statusCode, statusMessage, httpVersion } = parseResponseAndBuildTimeline(request.url, incomingMessage, internalRequestHeader);
+      timeline.map(t => timelineFileStreams.get(options.requestId)?.write(JSON.stringify(t) + '\n'));
       const responsePatch: Partial<Response> = {
         _id: responseId,
         parentId: request._id,
@@ -203,6 +247,7 @@ async function createWebSocketConnection(
       const settings = await models.settings.getOrCreate();
       models.response.create(responsePatch, settings.maxHistoryResponses);
       models.requestMeta.updateOrCreateByParentId(request._id, { activeResponseId: null });
+      deleteRequestMaps(request._id, `Unexpected response ${incomingMessage.statusCode}`);
     });
 
     ws.addEventListener('open', () => {
@@ -413,3 +458,31 @@ electron.app.on('window-all-closed', () => {
     ws.close();
   });
 });
+
+export function getAuthHeader(authentication: RequestAuthentication): RequestHeader | undefined {
+  if (!authentication || authentication.disabled) {
+    return;
+  }
+
+  switch (authentication.type) {
+    case 'basic': {
+      const { username, password, useISO88591 } = authentication;
+      const encoding = useISO88591 ? 'latin1' : 'utf8';
+      const header = getBasicAuthHeader(username, password, encoding);
+      return header;
+    }
+
+    case 'bearer': {
+      const { token, prefix } = authentication;
+      return getBearerAuthHeader(token, prefix);
+    }
+
+    case 'digest': {
+      return;
+    }
+
+    default: {
+      return;
+    }
+  }
+}