awa

src/main.rs

@@ -12,17 +12,23 @@ mod quic;
 
 use std::fs::File;
 use std::net::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 anyhow::{bail, Context};
+use quinn::crypto::rustls::QuicServerConfig;
+use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer};
 use tokio::{signal, sync::mpsc::Sender};
 use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
 use trust_dns_resolver::{Resolver, TokioAsyncResolver};
+use crate::quic::ALPN_QUIC_HTTP;
+use crate::quic::server::handle_connection;
 
 shadow!(build);
 
@@ -53,6 +59,29 @@ enum Commands {
     #[command(about = "Start client as GUI")]
     GuiClient,
     Devtest,
+    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 {
@@ -131,6 +160,97 @@ 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)?;
+            server_crypto.alpn_protocols = ALPN_QUIC_HTTP.iter().map(|&x| x.into()).collect();
+            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()?);
+
+            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

src/quic/client.rs

@@ -11,7 +11,7 @@ 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>> {
+) -> anyhow::Result<ClientConfig> {
     if let Some(server_certs) = server_certs {
         let mut certs = rustls::RootCertStore::empty();
         for cert in server_certs {
@@ -33,7 +33,7 @@ fn configure_client(
 pub fn make_client_endpoint(
     bind_addr: SocketAddr,
     server_certs: Option<&[&[u8]]>,
-) -> Result<Endpoint, Box<dyn Error + Send + Sync + 'static>> {
+) -> anyhow::Result<Endpoint> {
     let client_cfg = configure_client(server_certs)?;
     let mut endpoint = Endpoint::client(bind_addr)?;
     endpoint.set_default_client_config(client_cfg);

src/quic/mod.rs

@@ -1,2 +1,4 @@
 pub(crate) mod server;
-pub(crate) mod client;
+pub(crate) mod client;
+
+pub const ALPN_QUIC_HTTP: &[&[u8]] = &[b"hq-29"];

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.
@@ -45,4 +49,73 @@ fn configure_server(
     transport_config.max_concurrent_uni_streams(0_u8.into());
 
     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(())
+}
+
+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(())
 }