xds-server 0.1.0

//! Secret Discovery Service (SDS) implementation.
//!
//! SDS provides secret (TLS certificate) configuration to Envoy proxies.

use std::sync::Arc;

use async_trait::async_trait;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Request, Response, Status, Streaming};
use tracing::{debug, error, info, instrument};

use xds_cache::ShardedCache;
use xds_core::{NodeHash, ResourceRegistry, TypeUrl};

use crate::delta::{delta_response_to_proto, ClientResourceState, DeltaHandler};
use crate::sotw::SotwHandler;
use crate::stream::StreamContext;

// Re-export the data-plane-api types
pub use xds_types::envoy::service::discovery::v3::{DiscoveryRequest, DiscoveryResponse};
pub use xds_types::envoy::service::secret::v3::secret_discovery_service_server::{
    SecretDiscoveryService, SecretDiscoveryServiceServer,
};
pub use xds_types::envoy::service::discovery::v3::{
    DeltaDiscoveryRequest, DeltaDiscoveryResponse,
};

/// Secret Discovery Service.
#[derive(Debug, Clone)]
pub struct SdsService {
    /// Shared cache.
    cache: Arc<ShardedCache>,
    /// Resource registry.
    registry: Arc<ResourceRegistry>,
    /// SotW handler.
    sotw_handler: Arc<SotwHandler>,
    /// Delta handler.
    delta_handler: Arc<DeltaHandler>,
}

impl SdsService {
    /// Create a new SDS service.
    pub fn new(cache: Arc<ShardedCache>, registry: Arc<ResourceRegistry>) -> Self {
        let sotw_handler = Arc::new(SotwHandler::new(Arc::clone(&cache), Arc::clone(&registry)));
        let delta_handler = Arc::new(DeltaHandler::new(Arc::clone(&cache), Arc::clone(&registry)));
        Self {
            cache,
            registry,
            sotw_handler,
            delta_handler,
        }
    }

    /// Create a new SDS service with shared protocol handlers.
    ///
    /// This allows sharing handlers across multiple services for better
    /// resource efficiency.
    pub fn new_with_handlers(
        cache: Arc<ShardedCache>,
        registry: Arc<ResourceRegistry>,
        sotw_handler: Arc<SotwHandler>,
        delta_handler: Arc<DeltaHandler>,
    ) -> Self {
        Self {
            cache,
            registry,
            sotw_handler,
            delta_handler,
        }
    }

    /// Get the type URL for secrets.
    #[inline]
    pub fn type_url() -> &'static str {
        TypeUrl::SECRET
    }

    /// Get a reference to the cache.
    #[allow(dead_code)]
    pub fn cache(&self) -> &ShardedCache {
        &self.cache
    }

    /// Get a reference to the registry.
    #[allow(dead_code)]
    pub fn registry(&self) -> &ResourceRegistry {
        &self.registry
    }

    /// Convert this service into a tonic service for use with Server::add_service.
    pub fn into_service(self) -> SecretDiscoveryServiceServer<Self> {
        SecretDiscoveryServiceServer::new(self)
    }

    /// Convert a SotW response to a proto DiscoveryResponse.
    ///
    /// Returns an error if any resource fails to encode, rather than
    /// silently dropping it.
    fn convert_response(
        &self,
        response: crate::sotw::SotwResponse,
    ) -> Result<DiscoveryResponse, Status> {
        use xds_types::google::protobuf::Any;

        let resources: Vec<Any> = response
            .resources
            .iter()
            .map(|r| {
                r.encode().map(|encoded| Any {
                    type_url: encoded.type_url.clone(),
                    value: encoded.value.clone(),
                })
            })
            .collect::<Result<Vec<_>, _>>()
            .map_err(|e| Status::internal(format!("failed to encode resource: {}", e)))?;

        Ok(DiscoveryResponse {
            version_info: response.version_info,
            resources,
            type_url: TypeUrl::SECRET.to_string(),
            nonce: response.nonce,
            canary: false,
            control_plane: None,
            resource_errors: vec![],
        })
    }
}

/// Response stream type for SDS.
pub type SdsResponseStream = ReceiverStream<Result<DiscoveryResponse, Status>>;

/// Delta response stream type for SDS.
pub type SdsDeltaResponseStream = ReceiverStream<Result<DeltaDiscoveryResponse, Status>>;

#[async_trait]
impl SecretDiscoveryService for SdsService {
    type StreamSecretsStream = SdsResponseStream;

    #[instrument(skip(self, request), name = "sds_stream")]
    async fn stream_secrets(
        &self,
        request: Request<Streaming<DiscoveryRequest>>,
    ) -> Result<Response<Self::StreamSecretsStream>, Status> {
        let mut stream = request.into_inner();
        let (tx, rx) = mpsc::channel(16);

        let service = self.clone();
        let mut ctx = StreamContext::new();

        info!(stream = %ctx.id(), "SDS stream started");

        tokio::spawn(async move {
            let mut node_hash: Option<NodeHash> = None;

            while let Some(result) = tokio_stream::StreamExt::next(&mut stream).await {
                match result {
                    Ok(request) => {
                        // Validate type URL
                        if !request.type_url.is_empty() && request.type_url != TypeUrl::SECRET {
                            error!(
                                stream = %ctx.id(),
                                expected = TypeUrl::SECRET,
                                got = %request.type_url,
                                "invalid type URL for SDS"
                            );
                            continue;
                        }

                        // Extract node info
                        if node_hash.is_none() {
                            if let Some(ref node) = request.node {
                                let hash = NodeHash::from_id(&node.id);
                                ctx.set_node(node.id.clone(), hash);
                                node_hash = Some(hash);
                            }
                        }

                        let hash = match node_hash {
                            Some(h) => h,
                            None => {
                                // First request must include node information
                                error!(
                                    stream = %ctx.id(),
                                    "first request missing required node information"
                                );
                                let _ = tx.send(Err(Status::invalid_argument(
                                    "first request must include node information"
                                ))).await;
                                break;
                            }
                        };

                        // Process request
                        match service.sotw_handler.process_request(
                            &ctx,
                            TypeUrl::SECRET.into(),
                            &request.version_info,
                            &request.resource_names,
                            hash,
                        ) {
                            Ok(Some(response)) => {
                                match service.convert_response(response) {
                                    Ok(discovery_response) => {
                                        if tx.send(Ok(discovery_response)).await.is_err() {
                                            break;
                                        }
                                    }
                                    Err(e) => {
                                        error!(stream = %ctx.id(), error = %e, "failed to convert response");
                                        let _ = tx.send(Err(e)).await;
                                        break;
                                    }
                                }
                            }
                            Ok(None) => {}
                            Err(e) => {
                                error!(stream = %ctx.id(), error = %e, "SDS request failed");
                                break;
                            }
                        }
                    }
                    Err(e) => {
                        error!(stream = %ctx.id(), error = %e, "stream error");
                        break;
                    }
                }
            }

            info!(stream = %ctx.id(), "SDS stream ended");
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    type DeltaSecretsStream = SdsDeltaResponseStream;

    #[instrument(skip(self, request), name = "sds_delta_stream")]
    async fn delta_secrets(
        &self,
        request: Request<Streaming<DeltaDiscoveryRequest>>,
    ) -> Result<Response<Self::DeltaSecretsStream>, Status> {
        let mut stream = request.into_inner();
        let (tx, rx) = mpsc::channel(16);

        let service = self.clone();
        let mut ctx = StreamContext::new();
        info!(stream = %ctx.id(), "Delta SDS stream started");

        tokio::spawn(async move {
            let mut node_hash: Option<NodeHash> = None;
            let mut client_state = ClientResourceState::new();

            while let Some(result) = tokio_stream::StreamExt::next(&mut stream).await {
                match result {
                    Ok(request) => {
                        if !request.type_url.is_empty() && request.type_url != TypeUrl::SECRET {
                            error!(
                                stream = %ctx.id(),
                                expected = TypeUrl::SECRET,
                                got = %request.type_url,
                                "invalid type URL for delta SDS"
                            );
                            continue;
                        }

                        if node_hash.is_none() {
                            if let Some(ref node) = request.node {
                                let hash = NodeHash::from_id(&node.id);
                                ctx.set_node(node.id.clone(), hash);
                                node_hash = Some(hash);
                            }
                        }

                        let hash = match node_hash {
                            Some(h) => h,
                            None => {
                                error!(
                                    stream = %ctx.id(),
                                    "first delta request missing required node information"
                                );
                                let _ = tx
                                    .send(Err(Status::invalid_argument(
                                        "first request must include node information",
                                    )))
                                    .await;
                                break;
                            }
                        };

                        if !request.response_nonce.is_empty() {
                            if let Some(ref err) = request.error_detail {
                                service.delta_handler.handle_nack(
                                    &ctx,
                                    TypeUrl::SECRET.into(),
                                    &request.response_nonce,
                                    &err.message,
                                );
                            } else {
                                service.delta_handler.handle_ack(
                                    &ctx,
                                    TypeUrl::SECRET.into(),
                                    &request.response_nonce,
                                );
                            }
                        }

                        match service.delta_handler.process_request(
                            &ctx,
                            TypeUrl::SECRET.into(),
                            &mut client_state,
                            request.resource_names_subscribe,
                            request.resource_names_unsubscribe,
                            hash,
                        ) {
                            Ok(Some(response)) => match delta_response_to_proto(response) {
                                Ok(proto_response) => {
                                    if tx.send(Ok(proto_response)).await.is_err() {
                                        break;
                                    }
                                }
                                Err(e) => {
                                    error!(stream = %ctx.id(), error = %e, "failed to encode delta response");
                                    let _ = tx.send(Err(e)).await;
                                    break;
                                }
                            },
                            Ok(None) => {}
                            Err(e) => {
                                error!(stream = %ctx.id(), error = %e, "delta SDS request failed");
                                break;
                            }
                        }
                    }
                    Err(e) => {
                        error!(stream = %ctx.id(), error = %e, "delta stream error");
                        break;
                    }
                }
            }

            info!(stream = %ctx.id(), "Delta SDS stream ended");
            drop(tx);
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    #[instrument(skip(self, request), name = "sds_fetch")]
    async fn fetch_secrets(
        &self,
        request: Request<DiscoveryRequest>,
    ) -> Result<Response<DiscoveryResponse>, Status> {
        let request = request.into_inner();
        let ctx = StreamContext::new();

        let node_hash = request
            .node
            .as_ref()
            .map(|n| NodeHash::from_id(&n.id))
            .ok_or_else(|| Status::invalid_argument("node is required"))?;

        debug!(node = ?node_hash, "SDS fetch request");

        let response = self
            .sotw_handler
            .process_request(
                &ctx,
                TypeUrl::SECRET.into(),
                &request.version_info,
                &request.resource_names,
                node_hash,
            )
            .map_err(|e| Status::internal(e.to_string()))?
            .ok_or_else(|| Status::not_found("no secrets available"))?;

        Ok(Response::new(self.convert_response(response)?))
    }
}

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

    #[test]
    fn sds_service_creation() {
        let cache = Arc::new(ShardedCache::new());
        let registry = Arc::new(ResourceRegistry::new());
        let _service = SdsService::new(cache, registry);
    }

    #[test]
    fn sds_type_url() {
        assert_eq!(SdsService::type_url(), TypeUrl::SECRET);
    }
}