ZSTD on gateway

src/gateway/events/Connection.ts

@@ -28,7 +28,8 @@ import { Message } from "./Message";
 import { Deflate, Inflate } from "fast-zlib";
 import { URL } from "url";
 import { Config, ErlpackType } from "@spacebar/util";
-// import zlib from "node:zlib";
+import zlib from "node:zlib";
+import { Decoder, Encoder } from "@toondepauw/node-zstd";
 
 let erlpack: ErlpackType | null = null;
 try {
@@ -108,7 +109,7 @@ export async function Connection(
 			return socket.close(CLOSECODES.Decode_error);
 
 		if (socket.encoding === "etf" && !erlpack)
-			throw new Error("Erlpack is not installed: 'npm i erlpack'");
+			throw new Error("Erlpack is not installed: 'npm i @yukikaze-bot/erlpack'");
 
 		socket.version = Number(searchParams.get("version")) || 8;
 		if (socket.version != 8)
@@ -121,11 +122,8 @@ export async function Connection(
 				socket.deflate = new Deflate();
 				socket.inflate = new Inflate();
 			} else if (socket.compress === "zstd-stream") {
-				// TODO
-				return socket.close(
-					CLOSECODES.Decode_error
-				);
-
+				socket.zstdEncoder = new Encoder(6);
+				socket.zstdDecoder = new Decoder();
 			} else {
 				return socket.close(CLOSECODES.Decode_error);
 			}

src/gateway/events/Message.ts

@@ -45,9 +45,15 @@ export async function Message(this: WebSocket, buffer: WS.Data) {
 	) {
 		data = bigIntJson.parse(buffer.toString());
 	} else if (this.encoding === "json" && buffer instanceof Buffer) {
-		if (this.inflate) {
+		if (this.compress === "zlib-stream") {
 			try {
-				buffer = this.inflate.process(buffer);
+				buffer = this.inflate!.process(buffer);
+			} catch {
+				buffer = buffer.toString();
+			}
+		} else if (this.compress === "zstd-stream") {
+			try {
+				buffer = await this.zstdDecoder!.decode(buffer);
 			} catch {
 				buffer = buffer.toString();
 			}

src/gateway/util/Send.ts

@@ -43,22 +43,20 @@ const recurseJsonReplace = (json: any) => {
 	return json;
 };
 
-export function Send(socket: WebSocket, data: Payload) {
+export async function Send(socket: WebSocket, data: Payload) {
 	if (process.env.WS_VERBOSE)
 		console.log(`[Websocket] Outgoing message: ${JSON.stringify(data)}`);
 
 	if (process.env.WS_DUMP) {
 		const id = socket.session_id || "unknown";
 
-		(async () => {
-			await fs.mkdir(path.join("dump", id), {
-				recursive: true,
-			});
-			await fs.writeFile(
-				path.join("dump", id, `${Date.now()}.out.json`),
-				JSON.stringify(data, null, 2),
-			);
-		})();
+		await fs.mkdir(path.join("dump", id), {
+			recursive: true,
+		});
+		await fs.writeFile(
+			path.join("dump", id, `${Date.now()}.out.json`),
+			JSON.stringify(data, null, 2),
+		);
 	}
 
 	let buffer: Buffer | string;
@@ -71,9 +69,15 @@ export function Send(socket: WebSocket, data: Payload) {
 	else if (socket.encoding === "json")
 		buffer = JSON.stringify(data, JSONReplacer);
 	else return;
+
 	// TODO: compression
-	if (socket.deflate) {
-		buffer = socket.deflate.process(buffer) as Buffer;
+	if (socket.compress === "zlib-stream") {
+		buffer = socket.deflate!.process(buffer) as Buffer;
+	} else if (socket.compress === "zstd-stream") {
+		if (typeof(buffer) === "string")
+			buffer = Buffer.from(buffer as string);
+
+		buffer = await socket.zstdEncoder!.encode(buffer as Buffer) as Buffer;
 	}
 
 	return new Promise((res, rej) => {

src/gateway/util/WebSocket.ts

@@ -20,19 +20,24 @@ import { Intents, ListenEventOpts, Permissions } from "@spacebar/util";
 import WS from "ws";
 import { Deflate, Inflate } from "fast-zlib";
 import { Capabilities } from "./Capabilities";
+import { ZstdCompress } from "zlib";
+import { ZstdDecompress } from "node:zlib";
+import { Decoder, Encoder } from "@toondepauw/node-zstd";
 
 export interface WebSocket extends WS {
 	version: number;
 	user_id: string;
 	session_id: string;
 	encoding: "etf" | "json";
-	compress?: "zlib-stream";
+	compress?: "zlib-stream" | "zstd-stream";
 	ipAddress?: string;
 	userAgent?: string; // for cdn request signing
 	shard_count?: bigint;
 	shard_id?: bigint;
 	deflate?: Deflate;
 	inflate?: Inflate;
+	zstdEncoder?: Encoder;
+	zstdDecoder?: Decoder;
 	heartbeatTimeout: NodeJS.Timeout;
 	readyTimeout: NodeJS.Timeout;
 	intents: Intents;