pixelpass/src/common/control.rs

//! Friends control-plane protocol and service.
//!
//! This is the always-on presence channel that rides the [`CONTROL_ALPN`]
//! endpoint (bound with the persistent identity — see
//! [`super::endpoint::bind_control`]). It's how two peers' GUIs exchange friend
//! requests and pushed share-codes, independent of any video session.
//!
//! Wire shape: **one message per connection.** The sender opens a bi-stream,
//! writes the JSON-encoded [`ControlMsg`], and finishes its send side (EOF
//! delimits the message — no length framing needed). The receiver reads to EOF,
//! parses, hands the message up, then writes a one-byte [`ACK`] back so the
//! sender knows it was delivered *and* parsed. That delivery signal is what
//! lets the host-side code-push queue (a later phase) tell "sent" from "friend
//! was offline." A friend's *reply* (accept/decline) is a separate later
//! connection in the other direction, because acceptance can happen minutes
//! after the request — not a response on the same stream.

use std::time::Duration;

use anyhow::{Context, Result, bail};
use iroh::endpoint::{Incoming, VarInt};
use iroh::{Endpoint, EndpointAddr, EndpointId};
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;

use super::alpn::CONTROL_ALPN;

/// Upper bound on a single control message. Generous for a display name plus a
/// share-code ticket (~150 chars); rejects a peer trying to make us buffer a
/// huge blob.
const MAX_MSG: usize = 64 * 1024;

/// One-byte application acknowledgement the receiver returns once it has parsed
/// a message. ASCII ACK (0x06).
const ACK: &[u8] = b"\x06";

/// Bound on each phase of the send handshake, so a half-dead peer or relay
/// can't park a sender (or an inbound handler) forever.
const IO_TIMEOUT: Duration = Duration::from_secs(10);

/// A message on the friends control plane.
///
/// `#[serde(tag = "type")]` keeps the JSON self-describing and lets us add
/// variants without breaking older peers (an unknown tag fails to parse and is
/// logged, rather than being silently misread as another variant).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ControlMsg {
    /// "I'm online; here's my current display name." A presence/name refresh.
    Hello { name: String },
    /// Ask the recipient to become friends.
    FriendRequest { name: String },
    /// Accept a request the recipient previously sent us.
    FriendAccept { name: String },
    /// Decline a pending request, or cancel an outgoing one.
    FriendDecline,
    /// A host pushing a freshly generated share-code to an accepted friend.
    ShareCode { name: String, ticket: String },
}

/// A received control message, paired with the *authenticated* sender id (the
/// connection's verified remote public key — not a value the peer can spoof in
/// the payload, which is why no variant carries a sender id).
#[derive(Debug, Clone)]
pub struct Inbound {
    pub from: EndpointId,
    pub msg: ControlMsg,
}

fn encode(msg: &ControlMsg) -> Result<Vec<u8>> {
    serde_json::to_vec(msg).context("failed to encode control message")
}

fn decode(bytes: &[u8]) -> Result<ControlMsg> {
    serde_json::from_slice(bytes).context("failed to decode control message")
}

/// Deliver one message to `peer` over `endpoint`, returning once the recipient
/// has acknowledged it. An error means it was *not* delivered (peer offline,
/// unreachable, or rejected the stream) — the caller can queue and retry.
///
/// `peer` is usually a bare [`EndpointId`] — friends store only the stable id,
/// and n0 DNS discovery resolves it to a live address. The full [`EndpointAddr`]
/// form exists for callers that already hold one (and for hermetic tests).
pub async fn send(
    endpoint: &Endpoint,
    peer: impl Into<EndpointAddr>,
    msg: &ControlMsg,
) -> Result<()> {
    let payload = encode(msg)?;
    let conn = tokio::time::timeout(IO_TIMEOUT, endpoint.connect(peer, CONTROL_ALPN))
        .await
        .context("timed out connecting to peer")?
        .context("failed to connect to peer")?;

    let io = async {
        let (mut send, mut recv) = conn
            .open_bi()
            .await
            .context("failed to open control stream")?;
        send.write_all(&payload)
            .await
            .context("failed to write control message")?;
        send.finish().context("failed to finish control stream")?;
        // Read the peer's ACK. read_to_end returns once the peer finishes its
        // send side, so this also serves as "the peer is done with us."
        let ack = recv
            .read_to_end(ACK.len() + 1)
            .await
            .context("peer closed the control stream without acknowledging")?;
        if ack != ACK {
            bail!(
                "peer sent an unexpected acknowledgement ({} bytes)",
                ack.len()
            );
        }
        Ok(())
    };

    let result = tokio::time::timeout(IO_TIMEOUT, io)
        .await
        .context("timed out sending control message")?;
    // Clean close so the peer's `closed().await` returns promptly either way.
    conn.close(VarInt::from_u32(0), b"done");
    result
}

/// Run the control-plane accept loop, forwarding every received message to
/// `tx`. Returns when the endpoint stops accepting (i.e. it was closed).
pub async fn serve(endpoint: Endpoint, tx: mpsc::Sender<Inbound>) {
    while let Some(incoming) = endpoint.accept().await {
        let tx = tx.clone();
        tokio::spawn(async move {
            if let Err(e) = handle(incoming, &tx).await {
                tracing::warn!("control: inbound connection failed: {e:#}");
            }
        });
    }
    tracing::info!("control: endpoint stopped accepting");
}

async fn handle(incoming: Incoming, tx: &mpsc::Sender<Inbound>) -> Result<()> {
    let conn = incoming
        .await
        .context("inbound control connection failed")?;
    let from = conn.remote_id();

    let msg = async {
        let (mut send, mut recv) = conn
            .accept_bi()
            .await
            .context("failed to accept control stream")?;
        let bytes = recv
            .read_to_end(MAX_MSG)
            .await
            .context("failed to read control message")?;
        let msg = decode(&bytes)?;
        // ACK only after a successful parse, so the sender's delivery signal
        // means "received and understood."
        send.write_all(ACK).await.context("failed to write ack")?;
        send.finish().context("failed to finish ack stream")?;
        Ok::<_, anyhow::Error>(msg)
    };

    let msg = tokio::time::timeout(IO_TIMEOUT, msg)
        .await
        .context("timed out reading control message")??;

    // Hand the message up first, so it reaches the UI promptly even when the
    // sender is slow to close (a degraded link could otherwise delay a friend
    // request / pushed code by up to IO_TIMEOUT).
    tx.send(Inbound { from, msg })
        .await
        .map_err(|_| anyhow::anyhow!("control: receiver dropped"))?;

    // Then wait (briefly) for the sender's close so our ACK has flushed before
    // the connection is dropped at the end of this scope.
    let _ = tokio::time::timeout(IO_TIMEOUT, conn.closed()).await;
    Ok(())
}

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

    #[test]
    fn control_msg_round_trips() {
        let cases = [
            ControlMsg::Hello {
                name: "alice".into(),
            },
            ControlMsg::FriendRequest { name: "bob".into() },
            ControlMsg::FriendAccept {
                name: "carol".into(),
            },
            ControlMsg::FriendDecline,
            ControlMsg::ShareCode {
                name: "dave".into(),
                ticket: "endpointaa…".into(),
            },
        ];
        for msg in cases {
            let bytes = encode(&msg).unwrap();
            assert_eq!(decode(&bytes).unwrap(), msg);
        }
    }

    #[test]
    fn unknown_tag_is_rejected() {
        assert!(decode(br#"{"type":"nonsense"}"#).is_err());
    }

    /// Bind a control-plane endpoint with a *fresh* random key, so two of them
    /// in one test get distinct ids (two real machines each have their own
    /// persistent key; `bind_control` would give both the same one here, and
    /// iroh refuses "connecting to ourself").
    async fn bind_test_control() -> Endpoint {
        iroh::Endpoint::builder(iroh::endpoint::presets::N0)
            .secret_key(iroh::SecretKey::generate())
            .alpns(vec![CONTROL_ALPN.to_vec()])
            .bind()
            .await
            .unwrap()
    }

    /// End-to-end over two real iroh endpoints on this machine. Ignored by
    /// default — it binds endpoints and waits on the relay, so it's slow and
    /// network-dependent. Run with `cargo test -- --ignored control`.
    #[tokio::test]
    #[ignore = "binds real iroh endpoints; run on demand"]
    async fn loopback_delivers_and_acks() {
        let server = bind_test_control().await;
        let client = bind_test_control().await;
        // Connect by full addr so the test doesn't depend on DNS discovery.
        server.online().await;
        client.online().await;
        let server_addr = server.addr();

        let (tx, mut rx) = mpsc::channel(4);
        let server_ep = server.clone();
        let serve_task = tokio::spawn(async move { serve(server_ep, tx).await });

        let msg = ControlMsg::FriendRequest {
            name: "tester".into(),
        };
        // Full addr (not just the id) so the test doesn't depend on DNS discovery.
        send(&client, server_addr.clone(), &msg).await.unwrap();

        let got = tokio::time::timeout(Duration::from_secs(15), rx.recv())
            .await
            .expect("no inbound within 15s")
            .expect("channel closed");
        assert_eq!(got.msg, msg);
        assert_eq!(got.from, client.addr().id);

        server.close().await;
        client.close().await;
        serve_task.abort();
    }
}