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();
}
diff --git a/src/gateway/util/Send.ts b/src/gateway/util/Send.ts
index 9b1eef6a6..cee3aa8c9 100644
--- a/src/gateway/util/Send.ts
+++ b/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) => {
diff --git a/src/gateway/util/WebSocket.ts b/src/gateway/util/WebSocket.ts
index 2f8fa7cf3..09bae2d02 100644
--- a/src/gateway/util/WebSocket.ts
+++ b/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;
|