From 5431c2d9d3afc0b2d4c42fc58bc370003d7eda4d Mon Sep 17 00:00:00 2001
From: SeanOMik <seanomik@gmail.com>
Date: Tue, 18 Apr 2023 01:16:29 -0400
Subject: [PATCH] Implement check upload status endpoint

---
 src/api/uploads.rs | 25 ++++++++++++++++++++++---
 src/main.rs        |  6 ++++--
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src/api/uploads.rs b/src/api/uploads.rs
index 0104377..e8a0362 100644
--- a/src/api/uploads.rs
+++ b/src/api/uploads.rs
@@ -1,4 +1,4 @@
-use actix_web::{HttpResponse, HttpRequest, post, web, patch, put, delete};
+use actix_web::{HttpResponse, HttpRequest, post, web, patch, put, delete, get};
 use bytes::{BytesMut, Bytes, BufMut};
 use qstring::QString;
 use tracing::{debug};
@@ -84,8 +84,6 @@ pub async fn finish_chunked_upload(body: Bytes, path: web::Path<(String, String)
 
     database.replace_digest(&layer_uuid, &digest).await.unwrap();
 
-    //let first_digest = digest.first().unwrap().to_owned();
-
     HttpResponse::Created()
         .insert_header(("Location", format!("/v2/{}/blobs/{}", name, digest)))
         .insert_header(("Content-Length", 0))
@@ -103,4 +101,25 @@ pub async fn cancel_upload(path: web::Path<(String, String)>, state: web::Data<A
     // I'm not sure what this response should be, its not specified in the registry spec.
     HttpResponse::Ok()
         .finish()
+}
+
+#[get("/{uuid}")]
+pub async fn check_upload_status(path: web::Path<(String, String)>, state: web::Data<AppState>) -> HttpResponse {
+    let (name, layer_uuid) = (path.0.to_owned(), path.1.to_owned());
+    
+    let database = &state.database;
+    let ending = match database.get_digest(&layer_uuid).await.unwrap() {
+        Some(current_bytes) => {
+            current_bytes.len()
+        },
+        None => {
+            0
+        }
+    };
+
+    HttpResponse::Created()
+        .insert_header(("Location", format!("/v2/{}/blobs/uploads/{}", name, layer_uuid)))
+        .insert_header(("Range", format!("0-{}", ending)))
+        .insert_header(("Docker-Upload-Digest", layer_uuid))
+        .finish()
 }
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index 49ade06..fde28b3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,7 +15,7 @@ use database::Database;
 #[actix_web::main]
 async fn main() -> std::io::Result<()> {
     let pool = SqlitePoolOptions::new()
-        .max_connections(5)
+        .max_connections(15)
         .connect("test.db").await.unwrap();
 
     pool.create_schema().await.unwrap();
@@ -28,7 +28,8 @@ async fn main() -> std::io::Result<()> {
         .with_max_level(Level::DEBUG)
         .init();
 
-    let payload_config = web::PayloadConfig::new(31_460_000); // 30mb
+    // TODO: Make configurable by deployment
+    let payload_config = web::PayloadConfig::new(5 * 1024 * 1024 * 1024); // 5Gb 
 
     debug!("Starting http server...");
 
@@ -68,6 +69,7 @@ async fn main() -> std::io::Result<()> {
                                             .service(api::uploads::chunked_upload_layer)
                                             .service(api::uploads::finish_chunked_upload)
                                             .service(api::uploads::cancel_upload)
+                                            .service(api::uploads::check_upload_status)
                                             // TODO: Cross Repository Blob Mount
                                     )
                             )