Merge 9598657a81 into 50ee62172f

cmd/dashboard/controller/jwt.go

@@ -88,18 +88,21 @@ func authenticator() func(c *gin.Context) (interface{}, error) {
 		}
 
 		var user model.User
+		realip := c.GetString(model.CtxKeyRealIPStr)
 		if err := singleton.DB.Select("id", "password").Where("username = ?", loginVals.Username).First(&user).Error; err != nil {
 			if err == gorm.ErrRecordNotFound {
-				model.BlockIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr), model.WAFBlockReasonTypeLoginFail)
+				model.BlockIP(singleton.DB, realip, model.WAFBlockReasonTypeLoginFail, model.BlockIDUnknownUser)
 			}
 			return nil, jwt.ErrFailedAuthentication
 		}
 
 		if err := bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(loginVals.Password)); err != nil {
-			model.BlockIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr), model.WAFBlockReasonTypeLoginFail)
+			model.BlockIP(singleton.DB, realip, model.WAFBlockReasonTypeLoginFail, int64(user.ID))
 			return nil, jwt.ErrFailedAuthentication
 		}
 
+		model.ClearIP(singleton.DB, realip, model.BlockIDUnknownUser)
+		model.ClearIP(singleton.DB, realip, int64(user.ID))
 		return utils.Itoa(user.ID), nil
 	}
 }
@@ -169,10 +172,10 @@ func optionalAuthMiddleware(mw *jwt.GinJWTMiddleware) func(c *gin.Context) {
 		identity := mw.IdentityHandler(c)
 
 		if identity != nil {
-			model.ClearIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr))
+			model.ClearIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr), model.BlockIDToken)
 			c.Set(mw.IdentityKey, identity)
 		} else {
-			if err := model.BlockIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr), model.WAFBlockReasonTypeBruteForceToken); err != nil {
+			if err := model.BlockIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr), model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken); err != nil {
 				waf.ShowBlockPage(c, err)
 				return
 			}

cmd/dashboard/controller/user.go

@@ -151,6 +151,6 @@ func batchDeleteUser(c *gin.Context) (any, error) {
 		return nil, singleton.Localizer.ErrorT("can't delete yourself")
 	}
 
-	singleton.OnUserDelete(ids)
-	return nil, singleton.DB.Where("id IN (?)", ids).Delete(&model.User{}).Error
+	err := singleton.OnUserDelete(ids, newGormError)
+	return nil, err
 }

cmd/dashboard/controller/waf.go

@@ -1,6 +1,8 @@
 package controller
 
 import (
+	"strconv"
+
 	"github.com/gin-gonic/gin"
 
 	"github.com/nezhahq/nezha/model"
@@ -13,12 +15,24 @@ import (
 // @Schemes
 // @Description List server
 // @Tags auth required
+// @Param limit query uint false "Page limit"
+// @Param offset query uint false "Page offset"
 // @Produce json
 // @Success 200 {object} model.CommonResponse[[]model.WAFApiMock]
 // @Router /waf [get]
 func listBlockedAddress(c *gin.Context) ([]*model.WAF, error) {
+	limit, err := strconv.Atoi(c.Query("limit"))
+	if err != nil || limit < 1 {
+		limit = 25
+	}
+
+	offset, err := strconv.Atoi(c.Query("offset"))
+	if err != nil || offset < 1 {
+		offset = 1
+	}
+
 	var waf []*model.WAF
-	if err := singleton.DB.Find(&waf).Error; err != nil {
+	if err := singleton.DB.Limit(limit).Offset(offset).Find(&waf).Error; err != nil {
 		return nil, err
 	}
 

cmd/dashboard/controller/ws.go

@@ -107,6 +107,8 @@ func serverStream(c *gin.Context) (any, error) {
 		return nil, newWsError("%v", err)
 	}
 	defer conn.Close()
+	singleton.OnlineUsers.Add(1)
+	defer singleton.OnlineUsers.Add(^uint64(0))
 	count := 0
 	for {
 		stat, err := getServerStat(c, count == 0)
@@ -164,6 +166,7 @@ func getServerStat(c *gin.Context, withPublicNote bool) ([]byte, error) {
 
 		return utils.Json.Marshal(model.StreamServerData{
 			Now:     time.Now().Unix() * 1000,
+			Online:  singleton.OnlineUsers.Load(),
 			Servers: servers,
 		})
 	})

cmd/dashboard/controller/ws_test.go

@@ -0,0 +1,26 @@
+package controller
+
+import (
+	"sync/atomic"
+	"testing"
+)
+
+func TestWs(t *testing.T) {
+	onlineUsers := new(atomic.Uint64)
+	onlineUsers.Add(1)
+	if onlineUsers.Load() != 1 {
+		t.Error("onlineUsers.Add(1) failed")
+	}
+	onlineUsers.Add(1)
+	if onlineUsers.Load() != 2 {
+		t.Error("onlineUsers.Add(1) failed")
+	}
+	onlineUsers.Add(^uint64(0))
+	if onlineUsers.Load() != 1 {
+		t.Error("onlineUsers.Add(^uint64(0)) failed")
+	}
+	onlineUsers.Add(^uint64(0))
+	if onlineUsers.Load() != 0 {
+		t.Error("onlineUsers.Add(^uint64(0)) failed")
+	}
+}

model/server_api.go

@@ -16,6 +16,7 @@ type StreamServer struct {
 
 type StreamServerData struct {
 	Now     int64          `json:"now,omitempty"`
+	Online  uint64         `json:"online,omitempty"`
 	Servers []StreamServer `json:"servers,omitempty"`
 }
 

model/user.go

@@ -19,6 +19,10 @@ type User struct {
 }
 
 func (u *User) BeforeSave(tx *gorm.DB) error {
+	if u.AgentSecret != "" {
+		return nil
+	}
+
 	key, err := utils.GenerateRandomString(32)
 	if err != nil {
 		return err

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"`
+	ID              uint64 `json:"id,omitempty"`
+	IP              string `json:"ip,omitempty"`
+	BlockReason     uint8  `json:"block_reason,omitempty"`
+	BlockTimestamp  uint64 `json:"block_timestamp,omitempty"`
+	BlockIdentifier uint64 `json:"block_identifier,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"`
+	ID              uint64 `gorm:"primaryKey" json:"id,omitempty"`
+	IP              []byte `gorm:"type:binary(16);index:idx_block_identifier" json:"ip,omitempty"`
+	BlockReason     uint8  `json:"block_reason,omitempty"`
+	BlockTimestamp  uint64 `json:"block_timestamp,omitempty"`
+	BlockIdentifier int64  `gorm:"index:idx_block_identifier" json:"block_identifier,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{}).Select("block_timestamp").Order("id desc").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 int64
+	if err := db.Model(&WAF{}).Where("ip = ?", ipBinary).Count(&count).Error; err != nil {
+		return err
+	}
+
 	now := time.Now().Unix()
-	if powAdd(w.Count, 4, w.LastBlockTimestamp) > uint64(now) {
+	if powAdd(uint64(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,20 @@ 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,
+		BlockReason:     reason,
+		BlockTimestamp:  uint64(time.Now().Unix()),
+		BlockIdentifier: uid,
+	}
 	return db.Transaction(func(tx *gorm.DB) error {
-		if err := tx.Where(&w).Attrs(WAF{
-			LastBlockReason:    reason,
-			LastBlockTimestamp: uint64(time.Now().Unix()),
-		}).FirstOrCreate(&w).Error; err != nil {
-			return err
+		var lastRecord WAF
+		if err := tx.Model(&WAF{}).Order("id desc").Where("ip = ?", ipBinary).First(&lastRecord).Error; err != nil {
+			if !errors.Is(err, gorm.ErrRecordNotFound) {
+				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.Create(&w).Error
 	})
 }
 

service/rpc/auth.go

@@ -41,12 +41,12 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
 	userId, ok := singleton.AgentSecretToUserId[clientSecret]
 	if !ok && subtle.ConstantTimeCompare([]byte(clientSecret), []byte(singleton.Conf.AgentSecretKey)) != 1 {
 		singleton.UserLock.RUnlock()
-		model.BlockIP(singleton.DB, ip, model.WAFBlockReasonTypeAgentAuthFail)
+		model.BlockIP(singleton.DB, ip, model.WAFBlockReasonTypeAgentAuthFail, model.BlockIDgRPC)
 		return 0, status.Error(codes.Unauthenticated, "客户端认证失败")
 	}
 	singleton.UserLock.RUnlock()
 
-	model.ClearIP(singleton.DB, ip)
+	model.ClearIP(singleton.DB, ip, model.BlockIDgRPC)
 
 	var clientUUID string
 	if value, ok := md["client_uuid"]; ok {

service/singleton/frontend-templates.yaml

@@ -9,10 +9,10 @@
   name: "Official"
   repository: "https://github.com/hamster1963/nezha-dash-v1"
   author: "hamster1963"
-  version: "v1.7.4"
+  version: "v1.7.6"
   isofficial: true
 - path: "nazhua-dist"
   name: "Nazhua"
   repository: "https://github.com/hi2shark/nazhua"
   author: "hi2hi"
-  version: "v0.4.22"
+  version: "v0.4.23"

service/singleton/singleton.go

@@ -3,6 +3,7 @@ package singleton
 import (
 	_ "embed"
 	"log"
+	"sync/atomic"
 	"time"
 
 	"github.com/patrickmn/go-cache"
@@ -23,6 +24,7 @@ var (
 	Loc               *time.Location
 	FrontendTemplates []model.FrontendTemplate
 	DashboardBootTime = uint64(time.Now().Unix())
+	OnlineUsers       = new(atomic.Uint64)
 )
 
 //go:embed frontend-templates.yaml

service/singleton/user.go

@@ -39,50 +39,65 @@ func OnUserUpdate(u *model.User) {
 	AgentSecretToUserId[u.AgentSecret] = u.ID
 }
 
-func OnUserDelete(id []uint64) {
+func OnUserDelete(id []uint64, errorFunc func(string, ...interface{}) error) error {
 	UserLock.Lock()
 	defer UserLock.Unlock()
 
 	if len(id) < 1 {
-		return
+		return Localizer.ErrorT("user id not specified")
 	}
 
 	var (
-		cron   bool
-		server bool
+		cron, server   bool
+		crons, servers []uint64
 	)
 
 	for _, uid := range id {
-		secret := UserIdToAgentSecret[uid]
-		delete(AgentSecretToUserId, secret)
-		delete(UserIdToAgentSecret, uid)
+		err := DB.Transaction(func(tx *gorm.DB) error {
+			CronLock.RLock()
+			crons = model.FindUserID(CronList, uid)
+			CronLock.RUnlock()
 
-		CronLock.RLock()
-		crons := model.FindUserID(CronList, uid)
-		CronLock.RUnlock()
+			cron = len(crons) > 0
+			if cron {
+				if err := tx.Unscoped().Delete(&model.Cron{}, "id in (?)", crons).Error; err != nil {
+					return err
+				}
+			}
 
-		cron = len(crons) > 0
-		if cron {
-			DB.Unscoped().Delete(&model.Cron{}, "id in (?)", crons)
-			OnDeleteCron(crons)
-		}
+			SortedServerLock.RLock()
+			servers = model.FindUserID(SortedServerList, uid)
+			SortedServerLock.RUnlock()
 
-		SortedServerLock.RLock()
-		servers := model.FindUserID(SortedServerList, uid)
-		SortedServerLock.RUnlock()
-
-		server = len(servers) > 0
-		if server {
-			DB.Transaction(func(tx *gorm.DB) error {
+			server = len(servers) > 0
+			if server {
 				if err := tx.Unscoped().Delete(&model.Server{}, "id in (?)", servers).Error; err != nil {
 					return err
 				}
 				if err := tx.Unscoped().Delete(&model.ServerGroupServer{}, "server_id in (?)", servers).Error; err != nil {
 					return err
 				}
-				return nil
-			})
+			}
 
+			if err := tx.Unscoped().Delete(&model.Transfer{}, "server_id in (?)", servers).Error; err != nil {
+				return err
+			}
+
+			if err := tx.Where("id IN (?)", id).Delete(&model.User{}).Error; err != nil {
+				return err
+			}
+			return nil
+		})
+
+		if err != nil {
+			return errorFunc("%v", err)
+		}
+
+		if cron {
+			OnDeleteCron(crons)
+		}
+
+		if server {
 			AlertsLock.Lock()
 			for _, sid := range servers {
 				for _, alert := range Alerts {
@@ -93,10 +108,13 @@ func OnUserDelete(id []uint64) {
 					}
 				}
 			}
-			DB.Unscoped().Delete(&model.Transfer{}, "server_id in (?)", servers)
 			AlertsLock.Unlock()
 			OnServerDelete(servers)
 		}
+
+		secret := UserIdToAgentSecret[uid]
+		delete(AgentSecretToUserId, secret)
+		delete(UserIdToAgentSecret, uid)
 	}
 
 	if cron {
@@ -106,4 +124,6 @@ func OnUserDelete(id []uint64) {
 	if server {
 		ReSortServer()
 	}
+
+	return nil
 }