From d64b1e7ca0bf5167cfd598f86842ebee474f03c9 Mon Sep 17 00:00:00 2001
From: amigin <amigin@gmail.com>
Date: Fri, 26 Apr 2024 23:46:17 +0200
Subject: [PATCH] Made Timeout of LazyAcceptTcpStream

---
 src/http_control/startup.rs     |  2 +-
 src/http_server/https_server.rs | 71 +++++++++++++++++++++------------
 2 files changed, 46 insertions(+), 27 deletions(-)

diff --git a/src/http_control/startup.rs b/src/http_control/startup.rs
index 67179d3..9624abb 100644
--- a/src/http_control/startup.rs
+++ b/src/http_control/startup.rs
@@ -7,7 +7,7 @@ use crate::app::AppContext;
 const DEFAULT_PORT: u16 = 8000;
 
 pub fn start(app: &Arc<AppContext>) {
-    let http_port = if let Ok(result) = std::env::var("CONTROL_HTTP_PORT") {
+    let http_port = if let Ok(result) = std::env::var("HTTP_PORT") {
         match result.parse() {
             Ok(port) => port,
             Err(_) => DEFAULT_PORT,
diff --git a/src/http_server/https_server.rs b/src/http_server/https_server.rs
index 0768088..3fcbf7b 100644
--- a/src/http_server/https_server.rs
+++ b/src/http_server/https_server.rs
@@ -1,3 +1,4 @@
+use std::time::Duration;
 use std::{net::SocketAddr, sync::Arc};
 
 use hyper_util::rt::TokioIo;
@@ -39,34 +40,52 @@ async fn start_https_server_loop(addr: SocketAddr, app: Arc<AppContext>) {
 
         println!("Accepted connection from  {}", socket_addr);
 
-        let result = lazy_accept_tcp_stream(app.clone(), endpoint_port, tcp_stream).await;
+        let app = app.clone();
+        tokio::spawn(async move { handle_connection(app, endpoint_port, tcp_stream, socket_addr) });
+    }
+}
 
-        if let Err(err) = &result {
-            eprintln!("failed to perform tls handshake: {err:#}");
-            continue;
-        }
+async fn handle_connection(
+    app: Arc<AppContext>,
+    endpoint_port: u16,
+    tcp_stream: TcpStream,
+    socket_addr: SocketAddr,
+) {
+    let future = lazy_accept_tcp_stream(app.clone(), endpoint_port, tcp_stream);
 
-        let (tls_stream, endpoint_info, cn_user_name) = result.unwrap();
-
-        if endpoint_info.http_type.is_protocol_http1() {
-            kick_off_https1(
-                app.clone(),
-                socket_addr,
-                endpoint_info,
-                tls_stream,
-                cn_user_name,
-                endpoint_port,
-            );
-        } else {
-            kick_off_https2(
-                app.clone(),
-                socket_addr,
-                endpoint_info,
-                tls_stream,
-                cn_user_name,
-                endpoint_port,
-            );
-        }
+    let result = tokio::time::timeout(Duration::from_secs(10), future).await;
+
+    if result.is_err() {
+        println!("Timeout waiting for tls handshake from {}", socket_addr);
+    }
+
+    let result = result.unwrap();
+
+    if let Err(err) = &result {
+        eprintln!("failed to perform tls handshake: {err:#}");
+        return;
+    }
+
+    let (tls_stream, endpoint_info, cn_user_name) = result.unwrap();
+
+    if endpoint_info.http_type.is_protocol_http1() {
+        kick_off_https1(
+            app,
+            socket_addr,
+            endpoint_info,
+            tls_stream,
+            cn_user_name,
+            endpoint_port,
+        );
+    } else {
+        kick_off_https2(
+            app,
+            socket_addr,
+            endpoint_info,
+            tls_stream,
+            cn_user_name,
+            endpoint_port,
+        );
     }
 }