nordisk/gunnar

crates/gunnar-repair/src/proto.rs

//! **The verbs over gunnar's own generated clients** (gunnar-proto) and the one
//! credential signer (`gunnar_grpc::client_signer`): what gunnar-cli and
//! gunnar-server hand [`crate::pump::recharge_from_twin`] and [`crate::fence::fence_primary`].
//! Every call is signed with the caller's ssh key; this crate never stores one.

use gunnar_grpc::client_signer::{authed, ControlChannel};
use gunnar_proto::bootstrap_client::BootstrapClient;
use gunnar_proto::mirror_client::MirrorClient;

use crate::fence::{Fenceable, STANDBY};
use crate::grpc::{
    fault, restore_stream, TonicRestore, TonicStream, RESTORE_METHOD, RESTORE_PATH, RESTORE_SERVICE,
};
use crate::pump::{Fault, Main, Twin};
use crate::view::{AsMainStatus, AsTwinHead, MainStatus, TwinHead};

crate::wire!(gunnar_proto);

/// The byte ceiling of one `LogFetch` the pump asks for (the CLI's, 32 MiB).
pub const FETCH_MAX_BYTES: u64 = 32 << 20;

/// **The twin**, over its control plane.
#[derive(Clone)]
pub struct ProtoTwin {
    pub addr: String,
    pub channel: ControlChannel,
    pub key: ssh_key::PrivateKey,
}

impl Twin for ProtoTwin {
    type Chunk = gunnar_proto::LogChunk;
    type Stream = TonicStream<gunnar_proto::LogChunk>;

    fn addr(&self) -> String {
        self.addr.clone()
    }

    async fn head(&self) -> Result<TwinHead, Fault> {
        let mut c = MirrorClient::new(self.channel.clone());
        let r = c
            .log_head(authed(&self.key, gunnar_proto::LogHeadRequest::default()))
            .await
            .map_err(|s| fault(&s))?;
        Ok(r.into_inner().twin_head())
    }

    async fn snapshot(
        &self,
        from_item: u64,
        from_offset: u64,
        expect_id: &str,
    ) -> Result<Self::Stream, Fault> {
        let mut c = MirrorClient::new(self.channel.clone());
        let req = gunnar_proto::SnapshotFetchRequest {
            from_item,
            from_offset,
            expect_id: expect_id.to_owned(),
        };
        let s = c
            .snapshot_fetch(authed(&self.key, req))
            .await
            .map_err(|s| fault(&s))?;
        Ok(TonicStream(s.into_inner()))
    }

    async fn log(&self, from: u64, to: u64, from_offset: u64) -> Result<Self::Stream, Fault> {
        let mut c = MirrorClient::new(self.channel.clone());
        let req = gunnar_proto::LogFetchRequest {
            from,
            to,
            max_bytes: FETCH_MAX_BYTES,
            from_offset,
        };
        let s = c
            .log_fetch(authed(&self.key, req))
            .await
            .map_err(|s| fault(&s))?;
        Ok(TonicStream(s.into_inner()))
    }
}

/// **The main** (fresh or restoring). The Restore stream is signed with `key`, which
/// the box must have been fed as a restore key (`--restore-pubkey`, or an admin key
/// from the installer) — the same key that read the twin.
#[derive(Clone)]
pub struct ProtoMain {
    pub addr: String,
    pub channel: ControlChannel,
    pub key: ssh_key::PrivateKey,
}

impl Main<gunnar_proto::LogChunk> for ProtoMain {
    type Sink = TonicRestore<gunnar_proto::LogChunk, gunnar_proto::RestoreResponse>;

    async fn status(&self) -> Result<MainStatus, Fault> {
        let mut c = BootstrapClient::new(self.channel.clone());
        let r = c
            .status(gunnar_proto::BootstrapStatusRequest::default())
            .await
            .map_err(|s| fault(&s))?;
        Ok(r.into_inner().main_status())
    }

    async fn open_restore(&self) -> Result<Self::Sink, Fault> {
        let channel = self.channel.clone();
        let key = self.key.clone();
        Ok(restore_stream(move |body, codec| async move {
            let mut grpc = tonic::client::Grpc::new(channel);
            grpc.ready().await.map_err(|e| {
                tonic::Status::unavailable(format!("the main's channel was not ready: {e}"))
            })?;
            // A raw client_streaming call does not name its method the way a generated
            // client does; the signer binds the credential to exactly this.
            let mut request = authed(&key, body);
            request
                .extensions_mut()
                .insert(tonic::GrpcMethod::new(RESTORE_SERVICE, RESTORE_METHOD));
            grpc.client_streaming(
                request,
                tonic::codegen::http::uri::PathAndQuery::from_static(RESTORE_PATH),
                codec,
            )
            .await
        }))
    }
}

/// **The primary to fence**, signed with an admin key it holds.
#[derive(Clone)]
pub struct ProtoFence {
    pub addr: String,
    pub channel: ControlChannel,
    pub key: ssh_key::PrivateKey,
}

/// A `ReplicationRole` as its word.
pub fn role_word(role: i32) -> &'static str {
    match gunnar_proto::ReplicationRole::try_from(role) {
        Ok(gunnar_proto::ReplicationRole::Primary) => "primary",
        Ok(gunnar_proto::ReplicationRole::Standby) => STANDBY,
        Ok(gunnar_proto::ReplicationRole::Browse) => "browse",
        _ => "unknown",
    }
}

impl Fenceable for ProtoFence {
    fn addr(&self) -> String {
        self.addr.clone()
    }

    async fn set_standby(&self) -> Result<String, Fault> {
        let mut c = MirrorClient::new(self.channel.clone());
        let req = gunnar_proto::SetRoleRequest {
            role: gunnar_proto::ReplicationRole::Standby as i32,
        };
        let r = c
            .set_role(authed(&self.key, req))
            .await
            .map_err(|s| fault(&s))?;
        Ok(role_word(r.into_inner().role).to_owned())
    }

    async fn demote(&self) -> Result<usize, Fault> {
        let mut c = MirrorClient::new(self.channel.clone());
        let r = c
            .demote_replication(authed(&self.key, gunnar_proto::DemoteReplicationRequest {}))
            .await
            .map_err(|s| fault(&s))?;
        Ok(r.into_inner().targets.len())
    }
}