From 628f08617dc90a37a02f5e74dd8afe2054b3d495 Mon Sep 17 00:00:00 2001 From: Evan Thomas Date: Thu, 31 Oct 2024 16:39:18 +0100 Subject: [PATCH] Allow downloads of outputs --- src/config.rs | 3 ++ src/external/s3/models.rs | 8 ++-- src/main.rs | 4 ++ src/submissions/models.rs | 27 +++++++++++- src/submissions/views.rs | 90 ++++++++++++++++++++++++++++++++++++--- 5 files changed, 121 insertions(+), 11 deletions(-) diff --git a/src/config.rs b/src/config.rs index 58b43c8..5a792e3 100644 --- a/src/config.rs +++ b/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, diff --git a/src/external/s3/models.rs b/src/external/s3/models.rs index d761706..85b8b2b 100644 --- a/src/external/s3/models.rs +++ b/src/external/s3/models.rs @@ -24,9 +24,10 @@ impl From 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, - size_bytes: i64, + pub filename: String, + pub last_modified: DateTime, + pub size_bytes: i64, + pub url: Option, } impl From for OutputObjectResponse { @@ -37,6 +38,7 @@ impl From for OutputObjectResponse { last_modified: model.last_modified, filename: filename, size_bytes: model.size_bytes, + url: None, } } } diff --git a/src/main.rs b/src/main.rs index e81f292..8eb2552 100644 --- a/src/main.rs +++ b/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", diff --git a/src/submissions/models.rs b/src/submissions/models.rs index 0ee48cc..eb57288 100644 --- a/src/submissions/models.rs +++ b/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 = 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, +} diff --git a/src/submissions/views.rs b/src/submissions/views.rs index 16bb644..c40f27e 100644 --- a/src/submissions/views.rs +++ b/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::::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::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, (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, +) -> Result { + let config: crate::config::Config = crate::config::Config::from_env(); + let s3 = crate::external::s3::services::get_client(&config).await; + + let token_data = decode::( + &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` + )) +}