proxmox-mtls-gateway/src/main.rs

use std::collections::HashMap;

use async_trait::async_trait;
use clap::Parser;
use log::{debug, info, warn};
use pingora_core::server::configuration::Opt;
use pingora_core::Result;
use pingora_core::{protocols::ALPN, server::Server, upstreams::peer::Peer};
use pingora_core::{tls::ssl::SslVerifyMode, upstreams::peer::HttpPeer};
use pingora_http::{RequestHeader, ResponseHeader};
use pingora_proxy::{ProxyHttp, Session};
use reqwest::{Client, Url};
use serde::{Deserialize, Serialize};

/*
 * Configuration file
 */

#[derive(Debug, Deserialize)]
pub struct Config {
    listen_addr: String,
    tls: TlsConfig,
    proxmox: ProxmoxConfig,
    users: HashMap<String, UserConfig>,
}

#[derive(Debug, Deserialize)]
pub struct TlsConfig {
    ca_file: String,
    cert_file: String,
    key_file: String,
}

#[derive(Debug, Deserialize)]
pub struct ProxmoxConfig {
    host: String,
    port: u16,
}

#[derive(Debug, Deserialize, Clone)]
pub struct UserConfig {
    username: String,
    password: String,
    realm: String,
}

#[derive(Serialize)]
struct ProxmoxLoginReq<'a> {
    username: &'a str,
    password: &'a str,
    realm: &'a str,
}

/*
 * ProxmoxVE API
 */

#[derive(Deserialize, Debug)]
struct ProxmoxRes<T> {
    data: T,
}

#[derive(Deserialize, Debug, Clone)]
pub struct ProxmoxSession {
    pub username: String,
    #[serde(rename = "CSRFPreventionToken")]
    pub csrf_token: String,
    pub ticket: String,
}

pub struct MTLSGateway {
    config: Config,
}

pub struct Context {
    subject_name: String,
    user: UserConfig,
    session: Option<ProxmoxSession>,
}

#[async_trait]
impl ProxyHttp for MTLSGateway {
    type CTX = Context;
    fn new_ctx(&self) -> Self::CTX {
        Context {
            subject_name: "".to_string(),
            user: UserConfig {
                username: "".to_string(),
                password: "".to_string(),
                realm: "".to_string(),
            },
            session: None,
        }
    }

    async fn request_filter(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<bool> {
        if let Some(cert) = &session
            .downstream_session
            .digest()
            .and_then(|x| x.ssl_digest.as_ref())
            .map(|x| &x.certificate)
        {
            let subjects = cert
                .subject_name()
                .entries()
                .filter_map(|x| x.data().as_utf8().ok())
                .map(|x| x.to_string());
            if let Some((subject_name, user_config)) = subjects
                .filter_map(|x| self.config.users.get(&x).map(|y| (x, y)))
                .next()
            {
                ctx.user = user_config.clone();
                ctx.subject_name = subject_name;
                Ok(false)
            } else {
                session.respond_error(403 /* FORBIDDEN */).await?;
                return Ok(true);
            }
        } else {
            session.respond_error(401 /* UNAUTHORIZED */).await?;
            return Ok(true);
        }
    }

    async fn upstream_peer(
        &self,
        _session: &mut Session,
        _ctx: &mut Self::CTX,
    ) -> Result<Box<HttpPeer>> {
        let mut peer = HttpPeer::new(
            (self.config.proxmox.host.as_str(), self.config.proxmox.port),
            true,
            self.config.proxmox.host.clone(),
        );

        peer.get_mut_peer_options().unwrap().alpn = ALPN::H2;

        Ok(Box::new(peer))
    }

    async fn upstream_request_filter(
        &self,
        _session: &mut Session,
        upstream_request: &mut RequestHeader,
        ctx: &mut Self::CTX,
    ) -> Result<()> {
        if !upstream_request.headers.contains_key("Cookie") {
            info!(
                "Bootstrapping session for {} using {}@{}",
                ctx.subject_name, ctx.user.username, ctx.user.realm
            );
            let session = login(&ctx.user, &self.config).await.unwrap();
            upstream_request.insert_header(
                "Cookie",
                format!("PVEAuthCookie={}", urlencoding::encode(&session.ticket)),
            )?;
            ctx.session = Some(session);
        }

        Ok(())
    }

    async fn response_filter(
        &self,
        _session: &mut Session,
        upstream_response: &mut ResponseHeader,
        ctx: &mut Self::CTX,
    ) -> Result<()>
    where
        Self::CTX: Send + Sync,
    {
        // replace existing header if any
        upstream_response
            .insert_header("Server", "PVE mTLS Gateway")
            .unwrap();
        if let Some(session) = &ctx.session {
            upstream_response.append_header(
                "Set-Cookie",
                format!(
                    "PVEAuthCookie={};Path=/;SameSite=Strict;Secure",
                    urlencoding::encode(&session.ticket)
                ),
            )?;
        }

        // because we don't support h3
        upstream_response.remove_header("alt-svc");

        Ok(())
    }

    async fn logging(
        &self,
        session: &mut Session,
        _e: Option<&pingora_core::Error>,
        ctx: &mut Self::CTX,
    ) {
        let response_code = session
            .response_written()
            .map_or(0, |resp| resp.status.as_u16());
        debug!(
            "{} response code: {response_code}",
            self.request_summary(session, ctx)
        );
    }
}

fn main() {
    env_logger::init();

    let config = std::fs::read_to_string(std::env::var("CONFIG").expect("config file path"))
        .expect("user mapping file");
    let config: Config = toml::from_str(&config).expect("valid user mapping");
    let listen_addr = config.listen_addr.clone();

    // read command line arguments
    let opt = Opt::parse();
    let mut my_server = Server::new(Some(opt)).unwrap();
    my_server.bootstrap();

    let mut tls_settings = pingora_core::listeners::TlsSettings::intermediate(
        &config.tls.cert_file,
        &config.tls.key_file,
    )
    .unwrap();
    tls_settings.enable_h2();
    tls_settings.set_verify(SslVerifyMode::all());
    tls_settings.set_ca_file(&config.tls.ca_file).unwrap();

    let mut my_proxy =
        pingora_proxy::http_proxy_service(&my_server.configuration, MTLSGateway { config });
    my_proxy.add_tls_with_settings(&listen_addr, None, tls_settings);

    my_server.add_service(my_proxy);
    my_server.run_forever()
}

fn proxmox_url(path: &str, config: &Config) -> Result<Url, url::ParseError> {
    Url::parse(&format!(
        "https://{}:{}{}",
        config.proxmox.host, config.proxmox.port, path
    ))
}
async fn login(user: &UserConfig, config: &Config) -> Result<ProxmoxSession, reqwest::Error> {
    let client = Client::new();
    let res = client
        .post(proxmox_url("/api2/json/access/ticket", config).unwrap())
        .json(&ProxmoxLoginReq {
            username: &user.username,
            password: &user.password,
            realm: &user.realm,
        })
        .send()
        .await?
        .error_for_status()?;

    let sess = res.json::<ProxmoxRes<ProxmoxSession>>().await?.data;
    Ok(sess)
}