blueprint-manager-bridge 0.2.0-alpha.12

use crate::VSOCK_PORT;
use crate::blueprint_manager_bridge_client::BlueprintManagerBridgeClient;
use crate::{
    AddOwnerToServiceRequest, Error, PortRequest, RegisterBlueprintServiceProxyRequest,
    RemoveOwnerFromServiceRequest, ServiceOwner, TlsProfileConfig,
    UnregisterBlueprintServiceProxyRequest, UpdateBlueprintServiceTlsProfileRequest,
};
use blueprint_auth::models::{ServiceOwnerModel, TlsProfile};
use blueprint_core::debug;
use hyper_util::rt::TokioIo;
use std::path::Path;
use tokio::net::UnixStream;
use tokio_vsock::{VsockAddr, VsockStream};
use tonic::transport::Channel;

#[derive(Debug)]
pub struct Bridge {
    client: BlueprintManagerBridgeClient<Channel>,
}

impl Bridge {
    /// Connect to the blueprint manager bridge.
    ///
    /// NOTE: This should not be used directly in blueprints, see [`BlueprintEnvironment::bridge()`].
    ///
    /// # Errors
    ///
    /// * Unable to connect to the bridge (the host never created a server?)
    ///
    /// [`BlueprintEnvironment::bridge()`]: https://docs.rs/blueprint-runner/latest/blueprint_runner/config/struct.BlueprintEnvironment.html#method.bridge
    pub async fn connect(socket_path: Option<&Path>) -> Result<Self, Error> {
        let channel = match socket_path {
            Some(path) => {
                debug!("Connecting to UDS bridge at {}", path.display());

                let path = path.to_path_buf();
                Channel::from_static("http://[::]:50051")
                    .connect_with_connector(tower::service_fn(move |_| {
                        let path = path.clone();
                        async move {
                            Ok::<_, std::io::Error>(TokioIo::new(UnixStream::connect(&path).await?))
                        }
                    }))
                    .await?
            }
            None => {
                debug!("Connecting to VSOCK bridge at port {}", VSOCK_PORT);
                Channel::from_static("http://[::]:50051")
                    .connect_with_connector(tower::service_fn(|_| async {
                        Ok::<_, std::io::Error>(TokioIo::new(
                            VsockStream::connect(VsockAddr::new(
                                tokio_vsock::VMADDR_CID_HOST,
                                VSOCK_PORT,
                            ))
                            .await?,
                        ))
                    }))
                    .await?
            }
        };

        Ok(Self {
            client: BlueprintManagerBridgeClient::new(channel),
        })
    }

    /// Sends a Ping request to the blueprint manager bridge.
    ///
    /// This method is used to check the connectivity and responsiveness of the bridge.
    ///
    /// # Errors
    /// - Returns an error if the ping operation fails, for example, due to network issues or if the bridge service is not responding.
    pub async fn ping(&self) -> Result<(), Error> {
        self.client.clone().ping(()).await?;
        Ok(())
    }

    /// Requests a port from the blueprint manager bridge.
    ///
    /// The bridge will attempt to reserve the preferred port if provided,
    /// otherwise it will assign an available port.
    ///
    /// # Arguments
    /// * `preferred`: An optional `u16` specifying the preferred port number.
    ///
    /// # Returns
    /// A `Result` containing the `u16` port number assigned by the bridge,
    /// or an `Error` if the request fails.
    ///
    /// # Errors
    /// - Returns an error if the port request to the bridge fails, e.g., if the bridge cannot allocate a port or encounters an internal error.
    #[allow(clippy::cast_possible_truncation)]
    pub async fn request_port(&self, preferred: Option<u16>) -> Result<u16, Error> {
        let reply = self
            .client
            .clone()
            .request_port(PortRequest {
                preferred_port: u32::from(preferred.unwrap_or(0)),
            })
            .await?
            .into_inner();

        Ok(reply.port as u16)
    }

    /// Registers a blueprint service proxy with the blueprint manager bridge.
    ///
    /// This method allows a blueprint to make a service accessible via the proxy.
    ///
    /// # Arguments
    /// * `service_id`: A `u64` unique identifier for the service.
    /// * `api_key_prefix`: An optional string slice representing the prefix of the API key used for service authentication.
    /// * `upstream_url`: A string slice representing the URL of the upstream service.
    /// * `owners`: A slice of `ServiceOwnerModel` defining the owners authorized to use this service.
    /// * `tls_profile`: An optional `TlsProfile` defining the TLS configuration for this service.
    ///
    /// # Errors
    /// - Returns an error if registering the service proxy fails, such as issues with the provided parameters or bridge internal errors.
    pub async fn register_blueprint_service_proxy(
        &self,
        service_id: u64,
        api_key_prefix: Option<&str>,
        upstream_url: &str,
        owners: &[ServiceOwnerModel],
        tls_profile: Option<TlsProfile>,
    ) -> Result<(), Error> {
        let owners = owners
            .iter()
            .map(|owner| ServiceOwner {
                key_type: owner.key_type,
                key_bytes: owner.key_bytes.clone(),
            })
            .collect();

        let tls_profile = tls_profile.map(TlsProfileConfig::from);

        let request = RegisterBlueprintServiceProxyRequest {
            service_id,
            api_key_prefix: api_key_prefix.unwrap_or_default().to_owned(),
            upstream_url: upstream_url.to_string(),
            owners,
            tls_profile,
        };

        self.client
            .clone()
            .register_blueprint_service_proxy(request)
            .await?;

        Ok(())
    }

    /// Unregisters a blueprint service proxy from the blueprint manager bridge.
    ///
    /// This method is called by a blueprint to remove a service from the proxy.
    ///
    /// # Arguments
    /// * `service_id`: A `u64` unique identifier for the service to be unregistered.
    ///
    /// # Errors
    /// - Returns an error if unregistering the service proxy fails, for instance, if the service ID is not found or the bridge encounters an issue.
    pub async fn unregister_blueprint_service_proxy(&self, service_id: u64) -> Result<(), Error> {
        let request = UnregisterBlueprintServiceProxyRequest { service_id };

        self.client
            .clone()
            .unregister_blueprint_service_proxy(request)
            .await?;

        Ok(())
    }

    /// Adds an owner to a registered blueprint service.
    ///
    /// This method allows a blueprint to grant an additional owner access to a service.
    ///
    /// # Arguments
    /// * `service_id`: A `u64` unique identifier for the service.
    /// * `owner_to_add`: A `ServiceOwnerModel` representing the owner to be added.
    ///
    /// # Errors
    /// - Returns an error if adding an owner to the service fails, e.g., if the service ID is invalid or the owner cannot be added.
    pub async fn add_owner_to_service(
        &self,
        service_id: u64,
        owner_to_add: ServiceOwnerModel,
    ) -> Result<(), Error> {
        let request = AddOwnerToServiceRequest {
            service_id,
            owner_to_add: Some(ServiceOwner {
                key_type: owner_to_add.key_type,
                key_bytes: owner_to_add.key_bytes,
            }),
        };

        self.client.clone().add_owner_to_service(request).await?;

        Ok(())
    }

    /// Removes an owner from a registered blueprint service.
    ///
    /// This method allows a blueprint to revoke an owner's access to a service.
    ///
    /// # Arguments
    /// * `service_id`: A `u64` unique identifier for the service.
    /// * `owner_to_remove`: A `ServiceOwnerModel` representing the owner to be removed.
    ///
    /// # Errors
    /// - Returns an error if removing an owner from the service fails, for example, if the service ID or owner is not found.
    pub async fn remove_owner_from_service(
        &self,
        service_id: u64,
        owner_to_remove: ServiceOwnerModel,
    ) -> Result<(), Error> {
        let request = RemoveOwnerFromServiceRequest {
            service_id,
            owner_to_remove: Some(ServiceOwner {
                key_type: owner_to_remove.key_type,
                key_bytes: owner_to_remove.key_bytes,
            }),
        };

        self.client
            .clone()
            .remove_owner_from_service(request)
            .await?;

        Ok(())
    }

    /// Updates the TLS profile for a registered blueprint service.
    ///
    /// This method allows a blueprint to update the TLS configuration of a service after registration.
    /// TLS assets should already be envelope-encrypted before calling this method.
    ///
    /// # Arguments
    /// * `service_id`: A `u64` unique identifier for the service.
    /// * `tls_profile`: An optional `TlsProfile` defining the new TLS configuration.
    ///   If `None`, TLS will be disabled for the service.
    ///
    /// # Errors
    /// - Returns an error if updating the TLS profile fails, such as invalid service ID or bridge internal errors.
    pub async fn update_blueprint_service_tls_profile(
        &self,
        service_id: u64,
        tls_profile: Option<TlsProfile>,
    ) -> Result<(), Error> {
        let tls_profile = tls_profile.map(TlsProfileConfig::from);

        let request = UpdateBlueprintServiceTlsProfileRequest {
            service_id,
            tls_profile,
        };

        self.client
            .clone()
            .update_blueprint_service_tls_profile(request)
            .await?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use blueprint_auth::models::TlsProfile;
    use zerocopy::IntoBytes;

    #[test]
    fn test_tls_profile_conversion() {
        let auth_profile = TlsProfile {
            tls_enabled: true,
            require_client_mtls: true,
            encrypted_server_cert: b"server_cert".to_vec(),
            encrypted_server_key: b"server_key".to_vec(),
            encrypted_client_ca_bundle: b"client_ca".to_vec(),
            encrypted_upstream_ca_bundle: b"upstream_ca".to_vec(),
            encrypted_upstream_client_cert: b"upstream_cert".to_vec(),
            encrypted_upstream_client_key: b"upstream_key".to_vec(),
            client_cert_ttl_hours: 24,
            sni: Some("example.com".to_string()),
            subject_alt_name_template: Some("service-{id}.example.com".to_string()),
            allowed_dns_names: vec!["example.com".to_string(), "*.example.com".to_string()],
        };

        // Convert auth to protobuf
        let proto_config: TlsProfileConfig = auth_profile.clone().into();

        assert!(proto_config.tls_enabled);
        assert!(proto_config.require_client_mtls);
        assert_eq!(proto_config.encrypted_server_cert, b"server_cert");
        assert_eq!(proto_config.encrypted_server_key, b"server_key");
        assert_eq!(proto_config.encrypted_client_ca_bundle, b"client_ca");
        assert_eq!(proto_config.sni, Some("example.com".to_string()));

        // Convert protobuf back to auth
        let converted_auth: TlsProfile = proto_config.into();

        assert_eq!(converted_auth.tls_enabled, auth_profile.tls_enabled);
        assert_eq!(
            converted_auth.require_client_mtls,
            auth_profile.require_client_mtls
        );
        assert_eq!(
            converted_auth.encrypted_server_cert,
            auth_profile.encrypted_server_cert
        );
        assert_eq!(
            converted_auth.encrypted_server_key,
            auth_profile.encrypted_server_key
        );
        assert_eq!(
            converted_auth.client_cert_ttl_hours,
            auth_profile.client_cert_ttl_hours
        );
        assert_eq!(converted_auth.sni, auth_profile.sni);
        assert_eq!(
            converted_auth.allowed_dns_names,
            auth_profile.allowed_dns_names
        );
    }

    #[test]
    fn test_tls_profile_none_conversion() {
        let auth_profile = TlsProfile {
            tls_enabled: false,
            require_client_mtls: false,
            encrypted_server_cert: Vec::new(),
            encrypted_server_key: Vec::new(),
            encrypted_client_ca_bundle: Vec::new(),
            encrypted_upstream_ca_bundle: Vec::new(),
            encrypted_upstream_client_cert: Vec::new(),
            encrypted_upstream_client_key: Vec::new(),
            client_cert_ttl_hours: 0,
            sni: None,
            subject_alt_name_template: None,
            allowed_dns_names: Vec::new(),
        };

        let proto_config: TlsProfileConfig = auth_profile.clone().into();
        let converted_auth: TlsProfile = proto_config.into();

        assert!(!converted_auth.tls_enabled);
        assert!(!converted_auth.require_client_mtls);
        assert!(converted_auth.encrypted_server_cert.is_empty());
        assert!(converted_auth.sni.is_none());
        assert!(converted_auth.allowed_dns_names.is_empty());
    }

    #[test]
    fn test_tls_profile_partial_fields() {
        let auth_profile = TlsProfile {
            tls_enabled: true,
            require_client_mtls: false,
            encrypted_server_cert: b"cert".to_vec(),
            encrypted_server_key: b"key".to_vec(),
            encrypted_client_ca_bundle: Vec::new(),
            encrypted_upstream_ca_bundle: Vec::new(),
            encrypted_upstream_client_cert: Vec::new(),
            encrypted_upstream_client_key: Vec::new(),
            client_cert_ttl_hours: 48,
            sni: None,
            subject_alt_name_template: Some("*.example.com".to_string()),
            allowed_dns_names: vec!["example.com".to_string()],
        };

        let proto_config: TlsProfileConfig = auth_profile.clone().into();
        let converted_auth: TlsProfile = proto_config.into();

        assert!(converted_auth.tls_enabled);
        assert!(!converted_auth.require_client_mtls);
        assert_eq!(converted_auth.encrypted_server_cert, b"cert");
        assert_eq!(converted_auth.client_cert_ttl_hours, 48);
        assert_eq!(
            converted_auth.subject_alt_name_template,
            Some("*.example.com".to_string())
        );
        assert_eq!(
            converted_auth.allowed_dns_names,
            vec!["example.com".to_string()]
        );
    }

    #[test]
    fn test_tls_profile_upstream_only() {
        let auth_profile = TlsProfile {
            tls_enabled: true,
            require_client_mtls: false,
            encrypted_server_cert: Vec::new(),
            encrypted_server_key: Vec::new(),
            encrypted_client_ca_bundle: Vec::new(),
            encrypted_upstream_ca_bundle: b"upstream_ca".to_vec(),
            encrypted_upstream_client_cert: b"upstream_cert".to_vec(),
            encrypted_upstream_client_key: b"upstream_key".to_vec(),
            client_cert_ttl_hours: 0,
            sni: None,
            subject_alt_name_template: None,
            allowed_dns_names: Vec::new(),
        };

        let proto_config: TlsProfileConfig = auth_profile.clone().into();
        let converted_auth: TlsProfile = proto_config.into();

        assert!(converted_auth.tls_enabled);
        assert!(!converted_auth.require_client_mtls);
        assert!(converted_auth.encrypted_server_cert.is_empty());
        assert!(converted_auth.encrypted_server_key.is_empty());
        assert_eq!(converted_auth.encrypted_upstream_ca_bundle, b"upstream_ca");
        assert_eq!(
            converted_auth.encrypted_upstream_client_cert,
            b"upstream_cert"
        );
        assert_eq!(
            converted_auth.encrypted_upstream_client_key,
            b"upstream_key"
        );
    }

    #[test]
    fn test_tls_profile_server_validation_scenario() {
        // Test a typical server TLS scenario without client mTLS
        let auth_profile = TlsProfile {
            tls_enabled: true,
            require_client_mtls: false,
            encrypted_server_cert:
                b"-----BEGIN CERTIFICATE-----\nserver cert\n-----END CERTIFICATE-----"
                    .as_bytes()
                    .to_vec(),
            encrypted_server_key:
                b"-----BEGIN PRIVATE KEY-----\nserver key\n-----END PRIVATE KEY-----"
                    .as_bytes()
                    .to_vec(),
            encrypted_client_ca_bundle: Vec::new(),
            encrypted_upstream_ca_bundle: Vec::new(),
            encrypted_upstream_client_cert: Vec::new(),
            encrypted_upstream_client_key: Vec::new(),
            client_cert_ttl_hours: 24,
            sni: Some("api.example.com".to_string()),
            subject_alt_name_template: None,
            allowed_dns_names: vec!["api.example.com".to_string()],
        };

        let proto_config: TlsProfileConfig = auth_profile.clone().into();
        let converted_auth: TlsProfile = proto_config.into();

        assert!(converted_auth.tls_enabled);
        assert!(!converted_auth.require_client_mtls);
        assert!(!converted_auth.encrypted_server_cert.is_empty());
        assert!(!converted_auth.encrypted_server_key.is_empty());
        assert!(converted_auth.encrypted_client_ca_bundle.is_empty());
        assert_eq!(converted_auth.sni, Some("api.example.com".to_string()));
    }
}