initial client and server

Cargo.lock

@@ -92,7 +92,7 @@ dependencies = [
  "hashbrown 0.15.2",
  "paste",
  "static_assertions",
- "windows",
+ "windows 0.58.0",
  "windows-core 0.58.0",
 ]
 
@@ -1874,8 +1874,10 @@ dependencies = [
  "quinn",
  "rcgen",
  "rustls 0.23.23",
+ "rustls-pemfile 2.2.0",
  "shadow-rs",
  "tokio",
+ "tracing",
  "trust-dns-resolver",
 ]
 
@@ -1941,13 +1943,13 @@ dependencies = [
 
 [[package]]
 name = "hostname"
-version = "0.3.1"
+version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3c731c3e10504cc8ed35cfe2f1db4c9274c3d35fa486e3b31df46f068ef3e867"
+checksum = "f9c7c7c8ac16c798734b8a24560c1362120597c40d5e1459f09498f8f6c8f2ba"
 dependencies = [
+ "cfg-if",
  "libc",
- "match_cfg",
- "winapi",
+ "windows 0.52.0",
 ]
 
 [[package]]
@@ -2151,7 +2153,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b77d01e822461baa8409e156015a1d91735549f0f2c17691bd2d996bef238f7f"
 dependencies = [
  "byteorder-lite",
- "quick-error 2.0.1",
+ "quick-error",
 ]
 
 [[package]]
@@ -2452,12 +2454,6 @@ dependencies = [
  "libc",
 ]
 
-[[package]]
-name = "match_cfg"
-version = "0.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ffbee8634e0d45d258acb448e7eaab3fce7a0a467395d4d9f228e3c1f01fb2e4"
-
 [[package]]
 name = "memchr"
 version = "2.7.4"
@@ -3214,12 +3210,6 @@ version = "1.0.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "afbdc74edc00b6f6a218ca6a5364d6226a259d4b8ea1af4a0ea063f27e179f4d"
 
-[[package]]
-name = "quick-error"
-version = "1.2.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0"
-
 [[package]]
 name = "quick-error"
 version = "2.0.1"
@@ -3428,12 +3418,11 @@ checksum = "19b30a45b0cd0bcca8037f3d0dc3421eaf95327a17cad11964fb8179b4fc4832"
 
 [[package]]
 name = "resolv-conf"
-version = "0.7.0"
+version = "0.7.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "52e44394d2086d010551b14b53b1f24e31647570cd1deb0379e2c21b329aba00"
+checksum = "48375394603e3dd4b2d64371f7148fd8c7baa2680e28741f2cb8d23b59e3d4c4"
 dependencies = [
  "hostname",
- "quick-error 1.2.3",
 ]
 
 [[package]]
@@ -4970,7 +4959,7 @@ dependencies = [
  "wasm-bindgen",
  "web-sys",
  "wgpu-types",
- "windows",
+ "windows 0.58.0",
 ]
 
 [[package]]
@@ -4999,9 +4988,9 @@ dependencies = [
 
 [[package]]
 name = "widestring"
-version = "1.1.0"
+version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7219d36b6eac893fa81e84ebe06485e7dcbb616177469b142df14f1f4deb1311"
+checksum = "dd7cf3379ca1aac9eea11fba24fd7e315d621f8dfe35c8d7d2be8b793726e07d"
 
 [[package]]
 name = "winapi"
@@ -5034,6 +5023,16 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
 
+[[package]]
+name = "windows"
+version = "0.52.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e48a53791691ab099e5e2ad123536d0fff50652600abaf43bbf952894110d0be"
+dependencies = [
+ "windows-core 0.52.0",
+ "windows-targets 0.52.6",
+]
+
 [[package]]
 name = "windows"
 version = "0.58.0"

Cargo.toml

@@ -22,6 +22,8 @@ rcgen = { version = "0.13" }
 rustls = { version = "0.23", features = ["prefer-post-quantum"] }
 shadow-rs = "0.38"
 tokio = { version = "1.43", features = ["full"] }
+rustls-pemfile = "2.2.0"
+tracing = { version = "0.1.41", features = ["log-always"] }
 trust-dns-resolver = { version = "0.23.2", features = ["tokio", "tokio-rustls", "rustls", "dns-over-rustls"] }
 
 [build-dependencies]

src/main.rs

@@ -11,18 +11,26 @@ mod pman;
 mod quic;
 
 use std::fs::File;
-use std::net::SocketAddr;
+use std::io::Write;
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
+use std::path::{Path, PathBuf};
 use std::str::FromStr;
 use bunt::println;
 use clap::{Parser, Subcommand};
 use config::{Config, FileFormat, Source};
-use log::{debug, info};
+use log::{debug, error, info};
 use pman::{init_process_manager, ProcessCommand, ProcessManager};
 use shadow_rs::shadow;
 use std::sync::{Arc, Mutex, OnceLock};
+use std::time::{Duration, Instant};
+use anyhow::{anyhow, bail, Context, Error};
+use quinn::crypto::rustls::QuicServerConfig;
+use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer};
 use tokio::{signal, sync::mpsc::Sender};
+use tracing::{info_span, Instrument};
 use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
 use trust_dns_resolver::{Resolver, TokioAsyncResolver};
+use crate::quic::server::handle_connection;
 
 shadow!(build);
 
@@ -53,6 +61,30 @@ enum Commands {
     #[command(about = "Start client as GUI")]
     GuiClient,
     Devtest,
+    #[command(about = "Start server")]
+    Server {
+        /// file to log TLS keys to for debugging
+        #[clap(long = "keylog")]
+        keylog: bool,
+        /// TLS private key in PEM format
+        #[clap(short = 'k', long = "key", requires = "cert")]
+        key: Option<PathBuf>,
+        /// TLS certificate in PEM format
+        #[clap(short = 'c', long = "cert", requires = "key")]
+        cert: Option<PathBuf>,
+        /// Enable stateless retries
+        #[clap(long = "stateless-retry")]
+        stateless_retry: bool,
+        /// Address to listen on
+        #[clap(long = "listen", default_value = "[::1]:4433")]
+        listen: SocketAddr,
+        /// Client address to block
+        #[clap(long = "block")]
+        block: Option<SocketAddr>,
+        /// Maximum number of concurrent connections to allow
+        #[clap(long = "connection-limit")]
+        connection_limit: Option<usize>,
+    }
 }
 
 fn config() -> &'static Config {
@@ -92,6 +124,7 @@ async fn main() -> anyhow::Result<()> {
 
     match args.command {
         Commands::CliClient => {
+            let custom_certs_fut = read_opt_server_cert();
             let host = config().get_string("remote_host")?;
             let port = config().get::<u16>("remote_port")?;
             let resolver = TokioAsyncResolver::tokio(
@@ -100,10 +133,52 @@ async fn main() -> anyhow::Result<()> {
             );
             let response = resolver.lookup_ip(&host).await?;
             let ip_addr = response.iter().next().expect("no addresses returned!");
-            let ip_sock = SocketAddr::new(ip_addr, port);
-            let client_endpoint = quic::client::make_client_endpoint(ip_sock.clone(), None).unwrap();
-            let client = client_endpoint.connect(ip_sock, &host)?.await?;
+            info!("Remote host {} resolved to {:?}", host, ip_addr);
+            let server_sock = SocketAddr::new(ip_addr, port);
+            let client_sock = if server_sock.is_ipv6() {
+                SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0,0,0,0,0,0,0,1)), 0)
+            } else {
+                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127,0,0,1)), 0)
+            };
+            info!("Client address: {:?}", client_sock);
+            let custom_certs = custom_certs_fut.await?;
+            let client_endpoint = quic::client::make_client_endpoint(client_sock.clone(), custom_certs)?;
+            let client = client_endpoint.connect(server_sock, &host)?.await?;
             println!("[client] connected: addr={}", client.remote_address());
+            loop {
+                let stream = client.open_bi().await;
+                let (mut send, mut recv) = match stream {
+                    Err(quinn::ConnectionError::ApplicationClosed { .. }) => {
+                        info!("connection closed");
+                        return Ok(());
+                    }
+                    Err(e) => {
+                        return bail!(e);
+                    }
+                    Ok(s) => s,
+                };
+                send.write_all("ping".as_bytes()).await.expect("TODO: panic message");
+                send.finish()?;
+
+                let response_start = Instant::now();
+                let resp = recv
+                    .read_to_end(usize::MAX)
+                    .await
+                    .map_err(|e| anyhow!("failed to read response: {}", e))?;
+
+                let duration = response_start.elapsed();
+                eprintln!(
+                    "response received in {:?} - {} KiB/s",
+                    duration,
+                    resp.len() as f32 / (duration_secs(&duration) * 1024.0)
+                );
+                std::io::stdout().write_all(&resp).unwrap();
+                std::io::stdout().flush().unwrap();
+                client.close(0u32.into(), b"done");
+
+                // Give the server a fair chance to receive the close packet
+                client_endpoint.wait_idle().await;
+            }
         }
 
         Commands::Shadow => {
@@ -131,9 +206,120 @@ async fn main() -> anyhow::Result<()> {
                 }
             }
         }
+
+        Commands::Server {
+            keylog, key, cert, stateless_retry, listen, block, connection_limit
+        } => {
+            let (certs, key) = if let (Some(key_path), Some(cert_path)) = (&key, &cert) {
+                let key = std::fs::read(key_path).context("failed to read private key")?;
+                let key = if key_path.extension().is_some_and(|x| x == "der") {
+                    PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(key))
+                } else {
+                    rustls_pemfile::private_key(&mut &*key)
+                        .context("malformed PKCS #1 private key")?
+                        .ok_or_else(|| anyhow::Error::msg("no private keys found"))?
+                };
+                let cert_chain = std::fs::read(cert_path).context("failed to read certificate chain")?;
+                let cert_chain = if cert_path.extension().is_some_and(|x| x == "der") {
+                    vec![CertificateDer::from(cert_chain)]
+                } else {
+                    rustls_pemfile::certs(&mut &*cert_chain)
+                        .collect::<Result<_, _>>()
+                        .context("invalid PEM-encoded certificate")?
+                };
+
+                (cert_chain, key)
+            } else {
+                let default_dir = dirs::data_dir().unwrap();
+                let path = default_dir.join("hai-server");
+                let cert_path = path.join("cert.der");
+                let key_path = path.join("key.der");
+                let (cert, key) = match std::fs::read(&cert_path).and_then(|x| Ok((x, std::fs::read(&key_path)?))) {
+                    Ok((cert, key)) => (
+                        CertificateDer::from(cert),
+                        PrivateKeyDer::try_from(key).map_err(anyhow::Error::msg)?,
+                    ),
+                    Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
+                        info!("generating self-signed certificate");
+                        let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
+                        let key = PrivatePkcs8KeyDer::from(cert.key_pair.serialize_der());
+                        let cert = cert.cert.into();
+                        std::fs::create_dir_all(path).context("failed to create certificate directory")?;
+                        std::fs::write(&cert_path, &cert).context("failed to write certificate")?;
+                        std::fs::write(&key_path, key.secret_pkcs8_der())
+                            .context("failed to write private key")?;
+                        (cert, key.into())
+                    }
+                    Err(e) => {
+                        bail!("failed to read certificate: {}", e);
+                    }
+                };
+
+                (vec![cert], key)
+            };
+            let mut server_crypto = rustls::ServerConfig::builder()
+                //TODO change to auth
+                .with_no_client_auth()
+                .with_single_cert(certs, key)?;
+            if keylog {
+                server_crypto.key_log = Arc::new(rustls::KeyLogFile::new());
+            }
+
+            let mut server_config =
+                quinn::ServerConfig::with_crypto(Arc::new(QuicServerConfig::try_from(server_crypto)?));
+            let transport_config = Arc::get_mut(&mut server_config.transport).unwrap();
+            transport_config.max_concurrent_uni_streams(0_u8.into());
+
+            let endpoint = quinn::Endpoint::server(server_config, listen)?;
+            eprintln!("listening on {}", endpoint.local_addr().unwrap());
+
+            while let Some(conn) = endpoint.accept().await {
+                if connection_limit
+                    .is_some_and(|n| endpoint.open_connections() >= n)
+                {
+                    info!("refusing due to open connection limit");
+                    conn.refuse();
+                } else if Some(conn.remote_address()) == block {
+                    info!("refusing blocked client IP address");
+                    conn.refuse();
+                } else if stateless_retry && !conn.remote_address_validated() {
+                    info!("requiring connection to validate its address");
+                    conn.retry().unwrap();
+                } else {
+                    info!("accepting connection");
+                    let fut = handle_connection(conn);
+                    tokio::spawn(async move {
+                        if let Err(e) = fut.await {
+                            error!("connection failed: {reason}", reason = e.to_string())
+                        }
+                    });
+                }
+            }
+        }
     };
 
     //handling anything here for gui wont work
     println!("exit");
     Ok(())
 }
+
+async fn read_opt_server_cert() -> anyhow::Result<Option<Vec<u8>>> {
+    let cert_path_res = config().get_string("remote_cert");
+    if cert_path_res.is_err() {
+        return Ok(None);
+    }
+    let cert_path = cert_path_res?;
+    let res = tokio::fs::read(cert_path).await;
+    if let Err(e) = res {
+        error!("failed to read certificate: {reason}", reason = e.to_string());
+        Err(Error::from(e))
+    } else {
+        info!("read certificate successfully read");
+        Ok(Some(res?))
+    }
+}
+
+
+fn duration_secs(x: &Duration) -> f32 {
+    x.as_secs() as f32 + x.subsec_nanos() as f32 * 1e-9
+}

src/quic/client.rs

@@ -1,6 +1,8 @@
 use std::error::Error;
 use std::net::SocketAddr;
 use std::sync::Arc;
+use anyhow::anyhow;
+use log::info;
 use quinn::{ClientConfig, Endpoint};
 use rustls::pki_types::CertificateDer;
 
@@ -10,13 +12,11 @@ use rustls::pki_types::CertificateDer;
 ///
 /// - server_certs: a list of trusted certificates in DER format.
 fn configure_client(
-    server_certs: Option<&[&[u8]]>,
-) -> Result<ClientConfig, Box<dyn Error + Send + Sync + 'static>> {
+    server_certs: Option<Vec<u8>>,
+) -> anyhow::Result<ClientConfig> {
     if let Some(server_certs) = server_certs {
         let mut certs = rustls::RootCertStore::empty();
-        for cert in server_certs {
-            certs.add(CertificateDer::from(*cert))?;
-        }
+        certs.add(CertificateDer::from(server_certs.as_slice()))?;
 
         Ok(ClientConfig::with_root_certificates(Arc::new(certs))?)
     } else {
@@ -32,10 +32,35 @@ fn configure_client(
 #[allow(unused)]
 pub fn make_client_endpoint(
     bind_addr: SocketAddr,
-    server_certs: Option<&[&[u8]]>,
-) -> Result<Endpoint, Box<dyn Error + Send + Sync + 'static>> {
+    server_certs: Option<Vec<u8>>,
+) -> anyhow::Result<Endpoint> {
     let client_cfg = configure_client(server_certs)?;
     let mut endpoint = Endpoint::client(bind_addr)?;
     endpoint.set_default_client_config(client_cfg);
     Ok(endpoint)
 }
+
+pub(crate) async fn handle_request(
+    (mut send, mut recv): (quinn::SendStream, quinn::RecvStream),
+) -> anyhow::Result<()> {
+    let req = recv
+        .read_to_end(64 * 1024)
+        .await
+        .map_err(|e| anyhow!("failed reading request: {}", e))?;
+    let mut escaped = String::new();
+    for &x in &req[..] {
+        let part = std::ascii::escape_default(x).collect::<Vec<_>>();
+        escaped.push_str(std::str::from_utf8(&part)?);
+    }
+    tracing::info!(content = %escaped);
+    // Execute the request
+    let resp = "ping {}";
+    // Write the response
+    send.write_all(&resp.as_bytes())
+        .await
+        .map_err(|e| anyhow!("failed to send response: {}", e))?;
+    // Gracefully terminate the stream
+    send.finish()?;
+    info!("complete");
+    Ok(())
+}

src/quic/server.rs

@@ -1,9 +1,13 @@
 use std::error::Error;
+use std::fmt::Debug;
 use std::net::SocketAddr;
 use std::sync::Arc;
+use anyhow::anyhow;
+use log::{error, info};
 use quinn::{Endpoint, ServerConfig};
 use rustls::pki_types::{CertificateDer, PrivatePkcs8KeyDer};
 use rustls::pki_types::pem::PemObject;
+use tracing::{info_span, Instrument};
 
 /// Constructs a QUIC endpoint configured to listen for incoming connections on a certain address
 /// and port.
@@ -46,3 +50,72 @@ fn configure_server(
 
     Ok((server_config, cert_der))
 }
+
+pub async fn handle_connection(conn: quinn::Incoming) -> anyhow::Result<()> {
+    let connection = conn.await?;
+    let span = info_span!(
+        "connection",
+        remote = %connection.remote_address(),
+        protocol = %connection
+            .handshake_data()
+            .unwrap()
+            .downcast::<quinn::crypto::rustls::HandshakeData>().unwrap()
+            .protocol
+            .map_or_else(|| "<none>".into(), |x| String::from_utf8_lossy(&x).into_owned())
+    );
+    async {
+        info!("established");
+
+        // Each stream initiated by the client constitutes a new request.
+        loop {
+            let stream = connection.accept_bi().await;
+            let stream = match stream {
+                Err(quinn::ConnectionError::ApplicationClosed { .. }) => {
+                    info!("connection closed");
+                    return Ok(());
+                }
+                Err(e) => {
+                    return Err(e);
+                }
+                Ok(s) => s,
+            };
+            let fut = handle_request(stream);
+            tokio::spawn(
+                async move {
+                    if let Err(e) = fut.await {
+                        error!("failed: {reason}", reason = e.to_string());
+                    }
+                }
+                    .instrument(info_span!("request")),
+            );
+        }
+    }
+        .instrument(span)
+        .await?;
+    Ok(())
+}
+
+pub(crate) async fn handle_request(
+    (mut send, mut recv): (quinn::SendStream, quinn::RecvStream),
+) -> anyhow::Result<()> {
+    let req = recv
+        .read_to_end(64 * 1024)
+        .await
+        .map_err(|e| anyhow!("failed reading request: {}", e))?;
+    let mut escaped = String::new();
+    for &x in &req[..] {
+        let part = std::ascii::escape_default(x).collect::<Vec<_>>();
+        escaped.push_str(std::str::from_utf8(&part)?);
+    }
+    tracing::info!(content = %escaped);
+    // Execute the request
+    let resp = "Hello World {}";
+    // Write the response
+    send.write_all(&resp.as_bytes())
+        .await
+        .map_err(|e| anyhow!("failed to send response: {}", e))?;
+    // Gracefully terminate the stream
+    send.finish()?;
+    info!("complete");
+    Ok(())
+}