punchmole
Version:
A simple reverse proxy tunnel service which is easy deployable and can be easily customised and used with nodejs code
304 lines (294 loc) • 13.9 kB
JavaScript
import { EventEmitter } from 'node:events'
import { WebSocket } from 'ws'
import http from 'node:http'
const WEBSOCKET_HANDSHAKE_HEADERS = new Set([
'connection',
'upgrade',
'host',
'content-length',
'transfer-encoding',
'keep-alive',
'proxy-connection',
'proxy-authenticate',
'proxy-authorization',
'te',
'trailer',
'sec-websocket-key',
'sec-websocket-version',
'sec-websocket-extensions',
'sec-websocket-accept',
'sec-websocket-protocol',
])
function encodeWebsocketPayload(rawData) {
return Buffer.from(rawData).toString('base64')
}
function decodeWebsocketPayload(rawData) {
// Backward compatible with older servers that forwarded Buffer JSON / raw strings
if (typeof rawData === 'string') {
return Buffer.from(rawData, 'base64')
}
if (rawData && rawData.type === 'Buffer' && Array.isArray(rawData.data)) {
return Buffer.from(rawData.data)
}
return Buffer.from(rawData || '')
}
function buildUpstreamWebsocketOptions(headers = {}) {
const upstreamHeaders = {}
let protocols
for (const [key, value] of Object.entries(headers)) {
const lowerKey = key.toLowerCase()
if (lowerKey === 'sec-websocket-protocol') {
protocols = String(value)
.split(',')
.map((protocol) => protocol.trim())
.filter(Boolean)
continue
}
if (!WEBSOCKET_HANDSHAKE_HEADERS.has(lowerKey)) {
upstreamHeaders[key] = value
}
}
return { protocols, headers: upstreamHeaders }
}
export function PunchmoleClient(apiKey, domain, targetUrl, endpointUrl, log = console) {
const eventEmitter = new EventEmitter();
const ws = new WebSocket(endpointUrl);
ws.on('open', () => {
log.info(new Date(), 'connection with upstream server opened to forward public url', domain, 'to local url', targetUrl);
ws.send(JSON.stringify({'type': 'register', 'domain': domain, 'apiKey': apiKey}));
ws.isAlive = true;
})
const interval = setInterval(async () => {
try {
await ws.ping();
} catch(e) {
log.error(new Date(), 'error while sending ping to upstream server', e);
}
}, 10000)
ws.on('ping', () => {
log.debug('ping', ws.readyState);
// it can happen that we receive a ping but the connection breaks in between
try {
ws.pong()
} catch(e) {
// we ignore the error for ws.pong()
}
})
ws.on('pong', () => {
})
ws.on('close', () => {
log.info(new Date(), 'connection with upstream server closed')
clearInterval(interval)
eventEmitter.emit('close')
})
ws.on('error', (err) => {
log.error(new Date(), 'websocket error', err)
eventEmitter.emit('error', err)
})
const controllers = {}
const requests = {}
const websocketConnections = {}
// opening a new websocket may take some time and the remote server has already sent the first message
// for that purpose we push messages to a buffer (in order until the buffer is empty) which is sent then
// automatically once the websocket connection for one connection id (request id) is open,
let websocketMessageBuffer = []
// run through all buffered websocket messages in 30ms intervals and check if they're sendable
setInterval(() => {
for(let message of websocketMessageBuffer) {
if(
websocketConnections[message.request.id] &&
websocketConnections[message.request.id].readyState === WebSocket.OPEN
) {
log.debug('sending buffered message to finally opened websocket connection', message)
websocketConnections[message.request.id].send(
decodeWebsocketPayload(message.request.rawData),
{ binary: !!message.request.isBinary }
)
websocketMessageBuffer = websocketMessageBuffer.filter((v) => v.request.id !== message.request.id)
}
}
}, 30)
ws.on('message', async (rawMessage) => {
const request = JSON.parse(rawMessage)
let targetRequest = requests[request.id]
switch(request.type) {
case 'error':
log.error(new Date(), 'error received from upstream server', request)
eventEmitter.emit('error', request)
break
case 'registered':
log.info(new Date(), 'registration successful', request)
eventEmitter.emit('registered', request)
break
case 'request-end':
// if the foreign host stops the connection we need to stop streaming data (if we're streaming...)
log.debug(new Date(), 'request end received, stopping data if there is any', request.id)
// cleaning up status controllers
if(controllers[request.id]) {
controllers[request.id].controller.abort()
delete controllers[request.id].controller
delete controllers[request.id]
}
// cleaning up requests
if(requests[request.id]) {
requests[request.id].request.destroy()
delete requests[request.id]
}
eventEmitter.emit('request-end', request)
break
case 'request-start':
controllers[request.id] = {date: new Date(), controller: new AbortController()}
const url = targetUrl + request.url
log.debug(new Date(), 'received request', JSON.stringify(request))
requests[request.id] = {
responseStarted: false,
request: http.request(url, {
method: request.method,
headers: request.headers,
signal: controllers[request.id].controller.signal
}, (response) => {
requests[request.id].responseStarted = true
log.info(new Date(), 'request', request.headers['x-original-forwarded-for'], request.method, JSON.stringify(url), response.statusCode, JSON.stringify(response.statusMessage))
ws.send(JSON.stringify({
type: 'response-start',
id: request.id,
statusCode: response.statusCode,
statusMessage: response.statusMessage,
headers: response.headers
}))
response.on('data', data => {
ws.send(JSON.stringify({
type: 'data',
id: request.id,
data: data.toString('binary')
}))
})
response.on('end', async () => {
await ws.send(JSON.stringify({
type: 'data-end',
id: request.id
}))
})
controllers[request.id].controller.signal.addEventListener('abort', () => {
log.debug(new Date(), 'ending stream, remote client closed connection', request.id)
response.destroy()
}, { once: true })
})
}
requests[request.id].request.on('error', (err) => {
log.error(new Date(), 'request error', request.id, err);
if(!requests[request.id] || requests[request.id].responseStarted === false) {
ws.send(JSON.stringify({
type: 'response-start',
id: request.id,
statusCode: 503,
statusMessage: 'Service Unavailable',
headers: {}
}))
if(err.code === 'ECONNREFUSED') {
ws.send(JSON.stringify({
type: 'data',
id: request.id,
data: `remote service not available, check configuration/target url: ${err.code}`
}));
} else {
ws.send(JSON.stringify({
type: 'data',
id: request.id,
data: `${err.toString('binary')} (${err.code})`
}));
}
}
ws.send(JSON.stringify({
type: 'data-end',
id: request.id
}))
if(requests[request.id]) {
requests[request.id].request.destroy()
delete requests[request.id]
}
})
eventEmitter.emit('request', request)
break
case 'request-data':
if(targetRequest && targetRequest.request) {
log.debug(new Date(), 'request data has been received forwarding to remote service', request.id, Buffer.from(request.data, 'binary').length)
targetRequest.request.write(Buffer.from(request.data, 'binary'))
}
break
case 'request-data-end':
if(targetRequest && targetRequest.request) {
log.debug(new Date(), 'request has been finished, ending request to remote service', request.id)
targetRequest.request.end()
}
break
case 'websocket-connection':
log.info(new Date(), 'new websocket connection requested', request, targetUrl)
const websocketConnectionUrl = targetUrl.replace(/^http/, 'ws') + request.url
const upstream = buildUpstreamWebsocketOptions(request.headers)
log.debug(new Date(), '-> opening websocket connection...', websocketConnectionUrl, {
protocols: upstream.protocols,
headers: upstream.headers
})
// Pass subprotocols via the protocols argument (not headers). Forwarding
// sec-websocket-protocol as a raw header makes `ws` reject the handshake with
// "Server sent a subprotocol but none was requested" (breaks Vite/some HMR clients).
websocketConnections[request.id] = new WebSocket(
websocketConnectionUrl,
upstream.protocols,
{ headers: upstream.headers }
)
websocketConnections[request.id].on('message', (rawData, isBinary) => {
ws.send(JSON.stringify({
type: 'websocket-message',
id: request.id,
rawData: encodeWebsocketPayload(rawData),
isBinary
}))
})
websocketConnections[request.id].on('close', () => {
log.info(new Date(), 'service side closed websocket connection, sending forward and cleaning up...', request.id)
ws.send(JSON.stringify({
type: 'websocket-connection-closed',
id: request.id
}))
delete websocketConnections[request.id]
})
websocketConnections[request.id].on('error', (error) => {
log.warn(new Date(), 'websocket connected to service got an error', error)
})
break
case 'websocket-connection-closed':
if(websocketConnections[request.id]) {
log.debug(new Date(), 'remote client closed connection to websocket, closing websocket connection', request.id)
websocketConnections[request.id].close()
}
break
case 'websocket-message':
if(
websocketConnections[request.id] &&
websocketConnections[request.id].readyState === WebSocket.OPEN &&
// check if the websocket message buffer has messages for the given connection (request id) and if it
// has enqueue new messages rather then sending it directly to ensure the order is given
// (some protocols require separate authentication and/or handshake procedures which keep account for here)
!websocketMessageBuffer.find((v) => v.request.id === request.id)
) {
log.debug(new Date(), 'websocket message received and ready to send forward', request.id)
websocketConnections[request.id].send(
decodeWebsocketPayload(request.rawData),
{ binary: !!request.isBinary }
)
} else {
log.debug(new Date(), 'websocket message received but not ready to send forward yet, buffering...', request.id)
websocketMessageBuffer.push({
date: new Date(),
request: request,
})
}
}
})
eventEmitter.close = function close() {
ws.close()
}
return eventEmitter
}