diff --git a/go-backend/internal/ws/server.go b/go-backend/internal/ws/server.go index 82d5b2b..acde858 100644 --- a/go-backend/internal/ws/server.go +++ b/go-backend/internal/ws/server.go @@ -30,22 +30,23 @@ type broadcastMessage struct { Data string `json:"data"` } -type connWrap struct { - conn *websocket.Conn - mu sync.Mutex -} - type nodeSession struct { nodeID int64 secret string - conn *connWrap + conn *websocket.Conn crypto *security.AESCrypto // 缓存的 AES 加密器,避免每条消息重建 } type adminSession struct { userID int64 claims auth.Claims - conn *connWrap + conn *websocket.Conn +} + +type monitorSession struct { + userID int64 + claims auth.Claims + conn *websocket.Conn } type commandResponse struct { @@ -83,10 +84,11 @@ type Server struct { getUserAuthState func(userID int64) (*auth.UserAuthState, error) mu sync.RWMutex - admins map[*adminSession]struct{} - nodes map[int64]*nodeSession - byConn map[*websocket.Conn]*nodeSession - pending map[string]pendingRequest + admins map[*adminSession]struct{} + monitors map[*monitorSession]struct{} + nodes map[int64]*nodeSession + byConn map[*websocket.Conn]*nodeSession + pending map[string]pendingRequest } type SystemInfo struct { @@ -130,10 +132,11 @@ func NewServer(repo *repo.Repository, jwtSecret string) *Server { upgrader: websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true }, }, - admins: make(map[*adminSession]struct{}), - nodes: make(map[int64]*nodeSession), - byConn: make(map[*websocket.Conn]*nodeSession), - pending: make(map[string]pendingRequest), + admins: make(map[*adminSession]struct{}), + monitors: make(map[*monitorSession]struct{}), + nodes: make(map[int64]*nodeSession), + byConn: make(map[*websocket.Conn]*nodeSession), + pending: make(map[string]pendingRequest), } } @@ -172,15 +175,19 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { http.Error(w, "forbidden", http.StatusForbidden) return } - if claims.RoleID != 0 { + if claims.RoleID == 0 { + if !s.validateAdminSession(userID, claims) { + http.Error(w, "forbidden", http.StatusForbidden) + return + } + s.handleAdmin(w, r, userID, claims) + return + } + if !s.validateMonitorSession(userID, claims) { http.Error(w, "forbidden", http.StatusForbidden) return } - if !s.validateAdminSession(userID, claims) { - http.Error(w, "forbidden", http.StatusForbidden) - return - } - s.handleAdmin(w, r, userID, claims) + s.handleMonitor(w, r, userID, claims) return } @@ -192,14 +199,13 @@ func (s *Server) handleAdmin(w http.ResponseWriter, r *http.Request, userID int6 if err != nil { return } - cw := &connWrap{conn: conn} _ = conn.SetReadDeadline(time.Now().Add(wsPongWait)) conn.SetPongHandler(func(string) error { return conn.SetReadDeadline(time.Now().Add(wsPongWait)) }) done := make(chan struct{}) - session := &adminSession{userID: userID, claims: claims, conn: cw} - go startKeepalive(cw, done, func() bool { + session := &adminSession{userID: userID, claims: claims, conn: conn} + go startKeepalive(conn, done, func() bool { return s.validateAdminSession(session.userID, session.claims) }) @@ -222,18 +228,51 @@ func (s *Server) handleAdmin(w http.ResponseWriter, r *http.Request, userID int6 } } -func (s *Server) handleNode(w http.ResponseWriter, r *http.Request, nodeID int64, secret string) { +func (s *Server) handleMonitor(w http.ResponseWriter, r *http.Request, userID int64, claims auth.Claims) { conn, err := s.upgrader.Upgrade(w, r, nil) if err != nil { return } - cw := &connWrap{conn: conn} _ = conn.SetReadDeadline(time.Now().Add(wsPongWait)) conn.SetPongHandler(func(string) error { return conn.SetReadDeadline(time.Now().Add(wsPongWait)) }) done := make(chan struct{}) - go startKeepalive(cw, done, nil) + session := &monitorSession{userID: userID, claims: claims, conn: conn} + go startKeepalive(conn, done, func() bool { + return s.validateMonitorSession(session.userID, session.claims) + }) + + s.mu.Lock() + s.monitors[session] = struct{}{} + s.mu.Unlock() + + defer func() { + close(done) + s.mu.Lock() + delete(s.monitors, session) + s.mu.Unlock() + _ = conn.Close() + }() + + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } +} + +func (s *Server) handleNode(w http.ResponseWriter, r *http.Request, nodeID int64, secret string) { + conn, err := s.upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + _ = conn.SetReadDeadline(time.Now().Add(wsPongWait)) + conn.SetPongHandler(func(string) error { + return conn.SetReadDeadline(time.Now().Add(wsPongWait)) + }) + done := make(chan struct{}) + go startKeepalive(conn, done, nil) version := r.URL.Query().Get("version") httpVal := parseIntDefault(r.URL.Query().Get("http"), 0) @@ -242,15 +281,15 @@ func (s *Server) handleNode(w http.ResponseWriter, r *http.Request, nodeID int64 s.mu.Lock() if old, ok := s.nodes[nodeID]; ok { - _ = old.conn.conn.Close() - delete(s.byConn, old.conn.conn) + _ = old.conn.Close() + delete(s.byConn, old.conn) } // 初始化 AES 加密器并缓存(仅创建一次) var nodeCrypto *security.AESCrypto if strings.TrimSpace(secret) != "" { nodeCrypto, _ = security.NewAESCrypto(secret) } - ns := &nodeSession{nodeID: nodeID, secret: secret, conn: cw, crypto: nodeCrypto} + ns := &nodeSession{nodeID: nodeID, secret: secret, conn: conn, crypto: nodeCrypto} s.nodes[nodeID] = ns s.byConn[conn] = ns s.mu.Unlock() @@ -270,7 +309,7 @@ func (s *Server) handleNode(w http.ResponseWriter, r *http.Request, nodeID int64 needOfflineBroadcast := false s.mu.Lock() current, ok := s.nodes[nodeID] - if ok && current.conn.conn == conn { + if ok && current.conn == conn { delete(s.nodes, nodeID) needOfflineBroadcast = true } @@ -400,7 +439,7 @@ func (s *Server) SendCommand(nodeID int64, cmdType string, data interface{}, tim s.mu.RLock() ns, ok := s.nodes[nodeID] s.mu.RUnlock() - if !ok || ns == nil || ns.conn == nil || ns.conn.conn == nil { + if !ok || ns == nil || ns.conn == nil { return CommandResult{}, errors.New("节点不在线") } @@ -450,11 +489,9 @@ func (s *Server) SendCommand(nodeID int64, cmdType string, data interface{}, tim } } - ns.conn.mu.Lock() - _ = ns.conn.conn.SetWriteDeadline(time.Now().Add(wsWriteWait)) - err = ns.conn.conn.WriteMessage(websocket.TextMessage, messageData) - _ = ns.conn.conn.SetWriteDeadline(time.Time{}) - ns.conn.mu.Unlock() + _ = ns.conn.SetWriteDeadline(time.Now().Add(wsWriteWait)) + err = ns.conn.WriteMessage(websocket.TextMessage, messageData) + _ = ns.conn.SetWriteDeadline(time.Time{}) if err != nil { cleanup() return CommandResult{}, err @@ -570,38 +607,51 @@ func (s *Server) broadcastStatus(nodeID int64, status int) { "data": status, } raw, _ := json.Marshal(payload) - s.broadcastToAdmins(string(raw)) + s.broadcastToRealtime(string(raw)) } func (s *Server) broadcastInfo(nodeID int64, data string) { payload := broadcastMessage{ID: nodeID, Type: "info", Data: data} raw, _ := json.Marshal(payload) - s.broadcastToAdmins(string(raw)) + s.broadcastToRealtime(string(raw)) } func (s *Server) broadcastTyped(nodeID int64, msgType string, data string) { payload := broadcastMessage{ID: nodeID, Type: msgType, Data: data} raw, _ := json.Marshal(payload) - s.broadcastToAdmins(string(raw)) + s.broadcastToRealtime(string(raw)) } -func (s *Server) broadcastToAdmins(message string) { +func (s *Server) broadcastToRealtime(message string) { s.mu.RLock() admins := make([]*adminSession, 0, len(s.admins)) for c := range s.admins { admins = append(admins, c) } + monitors := make([]*monitorSession, 0, len(s.monitors)) + for c := range s.monitors { + monitors = append(monitors, c) + } s.mu.RUnlock() for _, c := range admins { - if c == nil || c.conn == nil || c.conn.conn == nil { + if c == nil || c.conn == nil { continue } - c.conn.mu.Lock() - _ = c.conn.conn.SetWriteDeadline(time.Now().Add(wsWriteWait)) - err := c.conn.conn.WriteMessage(websocket.TextMessage, []byte(message)) - _ = c.conn.conn.SetWriteDeadline(time.Time{}) - c.conn.mu.Unlock() + _ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait)) + err := c.conn.WriteMessage(websocket.TextMessage, []byte(message)) + _ = c.conn.SetWriteDeadline(time.Time{}) + if err != nil { + log.Printf("websocket broadcast failed: %v", err) + } + } + for _, c := range monitors { + if c == nil || c.conn == nil { + continue + } + _ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait)) + err := c.conn.WriteMessage(websocket.TextMessage, []byte(message)) + _ = c.conn.SetWriteDeadline(time.Time{}) if err != nil { log.Printf("websocket broadcast failed: %v", err) } @@ -655,8 +705,31 @@ func (s *Server) validateAdminSession(userID int64, claims auth.Claims) bool { return true } -func startKeepalive(cw *connWrap, done <-chan struct{}, validate func() bool) { - if cw == nil || cw.conn == nil { +func (s *Server) validateMonitorSession(userID int64, claims auth.Claims) bool { + if s == nil { + return false + } + if claims.Exp <= time.Now().Unix() { + return false + } + if s.getUserAuthState != nil { + state, err := s.getUserAuthState(userID) + if err != nil || state == nil || state.Status != 1 || state.RoleID != claims.RoleID || claims.IatMs <= state.PasswordChangedAt { + return false + } + } + if s.repo == nil { + return false + } + allowed, err := s.repo.HasMonitorPermission(userID) + if err != nil || !allowed { + return false + } + return true +} + +func startKeepalive(conn *websocket.Conn, done <-chan struct{}, validate func() bool) { + if conn == nil { return } ticker := time.NewTicker(wsPingPeriod) @@ -668,16 +741,14 @@ func startKeepalive(cw *connWrap, done <-chan struct{}, validate func() bool) { return case <-ticker.C: if validate != nil && !validate() { - _ = cw.conn.Close() + _ = conn.Close() return } - cw.mu.Lock() - _ = cw.conn.SetWriteDeadline(time.Now().Add(wsWriteWait)) - err := cw.conn.WriteMessage(websocket.PingMessage, nil) - _ = cw.conn.SetWriteDeadline(time.Time{}) - cw.mu.Unlock() + _ = conn.SetWriteDeadline(time.Now().Add(wsWriteWait)) + err := conn.WriteMessage(websocket.PingMessage, nil) + _ = conn.SetWriteDeadline(time.Time{}) if err != nil { - _ = cw.conn.Close() + _ = conn.Close() return } } diff --git a/go-backend/internal/ws/server_test.go b/go-backend/internal/ws/server_test.go index b1322ad..5cba6b8 100644 --- a/go-backend/internal/ws/server_test.go +++ b/go-backend/internal/ws/server_test.go @@ -4,9 +4,13 @@ import ( "net/http" "net/http/httptest" "net/url" + "strings" "testing" "go-backend/internal/auth" + "go-backend/internal/store/repo" + + "github.com/gorilla/websocket" ) func TestServeHTTPRejectsDisabledAdminToken(t *testing.T) { @@ -108,3 +112,40 @@ func TestValidateAdminSessionRejectsExpiredToken(t *testing.T) { t.Fatal("expected expired token to be rejected") } } + +func TestServeHTTPAllowsMonitorTokenWithPermission(t *testing.T) { + secret := "unit-test-secret" + token, err := auth.GenerateToken(2, "normal_user", 1, secret) + if err != nil { + t.Fatalf("generate token: %v", err) + } + + r, err := repo.Open(t.TempDir() + "/monitor.db") + if err != nil { + t.Fatalf("open repo: %v", err) + } + defer r.Close() + if err := r.InsertMonitorPermission(2, 123); err != nil { + t.Fatalf("insert permission: %v", err) + } + + server := NewServer(r, secret) + server.SetUserAuthStateLookup(func(userID int64) (*auth.UserAuthState, error) { + return &auth.UserAuthState{ID: userID, RoleID: 1, Status: 1, PasswordChangedAt: 0}, nil + }) + + ts := httptest.NewServer(server) + defer ts.Close() + + conn, resp, err := websocket.DefaultDialer.Dial( + "ws"+strings.TrimPrefix(ts.URL, "http")+"/system-info?type=0&secret="+url.QueryEscape(token), + nil, + ) + if err != nil { + if resp != nil { + t.Fatalf("dial websocket error = %v, status=%d", err, resp.StatusCode) + } + t.Fatalf("dial websocket error = %v", err) + } + _ = conn.Close() +}