Dev/partition allow accounts

etc/config.yaml

@@ -98,6 +98,8 @@ Partitions:
     DefaultMemPerCpu: 0
     # Optional maximum memory per cpu in MB, 0 indicates no limit
     MaxMemPerCpu: 0
+    AllowedAccounts: ac1, ac2
+    DeniedAccounts: ac1, ac2
 
 DefaultPartition: CPU
 

protos/Crane.proto

@@ -386,6 +386,19 @@ message BlockAccountOrUserReply {
   ErrCode reason = 2;
 }
 
+message ModifyPartitionAllowedOrDeniedAccountsRequest {
+  uint32 uid = 1;
+  string partition_name = 2;
+  bool is_modify_allowed = 3;
+  repeated string accounts = 4;
+}
+
+
+message ModifyPartitionAllowedOrDeniedAccountsReply {
+  bool ok = 1;
+  ErrCode err_code = 2;
+}
+
 message MigrateSshProcToCgroupRequest {
   int32 pid = 1;
   uint32 task_id = 2;
@@ -768,6 +781,7 @@ service CraneCtld {
   rpc QueryPartitionInfo(QueryPartitionInfoRequest) returns (QueryPartitionInfoReply);
   rpc ModifyTask(ModifyTaskRequest) returns (ModifyTaskReply);
   rpc ModifyNode(ModifyCranedStateRequest) returns (ModifyCranedStateReply);
+  rpc ModifyPartitionAllowedOrDeniedAccounts(ModifyPartitionAllowedOrDeniedAccountsRequest) returns (ModifyPartitionAllowedOrDeniedAccountsReply);
 
   /* RPCs called from cacctmgr */
   rpc AddAccount(AddAccountRequest) returns (AddAccountReply);

protos/PublicDefs.proto

@@ -285,6 +285,9 @@ message PartitionInfo {
   ResourceView res_total = 6;
   ResourceView res_avail = 7;
   ResourceView res_alloc = 8;
+
+  repeated string allowed_accounts = 9;
+  repeated string denied_accounts = 10;
 }
 
 message CranedInfo {

src/CraneCtld/AccountManager.cpp

@@ -976,8 +976,7 @@ std::expected<void, std::string> AccountManager::CheckAndApplyQosLimitOnTask(
   return {};
 }
 
-std::expected<void, std::string> AccountManager::CheckUidIsAdmin(
-    uint32_t uid) {
+std::expected<void, std::string> AccountManager::CheckUidIsAdmin(uint32_t uid) {
   util::read_lock_guard user_guard(m_rw_user_mutex_);
   auto user_result = GetUserInfoByUidNoLock_(uid);
   if (!user_result) {
@@ -1004,6 +1003,30 @@ AccountManager::CraneExpected<void> AccountManager::CheckIfUidHasPermOnUser(
   return CheckIfUserHasPermOnUserNoLock_(*op_user, user, read_only_priv);
 }
 
+  CraneErrCodeExpected<void> AccountManager::CheckModifyPartitionAllowedOrDeniedAccounts(
+        uint32_t uid, const std::string& partition_name,
+        const std::unordered_set<std::string>& accounts) {
+  CraneErrCodeExpected<void> result{};
+
+  util::read_lock_guard user_guard(m_rw_user_mutex_);
+  util::read_lock_guard account_guard(m_rw_account_mutex_);
+
+  auto user_result = GetUserInfoByUidNoLock_(uid);
+  if (!user_result) return std::unexpected(user_result.error());
+  const User* op_user = user_result.value();
+  result = CheckIfUserHasHigherPrivThan_(*op_user, User::None);
+  if (!result) return result;
+
+  for (const auto& account_name : accounts) {
+    const Account* account = GetAccountInfoNoLock_(account_name);
+    result =
+        CheckPartitionIsAllowedNoLock_(account, partition_name, false, false);
+    if (!result) return std::unexpected(result.error());
+  }
+
+  return result;
+}
+
 AccountManager::CraneExpected<void>
 AccountManager::CheckAddUserAllowedPartitionNoLock_(
     const User* user, const Account* account, const std::string& partition) {

src/CraneCtld/AccountManager.h

@@ -40,8 +40,6 @@ class AccountManager {
   using QosMapMutexSharedPtr = util::ScopeConstSharedPtr<
       std::unordered_map<std::string, std::unique_ptr<Qos>>, util::rw_mutex>;
 
-  using CraneErrCode = crane::grpc::ErrCode;
-
   template <typename T>
   using CraneExpected = std::expected<T, CraneErrCode>;
 
@@ -141,6 +139,10 @@ class AccountManager {
                                               const std::string& username,
                                               bool read_only_priv);
 
+  CraneErrCodeExpected<void> CheckModifyPartitionAllowedOrDeniedAccounts(
+      uint32_t uid, const std::string& partition_name,
+      const std::unordered_set<std::string>& accounts);
+
  private:
   void InitDataMap_();
 

src/CraneCtld/CraneCtld.cpp

@@ -458,6 +458,25 @@ void ParseConfig(int argc, char** argv) {
             }
           }
 
+          if (partition["AllowedAccounts"] &&
+              !partition["AllowedAccounts"].IsNull()) {
+            auto allowed_accounts_str =
+                partition["AllowedAccounts"].as<std::string>();
+            std::vector<std::string> allowed_accounts =
+                absl::StrSplit(absl::StripAsciiWhitespace(allowed_accounts_str).data(), ",");
+            for (const auto& account_name : allowed_accounts) {
+              part.allowed_accounts.insert(account_name);
+            }
+          }
+
+          if (partition["DeniedAccounts"] && !partition["DeniedAccounts"].IsNull()) {
+            auto denied_accounts_str = partition["DeniedAccounts"].as<std::string>();
+            std::vector<std::string> denied_accounts = absl::StrSplit(absl::StripAsciiWhitespace(denied_accounts_str).data(), ",");
+            for (const auto& account_name : denied_accounts) {
+              part.denied_accounts.insert(account_name);
+            }
+          }
+
           if (partition["DefaultMemPerCpu"] &&
               !partition["DefaultMemPerCpu"].IsNull()) {
             part.default_mem_per_cpu =

src/CraneCtld/CranedMetaContainer.cpp

@@ -283,9 +283,12 @@ void CranedMetaContainer::InitFromConfig(const Config& config) {
     part_meta.partition_global_meta.res_total_inc_dead = part_res;
     part_meta.partition_global_meta.node_cnt = part_meta.craned_ids.size();
     part_meta.partition_global_meta.nodelist_str = partition.nodelist_str;
+    part_meta.partition_global_meta.allowed_accounts = partition.allowed_accounts;
+    part_meta.partition_global_meta.denied_accounts = partition.denied_accounts;
 
     CRANE_DEBUG(
-        "partition [{}]'s Global resource now: (cpu: {}, mem: {}, gres: {}). "
+        "partition [{}]'s Global resource now: (cpu: {}, mem: {}, "
+        "gres: {}). "
         "It has {} craneds.",
         part_name,
         part_meta.partition_global_meta.res_total_inc_dead.CpuCount(),
@@ -346,7 +349,15 @@ CranedMetaContainer::QueryAllPartitionInfo() {
     part_info->set_total_nodes(part_meta->partition_global_meta.node_cnt);
     part_info->set_alive_nodes(
         part_meta->partition_global_meta.alive_craned_cnt);
-
+    auto* allowed_accounts = part_info->mutable_allowed_accounts();
+    for (const auto& account_name :
+         part_meta->partition_global_meta.allowed_accounts) {
+      allowed_accounts->Add()->assign(account_name);
+    }
+    auto* denied_accounts = part_info->mutable_denied_accounts();
+    for (const auto& account_name : part_meta->partition_global_meta.denied_accounts) {
+      denied_accounts->Add()->assign(account_name);
+    }
     *part_info->mutable_res_total() = static_cast<crane::grpc::ResourceView>(
         part_meta->partition_global_meta.res_total);
     *part_info->mutable_res_avail() = static_cast<crane::grpc::ResourceView>(
@@ -378,7 +389,15 @@ crane::grpc::QueryPartitionInfoReply CranedMetaContainer::QueryPartitionInfo(
   part_info->set_name(part_meta->partition_global_meta.name);
   part_info->set_total_nodes(part_meta->partition_global_meta.node_cnt);
   part_info->set_alive_nodes(part_meta->partition_global_meta.alive_craned_cnt);
-
+  auto* allowed_accounts = part_info->mutable_allowed_accounts();
+  for (const auto& account_name :
+       part_meta->partition_global_meta.allowed_accounts) {
+    allowed_accounts->Add()->assign(account_name);
+  }
+  auto* denied_accounts = part_info->mutable_denied_accounts();
+  for (const auto& account_name : part_meta->partition_global_meta.denied_accounts) {
+    denied_accounts->Add()->assign(account_name);
+  }
   if (part_meta->partition_global_meta.alive_craned_cnt > 0)
     part_info->set_state(crane::grpc::PartitionState::PARTITION_UP);
   else
@@ -573,6 +592,53 @@ crane::grpc::ModifyCranedStateReply CranedMetaContainer::ChangeNodeState(
   return reply;
 }
 
+  CraneErrCodeExpected<void> CranedMetaContainer::ModifyPartitionAllowedOrDeniedAccounts(
+        const std::string& partition_name,
+        bool is_modify_allowed,
+        const std::unordered_set<std::string>& accounts) {
+  CraneErrCodeExpected<void> result{};
+
+  auto part_meta_map = partition_metas_map_.GetMapSharedPtr();
+  if (!part_meta_map->contains(partition_name))
+    return std::unexpected(CraneErrCode::ERR_INVALID_PARTITION);
+
+  auto part_meta = part_meta_map->at(partition_name).GetExclusivePtr();
+  auto& allowed_accounts = part_meta->partition_global_meta.allowed_accounts;
+  auto& denied_accounts = part_meta->partition_global_meta.denied_accounts;
+
+  if (is_modify_allowed) {
+    allowed_accounts = accounts;
+  } else {
+    denied_accounts = accounts;
+  }
+
+  return result;
+}
+
+std::expected<void, std::string> CranedMetaContainer::CheckIfAccountIsAllowedInPartition(
+    const std::string& partition_name, const std::string& account_name) {
+  auto part_metas_map = partition_metas_map_.GetMapSharedPtr();
+
+  if (!part_metas_map->contains(partition_name)) return std::unexpected("Partition does not exist.");
+
+  auto part_meta = part_metas_map->at(partition_name).GetExclusivePtr();
+  const auto& allowed_accounts = part_meta->partition_global_meta.allowed_accounts;
+  if (!allowed_accounts.empty()) {
+    if (!allowed_accounts.contains(account_name))
+      return std::unexpected(
+        "The account is not in the AllowedAccounts of the partition "
+        "specified for the task, submission of the task is prohibited.");
+  } else {
+    const auto& denied_accounts = part_meta->partition_global_meta.denied_accounts;
+    if (denied_accounts.contains(account_name))
+      return std::unexpected(
+        "The account is in the DeniedAccounts of the partition "
+        "specified for the task, submission of the task is prohibited.");
+  }
+
+  return {};
+}
+
 void CranedMetaContainer::AddDedicatedResource(
     const CranedId& node_id, const DedicatedResourceInNode& resource) {
   if (!craned_meta_map_.Contains(node_id)) {

src/CraneCtld/CranedMetaContainer.h

@@ -90,6 +90,14 @@ class CranedMetaContainer final {
   crane::grpc::ModifyCranedStateReply ChangeNodeState(
       const crane::grpc::ModifyCranedStateRequest& request);
 
+  CraneErrCodeExpected<void> ModifyPartitionAllowedOrDeniedAccounts(
+      const std::string& partition_name,
+      bool is_modify_allowed,
+      const std::unordered_set<std::string>& accounts);
+
+  std::expected<void, std::string> CheckIfAccountIsAllowedInPartition(const std::string& partition_name,
+                                          const std::string& account_name);
+
   void CranedUp(const CranedId& craned_id);
 
   void CranedDown(const CranedId& craned_id);

src/CraneCtld/CtldGrpcServer.cpp

@@ -248,6 +248,40 @@ grpc::Status CraneCtldServiceImpl::ModifyNode(
   return grpc::Status::OK;
 }
 
+grpc::Status CraneCtldServiceImpl::ModifyPartitionAllowedOrDeniedAccounts(
+    grpc::ServerContext *context,
+    const crane::grpc::ModifyPartitionAllowedOrDeniedAccountsRequest *request,
+    crane::grpc::ModifyPartitionAllowedOrDeniedAccountsReply *response) {
+  CraneErrCodeExpected<void> result;
+
+  std::unordered_set<std::string> accounts;
+
+  for (const auto &account_name : request->accounts()) {
+    accounts.insert(account_name);
+  }
+
+  result = g_account_manager->CheckModifyPartitionAllowedOrDeniedAccounts(
+      request->uid(), request->partition_name(), accounts);
+
+  if (!result) {
+    response->set_ok(false);
+    response->set_err_code(result.error());
+    return grpc::Status::OK;
+  }
+
+  result = g_meta_container->ModifyPartitionAllowedOrDeniedAccounts(
+        request->partition_name(), request->is_modify_allowed(), accounts);
+
+  if (!result) {
+    response->set_ok(false);
+    response->set_err_code(result.error());
+  } else {
+    response->set_ok(true);
+  }
+
+  return grpc::Status::OK;
+}
+
 grpc::Status CraneCtldServiceImpl::QueryTasksInfo(
     grpc::ServerContext *context,
     const crane::grpc::QueryTasksInfoRequest *request,
@@ -379,7 +413,7 @@ grpc::Status CraneCtldServiceImpl::AddQos(
   int64_t sec = qos_info->max_time_limit_per_task();
   if (!CheckIfTimeLimitSecIsValid(sec)) {
     response->set_ok(false);
-    response->set_reason(AccountManager::CraneErrCode::ERR_TIME_LIMIT);
+    response->set_reason(CraneErrCode::ERR_TIME_LIMIT);
     return grpc::Status::OK;
   }
   qos.max_time_limit_per_task = absl::Seconds(sec);
@@ -980,6 +1014,11 @@ CtldServer::SubmitTaskToScheduler(std::unique_ptr<TaskInCtld> task) {
     return std::unexpected(enable_res.error());
   }
 
+  auto result = g_meta_container->CheckIfAccountIsAllowedInPartition(task->partition_id,
+                                                            task->account);
+  if (!result)
+    return std::unexpected(result.error());
+
   err = g_task_scheduler->AcquireTaskAttributes(task.get());
 
   if (err == CraneErr::kOk)

src/CraneCtld/CtldGrpcServer.h

@@ -227,6 +227,11 @@ class CraneCtldServiceImpl final : public crane::grpc::CraneCtld::Service {
       const crane::grpc::ModifyCranedStateRequest *request,
       crane::grpc::ModifyCranedStateReply *response) override;
 
+  grpc::Status ModifyPartitionAllowedOrDeniedAccounts(
+      grpc::ServerContext *context,
+      const crane::grpc::ModifyPartitionAllowedOrDeniedAccountsRequest *request,
+      crane::grpc::ModifyPartitionAllowedOrDeniedAccountsReply *response) override;
+
   grpc::Status AddAccount(grpc::ServerContext *context,
                           const crane::grpc::AddAccountRequest *request,
                           crane::grpc::AddAccountReply *response) override;

src/CraneCtld/CtldPublicDefs.h

@@ -80,7 +80,8 @@ struct Config {
     // optional, 0 indicates no limit
     uint64_t max_mem_per_cpu;
     std::unordered_set<std::string> nodes;
-    std::unordered_set<std::string> AllowAccounts;
+    std::unordered_set<std::string> allowed_accounts;
+    std::unordered_set<std::string> denied_accounts;
   };
 
   struct CraneCtldListenConf {
@@ -210,6 +211,10 @@ struct PartitionGlobalMeta {
 
   std::string name;
   std::string nodelist_str;
+
+  std::unordered_set<std::string> allowed_accounts;
+  std::unordered_set<std::string> denied_accounts;
+
   uint32_t node_cnt;
   uint32_t alive_craned_cnt;
 };

src/Craned/CranedPublicDefs.h

@@ -45,7 +45,7 @@ struct TaskInfoOfUid {
 
 struct Partition {
   std::unordered_set<std::string> nodes;
-  std::unordered_set<std::string> AllowAccounts;
+  std::unordered_set<std::string> AllowedAccounts;
 };
 
 struct Config {

src/Utilities/PublicHeader/include/crane/PublicHeader.h

@@ -61,9 +61,14 @@ enum class CraneErr : uint16_t {
   __ERR_SIZE  // NOLINT(bugprone-reserved-identifier)
 };
 
+using CraneErrCode = crane::grpc::ErrCode;
+
 template <typename T>
 using CraneExpected = std::expected<T, CraneErr>;
 
+template <typename T>
+using CraneErrCodeExpected = std::expected<T, CraneErrCode>;
+
 inline const char* kCtldDefaultPort = "10011";
 inline const char* kCranedDefaultPort = "10010";
 inline const char* kCforedDefaultPort = "10012";