openrtc 1.0.1

#![cfg(not(target_arch = "wasm32"))]

use anyhow::Result;
use async_trait::async_trait;
use iroh::{
    protocol::{DynProtocolHandler, RouterBuilder},
    Endpoint,
};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Clone)]
pub struct RuntimeContext {
    pub endpoint: Endpoint,
    pub node_id: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProtocolPluginKind {
    /// A plugin that is notified about runtime readiness and may open streams
    /// over already-established peer sessions, but does not contribute a new
    /// ALPN handler to the endpoint/router.
    SessionAdjacent,
    /// A plugin that contributes a custom ALPN protocol and expects the native
    /// runtime manager/router to register it before endpoint bind.
    Alpn,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProtocolPluginReadiness {
    Pending,
    Ready,
    Degraded(String),
}

impl ProtocolPluginReadiness {
    pub fn is_ready(&self) -> bool {
        matches!(self, Self::Ready)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProtocolPluginDescriptor {
    pub name: String,
    pub kind: ProtocolPluginKind,
    pub alpn: Option<Vec<u8>>,
    pub has_router_handler: bool,
    pub readiness: ProtocolPluginReadiness,
}

#[async_trait]
pub trait ProtocolPlugin: Send + Sync {
    /// Delegation rule:
    /// - Rust owns protocol registration and runtime lifecycle
    /// - plugins describe how they integrate with the runtime, but they do not
    ///   become a second lifecycle engine
    fn kind(&self) -> ProtocolPluginKind {
        if self.alpn().is_some() {
            ProtocolPluginKind::Alpn
        } else {
            ProtocolPluginKind::SessionAdjacent
        }
    }
    fn name(&self) -> &'static str;
    fn alpn(&self) -> Option<Vec<u8>>;

    /// Optional native router handler.
    ///
    /// Plugins that return a handler here can be wired directly into an iroh
    /// `RouterBuilder` through `ProtocolRegistry::apply_router_handlers()`.
    /// Session-adjacent plugins should leave this as `None`.
    fn router_handler(&self) -> Option<Box<dyn DynProtocolHandler>> {
        None
    }

    fn has_router_handler(&self) -> bool {
        false
    }

    fn readiness(&self) -> ProtocolPluginReadiness {
        ProtocolPluginReadiness::Ready
    }

    fn descriptor(&self) -> ProtocolPluginDescriptor {
        ProtocolPluginDescriptor {
            name: self.name().to_string(),
            kind: self.kind(),
            alpn: self.alpn(),
            has_router_handler: self.has_router_handler(),
            readiness: self.readiness(),
        }
    }

    async fn on_runtime_ready(&self, context: RuntimeContext) -> Result<()>;
}

#[derive(Default)]
pub struct ProtocolRegistry {
    plugins: Arc<RwLock<HashMap<String, Arc<dyn ProtocolPlugin>>>>,
}

impl ProtocolRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    pub async fn register(&self, plugin: Arc<dyn ProtocolPlugin>) {
        let mut guard = self.plugins.write().await;
        guard.insert(plugin.name().to_string(), plugin);
    }

    pub async fn register_checked(&self, plugin: Arc<dyn ProtocolPlugin>) -> Result<()> {
        let mut guard = self.plugins.write().await;
        let name = plugin.name().to_string();
        if guard.contains_key(&name) {
            return Err(anyhow::anyhow!(
                "protocol plugin '{}' is already registered",
                name
            ));
        }

        if plugin.kind() == ProtocolPluginKind::Alpn {
            let alpn = plugin.alpn().ok_or_else(|| {
                anyhow::anyhow!("ALPN protocol plugin '{}' must declare an ALPN", name)
            })?;
            let duplicate = guard.values().any(|existing| {
                existing.kind() == ProtocolPluginKind::Alpn
                    && existing.alpn().as_deref() == Some(alpn.as_slice())
            });
            if duplicate {
                return Err(anyhow::anyhow!(
                    "ALPN protocol plugin '{}' conflicts with an already registered ALPN",
                    name
                ));
            }
        }

        guard.insert(name, plugin);
        Ok(())
    }

    pub async fn unregister(&self, name: &str) -> Option<Arc<dyn ProtocolPlugin>> {
        let mut guard = self.plugins.write().await;
        guard.remove(name)
    }

    pub async fn contains(&self, name: &str) -> bool {
        self.plugins.read().await.contains_key(name)
    }

    pub async fn list(&self) -> Vec<String> {
        let guard = self.plugins.read().await;
        let mut names: Vec<String> = guard.keys().cloned().collect();
        names.sort();
        names
    }

    pub async fn extra_alpns(&self) -> Vec<Vec<u8>> {
        let guard = self.plugins.read().await;
        guard
            .values()
            .filter(|plugin| plugin.kind() == ProtocolPluginKind::Alpn)
            .filter_map(|plugin| plugin.alpn())
            .collect()
    }

    pub async fn descriptors(&self) -> Vec<ProtocolPluginDescriptor> {
        let guard = self.plugins.read().await;
        let mut descriptors: Vec<ProtocolPluginDescriptor> =
            guard.values().map(|plugin| plugin.descriptor()).collect();
        descriptors.sort_by(|a, b| a.name.cmp(&b.name));
        descriptors
    }

    pub async fn apply_router_handlers(&self, mut builder: RouterBuilder) -> Result<RouterBuilder> {
        let plugins: Vec<Arc<dyn ProtocolPlugin>> =
            self.plugins.read().await.values().cloned().collect();

        for plugin in plugins {
            if plugin.kind() != ProtocolPluginKind::Alpn {
                continue;
            }
            let Some(handler) = plugin.router_handler() else {
                continue;
            };
            let Some(alpn) = plugin.alpn() else {
                return Err(anyhow::anyhow!(
                    "ALPN protocol plugin '{}' returned a handler without an ALPN",
                    plugin.name()
                ));
            };
            builder = builder.accept(alpn, handler);
        }

        Ok(builder)
    }

    pub async fn notify_runtime_ready(&self, context: RuntimeContext) -> Result<()> {
        let plugins: Vec<Arc<dyn ProtocolPlugin>> =
            self.plugins.read().await.values().cloned().collect();

        for plugin in plugins {
            plugin.on_runtime_ready(context.clone()).await?;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use iroh::{
        endpoint::Connection,
        protocol::{AcceptError, ProtocolHandler},
    };
    use std::sync::atomic::{AtomicUsize, Ordering};

    struct TestPlugin {
        calls: Arc<AtomicUsize>,
    }

    #[async_trait]
    impl ProtocolPlugin for TestPlugin {
        fn name(&self) -> &'static str {
            "test"
        }

        fn alpn(&self) -> Option<Vec<u8>> {
            Some(b"pluto/test/1".to_vec())
        }

        fn readiness(&self) -> ProtocolPluginReadiness {
            ProtocolPluginReadiness::Ready
        }

        async fn on_runtime_ready(&self, _context: RuntimeContext) -> Result<()> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            Ok(())
        }
    }

    #[tokio::test]
    async fn registry_tracks_plugins_and_alpns() {
        let registry = ProtocolRegistry::new();
        let calls = Arc::new(AtomicUsize::new(0));

        registry
            .register(Arc::new(TestPlugin {
                calls: calls.clone(),
            }))
            .await;

        assert!(registry.contains("test").await);
        assert_eq!(registry.list().await, vec!["test".to_string()]);
        assert_eq!(registry.extra_alpns().await, vec![b"pluto/test/1".to_vec()]);
        assert_eq!(
            registry.descriptors().await,
            vec![ProtocolPluginDescriptor {
                name: "test".to_string(),
                kind: ProtocolPluginKind::Alpn,
                alpn: Some(b"pluto/test/1".to_vec()),
                has_router_handler: false,
                readiness: ProtocolPluginReadiness::Ready,
            }]
        );

        let endpoint = iroh::Endpoint::builder(iroh::endpoint::presets::N0)
            .bind()
            .await
            .expect("bind endpoint");
        let context = RuntimeContext {
            node_id: endpoint.id().to_string(),
            endpoint,
        };

        registry
            .notify_runtime_ready(context)
            .await
            .expect("notify plugins");

        assert_eq!(calls.load(Ordering::SeqCst), 1);
    }

    struct SessionPlugin;

    #[async_trait]
    impl ProtocolPlugin for SessionPlugin {
        fn kind(&self) -> ProtocolPluginKind {
            ProtocolPluginKind::SessionAdjacent
        }

        fn name(&self) -> &'static str {
            "session"
        }

        fn alpn(&self) -> Option<Vec<u8>> {
            Some(b"pluto/ignored/1".to_vec())
        }

        async fn on_runtime_ready(&self, _context: RuntimeContext) -> Result<()> {
            Ok(())
        }
    }

    #[tokio::test]
    async fn extra_alpns_only_includes_alpn_plugins() {
        let registry = ProtocolRegistry::new();
        registry.register(Arc::new(SessionPlugin)).await;
        let alpns = registry.extra_alpns().await;
        assert!(alpns.is_empty());
    }

    #[derive(Debug)]
    struct NoopHandler;

    impl ProtocolHandler for NoopHandler {
        async fn accept(&self, _connection: Connection) -> std::result::Result<(), AcceptError> {
            Ok(())
        }
    }

    struct HandlerPlugin;

    #[async_trait]
    impl ProtocolPlugin for HandlerPlugin {
        fn name(&self) -> &'static str {
            "handler"
        }

        fn alpn(&self) -> Option<Vec<u8>> {
            Some(b"pluto/handler/1".to_vec())
        }

        fn has_router_handler(&self) -> bool {
            true
        }

        fn router_handler(&self) -> Option<Box<dyn DynProtocolHandler>> {
            Some(Box::new(NoopHandler))
        }

        async fn on_runtime_ready(&self, _context: RuntimeContext) -> Result<()> {
            Ok(())
        }
    }

    #[tokio::test]
    async fn registry_applies_native_router_handlers() {
        let registry = ProtocolRegistry::new();
        registry.register(Arc::new(HandlerPlugin)).await;

        let endpoint = iroh::Endpoint::builder(iroh::endpoint::presets::N0)
            .alpns(vec![b"pluto/handler/1".to_vec()])
            .bind()
            .await
            .expect("bind endpoint");
        let builder = iroh::protocol::Router::builder(endpoint);
        let builder = registry
            .apply_router_handlers(builder)
            .await
            .expect("apply registered handlers");
        let router = builder.spawn();

        assert_eq!(
            registry.descriptors().await,
            vec![ProtocolPluginDescriptor {
                name: "handler".to_string(),
                kind: ProtocolPluginKind::Alpn,
                alpn: Some(b"pluto/handler/1".to_vec()),
                has_router_handler: true,
                readiness: ProtocolPluginReadiness::Ready,
            }]
        );

        router.shutdown().await.expect("shutdown router");
    }
}