koi-proxy 0.4.0

TLS-terminating reverse proxy with automatic certificate management
Documentation
use serde::{Deserialize, Serialize};

use koi_common::paths;

use crate::ProxyError;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, utoipa::ToSchema)]
pub struct ProxyEntry {
    pub name: String,
    pub listen_port: u16,
    pub backend: String,
    #[serde(default)]
    pub allow_remote: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
struct ProxySection {
    #[serde(default)]
    entries: Vec<ProxyEntry>,
}

pub fn config_path() -> std::path::PathBuf {
    paths::koi_data_dir().join("config.toml")
}

pub fn config_path_with_override(data_dir: Option<&std::path::Path>) -> std::path::PathBuf {
    paths::koi_data_dir_with_override(data_dir).join("config.toml")
}

pub fn load_entries() -> Result<Vec<ProxyEntry>, ProxyError> {
    load_entries_from(&config_path())
}

pub fn load_entries_with_data_dir(
    data_dir: Option<&std::path::Path>,
) -> Result<Vec<ProxyEntry>, ProxyError> {
    load_entries_from(&config_path_with_override(data_dir))
}

fn load_entries_from(path: &std::path::Path) -> Result<Vec<ProxyEntry>, ProxyError> {
    if !path.exists() {
        return Ok(Vec::new());
    }
    let raw = std::fs::read_to_string(path).map_err(|e| ProxyError::Io(e.to_string()))?;
    let value: toml::Value = raw
        .parse()
        .map_err(|e| ProxyError::Config(format!("Invalid config.toml: {e}")))?;
    let proxy = value
        .get("proxy")
        .cloned()
        .unwrap_or_else(|| toml::Value::Table(toml::map::Map::new()));
    let proxy: ProxySection = proxy
        .try_into()
        .map_err(|e| ProxyError::Config(format!("Invalid proxy section: {e}")))?;
    Ok(proxy.entries)
}

pub fn save_entries(entries: &[ProxyEntry]) -> Result<(), ProxyError> {
    save_entries_to(entries, &config_path())
}

fn save_entries_to(entries: &[ProxyEntry], path: &std::path::Path) -> Result<(), ProxyError> {
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent).map_err(|e| ProxyError::Io(e.to_string()))?;
    }

    let mut root = if path.exists() {
        let raw = std::fs::read_to_string(path).map_err(|e| ProxyError::Io(e.to_string()))?;
        raw.parse::<toml::Value>()
            .unwrap_or_else(|_| toml::Value::Table(toml::map::Map::new()))
    } else {
        toml::Value::Table(toml::map::Map::new())
    };

    let proxy = ProxySection {
        entries: entries.to_vec(),
    };
    let proxy_value = toml::Value::try_from(proxy)
        .map_err(|e| ProxyError::Config(format!("Proxy config serialize error: {e}")))?;

    if let toml::Value::Table(table) = &mut root {
        table.insert("proxy".to_string(), proxy_value);
    }

    let raw = toml::to_string_pretty(&root)
        .map_err(|e| ProxyError::Config(format!("Config serialize error: {e}")))?;
    std::fs::write(path, raw).map_err(|e| ProxyError::Io(e.to_string()))?;
    Ok(())
}

pub fn upsert_entry(entry: ProxyEntry) -> Result<Vec<ProxyEntry>, ProxyError> {
    upsert_entry_with_data_dir(entry, None)
}

pub fn upsert_entry_with_data_dir(
    entry: ProxyEntry,
    data_dir: Option<&std::path::Path>,
) -> Result<Vec<ProxyEntry>, ProxyError> {
    let path = config_path_with_override(data_dir);
    let mut entries = load_entries_from(&path)?;
    if let Some(existing) = entries.iter_mut().find(|e| e.name == entry.name) {
        *existing = entry;
    } else {
        entries.push(entry);
    }
    entries.sort_by(|a, b| a.name.cmp(&b.name));
    save_entries_to(&entries, &path)?;
    Ok(entries)
}

pub fn remove_entry(name: &str) -> Result<Vec<ProxyEntry>, ProxyError> {
    remove_entry_with_data_dir(name, None)
}

pub fn remove_entry_with_data_dir(
    name: &str,
    data_dir: Option<&std::path::Path>,
) -> Result<Vec<ProxyEntry>, ProxyError> {
    let path = config_path_with_override(data_dir);
    let mut entries = load_entries_from(&path)?;
    let before = entries.len();
    entries.retain(|e| e.name != name);
    if entries.len() == before {
        return Err(ProxyError::NotFound(name.to_string()));
    }
    save_entries_to(&entries, &path)?;
    Ok(entries)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn config_path_is_under_data_dir() {
        let _ = koi_common::test::ensure_data_dir("koi-proxy-config-tests");
        let path = config_path();
        assert!(path.ends_with("config.toml"));
    }

    #[test]
    fn proxy_entry_round_trip() {
        let entry = ProxyEntry {
            name: "grafana".to_string(),
            listen_port: 443,
            backend: "http://localhost:3000".to_string(),
            allow_remote: false,
        };
        let proxy = ProxySection {
            entries: vec![entry.clone()],
        };
        let value = toml::Value::try_from(proxy).unwrap();
        let decoded: ProxySection = value.try_into().unwrap();
        assert_eq!(decoded.entries[0], entry);
    }
}