continuum-proxy/src/printer/bambu_mqtt.rs

//! The real MQTT wiring for Bambu printers - the piece `connect`/
//! `send_gcode_line` in `bambu.rs` used to just stub out with a
//! `println!`. Two things make this more than "call `publish`, get a
//! `String` back":
//!
//! 1. MQTT v3.1.1 (BambuMQTT v1) and MQTT v5 (BambuMQTT v2) are two
//!    entirely separate rumqttc client types (`rumqttc::AsyncClient` vs
//!    `rumqttc::v5::AsyncClient`) with different event/error types -
//!    `BambuMqttClient` is the enum that hides which one a connection
//!    ended up using, same pattern as `PrinterHandle` hides which vendor
//!    struct you're holding.
//! 2. A response to a published request doesn't come back from that same
//!    call - it arrives later, asynchronously, on a *subscription* to
//!    `device/{sn}/report`, correlated only by matching `sequence_id`.
//!    That needs: a background task continuously polling the connection
//!    (rumqttc doesn't do this for you - nothing happens until you poll
//!    the `EventLoop` yourself, repeatedly, forever), and a shared table
//!    of "who's waiting for which `sequence_id`" that task completes as
//!    reports arrive. This is the first real use of `tokio::spawn` in
//!    this crate - flagged as deferred since `printer/mod.rs`'s very
//!    first doc comment, and this is why: it only makes sense once there's
//!    a long-running background job (polling forever) that has to run
//!    independently of whatever call site is waiting on a response.
//!
//! Unverified against a live printer (no broker to test against here):
//! the exact `ConnectReturnCode`/error shape a real Bambu printer's
//! embedded broker gives when it doesn't understand MQTT v5's CONNECT at
//! all (rather than a clean version-rejection CONNACK) - `open()` below
//! treats *any* failure of the v5 attempt as "fall back to v3.1.1", not
//! just a precisely-matched version-rejection code, since a broker that
//! can't parse a v5 CONNECT packet may just drop the connection instead
//! of replying with one.

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use rumqttc::v5::mqttbytes::v5::Packet as PacketV5;
use rumqttc::v5::mqttbytes::QoS as QoSV5;
use rumqttc::{Event, Packet, QoS, Transport, TlsConfiguration};
use serde_json::Value;
use tokio::sync::{oneshot, Mutex};

use super::bambu::BambuTls;
use super::PrinterError;
use super::BambuProtocolVersion;

/// Requests waiting on a report with a matching `sequence_id`. Shared
/// between whichever call registered one (in `request`) and the
/// background poll task that completes it when the report arrives.
type PendingTable = Arc<Mutex<HashMap<String, oneshot::Sender<Value>>>>;

const REQUEST_TIMEOUT: Duration = Duration::from_secs(10);

enum BambuMqttClient {
    V1(rumqttc::AsyncClient),
    V2(rumqttc::v5::AsyncClient),
}

/// One live MQTT connection to a Bambu printer - whichever protocol
/// version `open()` settled on.
pub struct BambuMqttSession {
    client: BambuMqttClient,
    pending: PendingTable,
    sn: String,
    next_sequence_id: AtomicU64,
}

impl BambuMqttSession {
    /// `ForceV1`/`ForceV2` connect with that protocol directly. `Auto`
    /// tries v5 first and falls back to v3.1.1 if that attempt fails -
    /// see this module's doc comment for why "any failure", not just a
    /// precisely-matched version-rejection code.
    pub async fn open(
        host: &str,
        access_code: &str,
        sn: &str,
        tls: &BambuTls,
        preference: super::bambu::BambuProtocolPreference,
    ) -> Result<(Self, BambuProtocolVersion), PrinterError> {
        use super::bambu::BambuProtocolPreference as Pref;

        match preference {
            Pref::ForceV1 => {
                let session = Self::open_v1(host, access_code, sn, tls).await?;
                Ok((session, BambuProtocolVersion::V1))
            }
            Pref::ForceV2 => {
                let session = Self::open_v2(host, access_code, sn, tls).await?;
                Ok((session, BambuProtocolVersion::V2))
            }
            Pref::Auto => match Self::open_v2(host, access_code, sn, tls).await {
                Ok(session) => Ok((session, BambuProtocolVersion::V2)),
                Err(_) => {
                    let session = Self::open_v1(host, access_code, sn, tls).await?;
                    Ok((session, BambuProtocolVersion::V1))
                }
            },
        }
    }

    async fn open_v1(host: &str, access_code: &str, sn: &str, tls: &BambuTls) -> Result<Self, PrinterError> {
        let connector = tls.build_connector().map_err(|err| PrinterError(format!("TLS setup failed: {err}")))?;

        let mut options = rumqttc::MqttOptions::new(format!("continuum-{sn}"), host, 8883);
        options.set_credentials("bblp", access_code);
        options.set_transport(Transport::tls_with_config(TlsConfiguration::NativeConnector(connector)));
        options.set_keep_alive(Duration::from_secs(30));

        let (client, mut event_loop) = rumqttc::AsyncClient::new(options, 10);
        client
            .subscribe(format!("device/{sn}/report"), QoS::AtMostOnce)
            .await
            .map_err(|err| PrinterError(format!("subscribe failed: {err}")))?;

        // Nothing actually happens on the wire until the eventloop is
        // polled - this first poll is what makes the TCP+TLS+CONNECT
        // handshake happen, and its result is the only way to know
        // whether it worked.
        match event_loop.poll().await {
            Ok(_) => {}
            Err(err) => return Err(PrinterError(format!("MQTT v3.1.1 connect failed: {err}"))),
        }

        let pending: PendingTable = Arc::new(Mutex::new(HashMap::new()));
        tokio::spawn(poll_loop_v1(event_loop, pending.clone()));

        Ok(Self { client: BambuMqttClient::V1(client), pending, sn: sn.to_string(), next_sequence_id: AtomicU64::new(1) })
    }

    async fn open_v2(host: &str, access_code: &str, sn: &str, tls: &BambuTls) -> Result<Self, PrinterError> {
        let connector = tls.build_connector().map_err(|err| PrinterError(format!("TLS setup failed: {err}")))?;

        let mut options = rumqttc::v5::MqttOptions::new(format!("continuum-{sn}"), host, 8883);
        options.set_credentials("bblp", access_code);
        options.set_transport(Transport::tls_with_config(TlsConfiguration::NativeConnector(connector)));
        options.set_keep_alive(Duration::from_secs(30));

        let (client, mut event_loop) = rumqttc::v5::AsyncClient::new(options, 10);
        client
            .subscribe(format!("device/{sn}/report"), QoSV5::AtMostOnce)
            .await
            .map_err(|err| PrinterError(format!("subscribe failed: {err}")))?;

        match event_loop.poll().await {
            Ok(_) => {}
            Err(err) => return Err(PrinterError(format!("MQTT v5 connect failed: {err}"))),
        }

        let pending: PendingTable = Arc::new(Mutex::new(HashMap::new()));
        tokio::spawn(poll_loop_v2(event_loop, pending.clone()));

        Ok(Self { client: BambuMqttClient::V2(client), pending, sn: sn.to_string(), next_sequence_id: AtomicU64::new(1) })
    }

    /// Publishes `body` under `device/{sn}/request`, wrapped as
    /// `{category: body}` (Bambu's top-level key names the *category*,
    /// e.g. "print"/"info" - `body`'s own `command` field says which
    /// specific one), and waits up to `REQUEST_TIMEOUT` for a report
    /// whose `sequence_id` matches. Overwrites whatever `sequence_id`
    /// `body` already had (continuum-schema-instance fills it with an
    /// empty placeholder, since it has no way to know a real one) with a
    /// fresh one from this session's own counter.
    pub async fn request(&self, category: &str, mut body: Value) -> Result<Value, PrinterError> {
        let sequence_id = self.next_sequence_id.fetch_add(1, Ordering::Relaxed).to_string();
        body["sequence_id"] = Value::String(sequence_id.clone());

        let (tx, rx) = oneshot::channel();
        self.pending.lock().await.insert(sequence_id.clone(), tx);

        if let Err(err) = self.publish(category, &body).await {
            self.pending.lock().await.remove(&sequence_id);
            return Err(err);
        }

        match tokio::time::timeout(REQUEST_TIMEOUT, rx).await {
            Ok(Ok(response)) => Ok(response),
            Ok(Err(_)) => Err(PrinterError(format!("{} closed before replying", self.sn))),
            Err(_) => {
                self.pending.lock().await.remove(&sequence_id);
                Err(PrinterError(format!("{} did not reply within {REQUEST_TIMEOUT:?}", self.sn)))
            }
        }
    }

    /// Publishes `body` under `device/{sn}/request`, wrapped as
    /// `{category: body}`, without waiting for a reply - for commands
    /// like `gcode_line` that only get acknowledged implicitly (the
    /// printer just executes it), not answered directly.
    pub async fn publish(&self, category: &str, body: &Value) -> Result<(), PrinterError> {
        self.publish_raw(serde_json::json!({ category: body }).to_string()).await
    }

    /// Publishes an already-serialized payload under `device/{sn}/request`
    /// as-is, with no category wrapping - the low-level primitive both
    /// `publish` above and `GenericPrinter::publish` (a raw pass-through,
    /// no envelope opinion at all) are built on.
    pub async fn publish_raw(&self, payload: String) -> Result<(), PrinterError> {
        let topic = format!("device/{}/request", self.sn);

        match &self.client {
            BambuMqttClient::V1(client) => client
                .publish(topic, QoS::AtLeastOnce, false, payload)
                .await
                .map_err(|err| PrinterError(format!("publish failed: {err}"))),
            BambuMqttClient::V2(client) => client
                .publish(topic, QoSV5::AtLeastOnce, false, payload)
                .await
                .map_err(|err| PrinterError(format!("publish failed: {err}"))),
        }
    }
}

/// Runs forever, polling the v3.1.1 eventloop and completing pending
/// requests as matching reports arrive. Ends (silently) once the
/// connection is gone for good - dropping the `BambuMqttSession` drops
/// its `AsyncClient`, which is what eventually makes that happen.
async fn poll_loop_v1(mut event_loop: rumqttc::EventLoop, pending: PendingTable) {
    loop {
        match event_loop.poll().await {
            Ok(Event::Incoming(Packet::Publish(publish))) => {
                complete_pending(&pending, &publish.payload).await;
            }
            Ok(_) => {}
            Err(_) => return,
        }
    }
}

async fn poll_loop_v2(mut event_loop: rumqttc::v5::EventLoop, pending: PendingTable) {
    loop {
        match event_loop.poll().await {
            Ok(rumqttc::v5::Event::Incoming(PacketV5::Publish(publish))) => {
                complete_pending(&pending, &publish.payload).await;
            }
            Ok(_) => {}
            Err(_) => return,
        }
    }
}

/// Parses an incoming report and, if its `sequence_id` matches something
/// `request` is waiting on, hands it the parsed JSON. A report with no
/// `sequence_id`, or one nothing's waiting on (an unsolicited status
/// push, not a reply to anything), is silently ignored - most Bambu
/// reports aren't replies to a specific request at all.
async fn complete_pending(pending: &PendingTable, payload: &[u8]) {
    let Ok(value) = serde_json::from_slice::<Value>(payload) else { return };
    let Some(sequence_id) = find_sequence_id(&value) else { return };

    if let Some(sender) = pending.lock().await.remove(&sequence_id) {
        let _ = sender.send(value);
    }
}

/// A report is `{category: {"sequence_id": "...", ...}}` - one key, same
/// shape as a request. Look inside whichever single category key is
/// present rather than requiring the caller to already know it.
fn find_sequence_id(report: &Value) -> Option<String> {
    let object = report.as_object()?;
    let (_, category_body) = object.iter().next()?;
    category_body.get("sequence_id")?.as_str().map(str::to_string)
}