From 931ce2b5adbece17d4f6d1fa5c5b99f612e77383 Mon Sep 17 00:00:00 2001
From: "Zhidong Peng (HE/HIM)" <zpeng@microsoft.com>
Date: Tue, 10 Feb 2026 16:46:55 -0800
Subject: [PATCH 1/2] Use SockRef.set_read_timeout without conversion

---
 Cargo.lock                            |  1 +
 proxy_agent/Cargo.toml                |  1 +
 proxy_agent/src/proxy/proxy_server.rs | 56 +++++++--------------------
 3 files changed, 17 insertions(+), 41 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0dab16fb..1df43fca 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -191,6 +191,7 @@ dependencies = [
  "serde-xml-rs",
  "serde_derive",
  "serde_json",
+ "socket2",
  "static_vcruntime",
  "sysinfo",
  "thiserror",
diff --git a/proxy_agent/Cargo.toml b/proxy_agent/Cargo.toml
index f7dd3a01..6b0987d6 100644
--- a/proxy_agent/Cargo.toml
+++ b/proxy_agent/Cargo.toml
@@ -28,6 +28,7 @@ tower-http = { version = "0.6.2", features = ["limit"] }
 clap = { version = "4.5.17", features =["derive"] } # Command Line Argument Parser 
 thiserror = "1.0.64"
 libc = "0.2.147"
+socket2 = "0.5"               # Set socket options without tokio/std conversion
 
 [dependencies.uuid]
 version = "1.3.0"
diff --git a/proxy_agent/src/proxy/proxy_server.rs b/proxy_agent/src/proxy/proxy_server.rs
index 666642a9..5aab37e6 100644
--- a/proxy_agent/src/proxy/proxy_server.rs
+++ b/proxy_agent/src/proxy/proxy_server.rs
@@ -253,24 +253,20 @@ impl ProxyServer {
         tokio::spawn({
             let cloned_proxy_server = self.clone();
             async move {
-                let (stream, _cloned_std_stream) =
-                    match Self::set_stream_read_time_out(stream, &mut tcp_connection_logger) {
-                        Ok((stream, cloned_std_stream)) => (stream, cloned_std_stream),
-                        Err(e) => {
-                            tcp_connection_logger.write(
-                                LoggerLevel::Error,
-                                format!("Failed to set stream read timeout: {e}"),
-                            );
-                            return;
-                        }
-                    };
+                // Get raw socket ID before any conversion (Windows only)
+                #[cfg(windows)]
+                let raw_socket_id = Self::get_stream_raw_socket_id(&stream);
+
+                // Set read timeout directly on the socket without conversion
+                Self::set_stream_read_time_out(&stream, &mut tcp_connection_logger);
+
                 let tcp_connection_context = TcpConnectionContext::new(
                     tcp_connection_id,
                     client_addr,
                     cloned_proxy_server.redirector_shared_state.clone(),
                     cloned_proxy_server.proxy_server_shared_state.clone(),
                     #[cfg(windows)]
-                    ProxyServer::get_stream_rocket_id(&_cloned_std_stream),
+                    raw_socket_id,
                 )
                 .await;
 
@@ -324,46 +320,24 @@ impl ProxyServer {
     }
 
     #[cfg(windows)]
-    fn get_stream_rocket_id(stream: &std::net::TcpStream) -> usize {
+    fn get_stream_raw_socket_id(stream: &TcpStream) -> usize {
         use std::os::windows::io::AsRawSocket;
         stream.as_raw_socket() as usize
     }
 
     // Set the read timeout for the stream
-    fn set_stream_read_time_out(
-        stream: TcpStream,
-        connection_logger: &mut ConnectionLogger,
-    ) -> Result<(TcpStream, std::net::TcpStream)> {
-        // Convert the stream to a std stream
-        let std_stream = stream.into_std().map_err(|e| {
-            Error::Io(
-                "Failed to convert Tokio stream into std equivalent".to_string(),
-                e,
-            )
-        })?;
+    // Uses socket2::SockRef to set socket options directly on the tokio stream
+    // socket2 crate alreasdy used by tokio internally, so it won't cause extra dependency
+    fn set_stream_read_time_out(stream: &TcpStream, connection_logger: &mut ConnectionLogger) {
+        use socket2::SockRef;
 
-        // Set the read timeout
-        if let Err(e) = std_stream.set_read_timeout(Some(std::time::Duration::from_secs(10))) {
+        let sock_ref = SockRef::from(stream);
+        if let Err(e) = sock_ref.set_read_timeout(Some(std::time::Duration::from_secs(10))) {
             connection_logger.write(
                 LoggerLevel::Warn,
                 format!("Failed to set read timeout: {e}"),
             );
         }
-
-        // Clone the stream for the service_fn
-        let cloned_std_stream = std_stream
-            .try_clone()
-            .map_err(|e| Error::Io("Failed to clone TCP stream".to_string(), e))?;
-
-        // Convert the std stream back
-        let tokio_tcp_stream = TcpStream::from_std(std_stream).map_err(|e| {
-            Error::Io(
-                "Failed to convert std stream into Tokio equivalent".to_string(),
-                e,
-            )
-        })?;
-
-        Ok((tokio_tcp_stream, cloned_std_stream))
     }
 
     async fn handle_new_http_request(

From 6a3c30d2c7871798ff2f406a0a3622ac9f7c351d Mon Sep 17 00:00:00 2001
From: "Zhidong Peng (HE/HIM)" <zpeng@microsoft.com>
Date: Tue, 10 Feb 2026 17:49:42 -0800
Subject: [PATCH 2/2] fix

---
 proxy_agent/src/proxy/proxy_server.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/proxy_agent/src/proxy/proxy_server.rs b/proxy_agent/src/proxy/proxy_server.rs
index 5aab37e6..9ecd8e90 100644
--- a/proxy_agent/src/proxy/proxy_server.rs
+++ b/proxy_agent/src/proxy/proxy_server.rs
@@ -327,7 +327,7 @@ impl ProxyServer {
 
     // Set the read timeout for the stream
     // Uses socket2::SockRef to set socket options directly on the tokio stream
-    // socket2 crate alreasdy used by tokio internally, so it won't cause extra dependency
+    // socket2 crate already used by tokio internally, so it won't cause extra dependency
     fn set_stream_read_time_out(stream: &TcpStream, connection_logger: &mut ConnectionLogger) {
         use socket2::SockRef;