path: root/src/transport/mycelium.rs

use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use base64::engine::{Engine, general_purpose::STANDARD as BASE64};
use reqwest::StatusCode;
use serde::{Deserialize, Serialize};
use std::time::Duration;
use thiserror::Error;
use url::Url;

use super::daemon::MyceliumSetup;
use crate::transport::wire;
use crate::transport::{EncodedMessage, MessageHandler, ParsedMessage, Transport};

const MYCELIUM_POLL_TIMEOUT_SECS: u64 = 300;

/// The one Mycelium topic all of continuum's messages travel on; what a message is, its
/// payload says (`wire`).
const TOPIC: &str = "continuum.v1";

const EMPTY_QUEUE_DELAY: Duration = Duration::from_millis(25);

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct InboundMessage {
    id: String,
    src_pk: String,
    topic: String,
    payload: String,
}

// Deserialize only for debugging
#[derive(Debug, Serialize, Deserialize)]
struct OutboundMessage {
    dst: Destination,
    topic: String,
    payload: String,
}

// Deserialize only for debugging
#[derive(Debug, Serialize, Deserialize)]
struct Destination {
    pk: String,
}

/// The part of Mycelium's `/api/v1/admin` answer we use.
#[derive(Deserialize, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct MyceliumAdminInfo {
    pub node_pubkey: String,
}

fn is_pubkey(s: &str) -> bool {
    s.len() == 64 && s.chars().all(|c| c.is_ascii_hexdigit())
}

pub struct MyceliumTransport {
    // Wrap your zero-allocation lifecycle supervisor directly inside the driver plane
    pub setup: tokio::sync::Mutex<MyceliumSetup>,
    pub client: reqwest::Client,
    pub poll_client: reqwest::Client,
    pub messages_endpoint: Url,
    pubkey_basepath: Url,
    local_id: String,
}

impl MyceliumTransport {
    pub async fn connect(setup: MyceliumSetup) -> Result<Self> {
        let rest_url = setup.rest_url.clone();

        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(10))
            .build()?;
        let poll_client = reqwest::Client::builder()
            .timeout(Duration::from_secs(MYCELIUM_POLL_TIMEOUT_SECS + 30))
            .build()?;

        let admin: MyceliumAdminInfo = client
            .get(rest_url.join("/api/v1/admin")?)
            .send()
            .await?
            .error_for_status()?
            .json()
            .await
            .context("Querying local Mycelium identity")?;

        let transport = Self {
            setup: tokio::sync::Mutex::new(setup),
            client,
            poll_client,
            messages_endpoint: rest_url.join("/api/v1/messages")?,
            pubkey_basepath: rest_url.join("/api/v1/pubkey/")?,
            local_id: admin.node_pubkey,
        };
        transport.discard_backlog().await;
        Ok(transport)
    }

    /// Drops what Mycelium queued for us before we started. A Mycelium that outlives us
    /// (an external one, or across our restart) keeps accepting messages meant for us:
    /// HELLO retries, keepalives, sync messages of sessions that ended with our previous
    /// run. Processing them would only replay stale handshakes. Everything that still
    /// matters is resent: peers retry their HELLOs, and sync starts over after the
    /// handshake. Runs before we send anything, so no answer to us is dropped.
    ///
    /// Bounded in time, not count: the backlog is whatever piled up while we were away,
    /// which a local loop drains in well under a second even for large swarms. Whatever
    /// is still arriving after that is current traffic (peers that don't know we're
    /// back), and it is processed normally, exactly as without this step.
    async fn discard_backlog(&self) {
        const DRAIN_FOR: Duration = Duration::from_secs(2);
        let deadline = std::time::Instant::now() + DRAIN_FOR;
        let mut discarded = 0u64;
        while std::time::Instant::now() < deadline {
            let response = self
                .client
                .get(self.messages_endpoint.clone())
                .query(&[("timeout", 0)])
                .query(&[("topic", encode_topic(TOPIC))])
                .send()
                .await;
            match response {
                Ok(response) if response.status() == StatusCode::NO_CONTENT => break,
                Ok(response) if response.status().is_success() => discarded += 1,
                Ok(response) => {
                    tracing::warn!(status = %response.status(), "Couldn't drain the Mycelium inbox");
                    break;
                }
                Err(error) => {
                    tracing::warn!(error = %error, "Couldn't drain the Mycelium inbox");
                    break;
                }
            }
        }
        if discarded > 0 {
            tracing::info!(
                discarded,
                "Discarded messages Mycelium queued for us before we started"
            );
        }
    }

    async fn lookup_pubkey(&self, ip: &std::net::Ipv6Addr) -> Result<String> {
        #[derive(Deserialize)]
        struct PublicKeyResponse {
            #[serde(rename = "publicKey")]
            public_key: String,
        }

        let response: PublicKeyResponse = self
            .client
            .get(self.pubkey_basepath.join(&ip.to_string())?)
            .send()
            .await?
            .error_for_status()?
            .json()
            .await?;
        Ok(response.public_key)
    }
}

#[async_trait]
impl Transport for MyceliumTransport {
    fn local_id(&self) -> &str {
        &self.local_id
    }

    async fn resolve_seed(&self, hint: &str) -> Result<String> {
        if is_pubkey(hint) {
            return Ok(hint.to_owned());
        }

        if let Ok(ip) = hint.parse::<std::net::Ipv6Addr>() {
            return self.lookup_pubkey(&ip).await;
        }

        let ip = tokio::time::timeout(Duration::from_secs(10), tokio::net::lookup_host((hint, 0)))
            .await
            .context("DNS lookup timed out")??
            .find_map(|addr| match addr.ip() {
                std::net::IpAddr::V6(ip) => Some(ip),
                _ => None,
            })
            .ok_or_else(|| anyhow!("{hint}: no AAAA record found"))?;

        self.lookup_pubkey(&ip).await
    }

    fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
        encode_message(peer, message)
    }

    async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> {
        tracing::trace!(payload_bytes = message.body.len(), "Sending message");
        let response = self
            .client
            .post(self.messages_endpoint.clone())
            .header(reqwest::header::CONTENT_TYPE, "application/json")
            .body(message.body.clone())
            .send()
            .await?;

        let status = response.status();
        if !status.is_success() {
            tracing::error!(status = %status, "Mycelium rejected outbound message");
            return Err(TransportError::MyceliumRejected {
                peer: message.peer.clone(),
                status,
            }
            .into());
        }

        // Mycelium's message ID, for GET /api/v1/messages/status/{id} when chasing
        // a message that never arrived.
        if tracing::enabled!(tracing::Level::DEBUG) {
            let Ok(body) = response
                .bytes()
                .await
                .map_err(anyhow::Error::from)
                .and_then(|bytes| Ok(serde_json::from_slice::<serde_json::Value>(&bytes)?))
            else {
                tracing::debug!(to = %message.peer, "Sent message (couldn't decode response from Mycelium REST API)");
                return Ok(());
            };
            let Ok(outbound) = serde_json::from_slice::<OutboundMessage>(&message.body) else {
                tracing::debug!(to = %message.peer, body_len = %message.body.len(), "Sent message (couldn't parse outbound message)");
                return Ok(());
            };
            let Ok(parsed) = decode_payload(&outbound.topic, &outbound.payload) else {
                tracing::debug!(to = %message.peer, message = ?outbound, "Sent message (couldn't decode outbound message)");
                return Ok(());
            };
            let Some(id) = body.get("id").and_then(|id| id.as_str()) else {
                tracing::debug!(to = %message.peer, message = ?parsed, "Sent message (couldn't read id from Mycelium REST API response)");
                return Ok(());
            };
            tracing::debug!(to = %message.peer, message = ?parsed, id = %id, "Sent message");
        }
        Ok(())
    }

    async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
        tracing::info!(
            endpoint = self.messages_endpoint.as_str(),
            "Starting Mycelium message consumer"
        );

        loop {
            tracing::trace!("Polling Mycelium inbox");

            let response = match self
                .poll_client
                .get(self.messages_endpoint.clone())
                .query(&[("timeout", MYCELIUM_POLL_TIMEOUT_SECS)])
                .query(&[("topic", encode_topic(TOPIC))])
                .send()
                .await
            {
                Ok(resp) => resp,
                Err(error) => {
                    tracing::warn!(error = %error, "Mycelium inbox poll failed");
                    tokio::time::sleep(Duration::from_millis(250)).await;
                    continue;
                }
            };

            let status = response.status();

            if status == StatusCode::NO_CONTENT {
                tracing::trace!("Mycelium inbox is empty");
                tokio::time::sleep(EMPTY_QUEUE_DELAY).await;
                continue;
            }

            if !status.is_success() {
                tracing::warn!(status = %status, "Mycelium inbox returned unsuccessful HTTP status");
                tokio::time::sleep(Duration::from_millis(500)).await;
                continue;
            }

            let raw_bytes = match response.bytes().await {
                Ok(bytes) => bytes,
                Err(error) => {
                    tracing::warn!(error = %error, "Failed to read raw Mycelium response stream");
                    tokio::time::sleep(Duration::from_millis(100)).await;
                    continue;
                }
            };

            let inbound_message = match serde_json::from_slice::<InboundMessage>(&raw_bytes) {
                Ok(msg) => msg,
                Err(error) => {
                    tracing::error!(
                        error = %error,
                        raw_payload = %String::from_utf8_lossy(&raw_bytes),
                        "Failed to decode inbound Mycelium message structure"
                    );
                    tokio::time::sleep(Duration::from_millis(500)).await;
                    continue;
                }
            };

            match parse_message(&inbound_message) {
                Ok((sender, parsed)) => {
                    tracing::debug!(
                        from = %inbound_message.src_pk,
                        message = ?parsed,
                        id = %inbound_message.id,
                        //payload_bytes = raw_bytes.len(),
                        "Received message"
                    );
                    let payload_len = base64_decoded_len(&inbound_message.payload);
                    if let Err(error) = handler(sender, parsed, payload_len).await {
                        tracing::warn!(error = %format!("{error:#}"), "Failed to process inbound message");
                    }
                }
                Err(transport_error) => {
                    tracing::error!(
                        error = %transport_error,
                        topic = %inbound_message.topic,
                        "Dropping poisoned packet: parsing failed gracefully"
                    );
                    // No sleep required: we successfully popped a malformed packet and want to process the next one immediately
                }
            }
        }
    }

    async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> {
        let mut guard = self.setup.lock().await;
        guard.wait_for_exit().await
    }
}

/// A message as the Mycelium REST API takes it (the JSON body of an outbound message).
/// Also the simulator's wire format (`transport::sim`), so tests go through the same
/// encoding.
pub(crate) fn encode_message(peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
    let (outbound, payload_len) = format_message(peer, message)?;
    Ok(EncodedMessage {
        peer: peer.to_owned(),
        body: serde_json::to_vec(&outbound)?.into(),
        kind: message.kind(),
        payload_len,
    })
}

/// The message in a body made by [`encode_message`].
#[cfg(test)]
pub(crate) fn decode_message(body: &[u8]) -> Result<ParsedMessage> {
    let outbound: OutboundMessage = serde_json::from_slice(body)?;
    Ok(decode_payload(&outbound.topic, &outbound.payload)?)
}

fn encode_topic(topic: &str) -> String {
    BASE64.encode(topic.as_bytes())
}

fn decode_topic(topic: &str) -> Result<String, TransportError> {
    let bytes = BASE64.decode(topic)?;
    Ok(String::from_utf8(bytes)?)
}

fn encode_payload(payload: &Vec<u8>) -> String {
    BASE64.encode(payload)
}

/// Length of what padded base64 `encoded` decodes to.
fn base64_decoded_len(encoded: &str) -> usize {
    let padding = encoded.bytes().rev().take_while(|&b| b == b'=').count();
    (encoded.len() / 4 * 3).saturating_sub(padding)
}

fn decode_payload(topic_encoded: &str, payload: &str) -> Result<ParsedMessage, TransportError> {
    let topic = decode_topic(topic_encoded)?;
    if topic != TOPIC {
        return Err(TransportError::UnknownTopic(topic));
    }
    Ok(wire::decode(&BASE64.decode(payload)?)?)
}

fn parse_message(message: &InboundMessage) -> Result<(String, ParsedMessage), TransportError> {
    let sender = message.src_pk.clone();

    if sender.is_empty() {
        return Err(TransportError::NoSender);
    }

    let parsed = decode_payload(&message.topic, &message.payload)?;

    Ok((sender, parsed))
}

/// The message for Mycelium's API, and the length of its payload.
fn format_message(
    recipient: &str,
    message: &ParsedMessage,
) -> Result<(OutboundMessage, usize), TransportError> {
    let payload_bytes = wire::encode(message)?;
    let outbound = OutboundMessage {
        dst: Destination {
            pk: recipient.to_owned(),
        },
        topic: encode_topic(TOPIC),
        payload: encode_payload(&payload_bytes),
    };
    Ok((outbound, payload_bytes.len()))
}

#[derive(Error, Debug)]
pub enum TransportError {
    #[error("Mycelium delivered a message without a sender public key")]
    NoSender,

    #[error("Base64 decoding failed: {0}")]
    Base64Decode(#[from] base64::DecodeError),

    #[error("UTF-8 decoding failed: {0}")]
    Utf8Decode(#[from] std::string::FromUtf8Error),

    #[error(transparent)]
    Wire(#[from] wire::WireError),

    #[error("Unknown topic: {0}")]
    UnknownTopic(String),

    #[error("Network HTTP request failed: {0}")]
    Http(#[from] reqwest::Error),

    #[error("Mycelium server rejected outbound message to {peer} (HTTP {status})")]
    MyceliumRejected {
        peer: String,
        status: reqwest::StatusCode,
    },
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::transport::{
        BlobBody, BlobId, BlobMessage, Epoch, HeadsSummary, Keepalive, MAX_LISTED_HEADS,
    };
    use automerge::ChangeHash;

    /// Mycelium splits messages into chunks of this size.
    const MYCELIUM_CHUNK: usize = 1300;

    #[test]
    fn largest_keepalive_fits_in_one_chunk() {
        let keepalive = Keepalive {
            heads: HeadsSummary::Heads(
                (0..MAX_LISTED_HEADS as u8)
                    .map(|n| ChangeHash([n; 32]))
                    .collect(),
            ),
            epoch: Epoch {
                generation: u64::MAX,
                id: [0xff; 32],
            },
            epoch_seq: u32::MAX,
            ready: true,
            syncs_with_you: true,
        };
        let payload = wire::encode(&ParsedMessage::Pong(keepalive.clone())).unwrap();
        // Mycelium's message data: a length byte, the topic, then our payload.
        let size = 1 + TOPIC.len() + payload.len();
        assert!(
            size + 50 <= MYCELIUM_CHUNK,
            "keepalive with {MAX_LISTED_HEADS} heads takes {size} bytes"
        );
        let ParsedMessage::Pong(decoded) = wire::decode(&payload).unwrap() else {
            panic!("not a PONG");
        };
        assert_eq!(decoded.heads, keepalive.heads);
    }

    #[test]
    fn blob_messages_round_trip() {
        let epoch = Epoch {
            generation: 3,
            id: [7; 32],
        };
        let id = BlobId::Change(ChangeHash([9; 32]));
        let decode_blob = |payload: &[u8]| match wire::decode(payload).unwrap() {
            ParsedMessage::Blob(blob) => blob,
            other => panic!("not a blob: {other:?}"),
        };
        let chunk = ParsedMessage::Blob(BlobMessage {
            epoch,
            epoch_seq: 4,
            id,
            body: BlobBody::Chunk {
                index: 5,
                data: vec![1u8; 64 * 1024].into(),
            },
        });
        let bytes = wire::encode(&chunk).unwrap();
        // A 64 KiB chunk costs only a few bytes of framing.
        assert!(bytes.len() < 64 * 1024 + 128, "{} bytes", bytes.len());
        let decoded = decode_blob(&bytes);
        assert_eq!(decoded.id, id);
        assert_eq!(decoded.epoch, epoch);
        assert_eq!(decoded.epoch_seq, 4);
        assert!(
            matches!(decoded.body, BlobBody::Chunk { index: 5, ref data } if data.len() == 64 * 1024)
        );

        let manifest = ParsedMessage::Blob(BlobMessage {
            epoch,
            epoch_seq: 4,
            id,
            body: BlobBody::Manifest {
                size: 100,
                chunk_size: 64,
                hashes: vec![[1; 32], [2; 32]],
            },
        });
        let decoded = decode_blob(&wire::encode(&manifest).unwrap());
        assert!(matches!(
            decoded.body,
            BlobBody::Manifest { size: 100, chunk_size: 64, ref hashes } if hashes.len() == 2
        ));
    }

    #[test]
    fn messages_round_trip_through_the_api_body() {
        let keepalive = ParsedMessage::Ping(Keepalive {
            heads: HeadsSummary::of(vec![ChangeHash([1; 32])]),
            epoch: Epoch::GENESIS,
            epoch_seq: 0,
            ready: true,
            syncs_with_you: true,
        });
        let encoded = encode_message("peer", &keepalive).unwrap();
        assert!(matches!(
            decode_message(&encoded.body).unwrap(),
            ParsedMessage::Ping(_)
        ));
    }

    #[test]
    fn listed_heads_equal_their_digest() {
        let heads: Vec<_> = (0..3u8).map(|n| ChangeHash([n; 32])).collect();
        let listed = HeadsSummary::of(heads.clone());
        let digest = HeadsSummary::Digest(listed.digest());
        assert_eq!(listed, digest);
        assert_ne!(listed, HeadsSummary::of(heads[..2].to_vec()));
    }
}