From 51cda94c1a216514712788a113078c58aff84636 Mon Sep 17 00:00:00 2001
From: Chewing_Bever <roosensjef@gmail.com>
Date: Wed, 12 Jul 2023 11:04:31 +0200
Subject: [PATCH] feat: streaming upload of archives

---
 server/Cargo.toml      |  3 +++
 server/src/repo/mod.rs | 46 ++++++++++++++++++++++++++++++++++++++----
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/server/Cargo.toml b/server/Cargo.toml
index c5714c1..e2b56b9 100644
--- a/server/Cargo.toml
+++ b/server/Cargo.toml
@@ -7,6 +7,9 @@ edition = "2021"
 
 [dependencies]
 axum = "0.6.18"
+futures = "0.3.28"
 tokio = { version = "1.29.1", features = ["full"] }
+tokio-util = { version = "0.7.8", features = ["io"] }
 tower = { version = "0.4.13", features = ["make"] }
 tower-http = { version = "0.4.1", features = ["fs"] }
+uuid = { version = "1.4.0", features = ["v4"] }
diff --git a/server/src/repo/mod.rs b/server/src/repo/mod.rs
index d58bc68..d02040c 100644
--- a/server/src/repo/mod.rs
+++ b/server/src/repo/mod.rs
@@ -1,12 +1,50 @@
-use axum::extract::{Path, State};
-use axum::routing::get_service;
+use axum::extract::{BodyStream, Path, State};
+use axum::http::StatusCode;
+use axum::response::IntoResponse;
+use axum::routing::{get_service, post};
 use axum::Router;
+use futures::StreamExt;
+use futures::TryFutureExt;
+use futures::TryStreamExt;
+use tokio::{fs, io, io::AsyncWriteExt};
 use tower_http::services::ServeDir;
+use uuid::Uuid;
 
 pub fn router(config: &crate::Config) -> Router<crate::Config> {
     // Try to serve packages by default, and try the database files instead if not found
-    let service = ServeDir::new(&config.pkg_dir).fallback(ServeDir::new(&config.repo_dir));
+    let serve_repos =
+        get_service(ServeDir::new(&config.pkg_dir).fallback(ServeDir::new(&config.repo_dir)));
     Router::new()
-        .route("/*path", get_service(service))
+        .route(
+            "/:repo",
+            post(post_package_archive).get(serve_repos.clone()),
+        )
+        .fallback(serve_repos)
         .with_state(config.clone())
 }
+
+async fn post_package_archive(
+    State(config): State<crate::Config>,
+    Path(repo): Path<String>,
+    body: BodyStream,
+) -> Result<(), StatusCode> {
+    // let mut body_reader = tokio_util::io::StreamReader::new(
+    //     body.map_err(|err| io::Error::new(io::ErrorKind::Other, err)),
+    // );
+    let mut body = body.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR);
+
+    // We first stream the uploaded file to disk
+    let uuid: uuid::fmt::Simple = Uuid::new_v4().into();
+    let path = config.pkg_dir.join(uuid.to_string());
+    let mut f = fs::File::create(&path)
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
+        .await?;
+
+    while let Some(chunk) = body.next().await {
+        f.write_all(&chunk?)
+            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
+            .await?;
+    }
+
+    Ok(())
+}