Add permission checking for kill queries. (vesoft-inc#3896)

* Add permission checking for kill queries.

* Fix typo.

* Add new line.

* Fix typo.

* Fix typo.

Co-authored-by: Sophie <84560950+Sophie-Xie@users.noreply.github.com>

src/graph/executor/admin/KillQueryExecutor.cpp

@@ -13,8 +13,6 @@ namespace graph {
 folly::Future<Status> KillQueryExecutor::execute() {
   SCOPED_TIMER(&execTime_);
 
-  // TODO: permission check
-
   QueriesMap toBeVerifiedQueries;
   QueriesMap killQueries;
   NG_RETURN_IF_ERROR(verifyTheQueriesByLocalCache(toBeVerifiedQueries, killQueries));
@@ -92,9 +90,15 @@ Status KillQueryExecutor::verifyTheQueriesByLocalCache(QueriesMap& toBeVerifiedQ
       auto sessionPtr = sessionMgr->findSessionFromCache(sessionId);
       if (sessionPtr == nullptr) {
         toBeVerifiedQueries[sessionId].emplace(epId);
-      } else if (!sessionPtr->findQuery(epId)) {
-        return Status::Error(
-            "ExecutionPlanId[%ld] does not exist in Session[%ld].", epId, sessionId);
+      } else {  // Sessions in current host
+        // Only GOD role could kill others' queries.
+        if (sessionPtr->user() != session->user() && !session->isGod()) {
+          return Status::PermissionError("Only GOD role could kill others' queries.");
+        }
+        if (!sessionPtr->findQuery(epId)) {
+          return Status::Error(
+              "ExecutionPlanId[%ld] does not exist in Session[%ld].", epId, sessionId);
+        }
       }
       killQueries[sessionId].emplace(epId);
     }
@@ -125,6 +129,7 @@ void KillQueryExecutor::killCurrentHostQueries(const QueriesMap& killQueries,
 
 Status KillQueryExecutor::verifyTheQueriesByMetaInfo(
     const QueriesMap& toBeVerifiedQueries, const std::vector<meta::cpp2::Session>& sessionsInMeta) {
+  auto* currentSession = qctx()->rctx()->session();
   for (auto& s : toBeVerifiedQueries) {
     auto sessionId = s.first;
     auto found = std::find_if(sessionsInMeta.begin(), sessionsInMeta.end(), [sessionId](auto& val) {
@@ -139,6 +144,10 @@ Status KillQueryExecutor::verifyTheQueriesByMetaInfo(
             "ExecutionPlanId[%ld] does not exist in Session[%ld].", epId, sessionId);
       }
     }
+    // Only GOD role could kill others' queries.
+    if (found->get_user_name() != currentSession->user() && !currentSession->isGod()) {
+      return Status::PermissionError("Only GOD role could kill others' queries.");
+    }
   }
   return Status::OK();
 }

src/graph/service/PermissionCheck.cpp

@@ -212,8 +212,10 @@ Status PermissionCheck::permissionCheck(ClientSession *session,
       // No permission checking for sequential sentence.
       return Status::OK();
     }
-    case Sentence::Kind::kShowQueries:
-    case Sentence::Kind::kKillQuery: {
+    case Sentence::Kind::kKillQuery:
+      // Only GOD could kill all queries, other roles only could kill own queries.
+      return Status::OK();
+    case Sentence::Kind::kShowQueries: {
       return Status::OK();
     }
   }

tests/Makefile

@@ -105,7 +105,9 @@ test: sess
 
 slow-query: currdir
 	$(test_j) tck/steps/test_kill_slow_query_via_same_service.py && \
-	$(test_j) tck/steps/test_kill_slow_query_via_different_service.py
+	$(test_j) tck/steps/test_kill_slow_query_via_different_service.py && \
+	$(test_j) tck/steps/test_kill_permission_via_same_service.py && \
+	$(test_j) tck/steps/test_kill_permission_via_different_service.py
 
 standalone-tck: jobs
 	$(test_j_sa) tck/steps/test_tck.py

tests/tck/slowquery/permissionViaDifferentService.feature

@@ -0,0 +1,44 @@
+# Copyright (c) 2022 vesoft inc. All rights reserved.
+#
+# This source code is licensed under Apache 2.0 License.
+Feature: Test kill queries permission from different services
+
+  Scenario: Setup slow query in service 1
+    # Set up a slow query which will be killed later.
+    Given a graph with space named "nba"
+    When executing query via graph 1:
+      """
+      USE nba;
+      GO 100000 STEPS FROM "Tim Duncan" OVER like YIELD like._dst
+      """
+    Then an ExecutionError should be raised at runtime: Execution had been killed
+
+  Scenario: Test permisson of kill queries from service 0
+    Given a graph with space named "nba"
+    When executing query:
+      """
+      CREATE USER IF NOT EXISTS test_permission WITH PASSWORD 'test';
+      GRANT ROLE USER ON nba TO test_permission;
+      """
+    Then the execution should be successful
+    And wait 3 seconds
+    When executing query with user test_permission with password test:
+      """
+      USE nba;
+      SHOW QUERIES
+      | YIELD $-.SessionID AS sid, $-.ExecutionPlanID AS eid, $-.DurationInUSec AS dur
+      WHERE $-.DurationInUSec > 1000000 AND $-.`Query` CONTAINS "GO"
+      | ORDER BY $-.dur
+      | KILL QUERY(session=$-.sid, plan=$-.eid)
+      """
+    Then an PermissionError should be raised at runtime: Only GOD role could kill others' queries.
+    When executing query with user root with password nebula:
+      """
+      USE nba;
+      SHOW QUERIES
+      | YIELD $-.SessionID AS sid, $-.ExecutionPlanID AS eid, $-.DurationInUSec AS dur
+      WHERE $-.DurationInUSec > 1000000 AND $-.`Query` CONTAINS "GO"
+      | ORDER BY $-.dur
+      | KILL QUERY(session=$-.sid, plan=$-.eid)
+      """
+    Then the execution should be successful

tests/tck/slowquery/permissionViaSameService.feature

@@ -0,0 +1,43 @@
+# Copyright (c) 2022 vesoft inc. All rights reserved.
+#
+# This source code is licensed under Apache 2.0 License.
+Feature: Test kill queries from same service
+
+  Scenario: Setup slow query
+    # Set up a slow query which will be killed later.
+    Given a graph with space named "nba"
+    When executing query:
+      """
+      GO 100000 STEPS FROM "Tim Duncan" OVER like YIELD like._dst
+      """
+    Then an ExecutionError should be raised at runtime: Execution had been killed
+
+  Scenario: Test permisson of kill queries
+    Given a graph with space named "nba"
+    When executing query:
+      """
+      CREATE USER IF NOT EXISTS test_permission WITH PASSWORD 'test';
+      GRANT ROLE USER ON nba TO test_permission;
+      """
+    Then the execution should be successful
+    And wait 3 seconds
+    When executing query with user test_permission with password test:
+      """
+      USE nba;
+      SHOW QUERIES
+      | YIELD $-.SessionID AS sid, $-.ExecutionPlanID AS eid, $-.DurationInUSec AS dur
+      WHERE $-.DurationInUSec > 1000000 AND $-.`Query` CONTAINS "GO"
+      | ORDER BY $-.dur
+      | KILL QUERY(session=$-.sid, plan=$-.eid)
+      """
+    Then an PermissionError should be raised at runtime: Only GOD role could kill others' queries.
+    When executing query with user root with password nebula:
+      """
+      USE nba;
+      SHOW QUERIES
+      | YIELD $-.SessionID AS sid, $-.ExecutionPlanID AS eid, $-.DurationInUSec AS dur
+      WHERE $-.DurationInUSec > 1000000 AND $-.`Query` CONTAINS "GO"
+      | ORDER BY $-.dur
+      | KILL QUERY(session=$-.sid, plan=$-.eid)
+      """
+    Then the execution should be successful

tests/tck/steps/test_kill_permission_via_different_service.py

@@ -0,0 +1,7 @@
+# Copyright (c) 2022 vesoft inc. All rights reserved.
+#
+# This source code is licensed under Apache 2.0 License.
+
+from pytest_bdd import scenarios
+
+scenarios('slowquery/permissionViaDifferentService.feature')

tests/tck/steps/test_kill_permission_via_same_service.py

@@ -0,0 +1,7 @@
+# Copyright (c) 2022 vesoft inc. All rights reserved.
+#
+# This source code is licensed under Apache 2.0 License.
+
+from pytest_bdd import scenarios
+
+scenarios('slowquery/permissionViaSameService.feature')