From 9896108b5a64e64621fd6c598e0e92dca70f2b4e Mon Sep 17 00:00:00 2001
From: synoet <teonys@nyu.edu>
Date: Sun, 5 May 2024 23:59:32 -0400
Subject: [PATCH] FIX ws upgrade

---
 theia/proxy-rs/src/main.rs | 2 +-
 theia/proxy-rs/src/ws.rs   | 4 +---
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/theia/proxy-rs/src/main.rs b/theia/proxy-rs/src/main.rs
index 105adb72..d40d87be 100644
--- a/theia/proxy-rs/src/main.rs
+++ b/theia/proxy-rs/src/main.rs
@@ -212,7 +212,7 @@ async fn handle(
 
     let host = format!("ws://{}:{}", cluster_address, port);
 
-    ws.on_upgrade(move |socket| ws::forward(&host, ws));
+    ws.on_upgrade(move |socket| ws::forward(host, socket));
     (StatusCode::OK, "authorized".to_string())
 }
 
diff --git a/theia/proxy-rs/src/ws.rs b/theia/proxy-rs/src/ws.rs
index 31ded65a..2f47d581 100644
--- a/theia/proxy-rs/src/ws.rs
+++ b/theia/proxy-rs/src/ws.rs
@@ -1,6 +1,4 @@
-use anyhow::Result;
 use axum::extract::ws::{CloseFrame, Message as AxumMessage, WebSocket};
-use axum::extract::WebSocketUpgrade;
 use futures_util::{SinkExt, StreamExt};
 use tokio_tungstenite::tungstenite;
 use tokio_tungstenite::{connect_async, tungstenite::protocol::Message as TsMessage};
@@ -68,7 +66,7 @@ impl WebSocketMessage {
     }
 }
 
-pub async fn forward(url: &str, client_ws: WebSocketUpgrade) {
+pub async fn forward(url: String, client_ws: WebSocket) {
     let server_ws = match connect_async(url).await {
         Ok((ws, _)) => ws,
         Err(e) => {