tcp-ws-proxy/index.ts

import { WebSocket, WebSocketServer } from "ws";
import { Socket, createServer } from "net";
import { Packet } from "./src/enum/Packet";
import { Endian, createReader, createWriter } from "bufferstuff";
import { existsSync, readFileSync } from "fs";
import IServerConfig from "./src/interface/IServerConfig";
import { AuthState } from "./src/enum/AuthState";
import IClientConfig from "./src/interface/IClientConfig";

const keepAliveMessage = createWriter(Endian.LE, 1).writeUByte(Packet.KeepAlive).toBuffer();

if (process.argv[2] === "server") {
	if (!existsSync("./config/server-config.json")) {
		console.error("server-config.json is missing!");
		process.exit(1);
	}
	const config:IServerConfig = JSON.parse(readFileSync("./config/server-config.json").toString());
	const server = new WebSocketServer({ port: config.port }, () => console.log(`Server started at ${config.port}`));

	server.on("connection", (socket) => {
		console.log("Connection");
		let queuedMessages = new Array<Buffer>();
		let connectedToServer = false;
		let connectingToServer = false;
		let authed = false;
		let client:Socket | null = null;
		const clientKeepAlive = setInterval(() => {
			socket.send(keepAliveMessage);
		}, 5000);

		socket.on("message", (data, isBinary) => {
			if (!isBinary) {
				return;
			}

			// NOTE: The types declarations for ws are really messed up >:(
			// @ts-ignore
			const packetData = createReader(Endian.LE, data);
			const packetId = packetData.readUByte();
			let tempReadLength = 0;
			
			switch (packetId) {
				case Packet.KeepAlive: break; // We don't really care, it's just so CF doesn't drop the connection.
				case Packet.Auth:
					try {
						if (packetData.readShortString() !== config.authKey) {
							socket.send(createWriter(Endian.LE, 2).writeUByte(Packet.Auth).writeUByte(AuthState.Bad).toBuffer());
							socket.close();
							return;
						}
						socket.send(createWriter(Endian.LE, 2).writeUByte(Packet.Auth).writeUByte(AuthState.Good).toBuffer());
						authed = true;
	
						client = new Socket();
	
						client.on("connect", () => {
							connectedToServer = true;
							if (queuedMessages.length > 0) {
								for (const message of queuedMessages) {
									console.log("Sent", message);
									client?.write(message);
								}
								queuedMessages.length = 0;
							}
						});
	
						client.on("data", (chunk) => {
							socket.send(createWriter(Endian.LE, 5)
								.writeUByte(Packet.Data)
								.writeUInt(chunk.length)
								.writeBuffer(chunk)
								.toBuffer());
						});

						function clientCloseOrError() {
							socket.close();
						}
						client.on("close", clientCloseOrError);
						client.on("error", clientCloseOrError);
	
						connectingToServer = true;
						client.connect({
							host: config.localHost,
							port: config.localPort
						});
					} catch (e) {
						client?.end();
						client = null;
						socket.close();
					}
					break;
				case Packet.Data:
					if (!authed) {
						return;
					}

					tempReadLength = packetData.readUInt();
					if (connectedToServer) {
						client?.write(packetData.readBuffer(tempReadLength));
					} else {
						queuedMessages.push(packetData.readBuffer(tempReadLength));
					}
					//console.log("[SERVER] Data:", data, " Length:", tempReadLength);
					break;
			}
		});

		function closeOrError() {
			clearInterval(clientKeepAlive);
			client?.end();
		}
		socket.on("close", closeOrError);
		socket.on("error", closeOrError);
	});
} else if (process.argv[2] === "client") {
	if (!existsSync("./config/client-config.json")) {
		console.error("client-config.json is missing!");
		process.exit(1);
	}
	const config:IClientConfig = JSON.parse(readFileSync("./config/client-config.json").toString());
	const server = createServer((socket) => {
		let authed = false;
		let queuedMessages = new Array<Buffer>();
		let txBytes = 0;
		let rxBytes = 0;
		const txrxInterval = setInterval(() => {
			console.log(`TX: ${(txBytes / 1024).toFixed(2)}KB/s  |  RX: ${(rxBytes / 1024).toFixed(2)}KB/s`);
			txBytes = rxBytes = 0;
		}, 1000);

		const client = new WebSocket(config.remoteAddress);
		client.on("open", () => {
			// Send off auth as soon as we connect
			client.send(createWriter(Endian.LE, 2 + config.authKey.length).writeUByte(Packet.Auth).writeShortString(config.authKey).toBuffer());
		});

		client.on("message", (data) => {
			// @ts-ignore
			const packetData = createReader(Endian.LE, data);
			const packetId = packetData.readUByte();
			let tempReadLength = 0;
			let bufferData:Buffer;

			switch (packetId) {
				case Packet.KeepAlive:
					client.send(keepAliveMessage);
					break;

				case Packet.Auth:
					if (authed) {
						return;
					}

					authed = packetData.readUByte() === AuthState.Good;
					if (authed && queuedMessages.length > 0) {
						for (const message of queuedMessages) {
							client.send(message);
						}
						queuedMessages.length = 0;
					}
					break;
				case Packet.Data:
					tempReadLength = packetData.readUInt();
					bufferData = packetData.readBuffer(tempReadLength);
					rxBytes += bufferData.length;
					socket.write(bufferData);
					break;
			}
		});

		function clientCloseOrError() {
			socket.end();
			clearInterval(txrxInterval);
		}
		client.on("close", clientCloseOrError);
		client.on("error", clientCloseOrError);

		socket.on("data", (chunk) => {
			const bufferData = createWriter(Endian.LE, 5).writeUByte(Packet.Data).writeUInt(chunk.length).writeBuffer(chunk).toBuffer();
			txBytes += chunk.length;
			if (authed) {
				client.send(bufferData);
			} else {
				queuedMessages.push(bufferData);
			}
		});

		function serverCloseOrError() {
			client.close();
			clearInterval(txrxInterval);
		}
		socket.on("close", serverCloseOrError);
		socket.on("error", serverCloseOrError);
	});

	server.listen(config.localPort, () => console.log(`Local server listening at ${config.localPort}`));
}