unb-server 2.0.3

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use unb_core::NodeIdentity;

use crate::handler::HandlerError;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VerifiedPeer {
    node_id: String,
    instance_id: String,
    epoch: u64,
    declared_only: bool,
}

impl VerifiedPeer {
    pub(crate) fn from_identity(identity: &NodeIdentity) -> Self {
        Self {
            node_id: identity.node_id.clone(),
            instance_id: identity.instance_id.clone(),
            epoch: identity.epoch,
            declared_only: false,
        }
    }

    pub fn node_id(&self) -> &str {
        &self.node_id
    }

    pub fn instance_id(&self) -> &str {
        &self.instance_id
    }

    pub fn epoch(&self) -> u64 {
        self.epoch
    }

    pub fn declared_only(&self) -> bool {
        self.declared_only
    }
}

pub struct PeerRequest {
    local: NodeIdentity,
    remote: NodeIdentity,
    extensions: http::Extensions,
}

impl PeerRequest {
    pub(crate) fn new(local: NodeIdentity, remote: NodeIdentity) -> PeerRequest {
        PeerRequest {
            local,
            remote,
            extensions: http::Extensions::new(),
        }
    }

    pub fn local(&self) -> &NodeIdentity {
        &self.local
    }

    pub fn remote(&self) -> &NodeIdentity {
        &self.remote
    }

    pub fn extensions(&self) -> &http::Extensions {
        &self.extensions
    }

    pub fn extensions_mut(&mut self) -> &mut http::Extensions {
        &mut self.extensions
    }

    pub fn accept(&mut self) {
        self.insert_verified(false);
    }

    pub fn accept_declared(&mut self) {
        self.insert_verified(true);
    }

    pub(crate) fn verified(&self) -> Option<VerifiedPeer> {
        self.extensions.get::<VerifiedPeer>().cloned()
    }

    fn insert_verified(&mut self, declared_only: bool) {
        self.extensions.insert(VerifiedPeer {
            node_id: self.remote.node_id.clone(),
            instance_id: self.remote.instance_id.clone(),
            epoch: self.remote.epoch,
            declared_only,
        });
    }
}

pub trait PeerLayer: Send + Sync + 'static {
    fn admit(
        &self,
        request: PeerRequest,
        next: PeerNext,
    ) -> Pin<Box<dyn Future<Output = Result<PeerRequest, HandlerError>> + Send + '_>>;
}

pub struct PeerNext {
    layers: Arc<[Arc<dyn PeerLayer>]>,
    index: usize,
}

impl PeerNext {
    pub(crate) fn root(layers: Arc<[Arc<dyn PeerLayer>]>) -> PeerNext {
        PeerNext { layers, index: 0 }
    }

    pub async fn admit(mut self, request: PeerRequest) -> Result<PeerRequest, HandlerError> {
        if self.index < self.layers.len() {
            let layer = self.layers[self.index].clone();
            self.index += 1;
            layer.admit(request, self).await
        } else {
            Ok(request)
        }
    }
}

pub struct InsecureAcceptDeclaredPeerIdentities;

impl PeerLayer for InsecureAcceptDeclaredPeerIdentities {
    fn admit(
        &self,
        mut request: PeerRequest,
        next: PeerNext,
    ) -> Pin<Box<dyn Future<Output = Result<PeerRequest, HandlerError>> + Send + '_>> {
        request.accept_declared();
        Box::pin(next.admit(request))
    }
}

pub struct PeerLayerFn<F>(F);

impl<F, Fut> PeerLayer for PeerLayerFn<F>
where
    F: Fn(PeerRequest, PeerNext) -> Fut + Send + Sync + 'static,
    Fut: Future<Output = Result<PeerRequest, HandlerError>> + Send + 'static,
{
    fn admit(
        &self,
        request: PeerRequest,
        next: PeerNext,
    ) -> Pin<Box<dyn Future<Output = Result<PeerRequest, HandlerError>> + Send + '_>> {
        Box::pin((self.0)(request, next))
    }
}

pub fn peer_layer_fn<F, Fut>(f: F) -> PeerLayerFn<F>
where
    F: Fn(PeerRequest, PeerNext) -> Fut + Send + Sync + 'static,
    Fut: Future<Output = Result<PeerRequest, HandlerError>> + Send + 'static,
{
    PeerLayerFn(f)
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::Value;
    use unb_core::NodeIdentity;

    fn identity(node: &str) -> NodeIdentity {
        NodeIdentity {
            node_id: node.into(),
            instance_id: format!("{node}-i"),
            epoch: 1,
            proof: Value::Null,
        }
    }

    fn chain(layer: impl PeerLayer) -> PeerNext {
        let layers: Arc<[Arc<dyn PeerLayer>]> =
            Arc::from(vec![Arc::new(layer) as Arc<dyn PeerLayer>]);
        PeerNext::root(layers)
    }

    #[tokio::test]
    async fn a_closure_that_accepts_and_continues_admits() {
        let next = chain(peer_layer_fn(
            |mut request: PeerRequest, next: PeerNext| async move {
                request.accept();
                next.admit(request).await
            },
        ));
        let admitted = next
            .admit(PeerRequest::new(identity("local"), identity("remote")))
            .await
            .unwrap();
        let verified = admitted.verified().unwrap();
        assert!(!verified.declared_only());
    }

    #[tokio::test]
    async fn a_closure_that_continues_without_accepting_produces_no_verified_peer() {
        let next = chain(peer_layer_fn(
            |request: PeerRequest, next: PeerNext| async move { next.admit(request).await },
        ));
        let admitted = next
            .admit(PeerRequest::new(identity("local"), identity("remote")))
            .await
            .unwrap();
        assert!(admitted.verified().is_none());
    }

    #[tokio::test]
    async fn a_closure_that_errors_rejects() {
        let next = chain(peer_layer_fn(
            |_request: PeerRequest, _next: PeerNext| async move {
                Err(HandlerError::new(crate::ErrorCode::Unauthorized, "denied"))
            },
        ));
        let result = next
            .admit(PeerRequest::new(identity("local"), identity("remote")))
            .await;
        assert!(result.is_err());
    }
}