From 59671e116f9f78406a5dc7742b5496b090bc6049 Mon Sep 17 00:00:00 2001
From: huerni <47264950+huerni@users.noreply.github.com>
Date: Tue, 17 Dec 2024 13:41:19 +0800
Subject: [PATCH] feat: When submitting a task, check whether the account is in
 the partition's AllowAccounts.

---
 src/CraneCtld/CranedMetaContainer.cpp | 12 ++++++++++++
 src/CraneCtld/CranedMetaContainer.h   |  3 +++
 src/CraneCtld/CtldGrpcServer.cpp      |  7 +++++++
 3 files changed, 22 insertions(+)

diff --git a/src/CraneCtld/CranedMetaContainer.cpp b/src/CraneCtld/CranedMetaContainer.cpp
index 1196e54ce..504c64dbb 100644
--- a/src/CraneCtld/CranedMetaContainer.cpp
+++ b/src/CraneCtld/CranedMetaContainer.cpp
@@ -600,6 +600,18 @@ CraneErrCodeExpected<void> CranedMetaContainer::ModifyPartitionAllowAccounts(
   return result;
 }
 
+bool CranedMetaContainer::CheckIfAccountIsAllowedInPartition(
+    const std::string& partition_name, const std::string& account_name) {
+  if (!partition_metas_map_.Contains(partition_name)) return false;
+
+  auto part_meta = partition_metas_map_.GetValueExclusivePtr(partition_name);
+
+  if (!part_meta->partition_global_meta.allow_accounts.contains(account_name))
+    return false;
+
+  return true;
+}
+
 void CranedMetaContainer::AddDedicatedResource(
     const CranedId& node_id, const DedicatedResourceInNode& resource) {
   if (!craned_meta_map_.Contains(node_id)) {
diff --git a/src/CraneCtld/CranedMetaContainer.h b/src/CraneCtld/CranedMetaContainer.h
index 090552375..5fccfd9c7 100644
--- a/src/CraneCtld/CranedMetaContainer.h
+++ b/src/CraneCtld/CranedMetaContainer.h
@@ -94,6 +94,9 @@ class CranedMetaContainer final {
       const std::string& partition_name,
       const std::unordered_set<std::string>& allow_accounts);
 
+  bool CheckIfAccountIsAllowedInPartition(const std::string& partition_name,
+                                          const std::string& account_name);
+
   void CranedUp(const CranedId& craned_id);
 
   void CranedDown(const CranedId& craned_id);
diff --git a/src/CraneCtld/CtldGrpcServer.cpp b/src/CraneCtld/CtldGrpcServer.cpp
index 191fa96f0..1b2d0965d 100644
--- a/src/CraneCtld/CtldGrpcServer.cpp
+++ b/src/CraneCtld/CtldGrpcServer.cpp
@@ -1013,6 +1013,13 @@ CtldServer::SubmitTaskToScheduler(std::unique_ptr<TaskInCtld> task) {
     return std::unexpected(enable_res.error());
   }
 
+  if (!g_meta_container->CheckIfAccountIsAllowedInPartition(task->partition_id,
+                                                            task->account))
+    return std::unexpected(
+        "The account is not in the AllowAccounts of the partition "
+        "specified "
+        "for the task, submission of the task is prohibited.");
+
   err = g_task_scheduler->AcquireTaskAttributes(task.get());
 
   if (err == CraneErr::kOk)