snap-control 0.5.2

// Copyright 2025 Anapaya Systems
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//! SNAP control plane API server.

use std::{net::SocketAddr, sync::Arc, time::Duration};

use axum::{BoxError, Router, error_handling::HandleErrorLayer};
use endhost_api::routes::nest_endhost_api;
use endhost_api_models::{
    SegmentsDiscovery,
    underlays::{ScionRouter, Underlays},
};
use http::StatusCode;
use scion_sdk_observability::info_trace_layer;
use sciparse::identifier::isd_asn::IsdAsn;
use tokio::net::TcpListener;
use tokio_util::sync::CancellationToken;
use tower::{ServiceBuilder, timeout::TimeoutLayer};
use url::Url;

use crate::{
    crpc_api::api_service::{
        model::{SnapDataPlaneResolver, SnapTunIdentityRegistry},
        nest_snap_control_api,
    },
    model::UnderlayDiscovery,
    server::{
        auth::AuthMiddlewareLayer,
        metrics::{Metrics, PrometheusMiddlewareLayer},
    },
};

pub mod auth;
pub mod identity_registry;
pub mod jwks_key_store;
pub mod metrics;
pub mod mock_segment_lister;
pub mod state;
pub mod token_verifier;

pub use token_verifier::SnapTokenVerifier;

const CONTROL_PLANE_API_TIMEOUT: Duration = Duration::from_secs(30);

/// Start the SNAP control plane API server.
pub async fn start<UD, SL, SR, IR>(
    cancellation_token: CancellationToken,
    listener: TcpListener,
    underlay_discovery: UD,
    segment_lister: SL,
    snap_resolver: SR,
    identity_registry: Arc<IR>,
    token_verifier: SnapTokenVerifier,
    metrics: Metrics,
) -> std::io::Result<()>
where
    UD: UnderlayDiscovery + 'static + Send + Sync,
    SL: SegmentsDiscovery + 'static + Send + Sync,
    SR: SnapDataPlaneResolver + 'static + Send + Sync,
    IR: SnapTunIdentityRegistry + 'static + Send + Sync,
{
    let router = Router::new();

    let dp_discovery = Arc::new(underlay_discovery);
    let segment_lister = Arc::new(segment_lister);
    let snap_resolver = Arc::new(snap_resolver);

    let snap_cp_addr = listener
        .local_addr()
        .map_err(|e| std::io::Error::other(format!("Failed to get own local address: {e}")))?;

    let snap_cp_api = match snap_cp_addr {
        SocketAddr::V4(addr) => {
            Url::parse(&format!("http://{addr}"))
                .expect("It is safe to format a SocketAddr as a URL")
        }
        SocketAddr::V6(addr) => {
            Url::parse(&format!("http://[{}]:{}", addr.ip(), addr.port()))
                .expect("It is safe to format a SocketAddr as a URL")
        }
    };

    let router = nest_endhost_api(
        router,
        Arc::new(UnderlayDiscoveryAdapter::new(
            dp_discovery.clone(),
            snap_cp_api,
        )),
        segment_lister.clone(),
    );

    let router = nest_snap_control_api(router, snap_resolver, identity_registry);

    let router = router.layer(
        ServiceBuilder::new()
            .layer(HandleErrorLayer::new(|err: BoxError| {
                async move {
                    tracing::error!(error=%err, "Control plane API error");

                    (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        format!("Unhandled error: {err}"),
                    )
                }
            }))
            .layer(info_trace_layer())
            .layer(TimeoutLayer::new(CONTROL_PLANE_API_TIMEOUT))
            .layer(PrometheusMiddlewareLayer::new(metrics))
            .layer(AuthMiddlewareLayer::new(token_verifier)),
    );

    tracing::info!(addr=%snap_cp_addr, "Starting control plane API");

    if let Err(e) = axum::serve(
        listener,
        router.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .with_graceful_shutdown(cancellation_token.cancelled_owned())
    .await
    {
        tracing::error!(error=%e, "Control plane API server unexpectedly stopped");
    }

    tracing::info!("Shutting down control plane API server");

    Ok(())
}

/// Adapter implementing UnderlayDiscovery for any DataPlaneDiscovery.
struct UnderlayDiscoveryAdapter<T: UnderlayDiscovery> {
    underlay_discovery: Arc<T>,
    snap_cp_api: Url,
}

impl<T: UnderlayDiscovery> UnderlayDiscoveryAdapter<T> {
    fn new(underlay_discovery: Arc<T>, snap_cp_api: Url) -> Self {
        Self {
            underlay_discovery,
            snap_cp_api,
        }
    }
}

impl<T: UnderlayDiscovery> endhost_api_models::UnderlayDiscovery for UnderlayDiscoveryAdapter<T> {
    fn list_underlays(&self, isd_as: IsdAsn) -> Underlays {
        let dps = self.underlay_discovery.list_udp_underlays();
        let mut udp_underlay = Vec::new();
        for dp in dps {
            for router_as in dp.isd_ases {
                if isd_as != IsdAsn::WILDCARD && router_as.isd_as != isd_as {
                    continue;
                };

                udp_underlay.push(ScionRouter {
                    isd_as: router_as.isd_as,
                    internal_interface: dp.endpoint,
                    interfaces: router_as.interfaces.clone(),
                });
            }
        }

        let sus = self.underlay_discovery.list_snap_underlays();
        if sus.is_empty() {
            return Underlays {
                udp_underlay,
                snap_underlay: Vec::new(),
            };
        }

        let mut snap_underlay = Vec::new();
        let all_ases: Vec<IsdAsn> = sus.iter().flat_map(|su| su.isd_ases.clone()).collect();
        if isd_as == IsdAsn::WILDCARD || all_ases.contains(&isd_as) {
            snap_underlay.push(endhost_api_models::underlays::Snap {
                address: self.snap_cp_api.clone(),
                isd_ases: all_ases,
            });
        }

        Underlays {
            udp_underlay,
            snap_underlay,
        }
    }
}