continuum-proxy/src/provision.rs

//! Adding and removing printers on a running gateway, on the console's say-so.
//!
//! The access code arrives in a `provision_printer` message, is checked by
//! actually connecting, and is then written to the local store (src/store.rs).
//! It is never logged and never sent anywhere. The cloud does not keep it -
//! it relayed it - so this is the only place it exists.
//!
//! Connect first, store second: a printer that cannot be reached with the
//! code the operator typed is a typo to report now, not a row to discover
//! later in a fleet that quietly never comes up.

use std::time::Duration;

use tracing::info;

use crate::fleet;
use crate::printer::GenericPrinter;
use crate::store::Store;
use crate::uplink::SharedFleet;
use crate::uplink::protocol::ProvisionPrinter;

/// Bambu's protocol auto-detection can spend several seconds on a first
/// attempt, and an operator is watching a spinner.
const CONNECT_TIMEOUT: Duration = Duration::from_secs(30);

/// Ids end up in MQTT topics, file names (pinned certificates) and log
/// lines, so they are held to a small safe alphabet.
fn valid_id(id: &str) -> bool {
    !id.is_empty() && id.len() <= 64 && id.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
}

pub async fn add(fleet: &SharedFleet, store: &Store, request: ProvisionPrinter) -> Result<(), String> {
    if !valid_id(&request.id) {
        return Err("printer id must be 1-64 letters, digits, '-' or '_'".into());
    }
    if fleet.lock().await.iter().any(|p| p.base().id == request.id) {
        return Err(format!("this gateway already has a printer with id '{}'", request.id));
    }

    let entry = fleet::PrinterEntry {
        id: request.id,
        name: request.name,
        vendor: request.vendor,
        host: request.host,
        access_code: request.access_code,
        sn: request.sn,
        api_key: request.api_key,
        port: None,
        baud_rate: 115_200,
        ca_cert_path: None,
        require_valid_cert: true,
    };

    // Built (which checks the required fields are present and non-blank) and
    // connected before the fleet lock is taken: connecting can take seconds,
    // and the lock is what telemetry waits on.
    let mut printer = fleet::build(entry.clone()).map_err(|err| err.to_string())?;
    match tokio::time::timeout(CONNECT_TIMEOUT, printer.connect()).await {
        Ok(Ok(())) => {}
        Ok(Err(err)) => return Err(format!("could not connect to the printer: {err}")),
        Err(_) => return Err("the printer did not answer in time - check the address and that it is on".into()),
    }

    store.save_printer(&entry).map_err(|err| format!("connected, but could not save it locally: {err}"))?;
    info!(printer = %entry.id, "printer added from the console");

    let mut fleet = fleet.lock().await;
    // Re-checked under the lock: two requests for one id can both pass the
    // check above while the first is still connecting.
    if fleet.iter().any(|p| p.base().id == entry.id) {
        return Err(format!("this gateway already has a printer with id '{}'", entry.id));
    }
    fleet.push(printer);
    Ok(())
}

pub async fn remove(fleet: &SharedFleet, store: &Store, printer_id: &str) -> Result<(), String> {
    let stored = store.remove_printer(printer_id).map_err(|err| err.to_string())?;
    let mut fleet = fleet.lock().await;
    let before = fleet.len();
    fleet.retain(|p| p.base().id != printer_id);
    if !stored && fleet.len() == before {
        return Err(format!("no printer with id '{printer_id}' on this gateway"));
    }
    info!(printer = %printer_id, "printer removed from the console");
    Ok(())
}

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

    #[test]
    fn ids_are_a_small_safe_alphabet() {
        assert!(valid_id("p1p"));
        assert!(valid_id("Printer_2-b"));
        assert!(!valid_id(""));
        assert!(!valid_id("../etc/passwd"));
        assert!(!valid_id("has space"));
        assert!(!valid_id(&"a".repeat(65)));
    }
}

#[cfg(test)]
mod behaviour {
    use std::sync::Arc;

    use super::*;

    fn request(id: &str) -> ProvisionPrinter {
        ProvisionPrinter {
            id: id.into(),
            name: "Bench P1".into(),
            vendor: "bambu".into(),
            // Nothing listens here, so a connect attempt is refused at once.
            host: Some("127.0.0.1".into()),
            access_code: Some("12345678".into()),
            sn: Some("01P00A000000001".into()),
            api_key: None,
        }
    }

    fn store() -> (Store, std::path::PathBuf) {
        let path = std::env::temp_dir().join(format!("continuum-provision-test-{}-{}.db", std::process::id(), rand_suffix()));
        (Store::open(&path).unwrap(), path)
    }

    fn rand_suffix() -> u128 {
        std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos()
    }

    fn empty_fleet() -> SharedFleet {
        Arc::new(tokio::sync::Mutex::new(Vec::new()))
    }

    #[tokio::test]
    async fn a_bad_id_is_refused_before_anything_happens() {
        let (store, path) = store();
        let err = add(&empty_fleet(), &store, request("../x")).await.unwrap_err();
        assert!(err.contains("printer id"), "{err}");
        assert!(store.printers().unwrap().is_empty());
        let _ = std::fs::remove_file(path);
    }

    #[tokio::test]
    async fn a_blank_credential_is_named_and_nothing_is_saved() {
        let (store, path) = store();
        let mut req = request("p1");
        req.access_code = Some("  ".into());
        let err = add(&empty_fleet(), &store, req).await.unwrap_err();
        assert!(err.contains("access_code"), "{err}");
        assert!(store.printers().unwrap().is_empty());
        let _ = std::fs::remove_file(path);
    }

    #[tokio::test]
    async fn a_printer_that_cannot_be_reached_is_not_saved() {
        let (store, path) = store();
        let fleet = empty_fleet();
        let err = add(&fleet, &store, request("p1")).await.unwrap_err();
        assert!(err.contains("could not connect") || err.contains("did not answer"), "{err}");
        assert!(store.printers().unwrap().is_empty(), "a failed add must not leave a stored printer");
        assert!(fleet.lock().await.is_empty());
        let _ = std::fs::remove_file(path);
    }

    #[tokio::test]
    async fn the_error_never_contains_the_access_code() {
        let (store, path) = store();
        let err = add(&empty_fleet(), &store, request("p1")).await.unwrap_err();
        assert!(!err.contains("12345678"), "{err}");
        let _ = std::fs::remove_file(path);
    }

    #[tokio::test]
    async fn removing_forgets_a_printer_in_both_places() {
        let (store, path) = store();
        let entry = crate::fleet::PrinterEntry {
            id: "p1".into(),
            name: "Bench".into(),
            vendor: "bambu".into(),
            host: Some("127.0.0.1".into()),
            access_code: Some("12345678".into()),
            sn: Some("01P00A000000001".into()),
            api_key: None,
            port: None,
            baud_rate: 115_200,
            ca_cert_path: None,
            require_valid_cert: true,
        };
        store.save_printer(&entry).unwrap();
        let fleet = empty_fleet();
        fleet.lock().await.push(crate::fleet::build(entry).unwrap());

        remove(&fleet, &store, "p1").await.unwrap();
        assert!(fleet.lock().await.is_empty());
        assert!(store.printers().unwrap().is_empty());

        let err = remove(&fleet, &store, "p1").await.unwrap_err();
        assert!(err.contains("no printer"), "{err}");
        let _ = std::fs::remove_file(path);
    }

    /// The success path, using the OctoPrint client because its `connect` is
    /// the one that succeeds without a real machine (the driver is a stub).
    /// What this proves is the plumbing, not a vendor: a connected printer is
    /// saved, joins the running fleet, and is still there after a "restart"
    /// that reopens the store and rebuilds the fleet from it.
    #[tokio::test]
    async fn an_added_printer_joins_the_fleet_and_survives_a_restart() {
        let (store, path) = store();
        let fleet = empty_fleet();
        let mut req = request("octo1");
        req.vendor = "octoprint".into();
        req.sn = None;
        req.access_code = None;
        req.api_key = Some("key-123".into());

        add(&fleet, &store, req.clone()).await.unwrap();
        assert!(fleet.lock().await.iter().any(|p| p.base().id == "octo1"));

        // Same id again: refused, and the first is untouched.
        let again = add(&fleet, &store, req).await.unwrap_err();
        assert!(again.contains("already has"), "{again}");
        assert_eq!(fleet.lock().await.len(), 1);

        // "Restart": reopen the store, rebuild what main would.
        drop(store);
        let reopened = Store::open(&path).unwrap();
        let rebuilt: Vec<_> = reopened.printers().unwrap().into_iter().map(|e| crate::fleet::build(e).unwrap()).collect();
        assert_eq!(rebuilt.len(), 1);
        assert_eq!(rebuilt[0].base().id, "octo1");
        let _ = std::fs::remove_file(path);
    }

    #[test]
    fn the_wire_message_parses_and_its_debug_hides_the_credential() {
        let json = r#"{"type":"provision_printer","request_id":"r1","printer":{"id":"p1","name":"Bench","vendor":"bambu","host":"10.0.0.5","access_code":"98765432","sn":"SN1"}}"#;
        let message: crate::uplink::ServerMessage = serde_json::from_str(json).unwrap();
        let shown = format!("{message:?}");
        assert!(shown.contains("p1") && shown.contains("<redacted>"), "{shown}");
        assert!(!shown.contains("98765432"), "the access code reached a Debug string: {shown}");
    }
}