From 39e856b7d1148dcbcd8a311e8f690efcbfc9d6f5 Mon Sep 17 00:00:00 2001
From: slhmy <slhmy.zzy@gmail.com>
Date: Sat, 27 Jul 2024 19:23:13 +0800
Subject: [PATCH] Support grant user role

---
 cmd/init/casbin.go    |  1 -
 cmd/init/main.go      |  7 +++++++
 services/user/user.go | 25 +++++++++++++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/cmd/init/casbin.go b/cmd/init/casbin.go
index 6823da5..2973729 100644
--- a/cmd/init/casbin.go
+++ b/cmd/init/casbin.go
@@ -9,7 +9,6 @@ func loadCasbinPolicies() {
 	enforcer := casbin_agent.GetDefaultCasbinEnforcer()
 
 	_, err := enforcer.AddGroupingPolicies([][]string{
-		{`user:root`, `role:super`, `system`},
 		{`role:super`, `role:admin`, `system`},
 	})
 	if err != nil {
diff --git a/cmd/init/main.go b/cmd/init/main.go
index 4b73b53..8dd852e 100644
--- a/cmd/init/main.go
+++ b/cmd/init/main.go
@@ -2,12 +2,19 @@ package main
 
 import (
 	"context"
+
+	user_service "github.com/oj-lab/oj-lab-platform/services/user"
 )
 
 func main() {
 	ctx := context.Background()
 	initDB()
 	loadCasbinPolicies()
+	err := user_service.GrantUserRole(ctx, "root", "super", "system")
+	if err != nil {
+		panic(err)
+	}
+
 	loadProblemPackages(ctx)
 	println("init success")
 }
diff --git a/services/user/user.go b/services/user/user.go
index 47754c7..44fc9d0 100644
--- a/services/user/user.go
+++ b/services/user/user.go
@@ -2,8 +2,10 @@ package user_service
 
 import (
 	"context"
+	"fmt"
 
 	user_model "github.com/oj-lab/oj-lab-platform/models/user"
+	casbin_agent "github.com/oj-lab/oj-lab-platform/modules/agent/casbin"
 	gorm_agent "github.com/oj-lab/oj-lab-platform/modules/agent/gorm"
 	auth_module "github.com/oj-lab/oj-lab-platform/modules/auth"
 	log_module "github.com/oj-lab/oj-lab-platform/modules/log"
@@ -49,6 +51,29 @@ func UpdateUser(ctx context.Context, user user_model.User) error {
 		auth_module.LoginSessionData{})
 }
 
+func GrantUserRole(ctx context.Context, account, role, domain string) error {
+	exist, err := CheckUserExist(ctx, account)
+	if err != nil {
+		return err
+	}
+	if !exist {
+		return fmt.Errorf("user not exist")
+	}
+
+	enforcer := casbin_agent.GetDefaultCasbinEnforcer()
+	account = casbin_agent.UserSubjectPrefix + account
+	role = casbin_agent.RoleSubjectPrefix + role
+	notDuplicated, err := enforcer.AddRoleForUserInDomain(account, role, domain)
+	if err != nil {
+		return err
+	}
+	if !notDuplicated {
+		return fmt.Errorf("role already granted")
+	}
+
+	return nil
+}
+
 func CheckUserExist(ctx context.Context, account string) (bool, error) {
 	getOptions := user_model.GetUserOptions{
 		AccountQuery: account,