Fix spec deviations and implement graceful shutdown drain

- Replace determine_if_https() with ProxyState.is_https field so X-Forwarded-Proto
  reflects the listener's protocol instead of guessing from the Host header
- Return ProxyError::BadGateway with host/upstream context for non-connect upstream
  errors instead of bare StatusCode::BAD_GATEWAY
- Implement InFlightCounter with RAII guard for tracking in-flight connections
- Add drain_in_flight() to wait for connections to complete on shutdown, with
  configurable timeout before forcing exit
- Mark review/core-components and review/integration-readiness as complete

src/main.rs

@@ -15,7 +15,7 @@ use reverse_proxy::health;
 use reverse_proxy::logging;
 use reverse_proxy::proxy::{build_router, create_http_client, create_https_client, ProxyState};
 use reverse_proxy::rate_limit::{start_eviction_task, RateLimiter};
-use reverse_proxy::server::serve_https_listener;
+use reverse_proxy::server::{drain_in_flight, serve_https_listener, InFlightCounter};
 use reverse_proxy::shutdown::GracefulShutdown;
 use reverse_proxy::tls::acceptor::{setup_tls, TlsMode};
 use reverse_proxy::tls::redirect;
@@ -74,6 +74,7 @@ async fn run_server(loaded_config: cli::LoadedConfig, config_path: &str) -> Resu
         config: config_arc.clone(),
         http_client,
         https_client,
+        is_https: true,
     });
 
     let reload_handle = Arc::new(ConfigReloadHandle::new(
@@ -197,6 +198,8 @@ async fn run_server(loaded_config: cli::LoadedConfig, config_path: &str) -> Resu
 
     let app = build_router(proxy_state.clone(), config_arc.clone(), rate_limiter);
 
+    let in_flight = InFlightCounter::new();
+
     let mut https_server_handles = Vec::new();
 
     for ((listener_config, tcp_listener), tls_acceptor) in
@@ -209,6 +212,7 @@ async fn run_server(loaded_config: cli::LoadedConfig, config_path: &str) -> Resu
             tls_acceptor,
             app.clone(),
             shutdown_rx,
+            in_flight.clone(),
         ));
 
         info!(
@@ -235,5 +239,15 @@ async fn run_server(loaded_config: cli::LoadedConfig, config_path: &str) -> Resu
         handle.abort();
     }
 
+    let remaining = drain_in_flight(&in_flight, shutdown.shutdown_timeout()).await;
+    if remaining > 0 {
+        warn!(
+            remaining = remaining,
+            "shutdown timeout expired, forcing exit"
+        );
+    } else {
+        info!("all in-flight requests completed");
+    }
+
     Ok(())
 }

src/proxy/handler.rs

@@ -22,6 +22,7 @@ pub struct ProxyState {
     pub config: Arc<ArcSwap<DynamicConfig>>,
     pub http_client: Client<HttpConnector, Body>,
     pub https_client: Client<hyper_rustls::HttpsConnector<HttpConnector>, Body>,
+    pub is_https: bool,
 }
 
 async fn health_handler() -> impl IntoResponse {
@@ -53,8 +54,8 @@ async fn proxy_handler(
         None => return ProxyError::UnknownHost.into_response(),
     };
 
-    let is_https = determine_if_https(host);
-    inject_proxy_headers(req.headers_mut(), remote_addr, is_https);
+    let host_owned = host.to_string();
+    inject_proxy_headers(req.headers_mut(), remote_addr, state.is_https);
     remove_hop_by_hop(req.headers_mut());
 
     let upstream_scheme = site.upstream_scheme.clone();
@@ -89,24 +90,18 @@ async fn proxy_handler(
             if e.is_connect() {
                 ProxyError::UpstreamConnection(e).into_response()
             } else {
-                warn!(error = %e, "upstream request failed");
-                StatusCode::BAD_GATEWAY.into_response()
+                let upstream_addr = format!("{}://{}", upstream_scheme, upstream);
+                ProxyError::BadGateway {
+                    host: host_owned,
+                    upstream: upstream_addr,
+                }
+                .into_response()
             }
         }
         Err(_) => ProxyError::UpstreamTimeout.into_response(),
     }
 }
 
-fn determine_if_https(host: &str) -> bool {
-    let port_str = host.split(':').nth(1);
-    if let Some(port) = port_str {
-        if let Ok(p) = port.parse::<u16>() {
-            return p == 443;
-        }
-    }
-    true
-}
-
 fn build_upstream_uri(scheme: &str, upstream: &str, original_uri: &Uri) -> Uri {
     let path = original_uri.path();
     let query = original_uri
@@ -200,6 +195,7 @@ mod tests {
             ))),
             http_client: create_http_client(),
             https_client: create_https_client(),
+            is_https: true,
         })
     }
 
@@ -298,26 +294,6 @@ mod tests {
         assert_eq!(resp.status(), StatusCode::NOT_FOUND);
     }
 
-    #[test]
-    fn test_determine_if_https_port_443() {
-        assert!(determine_if_https("example.com:443"));
-    }
-
-    #[test]
-    fn test_determine_if_https_port_80() {
-        assert!(!determine_if_https("example.com:80"));
-    }
-
-    #[test]
-    fn test_determine_if_https_no_port() {
-        assert!(determine_if_https("example.com"));
-    }
-
-    #[test]
-    fn test_determine_if_https_port_8443() {
-        assert!(!determine_if_https("example.com:8443"));
-    }
-
     #[test]
     fn test_build_upstream_uri_with_query() {
         let uri: Uri = "/path?foo=bar".parse().unwrap();

src/server.rs

@@ -1,4 +1,6 @@
 use std::net::SocketAddr;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
 
 use axum::extract::ConnectInfo;
 use axum::http::Request;
@@ -9,13 +11,46 @@ use hyper_util::service::TowerToHyperService;
 use tokio::net::TcpListener;
 use tokio_rustls::TlsAcceptor;
 use tower::Service;
-use tracing::{error, warn};
+use tracing::{error, info, warn};
+
+pub struct InFlightCounter {
+    count: AtomicUsize,
+}
+
+struct InFlightGuard(Arc<InFlightCounter>);
+
+impl Drop for InFlightGuard {
+    fn drop(&mut self) {
+        self.0.decrement();
+    }
+}
+
+impl InFlightCounter {
+    pub fn new() -> Arc<Self> {
+        Arc::new(Self {
+            count: AtomicUsize::new(0),
+        })
+    }
+
+    pub fn increment(&self) {
+        self.count.fetch_add(1, Ordering::SeqCst);
+    }
+
+    pub fn decrement(&self) {
+        self.count.fetch_sub(1, Ordering::SeqCst);
+    }
+
+    pub fn is_zero(&self) -> bool {
+        self.count.load(Ordering::SeqCst) == 0
+    }
+}
 
 pub async fn serve_https_listener(
     tcp_listener: TcpListener,
     tls_acceptor: TlsAcceptor,
     router: Router,
     mut shutdown_rx: tokio::sync::watch::Receiver<bool>,
+    in_flight: Arc<InFlightCounter>,
 ) {
     let local_addr = tcp_listener.local_addr();
 
@@ -32,8 +67,11 @@ pub async fn serve_https_listener(
 
                 let tls_acceptor = tls_acceptor.clone();
                 let router = router.clone();
+                let in_flight = in_flight.clone();
 
                 tokio::spawn(async move {
+                    let _guard = InFlightGuard(in_flight.clone());
+
                     let tls_stream = match tls_acceptor.accept(tcp_stream).await {
                         Ok(stream) => stream,
                         Err(e) => {
@@ -63,7 +101,7 @@ pub async fn serve_https_listener(
             }
             _ = shutdown_rx.changed() => {
                 if let Ok(addr) = local_addr {
-                    tracing::info!(addr = %addr, "HTTPS listener shutting down");
+                    info!(addr = %addr, "HTTPS listener shutting down");
                 }
                 break;
             }
@@ -71,6 +109,24 @@ pub async fn serve_https_listener(
     }
 }
 
+/// Wait for in-flight connections to drain, with a timeout.
+/// Returns the number of connections still in-flight when the timeout expired (0 if all drained).
+pub async fn drain_in_flight(
+    in_flight: &Arc<InFlightCounter>,
+    timeout: std::time::Duration,
+) -> usize {
+    let start = std::time::Instant::now();
+    loop {
+        if in_flight.is_zero() {
+            return 0;
+        }
+        if start.elapsed() >= timeout {
+            return in_flight.count.load(Ordering::SeqCst);
+        }
+        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+    }
+}
+
 #[derive(Clone)]
 struct ConnectInfoService<S> {
     inner: S,