path: root/src/main.rs

#![feature(try_trait_v2)]
#![feature(exclusive_range_pattern)]

pub mod config;
pub mod error;
pub mod files;
pub mod proxy;

use crate::{
    config::{Config, HostConfig},
    files::serve_files,
    proxy::proxy_request,
};
use anyhow::{anyhow, bail, Context, Result};
use error::ServiceError;
use http_body_util::{combinators::BoxBody, BodyExt};
use hyper::{
    body::Incoming,
    header::{CONTENT_TYPE, HOST, SERVER},
    http::HeaderValue,
    server::conn::http1,
    service::service_fn,
    Request, Response, StatusCode,
};
use log::{debug, error, info, warn};
use std::{
    fs::File,
    io::BufReader,
    net::SocketAddr,
    path::Path,
    process::exit,
    sync::{
        atomic::{AtomicUsize, Ordering::Relaxed},
        Arc,
    },
};
use tokio::{
    io::{AsyncRead, AsyncWrite},
    net::TcpListener,
    signal::ctrl_c,
};
use tokio_rustls::TlsAcceptor;

pub struct State {
    pub config: Config,
    pub total_connection: AtomicUsize,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    env_logger::init_from_env("LOG");

    let config_path = std::env::args().skip(1).next().ok_or(anyhow!(
        "first argument is expected to be the configuration file"
    ))?;

    let config = match Config::load(&config_path) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("error {e:?}");
            exit(1);
        }
    };
    let state = Arc::new(State {
        config,
        total_connection: AtomicUsize::new(0),
    });

    {
        let state = state.clone();
        tokio::spawn(async move {
            if let Err(e) = serve_http(state).await {
                error!("{e}");
                exit(1)
            }
        });
    }
    {
        let state = state.clone();
        tokio::spawn(async move {
            if let Err(e) = serve_https(state).await {
                error!("{e}");
                exit(1)
            }
        });
    }

    ctrl_c().await.unwrap();
    Ok(())
}

async fn serve_http(state: Arc<State>) -> Result<()> {
    let http_config = match &state.config.http {
        Some(n) => n,
        None => return Ok(()),
    };
    let listener = TcpListener::bind(http_config.bind).await?;
    info!("serving http");
    loop {
        let (stream, addr) = listener.accept().await.context("accepting connection")?;
        debug!("connection from {addr}");
        let config = state.clone();
        tokio::spawn(async move { serve_stream(config, stream, addr).await });
    }
}

async fn serve_https(state: Arc<State>) -> Result<()> {
    let https_config = match &state.config.https {
        Some(n) => n,
        None => return Ok(()),
    };
    let tls_config = {
        let certs = load_certs(&https_config.tls_cert)?;
        let key = load_private_key(&https_config.tls_key)?;
        let mut cfg = rustls::ServerConfig::builder()
            .with_safe_defaults()
            .with_no_client_auth()
            .with_single_cert(certs, key)?;
        cfg.alpn_protocols = vec![
            //b"h2".to_vec(),
            b"http/1.1".to_vec(),
        ];
        Arc::new(cfg)
    };
    let listener = TcpListener::bind(https_config.bind).await?;
    let tls_acceptor = Arc::new(TlsAcceptor::from(tls_config));

    info!("serving https");
    loop {
        let (stream, addr) = listener.accept().await.context("accepting connection")?;
        let state = state.clone();
        let tls_acceptor = tls_acceptor.clone();
        tokio::task::spawn(async move {
            debug!("connection from {addr}");
            match tls_acceptor.accept(stream).await {
                Ok(stream) => serve_stream(state, stream, addr).await,
                Err(e) => warn!("error accepting tls: {e}"),
            };
        });
    }
}

pub async fn serve_stream<T: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
    state: Arc<State>,
    stream: T,
    addr: SocketAddr,
) {
    let conns = state.total_connection.fetch_add(1, Relaxed);

    if conns >= state.config.limits.max_incoming_connections {
        let conn = http1::Builder::new()
            .serve_connection(
                stream,
                service_fn(|req| {
                    let state = state.clone();
                    async move {
                        match service(state, req, addr).await {
                            Ok(r) => Ok(r),
                            Err(ServiceError::Hyper(e)) => Err(e),
                            Err(error) => Ok({
                                let mut resp =
                                    Response::new(format!("gnix encountered an issue: {error}"));
                                *resp.status_mut() = StatusCode::BAD_REQUEST;
                                resp.headers_mut()
                                    .insert(CONTENT_TYPE, HeaderValue::from_static("text/plain"));
                                resp
                            }
                            .map(|b| b.map_err(|e| match e {}).boxed())),
                        }
                    }
                }),
            )
            .with_upgrades();
        if let Err(err) = conn.await {
            warn!("error: {:?}", err);
        }
    } else {
        warn!("connection dropped: too many incoming");
    }

    state.total_connection.fetch_sub(1, Relaxed);
}

fn load_certs(path: &Path) -> anyhow::Result<Vec<rustls::Certificate>> {
    let mut reader = BufReader::new(File::open(path).context("reading tls certs")?);
    let certs = rustls_pemfile::certs(&mut reader).context("parsing tls certs")?;
    Ok(certs.into_iter().map(rustls::Certificate).collect())
}
fn load_private_key(path: &Path) -> anyhow::Result<rustls::PrivateKey> {
    let mut reader = BufReader::new(File::open(path).context("reading tls private key")?);
    let keys =
        rustls_pemfile::pkcs8_private_keys(&mut reader).context("parsing tls private key")?;
    if keys.len() != 1 {
        bail!("expected a single private key, found {}", keys.len())
    }
    Ok(rustls::PrivateKey(keys[0].clone()))
}

async fn service(
    state: Arc<State>,
    req: Request<Incoming>,
    addr: SocketAddr,
) -> Result<hyper::Response<BoxBody<bytes::Bytes, ServiceError>>, ServiceError> {
    debug!("{addr} ~> {:?} {}", req.headers().get(HOST), req.uri());

    let route = state
        .config
        .hosts
        .get(remove_port(
            &req.headers()
                .get(HOST)
                .and_then(|e| e.to_str().ok())
                .map(String::from)
                .unwrap_or(String::from("")),
        ))
        .ok_or(ServiceError::NoHost)?;

    let mut resp = match route {
        HostConfig::Backend { backend } => proxy_request(req, addr, backend).await,
        HostConfig::Files { files } => serve_files(req, files).await,
    }?;

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

    return Ok(resp);
}

pub fn remove_port(s: &str) -> &str {
    s.split_once(":").map(|(s, _)| s).unwrap_or(s)
}