rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-15 07:17:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(auth): secure MySQL and PostgreSQL database selection (#10382)

Browse Source 
* fix(mysql): authorize the default database at login

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(auth): validate MySQL handshake scope and SQL USE

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(auth): authorize PostgreSQL schema selection

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(servers): defer PostgreSQL database lookup until authorized

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* test(servers): propagate catalog lookup errors in protocol tests

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(postgres): skip unusable search path candidates

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(auth): log database authorization and selection outcomes

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(auth): preserve search path responses and schema boundaries

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(auth): clarify database selection diagnostics

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(mysql): format handshake username in structured logs

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(auth): secure database-less sessions and failed startup

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(mysql): retain authorized default database selection

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* refactor(auth): separate database selection from request dispatch

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* test(mysql): align schema validation with database-less login

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* docs(auth): explain database-less session restrictions

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

---------

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
dennis zhuang
2026-11-13 08:07:57 +00:00
committed by GitHub
parent f88af153ec
commit 6e025eef61
16 changed files with 1375 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/function/src/permission.rs
+11
View File
@@ -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,
src/common/function/src/system/database.rs
+3 -1
View File
@@ -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),
)))
}
}
};
src/common/function/src/system/pg_catalog.rs
+9 -3
View File
@@ -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<ColumnarValue> {
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);
src/frontend/AGENTS.md
+4
View File
@@ -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.
src/frontend/src/instance.rs
+450 -1
View File
@@ -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::<PermissionCheckerRef>()
.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<Output> {
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::<SqlQueryInterceptorRef<Error>>();
let query_interceptor = query_interceptor_opt.as_ref();
@@ -993,6 +1017,8 @@ impl Instance {
query_ctx: QueryContextRef,
) -> Result<Option<DescribeResult>> {
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<UserInfoRef> {
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::<auth::UserProviderRef>(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::<auth::UserProviderRef>(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();
src/frontend/src/instance/schema_selection.rs
+218
View File
@@ -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<String>,
query_ctx: QueryContextRef,
) -> Result<Output> {
let mut authorized = Vec::with_capacity(schemas.len());
let provider = self.plugins.get::<auth::UserProviderRef>();
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::<auth::UserProviderRef>() 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<HashSet<(String, String)>> {
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<Option<(String, String)>> {
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 {},
))
}
src/operator/src/expr_helper.rs
+55
View File
@@ -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"
src/operator/src/statement.rs
+40 -10
View File
@@ -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<String>,
query_ctx: QueryContextRef,
) -> Result<Output> {
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<Output> {
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
src/servers/src/mysql/handler.rs
+285 -85
View File
@@ -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<auth::UserInfoRef> {
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<String> {
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<W: AsyncWrite + Send + Sync + Unpin> AsyncMysqlShim<W> 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<W: AsyncWrite + Send + Sync + Unpin> AsyncMysqlShim<W> 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<W: AsyncWrite + Send + Sync + Unpin> AsyncMysqlShim<W> 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<W: AsyncWrite + Send + Sync + Unpin> AsyncMysqlShim<W> 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<W: AsyncWrite + Send + Sync + Unpin> AsyncMysqlShim<W> 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<bool> {
Ok(true)
async fn is_valid_schema(&self, _: &str, schema: &str) -> Result<bool> {
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<auth::UserInfoRef> {
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 =
<MysqlInstanceShim as AsyncMysqlShim<tokio::io::Sink>>::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!(
!<MysqlInstanceShim as AsyncMysqlShim<tokio::io::Sink>>::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};
src/servers/src/mysql/writer.rs
+22 -1
View File
@@ -151,10 +151,31 @@ pub async fn write_output<W: AsyncWrite + Send + Sync + Unpin>(
Ok(())
}
pub(super) fn is_no_database_selected(error: &impl ErrorExt) -> bool {
matches!(
error.as_any().downcast_ref::<session::SessionError>(),
Some(session::SessionError::NoDatabaseSelected {})
) || std::iter::successors(error.source(), |error| error.source()).any(|error| {
let session_error = error.downcast_ref::<session::SessionError>().or_else(|| {
error
.downcast_ref::<common_error::ext::BoxedError>()
.and_then(|error| error.as_any().downcast_ref::<session::SessionError>())
});
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);
src/servers/src/postgres/auth_handler.rs
+55 -36
View File
@@ -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<C>(
async fn check_database_exists<C>(
client: &mut C,
query_handler: ServerSqlQueryHandlerRef,
) -> PgWireResult<DbResolution>
query_handler: &ServerSqlQueryHandlerRef,
) -> PgWireResult<bool>
where
C: ClientInfo + Unpin + Send,
C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send,
C::Error: Debug,
PgWireError: From<<C as Sink<PgWireBackendMessage>>::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)
}
src/servers/tests/mod.rs
+3 -3
View File
@@ -233,13 +233,13 @@ impl SqlQueryHandler for DummyInstance {
}
async fn is_valid_schema(&self, catalog: &str, schema: &str) -> Result<bool> {
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)
}
}
src/servers/tests/mysql/mysql_server_test.rs
+34 -4
View File
@@ -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<Option<String>> =
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<String> = 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());
src/servers/tests/postgres/mod.rs
+126
View File
@@ -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 =
src/session/src/lib.rs
+20
View File
@@ -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;
src/sql/src/statements/set_variables.rs
+40 -15
View File
@@ -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<Vec<String>> {
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::<Vec<_>>()),
set.search_path()
);
}
}
}