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

use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::net::{SocketAddrV4, SocketAddrV6};
use std::path::{Path, PathBuf};
use std::time::Duration;

pub const ALPN_CONTROL: &[u8] = b"/geth/control/1";
pub const ALPN_KV: &[u8] = b"/geth/kv/1";
pub const ALPN_CAS: &[u8] = b"/geth/cas/1";
pub const ALPN_PUBSUB: &[u8] = b"/geth/pubsub/1";
pub const ALPN_PIPE: &[u8] = b"/geth/pipe/1";
pub const ALPN_DB: &[u8] = b"/geth/db/1";
pub const ALPN_DOCUMENT: &[u8] = b"/geth/document/1";
pub const ALPN_SSH_PROXY: &[u8] = b"/geth/ssh-proxy/1";

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct GethIrohConfig {
    pub secret_key_path: PathBuf,
    pub relay_mode: GethRelayMode,
    pub local_discovery: bool,
    pub bind_ipv4: Option<SocketAddrV4>,
    pub bind_ipv6: Option<SocketAddrV6>,
    pub alpns: Vec<Vec<u8>>,
}

impl GethIrohConfig {
    #[must_use]
    pub fn local(secret_key_path: impl Into<PathBuf>) -> Self {
        Self::local_with_relay(secret_key_path, GethRelayMode::Default)
    }

    #[must_use]
    pub fn local_with_relay(
        secret_key_path: impl Into<PathBuf>,
        relay_mode: GethRelayMode,
    ) -> Self {
        Self {
            secret_key_path: secret_key_path.into(),
            relay_mode,
            local_discovery: true,
            bind_ipv4: None,
            bind_ipv6: None,
            alpns: default_protocol_router().alpns(),
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum ProtocolKind {
    Control,
    Kv,
    Cas,
    Pubsub,
    Pipe,
    Db,
    Document,
    SshProxy,
}

impl ProtocolKind {
    #[must_use]
    pub fn name(self) -> &'static str {
        match self {
            Self::Control => "control",
            Self::Kv => "kv",
            Self::Cas => "cas",
            Self::Pubsub => "pubsub",
            Self::Pipe => "pipe",
            Self::Db => "db",
            Self::Document => "document",
            Self::SshProxy => "ssh-proxy",
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProtocolDescriptor {
    pub kind: ProtocolKind,
    pub name: String,
    pub alpn: Vec<u8>,
}

impl ProtocolDescriptor {
    #[must_use]
    pub fn new(kind: ProtocolKind, alpn: &'static [u8]) -> Self {
        Self {
            kind,
            name: kind.name().to_owned(),
            alpn: alpn.to_vec(),
        }
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ProtocolRouter {
    protocols: BTreeMap<Vec<u8>, ProtocolDescriptor>,
}

impl ProtocolRouter {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    pub fn register(&mut self, descriptor: ProtocolDescriptor) -> Result<(), RouterError> {
        if self.protocols.contains_key(&descriptor.alpn) {
            return Err(RouterError::DuplicateAlpn {
                alpn: display_alpn(&descriptor.alpn),
            });
        }
        self.protocols.insert(descriptor.alpn.clone(), descriptor);
        Ok(())
    }

    pub fn require(&self, alpn: &[u8]) -> Result<&ProtocolDescriptor, RouterError> {
        self.protocols
            .get(alpn)
            .ok_or_else(|| RouterError::UnknownAlpn {
                alpn: display_alpn(alpn),
            })
    }

    #[must_use]
    pub fn alpns(&self) -> Vec<Vec<u8>> {
        self.protocols.keys().cloned().collect()
    }

    #[must_use]
    pub fn descriptors(&self) -> Vec<ProtocolDescriptor> {
        self.protocols.values().cloned().collect()
    }
}

#[must_use]
pub fn default_protocol_router() -> ProtocolRouter {
    let mut router = ProtocolRouter::new();
    for descriptor in default_protocol_descriptors() {
        router
            .register(descriptor)
            .expect("default geth ALPNs are unique");
    }
    router
}

#[must_use]
pub fn default_protocol_descriptors() -> Vec<ProtocolDescriptor> {
    vec![
        ProtocolDescriptor::new(ProtocolKind::Control, ALPN_CONTROL),
        ProtocolDescriptor::new(ProtocolKind::Kv, ALPN_KV),
        ProtocolDescriptor::new(ProtocolKind::Cas, ALPN_CAS),
        ProtocolDescriptor::new(ProtocolKind::Pubsub, ALPN_PUBSUB),
        ProtocolDescriptor::new(ProtocolKind::Pipe, ALPN_PIPE),
        ProtocolDescriptor::new(ProtocolKind::Db, ALPN_DB),
        ProtocolDescriptor::new(ProtocolKind::Document, ALPN_DOCUMENT),
        ProtocolDescriptor::new(ProtocolKind::SshProxy, ALPN_SSH_PROXY),
    ]
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum GethRelayMode {
    Disabled,
    Default,
    Staging,
    Custom {
        name: String,
        relay_urls: Vec<String>,
    },
}

impl GethRelayMode {
    fn to_iroh(&self) -> Result<iroh::RelayMode, IrohError> {
        match self {
            Self::Disabled => Ok(iroh::RelayMode::Disabled),
            Self::Default => Ok(iroh::RelayMode::Default),
            Self::Staging => Ok(iroh::RelayMode::Staging),
            Self::Custom { relay_urls, .. } => {
                let relay_urls = relay_urls
                    .iter()
                    .map(|url| {
                        url.parse::<iroh::RelayUrl>()
                            .map_err(|error| IrohError::InvalidRelayUrl {
                                url: url.clone(),
                                message: error.to_string(),
                            })
                    })
                    .collect::<Result<Vec<_>, _>>()?;
                Ok(iroh::RelayMode::Custom(iroh::RelayMap::from_iter(
                    relay_urls,
                )))
            }
        }
    }

    #[must_use]
    pub fn label(&self) -> String {
        match self {
            Self::Disabled => "disabled".to_owned(),
            Self::Default => "default".to_owned(),
            Self::Staging => "staging".to_owned(),
            Self::Custom { name, .. } => format!("custom:{name}"),
        }
    }
}

#[derive(Clone, Debug)]
pub struct GethIrohEndpoint {
    endpoint: iroh::Endpoint,
    status: EndpointStatus,
}

impl GethIrohEndpoint {
    #[must_use]
    pub fn status(&self) -> EndpointStatus {
        self.status.clone()
    }

    #[must_use]
    pub fn node_id(&self) -> String {
        self.endpoint.node_id().to_string()
    }

    #[must_use]
    pub fn endpoint(&self) -> iroh::Endpoint {
        self.endpoint.clone()
    }

    pub async fn node_addr_snapshot(&self) -> Result<GethNodeAddr, IrohError> {
        use n0_watcher::Watcher;

        let mut watcher = self.endpoint.node_addr();
        let node_addr = tokio::time::timeout(Duration::from_secs(2), watcher.initialized())
            .await
            .map_err(|_| IrohError::NodeAddrTimeout)?
            .map_err(|error| IrohError::NodeAddrUnavailable(error.to_string()))?;
        Ok(GethNodeAddr {
            endpoint_id: node_addr.node_id.to_string(),
            relay_url: node_addr.relay_url.map(|url| url.to_string()),
            direct_addresses: node_addr
                .direct_addresses
                .into_iter()
                .map(|addr| addr.to_string())
                .collect(),
        })
    }

    pub async fn shutdown(&self) {
        self.endpoint.close().await;
    }
}

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

pub async fn start_endpoint(config: &GethIrohConfig) -> Result<GethIrohEndpoint, IrohError> {
    let secret_key = load_or_create_secret_key(&config.secret_key_path)?;
    let mut builder = iroh::Endpoint::builder()
        .secret_key(secret_key)
        .relay_mode(config.relay_mode.to_iroh()?)
        .alpns(config.alpns.clone());

    if config.local_discovery {
        builder = builder.discovery_local_network();
    }

    if let Some(bind_ipv4) = config.bind_ipv4 {
        builder = builder.bind_addr_v4(bind_ipv4);
    }
    if let Some(bind_ipv6) = config.bind_ipv6 {
        builder = builder.bind_addr_v6(bind_ipv6);
    }

    let endpoint = builder.bind().await.map_err(IrohError::from)?;
    let status = EndpointStatus {
        enabled: true,
        endpoint_id: Some(endpoint.node_id().to_string()),
        relay_mode: config.relay_mode.label(),
        local_discovery: config.local_discovery,
        note: format!(
            "Iroh endpoint started with iroh 0.90.0; relay mode: {:?}; local discovery: {}",
            config.relay_mode, config.local_discovery
        ),
    };
    Ok(GethIrohEndpoint { endpoint, status })
}

pub fn load_or_create_secret_key(path: &Path) -> Result<iroh::SecretKey, IrohError> {
    if path.exists() {
        return load_secret_key(path);
    }

    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    let secret_key = iroh::SecretKey::generate(rand_core::OsRng);
    let tmp = path.with_extension("tmp");
    std::fs::write(&tmp, hex::encode(secret_key.to_bytes()))?;
    std::fs::rename(tmp, path)?;
    Ok(secret_key)
}

pub fn load_secret_key(path: &Path) -> Result<iroh::SecretKey, IrohError> {
    let encoded = std::fs::read_to_string(path)?;
    let bytes = hex::decode(encoded.trim())?;
    let bytes: [u8; 32] = bytes
        .try_into()
        .map_err(|_| IrohError::InvalidSecretKeyLength)?;
    Ok(iroh::SecretKey::from_bytes(&bytes))
}

#[must_use]
pub fn all_alpns() -> Vec<Vec<u8>> {
    default_protocol_router().alpns()
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct EndpointStatus {
    pub enabled: bool,
    pub endpoint_id: Option<String>,
    pub relay_mode: String,
    pub local_discovery: bool,
    pub note: String,
}

impl EndpointStatus {
    #[must_use]
    pub fn scaffolded() -> Self {
        Self {
            enabled: false,
            endpoint_id: None,
            relay_mode: "not-started".to_owned(),
            local_discovery: false,
            note: "Iroh endpoint is not running outside daemon mode".to_owned(),
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum IrohError {
    #[error("io error: {0}")]
    Io(#[from] std::io::Error),
    #[error("invalid hex iroh secret key: {0}")]
    Hex(#[from] hex::FromHexError),
    #[error("invalid iroh secret key length")]
    InvalidSecretKeyLength,
    #[error("invalid iroh relay URL `{url}`: {message}")]
    InvalidRelayUrl { url: String, message: String },
    #[error("failed to bind iroh endpoint: {0}")]
    Bind(Box<iroh::endpoint::BindError>),
    #[error("timed out waiting for iroh node address")]
    NodeAddrTimeout,
    #[error("iroh node address watcher is unavailable: {0}")]
    NodeAddrUnavailable(String),
}

#[derive(Debug, thiserror::Error)]
pub enum RouterError {
    #[error("duplicate geth ALPN registration: {alpn}")]
    DuplicateAlpn { alpn: String },
    #[error("unknown geth ALPN: {alpn}")]
    UnknownAlpn { alpn: String },
}

impl From<iroh::endpoint::BindError> for IrohError {
    fn from(error: iroh::endpoint::BindError) -> Self {
        Self::Bind(Box::new(error))
    }
}

fn display_alpn(alpn: &[u8]) -> String {
    String::from_utf8(alpn.to_vec()).unwrap_or_else(|_| format!("0x{}", hex::encode(alpn)))
}

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

    #[test]
    fn all_alpns_contains_expected_protocols() {
        let alpns = all_alpns();
        assert!(alpns.contains(&ALPN_CONTROL.to_vec()));
        assert!(alpns.contains(&ALPN_CAS.to_vec()));
        assert!(alpns.contains(&ALPN_SSH_PROXY.to_vec()));
        assert_eq!(alpns.len(), 8);
    }

    #[test]
    fn default_router_registers_all_protocols() {
        let router = default_protocol_router();
        assert_eq!(
            router.require(ALPN_CONTROL).expect("control").kind,
            ProtocolKind::Control
        );
        assert_eq!(
            router.require(ALPN_SSH_PROXY).expect("ssh proxy").kind,
            ProtocolKind::SshProxy
        );
        assert_eq!(router.descriptors().len(), 8);
    }

    #[test]
    fn router_rejects_duplicate_alpns() {
        let mut router = ProtocolRouter::new();
        router
            .register(ProtocolDescriptor::new(ProtocolKind::Control, ALPN_CONTROL))
            .expect("first registration");
        let error = router
            .register(ProtocolDescriptor::new(ProtocolKind::Cas, ALPN_CONTROL))
            .expect_err("duplicate error");
        assert!(matches!(error, RouterError::DuplicateAlpn { .. }));
    }

    #[test]
    fn router_rejects_unknown_alpn() {
        let router = default_protocol_router();
        let error = router
            .require(b"/geth/unknown/1")
            .expect_err("unknown error");
        assert!(matches!(error, RouterError::UnknownAlpn { .. }));
    }

    #[test]
    fn iroh_secret_key_persists() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("iroh.ed25519");
        let first = load_or_create_secret_key(&path).expect("create key");
        let second = load_or_create_secret_key(&path).expect("load key");
        assert_eq!(first.public(), second.public());
    }

    #[test]
    fn local_config_selects_relay_mode() {
        let config = GethIrohConfig::local_with_relay("iroh.ed25519", GethRelayMode::Staging);
        assert_eq!(config.relay_mode, GethRelayMode::Staging);
        assert!(matches!(
            config.relay_mode.to_iroh().expect("iroh relay mode"),
            iroh::RelayMode::Staging
        ));
    }

    #[test]
    fn custom_relay_mode_builds_iroh_relay_map() {
        let mode = GethRelayMode::Custom {
            name: "home".to_owned(),
            relay_urls: vec!["https://relay.example.com".to_owned()],
        };
        let relay_mode = mode.to_iroh().expect("custom relay mode");
        match relay_mode {
            iroh::RelayMode::Custom(map) => {
                assert_eq!(map.len(), 1);
                assert_eq!(mode.label(), "custom:home");
            }
            other => panic!("unexpected relay mode: {other:?}"),
        }
    }

    #[tokio::test]
    async fn endpoint_status_tracks_node_id_when_bind_is_available() {
        let dir = tempfile::tempdir().expect("tempdir");
        let mut config = GethIrohConfig::local_with_relay(
            dir.path().join("iroh.ed25519"),
            GethRelayMode::Disabled,
        );
        config.local_discovery = false;
        match start_endpoint(&config).await {
            Ok(endpoint) => {
                let status = endpoint.status();
                assert!(status.enabled);
                assert_eq!(status.endpoint_id, Some(endpoint.node_id()));
                assert_eq!(status.relay_mode, "disabled");
                assert!(!status.local_discovery);
                let node_addr = endpoint.node_addr_snapshot().await.expect("node addr");
                assert_eq!(node_addr.endpoint_id, endpoint.node_id());
                assert!(!node_addr.direct_addresses.is_empty());
                endpoint.shutdown().await;
            }
            Err(IrohError::Bind(error)) => {
                eprintln!("skipping endpoint bind assertion; UDP unavailable: {error}");
            }
            Err(error) => panic!("unexpected iroh startup error: {error}"),
        }
    }
}