Brought Example up to spec

example/server/src/main.rs

@@ -2,14 +2,19 @@ mod clients;
 mod handlers;
 mod keys;
 mod tls;
+#[path = "web-server.rs"]
+mod web_server;
 
-use mtp::host::{AuthenticationPolicy, HostConfig, MTPHost};
+use mtp::host::HostConfig;
 use mtp::type_map::TypeMap;
 use std::future::Future;
 use std::path::Path;
 use std::pin::Pin;
 use std::sync::Arc;
 
+const CONNECTION_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
+const MAX_MESSAGES_PER_CONNECTION: u64 = 10_000;
+
 fn dev_cert_paths() -> (String, String) {
     let cert = std::env::var("MTP_DEV_CERT").unwrap_or_else(|_| {
         if Path::new("example/dev-cert/cert.pem").exists() {
@@ -29,62 +34,81 @@ fn dev_cert_paths() -> (String, String) {
 }
 
 async fn handle_pipe_loopback(
-    conn: &mtp::host::MTPConnection,
-    req: mtp::host::PipeRequest,
+    conn: &mtp::webserver::WebMTPConnection,
+    request: mtp::host::PipeRequest<
+        mtp::webserver::WebMtpSender,
+        mtp::webserver::H3TransportReceiver,
+    >,
 ) -> Result<(), Box<dyn std::error::Error>> {
-    let pipe_id = req.id();
+    let pipe_id = request.id();
+    println!("  [loopback] Accepting pipe {pipe_id} ...");
+    let mut reader = request.accept().await?;
+
+    let return_pipe = conn.create_pipe("loopback").await?;
     println!(
-        "  [loopback] Pipe request: id={pipe_id} description={:?}",
-        req.description()
+        "  [loopback] Requested return pipe {}; waiting for client acceptance ...",
+        return_pipe.pipe_id()
     );
+    let Some(mut writer) = return_pipe.wait().await? else {
+        return Err("client denied the return pipe".into());
+    };
 
-    println!("  [loopback] Calling accept() for pipe {pipe_id} ...");
-    let mut reader = req.accept().await?;
-    println!("  [loopback] Pipe {pipe_id} accepted, reading data ...");
-
-    let handle = conn.create_pipe("loopback").await?;
-    println!(
-        "  [loopback] Return pipe created (id={}), waiting for client ...",
-        handle.pipe_id()
-    );
-
-    match handle.wait().await? {
-        Some(mut writer) => {
-            println!("  [loopback] Client accepted return pipe, echoing incoming bytes ...");
-            let mut total = 0usize;
-            let mut buf = [0u8; 16 * 1024];
-            loop {
-                let n = tokio::io::AsyncReadExt::read(&mut reader, &mut buf).await?;
-                if n == 0 {
-                    break;
-                }
-                total += n;
-                tokio::io::AsyncWriteExt::write_all(&mut writer, &buf[..n]).await?;
-            }
-            writer.finish().await?;
-            println!(
-                "  [loopback] Pipe {pipe_id} loopback complete ({} bytes)",
-                total
-            );
-        }
-        None => {
-            eprintln!("  [loopback] Return pipe denied by client for pipe {pipe_id}");
-        }
-    }
-
+    let copied = tokio::io::copy(&mut reader, &mut writer).await?;
+    writer.finish_async().await?;
+    println!("  [loopback] Pipe {pipe_id} complete ({copied} bytes)");
     Ok(())
 }
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    tracing_subscriber::fmt::init();
     let (cert_path, key_path) = dev_cert_paths();
-    let (cert_pem, key_pem) = tls::load_or_generate_tls(&cert_path, &key_path)?;
-    let cert_hash = tls::certificate_sha256_hex(&cert_pem)?;
-    tls::export_webtransport_cert_hash(&cert_hash)?;
+    let (cert_pem, key_pem) = tls::load_or_generate_tls(&cert_path, &key_path).await?;
+    let cert_hash = tls::certificate_sha256_hex(&cert_pem).await?;
+    tls::export_webtransport_cert_hash(&cert_hash).await?;
     println!("WebTransport certificate sha256: {cert_hash}");
 
     let (_host_id, host_keyring) = keys::load_or_generate_host_keys("host.mk")?;
-    keys::export_host_public_keys(&host_keyring)?;
+    keys::export_host_public_keys(&host_keyring).await?;
+
+    let (clients, next_id) = clients::load_client_db("clients.json").await?;
+    let clients_for_get = clients.clone();
+    let get_existing_client = move |id: u64, _description: Option<String>| {
+        let clients = clients_for_get.clone();
+        Box::pin(async move { clients.lock().ok()?.get(&id).cloned() })
+            as Pin<Box<dyn Future<Output = Option<mtp::crypto::PublicKeyBundle>> + Send>>
+    };
+    let clients_for_register = clients.clone();
+    let next_id_for_register = next_id.clone();
+    let complete_register = move |bundle: mtp::crypto::PublicKeyBundle,
+                                  _description: Option<String>| {
+        let clients = clients_for_register.clone();
+        let next_id = next_id_for_register.clone();
+
+        Box::pin(async move {
+            let id = {
+                let mut next = next_id.lock().expect("client id mutex poisoned");
+                let id = *next;
+                *next += 1;
+                id
+            };
+
+            let json = {
+                let mut db = clients.lock().expect("client database mutex poisoned");
+                db.insert(id, bundle);
+                serde_json::to_string_pretty(&*db).ok()
+            };
+
+            if let Some(json) = json {
+                if let Err(error) = tokio::fs::write("clients.json", json).await {
+                    eprintln!("Failed to persist clients.json: {error}");
+                }
+            }
+
+            println!("Registered new client with ID: {id}");
+            id
+        }) as Pin<Box<dyn Future<Output = u64> + Send>>
+    };
 
     let decrypt_keyring = Arc::new(
         match mtp::crypto::Keyring::from_bytes(&host_keyring.to_bytes()) {
@@ -95,51 +119,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         },
     );
 
-    let (clients, next_id) = clients::load_client_db("clients.json")?;
-
-    let clients_for_get = clients.clone();
-    let get_existing_user = move |id: u64, _description: Option<String>| {
-        let clients = clients_for_get.clone();
-        Box::pin(async move {
-            let result = clients.lock().unwrap().get(&id).cloned();
-            if result.is_some() {
-                println!("Auth lookup: client ID {id} found");
-            } else {
-                eprintln!("Auth lookup: unknown client ID {id}");
-            }
-            result
-        }) as Pin<Box<dyn Future<Output = Option<mtp::crypto::PublicKeyBundle>> + Send>>
-    };
-
-    let clients_for_register = clients.clone();
-    let next_id_for_register = next_id.clone();
-    let clients_path = "clients.json".to_string();
-    let complete_register = move |bundle: mtp::crypto::PublicKeyBundle,
-                                  _description: Option<String>| {
-        let db_arc = clients_for_register.clone();
-        let nid_arc = next_id_for_register.clone();
-        let path = clients_path.clone();
-        Box::pin(async move {
-            let mut db = db_arc.lock().unwrap();
-            let mut nid = nid_arc.lock().unwrap();
-            let id = *nid;
-            *nid += 1;
-            db.insert(id, bundle);
-            match serde_json::to_string_pretty(&*db) {
-                Ok(json) => match std::fs::write(&path, json) {
-                    Ok(()) => {}
-                    Err(e) => eprintln!("Failed to persist client database to {path}: {e}"),
-                },
-                Err(e) => {
-                    eprintln!("Failed to serialize client database after registering {id}: {e}")
-                }
-            }
-            println!("Registered new client with ID: {}", id);
-            id
-        }) as Pin<Box<dyn Future<Output = u64> + Send>>
-    };
-
-    println!("Starting MTP server on port 8080 ...");
+    println!("Starting integrated MTP web server on port 8080 ...");
 
     let config = HostConfig::new(
         std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
@@ -149,13 +129,21 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     )
     .with_authentication(
         host_keyring,
-        Box::new(get_existing_user),
+        Box::new(get_existing_client),
         Box::new(complete_register),
-    )
-    .with_authentication_policy(AuthenticationPolicy::ForceAuthentication);
+    );
 
-    let mut host = MTPHost::new(config).await?;
-    println!("Server listening on {}", host.local_addr());
+    let _https = web_server::spawn_https(
+        std::net::SocketAddr::new(config.ip, config.port),
+        &config.tls_fullchain,
+        &config.tls_key,
+    )
+    .await?;
+    let mut host = mtp::webserver::MTPWebServer::new(config, web_server::config()?).await?;
+    println!(
+        "Server listening on https://{} (TCP HTTPS + UDP WebTransport)",
+        host.local_addr()
+    );
 
     while let Some(conn) = host.accept().await? {
         let decrypt_keyring = Arc::clone(&decrypt_keyring);
@@ -170,56 +158,81 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             let tm: &TypeMap = conn.codec.registry().get(&conn.version).unwrap();
 
             println!("Waiting for messages / pipe requests ...");
-            loop {
-                tokio::select! {
-                    biased;
+            let mut pipe_open = true;
+            let mut message_open = true;
+            let mut messages_received = 0_u64;
 
-                    pipe_req = conn.receive_pipe() => {
-                        match pipe_req {
-                            Ok(req) => {
-                                println!("  Pipe request: id={} desc={:?}", req.id(), req.description());
-                                if let Err(e) = handle_pipe_loopback(&conn, req).await {
-                                    eprintln!("  Pipe loopback error: {e}");
-                                }
-                            }
-                            Err(e) => {
-                                println!("Pipe channel closed: {e}");
-                                break;
-                            }
-                        }
-                    }
-                    msg = conn.receive() => {
-                        match msg {
-                            Ok(msg) => {
-                                println!("Received: {msg}");
-                                let response = match handlers::process_and_respond(
-                                    &msg,
-                                    tm,
-                                    conn.client_public_key.as_ref(),
-                                    &decrypt_keyring,
-                                ) {
-                                    Ok(response) => response,
-                                    Err(e) => {
-                                        eprintln!("Failed to build response: {e}");
-                                        continue;
+            while pipe_open || message_open {
+                let activity = tokio::time::timeout(CONNECTION_IDLE_TIMEOUT, async {
+                    tokio::select! {
+                        biased;
+                        pipe_request = conn.receive_pipe(), if pipe_open => {
+                            match pipe_request {
+                                Ok(request) => {
+                                    if let Err(error) = handle_pipe_loopback(&conn, request).await {
+                                        eprintln!("  [loopback] Pipe error: {error}");
                                     }
-                                };
-                                println!("Sending: {response}");
-                                if let Err(e) = conn.sender.send(&response).await {
-                                    eprintln!("Send error: {e}");
-                                    break;
+                                }
+                                Err(mtp::common::CommunicationError::StreamClosed)
+                                | Err(mtp::common::CommunicationError::ClosedByPeer) => {
+                                    println!("Pipe channel closed normally");
+                                    pipe_open = false;
+                                }
+                                Err(error) => {
+                                    println!("Pipe channel closed: {error}");
+                                    pipe_open = false;
                                 }
                             }
-                            Err(e) => {
-                                println!("Connection ended: {e}");
-                                break;
+                        }
+                        message = conn.receive(), if message_open => {
+                            match message {
+                                Ok(message) => {
+                                    messages_received += 1;
+                                    println!("Received: {message}");
+                                    match handlers::process_and_respond(
+                                        &message,
+                                        tm,
+                                        conn.client_public_key.as_ref(),
+                                        &decrypt_keyring,
+                                    ) {
+                                        Ok(response) => {
+                                            println!("Sending: {response}");
+                                            if let Err(error) = conn.sender.send(&response).await {
+                                                eprintln!("Send error: {error}");
+                                                pipe_open = false;
+                                                message_open = false;
+                                            }
+                                        }
+                                        Err(error) => {
+                                            eprintln!("Failed to build response: {error}");
+                                        }
+                                    }
+                                }
+                                Err(mtp::common::CommunicationError::StreamClosed)
+                                | Err(mtp::common::CommunicationError::ClosedByPeer) => {
+                                    println!("Message channel closed normally");
+                                    message_open = false;
+                                }
+                                Err(error) => {
+                                    println!("Message channel closed: {error}");
+                                    message_open = false;
+                                }
                             }
                         }
                     }
+                })
+                .await;
+
+                if activity.is_err() {
+                    println!("Connection idle timeout reached");
+                    break;
+                }
+                if messages_received >= MAX_MESSAGES_PER_CONNECTION {
+                    println!("Connection message limit reached");
+                    break;
                 }
             }
 
-            conn.sender.close();
             println!("Connection closed\n");
         });
     }