From e190721ece209cd0208fe99f3d1627cc171ecba9 Mon Sep 17 00:00:00 2001 From: Michael Chung Date: Mon, 11 Dec 2023 15:11:28 -0500 Subject: [PATCH] Ok inital commit for server side need to run inference on the specific model once downloaded to verify that the model is compatible with our infrastructure. --- crates/lib/deprecated/config/src/bad_urls.rs | 5 + .../drain_multipart_weights_request.rs | 100 ----- .../endpoints/weights/upload_weight.rs | 385 +++++++----------- 3 files changed, 158 insertions(+), 332 deletions(-) delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/weights/drain_multipart_weights_request.rs diff --git a/crates/lib/deprecated/config/src/bad_urls.rs b/crates/lib/deprecated/config/src/bad_urls.rs index 13be46019b..280ed1d3bb 100644 --- a/crates/lib/deprecated/config/src/bad_urls.rs +++ b/crates/lib/deprecated/config/src/bad_urls.rs @@ -3,11 +3,16 @@ use crate::is_bad_download_url::is_bad_download_url; /// Reports if the URL is inappropriate for downloading as a TTS model, vocoder, etc. +/// DEPRECATE: pub fn is_bad_tts_model_download_url(url: &str) -> anyhow::Result { // TODO: Just use this. is_bad_download_url(url) } +pub fn is_bad_model_weights_download_url(url: &str) -> anyhow::Result { + is_bad_download_url(url) +} + #[cfg(test)] mod tests { use crate::bad_urls::is_bad_tts_model_download_url; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/drain_multipart_weights_request.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/drain_multipart_weights_request.rs deleted file mode 100644 index c87a9d82c5..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/drain_multipart_weights_request.rs +++ /dev/null @@ -1,100 +0,0 @@ -use actix_multipart::Multipart; -use actix_web::web::BytesMut; -use futures::TryStreamExt; -use log::warn; - -use errors::AnyhowResult; -use tokens::tokens::weights::ModelWeightToken; - -use crate::http_server::web_utils::read_multipart_field_bytes::{checked_read_multipart_bytes, read_multipart_field_as_text}; - -pub struct WeightUploadRequest { - pub uuid_idempotency_token: Option, - pub model_weights_token: Option, - pub file_name: Option, - pub file_bytes: Option, - pub media_source: MediaSource, -} - -/// Where the frontend tells us the file came from. -#[derive(Eq, PartialEq, Clone, Copy)] -pub enum MediaSource { - Unknown, - UserFile, - /// Eg. the web audio API in Javascript. - UserDeviceApi, -} - -/// Pull common parts out of multipart media HTTP requests, typically for handling file uploads. -pub async fn drain_multipart_request(mut multipart_payload: Multipart) -> AnyhowResult { - let mut uuid_idempotency_token = None; - let mut model_weight_token = None; - let mut file_bytes = None; - let mut file_name = None; - let mut media_source = None; - - while let Ok(Some(mut field)) = multipart_payload.try_next().await { - let mut field_name = None; - let mut field_filename = None; - - if let content_disposition = field.content_disposition() { - field_name = content_disposition.get_name() - .map(|s| s.to_string()); - field_filename = content_disposition.get_filename() // NB: Only used for the file bytes. - .map(|s| s.to_string()); - } - - match field_name.as_deref() { - Some("uuid_idempotency_token") => { - uuid_idempotency_token = read_multipart_field_as_text(&mut field).await - .map_err(|e| { - warn!("Error reading idempotency token: {:}", e); - e - })?; - }, - Some("dataset_token") => { - model_weight_token = read_multipart_field_as_text(&mut field).await - .map_err(|e| { - warn!("Error reading dataset token: {:}", e); - e - })?; - }, - Some("file") => { - file_name = field_filename.clone(); - file_bytes = checked_read_multipart_bytes(&mut field).await - .map_err(|e| { - warn!("Error reading audio upload: {:}", e); - e - })?; - }, - Some("source") => { - media_source = read_multipart_field_as_text(&mut field).await - .map_err(|e| { - warn!("Error reading source: {:}", e); - e - })?; - }, - _ => continue, - } - } - - let media_source = match media_source.as_deref() { - Some("device") => MediaSource::UserDeviceApi, - Some("file") => MediaSource::UserFile, - _ => { - if file_name.as_deref() == Some("blob") { - MediaSource::UserDeviceApi - } else { - MediaSource::Unknown - } - }, - }; - - Ok(UploadSampleRequest { - uuid_idempotency_token, - dataset_token, - file_name, - file_bytes, - media_source, - }) -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/upload_weight.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/upload_weight.rs index 2d66aa7684..8d2aa5fc5c 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/upload_weight.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/upload_weight.rs @@ -2,10 +2,10 @@ use std::collections::HashSet; use std::sync::Arc; use actix_multipart::Multipart; -use actix_web::{HttpRequest, HttpResponse, web}; +use actix_web::{ HttpRequest, HttpResponse, web }; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; -use mysql_queries::queries::model_weights::create_weight::{self, create_weight}; +use mysql_queries::queries::model_weights::create_weight::{ self, create_weight }; use once_cell::sync::Lazy; use tokens::tokens::model_weights::ModelWeightToken; @@ -13,259 +13,180 @@ use tokens::tokens::model_weights::ModelWeightToken; use crate::http_server::endpoints::media_uploads::common::upload_error::UploadError; use crate::server_state::ServerState; -use log::{error, info, warn}; +use log::{ error, info, warn }; use buckets::public::weight_uploads::original_file::WeightUploadOriginalFilePath; 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 mysql_queries::queries::media_uploads::insert_media_upload::{ Args, insert_media_upload }; -use crate::http_server::endpoints::media_uploads::common::drain_multipart_request::{drain_multipart_request, MediaSource}; +use crate::http_server::endpoints::media_uploads::common::drain_multipart_request::{ + drain_multipart_request, + MediaSource, +}; use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format; use utoipa::ToSchema; +use config::bad_urls::is_bad_model_weights_download_url; + +use enums::by_table::model_weights::{ + weights_category::WeightsCategory, + weights_types::WeightsType, +}; +use enums::common::visibility::Visibility; +use utoipa::ToSchema; + +#[derive(Deserialize, ToSchema)] +pub struct UploadModelWeightsRequest { + idempotency_token: String, + title: String, + download_url: String, + model_weights_type: WeightsType, + model_weights_category: WeightsCategory, + creator_set_visibility: Option, +} + +/// Tell the frontend how to deal with the download queue. +#[derive(Serialize, ToSchema)] +pub enum DownloadJobType { + /// Modern shared download type (everything will use this in the + /// future, and this endpoint will die.) + #[serde(rename = "generic")] + Generic, +} + +#[derive(Serialize, ToSchema)] +pub struct UploadModelWeightsSuccessResponse { + pub success: bool, + /// This is how frontend clients can request the job execution status. + pub job_token: String, + /// This is a transitional field to tell the frontend how to process the job checking. + pub job_type: DownloadJobType, +} + +#[derive(Debug)] +pub enum UploadModelWeightsError { + BadInput(String), + MustBeLoggedIn, + ServerError, + RateLimited, +} + +impl ResponseError for UploadModelWeightsError { + fn status_code(&self) -> StatusCode { + match *self { + UploadModelWeightsError::BadInput(_) => StatusCode::BAD_REQUEST, + UploadModelWeightsError::MustBeLoggedIn => StatusCode::UNAUTHORIZED, + UploadModelWeightsError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + UploadModelWeightsError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + UploadModelWeightsError::BadInput(reason) => reason.to_string(), + UploadModelWeightsError::MustBeLoggedIn => "user must be logged in".to_string(), + UploadModelWeightsError::ServerError => "server error".to_string(), + UploadModelWeightsError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + 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?; - - // send a body with the required payload along side the multipart upload - - let weight_upload_token = response.to_model_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: ModelWeightToken, - }, - WeightSuccessfullyUploaded { - upload_token: ModelWeightToken, + mut multipart_payload: Multipart +) -> Result { + // Rate limiter check + if + let Err(_err) = server_state.redis_rate_limiters.model_upload.rate_limit_request( + &http_request + ) + { + return Err(UploadModelWeightsError::RateLimited); } - } - - impl SuccessCase { - pub fn to_model_token(self) -> ModelWeightToken { - 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 + + let maybe_user_session = server_state.session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool).await .map_err(|e| { - error!("Session checker error: {:?}", e); - UploadError::ServerError + warn!("Session checker error: {:?}", e); + UploadModelWeightsError::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); + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(UploadModelWeightsError::MustBeLoggedIn); } - &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(reason) = validate_idempotency_token(&request.idempotency_token) { + return Err(UploadModelWeightsError::BadInput(reason)); } - - if let Err(_err) = server_state.redis_rate_limiters.model_upload.rate_limit_request(&http_request) { - return Err(UploadError::RateLimited); + + if let Err(reason) = validate_model_title(&request.title) { + return Err(UploadModelWeightsError::BadInput(reason)); } - - // ==================== 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. - } + + match is_bad_model_weights_download_url(&download_url) { + Ok(false) => {} // Ok case + Ok(true) => { + return Err(UploadModelWeightsError::BadInput("Bad model download URL".to_string())); + } + Err(err) => { + warn!("Error parsing url: {:?}", err); + return Err(UploadModelWeightsError::BadInput("Bad model download URL".to_string())); + } } - - 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 supported_model_weights_type = request.model_weights_type; + + let job_token; + let download_job_type; + + let (download_job_token, _record_id) = insert_generic_download_job( + InsertGenericDownloadJobArgs { + uuid_idempotency_token: &uuid, + download_type: GenericDownloadType::Vits, + download_url: &download_url, + title: &title, + creator_user_token: &user_session.user_token, + creator_ip_address: &ip_address, + creator_set_visibility: visibility, // TODO: Creator set preference + mysql_pool: &server_state.mysql_pool, + } + ).await.map_err(|err| { + warn!("New generic download creation DB error (for generic model): {:?}", err); + UploadModelWeightsError::ServerError + })?; + + job_token = download_job_token.to_string(); + download_job_type = DownloadJobType::Generic; + + let response = UploadModelWeightsSuccessResponse { + success: true, + job_token: job_token.to_string(), + job_type: download_job_type, }; - - 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))); - } - media_upload_type = match mimetype { - "binary/octet-stream" /* .bin */ => Some(MediaUploadType::Binary), - _ => None, - }; + + let body = serde_json::to_string(&response).map_err(|_e| UploadModelWeightsError::ServerError)?; + + Ok(HttpResponse::Ok().content_type("application/json").body(body)) +} + +fn validate_idempotency_token(token: &str) -> Result<(), String> { + if token.len() != 36 { + return Err("idempotency token should be 36 characters".to_string()); } - - // This is fine for now we just don't expose the path! - let public_upload_path = WeightUploadOriginalFilePath::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: &ModelWeightToken::generate(), - weights_type: WeightsType::LoRA, // to do parse multi part for this - weights_category: WeightsCategory::ImageGeneration, - title: upload_weights_request.file_name.as_deref(), - maybe_thumbnail_token: &ThumbnailToken::generate(), - description_markdown: "", - description_rendered_html: "", - creator_user_token: maybe_user_token.as_str(), - creator_ip_address: &ip_address, - creator_set_visibility, - maybe_last_update_user_token: "", - original_download_url: "", - 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(|e| { - error!("Error creating weights upload: {:?}", e); - UploadError::ServerError - })?; - - info!("new weights upload id: {} token: {:?}", record_id, &token); - - Ok(SuccessCase::WeightSuccessfullyUploaded { - upload_token: token, - }) - } - \ No newline at end of file + + Ok(()) +}