From 72fdd75568cfe551b91351fb530decfd33152c10 Mon Sep 17 00:00:00 2001 From: Michael Chung Date: Thu, 30 Nov 2023 17:33:25 -0500 Subject: [PATCH] added upload weights / create weight endpoint original file for the weights weights file type for safetensors or checkpoint Gotta add one more thing for record creation --- crates/_types/buckets/src/public/mod.rs | 3 +- .../src/public/weight_uploads/directory.rs | 77 ++++ .../buckets/src/public/weight_uploads/mod.rs | 2 + .../public/weight_uploads/original_file.rs | 104 ++++++ .../enums/src/by_table/model_weights/mod.rs | 1 + .../model_weights/weights_file_type.rs | 68 ++++ .../tokens/src/tokens/weight_uploads.rs | 14 + .../endpoints/weights/create_weight.rs | 335 ++++++++++++++++++ .../src/http_server/endpoints/weights/mod.rs | 3 +- 9 files changed, 605 insertions(+), 2 deletions(-) create mode 100644 crates/_types/buckets/src/public/weight_uploads/directory.rs create mode 100644 crates/_types/buckets/src/public/weight_uploads/mod.rs create mode 100644 crates/_types/buckets/src/public/weight_uploads/original_file.rs create mode 100644 crates/_types/enums/src/by_table/model_weights/weights_file_type.rs create mode 100644 crates/_types/tokens/src/tokens/weight_uploads.rs diff --git a/crates/_types/buckets/src/public/mod.rs b/crates/_types/buckets/src/public/mod.rs index d85de56720..dc35c40ed4 100644 --- a/crates/_types/buckets/src/public/mod.rs +++ b/crates/_types/buckets/src/public/mod.rs @@ -2,4 +2,5 @@ pub mod media_files; pub mod media_uploads; pub mod public_path; pub mod voice_conversion_results; -pub mod zs_voices; \ No newline at end of file +pub mod zs_voices; +pub mod weight_uploads; \ No newline at end of file diff --git a/crates/_types/buckets/src/public/weight_uploads/directory.rs b/crates/_types/buckets/src/public/weight_uploads/directory.rs new file mode 100644 index 0000000000..fcedda62dd --- /dev/null +++ b/crates/_types/buckets/src/public/weight_uploads/directory.rs @@ -0,0 +1,77 @@ +use crockford::crockford_entropy_lower; +// TODO Thoughts NOTE: +// I am wondering if it be easier if we just had a single struct and impl to handle downloads from public and private buckets. +// essentially, the only difference is the prefix and file postfix and the path of where the file goes and the path will be constructed from there. +// then you can tie in the downloader as a single object that can be used for both public and private buckets. + +use crate::public::public_path::PublicPath; +use crate::util::hashed_directory_path_long_string::hashed_directory_path_long_string; + +const WEIGHT_UPLOAD_DIRECTORY : &str = "/weight_upload"; + +/// Directory for user media uploads. +/// Each uploaded file gets its own directory so that we can store the original +/// file alongside re-processed derivatives. +#[derive(Clone)] +pub struct WeightUploadDirectory { + object_hash: String, + directory: String, +} + +impl PublicPath for WeightUploadDirectory {} + +impl WeightUploadDirectory { + + pub fn generate_new() -> Self { + let entropy = crockford_entropy_lower(32); + Self::from_object_hash(&entropy) + } + + pub fn from_object_hash(object_hash: &str) -> Self { + // TODO: Path construction could be cleaner. + let middle = hashed_directory_path_long_string(object_hash); + let directory = format!("{}/{}{}", WEIGHT_UPLOAD_DIRECTORY, middle, object_hash); + Self { + object_hash: object_hash.to_string(), + directory, + } + } + + pub fn get_directory_path_str(&self) -> &str { + &self.directory + } + + pub fn get_object_hash(&self) -> &str { + &self.object_hash + } +} + +#[cfg(test)] +mod tests { + use crate::public::weight_uploads::directory::WeightUploadDirectory; + + #[test] + pub fn generate_new_entropy() { + let directory = WeightUploadDirectory::generate_new(); + assert_eq!(directory.get_object_hash().len(), 32); + } + + #[test] + pub fn get_directory_path_str() { + let directory = WeightUploadDirectory::from_object_hash("abcdefghijk"); + assert_eq!(directory.get_directory_path_str(), "/weight_upload/a/b/c/d/e/abcdefghijk"); + } + + #[test] + pub fn get_directory_path_str_short_name() { + let directory = WeightUploadDirectory::from_object_hash("foo"); + assert_eq!(directory.get_directory_path_str(), "/weight_upload/f/o/foo"); + } + + #[test] + pub fn get_object_hash() { + let hash = "abcdefghijk"; + let directory = WeightUploadDirectory::from_object_hash(hash); + assert_eq!(directory.get_object_hash(), hash); + } +} diff --git a/crates/_types/buckets/src/public/weight_uploads/mod.rs b/crates/_types/buckets/src/public/weight_uploads/mod.rs new file mode 100644 index 0000000000..42ede393cf --- /dev/null +++ b/crates/_types/buckets/src/public/weight_uploads/mod.rs @@ -0,0 +1,2 @@ +pub mod directory; +pub mod original_file; \ No newline at end of file diff --git a/crates/_types/buckets/src/public/weight_uploads/original_file.rs b/crates/_types/buckets/src/public/weight_uploads/original_file.rs new file mode 100644 index 0000000000..7f0a40a6b7 --- /dev/null +++ b/crates/_types/buckets/src/public/weight_uploads/original_file.rs @@ -0,0 +1,104 @@ +use std::path::PathBuf; + +use crockford::crockford_entropy_lower; + +use crate::public::media_uploads::directory::MediaUploadDirectory; +use crate::public::public_path::PublicPath; + +const ORIGINAL_FILE_BASENAME : &str = "original_upload.bin"; + +// TODO: Generate these from a macro. + +/// The original user upload file. +/// It may have derivative files (down samples, crops, etc.) that live alongside it. +#[derive(Clone)] +pub struct MediaUploadOriginalFilePath { + directory: MediaUploadDirectory, + full_object_path: String, +} + +impl PublicPath for MediaUploadOriginalFilePath {} + +impl MediaUploadOriginalFilePath { + + pub fn generate_new() -> Self { + let entropy = crockford_entropy_lower(32); + Self::from_object_hash(&entropy) + } + + pub fn from_object_hash(hash: &str) -> Self { + // TODO: Path construction could be cleaner. + let directory = MediaUploadDirectory::from_object_hash(hash); + let full_object_path = format!("{}/{}", directory.get_directory_path_str(), ORIGINAL_FILE_BASENAME); + Self { + directory, + full_object_path, + } + } + + pub fn get_full_object_path_str(&self) -> &str { + &self.full_object_path + } + + pub fn to_full_object_pathbuf(&self) -> PathBuf { + PathBuf::from(&self.full_object_path) + } + + pub fn get_directory(&self) -> &MediaUploadDirectory { + &self.directory + } + + pub fn get_object_hash(&self) -> &str { + &self.directory.get_object_hash() + } + + pub fn get_basename() -> &'static str { + ORIGINAL_FILE_BASENAME + } +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use crate::public::media_uploads::original_file::MediaUploadOriginalFilePath; + + #[test] + pub fn generate_new_entropy() { + let file = MediaUploadOriginalFilePath::generate_new(); + assert_eq!(file.get_object_hash().len(), 32); + assert_eq!(file.get_directory().get_object_hash().len(), 32); + } + + #[test] + pub fn get_full_object_path_str() { + let file = MediaUploadOriginalFilePath::from_object_hash("abcdefghijk"); + assert_eq!(file.get_full_object_path_str(), "/media_upload/a/b/c/d/e/abcdefghijk/original_upload.bin"); + } + + #[test] + pub fn to_full_object_pathbuf() { + let file = MediaUploadOriginalFilePath::from_object_hash("abcdefghijk"); + assert_eq!(file.to_full_object_pathbuf(), PathBuf::from("/media_upload/a/b/c/d/e/abcdefghijk/original_upload.bin")); + } + + #[test] + pub fn get_full_object_path_str_short_name() { + let file = MediaUploadOriginalFilePath::from_object_hash("foo"); + assert_eq!(file.get_full_object_path_str(), "/media_upload/f/o/foo/original_upload.bin"); + } + + #[test] + pub fn get_full_object_path_str_starts_with_directory() { + let file = MediaUploadOriginalFilePath::from_object_hash("abcdefghijk"); + let full_path = file.get_full_object_path_str(); + assert!(full_path.starts_with(file.get_directory().get_directory_path_str())) + } + + #[test] + pub fn get_object_hash() { + let hash = "abcdefghijk"; + let file = MediaUploadOriginalFilePath::from_object_hash(hash); + assert_eq!(file.get_object_hash(), hash); + } +} diff --git a/crates/_types/enums/src/by_table/model_weights/mod.rs b/crates/_types/enums/src/by_table/model_weights/mod.rs index 5217cf739e..0755131487 100644 --- a/crates/_types/enums/src/by_table/model_weights/mod.rs +++ b/crates/_types/enums/src/by_table/model_weights/mod.rs @@ -1,2 +1,3 @@ pub mod weights_types; pub mod weights_category; +pub mod weights_file_type; \ No newline at end of file diff --git a/crates/_types/enums/src/by_table/model_weights/weights_file_type.rs b/crates/_types/enums/src/by_table/model_weights/weights_file_type.rs new file mode 100644 index 0000000000..a7adeb72cd --- /dev/null +++ b/crates/_types/enums/src/by_table/model_weights/weights_file_type.rs @@ -0,0 +1,68 @@ +use std::collections::BTreeSet; + +#[cfg(test)] +use strum::EnumCount; +#[cfg(test)] +use strum::EnumIter; +// TODO we will need to scan the checkpoints for malicious code. We can't just trust the file extension. +#[cfg_attr(test, derive(EnumIter, EnumCount))] +#[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize)] +pub enum WeightsFileType { + #[serde(rename = "checkpoint")] + Checkpoint, + #[serde(rename = "safetensors")] + SafeTensors, +} + +impl WeightsFileType { + pub fn to_str(&self) -> &'static str { + match self { + Self::Checkpoint => "checkpoint", + Self::SafeTensors => "safetensors", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "checkpoint" => Ok(Self::Checkpoint), + "safetensors" => Ok(Self::SafeTensors), + _ => Err(format!("invalid value: {:?}", value)), + } + } + + pub fn all_variants() -> BTreeSet { + BTreeSet::from([ + Self::Checkpoint, + Self::SafeTensors, + ]) + } +} + +impl_enum_display_and_debug_using_to_str!(WeightsFileType); +impl_mysql_enum_coders!(WeightsFileType); + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_to_str() { + assert_eq!(WeightsFileType::Checkpoint.to_str(), "checkpoint"); + assert_eq!(WeightsFileType::SafeTensors.to_str(), "safetensors"); + } + + #[test] + fn test_from_str() { + assert_eq!(WeightsFileType::from_str("checkpoint").unwrap(), WeightsFileType::Checkpoint); + assert_eq!(WeightsFileType::from_str("safetensors").unwrap(), WeightsFileType::SafeTensors); + assert!(WeightsFileType::from_str("invalid").is_err()); + } + + #[test] + fn test_all_variants() { + let variants = WeightsFileType::all_variants(); + assert_eq!(variants.len(), 2); + assert!(variants.contains(&WeightsFileType::Checkpoint)); + assert!(variants.contains(&WeightsFileType::SafeTensors)); + } +} \ No newline at end of file diff --git a/crates/_types/tokens/src/tokens/weight_uploads.rs b/crates/_types/tokens/src/tokens/weight_uploads.rs new file mode 100644 index 0000000000..cd7c046fc8 --- /dev/null +++ b/crates/_types/tokens/src/tokens/weight_uploads.rs @@ -0,0 +1,14 @@ +use std::fmt::Debug; + +use serde::Deserialize; +use serde::Serialize; + +use crate::prefixes::TokenPrefix; + +/// The primary key for user weight uploads +#[derive(Clone, PartialEq, Eq, sqlx::Type, Debug, Serialize, Deserialize)] +#[sqlx(transparent)] +pub struct WeightUploadToken(pub String); + +impl_string_token!(WeightUploadToken); +impl_crockford_generator!(WeightUploadToken, 32usize, TokenPrefix::WeightUpload, CrockfordLower); \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/create_weight.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/create_weight.rs index e69de29bb2..1db5217ba1 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/create_weight.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/create_weight.rs @@ -0,0 +1,335 @@ +use std::collections::HashSet; +use std::sync::Arc; + +use actix_multipart::Multipart; +use actix_web::{HttpRequest, HttpResponse, web}; +use mysql_queries::queries::model_weights::create_weight::{self, create_weight}; +use once_cell::sync::Lazy; + +use tokens::tokens::media_uploads::MediaUploadToken; + +use crate::http_server::endpoints::media_uploads::common::upload_error::UploadError; +use crate::server_state::ServerState; + +use log::{error, info, warn}; + +use buckets::public::media_uploads::original_file::MediaUploadOriginalFilePath; +use enums::by_table::media_uploads::media_upload_source::MediaUploadSource; +use enums::by_table::media_uploads::media_upload_type::MediaUploadType; +use enums::common::visibility::Visibility; +use hashing::sha256::sha256_hash_bytes::sha256_hash_bytes; +use http_server_common::request::get_request_ip::get_request_ip; +use media::decode_basic_audio_info::decode_basic_audio_bytes_info; +use mimetypes::mimetype_for_bytes::get_mimetype_for_bytes; +use mysql_queries::queries::media_uploads::get_media_upload_by_uuid::get_media_upload_by_uuid_with_connection; +use mysql_queries::queries::media_uploads::insert_media_upload::{Args, insert_media_upload}; +use tokens::tokens::media_uploads::MediaUploadToken; + +use crate::http_server::endpoints::media_uploads::common::drain_multipart_request::{drain_multipart_request, MediaSource}; +use crate::http_server::endpoints::media_uploads::common::upload_error::UploadError; +use crate::server_state::ServerState; +use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format; + +#[derive(Serialize)] +pub struct UploadWeightsSuccessResponse { + pub success: bool, + pub upload_token: MediaUploadToken, +} + +static ALLOWED_MIME_TYPES : Lazy> = Lazy::new(|| { + HashSet::from([ + "binary/octet-stream", + ]) +}); + +pub async fn upload_weights_handler( + http_request: HttpRequest, + server_state: web::Data>, + mut multipart_payload: Multipart, +) -> Result { + + let response = handle_weight_upload( + &http_request, + &server_state, + multipart_payload, + &ALLOWED_MIME_TYPES).await?; + + let weight_upload_token = response.to_media_token(); + + let response: UploadWeightsSuccessResponse = UploadWeightsSuccessResponse { + success: true, + upload_token: weight_upload_token, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| UploadError::ServerError)?; + + return Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)); +} + + +pub enum SuccessCase { + WeightAlreadyUploaded { + existing_upload_token: WeightUploadToken, + }, + WeightSuccessfullyUploaded { + upload_token: WeightUploadToken, + } +} + +impl SuccessCase { + pub fn to_media_token(self) -> WeightUploadToken { + match self { + SuccessCase::WeightAlreadyUploaded { existing_upload_token } => existing_upload_token, + SuccessCase::WeightSuccessfullyUploaded { upload_token } => upload_token, + } + } +} + +pub async fn handle_weight_upload( + http_request: &HttpRequest, + server_state: &web::Data>, + mut multipart_payload: Multipart, + allowed_mimetypes: &HashSet<&'static str>, +) -> Result { + + let mut mysql_connection = server_state.mysql_pool + .acquire() + .await + .map_err(|err| { + error!("MySql pool error: {:?}", err); + UploadError::ServerError + })?; + + // ==================== READ SESSION ==================== // + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session_from_connection(http_request, &mut mysql_connection) + .await + .map_err(|e| { + error!("Session checker error: {:?}", e); + UploadError::ServerError + })?; + + // ==================== RATE LIMIT ==================== // + + let rate_limiter = match maybe_user_session { + None => &server_state.redis_rate_limiters.logged_out, + Some(ref user) => { + if user.is_banned { + return Err(UploadError::NotAuthorized); + } + &server_state.redis_rate_limiters.logged_in + }, + }; + + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { + return Err(UploadError::RateLimited); + } + + if let Err(_err) = server_state.redis_rate_limiters.model_upload.rate_limit_request(&http_request) { + return Err(UploadError::RateLimited); + } + + // ==================== READ MULTIPART REQUEST ==================== // + + let upload_weights_request = drain_multipart_request(multipart_payload) + .await + .map_err(|e| { + // TODO: Error handling could be nicer. + UploadError::BadInput("bad request".to_string()) + })?; + + let uuid_idempotency_token = upload_weights_request.uuid_idempotency_token + .ok_or(UploadError::BadInput("no uuid".to_string()))?; + + let maybe_existing_upload = + get_media_upload_by_uuid_with_connection(&uuid_idempotency_token, &mut mysql_connection) + .await; + + match maybe_existing_upload { + Err(err) => { + error!("Error checking for previous upload: {:?}", err); + return Err(UploadError::ServerError); + } + Ok(Some(upload)) => { + // File already uploaded and frontend didn't protect us. + return Ok(SuccessCase::WeightAlreadyUploaded { + existing_upload_token: upload.token, + }); + } + Ok(None) => { + // Proceed. + } + } + + if let Err(reason) = validate_idempotency_token_format(&uuid_idempotency_token) { + return Err(UploadError::BadInput(reason)); + } + + let creator_set_visibility = maybe_user_session + .as_ref() + .map(|user_session| user_session.preferred_tts_result_visibility) // TODO: We need a new type of visibility control. + .unwrap_or(Visibility::default()); + + let ip_address = get_request_ip(&http_request); + + let maybe_user_token = maybe_user_session + .map(|session| session.get_strongly_typed_user_token()); + + let maybe_file_size_bytes = upload_weights_request.file_bytes + .as_ref() + .map(|bytes| bytes.len()); + + info!("Upload maybe filesize: {:?}", maybe_file_size_bytes); + + let maybe_mimetype = upload_weights_request.file_bytes + .as_ref() + .map(|bytes| get_mimetype_for_bytes(bytes.as_ref())) + .flatten(); + + let bytes = match upload_weights_request.file_bytes { + None => return Err(UploadError::BadInput("missing file contents".to_string())), + Some(bytes) => bytes, + }; + + let hash = sha256_hash_bytes(&bytes) + .map_err(|io_error| { + error!("Problem hashing bytes: {:?}", io_error); + UploadError::ServerError + })?; + + let file_size_bytes = bytes.len(); + + let mut maybe_duration_millis = None; + let mut maybe_codec_name = None; + let mut media_upload_type = None; + + if let Some(mimetype) = maybe_mimetype.as_deref() { + + if !allowed_mimetypes.contains(mimetype) { + // NB: Don't let our error message inject malicious strings + let filtered_mimetype = mimetype + .chars() + .filter(|c| c.is_ascii()) + .filter(|c| c.is_alphanumeric() || *c == '/') + .collect::(); + return Err(UploadError::BadInput(format!("unpermitted mime type: {}", &filtered_mimetype))); + } + + // NB: .aiff (audio/aiff) isn't supported by Symphonia: + // It contains uncompressed PCM-encoded audio similar to wav. + // See: https://github.com/pdeljanov/Symphonia/issues/75 + // NB: The following formats are not supported by Symphonia and + // do not have any open issues filed. They may simply be too old: + // - .wma (audio/x-ms-wma) + // - .avi (video/x-msvideo) + media_upload_type = match mimetype { + // Audio + "audio/aac" /* .aac */ => Some(MediaUploadType::Audio), + _ => None, + }; + + + let media_upload_type = match media_upload_type { + Some(m) => m, + None => { + warn!("Invalid mimetype: {:?}", maybe_mimetype); + return Err(UploadError::BadInput(format!("unknown mimetype: {:?}", maybe_mimetype))) + }, + }; + + // probably don't need this + let media_upload_source = match upload_weights_request.media_source { + MediaSource::Unknown => MediaUploadSource::Unknown, + MediaSource::UserFile => MediaUploadSource::File, + MediaSource::UserDeviceApi => MediaUploadSource::DeviceApi, + }; + + // This is fine for now we just don't expose the path! + let public_upload_path = MediaUploadOriginalFilePath::generate_new(); + + info!("Uploading media to bucket path: {}", public_upload_path.get_full_object_path_str()); + + // need to upload based on visibility + server_state.public_bucket_client.upload_file_with_content_type( + public_upload_path.get_full_object_path_str(), + bytes.as_ref(), + mime_type) + .await + .map_err(|e| { + warn!("Upload media bytes to bucket error: {:?}", e); + UploadError::ServerError + })?; + + let token = create_weight(CreateModelWeightsArgs { + token: &MediaUploadToken::generate(), + weights_type: media_upload_type, + weights_category: media_upload_category, + title: upload_weights_request.file_name.as_deref(), + maybe_thumbnail_token: None, + description_markdown: None, + description_rendered_html: None, + creator_user_token: maybe_user_token.as_deref(), + creator_ip_address: &ip_address, + creator_set_visibility, + maybe_last_update_user_token: None, + original_download_url: None, + original_filename: upload_weights_request.file_name.as_deref(), + file_size_bytes: file_size_bytes as i32, + file_checksum_sha2: hash, + private_bucket_hash: private_bucket_hash, + maybe_private_bucket_prefix: None, + maybe_private_bucket_extension: None, + cached_user_ratings_total_count: 0, + cached_user_ratings_positive_count: 0, + cached_user_ratings_negative_count: 0, + maybe_cached_user_ratings_ratio: None, + version: 0, + mysql_pool: &server_state.mysql_pool, + }) + .await + .map_err(|err| { + warn!("New generic download creation DB error: {:?}", err); + UploadError::ServerError + })?; + }).await?; + + + + let (token, record_id) = insert_media_upload(Args{ + uuid_idempotency_token: &uuid_idempotency_token, + media_type: media_upload_type, + media_source: media_upload_source, + maybe_original_filename: upload_weights_request.file_name.as_deref(), + original_file_size_bytes: file_size_bytes as u64, + maybe_original_duration_millis: maybe_duration_millis, + maybe_original_mime_type: maybe_mimetype, + maybe_original_audio_encoding: maybe_codec_name.as_deref(), + maybe_original_video_encoding: None, + maybe_original_frame_width: None, // TODO + maybe_original_frame_height: None, // TODO + checksum_sha2: &hash, + public_upload_path: &public_upload_path, + maybe_extra_file_modification_info: None, + maybe_creator_user_token: maybe_user_token.as_ref(), + maybe_creator_anonymous_visitor_token: None, + creator_ip_address: &ip_address, + creator_set_visibility, + mysql_pool: &server_state.mysql_pool, + }) + .await + .map_err(|err| { + warn!("New generic download creation DB error: {:?}", err); + UploadError::ServerError + })?; + + info!("new media upload id: {} token: {:?}", record_id, &token); + + Ok(SuccessCase::WeightSuccessfullyUploaded { + upload_token: token, + }) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/mod.rs index 66601e374d..b9e2bd05c2 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/mod.rs @@ -2,4 +2,5 @@ pub mod get_weight; pub mod delete_weight; pub mod update_weight; pub mod list_weights_by_user; -pub mod list_availible_weights; \ No newline at end of file +pub mod list_availible_weights; +pub mod create_weight; \ No newline at end of file