From e4115d13ef239c2aade5116e58ecc2693092c8e9 Mon Sep 17 00:00:00 2001
From: Flam3rboy <34555296+Flam3rboy@users.noreply.github.com>
Date: Tue, 6 Apr 2021 04:06:32 +0200
Subject: [PATCH] :zap: zlib compression

---
 src/events/Connection.ts | 11 +++++++++--
 src/util/Send.ts         |  6 ++++++
 src/util/WebSocket.ts    |  3 +++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/events/Connection.ts b/src/events/Connection.ts
index 0c4dd2b8..34274d9c 100644
--- a/src/events/Connection.ts
+++ b/src/events/Connection.ts
@@ -5,6 +5,7 @@ import { Message } from "./Message";
 import { setHeartbeat } from "../util/setHeartbeat";
 import { Send } from "../util/Send";
 import { CLOSECODES, OPCODES } from "../util/Constants";
+import { createDeflate } from "zlib";
 
 // TODO: check rate limit
 // TODO: specify rate limit in config
@@ -24,8 +25,14 @@ export async function Connection(this: Server, socket: WebSocket, request: Incom
 		if (socket.version != 8) return socket.close(CLOSECODES.Invalid_API_version);
 
 		// @ts-ignore
-		socket.compression = searchParams.get("compress") || "";
-		// TODO: compression
+		socket.compress = searchParams.get("compress") || "";
+		if (socket.compress) {
+			if (socket.compress !== "zlib-stream") return socket.close(CLOSECODES.Decode_error);
+			socket.deflate = createDeflate({ chunkSize: 65535 });
+			socket.deflate.on("data", (chunk) => socket.send(chunk));
+		}
+
+		socket.sequence = 0;
 
 		setHeartbeat(socket);
 
diff --git a/src/util/Send.ts b/src/util/Send.ts
index d38865b3..b691b525 100644
--- a/src/util/Send.ts
+++ b/src/util/Send.ts
@@ -1,6 +1,7 @@
 import erlpack from "erlpack";
 import { promisify } from "util";
 import { Payload } from "../util/Constants";
+import { deflateSync } from "zlib";
 
 import WebSocket from "./WebSocket";
 
@@ -11,6 +12,11 @@ export async function Send(socket: WebSocket, data: Payload) {
 	else if (socket.encoding === "json") buffer = JSON.stringify(data);
 
 	// TODO: compression
+	if (socket.deflate) {
+		socket.deflate.write(buffer);
+		socket.deflate.flush();
+		return;
+	}
 
 	return new Promise((res, rej) => {
 		socket.send(buffer, (err) => {
diff --git a/src/util/WebSocket.ts b/src/util/WebSocket.ts
index 5d107d9c..4f866ad5 100644
--- a/src/util/WebSocket.ts
+++ b/src/util/WebSocket.ts
@@ -1,14 +1,17 @@
 import { Intents } from "fosscord-server-util";
 import WS, { Server, Data } from "ws";
+import { Deflate } from "zlib";
 
 interface WebSocket extends WS {
 	version: number;
 	user_id: bigint;
 	encoding: "etf" | "json";
 	compress?: "zlib-stream";
+	deflate?: Deflate;
 	heartbeatTimeout: NodeJS.Timeout;
 	readyTimeout: NodeJS.Timeout;
 	intents: Intents;
+	sequence: number;
 }
 
 export default WebSocket;