Add WebSocket transform perspective support

Author: pimterry

src/components/view/http/http-details-pane.tsx

@@ -227,7 +227,7 @@ export class HttpDetailsPane extends React.Component<{
         } else if (
             expandedCard === 'webSocketMessages' &&
             exchange.isWebSocket() &&
-            exchange.wasAccepted()
+            exchange.wasAccepted
         ) {
             return this.renderWebSocketMessages(exchange);
         } else {
@@ -363,7 +363,7 @@ export class HttpDetailsPane extends React.Component<{
         }
 
         if (exchange.isWebSocket()) {
-            if (exchange.wasAccepted()) {
+            if (exchange.wasAccepted) {
                 cards.push(this.renderWebSocketMessages(exchange));
 
                 if (exchange.closeState) {

src/components/view/http/http-request-card.tsx

@@ -1,7 +1,7 @@
 import * as React from 'react';
 import { observer } from 'mobx-react';
 
-import { HttpExchange, HtkRequest, HttpVersion, HttpExchangeView } from '../../../types';
+import { HtkRequest, HttpVersion, HttpExchangeView } from '../../../types';
 
 import { getSummaryColor } from '../../../model/events/categorization';
 import { getMethodDocs } from '../../../model/http/http-docs';
@@ -29,7 +29,7 @@ import { HttpVersionPill } from '../../common/http-version-pill';
 import { HeaderDetails } from './header-details';
 import { UrlBreakdown } from '../url-breakdown';
 import { HandlerClassKey } from '../../../model/rules/rules';
-import { MatchedRulePill } from './matched-rule-pill';
+import { MatchedRulePill, shouldShowRuleDetails } from './matched-rule-pill';
 
 const RawRequestDetails = (p: {
     request: HtkRequest,
@@ -104,14 +104,9 @@ export const HttpRequestCard = observer((props: HttpRequestCardProps) => {
     const { exchange, matchedRuleData, onRuleClicked } = props;
     const { request } = exchange;
 
-    // We consider passthrough as a no-op, and so don't show anything in that case.
-    const noopRule = matchedRuleData?.stepTypes.every(
-        type => type === 'passthrough' || type === 'ws-passthrough'
-    )
-
     return <CollapsibleCard {...props} direction='right'>
         <header>
-            { matchedRuleData && !noopRule &&
+            { shouldShowRuleDetails(matchedRuleData) &&
                 <MatchedRulePill
                     ruleData={matchedRuleData}
                     onClick={onRuleClicked}

src/components/view/http/matched-rule-pill.tsx

@@ -12,13 +12,25 @@ import { getSummaryColor } from '../../../model/events/categorization';
 import { aOrAn, uppercaseFirst } from '../../../util/text';
 import { PillButton } from '../../common/pill';
 
+export interface MatchedRuleData {
+    stepTypes: HandlerClassKey[];
+    status: 'unchanged' | 'modified-types' | 'deleted';
+}
+
+export const shouldShowRuleDetails = (
+    matchedRuleData: MatchedRuleData | undefined
+): matchedRuleData is MatchedRuleData => {
+    // We never bother showing rule details for pure-passthrough rules
+    return !!matchedRuleData?.stepTypes.length &&
+        !matchedRuleData?.stepTypes.every(
+            type => type === 'passthrough' || type === 'ws-passthrough'
+        );
+}
+
 export const MatchedRulePill = styled(inject('uiStore')((p: {
     className?: string,
     uiStore?: UiStore,
-    ruleData: {
-        stepTypes: HandlerClassKey[],
-        status: 'unchanged' | 'modified-types' | 'deleted'
-    },
+    ruleData: MatchedRuleData,
     onClick: () => void
 }) => {
     const { stepTypes } = p.ruleData;

src/components/view/http/transform-card.tsx

@@ -5,11 +5,10 @@ import { observer } from 'mobx-react-lite';
 import { styled } from '../../../styles';
 
 import { ContentPerspective, UiStore } from '../../../model/ui/ui-store';
-import { HandlerClassKey } from '../../../model/rules/rules';
 
 import { PillSelect } from '../../common/pill';
 import { MediumCard } from '../../common/card';
-import { MatchedRulePill } from './matched-rule-pill';
+import { MatchedRuleData, MatchedRulePill, shouldShowRuleDetails } from './matched-rule-pill';
 
 const DropdownContainer = styled.div`
     display: inline-block;
@@ -41,20 +40,11 @@ const PerspectiveSelector = observer((p: {
 });
 
 export const TransformCard = (p: {
-    matchedRuleData?: {
-        stepTypes: HandlerClassKey[],
-        status: 'unchanged' | 'modified-types' | 'deleted'
-    } | undefined,
+    matchedRuleData?: MatchedRuleData | undefined,
     onRuleClicked: () => void,
     uiStore: UiStore
 }) => {
-    // We consider passthrough as a no-op, and so don't show anything in that case.
-    const noopRule = p.matchedRuleData?.stepTypes.every(
-        type => type === 'passthrough' || type === 'ws-passthrough'
-    );
-
-    // // Show nothing when there's no rule data available or the traffic was no-op'd
-    if (noopRule || !p.matchedRuleData) return null;
+    if (!shouldShowRuleDetails(p.matchedRuleData)) return null;
 
     return <MediumCard>
         <MatchedRulePill

src/components/view/view-page.tsx

@@ -294,7 +294,7 @@ class ViewPage extends React.Component<ViewPageProps> {
                 ) ||
                 (
                     expandedViewCard === 'webSocketMessages' &&
-                    !(selectedHttpExchange.isWebSocket() && selectedHttpExchange.wasAccepted())
+                    !(selectedHttpExchange.isWebSocket() && selectedHttpExchange.wasAccepted)
                 )
             ) {
                 runInAction(() => {

src/model/events/events-store.ts

@@ -520,6 +520,12 @@ export class EventsStore {
             case 'passthrough-abort':
                 exchange.updateFromUpstreamAbort(event.eventData);
                 break;
+
+            case 'passthrough-websocket-connect':
+                if (!exchange.isWebSocket()) throw new Error('Received WS connect event for non-WS');
+                const webSocket = exchange as WebSocketStream;
+                webSocket.updateWithUpstreamConnect(event.eventData);
+                break;
         }
     }
 

src/model/http/http-exchange.ts

@@ -117,7 +117,7 @@ export class HttpExchange extends HTKEventBase implements HttpExchangeView {
 
     constructor(
         request: InputRequest,
-        private readonly apiStore: ApiStore
+        protected readonly apiStore: ApiStore
     ) {
         super();
 

src/model/websockets/upstream-websocket.ts

@@ -0,0 +1,46 @@
+import { ApiStore } from "../api/api-store";
+import { StreamMessage } from "../events/stream-message";
+
+import { UpstreamHttpExchange } from "../http/upstream-exchange";
+import { WebSocketStream } from "./websocket-stream";
+import { WebSocketView } from "./websocket-views";
+
+/**
+ * This represents the upstream side of a proxied WebSocket connection. In the websocket
+ * case, at the time of writing, the only modifications made during proxying are redirection
+ * so this really just proxies through to the downstream side except for the initial
+ * upstream connection parameters.
+ *
+ * By and large this is a minimal PoC & structure for future development - there's very little
+ * real usage of this at the moment until we have more WebSocket transformations available.
+ */
+export class UpstreamWebSocket extends UpstreamHttpExchange implements WebSocketView {
+
+    constructor(downstream: WebSocketStream, apiStore: ApiStore) {
+        super(downstream, apiStore);
+    }
+
+    isWebSocket() {
+        return true;
+    }
+
+    declare public readonly downstream: WebSocketStream;
+
+    get upstream(): UpstreamWebSocket {
+        return this;
+    }
+
+    get original(): WebSocketView {
+        return this.downstream.original;
+    }
+
+    get transformed(): WebSocketView {
+        return this.downstream.transformed;
+    }
+
+    get wasAccepted() { return this.downstream.wasAccepted; }
+    get selectedSubprotocol() { return this.downstream.selectedSubprotocol; }
+    get messages(): readonly StreamMessage[] { return this.downstream.messages; }
+    get closeState() { return this.downstream.closeState; }
+
+}

src/model/websockets/websocket-stream.ts

@@ -1,22 +1,27 @@
 import * as _ from 'lodash';
-import { observable, action } from 'mobx';
+import { observable, action, computed } from 'mobx';
 
 import {
     InputRequest,
     InputResponse,
     InputWebSocketMessage,
     InputWebSocketClose,
-    InputFailedRequest
+    InputFailedRequest,
+    InputRuleEventDataMap
 } from '../../types';
 
 import { ApiStore } from '../api/api-store';
 import { StreamMessage } from '../events/stream-message';
 import { HttpExchange } from '../http/http-exchange';
 
+import { WebSocketOriginalView, WebSocketTransformedView, WebSocketView } from './websocket-views';
+import { UpstreamWebSocket } from './upstream-websocket';
+
 // A websocket stream is an HTTP exchange (the initial setup, or even rejection), but
 // may include a series of many ongoing messages and a final websocket close event,
 // if the initial websocket connection is successful.
-export class WebSocketStream extends HttpExchange {
+export class WebSocketStream extends HttpExchange implements WebSocketView {
+
     constructor(request: InputRequest, apiStore: ApiStore) {
         super(request, apiStore);
         this.searchIndex += '\nwebsocket';
@@ -26,11 +31,43 @@ export class WebSocketStream extends HttpExchange {
         return true;
     }
 
+    declare public upstream: UpstreamWebSocket | undefined;
+
+    // These are the same as HttpExchangeViewBase, but need to be copied here (because we're not a _view_,
+    // we're original, and TS has no proper mixin support).
+    @computed
+    get original(): WebSocketView {
+        if (!this.upstream) return this;
+
+        // If the request is original, then upstream data matches original data
+        // I.e. only possible transform was after all upstream data
+        if (!this.upstream.wasRequestTransformed) {
+            return this.upstream;
+        } else {
+            return new WebSocketOriginalView(this.downstream, this.apiStore);
+        }
+    }
+
+    @computed
+    get transformed(): WebSocketView {
+        if (!this.upstream) return this;
+
+        // If the response is original, then upstream data matches transformed data
+        // I.e. all transforms happened before any upstream data
+        if (!this.upstream?.wasResponseTransformed) {
+            return this.upstream;
+        } else {
+            return new WebSocketTransformedView(this.downstream, this.apiStore);
+        }
+    }
+
     @observable
     private accepted = false;
+    get wasAccepted() { return this.accepted; }
 
     @observable
     private subprotocol: string | undefined;
+    get selectedSubprotocol() { return this.subprotocol; }
 
     @action
     setAccepted(response: InputResponse) {
@@ -41,14 +78,6 @@ export class WebSocketStream extends HttpExchange {
         Object.assign(this.timingEvents, response.timingEvents);
     }
 
-    wasAccepted() {
-        return this.accepted;
-    }
-
-    get selectedSubprotocol() {
-        return this.subprotocol;
-    }
-
     @observable
     readonly messages: Array<StreamMessage> = [];
 
@@ -73,7 +102,7 @@ export class WebSocketStream extends HttpExchange {
     }
 
     markAborted(request: InputFailedRequest) {
-        if (!this.wasAccepted()) {
+        if (!this.wasAccepted) {
             // An abort before accept acts exactly as in normal HTTP
             return super.markAborted(request);
         } else {
@@ -88,6 +117,16 @@ export class WebSocketStream extends HttpExchange {
         }
     }
 
+    // The assorted normal upstreamFromUpstream... methods will never be called, as the server side
+    // is completely different, so UpstreamHttpExchange will never be populated by the base class.
+
+    updateWithUpstreamConnect(params: InputRuleEventDataMap['passthrough-websocket-connect']) {
+        if (!this.upstream) {
+            this.upstream = new UpstreamWebSocket(this, this.apiStore);
+        }
+        this.upstream.updateWithRequestHead(params);
+    }
+
     cleanup() {
         super.cleanup();
 

src/model/websockets/websocket-views.ts

@@ -0,0 +1,77 @@
+import { InputWebSocketClose, WebSocketStream } from '../../types';
+import { ApiStore } from '../api/api-store';
+
+import { StreamMessage } from '../events/stream-message';
+import {
+    HttpExchangeView,
+    HttpExchangeOriginalView,
+    HttpExchangeTransformedView
+} from '../http/http-exchange-views';
+import { UpstreamWebSocket } from './upstream-websocket';
+
+export interface WebSocketView extends HttpExchangeView {
+
+    get downstream(): WebSocketStream;
+    get upstream(): UpstreamWebSocket | undefined;
+
+    get original(): WebSocketView;
+    get transformed(): WebSocketView;
+
+    get wasAccepted(): boolean;
+    get selectedSubprotocol(): string | undefined;
+    get messages(): ReadonlyArray<StreamMessage>;
+    get closeState(): InputWebSocketClose | 'aborted' | undefined;
+
+}
+
+export class WebSocketOriginalView extends HttpExchangeOriginalView implements WebSocketView {
+
+    constructor(downstreamWebSocket: WebSocketStream, apiStore: ApiStore) {
+        super(downstreamWebSocket, apiStore);
+    }
+
+    declare public readonly downstream: WebSocketStream;
+
+    isWebSocket() {
+        return true;
+    }
+
+    get upstream(): UpstreamWebSocket | undefined {
+        return this.downstream.upstream;
+    }
+
+    get original(): WebSocketView { return this; }
+    get transformed(): WebSocketView { return this.downstream.transformed; }
+
+    get wasAccepted() { return this.downstream.wasAccepted; }
+    get selectedSubprotocol() { return this.downstream.selectedSubprotocol; }
+    get messages() { return this.downstream.messages; }
+    get closeState() { return this.downstream.closeState; }
+
+}
+
+export class WebSocketTransformedView extends HttpExchangeTransformedView implements WebSocketView {
+
+    constructor(exchange: WebSocketStream, apiStore: ApiStore) {
+        super(exchange, apiStore);
+    }
+
+    declare public readonly downstream: WebSocketStream;
+
+    isWebSocket() {
+        return true;
+    }
+
+    get upstream(): UpstreamWebSocket | undefined {
+        return this.downstream.upstream;
+    }
+
+    get original(): WebSocketView { return this.downstream.original; }
+    get transformed(): WebSocketView { return this; }
+
+    get wasAccepted() { return this.downstream.wasAccepted; }
+    get selectedSubprotocol() { return this.downstream.selectedSubprotocol; }
+    get messages() { return this.downstream.messages; }
+    get closeState() { return this.downstream.closeState; }
+
+}