OpenPXE/crates/core/src/branding.rs

//! Operator-controlled branding overrides.
//!
//! The browser tab's logo (`/assets/logo.svg`) defaults to the bundled
//! rainbow-horizon mark. Operators who deploy OpenPXE behind their own
//! branding can upload a replacement that lives at
//! `<work_dir>/branding/logo.<ext>` and is served in preference to the
//! bundled SVG when present. Borrowed-from-FleetDM: tenant chrome, same
//! product.
//!
//! Storage policy mirrors `HostBindings` / `BootLog`: in-memory cache is
//! authoritative for the current process, disk is the source of truth on
//! restart, and a corrupt cache file falls back to the bundled default
//! rather than blocking startup.

use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::sync::Arc;

/// Allowed MIME types for an uploaded logo. We deliberately keep this
/// narrow — anything that can be `<img src="...">`'d into the brand
/// block, no scripts. SVG carries the obvious XSS risk for raw inline
/// HTML; we always serve the bytes as a separate asset with a strict
/// content-type rather than inlining, so SVG is safe.
pub const ALLOWED_LOGO_MIMES: &[&str] = &[
    "image/svg+xml",
    "image/png",
    "image/jpeg",
    "image/webp",
    "image/gif",
];

/// Disk cap for an uploaded logo. PXE WebUIs are operator-facing — even
/// a generous 2 MB cap is comfortable for any reasonable brand mark and
/// puts a clear bound on memory + serialization cost.
pub const MAX_LOGO_BYTES: usize = 2 * 1024 * 1024;

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
struct Inner {
    /// File name (relative to the branding dir) for the active logo, if
    /// any. Always under `<work_dir>/branding/`; never an absolute path
    /// from the operator.
    logo_filename: Option<String>,
    /// MIME of the active logo, mirroring `logo_filename`. Cached here
    /// so the HTTP layer can set Content-Type without re-sniffing.
    logo_mime: Option<String>,
    /// Monotonic counter bumped on every set/clear. Surfaces as a
    /// cache-bust token (`/assets/logo.svg?r=<rev>`) so the browser
    /// fetches the new bytes the moment the operator swaps the logo —
    /// the app version alone can't do this since it doesn't change on
    /// upload. Persisted so the token stays stable across restarts and
    /// keeps climbing across multiple swaps.
    #[serde(default)]
    rev: u64,
}

/// In-memory + on-disk override registry. Cheap to clone; locks are
/// brief. The `branding.json` cache lives alongside the active asset
/// inside `<work_dir>/branding/`.
#[derive(Debug, Clone)]
pub struct BrandingStore {
    /// Root directory: `<work_dir>/branding/`. Created on first write.
    dir: Arc<PathBuf>,
    inner: Arc<RwLock<Inner>>,
}

impl BrandingStore {
    /// Load (or initialise empty) from `<work_dir>/branding/`. Tolerates
    /// missing directories, partial state, and corrupt JSON — a bad
    /// cache should never block PXE for the network.
    #[must_use]
    pub fn load_or_default(work_dir: &Path) -> Self {
        let dir = work_dir.join("branding");
        let path = dir.join("branding.json");
        let mut inner = Inner::default();
        if let Ok(text) = std::fs::read_to_string(&path) {
            match serde_json::from_str::<Inner>(&text) {
                Ok(parsed) => {
                    // Sanity: if the JSON says we have a logo but the
                    // file is gone, clear the in-memory pointer so
                    // /assets/logo.svg falls back to the bundled SVG
                    // rather than 500ing on a missing file.
                    if let Some(name) = parsed.logo_filename.as_deref() {
                        if dir.join(name).is_file() {
                            inner = parsed;
                        } else {
                            tracing::warn!(
                                target: "openpxe::branding",
                                file = %name,
                                "branding.json points at missing file; clearing"
                            );
                        }
                    } else {
                        inner = parsed;
                    }
                }
                Err(e) => {
                    tracing::warn!(
                        target: "openpxe::branding",
                        "branding.json present but unreadable ({e}); starting empty"
                    );
                }
            }
        }
        Self {
            dir: Arc::new(dir),
            inner: Arc::new(RwLock::new(inner)),
        }
    }

    /// Absolute path to the active logo, if one is set and present on
    /// disk. `None` means the HTTP layer should serve the bundled SVG.
    #[must_use]
    pub fn logo_path(&self) -> Option<PathBuf> {
        let g = self.inner.read();
        g.logo_filename.as_deref().map(|n| self.dir.join(n))
    }

    /// MIME of the active logo, if any. The HTTP layer pairs this with
    /// the bytes returned by [`Self::logo_path`].
    #[must_use]
    pub fn logo_mime(&self) -> Option<String> {
        self.inner.read().logo_mime.clone()
    }

    /// Replace the active logo. Returns the chosen on-disk filename so
    /// the caller can echo it back in the API response. Old logos are
    /// removed best-effort.
    pub fn set_logo(&self, mime: &str, ext: &str, bytes: &[u8]) -> std::io::Result<String> {
        std::fs::create_dir_all(self.dir.as_path())?;
        // Single canonical filename per upload — overwriting the old one
        // (after clearing it) keeps the directory tidy and avoids any
        // path-traversal concern: the operator never supplies the name.
        let safe_ext = sanitize_ext(ext);
        let filename = format!("logo.{safe_ext}");
        let final_path = self.dir.join(&filename);
        // Atomic write: tmp -> rename. Guarantees the file is either
        // entirely the old logo or entirely the new one.
        let tmp = final_path.with_extension(format!("{safe_ext}.tmp"));
        std::fs::write(&tmp, bytes)?;
        std::fs::rename(&tmp, &final_path)?;
        // Clean up any sibling logo.<otherext> so there's exactly one
        // canonical file at any time.
        if let Ok(entries) = std::fs::read_dir(self.dir.as_path()) {
            for e in entries.flatten() {
                let p = e.path();
                let name = p
                    .file_name()
                    .and_then(|s| s.to_str())
                    .unwrap_or("");
                if name.starts_with("logo.") && name != filename {
                    let _ = std::fs::remove_file(&p);
                }
            }
        }

        {
            let mut g = self.inner.write();
            g.logo_filename = Some(filename.clone());
            g.logo_mime = Some(mime.to_string());
            g.rev = g.rev.wrapping_add(1);
        }
        self.persist();
        tracing::info!(
            target: "openpxe::branding",
            file = %filename, mime = %mime, size = bytes.len(),
            "custom logo installed"
        );
        Ok(filename)
    }

    /// Drop the override and return to the bundled SVG.
    pub fn clear_logo(&self) -> std::io::Result<()> {
        let removed = {
            let mut g = self.inner.write();
            let removed = g.logo_filename.take();
            g.logo_mime = None;
            g.rev = g.rev.wrapping_add(1);
            removed
        };
        if let Some(name) = removed {
            let p = self.dir.join(&name);
            let _ = std::fs::remove_file(&p);
            tracing::info!(target: "openpxe::branding", file = %name, "custom logo cleared");
        }
        self.persist();
        Ok(())
    }

    /// Convenience: true if a custom logo is configured. Surfaces on
    /// `/api/status` so the WebUI can show "Custom logo: yes" without
    /// fetching the asset itself.
    #[must_use]
    pub fn has_logo(&self) -> bool {
        self.inner.read().logo_filename.is_some()
    }

    /// Cache-bust token for the logo asset URL. Changes on every
    /// set/clear so `/assets/logo.svg?r=<rev>` resolves to a fresh URL
    /// whenever the operator swaps the brand mark. Stable otherwise.
    #[must_use]
    pub fn logo_rev(&self) -> u64 {
        self.inner.read().rev
    }

    fn persist(&self) {
        let snap = self.inner.read().clone();
        let body = match serde_json::to_vec_pretty(&snap) {
            Ok(b) => b,
            Err(e) => {
                tracing::warn!(target: "openpxe::branding", "serialize branding.json: {e}");
                return;
            }
        };
        if let Err(e) = std::fs::create_dir_all(self.dir.as_path()) {
            tracing::warn!(target: "openpxe::branding", "mkdir branding/: {e}");
            return;
        }
        let path = self.dir.join("branding.json");
        let tmp = path.with_extension("json.tmp");
        if let Err(e) = std::fs::write(&tmp, body) {
            tracing::warn!(target: "openpxe::branding", "write branding.json tmp: {e}");
            return;
        }
        if let Err(e) = std::fs::rename(&tmp, &path) {
            tracing::warn!(target: "openpxe::branding", "rename branding.json: {e}");
        }
    }
}

/// Trim arbitrary operator-supplied extension strings to a small, safe
/// alphanumeric form. Anything weird collapses to `bin`. We never let
/// the extension affect the path beyond the final segment of `logo.<x>`.
fn sanitize_ext(ext: &str) -> String {
    let lc: String = ext
        .chars()
        .filter(char::is_ascii_alphanumeric)
        .map(|c| c.to_ascii_lowercase())
        .collect();
    if lc.is_empty() || lc.len() > 5 {
        "bin".into()
    } else {
        lc
    }
}

/// Pick a safe filesystem extension from a MIME type. Returns `None`
/// if the MIME isn't on the [`ALLOWED_LOGO_MIMES`] allowlist.
#[must_use]
pub fn ext_for_mime(mime: &str) -> Option<&'static str> {
    match mime {
        "image/svg+xml" => Some("svg"),
        "image/png" => Some("png"),
        "image/jpeg" => Some("jpg"),
        "image/webp" => Some("webp"),
        "image/gif" => Some("gif"),
        _ => None,
    }
}

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

    #[test]
    fn empty_after_load_when_no_branding_dir() {
        let dir = tempdir().unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        assert!(!b.has_logo());
        assert!(b.logo_path().is_none());
        assert!(b.logo_mime().is_none());
    }

    #[test]
    fn set_clear_round_trip_persists() {
        let dir = tempdir().unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        let name = b.set_logo("image/png", "png", b"\x89PNG\r\n\x1a\nfake").unwrap();
        assert_eq!(name, "logo.png");
        assert!(b.has_logo());
        assert_eq!(b.logo_mime().as_deref(), Some("image/png"));
        let p = b.logo_path().unwrap();
        assert!(p.is_file());

        // Re-open and confirm the override survives a restart.
        drop(b);
        let b2 = BrandingStore::load_or_default(dir.path());
        assert!(b2.has_logo());
        assert_eq!(b2.logo_mime().as_deref(), Some("image/png"));

        // Clear; the file goes away and has_logo flips off.
        b2.clear_logo().unwrap();
        assert!(!b2.has_logo());
        assert!(!p.exists());
    }

    #[test]
    fn replacing_logo_removes_old_extension_sibling() {
        // PNG then SVG; only the SVG should remain on disk.
        let dir = tempdir().unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        b.set_logo("image/png", "png", b"\x89PNG\r\n\x1a\nfake").unwrap();
        b.set_logo("image/svg+xml", "svg", br#"<svg xmlns="http://www.w3.org/2000/svg"/>"#).unwrap();
        let entries: Vec<_> = std::fs::read_dir(dir.path().join("branding"))
            .unwrap()
            .filter_map(|e| e.ok().map(|e| e.file_name().to_string_lossy().into_owned()))
            .collect();
        assert!(entries.iter().any(|n| n == "logo.svg"), "got {entries:?}");
        assert!(!entries.iter().any(|n| n == "logo.png"), "stale PNG left over: {entries:?}");
    }

    #[test]
    fn logo_rev_bumps_on_each_set_and_clear() {
        let dir = tempdir().unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        assert_eq!(b.logo_rev(), 0);
        b.set_logo("image/png", "png", b"\x89PNG\r\n\x1a\nfake").unwrap();
        assert_eq!(b.logo_rev(), 1);
        b.set_logo("image/png", "png", b"\x89PNG\r\n\x1a\nfake2").unwrap();
        assert_eq!(b.logo_rev(), 2);
        b.clear_logo().unwrap();
        assert_eq!(b.logo_rev(), 3);
        // Survives a restart.
        drop(b);
        let b2 = BrandingStore::load_or_default(dir.path());
        assert_eq!(b2.logo_rev(), 3);
    }

    #[test]
    fn sanitize_ext_strips_separators_and_path_chars() {
        assert_eq!(sanitize_ext("svg"), "svg");
        // Path separators and non-alphanumerics filter out, leaving just
        // letters. The remaining "etcpasswd" exceeds the 5-char cap so
        // it collapses to `bin` rather than producing `etcpa`.
        assert_eq!(sanitize_ext("../etc/passwd"), "bin");
        // Short alphanumeric strip-through stays itself.
        assert_eq!(sanitize_ext("../svg"), "svg");
        assert_eq!(sanitize_ext(""), "bin");
        assert_eq!(sanitize_ext("PNG"), "png");
        // Anything past five chars is suspicious — collapse to `bin`.
        assert_eq!(sanitize_ext("svgvvvv"), "bin");
    }

    #[test]
    fn missing_file_referenced_by_json_resolves_to_empty() {
        // If the operator nukes the file out from under the JSON cache,
        // we should silently fall back to no-override rather than
        // hanging on to a bogus path.
        let dir = tempdir().unwrap();
        let brand_dir = dir.path().join("branding");
        std::fs::create_dir_all(&brand_dir).unwrap();
        // Hand-write a branding.json claiming logo.png exists.
        let inner = Inner {
            logo_filename: Some("logo.png".into()),
            logo_mime: Some("image/png".into()),
            rev: 0,
        };
        std::fs::write(
            brand_dir.join("branding.json"),
            serde_json::to_vec_pretty(&inner).unwrap(),
        )
        .unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        assert!(!b.has_logo(), "should fall back when referenced file is missing");
    }

    #[test]
    fn ext_for_mime_only_accepts_known_types() {
        assert_eq!(ext_for_mime("image/png"), Some("png"));
        assert_eq!(ext_for_mime("image/svg+xml"), Some("svg"));
        assert_eq!(ext_for_mime("application/octet-stream"), None);
        assert_eq!(ext_for_mime("text/html"), None);
    }
}