From 814791f9f165b7e95f235036a70badea0971302a Mon Sep 17 00:00:00 2001 From: David Marby Date: Thu, 22 Sep 2022 10:48:43 +0200 Subject: [PATCH] Send and store cookies for websocket requests (#5205) * Send and store cookies for websocket requests * Lowercase for consistency --- .../insomnia-smoke-test/server/websocket.ts | 6 +++ .../insomnia/src/main/network/websocket.ts | 38 ++++++++++++++++++- .../insomnia/src/models/websocket-request.ts | 4 ++ packages/insomnia/src/network/network.ts | 2 +- .../ui/components/websockets/action-bar.tsx | 6 ++- .../websockets/websocket-response-pane.tsx | 5 +-- 6 files changed, 54 insertions(+), 7 deletions(-) diff --git a/packages/insomnia-smoke-test/server/websocket.ts b/packages/insomnia-smoke-test/server/websocket.ts index ff127439d..ec14e471d 100644 --- a/packages/insomnia-smoke-test/server/websocket.ts +++ b/packages/insomnia-smoke-test/server/websocket.ts @@ -13,6 +13,12 @@ export function startWebSocketServer(server: Server) { }); wsServer.on('connection', handleConnection); + + wsServer.on('headers', (headers, request) => { + if (request.url === '/cookies') { + headers.push('Set-Cookie: insomnia-websocket-test-cookie=foo'); + } + }); } const handleConnection = (ws: WebSocket, req: IncomingMessage) => { diff --git a/packages/insomnia/src/main/network/websocket.ts b/packages/insomnia/src/main/network/websocket.ts index 3cf3b2d9c..69723c4a3 100644 --- a/packages/insomnia/src/main/network/websocket.ts +++ b/packages/insomnia/src/main/network/websocket.ts @@ -1,6 +1,7 @@ import electron, { ipcMain } from 'electron'; import fs from 'fs'; import { IncomingMessage } from 'http'; +import { jarFromCookies } from 'insomnia-cookies'; import { setDefaultProtocol } from 'insomnia-url'; import mkdirp from 'mkdirp'; import path from 'path'; @@ -15,15 +16,17 @@ import { } from 'ws'; import { AUTH_BASIC, AUTH_BEARER } from '../../common/constants'; -import { generateId } from '../../common/misc'; +import { generateId, getSetCookieHeaders } from '../../common/misc'; import { webSocketRequest } from '../../models'; import * as models from '../../models'; +import { CookieJar } from '../../models/cookie-jar'; import { Environment } from '../../models/environment'; import { RequestAuthentication, RequestHeader } from '../../models/request'; import { BaseWebSocketRequest } from '../../models/websocket-request'; import type { WebSocketResponse } from '../../models/websocket-response'; import { getBasicAuthHeader } from '../../network/basic-auth/get-header'; import { getBearerAuthHeader } from '../../network/bearer-auth/get-header'; +import { addSetCookiesToToughCookieJar } from '../../network/network'; import { urlMatchesCertHost } from '../../network/url-matches-cert-host'; export interface WebSocketConnection extends WebSocket { @@ -97,6 +100,7 @@ const createWebSocketConnection = async ( url: string; headers: RequestHeader[]; authentication: RequestAuthentication; + cookieJar: CookieJar; } ): Promise => { const existingConnection = WebSocketConnections.get(options.requestId); @@ -172,6 +176,14 @@ const createWebSocketConnection = async ( } }); + if (request.settingSendCookies && options.cookieJar.cookies.length) { + const jar = jarFromCookies(options.cookieJar.cookies); + const cookieHeader = jar.getCookieStringSync(options.url); + if (cookieHeader) { + lowerCasedEnabledHeaders['cookie'] = cookieHeader; + } + } + const ws = new WebSocket(options.url, { headers: lowerCasedEnabledHeaders, cert: pemCertificates, @@ -186,7 +198,6 @@ const createWebSocketConnection = async ( // @ts-expect-error -- private property const internalRequestHeader = ws._req._header; const { timeline, responseHeaders, statusCode, statusMessage, httpVersion } = parseResponseAndBuildTimeline(options.url, incomingMessage, internalRequestHeader); - timeline.map(t => timelineFileStreams.get(options.requestId)?.write(JSON.stringify(t) + '\n')); const responsePatch: Partial = { _id: responseId, parentId: request._id, @@ -199,10 +210,30 @@ const createWebSocketConnection = async ( elapsedTime: performance.now() - start, timelinePath, eventLogPath: responseBodyPath, + settingSendCookies: request.settingSendCookies, + settingStoreCookies: request.settingStoreCookies, }; + const settings = await models.settings.getOrCreate(); models.webSocketResponse.create(responsePatch, settings.maxHistoryResponses); models.requestMeta.updateOrCreateByParentId(request._id, { activeResponseId: null }); + + if (request.settingStoreCookies) { + const setCookieStrings: string[] = getSetCookieHeaders(responseHeaders).map(h => h.value); + const totalSetCookies = setCookieStrings.length; + if (totalSetCookies) { + const currentUrl = request.url; + const { cookies, rejectedCookies } = await addSetCookiesToToughCookieJar({ setCookieStrings, currentUrl, cookieJar: options.cookieJar }); + rejectedCookies.forEach(errorMessage => timeline.push({ value: `Rejected cookie: ${errorMessage}`, name: 'Text', timestamp: Date.now() })); + const hasCookiesToPersist = totalSetCookies > rejectedCookies.length; + if (hasCookiesToPersist) { + await models.cookieJar.update(options.cookieJar, { cookies }); + timeline.push({ value: `Saved ${totalSetCookies} cookies`, name: 'Text', timestamp: Date.now() }); + } + } + } + + timeline.map(t => timelineFileStreams.get(options.requestId)?.write(JSON.stringify(t) + '\n')); }); ws.on('unexpected-response', async (clientRequest, incomingMessage) => { incomingMessage.on('data', chunk => { @@ -224,6 +255,8 @@ const createWebSocketConnection = async ( elapsedTime: performance.now() - start, timelinePath, eventLogPath: responseBodyPath, + settingSendCookies: request.settingSendCookies, + settingStoreCookies: request.settingStoreCookies, }; const settings = await models.settings.getOrCreate(); models.webSocketResponse.create(responsePatch, settings.maxHistoryResponses); @@ -402,6 +435,7 @@ export interface WebSocketBridgeAPI { url: string; headers: RequestHeader[]; authentication: RequestAuthentication; + cookieJar: CookieJar; }) => void; close: typeof closeWebSocketConnection; closeAll: typeof closeAllWebSocketConnections; diff --git a/packages/insomnia/src/models/websocket-request.ts b/packages/insomnia/src/models/websocket-request.ts index f72033864..3ba158250 100644 --- a/packages/insomnia/src/models/websocket-request.ts +++ b/packages/insomnia/src/models/websocket-request.ts @@ -20,6 +20,8 @@ export interface BaseWebSocketRequest { authentication: RequestAuthentication; parameters: RequestParameter[]; settingEncodeUrl: boolean; + settingStoreCookies: boolean; + settingSendCookies: boolean; } export type WebSocketRequest = BaseModel & BaseWebSocketRequest & { type: typeof type }; @@ -40,6 +42,8 @@ export const init = (): BaseWebSocketRequest => ({ authentication: {}, parameters: [], settingEncodeUrl: true, + settingStoreCookies: true, + settingSendCookies: true, }); export const migrate = (doc: WebSocketRequest) => doc; diff --git a/packages/insomnia/src/network/network.ts b/packages/insomnia/src/network/network.ts index 8e71c079d..ab1a45b92 100644 --- a/packages/insomnia/src/network/network.ts +++ b/packages/insomnia/src/network/network.ts @@ -203,7 +203,7 @@ export const getCurrentUrl = ({ headerResults, finalUrl }: { headerResults: any; } }; -const addSetCookiesToToughCookieJar = async ({ setCookieStrings, currentUrl, cookieJar }: any) => { +export const addSetCookiesToToughCookieJar = async ({ setCookieStrings, currentUrl, cookieJar }: any) => { const rejectedCookies: string[] = []; const jar = jarFromCookies(cookieJar.cookies); for (const setCookieStr of setCookieStrings) { diff --git a/packages/insomnia/src/ui/components/websockets/action-bar.tsx b/packages/insomnia/src/ui/components/websockets/action-bar.tsx index 928e3f33b..8004f7cc0 100644 --- a/packages/insomnia/src/ui/components/websockets/action-bar.tsx +++ b/packages/insomnia/src/ui/components/websockets/action-bar.tsx @@ -5,6 +5,7 @@ import styled from 'styled-components'; import { hotKeyRefs } from '../../../common/hotkeys'; import { executeHotKey } from '../../../common/hotkeys-listener'; import { getRenderContext, render, RENDER_PURPOSE_SEND } from '../../../common/render'; +import { cookieJar } from '../../../models'; import { WebSocketRequest } from '../../../models/websocket-request'; import { ReadyState } from '../../context/websocket-client/use-ws-ready-state'; import { OneLineEditor } from '../codemirror/one-line-editor'; @@ -81,12 +82,14 @@ export const WebSocketActionBar: FC = ({ request, workspaceId, e try { const renderContext = await getRenderContext({ request, environmentId, purpose: RENDER_PURPOSE_SEND }); const { url: rawUrl, headers, authentication, parameters } = request; - // Render any nunjucks tags in the url/headers/authentication settings + // Render any nunjucks tags in the url/headers/authentication settings/cookies + const workspaceCookieJar = await cookieJar.getOrCreateForParentId(workspaceId); const rendered = await render({ url: rawUrl, headers, authentication, parameters, + workspaceCookieJar, }, renderContext); const queryString = buildQueryStringFromParams(rendered.parameters); const url = joinUrlAndQueryString(rendered.url, queryString); @@ -96,6 +99,7 @@ export const WebSocketActionBar: FC = ({ request, workspaceId, e url, headers: rendered.headers, authentication: rendered.authentication, + cookieJar: rendered.workspaceCookieJar, }); } catch (err) { if (err.type === 'render') { diff --git a/packages/insomnia/src/ui/components/websockets/websocket-response-pane.tsx b/packages/insomnia/src/ui/components/websockets/websocket-response-pane.tsx index bf3d8c89f..18fd61b6e 100644 --- a/packages/insomnia/src/ui/components/websockets/websocket-response-pane.tsx +++ b/packages/insomnia/src/ui/components/websockets/websocket-response-pane.tsx @@ -180,9 +180,8 @@ const WebSocketActiveResponsePane: FC<{ requestId: string; response: WebSocketRe