From 49acdef16bf9a13799cedb300ff925d2d553d648 Mon Sep 17 00:00:00 2001 From: Ashish Kolhe Date: Fri, 23 Oct 2026 16:38:17 +0530 Subject: [PATCH] feat(license): licence-gate the search routes that skipped the search check --- Cargo.lock | 1 + src/api/http/Cargo.toml | 3 + .../handler/http/router/middlewares/mod.rs | 2 + .../http/router/middlewares/search_license.rs | 188 ++++++ src/api/http/src/handler/http/router/mod.rs | 209 +++--- .../http/src/handler/http/router/openapi.rs | 24 +- .../handler/http/router/search_gate_tests.rs | 618 ++++++++++++++++++ 7 files changed, 942 insertions(+), 103 deletions(-) create mode 100644 src/api/http/src/handler/http/router/middlewares/search_license.rs create mode 100644 src/api/http/src/handler/http/router/search_gate_tests.rs diff --git a/Cargo.lock b/Cargo.lock index a3aa5a012e..4a842089a5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8139,6 +8139,7 @@ dependencies = [ "reqwest 0.13.4", "sea-orm", "search", + "search_service", "serde", "serde_json", "snap", diff --git a/src/api/http/Cargo.toml b/src/api/http/Cargo.toml index 4e4f761924..afb6fbe281 100644 --- a/src/api/http/Cargo.toml +++ b/src/api/http/Cargo.toml @@ -21,6 +21,7 @@ enterprise = [ "openobserve-api-search/enterprise", "openobserve-api-common/enterprise", "openobserve-core/enterprise", + "search_service/enterprise", "stream/enterprise", ] cloud = [ @@ -54,6 +55,7 @@ openobserve-tls.workspace = true ingestion-common.workspace = true stream.workspace = true search.workspace = true +search_service.workspace = true openobserve-api-ingest.workspace = true openobserve-api-management.workspace = true openobserve-api-pipelines.workspace = true @@ -98,6 +100,7 @@ futures-util.workspace = true base64.workspace = true sea-orm.workspace = true svix-ksuid.workspace = true +o2_enterprise = { workspace = true, default-features = false, features = ["test-utils"] } tracing.workspace = true tracing-subscriber.workspace = true zstd.workspace = true diff --git a/src/api/http/src/handler/http/router/middlewares/mod.rs b/src/api/http/src/handler/http/router/middlewares/mod.rs index de9a39c3a7..bd7645069a 100644 --- a/src/api/http/src/handler/http/router/middlewares/mod.rs +++ b/src/api/http/src/handler/http/router/middlewares/mod.rs @@ -17,8 +17,10 @@ mod ingest_rejections; mod org_blocking; mod password_policy; mod root_only; +mod search_license; pub use ingest_rejections::ingest_rejections_middleware; pub use org_blocking::{blocked_org_param, blocked_orgs_middleware}; pub use password_policy::password_policy_middleware; pub use root_only::root_only_middleware; +pub use search_license::search_license_middleware; diff --git a/src/api/http/src/handler/http/router/middlewares/search_license.rs b/src/api/http/src/handler/http/router/middlewares/search_license.rs new file mode 100644 index 0000000000..a1a9de74fd --- /dev/null +++ b/src/api/http/src/handler/http/router/middlewares/search_license.rs @@ -0,0 +1,188 @@ +// Copyright 2026 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 axum::{ + extract::{FromRequestParts, RawPathParams, Request}, + http::request::Parts, + middleware::Next, + response::Response, +}; +use common::meta::http::HttpResponse as MetaHttpResponse; +use config::{META_ORG_ID, utils::base64}; + +/// Answers the search handlers' licence 429 before a gated route reads any data. +pub async fn search_license_middleware(request: Request, next: Next) -> Response { + let (mut parts, body) = request.into_parts(); + if let Some(refused) = refusal(&mut parts).await { + return refused; + } + next.run(Request::from_parts(parts, body)).await +} + +async fn refusal(parts: &mut Parts) -> Option { + // Unreadable params fall through to the closed-gate answer for an unknown org and stream. + let (org_id, streams) = match RawPathParams::from_request_parts(parts, &()).await { + Ok(params) => ( + param(¶ms, "org_id").unwrap_or_default().to_string(), + path_streams(¶ms), + ), + Err(_) => (String::new(), Vec::new()), + }; + refusal_for(&org_id, &streams, search_service::check_search_allowed) +} + +fn param<'a>(params: &'a RawPathParams, name: &str) -> Option<&'a str> { + params + .iter() + .find_map(|(key, value)| (key == name).then_some(value)) +} + +fn path_streams(params: &RawPathParams) -> Vec { + if let Some(stream) = param(params, "stream_name") { + return vec![stream.to_string()]; + } + param(params, "stream_names") + .and_then(|encoded| base64::decode_url(encoded).ok()) + .map(|names| names.split(',').map(str::to_string).collect()) + .unwrap_or_default() +} + +fn refusal_for( + org_id: &str, + streams: &[String], + gate: impl Fn(&str, Option<&str>) -> Result<(), infra::errors::Error>, +) -> Option { + if org_id == META_ORG_ID { + return None; + } + let verdict = if streams.is_empty() { + gate(org_id, None) + } else { + streams + .iter() + .try_for_each(|stream| gate(org_id, Some(stream))) + }; + verdict.err().map(MetaHttpResponse::too_many_requests) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use axum::{Router, body::Body, http::StatusCode, routing::get}; + use infra::errors::Error; + use tower::ServiceExt; + + use super::*; + + const BLOCKED: &str = "Search is temporarily disabled due to exceeding allotted ingestion limit. Please contact your administrator."; + + fn open(_: &str, _: Option<&str>) -> Result<(), Error> { + Ok(()) + } + + fn closed(_: &str, _: Option<&str>) -> Result<(), Error> { + Err(Error::Message(BLOCKED.to_string())) + } + + fn names(streams: &[&str]) -> Vec { + streams.iter().map(|s| s.to_string()).collect() + } + + async fn body_bytes(response: Response) -> Vec { + axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap() + .to_vec() + } + + async fn streams_at(encoded: &str) -> String { + let app = Router::new().route( + "/{org_id}/{stream_names}/_around_multi", + get(|params: RawPathParams| async move { path_streams(¶ms).join("|") }), + ); + let request = axum::http::Request::get(format!("/acme/{encoded}/_around_multi")) + .body(Body::empty()) + .unwrap(); + String::from_utf8(body_bytes(app.oneshot(request).await.unwrap()).await).unwrap() + } + + #[test] + fn open_gate_refuses_nothing() { + assert!(refusal_for("acme", &[], open).is_none()); + assert!(refusal_for("acme", &names(&["logs"]), open).is_none()); + } + + #[tokio::test] + async fn closed_gate_answers_the_search_handler_429() { + let refused = refusal_for("acme", &[], closed).unwrap(); + assert_eq!(refused.status(), StatusCode::TOO_MANY_REQUESTS); + assert!(refused.headers().get("retry-after").is_none()); + let body = body_bytes(refused).await; + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!( + json, + serde_json::json!({"code": 429, "message": format!("Error# {BLOCKED}")}) + ); + let handler_body = body_bytes(MetaHttpResponse::too_many_requests(Error::Message( + BLOCKED.to_string(), + ))) + .await; + assert_eq!(body, handler_body); + } + + #[test] + fn meta_org_is_never_refused_and_never_asks_the_gate() { + let calls = AtomicUsize::new(0); + let counting = |org: &str, stream: Option<&str>| { + calls.fetch_add(1, Ordering::SeqCst); + closed(org, stream) + }; + assert!(refusal_for(META_ORG_ID, &names(&["logs"]), counting).is_none()); + assert!(refusal_for(META_ORG_ID, &[], counting).is_none()); + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + + #[test] + fn each_path_stream_reaches_the_gate() { + let only_triggers = |_: &str, stream: Option<&str>| match stream { + Some("triggers") => Ok(()), + _ => Err(Error::Message(BLOCKED.to_string())), + }; + assert!(refusal_for("acme", &names(&["triggers"]), only_triggers).is_none()); + assert!(refusal_for("acme", &names(&["logs"]), only_triggers).is_some()); + assert!(refusal_for("acme", &names(&["triggers", "logs"]), only_triggers).is_some()); + } + + #[tokio::test] + async fn around_multi_stream_names_are_decoded() { + assert_eq!( + streams_at(&base64::encode_url("triggers,logs")).await, + "triggers|logs" + ); + } + + #[tokio::test] + async fn undecodable_stream_names_are_gated_as_unknown() { + assert_eq!(streams_at("!!not-base64!!").await, ""); + let asked = std::sync::Mutex::new(Vec::new()); + let recording = |_: &str, stream: Option<&str>| { + asked.lock().unwrap().push(stream.map(str::to_string)); + closed("", None) + }; + assert!(refusal_for("acme", &[], recording).is_some()); + assert_eq!(*asked.lock().unwrap(), vec![None]); + } +} diff --git a/src/api/http/src/handler/http/router/mod.rs b/src/api/http/src/handler/http/router/mod.rs index 1d0f1d3ef6..53680486d1 100644 --- a/src/api/http/src/handler/http/router/mod.rs +++ b/src/api/http/src/handler/http/router/mod.rs @@ -85,7 +85,7 @@ use crate::{ }, router::middlewares::{ blocked_org_param, blocked_orgs_middleware, ingest_rejections_middleware, - password_policy_middleware, root_only_middleware, + password_policy_middleware, root_only_middleware, search_license_middleware, }, }, }; @@ -99,6 +99,8 @@ mod one_ack_tests; pub mod openapi; #[cfg(test)] mod openapi_coverage; +#[cfg(test)] +mod search_gate_tests; use common::meta::http::ERROR_HEADER; @@ -1292,9 +1294,13 @@ pub fn service_routes() -> Router { #[cfg(not(feature = "enterprise"))] let server = cfg.common.instance_name_short.to_string(); - // Downtimes, Slack and service streams live in their own tables to keep this function under the line limit. - let mut router = - public_dashboard_admin_routes(Router::new().merge(downtime_routes()).merge(slack_routes())); + // Downtimes, Slack, search and service streams live in their own tables to keep this function under the line limit. + let mut router = public_dashboard_admin_routes( + Router::new() + .merge(downtime_routes()) + .merge(slack_routes()) + .merge(search_routes()), + ); router = router.merge(service_streams_routes()); // Full UI configuration — authenticated counterpart of the unauthenticated // `/config` bootstrap in config_routes() @@ -1402,7 +1408,6 @@ pub fn service_routes() -> Router { .route("/{org_id}/{stream_name}/traces/session/details", get(traces::session::get_session_details)) .route("/{org_id}/{stream_name}/traces/user", get(traces::user::get_latest_users)) .route("/{org_id}/{stream_name}/traces/time_range", get(traces::time_index::get_trace_time_range)) - .route("/{org_id}/{stream_name}/traces/{trace_id}/details", get(traces::details::get_trace_details)) .route("/{org_id}/{stream_name}/traces/{trace_id}/dag", get(traces::dag::get_trace_dag)) // Profiles query (ingest remains under /v1/profiles) @@ -1412,43 +1417,6 @@ pub fn service_routes() -> Router { .route("/{org_id}/{stream_name}/profiles/merge", get(profiles_query::merge_profiles_get).post(profiles_query::merge_profiles)) .route("/{org_id}/{stream_name}/profiles/diff", post(profiles_query::diff_profiles)) - // Database Monitoring — its own top-level module, not under /traces. - // The path segment must stay `db_monitoring` to match the OFGA resource - // these routes authorize against (`db_monitoring:{org}`, see o2_openfga - // ROUTE_PERMISSIONS). - // - // Every route below is registered in both builds; the build-type gate - // lives in the handlers, three of which (deadlocks/blocking/ - // table_health) are dual-implemented and answer 403 on OSS. Registering - // them in the enterprise-gated block instead would 404 them, losing the - // distinction between "not licensed" and "no such endpoint". Runtime - // off-switch is ZO_DB_MONITORING_ENABLED. - .route("/{org_id}/db_monitoring/databases", get(db_monitoring::handler::get_dbm_databases)) - .route("/{org_id}/db_monitoring/queries", get(db_monitoring::handler::get_dbm_queries)) - .route("/{org_id}/db_monitoring/query/history", get(db_monitoring::handler::get_dbm_query_history)) - .route("/{org_id}/db_monitoring/query/endpoints", get(db_monitoring::handler::get_dbm_query_endpoints)) - .route("/{org_id}/db_monitoring/samples", get(db_monitoring::handler::get_dbm_samples)) - // Server-vantage events (read the canonical o2_dbm_* columns) - .route("/{org_id}/db_monitoring/deadlocks", get(db_monitoring::handler::get_dbm_deadlocks)) - .route("/{org_id}/db_monitoring/blocking", get(db_monitoring::handler::get_dbm_blocking)) - .route("/{org_id}/db_monitoring/activity", get(db_monitoring::handler::get_dbm_activity)) - // `query/insights` returns the query-detail page's Logs-side pair in one - // round trip; `query/plans` and `query/server_metrics` are superseded by - // it (same sections, same envelopes) and stay registered for - // compatibility. - .route("/{org_id}/db_monitoring/query/insights", get(db_monitoring::handler::get_dbm_query_insights)) - .route("/{org_id}/db_monitoring/query/plans", get(db_monitoring::handler::get_dbm_query_plans)) - .route("/{org_id}/db_monitoring/query/server_metrics", get(db_monitoring::handler::get_dbm_query_server_metrics)) - .route("/{org_id}/db_monitoring/server_queries", get(db_monitoring::handler::get_dbm_server_queries)) - .route("/{org_id}/db_monitoring/server_samples", get(db_monitoring::handler::get_dbm_server_samples)) - .route("/{org_id}/db_monitoring/table_health", get(db_monitoring::handler::get_dbm_table_health)) - .route("/{org_id}/db_monitoring/instances", get(db_monitoring::handler::get_dbm_instances)) - .route( - "/{org_id}/db_monitoring/instance_metrics", - get(db_monitoring::handler::get_dbm_instance_metrics), - ) - .route("/{org_id}/db_monitoring/badges", get(db_monitoring::handler::get_dbm_badges)) - // LLM Model Pricing .route("/{org_id}/llm/models", get(model_pricing::list).post(model_pricing::create)) // NOTE: named routes MUST be registered before {model_id} or they match as a model ID @@ -1471,7 +1439,6 @@ pub fn service_routes() -> Router { .route("/{org_id}/prometheus/api/v1/metadata", get(promql::metadata)) .route("/{org_id}/prometheus/api/v1/series", get(promql::series_get).post(promql::series_post)) .route("/{org_id}/prometheus/api/v1/labels", get(promql::labels_get).post(promql::labels_post)) - .route("/{org_id}/prometheus/api/v1/label/{label_name}/values", get(promql::label_values)) .route("/{org_id}/prometheus/api/v1/format_query", get(promql::format_query_get).post(promql::format_query_post)) .route("/{org_id}/prometheus/api/v1/parse_tree", post(promql::parse_tree)) @@ -1484,12 +1451,6 @@ pub fn service_routes() -> Router { .route("/{org_id}/_search_history", post(search::search_history)) .route("/{org_id}/result_schema", post(search::result_schema)) - // Multi-stream search - .route("/{org_id}/_search_multi", post(search::multi_streams::search_multi)) - .route("/{org_id}/_search_multi_stream", post(search::multi_streams::search_multi_stream)) - .route("/{org_id}/_search_partition_multi", post(search::multi_streams::_search_partition_multi)) - .route("/{org_id}/{stream_names}/_around_multi", get(search::multi_streams::around_multi)) - // HTTP/2 streaming .route("/{org_id}/_search_stream", post(search::search_stream::search_http2_stream)) .route("/{org_id}/_values_stream", post(search::search_stream::values_http2_stream)) @@ -1643,8 +1604,7 @@ pub fn service_routes() -> Router { // Anomaly Detection .route("/{org_id}/anomaly_detection", get(anomaly_detection::list_configs).post(anomaly_detection::create_config)) .route("/{org_id}/anomaly_detection/{config_id}", get(anomaly_detection::get_config).put(anomaly_detection::update_config).delete(anomaly_detection::delete_config)) - .route("/{org_id}/anomaly_detection/{config_id}/train", post(anomaly_detection::train_model).delete(anomaly_detection::cancel_training)) - .route("/{org_id}/anomaly_detection/{config_id}/detect", post(anomaly_detection::detect_anomalies)) + .route("/{org_id}/anomaly_detection/{config_id}/train", delete(anomaly_detection::cancel_training)) .route("/{org_id}/anomaly_detection/{config_id}/history", get(anomaly_detection::get_detection_history)) .route("/{org_id}/anomaly_detection/history", get(alerts::history::get_all_anomaly_history)); router = alert_calibration_routes(router); @@ -1755,21 +1715,11 @@ pub fn service_routes() -> Router { ) .route( "/{org_id}/annotation_queues/{queue_id}/items", - post(annotation_queues::enqueue_annotation_queue_item) - .delete(annotation_queues::clear_annotation_queue_items), - ) - .route( - "/{org_id}/annotation_queues/{queue_id}/items/{queue_item_id}", - get(annotation_queues::get_annotation_queue_item), + delete(annotation_queues::clear_annotation_queue_items), ) .route( "/{org_id}/annotation_queues/{queue_id}/items/{queue_item_id}/reviews", - get(annotation_queues::list_annotation_queue_item_reviews) - .post(annotation_queues::review_annotation_queue_item), - ) - .route( - "/{org_id}/annotation_queues/{queue_id}/items/{queue_item_id}/push_to_dataset", - post(datasets::push_annotation_queue_item_to_dataset), + post(annotation_queues::review_annotation_queue_item), ) .route( "/{org_id}/annotation_queues/{queue_id}/items/archive", @@ -1791,9 +1741,7 @@ pub fn service_routes() -> Router { ) .route( "/{org_id}/datasets/{dataset_id}/items", - get(datasets::list_dataset_items) - .post(datasets::push_dataset_item) - .put(datasets::upsert_dataset_items), + get(datasets::list_dataset_items).put(datasets::upsert_dataset_items), ) .route( "/{org_id}/datasets/{dataset_id}/rows", @@ -1819,27 +1767,15 @@ pub fn service_routes() -> Router { "/{org_id}/experiments", get(experiments::list_experiments).post(experiments::create_experiment), ) - .route( - "/{org_id}/experiments/compare", - get(experiments::compare_experiments), - ) .route( "/{org_id}/experiments/{experiment_id}", - get(experiments::get_experiment).delete(experiments::delete_experiment), + delete(experiments::delete_experiment), ) .route( "/{org_id}/experiments/{experiment_id}/baseline", put(experiments::set_experiment_baseline) .delete(experiments::clear_experiment_baseline), ) - .route( - "/{org_id}/experiments/{experiment_id}/rows", - get(experiments::list_experiment_result_rows), - ) - .route( - "/{org_id}/experiments/{experiment_id}/rows/{row_id}", - get(experiments::get_experiment_row), - ) .route( "/{org_id}/experiments/{experiment_id}/rows/{row_id}/trials/{trial_index}/retry", post(experiments::retry_experiment_slot), @@ -1872,9 +1808,6 @@ pub fn service_routes() -> Router { // On-demand human annotation from Discovery .route("/{org_id}/annotations", post(annotations::annotate_target)) - // LLM score-based Discovery - .route("/{org_id}/discovery", get(discovery::list_discovery_items)) - // LLM Providers (Online Eval Phase 2) .route("/{org_id}/providers", get(providers::list_providers).post(providers::create_provider)) .route("/{org_id}/providers/test", post(providers::test_provider_config)) @@ -1892,7 +1825,6 @@ pub fn service_routes() -> Router { .route("/{org_id}/tasks/{entity_id}/signing/end_grace", post(remote_tasks::end_remote_task_signing_grace)) .route("/{org_id}/tasks/{entity_id}/signing", get(remote_tasks::get_remote_task_signing_status).delete(remote_tasks::revoke_remote_task_signing_secret)) .route("/{org_id}/tasks/{entity_id}/versions", get(remote_tasks::list_remote_task_versions)) - .route("/{org_id}/tasks/{entity_id}/stats", get(remote_tasks::get_remote_task_stats)) .route("/{org_id}/tasks/{entity_id}/draft", get(remote_tasks::get_remote_task_draft).delete(remote_tasks::discard_remote_task_draft)) .route("/{org_id}/tasks/{entity_id}/test_connection", post(remote_tasks::publish_remote_task)) .route("/{org_id}/tasks/{entity_id}/test_run", post(remote_tasks::test_run_remote_task)) @@ -1934,11 +1866,9 @@ pub fn service_routes() -> Router { { router = fleet_routes(router) // Search jobs - .route("/{org_id}/search_jobs", get(search::search_job::list_status).post(search::search_job::submit_job)) + .route("/{org_id}/search_jobs", get(search::search_job::list_status)) .route("/{org_id}/search_jobs/{job_id}", get(search::search_job::get_status).delete(search::search_job::delete_job)) - .route("/{org_id}/search_jobs/{job_id}/result", get(search::search_job::get_job_result)) .route("/{org_id}/search_jobs/{job_id}/cancel", post(search::search_job::cancel_job)) - .route("/{org_id}/search_jobs/{job_id}/retry", post(search::search_job::retry_job)) // Query manager .route("/{org_id}/query_manager/status", get(search::query_manager::query_status)) @@ -1976,16 +1906,6 @@ pub fn service_routes() -> Router { .route("/license", get(license::get_license_info).post(license::store_license)) .route("/license/refresh", post(license::refresh_license_limits)) - // Topology - .route("/{org_id}/traces/service_graph/topology/current", get(traces::get_current_topology)) - .route("/{org_id}/traces/service_graph/edge/history", get(traces::get_edge_history)) - // Agent behavior signals (loop / failure / cost) — reads the _agent_signals stream - .route("/{org_id}/traces/agent_signals", get(traces::get_agent_signals)) - .route("/{org_id}/traces/agent_signals/compare", post(traces::compare_agent_versions)) - - // Patterns - .route("/{org_id}/streams/{stream_name}/patterns/extract", post(patterns::extract_patterns)) - // Service streams .route("/{org_id}/service_streams", get(service_streams::list_services)) .route("/{org_id}/service_streams/config/identity", get(service_streams::get_identity_config).put(service_streams::save_identity_config)) @@ -2378,8 +2298,6 @@ pub fn downtime_routes() -> Router { Router::new() .route("/v2/{org_id}/downtimes", get(downtimes::list_downtimes).post(downtimes::create_downtime)) .route("/v2/{org_id}/downtimes/preview", post(downtimes::preview_downtime)) - .route("/v2/{org_id}/downtimes/resources", post(downtimes::downtime_resources)) - .route("/v2/{org_id}/downtimes/values", post(downtimes::downtime_values)) .route("/v2/{org_id}/downtimes/move", patch(downtimes::move_downtimes)) .route("/v2/{org_id}/downtimes/{downtime_id}", get(downtimes::get_downtime).put(downtimes::update_downtime).delete(downtimes::delete_downtime)) .route("/v2/{org_id}/downtimes/{downtime_id}/cancel", post(downtimes::cancel_downtime)) @@ -2405,6 +2323,101 @@ fn report_routes() -> Router { .route("/v2/{org_id}/reports/{report_id}/trigger", put(dashboards::reports::trigger_report_v2)) } +/// Routes that read customer data through search; the licence gate answers 429 while search is blocked. +#[rustfmt::skip] +fn search_routes() -> Router { + #[cfg(feature = "enterprise")] + let llm_eval = get_o2_config().llm_eval_config.enabled; + #[cfg(not(feature = "enterprise"))] + let llm_eval = false; + search_routes_with(llm_eval) +} + +#[rustfmt::skip] +#[cfg_attr(not(feature = "enterprise"), allow(unused_variables))] +fn search_routes_with(llm_eval: bool) -> Router { + let router = Router::new() + // Multi-stream search + .route("/{org_id}/_search_multi", post(search::multi_streams::search_multi)) + .route("/{org_id}/_search_multi_stream", post(search::multi_streams::search_multi_stream)) + .route("/{org_id}/_search_partition_multi", post(search::multi_streams::_search_partition_multi)) + .route("/{org_id}/{stream_names}/_around_multi", get(search::multi_streams::around_multi)) + .route("/{org_id}/{stream_name}/traces/{trace_id}/details", get(traces::details::get_trace_details)) + .route("/{org_id}/prometheus/api/v1/label/{label_name}/values", get(promql::label_values)) + .route("/v2/{org_id}/downtimes/resources", post(downtimes::downtime_resources)) + .route("/v2/{org_id}/downtimes/values", post(downtimes::downtime_values)) + + // Database Monitoring — its own top-level module, not under /traces. + // The path segment must stay `db_monitoring` to match the OFGA resource + // these routes authorize against (`db_monitoring:{org}`, see o2_openfga + // ROUTE_PERMISSIONS). + // + // Every route below is registered in both builds; the build-type gate + // lives in the handlers, three of which (deadlocks/blocking/ + // table_health) are dual-implemented and answer 403 on OSS. Registering + // them in the enterprise-gated block instead would 404 them, losing the + // distinction between "not licensed" and "no such endpoint". Runtime + // off-switch is ZO_DB_MONITORING_ENABLED. + .route("/{org_id}/db_monitoring/databases", get(db_monitoring::handler::get_dbm_databases)) + .route("/{org_id}/db_monitoring/queries", get(db_monitoring::handler::get_dbm_queries)) + .route("/{org_id}/db_monitoring/query/history", get(db_monitoring::handler::get_dbm_query_history)) + .route("/{org_id}/db_monitoring/query/endpoints", get(db_monitoring::handler::get_dbm_query_endpoints)) + .route("/{org_id}/db_monitoring/samples", get(db_monitoring::handler::get_dbm_samples)) + // Server-vantage events (read the canonical o2_dbm_* columns) + .route("/{org_id}/db_monitoring/deadlocks", get(db_monitoring::handler::get_dbm_deadlocks)) + .route("/{org_id}/db_monitoring/blocking", get(db_monitoring::handler::get_dbm_blocking)) + .route("/{org_id}/db_monitoring/activity", get(db_monitoring::handler::get_dbm_activity)) + // `query/insights` returns the query-detail page's Logs-side pair in one + // round trip; `query/plans` and `query/server_metrics` are superseded by + // it (same sections, same envelopes) and stay registered for + // compatibility. + .route("/{org_id}/db_monitoring/query/insights", get(db_monitoring::handler::get_dbm_query_insights)) + .route("/{org_id}/db_monitoring/query/plans", get(db_monitoring::handler::get_dbm_query_plans)) + .route("/{org_id}/db_monitoring/query/server_metrics", get(db_monitoring::handler::get_dbm_query_server_metrics)) + .route("/{org_id}/db_monitoring/server_queries", get(db_monitoring::handler::get_dbm_server_queries)) + .route("/{org_id}/db_monitoring/server_samples", get(db_monitoring::handler::get_dbm_server_samples)) + .route("/{org_id}/db_monitoring/table_health", get(db_monitoring::handler::get_dbm_table_health)) + .route("/{org_id}/db_monitoring/instances", get(db_monitoring::handler::get_dbm_instances)) + .route( + "/{org_id}/db_monitoring/instance_metrics", + get(db_monitoring::handler::get_dbm_instance_metrics), + ) + .route("/{org_id}/db_monitoring/badges", get(db_monitoring::handler::get_dbm_badges)); + #[cfg(feature = "enterprise")] + let router = { + let router = router + .route("/{org_id}/anomaly_detection/{config_id}/train", post(anomaly_detection::train_model)) + .route("/{org_id}/anomaly_detection/{config_id}/detect", post(anomaly_detection::detect_anomalies)) + .route("/{org_id}/search_jobs", post(search::search_job::submit_job)) + .route("/{org_id}/search_jobs/{job_id}/result", get(search::search_job::get_job_result)) + .route("/{org_id}/search_jobs/{job_id}/retry", post(search::search_job::retry_job)) + .route("/{org_id}/traces/service_graph/topology/current", get(traces::get_current_topology)) + .route("/{org_id}/traces/service_graph/edge/history", get(traces::get_edge_history)) + // Agent behavior signals (loop / failure / cost) — reads the _agent_signals stream + .route("/{org_id}/traces/agent_signals", get(traces::get_agent_signals)) + .route("/{org_id}/traces/agent_signals/compare", post(traces::compare_agent_versions)) + .route("/{org_id}/streams/{stream_name}/patterns/extract", post(patterns::extract_patterns)); + if llm_eval { + router + .route("/{org_id}/annotation_queues/{queue_id}/items", post(annotation_queues::enqueue_annotation_queue_item)) + .route("/{org_id}/annotation_queues/{queue_id}/items/{queue_item_id}", get(annotation_queues::get_annotation_queue_item)) + .route("/{org_id}/annotation_queues/{queue_id}/items/{queue_item_id}/reviews", get(annotation_queues::list_annotation_queue_item_reviews)) + .route("/{org_id}/annotation_queues/{queue_id}/items/{queue_item_id}/push_to_dataset", post(datasets::push_annotation_queue_item_to_dataset)) + .route("/{org_id}/datasets/{dataset_id}/items", post(datasets::push_dataset_item)) + .route("/{org_id}/experiments/compare", get(experiments::compare_experiments)) + .route("/{org_id}/experiments/{experiment_id}", get(experiments::get_experiment)) + .route("/{org_id}/experiments/{experiment_id}/rows", get(experiments::list_experiment_result_rows)) + .route("/{org_id}/experiments/{experiment_id}/rows/{row_id}", get(experiments::get_experiment_row)) + .route("/{org_id}/discovery", get(discovery::list_discovery_items)) + .route("/{org_id}/tasks/{entity_id}/stats", get(remote_tasks::get_remote_task_stats)) + } else { + router + } + }; + // A route layer, not a layer: it must not reach the ungated methods merged onto the same paths. + router.route_layer(middleware::from_fn(search_license_middleware)) +} + /// On-call management routes; the caller mounts them only when on-call is enabled. #[cfg(feature = "enterprise")] fn oncall_routes(router: Router) -> Router { diff --git a/src/api/http/src/handler/http/router/openapi.rs b/src/api/http/src/handler/http/router/openapi.rs index 7d543fd8b2..018cae1925 100644 --- a/src/api/http/src/handler/http/router/openapi.rs +++ b/src/api/http/src/handler/http/router/openapi.rs @@ -1683,6 +1683,9 @@ mod tests { #[cfg(feature = "enterprise")] const SERVICE_ROUTES_SIGNATURE: &str = "pub fn service_routes("; + #[cfg(feature = "enterprise")] + const SEARCH_ROUTES_SIGNATURE: &str = "fn search_routes_with("; + #[cfg(feature = "enterprise")] const HTTP_METHODS: &[&str] = &[ "get", "post", "put", "delete", "patch", "head", "options", "trace", @@ -1909,13 +1912,18 @@ mod tests { // Ends at the column-0 close brace, so an item after the body is never scraped into it. #[cfg(feature = "enterprise")] fn service_routes_body() -> &'static str { + router_fn_body(SERVICE_ROUTES_SIGNATURE) + } + + #[cfg(feature = "enterprise")] + fn router_fn_body(signature: &str) -> &'static str { let start = ROUTER_SOURCE - .find(SERVICE_ROUTES_SIGNATURE) - .expect("the router must define service_routes()"); - let after_signature = &ROUTER_SOURCE[start + SERVICE_ROUTES_SIGNATURE.len()..]; + .find(signature) + .unwrap_or_else(|| panic!("the router must define {signature}")); + let after_signature = &ROUTER_SOURCE[start + signature.len()..]; let end = after_signature .find("\n}\n") - .expect("service_routes() must end at a column-0 close brace"); + .unwrap_or_else(|| panic!("{signature} must end at a column-0 close brace")); &after_signature[..end] } @@ -2014,7 +2022,13 @@ mod tests { #[cfg(feature = "enterprise")] fn registered_gated_routes() -> BTreeSet<(String, String)> { - let live = normalize(&strip_comments(service_routes_body())); + // The licence-gated search routes are a separate group that service_routes() merges in. + let routes = [ + service_routes_body(), + router_fn_body(SEARCH_ROUTES_SIGNATURE), + ] + .concat(); + let live = normalize(&strip_comments(&routes)); let mut routes = BTreeSet::new(); for call in live.split(".route(").skip(1) { let Some(arguments) = balanced_arguments(call) else { diff --git a/src/api/http/src/handler/http/router/search_gate_tests.rs b/src/api/http/src/handler/http/router/search_gate_tests.rs new file mode 100644 index 0000000000..514b1fb934 --- /dev/null +++ b/src/api/http/src/handler/http/router/search_gate_tests.rs @@ -0,0 +1,618 @@ +// Copyright 2026 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::BTreeSet; + +use axum::body::Body; +use tower::ServiceExt; + +use super::*; + +/// Every route `search_routes_with` registers, as (method, template, handler). +#[rustfmt::skip] +const GATED: [(&str, &str, &str); 46] = [ + ("POST", "/{org_id}/_search_multi", "search::multi_streams::search_multi"), + ("POST", "/{org_id}/_search_multi_stream", "search::multi_streams::search_multi_stream"), + ("POST", "/{org_id}/_search_partition_multi", "search::multi_streams::_search_partition_multi"), + ("GET", "/{org_id}/{stream_names}/_around_multi", "search::multi_streams::around_multi"), + ("GET", "/{org_id}/{stream_name}/traces/{trace_id}/details", "traces::details::get_trace_details"), + ("GET", "/{org_id}/prometheus/api/v1/label/{label_name}/values", "promql::label_values"), + ("POST", "/v2/{org_id}/downtimes/resources", "downtimes::downtime_resources"), + ("POST", "/v2/{org_id}/downtimes/values", "downtimes::downtime_values"), + ("GET", "/{org_id}/db_monitoring/databases", "db_monitoring::handler::get_dbm_databases"), + ("GET", "/{org_id}/db_monitoring/queries", "db_monitoring::handler::get_dbm_queries"), + ("GET", "/{org_id}/db_monitoring/query/history", "db_monitoring::handler::get_dbm_query_history"), + ("GET", "/{org_id}/db_monitoring/query/endpoints", "db_monitoring::handler::get_dbm_query_endpoints"), + ("GET", "/{org_id}/db_monitoring/samples", "db_monitoring::handler::get_dbm_samples"), + ("GET", "/{org_id}/db_monitoring/deadlocks", "db_monitoring::handler::get_dbm_deadlocks"), + ("GET", "/{org_id}/db_monitoring/blocking", "db_monitoring::handler::get_dbm_blocking"), + ("GET", "/{org_id}/db_monitoring/activity", "db_monitoring::handler::get_dbm_activity"), + ("GET", "/{org_id}/db_monitoring/query/insights", "db_monitoring::handler::get_dbm_query_insights"), + ("GET", "/{org_id}/db_monitoring/query/plans", "db_monitoring::handler::get_dbm_query_plans"), + ("GET", "/{org_id}/db_monitoring/query/server_metrics", "db_monitoring::handler::get_dbm_query_server_metrics"), + ("GET", "/{org_id}/db_monitoring/server_queries", "db_monitoring::handler::get_dbm_server_queries"), + ("GET", "/{org_id}/db_monitoring/server_samples", "db_monitoring::handler::get_dbm_server_samples"), + ("GET", "/{org_id}/db_monitoring/table_health", "db_monitoring::handler::get_dbm_table_health"), + ("GET", "/{org_id}/db_monitoring/instances", "db_monitoring::handler::get_dbm_instances"), + ("GET", "/{org_id}/db_monitoring/instance_metrics", "db_monitoring::handler::get_dbm_instance_metrics"), + ("GET", "/{org_id}/db_monitoring/badges", "db_monitoring::handler::get_dbm_badges"), + ("POST", "/{org_id}/anomaly_detection/{config_id}/train", "anomaly_detection::train_model"), + ("POST", "/{org_id}/anomaly_detection/{config_id}/detect", "anomaly_detection::detect_anomalies"), + ("POST", "/{org_id}/search_jobs", "search::search_job::submit_job"), + ("GET", "/{org_id}/search_jobs/{job_id}/result", "search::search_job::get_job_result"), + ("POST", "/{org_id}/search_jobs/{job_id}/retry", "search::search_job::retry_job"), + ("GET", "/{org_id}/traces/service_graph/topology/current", "traces::get_current_topology"), + ("GET", "/{org_id}/traces/service_graph/edge/history", "traces::get_edge_history"), + ("GET", "/{org_id}/traces/agent_signals", "traces::get_agent_signals"), + ("POST", "/{org_id}/traces/agent_signals/compare", "traces::compare_agent_versions"), + ("POST", "/{org_id}/streams/{stream_name}/patterns/extract", "patterns::extract_patterns"), + ("POST", "/{org_id}/annotation_queues/{queue_id}/items", "annotation_queues::enqueue_annotation_queue_item"), + ("GET", "/{org_id}/annotation_queues/{queue_id}/items/{queue_item_id}", "annotation_queues::get_annotation_queue_item"), + ("GET", "/{org_id}/annotation_queues/{queue_id}/items/{queue_item_id}/reviews", "annotation_queues::list_annotation_queue_item_reviews"), + ("POST", "/{org_id}/annotation_queues/{queue_id}/items/{queue_item_id}/push_to_dataset", "datasets::push_annotation_queue_item_to_dataset"), + ("POST", "/{org_id}/datasets/{dataset_id}/items", "datasets::push_dataset_item"), + ("GET", "/{org_id}/experiments/compare", "experiments::compare_experiments"), + ("GET", "/{org_id}/experiments/{experiment_id}", "experiments::get_experiment"), + ("GET", "/{org_id}/experiments/{experiment_id}/rows", "experiments::list_experiment_result_rows"), + ("GET", "/{org_id}/experiments/{experiment_id}/rows/{row_id}", "experiments::get_experiment_row"), + ("GET", "/{org_id}/discovery", "discovery::list_discovery_items"), + ("GET", "/{org_id}/tasks/{entity_id}/stats", "remote_tasks::get_remote_task_stats"), +]; + +/// Modules whose every registered handler reads data through search. +const FULLY_GATED_MODULES: [&str; 5] = [ + "search::multi_streams::", + "db_monitoring::handler::", + "patterns::", + "discovery::", + "traces::details::", +]; + +/// Modules that hold gated and ungated handlers side by side. +const MIXED_MODULES: [&str; 9] = [ + "search::search_job", + "anomaly_detection", + "experiments", + "annotation_queues", + "datasets", + "remote_tasks", + "downtimes", + "promql", + "traces", +]; + +/// Each ungated handler of a mixed module, with why it stays open. +#[rustfmt::skip] +const UNGATED_SIBLINGS: &[(&str, &str)] = &[ + ("annotation_queues::archive_annotation_queue_items", "queue bookkeeping, no search"), + ("annotation_queues::clear_annotation_queue_items", "queue bookkeeping, no search"), + ("annotation_queues::create_annotation_queue", "queue bookkeeping, no search"), + ("annotation_queues::delete_annotation_queue", "queue bookkeeping, no search"), + ("annotation_queues::get_annotation_queue", "queue bookkeeping, no search"), + ("annotation_queues::list_annotation_queue_items", "stored queue items, no search"), + ("annotation_queues::list_annotation_queues", "queue bookkeeping, no search"), + ("annotation_queues::review_annotation_queue_item", "stores a human review, no search"), + ("annotation_queues::update_annotation_queue", "queue bookkeeping, no search"), + ("anomaly_detection::cancel_training", "stops work, no search"), + ("anomaly_detection::create_config", "create-time training waits in the scheduler while blocked"), + ("anomaly_detection::delete_config", "config bookkeeping, no search"), + ("anomaly_detection::get_config", "config bookkeeping, no search"), + ("anomaly_detection::get_detection_history", "operational history of the org's own detections"), + ("anomaly_detection::list_configs", "config bookkeeping, no search"), + ("anomaly_detection::update_config", "config bookkeeping, no search"), + ("datasets::create_dataset", "dataset bookkeeping, no search"), + ("datasets::delete_dataset", "dataset bookkeeping, no search"), + ("datasets::delete_dataset_item", "dataset bookkeeping, no search"), + ("datasets::get_dataset", "dataset bookkeeping, no search"), + ("datasets::get_dataset_item_versions", "stored dataset rows, no search"), + ("datasets::get_dataset_snapshot_rows", "stored dataset rows, no search"), + ("datasets::import_dataset_items", "rows come from the uploaded CSV, not telemetry"), + ("datasets::list_dataset_items", "stored dataset rows, no search"), + ("datasets::list_datasets", "dataset bookkeeping, no search"), + ("datasets::update_dataset", "dataset bookkeeping, no search"), + ("datasets::update_dataset_item", "rows come from the caller, not telemetry"), + ("datasets::upsert_dataset_items", "rows come from the caller, not telemetry"), + ("downtimes::cancel_downtime", "downtime bookkeeping, no search"), + ("downtimes::create_downtime", "downtime bookkeeping, no search"), + ("downtimes::delete_downtime", "downtime bookkeeping, no search"), + ("downtimes::extend_downtime", "downtime bookkeeping, no search"), + ("downtimes::get_downtime", "downtime bookkeeping, no search"), + ("downtimes::list_downtimes", "downtime bookkeeping, no search"), + ("downtimes::move_downtimes", "downtime bookkeeping, no search"), + ("downtimes::preview_downtime", "inventory match only"), + ("downtimes::update_downtime", "downtime bookkeeping, no search"), + ("experiments::cancel_experiment", "stops work, no search"), + ("experiments::clear_experiment_baseline", "experiment bookkeeping, no search"), + ("experiments::clone_experiment", "experiment bookkeeping; runs score as background eval work"), + ("experiments::create_experiment", "runs score as background eval work"), + ("experiments::delete_experiment", "experiment bookkeeping, no search"), + ("experiments::finalize_experiment", "experiment bookkeeping, no search"), + ("experiments::list_experiment_slots", "slot bookkeeping, no search"), + ("experiments::list_experiments", "summary skipped while blocked (S4)"), + ("experiments::preview_experiment", "previews the dataset, not telemetry"), + ("experiments::retry_experiment", "runs score as background eval work"), + ("experiments::retry_experiment_slot", "runs score as background eval work"), + ("experiments::set_experiment_baseline", "experiment bookkeeping, no search"), + ("experiments::submit_experiment_records", "records come from the caller"), + ("promql::format_query_get", "formats the query text, no data"), + ("promql::format_query_post", "formats the query text, no data"), + ("promql::labels_get", "schema only"), + ("promql::labels_post", "schema only"), + ("promql::metadata", "schema only"), + ("promql::parse_tree", "parses the query text, no data"), + ("promql::query_exemplars_get", "checks the licence in the handler (query_range)"), + ("promql::query_exemplars_post", "checks the licence in the handler (query_range)"), + ("promql::query_get", "checks the licence in the handler"), + ("promql::query_post", "checks the licence in the handler"), + ("promql::query_range_get", "checks the licence in the handler"), + ("promql::query_range_post", "checks the licence in the handler"), + ("promql::remote_write", "ingestion"), + ("promql::series_get", "checks the licence in the handler"), + ("promql::series_post", "checks the licence in the handler"), + ("remote_tasks::activate_remote_task_signing_candidate", "task bookkeeping, no search"), + ("remote_tasks::create_remote_task", "task bookkeeping, no search"), + ("remote_tasks::delete_remote_task", "task bookkeeping, no search"), + ("remote_tasks::discard_remote_task_draft", "task bookkeeping, no search"), + ("remote_tasks::end_remote_task_signing_grace", "task bookkeeping, no search"), + ("remote_tasks::get_remote_task", "task bookkeeping, no search"), + ("remote_tasks::get_remote_task_draft", "task bookkeeping, no search"), + ("remote_tasks::get_remote_task_signing_status", "task bookkeeping, no search"), + ("remote_tasks::list_remote_task_versions", "task bookkeeping, no search"), + ("remote_tasks::list_remote_tasks", "task bookkeeping, no search"), + ("remote_tasks::publish_remote_task", "calls the remote endpoint, no search"), + ("remote_tasks::replace_remote_task_auth_secret", "task bookkeeping, no search"), + ("remote_tasks::replace_remote_task_header_secret", "task bookkeeping, no search"), + ("remote_tasks::revoke_remote_task_auth_secret", "task bookkeeping, no search"), + ("remote_tasks::revoke_remote_task_header_secret", "task bookkeeping, no search"), + ("remote_tasks::revoke_remote_task_signing_secret", "task bookkeeping, no search"), + ("remote_tasks::rotate_remote_task_signing_secret", "task bookkeeping, no search"), + ("remote_tasks::save_remote_task_draft", "task bookkeeping, no search"), + ("remote_tasks::test_remote_task", "calls the remote endpoint, no search"), + ("remote_tasks::test_remote_task_signing_candidate", "calls the remote endpoint, no search"), + ("remote_tasks::test_run_remote_task", "calls the remote endpoint on caller input"), + ("search::search_job::cancel_job", "job bookkeeping, no search"), + ("search::search_job::delete_job", "job bookkeeping, no search"), + ("search::search_job::get_status", "job bookkeeping, no search"), + ("search::search_job::list_status", "job bookkeeping, no search"), + ("traces::get_latest_traces", "checks the licence in the handler"), + ("traces::get_latest_traces_stream", "checks the licence in the handler"), + ("traces::traces_write", "ingestion"), +]; + +/// Pages an over-limit customer needs to understand and fix their state; none may be gated. +const OPERATIONAL_READS: [&str; 22] = [ + "/v2/{org_id}/alerts/history", + "/v2/{org_id}/reports/history", + "/{org_id}/pipelines/history", + "/{org_id}/workflows/{id}/history", + "/v2/{org_id}/alert_calibration/digests", + "/v2/{org_id}/alert_calibration/digests/alerts", + "/v2/{org_id}/alert_calibration/digests/records", + "/v2/{org_id}/alert_calibration/digests/rules", + "/v2/{org_id}/alert_calibration/digests/run", + "/{org_id}/ai/chats", + "/{org_id}/ai/chats/{session_id}", + "/{org_id}/synthetics", + "/{org_id}/synthetics/{id}/runs", + "/{org_id}/_search_history", + "/{org_id}/search/profile", + "/license", + "/license/refresh", + "/{org_id}/config", + "/{org_id}/anomaly_detection/{config_id}/history", + "/{org_id}/anomaly_detection/history", + "/v2/{org_id}/alerts/incidents/{incident_id}", + "/{org_id}/search_jobs/{job_id}", +]; + +const SEARCH_ENTRY: &str = r"(?:search_service|SearchService|crate::search|core_promql::search)::(?:cache::)?(?:search|search_multi|search_partition|search_partition_multi)\(|\bgrpc_search(?:_partition)?\(|process_search_stream_request(?:_multi)?\(|eval_jobs::search::execute|execute_training_query\("; + +/// Every production file that calls a search entry point, and how its search is licence-gated. +#[rustfmt::skip] +const SEARCH_CALLERS: [(&str, Caller); 39] = [ + ("api/management/src/request/alert_calibration/digests.rs", Caller::Operational), + ("api/management/src/request/alerts/history.rs", Caller::Exempt("the org's own alert and anomaly run history; its routes stay open")), + ("api/management/src/request/dashboards/reports_history.rs", Caller::Exempt("the org's own report run history; its route stays open")), + ("api/management/src/request/db_monitoring/service/common.rs", Caller::Routed("db_monitoring::handler::")), + ("api/management/src/request/db_monitoring/service/instance_metrics.rs", Caller::Routed("db_monitoring::handler::")), + ("api/management/src/request/synthetics/last_check.rs", Caller::Exempt("last check on the synthetics list; its route stays open")), + ("api/management/src/request/workflows/mod.rs", Caller::Exempt("the org's own workflow run history; its route stays open")), + ("api/pipelines/src/request/pipelines/history.rs", Caller::Exempt("the org's own pipeline run history; its route stays open")), + ("api/search/src/profiles/mod.rs", Caller::InHandler), + ("api/search/src/promql/mod.rs", Caller::InHandler), + ("api/search/src/rum/lookup/query.rs", Caller::Helper("called by rum lookup, which checks the licence in rum/lookup/mod.rs")), + ("api/search/src/search/around.rs", Caller::Helper("around v1/v2 check in search/mod.rs; around_multi is in the gated group")), + ("api/search/src/search/mod.rs", Caller::InHandler), + ("api/search/src/search/multi_streams.rs", Caller::Routed("search::multi_streams::")), + ("api/search/src/search/patterns.rs", Caller::Routed("patterns::")), + ("api/search/src/search/search_inspector.rs", Caller::Exempt("search/profile stays open (D13)")), + ("api/search/src/search/search_stream.rs", Caller::InHandler), + ("api/search/src/traces/dag.rs", Caller::InHandler), + ("api/search/src/traces/details.rs", Caller::Routed("traces::details::")), + ("api/search/src/traces/list/query.rs", Caller::Helper("latest traces; traces/mod.rs checks the licence")), + ("api/search/src/traces/list/stream.rs", Caller::Helper("latest traces; traces/mod.rs checks the licence")), + ("api/search/src/traces/session.rs", Caller::InHandler), + ("api/search/src/traces/time_index/mod.rs", Caller::InHandler), + ("api/search/src/traces/user.rs", Caller::InHandler), + ("core/src/ai_chat.rs", Caller::Operational), + ("core/src/alerts/alert_calibration/mod.rs", Caller::Background("calibration runs")), + ("core/src/alerts/mod.rs", Caller::Background("alert evaluation")), + ("core/src/alerts/notifications/chart/mod.rs", Caller::Background("notification charts")), + ("core/src/anomaly_detection/partitioned_search.rs", Caller::Background("anomaly detection runs")), + ("core/src/db_monitoring/rollup.rs", Caller::Background("DBM rollup")), + ("core/src/downtimes/values.rs", Caller::Routed("downtimes::")), + ("core/src/metrics/prom.rs", Caller::Routed("promql::label_values")), + ("core/src/self_reporting/search.rs", Caller::Exempt("licence usage from the meta org")), + ("core/src/slo/job.rs", Caller::Background("SLO job")), + ("core/src/slo/reconcile.rs", Caller::Background("SLO reconcile")), + ("core/src/traces/agent_signals/api.rs", Caller::Routed("traces::get_agent_signals")), + ("core/src/traces/service_graph/mod.rs", Caller::Routed("traces::get_current_topology")), + ("jobs/src/job/mod.rs", Caller::Background("eval-job scoring")), + ("jobs/src/job/red_insights.rs", Caller::Background("red_insights")), +]; + +/// How a search call site is licence-gated. +enum Caller { + /// The file calls `check_search_allowed` itself. + InHandler, + /// The file serves handlers under this prefix, which sit in the gated group. + Routed(&'static str), + /// The file reads only an operational stream. + Operational, + Background(&'static str), + Helper(&'static str), + Exempt(&'static str), +} + +type Registration = (String, String, String); + +fn production_source() -> &'static str { + include_str!("mod.rs") + .split_once("\n#[cfg(test)]\nmod tests {") + .unwrap() + .0 +} + +fn fn_body<'a>(source: &'a str, signature: &str) -> &'a str { + let start = source.find(signature).unwrap(); + &source[start..start + source[start..].find("\n}\n").unwrap()] +} + +/// The `/api` tree the router serves, minus the gated group. +fn api_span_outside_the_group() -> String { + let source = production_source(); + let section = |from: &str, to: &str| { + let start = source.find(from).unwrap(); + &source[start..start + source[start..].find(to).unwrap()] + }; + let api = section("fn ai_chat_routes(", "pub fn other_service_routes("); + let group = fn_body(source, "fn search_routes_with("); + format!( + "{}{}", + api.replace(group, ""), + section("fn es_compat_routes(", "\n}\n") + ) +} + +/// Every (method, template, handler) the `.route(` calls in `text` register. +fn registrations(text: &str) -> BTreeSet { + let handler = regex::Regex::new( + r"(?:^|[\s(.,])(get|post|put|delete|patch|head)\(([A-Za-z_][A-Za-z0-9_:]*)\)", + ) + .unwrap(); + let mut found = BTreeSet::new(); + for (at, _) in text.match_indices(".route(") { + let args = &text[at + ".route(".len()..]; + let mut depth = 1; + let end = args + .char_indices() + .find(|(_, c)| { + depth += match c { + '(' => 1, + ')' => -1, + _ => 0, + }; + depth == 0 + }) + .unwrap() + .0; + let args = &args[..end]; + let template = args.split('"').nth(1).unwrap().to_string(); + for c in handler.captures_iter(args) { + found.insert((c[1].to_uppercase(), template.clone(), c[2].to_string())); + } + } + found +} + +fn gated() -> BTreeSet { + GATED + .iter() + .map(|(m, t, h)| (m.to_string(), t.to_string(), h.to_string())) + .collect() +} + +fn module_of(handler: &str) -> &str { + handler.rsplit_once("::").map_or("", |(module, _)| module) +} + +fn search_call_files() -> BTreeSet { + let src = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut roots: Vec<_> = std::fs::read_dir(src.join("api")) + .unwrap() + .map(|e| e.unwrap().path().join("src")) + .collect(); + roots.extend(["core/src", "mcp/src", "synthetics/src", "jobs/src"].map(|d| src.join(d))); + let entry = regex::Regex::new(SEARCH_ENTRY).unwrap(); + let mut hits = BTreeSet::new(); + let mut pending = roots; + while let Some(path) = pending.pop() { + if path.is_dir() { + pending.extend(std::fs::read_dir(&path).unwrap().map(|e| e.unwrap().path())); + continue; + } + let rel = path + .strip_prefix(&src) + .unwrap() + .to_string_lossy() + .replace('\\', "/"); + // This file quotes the entry pattern, so scanning it would find itself. + if !rel.ends_with(".rs") || rel.ends_with("router/search_gate_tests.rs") { + continue; + } + let text = std::fs::read_to_string(&path).unwrap(); + let production = text.split("\n#[cfg(test)]\nmod tests").next().unwrap(); + if entry.is_match(production) { + hits.insert(rel); + } + } + hits +} + +fn caller_source(file: &str) -> String { + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join(file); + let text = std::fs::read_to_string(path).unwrap(); + text.split("\n#[cfg(test)]\nmod tests") + .next() + .unwrap() + .to_string() +} + +#[cfg(feature = "enterprise")] +fn filled(template: &str) -> String { + template + .split('/') + .map(|segment| match segment { + "{org_id}" => "acme".to_string(), + "{stream_name}" => "s1".to_string(), + "{job_id}" => "j1".to_string(), + "{stream_names}" => config::utils::base64::encode_url("s1"), + s if s.starts_with('{') => "p1".to_string(), + s => s.to_string(), + }) + .collect::>() + .join("/") +} + +#[test] +fn search_routes_hold_exactly_the_gated_routes() { + let group = registrations(fn_body(production_source(), "fn search_routes_with(")); + assert_eq!(group, gated()); +} + +#[test] +fn no_gated_handler_is_registered_outside_search_routes() { + let outside = registrations(&api_span_outside_the_group()); + let gated_handlers: BTreeSet = GATED.iter().map(|(_, _, h)| h.to_string()).collect(); + let siblings: BTreeSet<&str> = UNGATED_SIBLINGS.iter().map(|(h, _)| *h).collect(); + let mut seen_siblings = BTreeSet::new(); + for (method, template, handler) in &outside { + let route = format!("{method} {template} {handler}"); + assert!( + !gated_handlers.contains(handler), + "gated handler outside the group: {route}" + ); + assert!( + !FULLY_GATED_MODULES.iter().any(|m| handler.starts_with(m)), + "a fully gated module registers outside the group: {route}" + ); + if MIXED_MODULES.contains(&module_of(handler)) { + assert!( + siblings.contains(handler.as_str()), + "unclassified handler: {route}" + ); + seen_siblings.insert(handler.as_str()); + } + } + assert_eq!( + seen_siblings, siblings, + "UNGATED_SIBLINGS lists a handler no route registers" + ); + assert!( + UNGATED_SIBLINGS + .iter() + .all(|(_, reason)| !reason.is_empty()) + ); +} + +#[test] +fn operational_reads_stay_outside_the_gate() { + let group: BTreeSet = + registrations(fn_body(production_source(), "fn search_routes_with(")) + .into_iter() + .map(|(_, template, _)| template) + .collect(); + let outside: BTreeSet = registrations(&api_span_outside_the_group()) + .into_iter() + .map(|(_, template, _)| template) + .collect(); + for template in OPERATIONAL_READS { + assert!( + outside.contains(template), + "{template} is no longer registered" + ); + assert!( + !group.contains(template), + "{template} must stay outside the gate" + ); + } +} + +#[test] +fn the_gate_is_a_route_layer_on_the_group_only() { + let source = production_source(); + assert_eq!( + source.matches("from_fn(search_license_middleware)").count(), + 1 + ); + assert!( + fn_body(source, "fn search_routes_with(") + .contains(".route_layer(middleware::from_fn(search_license_middleware))") + ); + assert!(!source.contains(".layer(middleware::from_fn(search_license_middleware")); +} + +#[tokio::test] +async fn route_layer_does_not_reach_methods_merged_on_the_same_path() { + const PATH: &str = "/v2/{org_id}/downtimes/values"; + let _built = Router::new() + .merge(search_routes_with(false)) + .route(PATH, get(|| async { "stub" })); + let closed = middleware::from_fn(|_: Request, _: Next| async { + common::meta::http::HttpResponse::too_many_requests("blocked") + }); + let app = Router::new() + .merge( + Router::new() + .route(PATH, post(|| async { "gated" })) + .route_layer(closed), + ) + .route(PATH, get(|| async { "open" })); + let send = |method: Method| { + app.clone().oneshot( + axum::http::Request::builder() + .method(method) + .uri("/v2/acme/downtimes/values") + .body(Body::empty()) + .unwrap(), + ) + }; + assert_eq!(send(Method::GET).await.unwrap().status(), StatusCode::OK); + assert_eq!( + send(Method::POST).await.unwrap().status(), + StatusCode::TOO_MANY_REQUESTS + ); +} + +#[test] +fn every_search_call_site_is_classified() { + let listed: BTreeSet = SEARCH_CALLERS.iter().map(|(f, _)| f.to_string()).collect(); + assert_eq!( + search_call_files(), + listed, + "classify every search call site in SEARCH_CALLERS" + ); + let operational = ["DIGEST_STREAM", "AI_CHAT_EVENTS_STREAM"] + .into_iter() + .chain(search_service::OPERATIONAL_STREAMS); + let operational: Vec<&str> = operational.collect(); + for (file, class) in SEARCH_CALLERS { + match class { + Caller::InHandler => { + assert!( + caller_source(file).contains("check_search_allowed("), + "{file}" + ) + } + Caller::Routed(prefix) => { + assert!( + GATED.iter().any(|(_, _, h)| h.starts_with(prefix)), + "{file}" + ) + } + Caller::Operational => { + let text = caller_source(file); + assert!(operational.iter().any(|name| text.contains(name)), "{file}") + } + Caller::Background(reason) | Caller::Helper(reason) | Caller::Exempt(reason) => { + assert!(!reason.is_empty(), "{file}") + } + } + } +} + +#[cfg(feature = "enterprise")] +#[tokio::test] +async fn closed_gate_refuses_every_gated_route() { + use o2_enterprise::enterprise::license::{TEST_STATICS_LOCK, close_search_gate_for_tests}; + + let _gate = close_search_gate_for_tests(TEST_STATICS_LOCK.lock().await); + let app = search_routes_with(true); + for (method, template, handler) in GATED { + let request = axum::http::Request::builder() + .method(method) + .uri(filled(template)) + .body(Body::empty()) + .unwrap(); + let response = app.clone().oneshot(request).await.unwrap(); + assert_eq!( + response.status(), + StatusCode::TOO_MANY_REQUESTS, + "{method} {template} {handler}" + ); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(json["code"], 429, "{template}"); + let message = json["message"].as_str().unwrap(); + assert!( + message.starts_with( + "Error# Search is temporarily disabled due to exceeding allotted ingestion limit." + ) && message.ends_with("Please contact your administrator."), + "{template}: {message}" + ); + } +} + +#[cfg(feature = "enterprise")] +#[tokio::test] +async fn meta_org_passes_the_closed_gate() { + use o2_enterprise::enterprise::license::{TEST_STATICS_LOCK, close_search_gate_for_tests}; + + let _gate = close_search_gate_for_tests(TEST_STATICS_LOCK.lock().await); + let request = axum::http::Request::post(format!("/{}/_search_multi", config::META_ORG_ID)) + .body(Body::empty()) + .unwrap(); + let response = search_routes_with(true).oneshot(request).await.unwrap(); + assert_ne!(response.status(), StatusCode::TOO_MANY_REQUESTS); +} + +#[cfg(feature = "enterprise")] +#[tokio::test] +async fn open_gate_lets_a_gated_route_through() { + let _held = o2_enterprise::enterprise::license::TEST_STATICS_LOCK + .lock() + .await; + let app = Router::new() + .route("/{org_id}/{stream_name}/probe", get(|| async { "ok" })) + .route_layer(middleware::from_fn(search_license_middleware)); + let request = axum::http::Request::get("/acme/logs/probe") + .body(Body::empty()) + .unwrap(); + assert_eq!(app.oneshot(request).await.unwrap().status(), StatusCode::OK); +}