myxce/src/frontend/apiserver/store/stores.go

package store

import "apiserver/models"

// Global store instances.
var (
	Jobs              *DataStore[uint64, models.JobInfo]
	ArrayJobs         *DataStore[uint64, models.ArrayJobInfo]
	Queues            *DataStore[string, models.QueueInfo]
	Hosts             *DataStore[string, models.HostInfo]
	Resources         *DataStore[string, models.ResourceDefinitionInfo]
	ResourceLocations *DataStore[string, models.ResourceDistributionInfo]
	Clusters          *DataStore[string, models.ClusterInfo]
	JobExecHosts      *DataStore[int64, models.JobExecHosts]
	Roles             *DataStore[string, RoleEntry]
	UserRoles         *DataStore[string, []string] // user_name → role_names
)

// RoleEntry holds a role's permissions.
type RoleEntry struct {
	RoleName    string
	Permissions []RolePermission
}

// RolePermission holds a single permission entry for a role.
type RolePermission struct {
	Action       string
	DomainType   string
	DomainValues []string
}

// InitStores creates all global store instances.
func InitStores() {
	Jobs = NewDataStore[uint64, models.JobInfo]()
	ArrayJobs = NewDataStore[uint64, models.ArrayJobInfo]()
	Queues = NewDataStore[string, models.QueueInfo]()
	Hosts = NewDataStore[string, models.HostInfo]()
	Resources = NewDataStore[string, models.ResourceDefinitionInfo]()
	ResourceLocations = NewDataStore[string, models.ResourceDistributionInfo]()
	Clusters = NewDataStore[string, models.ClusterInfo]()
	JobExecHosts = NewDataStore[int64, models.JobExecHosts]()
	Roles = NewDataStore[string, RoleEntry]()
	UserRoles = NewDataStore[string, []string]()
}

// Helper functions for RBAC.

func GetUserRoleNames(username string) []string {
	roles, ok := UserRoles.Get(username)
	if !ok {
		return []string{}
	}
	return roles
}

func GetRolePermissions(roleName string) []RolePermission {
	role, ok := Roles.Get(roleName)
	if !ok {
		return []RolePermission{}
	}
	return role.Permissions
}

func AddRoleForUser(username string, roleName string) {
	roles, _ := UserRoles.Get(username)
	for _, r := range roles {
		if r == roleName {
			return // already has this role
		}
	}
	roles = append(roles, roleName)
	UserRoles.Put(username, roles)
}

func DeleteRoleForUser(username string, roleName string) {
	roles, ok := UserRoles.Get(username)
	if !ok {
		return
	}
	newRoles := make([]string, 0, len(roles))
	for _, r := range roles {
		if r != roleName {
			newRoles = append(newRoles, r)
		}
	}
	if len(newRoles) == 0 {
		UserRoles.Delete(username)
	} else {
		UserRoles.Put(username, newRoles)
	}
}

func AddPermissionsForRole(roleName string, perms []RolePermission) {
	role, ok := Roles.Get(roleName)
	if !ok {
		role = RoleEntry{RoleName: roleName}
	}
	role.Permissions = append(role.Permissions, perms...)
	Roles.Put(roleName, role)
}

// GetAllQueueNames returns all queue names in the store.
func GetAllQueueNames() []string {
	return Queues.Keys()
}

// GetAllHostNames returns all host names in the store.
func GetAllHostNames() []string {
	return Hosts.Keys()
}

// GetAllRoleNames returns all role names.
func GetAllRoleNames(includeBuiltin bool) []string {
	roleMap := make(map[string]bool)
	for _, k := range Roles.Keys() {
		roleMap[k] = true
	}
	if includeBuiltin {
		for _, name := range models.BUILTIN_ROLE_NAMES {
			roleMap[name] = true
		}
	}
	names := make([]string, 0, len(roleMap))
	for k := range roleMap {
		names = append(names, k)
	}
	return names
}

// GetUsersUnderRoleName returns all user names that have the given role.
func GetUsersUnderRoleName(roleName string) []string {
	var users []string
	UserRoles.Range(func(userName string, roles []string) bool {
		for _, r := range roles {
			if r == roleName {
				users = append(users, userName)
				break
			}
		}
		return true
	})
	return users
}

// GetRoleAllPermissions returns permission action names for a role.
func GetRoleAllPermissions(roleName string) []string {
	role, ok := Roles.Get(roleName)
	if !ok {
		return []string{}
	}
	perms := make([]string, 0, len(role.Permissions))
	for _, p := range role.Permissions {
		perms = append(perms, p.Action)
	}
	return perms
}

// AddRoleNames ensures a role entry exists (creates empty if not).
func AddRoleNames(roleName string) {
	if !Roles.Has(roleName) {
		Roles.Put(roleName, RoleEntry{RoleName: roleName})
	}
}

// IsQueueExist checks if all queue names exist in the store.
func IsQueueExist(queueNames []string) bool {
	if len(queueNames) == 0 {
		return false
	}
	for _, q := range queueNames {
		if !Queues.Has(q) {
			return false
		}
	}
	return true
}