rpki/src/rtr/server/service.rs

use std::net::SocketAddr;
use std::path::Path;
use std::sync::{
    Arc,
    atomic::{AtomicU64, AtomicUsize, Ordering},
};

use tokio::sync::{Semaphore, broadcast, watch};
use tokio::task::JoinHandle;
use tracing::{error, warn};

use crate::rtr::cache::SharedRtrCache;
use crate::rtr::server::config::RtrServiceConfig;
use crate::rtr::server::connection::{
    RtrTransportConnectionCounters, RtrTransportConnectionCounts,
};
use crate::rtr::server::listener::RtrServer;
use crate::rtr::server::notifier::RtrNotifier;
use crate::rtr::server::ssh::{SshAuthMode, load_rtr_ssh_runtime_config};

pub struct RtrService {
    cache: SharedRtrCache,
    notify_tx: broadcast::Sender<()>,
    shutdown_tx: watch::Sender<bool>,
    connection_limiter: Arc<Semaphore>,
    handshake_limiter: Arc<Semaphore>,
    active_connections: Arc<AtomicUsize>,
    transport_connections: Arc<RtrTransportConnectionCounters>,
    connection_change_tx: watch::Sender<u64>,
    connection_change_seq: Arc<AtomicU64>,
    config: RtrServiceConfig,
}

#[derive(Clone)]
pub struct RtrServiceStats {
    active_connections: Arc<AtomicUsize>,
    transport_connections: Arc<RtrTransportConnectionCounters>,
    connection_change_tx: watch::Sender<u64>,
    max_connections: usize,
}

impl RtrServiceStats {
    pub fn active_connections(&self) -> usize {
        self.active_connections.load(Ordering::Relaxed)
    }

    pub fn transport_connections(&self) -> RtrTransportConnectionCounts {
        self.transport_connections.snapshot()
    }

    pub fn max_connections(&self) -> usize {
        self.max_connections
    }

    pub fn subscribe_connection_changes(&self) -> watch::Receiver<u64> {
        self.connection_change_tx.subscribe()
    }
}

impl RtrService {
    pub fn new(cache: SharedRtrCache) -> Self {
        Self::with_config(cache, RtrServiceConfig::default())
    }

    pub fn with_config(cache: SharedRtrCache, config: RtrServiceConfig) -> Self {
        let (notify_tx, _) = broadcast::channel(config.notify_queue_size);
        let (shutdown_tx, _) = watch::channel(false);
        let (connection_change_tx, _) = watch::channel(0);

        Self {
            cache,
            notify_tx,
            shutdown_tx,
            connection_limiter: Arc::new(Semaphore::new(config.max_connections)),
            handshake_limiter: Arc::new(Semaphore::new(config.max_concurrent_handshakes)),
            active_connections: Arc::new(AtomicUsize::new(0)),
            transport_connections: Arc::new(RtrTransportConnectionCounters::default()),
            connection_change_tx,
            connection_change_seq: Arc::new(AtomicU64::new(0)),
            config,
        }
    }

    pub fn cache(&self) -> SharedRtrCache {
        self.cache.clone()
    }

    pub fn notifier(&self) -> RtrNotifier {
        RtrNotifier::new(self.notify_tx.clone())
    }

    pub fn notify_cache_updated(&self) {
        let _ = self.notify_tx.send(());
    }

    pub fn active_connections(&self) -> usize {
        self.active_connections.load(Ordering::Relaxed)
    }

    pub fn transport_connections(&self) -> RtrTransportConnectionCounts {
        self.transport_connections.snapshot()
    }

    pub fn max_connections(&self) -> usize {
        self.config.max_connections
    }

    pub fn stats(&self) -> RtrServiceStats {
        RtrServiceStats {
            active_connections: self.active_connections.clone(),
            transport_connections: self.transport_connections.clone(),
            connection_change_tx: self.connection_change_tx.clone(),
            max_connections: self.config.max_connections,
        }
    }

    pub fn tcp_server(&self, bind_addr: SocketAddr) -> RtrServer {
        RtrServer::new(
            bind_addr,
            self.cache.clone(),
            self.notify_tx.clone(),
            self.shutdown_tx.clone(),
            self.connection_limiter.clone(),
            self.handshake_limiter.clone(),
            self.active_connections.clone(),
            self.transport_connections.clone(),
            self.connection_change_tx.clone(),
            self.connection_change_seq.clone(),
            self.config.clone(),
        )
    }

    pub fn tls_server(&self, bind_addr: SocketAddr) -> RtrServer {
        RtrServer::new(
            bind_addr,
            self.cache.clone(),
            self.notify_tx.clone(),
            self.shutdown_tx.clone(),
            self.connection_limiter.clone(),
            self.handshake_limiter.clone(),
            self.active_connections.clone(),
            self.transport_connections.clone(),
            self.connection_change_tx.clone(),
            self.connection_change_seq.clone(),
            self.config.clone(),
        )
    }

    pub fn ssh_server(&self, bind_addr: SocketAddr) -> RtrServer {
        RtrServer::new(
            bind_addr,
            self.cache.clone(),
            self.notify_tx.clone(),
            self.shutdown_tx.clone(),
            self.connection_limiter.clone(),
            self.handshake_limiter.clone(),
            self.active_connections.clone(),
            self.transport_connections.clone(),
            self.connection_change_tx.clone(),
            self.connection_change_seq.clone(),
            self.config.clone(),
        )
    }

    pub fn spawn_tcp(&self, bind_addr: SocketAddr) -> JoinHandle<()> {
        if self.config.warn_insecure_tcp {
            warn!(
                "starting plain TCP RTR service on {}. Per draft-ietf-sidrops-8210bis-25 Section 9, unsecured TCP must only be used on a trusted and controlled network",
                bind_addr
            );
        }
        let server = self.tcp_server(bind_addr);
        tokio::spawn(async move {
            if let Err(err) = server.run_tcp().await {
                error!("RTR TCP server {} exited with error: {:?}", bind_addr, err);
            }
        })
    }

    pub fn spawn_tls_from_pem(
        &self,
        bind_addr: SocketAddr,
        cert_path: impl AsRef<Path>,
        key_path: impl AsRef<Path>,
        client_ca_path: impl AsRef<Path>,
    ) -> JoinHandle<()> {
        let cert_path = cert_path.as_ref().to_path_buf();
        let key_path = key_path.as_ref().to_path_buf();
        let client_ca_path = client_ca_path.as_ref().to_path_buf();
        let server = self.tls_server(bind_addr);

        tokio::spawn(async move {
            if let Err(err) = server
                .run_tls_from_pem(cert_path, key_path, client_ca_path)
                .await
            {
                error!("RTR TLS server {} exited with error: {:?}", bind_addr, err);
            }
        })
    }

    pub fn spawn_tcp_and_tls_from_pem(
        &self,
        tcp_bind_addr: SocketAddr,
        tls_bind_addr: SocketAddr,
        cert_path: impl AsRef<Path>,
        key_path: impl AsRef<Path>,
        client_ca_path: impl AsRef<Path>,
    ) -> RunningRtrService {
        let tcp_handle = self.spawn_tcp(tcp_bind_addr);
        let tls_handle =
            self.spawn_tls_from_pem(tls_bind_addr, cert_path, key_path, client_ca_path);

        RunningRtrService {
            shutdown_tx: self.shutdown_tx.clone(),
            handles: vec![tcp_handle, tls_handle],
        }
    }

    #[allow(clippy::too_many_arguments)]
    pub fn spawn_ssh_from_openssh(
        &self,
        bind_addr: SocketAddr,
        host_key_path: impl AsRef<Path>,
        authorized_keys_path: impl AsRef<Path>,
        username: &str,
        subsystem_name: &str,
        auth_mode: SshAuthMode,
        password: Option<&str>,
    ) -> JoinHandle<()> {
        let host_key_path = host_key_path.as_ref().to_path_buf();
        let authorized_keys_path = authorized_keys_path.as_ref().to_path_buf();
        let username = username.to_string();
        let subsystem_name = subsystem_name.to_string();
        let password = password.map(ToString::to_string);
        let inactivity_timeout = Some(std::time::Duration::from_secs(3600));
        let keepalive_interval = self.config.tcp_keepalive;
        let server = self.ssh_server(bind_addr);

        tokio::spawn(async move {
            let runtime_config = match load_rtr_ssh_runtime_config(
                &host_key_path,
                &authorized_keys_path,
                &username,
                &subsystem_name,
                auth_mode,
                password.as_deref(),
                inactivity_timeout,
                keepalive_interval,
            ) {
                Ok(cfg) => Arc::new(cfg),
                Err(err) => {
                    error!(
                        "RTR SSH server {} failed to load configuration: {:?}",
                        bind_addr, err
                    );
                    return;
                }
            };

            if let Err(err) = server.run_ssh(runtime_config).await {
                error!("RTR SSH server {} exited with error: {:?}", bind_addr, err);
            }
        })
    }

    #[allow(clippy::too_many_arguments)]
    pub fn spawn_tcp_and_ssh_from_openssh(
        &self,
        tcp_bind_addr: SocketAddr,
        ssh_bind_addr: SocketAddr,
        host_key_path: impl AsRef<Path>,
        authorized_keys_path: impl AsRef<Path>,
        username: &str,
        subsystem_name: &str,
        auth_mode: SshAuthMode,
        password: Option<&str>,
    ) -> RunningRtrService {
        let tcp_handle = self.spawn_tcp(tcp_bind_addr);
        let ssh_handle = self.spawn_ssh_from_openssh(
            ssh_bind_addr,
            host_key_path,
            authorized_keys_path,
            username,
            subsystem_name,
            auth_mode,
            password,
        );

        RunningRtrService {
            shutdown_tx: self.shutdown_tx.clone(),
            handles: vec![tcp_handle, ssh_handle],
        }
    }

    #[allow(clippy::too_many_arguments)]
    pub fn spawn_tcp_tls_and_ssh_from_pem_and_openssh(
        &self,
        tcp_bind_addr: SocketAddr,
        tls_bind_addr: SocketAddr,
        ssh_bind_addr: SocketAddr,
        cert_path: impl AsRef<Path>,
        key_path: impl AsRef<Path>,
        client_ca_path: impl AsRef<Path>,
        host_key_path: impl AsRef<Path>,
        authorized_keys_path: impl AsRef<Path>,
        username: &str,
        subsystem_name: &str,
        auth_mode: SshAuthMode,
        password: Option<&str>,
    ) -> RunningRtrService {
        let tcp_handle = self.spawn_tcp(tcp_bind_addr);
        let tls_handle =
            self.spawn_tls_from_pem(tls_bind_addr, cert_path, key_path, client_ca_path);
        let ssh_handle = self.spawn_ssh_from_openssh(
            ssh_bind_addr,
            host_key_path,
            authorized_keys_path,
            username,
            subsystem_name,
            auth_mode,
            password,
        );

        RunningRtrService {
            shutdown_tx: self.shutdown_tx.clone(),
            handles: vec![tcp_handle, tls_handle, ssh_handle],
        }
    }

    pub fn spawn_tcp_only(&self, tcp_bind_addr: SocketAddr) -> RunningRtrService {
        let tcp_handle = self.spawn_tcp(tcp_bind_addr);

        RunningRtrService {
            shutdown_tx: self.shutdown_tx.clone(),
            handles: vec![tcp_handle],
        }
    }
}

pub struct RunningRtrService {
    shutdown_tx: watch::Sender<bool>,
    handles: Vec<JoinHandle<()>>,
}

impl RunningRtrService {
    pub fn shutdown(&self) {
        let _ = self.shutdown_tx.send(true);
    }

    pub async fn wait(self) {
        for handle in self.handles {
            let _ = handle.await;
        }
    }
}