Skip to content

Commit

Permalink
Add a 'Forward to host' websocket rule
Browse files Browse the repository at this point in the history
  • Loading branch information
@pimterry
pimterry committed Aug 11, 2023
1 parent 702c95b commit bd9d252
Show file tree
Hide file tree
Showing 4 changed files with 107 additions and 29 deletions.
87 changes: 64 additions & 23 deletions src/components/mock/handler-config.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,8 @@ import {
WebSocketPassThroughHandler,
EchoWebSocketHandlerDefinition,
RejectWebSocketHandlerDefinition,
ListenWebSocketHandlerDefinition
ListenWebSocketHandlerDefinition,
WebSocketForwardToHostHandler
} from '../../model/rules/definitions/websocket-rule-definitions';
import {
EthereumCallResultHandler,
Expand Down Expand Up @@ -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} />;
Expand Down Expand Up @@ -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;
Expand All @@ -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
Expand All @@ -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
`}
Expand All @@ -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.{
Expand All @@ -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);
Expand Down
9 changes: 6 additions & 3 deletions src/components/mock/handler-selection.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,8 @@ import {
WebSocketPassThroughHandler,
EchoWebSocketHandlerDefinition,
RejectWebSocketHandlerDefinition,
ListenWebSocketHandlerDefinition
ListenWebSocketHandlerDefinition,
WebSocketForwardToHostHandler
} from '../../model/rules/definitions/websocket-rule-definitions';
import {
EthereumCallResultHandler,
Expand Down Expand Up @@ -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':
Expand All @@ -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':
Expand Down
33 changes: 32 additions & 1 deletion src/model/rules/definitions/websocket-rule-definitions.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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];
Expand Down
7 changes: 5 additions & 2 deletions src/model/rules/rule-descriptions.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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':
Expand Down Expand Up @@ -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':
Expand All @@ -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':
Expand Down

0 comments on commit bd9d252

Please sign in to comment.