trz-gateway-client 0.2.23

//! The Terrazzo Gateway [Client].

use std::ops::ControlFlow;
use std::sync::Arc;
use std::sync::Mutex;

use autoclone::autoclone;
use connect::ConnectError;
use futures::FutureExt as _;
use futures::future::BoxFuture;
use futures::future::Either;
use nameth::NamedEnumValues as _;
use nameth::nameth;
use tokio::sync::oneshot;
use tracing::Instrument;
use tracing::info;
use tracing::info_span;
use trz_gateway_common::declare_identifier;
use trz_gateway_common::handle::ServerHandle;
use trz_gateway_common::id::ClientId;
use trz_gateway_common::id::ClientName;
use trz_gateway_common::retry_strategy::RetryStrategy;
use trz_gateway_common::security_configuration::certificate::CertificateConfig;
use trz_gateway_common::security_configuration::certificate::tls_server::ToTlsServer as _;
use trz_gateway_common::security_configuration::certificate::tls_server::ToTlsServerError;
use trz_gateway_common::security_configuration::custom_server_certificate_verifier::ChainOnlyServerCertificateVerifier;
use trz_gateway_common::security_configuration::trusted_store::TrustedStoreConfig;
use trz_gateway_common::security_configuration::trusted_store::tls_client::ToTlsClient as _;
use trz_gateway_common::security_configuration::trusted_store::tls_client::ToTlsClientError;
use uuid::Uuid;

use self::config::ClientTransport;
use self::config::SniOverrideError;
use self::service::ClientService;
use crate::tunnel_config::TunnelConfig;

pub mod certificate;
pub mod config;
pub mod connect;
mod connection;
mod health;
pub mod service;
mod transport_stream;

/// The [Client].
///
/// It creates a WebSocket tunnel with the Terrazzo Gateway, and then runs a
/// gRPC server that listens to requests sent or forwarded by the Terrazzo
/// Gateway over the WebSocket tunnel.
pub struct Client {
    /// The client name for troubleshooting purposes
    pub client_name: ClientName,

    gateway_client: GatewayClient,
    client_api_server: ClientApiServer,
}

/// Configuration for how to create tunnels to the Terrazzo Gateway
struct GatewayClient {
    /// The URL used to open the TCP connection.
    gateway_uri: String,

    /// The TLS server name to validate when it differs from [Self::uri].
    gateway_sni_override: Option<String>,

    /// How the Gateway is reached before applying its TLS protocol.
    transport: ClientTransport,

    /// The TLS client is used to create the secure WebSocket tunnel.
    ///
    /// (without client certificate auth)
    gateway_tls_connector: tokio_tungstenite::Connector,
}

/// Configuration for the server running through tunnels to the Terrazzo Gateway.
struct ClientApiServer {
    /// The TLS server is used to accept gRPC connections through the tunnel.
    ///
    /// The client uses its certiticate to authenticate the server side of the connection.
    client_api_acceptor: tokio_rustls::TlsAcceptor,

    /// A callback to configure the [tonic gRPC server](tonic::transport::Server).
    client_service: Arc<dyn ClientService>,

    /// The strategy to retry creating the WebSocket tunnel when it fails
    retry_strategy: RetryStrategy,

    /// A global mutable variable that holds the [AuthCode].
    ///
    /// This is used to periodically renew the certificate.
    current_auth_code: Arc<Mutex<AuthCode>>,
}

declare_identifier!(AuthCode);

impl Client {
    /// Creates a new [Client].
    pub fn new<C: TunnelConfig>(config: C) -> Result<Arc<Self>, NewClientError<C>> {
        let tls_client = config
            .gateway_pki()
            .to_tls_client(ChainOnlyServerCertificateVerifier)?;
        let tls_server = config.client_certificate().to_tls_server()?;
        let client_name = config.client_name();
        let tunnel_path = format!("/remote/tunnel/{client_name}");
        Ok(Arc::new(Client {
            client_name,
            gateway_client: GatewayClient {
                gateway_uri: config.url(&tunnel_path)?.to_string(),
                gateway_sni_override: config.gateway_sni_override().map(ToOwned::to_owned),
                transport: config.transport(),
                gateway_tls_connector: tokio_tungstenite::Connector::Rustls(tls_client.into()),
            },
            client_api_server: ClientApiServer {
                client_api_acceptor: tokio_rustls::TlsAcceptor::from(tls_server),
                client_service: Arc::new(config.client_service()),
                retry_strategy: config.retry_strategy(),
                current_auth_code: config.current_auth_code(),
            },
        }))
    }

    /// Runs the client and returns a handle to stop the client.
    pub async fn run(self: &Arc<Self>) -> Result<ServerHandle<()>, ConnectError> {
        let this = self.clone();
        let client_name = &this.client_name;
        let span = info_span!("Run", %client_name);
        async move {
            let client_id = ClientId::from(Uuid::new_v4().to_string());
            info!(%client_id, "Allocated new client id");
            let (shutdown_rx, terminated_tx, handle) = ServerHandle::new("Client");
            let (serving_tx, serving_rx) = oneshot::channel();
            let task = run_impl(this, client_id, serving_tx, shutdown_rx, terminated_tx);
            tokio::spawn(task.in_current_span());
            let _ = serving_rx.await;
            Ok(handle)
        }
        .instrument(span)
        .await
    }
}

#[autoclone]
async fn run_impl(
    this: Arc<Client>,
    client_id: ClientId,

    // Set when the client is serving connections
    serving_tx: oneshot::Sender<()>,

    // Set when the client should start shutting down
    shutdown_rx: impl Future<Output = ()> + Send + 'static,

    // Set when the client has shut down
    terminated_tx: oneshot::Sender<()>,
) {
    scopeguard::defer! { let _ = terminated_tx.send(()); };
    let retry_strategy = this.client_api_server.retry_strategy.clone();
    let shutdown_rx: BoxFuture<()> = Box::pin(shutdown_rx);
    let shutdown_rx = shutdown_rx.shared();

    let serving_tx = Arc::new(tokio::sync::Mutex::new(Some(serving_tx)));
    let connect_retry_delay = retry_strategy.peek() / 2;
    let mut attempt = 0usize;
    let connect_task = retry_strategy.process(|| {
        let current_attempt = attempt;
        attempt += 1;
        async move {
            autoclone!(this, client_id, shutdown_rx, serving_tx);
            let mut serving_tx = serving_tx.lock().await;
            let result = this
                .connect(client_id, shutdown_rx, connect_retry_delay, &mut serving_tx)
                .instrument(info_span!("Connect", current_attempt))
                .await;
            ControlFlow::<(), _>::Continue(match result {
                Ok(()) => ConnectRetryError::Closed,
                Err(error) => ConnectRetryError::Failed(error),
            })
        }
    });
    let connect_task = std::pin::pin!(connect_task);
    match futures::future::select(connect_task, shutdown_rx.clone()).await {
        Either::Left(((), _shutdown_rx)) => {}
        Either::Right(((), _connect_task)) => {}
    }
}

#[derive(Debug, thiserror::Error)]
enum ConnectRetryError {
    #[error("Connection closed")]
    Closed,
    #[error("Connection failed: {0}")]
    Failed(ConnectError),
}

#[nameth]
#[derive(thiserror::Error, Debug)]
pub enum NewClientError<C: TunnelConfig> {
    #[error("[{n}] {0}", n = self.name())]
    SniOverride(#[from] SniOverrideError),

    #[error("[{n}] {0}", n = self.name())]
    ToTlsClient(#[from] ToTlsClientError<<C::GatewayPki as TrustedStoreConfig>::Error>),

    #[error("[{n}] {0}", n = self.name())]
    ToTlsServer(#[from] ToTlsServerError<<C::ClientCertificate as CertificateConfig>::Error>),
}