[Add] TCP server to core MTP (HTTP/1.1 & HTTP/2) compatibility

mtp-webserver/tests/integration.rs

@@ -1,9 +1,11 @@
-use http::StatusCode;
+use http::{Method, StatusCode};
 use mtp_webserver::{MTPWebServer, WebServerConfig, WebServerError, WebServerMetrics};
-use std::net::IpAddr;
+use rustls::pki_types::{CertificateDer, ServerName, pem::PemObject};
+use std::net::{IpAddr, SocketAddr};
 use std::sync::Arc;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::time::Duration;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
 fn generate_self_signed_cert() -> (Vec<u8>, Vec<u8>) {
     let key_pair = rcgen::KeyPair::generate().expect("failed to generate self-signed key pair");
@@ -18,11 +20,52 @@ fn generate_self_signed_cert() -> (Vec<u8>, Vec<u8>) {
     )
 }
 
+fn host_config(port: u16, cert: Vec<u8>, key: Vec<u8>) -> mtp_host::HostConfig {
+    mtp_host::HostConfig::new(IpAddr::V4(std::net::Ipv4Addr::LOCALHOST), port, cert, key)
+}
+
+async fn tls_connect(
+    addr: SocketAddr,
+    cert_pem: &[u8],
+    alpn: Vec<Vec<u8>>,
+) -> tokio_rustls::client::TlsStream<tokio::net::TcpStream> {
+    let certs = CertificateDer::pem_slice_iter(cert_pem)
+        .collect::<Result<Vec<_>, _>>()
+        .unwrap();
+    let mut roots = rustls::RootCertStore::empty();
+    for cert in certs {
+        roots.add(cert).unwrap();
+    }
+    let mut config = rustls::ClientConfig::builder()
+        .with_root_certificates(roots)
+        .with_no_client_auth();
+    config.alpn_protocols = alpn;
+    let connector = tokio_rustls::TlsConnector::from(Arc::new(config));
+    connector
+        .connect(
+            ServerName::try_from("localhost").unwrap().to_owned(),
+            tokio::net::TcpStream::connect(addr).await.unwrap(),
+        )
+        .await
+        .unwrap()
+}
+
+async fn http1_request(addr: SocketAddr, cert_pem: &[u8], request: &str) -> Vec<u8> {
+    let mut stream = tls_connect(addr, cert_pem, vec![b"http/1.1".to_vec()]).await;
+    stream.write_all(request.as_bytes()).await.unwrap();
+    let mut response = Vec::new();
+    stream.read_to_end(&mut response).await.unwrap();
+    response
+}
+
 #[test]
 fn config_builder_defaults() {
     let config = WebServerConfig::new();
     assert_eq!(config.max_request_body, 4 * 1024 * 1024);
     assert_eq!(config.max_connections, 256);
+    assert!(config.serve_tcp_https);
+    assert_eq!(config.max_tcp_connections, 256);
+    assert_eq!(config.tls_handshake_timeout, Duration::from_secs(10));
     assert_eq!(config.request_timeout, Duration::from_secs(30));
 }
 
@@ -30,10 +73,13 @@ fn config_builder_defaults() {
 fn config_builder_chain() {
     let config = WebServerConfig::new()
         .max_connections(64)
+        .max_tcp_connections(32)
+        .tls_handshake_timeout(Duration::from_secs(2))
         .max_request_body(1024)
         .request_timeout(Duration::from_secs(5))
         .mtp_path("/ws");
     assert_eq!(config.max_connections, 64);
+    assert_eq!(config.max_tcp_connections, 32);
     assert_eq!(config.max_request_body, 1024);
     assert_eq!(config.request_timeout, Duration::from_secs(5));
 }
@@ -201,5 +247,301 @@ async fn graceful_shutdown_completes() {
     let server = MTPWebServer::new(host_config, WebServerConfig::new())
         .await
         .unwrap();
+    let addr = server.local_addr();
     server.shutdown().await;
+    assert!(tokio::net::TcpStream::connect(addr).await.is_err());
+    assert!(std::net::UdpSocket::bind(addr).is_ok());
+}
+
+#[tokio::test]
+async fn tcp_and_udp_share_port_zero_assignment() {
+    let (cert, key) = generate_self_signed_cert();
+    let server = MTPWebServer::new(host_config(0, cert, key), WebServerConfig::new())
+        .await
+        .unwrap();
+    let addr = server.local_addr();
+    assert!(addr.port() > 0);
+    assert!(tokio::net::TcpStream::connect(addr).await.is_ok());
+    assert!(std::net::UdpSocket::bind(addr).is_err());
+    server.close().await;
+}
+
+#[tokio::test]
+async fn tcp_conflict_fails_and_udp_only_does_not_claim_tcp() {
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let (cert, key) = generate_self_signed_cert();
+    assert!(
+        MTPWebServer::new(host_config(addr.port(), cert, key), WebServerConfig::new())
+            .await
+            .is_err()
+    );
+    drop(listener);
+
+    let (cert, key) = generate_self_signed_cert();
+    let server = MTPWebServer::new(
+        host_config(addr.port(), cert, key),
+        WebServerConfig::new().serve_tcp_https(false),
+    )
+    .await
+    .unwrap();
+    let tcp = tokio::net::TcpListener::bind(addr).await.unwrap();
+    drop(tcp);
+    server.close().await;
+}
+
+#[tokio::test]
+async fn close_and_drop_release_tcp_listener() {
+    let (cert, key) = generate_self_signed_cert();
+    let server = MTPWebServer::new(host_config(0, cert, key), WebServerConfig::new())
+        .await
+        .unwrap();
+    let addr = server.local_addr();
+    server.close().await;
+    assert!(tokio::net::TcpListener::bind(addr).await.is_ok());
+
+    let (cert, key) = generate_self_signed_cert();
+    let server = MTPWebServer::new(host_config(0, cert, key), WebServerConfig::new())
+        .await
+        .unwrap();
+    let addr = server.local_addr();
+    drop(server);
+    tokio::task::yield_now().await;
+    assert!(tokio::net::TcpListener::bind(addr).await.is_ok());
+}
+
+#[tokio::test]
+async fn http1_routes_bodies_chunks_and_timeout() {
+    let (cert, key) = generate_self_signed_cert();
+    let config = WebServerConfig::new()
+        .max_request_body(8)
+        .request_timeout(Duration::from_millis(20))
+        .route_method(Method::POST, "/exact", |request, response| async move {
+            response.body(request.body.unwrap()).body("-tail")
+        })
+        .unwrap()
+        .route_pattern("/users/{user}", |request, response, params| async move {
+            response.body(format!(
+                "{}:{}",
+                params["user"],
+                request.uri.query().unwrap_or("")
+            ))
+        })
+        .unwrap()
+        .route("/slow", |_, response| async move {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            response
+        })
+        .unwrap()
+        .fallback(|_, response| async move { response.status(StatusCode::IM_A_TEAPOT) })
+        .unwrap();
+    let server = MTPWebServer::new(host_config(0, cert.clone(), key), config)
+        .await
+        .unwrap();
+    let addr = server.local_addr();
+
+    let response = http1_request(addr, &cert, "POST /exact HTTP/1.1\r\nHost: localhost\r\nContent-Length: 4\r\nConnection: close\r\n\r\ndata").await;
+    assert!(String::from_utf8_lossy(&response).ends_with("data-tail"));
+    let response = http1_request(
+        addr,
+        &cert,
+        "GET /users/alice%2Dsmith?full=1 HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n",
+    )
+    .await;
+    assert!(String::from_utf8_lossy(&response).ends_with("alice-smith:full=1"));
+    let response = http1_request(
+        addr,
+        &cert,
+        "GET /exact HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n",
+    )
+    .await;
+    assert!(String::from_utf8_lossy(&response).starts_with("HTTP/1.1 418"));
+    let response = http1_request(addr, &cert, "POST /exact HTTP/1.1\r\nHost: localhost\r\nContent-Length: 9\r\nConnection: close\r\n\r\n123456789").await;
+    assert!(String::from_utf8_lossy(&response).starts_with("HTTP/1.1 413"));
+    let response = http1_request(
+        addr,
+        &cert,
+        "GET /slow HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n",
+    )
+    .await;
+    assert!(String::from_utf8_lossy(&response).starts_with("HTTP/1.1 408"));
+    server.close().await;
+}
+
+#[tokio::test]
+async fn tls_negotiates_h2() {
+    let (cert, key) = generate_self_signed_cert();
+    let config = WebServerConfig::new()
+        .route("/h2", |_, response| async move { response.body("over-h2") })
+        .unwrap();
+    let server = MTPWebServer::new(host_config(0, cert.clone(), key), config)
+        .await
+        .unwrap();
+    let stream = tls_connect(server.local_addr(), &cert, vec![b"h2".to_vec()]).await;
+    assert_eq!(stream.get_ref().1.alpn_protocol(), Some(b"h2".as_slice()));
+    let (mut sender, connection) = hyper::client::conn::http2::handshake(
+        hyper_util::rt::TokioExecutor::new(),
+        hyper_util::rt::TokioIo::new(stream),
+    )
+    .await
+    .unwrap();
+    tokio::spawn(async move {
+        let _ = connection.await;
+    });
+    let request = http::Request::builder()
+        .uri("https://localhost/h2")
+        .body(http_body_util::Empty::<bytes::Bytes>::new())
+        .unwrap();
+    let response = sender.send_request(request).await.unwrap();
+    assert_eq!(response.status(), StatusCode::OK);
+    let body = http_body_util::BodyExt::collect(response.into_body())
+        .await
+        .unwrap()
+        .to_bytes();
+    assert_eq!(body, "over-h2");
+    server.close().await;
+}
+
+#[tokio::test]
+async fn tcp_metrics_report_completion_and_errors() {
+    let (cert, key) = generate_self_signed_cert();
+    let metrics = Arc::new(TestMetrics::new());
+    let config = WebServerConfig::new()
+        .request_timeout(Duration::from_millis(10))
+        .with_metrics(metrics.clone())
+        .route("/slow", |_, response| async move {
+            tokio::time::sleep(Duration::from_millis(50)).await;
+            response
+        })
+        .unwrap();
+    let server = MTPWebServer::new(host_config(0, cert.clone(), key), config)
+        .await
+        .unwrap();
+    let response = http1_request(
+        server.local_addr(),
+        &cert,
+        "GET /slow HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n",
+    )
+    .await;
+    assert!(String::from_utf8_lossy(&response).starts_with("HTTP/1.1 408"));
+    assert_eq!(metrics.requests_started.load(Ordering::SeqCst), 1);
+    assert_eq!(metrics.requests_completed.load(Ordering::SeqCst), 1);
+    assert_eq!(metrics.errors.load(Ordering::SeqCst), 1);
+    server.close().await;
+}
+
+#[tokio::test]
+async fn http1_keep_alive_and_streaming_work() {
+    let (cert, key) = generate_self_signed_cert();
+    let config = WebServerConfig::new()
+        .route("/one", |_, response| async move { response.body("one") })
+        .unwrap()
+        .route("/stream", |_, response| async move {
+            let (tx, rx) = tokio::sync::mpsc::channel(1);
+            tokio::spawn(async move {
+                tx.send(bytes::Bytes::from_static(b"first-")).await.unwrap();
+                tokio::time::sleep(Duration::from_millis(10)).await;
+                let _ = tx.send(bytes::Bytes::from_static(b"second")).await;
+            });
+            response.stream(rx)
+        })
+        .unwrap();
+    let server = MTPWebServer::new(host_config(0, cert.clone(), key), config)
+        .await
+        .unwrap();
+    let addr = server.local_addr();
+    let mut stream = tls_connect(addr, &cert, vec![b"http/1.1".to_vec()]).await;
+    stream.write_all(b"GET /one HTTP/1.1\r\nHost: localhost\r\n\r\nGET /one HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n").await.unwrap();
+    let mut response = Vec::new();
+    stream.read_to_end(&mut response).await.unwrap();
+    assert_eq!(
+        String::from_utf8_lossy(&response)
+            .matches("HTTP/1.1 200")
+            .count(),
+        2
+    );
+
+    let response = http1_request(
+        addr,
+        &cert,
+        "GET /stream HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n",
+    )
+    .await;
+    let response = String::from_utf8_lossy(&response);
+    assert!(response.contains("first-"));
+    assert!(response.contains("second"));
+    assert!(!response.to_ascii_lowercase().contains("content-length:"));
+    server.close().await;
+}
+
+#[tokio::test]
+async fn shutdown_allows_active_request_within_drain_period() {
+    let (cert, key) = generate_self_signed_cert();
+    let entered = Arc::new(tokio::sync::Notify::new());
+    let handler_entered = Arc::clone(&entered);
+    let config = WebServerConfig::new()
+        .drain_timeout(Duration::from_millis(250))
+        .route("/work", move |_, response| {
+            let entered = Arc::clone(&handler_entered);
+            async move {
+                entered.notify_one();
+                tokio::time::sleep(Duration::from_millis(30)).await;
+                response.body("done")
+            }
+        })
+        .unwrap();
+    let server = MTPWebServer::new(host_config(0, cert.clone(), key), config)
+        .await
+        .unwrap();
+    let addr = server.local_addr();
+    let request = tokio::spawn(async move {
+        http1_request(
+            addr,
+            &cert,
+            "GET /work HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n",
+        )
+        .await
+    });
+    entered.notified().await;
+    server.shutdown().await;
+    assert!(String::from_utf8_lossy(&request.await.unwrap()).ends_with("done"));
+}
+
+#[tokio::test]
+async fn shutdown_terminates_request_after_drain_period() {
+    let (cert, key) = generate_self_signed_cert();
+    let entered = Arc::new(tokio::sync::Notify::new());
+    let handler_entered = Arc::clone(&entered);
+    let config = WebServerConfig::new()
+        .drain_timeout(Duration::from_millis(20))
+        .route("/stuck", move |_, response| {
+            let entered = Arc::clone(&handler_entered);
+            async move {
+                entered.notify_one();
+                std::future::pending::<()>().await;
+                response
+            }
+        })
+        .unwrap();
+    let server = MTPWebServer::new(host_config(0, cert.clone(), key), config)
+        .await
+        .unwrap();
+    let addr = server.local_addr();
+    let request = tokio::spawn(async move {
+        let mut stream = tls_connect(addr, &cert, vec![b"http/1.1".to_vec()]).await;
+        stream
+            .write_all(b"GET /stuck HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n")
+            .await
+            .unwrap();
+        let mut response = Vec::new();
+        let _ = stream.read_to_end(&mut response).await;
+        response
+    });
+    entered.notified().await;
+    server.shutdown().await;
+    let response = tokio::time::timeout(Duration::from_millis(250), request)
+        .await
+        .expect("connection task survived drain timeout")
+        .unwrap();
+    assert!(!String::from_utf8_lossy(&response).contains("HTTP/1.1 200"));
 }