Last active
March 17, 2024 13:31
-
-
Save digitaltembo/10d2df56ce1843ac468f8b6306149b0b to your computer and use it in GitHub Desktop.
Minimal WebSocket Server in NodeJS/TypeScript, no dependencies
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import { createHash } from "crypto"; | |
import { createServer, IncomingMessage } from "http"; | |
import { Duplex } from "stream"; | |
type WsFn = (val: string | Uint8Array) => void; | |
// Create this object whenever a websocket connection has been established | |
type WebSocket = { | |
write: WsFn; | |
addListener: (listener: WsFn) => void; | |
removeListener: (listener: WsFn) => void; | |
addCloseListener: (listener: () => void) => void; | |
close: () => void; | |
}; | |
// creates a server that simply echos what is sent to it | |
makeWsServer((ws: WebSocket) => { | |
ws.addListener((val: string | Uint8Array) => { | |
console.log("Got message!!!", val); | |
ws.write(val); | |
}); | |
}); | |
/** | |
* Data Frame Parsing | |
* | |
* Data passed in WebSockets takes place within a data frame containing | |
* - fin: a boolean indicating that the data frameis the end of a complete message | |
* - opCode: one of 5 reserved operations | |
* - hasMask: a boolean containing whether the payload should be masked by being | |
* XOR'd with a 4-byte repeating mask. The server does not need to transmit | |
* masked payloads, and the client is required to do so. | |
* - payloadLength: 7, 16, or 64-bit number indicating the length of the payload | |
* - mask: 4-byte number for masking payload, provided by client | |
* - payload: the good stuff | |
* | |
* This section deals with reading and writing the data frames | |
* | |
*/ | |
type DataFrame = { | |
fin: boolean; | |
opCode: WsOpCode; | |
hasMask?: boolean; | |
payload: Uint8Array; | |
}; | |
enum WsOpCode { | |
CONTINUATION = 0, | |
TEXT = 1, | |
BINARY = 2, | |
CONNECTION_CLOSE = 8, | |
PING = 9, | |
PONG = 10, | |
UNDEFINED = 0xff, | |
} | |
function findPayloadLen(data: Buffer, offset: number): [number, number] { | |
const first = data[offset] & 0x7f; | |
if (first === 126) { | |
return [data.readUInt16BE(offset + 1), 2]; | |
} else if (first === 127) { | |
const val = data.readBigUInt64BE(offset + 1); | |
if (val > Number.MAX_SAFE_INTEGER) { | |
throw new Error("Can't handle that kind of payload"); | |
} | |
return [Number(val), 4]; | |
} | |
return [first, 1]; | |
} | |
function parseDataFrame(data: Buffer) { | |
let offset = 0; | |
const fin = Boolean((data[offset] >> 7) & 1); | |
const opCodeInt = data[offset] & 0xf; | |
const opCode: WsOpCode | null = | |
opCodeInt in WsOpCode ? opCodeInt : WsOpCode.UNDEFINED; | |
offset++; | |
const hasMask = Boolean((data[offset] >> 7) & 1); | |
const [payloadLen, payloadWidth] = findPayloadLen(data, offset); | |
offset += payloadWidth; | |
const mask: Uint8Array | null = hasMask | |
? data.subarray(offset, offset + 4) | |
: null; | |
offset += hasMask ? 4 : 0; | |
const payload = | |
mask !== null | |
? Uint8Array.from( | |
data.subarray(offset, offset + payloadLen), | |
(maskedByte, index) => maskedByte ^ mask[index % 4] | |
) | |
: Uint8Array.from(data.subarray(offset, offset + payloadLen)); | |
return { fin, opCode, hasMask, payloadLen, mask, payload }; | |
} | |
function sendDataFrame(df: DataFrame, socket: Duplex) { | |
// payload length is encoded in 1, 2, or 4 byte values depending on size | |
const sizeBytes = | |
df.payload.length < 126 | |
? [df.payload.length] | |
: df.payload.length < 65536 | |
? [126, df.payload.length >> 8, df.payload.length & 0xff] | |
: [ | |
127, | |
df.payload.length >> 24, | |
(df.payload.length >> 16) & 0xff, | |
(df.payload.length >> 8) & 0xff, | |
]; | |
socket.write( | |
Uint8Array.from([ | |
(Number(df.fin) << 7) | df.opCode, | |
...sizeBytes, | |
...df.payload, | |
]) | |
); | |
} | |
// the PONG message is the PING message with a PONG opcode | |
function sendPong(ping: DataFrame, socket: Duplex) { | |
sendDataFrame({ ...ping, opCode: WsOpCode.PONG }, socket); | |
} | |
/** | |
* Upgrade Response | |
* | |
* The Http Server must respond to a Connection: Upgrade request with the following HTTPish headers | |
* in order to signal to the client that it can handle websocket connections | |
* | |
* Technically there are extensions and subprotocols that could also be referenced here | |
*/ | |
function upgradeConnection(req: IncomingMessage, socket: Duplex) { | |
const clientKey = req.headers["sec-websocket-key"]; | |
const reponseHeaders = [ | |
"HTTP/1.1 101 Switching Protocols", | |
"Upgrade: websocket", | |
`Sec-WebSocket-Accept: ${createHash("sha1") | |
.update(clientKey + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11") | |
.digest("base64")}`, | |
"Connection: Upgrade", | |
"\r\n", | |
].join("\r\n"); | |
socket.write(reponseHeaders); | |
} | |
const encoder = new TextEncoder(); | |
const decoder = new TextDecoder(); | |
function listenToSocket(socket: Duplex, listeners: WsFn[]) { | |
let strBuffer = ""; | |
let bufOffset = 0; | |
let buffer = new Uint8Array(2048); | |
socket.on("data", (data: Buffer) => { | |
const parsedFrame = parseDataFrame(data); | |
if (parsedFrame.opCode === WsOpCode.PING) { | |
sendPong(parsedFrame, socket); | |
return; | |
} | |
if (parsedFrame.opCode === WsOpCode.TEXT) { | |
strBuffer += decoder.decode(parsedFrame.payload); | |
if (parsedFrame.fin) { | |
listeners.forEach((listener) => listener(strBuffer)); | |
strBuffer = ""; | |
} | |
} else if (parsedFrame.opCode === WsOpCode.BINARY) { | |
// somewhat efficiently grow the underlying buffer | |
if (parsedFrame.payloadLen + bufOffset > buffer.length) { | |
let newSize = buffer.length * 2; | |
while (parsedFrame.payloadLen + bufOffset > newSize) { | |
newSize *= 2; | |
} | |
const newBuffer = new Uint8Array(newSize); | |
newBuffer.set(buffer); | |
buffer = newBuffer; | |
} | |
buffer.set(parsedFrame.payload, bufOffset); | |
bufOffset += parsedFrame.payloadLen; | |
if (parsedFrame.fin) { | |
listeners.forEach((listener) => listener(buffer)); | |
bufOffset = 0; | |
} | |
} else if (parsedFrame.opCode === WsOpCode.CONNECTION_CLOSE) { | |
socket.end(); | |
} | |
}); | |
} | |
export function makeWsServer(handleUpgrade: (handle: WebSocket) => void) { | |
const srv = createServer(); | |
srv.on("upgrade", (req: IncomingMessage, socket: Duplex, head: Buffer) => { | |
console.log("upgrading"); | |
upgradeConnection(req, socket); | |
const closeListeners: (() => void)[] = []; | |
const addCloseListener = (listener: () => void) => listeners.push(listener); | |
// note that if the socket is closed real quick none of the listeners will trigger | |
socket.on("close", () => closeListeners.forEach((closer) => closer())); | |
let listeners: WsFn[] = []; | |
const addListener = (listener: WsFn) => listeners.push(listener); | |
const removeListener = (removed: WsFn) => | |
(listeners = listeners.filter((listener) => listener !== removed)); | |
// write string/binary data via a single dataframe | |
const write = (val: string | Uint8Array) => { | |
// Don't bother fragmenting | |
if (typeof val === "string") { | |
const payload = encoder.encode(val); | |
sendDataFrame( | |
{ | |
fin: true, | |
opCode: WsOpCode.TEXT, | |
payload, | |
}, | |
socket | |
); | |
} else { | |
sendDataFrame( | |
{ | |
fin: true, | |
opCode: WsOpCode.BINARY, | |
payload: val, | |
}, | |
socket | |
); | |
} | |
}; | |
// a WebSocket connection closing down should be preceeded by the closing side | |
// first broadcasting one last CONNECTION_CLOSE dataframe | |
const close = () => { | |
sendDataFrame( | |
{ | |
fin: true, | |
opCode: WsOpCode.CONNECTION_CLOSE, | |
hasMask: false, | |
payload: Uint8Array.from([]), | |
}, | |
socket | |
); | |
}; | |
listenToSocket(socket, listeners); | |
handleUpgrade({ | |
write, | |
addListener, | |
removeListener, | |
close, | |
addCloseListener, | |
}); | |
}); | |
srv.listen(1337, "127.0.0.1", () => | |
console.log("listening on 127.0.0.1:1337") | |
); | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment