From 6e025eef61f0cdae3d0d212caf02f578762bdfae Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 13 Nov 2026 08:07:57 +0000 Subject: [PATCH] fix(auth): secure MySQL and PostgreSQL database selection (#10382) * fix(mysql): authorize the default database at login Signed-off-by: Dennis Zhuang * fix(auth): validate MySQL handshake scope and SQL USE Signed-off-by: Dennis Zhuang * fix(auth): authorize PostgreSQL schema selection Signed-off-by: Dennis Zhuang * fix(servers): defer PostgreSQL database lookup until authorized Signed-off-by: Dennis Zhuang * test(servers): propagate catalog lookup errors in protocol tests Signed-off-by: Dennis Zhuang * fix(postgres): skip unusable search path candidates Signed-off-by: Dennis Zhuang * fix(auth): log database authorization and selection outcomes Signed-off-by: Dennis Zhuang * fix(auth): preserve search path responses and schema boundaries Signed-off-by: Dennis Zhuang * fix(auth): clarify database selection diagnostics Signed-off-by: Dennis Zhuang * fix(mysql): format handshake username in structured logs Signed-off-by: Dennis Zhuang * fix(auth): secure database-less sessions and failed startup Signed-off-by: Dennis Zhuang * fix(mysql): retain authorized default database selection Signed-off-by: Dennis Zhuang * refactor(auth): separate database selection from request dispatch Signed-off-by: Dennis Zhuang * test(mysql): align schema validation with database-less login Signed-off-by: Dennis Zhuang * docs(auth): explain database-less session restrictions Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/function/src/permission.rs | 11 + src/common/function/src/system/database.rs | 4 +- src/common/function/src/system/pg_catalog.rs | 12 +- src/frontend/AGENTS.md | 4 + src/frontend/src/instance.rs | 451 +++++++++++++++++- src/frontend/src/instance/schema_selection.rs | 218 +++++++++ src/operator/src/expr_helper.rs | 55 +++ src/operator/src/statement.rs | 50 +- src/servers/src/mysql/handler.rs | 370 ++++++++++---- src/servers/src/mysql/writer.rs | 23 +- src/servers/src/postgres/auth_handler.rs | 91 ++-- src/servers/tests/mod.rs | 6 +- src/servers/tests/mysql/mysql_server_test.rs | 38 +- src/servers/tests/postgres/mod.rs | 126 +++++ src/session/src/lib.rs | 20 + src/sql/src/statements/set_variables.rs | 55 ++- 16 files changed, 1375 insertions(+), 159 deletions(-) create mode 100644 src/frontend/src/instance/schema_selection.rs diff --git a/src/common/function/src/permission.rs b/src/common/function/src/permission.rs index 1a8392a0585..f3f5c6731a4 100644 --- a/src/common/function/src/permission.rs +++ b/src/common/function/src/permission.rs @@ -27,6 +27,17 @@ pub fn ensure_admin_privilege( func_ctx: &FunctionContext, fn_name: &str, ) -> datafusion_common::Result<()> { + // Database-less MySQL sessions have no authorized default scope. Admin targets + // are runtime function arguments, so frontend relation checks cannot authorize + // them, even when qualified. Require database selection before invoking a handler. + if func_ctx.query_ctx.channel() == session::context::Channel::Mysql + && func_ctx.query_ctx.current_schema().is_empty() + { + return Err(DataFusionError::External(Box::new(BoxedError::new( + session::SessionError::NoDatabaseSelected {}, + )))); + } + let admin_stmt = Statement::Admin(Admin::Func(sql::ast::Function { name: ObjectName::from(vec![Ident::new(fn_name)]), uses_odbc_syntax: false, diff --git a/src/common/function/src/system/database.rs b/src/common/function/src/system/database.rs index 1884c5525f0..b9ad7ab7178 100644 --- a/src/common/function/src/system/database.rs +++ b/src/common/function/src/system/database.rs @@ -59,7 +59,9 @@ macro_rules! impl_current_schema_function { let func_ctx = find_function_context(&args)?; let db = func_ctx.query_ctx.current_schema(); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(db)))) + Ok(ColumnarValue::Scalar(ScalarValue::Utf8View( + (!db.is_empty()).then_some(db), + ))) } } }; diff --git a/src/common/function/src/system/pg_catalog.rs b/src/common/function/src/system/pg_catalog.rs index 3029361c202..bc311d84164 100644 --- a/src/common/function/src/system/pg_catalog.rs +++ b/src/common/function/src/system/pg_catalog.rs @@ -90,7 +90,9 @@ impl Function for CurrentSchemaFunction { let func_ctx = find_function_context(&args)?; let db = func_ctx.query_ctx.current_schema(); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(db)))) + Ok(ColumnarValue::Scalar(ScalarValue::Utf8View( + (!db.is_empty()).then_some(db), + ))) } } @@ -158,11 +160,15 @@ impl Function for CurrentSchemasFunction { &self, args: ScalarFunctionArgs, ) -> datafusion_common::Result { + let func_ctx = find_function_context(&args)?; + let schema = func_ctx.query_ctx.current_schema(); let args = ColumnarValue::values_to_arrays(&args.args)?; let input = as_boolean_array(&args[0]); - // Create a UTF8 array with a single value - let mut values = vec!["public"]; + let mut values = Vec::new(); + if !schema.is_empty() { + values.push(schema.as_str()); + } // include implicit schemas if input.value(0) { values.push(INFORMATION_SCHEMA_NAME); diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index 7072431a5ec..93354feb352 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -27,6 +27,7 @@ remote datanodes via `operator`/`client`. | `instance` | `src/frontend/src/instance.rs` | `Instance`: the core handler; implements `SqlQueryHandler`, `PrometheusHandler`, etc. | | `instance/builder` | `src/frontend/src/instance/builder.rs` | `FrontendBuilder` assembles `Instance` from its dependencies | | `instance/grpc` | `src/frontend/src/instance/grpc.rs` | `GrpcQueryHandler`: insert/delete/query/promql over gRPC | +| `instance/schema_selection` | `src/frontend/src/instance/schema_selection.rs` | SQL database selection and authorization for sessions without a selected schema | | `instance/region_query` | `src/frontend/src/instance/region_query.rs` | Routes distributed region reads to datanodes | | `instance/exemplar` | `src/frontend/src/instance/exemplar.rs` | Resolves and authorizes selector branches before bounded exemplar-table scans | | `instance/*` | `src/frontend/src/instance/` | Per-protocol handlers (`influxdb.rs`, `promql.rs`, `otlp/`, `jaeger.rs`, `logs.rs`, `prom_store.rs`, ...) | @@ -43,6 +44,9 @@ including caches owned by pipeline and protocol handlers. - **SQL query** (`instance.rs`): `do_query_inner` handles parsing, interceptors, permission checks, timeout/cancellation, and delegates planning/execution to `StatementExecutor`. Distributed scans enter through `region_query.rs`. + MySQL sessions without a selected database validate explicit scopes before + planning or executing SQL, including prepared statements; unqualified targets + fail without falling back to `public`. - **Exemplar query** (`instance/exemplar.rs`): authorize all selector branches before calling `PromQueryInterceptor::pre_execute_exemplars` for each resolved exemplar table, with that table's schema context, before scanning it. diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index f881846818b..e2d061167a7 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -30,6 +30,7 @@ mod otlp; pub mod prom_store; mod promql; mod region_query; +mod schema_selection; use std::collections::HashSet; use std::pin::Pin; @@ -396,6 +397,16 @@ impl Instance { } _ => { query_interceptor.pre_execute(Some(&stmt), None, query_ctx.clone())?; + if let Statement::SetVariables(set) = &stmt + && query_ctx.channel() == Channel::Postgres + && let Some(schemas) = set.search_path() + { + return self.set_search_path(schemas, query_ctx).await; + } + if let Statement::Use(schema) = &stmt { + self.authorize_database(query_ctx.current_catalog(), schema, &query_ctx) + .await?; + } if let Statement::Copy(sql::statements::copy::Copy::CopyDatabase( CopyDatabase::From(arg), )) = &stmt @@ -561,6 +572,8 @@ impl Instance { stmt: &Statement, query_ctx: &QueryContextRef, ) -> Result<()> { + self.check_unselected_mysql_database(stmt, query_ctx) + .await?; self.plugins .get::() .as_ref() @@ -782,7 +795,9 @@ impl Instance { } Err(e) => { if e.status_code().should_log_error() { - error!(e; "Failed to execute query: {stmt}"); + error!(e; user = query_ctx.current_user().username(), protocol = %query_ctx.channel(), + catalog = query_ctx.current_catalog(), schema = %query_ctx.current_schema(), + process_id = query_ctx.process_id(), "Failed to execute query: {stmt}"); } else { debug!("Failed to execute query: {stmt}, {e}"); } @@ -837,6 +852,15 @@ impl Instance { query_ctx: QueryContextRef, ) -> Result { ensure!(!self.is_suspended(), error::SuspendedSnafu); + if let Some(stmt) = &stmt { + self.check_unselected_mysql_database(stmt, &query_ctx) + .await?; + } else if query_ctx.channel() == Channel::Mysql && query_ctx.current_schema().is_empty() { + return Err(BoxedError::new( + session::SessionError::NoDatabaseSelected {}, + )) + .context(ExternalSnafu); + } let query_interceptor_opt = self.plugins.get::>(); let query_interceptor = query_interceptor_opt.as_ref(); @@ -993,6 +1017,8 @@ impl Instance { query_ctx: QueryContextRef, ) -> Result> { ensure!(!self.is_suspended(), error::SuspendedSnafu); + self.check_unselected_mysql_database(&stmt, &query_ctx) + .await?; // EXPLAIN / EXPLAIN ANALYZE wrap an inner statement; describe them when the // wrapped statement is something we already plan (so that bind parameters @@ -3476,6 +3502,429 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_sql_schema_selection_authorizes_before_mutation() -> TestResult<()> { + struct SchemaProvider; + #[async_trait::async_trait] + impl auth::UserProvider for SchemaProvider { + fn name(&self) -> &str { + "schema-test" + } + async fn authenticate( + &self, + _: auth::Identity<'_>, + _: auth::Password<'_>, + ) -> auth::error::Result { + unreachable!("the request already has an authenticated user") + } + async fn authorize( + &self, + _: &str, + schema: &str, + _: &UserInfoRef, + ) -> auth::error::Result<()> { + if schema == "backend_error" { + return auth::error::InternalStateSnafu { + msg: "backend unavailable", + } + .fail(); + } + if matches!(schema, "public" | "allowed" | "allowed_missing" | "admin") { + Ok(()) + } else { + auth::error::PermissionDeniedSnafu.fail() + } + } + } + let plugins = Plugins::new(); + plugins.insert::(Arc::new(SchemaProvider)); + let catalog = + catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?); + for schema in ["denied", "allowed", "admin"] { + catalog + .register_schema_sync(catalog::RegisterSchemaRequest { + catalog: "greptime".to_string(), + schema: schema.to_string(), + }) + .unwrap(); + } + let instance = test_instance_with_catalog_manager( + catalog, + test_table(1025, "target")?, + plugins, + None, + Arc::new(NoopProcedureExecutor), + ) + .await?; + for sql in [ + "Use denied", + " USE denied", + "/* switch */ USE denied", + "USE public; Use denied", + "USE missing", + ] { + let ctx = Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::Postgres, + )); + let results = instance.do_query_inner(sql, ctx.clone()).await; + assert_eq!( + StatusCode::PermissionDenied, + results.last().unwrap().as_ref().unwrap_err().status_code(), + "{sql}" + ); + assert_eq!("public", ctx.current_schema()); + } + for sql in ["USE allowed", "SET search_path TO allowed"] { + let ctx = Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::Postgres, + )); + assert!( + instance + .do_query_inner(sql, ctx.clone()) + .await + .remove(0) + .is_ok() + ); + assert_eq!("allowed", ctx.current_schema()); + } + let ctx = Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::Postgres, + )); + let error = instance + .do_query_inner("USE allowed_missing", ctx.clone()) + .await + .remove(0) + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + assert_eq!("public", ctx.current_schema()); + for (sql, expected) in [ + ("SET search_path TO denied, public", "public"), + ("SET search_path TO DEFAULT", "admin"), + ("SET search_path TO allowed_missing, allowed", "allowed"), + ( + "SET search_path TO denied, allowed_missing, public", + "public", + ), + (r#"SET search_path TO "$user", public"#, "admin"), + (r#"SET "search_path" TO denied, allowed"#, "allowed"), + ("SET SEARCH_PATH TO 'denied'", ""), + ("SET search_path TO allowed_missing", ""), + ("SET search_path TO ''", ""), + ("SET search_path TO allowed; SET search_path TO denied", ""), + ] { + let ctx = Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::Postgres, + )); + ctx.set_current_user(Arc::new(AdminUserInfo)); + let results = instance.do_query_inner(sql, ctx.clone()).await; + assert!( + results.iter().all(|result| matches!( + result, + Ok(Output { + data: OutputData::AffectedRows(0), + .. + }) + )), + "{sql}: {results:?}" + ); + assert_eq!(expected, ctx.current_schema(), "{sql}"); + } + + let ctx = Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::Postgres, + )); + // The default user's schema is absent or unauthorized, so public remains usable. + assert!( + instance + .do_query_inner(r#"SET search_path TO "$user", public"#, ctx.clone()) + .await + .remove(0) + .is_ok() + ); + assert_eq!("public", ctx.current_schema()); + assert!( + instance + .do_query_inner("SET search_path TO DEFAULT", ctx.clone()) + .await + .remove(0) + .is_ok() + ); + assert_eq!("public", ctx.current_schema()); + let error = instance + .do_query_inner("SET search_path TO backend_error, allowed", ctx.clone()) + .await + .remove(0) + .unwrap_err(); + assert_eq!(StatusCode::Unexpected, error.status_code()); + assert_eq!("public", ctx.current_schema()); + Ok(()) + } + + #[tokio::test] + async fn test_mysql_without_selected_database() -> TestResult<()> { + let plugins = Plugins::new(); + plugins.insert::(Arc::new(auth::tests::MockUserProvider::default())); + let instance = test_instance_with_plugins( + test_table(1024, "source")?, + test_table(1025, "target")?, + plugins, + ) + .await?; + let ctx = Arc::new(QueryContext::with_channel("greptime", "", Channel::Mysql)); + for sql in [ + "SELECT * FROM source", + "INSERT INTO source VALUES (1, 0)", + "CREATE TABLE t (ts TIMESTAMP TIME INDEX)", + "SHOW TABLES", + "TRUNCATE TABLE source", + "SHOW CREATE TABLE source", + "COPY source TO '/tmp/out'", + "SELECT flush_table('public.source')", + ] { + let result = instance.do_query_inner(sql, ctx.clone()).await.remove(0); + let error = match result { + Err(error) => Err::<(), _>(BoxedError::new(error)) + .context(servers::error::ExecuteQuerySnafu) + .unwrap_err(), + Ok(output) => { + let OutputData::Stream(stream) = output.data else { + panic!("Expected a query stream: {sql}"); + }; + common_recordbatch::util::collect(stream) + .await + .context(servers::error::CollectRecordbatchSnafu) + .unwrap_err() + } + }; + let (kind, message) = servers::mysql::writer::handle_err(error, ctx.clone()); + assert_eq!(1046, kind as u16, "{sql}: {message}"); + assert!(message.contains("No database selected"), "{message}"); + } + let output = instance + .do_query_inner( + "SELECT 1, DATABASE() IS NULL AS no_database, SCHEMA() IS NULL AS no_schema", + ctx.clone(), + ) + .await + .remove(0) + .unwrap(); + assert_eq!(2, output.data.pretty_print().await.matches("true").count()); + for sql in [ + "SELECT * FROM denied.source", + "SHOW TABLES FROM denied", + "USE denied", + ] { + let error = instance + .do_query_inner(sql, ctx.clone()) + .await + .remove(0) + .unwrap_err(); + assert_eq!( + StatusCode::AccessDenied, + error.status_code(), + "{sql}: {error}" + ); + assert!(ctx.current_schema().is_empty()); + } + for sql in [ + "SELECT * FROM public.source", + "SHOW TABLES FROM public", + "SET time_zone = '+00:00'", + ] { + assert!( + instance + .do_query_inner(sql, ctx.clone()) + .await + .remove(0) + .is_ok(), + "{sql}" + ); + } + let plan = instance + .do_describe_inner(parse_one_sql("SELECT * FROM public.source"), ctx.clone()) + .await + .unwrap() + .unwrap() + .logical_plan; + instance + .do_exec_plan_inner( + plan.clone(), + Some(parse_one_sql("SELECT * FROM public.source")), + ctx.clone(), + ) + .await + .unwrap(); + let error = instance + .do_describe_inner(parse_one_sql("SELECT * FROM denied.source"), ctx.clone()) + .await + .unwrap_err(); + assert_eq!(StatusCode::AccessDenied, error.status_code()); + assert!(ctx.current_schema().is_empty()); + let error = instance + .do_exec_plan_inner(plan.clone(), None, ctx.clone()) + .await + .unwrap_err(); + assert_eq!(StatusCode::DatabaseNotFound, error.status_code()); + let current_user = ctx.current_user(); + ctx.set_current_user(Arc::new(AdminUserInfo)); + let error = instance + .do_exec_plan_inner( + plan, + Some(parse_one_sql("SELECT * FROM public.source")), + ctx.clone(), + ) + .await + .unwrap_err(); + assert_eq!(StatusCode::AccessDenied, error.status_code()); + ctx.set_current_user(current_user); + instance + .do_query_inner("USE public", ctx.clone()) + .await + .remove(0) + .unwrap(); + assert_eq!("public", ctx.current_schema()); + instance + .do_query_inner("SELECT * FROM source", ctx) + .await + .remove(0) + .unwrap(); + Ok(()) + } + + #[tokio::test] + async fn test_empty_search_path_statements() -> TestResult<()> { + let instance = test_instance_with_plugins( + test_table(1024, "source")?, + test_table(1025, "target")?, + Plugins::new(), + ) + .await?; + let ctx = Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::Postgres, + )); + instance + .do_query_inner("SET search_path TO ''", ctx.clone()) + .await + .remove(0) + .unwrap(); + let tables = instance + .do_query_inner("SHOW TABLES", ctx.clone()) + .await + .remove(0) + .unwrap() + .data + .pretty_print() + .await; + assert!( + !tables.contains("source") && !tables.contains("target"), + "{tables}" + ); + let error = instance + .do_query_inner("INSERT INTO source VALUES (1, 0)", ctx.clone()) + .await + .remove(0) + .unwrap_err(); + assert!( + matches!(error, Error::TableOperation { source: operator::error::Error::TableNotFound { ref table_name, .. }, .. } if table_name == "greptime..source"), + "{error:?}" + ); + for sql in [ + "CREATE TABLE no_schema (ts TIMESTAMP TIME INDEX)", + "ALTER TABLE source ADD COLUMN no_schema_column INT", + "CREATE VIEW no_schema_view AS SELECT 1 AS n", + "CREATE FLOW no_schema_flow SINK TO sink AS SELECT * FROM public.source", + "CREATE FLOW no_schema_flow SINK TO public.sink AS SELECT * FROM source", + ] { + let error = instance + .do_query_inner(sql, ctx.clone()) + .await + .remove(0) + .unwrap_err(); + assert!( + matches!(error, Error::TableOperation { source: operator::error::Error::SchemaNotFound { ref schema_info, .. }, .. } if schema_info.is_empty()), + "{sql}: {error:?}" + ); + } + + for sql in ["SELECT * FROM source", "DELETE FROM source"] { + assert!( + instance + .do_query_inner(sql, ctx.clone()) + .await + .remove(0) + .is_err(), + "{sql}" + ); + } + let output = instance + .do_query_inner( + "SELECT current_schema() IS NULL AS empty, current_schemas(false) AS schemas", + ctx.clone(), + ) + .await + .remove(0) + .unwrap(); + let printed = output.data.pretty_print().await; + assert!( + printed.contains("true") && printed.contains("[]"), + "{printed}" + ); + let output = instance + .do_query_inner("DROP TABLE IF EXISTS source", ctx.clone()) + .await + .remove(0) + .unwrap(); + assert!(matches!(output.data, OutputData::AffectedRows(0))); + instance + .do_query_inner("USE public", ctx.clone()) + .await + .remove(0) + .unwrap(); + assert!( + instance + .do_query_inner("SELECT * FROM source", ctx) + .await + .remove(0) + .is_ok() + ); + Ok(()) + } + + #[tokio::test] + async fn test_operator_rejects_search_path_without_frontend() -> TestResult<()> { + let instance = test_instance_with_plugins( + test_table(1024, "source")?, + test_table(1025, "target")?, + Plugins::new(), + ) + .await?; + let ctx = Arc::new(QueryContext::with_channel( + "greptime", + "public", + Channel::Postgres, + )); + let error = instance + .statement_executor + .execute_sql(parse_one_sql("SET search_path TO denied"), ctx.clone()) + .await + .unwrap_err(); + assert!(matches!(error, operator::error::Error::NotSupported { .. })); + assert_eq!("public", ctx.current_schema()); + Ok(()) + } + #[tokio::test] async fn test_write_only_grpc_sql_is_checked_after_parsing() -> TestResult<()> { let plugins = Plugins::new(); diff --git a/src/frontend/src/instance/schema_selection.rs b/src/frontend/src/instance/schema_selection.rs new file mode 100644 index 00000000000..af0dbc33984 --- /dev/null +++ b/src/frontend/src/instance/schema_selection.rs @@ -0,0 +1,218 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashSet; + +use common_error::ext::{BoxedError, ErrorExt}; +use common_error::status_code::StatusCode; +use common_query::Output; +use common_telemetry::{debug, warn}; +use session::context::{Channel, QueryContextRef}; +use session::table_name::table_idents_to_full_name; +use snafu::{IntoError, ResultExt}; +use sql::ast::ObjectNamePartExt; +use sql::statements::copy::{Copy, CopyDatabase}; +use sql::statements::statement::Statement; +use sql::util::extract_tables_from_statement_checked; + +use crate::error::{ + Error, ExternalSnafu, InvalidSqlSnafu, PermissionSnafu, Result, TableOperationSnafu, +}; +use crate::instance::Instance; + +impl Instance { + /// Filters unauthorized candidates before selecting an existing schema. + pub(crate) async fn set_search_path( + &self, + schemas: Vec, + query_ctx: QueryContextRef, + ) -> Result { + let mut authorized = Vec::with_capacity(schemas.len()); + let provider = self.plugins.get::(); + let user = query_ctx.current_user(); + for schema in schemas { + let schema = if schema == "$user" { + user.username().to_string() + } else { + schema + }; + if schema.is_empty() { + continue; + } + if let Some(provider) = &provider { + match provider + .authorize(query_ctx.current_catalog(), &schema, &user) + .await + { + Ok(()) => {} + Err(error) + if matches!( + error.status_code(), + StatusCode::PermissionDenied | StatusCode::AccessDenied + ) => + { + debug!(user = user.username(), protocol = %query_ctx.channel(), + catalog = query_ctx.current_catalog(), schema = %schema, + process_id = query_ctx.process_id(), reason = %error.status_code(), + "Skipping unauthorized search_path candidate"); + continue; + } + Err(error) => return Err(error).context(PermissionSnafu), + } + } + authorized.push(schema); + } + self.statement_executor + .set_search_path(authorized, query_ctx) + .await + .context(TableOperationSnafu) + } + + /// Authorizes an explicit database access without changing the session. + pub(crate) async fn authorize_database( + &self, + catalog: &str, + schema: &str, + query_ctx: &QueryContextRef, + ) -> Result<()> { + let Some(provider) = self.plugins.get::() else { + return Ok(()); + }; + let user = query_ctx.current_user(); + if let Err(error) = provider.authorize(catalog, schema, &user).await { + if matches!( + error.status_code(), + StatusCode::PermissionDenied | StatusCode::AccessDenied + ) { + warn!(user = user.username(), protocol = %query_ctx.channel(), catalog, schema, + previous_schema = %query_ctx.current_schema(), process_id = query_ctx.process_id(), + reason = %error.status_code(), "Database access denied"); + } + return Err(error).context(PermissionSnafu); + } + Ok(()) + } + + /// Checks explicit targets before planning when no default scope is authorized. + pub(crate) async fn check_unselected_mysql_database( + &self, + stmt: &Statement, + query_ctx: &QueryContextRef, + ) -> Result<()> { + if query_ctx.channel() != Channel::Mysql || !query_ctx.current_schema().is_empty() { + return Ok(()); + } + for (catalog, schema) in resolve_mysql_database_scopes(stmt, query_ctx)? { + self.authorize_database(&catalog, &schema, query_ctx) + .await?; + } + Ok(()) + } +} + +fn resolve_mysql_database_scopes( + stmt: &Statement, + query_ctx: &QueryContextRef, +) -> Result> { + let mut stmt = stmt; + while let Statement::Explain(explain) = stmt { + stmt = &explain.statement; + } + let mut scopes = HashSet::new(); + let tables = extract_tables_from_statement_checked(stmt).ok_or_else(no_database_selected)?; + for table in tables { + let (catalog, schema, _) = table_idents_to_full_name(&table, query_ctx) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + if schema.is_empty() { + return Err(no_database_selected()); + } + scopes.insert((catalog, schema)); + } + if let Some(scope) = resolve_statement_database_scope(stmt, query_ctx)? { + scopes.insert(scope); + } + if scopes.is_empty() + && !matches!( + stmt, + Statement::Query(_) + | Statement::Use(_) + | Statement::SetVariables(_) + | Statement::ShowVariables(_) + | Statement::ShowDatabases(_) + | Statement::ShowCharset(_) + | Statement::ShowCollation(_) + | Statement::ShowStatus(_) + | Statement::ShowProcesslist(_) + | Statement::Kill(_) + ) + { + return Err(no_database_selected()); + } + Ok(scopes) +} + +fn resolve_statement_database_scope( + stmt: &Statement, + query_ctx: &QueryContextRef, +) -> Result> { + let database = match stmt { + Statement::ShowTables(show) => Some(show.database.as_deref()), + Statement::ShowTableStatus(show) => Some(show.database.as_deref()), + Statement::ShowViews(show) => Some(show.database.as_deref()), + _ => None, + }; + if let Some(database) = database { + let schema = database + .filter(|schema| !schema.is_empty()) + .ok_or_else(no_database_selected)?; + return Ok(Some(( + query_ctx.current_catalog().to_string(), + schema.to_string(), + ))); + } + let name = match stmt { + Statement::CreateDatabase(create) => &create.name, + Statement::AlterDatabase(alter) => alter.database_name(), + Statement::DropDatabase(drop) => drop.name(), + Statement::ShowCreateDatabase(show) => &show.database_name, + Statement::Copy(Copy::CopyDatabase(copy)) => match copy { + CopyDatabase::To(args) | CopyDatabase::From(args) => &args.database_name, + }, + _ => return Ok(None), + }; + let scope = match &name.0[..] { + [schema] => ( + query_ctx.current_catalog().to_string(), + schema.to_string_unquoted(), + ), + [catalog, schema] => (catalog.to_string_unquoted(), schema.to_string_unquoted()), + _ => { + return InvalidSqlSnafu { + err_msg: format!("Invalid database name: {name}"), + } + .fail(); + } + }; + if scope.1.is_empty() { + return Err(no_database_selected()); + } + Ok(Some(scope)) +} + +fn no_database_selected() -> Error { + ExternalSnafu.into_error(BoxedError::new( + session::SessionError::NoDatabaseSelected {}, + )) +} diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index af6c960556c..359ab9722b5 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -836,6 +836,14 @@ pub(crate) fn to_alter_table_expr( .map_err(BoxedError::new) .context(ExternalSnafu)?; + // SQL must not reach the gRPC entry point's default-schema fallback. + ensure!( + !schema_name.is_empty(), + crate::error::SchemaNotFoundSnafu { + schema_info: &schema_name + } + ); + let kind = match alter_table.alter_operation { AlterTableOperation::AddConstraint(_) => { return NotSupportedSnafu { @@ -1130,6 +1138,13 @@ pub fn to_create_flow_task_expr( .map(|s| s.to_owned()) .unwrap_or(query_ctx.current_schema()); + ensure!( + !schema.is_empty(), + crate::error::SchemaNotFoundSnafu { + schema_info: &schema + } + ); + let sink_table_name = TableName { catalog_name: catalog, schema_name: schema, @@ -1153,6 +1168,13 @@ pub fn to_create_flow_task_expr( .map(|s| s.to_string()) .unwrap_or(query_ctx.current_schema()); + ensure!( + !schema.is_empty(), + crate::error::SchemaNotFoundSnafu { + schema_info: &schema + } + ); + let table_name = TableName { catalog_name: catalog, schema_name: schema, @@ -1453,6 +1475,39 @@ SELECT * FROM tql; ); } + #[test] + fn test_create_flow_requires_schema_for_unqualified_tables() { + let ctx = QueryContext::arc(); + ctx.set_current_schema(""); + for (sink, source, valid) in [ + ("sink", "public.source", false), + ("public.sink", "source", false), + ("public.sink", "public.source", true), + ] { + let sql = format!("CREATE FLOW test_flow SINK TO {sink} AS SELECT * FROM {source}"); + let stmt = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .pop() + .unwrap(); + let Statement::CreateFlow(flow) = stmt else { + unreachable!() + }; + let result = to_create_flow_task_expr(flow, &ctx); + if valid { + assert!(result.is_ok(), "{result:?}"); + } else { + assert!( + matches!(result, Err(crate::error::Error::SchemaNotFound { ref schema_info, .. }) if schema_info.is_empty()), + "{result:?}" + ); + } + } + } + #[test] fn test_create_flow_expr() { let sql = r" diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index b52cb814cbb..feab3b15cba 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -522,11 +522,48 @@ impl StatementExecutor { SchemaNotFoundSnafu { schema_info: &db } ); + let previous_schema = query_ctx.current_schema(); query_ctx.set_current_schema(&db); + debug!(user = query_ctx.current_user().username(), protocol = %query_ctx.channel(), + catalog, schema = %db, previous_schema, process_id = query_ctx.process_id(), + "Database selected"); Ok(Output::new_with_record_batches(RecordBatches::empty())) } + /// Selects the first existing schema from candidates authorized by the caller. + pub async fn set_search_path( + &self, + schemas: Vec, + query_ctx: QueryContextRef, + ) -> Result { + let mut selected = String::new(); + for schema in schemas { + if !schema.is_empty() + && self + .catalog_manager + .schema_exists(query_ctx.current_catalog(), &schema, Some(&query_ctx)) + .await + .context(CatalogSnafu)? + { + selected = schema; + break; + } + debug!(user = query_ctx.current_user().username(), protocol = %query_ctx.channel(), + catalog = query_ctx.current_catalog(), schema = %schema, + process_id = query_ctx.process_id(), reason = "schema_not_found", + "Skipping unavailable search_path candidate"); + } + // An empty effective path must not retain access to the previous schema. + let previous_schema = query_ctx.current_schema(); + query_ctx.set_current_schema(&selected); + debug!(user = query_ctx.current_user().username(), protocol = %query_ctx.channel(), + catalog = query_ctx.current_catalog(), schema = %selected, previous_schema, + process_id = query_ctx.process_id(), empty_path = selected.is_empty(), + "Search path schema selected"); + Ok(Output::new_with_affected_rows(0)) + } + fn set_variables(&self, set_var: SetVariables, query_ctx: QueryContextRef) -> Result { let variable = sql::ast::object_name_to_string_unquoted(&set_var.variable); let var_name = variable.to_uppercase(); @@ -594,17 +631,10 @@ impl StatementExecutor { } }, "SEARCH_PATH" => { - if query_ctx.channel() == Channel::Postgres { - let search_path = set_var.search_path().context(NotSupportedSnafu { - feat: "Unsupported search path in set variable statement", - })?; - query_ctx.set_current_schema(search_path); - } else { - return NotSupportedSnafu { - feat: format!("Unsupported set variable {}", var_name), - } - .fail(); + return NotSupportedSnafu { + feat: "Unsupported search_path assignment", } + .fail(); } _ => { if query_ctx.channel() == Channel::Postgres || query_ctx.channel() == Channel::Mysql diff --git a/src/servers/src/mysql/handler.rs b/src/servers/src/mysql/handler.rs index bdfc7a72057..5d202c77af1 100644 --- a/src/servers/src/mysql/handler.rs +++ b/src/servers/src/mysql/handler.rs @@ -151,6 +151,124 @@ impl MysqlInstanceShim { } } + async fn authenticate_handshake_user( + &self, + auth_plugin: &str, + username: &str, + salt: &[u8], + auth_data: &[u8], + catalog: &str, + ) -> Option { + let Some(user_provider) = &self.user_provider else { + return Some(auth::userinfo_by_name(Some(username.to_string()))); + }; + let addr = self + .session + .conn_info() + .client_addr + .map(|addr| addr.to_string()); + let result = if username == BEARER_TOKEN_USER { + if auth_plugin != MysqlAuthMethod::CLEAR_PASSWORD_PLUGIN { + warn!("Bearer-token MySQL authentication requires mysql_clear_password"); + return None; + } + let token = auth_data.strip_suffix(&[0]).unwrap_or(auth_data); + let Ok(token) = std::str::from_utf8(token) else { + warn!("Bearer token is not valid UTF-8"); + return None; + }; + user_provider + .authenticate_bearer_token(token, catalog) + .await + } else { + let user_id = Identity::UserId(username, addr.as_deref()); + let password = match auth_plugin { + MysqlAuthMethod::NATIVE_PASSWORD_PLUGIN => { + Password::MysqlNativePassword(auth_data, salt) + } + MysqlAuthMethod::CLEAR_PASSWORD_PLUGIN => { + // The raw bytes received could be represented in C-like string, ended in '\0'. + // We must "trim" it to get the real password string. + let password = auth_data.strip_suffix(&[0]).unwrap_or(auth_data); + Password::PlainText(String::from_utf8_lossy(password).to_string().into()) + } + other => { + error!("Unsupported mysql auth plugin: {}", other); + return None; + } + }; + user_provider.authenticate(user_id, password).await + }; + match result { + Ok(user_info) => Some(user_info), + Err(error) => { + METRIC_AUTH_FAILURE + .with_label_values(&[error.status_code().as_ref()]) + .inc(); + warn!(error; "Failed to auth"); + None + } + } + } + + // on_init can write an ERR without ending the handshake. Validate here before + // publishing the authenticated session or entering opensrv's command loop. + async fn select_handshake_database( + &self, + catalog: &str, + schema: String, + user_info: &auth::UserInfoRef, + explicit_database: bool, + ) -> Option { + if let Some(provider) = &self.user_provider + && let Err(error) = provider.authorize(catalog, &schema, user_info).await + { + let denied = matches!( + error.status_code(), + common_error::status_code::StatusCode::AccessDenied + | common_error::status_code::StatusCode::PermissionDenied + ); + if !explicit_database && denied { + debug!(user = user_info.username(), protocol = "mysql", catalog, schema = %schema, + reason = %error.status_code(), + "Default database not authorized; starting without a selected database"); + return Some(String::new()); + } + METRIC_AUTH_FAILURE + .with_label_values(&[error.status_code().as_ref()]) + .inc(); + if denied { + warn!(user = user_info.username(), protocol = "mysql", catalog, schema = %schema, + reason = %error.status_code(), "Handshake database authorization denied"); + } else { + error!(error; user = user_info.username(), protocol = "mysql", catalog, schema = %schema, + "Handshake database authorization failed"); + } + return None; + } + if !explicit_database && schema.is_empty() { + return Some(schema); + } + match self.query_handler.is_valid_schema(catalog, &schema).await { + Ok(true) => Some(schema), + Ok(false) if !explicit_database => { + debug!(user = user_info.username(), protocol = "mysql", catalog, schema = %schema, + "Default database does not exist; starting without a selected database"); + Some(String::new()) + } + Ok(false) => { + warn!(user = user_info.username(), protocol = "mysql", catalog, schema = %schema, + "Unknown MySQL handshake database"); + None + } + Err(error) => { + error!(error; user = user_info.username(), protocol = "mysql", catalog, schema = %schema, + "Failed to validate MySQL handshake database"); + None + } + } + } + /// Describe the statement async fn do_describe( &self, @@ -432,76 +550,40 @@ impl AsyncMysqlShim for MysqlInstanceShi auth_data: &[u8], database: Option<&[u8]>, ) -> bool { + let database = match database.map(std::str::from_utf8).transpose() { + Ok(database) => database, + Err(_) => { + warn!("MySQL database is not valid UTF-8"); + return false; + } + }; + let (catalog, schema) = if let Some(database) = database { + let (catalog, schema) = parse_optional_catalog_and_schema_from_db_string(database); + (catalog.unwrap_or_else(|| self.session.catalog()), schema) + } else { + (self.session.catalog(), self.session.schema()) + }; // if not specified then **greptime** will be used let username = String::from_utf8_lossy(username); - let mut user_info = None; - let addr = self - .session - .conn_info() - .client_addr - .map(|addr| addr.to_string()); - if let Some(user_provider) = &self.user_provider { - let result = if username.as_ref() == BEARER_TOKEN_USER { - if auth_plugin != MysqlAuthMethod::CLEAR_PASSWORD_PLUGIN { - warn!("Bearer-token MySQL authentication requires mysql_clear_password"); - return false; - } - let token = auth_data.strip_suffix(&[0]).unwrap_or(auth_data); - let Ok(token) = std::str::from_utf8(token) else { - warn!("Bearer token is not valid UTF-8"); - return false; - }; - let catalog = if let Some(database) = database { - let Ok(database) = std::str::from_utf8(database) else { - warn!("MySQL database is not valid UTF-8"); - return false; - }; - parse_optional_catalog_and_schema_from_db_string(database) - .0 - .unwrap_or_else(|| self.session.catalog()) - } else { - self.session.catalog() - }; - user_provider - .authenticate_bearer_token(token, &catalog) - .await - } else { - let user_id = Identity::UserId(&username, addr.as_deref()); - let password = match auth_plugin { - MysqlAuthMethod::NATIVE_PASSWORD_PLUGIN => { - Password::MysqlNativePassword(auth_data, salt) - } - MysqlAuthMethod::CLEAR_PASSWORD_PLUGIN => { - // The raw bytes received could be represented in C-like string, ended in '\0'. - // We must "trim" it to get the real password string. - let password = auth_data.strip_suffix(&[0]).unwrap_or(auth_data); - Password::PlainText(String::from_utf8_lossy(password).to_string().into()) - } - other => { - error!("Unsupported mysql auth plugin: {}", other); - return false; - } - }; - user_provider.authenticate(user_id, password).await - }; - match result { - Ok(userinfo) => { - user_info = Some(userinfo); - } - Err(e) => { - METRIC_AUTH_FAILURE - .with_label_values(&[e.status_code().as_ref()]) - .inc(); - warn!(e; "Failed to auth"); - return false; - } - }; - } - let user_info = - user_info.unwrap_or_else(|| auth::userinfo_by_name(Some(username.to_string()))); - + let Some(user_info) = self + .authenticate_handshake_user(auth_plugin, &username, salt, auth_data, &catalog) + .await + else { + return false; + }; + let Some(schema) = self + .select_handshake_database(&catalog, schema, &user_info, database.is_some()) + .await + else { + return false; + }; + self.session.set_catalog(catalog); + self.session.set_schema(schema); self.session.set_user_info(user_info); + debug!(user = self.session.user_info().username(), protocol = "mysql", + catalog = %self.session.catalog(), schema = %self.session.schema(), + client_addr = ?self.session.conn_info().client_addr, "Handshake database selected"); true } @@ -610,9 +692,12 @@ impl AsyncMysqlShim for MysqlInstanceShi return Ok(()); } Err(e) => { - writer - .error(ErrorKind::ER_SP_BADSTATEMENT, e.output_msg().as_bytes()) - .await?; + let kind = if writer::is_no_database_selected(&e) { + ErrorKind::ER_NO_DB_ERROR + } else { + ErrorKind::ER_SP_BADSTATEMENT + }; + writer.error(kind, e.output_msg().as_bytes()).await?; return Ok(()); } } @@ -684,20 +769,6 @@ impl AsyncMysqlShim for MysqlInstanceShi self.session.catalog() }; - if !self - .query_handler - .is_valid_schema(&catalog, &schema) - .await? - { - return w - .error( - ErrorKind::ER_WRONG_DB_NAME, - format!("Unknown database '{}'", database).as_bytes(), - ) - .await - .map_err(|e| e.into()); - } - let user_info = &self.session.user_info(); if let Some(schema_validator) = &self.user_provider @@ -708,6 +779,15 @@ impl AsyncMysqlShim for MysqlInstanceShi METRIC_AUTH_FAILURE .with_label_values(&[e.status_code().as_ref()]) .inc(); + if e.status_code().should_log_error() { + error!(e; user = user_info.username(), protocol = "mysql", catalog = %catalog, schema = %schema, + client_addr = ?self.session.conn_info().client_addr, reason = %e.status_code(), + "Database selection authorization failed"); + } else { + warn!(user = user_info.username(), protocol = "mysql", catalog = %catalog, schema = %schema, + previous_schema = %self.session.schema(), client_addr = ?self.session.conn_info().client_addr, + reason = %e.status_code(), "Database selection denied"); + } return w .error( ErrorKind::ER_DBACCESS_DENIED_ERROR, @@ -717,10 +797,32 @@ impl AsyncMysqlShim for MysqlInstanceShi .map_err(|e| e.into()); } + if !self + .query_handler + .is_valid_schema(&catalog, &schema) + .await? + { + debug!(user = user_info.username(), protocol = "mysql", catalog = %catalog, schema = %schema, + client_addr = ?self.session.conn_info().client_addr, reason = "schema_not_found", + "Database selection failed"); + return w + .error( + ErrorKind::ER_WRONG_DB_NAME, + format!("Unknown database '{}'", database).as_bytes(), + ) + .await + .map_err(|e| e.into()); + } + + let previous_catalog = self.session.catalog(); + let previous_schema = self.session.schema(); if catalog_from_db.is_some() { self.session.set_catalog(catalog) } self.session.set_schema(schema); + debug!(user = user_info.username(), protocol = "mysql", catalog = %self.session.catalog(), + schema = %self.session.schema(), previous_catalog, previous_schema, + client_addr = ?self.session.conn_info().client_addr, "Database selected"); w.ok().await.map_err(|e| e.into()) } @@ -996,6 +1098,7 @@ mod tests { use std::sync::Arc; use async_trait::async_trait; + use auth::UserProvider; use common_query::Output; use datafusion_expr::LogicalPlan; use query::parser::PromQuery; @@ -1040,8 +1143,14 @@ mod tests { unimplemented!() } - async fn is_valid_schema(&self, _: &str, _: &str) -> Result { - Ok(true) + async fn is_valid_schema(&self, _: &str, schema: &str) -> Result { + if schema == "backend_error" { + return crate::error::InternalSnafu { + err_msg: "Schema backend unavailable", + } + .fail(); + } + Ok(schema != "missing") } } @@ -1055,6 +1164,97 @@ mod tests { ) } + struct HandshakeUserProvider { + inner: auth::tests::MockUserProvider, + fail_authorization: bool, + } + + #[async_trait] + impl UserProvider for HandshakeUserProvider { + fn name(&self) -> &str { + "handshake_test" + } + + async fn authenticate( + &self, + identity: Identity<'_>, + password: Password<'_>, + ) -> auth::error::Result { + self.inner.authenticate(identity, password).await + } + + async fn authorize( + &self, + catalog: &str, + schema: &str, + user_info: &auth::UserInfoRef, + ) -> auth::error::Result<()> { + if self.fail_authorization { + auth::error::InternalStateSnafu { + msg: "Authorization backend unavailable", + } + .fail() + } else { + self.inner.authorize(catalog, schema, user_info).await + } + } + } + + #[tokio::test] + async fn test_handshake_database_selection() { + for (allowed, default_schema, database, fail_authorization, expected_schema) in [ + ("allowed", "public", None, false, Some("")), + ("allowed", "public", Some("denied"), false, None), + ("allowed", "public", Some("allowed"), false, Some("allowed")), + ("missing", "public", Some("missing"), false, None), + ("public", "public", None, false, Some("public")), + ("missing", "missing", None, false, Some("")), + ("backend_error", "backend_error", None, false, None), + ("public", "public", None, true, None), + ] { + let mut shim = create_shim(); + shim.session.set_schema(default_schema.to_string()); + shim.user_provider = Some(Arc::new(HandshakeUserProvider { + inner: auth::tests::MockUserProvider { + schema: allowed.to_string(), + ..Default::default() + }, + fail_authorization, + })); + let result = + >::authenticate_with_database( + &shim, + MysqlAuthMethod::CLEAR_PASSWORD_PLUGIN, + b"greptime", + &[], + b"greptime\0", + database.map(str::as_bytes), + ) + .await; + assert_eq!(expected_schema.is_some(), result, "{allowed}: {database:?}"); + assert_eq!( + expected_schema.unwrap_or(default_schema), + shim.session.schema() + ); + } + } + + #[tokio::test] + async fn test_reject_invalid_handshake_database() { + let shim = create_shim(); + assert!( + !>::authenticate_with_database( + &shim, + MysqlAuthMethod::NATIVE_PASSWORD_PLUGIN, + b"greptime", + &[], + &[], + Some(&[0xff]), + ) + .await + ); + } + #[test] fn test_prepared_statement_cache_metrics() { use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; diff --git a/src/servers/src/mysql/writer.rs b/src/servers/src/mysql/writer.rs index c485a5189fc..cde1e082d07 100644 --- a/src/servers/src/mysql/writer.rs +++ b/src/servers/src/mysql/writer.rs @@ -151,10 +151,31 @@ pub async fn write_output( Ok(()) } +pub(super) fn is_no_database_selected(error: &impl ErrorExt) -> bool { + matches!( + error.as_any().downcast_ref::(), + Some(session::SessionError::NoDatabaseSelected {}) + ) || std::iter::successors(error.source(), |error| error.source()).any(|error| { + let session_error = error.downcast_ref::().or_else(|| { + error + .downcast_ref::() + .and_then(|error| error.as_any().downcast_ref::()) + }); + matches!( + session_error, + Some(session::SessionError::NoDatabaseSelected {}) + ) + }) +} + /// Handle GreptimeDB error, convert it to MySQL error pub fn handle_err(e: impl ErrorExt, query_ctx: QueryContextRef) -> (ErrorKind, String) { let status_code = e.status_code(); - let kind = mysql_error_kind(&status_code); + let kind = if is_no_database_selected(&e) { + ErrorKind::ER_NO_DB_ERROR + } else { + mysql_error_kind(&status_code) + }; if status_code.should_log_error() { let root_error = e.root_cause().unwrap_or(&e); diff --git a/src/servers/src/postgres/auth_handler.rs b/src/servers/src/postgres/auth_handler.rs index 169ce594f5e..1c1211490b8 100644 --- a/src/servers/src/postgres/auth_handler.rs +++ b/src/servers/src/postgres/auth_handler.rs @@ -24,6 +24,7 @@ use async_trait::async_trait; use common_catalog::parse_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; +use common_telemetry::{debug, error, warn}; use common_time::Timezone; use futures::{Sink, SinkExt}; use pgwire::api::auth::StartupHandler; @@ -88,8 +89,7 @@ impl PgUserProviderSource { } /// (user, catalog, schema) of this login. Catalog/schema are parsed from - /// the `database` startup parameter, which was validated at connection - /// start. + /// the `database` startup parameter. Existence is checked after authorization. fn resolve(login: &WireLoginInfo<'_>) -> Option<(String, String, String)> { let user = login.user()?.to_string(); let (catalog, schema) = parse_catalog_and_schema_from_db_string(login.database()?); @@ -136,6 +136,13 @@ impl PgUserProviderSource { { Err(e) => { record_auth_failure(e.status_code()); + if e.status_code().should_log_error() { + error!(e; user = user_info.username(), protocol = "postgres", catalog = %catalog, + schema = %schema, client_addr = login.host(), reason = %e.status_code(), "Startup database authorization failed"); + } else { + warn!(user = user_info.username(), protocol = "postgres", catalog = %catalog, + schema = %schema, client_addr = login.host(), reason = %e.status_code(), "Startup database authorization denied"); + } Err(PgWireError::ApiError(Box::new(AuthSnafu.into_error(e)))) } Ok(()) => Ok(()), @@ -233,6 +240,9 @@ where if let Some(current_schema) = client.metadata().get(super::METADATA_SCHEMA) { session.set_schema(current_schema.clone()); } + debug!(user = session.user_info().username(), protocol = "postgres", + catalog = %session.catalog(), schema = %session.schema(), + client_addr = ?session.conn_info().client_addr, "Handshake database selected"); } enum MechanismSelection { @@ -291,18 +301,18 @@ impl StartupHandler for PostgresServerHandlerInner { auth::protocol_negotiation(client, startup).await?; auth::save_startup_parameters_to_metadata(client, startup); - // check if db is valid - match resolve_db_info(client, self.query_handler.clone()).await? { - DbResolution::Resolved(catalog, schema) => { - let metadata = client.metadata_mut(); - let _ = metadata.insert(super::METADATA_CATALOG.to_owned(), catalog); - let _ = metadata.insert(super::METADATA_SCHEMA.to_owned(), schema); - } - DbResolution::NotFound(msg) => { - send_error(client, PgErrorCode::Ec3D000.to_err_info(msg)).await?; - return Ok(()); - } - } + let Some(db) = client.metadata().get(super::METADATA_DATABASE) else { + send_error( + client, + PgErrorCode::Ec3D000.to_err_info("Database not specified".to_owned()), + ) + .await?; + return Ok(()); + }; + let (catalog, schema) = parse_catalog_and_schema_from_db_string(db); + let metadata = client.metadata_mut(); + let _ = metadata.insert(super::METADATA_CATALOG.to_owned(), catalog); + let _ = metadata.insert(super::METADATA_SCHEMA.to_owned(), schema); // try to set TimeZone if let Some(tz) = client.metadata().get("TimeZone") { @@ -323,6 +333,9 @@ impl StartupHandler for PostgresServerHandlerInner { match self.select_mechanism(client).await { MechanismSelection::NoAuth => { + if !check_database_exists(client, &self.query_handler).await? { + return Ok(()); + } self.session.set_user_info(userinfo_by_name( client.metadata().get(super::METADATA_USER).cloned(), )); @@ -387,6 +400,9 @@ impl StartupHandler for PostgresServerHandlerInner { return send_password_authentication_failed(client).await; } + if !check_database_exists(client, &self.query_handler).await? { + return Ok(()); + } self.session.set_user_info(user_info); set_client_info(client, &self.session); self.auth_driver @@ -437,32 +453,35 @@ where .await } -enum DbResolution { - Resolved(String, String), - NotFound(String), -} - -/// A function extracted to resolve lifetime and readability issues: -async fn resolve_db_info( +async fn check_database_exists( client: &mut C, - query_handler: ServerSqlQueryHandlerRef, -) -> PgWireResult + query_handler: &ServerSqlQueryHandlerRef, +) -> PgWireResult where - C: ClientInfo + Unpin + Send, + C: ClientInfo + Sink + Unpin + Send, + C::Error: Debug, + PgWireError: From<>::Error>, { - let db_ref = client.metadata().get(super::METADATA_DATABASE); - if let Some(db) = db_ref { + let db = client.metadata().get(super::METADATA_DATABASE); + let error = if let Some(db) = db { let (catalog, schema) = parse_catalog_and_schema_from_db_string(db); - if query_handler - .is_valid_schema(&catalog, &schema) - .await - .map_err(convert_err)? - { - Ok(DbResolution::Resolved(catalog, schema)) - } else { - Ok(DbResolution::NotFound(format!("Database not found: {db}"))) + match query_handler.is_valid_schema(&catalog, &schema).await { + Ok(true) => return Ok(true), + Ok(false) => {} + Err(error) => { + let mut error_info: ErrorInfo = convert_err(error).into(); + error_info.severity = "FATAL".to_string(); + send_error(client, error_info).await?; + return Ok(false); + } } + debug!(protocol = "postgres", user = ?client.metadata().get(super::METADATA_USER), + catalog = %catalog, schema = %schema, reason = "schema_not_found", + "Startup database validation failed"); + format!("Database not found: {db}") } else { - Ok(DbResolution::NotFound("Database not specified".to_owned())) - } + "Database not specified".to_owned() + }; + send_error(client, PgErrorCode::Ec3D000.to_err_info(error)).await?; + Ok(false) } diff --git a/src/servers/tests/mod.rs b/src/servers/tests/mod.rs index d78f57b7b83..84d1d27255a 100644 --- a/src/servers/tests/mod.rs +++ b/src/servers/tests/mod.rs @@ -233,13 +233,13 @@ impl SqlQueryHandler for DummyInstance { } async fn is_valid_schema(&self, catalog: &str, schema: &str) -> Result { - Ok(self - .query_engine + self.query_engine .engine_state() .catalog_manager() .schema_exists(catalog, schema, None) .await - .unwrap()) + .map_err(common_error::ext::BoxedError::new) + .context(servers::error::CheckDatabaseValiditySnafu) } } diff --git a/src/servers/tests/mysql/mysql_server_test.rs b/src/servers/tests/mysql/mysql_server_test.rs index fdbdd4fd8f9..c1c6eb5e739 100644 --- a/src/servers/tests/mysql/mysql_server_test.rs +++ b/src/servers/tests/mysql/mysql_server_test.rs @@ -278,6 +278,13 @@ async fn test_schema_validation() -> Result<()> { let _ = create_connection_default_db_name(server_port, false) .await .unwrap(); + let mut connection = create_connection(server_port, None, false).await.unwrap(); + for schema in ["information_schema", "missing"] { + match connection.query_drop(format!("USE {schema}")).await { + Err(mysql_async::Error::Server(error)) => assert_eq!(1044, error.code), + result => panic!("expected database access denied for {schema}, got {result:?}"), + } + } mysql_server.shutdown().await.unwrap(); // change to another username @@ -290,6 +297,15 @@ async fn test_schema_validation() -> Result<()> { let fail = create_connection_default_db_name(server_port, false).await; assert!(fail.is_err()); + let mut connection = create_connection(server_port, None, false).await.unwrap(); + match connection.query_drop("USE public").await { + Err(mysql_async::Error::Server(error)) => assert_eq!(1044, error.code), + result => panic!("expected database access denied, got {result:?}"), + } + let database: Option> = + connection.query_first("SELECT DATABASE()").await.unwrap(); + assert_eq!(Some(None), database); + connection.disconnect().await.unwrap(); mysql_server.shutdown().await.unwrap(); Ok(()) @@ -461,7 +477,7 @@ async fn test_bearer_token_auth_over_clear_password() -> Result<()> { assert_eq!(1, value); assert_eq!(1, provider.authentications.load(Ordering::Relaxed)); assert_eq!( - usize::from(database.is_some()), + 1 + usize::from(database.is_some()), provider.authorizations.load(Ordering::Relaxed) ); connection.disconnect().await.unwrap(); @@ -579,9 +595,23 @@ async fn test_db_name() -> Result<()> { mysql_server.start(listening).await.unwrap(); let server_addr = mysql_server.bind_addr().unwrap(); - // None actually uses default database name - let r = create_connection_default_db_name(server_addr.port(), client_tls).await; - let _ = r.unwrap(); + let mut connection = create_connection(server_addr.port(), None, client_tls) + .await + .unwrap(); + let database: Option = connection + .query_first("SELECT DATABASE()") + .await + .unwrap() + .flatten(); + assert_eq!(Some("public".to_string()), database); + connection.query_drop("USE public").await.unwrap(); + let database: String = connection + .query_first("SELECT DATABASE()") + .await + .unwrap() + .unwrap(); + assert_eq!("public", database); + connection.disconnect().await.unwrap(); let r = create_connection(server_addr.port(), Some("tomcat"), client_tls).await; assert!(r.is_err()); diff --git a/src/servers/tests/postgres/mod.rs b/src/servers/tests/postgres/mod.rs index 1bc5e79c236..a47da20c177 100644 --- a/src/servers/tests/postgres/mod.rs +++ b/src/servers/tests/postgres/mod.rs @@ -346,6 +346,132 @@ async fn test_schema_validating() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_database_existence_is_checked_after_authentication_and_authorization() -> Result<()> { + let provider = Arc::new(MixedAuthProvider { + scram_verifier: PgScramSha256Verifier::from_password( + MIXED_AUTH_SCRAM_PASSWORD, + b"database-auth-salt", + 4096, + ) + .unwrap(), + selections: Mutex::new(HashMap::new()), + authorize_calls: Mutex::new(HashMap::new()), + }); + let (server, port) = start_test_server_with_user_provider(provider, Default::default()).await?; + for database in [DEFAULT_SCHEMA_NAME, "missing", "missing_catalog-public"] { + for (user, password, authenticated) in [ + (MIXED_AUTH_CLEARTEXT_USER, "wrong", false), + (MIXED_AUTH_SCRAM_USER, "wrong", false), + ("unknown", "wrong", false), + (MIXED_AUTH_DENIED_USER, "scram-password", false), + ( + MIXED_AUTH_CLEARTEXT_USER, + MIXED_AUTH_CLEARTEXT_PASSWORD, + true, + ), + (MIXED_AUTH_SCRAM_USER, "scram-password", true), + ] { + let result = tokio_postgres::Config::new() + .host("127.0.0.1") + .port(port) + .user(user) + .password(password) + .dbname(database) + .connect_timeout(Duration::from_secs(2)) + .connect(NoTls) + .await; + if authenticated && database == DEFAULT_SCHEMA_NAME { + let (client, connection) = result.unwrap(); + let task = tokio::spawn(connection); + client.simple_query("SELECT 1").await.unwrap(); + drop(client); + task.await.unwrap().unwrap(); + } else { + let error = match result { + Err(error) => error, + Ok(_) => panic!("unexpected successful login for {user} in {database}"), + }; + assert_eq!( + Some(if !authenticated { + &SqlState::INVALID_PASSWORD + } else if database == "missing_catalog-public" { + &SqlState::INTERNAL_ERROR + } else { + &SqlState::INVALID_CATALOG_NAME + }), + error.code(), + "{user} in {database}: {error}" + ); + } + } + } + server.shutdown().await.unwrap(); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_startup_database_backend_error_closes_connection() -> Result<()> { + let provider = Arc::new(MixedAuthProvider { + scram_verifier: PgScramSha256Verifier::from_password( + MIXED_AUTH_SCRAM_PASSWORD, + b"database-auth-salt", + 4096, + ) + .unwrap(), + selections: Mutex::new(HashMap::new()), + authorize_calls: Mutex::new(HashMap::new()), + }); + let (server, port) = start_test_server_with_user_provider(provider, Default::default()).await?; + let mut stream = TcpStream::connect(("127.0.0.1", port)).await.unwrap(); + // The in-memory catalog reports an error for an unknown catalog, exercising + // the backend-error path after successful authentication and authorization. + let mut startup = 196608i32.to_be_bytes().to_vec(); + startup.extend_from_slice( + format!("user\0{MIXED_AUTH_CLEARTEXT_USER}\0database\0missing_catalog-public\0\0") + .as_bytes(), + ); + stream + .write_all(&((startup.len() + 4) as i32).to_be_bytes()) + .await + .unwrap(); + stream.write_all(&startup).await.unwrap(); + let mut header = [0; 5]; + stream.read_exact(&mut header).await.unwrap(); + assert_eq!(b'R', header[0]); + let mut body = vec![0; i32::from_be_bytes(header[1..].try_into().unwrap()) as usize - 4]; + stream.read_exact(&mut body).await.unwrap(); + assert_eq!(3i32.to_be_bytes(), body.as_slice()); + let password = format!("{MIXED_AUTH_CLEARTEXT_PASSWORD}\0"); + let mut message = vec![b'p']; + message.extend_from_slice(&((password.len() + 4) as i32).to_be_bytes()); + message.extend_from_slice(password.as_bytes()); + stream.write_all(&message).await.unwrap(); + stream.read_exact(&mut header).await.unwrap(); + assert_eq!(b'E', header[0]); + let mut body = vec![0; i32::from_be_bytes(header[1..].try_into().unwrap()) as usize - 4]; + stream.read_exact(&mut body).await.unwrap(); + assert!(body.windows(7).any(|part| part == b"SFATAL\0")); + assert!(body.windows(7).any(|part| part == b"CXX000\0")); + // A client that ignores the startup error must not enter the query loop. + let query = b"SELECT 1\0"; + let mut message = vec![b'Q']; + message.extend_from_slice(&((query.len() + 4) as i32).to_be_bytes()); + message.extend_from_slice(query); + let _ = stream.write_all(&message).await; + let mut response = [0]; + let read = tokio::time::timeout(Duration::from_secs(2), stream.read(&mut response)) + .await + .unwrap(); + assert!( + matches!(read, Ok(0)) + || matches!(read, Err(ref error) if error.kind() == std::io::ErrorKind::ConnectionReset), + "{read:?}" + ); + server.shutdown().await.unwrap(); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_pg_scram_sha256_auth() -> Result<()> { let verifier = diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index fe4e7bdcc01..ae6bae9906c 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -38,6 +38,26 @@ use derive_more::Debug; use crate::context::{Channel, ConnInfo, QueryContextRef, dialect_for_channel}; +/// Errors caused by the current connection state. +#[derive(snafu::Snafu)] +#[common_macro::stack_trace_debug] +pub enum SessionError { + #[snafu(display("No database selected"))] + NoDatabaseSelected {}, +} + +impl common_error::ext::ErrorExt for SessionError { + fn status_code(&self) -> common_error::status_code::StatusCode { + // Keep the shared database-error category; the MySQL writer inspects this + // typed cause through error wrappers to emit ER_NO_DB_ERROR (1046). + common_error::status_code::StatusCode::DatabaseNotFound + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + /// Maximum number of warnings to store per session (similar to MySQL's max_error_count) const MAX_WARNINGS: usize = 64; diff --git a/src/sql/src/statements/set_variables.rs b/src/sql/src/statements/set_variables.rs index 346f1762e62..8b70e3e7fa8 100644 --- a/src/sql/src/statements/set_variables.rs +++ b/src/sql/src/statements/set_variables.rs @@ -26,23 +26,36 @@ pub struct SetVariables { } impl SetVariables { - /// Returns the first supported `search_path` value. - pub fn search_path(&self) -> Option<&str> { + /// Returns the schema candidates for a supported `search_path` assignment. + pub fn search_path(&self) -> Option> { let [ObjectNamePart::Identifier(variable)] = self.variable.0.as_slice() else { return None; }; - if variable.quote_style.is_some() || !variable.value.eq_ignore_ascii_case("search_path") { + if !variable.value.eq_ignore_ascii_case("search_path") { return None; } - match self.value.first()? { - Expr::Value(ValueWithSpan { - value: Value::SingleQuotedString(value) | Value::DoubleQuotedString(value), - .. - }) => Some(value), - Expr::Identifier(value) => Some(&value.value), - _ => None, + if let [Expr::Identifier(value)] = self.value.as_slice() + && value.quote_style.is_none() + && value.value.eq_ignore_ascii_case("default") + { + return Some(vec!["$user".to_string(), "public".to_string()]); } + + self.value + .iter() + .map(|value| match value { + Expr::Value(ValueWithSpan { + value: Value::SingleQuotedString(value) | Value::DoubleQuotedString(value), + .. + }) => Some(value.clone()), + Expr::Identifier(value) if value.quote_style.is_none() => { + Some(value.value.to_lowercase()) + } + Expr::Identifier(value) => Some(value.value.clone()), + _ => None, + }) + .collect() } } @@ -95,11 +108,20 @@ SET delayed_insert_timeout = 300"#, #[test] fn test_search_path() { for (sql, expected) in [ - ("SET search_path TO public", Some("public")), - ("SET SEARCH_PATH TO 'private'", Some("private")), - ("SET search_path TO \"quoted\"", Some("quoted")), - ("SET \"search_path\" TO private", None), + ("SET search_path TO public", Some(vec!["public"])), + ("SET search_path TO DEFAULT", Some(vec!["$user", "public"])), + (r#"SET search_path TO "default""#, Some(vec!["default"])), + ("SET search_path TO 'default'", Some(vec!["default"])), + ("SET SEARCH_PATH TO 'private'", Some(vec!["private"])), + ("SET search_path TO \"quoted\"", Some(vec!["quoted"])), + ("SET \"search_path\" TO private", Some(vec!["private"])), ("SET timezone TO 'UTC'", None), + ( + r#"SET search_path TO "$user", MISSING, "Mixed", 'a,b'"#, + Some(vec!["$user", "missing", "Mixed", "a,b"]), + ), + ("SET search_path TO ''", Some(vec![""])), + ("SET search_path TO public, 42", None), ] { let statements = ParserContext::create_with_dialect( sql, @@ -110,7 +132,10 @@ SET delayed_insert_timeout = 300"#, let [Statement::SetVariables(set)] = statements.as_slice() else { unreachable!() }; - assert_eq!(expected, set.search_path()); + assert_eq!( + expected.map(|v| v.into_iter().map(str::to_owned).collect::>()), + set.search_path() + ); } } }