diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/create_beta_keys_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/create_beta_keys_handler.rs index 9f967f08ad..e1ab2a9d02 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/create_beta_keys_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/create_beta_keys_handler.rs @@ -27,7 +27,7 @@ use user_input_common::check_for_slurs::contains_slurs; use user_input_common::markdown_to_html::markdown_to_html; use crate::http_server::endpoints::moderation::user_feature_flags::edit_user_feature_flags_handler::EditUserFeatureFlagsError; -use crate::http_server::web_utils::require_moderator::{require_moderator, RequireModeratorError}; +use crate::http_server::web_utils::require_moderator::{require_moderator, RequireModeratorError, UseDatabase}; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; @@ -108,7 +108,7 @@ pub async fn create_beta_keys_handler( server_state: web::Data>, ) -> Result { - let user_session = require_moderator(&http_request, &server_state) + let user_session = require_moderator(&http_request, &server_state, UseDatabase::Implicit) .await .map_err(|err| match err { RequireModeratorError::ServerError => CreateBetaKeysError::ServerError, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_pmx/upload_pmx_media_file_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_pmx/upload_pmx_media_file_handler.rs index 1d30d4eb44..199b1a9c0e 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_pmx/upload_pmx_media_file_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_pmx/upload_pmx_media_file_handler.rs @@ -27,10 +27,12 @@ use mysql_queries::queries::media_files::create::insert_media_file_from_file_upl use tokens::tokens::media_files::MediaFileToken; use videos::get_mp4_info::{get_mp4_info, get_mp4_info_for_bytes, get_mp4_info_for_bytes_and_len}; +use crate::http_server::endpoints::beta_keys::create_beta_keys_handler::CreateBetaKeysError; use crate::http_server::endpoints::media_files::upload::upload_engine_asset::drain_multipart_request::drain_multipart_request; use crate::http_server::endpoints::media_files::upload::upload_error::MediaFileUploadError; use crate::http_server::endpoints::media_files::upload::upload_new_scene_media_file_handler::UploadNewSceneMediaFileForm; use crate::http_server::endpoints::media_files::upload::upload_pmx::extract_and_upload_pmx_files::{extract_and_upload_pmx_files, PmxError}; +use crate::http_server::web_utils::require_moderator::{require_moderator, RequireModeratorError, UseDatabase}; use crate::server_state::ServerState; use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format; @@ -122,37 +124,47 @@ pub async fn upload_pmx_media_file_handler( // ==================== READ SESSION ==================== // - let maybe_user_session = server_state - .session_checker - .maybe_get_user_session_from_connection(&http_request, &mut mysql_connection) + // NB: We require a moderator to upload PMX files. + let user_session = require_moderator(&http_request, &server_state, UseDatabase::FromPool(&mut mysql_connection)) .await - .map_err(|e| { - error!("Session checker error: {:?}", e); - MediaFileUploadError::ServerError + .map_err(|err| match err { + RequireModeratorError::ServerError => MediaFileUploadError::ServerError, + RequireModeratorError::NotAuthorized => MediaFileUploadError::NotAuthorized, })?; let maybe_avt_token = server_state .avt_cookie_manager .get_avt_token_from_request(&http_request); - // ==================== BANNED USERS ==================== // +// // ==================== 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); +// MediaFileUploadError::ServerError +// })?; +// +// // ==================== BANNED USERS ==================== // +// +// if let Some(ref user) = maybe_user_session { +// if user.is_banned { +// return Err(MediaFileUploadError::NotAuthorized); +// } +// } - if let Some(ref user) = maybe_user_session { - if user.is_banned { - return Err(MediaFileUploadError::NotAuthorized); - } - } - - // ==================== RATE LIMIT ==================== // - - let rate_limiter = match maybe_user_session { - None => &server_state.redis_rate_limiters.file_upload_logged_out, - Some(ref _session) => &server_state.redis_rate_limiters.file_upload_logged_in, - }; - - if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { - return Err(MediaFileUploadError::RateLimited); - } +// // ==================== RATE LIMIT ==================== // +// +// let rate_limiter = match maybe_user_session { +// None => &server_state.redis_rate_limiters.file_upload_logged_out, +// Some(ref _session) => &server_state.redis_rate_limiters.file_upload_logged_in, +// }; +// +// if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { +// return Err(MediaFileUploadError::RateLimited); +// } // ==================== HANDLE IDEMPOTENCY ==================== // @@ -192,9 +204,10 @@ pub async fn upload_pmx_media_file_handler( let creator_set_visibility = form.maybe_visibility .map(|visibility| visibility.0) .or_else(|| { - maybe_user_session - .as_ref() - .map(|user_session| user_session.preferred_tts_result_visibility) + //maybe_user_session + // .as_ref() + // .map(|user_session| user_session.preferred_tts_result_visibility) + Some(user_session.preferred_tts_result_visibility) }) .unwrap_or(Visibility::default()); @@ -202,8 +215,8 @@ pub async fn upload_pmx_media_file_handler( 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_user_token = maybe_user_session + // .map(|session| session.get_strongly_typed_user_token()); // ==================== FILE DATA ==================== // @@ -254,7 +267,7 @@ pub async fn upload_pmx_media_file_handler( let (token, record_id) = insert_media_file_from_file_upload(InsertMediaFileFromUploadArgs { maybe_media_class: Some(MediaFileClass::Dimensional), media_file_type: MediaFileType::Pmx, - maybe_creator_user_token: maybe_user_token.as_ref(), + maybe_creator_user_token: Some(&user_session.get_strongly_typed_user_token()), maybe_creator_anonymous_visitor_token: maybe_avt_token.as_ref(), creator_ip_address: &ip_address, creator_set_visibility, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/moderation/info/moderator_token_info_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/moderation/info/moderator_token_info_handler.rs index 384517a670..a7de94f4a9 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/moderation/info/moderator_token_info_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/moderation/info/moderator_token_info_handler.rs @@ -26,7 +26,7 @@ use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::prompts::PromptToken; use tokens::tokens::users::UserToken; -use crate::http_server::web_utils::require_moderator::{require_moderator, RequireModeratorError}; +use crate::http_server::web_utils::require_moderator::{require_moderator, RequireModeratorError, UseDatabase}; use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; use crate::server_state::ServerState; @@ -87,7 +87,7 @@ pub async fn moderator_get_token_info_handler( server_state: web::Data> ) -> Result { - let user_session = require_moderator(&http_request, &server_state) + let user_session = require_moderator(&http_request, &server_state, UseDatabase::Implicit) .await .map_err(|err| match err { RequireModeratorError::ServerError => ModeratorTokenInfoError::ServerError, diff --git a/crates/service/web/storyteller_web/src/http_server/web_utils/require_moderator.rs b/crates/service/web/storyteller_web/src/http_server/web_utils/require_moderator.rs index 7dffe18b2c..b25247a6df 100644 --- a/crates/service/web/storyteller_web/src/http_server/web_utils/require_moderator.rs +++ b/crates/service/web/storyteller_web/src/http_server/web_utils/require_moderator.rs @@ -3,6 +3,9 @@ use std::fmt::{Display, Formatter}; use actix_web::HttpRequest; use log::warn; +use r2d2_redis::r2d2::PooledConnection; +use sqlx::MySql; +use sqlx::pool::PoolConnection; use mysql_queries::queries::users::user_sessions::get_user_session_by_token::SessionUserRecord; @@ -26,22 +29,40 @@ impl Display for RequireModeratorError { impl Error for RequireModeratorError {} +pub enum UseDatabase<'a> { + Implicit, + FromPool(&'a mut PoolConnection), +} + pub async fn require_moderator( http_request: &HttpRequest, server_state: &ServerState, + database: UseDatabase<'_>, ) -> Result { + // NB: Save a reference to a connection we open in a branch until the function ends. + let mut saved_connection = None; + + let mysql_connection = match database { + UseDatabase::Implicit => { + let mut connection = server_state.mysql_pool + .acquire() + .await + .map_err(|err| { + warn!("MySql pool error: {:?}", err); + RequireModeratorError::ServerError + })?; + + saved_connection = Some(connection); + + saved_connection.as_mut().expect("this should be safe - we just saved the connection") + }, + UseDatabase::FromPool(pool) => pool, + }; - let mut mysql_connection = server_state.mysql_pool - .acquire() - .await - .map_err(|err| { - warn!("MySql pool error: {:?}", err); - RequireModeratorError::ServerError - })?; let maybe_user_session = server_state .session_checker - .maybe_get_user_session_from_connection(&http_request, &mut mysql_connection) + .maybe_get_user_session_from_connection(&http_request, mysql_connection) .await .map_err(|e| { warn!("Session checker error: {:?}", e);