Allow downloads of outputs

src/config.rs

@@ -25,6 +25,7 @@ pub struct Config {
     pub interval_external_services: u64,
     pub submission_base_image: String,
     pub submission_base_image_tag: String,
+    pub serializer_secret_key: String,
 
     pub s3_prefix: String,  // Prefix within the bucket, ie. labcaller-dev
     pub pod_prefix: String, // What is prefixed to the pod name, ie. labcaller-dev}
@@ -80,6 +81,8 @@ impl Config {
                 .expect("SUBMISSION_BASE_IMAGE must be set"),
             submission_base_image_tag: env::var("SUBMISSION_BASE_IMAGE_TAG")
                 .expect("SUBMISSION_BASE_IMAGE_TAG must be set"),
+            serializer_secret_key: env::var("SERIALIZER_SECRET_KEY")
+                .expect("SERIALIZER_SECRET_KEY must be set"),
             db_prefix,
             db_url,
             s3_prefix,

src/external/s3/models.rs

@@ -24,9 +24,10 @@ impl From<aws_sdk_s3::types::Object> for OutputObject {
 #[derive(ToSchema, Serialize, FromQueryResult, Debug)]
 pub struct OutputObjectResponse {
     // Let's not show the full key path through the API
-    filename: String,
-    last_modified: DateTime<Utc>,
-    size_bytes: i64,
+    pub filename: String,
+    pub last_modified: DateTime<Utc>,
+    pub size_bytes: i64,
+    pub url: Option<String>,
 }
 
 impl From<OutputObject> for OutputObjectResponse {
@@ -37,6 +38,7 @@ impl From<OutputObject> for OutputObjectResponse {
             last_modified: model.last_modified,
             filename: filename,
             size_bytes: model.size_bytes,
+            url: None,
         }
     }
 }

src/main.rs

@@ -52,6 +52,10 @@ async fn main() {
         .route("/healthz", get(common::views::healthz))
         .route("/api/config", get(common::views::get_ui_config))
         .route("/api/status", get(common::views::get_status))
+        .route(
+            "/api/submissions/download/:token",
+            get(submissions::views::download_file),
+        )
         .with_state(db.clone())
         .nest(
             "/api/submissions",

src/submissions/models.rs

@@ -58,7 +58,19 @@ impl
         let submission = model_tuple.0;
         let uploads = model_tuple.1;
         let status = model_tuple.2;
-        let outputs = model_tuple.3;
+        let mut outputs: Vec<crate::external::s3::models::OutputObjectResponse> = model_tuple
+            .3
+            .into_iter()
+            .map(|output| output.into())
+            .collect();
+
+        // Set the url for each output object
+        for output in outputs.iter_mut() {
+            output.url = Some(format!(
+                "/api/submissions/{}/{}",
+                submission.id, output.filename
+            ));
+        }
         Self {
             id: submission.id,
             name: submission.name,
@@ -73,7 +85,7 @@ impl
                 .map(|association| association.into())
                 .collect(),
             status: status.into_iter().map(|status| status.into()).collect(),
-            outputs: outputs.into_iter().map(|output| output.into()).collect(),
+            outputs: outputs,
         }
     }
 }
@@ -145,3 +157,14 @@ impl SubmissionUpdate {
         model
     }
 }
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct DownloadToken {
+    pub token: String,
+}
+#[derive(Debug, Serialize, Deserialize)]
+pub(super) struct Claims {
+    pub(super) submission_id: Uuid,
+    pub(super) filename: String,
+    pub(super) exp: usize,
+}

src/submissions/views.rs

@@ -9,15 +9,18 @@ use crate::external::k8s::crd::{
 use anyhow::Result;
 use aws_sdk_s3::Client as S3Client;
 use axum::{
+    body::Bytes,
     debug_handler,
     extract::{Path, Query, State},
-    http::StatusCode,
+    http::{header, StatusCode},
     response::IntoResponse,
     routing, Json, Router,
 };
 use axum_keycloak_auth::{
     instance::KeycloakAuthInstance, layer::KeycloakAuthLayer, PassthroughMode,
 };
+use chrono::{Duration, Utc};
+use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation};
 use kube::{api::PostParams, Api};
 use rand::Rng;
 use sea_orm::{
@@ -41,6 +44,7 @@ pub fn router(
                 .delete(delete_one)
                 .post(execute_workflow),
         )
+        .route("/:id/:filename", routing::get(generate_download_token))
         .with_state((db, s3))
         .layer(
             KeycloakAuthLayer::<Role>::builder()
@@ -170,11 +174,9 @@ pub async fn get_one(
         Ok(obj) => obj.unwrap(),
         _ => return Err((StatusCode::NOT_FOUND, Json("Not Found".to_string()))),
     };
-    let outputs: Vec<crate::external::s3::models::OutputObject> =
-        crate::external::s3::services::get_outputs_from_id(s3, obj.id)
-            .await
-            .unwrap();
-
+    let outputs = crate::external::s3::services::get_outputs_from_id(s3, obj.id)
+        .await
+        .unwrap();
     let uploads = match obj.find_related(crate::uploads::db::Entity).all(&db).await {
         // Return all or none. If any fail, return an error
         Ok(uploads) => Some(uploads),
@@ -354,3 +356,79 @@ pub async fn execute_workflow(
         }
     }
 }
+
+pub async fn generate_download_token(
+    Path((submission_id, filename)): Path<(Uuid, String)>,
+) -> Result<Json<super::models::DownloadToken>, (StatusCode, String)> {
+    let config: crate::config::Config = crate::config::Config::from_env();
+    let expiration = Utc::now() + Duration::hours(1); // Token expiry in 1 hour
+    let claims = super::models::Claims {
+        submission_id,
+        filename: filename.clone(),
+        exp: expiration.timestamp() as usize,
+    };
+
+    let token = encode(
+        &Header::default(),
+        &claims,
+        &EncodingKey::from_secret(config.serializer_secret_key.as_ref()),
+    )
+    .map_err(|_| {
+        (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            "Token creation failed".to_string(),
+        )
+    })?;
+
+    Ok(Json(super::models::DownloadToken { token }))
+}
+
+pub async fn download_file(
+    Path(token): Path<String>,
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    let config: crate::config::Config = crate::config::Config::from_env();
+    let s3 = crate::external::s3::services::get_client(&config).await;
+
+    let token_data = decode::<super::models::Claims>(
+        &token,
+        &DecodingKey::from_secret(config.serializer_secret_key.as_ref()),
+        &Validation::default(),
+    )
+    .map_err(|_| (StatusCode::UNAUTHORIZED, "Invalid token".to_string()))?;
+
+    let claims = token_data.claims;
+    if claims.exp < Utc::now().timestamp() as usize {
+        return Err((StatusCode::UNAUTHORIZED, "Token expired".to_string()));
+    }
+
+    let key = format!(
+        "{}/outputs/{}/{}",
+        config.s3_prefix, claims.submission_id, claims.filename
+    );
+
+    println!("Key: {}", key);
+    println!("Token data: {:?}", claims);
+    let object = s3
+        .get_object()
+        .bucket(config.s3_bucket)
+        .key(key)
+        .send()
+        .await
+        .map_err(|_| (StatusCode::NOT_FOUND, "File not found".to_string()))?;
+
+    let body = object.body.collect().await.map_err(|_| {
+        (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            "Failed to read file".to_string(),
+        )
+    })?;
+
+    Ok((
+        StatusCode::OK,
+        [(
+            header::CONTENT_DISPOSITION,
+            format!("attachment; filename=\"{}\"", claims.filename),
+        )],
+        Bytes::from(body.into_bytes()), // Ensures compatibility with `IntoResponse`
+    ))
+}