From 037b2017c5ef02dad71e3ddb48bf220a74e51bb4 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 6 Mar 2024 03:00:27 -0500 Subject: [PATCH] feature flag integrations --- .../media_files/list/list_media_files.rs | 5 +- .../endpoints/workflows/enqueue_comfy_ui.rs | 8 +-- .../service/web/storyteller_web/src/main.rs | 1 + .../web/storyteller_web/src/server_state.rs | 5 ++ .../src/util/allowed_explore_media_access.rs | 52 ++++++++++++++++++ .../allowed_video_style_transfer_access.rs | 55 +++++++++++++++++++ .../web/storyteller_web/src/util/mod.rs | 4 +- 7 files changed, 123 insertions(+), 7 deletions(-) create mode 100644 crates/service/web/storyteller_web/src/util/allowed_explore_media_access.rs create mode 100644 crates/service/web/storyteller_web/src/util/allowed_video_style_transfer_access.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list/list_media_files.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list/list_media_files.rs index 833b7d26ae..ccd87d5a3d 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list/list_media_files.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list/list_media_files.rs @@ -24,6 +24,7 @@ use crate::http_server::common_responses::pagination_cursors::PaginationCursors; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; +use crate::util::allowed_explore_media_access::allowed_explore_media_access; use crate::util::allowed_studio_access::allowed_studio_access; #[derive(Deserialize, ToSchema, IntoParams)] @@ -136,8 +137,8 @@ pub async fn list_media_files_handler( // ==================== FEATURE FLAG CHECK ==================== // - if !allowed_studio_access(maybe_user_session.as_ref(), &server_state.flags) { - warn!("Storyteller Studio access is not permitted for user"); + if !allowed_explore_media_access(maybe_user_session.as_ref()) { + warn!("Explore media access is not permitted for user"); return Err(ListMediaFilesError::NotAuthorized); } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs index 3457b9cf63..3eb1bf8099 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs @@ -27,10 +27,10 @@ use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; -use crate::configs::plans::plan_category::{PlanCategory}; +use crate::configs::plans::plan_category::PlanCategory; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; -use crate::util::allowed_studio_access::allowed_studio_access; +use crate::util::allowed_video_style_transfer_access::allowed_video_style_transfer_access; use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format; /// Debug requests can get routed to special "debug-only" workers, which can @@ -144,8 +144,8 @@ pub async fn enqueue_comfy_ui_handler( // ==================== FEATURE FLAG CHECK ==================== // - if !allowed_studio_access(maybe_user_session.as_ref(), &server_state.flags) { - warn!("Storyteller Studio access is not permitted for user"); + if !allowed_video_style_transfer_access(maybe_user_session.as_ref(), &server_state.flags) { + warn!("Video style transfer access is not permitted for user"); return Err(EnqueueComfyError::NotAuthorized); } diff --git a/crates/service/web/storyteller_web/src/main.rs b/crates/service/web/storyteller_web/src/main.rs index 88597e2f3e..2c5a809d66 100644 --- a/crates/service/web/storyteller_web/src/main.rs +++ b/crates/service/web/storyteller_web/src/main.rs @@ -361,6 +361,7 @@ async fn main() -> AnyhowResult<()> { switch_voice_conversion_to_model_weights: easyenv::get_env_bool_or_default("FF_SWITCH_VOICE_CONVERSION_TO_MODEL_WEIGHTS", false), switch_tts_to_model_weights: easyenv::get_env_bool_or_default("FF_SWITCH_TTS_TO_MODEL_WEIGHTS", false), force_session_studio_flags: easyenv::get_env_bool_or_default("FF_FORCE_SESSION_STUDIO_FLAG", false), + force_session_video_style_transfer_flags: easyenv::get_env_bool_or_default("FF_FORCE_SESSION_VST_FLAG", false), }; let third_party_url_redirector = ThirdPartyUrlRedirector::new(server_environment); diff --git a/crates/service/web/storyteller_web/src/server_state.rs b/crates/service/web/storyteller_web/src/server_state.rs index d4f7ed9249..e1c59c252c 100644 --- a/crates/service/web/storyteller_web/src/server_state.rs +++ b/crates/service/web/storyteller_web/src/server_state.rs @@ -273,6 +273,11 @@ pub struct StaticFeatureFlags { /// TEMPORARY: Require users to have a studio flag on their `users` table record to use studio features. /// This will also prevent anonymous users from using the service APIs marked as studio-only. pub force_session_studio_flags: bool, + + // TODO(2024-03-05): Remove when done. + /// TEMPORARY: Require users to have a flag on their `users` table record to use video style transfer. + /// This will also prevent anonymous users from using the service APIs marked as VST-only. + pub force_session_video_style_transfer_flags: bool, } /// Instead of top level service denial, these are bans against entities that instead return diff --git a/crates/service/web/storyteller_web/src/util/allowed_explore_media_access.rs b/crates/service/web/storyteller_web/src/util/allowed_explore_media_access.rs new file mode 100644 index 0000000000..b9d801edc4 --- /dev/null +++ b/crates/service/web/storyteller_web/src/util/allowed_explore_media_access.rs @@ -0,0 +1,52 @@ +use enums::by_table::users::user_feature_flag::UserFeatureFlag; +use mysql_queries::queries::users::user_sessions::get_user_session_by_token::SessionUserRecord; +use users_component::utils::user_session_extended::UserSessionExtended; +use users_component::utils::user_session_feature_flags::UserSessionFeatureFlags; + +use crate::server_state::StaticFeatureFlags; + +/// Check whether we should allow the request access to explore media. +pub fn allowed_explore_media_access(maybe_session: Option) -> bool { + // NB: We're more careful here since this is sensitive until we have NLP checks. + maybe_session + .map(|ref session| session.can_access_explore_media()) + .unwrap_or(false) +} + +pub trait UserSessionExploreMediaFlag { + fn can_access_explore_media(&self) -> bool; +} + +impl UserSessionExploreMediaFlag for UserSessionExtended { + fn can_access_explore_media(&self) -> bool { + self.feature_flags.has_permission_unoptimized(UserFeatureFlag::ExploreMedia) + } +} + +impl UserSessionExploreMediaFlag for &UserSessionExtended { + fn can_access_explore_media(&self) -> bool { + self.feature_flags.has_permission_unoptimized(UserFeatureFlag::ExploreMedia) + } +} + +impl UserSessionExploreMediaFlag for SessionUserRecord { + fn can_access_explore_media(&self) -> bool { + user_session_has_feature(self) + } +} + +impl UserSessionExploreMediaFlag for &SessionUserRecord { + fn can_access_explore_media(&self) -> bool { + user_session_has_feature(self) + } +} + +fn user_session_has_feature(user_session: &SessionUserRecord) -> bool { + // TODO(bt, 2024-03-05): this is horrible. + // There should be a wrapper class between the query and the caller. + let flags = UserSessionFeatureFlags { + maybe_feature_flags: user_session.maybe_feature_flags.clone(), + }; + + flags.has_permission_unoptimized(UserFeatureFlag::ExploreMedia) +} diff --git a/crates/service/web/storyteller_web/src/util/allowed_video_style_transfer_access.rs b/crates/service/web/storyteller_web/src/util/allowed_video_style_transfer_access.rs new file mode 100644 index 0000000000..2203d9a327 --- /dev/null +++ b/crates/service/web/storyteller_web/src/util/allowed_video_style_transfer_access.rs @@ -0,0 +1,55 @@ +use enums::by_table::users::user_feature_flag::UserFeatureFlag; +use mysql_queries::queries::users::user_sessions::get_user_session_by_token::SessionUserRecord; +use users_component::utils::user_session_extended::UserSessionExtended; +use users_component::utils::user_session_feature_flags::UserSessionFeatureFlags; + +use crate::server_state::StaticFeatureFlags; + +/// Check whether we should allow the request access to video style transfer +pub fn allowed_video_style_transfer_access(maybe_session: Option, flags: &StaticFeatureFlags) -> bool { + if !flags.force_session_video_style_transfer_flags { + return true; + } + + maybe_session + .map(|ref session| session.can_access_video_style_transfer()) + .unwrap_or(false) +} + +pub trait UserSessionVideoStyleTransferFlag { + fn can_access_video_style_transfer(&self) -> bool; +} + +impl UserSessionVideoStyleTransferFlag for UserSessionExtended { + fn can_access_video_style_transfer(&self) -> bool { + self.feature_flags.has_permission_unoptimized(UserFeatureFlag::VideoStyleTransfer) + } +} + +impl UserSessionVideoStyleTransferFlag for &UserSessionExtended { + fn can_access_video_style_transfer(&self) -> bool { + self.feature_flags.has_permission_unoptimized(UserFeatureFlag::VideoStyleTransfer) + } +} + +impl UserSessionVideoStyleTransferFlag for SessionUserRecord { + fn can_access_video_style_transfer(&self) -> bool { + user_session_has_feature(self) + } +} + +impl UserSessionVideoStyleTransferFlag for &SessionUserRecord { + fn can_access_video_style_transfer(&self) -> bool { + user_session_has_feature(self) + } +} + +fn user_session_has_feature(user_session: &SessionUserRecord) -> bool { + // TODO(bt, 2024-03-05): this is horrible. + // There should be a wrapper class between the query and the caller. + let flags = UserSessionFeatureFlags { + maybe_feature_flags: user_session.maybe_feature_flags.clone(), + }; + + flags.has_permission_unoptimized(UserFeatureFlag::VideoStyleTransfer) +} diff --git a/crates/service/web/storyteller_web/src/util/mod.rs b/crates/service/web/storyteller_web/src/util/mod.rs index 13b83e2474..1b4f55b6ed 100644 --- a/crates/service/web/storyteller_web/src/util/mod.rs +++ b/crates/service/web/storyteller_web/src/util/mod.rs @@ -1,5 +1,7 @@ +pub mod allowed_explore_media_access; pub mod allowed_studio_access; +pub mod allowed_video_style_transfer_access; pub mod delete_role_disambiguation; pub mod encrypted_sort_id; pub mod placeholder_images; -pub mod troll_user_bans; +pub mod troll_user_bans; \ No newline at end of file