Add ability to download tar.gz of commit/branch/tag

Author: w4

Cargo.lock

@@ -33,10 +33,13 @@ serde = { version = "1.0", features = ["derive", "rc"] }
 sha2 = "0.10"
 syntect = "5"
 sled = { version = "0.34", features = ["compression"] }
+tar = "0.4"
+flate2 = "1.0"
 time = { version = "0.3", features = ["serde"] }
 timeago = { version = "0.4.2", default-features = false }
 tokio = { version = "1.19", features = ["full"] }
 tokio-util = { version = "0.7.3", features = ["io"] }
+tokio-stream = "0.1"
 tower = "0.4"
 tower-service = "0.3"
 tower-layer = "0.3"

Cargo.toml

@@ -7,12 +7,12 @@ use std::{
     time::Duration,
 };
 
-use anyhow::{Context, Result};
-use bytes::{Bytes, BytesMut};
+use anyhow::{anyhow, Context, Result};
+use bytes::{BufMut, Bytes, BytesMut};
 use comrak::{ComrakOptions, ComrakPlugins};
 use git2::{
     DiffFormat, DiffLineType, DiffOptions, DiffStatsFormat, Email, EmailCreateOptions, ObjectType,
-    Oid, Signature,
+    Oid, Signature, TreeWalkResult,
 };
 use moka::future::Cache;
 use parking_lot::Mutex;
@@ -22,7 +22,7 @@ use syntect::{
     util::LinesWithEndings,
 };
 use time::OffsetDateTime;
-use tracing::instrument;
+use tracing::{error, instrument, warn};
 
 use crate::syntax_highlight::ComrakSyntectAdapter;
 
@@ -272,6 +272,16 @@ impl OpenRepository {
             .await
     }
 
+    pub async fn default_branch(self: Arc<Self>) -> Result<Option<String>> {
+        tokio::task::spawn_blocking(move || {
+            let repo = self.repo.lock();
+            let head = repo.head().context("Couldn't find HEAD of repository")?;
+            Ok(head.shorthand().map(ToString::to_string))
+        })
+        .await
+        .context("Failed to join Tokio task")?
+    }
+
     #[instrument(skip(self))]
     pub async fn latest_commit(self: Arc<Self>) -> Result<Commit> {
         tokio::task::spawn_blocking(move || {
@@ -299,6 +309,87 @@ impl OpenRepository {
         .context("Failed to join Tokio task")?
     }
 
+    #[instrument(skip_all)]
+    pub async fn archive(
+        self: Arc<Self>,
+        res: tokio::sync::mpsc::Sender<Result<Bytes, anyhow::Error>>,
+        cont: tokio::sync::oneshot::Sender<()>,
+        commit: Option<&str>,
+    ) -> Result<(), anyhow::Error> {
+        const BUFFER_CAP: usize = 512 * 1024;
+
+        let commit = commit
+            .map(Oid::from_str)
+            .transpose()
+            .context("failed to build oid")?;
+
+        tokio::task::spawn_blocking(move || {
+            let buffer = BytesMut::with_capacity(BUFFER_CAP + 1024);
+
+            let flate = flate2::write::GzEncoder::new(buffer.writer(), flate2::Compression::fast());
+            let mut archive = tar::Builder::new(flate);
+
+            let repo = self.repo.lock();
+
+            let tree = if let Some(commit) = commit {
+                repo.find_commit(commit)?.tree()?
+            } else if let Some(reference) = &self.branch {
+                repo.resolve_reference_from_short_name(reference)?
+                    .peel_to_tree()?
+            } else {
+                repo.head()
+                    .context("Couldn't find HEAD of repository")?
+                    .peel_to_tree()?
+            };
+
+            // tell the web server it can send response headers to the requester
+            if cont.send(()).is_err() {
+                return Err(anyhow!("requester gone"));
+            }
+
+            let mut callback = |root: &str, entry: &git2::TreeEntry| -> TreeWalkResult {
+                if let Ok(blob) = entry.to_object(&repo).unwrap().peel_to_blob() {
+                    let path =
+                        Path::new(root).join(String::from_utf8_lossy(entry.name_bytes()).as_ref());
+
+                    let mut header = tar::Header::new_gnu();
+                    if let Err(error) = header.set_path(&path) {
+                        warn!(%error, "Attempted to write invalid path to archive");
+                        return TreeWalkResult::Skip;
+                    }
+                    header.set_size(blob.size() as u64);
+                    #[allow(clippy::cast_sign_loss)]
+                    header.set_mode(entry.filemode() as u32);
+                    header.set_cksum();
+
+                    if let Err(error) = archive.append(&header, blob.content()) {
+                        error!(%error, "Failed to write blob to archive");
+                        return TreeWalkResult::Abort;
+                    }
+                }
+
+                if archive.get_ref().get_ref().get_ref().len() >= BUFFER_CAP {
+                    let b = archive.get_mut().get_mut().get_mut().split().freeze();
+                    if let Err(error) = res.blocking_send(Ok(b)) {
+                        error!(%error, "Failed to send buffer to client");
+                        return TreeWalkResult::Abort;
+                    }
+                }
+
+                TreeWalkResult::Ok
+            };
+
+            tree.walk(git2::TreeWalkMode::PreOrder, &mut callback)?;
+
+            res.blocking_send(Ok(archive.into_inner()?.finish()?.into_inner().freeze()))?;
+
+            Ok::<_, anyhow::Error>(())
+        })
+        .await??;
+
+        Ok(())
+    }
+
     #[instrument(skip(self))]
     pub async fn commit(self: Arc<Self>, commit: &str) -> Result<Arc<Commit>, Arc<anyhow::Error>> {
         let commit = Oid::from_str(commit)

src/git.rs

@@ -17,6 +17,8 @@ pub struct View {
     pub repo: Repository,
     pub commit: Arc<Commit>,
     pub branch: Option<Arc<str>>,
+    pub dl_branch: Arc<str>,
+    pub id: Option<String>,
 }
 
 #[derive(Deserialize)]
@@ -33,8 +35,23 @@ pub async fn handle(
     Query(query): Query<UriQuery>,
 ) -> Result<Response> {
     let open_repo = git.repo(repository_path, query.branch.clone()).await?;
-    let commit = if let Some(commit) = query.id {
-        open_repo.commit(&commit).await?
+
+    let dl_branch = if let Some(branch) = query.branch.clone() {
+        branch
+    } else {
+        Arc::from(
+            open_repo
+                .clone()
+                .default_branch()
+                .await
+                .ok()
+                .flatten()
+                .unwrap_or_else(|| "master".to_string()),
+        )
+    };
+
+    let commit = if let Some(commit) = query.id.as_deref() {
+        open_repo.commit(commit).await?
     } else {
         Arc::new(open_repo.latest_commit().await?)
     };
@@ -43,5 +60,7 @@ pub async fn handle(
         repo,
         commit,
         branch: query.branch,
+        id: query.id,
+        dl_branch,
     }))
 }

src/methods/repo/commit.rs

@@ -4,6 +4,7 @@ mod diff;
 mod log;
 mod refs;
 mod smart_git;
+mod snapshot;
 mod summary;
 mod tag;
 mod tree;
@@ -32,6 +33,7 @@ use self::{
     log::handle as handle_log,
     refs::handle as handle_refs,
     smart_git::handle as handle_smart_git,
+    snapshot::handle as handle_snapshot,
     summary::handle as handle_summary,
     tag::handle as handle_tag,
     tree::handle as handle_tree,
@@ -89,6 +91,7 @@ where
         Some("diff") => h!(handle_diff),
         Some("patch") => h!(handle_patch),
         Some("tag") => h!(handle_tag),
+        Some("snapshot") => h!(handle_snapshot),
         Some(v) => {
             uri_parts.push(v);
 

src/methods/repo/mod.rs

@@ -0,0 +1,86 @@
+use std::sync::Arc;
+
+use anyhow::{anyhow, Context};
+use axum::{
+    body::{boxed, Body, BoxBody},
+    extract::Query,
+    http::Response,
+    Extension,
+};
+use serde::Deserialize;
+use tokio_stream::wrappers::ReceiverStream;
+use tracing::{error, info_span, Instrument};
+
+use super::{RepositoryPath, Result};
+use crate::git::Git;
+
+#[derive(Deserialize)]
+pub struct UriQuery {
+    #[serde(rename = "h")]
+    branch: Option<Arc<str>>,
+    id: Option<Arc<str>>,
+}
+
+pub async fn handle(
+    Extension(RepositoryPath(repository_path)): Extension<RepositoryPath>,
+    Extension(git): Extension<Arc<Git>>,
+    Query(query): Query<UriQuery>,
+) -> Result<Response<BoxBody>> {
+    let open_repo = git.repo(repository_path, query.branch.clone()).await?;
+
+    // byte stream back to the client
+    let (send, recv) = tokio::sync::mpsc::channel(1);
+
+    // channel for `archive` to tell us we can send headers etc back to
+    // the user so it has time to return an error
+    let (send_cont, recv_cont) = tokio::sync::oneshot::channel();
+
+    let id = query.id.clone();
+
+    let res = tokio::spawn(
+        async move {
+            if let Err(error) = open_repo
+                .archive(send.clone(), send_cont, id.as_deref())
+                .await
+            {
+                error!(%error, "Failed to build archive for client");
+                let _res = send.send(Err(anyhow!("archive builder failed"))).await;
+                return Err(error);
+            }
+
+            Ok(())
+        }
+        .instrument(info_span!("sender")),
+    );
+
+    // don't send any headers until `archive` has told us we're good
+    // to continue
+    if recv_cont.await.is_err() {
+        // sender disappearing means `archive` hit an issue during init, lets
+        // wait for the error back from the spawned tokio task to return to
+        // the client
+        res.await
+            .context("Tokio task failed")?
+            .context("Failed to build archive")?;
+
+        // ok, well this isn't ideal. the sender disappeared but we never got
+        // an error. this shouldn't be possible, i guess lets just return an
+        // internal error
+        return Err(anyhow!("Ran into inconsistent error state whilst building archive, please file an issue at https://github.com/w4/rgit/issues").into());
+    }
+
+    let file_name = query
+        .id
+        .as_deref()
+        .or(query.branch.as_deref())
+        .unwrap_or("main");
+
+    Ok(Response::builder()
+        .header("Content-Type", "application/gzip")
+        .header(
+            "Content-Disposition",
+            format!("attachment; filename=\"{file_name}.tar.gz\""),
+        )
+        .body(boxed(Body::wrap_stream(ReceiverStream::new(recv))))
+        .context("failed to build response")?)
+}

src/methods/repo/snapshot.rs

@@ -35,6 +35,10 @@
         <td colspan="2"><pre><a href="/{{ repo.display() }}/commit?id={{ parent }}{% call link::maybe_branch_suffix(branch) %}" class="no-style">{{ parent }}</a></pre></td>
     </tr>
     {%- endfor %}
+    <tr>
+        <th>download (tar.gz)</th>
+        <td colspan="2"><pre><a href="/{{ repo.display() }}/snapshot?{% if let Some(id) = id %}id={{ id }}{% else %}h={{ dl_branch }}{% endif %}">{{ id.as_deref().unwrap_or(dl_branch.as_ref()) }}</a></pre></td>
+    </tr>
     </tbody>
 </table>
 

templates/repo/commit.html

@@ -31,6 +31,10 @@
             </td>
         </tr>
     {% endif %}
+    <tr>
+        <th>download (tar.gz)</th>
+        <td colspan="2"><pre><a href="/{{ repo.display() }}/snapshot?h={{ tag.name }}">{{ tag.name }}</a></pre></td>
+    </tr>
     </tbody>
 </table>