From 871805a32c7ea81eadfc477dd8668941c71da892 Mon Sep 17 00:00:00 2001 From: "Zhao.wei" Date: Thu, 3 Apr 2025 23:47:18 +0800 Subject: [PATCH] feat: impl ratelimit feature (#6072) --- src/common/infra/config.rs | 3 + src/common/infra/ofga/mod.rs | 35 +- src/common/utils/auth.rs | 8 +- src/common/utils/auth_tests.rs | 4 + src/config/src/config.rs | 41 +- src/config/src/lib.rs | 1 + src/config/src/meta/mod.rs | 1 + src/config/src/meta/ratelimit.rs | 160 ++ src/config/src/router.rs | 92 + src/handler/http/request/actions/action.rs | 810 ++++---- .../http/request/actions/operations.rs | 26 +- src/handler/http/request/alerts/deprecated.rs | 16 + .../http/request/alerts/destinations.rs | 10 + src/handler/http/request/alerts/mod.rs | 16 + src/handler/http/request/alerts/templates.rs | 10 + src/handler/http/request/authz/fga.rs | 387 +++- src/handler/http/request/clusters/mod.rs | 2 + src/handler/http/request/dashboards/mod.rs | 15 + .../http/request/dashboards/reports.rs | 14 + .../request/dashboards/timed_annotations.rs | 10 + .../http/request/enrichment_table/mod.rs | 2 + src/handler/http/request/folders.rs | 24 + src/handler/http/request/functions/mod.rs | 10 + src/handler/http/request/keys/mod.rs | 60 +- src/handler/http/request/kv/mod.rs | 8 + src/handler/http/request/mod.rs | 3 +- src/handler/http/request/organization/org.rs | 16 + .../http/request/organization/settings.rs | 4 + src/handler/http/request/pipeline.rs | 12 + src/handler/http/request/promql/mod.rs | 16 + src/handler/http/request/ratelimit/mod.rs | 517 ++++++ src/handler/http/request/search/mod.rs | 15 +- .../http/request/search/query_manager.rs | 110 +- src/handler/http/request/search/saved_view.rs | 35 +- src/handler/http/request/search/search_job.rs | 748 +++++--- src/handler/http/request/search/utils.rs | 19 +- .../http/request/service_accounts/mod.rs | 10 + src/handler/http/request/short_url/mod.rs | 4 + src/handler/http/request/stream/mod.rs | 15 + src/handler/http/request/syslog/mod.rs | 8 + src/handler/http/request/traces/mod.rs | 2 + src/handler/http/request/users/mod.rs | 16 +- src/handler/http/router/mod.rs | 4 + src/handler/http/router/openapi.rs | 99 +- src/infra/src/table/entity/mod.rs | 1 + .../src/table/entity/rate_limit_rules.rs | 23 + .../m20250217_115548_ratelimit_table.rs | 136 ++ src/infra/src/table/migration/mod.rs | 2 + src/infra/src/table/mod.rs | 1 + src/infra/src/table/ratelimit.rs | 779 ++++++++ src/job/mod.rs | 28 +- src/main.rs | 25 +- src/router/http/mod.rs | 79 +- src/router/mod.rs | 2 + src/router/ratelimit/mod.rs | 16 + src/router/ratelimit/resource_extractor.rs | 615 ++++++ src/service/mod.rs | 2 + src/service/ratelimit/mod.rs | 16 + src/service/ratelimit/rule.rs | 162 ++ src/service/search/cluster/cache_multi.rs | 2 +- src/service/search/cluster/cacher.rs | 4 +- src/service/search/grpc_search.rs | 8 +- src/service/users.rs | 56 + src/super_cluster_queue/mod.rs | 2 + src/super_cluster_queue/ratelimit.rs | 83 + web/src/assets/images/common/hard-drive.svg | 11 + .../images/common/selectOrganization.svg | 39 + web/src/components/iam/quota/Quota.vue | 1641 +++++++++++++++++ .../components/iam/roles/PermissionsTable.vue | 20 +- .../shared/grid/NoOrganizationSelected.vue | 57 + .../composables/shared/useEnterpriseRoutes.ts | 10 + web/src/composables/useIsMetaOrg.ts | 2 +- web/src/composables/useRateLimiter.ts | 139 ++ web/src/locales/languages/en.json | 23 +- web/src/services/rate_limit.ts | 61 + web/src/stores/index.ts | 21 + web/src/views/IdentityAccessManagement.vue | 32 +- 77 files changed, 6671 insertions(+), 845 deletions(-) create mode 100644 src/config/src/meta/ratelimit.rs create mode 100644 src/config/src/router.rs create mode 100644 src/handler/http/request/ratelimit/mod.rs create mode 100644 src/infra/src/table/entity/rate_limit_rules.rs create mode 100644 src/infra/src/table/migration/m20250217_115548_ratelimit_table.rs create mode 100644 src/infra/src/table/ratelimit.rs create mode 100644 src/router/ratelimit/mod.rs create mode 100644 src/router/ratelimit/resource_extractor.rs create mode 100644 src/service/ratelimit/mod.rs create mode 100644 src/service/ratelimit/rule.rs create mode 100644 src/super_cluster_queue/ratelimit.rs create mode 100644 web/src/assets/images/common/hard-drive.svg create mode 100644 web/src/assets/images/common/selectOrganization.svg create mode 100644 web/src/components/iam/quota/Quota.vue create mode 100644 web/src/components/shared/grid/NoOrganizationSelected.vue create mode 100644 web/src/composables/useRateLimiter.ts create mode 100644 web/src/services/rate_limit.ts diff --git a/src/common/infra/config.rs b/src/common/infra/config.rs index 448879e05a..7a2477695c 100644 --- a/src/common/infra/config.rs +++ b/src/common/infra/config.rs @@ -24,6 +24,7 @@ use config::{ folder::Folder, function::Transform, promql::ClusterLeader, + ratelimit::CachedUserRoles, stream::StreamParams, }, }; @@ -96,5 +97,7 @@ pub static USER_SESSIONS: Lazy> = Lazy::new(Default::d pub static SHORT_URLS: Lazy> = Lazy::new(DashMap::default); // TODO: Implement rate limiting for maximum number of sessions pub static WS_SESSIONS: Lazy> = Lazy::new(DashMap::default); +pub static USER_ROLES_CACHE: Lazy> = + Lazy::new(Default::default); // Global registry for search requests by `trace_id` pub static WS_SEARCH_REGISTRY: Lazy> = Lazy::new(DashMap::new); diff --git a/src/common/infra/ofga/mod.rs b/src/common/infra/ofga/mod.rs index 50df11d954..ccea4cb1c4 100644 --- a/src/common/infra/ofga/mod.rs +++ b/src/common/infra/ofga/mod.rs @@ -39,7 +39,7 @@ use crate::{ }; pub async fn init() -> Result<(), anyhow::Error> { - use o2_openfga::{authorizer::authz::get_tuple_for_new_index, get_all_init_tuples}; + use o2_openfga::get_all_init_tuples; let mut init_tuples = vec![]; let mut migrate_native_objects = false; @@ -48,6 +48,7 @@ pub async fn init() -> Result<(), anyhow::Error> { let mut need_cipher_keys_migration = false; let mut need_action_scripts_migration = false; let mut need_alert_folders_migration = false; + let mut need_ratelimit_migration = false; let mut existing_meta: Option = match db::ofga::get_ofga_model().await { Ok(Some(model)) => Some(model), @@ -138,6 +139,8 @@ pub async fn init() -> Result<(), anyhow::Error> { let v0_0_10 = version_compare::Version::from("0.0.10").unwrap(); let v0_0_12 = version_compare::Version::from("0.0.12").unwrap(); let v0_0_13 = version_compare::Version::from("0.0.13").unwrap(); + let v0_0_15 = version_compare::Version::from("0.0.15").unwrap(); + let v0_0_16 = version_compare::Version::from("0.0.16").unwrap(); if meta_version > v0_0_4 && existing_model_version < v0_0_5 { need_migrate_index_streams = true; } @@ -154,6 +157,11 @@ pub async fn init() -> Result<(), anyhow::Error> { log::info!("[OFGA:Local] Alert folders migration needed"); need_alert_folders_migration = true; } + + if meta_version > v0_0_15 && existing_model_version < v0_0_16 { + log::info!("[OFGA:Local] Ratelimit migration needed"); + need_ratelimit_migration = true; + } } // 1. create a cluster lock @@ -168,26 +176,13 @@ pub async fn init() -> Result<(), anyhow::Error> { o2_openfga::config::OFGA_STORE_ID.insert("store_id".to_owned(), store_id); let mut tuples = vec![]; - let r = infra::schema::STREAM_SCHEMAS.read().await; let mut orgs = HashSet::new(); - for key in r.keys() { - if !key.contains('/') { - continue; - } - let split_key = key.split('/').collect::>(); - let org_name = split_key[0]; - orgs.insert(org_name); - if need_migrate_index_streams - && split_key.len() > 2 - && split_key[1] == "index" - && !migrate_native_objects - { - get_tuple_for_new_index(org_name, split_key[2], &mut tuples); - } + for user in USERS.iter() { + orgs.insert(user.value().org.to_string()); } log::info!("[OFGA:Local] Migrating native objects"); if migrate_native_objects { - for org_name in orgs { + for org_name in orgs.iter() { get_org_creation_tuples( org_name, &mut tuples, @@ -230,7 +225,7 @@ pub async fn init() -> Result<(), anyhow::Error> { } } else { log::info!("[OFGA:Local] Migrating index streams"); - for org_name in orgs { + for org_name in orgs.iter() { if need_migrate_index_streams { get_index_creation_tuples(org_name, &mut tuples).await; } @@ -261,6 +256,9 @@ pub async fn init() -> Result<(), anyhow::Error> { get_ownership_all_org_tuple(org_name, "alert_folders", &mut tuples); get_ownership_tuple(org_name, "alert_folders", DEFAULT_FOLDER, &mut tuples); } + if need_ratelimit_migration { + get_ownership_all_org_tuple(org_name, "ratelimit", &mut tuples); + } } if need_alert_folders_migration { match migrations::migrate_alert_folders().await { @@ -298,7 +296,6 @@ pub async fn init() -> Result<(), anyhow::Error> { } } } - drop(r); } Err(e) => { log::error!("Error setting OFGA model: {:?}", e); diff --git a/src/common/utils/auth.rs b/src/common/utils/auth.rs index 4a7c13e8ae..acdac472df 100644 --- a/src/common/utils/auth.rs +++ b/src/common/utils/auth.rs @@ -358,7 +358,7 @@ impl FromRequest for AuthExtractor { OFGA_MODELS.get("streams").unwrap().key, path_columns[1] ) - } else if method.eq("PUT") + } else if (method.eq("PUT") && !path_columns[1].starts_with("ratelimit")) || method.eq("DELETE") || path_columns[1].eq("reports") || path_columns[1].eq("savedviews") @@ -396,6 +396,12 @@ impl FromRequest for AuthExtractor { } else { // for things like dashboards and folders etc, // this will take form org:dashboard or org:folders + + // handle ratelimit:org + if method.eq("GET") && path_columns[1].starts_with("ratelimit") { + method = "LIST".to_string(); + } + format!( "{}:{}", OFGA_MODELS diff --git a/src/common/utils/auth_tests.rs b/src/common/utils/auth_tests.rs index ef0d3b464d..f3eac11266 100644 --- a/src/common/utils/auth_tests.rs +++ b/src/common/utils/auth_tests.rs @@ -3432,6 +3432,10 @@ mod tests { algorithm: String::default(), master_key: String::default(), }, + ratelimit: config::RateLimit { + ratelimit_rule_refresh_interval: 0, + ratelimit_enabled: false, + }, } } } diff --git a/src/config/src/config.rs b/src/config/src/config.rs index 897a9ef148..cc216e7812 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -33,7 +33,7 @@ use lettre::{ use once_cell::sync::Lazy; use crate::{ - meta::cluster, + meta::{cluster, meta_store::MetaStore}, utils::{file::get_file_meta, sysinfo}, }; @@ -416,6 +416,7 @@ pub struct Config { pub pipeline: Pipeline, pub health_check: HealthCheck, pub encryption: Encryption, + pub ratelimit: RateLimit, } #[derive(EnvConfig)] @@ -1773,6 +1774,21 @@ pub struct HealthCheck { pub failed_times: usize, } +#[derive(EnvConfig)] +pub struct RateLimit { + #[env_config( + name = "ZO_RATELIMIT_ENABLED", + default = false, + help = "ratelimit enabled" + )] + pub ratelimit_enabled: bool, + #[env_config( + name = "ZO_RATELIMIT_RULE_REFRESH_INTERVAL", + default = 10, + help = "unit: seconds, refresh interval for rate limit rules" + )] + pub ratelimit_rule_refresh_interval: usize, +} pub fn init() -> Config { dotenv_override().ok(); let mut cfg = Config::init().expect("config init error"); @@ -1853,6 +1869,11 @@ pub fn init() -> Config { panic!("pipeline config error: {e}"); } + // check ratelimit config + if let Err(e) = check_ratelimit_config(&mut cfg) { + panic!("ratelimit config error: {e}"); + } + cfg } @@ -2571,6 +2592,24 @@ fn check_pipeline_config(cfg: &mut Config) -> Result<(), anyhow::Error> { Ok(()) } +fn check_ratelimit_config(cfg: &mut Config) -> Result<(), anyhow::Error> { + if cfg.ratelimit.ratelimit_enabled { + let meta_store: MetaStore = cfg.common.queue_store.as_str().into(); + if meta_store != MetaStore::Nats { + return Err(anyhow::anyhow!( + "ZO_QUEUE_STORE must be nats when ratelimit is enabled" + )); + } + } + + if cfg.ratelimit.ratelimit_rule_refresh_interval < 2 { + return Err(anyhow::anyhow!( + "ratelimit rules refresh interval must be greater than or equal to 2 seconds" + )); + } + Ok(()) +} + fn check_health_check_config(cfg: &mut Config) -> Result<(), anyhow::Error> { if cfg.health_check.timeout == 0 { cfg.health_check.timeout = 5; diff --git a/src/config/src/lib.rs b/src/config/src/lib.rs index ab05aef6ea..0f82b5ea24 100644 --- a/src/config/src/lib.rs +++ b/src/config/src/lib.rs @@ -18,6 +18,7 @@ pub mod config; pub mod ider; pub mod meta; pub mod metrics; +pub mod router; pub mod utils; pub use config::*; diff --git a/src/config/src/meta/mod.rs b/src/config/src/meta/mod.rs index a179a1708f..0b6eb373f5 100644 --- a/src/config/src/meta/mod.rs +++ b/src/config/src/meta/mod.rs @@ -27,6 +27,7 @@ pub mod meta_store; pub mod otlp; pub mod pipeline; pub mod promql; +pub mod ratelimit; pub mod search; pub mod self_reporting; pub mod short_url; diff --git a/src/config/src/meta/ratelimit.rs b/src/config/src/meta/ratelimit.rs new file mode 100644 index 0000000000..e331a440d4 --- /dev/null +++ b/src/config/src/meta/ratelimit.rs @@ -0,0 +1,160 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::{collections::HashMap, time::Instant}; + +use anyhow::Context; +use bytes::Bytes; +use sea_orm::FromQueryResult; +use serde::{Deserialize, Serialize}; +use utoipa::ToSchema; + +pub struct CachedUserRoles { + pub roles: Vec, + pub expires_at: Instant, +} + +#[derive(FromQueryResult, Default, Debug, Clone, Serialize, Deserialize, ToSchema)] +pub struct RatelimitRule { + pub org: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub rule_type: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub rule_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub user_role: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub user_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub api_group_name: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub api_group_operation: Option, + pub threshold: i32, +} + +impl TryFrom<&Bytes> for RatelimitRule { + type Error = anyhow::Error; + + fn try_from(bytes: &Bytes) -> Result { + serde_json::from_slice(bytes).context("Failed to deserialize RatelimitRule") + } +} + +#[derive(Debug)] +pub struct RatelimitRuleUpdater(pub HashMap>); + +impl RatelimitRule { + pub fn from_updater( + org: &str, + role: &str, + updater: &RatelimitRuleUpdater, + ) -> Vec { + updater + .0 + .iter() + .flat_map(|(group_name, operations)| { + operations + .iter() + .map(|(operation, threshold)| RatelimitRule { + org: org.to_string(), + rule_type: Some(RatelimitRuleType::Exact.to_string()), + rule_id: Some(crate::ider::generate()), + user_role: Some(role.to_string()), + user_id: Some(".*".to_string()), + api_group_name: Some(group_name.clone()), + api_group_operation: Some(operation.clone().to_lowercase()), + threshold: *threshold, + }) + }) + .collect() + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum RatelimitRuleType { + Exact, + Regex, +} + +impl From<&str> for RatelimitRuleType { + fn from(s: &str) -> Self { + match s { + "exact" => RatelimitRuleType::Exact, + "regex" => RatelimitRuleType::Regex, + _ => RatelimitRuleType::Exact, // Default to Exact for invalid input + } + } +} + +impl std::fmt::Display for RatelimitRuleType { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + RatelimitRuleType::Exact => write!(f, "exact"), + RatelimitRuleType::Regex => write!(f, "regex"), + } + } +} + +impl RatelimitRule { + pub fn get_resource(&self) -> String { + get_resource_from_params( + self.get_org(), + self.get_rule_type(), + self.get_user_role(), + self.get_api_group_name(), + self.get_api_group_operation(), + self.get_user_id(), + ) + } + + pub fn get_rule_id(&self) -> &str { + self.rule_id.as_deref().unwrap_or_default() + } + pub fn get_org(&self) -> &str { + &self.org + } + pub fn get_rule_type(&self) -> &str { + self.rule_type.as_deref().unwrap_or_default() + } + pub fn get_user_role(&self) -> &str { + self.user_role.as_deref().unwrap_or_default() + } + pub fn get_user_id(&self) -> &str { + self.user_id.as_deref().unwrap_or_default() + } + pub fn get_api_group_name(&self) -> &str { + self.api_group_name.as_deref().unwrap_or_default() + } + pub fn get_api_group_operation(&self) -> &str { + self.api_group_operation.as_deref().unwrap_or_default() + } + pub fn get_threshold(&self) -> i32 { + self.threshold + } +} + +pub fn get_resource_from_params( + org: &str, + rule_type: &str, + user_role: &str, + api_group_name: &str, + api_group_operation: &str, + user_id: &str, +) -> String { + format!( + "{}:{}:{}:{}:{}:{}", + org, rule_type, user_role, api_group_name, api_group_operation, user_id + ) +} diff --git a/src/config/src/router.rs b/src/config/src/router.rs new file mode 100644 index 0000000000..7418a64662 --- /dev/null +++ b/src/config/src/router.rs @@ -0,0 +1,92 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +/// usize indicates the number of parts to skip based on their actual paths. +const QUERIER_ROUTES: [(&str, usize); 21] = [ + ("config", 0), // /config + ("summary", 2), // /api/{org_id}/summary + ("organizations", 1), // /api/organizations + ("settings", 2), // /api/{org_id}/settings/... + ("schema", 3), // /api/{org_id}/streams/{stream_name}/schema + ("streams", 2), // /api/{org_id}/streams/... + ("traces/latest", 3), // /api/{org_id}/{stream_name}/traces/latest + ("clusters", 1), // /api/clusters + ("query_manager", 2), // /api/{org_id}/query_manager/... + ("ws", 2), // /api/{org_id}/ws + ("_search", 2), // /api/{org_id}/_search + ("_around", 3), // /api/{org_id}/{stream_name}/_around + ("_values", 3), // /api/{org_id}/{stream_name}/_values + ("functions?page_num=", 2), // /api/{org_id}/functions + ("prometheus/api/v1/series", 2), // /api/{org_id}/prometheus/api/v1/series + ("prometheus/api/v1/query", 2), // /api/{org_id}/prometheus/api/v1/query + ("prometheus/api/v1/query_range", 2), // /api/{org_id}/prometheus/api/v1/query_range + ("prometheus/api/v1/query_exemplars", 2), // /api/{org_id}/prometheus/api/v1/query_exemplars + ("prometheus/api/v1/metadata", 2), // /api/{org_id}/prometheus/api/v1/metadata + ("prometheus/api/v1/labels", 2), // /api/{org_id}/prometheus/api/v1/labels + ("prometheus/api/v1/label/", 2), /* /api/{org_id}/prometheus/api/v1/label/ + * {label_name}/ + * values */ +]; +const QUERIER_ROUTES_BY_BODY: [&str; 2] = [ + "/prometheus/api/v1/query_range", + "/prometheus/api/v1/query_exemplars", +]; +const FIXED_QUERIER_ROUTES: [&str; 3] = ["/summary", "/schema", "/streams"]; +pub const INGESTER_ROUTES: [&str; 11] = [ + "/_bulk", + "/_multi", + "/_json", + "/_kinesis_firehose", + "/_sub", + "/v1/logs", + "/ingest/metrics/_json", + "/v1/metrics", + "/traces", + "/v1/traces", + "/traces/latest", +]; + +#[inline] +pub fn is_querier_route(path: &str) -> bool { + QUERIER_ROUTES.iter().any(|(route, skip_segments)| { + if path.contains(route) { + let segments = path + .split('/') + .filter(|s| !s.is_empty()) + .collect::>(); + // check if we have enough segments + if segments.len() <= *skip_segments { + return false; + } + let route_part = segments[*skip_segments..].join("/"); + route_part.starts_with(route) + && INGESTER_ROUTES + .iter() + .all(|ingest_route| !route_part.ends_with(ingest_route)) + } else { + false + } + }) +} + +#[inline] +pub fn is_querier_route_by_body(path: &str) -> bool { + QUERIER_ROUTES_BY_BODY.iter().any(|x| path.contains(x)) +} + +#[inline] +pub fn is_fixed_querier_route(path: &str) -> bool { + FIXED_QUERIER_ROUTES.iter().any(|x| path.contains(x)) +} diff --git a/src/handler/http/request/actions/action.rs b/src/handler/http/request/actions/action.rs index cce028f652..2a2402b298 100644 --- a/src/handler/http/request/actions/action.rs +++ b/src/handler/http/request/actions/action.rs @@ -13,38 +13,46 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::{collections::HashMap, io::Error, str::FromStr}; +use std::io::Error; use actix_multipart::Multipart; use actix_web::{HttpRequest, HttpResponse, delete, get, post, put, web}; -use bytes::Bytes; -use config::meta::actions::action::{Action, ExecutionDetailsType, UpdateActionDetailsRequest}; -use futures::{StreamExt, TryStreamExt}; -use futures_util::stream::{self}; -use infra::table::action_scripts; -use o2_enterprise::enterprise::actions::action_manager::{ - delete_app_from_target_cluster, get_action_details, get_actions, register_app, - serve_file_from_s3, update_app_on_target_cluster, -}; -use once_cell::sync::Lazy; -use regex::Regex; -use serde_json; +use config::meta::actions::action::UpdateActionDetailsRequest; use svix_ksuid::Ksuid; - -use crate::{ - common::{ - meta::{authz::Authz, http::HttpResponse as MetaHttpResponse}, - utils::auth::{UserEmail, check_permissions, remove_ownership, set_ownership}, +#[cfg(feature = "enterprise")] +use { + crate::{ + common::{ + meta::{authz::Authz, http::HttpResponse as MetaHttpResponse}, + utils::auth::{check_permissions, remove_ownership, set_ownership}, + }, + handler::http::models::action::{GetActionDetailsResponse, GetActionInfoResponse}, + service::organization::get_passcode, }, - handler::http::models::action::{GetActionDetailsResponse, GetActionInfoResponse}, - service::organization::get_passcode, + bytes::Bytes, + config::meta::actions::action::{Action, ExecutionDetailsType}, + futures::{StreamExt, TryStreamExt}, + futures_util::stream::{self}, + infra::table::action_scripts, + o2_enterprise::enterprise::actions::action_manager::{ + delete_app_from_target_cluster, get_action_details, get_actions, register_app, + serve_file_from_s3, update_app_on_target_cluster, + }, + once_cell::sync::Lazy, + regex::Regex, + serde_json, + std::collections::HashMap, + std::str::FromStr, }; +use crate::common::utils::auth::UserEmail; + +#[cfg(feature = "enterprise")] const MANDATORY_FIELDS_FOR_ACTION_CREATION: [&str; 5] = ["name", "owner", "file", "filename", "execution_details"]; - +#[cfg(feature = "enterprise")] static ENV_VAR_REGEX: Lazy = Lazy::new(|| Regex::new(r"^[A-Z][A-Z0-9_]*$").unwrap()); - +#[cfg(feature = "enterprise")] fn validate_environment_variables(env_vars: &HashMap) -> Result<(), String> { for key in env_vars.keys() { if !ENV_VAR_REGEX.is_match(key) { @@ -55,6 +63,8 @@ fn validate_environment_variables(env_vars: &HashMap) -> Result< } /// Delete Action +/// +/// #{"ratelimit_module":"Actions", "ratelimit_module_operation":"delete"}# #[utoipa::path( context_path = "/api", tag = "Actions", @@ -73,17 +83,28 @@ fn validate_environment_variables(env_vars: &HashMap) -> Result< )] #[delete("/{org_id}/actions/{ksuid}")] pub async fn delete_action(path: web::Path<(String, Ksuid)>) -> Result { - let (org_id, ksuid) = path.into_inner(); - match delete_app_from_target_cluster(&org_id, ksuid).await { - Ok(_) => { - remove_ownership(&org_id, "actions", Authz::new(&ksuid.to_string())).await; - Ok(MetaHttpResponse::ok("Action deleted")) + #[cfg(feature = "enterprise")] + { + let (org_id, ksuid) = path.into_inner(); + match delete_app_from_target_cluster(&org_id, ksuid).await { + Ok(_) => { + remove_ownership(&org_id, "actions", Authz::new(&ksuid.to_string())).await; + Ok(MetaHttpResponse::ok("Action deleted")) + } + Err(e) => Ok(MetaHttpResponse::bad_request(e)), } - Err(e) => Ok(MetaHttpResponse::bad_request(e)), + } + + #[cfg(not(feature = "enterprise"))] + { + drop(path); + Ok(HttpResponse::Forbidden().json("Not Supported")) } } /// Serve Action zip file +/// +/// #{"ratelimit_module":"Actions", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Actions", @@ -102,23 +123,34 @@ pub async fn delete_action(path: web::Path<(String, Ksuid)>) -> Result) -> Result { - let (org_id, ksuid) = path.into_inner(); - match serve_file_from_s3(&org_id, ksuid).await { - Ok((bytes, file_name)) => { - let resp = HttpResponse::Ok() - .insert_header(( - "Content-Disposition", - format!("attachment; filename=\"{}\"", file_name), - )) - .content_type("application/zip") - .streaming(stream::once(async { Ok::(bytes) })); - Ok(resp) + #[cfg(feature = "enterprise")] + { + let (org_id, ksuid) = path.into_inner(); + match serve_file_from_s3(&org_id, ksuid).await { + Ok((bytes, file_name)) => { + let resp = HttpResponse::Ok() + .insert_header(( + "Content-Disposition", + format!("attachment; filename=\"{}\"", file_name), + )) + .content_type("application/zip") + .streaming(stream::once(async { Ok::(bytes) })); + Ok(resp) + } + Err(e) => Ok(MetaHttpResponse::bad_request(e)), } - Err(e) => Ok(MetaHttpResponse::bad_request(e)), + } + + #[cfg(not(feature = "enterprise"))] + { + drop(path); + Ok(HttpResponse::Forbidden().json("Not Supported")) } } /// Update Action +/// +/// #{"ratelimit_module":"Actions", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Actions", @@ -140,40 +172,57 @@ pub async fn update_action_details( path: web::Path<(String, Ksuid)>, req: web::Json, ) -> Result { - let (org_id, ksuid) = path.into_inner(); - let mut req = req.into_inner(); + #[cfg(feature = "enterprise")] + { + let (org_id, ksuid) = path.into_inner(); + let mut req = req.into_inner(); - // Validate environment variables if they are being updated - if let Some(ref env_vars) = req.environment_variables { - if let Err(e) = validate_environment_variables(env_vars) { - return Ok(MetaHttpResponse::bad_request(e)); + // Validate environment variables if they are being updated + if let Some(ref env_vars) = req.environment_variables { + if let Err(e) = + crate::handler::http::request::actions::action::validate_environment_variables( + env_vars, + ) + { + return Ok(MetaHttpResponse::bad_request(e)); + } + } + + let sa = match req.service_account.clone() { + None => { + if let Ok(action) = action_scripts::get(&ksuid.to_string(), &org_id).await { + action.service_account + } else { + return Ok(MetaHttpResponse::bad_request("Failed to fetch action")); + } + } + Some(sa) => sa, + }; + let passcode = + if let Ok(res) = crate::service::organization::get_passcode(Some(&org_id), &sa).await { + res.passcode + } else { + return Ok(MetaHttpResponse::bad_request("Failed to fetch passcode")); + }; + + req.service_account = Some(sa); + match update_app_on_target_cluster(&org_id, ksuid, req, &passcode).await { + Ok(uuid) => Ok(MetaHttpResponse::json(serde_json::json!({"uuid":uuid}))), + Err(e) => Ok(MetaHttpResponse::bad_request(e)), } } - let sa = match req.service_account.clone() { - None => { - if let Ok(action) = action_scripts::get(&ksuid.to_string(), &org_id).await { - action.service_account - } else { - return Ok(MetaHttpResponse::bad_request("Failed to fetch action")); - } - } - Some(sa) => sa, - }; - let passcode = if let Ok(res) = get_passcode(Some(&org_id), &sa).await { - res.passcode - } else { - return Ok(HttpResponse::BadRequest().body("Failed to fetch passcode")); - }; - - req.service_account = Some(sa); - match update_app_on_target_cluster(&org_id, ksuid, req, &passcode).await { - Ok(uuid) => Ok(MetaHttpResponse::json(serde_json::json!({"uuid":uuid}))), - Err(e) => Ok(MetaHttpResponse::bad_request(e)), + #[cfg(not(feature = "enterprise"))] + { + drop(path); + drop(req); + Ok(HttpResponse::Forbidden().json("Not Supported")) } } /// List Actions +/// +/// #{"ratelimit_module":"Actions", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Actions", @@ -192,24 +241,35 @@ body = HttpResponse), ) )] #[get("/{org_id}/actions")] pub async fn list_actions(path: web::Path) -> Result { - let org_id = path.into_inner(); + #[cfg(feature = "enterprise")] + { + let org_id = path.into_inner(); - if let Ok(list) = get_actions(&org_id).await { - if let Ok(list) = list - .into_iter() - .map(GetActionInfoResponse::try_from) - .collect::, _>>() - { - Ok(MetaHttpResponse::json(list)) + if let Ok(list) = get_actions(&org_id).await { + if let Ok(list) = list + .into_iter() + .map(GetActionInfoResponse::try_from) + .collect::, _>>() + { + Ok(MetaHttpResponse::json(list)) + } else { + Ok(MetaHttpResponse::bad_request("Failed to transform actions")) + } } else { - Ok(MetaHttpResponse::bad_request("Failed to transform actions")) + Ok(MetaHttpResponse::bad_request("Failed to fetch actions")) } - } else { - Ok(MetaHttpResponse::bad_request("Failed to fetch actions")) + } + + #[cfg(not(feature = "enterprise"))] + { + drop(path); + Ok(HttpResponse::Forbidden().json("Not Supported")) } } /// Get single Action +/// +/// #{"ratelimit_module":"Actions", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Actions", @@ -229,24 +289,36 @@ body = HttpResponse), ) )] #[get("/{org_id}/actions/{action_id}")] pub async fn get_action_from_id(path: web::Path<(String, String)>) -> Result { - let (org_id, action_id) = path.into_inner(); - match get_action_details(&org_id, &action_id).await { - Ok(action) => { - if let Ok(resp) = GetActionDetailsResponse::try_from(action) { - Ok(MetaHttpResponse::json(resp)) - } else { - Ok(MetaHttpResponse::bad_request( - "Failed to fetch action details", - )) + #[cfg(feature = "enterprise")] + { + let (org_id, action_id) = path.into_inner(); + match get_action_details(&org_id, &action_id).await { + Ok(action) => { + if let Ok(resp) = GetActionDetailsResponse::try_from(action) { + Ok(MetaHttpResponse::json(resp)) + } else { + Ok(MetaHttpResponse::bad_request( + "Failed to fetch action details", + )) + } } + Err(e) => Ok(MetaHttpResponse::bad_request(e)), } - Err(e) => Ok(MetaHttpResponse::bad_request(e)), + } + + #[cfg(not(feature = "enterprise"))] + { + drop(path); + Ok(HttpResponse::Forbidden().json("Not Supported")) } } +/// UploadZippedAction +/// /// Upload a zipped action file and process it. /// This endpoint allows uploading a ZIP file containing an action, which will be extracted, /// processed, and executed. +/// #{"ratelimit_module":"Actions", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Actions", @@ -266,318 +338,334 @@ pub async fn get_action_from_id(path: web::Path<(String, String)>) -> Result, - mut payload: Multipart, + #[cfg_attr(not(feature = "enterprise"), allow(unused_mut))] mut payload: Multipart, req: HttpRequest, user_email: UserEmail, ) -> Result { - let org_id = path.into_inner(); - let mut file_data = Vec::new(); - let mut is_update = false; - let mut action = Action { - org_id: org_id.clone(), - ..Default::default() - }; + #[cfg(feature = "enterprise")] + { + let org_id = path.into_inner(); + let mut file_data = Vec::new(); + let mut is_update = false; + let mut action = Action { + org_id: org_id.clone(), + ..Default::default() + }; - // Validate Content-Type - if let Some(content_type) = req.headers().get("Content-Type") { - match content_type.to_str() { - Ok(content_type_str) => { - if !content_type_str.contains("multipart/form-data") { - return Ok(HttpResponse::BadRequest().body("Invalid Content-Type")); + // Validate Content-Type + if let Some(content_type) = req.headers().get("Content-Type") { + match content_type.to_str() { + Ok(content_type_str) => { + if !content_type_str.contains("multipart/form-data") { + return Ok(HttpResponse::BadRequest().body("Invalid Content-Type")); + } + } + Err(_) => { + return Ok(HttpResponse::BadRequest().body("Invalid Content-Type header value")); } - } - Err(_) => { - return Ok(HttpResponse::BadRequest().body("Invalid Content-Type header value")); } } - } - let mut received_fields = Vec::new(); - while let Ok(Some(mut field)) = payload.try_next().await { - match field.name().unwrap_or("") { - "description" => { - received_fields.push("description"); - let mut description = Vec::new(); - while let Some(chunk) = field.next().await { - match chunk { - Ok(bytes) => description.extend_from_slice(&bytes), - Err(_) => { - return Ok( - HttpResponse::BadRequest().body("Failed to read description field") - ); + let mut received_fields = Vec::new(); + while let Ok(Some(mut field)) = payload.try_next().await { + match field.name().unwrap_or("") { + "description" => { + received_fields.push("description"); + let mut description = Vec::new(); + while let Some(chunk) = field.next().await { + match chunk { + Ok(bytes) => description.extend_from_slice(&bytes), + Err(_) => { + return Ok(HttpResponse::BadRequest() + .body("Failed to read description field")); + } } } - } - if let Ok(description) = String::from_utf8(description) { - action.description = Some(description); - } else { - return Ok(HttpResponse::BadRequest() - .body("Description field contains invalid UTF-8 data")); - } - } - "file" => { - received_fields.push("file"); - if let Some(name) = field.content_disposition() { - action.zip_file_name = name.to_string(); - } - - while let Some(chunk) = field.next().await { - match chunk { - Ok(bytes) => file_data.extend_from_slice(&bytes), - Err(_) => return Ok(HttpResponse::BadRequest().body("Failed to read file")), - } - } - if file_data.is_empty() { - return Ok(HttpResponse::BadRequest().body("File is missing or empty")); - } - } - "filename" => { - received_fields.push("filename"); - let mut filename = Vec::new(); - while let Some(chunk) = field.next().await { - match chunk { - Ok(bytes) => filename.extend_from_slice(&bytes), - Err(_) => { - return Ok( - HttpResponse::BadRequest().body("Failed to read filename field") - ); - } - } - } - if filename.is_empty() { - return Ok( - HttpResponse::BadRequest().body("Filename field is missing or empty") - ); - } - if let Ok(filename) = String::from_utf8(filename) { - action.zip_file_name = filename; - } else { - return Ok(HttpResponse::BadRequest() - .body("Filename field contains invalid UTF-8 data")); - } - } - "name" => { - received_fields.push("name"); - let mut name = Vec::new(); - while let Some(chunk) = field.next().await { - match chunk { - Ok(bytes) => name.extend_from_slice(&bytes), - Err(_) => { - return Ok(HttpResponse::BadRequest().body("Failed to read name field")); - } - } - } - if name.is_empty() { - return Ok(HttpResponse::BadRequest().body("Name field is missing or empty")); - } - if let Ok(name) = String::from_utf8(name) { - action.name = name; - } else { - return Ok( - HttpResponse::BadRequest().body("Name field contains invalid UTF-8 data") - ); - } - } - "execution_details" => { - received_fields.push("execution_details"); - let mut details = Vec::new(); - while let Some(chunk) = field.next().await { - match chunk { - Ok(bytes) => details.extend_from_slice(&bytes), - Err(_) => { - return Ok(HttpResponse::BadRequest() - .body("Failed to read execution_details field")); - } - } - } - if details.is_empty() { - return Ok(HttpResponse::BadRequest() - .body("Execution details field is missing or empty")); - } - if let Ok(exec_details) = std::str::from_utf8(&details) { - if let Ok(exec_details_type) = ExecutionDetailsType::try_from(exec_details) { - action.execution_details = exec_details_type; + if let Ok(description) = String::from_utf8(description) { + action.description = Some(description); } else { - return Ok(HttpResponse::BadRequest().body("Invalid execution details")); + return Ok(HttpResponse::BadRequest() + .body("Description field contains invalid UTF-8 data")); } - } else { - return Ok(HttpResponse::BadRequest() - .body("Execution details field contains invalid UTF-8 data")); } - } - "cron_expr" => { - received_fields.push("cron_expr"); - let mut cron = Vec::new(); - while let Some(chunk) = field.next().await { - match chunk { - Ok(bytes) => cron.extend_from_slice(&bytes), - Err(_) => { - return Ok( - HttpResponse::BadRequest().body("Failed to read cron_expr field") - ); + "file" => { + received_fields.push("file"); + if let Some(name) = field.content_disposition() { + action.zip_file_name = name.to_string(); + } + + while let Some(chunk) = field.next().await { + match chunk { + Ok(bytes) => file_data.extend_from_slice(&bytes), + Err(_) => { + return Ok(HttpResponse::BadRequest().body("Failed to read file")); + } } } + if file_data.is_empty() { + return Ok(HttpResponse::BadRequest().body("File is missing or empty")); + } } - if let Ok(cron) = String::from_utf8(cron) { - action.cron_expr = Some(cron); - } else { - return Ok(HttpResponse::BadRequest() - .body("Cron expression contains invalid UTF-8 data")); + "filename" => { + received_fields.push("filename"); + let mut filename = Vec::new(); + while let Some(chunk) = field.next().await { + match chunk { + Ok(bytes) => filename.extend_from_slice(&bytes), + Err(_) => { + return Ok(HttpResponse::BadRequest() + .body("Failed to read filename field")); + } + } + } + if filename.is_empty() { + return Ok( + HttpResponse::BadRequest().body("Filename field is missing or empty") + ); + } + if let Ok(filename) = String::from_utf8(filename) { + action.zip_file_name = filename; + } else { + return Ok(HttpResponse::BadRequest() + .body("Filename field contains invalid UTF-8 data")); + } } - } - "environment_variables" => { - received_fields.push("environment_variables"); - let mut env_vars = Vec::new(); - while let Some(chunk) = field.next().await { - match chunk { - Ok(bytes) => env_vars.extend_from_slice(&bytes), - Err(_) => { + "name" => { + received_fields.push("name"); + let mut name = Vec::new(); + while let Some(chunk) = field.next().await { + match chunk { + Ok(bytes) => name.extend_from_slice(&bytes), + Err(_) => { + return Ok( + HttpResponse::BadRequest().body("Failed to read name field") + ); + } + } + } + if name.is_empty() { + return Ok( + HttpResponse::BadRequest().body("Name field is missing or empty") + ); + } + if let Ok(name) = String::from_utf8(name) { + action.name = name; + } else { + return Ok(HttpResponse::BadRequest() + .body("Name field contains invalid UTF-8 data")); + } + } + "execution_details" => { + received_fields.push("execution_details"); + let mut details = Vec::new(); + while let Some(chunk) = field.next().await { + match chunk { + Ok(bytes) => details.extend_from_slice(&bytes), + Err(_) => { + return Ok(HttpResponse::BadRequest() + .body("Failed to read execution_details field")); + } + } + } + if details.is_empty() { + return Ok(HttpResponse::BadRequest() + .body("Execution details field is missing or empty")); + } + if let Ok(exec_details) = std::str::from_utf8(&details) { + if let Ok(exec_details_type) = ExecutionDetailsType::try_from(exec_details) + { + action.execution_details = exec_details_type; + } else { + return Ok(HttpResponse::BadRequest().body("Invalid execution details")); + } + } else { + return Ok(HttpResponse::BadRequest() + .body("Execution details field contains invalid UTF-8 data")); + } + } + "cron_expr" => { + received_fields.push("cron_expr"); + let mut cron = Vec::new(); + while let Some(chunk) = field.next().await { + match chunk { + Ok(bytes) => cron.extend_from_slice(&bytes), + Err(_) => { + return Ok(HttpResponse::BadRequest() + .body("Failed to read cron_expr field")); + } + } + } + if let Ok(cron) = String::from_utf8(cron) { + action.cron_expr = Some(cron); + } else { + return Ok(HttpResponse::BadRequest() + .body("Cron expression contains invalid UTF-8 data")); + } + } + "environment_variables" => { + received_fields.push("environment_variables"); + let mut env_vars = Vec::new(); + while let Some(chunk) = field.next().await { + match chunk { + Ok(bytes) => env_vars.extend_from_slice(&bytes), + Err(_) => { + return Ok(HttpResponse::BadRequest() + .body("Failed to read environment_variables field")); + } + } + } + if let Ok(env_vars) = String::from_utf8(env_vars) { + if let Ok(env_vars) = serde_json::from_str(&env_vars) { + // Validate environment variables before assigning + if let Err(e) = validate_environment_variables(&env_vars) { + return Ok(MetaHttpResponse::bad_request(e)); + } + action.environment_variables = env_vars; + } else { return Ok(HttpResponse::BadRequest() - .body("Failed to read environment_variables field")); + .body("Invalid JSON in environment variables")); } - } - } - if let Ok(env_vars) = String::from_utf8(env_vars) { - if let Ok(env_vars) = serde_json::from_str(&env_vars) { - // Validate environment variables before assigning - if let Err(e) = validate_environment_variables(&env_vars) { - return Ok(MetaHttpResponse::bad_request(e)); - } - action.environment_variables = env_vars; } else { return Ok(HttpResponse::BadRequest() .body("Invalid JSON in environment variables")); } - } else { - return Ok( - HttpResponse::BadRequest().body("Invalid JSON in environment variables") - ); } - } - "owner" => { - received_fields.push("owner"); - let mut owner = Vec::new(); - while let Some(chunk) = field.next().await { - match chunk { - Ok(bytes) => owner.extend_from_slice(&bytes), - Err(_) => { - return Ok( - HttpResponse::BadRequest().body("Failed to read owner field") - ); + "owner" => { + received_fields.push("owner"); + let mut owner = Vec::new(); + while let Some(chunk) = field.next().await { + match chunk { + Ok(bytes) => owner.extend_from_slice(&bytes), + Err(_) => { + return Ok( + HttpResponse::BadRequest().body("Failed to read owner field") + ); + } } } - } - if let Ok(owner) = String::from_utf8(owner) { - action.created_by = owner; - } else { - return Ok( - HttpResponse::BadRequest().body("Owner field contains invalid UTF-8 data") - ); - } - } - "id" => { - received_fields.push("id"); - let mut id = Vec::new(); - while let Some(chunk) = field.next().await { - match chunk { - Ok(bytes) => id.extend_from_slice(&bytes), - Err(_) => { - return Ok(HttpResponse::BadRequest().body("Failed to read id field")); - } + if let Ok(owner) = String::from_utf8(owner) { + action.created_by = owner; + } else { + return Ok(HttpResponse::BadRequest() + .body("Owner field contains invalid UTF-8 data")); } } - if let Ok(id) = String::from_utf8(id) { - if let Ok(id) = Ksuid::from_str(&id) { - action.id = Some(id); - is_update = true; + "id" => { + received_fields.push("id"); + let mut id = Vec::new(); + while let Some(chunk) = field.next().await { + match chunk { + Ok(bytes) => id.extend_from_slice(&bytes), + Err(_) => { + return Ok( + HttpResponse::BadRequest().body("Failed to read id field") + ); + } + } + } + if let Ok(id) = String::from_utf8(id) { + if let Ok(id) = Ksuid::from_str(&id) { + action.id = Some(id); + is_update = true; + } else { + return Ok(HttpResponse::BadRequest().body("Invalid ID")); + } } else { return Ok(HttpResponse::BadRequest().body("Invalid ID")); } - } else { - return Ok(HttpResponse::BadRequest().body("Invalid ID")); } - } - "service_account" => { - received_fields.push("service_account"); - let mut service_account = Vec::new(); - while let Some(chunk) = field.next().await { - match chunk { - Ok(bytes) => service_account.extend_from_slice(&bytes), - Err(_) => { - return Ok(HttpResponse::BadRequest() - .body("Failed to read service_account field")); + "service_account" => { + received_fields.push("service_account"); + let mut service_account = Vec::new(); + while let Some(chunk) = field.next().await { + match chunk { + Ok(bytes) => service_account.extend_from_slice(&bytes), + Err(_) => { + return Ok(HttpResponse::BadRequest() + .body("Failed to read service_account field")); + } } } - } - if let Ok(service_account) = String::from_utf8(service_account) { - action.service_account = service_account; - } else { - return Ok(HttpResponse::BadRequest() - .body("Service account field contains invalid UTF-8 data")); - } - } - _ => {} - } - } - - if !MANDATORY_FIELDS_FOR_ACTION_CREATION - .iter() - .all(|field| received_fields.contains(field)) - { - return Ok(HttpResponse::BadRequest().body("Missing mandatory fields")); - } - - // If action ID is present as field then we know its an update request - // Hence we treat it as a PUT request and check for permissions - let method = match action.id { - Some(_) => "PUT", - None => "POST", - }; - if !check_permissions( - action.id.map(|ksuid| ksuid.to_string()), - &org_id, - &user_email.user_id, - "actions", - method, - "", - ) - .await - { - return Ok(HttpResponse::Forbidden().body("Unauthorized Access")); - } - - if file_data.is_empty() { - return Ok(HttpResponse::BadRequest().body("Uploaded file is empty")); - } - - let file_path = format!("files/{}/actions/{}", org_id, action.zip_file_name); - - let passcode = if let Ok(res) = get_passcode(Some(&org_id), &action.service_account).await { - res.passcode - } else { - return Ok(HttpResponse::BadRequest().body("Failed to fetch passcode")); - }; - - // Attempt to read the uploaded data as a ZIP file - match zip::read::ZipArchive::new(std::io::Cursor::new(file_data)) { - Ok(archive) => { - log::info!("Successfully read ZIP archive with {} files", archive.len()); - match register_app(action, archive, &file_path, &passcode).await { - Ok(uuid) => { - if !is_update { - set_ownership(&org_id, "actions", Authz::new(&uuid.to_string())).await; + if let Ok(service_account) = String::from_utf8(service_account) { + action.service_account = service_account; + } else { + return Ok(HttpResponse::BadRequest() + .body("Service account field contains invalid UTF-8 data")); } - Ok(MetaHttpResponse::json(serde_json::json!({"uuid":uuid}))) } - Err(e) => Ok(HttpResponse::BadRequest().json( - serde_json::json!({"message":format!("Failed to process action:{}", e)}), - )), + _ => {} } } - Err(e) => { - log::error!("Error reading ZIP file: {:?}", e); - Ok(HttpResponse::BadRequest() - .json(serde_json::json!({"message":format!("Error reading ZIP file: {:?}", e)}))) + + if !MANDATORY_FIELDS_FOR_ACTION_CREATION + .iter() + .all(|field| received_fields.contains(field)) + { + return Ok(HttpResponse::BadRequest().body("Missing mandatory fields")); } + + // If action ID is present as field then we know its an update request + // Hence we treat it as a PUT request and check for permissions + let method = match action.id { + Some(_) => "PUT", + None => "POST", + }; + if !check_permissions( + action.id.map(|ksuid| ksuid.to_string()), + &org_id, + &user_email.user_id, + "actions", + method, + "", + ) + .await + { + return Ok(HttpResponse::Forbidden().body("Unauthorized Access")); + } + + if file_data.is_empty() { + return Ok(HttpResponse::BadRequest().body("Uploaded file is empty")); + } + + let file_path = format!("files/{}/actions/{}", org_id, action.zip_file_name); + + let passcode = if let Ok(res) = get_passcode(Some(&org_id), &action.service_account).await { + res.passcode + } else { + return Ok(HttpResponse::BadRequest().body("Failed to fetch passcode")); + }; + + // Attempt to read the uploaded data as a ZIP file + match zip::read::ZipArchive::new(std::io::Cursor::new(file_data)) { + Ok(archive) => { + log::info!("Successfully read ZIP archive with {} files", archive.len()); + match register_app(action, archive, &file_path, &passcode).await { + Ok(uuid) => { + if !is_update { + set_ownership(&org_id, "actions", Authz::new(&uuid.to_string())).await; + } + Ok(MetaHttpResponse::json(serde_json::json!({"uuid":uuid}))) + } + Err(e) => Ok(HttpResponse::BadRequest().json( + serde_json::json!({"message":format!("Failed to process action:{}", e)}), + )), + } + } + Err(e) => { + log::error!("Error reading ZIP file: {:?}", e); + Ok(HttpResponse::BadRequest().json( + serde_json::json!({"message":format!("Error reading ZIP file: {:?}", e)}), + )) + } + } + } + + #[cfg(not(feature = "enterprise"))] + { + drop(path); + let _ = payload; + drop(req); + drop(user_email); + Ok(HttpResponse::Forbidden().json("Not Supported")) } } diff --git a/src/handler/http/request/actions/operations.rs b/src/handler/http/request/actions/operations.rs index 181253a235..22c6ecaebd 100644 --- a/src/handler/http/request/actions/operations.rs +++ b/src/handler/http/request/actions/operations.rs @@ -14,20 +14,24 @@ use std::io::Error; -use actix_http::header::HeaderMap; -use actix_web::{HttpResponse, get, post, web}; -use itertools::Itertools; -use o2_enterprise::enterprise::actions::{action_manager::trigger_action, meta::TriggerSource}; -use tracing::{Level, span}; - -use crate::{ - common::utils::{ - auth::{UserEmail, check_permissions}, - http::get_or_create_trace_id, +use actix_web::{HttpResponse, get, web}; +#[cfg(feature = "enterprise")] +use { + crate::{ + common::utils::{ + auth::{UserEmail, check_permissions}, + http::get_or_create_trace_id, + }, + handler::http::models::action::TestActionRequest, }, - handler::http::models::action::TestActionRequest, + actix_http::header::HeaderMap, + actix_web::post, + itertools::Itertools, + o2_enterprise::enterprise::actions::{action_manager::trigger_action, meta::TriggerSource}, + tracing::{Level, span}, }; +#[cfg(feature = "enterprise")] /// Test Action #[utoipa::path( context_path = "/api", diff --git a/src/handler/http/request/alerts/deprecated.rs b/src/handler/http/request/alerts/deprecated.rs index 3ecced6cbf..6b4a9cb1a2 100644 --- a/src/handler/http/request/alerts/deprecated.rs +++ b/src/handler/http/request/alerts/deprecated.rs @@ -38,6 +38,8 @@ use crate::{ }; /// CreateAlert +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"create"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -82,6 +84,8 @@ pub async fn save_alert( } /// UpdateAlert +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"update"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -121,6 +125,8 @@ pub async fn update_alert( } /// ListStreamAlerts +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"list"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -175,6 +181,8 @@ async fn list_stream_alerts(path: web::Path<(String, String)>, req: HttpRequest) } /// ListAlerts +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"list"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -279,6 +287,8 @@ async fn list_alerts(path: web::Path, req: HttpRequest) -> HttpResponse } /// GetAlertByName +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"get"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -328,6 +338,8 @@ async fn get_alert(path: web::Path<(String, String, String)>, req: HttpRequest) } /// DeleteAlert +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"delete"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -359,6 +371,8 @@ async fn delete_alert(path: web::Path<(String, String, String)>, req: HttpReques } /// EnableAlert +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"update"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -397,6 +411,8 @@ async fn enable_alert(path: web::Path<(String, String, String)>, req: HttpReques } /// TriggerAlert +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"update"}# #[deprecated] #[utoipa::path( context_path = "/api", diff --git a/src/handler/http/request/alerts/destinations.rs b/src/handler/http/request/alerts/destinations.rs index 3d909f8c95..5aab116ba5 100644 --- a/src/handler/http/request/alerts/destinations.rs +++ b/src/handler/http/request/alerts/destinations.rs @@ -36,6 +36,8 @@ impl From for HttpResponse { } /// CreateDestination +/// +/// #{"ratelimit_module":"Destinations", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Alerts", @@ -70,6 +72,8 @@ pub async fn save_destination( } /// UpdateDestination +/// +/// #{"ratelimit_module":"Destinations", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Alerts", @@ -104,6 +108,8 @@ pub async fn update_destination( } /// GetDestination +/// +/// #{"ratelimit_module":"Destinations", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Alerts", @@ -130,6 +136,8 @@ async fn get_destination(path: web::Path<(String, String)>) -> Result for HttpResponse { } /// CreateAlert +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Alerts", @@ -124,6 +126,8 @@ pub async fn create_alert( } /// GetAlert +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Alerts", @@ -159,6 +163,8 @@ async fn get_alert(path: web::Path<(String, Ksuid)>) -> HttpResponse { } /// UpdateAlert +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Alerts", @@ -197,6 +203,8 @@ pub async fn update_alert( } /// DeleteAlert +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"delete"}# #[utoipa::path( context_path = "/api", tag = "Alerts", @@ -225,6 +233,8 @@ async fn delete_alert(path: web::Path<(String, Ksuid)>) -> HttpResponse { } /// ListAlerts +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Alerts", @@ -284,6 +294,8 @@ async fn list_alerts(path: web::Path, req: HttpRequest) -> HttpResponse } /// EnableAlert +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Alerts", @@ -323,6 +335,8 @@ async fn enable_alert(path: web::Path<(String, Ksuid)>, req: HttpRequest) -> Htt } /// TriggerAlert +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Alerts", @@ -352,6 +366,8 @@ async fn trigger_alert(path: web::Path<(String, Ksuid)>) -> HttpResponse { } /// MoveAlerts +/// +/// #{"ratelimit_module":"Alerts", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Alerts", diff --git a/src/handler/http/request/alerts/templates.rs b/src/handler/http/request/alerts/templates.rs index 67370aeba3..ca7cbaf0af 100644 --- a/src/handler/http/request/alerts/templates.rs +++ b/src/handler/http/request/alerts/templates.rs @@ -39,6 +39,8 @@ impl From for HttpResponse { } /// CreateTemplate +/// +/// #{"ratelimit_module":"Templates", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Templates", @@ -69,6 +71,8 @@ pub async fn save_template( } /// UpdateTemplate +/// +/// #{"ratelimit_module":"Templates", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Templates", @@ -100,6 +104,8 @@ pub async fn update_template( } /// GetTemplateByName +/// +/// #{"ratelimit_module":"Templates", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Templates", @@ -126,6 +132,8 @@ async fn get_template(path: web::Path<(String, String)>) -> Result, _req: HttpRequest) -> Result, @@ -36,6 +55,22 @@ pub async fn create_role( } #[cfg(not(feature = "enterprise"))] +#[utoipa::path( + context_path = "/api", + tag = "Roles", + operation_id = "CreateRoles", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ), + request_body(content = UserRoleRequest, description = "UserRoleRequest", content_type = "application/json"), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[post("/{org_id}/roles")] pub async fn create_role( _org_id: web::Path, @@ -45,6 +80,25 @@ pub async fn create_role( } #[cfg(feature = "enterprise")] +/// DeleteRole +/// +/// #{"ratelimit_module":"Roles", "ratelimit_module_operation":"delete"}# +#[utoipa::path( + context_path = "/api", + tag = "Roles", + operation_id = "DeleteRole", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("role_id" = String, Path, description = "Role Id"), + ), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[delete("/{org_id}/roles/{role_id}")] pub async fn delete_role(path: web::Path<(String, String)>) -> Result { let (org_id, role_name) = path.into_inner(); @@ -56,12 +110,46 @@ pub async fn delete_role(path: web::Path<(String, String)>) -> Result) -> Result { Ok(HttpResponse::Forbidden().json("Not Supported")) } #[cfg(feature = "enterprise")] +/// ListRoles +/// +/// #{"ratelimit_module":"Roles", "ratelimit_module_operation":"list"}# +#[utoipa::path( + context_path = "/api", + tag = "Roles", + operation_id = "ListRoles", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ), + responses( + (status = 200, description = "Success", content_type = "application/json", body = Vec), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[get("/{org_id}/roles")] pub async fn get_roles(org_id: web::Path, req: HttpRequest) -> Result { let org_id = org_id.into_inner(); @@ -109,6 +197,21 @@ pub async fn get_roles(org_id: web::Path, req: HttpRequest) -> Result), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[get("/{org_id}/roles")] pub async fn get_roles( _org_id: web::Path, @@ -118,6 +221,26 @@ pub async fn get_roles( } #[cfg(feature = "enterprise")] +/// UpdateRoles +/// +/// #{"ratelimit_module":"Roles", "ratelimit_module_operation":"update"}# +#[utoipa::path( + context_path = "/api", + tag = "Roles", + operation_id = "UpdateRoles", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("role_id" = String, Path, description = "Role Id"), + ), + request_body(content = RoleRequest, description = "RoleRequest", content_type = "application/json"), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[put("/{org_id}/roles/{role_id}")] pub async fn update_role( path: web::Path<(String, String)>, @@ -142,6 +265,23 @@ pub async fn update_role( } #[cfg(not(feature = "enterprise"))] +#[utoipa::path( + context_path = "/api", + tag = "Roles", + operation_id = "UpdateRoles", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("role_id" = String, Path, description = "Role Id"), + ), + request_body(content = RoleRequest, description = "RoleRequest", content_type = "application/json"), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[post("/{org_id}/roles/{role_id}")] pub async fn update_role( _path: web::Path<(String, String)>, @@ -151,6 +291,26 @@ pub async fn update_role( } #[cfg(feature = "enterprise")] +/// GetResourcePermission +/// +/// #{"ratelimit_module":"Roles", "ratelimit_module_operation":"get"}# +#[utoipa::path( + context_path = "/api", + tag = "Roles", + operation_id = "GetResourcePermission", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("role_id" = String, Path, description = "Role Id"), + ("resource" = String, Path, description = "resource"), + ), + responses( + (status = 200, description = "Success", content_type = "application/json", body = Vec), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[get("/{org_id}/roles/{role_id}/permissions/{resource}")] pub async fn get_role_permissions( path: web::Path<(String, String, String)>, @@ -163,14 +323,50 @@ pub async fn get_role_permissions( } #[cfg(not(feature = "enterprise"))] +#[utoipa::path( + context_path = "/api", + tag = "Roles", + operation_id = "GetResourcePermission", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("role_id" = String, Path, description = "Role Id"), + ("resource" = String, Path, description = "resource"), + ), + responses( + (status = 200, description = "Success", content_type = "application/json", body = Vec), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[get("/{org_id}/roles/{role_id}/permissions/{resource}")] pub async fn get_role_permissions( - _path: web::Path<(String, String)>, + _path: web::Path<(String, String, String)>, ) -> Result { Ok(HttpResponse::Forbidden().json("Not Supported")) } #[cfg(feature = "enterprise")] +/// GetRoleUsers +/// +/// #{"ratelimit_module":"Roles", "ratelimit_module_operation":"list"}# +#[utoipa::path( + context_path = "/api", + tag = "Roles", + operation_id = "GetRoleUsers", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("role_id" = String, Path, description = "Role Id"), + ), + responses( + (status = 200, description = "Success", content_type = "application/json", body = Vec), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[get("/{org_id}/roles/{role_id}/users")] pub async fn get_users_with_role(path: web::Path<(String, String)>) -> Result { let (org_id, role_id) = path.into_inner(); @@ -181,12 +377,47 @@ pub async fn get_users_with_role(path: web::Path<(String, String)>) -> Result), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[get("/{org_id}/roles/{role_id}/users")] pub async fn get_users_with_role(_org_id: web::Path) -> Result { Ok(HttpResponse::Forbidden().json("Not Supported")) } #[cfg(feature = "enterprise")] +/// CreateGroup +/// +/// #{"ratelimit_module":"Groups", "ratelimit_module_operation":"create"}# +#[utoipa::path( + context_path = "/api", + tag = "Groups", + operation_id = "CreateGroup", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ), + request_body(content = UserGroup, description = "UserGroup", content_type = "application/json"), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[post("/{org_id}/groups")] pub async fn create_group( org_id: web::Path, @@ -208,6 +439,22 @@ pub async fn create_group( } #[cfg(not(feature = "enterprise"))] +#[utoipa::path( + context_path = "/api", + tag = "Groups", + operation_id = "CreateGroup", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ), + request_body(content = UserGroup, description = "UserGroup", content_type = "application/json"), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[post("/{org_id}/groups")] pub async fn create_group( _org_id: web::Path, @@ -217,6 +464,26 @@ pub async fn create_group( } #[cfg(feature = "enterprise")] +/// UpdateGroup +/// +/// #{"ratelimit_module":"Groups", "ratelimit_module_operation":"update"}# +#[utoipa::path( + context_path = "/api", + tag = "Groups", + operation_id = "UpdateGroup", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("group_name" = String, Path, description = "Group name"), + ), + request_body(content = UserGroupRequest, description = "UserGroupRequest", content_type = "application/json"), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[put("/{org_id}/groups/{group_name}")] pub async fn update_group( path: web::Path<(String, String)>, @@ -241,6 +508,23 @@ pub async fn update_group( } #[cfg(not(feature = "enterprise"))] +#[utoipa::path( + context_path = "/api", + tag = "Groups", + operation_id = "UpdateGroup", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("group_name" = String, Path, description = "Group name"), + ), + request_body(content = UserGroupRequest, description = "UserGroupRequest", content_type = "application/json"), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[put("/{org_id}/groups/{group_name}")] pub async fn update_group( _org_id: web::Path, @@ -250,6 +534,24 @@ pub async fn update_group( } #[cfg(feature = "enterprise")] +/// ListGroups +/// +/// #{"ratelimit_module":"Groups", "ratelimit_module_operation":"list"}# +#[utoipa::path( + context_path = "/api", + tag = "Groups", + operation_id = "ListGroups", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ), + responses( + (status = 200, description = "Success", content_type = "application/json", body = Vec), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[get("/{org_id}/groups")] pub async fn get_groups(path: web::Path, req: HttpRequest) -> Result { let org_id = path.into_inner(); @@ -298,12 +600,45 @@ pub async fn get_groups(path: web::Path, req: HttpRequest) -> Result), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[get("/{org_id}/groups")] pub async fn get_groups(_path: web::Path) -> Result { Ok(HttpResponse::Forbidden().json("Not Supported")) } #[cfg(feature = "enterprise")] +/// GetGroup +/// +/// #{"ratelimit_module":"Groups", "ratelimit_module_operation":"get"}# +#[utoipa::path( + context_path = "/api", + tag = "Groups", + operation_id = "GetGroup", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ), + responses( + (status = 200, description = "Success", content_type = "application/json", body = UserGroup), + (status = 500, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] #[get("/{org_id}/groups/{group_name}")] pub async fn get_group_details(path: web::Path<(String, String)>) -> Result { let (org_id, group_name) = path.into_inner(); @@ -315,6 +650,21 @@ pub async fn get_group_details(path: web::Path<(String, String)>) -> Result) -> Result { Ok(HttpResponse::Forbidden().json("Not Supported")) @@ -337,6 +687,25 @@ pub async fn get_resources(_org_id: web::Path) -> Result) -> Result { let (org_id, group_name) = path.into_inner(); @@ -348,6 +717,22 @@ pub async fn delete_group(path: web::Path<(String, String)>) -> Result) -> Result { Ok(HttpResponse::Forbidden().json("Not Supported")) diff --git a/src/handler/http/request/clusters/mod.rs b/src/handler/http/request/clusters/mod.rs index fcddd13657..e8bc35e6aa 100644 --- a/src/handler/http/request/clusters/mod.rs +++ b/src/handler/http/request/clusters/mod.rs @@ -24,6 +24,8 @@ use { }; /// ListClusters +/// +/// #{"ratelimit_module":"Clusters", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Clusters", diff --git a/src/handler/http/request/dashboards/mod.rs b/src/handler/http/request/dashboards/mod.rs index 6c0244935e..7d10b81ef5 100644 --- a/src/handler/http/request/dashboards/mod.rs +++ b/src/handler/http/request/dashboards/mod.rs @@ -68,6 +68,8 @@ impl From for HttpResponse { } /// CreateDashboard +/// +/// #{"ratelimit_module":"Dashboards", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Dashboards", @@ -117,6 +119,8 @@ pub async fn create_dashboard( } /// UpdateDashboard +/// +/// #{"ratelimit_module":"Dashboards", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Dashboards", @@ -169,6 +173,8 @@ async fn update_dashboard( } /// ListDashboards +/// +/// #{"ratelimit_module":"Dashboards", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Dashboards", @@ -202,6 +208,8 @@ async fn list_dashboards(org_id: web::Path, req: HttpRequest) -> impl Re } /// GetDashboard +/// +/// #{"ratelimit_module":"Dashboards", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Dashboards", @@ -230,6 +238,8 @@ async fn get_dashboard(path: web::Path<(String, String)>) -> impl Responder { } /// DeleteDashboard +/// +/// #{"ratelimit_module":"Dashboards", "ratelimit_module_operation":"delete"}# #[utoipa::path( context_path = "/api", tag = "Dashboards", @@ -260,6 +270,8 @@ async fn delete_dashboard(path: web::Path<(String, String)>) -> impl Responder { } /// MoveDashboard +/// +/// #{"ratelimit_module":"Dashboards", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Dashboards", @@ -299,6 +311,9 @@ async fn move_dashboard( } } +/// MoveDashboards +/// +/// #{"ratelimit_module":"Dashboards", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Dashboards", diff --git a/src/handler/http/request/dashboards/reports.rs b/src/handler/http/request/dashboards/reports.rs index 32d689a703..5ceac5dc3e 100644 --- a/src/handler/http/request/dashboards/reports.rs +++ b/src/handler/http/request/dashboards/reports.rs @@ -24,6 +24,8 @@ use crate::{ }; /// CreateReport +/// +/// #{"ratelimit_module":"Reports", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Reports", @@ -66,6 +68,8 @@ pub async fn create_report( } /// UpdateReport +/// +/// #{"ratelimit_module":"Reports", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Reports", @@ -103,6 +107,8 @@ async fn update_report( } /// ListReports +/// +/// #{"ratelimit_module":"Reports", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Reports", @@ -165,6 +171,8 @@ async fn list_reports(org_id: web::Path, req: HttpRequest) -> Result) -> Result) -> Result for HttpResponse { } /// CreateFolder +/// +/// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Folders", @@ -92,6 +94,8 @@ pub async fn create_folder( } /// UpdateFolder +/// +/// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Folders", @@ -131,6 +135,8 @@ pub async fn update_folder( } /// ListFolders +/// +/// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Folders", @@ -172,6 +178,8 @@ pub async fn list_folders( } /// GetFolder +/// +/// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Folders", @@ -202,6 +210,8 @@ pub async fn get_folder(path: web::Path<(String, FolderType, String)>) -> impl R } /// GetFolderByName +/// +/// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Folders", @@ -232,6 +242,8 @@ pub async fn get_folder_by_name(path: web::Path<(String, FolderType, String)>) - } /// DeleteFolder +/// +/// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"delete"}# #[utoipa::path( context_path = "/api", tag = "Folders", @@ -264,6 +276,8 @@ pub mod deprecated { use super::*; /// CreateFolder + /// + /// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"create"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -306,6 +320,8 @@ pub mod deprecated { } /// UpdateFolder + /// + /// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"update"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -346,6 +362,8 @@ pub mod deprecated { } /// ListFolders + /// + /// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"list"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -385,6 +403,8 @@ pub mod deprecated { } /// GetFolder + /// + /// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"get"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -416,6 +436,8 @@ pub mod deprecated { } /// GetFolderByName + /// + /// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"get"}# #[deprecated] #[utoipa::path( context_path = "/api", @@ -447,6 +469,8 @@ pub mod deprecated { } /// DeleteFolder + /// + /// #{"ratelimit_module":"Folders", "ratelimit_module_operation":"delete"}# #[deprecated] #[utoipa::path( context_path = "/api", diff --git a/src/handler/http/request/functions/mod.rs b/src/handler/http/request/functions/mod.rs index 3ebea04ae5..65fa49354e 100644 --- a/src/handler/http/request/functions/mod.rs +++ b/src/handler/http/request/functions/mod.rs @@ -19,6 +19,8 @@ use actix_web::{HttpRequest, HttpResponse, delete, get, post, put, web}; use config::meta::function::{TestVRLRequest, Transform}; /// CreateFunction +/// +/// #{"ratelimit_module":"Functions", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Functions", @@ -48,6 +50,8 @@ pub async fn save_function( } /// ListFunctions +/// +/// #{"ratelimit_module":"Functions", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Functions", @@ -96,6 +100,8 @@ async fn list_functions( } /// DeleteFunction +/// +/// #{"ratelimit_module":"Functions", "ratelimit_module_operation":"delete"}# #[utoipa::path( context_path = "/api", tag = "Functions", @@ -120,6 +126,8 @@ async fn delete_function(path: web::Path<(String, String)>) -> Result, - in_req: HttpRequest, body: web::Bytes, + in_req: HttpRequest, ) -> Result { #[cfg(feature = "enterprise")] { @@ -112,6 +115,9 @@ pub async fn save( } #[cfg(not(feature = "enterprise"))] { + drop(org_id); + drop(in_req); + drop(body); Ok(MetaHttpResponse::forbidden("not supported")) } } @@ -135,10 +141,7 @@ pub async fn save( tag = "Key" )] #[get("/{org_id}/cipher_keys/{key_name}")] -pub async fn get( - _req: HttpRequest, - path: web::Path<(String, String)>, -) -> Result { +pub async fn get(path: web::Path<(String, String)>) -> Result { #[cfg(feature = "enterprise")] { let (org_id, key_name) = path.into_inner(); @@ -169,7 +172,10 @@ pub async fn get( Ok(HttpResponse::Ok().json(res)) } #[cfg(not(feature = "enterprise"))] - Ok(MetaHttpResponse::forbidden("not supported")) + { + drop(path); + Ok(MetaHttpResponse::forbidden("not supported")) + } } /// list all keys for given org @@ -187,7 +193,7 @@ pub async fn get( tag = "Key" )] #[get("/{org_id}/cipher_keys")] -pub async fn list(_req: HttpRequest, path: web::Path) -> Result { +pub async fn list(path: web::Path) -> Result { #[cfg(feature = "enterprise")] { let org_id = path.into_inner(); @@ -217,7 +223,10 @@ pub async fn list(_req: HttpRequest, path: web::Path) -> Result) -> Result, -) -> Result { +pub async fn delete(path: web::Path<(String, String)>) -> Result { #[cfg(feature = "enterprise")] { let (org_id, key_name) = path.into_inner(); @@ -263,7 +269,10 @@ pub async fn delete( } } #[cfg(not(feature = "enterprise"))] - Ok(MetaHttpResponse::forbidden("not supported")) + { + drop(path); + Ok(MetaHttpResponse::forbidden("not supported")) + } } /// update the credentials for given key @@ -291,9 +300,9 @@ pub async fn delete( )] #[put("/{org_id}/cipher_keys/{name}")] pub async fn update( - in_req: HttpRequest, - body: web::Bytes, path: web::Path<(String, String)>, + body: web::Bytes, + in_req: HttpRequest, ) -> Result { #[cfg(feature = "enterprise")] { @@ -378,6 +387,9 @@ pub async fn update( } #[cfg(not(feature = "enterprise"))] { + drop(in_req); + drop(path); + drop(body); Ok(MetaHttpResponse::forbidden("not supported")) } } diff --git a/src/handler/http/request/kv/mod.rs b/src/handler/http/request/kv/mod.rs index 6bf4b52f59..65ecaffe38 100644 --- a/src/handler/http/request/kv/mod.rs +++ b/src/handler/http/request/kv/mod.rs @@ -21,6 +21,8 @@ use hashbrown::HashMap; use crate::{common::meta::http::HttpResponse as MetaHttpResponse, service::kv}; /// GetValue +/// +/// #{"ratelimit_module":"Key Values", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "KV", @@ -51,6 +53,8 @@ pub async fn get(path: web::Path<(String, String)>) -> Result) -> Result. -#[cfg(feature = "enterprise")] pub mod actions; pub mod alerts; pub mod authz; @@ -23,7 +22,6 @@ pub mod enrichment_table; #[allow(deprecated)] pub mod folders; pub mod functions; -#[cfg(feature = "enterprise")] pub mod keys; pub mod kv; pub mod logs; @@ -31,6 +29,7 @@ pub mod metrics; pub mod organization; pub mod pipeline; pub mod promql; +pub mod ratelimit; pub mod rum; #[cfg(feature = "enterprise")] pub mod script_server; diff --git a/src/handler/http/request/organization/org.rs b/src/handler/http/request/organization/org.rs index f443af846b..3f2003b41e 100644 --- a/src/handler/http/request/organization/org.rs +++ b/src/handler/http/request/organization/org.rs @@ -38,6 +38,8 @@ use crate::{ }; /// GetOrganizations +/// +/// #{"ratelimit_module":"Organizations", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Organizations", @@ -132,6 +134,8 @@ pub async fn organizations(user_email: UserEmail) -> Result } /// GetOrganizationSummary +/// +/// #{"ratelimit_module":"Summary", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Organizations", @@ -154,6 +158,8 @@ async fn org_summary(org_id: web::Path) -> Result { } /// GetIngestToken +/// +/// #{"ratelimit_module":"Ingestion Token", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Organizations", @@ -190,6 +196,8 @@ async fn get_user_passcode( } /// UpdateIngestToken +/// +/// #{"ratelimit_module":"Ingestion Token", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Organizations", @@ -226,6 +234,8 @@ async fn update_user_passcode( } /// GetRumIngestToken +/// +/// #{"ratelimit_module":"Rumtokens", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Organizations", @@ -262,6 +272,8 @@ async fn get_user_rumtoken( } /// UpdateRumIngestToken +/// +/// #{"ratelimit_module":"Rumtokens", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Organizations", @@ -298,6 +310,8 @@ async fn update_user_rumtoken( } /// CreateRumIngestToken +/// +/// #{"ratelimit_module":"Rumtokens", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Organizations", @@ -334,6 +348,8 @@ async fn create_user_rumtoken( } /// CreateOrganization +/// +/// #{"ratelimit_module":"Organizations", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Organizations", diff --git a/src/handler/http/request/organization/settings.rs b/src/handler/http/request/organization/settings.rs index ba3c722bba..f944179989 100644 --- a/src/handler/http/request/organization/settings.rs +++ b/src/handler/http/request/organization/settings.rs @@ -36,6 +36,8 @@ use crate::{ }; /// Organization specific settings +/// +/// #{"ratelimit_module":"Settings", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Organizations", @@ -112,6 +114,8 @@ async fn create( } /// Retrieve organization specific settings +/// +/// #{"ratelimit_module":"Settings", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Organizations", diff --git a/src/handler/http/request/pipeline.rs b/src/handler/http/request/pipeline.rs index a79eb8698b..84a9ea2fc4 100644 --- a/src/handler/http/request/pipeline.rs +++ b/src/handler/http/request/pipeline.rs @@ -36,6 +36,8 @@ impl From for HttpResponse { } /// CreatePipeline +/// +/// #{"ratelimit_module":"Pipeline", "ratelimit_module_operation":"create"}# #[utoipa::path( context_path = "/api", tag = "Pipeline", @@ -73,6 +75,8 @@ pub async fn save_pipeline( } /// ListPipelines +/// +/// #{"ratelimit_module":"Pipeline", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Pipelines", @@ -128,6 +132,8 @@ async fn list_pipelines( } /// GetStreamsWithPipeline +/// +/// #{"ratelimit_module":"Pipeline", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Pipelines", @@ -155,6 +161,8 @@ async fn list_streams_with_pipeline( } /// DeletePipeline +/// +/// #{"ratelimit_module":"Pipeline", "ratelimit_module_operation":"delete"}# #[utoipa::path( context_path = "/api", tag = "Pipelines", @@ -187,6 +195,8 @@ async fn delete_pipeline( } /// UpdatePipeline +/// +/// #{"ratelimit_module":"Pipeline", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Pipelines", @@ -219,6 +229,8 @@ pub async fn update_pipeline( } /// EnablePipeline +/// +/// #{"ratelimit_module":"Pipeline", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/api", tag = "Pipelines", diff --git a/src/handler/http/request/promql/mod.rs b/src/handler/http/request/promql/mod.rs index c88fca04e9..dc1a225337 100644 --- a/src/handler/http/request/promql/mod.rs +++ b/src/handler/http/request/promql/mod.rs @@ -69,6 +69,8 @@ pub async fn remote_write( } /// prometheus instant queries +/// +/// #{"ratelimit_module":"Metrics", "ratelimit_module_operation":"get"}# // refer: https://prometheus.io/docs/prometheus/latest/querying/api/#instant-queries #[utoipa::path( context_path = "/api", @@ -226,6 +228,8 @@ async fn query( } /// prometheus range queries +/// +/// #{"ratelimit_module":"Metrics", "ratelimit_module_operation":"get"}# // refer: https://prometheus.io/docs/prometheus/latest/querying/api/#range-queries #[utoipa::path( context_path = "/api", @@ -303,6 +307,8 @@ pub async fn query_range_post( } /// prometheus query exemplars +/// +/// #{"ratelimit_module":"Metrics", "ratelimit_module_operation":"get"}# // refer: https://prometheus.io/docs/prometheus/latest/querying/api/#querying-exemplars #[utoipa::path( context_path = "/api", @@ -521,6 +527,8 @@ async fn query_range( } /// prometheus query metric metadata +/// +/// #{"ratelimit_module":"Metrics", "ratelimit_module_operation":"get"}# // refer: https://prometheus.io/docs/prometheus/latest/querying/api/#querying-metric-metadata #[utoipa::path( context_path = "/api", @@ -581,6 +589,8 @@ pub async fn metadata( } /// prometheus finding series by label matchers +/// +/// #{"ratelimit_module":"Metrics", "ratelimit_module_operation":"get"}# // refer: https://prometheus.io/docs/prometheus/latest/querying/api/#finding-series-by-label-matchers #[utoipa::path( context_path = "/api", @@ -726,6 +736,8 @@ async fn series( } /// prometheus getting label names +/// +/// #{"ratelimit_module":"Metrics", "ratelimit_module_operation":"get"}# // refer: https://prometheus.io/docs/prometheus/latest/querying/api/#getting-label-names #[utoipa::path( context_path = "/api", @@ -822,6 +834,8 @@ async fn labels( } /// prometheus query label values +/// +/// #{"ratelimit_module":"Metrics", "ratelimit_module_operation":"get"}# // refer: https://prometheus.io/docs/prometheus/latest/querying/api/#querying-label-values #[utoipa::path( context_path = "/api", @@ -934,6 +948,8 @@ fn validate_metadata_params( } /// prometheus formatting query expressions +/// +/// #{"ratelimit_module":"Metrics", "ratelimit_module_operation":"get"}# // refer: https://prometheus.io/docs/prometheus/latest/querying/api/#formatting-query-expressions #[utoipa::path( context_path = "/api", diff --git a/src/handler/http/request/ratelimit/mod.rs b/src/handler/http/request/ratelimit/mod.rs new file mode 100644 index 0000000000..1f868caa49 --- /dev/null +++ b/src/handler/http/request/ratelimit/mod.rs @@ -0,0 +1,517 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::io::Error; + +use actix_web::{HttpResponse, get, put, web}; +use serde::Deserialize; +#[cfg(feature = "enterprise")] +use { + crate::{ + common::meta::http::HttpResponse as MetaHttpResponse, + service::{ratelimit, ratelimit::rule::RatelimitError}, + }, + actix_web::http, + config::meta::ratelimit::{RatelimitRule, RatelimitRuleUpdater}, + futures_util::StreamExt, + infra::table::ratelimit::RuleEntry, + o2_ratelimit::dataresource::{ + default_rules::{ + ApiGroupOperation, DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER, + DEFAULT_THRESHOLD_NO_LIMIT_FRONTEND, RATELIMIT_API_MAPPING, + get_ratelimit_global_default_api_info, has_group_name_and_operation, + }, + default_rules_provider::get_default_rules, + }, + std::collections::HashMap, +}; +#[cfg(feature = "enterprise")] +const QUOTA_PAGE_REQUIRED_ORG: &str = "_meta"; +#[cfg(feature = "enterprise")] +pub const QUOTA_PAGE_GLOBAL_RULES_ORG: &str = "global_rules"; + +#[cfg(feature = "enterprise")] +#[derive(Debug, Clone, Deserialize)] +struct QueryParams { + org_id: String, + update_type: Option, + user_role: Option, +} +#[cfg(feature = "enterprise")] +impl QueryParams { + fn get_org_id(&self) -> String { + match self.org_id.to_lowercase().as_str() { + QUOTA_PAGE_GLOBAL_RULES_ORG => DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER.to_string(), + _ => self.org_id.to_string(), + } + } + + fn get_update_type(&self) -> String { + self.update_type.clone().unwrap_or_default() + } +} + +#[cfg(not(feature = "enterprise"))] +#[derive(Debug, Clone, Deserialize)] +struct QueryParams(); + +#[cfg(feature = "enterprise")] +async fn validate_ratelimit_updater( + org_id: &str, + update_type: &str, + rules: &RatelimitRuleUpdater, +) -> Result<(), anyhow::Error> { + let global_default_rules = get_default_rules() + .await + .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; + + // module-level is org-level + let org_level_rules = infra::table::ratelimit::fetch_rules( + global_default_rules.clone(), + Some(org_id.to_string()), + Some(DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER.to_string()), + ) + .await + .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; + + let org_level_thresholds: HashMap<(String, ApiGroupOperation), i32> = org_level_rules + .into_iter() + .filter_map(|rule| { + let group_name = rule.api_group_name?; + let operation = ApiGroupOperation::try_from(rule.api_group_operation?.as_str()).ok()?; + Some(((group_name, operation), rule.threshold)) + }) + .collect(); + + for (api_group_name, operations) in rules.0.iter() { + for (operation, threshold) in operations.iter() { + // Validate operation format + let operation = ApiGroupOperation::try_from(operation.as_str()); + if operation.is_err() { + return Err(anyhow::anyhow!( + "rule operation_name is err, {:?}", + operation + )); + } + let operation = operation.unwrap(); + + // Validate group name and operation combination + if !has_group_name_and_operation(api_group_name, operation).await { + return Err(anyhow::anyhow!( + "rule module_name and operation_name is not found, module_name: {:?} operation_name: {:?}", + api_group_name, + operation, + )); + } + + // Validate threshold + if *threshold < DEFAULT_THRESHOLD_NO_LIMIT_FRONTEND { + return Err(anyhow::anyhow!( + "threshold must be greater than or equal to {}, got {}", + DEFAULT_THRESHOLD_NO_LIMIT_FRONTEND, + threshold + )); + } + + if update_type == "role" { + let compare_threshold = + org_level_thresholds.get(&(api_group_name.to_string(), operation)); + if compare_threshold.is_some() && *compare_threshold.unwrap() < *threshold { + let compare_threshold = compare_threshold.unwrap(); + return Err(anyhow::anyhow!( + "{}:{} threshold must be lower than or equal to {:?}, got {}, because module-level rule limit {}:{} is {:?}", + api_group_name, + operation.as_str(), + compare_threshold, + threshold, + api_group_name, + operation.as_str(), + compare_threshold, + )); + } + } + } + } + + Ok(()) +} + +#[cfg(feature = "enterprise")] +impl From for HttpResponse { + fn from(value: RatelimitError) -> Self { + match value { + RatelimitError::NotFound(_) => MetaHttpResponse::not_found(value), + error => MetaHttpResponse::bad_request(error), + } + } +} + +/// listApiModule +/// +/// #{"ratelimit_module":"Ratelimit", "ratelimit_module_operation":"list"}# +#[utoipa::path( + context_path = "/api", + tag = "Ratelimit", + operation_id = "listApiModule", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization Name"), + ), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 400, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] +#[get("/{org_id}/ratelimit/api_modules")] +pub async fn api_modules(path: web::Path) -> Result { + #[cfg(feature = "enterprise")] + { + let org_id = path.into_inner(); + if org_id != QUOTA_PAGE_REQUIRED_ORG { + return Ok(MetaHttpResponse::bad_request(format!( + "org_id: {} has no access", + org_id, + ))); + } + + if RATELIMIT_API_MAPPING.read().await.is_empty() { + let openapi_info = crate::handler::http::router::openapi::openapi_info().await; + o2_ratelimit::dataresource::default_rules::init_ratelimit_api_mapping(openapi_info) + .await; + } + + let info = get_ratelimit_global_default_api_info().await; + let all_group = info.get_all_groups(); + Ok(HttpResponse::Ok().json(all_group)) + } + + #[cfg(not(feature = "enterprise"))] + { + drop(path); + Ok(HttpResponse::Forbidden().json("Not Supported")) + } +} + +/// listModuleRatelimitRule +/// +/// #{"ratelimit_module":"Ratelimit", "ratelimit_module_operation":"list"}# +#[utoipa::path( + context_path = "/api", + tag = "Ratelimit", + operation_id = "listModuleRatelimitRule", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization Name"), + ("org_id" = String, Query, description = "Organization Name"), + ), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 400, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] +#[get("/{org_id}/ratelimit/module_list")] +pub async fn list_module_ratelimit( + path: web::Path, + query: web::Query, +) -> Result { + #[cfg(feature = "enterprise")] + { + let org_id = path.into_inner(); + if org_id != QUOTA_PAGE_REQUIRED_ORG { + return Ok(MetaHttpResponse::bad_request(format!( + "org_id: {} has no access", + org_id, + ))); + } + + let query = query.into_inner(); + let org_id = query.get_org_id(); + if RATELIMIT_API_MAPPING.read().await.is_empty() { + let openapi_info = crate::handler::http::router::openapi::openapi_info().await; + o2_ratelimit::dataresource::default_rules::init_ratelimit_api_mapping(openapi_info) + .await; + } + + let global_default_rules = get_default_rules() + .await + .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; + // module-level is org-level + let all_rules = infra::table::ratelimit::fetch_rules( + global_default_rules.clone(), + Some(org_id.clone()), + Some(DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER.to_string()), + ) + .await + .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; + + let info = get_ratelimit_global_default_api_info().await; + let api_group_info = info.api_groups(Some(org_id.as_str()), all_rules).await; + + Ok(HttpResponse::Ok().json(api_group_info)) + } + + #[cfg(not(feature = "enterprise"))] + { + drop(path); + let _ = query; + Ok(HttpResponse::Forbidden().json("Not Supported")) + } +} + +/// listRoleRatelimitRule +/// +/// #{"ratelimit_module":"Ratelimit", "ratelimit_module_operation":"list"}# +#[utoipa::path( + context_path = "/api", + tag = "Ratelimit", + operation_id = "listRoleRatelimitRule", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization Name"), + ("org_id" = String, Query, description = "Organization Name"), + ("user_role" = String, Query, description = "User Role Name"), + ), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 400, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] +#[get("/{org_id}/ratelimit/role_list")] +pub async fn list_role_ratelimit( + path: web::Path, + query: web::Query, +) -> Result { + #[cfg(feature = "enterprise")] + { + let org_id = path.into_inner(); + if org_id != QUOTA_PAGE_REQUIRED_ORG { + return Ok(MetaHttpResponse::bad_request(format!( + "org_id: {} has no access", + org_id, + ))); + } + + let query = query.into_inner(); + let org_id = query.get_org_id(); + let user_role = query.user_role.unwrap_or_default(); + if user_role.is_empty() { + return Ok(MetaHttpResponse::bad_request( + "user_role param is required".to_string(), + )); + } + + if RATELIMIT_API_MAPPING.read().await.is_empty() { + let openapi_info = crate::handler::http::router::openapi::openapi_info().await; + o2_ratelimit::dataresource::default_rules::init_ratelimit_api_mapping(openapi_info) + .await; + } + + let global_default_rules = get_default_rules().await.unwrap(); + + // module-level is org-level + let org_level_rules = infra::table::ratelimit::fetch_rules( + global_default_rules.clone(), + Some(org_id.clone()), + Some(DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER.to_string()), + ) + .await + .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; + + let role_level_rules = infra::table::ratelimit::fetch_rules( + org_level_rules, + Some(org_id.clone()), + Some(user_role), + ) + .await + .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; + + let info = get_ratelimit_global_default_api_info().await; + let api_group_info = info + .api_groups(Some(org_id.as_str()), role_level_rules) + .await; + + Ok(HttpResponse::Ok().json(api_group_info)) + } + + #[cfg(not(feature = "enterprise"))] + { + drop(path); + let _ = query; + Ok(HttpResponse::Forbidden().json("Not Supported")) + } +} + +/// UpdateRatelimitRule +/// +/// #{"ratelimit_module":"Ratelimit", "ratelimit_module_operation":"update"}# +#[utoipa::path( + context_path = "/api", + tag = "Ratelimit", + operation_id = "UpdateRatelimitRule", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("org_id" = String, Query, description = "Organization name"), + ("update_type" = String, Query, description = "Update Type"), + ("user_role" = Option, Query, description = "UserRole name"), + ), + request_body(content = String, description = "json array", content_type = "application/json", example = json!({"Key Values": {"list": 0,"create": 0,"delete": 0,"get": 0},"Service Accounts": {"update": 0,"list": 0,"create": 0,"delete": 0,"get": 0}})), + responses( + (status = 200, description = "Success", content_type = "application/json", body = HttpResponse), + (status = 400, description = "Failure", content_type = "application/json", body = HttpResponse), + ) +)] +#[put("/{org_id}/ratelimit/update")] +pub async fn update_ratelimit( + path: web::Path, + query: web::Query, + #[cfg_attr(not(feature = "enterprise"), allow(unused_mut))] mut payload: web::Payload, +) -> Result { + #[cfg(feature = "enterprise")] + { + let org_id = path.into_inner(); + if org_id != QUOTA_PAGE_REQUIRED_ORG { + return Ok(MetaHttpResponse::bad_request(format!( + "org_id: {} has no access", + org_id, + ))); + } + + let mut bytes = web::BytesMut::new(); + while let Some(item) = payload.next().await { + bytes.extend_from_slice( + &item.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?, + ); + } + + let updater = match parse_and_validate_ratelimit_payload(bytes).await { + Ok(updater) => updater, + Err(e) => { + return Ok(MetaHttpResponse::bad_request(e.to_string())); + } + }; + + let query = query.into_inner(); + let org_id = query.get_org_id(); + + if RATELIMIT_API_MAPPING.read().await.is_empty() { + let openapi_info = crate::handler::http::router::openapi::openapi_info().await; + o2_ratelimit::dataresource::default_rules::init_ratelimit_api_mapping(openapi_info) + .await; + } + + if let Err(e) = + validate_ratelimit_updater(org_id.as_str(), &query.get_update_type(), &updater).await + { + return Ok(MetaHttpResponse::bad_request(format!( + "validate ratelimit updater error: {}", + e, + ))); + } + + let user_role = match query.get_update_type().as_str() { + "module" => DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER.to_string(), + "role" => { + if query.user_role.is_none() || query.user_role.eq(&Some("".to_string())) { + return Ok(MetaHttpResponse::bad_request( + "user_role param is required when update_type=role".to_string(), + )); + } else { + query.user_role.clone().unwrap() + } + } + _ => { + return Ok(MetaHttpResponse::bad_request(format!( + "update_type is incorrect: {}, only support module or role", + query.get_update_type(), + ))); + } + }; + + let rules = RatelimitRule::from_updater(org_id.as_str(), user_role.as_str(), &updater); + log::debug!("RatelimitRule::from_updater rules: {:?}", rules); + match ratelimit::rule::update(RuleEntry::UpsertBatch(rules)).await { + Ok(()) => Ok(HttpResponse::Ok().json(MetaHttpResponse::message( + http::StatusCode::OK.into(), + "Ratelimit rule updated successfully".to_string(), + ))), + Err(e) => Ok(e.into()), + } + } + + #[cfg(not(feature = "enterprise"))] + { + drop(path); + let _ = query; + drop(payload); + Ok(HttpResponse::Forbidden().json("Not Supported")) + } +} + +#[cfg(feature = "enterprise")] +async fn parse_and_validate_ratelimit_payload( + bytes: web::BytesMut, +) -> Result { + let value: serde_json::Value = + serde_json::from_slice(&bytes).map_err(|e| format!("Failed to parse JSON: {}", e))?; + + let obj = value + .as_object() + .ok_or_else(|| "Root must be an object".to_string())?; + + let mut result = HashMap::new(); + + for (outer_key, outer_value) in obj { + let inner_obj = outer_value + .as_object() + .ok_or_else(|| format!("Value for key '{}' must be an object", outer_key))?; + + let mut inner_map = HashMap::new(); + + for (inner_key, inner_value) in inner_obj { + let number = inner_value.as_i64().ok_or_else(|| { + format!( + "Value for key '{}' in group '{}' must be an i32", + inner_key, outer_key + ) + })?; + + if number > i32::MAX as i64 || number < 0 { + return Err(format!( + "The Rate limit value of {} for the key '{}' in group '{}' exceeds the allowed integer range of [{}, {}]", + number, + inner_key, + outer_key, + 0, + i32::MAX + )); + } + + inner_map.insert(inner_key.clone(), number as i32); + } + + result.insert(outer_key.clone(), inner_map); + } + + Ok(RatelimitRuleUpdater(result)) +} diff --git a/src/handler/http/request/search/mod.rs b/src/handler/http/request/search/mod.rs index 117b04cc49..4bc252d2cb 100644 --- a/src/handler/http/request/search/mod.rs +++ b/src/handler/http/request/search/mod.rs @@ -62,12 +62,9 @@ use crate::{ pub(crate) mod around; pub(crate) mod error_utils; pub mod multi_streams; -#[cfg(feature = "enterprise")] pub mod query_manager; pub mod saved_view; -#[cfg(feature = "enterprise")] pub mod search_job; -#[cfg(feature = "enterprise")] pub(crate) mod utils; async fn can_use_distinct_stream( @@ -134,6 +131,8 @@ async fn can_use_distinct_stream( } /// SearchStreamData +/// +/// #{"ratelimit_module":"Search", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Search", @@ -367,6 +366,8 @@ pub async fn search( } /// SearchAround +/// +/// #{"ratelimit_module":"Search", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Search", @@ -461,6 +462,8 @@ pub async fn around_v1( } /// SearchAroundV2 +/// +/// #{"ratelimit_module":"Search", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Search", @@ -565,6 +568,8 @@ pub async fn around_v2( } /// SearchTopNValues +/// +/// #{"ratelimit_module":"Search", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Search", @@ -971,6 +976,8 @@ async fn values_v1( } /// SearchStreamPartition +/// +/// #{"ratelimit_module":"Search", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Search", @@ -1077,6 +1084,8 @@ pub async fn search_partition( } /// Search History +/// +/// #{"ratelimit_module":"Search", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Search", diff --git a/src/handler/http/request/search/query_manager.rs b/src/handler/http/request/search/query_manager.rs index ce193daa52..2b85ab430d 100644 --- a/src/handler/http/request/search/query_manager.rs +++ b/src/handler/http/request/search/query_manager.rs @@ -16,15 +16,26 @@ use std::io::Error; use actix_web::{HttpResponse, delete, get, put, web}; -use o2_enterprise::enterprise::common::infra::config::get_config as get_o2_config; - -use crate::common::meta::http::HttpResponse as MetaHttpResponse; +#[cfg(feature = "enterprise")] +use { + crate::common::meta::http::HttpResponse as MetaHttpResponse, + o2_enterprise::enterprise::common::infra::config::get_config as get_o2_config, +}; #[delete("/{org_id}/query_manager/{trace_id}")] pub async fn cancel_query(params: web::Path<(String, String)>) -> Result { - let (org_id, trace_id) = params.into_inner(); - let trace_ids = trace_id.split(',').collect::>(); - cancel_query_inner(&org_id, &trace_ids).await + #[cfg(feature = "enterprise")] + { + let (org_id, trace_id) = params.into_inner(); + let trace_ids = trace_id.split(',').collect::>(); + cancel_query_inner(&org_id, &trace_ids).await + } + + #[cfg(not(feature = "enterprise"))] + { + drop(params); + Ok(HttpResponse::Forbidden().json("Not Supported")) + } } #[put("/{org_id}/query_manager/cancel")] @@ -32,44 +43,73 @@ pub async fn cancel_multiple_query( params: web::Path, body: web::Bytes, ) -> Result { - let org_id = params.into_inner(); - let trace_ids: Vec = match config::utils::json::from_slice(&body) { - Ok(v) => v, - Err(e) => { - return Ok(MetaHttpResponse::bad_request(e)); - } - }; - let trace_ids = trace_ids.iter().map(|s| s.as_str()).collect::>(); - cancel_query_inner(&org_id, &trace_ids).await + #[cfg(feature = "enterprise")] + { + let org_id = params.into_inner(); + let trace_ids: Vec = match config::utils::json::from_slice(&body) { + Ok(v) => v, + Err(e) => { + return Ok(MetaHttpResponse::bad_request(e)); + } + }; + let trace_ids = trace_ids.iter().map(|s| s.as_str()).collect::>(); + cancel_query_inner(&org_id, &trace_ids).await + } + + #[cfg(not(feature = "enterprise"))] + { + drop(params); + drop(body); + Ok(HttpResponse::Forbidden().json("Not Supported")) + } } #[get("/{org_id}/query_manager/status")] pub async fn query_status(_params: web::Path) -> Result { - let res = crate::service::search::query_status().await; - match res { - Ok(query_status) => Ok(HttpResponse::Ok().json(query_status)), - Err(e) => Ok(MetaHttpResponse::bad_request(e)), + #[cfg(feature = "enterprise")] + { + let res = crate::service::search::query_status().await; + match res { + Ok(query_status) => Ok(HttpResponse::Ok().json(query_status)), + Err(e) => Ok(MetaHttpResponse::bad_request(e)), + } + } + + #[cfg(not(feature = "enterprise"))] + { + Ok(HttpResponse::Forbidden().json("Not Supported")) } } pub async fn cancel_query_inner(org_id: &str, trace_ids: &[&str]) -> Result { - if trace_ids.is_empty() { - return Ok(HttpResponse::BadRequest().json("Invalid trace_id")); - } - let mut res = Vec::with_capacity(trace_ids.len()); - for trace_id in trace_ids { - if trace_id.is_empty() { - continue; + #[cfg(feature = "enterprise")] + { + if trace_ids.is_empty() { + return Ok(HttpResponse::BadRequest().json("Invalid trace_id")); } - let ret = if get_o2_config().super_cluster.enabled { - o2_enterprise::enterprise::super_cluster::search::cancel_query(org_id, trace_id).await - } else { - crate::service::search::cancel_query(org_id, trace_id).await - }; - match ret { - Ok(status) => res.push(status), - Err(e) => return Ok(MetaHttpResponse::bad_request(e)), + let mut res = Vec::with_capacity(trace_ids.len()); + for trace_id in trace_ids { + if trace_id.is_empty() { + continue; + } + let ret = if get_o2_config().super_cluster.enabled { + o2_enterprise::enterprise::super_cluster::search::cancel_query(org_id, trace_id) + .await + } else { + crate::service::search::cancel_query(org_id, trace_id).await + }; + match ret { + Ok(status) => res.push(status), + Err(e) => return Ok(MetaHttpResponse::bad_request(e)), + } } + Ok(HttpResponse::Ok().json(res)) + } + + #[cfg(not(feature = "enterprise"))] + { + let _ = org_id; + let _ = trace_ids; + Ok(HttpResponse::Forbidden().json("Not Supported")) } - Ok(HttpResponse::Ok().json(res)) } diff --git a/src/handler/http/request/search/saved_view.rs b/src/handler/http/request/search/saved_view.rs index cd571f88e8..2bd285bad9 100644 --- a/src/handler/http/request/search/saved_view.rs +++ b/src/handler/http/request/search/saved_view.rs @@ -31,10 +31,9 @@ use crate::{ service::db::saved_view, }; -// GetSavedView -// -// Retrieve a single saved view associated with this org. -// +/// GetSavedView - Retrieve a single saved view associated with this org. +/// +/// #{"ratelimit_module":"Saved Views", "ratelimit_module_operation":"get"}# #[utoipa::path( context_path = "/api", tag = "Saved Views", @@ -70,10 +69,9 @@ pub async fn get_view(path: web::Path<(String, String)>) -> Result) -> Result { } } -// DeleteSavedViews -// -// Delete a view associated with this given org. -// +/// DeleteSavedViews - Delete a view associated with this given org. +/// +/// #{"ratelimit_module":"Saved Views", "ratelimit_module_operation":"delete"}# #[utoipa::path( context_path = "/api", tag = "Saved Views", @@ -143,10 +140,9 @@ pub async fn delete_view(path: web::Path<(String, String)>) -> Result, in_req: HttpRequest, body: web::Bytes, ) -> Result { - let cfg = get_config(); + #[cfg(feature = "enterprise")] + { + let cfg = get_config(); - let org_id = org_id.into_inner(); - let http_span = if cfg.common.tracing_search_enabled || cfg.common.tracing_enabled { - tracing::info_span!("/api/{org_id}/_search", org_id = org_id.clone()) - } else { - Span::none() - }; + let org_id = org_id.into_inner(); + let http_span = if cfg.common.tracing_search_enabled || cfg.common.tracing_enabled { + tracing::info_span!("/api/{org_id}/_search", org_id = org_id.clone()) + } else { + Span::none() + }; - let trace_id = get_or_create_trace_id(in_req.headers(), &http_span); - let user_id = in_req - .headers() - .get("user_id") - .and_then(|v| v.to_str().ok()) - .unwrap_or("") - .to_string(); + let trace_id = get_or_create_trace_id(in_req.headers(), &http_span); + let user_id = in_req + .headers() + .get("user_id") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); - let query = web::Query::>::from_query(in_req.query_string()).unwrap(); - let stream_type = get_stream_type_from_request(&query).unwrap_or_default(); + let query = match web::Query::>::from_query(in_req.query_string()) { + Ok(q) => q, + Err(e) => return Ok(MetaHttpResponse::bad_request(e)), + }; + let stream_type = get_stream_type_from_request(&query).unwrap_or_default(); - let use_cache = cfg.common.result_cache_enabled && get_use_cache_from_request(&query); - // handle encoding for query and aggs - let mut req: config::meta::search::Request = match json::from_slice(&body) { - Ok(v) => v, - Err(e) => return Ok(MetaHttpResponse::bad_request(e)), - }; - if let Err(e) = req.decode() { - return Ok(MetaHttpResponse::bad_request(e)); - } - req.use_cache = Some(use_cache); - - // update timeout - if req.timeout == 0 { - req.timeout = cfg.limit.search_job_timeout; - } - - // set search event type - req.search_type = Some(SearchEventType::SearchJob); - if req.search_event_context.is_none() { - req.search_event_context = req - .search_type - .as_ref() - .and_then(|event_type| get_search_event_context_from_request(event_type, &query)); - } - - // get stream name - let stream_names = match resolve_stream_names(&req.query.sql) { - Ok(v) => v.clone(), - Err(e) => { - return Ok( - HttpResponse::InternalServerError().json(meta::http::HttpResponse::error( - StatusCode::INTERNAL_SERVER_ERROR.into(), - e.to_string(), - )), - ); + let use_cache = cfg.common.result_cache_enabled && get_use_cache_from_request(&query); + // handle encoding for query and aggs + let mut req: config::meta::search::Request = match json::from_slice(&body) { + Ok(v) => v, + Err(e) => return Ok(MetaHttpResponse::bad_request(e)), + }; + if let Err(e) = req.decode() { + return Ok(MetaHttpResponse::bad_request(e)); } - }; + req.use_cache = Some(use_cache); - // Check permissions on stream - for stream_name in stream_names.iter() { - if let Some(res) = - check_stream_permissions(stream_name, &org_id, &user_id, &stream_type).await - { - return Ok(res); + // update timeout + if req.timeout == 0 { + req.timeout = cfg.limit.search_job_timeout; + } + + // set search event type + req.search_type = Some(SearchEventType::SearchJob); + if req.search_event_context.is_none() { + req.search_event_context = req + .search_type + .as_ref() + .and_then(|event_type| get_search_event_context_from_request(event_type, &query)); + } + + // get stream name + let stream_names = match resolve_stream_names(&req.query.sql) { + Ok(v) => v.clone(), + Err(e) => { + return Ok(HttpResponse::InternalServerError().json( + meta::http::HttpResponse::error( + StatusCode::INTERNAL_SERVER_ERROR.into(), + e.to_string(), + ), + )); + } + }; + + // Check permissions on stream + for stream_name in stream_names.iter() { + if let Some(res) = + check_stream_permissions(stream_name, &org_id, &user_id, &stream_type).await + { + return Ok(res); + } + } + + // add stream_names for rbac + let stream_names = json::to_string(&stream_names).unwrap(); + + // submit query to db + let res = submit( + &trace_id, + &org_id, + &user_id, + stream_type.as_str(), + &stream_names, + &json::to_string(&req).unwrap(), + req.query.start_time, + req.query.end_time, + ) + .await; + + match res { + Ok(job_id) => Ok(MetaHttpResponse::ok(format!( + "[Job_Id: {job_id}] Search Job submitted successfully." + ))), + Err(err) => { + log::error!("[trace_id {trace_id}] sumbit query error: {}", err); + Ok(MetaHttpResponse::internal_error(err.to_string())) + } } } - // add stream_names for rbac - let stream_names = json::to_string(&stream_names).unwrap(); - - // submit query to db - let res = submit( - &trace_id, - &org_id, - &user_id, - stream_type.as_str(), - &stream_names, - &json::to_string(&req).unwrap(), - req.query.start_time, - req.query.end_time, - ) - .await; - - match res { - Ok(job_id) => Ok(MetaHttpResponse::ok(format!( - "[Job_Id: {job_id}] Search Job submitted successfully." - ))), - Err(err) => { - log::error!("[trace_id {trace_id}] sumbit query error: {}", err); - Ok(MetaHttpResponse::internal_error(err.to_string())) - } + #[cfg(not(feature = "enterprise"))] + { + drop(org_id); + drop(in_req); + drop(body); + Ok(HttpResponse::Forbidden().json("Not Supported")) } } // 2. status_all +/// ListSearchJobs +/// +/// #{"ratelimit_module":"Search Jobs", "ratelimit_module_operation":"list"}# +#[utoipa::path( + context_path = "/api", + tag = "Search Jobs", + operation_id = "ListSearchJobs", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name") + ), + responses( + (status = 200, description = "List of search jobs", body = Vec, example = json!([{ + "id": "abc123", + "trace_id": "xyz789", + "org_id": "default", + "user_id": "user1", + "stream_type": "logs", + "stream_names": "[\"stream1\"]", + "payload": "...", + "start_time": 1675182660872049i64, + "end_time": 1675185660872049i64, + "created_at": 1675182660872049i64, + "updated_at": 1675182660872049i64, + "status": 1, + "cluster": "cluster1", + "result_path": "/path/to/result" + }])), + (status = 400, description = "Bad Request", body = MetaHttpResponse) + ) +)] #[get("/{org_id}/search_jobs")] pub async fn list_status(org_id: web::Path) -> Result { - let res = list_status_by_org_id(&org_id).await; - let res = match res { - Ok(res) => res, - Err(e) => return Ok(MetaHttpResponse::bad_request(e)), - }; + #[cfg(feature = "enterprise")] + { + let res = list_status_by_org_id(&org_id).await; + let res = match res { + Ok(res) => res, + Err(e) => return Ok(MetaHttpResponse::bad_request(e)), + }; - Ok(HttpResponse::Ok().json(res)) + Ok(HttpResponse::Ok().json(res)) + } + + #[cfg(not(feature = "enterprise"))] + { + drop(org_id); + Ok(HttpResponse::Forbidden().json("Not Supported")) + } } // 3. status +/// GetSearchJobStatus +/// +/// #{"ratelimit_module":"Search Jobs", "ratelimit_module_operation":"get"}# +#[utoipa::path( + context_path = "/api", + tag = "Search Jobs", + operation_id = "GetSearchJobStatus", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("job_id" = String, Path, description = "Search job ID") + ), + responses( + (status = 200, description = "Search job status", body = JobModel, example = json!({ + "id": "abc123", + "trace_id": "xyz789", + "org_id": "default", + "user_id": "user1", + "stream_type": "logs", + "stream_names": "[\"stream1\"]", + "payload": "...", + "start_time": 1675182660872049i64, + "end_time": 1675185660872049i64, + "created_at": 1675182660872049i64, + "updated_at": 1675182660872049i64, + "status": 1, + "cluster": "cluster1", + "result_path": "/path/to/result" + })), + (status = 400, description = "Bad Request", body = MetaHttpResponse) + ) +)] #[get("/{org_id}/search_jobs/{job_id}/status")] pub async fn get_status( path: web::Path<(String, String)>, in_req: HttpRequest, ) -> Result { - let user_id = in_req - .headers() - .get("user_id") - .and_then(|v| v.to_str().ok()) - .unwrap_or("") - .to_string(); + #[cfg(feature = "enterprise")] + { + let user_id = in_req + .headers() + .get("user_id") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); - let org_id = path.0.clone(); - let job_id = path.1.clone(); - let res = get(&job_id, &org_id).await; - let model = match res { - Ok(res) => res, - Err(e) => return Ok(MetaHttpResponse::bad_request(e)), - }; + let org_id = path.0.clone(); + let job_id = path.1.clone(); + let res = get(&job_id, &org_id).await; + let model = match res { + Ok(res) => res, + Err(e) => return Ok(MetaHttpResponse::bad_request(e)), + }; - // check permissions - if let Some(res) = check_permissions(&model, &org_id, &user_id).await { - return Ok(res); + // check permissions + if let Some(res) = check_permissions(&model, &org_id, &user_id).await { + return Ok(res); + } + Ok(HttpResponse::Ok().json(model)) + } + + #[cfg(not(feature = "enterprise"))] + { + drop(path); + drop(in_req); + Ok(HttpResponse::Forbidden().json("Not Supported")) } - Ok(HttpResponse::Ok().json(model)) } // 4. cancel +/// CancelSearchJob +/// +/// #{"ratelimit_module":"Search Jobs", "ratelimit_module_operation":"update"}# +#[utoipa::path( + context_path = "/api", + tag = "Search Jobs", + operation_id = "CancelSearchJob", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("job_id" = String, Path, description = "Search job ID") + ), + responses( + (status = 200, description = "Search job cancelled successfully", body = MetaHttpResponse, example = json!({ + "code": 200, + "message": "[Job_Id: abc123] Running Search Job cancelled successfully." + })), + (status = 400, description = "Bad Request", body = MetaHttpResponse) + ) +)] #[post("/{org_id}/search_jobs/{job_id}/cancel")] pub async fn cancel_job( path: web::Path<(String, String)>, in_req: HttpRequest, ) -> Result { - let user_id = in_req - .headers() - .get("user_id") - .and_then(|v| v.to_str().ok()) - .unwrap_or("") - .to_string(); + #[cfg(feature = "enterprise")] + { + let user_id = in_req + .headers() + .get("user_id") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); - let org_id = path.0.clone(); - let job_id = path.1.clone(); - match cancel_job_inner(&org_id, &job_id, &user_id).await { - Ok(res) if res.status() != StatusCode::OK => Ok(res), - Err(e) => Ok(MetaHttpResponse::bad_request(e)), - Ok(_) => Ok(MetaHttpResponse::ok(format!( - "[Job_Id: {job_id}] Running Search Job cancelled successfully." - ))), + let org_id = path.0.clone(); + let job_id = path.1.clone(); + match cancel_job_inner(&org_id, &job_id, &user_id).await { + Ok(res) if res.status() != StatusCode::OK => Ok(res), + Err(e) => Ok(MetaHttpResponse::bad_request(e)), + Ok(_) => Ok(MetaHttpResponse::ok(format!( + "[Job_Id: {job_id}] Running Search Job cancelled successfully." + ))), + } + } + + #[cfg(not(feature = "enterprise"))] + { + drop(path); + drop(in_req); + Ok(HttpResponse::Forbidden().json("Not Supported")) } } // 5. get +/// GetSearchJobResult +/// +/// #{"ratelimit_module":"Search Jobs", "ratelimit_module_operation":"get"}# +#[utoipa::path( + context_path = "/api", + tag = "Search Jobs", + operation_id = "GetSearchJobResult", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("job_id" = String, Path, description = "Search job ID"), + ("from" = Option, Query, description = "Pagination start offset"), + ("size" = Option, Query, description = "Number of results to return") + ), + responses( + (status = 200, description = "Search job results", body = Response, example = json!({ + "took": 155, + "hits": [], + "total": 27179431, + "from": 0, + "size": 100 + })), + (status = 400, description = "Bad Request", body = MetaHttpResponse), + (status = 404, description = "Not Found", body = MetaHttpResponse) + ) +)] #[get("/{org_id}/search_jobs/{job_id}/result")] pub async fn get_job_result( path: web::Path<(String, String)>, req: web::Query, in_req: HttpRequest, ) -> Result { - let user_id = in_req - .headers() - .get("user_id") - .and_then(|v| v.to_str().ok()) - .unwrap_or("") - .to_string(); + #[cfg(feature = "enterprise")] + { + let user_id = in_req + .headers() + .get("user_id") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); - let from = req.from.unwrap_or(0); - let size = req.size.unwrap_or(100); + let from = req.from.unwrap_or(0); + let size = req.size.unwrap_or(100); - let org_id = path.0.clone(); - let job_id = path.1.clone(); - let res = get(&job_id, &org_id).await; - let model = match res { - Ok(res) => res, - Err(e) => return Ok(MetaHttpResponse::bad_request(e)), - }; + let org_id = path.0.clone(); + let job_id = path.1.clone(); + let res = get(&job_id, &org_id).await; + let model = match res { + Ok(res) => res, + Err(e) => return Ok(MetaHttpResponse::bad_request(e)), + }; - // check permissions - if let Some(res) = check_permissions(&model, &org_id, &user_id).await { - return Ok(res); + // check permissions + if let Some(res) = check_permissions(&model, &org_id, &user_id).await { + return Ok(res); + } + + if model.error_message.is_some() { + Ok(MetaHttpResponse::ok(format!( + "job_id: {job_id} error: {}", + model.error_message.unwrap() + ))) + } else if model.status == 1 && model.partition_num != Some(1) { + let response = get_partition_result(&model, from, size).await; + Ok(response) + } else if model.result_path.is_none() || model.cluster.is_none() { + Ok(MetaHttpResponse::not_found(format!( + "[Job_Id: {job_id}] don't have result_path or cluster" + ))) + } else { + let path = model.result_path.clone().unwrap(); + let cluster = model.cluster.clone().unwrap(); + let response = get_result(&path, &cluster, from, size).await; + if let Err(e) = response { + return Ok(MetaHttpResponse::internal_error(e)); + } + Ok(HttpResponse::Ok().json(response.unwrap())) + } } - if model.error_message.is_some() { - Ok(MetaHttpResponse::ok(format!( - "job_id: {job_id} error: {}", - model.error_message.unwrap() - ))) - } else if model.status == 1 && model.partition_num != Some(1) { - let response = get_partition_result(&model, from, size).await; - Ok(response) - } else if model.result_path.is_none() || model.cluster.is_none() { - Ok(MetaHttpResponse::not_found(format!( - "[Job_Id: {job_id}] don't have result_path or cluster" - ))) - } else { - let path = model.result_path.clone().unwrap(); - let cluster = model.cluster.clone().unwrap(); - let response = get_result(&path, &cluster, from, size).await; - if let Err(e) = response { - return Ok(MetaHttpResponse::internal_error(e)); - } - Ok(HttpResponse::Ok().json(response.unwrap())) + #[cfg(not(feature = "enterprise"))] + { + drop(path); + let _ = req; + drop(in_req); + Ok(HttpResponse::Forbidden().json("Not Supported")) } } // 6. delete +/// DeleteSearchJob +/// +/// #{"ratelimit_module":"Search Jobs", "ratelimit_module_operation":"delete"}# +#[utoipa::path( + context_path = "/api", + tag = "Search Jobs", + operation_id = "DeleteSearchJob", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("job_id" = String, Path, description = "Search job ID") + ), + responses( + (status = 200, description = "Search job deleted successfully", body = MetaHttpResponse, example = json!({ + "code": 200, + "message": "[Job_Id: abc123] Running Search Job deleted successfully." + })), + (status = 400, description = "Bad Request", body = MetaHttpResponse), + (status = 404, description = "Not Found", body = MetaHttpResponse) + ) +)] #[delete("/{org_id}/search_jobs/{job_id}")] pub async fn delete_job( path: web::Path<(String, String)>, in_req: HttpRequest, ) -> Result { - let user_id = in_req - .headers() - .get("user_id") - .and_then(|v| v.to_str().ok()) - .unwrap_or("") - .to_string(); + #[cfg(feature = "enterprise")] + { + let user_id = in_req + .headers() + .get("user_id") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); - let org_id = path.0.clone(); - let job_id = path.1.clone(); + let org_id = path.0.clone(); + let job_id = path.1.clone(); - // 1. cancel the query - match cancel_job_inner(&org_id, &job_id, &user_id).await { - Ok(res) if res.status() != StatusCode::OK && res.status() != StatusCode::BAD_REQUEST => { - return Ok(res); + // 1. cancel the query + match cancel_job_inner(&org_id, &job_id, &user_id).await { + Ok(res) + if res.status() != StatusCode::OK && res.status() != StatusCode::BAD_REQUEST => + { + return Ok(res); + } + Err(e) => return Ok(MetaHttpResponse::bad_request(e)), + _ => {} + }; + + // 2. make the job_id in search_job table delete + match set_job_deleted(job_id.as_str()).await { + Ok(true) => Ok(MetaHttpResponse::ok(format!( + "[Job_Id: {job_id}] Running Search Job deleted successfully." + ))), + Ok(false) => Ok(MetaHttpResponse::not_found(format!( + "[Job_Id: {job_id}] Search Job not found" + ))), + Err(e) => Ok(MetaHttpResponse::bad_request(e)), } - Err(e) => return Ok(MetaHttpResponse::bad_request(e)), - _ => {} - }; + } - // 2. make the job_id in search_job table delete - match set_job_deleted(job_id.as_str()).await { - Ok(true) => Ok(MetaHttpResponse::ok(format!( - "[Job_Id: {job_id}] Running Search Job deleted successfully." - ))), - Ok(false) => Ok(MetaHttpResponse::not_found(format!( - "[Job_Id: {job_id}] Search Job not found" - ))), - Err(e) => Ok(MetaHttpResponse::bad_request(e)), + #[cfg(not(feature = "enterprise"))] + { + drop(path); + drop(in_req); + Ok(HttpResponse::Forbidden().json("Not Supported")) } } // 7. retry +/// RetrySearchJob +/// +/// #{"ratelimit_module":"Search Jobs", "ratelimit_module_operation":"update"}# +#[utoipa::path( + context_path = "/api", + tag = "Search Jobs", + operation_id = "RetrySearchJob", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("job_id" = String, Path, description = "Search job ID") + ), + responses( + (status = 200, description = "Search job retry initiated successfully", body = MetaHttpResponse, example = json!({ + "code": 200, + "message": "[Job_Id: abc123] Search Job retry successfully." + })), + (status = 400, description = "Bad Request", body = MetaHttpResponse), + (status = 403, description = "Forbidden - Job cannot be retried", body = MetaHttpResponse) + ) +)] #[post("/{org_id}/search_jobs/{job_id}/retry")] pub async fn retry_job( path: web::Path<(String, String)>, in_req: HttpRequest, ) -> Result { - let user_id = in_req - .headers() - .get("user_id") - .and_then(|v| v.to_str().ok()) - .unwrap_or("") - .to_string(); + #[cfg(feature = "enterprise")] + { + let user_id = in_req + .headers() + .get("user_id") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); - let org_id = path.0.clone(); - let job_id = path.1.clone(); + let org_id = path.0.clone(); + let job_id = path.1.clone(); - // 1. check the status of the job, only cancel, finish can be retry - let status = get(&job_id, &org_id).await; - let model = match status { - Ok(v) => v, - Err(e) => return Ok(MetaHttpResponse::bad_request(e)), - }; + // 1. check the status of the job, only cancel, finish can be retry + let status = get(&job_id, &org_id).await; + let model = match status { + Ok(v) => v, + Err(e) => return Ok(MetaHttpResponse::bad_request(e)), + }; - // check permissions - if let Some(res) = check_permissions(&model, &org_id, &user_id).await { - return Ok(res); + // check permissions + if let Some(res) = check_permissions(&model, &org_id, &user_id).await { + return Ok(res); + } + + if model.status != 2 && model.status != 3 { + return Ok(MetaHttpResponse::forbidden(format!( + "[Job_Id: {job_id}] Only canceled, finished search job can be retry" + ))); + } + + // 2. make the job_id as pending in search_job table + let res = retry_search_job(&job_id).await; + match res { + Ok(_) => Ok(MetaHttpResponse::ok(format!( + "[Job_Id: {job_id}] Search Job retry successfully." + ))), + Err(e) => Ok(MetaHttpResponse::bad_request(e)), + } } - if model.status != 2 && model.status != 3 { - return Ok(MetaHttpResponse::forbidden(format!( - "[Job_Id: {job_id}] Only canceled, finished search job can be retry" - ))); - } - - // 2. make the job_id as pending in search_job table - let res = retry_search_job(&job_id).await; - match res { - Ok(_) => Ok(MetaHttpResponse::ok(format!( - "[Job_Id: {job_id}] Search Job retry successfully." - ))), - Err(e) => Ok(MetaHttpResponse::bad_request(e)), + #[cfg(not(feature = "enterprise"))] + { + drop(in_req); + drop(path); + Ok(HttpResponse::Forbidden().json("Not Supported")) } } - +#[cfg(feature = "enterprise")] async fn cancel_job_inner( org_id: &str, job_id: &str, @@ -382,6 +650,7 @@ async fn cancel_job_inner( cancel_query_inner(org_id, &[&job.trace_id]).await } +#[cfg(feature = "enterprise")] async fn get_partition_result(job: &JobModel, from: i64, size: i64) -> HttpResponse { let req: Result = json::from_str(&job.payload); if let Err(e) = req { @@ -406,7 +675,7 @@ async fn get_partition_result(job: &JobModel, from: i64, size: i64) -> HttpRespo let response = response.unwrap(); apply_pagination(response, from, size) } - +#[cfg(feature = "enterprise")] fn apply_pagination(response: Response, from: i64, size: i64) -> HttpResponse { let mut res = response; res.pagination(from, size); @@ -414,6 +683,7 @@ fn apply_pagination(response: Response, from: i64, size: i64) -> HttpResponse { } // check permissions +#[cfg(feature = "enterprise")] async fn check_permissions(job: &JobModel, org_id: &str, user_id: &str) -> Option { let stream_type = StreamType::from(job.stream_type.as_str()); let stream_names: Vec = json::from_str(&job.stream_names).unwrap(); diff --git a/src/handler/http/request/search/utils.rs b/src/handler/http/request/search/utils.rs index be38ec3468..04fadcd0a1 100644 --- a/src/handler/http/request/search/utils.rs +++ b/src/handler/http/request/search/utils.rs @@ -13,17 +13,20 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use actix_web::HttpResponse; -use config::meta::stream::StreamType; -use o2_openfga::meta::mapping::OFGA_MODELS; - -use crate::common::{ - infra::config::USERS, - meta::{self, http::HttpResponse as MetaHttpResponse}, - utils::auth::{AuthExtractor, is_root_user}, +#[cfg(feature = "enterprise")] +use { + crate::common::{ + infra::config::USERS, + meta::{self, http::HttpResponse as MetaHttpResponse}, + utils::auth::{AuthExtractor, is_root_user}, + }, + actix_web::HttpResponse, + config::meta::stream::StreamType, + o2_openfga::meta::mapping::OFGA_MODELS, }; // Check permissions on stream +#[cfg(feature = "enterprise")] pub async fn check_stream_permissions( stream_name: &str, org_id: &str, diff --git a/src/handler/http/request/service_accounts/mod.rs b/src/handler/http/request/service_accounts/mod.rs index d1cc6d43a2..3e071b3121 100644 --- a/src/handler/http/request/service_accounts/mod.rs +++ b/src/handler/http/request/service_accounts/mod.rs @@ -37,6 +37,8 @@ use crate::{ }; /// ListServiceAccounts +/// +/// #{"ratelimit_module":"Service Accounts", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "ServiceAccounts", @@ -89,6 +91,8 @@ pub async fn list(org_id: web::Path, req: HttpRequest) -> Result, body: web::Bytes) -> Result, req: HttpRequest) -> impl Responder { })) } +/// StreamDeleteCache +/// +/// #{"ratelimit_module":"Streams", "ratelimit_module_operation":"delete"}# #[utoipa::path( context_path = "/api", tag = "Streams", diff --git a/src/handler/http/request/syslog/mod.rs b/src/handler/http/request/syslog/mod.rs index e0e9669851..eb6c0092c5 100644 --- a/src/handler/http/request/syslog/mod.rs +++ b/src/handler/http/request/syslog/mod.rs @@ -30,6 +30,8 @@ pub async fn toggle_state(details: web::Json) -> Result) -> Result impl Responder { } /// DeleteSyslogRoute +/// +/// #{"ratelimit_module":"Syslog Routes", "ratelimit_module_operation":"delete"}# #[utoipa::path( context_path = "/api", tag = "Syslog Routes", diff --git a/src/handler/http/request/traces/mod.rs b/src/handler/http/request/traces/mod.rs index fa30ed8cd2..6e24af19e7 100644 --- a/src/handler/http/request/traces/mod.rs +++ b/src/handler/http/request/traces/mod.rs @@ -88,6 +88,8 @@ async fn handle_req( } /// GetLatestTraces +/// +/// #{"ratelimit_module":"Traces", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Traces", diff --git a/src/handler/http/request/users/mod.rs b/src/handler/http/request/users/mod.rs index 50cadaae82..1ed3cdcba0 100644 --- a/src/handler/http/request/users/mod.rs +++ b/src/handler/http/request/users/mod.rs @@ -53,6 +53,8 @@ use crate::{ pub mod service_accounts; /// ListUsers +/// +/// #{"ratelimit_module":"Users", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Users", @@ -91,6 +93,8 @@ pub async fn list(org_id: web::Path, user_email: UserEmail) -> Result>( auth_cookie } /// RemoveUserFromOrganization +/// +/// #{"ratelimit_module":"Users", "ratelimit_module_operation":"delete"}# #[utoipa::path( context_path = "/api", tag = "Users", @@ -254,6 +264,8 @@ pub async fn delete( } /// AuthenticateUser +/// +/// #{"ratelimit_module":"Users", "ratelimit_module_operation":"update"}# #[utoipa::path( context_path = "/auth", tag = "Auth", @@ -639,7 +651,9 @@ pub async fn get_auth(_req: HttpRequest) -> Result { } } -/// ListUsers +/// ListUserRoles +/// +/// #{"ratelimit_module":"Users", "ratelimit_module_operation":"list"}# #[utoipa::path( context_path = "/api", tag = "Users", diff --git a/src/handler/http/router/mod.rs b/src/handler/http/router/mod.rs index 79bf573934..022e894fc0 100644 --- a/src/handler/http/router/mod.rs +++ b/src/handler/http/router/mod.rs @@ -571,6 +571,10 @@ pub fn get_service_routes(svc: &mut web::ServiceConfig) { .service(actions::action::update_action_details) .service(actions::action::serve_action_zip) .service(actions::action::delete_action) + .service(ratelimit::list_module_ratelimit) + .service(ratelimit::list_role_ratelimit) + .service(ratelimit::update_ratelimit) + .service(ratelimit::api_modules) .service(actions::operations::test_action); svc.service(service); diff --git a/src/handler/http/router/openapi.rs b/src/handler/http/router/openapi.rs index bf4ed7b7ba..2936406f81 100644 --- a/src/handler/http/router/openapi.rs +++ b/src/handler/http/router/openapi.rs @@ -14,6 +14,8 @@ // along with this program. If not, see . use config::{get_config, meta::stream::StreamType}; +#[cfg(feature = "enterprise")] +use o2_ratelimit::dataresource::default_rules::OpenapiInfo; use utoipa::{Modify, OpenApi, openapi::security::SecurityScheme}; use crate::{common::meta, handler::http::request}; @@ -138,6 +140,56 @@ use crate::{common::meta, handler::http::request}; request::clusters::list_clusters, request::short_url::shorten, request::short_url::retrieve, + request::ratelimit::list_module_ratelimit, + request::ratelimit::list_role_ratelimit, + request::ratelimit::update_ratelimit, + request::service_accounts::list, + request::service_accounts::save, + request::service_accounts::update, + request::service_accounts::delete, + request::service_accounts::get_api_token, + request::pipeline::save_pipeline, + request::pipeline::list_pipelines, + request::pipeline::list_streams_with_pipeline, + request::pipeline::delete_pipeline, + request::pipeline::update_pipeline, + request::pipeline::enable_pipeline, + request::dashboards::reports::create_report, + request::dashboards::reports::update_report, + request::dashboards::reports::list_reports, + request::dashboards::reports::get_report, + request::dashboards::reports::delete_report, + request::dashboards::reports::enable_report, + request::dashboards::reports::trigger_report, + request::actions::action::upload_zipped_action, + request::actions::action::delete_action, + request::actions::action::serve_action_zip, + request::actions::action::update_action_details, + request::actions::action::list_actions, + request::actions::action::get_action_from_id, + request::authz::fga::create_role, + request::authz::fga::delete_role, + request::authz::fga::get_roles, + request::authz::fga::update_role, + request::authz::fga::get_role_permissions, + request::authz::fga::get_users_with_role, + request::authz::fga::create_group, + request::authz::fga::update_group, + request::authz::fga::get_groups, + request::authz::fga::get_group_details, + request::authz::fga::delete_group, + request::keys::save, + request::keys::get, + request::keys::list, + request::keys::delete, + request::keys::update, + request::search::search_job::submit_job, + request::search::search_job::list_status, + request::search::search_job::get_status, + request::search::search_job::cancel_job, + request::search::search_job::get_job_result, + request::search::search_job::delete_job, + request::search::search_job::retry_job, ), components( schemas( @@ -283,7 +335,6 @@ use crate::{common::meta, handler::http::request}; config::meta::promql::Metadata, config::meta::promql::MetricType, // Functions - ), ), modifiers(&SecurityAddon), @@ -305,6 +356,7 @@ use crate::{common::meta, handler::http::request}; (name = "Syslog Routes", description = "Syslog Routes retrieval & management operations"), (name = "Clusters", description = "Super cluster operations"), (name = "Short Url", description = "Short Url Service"), + (name = "Ratelimit", description = "Ratelimit operations"), ), info( description = "OpenObserve API documents [https://openobserve.ai/docs/](https://openobserve.ai/docs/)", @@ -330,3 +382,48 @@ impl Modify for SecurityAddon { ); } } + +#[cfg(feature = "enterprise")] +pub async fn openapi_info() -> OpenapiInfo { + let api = ApiDoc::openapi(); + + // Group endpoints by tags with full operation details + let mut tag_operations: OpenapiInfo = std::collections::HashMap::new(); + + for (path, path_item) in &api.paths.paths { + for (method, operation) in path_item.operations.clone() { + let tags = operation + .tags + .clone() + .unwrap_or_else(|| vec!["untagged".to_string()]); + + let method = match method { + utoipa::openapi::PathItemType::Get => "GET", + utoipa::openapi::PathItemType::Post => "POST", + utoipa::openapi::PathItemType::Put => "PUT", + utoipa::openapi::PathItemType::Delete => "DELETE", + utoipa::openapi::PathItemType::Patch => "PATCH", + utoipa::openapi::PathItemType::Head => "HEAD", + utoipa::openapi::PathItemType::Options => "OPTIONS", + utoipa::openapi::PathItemType::Trace => "TRACE", + utoipa::openapi::PathItemType::Connect => "CONNECT", + }; + + let operation_info = ( + method.to_string(), + path.clone(), + operation.clone().description, + tags.clone(), + ); + + for tag in tags { + tag_operations + .entry(tag) + .or_default() + .push(operation_info.clone()); + } + } + } + + tag_operations +} diff --git a/src/infra/src/table/entity/mod.rs b/src/infra/src/table/entity/mod.rs index cc9d738b9b..913f90eb76 100644 --- a/src/infra/src/table/entity/mod.rs +++ b/src/infra/src/table/entity/mod.rs @@ -9,6 +9,7 @@ pub mod dashboards; pub mod destinations; pub mod distinct_value_fields; pub mod folders; +pub mod rate_limit_rules; pub mod search_job_partitions; pub mod search_job_results; pub mod search_jobs; diff --git a/src/infra/src/table/entity/rate_limit_rules.rs b/src/infra/src/table/entity/rate_limit_rules.rs new file mode 100644 index 0000000000..2f24a4f109 --- /dev/null +++ b/src/infra/src/table/entity/rate_limit_rules.rs @@ -0,0 +1,23 @@ +//! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.0 + +use sea_orm::entity::prelude::*; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel)] +#[sea_orm(table_name = "rate_limit_rules")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub rule_id: String, + pub org: String, + pub rule_type: String, + pub user_role: String, + pub user_id: String, + pub api_group_name: String, + pub api_group_operation: String, + pub threshold: i32, + pub created_at: i64, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/src/infra/src/table/migration/m20250217_115548_ratelimit_table.rs b/src/infra/src/table/migration/m20250217_115548_ratelimit_table.rs new file mode 100644 index 0000000000..e13d707a61 --- /dev/null +++ b/src/infra/src/table/migration/m20250217_115548_ratelimit_table.rs @@ -0,0 +1,136 @@ +use sea_orm::DbBackend; +use sea_orm_migration::prelude::*; + +const RESOURCE_KEY: &str = "rkey_idx"; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .create_table(create_ratelimit_table_statement()) + .await?; + + match manager.get_database_backend() { + DbBackend::MySql => { + manager + .create_index(create_ratelimit_resource_key_idx_stmnt_mysql()) + .await?; + } + _ => { + manager + .create_index(create_ratelimit_resource_key_idx_stmnt()) + .await?; + } + } + + Ok(()) + } + + async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .drop_index( + Index::drop() + .name(RESOURCE_KEY) + .table(RateLimitRules::Table) + .to_owned(), + ) + .await?; + manager + .drop_table(Table::drop().table(RateLimitRules::Table).to_owned()) + .await + } +} + +fn create_ratelimit_table_statement() -> TableCreateStatement { + Table::create() + .table(RateLimitRules::Table) + .if_not_exists() + .col( + ColumnDef::new(RateLimitRules::RuleId) + .string_len(64) + .primary_key() + .not_null(), + ) + .col( + ColumnDef::new(RateLimitRules::Org) + .string_len(100) + .not_null(), + ) + .col( + ColumnDef::new(RateLimitRules::RuleType) + .string_len(50) + .extra("DEFAULT 'exact'") + .not_null(), + ) + .col( + ColumnDef::new(RateLimitRules::UserRole) + .string_len(256) + .not_null(), + ) + .col( + ColumnDef::new(RateLimitRules::UserId) + .string_len(256) + .not_null(), + ) + .col( + ColumnDef::new(RateLimitRules::ApiGroupName) + .string_len(256) + .not_null(), + ) + .col( + ColumnDef::new(RateLimitRules::ApiGroupOperation) + .string_len(128) + .not_null(), + ) + .col( + ColumnDef::new(RateLimitRules::Threshold) + .integer() + .not_null(), + ) + .col( + ColumnDef::new(RateLimitRules::CreatedAt) + .big_integer() + .not_null(), + ) + .to_owned() +} + +fn create_ratelimit_resource_key_idx_stmnt() -> IndexCreateStatement { + sea_query::Index::create() + .name(RESOURCE_KEY) + .table(RateLimitRules::Table) + .col(RateLimitRules::Org) + .col(RateLimitRules::UserRole) + .col(RateLimitRules::UserId) + .col(RateLimitRules::ApiGroupName) + .col(RateLimitRules::ApiGroupOperation) + .unique() + .to_owned() +} + +fn create_ratelimit_resource_key_idx_stmnt_mysql() -> IndexCreateStatement { + sea_query::Index::create() + .name(RESOURCE_KEY) + .table(RateLimitRules::Table) + .col(RateLimitRules::Org) + .col(RateLimitRules::UserRole) + .col(RateLimitRules::UserId) + .to_owned() +} + +#[derive(DeriveIden)] +enum RateLimitRules { + Table, + Org, + RuleId, + RuleType, + UserRole, + UserId, + ApiGroupName, + ApiGroupOperation, + Threshold, + CreatedAt, +} diff --git a/src/infra/src/table/migration/mod.rs b/src/infra/src/table/migration/mod.rs index 977f3b60d9..554e15f371 100644 --- a/src/infra/src/table/migration/mod.rs +++ b/src/infra/src/table/migration/mod.rs @@ -44,6 +44,7 @@ mod m20250125_133700_populate_destinations_table; mod m20250125_153005_delete_metas_destinations; mod m20250125_172300_delete_metas_templates; mod m20250213_000001_add_dashboard_updated_at; +mod m20250217_115548_ratelimit_table; mod m20250320_000001_remove_alert_name_unique_constraint; pub struct Migrator; @@ -80,6 +81,7 @@ impl MigratorTrait for Migrator { Box::new(m20250125_133700_populate_destinations_table::Migration), Box::new(m20250125_153005_delete_metas_destinations::Migration), Box::new(m20250213_000001_add_dashboard_updated_at::Migration), + Box::new(m20250217_115548_ratelimit_table::Migration), Box::new(m20250320_000001_remove_alert_name_unique_constraint::Migration), ] } diff --git a/src/infra/src/table/mod.rs b/src/infra/src/table/mod.rs index 679059a026..57d802953a 100644 --- a/src/infra/src/table/mod.rs +++ b/src/infra/src/table/mod.rs @@ -32,6 +32,7 @@ pub mod distinct_values; pub mod entity; pub mod folders; mod migration; +pub mod ratelimit; pub mod search_job; pub mod search_queue; pub mod short_urls; diff --git a/src/infra/src/table/ratelimit.rs b/src/infra/src/table/ratelimit.rs new file mode 100644 index 0000000000..24b77f464e --- /dev/null +++ b/src/infra/src/table/ratelimit.rs @@ -0,0 +1,779 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use anyhow::{Context, anyhow}; +use bytes::Bytes; +use config::{ + meta::ratelimit::{RatelimitRule, RatelimitRuleType}, + utils::time::now, +}; +use sea_orm::{ + ActiveModelTrait, ActiveValue::Set, ColumnTrait, EntityTrait, QueryFilter, QueryOrder, + QuerySelect, TransactionTrait, +}; +use serde::{Deserialize, Serialize}; + +use crate::{ + db::{ORM_CLIENT, connect_to_orm}, + orm_err, + table::{ + entity::rate_limit_rules::{ActiveModel, Column, Entity}, + migration::{Expr, Order}, + }, +}; + +pub const RULE_EXISTS: &str = "Rule already exists"; +pub const RULE_NOT_FOUND: &str = "Rule not found"; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum RuleEntry { + Single(RatelimitRule), + Batch(Vec), + UpsertBatch(Vec), +} + +impl TryFrom<&Bytes> for RuleEntry { + type Error = anyhow::Error; + + fn try_from(bytes: &Bytes) -> Result { + serde_json::from_slice(bytes).context("Failed to deserialize RatelimitRule") + } +} + +pub async fn fetch_rules( + mut default_rules: Vec, + org_id: Option, + user_role: Option, +) -> Result, anyhow::Error> { + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let mut res = Entity::find() + .select_only() + .column(Column::Org) + .column(Column::RuleId) + .column(Column::RuleType) + .column(Column::UserRole) + .column(Column::ApiGroupName) + .column(Column::ApiGroupOperation) + .column(Column::UserId) + .column(Column::Threshold) + .order_by(Column::Threshold, Order::Asc); + + if let Some(org_id) = org_id { + res = res.filter(Column::Org.eq(org_id)); + }; + + if let Some(user_role) = user_role { + res = res.filter(Column::UserRole.eq(user_role)); + }; + + let records = res + .into_model::() + .all(client) + .await + .map_err(|e| anyhow!("DbError# {e}"))?; + + default_rules.extend(records); + Ok(default_rules) +} + +pub async fn fetch_rules_by_id(rule_id: &str) -> Result, anyhow::Error> { + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let res = Entity::find() + .select_only() + .column(Column::Org) + .column(Column::RuleId) + .column(Column::RuleType) + .column(Column::UserRole) + .column(Column::ApiGroupName) + .column(Column::ApiGroupOperation) + .column(Column::UserId) + .column(Column::Threshold) + .filter(Column::RuleId.eq(rule_id)); + let record = res + .into_model::() + .one(client) + .await + .map_err(|e| anyhow!("DbError# {e}"))?; + Ok(record) +} + +pub async fn add(rule: RuleEntry) -> Result<(), anyhow::Error> { + match rule { + RuleEntry::Single(rule) => add_single(rule).await, + RuleEntry::Batch(rules) => add_batch(rules).await, + RuleEntry::UpsertBatch(rules) => add_upsert_batch(rules).await, + } +} + +async fn add_batch(rules: Vec) -> Result<(), anyhow::Error> { + if rules.is_empty() { + return Ok(()); + } + + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + + // Collect all rule IDs and validate + let rule_ids: Vec = rules + .iter() + .map(|rule| { + rule.rule_id + .clone() + .ok_or_else(|| anyhow!("Rule ID is required for batch operations")) + }) + .collect::, _>>()?; + + // Check for existing rules + let existing_rules = Entity::find() + .filter(Column::RuleId.is_in(rule_ids.clone())) + .all(client) + .await + .map_err(|e| anyhow!("DbError# Failed to fetch existing rules: {}", e))?; + + // If any rule exists, return error with details + if !existing_rules.is_empty() { + let existing_ids: Vec = existing_rules + .iter() + .map(|rule| rule.rule_id.clone()) + .collect(); + + return Err(anyhow!( + "Cannot add batch: following rules already exist: {}", + existing_ids.join(", ") + )); + } + + // Prepare batch insert + let current_timestamp = now().timestamp(); + let active_rules: Vec = rules + .into_iter() + .map(|rule| ActiveModel { + org: Set(rule.org), + rule_id: Set(rule.rule_id.unwrap()), + rule_type: Set(rule + .rule_type + .unwrap_or(RatelimitRuleType::Exact.to_string())), + user_role: Set(rule.user_role.unwrap_or_default()), + user_id: Set(rule.user_id.unwrap_or_default()), + api_group_name: Set(rule.api_group_name.unwrap_or_default()), + api_group_operation: Set(rule.api_group_operation.unwrap_or_default()), + threshold: Set(rule.threshold), + created_at: Set(current_timestamp), + }) + .collect(); + + // Execute batch insert in a transaction + let txn = client + .begin() + .await + .map_err(|e| anyhow!("DbError# Failed to start transaction: {}", e))?; + + match Entity::insert_many(active_rules) + .exec(&txn) + .await + .map_err(|e| anyhow!("DbError# Failed to insert rules: {}", e)) + { + Ok(_) => { + txn.commit() + .await + .map_err(|e| anyhow!("DbError# Failed to commit transaction: {}", e))?; + } + Err(e) => { + txn.rollback() + .await + .map_err(|e| anyhow!("DbError# Failed to rollback transaction: {}", e))?; + return Err(e); + } + } + + Ok(()) +} + +async fn add_upsert_batch(rules: Vec) -> Result<(), anyhow::Error> { + if rules.is_empty() { + return Ok(()); + } + + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let current_time = config::utils::time::now_micros(); + + // Start transaction + let txn = client + .begin() + .await + .map_err(|e| anyhow!("DbError# Failed to start transaction: {}", e))?; + + for rule in rules { + // Check if rule exists + let existing_rule = Entity::find() + .filter(Column::Org.eq(rule.org.clone())) + .filter(Column::UserRole.eq(rule.user_role.clone().unwrap_or_default())) + .filter(Column::UserId.eq(rule.user_id.clone().unwrap_or_default())) + .filter(Column::ApiGroupName.eq(rule.api_group_name.clone().unwrap_or_default())) + .filter( + Column::ApiGroupOperation.eq(rule.api_group_operation.clone().unwrap_or_default()), + ) + .one(&txn) + .await + .map_err(|e| anyhow!("DbError# Failed to fetch existing rule: {}", e))?; + + match existing_rule { + Some(existing) => { + // Update existing rule using the Model's primary key + let mut active_model: ActiveModel = existing.clone().into(); + + // Update only the fields that should change + active_model.org = Set(rule.org); + active_model.rule_id = Set(existing.rule_id.clone()); + active_model.rule_type = Set(rule + .rule_type + .unwrap_or(RatelimitRuleType::Exact.to_string())); + active_model.user_role = Set(rule.user_role.unwrap_or_default()); + active_model.user_id = Set(rule.user_id.unwrap_or_default()); + active_model.api_group_name = Set(rule.api_group_name.unwrap_or_default()); + active_model.api_group_operation = + Set(rule.api_group_operation.unwrap_or_default()); + active_model.threshold = Set(rule.threshold); + active_model.created_at = Set(existing.created_at); + + if let Err(e) = active_model.save(&txn).await { + txn.rollback() + .await + .map_err(|e| anyhow!("DbError# Failed to rollback transaction: {}", e))?; + + return Err(anyhow!( + "DbError# Failed to update rule: {:?}, {}", + existing, + e + )); + } + } + None => { + // if threshold <= 0 (which means no limit) and no record in the table, + // we should not insert a new record + if rule.threshold > 0 { + // Insert new rule + let active_model = ActiveModel { + org: Set(rule.org), + rule_id: Set(rule.rule_id.unwrap()), + rule_type: Set(rule + .rule_type + .unwrap_or(RatelimitRuleType::Exact.to_string())), + user_role: Set(rule.user_role.unwrap_or_default()), + user_id: Set(rule.user_id.unwrap_or_default()), + api_group_name: Set(rule.api_group_name.unwrap_or_default()), + api_group_operation: Set(rule.api_group_operation.unwrap_or_default()), + threshold: Set(rule.threshold), + created_at: Set(current_time), + }; + log::debug!("add_upsert_batch insert active_model: {:?} ", active_model); + if let Err(e) = Entity::insert(active_model).exec(&txn).await { + txn.rollback().await.map_err(|e| { + anyhow!("DbError# Failed to rollback transaction: {}", e) + })?; + + return Err(anyhow!("DbError# Failed to insert rule: {}", e)); + } + } + } + } + // wait for a while to reduce the databases load + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + } + + // Commit transaction + txn.commit() + .await + .map_err(|e| anyhow!("DbError# Failed to commit transaction: {}", e))?; + + Ok(()) +} + +async fn add_single(rule: RatelimitRule) -> Result<(), anyhow::Error> { + let rule_id = rule.rule_id.clone(); + match fetch_rules_by_id(rule_id.unwrap().as_str()).await { + Ok(None) => { + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let active_rule = ActiveModel { + org: Set(rule.org), + rule_id: Set(rule.rule_id.unwrap()), + rule_type: Set(rule + .rule_type + .unwrap_or(RatelimitRuleType::Exact.to_string())), + user_role: Set(rule.user_role.unwrap_or("".to_string())), + user_id: Set(rule.user_id.unwrap_or_default()), + api_group_name: Set(rule.api_group_name.unwrap_or_default()), + api_group_operation: Set(rule.api_group_operation.unwrap_or_default()), + threshold: Set(rule.threshold), + created_at: Set(now().timestamp()), + }; + match Entity::insert(active_rule) + .exec(client) + .await + .map_err(|e| anyhow!("DbError# {e}")) + { + Ok(_) => Ok(()), + Err(e) => orm_err!(format!("Add ratelimit rule error: {e}")) + .map_err(|e| anyhow!("DbError# {e}"))?, + } + } + Ok(Some(_)) => Err(anyhow::anyhow!(RULE_EXISTS)), + Err(e) => { + log::error!("Add Rule Error fetching rule: {:?}", e.to_string()); + Err(anyhow::anyhow!(e.to_string())) + } + } +} + +pub async fn update(rule: RuleEntry) -> Result<(), anyhow::Error> { + match rule { + RuleEntry::Single(rule) => update_single(rule).await, + RuleEntry::Batch(rules) => update_batch(rules).await, + RuleEntry::UpsertBatch(rules) => add_upsert_batch(rules).await, + } +} + +async fn update_single(rule: RatelimitRule) -> Result<(), anyhow::Error> { + let rule_id = rule + .rule_id + .clone() + .ok_or_else(|| anyhow!("Rule ID is required"))?; + + match fetch_rules_by_id(rule_id.as_str()).await { + Ok(Some(_)) => { + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + match Entity::update_many() + .col_expr(Column::Org, Expr::value(&rule.org)) + .col_expr( + Column::RuleType, + Expr::value( + rule.rule_type + .unwrap_or(RatelimitRuleType::Exact.to_string()), + ), + ) + .col_expr(Column::Threshold, Expr::value(rule.threshold)) + .col_expr( + Column::UserRole, + Expr::value(rule.user_role.unwrap_or_default()), + ) + .col_expr( + Column::ApiGroupName, + Expr::value(rule.api_group_name.unwrap_or_default()), + ) + .col_expr( + Column::ApiGroupOperation, + Expr::value(rule.api_group_operation.unwrap_or_default()), + ) + .filter(Column::RuleId.eq(rule_id)) + .exec(client) + .await + .map_err(|e| anyhow!("DbError# {e}")) + { + Ok(_) => Ok(()), + Err(e) => orm_err!(format!("update ratelimit rule error: {e}")) + .map_err(|e| anyhow!("DbError# {e}"))?, + } + } + Ok(None) => Err(anyhow::anyhow!(RULE_NOT_FOUND)), + Err(e) => { + log::error!("Update rule Error fetching rule: {:?}", e.to_string()); + Err(anyhow::anyhow!(e.to_string())) + } + } +} + +async fn update_batch(rules: Vec) -> Result<(), anyhow::Error> { + if rules.is_empty() { + return Ok(()); + } + + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let txn = client + .begin() + .await + .map_err(|e| anyhow!("DbError# Failed to start transaction: {}", e))?; + + // Collect all rule IDs + let rule_ids: Vec = rules + .iter() + .map(|rule| { + rule.rule_id + .clone() + .ok_or_else(|| anyhow!("Rule ID is required for batch operations")) + }) + .collect::, _>>()?; + + // Check if all rules exist + let existing_rules = Entity::find() + .filter(Column::RuleId.is_in(rule_ids.clone())) + .all(&txn) + .await + .map_err(|e| anyhow!("DbError# Failed to fetch existing rules: {}", e))?; + + if existing_rules.len() != rule_ids.len() { + let existing_ids: std::collections::HashSet = existing_rules + .iter() + .map(|rule| rule.rule_id.clone()) + .collect(); + let missing_ids: Vec = rule_ids + .iter() + .filter(|id| !existing_ids.contains(*id)) + .cloned() + .collect(); + + return Err(anyhow!( + "The following rules do not exist: rules id {}", + missing_ids.join(", ") + )); + } + + // Update each rule in transaction + for rule in rules { + let rule_id = rule.rule_id.clone().unwrap(); + match Entity::update_many() + .col_expr(Column::Org, Expr::value(&rule.org)) + .col_expr( + Column::RuleType, + Expr::value( + rule.rule_type + .unwrap_or(RatelimitRuleType::Exact.to_string()), + ), + ) + .col_expr(Column::Threshold, Expr::value(rule.threshold)) + .col_expr( + Column::UserRole, + Expr::value(rule.user_role.unwrap_or_default()), + ) + .col_expr( + Column::ApiGroupName, + Expr::value(rule.api_group_name.unwrap_or_default()), + ) + .col_expr( + Column::ApiGroupOperation, + Expr::value(rule.api_group_operation.unwrap_or_default()), + ) + .filter(Column::RuleId.eq(rule_id)) + .exec(&txn) + .await + { + Ok(_) => continue, + Err(e) => { + txn.rollback() + .await + .map_err(|e| anyhow!("DbError# Failed to rollback transaction: {}", e))?; + return Err(anyhow!("DbError# Failed to update rule: {}", e)); + } + } + } + + // Commit transaction + txn.commit() + .await + .map_err(|e| anyhow!("DbError# Failed to commit transaction: {}", e))?; + + Ok(()) +} + +pub async fn delete(rule_id: String) -> Result<(), anyhow::Error> { + match fetch_rules_by_id(rule_id.as_str()).await { + Ok(Some(_)) => { + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + match Entity::delete_many() + .filter(Column::RuleId.eq(rule_id)) + .exec(client) + .await + .map_err(|e| anyhow!("DbError# {e}")) + { + Ok(_) => Ok(()), + Err(e) => orm_err!(format!("Delete ratelimit rule error: {e}")) + .map_err(|e| anyhow!("DbError# {e}"))?, + } + } + Ok(None) => Err(anyhow::anyhow!(RULE_NOT_FOUND)), + Err(e) => { + log::error!("Delete rule Error fetching rule: {:?}", e.to_string()); + Err(anyhow::anyhow!(e.to_string())) + } + } +} + +pub async fn list( + org: &str, + api: Option<&str>, + user_id: Option<&str>, +) -> Result, anyhow::Error> { + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let mut query = Entity::find() + .select_only() + .column(Column::Org) + .column(Column::RuleId) + .column(Column::RuleType) + .column(Column::UserRole) + .column(Column::ApiGroupName) + .column(Column::ApiGroupOperation) + .column(Column::Threshold) + .filter(Column::Org.eq(org)); + + if let Some(api) = api { + query = query.filter(Column::ApiGroupName.eq(api)); + } + + if let Some(user_id) = user_id { + query = query.filter(Column::UserRole.eq(user_id)); + } + + let records = query + .into_model::() + .all(client) + .await + .map_err(|e| anyhow!("DbError# {e}"))?; + + Ok(records) +} + +#[cfg(test)] +mod tests { + use std::sync::Once; + + use bytes::Bytes; + use sea_orm::{DatabaseBackend, DbErr, MockDatabase, MockExecResult}; + use serde_json::json; + + use super::*; + use crate::table::entity::rate_limit_rules::Model; + + fn create_test_rule() -> RatelimitRule { + RatelimitRule { + org: "test_org".to_string(), + rule_type: Some("exact".to_string()), + rule_id: Some("test_rule".to_string()), + user_role: Some("test_user_role".to_string()), + user_id: Some("test_user_id".to_string()), + api_group_name: Some("test_group_name".to_string()), + api_group_operation: Some("test_operation".to_string()), + threshold: 100, + } + } + + fn create_test_model() -> Model { + Model { + org: "test_org".to_string(), + rule_type: "exact".to_string(), + rule_id: "test_rule".to_string(), + user_role: "test_user_role".to_string(), + user_id: "test_user_id".to_string(), + api_group_name: "api_group_name".to_string(), + api_group_operation: "operation".to_string(), + threshold: 100, + created_at: 0, + } + } + + #[tokio::test] + async fn test_fetch_rules() -> Result<(), DbErr> { + let db = MockDatabase::new(DatabaseBackend::Postgres) + .append_query_results(vec![vec![create_test_model()]]) + .into_connection(); + + // Override the ORM_CLIENT for testing + let rules = Entity::find() + .select_only() + .column(Column::Org) + .column(Column::RuleId) + .column(Column::RuleType) + .column(Column::UserRole) + .column(Column::ApiGroupName) + .column(Column::Threshold) + .order_by(Column::Threshold, Order::Asc) + .into_model::() + .all(&db) + .await?; + + assert_eq!(rules.len(), 1); + let rule = &rules[0]; + assert_eq!(rule.org, "test_org"); + assert_eq!(rule.rule_type, Some("exact".to_string())); + assert_eq!(rule.rule_id, Some("test_rule".to_string())); + assert_eq!(rule.user_role, Some("test_user_role".to_string())); + assert_eq!(rule.api_group_name, Some("api_group_name".to_string())); + assert_eq!(rule.threshold, 100); + + Ok(()) + } + + #[tokio::test] + async fn test_fetch_rules_by_id() { + let test_rule = create_test_model(); + + let db = MockDatabase::new(DatabaseBackend::Postgres) + .append_query_results(vec![vec![test_rule.clone()], vec![]]) + .into_connection(); + + let result = Entity::find() + .select_only() + .column(Column::Org) + .column(Column::RuleId) + .column(Column::RuleType) + .column(Column::UserRole) + .column(Column::Threshold) + .filter(Column::RuleId.eq("test_rule")) + .into_model::() + .one(&db) + .await; + + assert!(result.is_ok()); + let rule = result.unwrap().unwrap(); + assert_eq!(rule.rule_id, Some("test_rule".to_string())); + assert_eq!(rule.org, "test_org"); + + let result = Entity::find() + .select_only() + .column(Column::Org) + .column(Column::RuleId) + .column(Column::RuleType) + .column(Column::UserRole) + .column(Column::Threshold) + .filter(Column::RuleId.eq("non_existent")) + .into_model::() + .one(&db) + .await; + + assert!(result.is_ok()); + assert!(result.unwrap().is_none()); + } + + #[tokio::test] + async fn test_update_rule() -> Result<(), DbErr> { + let db = MockDatabase::new(DatabaseBackend::Postgres) + .append_exec_results(vec![MockExecResult { + last_insert_id: 0, + rows_affected: 1, + }]) + .into_connection(); + + let rule = create_test_rule(); + + // Test update + let result = Entity::update_many() + .col_expr(Column::Org, Expr::value(&rule.org)) + .col_expr(Column::RuleType, Expr::value(&rule.rule_type.unwrap())) + .col_expr(Column::UserRole, Expr::value(&rule.user_role.unwrap())) + .col_expr(Column::Threshold, Expr::value(rule.threshold)) + .filter(Column::RuleId.eq(&rule.rule_id.unwrap())) + .exec(&db) + .await; + + assert!(result.is_ok()); + assert_eq!(result.unwrap().rows_affected, 1); + + Ok(()) + } + + #[test] + fn test_try_from_bytes() { + // Test valid JSON + let json_data = json!({ + "org": "test_org", + "rule_type": "exact", + "rule_id": "test_rule_id", + "user_role": "user_role", + "user_id": "test_user_id", + "api": "test_api", + "threshold": 100, + }); + let bytes = Bytes::from(json_data.to_string()); + + let result = RatelimitRule::try_from(&bytes); + assert!(result.is_ok()); + + let rule = result.unwrap(); + assert_eq!(rule.org, "test_org"); + assert_eq!(rule.rule_type, Some("exact".to_string())); + assert_eq!(rule.rule_id, Some("test_rule_id".to_string())); + assert_eq!(rule.user_role, Some("user_role".to_string())); + assert_eq!(rule.threshold, 100); + + // Test invalid JSON + let invalid_bytes = Bytes::from("invalid json"); + let result = RatelimitRule::try_from(&invalid_bytes); + assert!(result.is_err()); + + // Test missing fields + let invalid_json = json!({ + "org": "test_org" + }); + let bytes = Bytes::from(invalid_json.to_string()); + let result = RatelimitRule::try_from(&bytes); + assert!(result.is_err()); + } + + #[test] + fn test_ratelimit_rule_type() { + assert!(matches!( + RatelimitRuleType::try_from("exact"), + Ok(RatelimitRuleType::Exact) + )); + assert!(matches!( + RatelimitRuleType::try_from("regex"), + Ok(RatelimitRuleType::Regex) + )); + assert_eq!( + RatelimitRuleType::try_from("invalid").unwrap().to_string(), + RatelimitRuleType::Exact.to_string() + ); + assert_eq!(RatelimitRuleType::Exact.to_string(), "exact"); + assert_eq!(RatelimitRuleType::Regex.to_string(), "regex"); + } + + static INIT: Once = Once::new(); + + async fn setup_test_db(mock_db: sea_orm::DatabaseConnection) { + INIT.call_once(|| { + // Initialize ORM_CLIENT only once + ORM_CLIENT + .set(mock_db) + .expect("Failed to set mock database"); + }); + } + + #[tokio::test] + async fn test_add_batch_empty() { + let result = add_batch(vec![]).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_add_batch_missing_rule_id() { + let rules = vec![RatelimitRule { + org: "test_org".to_string(), + rule_id: None, + ..Default::default() + }]; + + // Create mock database + let db = MockDatabase::new(DatabaseBackend::Postgres).into_connection(); + + // Setup test database + setup_test_db(db).await; + + let result = add_batch(rules).await; + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("Rule ID is required") + ); + } +} diff --git a/src/job/mod.rs b/src/job/mod.rs index dd20eeb877..b1531d3e06 100644 --- a/src/job/mod.rs +++ b/src/job/mod.rs @@ -99,6 +99,18 @@ pub async fn init() -> Result<(), anyhow::Error> { // Auth auditing should be done by router also #[cfg(feature = "enterprise")] tokio::task::spawn(async move { self_reporting::run_audit_publish().await }); + #[cfg(feature = "enterprise")] + { + tokio::task::spawn(async move { db::ofga::watch().await }); + db::ofga::cache().await.expect("ofga model cache failed"); + o2_openfga::authorizer::authz::init_open_fga().await; + // RBAC model + if get_openfga_config().enabled { + if let Err(e) = crate::common::infra::ofga::init().await { + log::error!("OFGA init failed: {}", e); + } + } + } tokio::task::spawn(async move { promql_self_consume::run().await }); // Router doesn't need to initialize job @@ -131,8 +143,6 @@ pub async fn init() -> Result<(), anyhow::Error> { tokio::task::spawn(async move { db::dashboards::reports::watch().await }); tokio::task::spawn(async move { db::organization::watch().await }); tokio::task::spawn(async move { db::pipeline::watch().await }); - #[cfg(feature = "enterprise")] - tokio::task::spawn(async move { db::ofga::watch().await }); #[cfg(feature = "enterprise")] if LOCAL_NODE.is_ingester() || LOCAL_NODE.is_querier() { @@ -180,9 +190,6 @@ pub async fn init() -> Result<(), anyhow::Error> { infra_file_list::create_table_index().await?; infra_file_list::LOCAL_CACHE.create_table_index().await?; - #[cfg(feature = "enterprise")] - db::ofga::cache().await.expect("ofga model cache failed"); - #[cfg(feature = "enterprise")] if !LOCAL_NODE.is_compactor() { db::session::cache() @@ -215,22 +222,11 @@ pub async fn init() -> Result<(), anyhow::Error> { async move { o2_enterprise::enterprise::pipeline::pipeline_job::run().await }, ); - #[cfg(feature = "enterprise")] - o2_openfga::authorizer::authz::init_open_fga().await; - #[cfg(feature = "enterprise")] tokio::task::spawn(async move { cipher::run().await }); #[cfg(feature = "enterprise")] tokio::task::spawn(async move { db::keys::watch().await }); - // RBAC model - #[cfg(feature = "enterprise")] - if get_openfga_config().enabled { - if let Err(e) = crate::common::infra::ofga::init().await { - log::error!("OFGA init failed: {}", e); - } - } - // Shouldn't serve request until initialization finishes log::info!("Job initialization complete"); diff --git a/src/main.rs b/src/main.rs index 014262e3d0..1057295c28 100644 --- a/src/main.rs +++ b/src/main.rs @@ -685,10 +685,18 @@ async fn init_http_server() -> Result<(), anyhow::Error> { if config::cluster::LOCAL_NODE.is_router() { let http_client = router::http::create_http_client().expect("Failed to create http tls client"); + let factory = web::scope(&cfg.common.base_uri); + #[cfg(feature = "enterprise")] + let factory = factory.wrap( + o2_ratelimit::middleware::RateLimitController::new_with_extractor(Some( + router::ratelimit::resource_extractor::default_extractor, + )), + ); + app = app .service( // if `cfg.common.base_uri` is empty, scope("") still works as expected. - web::scope(&cfg.common.base_uri) + factory .wrap(middlewares::SlowLog::new( cfg.limit.http_slow_log_threshold, cfg.limit.circuit_breaker_enabled, @@ -797,10 +805,18 @@ async fn init_http_server_without_tracing() -> Result<(), anyhow::Error> { if config::cluster::LOCAL_NODE.is_router() { let http_client = router::http::create_http_client().expect("Failed to create http tls client"); + let factory = web::scope(&cfg.common.base_uri); + #[cfg(feature = "enterprise")] + let factory = factory.wrap( + o2_ratelimit::middleware::RateLimitController::new_with_extractor(Some( + router::ratelimit::resource_extractor::default_extractor, + )), + ); + app = app .service( // if `cfg.common.base_uri` is empty, scope("") still works as expected. - web::scope(&cfg.common.base_uri) + factory .wrap(middlewares::SlowLog::new( cfg.limit.http_slow_log_threshold, cfg.limit.circuit_breaker_enabled, @@ -1176,5 +1192,10 @@ async fn init_enterprise() -> Result<(), anyhow::Error> { } o2_enterprise::enterprise::pipeline::pipeline_file_server::PipelineFileServer::run().await?; + if config::get_config().ratelimit.ratelimit_enabled && o2_openfga::config::get_config().enabled + { + o2_ratelimit::init(openobserve::handler::http::router::openapi::openapi_info().await) + .await?; + } Ok(()) } diff --git a/src/router/http/mod.rs b/src/router/http/mod.rs index 42378fd50d..8aef4cd5ca 100644 --- a/src/router/http/mod.rs +++ b/src/router/http/mod.rs @@ -19,6 +19,7 @@ use ::config::{ cluster::{Role, RoleGroup}, promql::RequestRangeQuery, }, + router::{is_fixed_querier_route, is_querier_route, is_querier_route_by_body}, utils::rand::get_rand_element, }; use actix_web::{ @@ -34,51 +35,6 @@ use crate::common::{infra::cluster, utils::http::get_search_type_from_request}; mod ws; pub(crate) mod ws_v2; -/// usize indicates the number of parts to skip based on their actual paths. -const QUERIER_ROUTES: [(&str, usize); 21] = [ - ("config", 0), // /config - ("summary", 2), // /api/{org_id}/summary - ("organizations", 1), // /api/organizations - ("settings", 2), // /api/{org_id}/settings/... - ("schema", 3), // /api/{org_id}/streams/{stream_name}/schema - ("streams", 2), // /api/{org_id}/streams/... - ("traces/latest", 3), // /api/{org_id}/{stream_name}/traces/latest - ("clusters", 1), // /api/clusters - ("query_manager", 2), // /api/{org_id}/query_manager/... - ("ws", 2), // /api/{org_id}/ws - ("_search", 2), // /api/{org_id}/_search - ("_around", 3), // /api/{org_id}/{stream_name}/_around - ("_values", 3), // /api/{org_id}/{stream_name}/_values - ("functions?page_num=", 2), // /api/{org_id}/functions - ("prometheus/api/v1/series", 2), // /api/{org_id}/prometheus/api/v1/series - ("prometheus/api/v1/query", 2), // /api/{org_id}/prometheus/api/v1/query - ("prometheus/api/v1/query_range", 2), // /api/{org_id}/prometheus/api/v1/query_range - ("prometheus/api/v1/query_exemplars", 2), // /api/{org_id}/prometheus/api/v1/query_exemplars - ("prometheus/api/v1/metadata", 2), // /api/{org_id}/prometheus/api/v1/metadata - ("prometheus/api/v1/labels", 2), // /api/{org_id}/prometheus/api/v1/labels - ("prometheus/api/v1/label/", 2), /* /api/{org_id}/prometheus/api/v1/label/ - * {label_name}/ - * values */ -]; -const QUERIER_ROUTES_BY_BODY: [&str; 2] = [ - "/prometheus/api/v1/query_range", - "/prometheus/api/v1/query_exemplars", -]; -const FIXED_QUERIER_ROUTES: [&str; 3] = ["/summary", "/schema", "/streams"]; -const INGESTER_ROUTES: [&str; 11] = [ - "/_bulk", - "/_multi", - "/_json", - "/_kinesis_firehose", - "/_sub", - "/v1/logs", - "/ingest/metrics/_json", - "/v1/metrics", - "/traces", - "/v1/traces", - "/traces/latest", -]; - struct URLDetails { is_error: bool, error: Option, @@ -87,39 +43,6 @@ struct URLDetails { node_addr: String, } -#[inline] -fn is_querier_route(path: &str) -> bool { - QUERIER_ROUTES.iter().any(|(route, skip_segments)| { - if path.contains(route) { - let segments = path - .split('/') - .filter(|s| !s.is_empty()) - .collect::>(); - // check if we have enough segments - if segments.len() <= *skip_segments { - return false; - } - let route_part = segments[*skip_segments..].join("/"); - route_part.starts_with(route) - && INGESTER_ROUTES - .iter() - .all(|ingest_route| !route_part.ends_with(ingest_route)) - } else { - false - } - }) -} - -#[inline] -fn is_querier_route_by_body(path: &str) -> bool { - QUERIER_ROUTES_BY_BODY.iter().any(|x| path.contains(x)) -} - -#[inline] -fn is_fixed_querier_route(path: &str) -> bool { - FIXED_QUERIER_ROUTES.iter().any(|x| path.contains(x)) -} - #[route( "/config", method = "GET", diff --git a/src/router/mod.rs b/src/router/mod.rs index d540e4faa1..b5836efee9 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -15,3 +15,5 @@ pub mod grpc; pub mod http; +#[cfg(feature = "enterprise")] +pub mod ratelimit; diff --git a/src/router/ratelimit/mod.rs b/src/router/ratelimit/mod.rs new file mode 100644 index 0000000000..38ee9cf12a --- /dev/null +++ b/src/router/ratelimit/mod.rs @@ -0,0 +1,16 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +pub mod resource_extractor; diff --git a/src/router/ratelimit/resource_extractor.rs b/src/router/ratelimit/resource_extractor.rs new file mode 100644 index 0000000000..52cca6cd0c --- /dev/null +++ b/src/router/ratelimit/resource_extractor.rs @@ -0,0 +1,615 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::collections::HashMap; + +use actix_http::Method; +use actix_web::dev::ServiceRequest; +use config::meta::ratelimit::{RatelimitRule, RatelimitRuleType, get_resource_from_params}; +use futures_util::future::BoxFuture; +use o2_ratelimit::{ + dataresource::{ + db::RATELIMIT_RULES_CACHE, + default_rules::{ + DEFAULT_GLOBAL_ORG_IDENTIFIER, DEFAULT_GLOBAL_USER_ID_IDENTIFIER, + DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER, find_group_by_openapi, + }, + }, + middleware::{ExtractorRule, ExtractorRuleResult}, +}; +use regex::Regex; +use utoipa::OpenApi; + +use crate::{ + common::utils::auth::extract_auth_str, + handler::http::{ + auth::validator::get_user_email_from_auth_str, + request::ratelimit::QUOTA_PAGE_GLOBAL_RULES_ORG, router::openapi::ApiDoc, + }, +}; + +fn extract_org_id(path: &str) -> String { + let path_columns = path.split('/').collect::>(); + if path_columns.len() <= 1 { + return QUOTA_PAGE_GLOBAL_RULES_ORG.to_string(); + } + + if path.starts_with("v1/") || path.starts_with("v2/") || path.starts_with("v3/") { + path_columns + .get(1) + .map(|s| s.to_string()) + .unwrap_or_default() + } else { + path_columns + .first() + .map(|s| s.to_string()) + .unwrap_or_default() + } +} +pub fn default_extractor(req: &ServiceRequest) -> BoxFuture<'_, ExtractorRuleResult> { + let auth_str = extract_auth_str(req.request()); + let local_path = req.path().to_string(); + let path = match local_path + .strip_prefix(format!("{}/api/", config::get_config().common.base_uri).as_str()) + { + Some(path) => path, + None => &local_path, + }; + + let (path, org_id) = (path.to_string(), extract_org_id(path)); + + let openapi_path = extract_openapi_path(req); + Box::pin(async move { + let user_email = get_user_email_from_auth_str(&auth_str) + .await + .unwrap_or_default(); + let user_roles = + crate::service::users::get_user_roles(user_email.as_str(), Some(&org_id)).await; + log::debug!("found user_roles: {:?}", &user_roles); + let openapi_path = openapi_path.unwrap_or(path); + // find the group of the openapi_path + match find_group_by_openapi(&openapi_path).await { + Some(group) => { + let (api_group_name, api_group_operation) = group; + let default_resource = get_resource_from_params( + DEFAULT_GLOBAL_ORG_IDENTIFIER, + RatelimitRuleType::Exact.to_string().as_str(), + DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER, + api_group_name.as_str(), + api_group_operation.as_str(), + DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER, + ); + let (default_rule, custom_rules) = find_default_and_custom_rules( + &org_id, + user_roles.clone(), + api_group_name.as_str(), + api_group_operation.as_str(), + &user_email, + default_resource.as_str(), + ) + .await; + log::debug!( + "org_id{}, user_roles:{:?}, api_group_name:{}, api_group_operation:{}, user_email:{}, default_resource:{}, found default_rule: {:?}, custom_rules: {:?}", + &org_id, + user_roles, + api_group_name, + api_group_operation.as_str(), + user_email, + default_resource, + &default_rule, + &custom_rules + ); + // Select the final rule based on priority + select_final_rule_resource(default_rule, custom_rules, user_email.as_str()) + } + None => ( + ExtractorRule::DefaultOrgGlobal(None), + ExtractorRule::OrgLevel(None), + ExtractorRule::UserRole(None), + ExtractorRule::UserId(None), + ), + } + }) +} + +async fn find_default_and_custom_rules( + org_id: &str, + user_roles: Vec, + api_group_name: &str, + api_group_operation: &str, + user_id: &str, + rule_resource: &str, +) -> (Option, Vec) { + let default_rule = find_matching_rule(org_id, rule_resource, true).await; + + // custom_rules will includes custom org rules and custom role rules all + let mut custom_rules = vec![]; + + // Find role-level and user-level rules + for role in user_roles.iter() { + // user-level + let resource = get_resource_from_params( + org_id, + RatelimitRuleType::Exact.to_string().as_str(), + role, + api_group_name, + api_group_operation, + user_id, + ); + let custom_rule = find_matching_rule(org_id, &resource, false).await; + if let Some(rule) = custom_rule { + custom_rules.push(rule) + } + + // role-level + let resource = get_resource_from_params( + org_id, + RatelimitRuleType::Exact.to_string().as_str(), + role, + api_group_name, + api_group_operation, + DEFAULT_GLOBAL_USER_ID_IDENTIFIER, + ); + let custom_rule = find_matching_rule(org_id, &resource, false).await; + if let Some(rule) = custom_rule { + custom_rules.push(rule) + } + } + + // Find org-level rule + let org_level_resource = get_resource_from_params( + org_id, + RatelimitRuleType::Exact.to_string().as_str(), + DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER, + api_group_name, + api_group_operation, + DEFAULT_GLOBAL_USER_ID_IDENTIFIER, + ); + if let Some(org_level_rule) = + find_matching_rule(org_id, org_level_resource.as_str(), false).await + { + custom_rules.push(org_level_rule) + } + + (default_rule, custom_rules) +} + +fn select_final_rule_resource( + default_rule: Option, + custom_rules: Vec, + user_email: &str, +) -> ExtractorRuleResult { + if custom_rules.is_empty() { + return ( + ExtractorRule::DefaultOrgGlobal(default_rule), + ExtractorRule::OrgLevel(None), + ExtractorRule::UserRole(None), + ExtractorRule::UserId(None), + ); + } + + // Find org-level and role-level rules + let org_level_rule = custom_rules + .iter() + .find(|rule| { + rule.user_role + .eq(&Some(DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER.to_string())) + && rule + .user_id + .eq(&Some(DEFAULT_GLOBAL_USER_ID_IDENTIFIER.to_string())) + }) + .cloned(); + + let user_role_rule = custom_rules + .iter() + .filter(|rule| { + rule.user_role + .ne(&Some(DEFAULT_GLOBAL_USER_ROLE_IDENTIFIER.to_string())) + && rule + .user_id + .eq(&Some(DEFAULT_GLOBAL_USER_ID_IDENTIFIER.to_string())) + }) + .max_by_key(|rule| rule.threshold) + .cloned(); + + let user_role = if let Some(user_role_rule) = &user_role_rule { + user_role_rule.user_role.clone() + } else { + None + }; + + let user_id_rules = custom_rules + .iter() + .find(|rule| { + user_role_rule.is_some() + && rule.user_role.eq(&user_role) + && rule.user_id.eq(&Some(user_email.to_string())) + }) + .cloned(); + + ( + ExtractorRule::DefaultOrgGlobal(default_rule), + ExtractorRule::OrgLevel(org_level_rule), + ExtractorRule::UserRole(user_role_rule), + ExtractorRule::UserId(user_id_rules), + ) +} + +async fn find_matching_rule( + org: &str, + resource: &str, + search_default: bool, +) -> Option { + let rules = RATELIMIT_RULES_CACHE.read().await; + + let filter_rules: Vec = if search_default { + rules.iter().filter(|r| r.org == ".*").cloned().collect() + } else { + rules.iter().filter(|r| r.org == org).cloned().collect() + }; + drop(rules); + + for rule in filter_rules { + match rule.rule_type.as_deref() { + Some(type_str) if type_str == RatelimitRuleType::Exact.to_string() => { + if resource == rule.get_resource() { + log::debug!( + "Matching resource: {} with exact rule: {}", + resource, + rule.get_resource() + ); + return Some(rule); + } + } + Some(type_str) if type_str == RatelimitRuleType::Regex.to_string() => { + if let Ok(re) = Regex::new(&rule.get_resource()) { + if re.is_match(resource) { + log::debug!( + "Matching resource: {} with regex rule: {}", + resource, + rule.get_resource() + ); + return Some(rule); + } + } + } + _ => continue, + } + } + + None +} + +fn path_to_regex(path: &str) -> Regex { + let re = Regex::new(r"\{[^}]+\}").unwrap(); + let replaced = re.replace_all(path, r"([^/]+)"); + Regex::new(&format!("^{}$", replaced)).unwrap() +} + +fn extract_openapi_path(req: &ServiceRequest) -> Option { + let api = ApiDoc::openapi(); + let path_patterns: HashMap = { + api.paths + .paths + .keys() + .map(|path| { + let pattern = path_to_regex(path); + (path.clone(), pattern) + }) + .collect() + }; + + let path = req.path(); + let method = req.method(); + + // Find matching path + for (openapi_path, path_item) in &api.paths.paths { + // Check if path pattern matches + if let Some(pattern) = path_patterns.get(openapi_path) { + if !pattern.is_match(path) { + continue; + } + + // Check if method is supported for this path + let method_exists = path_item.operations.iter().any(|(op_method, _)| { + matches!( + (method, op_method), + (&Method::GET, utoipa::openapi::PathItemType::Get) + | (&Method::POST, utoipa::openapi::PathItemType::Post) + | (&Method::PUT, utoipa::openapi::PathItemType::Put) + | (&Method::DELETE, utoipa::openapi::PathItemType::Delete) + | (&Method::PATCH, utoipa::openapi::PathItemType::Patch) + | (&Method::HEAD, utoipa::openapi::PathItemType::Head) + | (&Method::OPTIONS, utoipa::openapi::PathItemType::Options) + | (&Method::TRACE, utoipa::openapi::PathItemType::Trace) + ) + }); + + if method_exists { + log::debug!( + "Matched OpenAPI path: {} for request path: {} and method: {}", + openapi_path, + path, + method + ); + return Some(format!("{}:{}", openapi_path, method)); + } else { + log::debug!( + "Path matched but method {} not supported for OpenAPI path: {}", + method, + openapi_path + ); + } + } + } + + None +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use utoipa::{OpenApi, openapi::path::Operation}; + + use super::*; + use crate::handler::http::router::openapi::ApiDoc; + + #[tokio::test] + async fn print_openapi_info_with() { + let api = ApiDoc::openapi(); + + println!("\nOpenAPI Documentation Info:"); + println!("=========================="); + + // Print general info + println!("Title: {}", api.info.title); + println!("Version: {}", api.info.version); + if let Some(description) = &api.info.description { + println!("Description: {}", description); + } + + // Print all tags with descriptions + if let Some(tags) = &api.tags { + println!("\nGlobal Tags:"); + println!("============"); + for tag in tags { + println!("Tag: {}", tag.name); + if let Some(desc) = &tag.description { + println!("Description: {}", desc); + } + if let Some(ext_docs) = &tag.external_docs { + println!( + "External Docs: {} ({})", + ext_docs.url, + ext_docs.description.as_deref().unwrap_or("") + ); + } + println!(); + } + } + + println!("\nAPI Endpoints:"); + println!("=============="); + + // Group endpoints by tags with full operation details + let mut tag_operations: HashMap)>> = + HashMap::new(); + + for (path, path_item) in &api.paths.paths { + for (method, operation) in path_item.operations.clone() { + let tags = operation + .tags + .clone() + .unwrap_or_else(|| vec!["untagged".to_string()]); + + let method = match method { + utoipa::openapi::PathItemType::Get => "GET", + utoipa::openapi::PathItemType::Post => "POST", + utoipa::openapi::PathItemType::Put => "PUT", + utoipa::openapi::PathItemType::Delete => "DELETE", + utoipa::openapi::PathItemType::Patch => "PATCH", + utoipa::openapi::PathItemType::Head => "HEAD", + utoipa::openapi::PathItemType::Options => "OPTIONS", + utoipa::openapi::PathItemType::Trace => "TRACE", + utoipa::openapi::PathItemType::Connect => "CONNECT", + }; + + let operation_info = ( + method.to_string(), + path.clone(), + operation.clone(), + tags.clone(), + ); + + for tag in tags { + tag_operations + .entry(tag) + .or_insert_with(Vec::new) + .push(operation_info.clone()); + } + } + } + + // Print grouped operations with detailed information + for (tag, operations) in tag_operations.iter() { + println!("\n[{}]", tag); + println!("{}", "=".repeat(tag.len() + 2)); + + for (method, path, operation, _) in operations { + println!("\n{} {}", method, path); + + // Print operation ID if available + if let Some(operation_id) = &operation.operation_id { + println!("Operation ID: {}", operation_id); + } + + // Print operation summary and description + if let Some(summary) = &operation.summary { + println!("Summary: {}", summary); + } + if let Some(description) = &operation.description { + println!("Description: {}", description); + } else { + println!("No description available"); + } + + println!("\n{}", "-".repeat(50)); + } + } + } + + fn create_rule( + org: &str, + user_role: Option<&str>, + user_id: Option<&str>, + threshold: i32, + ) -> RatelimitRule { + RatelimitRule { + org: org.to_string(), + rule_type: Some("exact".to_string()), + user_role: user_role.map(String::from), + user_id: user_id.map(String::from), + threshold, + ..Default::default() + } + } + + #[test] + fn test_select_final_rule_empty_rules() { + let result = select_final_rule_resource(None, vec![], "test@example.com"); + + assert!(matches!(result.0, ExtractorRule::DefaultOrgGlobal(None))); + assert!(matches!(result.1, ExtractorRule::OrgLevel(None))); + assert!(matches!(result.2, ExtractorRule::UserRole(None))); + assert!(matches!(result.3, ExtractorRule::UserId(None))); + } + + #[test] + fn test_select_final_rule_only_default() { + let default_rule = create_rule(".*", None, None, 100); + let result = + select_final_rule_resource(Some(default_rule.clone()), vec![], "test@example.com"); + + assert!(matches!(result.0, ExtractorRule::DefaultOrgGlobal(Some(r)) if r.threshold == 100)); + assert!(matches!(result.1, ExtractorRule::OrgLevel(None))); + assert!(matches!(result.2, ExtractorRule::UserRole(None))); + assert!(matches!(result.3, ExtractorRule::UserId(None))); + } + + #[test] + fn test_select_final_rule_org_level() { + let default_rule = create_rule(".*", None, None, 100); + let org_rule = create_rule("test_org", None, None, 50); + let custom_rules = vec![org_rule.clone()]; + + let result = select_final_rule_resource( + Some(default_rule.clone()), + custom_rules, + "test@example.com", + ); + + assert!(matches!(result.0, ExtractorRule::DefaultOrgGlobal(Some(r)) if r.threshold == 100)); + assert!(matches!(result.1, ExtractorRule::OrgLevel(Some(r)) if r.threshold == 50)); + assert!(matches!(result.2, ExtractorRule::UserRole(None))); + assert!(matches!(result.3, ExtractorRule::UserId(None))); + } + + #[test] + fn test_select_final_rule_user_role() { + let default_rule = create_rule(".*", None, None, 100); + let role_rule = create_rule("test_org", Some("admin"), Some(".*"), 75); + let custom_rules = vec![role_rule.clone()]; + + let result = select_final_rule_resource( + Some(default_rule.clone()), + custom_rules, + "test@example.com", + ); + + assert!(matches!(result.0, ExtractorRule::DefaultOrgGlobal(Some(r)) if r.threshold == 100)); + assert!(matches!(result.1, ExtractorRule::OrgLevel(None))); + assert!(matches!(result.2, ExtractorRule::UserRole(Some(r)) if r.threshold == 75)); + assert!(matches!(result.3, ExtractorRule::UserId(None))); + } + + #[test] + fn test_select_final_rule_user_id() { + let default_rule = create_rule(".*", None, None, 100); + let role_rule = create_rule("test_org", Some("admin"), Some(".*"), 75); + let user_rule = create_rule("test_org", Some("admin"), Some("test@example.com"), 25); + let custom_rules = vec![role_rule.clone(), user_rule.clone()]; + + let result = select_final_rule_resource( + Some(default_rule.clone()), + custom_rules, + "test@example.com", + ); + + assert!(matches!(result.0, ExtractorRule::DefaultOrgGlobal(Some(r)) if r.threshold == 100)); + assert!(matches!(result.1, ExtractorRule::OrgLevel(None))); + assert!(matches!(result.2, ExtractorRule::UserRole(Some(r)) if r.threshold == 75)); + assert!(matches!(result.3, ExtractorRule::UserId(Some(r)) if r.threshold == 25)); + } + + #[test] + fn test_select_final_rule_all_levels() { + let default_rule = create_rule(".*", None, None, 100); + let org_rule = create_rule("test_org", None, None, 80); + let role_rule = create_rule("test_org", Some("admin"), Some(".*"), 60); + let user_rule = create_rule("test_org", Some("admin"), Some("test@example.com"), 40); + let custom_rules = vec![org_rule.clone(), role_rule.clone(), user_rule.clone()]; + + let result = select_final_rule_resource( + Some(default_rule.clone()), + custom_rules, + "test@example.com", + ); + + assert!(matches!(result.0, ExtractorRule::DefaultOrgGlobal(Some(r)) if r.threshold == 100)); + assert!(matches!(result.1, ExtractorRule::OrgLevel(Some(r)) if r.threshold == 80)); + assert!(matches!(result.2, ExtractorRule::UserRole(Some(r)) if r.threshold == 60)); + assert!(matches!(result.3, ExtractorRule::UserId(Some(r)) if r.threshold == 40)); + } + + #[test] + fn test_select_final_rule_multiple_role_rules() { + let role_rule1 = create_rule("test_org", Some("admin"), Some(".*"), 75); + let role_rule2 = create_rule("test_org", Some("user"), Some(".*"), 50); + let custom_rules = vec![role_rule1.clone(), role_rule2.clone()]; + + let result = select_final_rule_resource(None, custom_rules, "test@example.com"); + + // Should select the rule with higher threshold + assert!(matches!(result.0, ExtractorRule::DefaultOrgGlobal(None))); + assert!(matches!(result.1, ExtractorRule::OrgLevel(None))); + assert!(matches!(result.2, ExtractorRule::UserRole(Some(r)) if r.threshold == 75)); + assert!(matches!(result.3, ExtractorRule::UserId(None))); + } + + #[test] + fn test_select_final_rule_non_matching_user_id() { + let role_rule = create_rule("test_org", Some("admin"), Some(".*"), 75); + let user_rule = create_rule("test_org", Some("admin"), Some("other@example.com"), 25); + let custom_rules = vec![role_rule.clone(), user_rule.clone()]; + + let result = select_final_rule_resource(None, custom_rules, "test@example.com"); + + assert!(matches!(result.0, ExtractorRule::DefaultOrgGlobal(None))); + assert!(matches!(result.1, ExtractorRule::OrgLevel(None))); + assert!(matches!(result.2, ExtractorRule::UserRole(Some(r)) if r.threshold == 75)); + assert!(matches!(result.3, ExtractorRule::UserId(None))); + } +} diff --git a/src/service/mod.rs b/src/service/mod.rs index 89616206d8..81d02b4a56 100644 --- a/src/service/mod.rs +++ b/src/service/mod.rs @@ -37,6 +37,8 @@ pub mod node; pub mod organization; pub mod pipeline; pub mod promql; +#[cfg(feature = "enterprise")] +pub mod ratelimit; pub mod schema; pub mod search; pub mod websocket_events; diff --git a/src/service/ratelimit/mod.rs b/src/service/ratelimit/mod.rs new file mode 100644 index 0000000000..f3c58ca78a --- /dev/null +++ b/src/service/ratelimit/mod.rs @@ -0,0 +1,16 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +pub mod rule; diff --git a/src/service/ratelimit/rule.rs b/src/service/ratelimit/rule.rs new file mode 100644 index 0000000000..5d8690a35c --- /dev/null +++ b/src/service/ratelimit/rule.rs @@ -0,0 +1,162 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::future::Future; + +use anyhow::Context; +use config::meta::ratelimit::RatelimitRule; +use infra::table::ratelimit::RuleEntry; +use o2_enterprise::enterprise::super_cluster::queue::ratelimit::{ + SUPER_CLUSTER_RATELIMIT_KEY_PREFIX, ratelimit_rule_delete, ratelimit_rule_put, + ratelimit_rule_update, +}; + +#[derive(Debug, thiserror::Error)] +pub enum RatelimitError { + #[error("ratelimit rule with ID {0} not found.")] + NotFound(String), + #[error("ratelimit rule entry type not supported: {0}")] + NotSupportRuleEntry(String), + #[error(transparent)] + DbError(#[from] anyhow::Error), +} + +#[derive(Debug)] +enum RuleOperation { + Save, + Update, + Delete, +} + +impl RuleOperation { + fn as_str(&self) -> &'static str { + match self { + RuleOperation::Save => "add", + RuleOperation::Update => "update", + RuleOperation::Delete => "delete", + } + } + + async fn execute_cluster_operation( + &self, + key: String, + value: bytes::Bytes, + ) -> Result<(), anyhow::Error> { + match self { + RuleOperation::Save => ratelimit_rule_put(key, value).await, + RuleOperation::Update => ratelimit_rule_update(key, value).await, + RuleOperation::Delete => ratelimit_rule_delete(key, value).await, + } + } +} + +async fn handle_rule_operation( + rule: RuleEntry, + operation: RuleOperation, + db_operation: impl Future>, +) -> Result<(), RatelimitError> { + db_operation.await.map_err(RatelimitError::DbError)?; + + #[cfg(feature = "enterprise")] + if o2_enterprise::enterprise::common::infra::config::get_config() + .super_cluster + .enabled + { + if let Err(e) = sync_to_super_cluster(&rule, &operation).await { + log::error!( + "[Ratelimit] error triggering super cluster event to {} rule: {}", + operation.as_str(), + e + ); + } + } + + Ok(()) +} + +async fn sync_to_super_cluster( + rule: &RuleEntry, + operation: &RuleOperation, +) -> Result<(), anyhow::Error> { + let value_vec = config::utils::json::to_vec(&rule) + .context("[Ratelimit] error serializing ratelimit rule for super_cluster event")?; + + let key = format!( + "/{SUPER_CLUSTER_RATELIMIT_KEY_PREFIX}/{}", + operation.as_str(), + ); + + operation + .execute_cluster_operation(key, value_vec.into()) + .await +} + +pub async fn save(rule: RuleEntry) -> Result<(), RatelimitError> { + handle_rule_operation( + rule.clone(), + RuleOperation::Save, + infra::table::ratelimit::add(rule), + ) + .await +} + +pub async fn save_batch(rules: RuleEntry) -> Result<(), RatelimitError> { + let RuleEntry::UpsertBatch(_) = &rules else { + return Err(RatelimitError::NotSupportRuleEntry( + "Expected upsert batch rule entry".to_string(), + )); + }; + + handle_rule_operation( + rules.clone(), + RuleOperation::Save, + infra::table::ratelimit::add(rules), + ) + .await +} + +pub async fn update(rule: RuleEntry) -> Result<(), RatelimitError> { + handle_rule_operation( + rule.clone(), + RuleOperation::Update, + infra::table::ratelimit::update(rule), + ) + .await +} + +pub async fn delete(rule: RuleEntry) -> Result<(), RatelimitError> { + let RuleEntry::Single(r) = &rule else { + return Err(RatelimitError::NotSupportRuleEntry( + "Batch rule entry not supported".to_string(), + )); + }; + + handle_rule_operation( + rule.clone(), + RuleOperation::Delete, + infra::table::ratelimit::delete(r.rule_id.clone().unwrap()), + ) + .await +} + +pub async fn list( + org: &str, + api: Option<&str>, + user_id: Option<&str>, +) -> Result, RatelimitError> { + infra::table::ratelimit::list(org, api, user_id) + .await + .map_err(RatelimitError::DbError) +} diff --git a/src/service/search/cluster/cache_multi.rs b/src/service/search/cluster/cache_multi.rs index e9e9b900f1..787f4c8e96 100644 --- a/src/service/search/cluster/cache_multi.rs +++ b/src/service/search/cluster/cache_multi.rs @@ -37,7 +37,7 @@ pub async fn get_cached_results( cache_req: CacheQueryRequest, ) -> Vec { // get nodes from cluster - let mut nodes = match infra_cluster::get_cached_online_query_nodes(None).await { + let mut nodes = match infra_cluster::get_cached_online_querier_nodes(None).await { Some(nodes) => nodes, None => { log::error!("[trace_id {trace_id}] get_cached_results: no querier node online"); diff --git a/src/service/search/cluster/cacher.rs b/src/service/search/cluster/cacher.rs index 0bb9f5f3e0..1b41d323e4 100644 --- a/src/service/search/cluster/cacher.rs +++ b/src/service/search/cluster/cacher.rs @@ -36,7 +36,7 @@ pub async fn get_cached_results( ) -> Option { let start = std::time::Instant::now(); // get nodes from cluster - let mut nodes = match infra_cluster::get_cached_online_query_nodes(None).await { + let mut nodes = match infra_cluster::get_cached_online_querier_nodes(None).await { Some(nodes) => nodes, None => { log::error!("[trace_id {trace_id}] get_cached_results: no querier node online"); @@ -257,7 +257,7 @@ pub async fn delete_cached_results(path: String) -> bool { let trace_id = path.clone(); let mut delete_response = true; // get nodes from cluster - let mut nodes = match infra_cluster::get_cached_online_query_nodes(None).await { + let mut nodes = match infra_cluster::get_cached_online_querier_nodes(None).await { Some(nodes) => nodes, None => { log::error!("[trace_id {trace_id}] delete_cached_results: no querier node online"); diff --git a/src/service/search/grpc_search.rs b/src/service/search/grpc_search.rs index 981923afde..a6a5b33998 100644 --- a/src/service/search/grpc_search.rs +++ b/src/service/search/grpc_search.rs @@ -37,9 +37,9 @@ pub async fn grpc_search( in_req: &search::Request, node_group: Option, ) -> Result { - let mut nodes = match infra_cluster::get_cached_online_query_nodes(node_group).await { - Some(nodes) => nodes, - None => { + let mut nodes = match infra_cluster::get_cached_online_querier_nodes(node_group).await { + Some(nodes) if !nodes.is_empty() => nodes, + _ => { log::error!("search->grpc: no querier node online"); return Err(server_internal_error("no querier node online")); } @@ -109,7 +109,7 @@ pub async fn grpc_search_partition( node_group: Option, skip_max_query_range: bool, ) -> Result { - let mut nodes = match infra_cluster::get_cached_online_query_nodes(node_group).await { + let mut nodes = match infra_cluster::get_cached_online_querier_nodes(node_group).await { Some(nodes) => nodes, None => { log::error!("search->grpc: no querier node online"); diff --git a/src/service/users.rs b/src/service/users.rs index 725218efa7..80e0ab0b5e 100644 --- a/src/service/users.rs +++ b/src/service/users.rs @@ -14,8 +14,12 @@ // along with this program. If not, see . use std::io::Error; +#[cfg(feature = "enterprise")] +use std::time::{Duration, Instant}; use actix_web::{HttpResponse, http}; +#[cfg(feature = "enterprise")] +use config::meta::ratelimit::CachedUserRoles; use config::{get_config, ider, utils::rand::generate_random_string}; #[cfg(feature = "enterprise")] use o2_openfga::{ @@ -23,6 +27,8 @@ use o2_openfga::{ }; use regex::Regex; +#[cfg(feature = "enterprise")] +use crate::common::infra::config::USER_ROLES_CACHE; use crate::{ common::{ infra::config::{ROOT_USER, USERS, USERS_RUM_TOKEN}, @@ -832,6 +838,56 @@ pub(crate) async fn create_root_user(org_id: &str, user_req: UserRequest) -> Res Ok(()) } +#[cfg(feature = "enterprise")] +pub async fn get_user_roles(user_email: &str, org_id: Option<&str>) -> Vec { + if let Some(roles) = check_cache(user_email).await { + return get_user_roles_by_org_id(roles, org_id); + } + + let roles = o2_openfga::authorizer::roles::get_roles_for_user(user_email).await; + update_cache(user_email, roles.clone()).await; + + get_user_roles_by_org_id(roles, org_id) +} +#[cfg(feature = "enterprise")] +fn get_user_roles_by_org_id(roles: Vec, org_id: Option<&str>) -> Vec { + match org_id { + Some(org_id) => roles + .iter() + .filter_map(|role| { + let parts: Vec<&str> = role.split('/').collect(); + if parts.first() == Some(&org_id) { + parts.get(1).map(|s| s.to_string()) + } else { + None + } + }) + .collect(), + None => roles, + } +} +#[cfg(feature = "enterprise")] +async fn check_cache(user_email: &str) -> Option> { + let cache = USER_ROLES_CACHE.read().await; + if let Some(cached) = cache.get(user_email) { + if cached.expires_at > Instant::now() { + return Some(cached.roles.clone()); + } + } + None +} +#[cfg(feature = "enterprise")] +async fn update_cache(user_email: &str, roles: Vec) { + let mut cache = USER_ROLES_CACHE.write().await; + cache.insert( + user_email.to_string(), + CachedUserRoles { + roles, + expires_at: Instant::now() + Duration::from_secs(60), + }, + ); +} + #[cfg(test)] mod tests { use infra::db as infra_db; diff --git a/src/super_cluster_queue/mod.rs b/src/super_cluster_queue/mod.rs index 8b3ae7c1e2..063e758c74 100644 --- a/src/super_cluster_queue/mod.rs +++ b/src/super_cluster_queue/mod.rs @@ -22,6 +22,7 @@ mod distinct_values; mod folders; mod meta; mod pipelines; +mod ratelimit; mod scheduler; mod schemas; mod search_job; @@ -49,6 +50,7 @@ pub async fn init() -> Result<(), anyhow::Error> { on_dashboard_msg: dashboards::process, on_pipeline_msg: pipelines::process, on_cipher_key_msg: cipher_keys::process, + on_rate_limit_msg: ratelimit::process, }; let schema_queue = SchemasQueue { on_schema_msg: schemas::process, diff --git a/src/super_cluster_queue/ratelimit.rs b/src/super_cluster_queue/ratelimit.rs new file mode 100644 index 0000000000..1a28319d7a --- /dev/null +++ b/src/super_cluster_queue/ratelimit.rs @@ -0,0 +1,83 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use infra::{ + errors::{Error, Result}, + table::ratelimit::{RULE_EXISTS, RULE_NOT_FOUND, RuleEntry}, +}; +use o2_enterprise::enterprise::super_cluster::queue::{Message, MessageType}; + +pub async fn process(msg: Message) -> Result<()> { + log::debug!( + "[SUPER_CLUSTER:RATELIMIT] LOCAL_NODE:{:?}, Processing message: {:?}", + config::cluster::LOCAL_NODE, + msg.key + ); + + let bytes = msg + .value + .ok_or(Error::Message("Message missing value".to_string()))?; + let rule = RuleEntry::try_from(&bytes).map_err(|e| Error::Message(e.to_string()))?; + + match msg.message_type { + MessageType::RatelimitAdd => match infra::table::ratelimit::add(rule.clone()).await { + Ok(_) => Ok(()), + Err(e) if e.to_string() == RULE_EXISTS => { + log::warn!( + "[SUPER_CLUSTER:RATELIMIT] Rule already exists, ignoring, rule: {:?}", + rule + ); + Ok(()) + } + Err(e) => Err(Error::Message(e.to_string())), + }, + MessageType::RatelimitUpdate => match infra::table::ratelimit::update(rule).await { + Ok(_) => Ok(()), + Err(e) if e.to_string() == RULE_NOT_FOUND => { + log::warn!("[SUPER_CLUSTER:RATELIMIT] Rule not found for update, ignoring"); + Ok(()) + } + Err(e) => Err(Error::Message(e.to_string())), + }, + MessageType::RatelimitDelete => { + let RuleEntry::Single(rule) = rule else { + log::error!("[SUPER_CLUSTER:RATELIMIT] Invalid message: {:?}", rule); + return Err(Error::Message("Invalid rule entry type".to_string())); + }; + + match infra::table::ratelimit::delete( + rule.rule_id + .ok_or_else(|| Error::Message("Missing rule_id".to_string()))?, + ) + .await + { + Ok(_) => Ok(()), + Err(e) if e.to_string() == RULE_NOT_FOUND => { + log::warn!("[SUPER_CLUSTER:RATELIMIT] Rule not found for deletion, ignoring"); + Ok(()) + } + Err(e) => Err(Error::Message(e.to_string())), + } + } + _ => { + log::error!( + "[SUPER_CLUSTER:RATELIMIT] Invalid message: type: {:?}, key: {}", + msg.message_type, + msg.key + ); + Err(Error::Message("Invalid message type".to_string())) + } + } +} diff --git a/web/src/assets/images/common/hard-drive.svg b/web/src/assets/images/common/hard-drive.svg new file mode 100644 index 0000000000..2e4cdbe272 --- /dev/null +++ b/web/src/assets/images/common/hard-drive.svg @@ -0,0 +1,11 @@ + + + + + + + + + + + diff --git a/web/src/assets/images/common/selectOrganization.svg b/web/src/assets/images/common/selectOrganization.svg new file mode 100644 index 0000000000..4628ef7837 --- /dev/null +++ b/web/src/assets/images/common/selectOrganization.svg @@ -0,0 +1,39 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/web/src/components/iam/quota/Quota.vue b/web/src/components/iam/quota/Quota.vue new file mode 100644 index 0000000000..0cd288f8eb --- /dev/null +++ b/web/src/components/iam/quota/Quota.vue @@ -0,0 +1,1641 @@ + + + + + + + \ No newline at end of file diff --git a/web/src/components/iam/roles/PermissionsTable.vue b/web/src/components/iam/roles/PermissionsTable.vue index 7240aa5ff5..b2285d3e17 100644 --- a/web/src/components/iam/roles/PermissionsTable.vue +++ b/web/src/components/iam/roles/PermissionsTable.vue @@ -297,7 +297,7 @@ const columns: any = [ { name: "display_name", field: "display_name", - label: t("common.name"), + label: t("quota.moduleName"), align: "left", sortable: true, }, @@ -335,15 +335,6 @@ const columns: any = [ slotName: "permission", style: { width: "80px" }, }, - { - name: "AllowDelete", - field: "permission", - label: t("iam.delete"), - align: "center", - slot: true, - slotName: "permission", - style: { width: "80px" }, - }, { name: "AllowPost", field: "permission", @@ -362,6 +353,15 @@ const columns: any = [ slotName: "permission", style: { width: "80px" }, }, + { + name: "AllowDelete", + field: "permission", + label: t("iam.delete"), + align: "center", + slot: true, + slotName: "permission", + style: { width: "80px" }, + } ]; const expandPermission = async (resource: any) => { diff --git a/web/src/components/shared/grid/NoOrganizationSelected.vue b/web/src/components/shared/grid/NoOrganizationSelected.vue new file mode 100644 index 0000000000..6acea92464 --- /dev/null +++ b/web/src/components/shared/grid/NoOrganizationSelected.vue @@ -0,0 +1,57 @@ + + + + + + + \ No newline at end of file diff --git a/web/src/composables/shared/useEnterpriseRoutes.ts b/web/src/composables/shared/useEnterpriseRoutes.ts index e91ed982cc..f889c198b4 100644 --- a/web/src/composables/shared/useEnterpriseRoutes.ts +++ b/web/src/composables/shared/useEnterpriseRoutes.ts @@ -28,6 +28,8 @@ const EditRole = () => import("@/components/iam/roles/EditRole.vue"); const EditGroup = () => import("@/components/iam/groups/EditGroup.vue"); +const Quota = () => import("@/components/iam/quota/Quota.vue"); + const Organizations = () => import("@/components/iam/organizations/AppOrganizations.vue"); @@ -117,6 +119,14 @@ const useEnterpriseRoutes = () => { routeGuard(to, from, next); }, }, + { + path: "quota", + name: "quota", + component: Quota , + beforeEnter(to: any, from: any, next: any) { + routeGuard(to, from, next); + }, + }, ] ); } diff --git a/web/src/composables/useIsMetaOrg.ts b/web/src/composables/useIsMetaOrg.ts index fbd47de9f5..2c30aaedce 100644 --- a/web/src/composables/useIsMetaOrg.ts +++ b/web/src/composables/useIsMetaOrg.ts @@ -6,7 +6,7 @@ export default function useIsMetaOrg() { const store = useStore(); const isMetaOrg = computed(() => { - return store.state.selectedOrganization.label === store.state.zoConfig.meta_org || store.state.selectedOrganization.identifier === store.state.zoConfig.meta_org; + return (store.state.selectedOrganization.label === store.state.zoConfig.meta_org || store.state.selectedOrganization.identifier === store.state.zoConfig.meta_org); }); return { diff --git a/web/src/composables/useRateLimiter.ts b/web/src/composables/useRateLimiter.ts new file mode 100644 index 0000000000..22ffb1cb9c --- /dev/null +++ b/web/src/composables/useRateLimiter.ts @@ -0,0 +1,139 @@ +// Copyright 2023 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +import { useStore } from "vuex"; +import { reactive } from "vue"; +import { ref } from "vue"; +import { useQuasar } from "quasar"; +import rateLimiterService from "@/services/rate_limit"; + const useRateLimiter = () => { + const store = useStore(); + const q = useQuasar(); + + const getApiLimitsByOrganization = async ( orgId: string) => { + try { + const response = await rateLimiterService.getApiLimits(orgId); + let transformedData: any = []; + + //predefined operation that we get from the api + const operations = ['list', 'get', 'create', 'update', 'delete']; + // Iterate over each module in api_group_info + Object.keys(response.data).forEach((moduleName) => { + const module = response.data[moduleName]; + + // Create an object to store the threshold values for each operation + let moduleThresholds: any = { + module_name: moduleName, + }; + + operations.forEach((operation) => { + // Check if the operation exists for the current module + if (module[operation.toLowerCase()] !== undefined) { + // If the operation exists, get the threshold value + moduleThresholds[operation.toLowerCase()] = module[operation.toLowerCase()]; + } else { + // If the operation doesn't exist, set it as '--' + moduleThresholds[operation.toLowerCase()] = '-'; + } + }); + // Add the transformed data to the array + transformedData.push(moduleThresholds); + }); + transformedData.sort((a: any, b: any) => a.module_name.localeCompare(b.module_name)); + store.dispatch("setApiLimitsByOrgId", { + ...store.state.allApiLimitsByOrgId, + [orgId]: transformedData, + }); + //this is done because once we update the api limits , we need to reset the role limits so that when we can fetch the latest roles limits from the api (updated one) + store.dispatch("setRoleLimitsByOrgIdByRole", { + ...store.state.allRoleLimitsByOrgIdByRole, + [orgId]: [], + }); + return transformedData; + } catch (error) { + + console.log(error); + } + }; + const getRoleLimitsByOrganization = async (orgId: string, rolename: string) => { + try { + const response = await rateLimiterService.getRoleLimits(orgId, rolename); + let transformedData: any = []; + + //predefined operation that we get from the api + const operations = ['list', 'get', 'create', 'update', 'delete']; + // Iterate over each module in api_group_info + Object.keys(response.data).forEach((moduleName) => { + const module = response.data[moduleName]; + + // Create an object to store the threshold values for each operation + let moduleThresholds: any = { + module_name: moduleName, + }; + + operations.forEach((operation) => { + // Check if the operation exists for the current module + if (module[operation.toLowerCase()] !== undefined) { + // If the operation exists, get the threshold value + moduleThresholds[operation.toLowerCase()] = module[operation.toLowerCase()]; + } else { + // If the operation doesn't exist, set it as '--' + moduleThresholds[operation.toLowerCase()] = '-'; + } + }); + // Add the transformed data to the array + transformedData.push(moduleThresholds); + + }); + transformedData.sort((a: any, b: any) => a.module_name.localeCompare(b.module_name)); + store.dispatch("setRoleLimitsByOrgIdByRole", { + ...store.state.allRoleLimitsByOrgIdByRole, + [orgId]: { + ...store.state.allRoleLimitsByOrgIdByRole[orgId], + [rolename]: transformedData, + }, + }); + return transformedData; + } catch (error) { + console.log(error); + } + }; + + const getModulesToDisplay = async (orgId: string) => { + try { + const response = await rateLimiterService.getModules(orgId); + const modulesToDisplay = response.data.map((role: any) => ({ + label: role, + value: role + })).sort((a: any, b: any) => a.label.localeCompare(b.label)); + //this is done to store the modules to display for the organization + store.dispatch("setModulesToDisplay", { + ...store.state.modulesToDisplay, + [orgId]: modulesToDisplay, + }); + return modulesToDisplay; + } catch (error) { + } + + + } + + return { + getApiLimitsByOrganization, + getRoleLimitsByOrganization, + getModulesToDisplay + } +} +export default useRateLimiter; diff --git a/web/src/locales/languages/en.json b/web/src/locales/languages/en.json index 3443bb6127..bbf682b35c 100644 --- a/web/src/locales/languages/en.json +++ b/web/src/locales/languages/en.json @@ -345,7 +345,8 @@ "searchRole": "Search Role", "searchGroup": "Search Group", "object": "Object", - "serviceAccounts": "Service Accounts" + "serviceAccounts": "Service Accounts", + "quota": "Quotas" }, "ticket": { "header": "Support Tickets", @@ -368,7 +369,9 @@ "updateTicket": "Update Ticket", "inputTextHolder": "Type here.", "selectHolder": "Select...", - "cancel": "Cancel" + "cancel": "Cancel", + "noOrganizationSelected": "No Organization Selected", + "selectOrganizationForQuota": "Select an Organization to set quotas" }, "organization": { "header": "Organizations", @@ -1174,6 +1177,22 @@ "mechanismType": "Mechanism Type", "actions": "Actions" }, + "quota":{ + "header": "Quotas", + "api-limits": "API Limits", + "role-limits": "Role Limits", + "api-search": "Search Module", + "role-search": "Search Role", + "moduleName": "Module Name", + "roleName": "Roles", + "listLimit": "List (Req/s)", + "getLimit": "Get (Req/s)", + "createLimit": "Create (Req/s)", + "updateLimit": "Update (Req/s)", + "deleteLimit": "Delete (Req/s)", + "table": "Table", + "json": "JSON" + }, "nodes": { "header": "All Nodes", "search": "Search Node", diff --git a/web/src/services/rate_limit.ts b/web/src/services/rate_limit.ts new file mode 100644 index 0000000000..d4c165fe31 --- /dev/null +++ b/web/src/services/rate_limit.ts @@ -0,0 +1,61 @@ +import http from "./http"; + +const rate_limit = { + getApiLimits: (org_identifier: string) => { + let url = `/api/_meta/ratelimit/module_list`; + if(org_identifier){ + url += `?org_id=${org_identifier}`; + } + return http().get(url); + }, + getRoleLimits: (org_identifier: string, role_name: string) => { + let url = `/api/_meta/ratelimit/role_list?user_role=${role_name}&org_id=${org_identifier}`; + return http().get(url); + }, + update_batch: (org_identifier: string, data: any, update_type?: string, user_role?: string) => { + let url = `/api/_meta/ratelimit/update`; + if(update_type){ + url += `?update_type=${update_type}`; + } + if(user_role){ + url += `&user_role=${user_role}`; + } + if(org_identifier){ + url += `&org_id=${org_identifier}`; + } + return http().put(url, data); + }, + download_template: (org_identifier: string) => { + let url = `/api/_meta/ratelimit/download_template`; + if(org_identifier){ + url += `?org_id=${org_identifier}`; + } + return http().get(url); + }, + upload_template: (org_identifier: string, file: any) => { + let url = `/api/_meta/ratelimit/upload`; + if(org_identifier){ + url += `?org_id=${org_identifier}`; + } + const formData = new FormData(); + formData.append('rules', JSON.stringify(file)); + return http().put(url, formData, { + headers: { + 'Content-Type': 'multipart/form-data', + } + }); + }, + getModules: (org_identifier: string) => { + let url = `/api/_meta/ratelimit/api_modules`; + if(org_identifier){ + url += `?org_id=${org_identifier}`; + } + return http().get(url); + } + +} + +export default rate_limit; + + + diff --git a/web/src/stores/index.ts b/web/src/stores/index.ts index 4069ff8902..7ea351d141 100644 --- a/web/src/stores/index.ts +++ b/web/src/stores/index.ts @@ -82,6 +82,9 @@ export default createStore({ hiddenMenus: [], sessionId: "", webSocketUrl: "", + allApiLimitsByOrgId: {}, + allRoleLimitsByOrgIdByRole: {}, + modulesToDisplay: {}, }, mutations: { login(state, payload) { @@ -211,6 +214,15 @@ export default createStore({ setHiddenMenus(state, payload) { state.hiddenMenus = payload; }, + setApiLimitsByOrgId(state, payload) { + state.allApiLimitsByOrgId = payload; + }, + setRoleLimitsByOrgIdByRole(state, payload) { + state.allRoleLimitsByOrgIdByRole = payload; + }, + setModulesToDisplay(state, payload) { + state.modulesToDisplay = payload; + }, }, actions: { login(context, payload) { @@ -339,6 +351,15 @@ export default createStore({ setHiddenMenus(context, payload) { context.commit("setHiddenMenus", payload); }, + setApiLimitsByOrgId(context, payload) { + context.commit("setApiLimitsByOrgId", payload); + }, + setRoleLimitsByOrgIdByRole(context, payload) { + context.commit("setRoleLimitsByOrgIdByRole", payload); + }, + setModulesToDisplay(context, payload) { + context.commit("setModulesToDisplay", payload); + }, }, modules: {}, }); diff --git a/web/src/views/IdentityAccessManagement.vue b/web/src/views/IdentityAccessManagement.vue index 0e9670d38e..0032408485 100644 --- a/web/src/views/IdentityAccessManagement.vue +++ b/web/src/views/IdentityAccessManagement.vue @@ -1,7 +1,7 @@