Merge pull request #5211 from junnplus/management-check

add access check for management mode

query/src/interpreters/access/management_mode_access.rs

@@ -0,0 +1,101 @@
+// Copyright 2022 Datafuse Labs.
+//
+// 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::sync::Arc;
+
+use common_exception::ErrorCode;
+use common_exception::Result;
+use common_planners::PlanNode;
+use common_planners::ShowPlan;
+
+use crate::sessions::QueryContext;
+
+pub struct ManagementModeAccess {
+    ctx: Arc<QueryContext>,
+}
+
+impl ManagementModeAccess {
+    pub fn create(ctx: Arc<QueryContext>) -> Self {
+        ManagementModeAccess { ctx }
+    }
+
+    // Check what we can do if in management mode.
+    pub fn check(&self, plan: &PlanNode) -> Result<()> {
+        // Allows for management-mode.
+        if self.ctx.get_config().query.management_mode {
+            return match plan {
+                PlanNode::Empty(_)
+
+                // Show.
+                | PlanNode::Show(ShowPlan::ShowDatabases(_))
+                | PlanNode::Show(ShowPlan::ShowTables(_))
+                | PlanNode::Show(ShowPlan::ShowEngines(_))
+                | PlanNode::Show(ShowPlan::ShowFunctions(_))
+                | PlanNode::Show(ShowPlan::ShowGrants(_))
+                | PlanNode::Show(ShowPlan::ShowSettings(_))
+                | PlanNode::Show(ShowPlan::ShowUsers(_))
+                | PlanNode::Show(ShowPlan::ShowRoles(_))
+
+                // Database.
+                | PlanNode::CreateDatabase(_)
+                | PlanNode::ShowCreateDatabase(_)
+                | PlanNode::DropDatabase(_)
+
+                // Table.
+                | PlanNode::CreateTable(_)
+                | PlanNode::DropTable(_)
+                | PlanNode::DescribeTable(_)
+                | PlanNode::ShowCreateTable(_)
+
+                // User.
+                | PlanNode::CreateUser(_)
+                | PlanNode::DropUser(_)
+                | PlanNode::AlterUser(_)
+
+                // Role.
+                | PlanNode::CreateRole(_)
+                | PlanNode::DropRole(_)
+
+                // Privilege.
+                | PlanNode::GrantPrivilege(_)
+                | PlanNode::RevokePrivilege(_)
+                | PlanNode::GrantRole(_)
+                | PlanNode::RevokeRole(_)
+
+                // Stage.
+                | PlanNode::CreateUserStage(_)
+                | PlanNode::DropUserStage(_)
+                | PlanNode::DescribeUserStage(_)
+
+                | PlanNode::List(_)
+
+                // UDF.
+                | PlanNode::CreateUserUDF(_)
+                | PlanNode::DropUserUDF(_)
+                | PlanNode::AlterUserUDF(_)
+
+                // USE.
+                | PlanNode::UseDatabase(_)
+
+                // CALL
+                | PlanNode::Call(_) => Ok(()),
+                _ => Err(ErrorCode::ManagementModePermissionDenied(format!(
+                    "Access denied for operation:{:?} in management-mode",
+                    plan.name()
+                ))),
+            };
+        };
+        Ok(())
+    }
+}

query/src/interpreters/access/mod.rs

@@ -0,0 +1,17 @@
+// Copyright 2022 Datafuse Labs.
+//
+// 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.
+
+mod management_mode_access;
+
+pub use management_mode_access::ManagementModeAccess;

query/src/interpreters/interpreter_factory_interceptor.rs

@@ -22,6 +22,7 @@ use common_streams::ErrorStream;
 use common_streams::ProgressStream;
 use common_streams::SendableDataBlockStream;
 
+use crate::interpreters::access::ManagementModeAccess;
 use crate::interpreters::Interpreter;
 use crate::interpreters::InterpreterPtr;
 use crate::interpreters::InterpreterQueryLog;
@@ -30,18 +31,22 @@ use crate::sessions::QueryContext;
 
 pub struct InterceptorInterpreter {
     ctx: Arc<QueryContext>,
+    plan: PlanNode,
     inner: InterpreterPtr,
     query_log: InterpreterQueryLog,
     source_pipe_builder: Mutex<Option<SourcePipeBuilder>>,
+    management_mode_access: ManagementModeAccess,
 }
 
 impl InterceptorInterpreter {
     pub fn create(ctx: Arc<QueryContext>, inner: InterpreterPtr, plan: PlanNode) -> Self {
         InterceptorInterpreter {
             ctx: ctx.clone(),
+            plan: plan.clone(),
             inner,
-            query_log: InterpreterQueryLog::create(ctx, Some(plan)),
+            query_log: InterpreterQueryLog::create(ctx.clone(), Some(plan)),
             source_pipe_builder: Mutex::new(None),
+            management_mode_access: ManagementModeAccess::create(ctx),
         }
     }
 }
@@ -56,6 +61,9 @@ impl Interpreter for InterceptorInterpreter {
         &self,
         input_stream: Option<SendableDataBlockStream>,
     ) -> Result<SendableDataBlockStream> {
+        // Management mode access check.
+        self.management_mode_access.check(&self.plan)?;
+
         let _ = self
             .inner
             .set_source_pipe_builder((*self.source_pipe_builder.lock()).clone());

query/src/interpreters/mod.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod access;
 mod interpreter;
 mod interpreter_call;
 mod interpreter_common;
@@ -71,6 +72,7 @@ mod interpreter_view_drop;
 mod plan_schedulers;
 mod stream;
 
+pub use access::ManagementModeAccess;
 pub use interpreter::Interpreter;
 pub use interpreter::InterpreterPtr;
 pub use interpreter_call::CallInterpreter;

query/tests/it/interpreters/access/management_mode_access.rs

@@ -0,0 +1,171 @@
+// Copyright 2022 Datafuse Labs.
+//
+// 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 common_base::tokio;
+use common_exception::Result;
+use databend_query::interpreters::InterpreterFactory;
+use databend_query::sql::PlanParser;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn test_management_mode_access() -> Result<()> {
+    struct TestGroup {
+        name: &'static str,
+        tests: Vec<Test>,
+    }
+
+    struct Test {
+        name: &'static str,
+        query: &'static str,
+        is_err: bool,
+    }
+
+    let groups: Vec<TestGroup> = vec![
+        TestGroup {
+            name: "show",
+            tests: vec![
+                Test {
+                    name: "show-databases",
+                    query: "SHOW DATABASES",
+                    is_err: false,
+                },
+                Test {
+                    name: "show-engines",
+                    query: "SHOW ENGINES",
+                    is_err: false,
+                },
+                Test {
+                    name: "show-functions",
+                    query: "SHOW FUNCTIONS",
+                    is_err: false,
+                },
+                Test {
+                    name: "show-grants",
+                    query: "SHOW GRANTS",
+                    is_err: false,
+                },
+                Test {
+                    name: "show-settings",
+                    query: "SHOW SETTINGS",
+                    is_err: false,
+                },
+                Test {
+                    name: "show-tables",
+                    query: "SHOW TABLES",
+                    is_err: false,
+                },
+                Test {
+                    name: "show-users",
+                    query: "SHOW USERS",
+                    is_err: false,
+                },
+            ],
+        },
+        TestGroup {
+            name: "database",
+            tests: vec![
+                Test {
+                    name: "db-create-access-passed",
+                    query: "CREATE DATABASE db1",
+                    is_err: false,
+                },
+                Test {
+                    name: "db-show-access-passed",
+                    query: "SHOW CREATE DATABASE db1",
+                    is_err: false,
+                },
+                Test {
+                    name: "db-drop-access-passed",
+                    query: "DROP DATABASE IF EXISTS db1",
+                    is_err: false,
+                },
+            ],
+        },
+        TestGroup {
+            name: "table",
+            tests: vec![
+                Test {
+                    name: "table-create-access-passed",
+                    query: "CREATE TABLE t1(a int)",
+                    is_err: false,
+                },
+                Test {
+                    name: "table-desc-access-passed",
+                    query: "DESC t1",
+                    is_err: false,
+                },
+                Test {
+                    name: "table-show-create-access-passed",
+                    query: "SHOW CREATE TABLE t1",
+                    is_err: false,
+                },
+                Test {
+                    name: "table-drop-access-passed",
+                    query: "DROP TABLE t1",
+                    is_err: false,
+                },
+            ],
+        },
+        TestGroup {
+            name: "denied",
+            tests: vec![
+                Test {
+                    name: "table-create-access-passed",
+                    query: "CREATE TABLE t1(a int)",
+                    is_err: false,
+                },
+                Test {
+                    name: "insert-denied",
+                    query: "insert into t1 values(1)",
+                    is_err: true,
+                },
+                Test {
+                    name: "select-denied",
+                    query: "SELECT * FROM t1",
+                    is_err: true,
+                },
+                Test {
+                    name: "show-processlist-denied",
+                    query: "SHOW PROCESSLIST",
+                    is_err: true,
+                },
+                Test {
+                    name: "show-metrics-denied",
+                    query: "SHOW METRICS",
+                    is_err: true,
+                },
+            ],
+        },
+    ];
+
+    let conf = crate::tests::ConfigBuilder::create()
+        .with_management_mode()
+        .config();
+    let ctx = crate::tests::create_query_context_with_config(conf.clone(), None).await?;
+
+    for group in groups {
+        for test in group.tests {
+            let plan = PlanParser::parse(ctx.clone(), test.query).await?;
+            let interpreter = InterpreterFactory::get(ctx.clone(), plan)?;
+            let res = interpreter.execute(None).await;
+            assert_eq!(
+                test.is_err,
+                res.is_err(),
+                "in test case:{:?}",
+                (group.name, test.name)
+            );
+        }
+    }
+
+    Ok(())
+}

query/tests/it/interpreters/access/mod.rs

@@ -0,0 +1,15 @@
+// Copyright 2022 Datafuse Labs.
+//
+// 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.
+
+mod management_mode_access;

query/tests/it/interpreters/mod.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod access;
 mod interpreter_call;
 mod interpreter_database_create;
 mod interpreter_database_drop;