git-remote-iroh 0.1.0

use std::path::PathBuf;

use anyhow::Context;
use clap::Parser;
use iroh::{
    Endpoint, SecretKey,
    protocol::{AcceptError, ProtocolHandler, Router},
};
use tokio::{io::AsyncReadExt, process::Command, task::JoinSet};
use tokio_util::sync::CancellationToken;
use tracing::{info, instrument, trace, warn};

use crate::{GIT_IROH_ALPN, GitOp, keys};

#[derive(Parser, Debug)]
#[command()]
pub struct ServerArgs {
    #[arg(env = "GIT_DIR")]
    git_dir: String,
}

pub struct Server {
    args: ServerArgs,
    key: SecretKey,
}

impl Server {
    pub fn new(args: ServerArgs) -> Result<Self, anyhow::Error> {
        let key = keys::load_or_create_key(&args.git_dir)?;
        Ok(Self { args, key })
    }

    pub async fn run(&self, cancel: CancellationToken) -> Result<(), anyhow::Error> {
        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 = GitIrohHandler::new(self.args.git_dir.as_str().into());
        let router = Router::builder(ep.clone())
            .accept(GIT_IROH_ALPN, handler)
            .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 GitIrohHandler {
    /// the git repo directory to serve
    repo_dir: PathBuf,
}

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

impl ProtocolHandler for GitIrohHandler {
    #[instrument(skip(conn), err)]
    async fn accept(&self, conn: iroh::endpoint::Connection) -> Result<(), AcceptError> {
        let remote = conn.remote_id();
        info!(%remote, "accept");

        let (mut send, mut recv) = conn.accept_bi().await?;

        // read operation byte
        let op_byte = recv
            .read_u8()
            .await
            .map_err(|e| std::io::Error::other(format!("read op: {e}")))?;

        let op = GitOp::from_byte(op_byte)
            .ok_or_else(|| std::io::Error::other(format!("unknown git op: {}", op_byte)))?;

        info!(?op, %remote);

        // 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 mut tasks: JoinSet<anyhow::Result<()>> = JoinSet::new();
        tasks.spawn(async move {
            tokio::io::copy(&mut recv, &mut child_stdin)
                .await
                .context("recv -> stdin")?;
            Ok(())
        });
        tasks.spawn(async move {
            tokio::io::copy(&mut child_stdout, &mut send)
                .await
                .context("stdout -> send")?;
            send.finish()?;
            Ok(())
        });
        loop {
            tokio::select! {
                _ = cancel.cancelled() => {
                    child.kill().await?;
                }
                res = tasks.join_next() => {
                    trace!(?res);
                    cancel.cancel();
                    tasks.abort_all();
                }
                res = child.wait() => {
                    let status = res?;
                    if !status.success() {
                        warn!(%op, %status, "process exit");
                    }
                    tasks.abort_all();
                    break;
                }
            }
        }

        info!(%op, %remote, "complete");
        Ok(())
    }
}