geth/crates/geth-discovery/src/lib.rs

use base64::Engine;
use geth_crypto::AgentKey;
use geth_types::{AgentId, NodeId, UnixMillis};
use serde::{Deserialize, Serialize};

pub const PEER_CARD_SIGNATURE_NAMESPACE: &str = "geth.peer-card.v1@geth.local";
pub const PEER_CARD_LAN_DISCOVERY_SERVICE: &str = "geth-peer-card";
pub const PEER_CARD_TXT_VERSION_KEY: &str = "geth";
pub const PEER_CARD_TXT_VERSION: &str = "peer-card-v1";
pub const PEER_CARD_TXT_CHUNKS_KEY: &str = "card-chunks";
pub const PEER_CARD_TXT_CHUNK_PREFIX: &str = "card-";
const PEER_CARD_TXT_CHUNK_BYTES: usize = 200;
const PEER_CARD_TXT_MAX_CHUNKS: usize = 32;

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct PeerCard {
    pub node_id: NodeId,
    pub agent_id: AgentId,
    pub endpoints: Vec<EndpointCandidate>,
    pub issued_at: UnixMillis,
    pub signature: SignatureMetadata,
}

impl PeerCard {
    pub fn signed(
        node_id: NodeId,
        agent_key: &AgentKey,
        endpoints: Vec<EndpointCandidate>,
        issued_at: UnixMillis,
    ) -> Result<Self, DiscoveryError> {
        let agent_id = agent_key.agent_id();
        let payload = PeerCardSigningPayload {
            node_id: node_id.clone(),
            agent_id: agent_id.clone(),
            endpoints: endpoints.clone(),
            issued_at,
        };
        let signature = agent_key.sign_canonical(PEER_CARD_SIGNATURE_NAMESPACE, &payload)?;
        Ok(Self {
            node_id,
            agent_id: agent_id.clone(),
            endpoints,
            issued_at,
            signature: SignatureMetadata {
                namespace: PEER_CARD_SIGNATURE_NAMESPACE.to_owned(),
                signer: agent_id.to_string(),
                public_key: agent_key.public_key_hex(),
                signature: hex::encode(signature),
            },
        })
    }

    pub fn validate_candidate(&self) -> Result<(), DiscoveryError> {
        if self.endpoints.is_empty() {
            return Err(DiscoveryError::MissingEndpoint);
        }
        self.verify_signature()
    }

    pub fn verify_signature(&self) -> Result<(), DiscoveryError> {
        if self.signature.namespace != PEER_CARD_SIGNATURE_NAMESPACE {
            return Err(DiscoveryError::InvalidSignatureNamespace(
                self.signature.namespace.clone(),
            ));
        }
        if self.signature.signer.is_empty()
            || self.signature.public_key.is_empty()
            || self.signature.signature.is_empty()
        {
            return Err(DiscoveryError::UnsignedPeerCard);
        }
        if self.signature.signer != self.agent_id.as_str() {
            return Err(DiscoveryError::SignerMismatch {
                signer: self.signature.signer.clone(),
                agent: self.agent_id.to_string(),
            });
        }
        let public_key = hex::decode(&self.signature.public_key)?;
        let fingerprint = geth_crypto::key_fingerprint(&public_key);
        if fingerprint != self.signature.signer {
            return Err(DiscoveryError::SignerPublicKeyMismatch {
                signer: self.signature.signer.clone(),
                fingerprint,
            });
        }
        let signature = hex::decode(&self.signature.signature)?;
        geth_crypto::verify_canonical(
            &public_key,
            PEER_CARD_SIGNATURE_NAMESPACE,
            &self.signing_payload(),
            &signature,
        )?;
        Ok(())
    }

    fn signing_payload(&self) -> PeerCardSigningPayload {
        PeerCardSigningPayload {
            node_id: self.node_id.clone(),
            agent_id: self.agent_id.clone(),
            endpoints: self.endpoints.clone(),
            issued_at: self.issued_at,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
struct PeerCardSigningPayload {
    node_id: NodeId,
    agent_id: AgentId,
    endpoints: Vec<EndpointCandidate>,
    issued_at: UnixMillis,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct EndpointCandidate {
    pub endpoint_id: String,
    pub relay_url: Option<String>,
    #[serde(default)]
    pub direct_addresses: Vec<String>,
    pub source: DiscoverySource,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SignatureMetadata {
    pub namespace: String,
    pub signer: String,
    #[serde(default)]
    pub public_key: String,
    pub signature: String,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum DiscoverySource {
    Manual,
    Mdns,
    PeerExchange,
    Imported,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct DiscoveredPeer {
    pub card: PeerCard,
    pub discovered_at: UnixMillis,
    pub source: DiscoverySource,
    pub trust_state: CandidateTrustState,
}

impl DiscoveredPeer {
    pub fn candidate(
        card: PeerCard,
        discovered_at: UnixMillis,
        source: DiscoverySource,
    ) -> Result<Self, DiscoveryError> {
        card.validate_candidate()?;
        Ok(Self {
            card,
            discovered_at,
            source,
            trust_state: CandidateTrustState::CandidateOnly,
        })
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum CandidateTrustState {
    CandidateOnly,
}

pub trait DiscoveryBackend {
    fn candidates(&self) -> Result<Vec<DiscoveredPeer>, DiscoveryError>;
}

pub fn peer_card_txt_attributes(
    card: &PeerCard,
) -> Result<Vec<(String, Option<String>)>, DiscoveryError> {
    card.validate_candidate()?;
    let json = serde_json::to_vec(card)?;
    let encoded = base64::engine::general_purpose::STANDARD_NO_PAD.encode(json);
    let chunks = encoded
        .as_bytes()
        .chunks(PEER_CARD_TXT_CHUNK_BYTES)
        .map(|chunk| {
            std::str::from_utf8(chunk)
                .expect("base64 is utf-8")
                .to_owned()
        })
        .collect::<Vec<_>>();
    if chunks.is_empty() || chunks.len() > PEER_CARD_TXT_MAX_CHUNKS {
        return Err(DiscoveryError::PeerCardTxtTooLarge {
            chunks: chunks.len(),
            max_chunks: PEER_CARD_TXT_MAX_CHUNKS,
        });
    }

    let mut attributes = vec![
        (
            PEER_CARD_TXT_VERSION_KEY.to_owned(),
            Some(PEER_CARD_TXT_VERSION.to_owned()),
        ),
        (
            PEER_CARD_TXT_CHUNKS_KEY.to_owned(),
            Some(chunks.len().to_string()),
        ),
    ];
    for (index, chunk) in chunks.into_iter().enumerate() {
        attributes.push((
            format!("{PEER_CARD_TXT_CHUNK_PREFIX}{index:02}"),
            Some(chunk),
        ));
    }
    Ok(attributes)
}

pub fn peer_card_from_txt_attributes<'a>(
    attributes: impl IntoIterator<Item = (&'a str, Option<&'a str>)>,
) -> Result<PeerCard, DiscoveryError> {
    let attributes = attributes.into_iter().collect::<Vec<_>>();
    let version = txt_value(&attributes, PEER_CARD_TXT_VERSION_KEY)
        .ok_or_else(|| DiscoveryError::MissingTxtAttribute(PEER_CARD_TXT_VERSION_KEY.to_owned()))?;
    if version != PEER_CARD_TXT_VERSION {
        return Err(DiscoveryError::UnsupportedTxtVersion(version.to_owned()));
    }
    let chunk_count = txt_value(&attributes, PEER_CARD_TXT_CHUNKS_KEY)
        .ok_or_else(|| DiscoveryError::MissingTxtAttribute(PEER_CARD_TXT_CHUNKS_KEY.to_owned()))?
        .parse::<usize>()
        .map_err(|_| DiscoveryError::InvalidTxtAttribute(PEER_CARD_TXT_CHUNKS_KEY.to_owned()))?;
    if chunk_count == 0 || chunk_count > PEER_CARD_TXT_MAX_CHUNKS {
        return Err(DiscoveryError::PeerCardTxtTooLarge {
            chunks: chunk_count,
            max_chunks: PEER_CARD_TXT_MAX_CHUNKS,
        });
    }

    let mut encoded = String::new();
    for index in 0..chunk_count {
        let key = format!("{PEER_CARD_TXT_CHUNK_PREFIX}{index:02}");
        let chunk = txt_value(&attributes, &key)
            .ok_or_else(|| DiscoveryError::MissingTxtAttribute(key.clone()))?;
        encoded.push_str(chunk);
    }
    let json = base64::engine::general_purpose::STANDARD_NO_PAD.decode(encoded)?;
    let card: PeerCard = serde_json::from_slice(&json)?;
    card.validate_candidate()?;
    Ok(card)
}

fn txt_value<'a>(attributes: &[(&'a str, Option<&'a str>)], key: &str) -> Option<&'a str> {
    attributes
        .iter()
        .find_map(|(candidate, value)| (*candidate == key).then_some(*value).flatten())
}

#[derive(Debug, thiserror::Error)]
pub enum DiscoveryError {
    #[error("peer card has no endpoint candidates")]
    MissingEndpoint,
    #[error("peer card is missing signature metadata")]
    UnsignedPeerCard,
    #[error("invalid peer card signature namespace: {0}")]
    InvalidSignatureNamespace(String),
    #[error("peer card signer {signer} does not match agent {agent}")]
    SignerMismatch { signer: String, agent: String },
    #[error("peer card signer {signer} does not match public key fingerprint {fingerprint}")]
    SignerPublicKeyMismatch { signer: String, fingerprint: String },
    #[error("invalid peer card hex: {0}")]
    Hex(#[from] hex::FromHexError),
    #[error("peer card signature error: {0}")]
    Crypto(#[from] geth_crypto::CryptoError),
    #[error("peer card TXT payload is too large: {chunks} chunks, max {max_chunks}")]
    PeerCardTxtTooLarge { chunks: usize, max_chunks: usize },
    #[error("missing peer card TXT attribute: {0}")]
    MissingTxtAttribute(String),
    #[error("invalid peer card TXT attribute: {0}")]
    InvalidTxtAttribute(String),
    #[error("unsupported peer card TXT version: {0}")]
    UnsupportedTxtVersion(String),
    #[error("peer card TXT base64 error: {0}")]
    Base64(#[from] base64::DecodeError),
    #[error("peer card JSON error: {0}")]
    Json(#[from] serde_json::Error),
}

#[must_use]
pub fn discovery_is_untrusted_note() -> &'static str {
    "discovery returns candidate peers only and never grants trust or authorization"
}

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

    fn signed_card() -> PeerCard {
        let key = AgentKey::generate();
        let agent_id = key.agent_id();
        PeerCard::signed(
            format!("node:{agent_id}").into(),
            &key,
            vec![EndpointCandidate {
                endpoint_id: "endpoint:iroh".to_owned(),
                relay_url: None,
                direct_addresses: vec!["127.0.0.1:12345".to_owned()],
                source: DiscoverySource::Manual,
            }],
            UnixMillis(1),
        )
        .expect("signed card")
    }

    #[test]
    fn signed_peer_card_is_valid_candidate() {
        signed_card()
            .validate_candidate()
            .expect("valid signed candidate");
    }

    #[test]
    fn discovered_peer_is_candidate_only() {
        let peer = DiscoveredPeer::candidate(signed_card(), UnixMillis(2), DiscoverySource::Mdns)
            .expect("candidate");
        assert_eq!(peer.trust_state, CandidateTrustState::CandidateOnly);
    }

    #[test]
    fn unsigned_peer_card_is_not_valid_candidate() {
        let mut card = signed_card();
        card.signature.signature.clear();
        assert!(matches!(
            card.validate_candidate(),
            Err(DiscoveryError::UnsignedPeerCard)
        ));
    }

    #[test]
    fn tampered_peer_card_signature_is_rejected() {
        let mut card = signed_card();
        card.endpoints[0].endpoint_id = "endpoint:tampered".to_owned();
        assert!(matches!(
            card.validate_candidate(),
            Err(DiscoveryError::Crypto(geth_crypto::CryptoError::Verify))
        ));
    }

    #[test]
    fn peer_card_txt_attributes_roundtrip_signed_payload() {
        let card = signed_card();
        let attributes = peer_card_txt_attributes(&card).expect("txt attributes");
        assert_eq!(
            attributes
                .iter()
                .find(|(key, _)| key == PEER_CARD_TXT_VERSION_KEY)
                .and_then(|(_, value)| value.as_deref()),
            Some(PEER_CARD_TXT_VERSION)
        );
        assert!(attributes.iter().any(|(key, _)| key == "card-00"));
        assert!(
            attributes
                .iter()
                .all(|(key, value)| key.len() + value.as_deref().unwrap_or("").len() <= 254)
        );

        let decoded = peer_card_from_txt_attributes(
            attributes
                .iter()
                .map(|(key, value)| (key.as_str(), value.as_deref())),
        )
        .expect("decode txt attributes");

        assert_eq!(decoded, card);
    }

    #[test]
    fn peer_card_txt_attributes_reject_missing_chunk() {
        let card = signed_card();
        let attributes = peer_card_txt_attributes(&card).expect("txt attributes");
        let without_first_chunk = attributes
            .iter()
            .filter(|(key, _)| key != "card-00")
            .map(|(key, value)| (key.as_str(), value.as_deref()));

        assert!(matches!(
            peer_card_from_txt_attributes(without_first_chunk),
            Err(DiscoveryError::MissingTxtAttribute(key)) if key == "card-00"
        ));
    }
}