Skip to main content

git_remote_iroh/
server.rs

1use std::path::PathBuf;
2
3use anyhow::Context;
4use clap::Parser;
5use iroh::{
6    Endpoint, SecretKey,
7    protocol::{AcceptError, ProtocolHandler, Router},
8};
9use tokio::{io::AsyncReadExt, process::Command, task::JoinSet};
10use tokio_util::sync::CancellationToken;
11use tracing::{info, instrument, trace, warn};
12
13use crate::{GIT_IROH_ALPN, GitOp, keys};
14
15#[derive(Parser, Debug)]
16#[command()]
17pub struct ServerArgs {
18    #[arg(env = "GIT_DIR")]
19    git_dir: String,
20}
21
22pub struct Server {
23    args: ServerArgs,
24    key: SecretKey,
25}
26
27impl Server {
28    pub fn new(args: ServerArgs) -> Result<Self, anyhow::Error> {
29        let key = keys::load_or_create_key(&args.git_dir)?;
30        Ok(Self { args, key })
31    }
32
33    pub async fn run(&self, cancel: CancellationToken) -> Result<(), anyhow::Error> {
34        let ep = Endpoint::builder(iroh::endpoint::presets::N0)
35            .secret_key(self.key.clone())
36            .alpns(vec![GIT_IROH_ALPN.to_vec()])
37            .bind()
38            .await
39            .context("failed to bind iroh endpoint")?;
40        ep.online().await;
41        let node_id = ep.id();
42        info!(addr = format!("iroh://{node_id}"), "server started");
43
44        let handler = GitIrohHandler::new(self.args.git_dir.as_str().into());
45        let router = Router::builder(ep.clone())
46            .accept(GIT_IROH_ALPN, handler)
47            .spawn();
48        cancel.cancelled().await;
49        router.shutdown().await?;
50        Ok(())
51    }
52}
53
54/// Git protocol handler for iroh.
55///
56/// Accepts incoming connections, reads the operation, spawns the appropriate
57/// command, connecting stdin/stdout to the iroh bidirectional stream.
58#[derive(Debug, Clone)]
59pub struct GitIrohHandler {
60    /// the git repo directory to serve
61    repo_dir: PathBuf,
62}
63
64impl GitIrohHandler {
65    pub fn new(repo_dir: PathBuf) -> Self {
66        Self { repo_dir }
67    }
68}
69
70impl ProtocolHandler for GitIrohHandler {
71    #[instrument(skip(conn), err)]
72    async fn accept(&self, conn: iroh::endpoint::Connection) -> Result<(), AcceptError> {
73        let remote = conn.remote_id();
74        info!(%remote, "accept");
75
76        let (mut send, mut recv) = conn.accept_bi().await?;
77
78        // read operation byte
79        let op_byte = recv
80            .read_u8()
81            .await
82            .map_err(|e| std::io::Error::other(format!("read op: {e}")))?;
83
84        let op = GitOp::from_byte(op_byte)
85            .ok_or_else(|| std::io::Error::other(format!("unknown git op: {}", op_byte)))?;
86
87        info!(?op, %remote);
88
89        // spawn the git command
90        let mut child = Command::new(op.to_string())
91            .arg(&self.repo_dir)
92            .stdin(std::process::Stdio::piped())
93            .stdout(std::process::Stdio::piped())
94            .stderr(std::process::Stdio::inherit())
95            .spawn()
96            .map_err(|e| std::io::Error::other(format!("spawn {op}: {e}")))?;
97
98        let mut child_stdin = child.stdin.take().unwrap();
99        let mut child_stdout = child.stdout.take().unwrap();
100
101        let cancel = CancellationToken::new();
102        let mut tasks: JoinSet<anyhow::Result<()>> = JoinSet::new();
103        tasks.spawn(async move {
104            tokio::io::copy(&mut recv, &mut child_stdin)
105                .await
106                .context("recv -> stdin")?;
107            Ok(())
108        });
109        tasks.spawn(async move {
110            tokio::io::copy(&mut child_stdout, &mut send)
111                .await
112                .context("stdout -> send")?;
113            send.finish()?;
114            Ok(())
115        });
116        loop {
117            tokio::select! {
118                _ = cancel.cancelled() => {
119                    child.kill().await?;
120                }
121                res = tasks.join_next() => {
122                    trace!(?res);
123                    cancel.cancel();
124                    tasks.abort_all();
125                }
126                res = child.wait() => {
127                    let status = res?;
128                    if !status.success() {
129                        warn!(%op, %status, "process exit");
130                    }
131                    tasks.abort_all();
132                    break;
133                }
134            }
135        }
136
137        info!(%op, %remote, "complete");
138        Ok(())
139    }
140}