remove feature flag, enable_compression arg, fix integration test

Author: haardikk21

crates/websocket-proxy/Cargo.toml

@@ -32,11 +32,10 @@ redis = "0.30.0"
 redis-test = { version = "0.10.0", optional = true }
 uuid = { version = "1.16.0", features = ["v4"] }
 tokio-util = "0.7.12"
-brotli = { version = "8.0.1", optional = true }
+brotli = "8.0.1"
 
 [dependencies.ring]
 version = "0.17.12"
 
 [features]
 integration = ["redis-test"]
-compression = ["brotli"]

crates/websocket-proxy/README.md

@@ -62,3 +62,8 @@ When Redis is enabled, the following features are available:
 
 If the Redis connection fails, the proxy will automatically fall back to in-memory rate limiting.
 
+### Brotli Compression
+
+The proxy supports compressing messages to downstream clients using Brotli.
+
+To enable this, pass the parameter `--enable-compression`

crates/websocket-proxy/src/integration.rs

@@ -27,7 +27,7 @@ mod test {
         server: Server,
         server_addr: SocketAddr,
         client_id_to_handle: HashMap<usize, JoinHandle<()>>,
-        sender: Sender<String>,
+        sender: Sender<Vec<u8>>,
     }
 
     impl TestHarness {
@@ -153,7 +153,10 @@ mod test {
         }
 
         fn send_messages(&mut self, messages: Vec<&str>) {
-            let messages: Vec<String> = messages.into_iter().map(String::from).collect();
+            let messages: Vec<Vec<u8>> = messages
+                .into_iter()
+                .map(|m| m.as_bytes().to_vec())
+                .collect();
 
             for message in messages.iter() {
                 match self.sender.send(message.clone()) {

crates/websocket-proxy/src/main.rs

@@ -18,7 +18,6 @@ use clap::Parser;
 use dotenvy::dotenv;
 use metrics_exporter_prometheus::PrometheusBuilder;
 use rate_limit::RedisRateLimit;
-#[cfg(feature = "compression")]
 use std::io::Write;
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -70,6 +69,13 @@ struct Args {
         help = "Maximum number of concurrently connected clients"
     )]
     per_ip_connections_limit: usize,
+    #[arg(
+        long,
+        env,
+        default_value = "false",
+        help = "Enable brotli compression on messages to downstream clients"
+    )]
+    enable_compression: bool,
 
     #[arg(
         long,
@@ -210,8 +216,7 @@ async fn main() {
             .active_connections
             .set((send.receiver_count() - 1) as f64);
 
-        #[cfg(feature = "compression")]
-        let message_data = {
+        let message_data = if args.enable_compression {
             let data_bytes = data.as_bytes();
             let mut compressed_data_bytes = Vec::new();
             {
@@ -220,11 +225,10 @@ async fn main() {
                 compressor.write_all(&data_bytes).unwrap();
             }
             compressed_data_bytes
+        } else {
+            data.into_bytes()
         };
 
-        #[cfg(not(feature = "compression"))]
-        let message_data = data.into_bytes();
-
         match send.send(message_data) {
             Ok(_) => (),
             Err(e) => error!(message = "failed to send data", error = e.to_string()),