diff --git a/etc/config.yaml b/etc/config.yaml index 4a1a8d6ec..7a868588b 100644 --- a/etc/config.yaml +++ b/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 diff --git a/protos/Crane.proto b/protos/Crane.proto index be2f822f0..bb8ff1586 100644 --- a/protos/Crane.proto +++ b/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); diff --git a/protos/PublicDefs.proto b/protos/PublicDefs.proto index a635191eb..89259e4be 100644 --- a/protos/PublicDefs.proto +++ b/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 { diff --git a/src/CraneCtld/AccountManager.cpp b/src/CraneCtld/AccountManager.cpp index fd4e8a239..2cab2258b 100644 --- a/src/CraneCtld/AccountManager.cpp +++ b/src/CraneCtld/AccountManager.cpp @@ -976,8 +976,7 @@ std::expected AccountManager::CheckAndApplyQosLimitOnTask( return {}; } -std::expected AccountManager::CheckUidIsAdmin( - uint32_t uid) { +std::expected 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 AccountManager::CheckIfUidHasPermOnUser( return CheckIfUserHasPermOnUserNoLock_(*op_user, user, read_only_priv); } + CraneErrCodeExpected AccountManager::CheckModifyPartitionAllowedOrDeniedAccounts( + uint32_t uid, const std::string& partition_name, + const std::unordered_set& accounts) { + CraneErrCodeExpected 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 AccountManager::CheckAddUserAllowedPartitionNoLock_( const User* user, const Account* account, const std::string& partition) { diff --git a/src/CraneCtld/AccountManager.h b/src/CraneCtld/AccountManager.h index 812acc5a5..d1d1ea392 100644 --- a/src/CraneCtld/AccountManager.h +++ b/src/CraneCtld/AccountManager.h @@ -40,8 +40,6 @@ class AccountManager { using QosMapMutexSharedPtr = util::ScopeConstSharedPtr< std::unordered_map>, util::rw_mutex>; - using CraneErrCode = crane::grpc::ErrCode; - template using CraneExpected = std::expected; @@ -141,6 +139,10 @@ class AccountManager { const std::string& username, bool read_only_priv); + CraneErrCodeExpected CheckModifyPartitionAllowedOrDeniedAccounts( + uint32_t uid, const std::string& partition_name, + const std::unordered_set& accounts); + private: void InitDataMap_(); diff --git a/src/CraneCtld/CraneCtld.cpp b/src/CraneCtld/CraneCtld.cpp index 480ca8b95..ebf296bc6 100644 --- a/src/CraneCtld/CraneCtld.cpp +++ b/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::vector 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::vector 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 = diff --git a/src/CraneCtld/CranedMetaContainer.cpp b/src/CraneCtld/CranedMetaContainer.cpp index c5b78d125..e2eb9bad4 100644 --- a/src/CraneCtld/CranedMetaContainer.cpp +++ b/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( part_meta->partition_global_meta.res_total); *part_info->mutable_res_avail() = static_cast( @@ -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 CranedMetaContainer::ModifyPartitionAllowedOrDeniedAccounts( + const std::string& partition_name, + bool is_modify_allowed, + const std::unordered_set& accounts) { + CraneErrCodeExpected 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 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)) { diff --git a/src/CraneCtld/CranedMetaContainer.h b/src/CraneCtld/CranedMetaContainer.h index 38d995d24..49456d185 100644 --- a/src/CraneCtld/CranedMetaContainer.h +++ b/src/CraneCtld/CranedMetaContainer.h @@ -90,6 +90,14 @@ class CranedMetaContainer final { crane::grpc::ModifyCranedStateReply ChangeNodeState( const crane::grpc::ModifyCranedStateRequest& request); + CraneErrCodeExpected ModifyPartitionAllowedOrDeniedAccounts( + const std::string& partition_name, + bool is_modify_allowed, + const std::unordered_set& accounts); + + std::expected 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 bc00b3152..d0126332d 100644 --- a/src/CraneCtld/CtldGrpcServer.cpp +++ b/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 result; + + std::unordered_set 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 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) diff --git a/src/CraneCtld/CtldGrpcServer.h b/src/CraneCtld/CtldGrpcServer.h index d35493fec..d9ea9404c 100644 --- a/src/CraneCtld/CtldGrpcServer.h +++ b/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; diff --git a/src/CraneCtld/CtldPublicDefs.h b/src/CraneCtld/CtldPublicDefs.h index 1cd237ed9..bd440474b 100644 --- a/src/CraneCtld/CtldPublicDefs.h +++ b/src/CraneCtld/CtldPublicDefs.h @@ -80,7 +80,8 @@ struct Config { // optional, 0 indicates no limit uint64_t max_mem_per_cpu; std::unordered_set nodes; - std::unordered_set AllowAccounts; + std::unordered_set allowed_accounts; + std::unordered_set denied_accounts; }; struct CraneCtldListenConf { @@ -210,6 +211,10 @@ struct PartitionGlobalMeta { std::string name; std::string nodelist_str; + + std::unordered_set allowed_accounts; + std::unordered_set denied_accounts; + uint32_t node_cnt; uint32_t alive_craned_cnt; }; diff --git a/src/Craned/CranedPublicDefs.h b/src/Craned/CranedPublicDefs.h index 3674aff39..bba01f36b 100644 --- a/src/Craned/CranedPublicDefs.h +++ b/src/Craned/CranedPublicDefs.h @@ -45,7 +45,7 @@ struct TaskInfoOfUid { struct Partition { std::unordered_set nodes; - std::unordered_set AllowAccounts; + std::unordered_set AllowedAccounts; }; struct Config { diff --git a/src/Utilities/PublicHeader/include/crane/PublicHeader.h b/src/Utilities/PublicHeader/include/crane/PublicHeader.h index fbb080ada..21462a75a 100644 --- a/src/Utilities/PublicHeader/include/crane/PublicHeader.h +++ b/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 using CraneExpected = std::expected; +template +using CraneErrCodeExpected = std::expected; + inline const char* kCtldDefaultPort = "10011"; inline const char* kCranedDefaultPort = "10010"; inline const char* kCforedDefaultPort = "10012";