ZSTD on gateway

2 files changed, 13 insertions, 9 deletions

diff --git a/src/gateway/events/Connection.ts b/src/gateway/events/Connection.ts

index b187019e9..bc74dee65 100644
--- a/src/gateway/events/Connection.ts
+++ b/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);
 			}
diff --git a/src/gateway/events/Message.ts b/src/gateway/events/Message.ts

index 6fb222b61..2987f464e 100644
--- a/src/gateway/events/Message.ts
+++ b/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();
 			}