trpha/src/main.rs

//! # Tuxcord Reverse Proxy Header Authenthication
#![feature(
    anonymous_lifetime_in_impl_trait,
    iterator_try_collect,
    rwlock_downgrade,
    try_blocks
)]
#![allow(clippy::missing_errors_doc)]

use std::{
    fs::File,
    io::{self, Read, Write},
    net::{self, Shutdown, TcpListener, TcpStream},
    sync::Arc,
    thread,
};

use anyhow::Context;
use clap::Parser;
use dashmap::{DashMap, mapref::one::Ref};
use thiserror::Error;

use crate::utils::headers::HeadersExt;

pub mod args;
pub mod config;
pub mod constants;
#[cfg(feature = "ipc")]
pub mod ipc;
pub mod utils;

use args::Args;

fn main() -> anyhow::Result<()> {
    let args = Args::parse();

    if !args.config.exists() {
        println!(
            "{:?} doesn't exist, creating a default config",
            &args.config
        );
        File::create(args.config)
            .context("failure creating the config file")?
            .write_all(constants::DEFAULT_CONFIG)
            .context("failure writing the contents to the config file")?;
        return Ok(());
    }

    let mut config_file = File::open(&args.config).context("failure opening the config file")?;
    let mut config = String::new();
    config_file
        .read_to_string(&mut config)
        .context("failure reading the config file")?;
    let config: config::Schema = toml::from_str(&config).context("invalid config file")?;

    #[cfg(feature = "ipc")]
    if let Some(args::Commands::Ipc) = args.command {
        ipc::start_client(config)?;
        return Ok(());
    }

    println!("config: {config:#?}");
    let listener = TcpListener::bind(config.listen_at).context("failure tcp listening")?;
    let config_arc = Arc::new(config); // will also serve as a counter

    #[cfg(feature = "ipc")]
    ipc::handle_daemon(config_arc.clone());

    for stream in listener.incoming() {
        match stream {
            Ok(mut client) => {
                let config_arc = config_arc.clone();
                thread::spawn(|| {
                    if let Err(err) = handle_client(&mut client, config_arc.as_ref()) {
                        eprintln!("err: invalid req head ({err}), closing...");
                        _ = client.shutdown(Shutdown::Both);
                    }
                    drop(client);
                    drop(config_arc);
                });
            }
            Err(err) => eprintln!("error with an incoming listener {err:#?}"),
        }
    }

    unreachable!("listener had to be killed unexpectedly");
}

#[derive(Error, Debug)]
pub enum ClientError {
    #[error("error reading request head: {0}")]
    HeadReadError(io::Error),
    #[error("failed to find \"host\" header")]
    HostHeaderNotFound,
    #[error("requested host {0:?}, but it's not registered")]
    HostNotRegistered(String),
    #[error("failed to connect to backend: {0}")]
    BackendConnectFail(io::Error),
    #[error("io error exchanging client and backend: {0}")]
    ExchangeIoError(io::Error),
    #[error("io error capturing {0} buffer: {0}")]
    BufferError(&'static str, io::Error),
    #[error("io error writing auth headers: {0}")]
    AuthHeadersIoError(io::Error),
}

// fn(req, (adrr, etc)) -> Action
// Action {
//  CaptureBody?
//  CaptureResponse?
// }
//
// cases:
// * GET login -> get response to access csrf token
// * POST login -> get body to handle authentication
//
// auth headers are just added depending on the ID cookie soo, keep a hashmap of IDs and user
// data??
//
// check authentication first of everything as it changes decisions on everything and give it to fn
// too
// fn(req, Option<AuthData>, (adrr, etc)) -> Action
//
// and also add an action to... hook in the response without even letting backend know (redirect
// ppl out and such), also in no case we need multiple actions. NOPE, POST login should capture
// body and hijack the response, I'll also just make this flexible too.

#[allow(clippy::type_complexity)]
struct ProxyAction<'a> {
    // combining like this eases a lot the implementation
    capture_body_and_hijack: Option<(
        fn(&httparse::Request, &[u8]) -> Result<(), ClientError>,
        fn(&httparse::Request, &mut TcpStream) -> Result<(), ClientError>,
    )>,
    capture_response: Option<fn(&httparse::Request, &[u8]) -> Result<(), ClientError>>,
    auth: Option<Ref<'a, String, AuthData>>,
}

struct AuthData {
    pub username: String,
    pub email: String,
    pub fullname: Option<String>,
}
impl AuthData {
    pub fn write_as_headers(&self, into: &mut impl Write) -> Result<(), ClientError> {
        use constants::HTTP_HEADERS_NEWLINE as NL;
        let r: Result<_, io::Error> = try {
            into.write_all(b"X-WEBAUTH-USER: ")?;
            into.write_all(self.username.as_bytes())?;
            into.write_all(NL)?;
            into.write_all(b"X-WEBAUTH-EMAIL: ")?;
            into.write_all(self.email.as_bytes())?;
            into.write_all(NL)?;
            if let Some(fullname) = self.fullname.as_ref() {
                into.write_all(b"X-WEBAUTH-FULLNAME: ")?;
                into.write_all(fullname.as_bytes())?;
                into.write_all(NL)?;
            }
        };
        r.map_err(ClientError::AuthHeadersIoError)
    }
}

fn what_to_do<'a>(
    req: &httparse::Request,
    id_auth: Option<(String, Option<Ref<'a, String, AuthData>>)>,
    _backend: (net::SocketAddr, String, String),
) -> ProxyAction<'a> {
    let is_authd = id_auth.as_ref().is_some_and(|id_auth| id_auth.1.is_some());
    match (req.path, req.method, is_authd) {
        (Some("/user/login"), Some("GET"), false) => todo!(),
        (Some("/user/login"), Some("POST"), false) => todo!(),
        _ => ProxyAction {
            capture_body_and_hijack: None,
            capture_response: None,
            auth: id_auth.and_then(|id_auth| id_auth.1),
        },
    }
}

// TODO: might wanna save this one between runtimes
static AUTH_MAP: std::sync::LazyLock<DashMap<String, AuthData>> =
    std::sync::LazyLock::new(DashMap::new);
static _CSRF_MAP: std::sync::LazyLock<DashMap<String, String>> =
    std::sync::LazyLock::new(DashMap::new);

fn handle_client(client: &mut TcpStream, config: &config::Schema) -> Result<(), ClientError> {
    use ClientError as E;
    use constants::DOUBLE_CRLF;

    let mut header_buf = [0u8; 1024 * 8];
    let mut read_pos = 0usize;

    let mut headers;
    let (pos, req) = loop {
        headers = [httparse::EMPTY_HEADER; 16];
        let mut req = httparse::Request::new(&mut headers);

        let n = client
            .read(&mut header_buf[read_pos..])
            .map_err(E::HeadReadError)?;
        read_pos += n;

        if let Ok(httparse::Status::Complete(n)) = req.parse(&header_buf[0..read_pos]) {
            break (n - DOUBLE_CRLF.len(), req);
        }
    };

    let theres_body = req.headers.has_any(["conten-length", "transfer-encoding"]);

    let host_header = req.headers.get("host").ok_or(E::HostHeaderNotFound)?;
    let cookies = req.headers.get("cookie");

    // Now find that header and pass everything
    let read_hosts = config.hosts.read();
    let (addr, id_keyname, csrf_keyname) = read_hosts
        .get(host_header.as_ref())
        .ok_or_else(|| E::HostNotRegistered(host_header.to_string()))?
        .clone();
    drop(read_hosts);

    // Try to find auth data
    let id_auth = if let Some(cookies) = cookies {
        if let Some(id_ck) = cookie::Cookie::split_parse(cookies)
            .filter_map(Result::ok)
            .find(|ck| ck.name() == id_keyname)
        {
            let id = id_ck.value().to_string();
            let val = AUTH_MAP.get(&id);
            Some((id, val))
        } else {
            None
        }
    } else {
        None
    };

    let action = what_to_do(&req, id_auth, (addr, id_keyname, csrf_keyname));
    if let Some((cb1, cb2)) = action.capture_body_and_hijack {
        let mut body = header_buf[pos..read_pos].to_vec();
        if theres_body {
            io::copy(client, &mut body).map_err(|err| E::BufferError("body", err))?;
        }
        cb1(&req, &body[..])?;
        cb2(&req, client)?;
    } else {
        let mut stream = TcpStream::connect(addr).map_err(E::BackendConnectFail)?;
        stream
            .write_all(&header_buf[0..pos])
            .map_err(E::ExchangeIoError)?;
        if let Some(auth) = action.auth {
            auth.write_as_headers(&mut stream)?;
        }

        // send our overhead
        stream
            .write_all(&header_buf[pos..read_pos])
            .map_err(E::ExchangeIoError)?;
        if theres_body {
            io::copy(client, &mut stream).map_err(E::ExchangeIoError)?;
        }
        if let Some(cb) = action.capture_response {
            let mut buf = Vec::new();
            stream
                .read_to_end(&mut buf)
                .map_err(|err| E::BufferError("response", err))?;
            cb(&req, &buf[..])?;
            client.write_all(&buf[..]).map_err(E::ExchangeIoError)?;
        } else {
            io::copy(&mut stream, client).map_err(E::ExchangeIoError)?;
        }
    }

    Ok(())
}