git-remote-iroh 0.2.1

use std::path::PathBuf;

use anyhow::Context;
use iroh::{Endpoint, SecretKey, protocol::Router};
use iroh_rings::RingGate;
use tokio::{io::AsyncWriteExt, process::Command};
use tokio_util::sync::CancellationToken;
use tracing::{Instrument, info, warn, warn_span};

use crate::{
    GIT_IROH_ALPN, GitOp,
    auth::{AuthAdmin, open_db},
    cli::{PeerCommands, Subcommands},
    keys,
};

pub struct Server {
    git_dir: PathBuf,
    key: SecretKey,
    config_dir: PathBuf,
}

impl Server {
    pub fn run_command(
        commands: &Subcommands,
        config_dir: &PathBuf,
        git_dir: &PathBuf,
    ) -> Result<(), anyhow::Error> {
        let admin = AuthAdmin::new(config_dir, git_dir)?;
        match commands {
            Subcommands::Whoami => {
                println!(
                    "{} {}",
                    admin.key.public().to_z32(),
                    admin.key_file.to_string_lossy()
                );
                Ok(())
            }
            &Subcommands::Peer {
                commands:
                    Some(PeerCommands::Add {
                        ref ring,
                        ref name,
                        ref id,
                    }),
                ..
            } => admin.add_peer(ring, name, id).context("peer add"),
            &Subcommands::Peer {
                commands: None,
                ref ring,
            } => admin
                .list_peers(ring.as_ref().unwrap())
                .context("peer list"),
            &Subcommands::Ring => admin.list_rings().context("ring list"),
        }
    }

    pub fn new(config_dir: PathBuf, git_dir: &PathBuf) -> Result<Self, anyhow::Error> {
        let (_, key) = keys::load_or_create_key(&config_dir, git_dir)?;
        Ok(Self {
            git_dir: git_dir.clone(),
            key,
            config_dir,
        })
    }

    pub async fn run(&self, cancel: CancellationToken) -> Result<(), anyhow::Error> {
        let reg = open_db(&self.config_dir)?;
        let ep = Endpoint::builder(iroh::endpoint::presets::N0)
            .secret_key(self.key.clone())
            .alpns(vec![GIT_IROH_ALPN.to_vec()])
            .bind()
            .await
            .context("failed to bind iroh endpoint")?;
        ep.online().await;
        let node_id = ep.id();
        info!(addr = format!("iroh://{node_id}"), "server started");

        let handler = GitIrohTransfer::new(self.git_dir.clone());
        let gate = RingGate::new(reg, handler);
        let router = Router::builder(ep.clone())
            .accept(GIT_IROH_ALPN, gate)
            .spawn();
        cancel.cancelled().await;
        router.shutdown().await?;
        Ok(())
    }
}

/// Git protocol handler for iroh.
///
/// Accepts incoming connections, reads the operation, spawns the appropriate
/// command, connecting stdin/stdout to the iroh bidirectional stream.
#[derive(Debug, Clone)]
pub struct GitIrohTransfer {
    /// the git repo directory to serve
    repo_dir: PathBuf,
}

impl GitIrohTransfer {
    pub fn new(repo_dir: PathBuf) -> Self {
        Self { repo_dir }
    }
}

impl iroh_rings::Transfer for GitIrohTransfer {
    async fn can_access(&self, peer: &iroh::EndpointId, resource_id: &[u8]) -> bool {
        // Allow public read access to any peer
        if let Ok(GitOp::ClientFetchServerPush) = GitOp::from_resource_id(resource_id) {
            info!(%peer, "allowing public read");
            true
        } else {
            false
        }
    }

    async fn transfer(
        &self,
        resource_id: &[u8],
        send: &mut iroh::endpoint::SendStream,
        recv: &mut iroh::endpoint::RecvStream,
    ) -> anyhow::Result<()> {
        let op = GitOp::from_resource_id(resource_id)?;
        info!(?op);

        // spawn the git command
        let mut child = Command::new(op.to_string())
            .arg(&self.repo_dir)
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::inherit())
            .spawn()
            .map_err(|e| std::io::Error::other(format!("spawn {op}: {e}")))?;

        let mut child_stdin = child.stdin.take().unwrap();
        let mut child_stdout = child.stdout.take().unwrap();

        let cancel = CancellationToken::new();
        {
            let cancel = cancel.clone();
            tokio::spawn(
                async move {
                    let status = child.wait().await?;
                    if !status.success() {
                        warn!(%op, %status);
                    }
                    cancel.cancel();
                    Ok::<(), anyhow::Error>(())
                }
                .instrument(warn_span!("process wait")),
            );
        }

        let (r1, r2) = tokio::join!(
            async {
                tokio::select! {
                    res = tokio::io::copy(recv, &mut child_stdin) => { res?; Ok(()) }
                    _ = cancel.cancelled() => { recv.stop(0u8.into())?; Ok::<(), anyhow::Error>(()) }
                }
            },
            async {
                tokio::select! {
                    res = tokio::io::copy(&mut child_stdout, send) => { res?; Ok(()) }
                    _ = cancel.cancelled() => { send.flush().await?; send.finish()?; Ok::<(), anyhow::Error>(()) }
                }
            },
        );
        r1.context("stdin -> send")?;
        r2.context("recv -> stdout")?;
        Ok(())
    }
}