From 71a10041b0df729659415201e64cd6084fe3a6bc Mon Sep 17 00:00:00 2001
From: cheesycod <cheesycod@outlook.com>
Date: Mon, 12 Feb 2024 09:03:12 +0530
Subject: [PATCH 1/2] add allow_compression to proxy

---
 src/config.rs  |  6 ++++++
 src/server.rs  | 12 ++++++++++--
 src/upgrade.rs |  4 ++--
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 71ecf7b..8903acb 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -42,6 +42,8 @@ pub struct Config {
     pub backpressure: usize,
     #[serde(default = "default_validate_token")]
     pub validate_token: bool,
+    #[serde(default = "default_allow_compression")]
+    pub allow_compression: bool,
     #[serde(default)]
     pub twilight_http_proxy: Option<String>,
     pub externally_accessible_url: String,
@@ -209,6 +211,10 @@ const fn default_validate_token() -> bool {
     true
 }
 
+const fn default_allow_compression() -> bool {
+    true
+}   
+
 pub enum Error {
     InvalidConfig(JsonError),
     NotFound(String),
diff --git a/src/server.rs b/src/server.rs
index d667905..606af8e 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -271,10 +271,18 @@ pub async fn handle_client<S: 'static + AsyncRead + AsyncWrite + Unpin + Send>(
 
                 trace!("[{addr}] Shard ID is {shard_id}");
 
+		let compress = {
+                    if CONFIG.allow_compression {
+			identify.d.compress
+		    } else {
+                        Some(false)
+                    } 
+                };
+
                 // Create a new session for this client
                 let session = Session {
                     shard_id,
-                    compress: identify.d.compress,
+                    compress,
                 };
                 let session_id = state.create_session(session);
 
@@ -291,7 +299,7 @@ pub async fn handle_client<S: 'static + AsyncRead + AsyncWrite + Unpin + Send>(
                         0,
                     )));
 
-                    let _res = sender.send(identify.d.compress);
+                    let _res = sender.send(compress);
                 }
             }
             6 => {
diff --git a/src/upgrade.rs b/src/upgrade.rs
index 527be15..588d225 100644
--- a/src/upgrade.rs
+++ b/src/upgrade.rs
@@ -15,7 +15,7 @@ use tracing::error;
 
 use std::net::SocketAddr;
 
-use crate::{server::handle_client, state::State};
+use crate::{server::handle_client, state::State, config::CONFIG};
 
 /// Websocket GUID constant as specified in RFC6455:
 /// <https://datatracker.ietf.org/doc/html/rfc6455#section-1.3>
@@ -36,7 +36,7 @@ pub fn server(
 
     // Track whether the client requested zlib encoding in the query
     // string parameters
-    let use_zlib = query.map_or(false, |q| q.contains("compress=zlib-stream"));
+    let use_zlib = CONFIG.allow_compression && query.map_or(false, |q| q.contains("compress=zlib-stream"));
 
     let mut response = Response::new(Full::default());
 

From afbee65ed416d9b0164b4c6b01e2e2a71e8be610 Mon Sep 17 00:00:00 2001
From: cheesycod <cheesycod@outlook.com>
Date: Wed, 14 Feb 2024 20:17:03 +0530
Subject: [PATCH 2/2] format code

---
 src/config.rs  |  2 +-
 src/server.rs  | 13 +++++--------
 src/upgrade.rs |  5 +++--
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 8903acb..62121b8 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -213,7 +213,7 @@ const fn default_validate_token() -> bool {
 
 const fn default_allow_compression() -> bool {
     true
-}   
+}
 
 pub enum Error {
     InvalidConfig(JsonError),
diff --git a/src/server.rs b/src/server.rs
index 606af8e..ac2e256 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -271,19 +271,16 @@ pub async fn handle_client<S: 'static + AsyncRead + AsyncWrite + Unpin + Send>(
 
                 trace!("[{addr}] Shard ID is {shard_id}");
 
-		let compress = {
+                let compress = {
                     if CONFIG.allow_compression {
-			identify.d.compress
-		    } else {
+                        identify.d.compress
+                    } else {
                         Some(false)
-                    } 
+                    }
                 };
 
                 // Create a new session for this client
-                let session = Session {
-                    shard_id,
-                    compress,
-                };
+                let session = Session { shard_id, compress };
                 let session_id = state.create_session(session);
 
                 // The client is connected to this shard, so prepare for sending commands to it
diff --git a/src/upgrade.rs b/src/upgrade.rs
index 588d225..b69baa0 100644
--- a/src/upgrade.rs
+++ b/src/upgrade.rs
@@ -15,7 +15,7 @@ use tracing::error;
 
 use std::net::SocketAddr;
 
-use crate::{server::handle_client, state::State, config::CONFIG};
+use crate::{config::CONFIG, server::handle_client, state::State};
 
 /// Websocket GUID constant as specified in RFC6455:
 /// <https://datatracker.ietf.org/doc/html/rfc6455#section-1.3>
@@ -36,7 +36,8 @@ pub fn server(
 
     // Track whether the client requested zlib encoding in the query
     // string parameters
-    let use_zlib = CONFIG.allow_compression && query.map_or(false, |q| q.contains("compress=zlib-stream"));
+    let use_zlib =
+        CONFIG.allow_compression && query.map_or(false, |q| q.contains("compress=zlib-stream"));
 
     let mut response = Response::new(Full::default());