path: root/src/main.rs

/*
    This file is part of gnix (https://codeberg.org/metamuffin/gnix)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2025 metamuffin <metamuffin.org>
*/
#![feature(
    try_trait_v2,
    slice_split_once,
    iterator_try_collect,
    never_type,
    string_from_utf8_lossy_owned
)]

pub mod certs;
pub mod config;
pub mod control_socket;
pub mod deser_helpers;
pub mod error;
pub mod generation;
pub mod h3_support;
pub mod modules;

use crate::{
    config::{Config, ConfigPackage},
    control_socket::{
        create_unix_listener, cs_client_reload, handle_control_socket_request,
        serve_control_socket, ControlSocketRequest,
    },
    generation::Generation,
};
use aes_gcm_siv::{aead::generic_array::GenericArray, Aes256GcmSiv, KeyInit};
use anyhow::{anyhow, Context, Result};
use bytes::Bytes;
use clap::Parser;
use error::ServiceError;
use h3::server::RequestStream;
use h3_quinn::SendStream;
use h3_support::H3RequestBody;
use http::header::{CONTENT_LENGTH, TRANSFER_ENCODING};
use http_body_util::{combinators::BoxBody, BodyExt};
use hyper::{
    body::Incoming,
    header::{CONTENT_TYPE, HOST, SERVER},
    http::HeaderValue,
    service::service_fn,
    Request, Response, Uri,
};
use hyper_util::rt::{TokioExecutor, TokioIo};
use log::{debug, error, info, warn, LevelFilter};
use modules::NodeContext;
use notify::{RecursiveMode, Watcher};
use std::{
    future::Future,
    net::{IpAddr, SocketAddr},
    path::PathBuf,
    process::exit,
    str::FromStr,
    sync::Arc,
};
use tokio::{
    net::{TcpListener, UdpSocket},
    signal::ctrl_c,
    spawn,
    sync::{mpsc::channel, RwLock, Semaphore},
};
use users::{get_user_by_name, switch::set_current_uid};

pub struct State {
    pub crypto_key: Aes256GcmSiv,
    pub generation: RwLock<Arc<Generation>>,
    pub l_incoming: Semaphore,
    pub l_outgoing: Semaphore,
    pub l_incoming_h3: Semaphore,
    pub quic_endpoints: RwLock<Vec<quinn::Endpoint>>,
}

/// a simple stupid reverse proxy
#[derive(Parser)]
struct Args {
    /// Switch to specified user after loading configuration and certificates.
    #[arg(short, long)]
    user: Option<String>,
    /// Create/Use a control socket to reload configuration with.
    #[arg(short = 's', long)]
    control_socket: Option<PathBuf>,
    /// Watch the configuration file and reload on changes. Can be used with --reload.
    #[arg(short, long)]
    watch: bool,
    /// Reload configuration of another running instance. Requires --control-socket option.
    #[arg(short, long)]
    reload: bool,
    /// Path to main configuration file in YAML format
    config: PathBuf,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let args = Args::parse();

    env_logger::Builder::new()
        .filter_level(LevelFilter::Info)
        .parse_env("LOG")
        .init();

    rustls::crypto::ring::default_provider()
        .install_default()
        .unwrap();

    info!("Reading configuration and certificates...");
    let config_package = ConfigPackage::new(&args.config)?;

    if args.reload {
        let cs_path = args
            .control_socket
            .ok_or(anyhow!("reload needs control socket path"))?;
        if args.watch {
            watch_config(args.config.clone(), || async {
                let config_package = ConfigPackage::new(&args.config)?;
                cs_client_reload(&cs_path, config_package).await?;
                info!("Remote configuration updated");
                Ok(())
            })
            .await?;
            exit(1);
        } else {
            match cs_client_reload(&cs_path, config_package).await {
                Ok(()) => {
                    info!("Remote configuration updated");
                    exit(0);
                }
                Err(e) => {
                    error!("Error: {e}");
                    exit(1);
                }
            }
        }
    }

    let mut http_listeners = Vec::new();
    let mut https_listeners = Vec::new();
    let mut https_listeners_h3 = Vec::new();
    let mut control_socket_listeners = Vec::new();

    let init_config: Config =
        serde_yml::from_str(&config_package.config).context("parsing config YAML")?;
    if let Some(path) = args.control_socket {
        control_socket_listeners.push(
            create_unix_listener(&path)
                .await
                .context(anyhow!("creating Unix listener for {path:?}"))?,
        );
        info!("Control Socket listener bound to {path:?}/unix");
    }

    if let Some(http_config) = &init_config.http {
        for addr in &http_config.bind {
            http_listeners.push(
                TcpListener::bind(addr)
                    .await
                    .context(anyhow!("creating TCP listener for {addr}"))?,
            );
            info!("HTTP listener bound to {addr}/tcp");
        }
    }
    if let Some(https_config) = &init_config.https {
        for addr in &https_config.bind {
            if !https_config.disable_h1 || !https_config.disable_h2 {
                https_listeners.push(
                    TcpListener::bind(addr)
                        .await
                        .context(anyhow!("creating TCP listener for {addr}"))?,
                );
                info!("HTTPS (h1+h2) listener bound to {addr}/tcp");
            }
            if !https_config.disable_h3 {
                https_listeners_h3.push(
                    UdpSocket::bind(addr)
                        .await
                        .context(anyhow!("creating UDP listener for {addr}"))?,
                );
                info!("HTTPS (h3) listener bound to {addr}/udp");
            }
        }
    }
    drop(init_config);

    if let Some(username) = args.user {
        info!("Switching user to {username:?}...");
        let user = get_user_by_name(&username).ok_or(anyhow!("user for setuid not found"))?;
        set_current_uid(user.uid()).context("setuid")?;
    }

    info!("Initializing request handlers...");
    let generation = Generation::new(config_package)?;
    let state = Arc::new(State {
        crypto_key: aes_gcm_siv::Aes256GcmSiv::new(GenericArray::from_slice(
            &generation.config.private_key,
        )),
        quic_endpoints: RwLock::new(Vec::new()),
        l_incoming: Semaphore::new(generation.config.limits.max_incoming_connections),
        l_incoming_h3: Semaphore::new(generation.config.limits.max_incoming_connections_h3),
        l_outgoing: Semaphore::new(generation.config.limits.max_outgoing_connections),
        generation: RwLock::new(Arc::new(generation)),
    });

    if args.watch {
        let state = state.clone();
        let config = args.config.clone();
        spawn(exit_on_error(
            "Config Watch",
            watch_config(args.config.clone(), move || {
                let config = config.clone();
                let state = state.clone();
                async move {
                    let config_package = ConfigPackage::new(&config)?;
                    handle_control_socket_request(
                        state.clone(),
                        ControlSocketRequest::Config(config_package),
                    )
                    .await?;
                    info!("Local configuration updated");
                    Ok(())
                }
            }),
        ));
    }

    for li in control_socket_listeners {
        spawn(exit_on_error(
            "Control Socket",
            serve_control_socket(state.clone(), li),
        ));
    }
    for li in http_listeners {
        spawn(exit_on_error("HTTP", serve_http(state.clone(), li)));
    }
    for li in https_listeners {
        spawn(exit_on_error(
            "HTTPS (h1+h2)",
            serve_https(state.clone(), li),
        ));
    }
    for li in https_listeners_h3 {
        spawn(exit_on_error(
            "HTTPS (h3)",
            serve_https_h3(state.clone(), li),
        ));
    }
    info!("Ready to accept connections");

    ctrl_c().await.unwrap();
    exit(0)
}

async fn exit_on_error(name: &str, fut: impl Future<Output = Result<()>>) {
    if let Err(e) = fut.await {
        error!("{name}: {e:#}");
        exit(1)
    }
}

async fn serve_http(state: Arc<State>, listener: TcpListener) -> Result<()> {
    let listen_addr = listener.local_addr()?;
    loop {
        let (stream, addr) = listener.accept().await.context("accepting connection")?;
        debug!("connection from {addr}");
        let stream = TokioIo::new(stream);
        let state = state.clone();
        spawn(async move { serve_stream(state, stream, addr, false, listen_addr).await });
    }
}

async fn serve_https(state: Arc<State>, listener: TcpListener) -> Result<()> {
    let listen_addr = listener.local_addr()?;
    loop {
        let (stream, addr) = listener.accept().await.context("accepting connection")?;
        let generation = state.generation.read().await.clone();
        let Some(tls_acceptor) = generation.tls_acceptor.clone() else {
            warn!("HTTPS (h1+h2) listener is missing TLS configuration");
            break;
        };
        let state = state.clone();
        spawn(async move {
            debug!("connection from {addr}");
            match tls_acceptor.accept(stream).await {
                Ok(stream) => {
                    serve_stream(state, TokioIo::new(stream), addr, true, listen_addr).await
                }
                Err(e) => warn!("error accepting tls: {e}"),
            };
        });
    }
    info!(
        "HTTPS (h1+h2) listener for {} shutting down",
        listener.local_addr().unwrap()
    );
    Ok(())
}

async fn serve_https_h3(state: Arc<State>, listener: UdpSocket) -> Result<()> {
    let Some(quic_config) = state.generation.read().await.quic_config.clone() else {
        warn!("HTTPS (h3) listener is missing TLS configuration");
        return Ok(());
    };
    let listen_addr = listener.local_addr()?;
    let endpoint = quinn::Endpoint::new(
        quinn::EndpointConfig::default(),
        Some(quic_config),
        listener.into_std().unwrap(),
        quinn::default_runtime().unwrap(),
    )?;
    state.quic_endpoints.write().await.push(endpoint.clone());
    while let Some(conn) = endpoint.accept().await {
        let state = state.clone();
        spawn(serve_stream_h3(conn, state, listen_addr));
    }
    Ok(())
}

pub async fn serve_stream<T: Unpin + Send + 'static + hyper::rt::Read + hyper::rt::Write>(
    state: Arc<State>,
    stream: T,
    addr: SocketAddr,
    secure: bool,
    listen_addr: SocketAddr,
) {
    if let Ok(_semaphore) = state.l_incoming.try_acquire() {
        let builder = hyper_util::server::conn::auto::Builder::new(TokioExecutor::new());
        let conn = builder.serve_connection_with_upgrades(
            stream,
            service_fn(|req| {
                let state = state.clone();
                async move {
                    let req = req.map(|body: Incoming| body.map_err(ServiceError::Hyper).boxed());
                    match service(state, req, addr, secure, listen_addr).await {
                        Ok(r) => Ok(r),
                        Err(ServiceError::Hyper(e)) => Err(e),
                        Err(error) => Ok(error_response(addr, error)),
                    }
                }
            }),
        );
        if let Err(err) = conn.await {
            warn!("error: {:?}", err);
        }
    } else {
        warn!("connection dropped: too many incoming");
    }
}

async fn serve_stream_h3(conn: quinn::Incoming, state: Arc<State>, listen_addr: SocketAddr) {
    let addr = conn.remote_address();
    debug!("h3 connection attempt from {addr}");
    let Ok(_sem) = state.l_incoming_h3.try_acquire() else {
        return conn.refuse();
    };
    let conn = match conn.accept() {
        Ok(conn) => conn,
        Err(e) => return warn!("quic accep failed: {e}"),
    };
    let conn = match conn.await {
        Ok(conn) => conn,
        Err(e) => return warn!("quic connection failed: {e}"),
    };
    let mut conn =
        match h3::server::Connection::<_, Bytes>::new(h3_quinn::Connection::new(conn)).await {
            Ok(conn) => conn,
            Err(e) => return warn!("h3 accept failed {e}"),
        };
    debug!("h3 stream from {addr}");
    let generation = state.generation.read().await;
    let max_par_requests = Semaphore::new(generation.config.limits.max_requests_per_connnection);
    loop {
        match conn.accept().await {
            Ok(Some(x)) => {
                let Ok((req, stream)) = x.resolve_request().await else {
                    warn!("h3 request accept failed");
                    continue;
                };
                let Ok(_sem_req) = max_par_requests.acquire().await else {
                    warn!("h3 par request semasphore closed");
                    return;
                };
                let state = state.clone();
                spawn(async move {
                    let (mut send, recv) = stream.split();
                    let req = req.map(|()| H3RequestBody(recv).boxed());

                    let resp = service(state.clone(), req, addr, true, listen_addr)
                        .await
                        .unwrap_or_else(|error| error_response(addr, error));

                    send_h3_response(resp, &mut send).await;
                });
                drop(_sem_req)
            }
            Ok(None) => break,
            Err(e) => {
                warn!("h3 connection error: {e}");
                break;
            }
        }
    }
    drop(_sem);
}

async fn send_h3_response(
    resp: Response<BoxBody<Bytes, ServiceError>>,
    send: &mut RequestStream<SendStream<Bytes>, Bytes>,
) {
    let (parts, mut body) = resp.into_parts();
    let mut resp = Response::from_parts(parts, ());

    resp.headers_mut().remove(TRANSFER_ENCODING); // TODO allow "trailers" options
    resp.headers_mut().remove(CONTENT_LENGTH);

    if let Err(e) = send.send_response(resp).await {
        debug!("h3 response send error: {e}");
        return;
    };
    while let Some(frame) = body.frame().await {
        match frame {
            Ok(frame) => {
                if frame.is_data() {
                    let data = frame.into_data().unwrap();
                    if let Err(e) = send.send_data(data).await {
                        debug!("h3 body send error: {e}");
                        return;
                    }
                } else if frame.is_trailers() {
                    let trailers = frame.into_trailers().unwrap();
                    if let Err(e) = send.send_trailers(trailers).await {
                        debug!("h3 trailers send error: {e}");
                        return;
                    }
                }
            }
            Err(_) => todo!(),
        }
    }
    if let Err(e) = send.finish().await {
        debug!("h3 response finish error: {e}");
    }
}

fn error_response(addr: SocketAddr, error: ServiceError) -> Response<BoxBody<Bytes, ServiceError>> {
    {
        warn!("service error {addr} {error:?}");
        let mut resp = Response::new(format!(
            "Sorry, we were unable to process your request: {error}"
        ));
        *resp.status_mut() = error.status_code();
        resp.headers_mut()
            .insert(CONTENT_TYPE, HeaderValue::from_static("text/plain"));
        resp.headers_mut()
            .insert(SERVER, HeaderValue::from_static("gnix"));
        resp
    }
    .map(|b| b.map_err(|e| match e {}).boxed())
}

async fn service(
    state: Arc<State>,
    mut request: Request<BoxBody<Bytes, ServiceError>>,
    mut addr: SocketAddr,
    secure: bool,
    listen_addr: SocketAddr,
) -> Result<hyper::Response<BoxBody<bytes::Bytes, ServiceError>>, ServiceError> {
    let generation = state.generation.read().await.clone();
    // move uri authority used in HTTP/2 to Host header field
    {
        let uri = request.uri_mut();
        if let Some(authority) = uri.authority() {
            let host =
                HeaderValue::from_str(authority.host()).map_err(|_| ServiceError::InvalidUri)?;

            let mut new_uri = http::uri::Parts::default();
            new_uri.path_and_query = uri.path_and_query().cloned();
            *uri = Uri::from_parts(new_uri).map_err(|_| ServiceError::InvalidUri)?;
            request.headers_mut().insert(HOST, host);
        }
    }

    if generation.config.source_ip_from_header {
        if let Some(x) = request.headers_mut().remove("x-real-ip") {
            addr = SocketAddr::new(
                IpAddr::from_str(x.to_str()?).map_err(|_| ServiceError::InvalidHeader)?,
                0,
            );
        } else {
            return Err(ServiceError::XRealIPMissing);
        }
    }

    debug!(
        "{addr} ~> {:?} {}",
        request.headers().get(HOST),
        request.uri()
    );

    let mut context = NodeContext {
        addr,
        state,
        secure,
        listen_addr,
    };
    let mut resp = generation.handler.handle(&mut context, request).await?;

    if !generation.config.disable_server_header {
        let server_header = resp.headers().get(SERVER).cloned();
        resp.headers_mut().insert(
            SERVER,
            if let Some(o) = server_header {
                HeaderValue::from_str(&format!(
                    "{} via gnix",
                    o.to_str().ok().unwrap_or("invalid")
                ))
                .unwrap()
            } else {
                HeaderValue::from_static("gnix")
            },
        );
    }

    Ok(resp)
}

async fn watch_config<R: Future<Output = Result<()>>>(
    path: PathBuf,
    mut reload: impl FnMut() -> R,
) -> Result<()> {
    let (tx, mut rx) = channel(1);
    let mut w = notify::recommended_watcher(move |r| {
        let _ = tx.blocking_send(r);
    })?;
    w.watch(path.parent().unwrap(), RecursiveMode::NonRecursive)?;

    let path = path.canonicalize()?;
    while let Some(r) = rx.recv().await {
        match r {
            Ok(ev) => {
                if matches!(
                    ev.kind,
                    notify::EventKind::Access(notify::event::AccessKind::Close(
                        notify::event::AccessMode::Write
                    )) | notify::EventKind::Modify(notify::event::ModifyKind::Data(_))
                ) && ev.paths.contains(&path)
                {
                    if let Err(e) = reload().await {
                        error!("Error: {e:#}")
                    }
                }
            }
            Err(e) => error!("watch error: {e:#}"),
        }
    }

    Ok(())
}