noq/interop/src/server.rs

use std::{
    ascii, cmp,
    ffi::OsStr,
    fs,
    net::{IpAddr, SocketAddr},
    path::PathBuf,
    pin::Pin,
    str, sync,
    task::{Context, Poll},
};

use anyhow::{anyhow, bail, Context as _, Result};
use bytes::Bytes;
use futures::{ready, AsyncReadExt, Future, StreamExt, TryFutureExt, future};
use http::{Response, StatusCode};
use http_body::Body as _;
use hyper::{body::HttpBody, service::{make_service_fn, service_fn}};
use structopt::{self, StructOpt};
use tokio::net::{TcpListener, TcpStream};
use tokio_rustls::{server::TlsStream, TlsAcceptor};
use tracing::{error, info, info_span};
use tracing_futures::Instrument as _;

use quinn::SendStream;
use quinn_h3::{self, server::RecvRequest, IntoBody};
use sync::Arc;

#[derive(StructOpt, Debug, Clone)]
#[structopt(name = "h3_server")]
struct Opt {
    /// TLS private key
    #[structopt(
        parse(from_os_str),
        short = "k",
        long = "key",
        requires = "cert",
        default_value = "key.der"
    )]
    key: PathBuf,
    /// TLS certificate
    #[structopt(
        parse(from_os_str),
        short = "c",
        long = "cert",
        requires = "key",
        default_value = "cert.der"
    )]
    cert: PathBuf,
    /// Address to listen on
    #[structopt(long = "listen", short = "l", default_value = "::")]
    listen: IpAddr,
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
            .finish(),
    )
    .unwrap();
    let opt = Opt::from_args();

    let key = fs::read(&opt.key).context("failed to read private key")?;
    let key = match opt.key.as_path().extension().and_then(OsStr::to_str) {
        Some("der") => quinn::PrivateKey::from_der(&key[..])?,
        _ => quinn::PrivateKey::from_pem(&key[..])?,
    };

    let cert_chain = fs::read(&opt.cert).context("failed to read certificate chain")?;
    let cert_chain = match opt.cert.as_path().extension().and_then(OsStr::to_str) {
        Some("der") => {
            quinn::CertificateChain::from_certs(vec![quinn::Certificate::from_der(&cert_chain)?])
        }
        _ => quinn::CertificateChain::from_pem(&cert_chain)?,
    };

    let mut server_config = quinn::ServerConfigBuilder::default();
    server_config.certificate(cert_chain, key)?;
    server_config.protocols(&[quinn_h3::ALPN, b"hq-27"]);

    let main = server(server_config.clone(), SocketAddr::new(opt.listen, 4433));
    let default = server(server_config.clone(), SocketAddr::new(opt.listen, 443));
    server_config.use_stateless_retry(true);
    let retry = server(server_config.clone(), SocketAddr::new(opt.listen, 4434));

    tokio::try_join!(main, default, retry, h2_server(server_config.clone()))?;

    Ok(())
}

async fn server(server_config: quinn::ServerConfigBuilder, addr: SocketAddr) -> Result<()> {
    let mut transport = quinn::TransportConfig::default();
    transport.initial_window(1024 * 1024);
    transport.send_window(1024 * 1024 * 3);
    transport.receive_window(1024 * 1024);
    let mut server_config = server_config.build();
    server_config.transport = Arc::new(transport);

    let mut endpoint_builder = quinn::Endpoint::builder();
    endpoint_builder.listen(server_config);
    let (_, mut incoming) = endpoint_builder.bind(&addr)?;

    println!("server listening on {}", addr);
    while let Some(connecting) = incoming.next().await {
        tokio::spawn(async move {
            let protos = connecting.authentication_data().protocol.unwrap();
            println!("server received connection");

            if protos == b"h3-27" {
                if let Err(e) = h3_handle_connection(connecting).await {
                    error!("handling connection failed: {:?}", e)
                }
            } else if protos == b"hq-27" {
                if let Err(e) = hq_handle_connection(connecting).await {
                    error!("handling connection failed: {:?}", e)
                }
            }
        });
    }
    Ok(())
}

async fn h3_handle_connection(connecting: quinn::Connecting) -> Result<()> {
    let connecting = quinn_h3::server::Connecting::from(connecting);
    let mut incoming = connecting.await.context("accept failed")?;
    tokio::spawn(async move {
        while let Some(request) = incoming.next().await {
            tokio::spawn(async move {
                if let Err(e) = h3_handle_request(request).await {
                    eprintln!("request error: {}", e)
                }
            });
        }
    });

    Ok(())
}

async fn h3_handle_request(recv_request: RecvRequest) -> Result<()> {
    let (mut request, sender) = recv_request.await?;
    println!("received request: {:?}", request);

    let body = request.body_mut().read_to_end().await?;
    println!("received body: {}", String::from_utf8_lossy(&body));

    if let Some(trailers) = request.body_mut().trailers().await? {
        println!("received trailers: {:?}", trailers);
    }

    match request.uri().path() {
        "/" => h3_home(sender).await?,
        path => match parse_size(path) {
            Ok(n) => h3_payload(sender, n).await?,
            Err(_) => h3_home(sender).await?,
        },
    };

    Ok(())
}

async fn h3_home(sender: quinn_h3::server::Sender) -> Result<()> {
    let response = Response::builder()
        .status(StatusCode::OK)
        .header("server", VERSION)
        .body(HOME.into_body())
        .expect("failed to build response");
    sender
        .send_response(response)
        .await
        .map_err(|e| anyhow!("failed to send response: {:?}", e))?;
    Ok(())
}

async fn h3_payload(sender: quinn_h3::server::Sender, len: usize) -> Result<()> {
    if len > 1_000_000_000 {
        let response = Response::builder()
            .status(StatusCode::BAD_REQUEST)
            .header("server", VERSION)
            .body(format!("requested {}: too large", len).into_body())
            .expect("failed to build response");
        sender.send_response(response).await?;
        return Ok(());
    }

    let mut buf = TEXT.repeat(len / TEXT.len() + 1);
    buf.truncate(len);

    let response = Response::builder()
        .status(StatusCode::OK)
        .header("server", VERSION)
        .body(buf.into_body())
        .expect("failed to build response");

    sender
        .send_response(response)
        .await
        .map_err(|e| anyhow!("failed to send response: {:?}", e))?;

    Ok(())
}

async fn hq_handle_connection(conn: quinn::Connecting) -> Result<()> {
    let quinn::NewConnection {
        connection,
        mut bi_streams,
        ..
    } = conn.await?;
    let span = info_span!(
        "connection",
        remote = %connection.remote_address(),
        protocol = %connection
            .authentication_data()
            .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.
        while let Some(stream) = bi_streams.next().await {
            let stream = match stream {
                Err(quinn::ConnectionError::ApplicationClosed { .. }) => {
                    info!("connection closed");
                    return Ok(());
                }
                Err(e) => {
                    return Err(e);
                }
                Ok(s) => s,
            };
            tokio::spawn(
                hq_handle_request(stream)
                    .unwrap_or_else(move |e| error!("failed: {reason}", reason = e.to_string()))
                    .instrument(info_span!("request")),
            );
        }
        Ok(())
    }
    .instrument(span)
    .await?;
    Ok(())
}

async fn hq_handle_request((send, recv): (quinn::SendStream, quinn::RecvStream)) -> 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 = ascii::escape_default(x).collect::<Vec<_>>();
        escaped.push_str(str::from_utf8(&part).unwrap());
    }
    info!(content = %escaped);
    // Execute the request
    hq_process_get(send, &req).await?;
    Ok(())
}

async fn hq_process_get(mut send: SendStream, x: &[u8]) -> Result<()> {
    if x.len() < 4 || &x[0..4] != b"GET " {
        bail!("missing GET");
    }
    if x[4..].len() < 2 || &x[x.len() - 2..] != b"\r\n" {
        bail!("missing \\r\\n");
    }
    let x = &x[4..x.len() - 2];
    let end = x.iter().position(|&c| c == b' ').unwrap_or_else(|| x.len());
    let path = str::from_utf8(&x[..end]).context("path is malformed UTF-8")?;

    // Write the response
    match parse_size(path) {
        Ok(n) if n <= 1_000_000_000 => {
            let mut remaining = n;
            while remaining > 0 {
                let size = cmp::min(remaining, TEXT.len());
                send.write_all(&TEXT[..size])
                    .await
                    .map_err(|e| anyhow!("failed to send response: {}", e))?;
                remaining -= size;
            }
        }
        Ok(_) | Err(_) => {
            send.write_all(&HOME.as_bytes())
                .await
                .map_err(|e| anyhow!("failed to send response: {}", e))?;
        }
    }

    // Gracefully terminate the stream
    send.finish()
        .await
        .map_err(|e| anyhow!("failed to shutdown stream: {}", e))?;

    Ok(())
}

fn parse_size(literal: &str) -> Result<usize> {
    if literal.is_empty() {
        return Err(anyhow!("path empty"));
    }
    let pos = literal[1..]
        .find(|c: char| !c.is_ascii_digit())
        .map(|p| p + 1)
        .unwrap_or_else(|| literal.len());
    let num: usize = literal[1..pos]
        .parse()
        .map_err(|_| anyhow!("parse failed"))?;
    let scale = match literal[pos..].to_uppercase().as_str() {
        "K" => 1000,
        "M" => 1_000_000,
        "G" => 1_000_000_000,
        _ => 1,
    };
    Ok(num * scale)
}

const ALT_SVC: &str = "h3-27=\":443\"";

fn h2_home() -> hyper::Response<hyper::Body> {
    Response::builder()
        .status(StatusCode::OK)
        .header("Alt-Svc", ALT_SVC)
        .body(HOME.into())
        .expect("failed to build response")
}

fn h2_payload(len: usize) -> hyper::Response<hyper::Body> {
    if len > 1_000_000_000 {
        let response = Response::builder()
            .status(StatusCode::BAD_REQUEST)
            .header("Alt-Svc", ALT_SVC)
            .body(Bytes::from(format!("requested {}: too large", len)).into())
            .expect("failed to build response");
        return response;
    }

    let mut buf = TEXT.repeat(len / TEXT.len() + 1);
    buf.truncate(len);
    Response::builder()
        .status(StatusCode::OK)
        .body(buf.into())
        .expect("failed to build response")
}

async fn h2_handle(request: hyper::Request<hyper::Body>) -> Result<hyper::Response<hyper::Body>> {
    Ok(match request.uri().path() {
        "/" => h2_home(),
        path => match parse_size(path) {
            Ok(n) => h2_payload(n),
            Err(_) => h2_home(),
        },
    })
}

async fn h2_server(server_config: quinn::ServerConfigBuilder) -> Result<()> {
    let mut tls_cfg = (*server_config.build().crypto).clone();
    tls_cfg.set_protocols(&[b"h2".to_vec(), b"http/1.1".to_vec()]);
    let tls_acceptor = TlsAcceptor::from(sync::Arc::new(tls_cfg));

    let tcp = TcpListener::bind(&SocketAddr::new([0, 0, 0, 0].into(), 443)).await?;

    let service = make_service_fn(|_conn| async { Ok::<_, anyhow::Error>(service_fn(h2_handle)) });
    let server = hyper::Server::builder(HyperAcceptor::new(tcp, tls_acceptor)).serve(service);

    if let Err(e) = server.await {
        error!("server error: {}", e);
    }
    Ok(())
}

struct HyperAcceptor {
    tcp: TcpListener,
    tls: TlsAcceptor,
    handshake: Option<tokio_rustls::Accept<TcpStream>>,
}

impl HyperAcceptor {
    pub fn new(tcp: TcpListener, tls: TlsAcceptor) -> Self {
        Self {
            tls,
            tcp,
            handshake: None,
        }
    }
}

impl hyper::server::accept::Accept for HyperAcceptor {
    type Conn = TlsStream<TcpStream>;
    type Error = anyhow::Error;

    fn poll_accept(
        mut self: Pin<&mut Self>,
        cx: &mut Context,
    ) -> Poll<Option<Result<Self::Conn, Self::Error>>> {
        loop {
            match self.handshake {
                Some(ref mut h) => {
                    let conn = ready!(Pin::new(h).poll(cx))?;
                    std::mem::replace(&mut self.handshake, None);
                    return Poll::Ready(Some(Ok(conn)));
                }
                None => {
                    let (stream, _) = ready!(self.tcp.poll_accept(cx))?;
                    self.handshake = Some(self.tls.accept(stream));
                }
            }
        }
    }
}

const TEXT: &[u8] =
    b"It would be different if we could not step back and reflect on the process,\n\
    but were merely led from impulse to impulse without self- consciousness. But human\n\
    beings do not act solely on impulse. They are prudent, they reflect, they weigh\n\
    consequences, they ask whether what they are doing is worth while. Not only are their\n\
    lives full of particular choices that hang together in larger activities with temporal\n\
    structure: they also decide in the broadest terms what to pursue and what to avoid, what\n\
    the priorities among their various aims should be, and what kind of people they want to\n\
    be or become. Some men are faced with such choices by the large decisions they make from\n\
    time to time; some merely by reflection on the course their lives are taking as the product\n\
    of countless small decisions. They decide whom to marry, what profession to follow, whether\n\
    to join the Country Club, or the Resistance; or they may just wonder why they go on being\n\
    salesmen or academics or taxi drivers, and then stop thinking about it after a certain period\n\
    of inconclusive reflection.";

const HOME: &str = r##"
<html xmlns="http://www.w3.org/1999/xhtml" xml:lang="en" lang="en">
  <head>
    <title>Quinn H3 interop server</title>
  </head>
  <body>
    <h1>Welcome to the quinn-h3 interop server.</h1>
    <p>
      <strong>Draft version:</strong> draft-24<br/>
      <strong>Available tests:</strong> VHDCRZSBU3
    </p>
    <p>
      Use '/{n}' to get <i>n</i> bytes of deep thoughts.<br/>
      For example <a href="/1000000">/1000000</a>
      to get 1MB. Limit: 1GB
    </p>
    <p>Checkout our project's <a href="https://github.com/djc/quinn">repository</a>.</p>
    <p>Say hi on quickdev slack workspace at `quinn`.</p>
  </body>
</html>
"##;

const VERSION: &str = "quinn-h3:0.0.1";