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

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

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 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 {
            secret_key_path: secret_key_path.into(),
            relay_mode: GethRelayMode::Disabled,
            bind_ipv4: None,
            bind_ipv6: None,
            alpns: all_alpns(),
        }
    }
}

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

impl GethRelayMode {
    fn to_iroh(&self) -> iroh::RelayMode {
        match self {
            Self::Disabled => iroh::RelayMode::Disabled,
            Self::Default => iroh::RelayMode::Default,
            Self::Staging => iroh::RelayMode::Staging,
        }
    }
}

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

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

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 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()),
        note: format!(
            "Iroh endpoint started with iroh 0.90.0; relay mode: {:?}",
            config.relay_mode
        ),
    };
    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>> {
    [
        ALPN_CONTROL,
        ALPN_KV,
        ALPN_CAS,
        ALPN_PUBSUB,
        ALPN_PIPE,
        ALPN_DB,
        ALPN_DOCUMENT,
        ALPN_SSH_PROXY,
    ]
    .into_iter()
    .map(<[u8]>::to_vec)
    .collect()
}

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

impl EndpointStatus {
    #[must_use]
    pub fn scaffolded() -> Self {
        Self {
            enabled: false,
            endpoint_id: None,
            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("failed to bind iroh endpoint: {0}")]
    Bind(Box<iroh::endpoint::BindError>),
}

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

#[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 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());
    }

    #[tokio::test]
    async fn endpoint_status_tracks_node_id_when_bind_is_available() {
        let dir = tempfile::tempdir().expect("tempdir");
        let config = GethIrohConfig::local(dir.path().join("iroh.ed25519"));
        match start_endpoint(&config).await {
            Ok(endpoint) => {
                let status = endpoint.status();
                assert!(status.enabled);
                assert_eq!(status.endpoint_id, Some(endpoint.node_id()));
                endpoint.shutdown().await;
            }
            Err(IrohError::Bind(error)) => {
                eprintln!("skipping endpoint bind assertion; UDP unavailable: {error}");
            }
            Err(error) => panic!("unexpected iroh startup error: {error}"),
        }
    }
}