path: root/src/transport/mycelium/mod.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 crate::transport::wire;
use crate::transport::{
    Addressed, EncodedMessage, HandlerFn, MessageHandler, ParsedMessage, Transport, hand_on_probe,
    is_probe,
};

pub(crate) mod daemon;

pub(crate) use daemon::{MyceliumSetup, MyceliumSpawnConfig};

/// The scheme of this transport's addresses (`transport::paths`).
pub(crate) const SCHEME: &str = "mycelium";

const MYCELIUM_POLL_TIMEOUT_SECS: u64 = 300;
/// How long Mycelium tries to deliver a message before dropping it.
const DELIVERY_LIMIT: Duration = Duration::from_secs(300);
/// The largest inbound message read (its JSON, base64 payload included: about 190 MiB of
/// payload). Anyone can send us anything; this bounds what one message costs before it
/// is even parsed. Generous, because a sync message can carry a whole history until
/// large deltas travel as blobs.
const MAX_INBOUND_BYTES: usize = 256 << 20;
/// Polls of the inbox in progress at once.
const INBOX_POLLS: usize = 2;
/// How long a request to our Mycelium may take, before what its body adds
/// (`post_timeout`).
const REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
/// The slowest our Mycelium is expected to take a message body, in bytes per second:
/// a large sync message gets more time than a STATUS, instead of failing (and being
/// generated and sent whole again) at a fixed limit.
const POST_MIN_RATE: u64 = 1 << 20;

/// The one Mycelium topic all of continuum's messages travel on; what a message is, its
/// payload says (`wire`).
const TOPIC: &str = "continuum";
/// How many senders on another wire version are remembered as already logged.
const MAX_OTHER_VERSIONS: usize = 1024;

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(crate) struct MyceliumAdminInfo {
    pub(crate) node_pubkey: String,
}

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

pub(crate) struct MyceliumTransport {
    /// The daemon, kept to wait for its exit.
    setup: tokio::sync::Mutex<MyceliumSetup>,
    client: reqwest::Client,
    /// For the long-polling inbox (its timeout is longer).
    poll_client: reqwest::Client,
    messages_endpoint: Url,
    pubkey_basepath: Url,
    local_id: String,
}

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

        let client = reqwest::Client::builder()
            .timeout(REQUEST_TIMEOUT)
            .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 inbox");
                    break;
                }
                Err(error) => {
                    tracing::warn!(error = %error, "Couldn't drain the 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
    }

    fn is_peer_id(&self, id: &str) -> bool {
        is_pubkey(id)
    }

    fn delivery_limit(&self) -> Duration {
        DELIVERY_LIMIT
    }

    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")
            .timeout(post_timeout(message.body.len()))
            .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 = %Addressed(SCHEME, &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 = %Addressed(SCHEME, &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 = %Addressed(SCHEME, &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 = %Addressed(SCHEME, &message.peer), message = ?parsed, "Sent message (couldn't read id from Mycelium REST API response)");
                return Ok(());
            };
            tracing::debug!(to = %Addressed(SCHEME, &message.peer), message = ?parsed, id = %id, "Sent message");
        }
        Ok(())
    }

    async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
        // Senders on another wire version, logged once each.
        let mut other_versions = std::collections::HashSet::new();
        tracing::info!(
            endpoint = self.messages_endpoint.as_str(),
            "Starting the consumer"
        );

        // Several polls at once, so one message's HTTP round trip overlaps the next one's
        // and handling them: the messages still go to the handler one at a time.
        let (tx, mut rx) = tokio::sync::mpsc::channel(INBOX_POLLS * 4);
        for _ in 0..INBOX_POLLS {
            tokio::spawn(poll_inbox(
                self.poll_client.clone(),
                self.messages_endpoint.clone(),
                tx.clone(),
            ));
        }
        drop(tx);

        // PINGs and PONGs aren't handled after the message before them (`is_probe`).
        let handler: std::sync::Arc<HandlerFn> = std::sync::Arc::from(handler);
        while let Some(inbound_message) = rx.recv().await {
            match parse_message(&inbound_message) {
                Ok((sender, parsed)) => {
                    tracing::debug!(
                        from = %Addressed(SCHEME, &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 is_probe(&parsed) {
                        hand_on_probe(SCHEME, &handler, sender, parsed, payload_len);
                    } else if let Err(error) = handler(sender, parsed, payload_len).await {
                        tracing::warn!(error = %format!("{error:#}"), "Failed to process inbound message");
                    }
                }
                Err(TransportError::Wire(wire::WireError::Version(version))) => {
                    // Bounded: anyone can send us anything.
                    if other_versions.len() >= MAX_OTHER_VERSIONS {
                        other_versions.clear();
                    }
                    if other_versions.insert(inbound_message.src_pk.clone()) {
                        tracing::warn!(
                            from = %Addressed(SCHEME, &inbound_message.src_pk),
                            version,
                            ours = wire::VERSION,
                            "Peer speaks another wire version (another build); ignoring its messages"
                        );
                    }
                }
                Err(transport_error) => {
                    // Who sent it and what it claims to be.
                    let kind = BASE64
                        .decode(&inbound_message.payload)
                        .ok()
                        .and_then(|payload| payload.get(1).copied());
                    tracing::error!(
                        error = %transport_error,
                        from = %Addressed(SCHEME, &inbound_message.src_pk),
                        kind = ?kind,
                        payload_bytes = base64_decoded_len(&inbound_message.payload),
                        topic = %inbound_message.topic,
                        "Dropping poisoned packet: parsing failed gracefully"
                    );
                }
            }
        }
        Ok(())
    }

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

/// Long-polls the inbox and passes what it pops to `tx` (`start_consumer`).
async fn poll_inbox(
    client: reqwest::Client,
    endpoint: Url,
    tx: tokio::sync::mpsc::Sender<InboundMessage>,
) {
    loop {
        tracing::trace!("Polling the inbox");

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

        let status = response.status();

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

        if !status.is_success() {
            tracing::warn!(status = %status, "Polling the inbox failed");
            tokio::time::sleep(Duration::from_millis(500)).await;
            continue;
        }

        let raw_bytes = match read_capped(response, MAX_INBOUND_BYTES).await {
            Ok(Some(bytes)) => bytes,
            Ok(None) => {
                tracing::warn!(
                    limit_bytes = MAX_INBOUND_BYTES,
                    "Dropping an inbound message larger than the limit"
                );
                continue;
            }
            Err(error) => {
                tracing::warn!(error = %error, "Couldn't read a message from the inbox");
                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),
                    "Dropping a message that doesn't decode"
                );
                tokio::time::sleep(Duration::from_millis(500)).await;
                continue;
            }
        };

        if tx.send(inbound_message).await.is_err() {
            return; // the consumer stopped
        }
    }
}

/// The body of `response`, unless it is larger than `limit` bytes (None): read no
/// further than that, whatever the sender claims.
async fn read_capped(
    mut response: reqwest::Response,
    limit: usize,
) -> reqwest::Result<Option<bytes::Bytes>> {
    if response
        .content_length()
        .is_some_and(|length| length > limit as u64)
    {
        return Ok(None);
    }
    let mut body = bytes::BytesMut::new();
    while let Some(chunk) = response.chunk().await? {
        if body.len() + chunk.len() > limit {
            return Ok(None);
        }
        body.extend_from_slice(&chunk);
    }
    Ok(Some(body.freeze()))
}

/// 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(),
        address: peer.to_owned(),
        body: serde_json::to_vec(&outbound)?.into(),
        kind: message.kind(),
        payload_len,
        message: None,
    })
}

/// 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)
}

/// How long posting a body of `bytes` may take: `REQUEST_TIMEOUT`, plus the body at
/// `POST_MIN_RATE`.
fn post_timeout(bytes: usize) -> Duration {
    REQUEST_TIMEOUT + Duration::from_millis(bytes as u64 * 1000 / POST_MIN_RATE)
}

/// 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(crate) 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, MAX_LISTED_HEADS, Status,
    };
    use automerge::ChangeHash;

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

    #[test]
    fn post_timeout_grows_with_the_body() {
        assert_eq!(post_timeout(0), REQUEST_TIMEOUT);
        assert_eq!(
            post_timeout(64 << 20),
            REQUEST_TIMEOUT + Duration::from_secs(64)
        );
    }

    #[test]
    fn largest_keepalive_fits_in_one_chunk() {
        // A digest with a full sample: the most a keepalive carries.
        let status = Status {
            heads: HeadsSummary::of(
                (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::StatusReply(status.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,
            "STATUS with {MAX_LISTED_HEADS} heads takes {size} bytes"
        );
        let ParsedMessage::StatusReply(decoded) = wire::decode(&payload).unwrap() else {
            panic!("not a STATUS_REPLY");
        };
        assert_eq!(decoded.heads, status.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 status = ParsedMessage::Status(Status {
            heads: HeadsSummary::of(vec![ChangeHash([1; 32])]),
            epoch: Epoch::GENESIS,
            epoch_seq: 0,
            ready: true,
            syncs_with_you: true,
        });
        let encoded = encode_message("peer", &status).unwrap();
        assert!(matches!(
            decode_message(&encoded.body).unwrap(),
            ParsedMessage::Status(_)
        ));
    }

    #[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 {
            digest: listed.digest(),
            sample: Vec::new(),
        };
        assert_eq!(listed, digest);
        assert_ne!(listed, HeadsSummary::of(heads[..2].to_vec()));
    }
}