feat: list & block online users

cmd/dashboard/controller/controller.go

@@ -132,6 +132,9 @@ func routers(r *gin.Engine, frontendDist fs.FS) {
 	auth.GET("/waf", pCommonHandler(listBlockedAddress))
 	auth.POST("/batch-delete/waf", adminHandler(batchDeleteBlockedAddress))
 
+	auth.GET("/online-user", pCommonHandler(listOnlineUser))
+	auth.GET("/online-user/batch-block", adminHandler(batchBlockOnlineUser))
+
 	auth.PATCH("/setting", adminHandler(updateConfig))
 
 	r.NoRoute(fallbackToFrontend(frontendDist))

cmd/dashboard/controller/setting.go

@@ -50,7 +50,7 @@ func listConfig(c *gin.Context) (model.SettingResponse, error) {
 // @Security BearerAuth
 // @Schemes
 // @Description Edit config
-// @Tags auth required
+// @Tags admin required
 // @Accept json
 // @Param body body model.SettingForm true "SettingForm"
 // @Produce json

cmd/dashboard/controller/user.go

@@ -2,6 +2,7 @@ package controller
 
 import (
 	"slices"
+	"strconv"
 
 	"github.com/gin-gonic/gin"
 	"golang.org/x/crypto/bcrypt"
@@ -76,7 +77,7 @@ func updateProfile(c *gin.Context) (any, error) {
 // @Security BearerAuth
 // @Schemes
 // @Description List user
-// @Tags auth required
+// @Tags admin required
 // @Produce json
 // @Success 200 {object} model.CommonResponse[[]model.User]
 // @Router /user [get]
@@ -93,7 +94,7 @@ func listUser(c *gin.Context) ([]model.User, error) {
 // @Security BearerAuth
 // @Schemes
 // @Description Create user
-// @Tags auth required
+// @Tags admin required
 // @Accept json
 // @param request body model.UserForm true "User Request"
 // @Produce json
@@ -135,7 +136,7 @@ func createUser(c *gin.Context) (uint64, error) {
 // @Security BearerAuth
 // @Schemes
 // @Description Batch delete users
-// @Tags auth required
+// @Tags admin required
 // @Accept json
 // @param request body []uint true "id list"
 // @Produce json
@@ -154,3 +155,59 @@ func batchDeleteUser(c *gin.Context) (any, error) {
 	err := singleton.OnUserDelete(ids, newGormError)
 	return nil, err
 }
+
+// List online users
+// @Summary List online users
+// @Security BearerAuth
+// @Schemes
+// @Description List online users
+// @Tags auth required
+// @Param limit query uint false "Page limit"
+// @Param offset query uint false "Page offset"
+// @Produce json
+// @Success 200 {object} model.PaginatedResponse[[]model.OnlineUser, model.OnlineUser]
+// @Router /online-user [get]
+func listOnlineUser(c *gin.Context) (*model.Value[[]*model.OnlineUser], 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 < 0 {
+		offset = 0
+	}
+
+	return &model.Value[[]*model.OnlineUser]{
+		Value: singleton.GetOnlineUsers(limit, offset),
+		Pagination: model.Pagination{
+			Offset: offset,
+			Limit:  limit,
+			Total:  int64(singleton.GetOnlineUserCount()),
+		},
+	}, nil
+}
+
+// Batch block online user
+// @Summary Batch block online user
+// @Security BearerAuth
+// @Schemes
+// @Description Batch block online user
+// @Tags admin required
+// @Accept json
+// @Param request body []string true "block list"
+// @Produce json
+// @Success 200 {object} model.CommonResponse[any]
+// @Router /online-user/batch-block [patch]
+func batchBlockOnlineUser(c *gin.Context) (any, error) {
+	var list []string
+	if err := c.ShouldBindJSON(&list); err != nil {
+		return nil, err
+	}
+
+	if err := singleton.BlockByIPs(list); err != nil {
+		return nil, newGormError("%v", err)
+	}
+
+	return nil, nil
+}

cmd/dashboard/controller/waf.go

@@ -56,7 +56,7 @@ func listBlockedAddress(c *gin.Context) (*model.Value[[]*model.WAF], error) {
 // @Security BearerAuth
 // @Schemes
 // @Description Edit server
-// @Tags auth required
+// @Tags admin required
 // @Accept json
 // @Param request body []string true "block list"
 // @Produce json

cmd/dashboard/controller/ws.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/gorilla/websocket"
+	"github.com/hashicorp/go-uuid"
 	"golang.org/x/sync/singleflight"
 
 	"github.com/nezhahq/nezha/model"
@@ -102,13 +103,29 @@ func checkSameOrigin(r *http.Request) bool {
 // @Success 200 {object} model.StreamServerData
 // @Router /ws/server [get]
 func serverStream(c *gin.Context) (any, error) {
+	connId, err := uuid.GenerateUUID()
+	if err != nil {
+		return nil, newWsError("%v", err)
+	}
+
 	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
 	if err != nil {
 		return nil, newWsError("%v", err)
 	}
 	defer conn.Close()
-	singleton.OnlineUsers.Add(1)
-	defer singleton.OnlineUsers.Add(^uint64(0))
+
+	userIp := c.GetString(model.CtxKeyRealIPStr)
+	if userIp == "" {
+		userIp = c.RemoteIP()
+	}
+
+	singleton.AddOnlineUser(connId, &model.OnlineUser{
+		IP:          userIp,
+		ConnectedAt: time.Now(),
+		Conn:        conn,
+	})
+	defer singleton.RemoveOnlineUser(connId)
+
 	count := 0
 	for {
 		stat, err := getServerStat(c, count == 0)
@@ -166,7 +183,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(),
+			Online:  singleton.GetOnlineUserCount(),
 			Servers: servers,
 		})
 	})

model/server_api.go

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

model/user.go

@@ -1,6 +1,9 @@
 package model
 
 import (
+	"time"
+
+	"github.com/gorilla/websocket"
 	"github.com/nezhahq/nezha/pkg/utils"
 	"gorm.io/gorm"
 )
@@ -20,7 +23,6 @@ type User struct {
 
 type UserInfo struct {
 	Role        uint8
-	_           [3]byte
 	AgentSecret string
 }
 
@@ -42,3 +44,11 @@ type Profile struct {
 	User
 	LoginIP string `json:"login_ip,omitempty"`
 }
+
+type OnlineUser struct {
+	UserID      uint64    `json:"user_id,omitempty"`
+	ConnectedAt time.Time `json:"connected_at,omitempty"`
+	IP          string    `json:"ip,omitempty"`
+
+	Conn *websocket.Conn `json:"-"`
+}

model/waf.go

@@ -14,12 +14,14 @@ const (
 	WAFBlockReasonTypeLoginFail
 	WAFBlockReasonTypeBruteForceToken
 	WAFBlockReasonTypeAgentAuthFail
+	WAFBlockReasonTypeManual
 )
 
 const (
 	BlockIDgRPC = -127 + iota
 	BlockIDToken
 	BlockIDUnknownUser
+	BlockIDManual
 )
 
 type WAFApiMock struct {

service/singleton/online_user.go

@@ -0,0 +1,68 @@
+package singleton
+
+import (
+	"slices"
+	"sync"
+
+	"github.com/nezhahq/nezha/model"
+)
+
+var (
+	OnlineUserMap     = make(map[string]*model.OnlineUser)
+	OnlineUserMapLock = new(sync.Mutex)
+)
+
+func AddOnlineUser(connId string, user *model.OnlineUser) {
+	OnlineUserMapLock.Lock()
+	defer OnlineUserMapLock.Unlock()
+	OnlineUserMap[connId] = user
+}
+
+func RemoveOnlineUser(connId string) {
+	OnlineUserMapLock.Lock()
+	defer OnlineUserMapLock.Unlock()
+	delete(OnlineUserMap, connId)
+}
+
+func BlockByIPs(ipList []string) error {
+	OnlineUserMapLock.Lock()
+	defer OnlineUserMapLock.Unlock()
+
+	for _, ip := range ipList {
+		if err := model.BlockIP(DB, ip, model.WAFBlockReasonTypeManual, model.BlockIDManual); err != nil {
+			return err
+		}
+		for _, user := range OnlineUserMap {
+			if user.IP == ip && user.Conn != nil {
+				user.Conn.Close()
+			}
+		}
+	}
+
+	return nil
+}
+
+func GetOnlineUsers(limit, offset int) []*model.OnlineUser {
+	OnlineUserMapLock.Lock()
+	defer OnlineUserMapLock.Unlock()
+	var users []*model.OnlineUser
+	for _, user := range OnlineUserMap {
+		users = append(users, user)
+	}
+	slices.SortFunc(users, func(i, j *model.OnlineUser) int {
+		return i.ConnectedAt.Compare(j.ConnectedAt)
+	})
+	if offset > len(users) {
+		return nil
+	}
+	if offset+limit > len(users) {
+		return users[offset:]
+	}
+	return users[offset : offset+limit]
+}
+
+func GetOnlineUserCount() int {
+	OnlineUserMapLock.Lock()
+	defer OnlineUserMapLock.Unlock()
+	return len(OnlineUserMap)
+}

service/singleton/singleton.go

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