Add a 'Forward to host' websocket rule

Author: pimterry

src/components/mock/handler-config.tsx

@@ -50,7 +50,8 @@ import {
     WebSocketPassThroughHandler,
     EchoWebSocketHandlerDefinition,
     RejectWebSocketHandlerDefinition,
-    ListenWebSocketHandlerDefinition
+    ListenWebSocketHandlerDefinition,
+    WebSocketForwardToHostHandler
 } from '../../model/rules/definitions/websocket-rule-definitions';
 import {
     EthereumCallResultHandler,
@@ -151,7 +152,11 @@ export function HandlerConfiguration(props: {
         case 'file':
             return <FromFileResponseHandlerConfig {...configProps} />;
         case 'forward-to-host':
-            return <ForwardToHostHandlerConfig {...configProps} />;
+        case 'ws-forward-to-host':
+            return <ForwardToHostHandlerConfig
+                {...configProps}
+                handlerKey={handlerKey}
+            />;
         case 'passthrough':
         case 'ws-passthrough':
             return <PassThroughHandlerConfig {...configProps} />;
@@ -657,9 +662,14 @@ const UrlInput = styled(TextInput)`
 
 @inject('rulesStore')
 @observer
-class ForwardToHostHandlerConfig extends HandlerConfig<ForwardToHostHandler, {
-    rulesStore?: RulesStore
-}> {
+class ForwardToHostHandlerConfig extends HandlerConfig<
+    | ForwardToHostHandler
+    | WebSocketForwardToHostHandler,
+    {
+        rulesStore?: RulesStore,
+        handlerKey: 'forward-to-host' | 'ws-forward-to-host'
+    }
+> {
 
     @observable
     private error: Error | undefined;
@@ -682,9 +692,19 @@ class ForwardToHostHandlerConfig extends HandlerConfig<ForwardToHostHandler, {
     }
 
     render() {
-        const { targetHost, updateHostHeader, error, onTargetChange, onUpdateHeaderChange } = this;
+        const {
+            targetHost,
+            updateHostHeader,
+            error,
+            onTargetChange,
+            onUpdateHeaderChange
+        } = this;
         const { targetHost: savedTargetHost } = this.props.handler.forwarding!;
 
+        const messageType = this.props.handlerKey === 'ws-forward-to-host'
+            ? 'WebSocket'
+            : 'request';
+
         return <ConfigContainer>
             <SectionLabel>Replacement host</SectionLabel>
             <UrlInput
@@ -699,7 +719,7 @@ class ForwardToHostHandlerConfig extends HandlerConfig<ForwardToHostHandler, {
                 value={updateHostHeader.toString()}
                 onChange={onUpdateHeaderChange}
                 title={dedent`
-                    Most servers will not accept requests that arrive
+                    Most servers will not accept ${messageType}s that arrive
                     with the wrong host header, so it's typically useful
                     to automatically change it to match the new host
                 `}
@@ -709,7 +729,7 @@ class ForwardToHostHandlerConfig extends HandlerConfig<ForwardToHostHandler, {
             </ConfigSelect>
             { savedTargetHost &&
                 <ConfigExplanation>
-                    All matching requests will be forwarded to {savedTargetHost},
+                    All matching {messageType}s will be forwarded to {savedTargetHost},
                     keeping their existing path{
                         !savedTargetHost.includes('://') ? ', protocol,' : ''
                     } and query string.{
@@ -726,26 +746,47 @@ class ForwardToHostHandlerConfig extends HandlerConfig<ForwardToHostHandler, {
         try {
             if (!this.targetHost) throw new Error('A target host is required');
 
-            const protocolMatch = this.targetHost.match(/^\w+:\/\//);
-            if (protocolMatch) {
-                const pathWithoutProtocol = this.targetHost.slice(protocolMatch[0].length);
+            let urlWithoutProtocol: string;
 
-                if (pathWithoutProtocol.includes('/')) {
-                    throw new Error('The replacement host shouldn\'t include a path, since it won\'t be used');
-                }
-                if (pathWithoutProtocol.includes('?')) {
-                    throw new Error('The replacement host shouldn\'t include a query string, since it won\'t be used');
+            const protocolMatch = this.targetHost.match(/^(\w+):\/\//);
+            if (protocolMatch) {
+                const validProtocols = this.props.handlerKey === 'ws-forward-to-host'
+                    ? ['ws', 'wss']
+                    : ['http', 'https'];
+
+                if (!validProtocols.includes(protocolMatch[1].toLowerCase())) {
+                    throw new Error(
+                        `The protocol must be either ${validProtocols[0]} or ${validProtocols[1]}`
+                    );
                 }
+
+                urlWithoutProtocol = this.targetHost.slice(protocolMatch[0].length);
             } else {
-                if (this.targetHost.includes('/')) {
-                    throw new Error('The replacement host shouldn\'t include a path, since it won\'t be used');
-                }
-                if (this.targetHost.includes('?')) {
-                    throw new Error('The replacement host shouldn\'t include a query string, since it won\'t be used');
-                }
+                urlWithoutProtocol = this.targetHost;
+            }
+
+            if (urlWithoutProtocol.includes('/')) {
+                throw new Error(
+                    'The replacement host shouldn\'t include a path, since it won\'t be used'
+                );
+            }
+            if (urlWithoutProtocol.includes('?')) {
+                throw new Error(
+                    'The replacement host shouldn\'t include a query string, since it won\'t be used'
+                );
             }
 
-            this.props.onChange(new ForwardToHostHandler(this.targetHost, this.updateHostHeader, this.props.rulesStore!));
+            const HandlerClass = this.props.handlerKey === 'ws-forward-to-host'
+                ? WebSocketForwardToHostHandler
+                : ForwardToHostHandler;
+
+            this.props.onChange(
+                new HandlerClass(
+                    this.targetHost,
+                    this.updateHostHeader,
+                    this.props.rulesStore!
+                )
+            );
             this.error = undefined;
         } catch (e) {
             console.log(e);

src/components/mock/handler-selection.tsx

@@ -33,7 +33,8 @@ import {
     WebSocketPassThroughHandler,
     EchoWebSocketHandlerDefinition,
     RejectWebSocketHandlerDefinition,
-    ListenWebSocketHandlerDefinition
+    ListenWebSocketHandlerDefinition,
+    WebSocketForwardToHostHandler
 } from '../../model/rules/definitions/websocket-rule-definitions';
 import {
     EthereumCallResultHandler,
@@ -97,8 +98,6 @@ const instantiateHandler = (
             return new FromFileResponseHandler(200, undefined, '');
         case 'passthrough':
             return new PassThroughHandler(rulesStore);
-        case 'ws-passthrough':
-            return new WebSocketPassThroughHandler(rulesStore);
         case 'forward-to-host':
             return new ForwardToHostHandler('', true, rulesStore);
         case 'req-res-transformer':
@@ -116,6 +115,10 @@ const instantiateHandler = (
         case 'reset-connection':
             return new ResetConnectionHandler();
 
+        case 'ws-passthrough':
+            return new WebSocketPassThroughHandler(rulesStore);
+        case 'ws-forward-to-host':
+            return new WebSocketForwardToHostHandler('', true, rulesStore);
         case 'ws-echo':
             return new EchoWebSocketHandlerDefinition();
         case 'ws-reject':

src/model/rules/definitions/websocket-rule-definitions.ts

@@ -11,6 +11,7 @@ import * as serializr from 'serializr';
 import { RulesStore } from '../rules-store';
 
 import { MethodNames } from '../../http/methods';
+import { serializeAsTag } from '../../serialization';
 import {
     HttpMatcherLookup,
     HttpMatcher,
@@ -66,6 +67,35 @@ serializr.createModelSchema(WebSocketPassThroughHandler, {
     type: serializr.primitive()
 }, (context) => new WebSocketPassThroughHandler(context.args.rulesStore));
 
+export class WebSocketForwardToHostHandler extends wsHandlers.PassThroughWebSocketHandlerDefinition {
+
+    readonly uiType = 'ws-forward-to-host';
+
+    constructor(forwardToLocation: string, updateHostHeader: boolean, rulesStore: RulesStore) {
+        super({
+            ...rulesStore.activePassthroughOptions,
+            forwarding: {
+                targetHost: forwardToLocation,
+                updateHostHeader: updateHostHeader
+            }
+        });
+    }
+
+}
+
+serializr.createModelSchema(WebSocketForwardToHostHandler, {
+    uiType: serializeAsTag(() => 'ws-forward-to-host'),
+    type: serializr.primitive(),
+    forwarding: serializr.map(serializr.primitive())
+}, (context) => {
+    const data = context.json;
+    return new WebSocketForwardToHostHandler(
+        data.forwarding.targetHost,
+        data.forwarding.updateHostHeader,
+        context.args.rulesStore
+    );
+});
+
 export const WebSocketMatcherLookup = {
     ..._.omit(HttpMatcherLookup, MethodNames),
     'method': WebSocketMethodMatcher, // Unlike HTTP rules, WS uses a single method matcher
@@ -81,7 +111,8 @@ export const WebSocketInitialMatcherClasses = [
 
 export const WebSocketHandlerLookup = {
     ...wsHandlers.WsHandlerDefinitionLookup,
-    'ws-passthrough': WebSocketPassThroughHandler
+    'ws-passthrough': WebSocketPassThroughHandler,
+    'ws-forward-to-host': WebSocketForwardToHostHandler
 };
 
 type WebSocketHandlerClass = typeof WebSocketHandlerLookup[keyof typeof WebSocketHandlerLookup];

src/model/rules/rule-descriptions.ts

@@ -109,6 +109,7 @@ export function nameHandlerClass(key: HandlerClassKey): string {
         case 'file':
             return "file response";
         case 'forward-to-host':
+        case 'ws-forward-to-host':
             return "forwarding";
         case 'passthrough':
         case 'ws-passthrough':
@@ -184,8 +185,6 @@ export function summarizeHandlerClass(key: HandlerClassKey): string {
             return "Forward the request to a different host";
         case 'passthrough':
             return "Pass the request on to its destination";
-        case 'ws-passthrough':
-            return "Pass the WebSocket through to its destination";
         case 'req-res-transformer':
             return "Transform the real request or response automatically";
         case 'request-breakpoint':
@@ -201,6 +200,10 @@ export function summarizeHandlerClass(key: HandlerClassKey): string {
         case 'reset-connection':
             return "Forcibly reset the connection";
 
+        case 'ws-passthrough':
+            return "Pass the WebSocket through to its destination";
+        case 'ws-forward-to-host':
+            return "Forward the WebSocket to a different host";
         case 'ws-reject':
             return "Reject the WebSocket setup request";
         case 'ws-listen':