feat: user roles (#852)

* [WIP] feat: user roles

* update

* update

* admin handler

* update

* feat: user-specific connection secret

* simplify some logics

* cleanup

* update waf

* update user api error handling

* update waf api

* fix codeql

* update waf table

* fix several problems

* add pagination for waf api

* update permission checks

* switch to runtime check

* 1

* cover?

* some changes

model/alertrule.go

@@ -62,10 +62,15 @@ func (r *AlertRule) Enabled() bool {
 }
 
 // Snapshot 对传入的Server进行该报警规则下所有type的检查 返回每项检查结果
-func (r *AlertRule) Snapshot(cycleTransferStats *CycleTransferStats, server *Server, db *gorm.DB) []bool {
-	point := make([]bool, 0, len(r.Rules))
-	for _, rule := range r.Rules {
-		point = append(point, rule.Snapshot(cycleTransferStats, server, db))
+func (r *AlertRule) Snapshot(cycleTransferStats *CycleTransferStats, server *Server, db *gorm.DB, role uint8) []bool {
+	point := make([]bool, len(r.Rules))
+
+	if r.UserID != server.UserID && role != RoleAdmin {
+		return point
+	}
+
+	for i, rule := range r.Rules {
+		point[i] = rule.Snapshot(cycleTransferStats, server, db)
 	}
 	return point
 }

model/api.go

@@ -15,6 +15,23 @@ type CommonResponse[T any] struct {
 	Error   string `json:"error,omitempty"`
 }
 
+type PaginatedResponse[S ~[]E, E any] struct {
+	Success bool      `json:"success,omitempty"`
+	Data    *Value[S] `json:"data,omitempty"`
+	Error   string    `json:"error,omitempty"`
+}
+
+type Value[T any] struct {
+	Value      T          `json:"value,omitempty"`
+	Pagination Pagination `json:"pagination,omitempty"`
+}
+
+type Pagination struct {
+	Offset int   `json:"offset,omitempty"`
+	Limit  int   `json:"limit,omitempty"`
+	Total  int64 `json:"total,omitempty"`
+}
+
 type LoginResponse struct {
 	Token  string `json:"token,omitempty"`
 	Expire string `json:"expire,omitempty"`

model/common.go

@@ -2,6 +2,8 @@ package model
 
 import (
 	"time"
+
+	"github.com/gin-gonic/gin"
 )
 
 const (
@@ -18,6 +20,47 @@ type Common struct {
 	UpdatedAt time.Time `gorm:"autoUpdateTime" json:"updated_at,omitempty"`
 	// Do not use soft deletion
 	// DeletedAt gorm.DeletedAt `gorm:"index" json:"deleted_at,omitempty"`
+
+	UserID uint64 `json:"-"`
+}
+
+func (c *Common) GetID() uint64 {
+	return c.ID
+}
+
+func (c *Common) GetUserID() uint64 {
+	return c.UserID
+}
+
+func (c *Common) HasPermission(ctx *gin.Context) bool {
+	auth, ok := ctx.Get(CtxKeyAuthorizedUser)
+	if !ok {
+		return false
+	}
+
+	user := *auth.(*User)
+	if user.Role == RoleAdmin {
+		return true
+	}
+
+	return user.ID == c.UserID
+}
+
+type CommonInterface interface {
+	GetID() uint64
+	GetUserID() uint64
+	HasPermission(*gin.Context) bool
+}
+
+func FindByUserID[S ~[]E, E CommonInterface](s S, uid uint64) []uint64 {
+	var list []uint64
+	for _, v := range s {
+		if v.GetUserID() == uid {
+			list = append(list, v.GetID())
+		}
+	}
+
+	return list
 }
 
 type Response struct {

model/user.go

@@ -1,9 +1,41 @@
 package model
 
+import (
+	"github.com/nezhahq/nezha/pkg/utils"
+	"gorm.io/gorm"
+)
+
+const (
+	RoleAdmin uint8 = iota
+	RoleMember
+)
+
 type User struct {
 	Common
-	Username string `json:"username,omitempty" gorm:"uniqueIndex"`
-	Password string `json:"password,omitempty" gorm:"type:char(72)"`
+	Username    string `json:"username,omitempty" gorm:"uniqueIndex"`
+	Password    string `json:"password,omitempty" gorm:"type:char(72)"`
+	Role        uint8  `json:"role,omitempty"`
+	AgentSecret string `json:"agent_secret,omitempty" gorm:"type:char(32)"`
+}
+
+type UserInfo struct {
+	Role        uint8
+	_           [3]byte
+	AgentSecret string
+}
+
+func (u *User) BeforeSave(tx *gorm.DB) error {
+	if u.AgentSecret != "" {
+		return nil
+	}
+
+	key, err := utils.GenerateRandomString(32)
+	if err != nil {
+		return err
+	}
+
+	u.AgentSecret = key
+	return nil
 }
 
 type Profile struct {

model/user_group.go

@@ -1,6 +0,0 @@
-package model
-
-type UserGroup struct {
-	Common
-	Name string `json:"name"`
-}

model/user_group_user.go

@@ -1,7 +0,0 @@
-package model
-
-type UserGroupUser struct {
-	Common
-	UserGroupId uint64 `json:"user_group_id"`
-	UserId      uint64 `json:"user_id"`
-}

model/waf.go

@@ -16,22 +16,30 @@ const (
 	WAFBlockReasonTypeAgentAuthFail
 )
 
+const (
+	BlockIDgRPC = -127 + iota
+	BlockIDToken
+	BlockIDUnknownUser
+)
+
 type WAFApiMock struct {
-	IP                 string `json:"ip,omitempty"`
-	Count              uint64 `json:"count,omitempty"`
-	LastBlockReason    uint8  `json:"last_block_reason,omitempty"`
-	LastBlockTimestamp uint64 `json:"last_block_timestamp,omitempty"`
+	IP              string `json:"ip,omitempty"`
+	BlockIdentifier int64  `json:"block_identifier,omitempty"`
+	BlockReason     uint8  `json:"block_reason,omitempty"`
+	BlockTimestamp  uint64 `json:"block_timestamp,omitempty"`
+	Count           uint64 `json:"count,omitempty"`
 }
 
 type WAF struct {
-	IP                 []byte `gorm:"type:binary(16);primaryKey" json:"ip,omitempty"`
-	Count              uint64 `json:"count,omitempty"`
-	LastBlockReason    uint8  `json:"last_block_reason,omitempty"`
-	LastBlockTimestamp uint64 `json:"last_block_timestamp,omitempty"`
+	IP              []byte `gorm:"type:binary(16);primaryKey" json:"ip,omitempty"`
+	BlockIdentifier int64  `gorm:"primaryKey" json:"block_identifier,omitempty"`
+	BlockReason     uint8  `json:"block_reason,omitempty"`
+	BlockTimestamp  uint64 `gorm:"index" json:"block_timestamp,omitempty"`
+	Count           uint64 `json:"count,omitempty"`
 }
 
 func (w *WAF) TableName() string {
-	return "waf"
+	return "nz_waf"
 }
 
 func CheckIP(db *gorm.DB, ip string) error {
@@ -42,22 +50,31 @@ func CheckIP(db *gorm.DB, ip string) error {
 	if err != nil {
 		return err
 	}
-	var w WAF
-	result := db.Limit(1).Find(&w, "ip = ?", ipBinary)
+
+	var blockTimestamp uint64
+	result := db.Model(&WAF{}).Order("block_timestamp desc").Select("block_timestamp").Where("ip = ?", ipBinary).Limit(1).Find(&blockTimestamp)
 	if result.Error != nil {
 		return result.Error
 	}
-	if result.RowsAffected == 0 { // 检查是否未找到记录
+
+	// 检查是否未找到记录
+	if result.RowsAffected < 1 {
 		return nil
 	}
+
+	var count uint64
+	if err := db.Model(&WAF{}).Select("SUM(count)").Where("ip = ?", ipBinary).Scan(&count).Error; err != nil {
+		return err
+	}
+
 	now := time.Now().Unix()
-	if powAdd(w.Count, 4, w.LastBlockTimestamp) > uint64(now) {
+	if powAdd(count, 4, blockTimestamp) > uint64(now) {
 		return errors.New("you are blocked by nezha WAF")
 	}
 	return nil
 }
 
-func ClearIP(db *gorm.DB, ip string) error {
+func ClearIP(db *gorm.DB, ip string, uid int64) error {
 	if ip == "" {
 		return nil
 	}
@@ -65,7 +82,7 @@ func ClearIP(db *gorm.DB, ip string) error {
 	if err != nil {
 		return err
 	}
-	return db.Unscoped().Delete(&WAF{}, "ip = ?", ipBinary).Error
+	return db.Unscoped().Delete(&WAF{}, "ip = ? and block_identifier = ?", ipBinary, uid).Error
 }
 
 func BatchClearIP(db *gorm.DB, ip []string) error {
@@ -83,7 +100,7 @@ func BatchClearIP(db *gorm.DB, ip []string) error {
 	return db.Unscoped().Delete(&WAF{}, "ip in (?)", ips).Error
 }
 
-func BlockIP(db *gorm.DB, ip string, reason uint8) error {
+func BlockIP(db *gorm.DB, ip string, reason uint8, uid int64) error {
 	if ip == "" {
 		return nil
 	}
@@ -91,16 +108,19 @@ func BlockIP(db *gorm.DB, ip string, reason uint8) error {
 	if err != nil {
 		return err
 	}
-	var w WAF
-	w.IP = ipBinary
+	w := WAF{
+		IP:              ipBinary,
+		BlockIdentifier: uid,
+	}
+	now := uint64(time.Now().Unix())
 	return db.Transaction(func(tx *gorm.DB) error {
 		if err := tx.Where(&w).Attrs(WAF{
-			LastBlockReason:    reason,
-			LastBlockTimestamp: uint64(time.Now().Unix()),
+			BlockReason:    reason,
+			BlockTimestamp: now,
 		}).FirstOrCreate(&w).Error; err != nil {
 			return err
 		}
-		return tx.Exec("UPDATE waf SET count = count + 1, last_block_reason = ?, last_block_timestamp = ? WHERE ip = ?", reason, uint64(time.Now().Unix()), ipBinary).Error
+		return tx.Exec("UPDATE nz_waf SET count = count + 1, block_reason = ?, block_timestamp = ? WHERE ip = ? and block_identifier = ?", reason, now, ipBinary, uid).Error
 	})
 }