From 312c9a9c5c416330a7cd570ebe85cc8dc72defcd Mon Sep 17 00:00:00 2001 From: sagitchu Date: Fri, 15 May 2026 23:16:14 +0800 Subject: [PATCH 1/2] fix: harden panel self-upgrade --- .../internal/http/handler/system_upgrade.go | 97 ++++++++++++++++--- .../http/handler/system_upgrade_test.go | 81 ++++++++++++---- 2 files changed, 141 insertions(+), 37 deletions(-) diff --git a/go-backend/internal/http/handler/system_upgrade.go b/go-backend/internal/http/handler/system_upgrade.go index 04c53de..905c56c 100644 --- a/go-backend/internal/http/handler/system_upgrade.go +++ b/go-backend/internal/http/handler/system_upgrade.go @@ -31,6 +31,10 @@ const ( var safeBackendContainerPattern = regexp.MustCompile(`^[A-Za-z0-9_.-]+$`) var enableIPv6ComposePattern = regexp.MustCompile(`(?im)^\s*enable_ipv6\s*:\s*['"]?true['"]?\s*(?:#.*)?$`) var systemUpgradeReleaseBaseURL = githubHTMLBase +var systemUpgradeAPIBaseURL = githubAPIBase +var systemUpgradeHTTPGet = func(client *http.Client, url string) (*http.Response, error) { + return client.Get(url) +} type systemUpgradeExecutor struct { deployDir string @@ -199,7 +203,7 @@ if [ ! -f .env ]; then fi log "拉取新镜像..." -if ! docker compose pull backend frontend 2>&1 | tee -a "$LOGFILE"; then +if ! docker compose pull backend frontend >> "$LOGFILE" 2>&1; then log "错误: 拉取镜像失败" exit 1 fi @@ -208,7 +212,7 @@ log "等待旧容器释放资源..." sleep 3 log "重启服务（force-recreate）..." -if ! docker compose up -d --force-recreate --remove-orphans backend frontend 2>&1 | tee -a "$LOGFILE"; then +if ! docker compose up -d --force-recreate --remove-orphans backend frontend >> "$LOGFILE" 2>&1; then log "错误: 重启服务失败" exit 1 fi @@ -317,6 +321,68 @@ func (e *systemUpgradeExecutor) replaceCompose(path string, data []byte) error { return writeFileWithMode(path, data, mode) } +func (h *Handler) buildSystemUpgradeDownloadURL(version, filename string) string { + enabled, proxyURL := h.getGithubProxyConfig() + base := fmt.Sprintf("%s/%s/releases/download/%s/%s", strings.TrimRight(systemUpgradeReleaseBaseURL, "/"), githubRepo, version, filename) + if enabled { + return fmt.Sprintf("%s/%s", proxyURL, base) + } + return base +} + +func (h *Handler) fetchSystemUpgradeReleases(perPage int) ([]githubRelease, error) { + if perPage <= 0 { + perPage = 20 + } + + client := &http.Client{Timeout: 15 * time.Second} + url := fmt.Sprintf("%s/repos/%s/releases?per_page=%d", strings.TrimRight(systemUpgradeAPIBaseURL, "/"), githubRepo, perPage) + if enabled, proxyURL := h.getGithubProxyConfig(); enabled { + url = fmt.Sprintf("%s/%s", proxyURL, url) + } + + resp, err := systemUpgradeHTTPGet(client, url) + if err != nil { + return nil, fmt.Errorf("请求GitHub API失败: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) + return nil, fmt.Errorf("GitHub API返回 %d: %s", resp.StatusCode, string(body)) + } + + var releases []githubRelease + if err := json.NewDecoder(resp.Body).Decode(&releases); err != nil { + return nil, fmt.Errorf("解析GitHub API响应失败: %v", err) + } + + return releases, nil +} + +func (h *Handler) resolveSystemUpgradeLatestReleaseByChannel(channel string) (string, error) { + normalizedChannel := normalizeReleaseChannel(channel) + releases, err := h.fetchSystemUpgradeReleases(50) + if err != nil { + return "", err + } + + for _, r := range releases { + if r.Draft { + continue + } + tag := strings.TrimSpace(r.TagName) + if tag == "" { + continue + } + if releaseChannelFromTag(tag) == normalizedChannel { + return tag, nil + } + } + + return "", fmt.Errorf("未找到%s版本号", releaseChannelLabel(normalizedChannel)) +} + func fileModeOrDefault(path string, fallback os.FileMode) (os.FileMode, error) { info, err := os.Stat(path) if err != nil { @@ -367,9 +433,9 @@ func (e *systemUpgradeExecutor) startHelper(ctx context.Context, imageID, helper } func (h *Handler) downloadReleaseAsset(version, filename string) ([]byte, error) { - url := fmt.Sprintf("%s/%s/releases/download/%s/%s", strings.TrimRight(systemUpgradeReleaseBaseURL, "/"), githubRepo, version, filename) + url := h.buildSystemUpgradeDownloadURL(version, filename) client := &http.Client{Timeout: 60 * time.Second} - resp, err := client.Get(url) + resp, err := systemUpgradeHTTPGet(client, url) if err != nil { return nil, fmt.Errorf("下载%s失败: %v", filename, err) } @@ -457,7 +523,7 @@ func (h *Handler) systemVersion(w http.ResponseWriter, r *http.Request) { current := currentPanelVersion() exec := newSystemUpgradeExecutor() capability := exec.capability(r.Context()) - latest, err := resolveLatestReleaseByChannel(channel) + latest, err := h.resolveSystemUpgradeLatestReleaseByChannel(channel) response.WriteJSON(w, response.OK(systemUpgradeVersionResponse(current, channel, latest, err, capability))) } @@ -477,7 +543,7 @@ func (h *Handler) systemCheckUpdates(w http.ResponseWriter, r *http.Request) { exec := newSystemUpgradeExecutor() capability := exec.capability(r.Context()) - githubReleases, err := fetchGitHubReleases(50) + githubReleases, err := h.fetchSystemUpgradeReleases(50) if err != nil { response.WriteJSON(w, response.Err(-2, fmt.Sprintf("获取版本列表失败: %v", err))) return @@ -514,22 +580,21 @@ func (h *Handler) systemUpgrade(w http.ResponseWriter, r *http.Request) { return } channel := normalizeReleaseChannel(req.Channel) - version := strings.TrimSpace(req.Version) - if version == "" { - var err error - version, err = resolveLatestReleaseByChannel(channel) - if err != nil { - response.WriteJSON(w, response.Err(-2, fmt.Sprintf("获取最新%s失败: %v", releaseChannelLabel(channel), err))) - return - } - } - exec := newSystemUpgradeExecutor() capability := exec.capability(r.Context()) if !capability.Capable { response.WriteJSON(w, response.ErrDefault("当前环境不支持面板自升级: "+strings.Join(capability.Reasons, "; "))) return } + version := strings.TrimSpace(req.Version) + if version == "" { + var err error + version, err = h.resolveSystemUpgradeLatestReleaseByChannel(channel) + if err != nil { + response.WriteJSON(w, response.Err(-2, fmt.Sprintf("获取最新%s失败: %v", releaseChannelLabel(channel), err))) + return + } + } imageID, err := exec.currentBackendImage(r.Context()) if err != nil { response.WriteJSON(w, response.Err(-2, err.Error())) diff --git a/go-backend/internal/http/handler/system_upgrade_test.go b/go-backend/internal/http/handler/system_upgrade_test.go index 8222c54..3eff141 100644 --- a/go-backend/internal/http/handler/system_upgrade_test.go +++ b/go-backend/internal/http/handler/system_upgrade_test.go @@ -3,6 +3,7 @@ package handler import ( "bytes" "encoding/json" + "io" "net/http" "net/http/httptest" "os" @@ -10,6 +11,9 @@ import ( "reflect" "strings" "testing" + "time" + + "go-backend/internal/store/repo" ) func TestSelectComposeAssetUsesIPv6Template(t *testing.T) { @@ -21,19 +25,39 @@ func TestSelectComposeAssetUsesIPv6Template(t *testing.T) { } } -func TestDownloadReleaseAssetUsesDirectReleaseURL(t *testing.T) { - var gotPath string - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotPath = r.URL.Path - _, _ = w.Write([]byte("services:\n backend:\n image: test\n")) - })) - defer server.Close() +func TestDownloadReleaseAssetUsesGithubProxyWhenEnabled(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "test.db") + repoStore, err := repo.Open(dbPath) + if err != nil { + t.Fatalf("repo.Open() error = %v", err) + } + defer repoStore.Close() + h := &Handler{repo: repoStore} originalBase := systemUpgradeReleaseBaseURL - systemUpgradeReleaseBaseURL = server.URL + systemUpgradeReleaseBaseURL = "https://example.invalid" t.Cleanup(func() { systemUpgradeReleaseBaseURL = originalBase }) - h := &Handler{} + originalGet := systemUpgradeHTTPGet + defer func() { systemUpgradeHTTPGet = originalGet }() + + var gotURL string + systemUpgradeHTTPGet = func(client *http.Client, url string) (*http.Response, error) { + gotURL = url + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("services:\n backend:\n image: test\n")), + }, nil + } + + now := time.Now().UnixMilli() + if err := repoStore.UpsertConfig("github_proxy_enabled", "true", now); err != nil { + t.Fatalf("UpsertConfig() github_proxy_enabled error = %v", err) + } + if err := repoStore.UpsertConfig("github_proxy_url", "https://proxy.example.com", now); err != nil { + t.Fatalf("UpsertConfig() github_proxy_url error = %v", err) + } + data, err := h.downloadReleaseAsset("2.1.9", "docker-compose-v4.yml") if err != nil { t.Fatalf("downloadReleaseAsset() error = %v", err) @@ -42,24 +66,39 @@ func TestDownloadReleaseAssetUsesDirectReleaseURL(t *testing.T) { t.Fatalf("downloadReleaseAsset() data = %q, want compose data", string(data)) } - wantPath := "/" + githubRepo + "/releases/download/2.1.9/docker-compose-v4.yml" - if gotPath != wantPath { - t.Fatalf("download path = %q, want %q", gotPath, wantPath) + wantURL := "https://proxy.example.com/https://example.invalid/Sagit-chu/flvx/releases/download/2.1.9/docker-compose-v4.yml" + if gotURL != wantURL { + t.Fatalf("download URL = %q, want %q", gotURL, wantURL) } } func TestDownloadReleaseAssetRejectsOversizedBody(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _, _ = w.Write(bytes.Repeat([]byte("a"), maxSystemUpgradeComposeAssetBytes+1)) - })) - defer server.Close() + dbPath := filepath.Join(t.TempDir(), "test.db") + repoStore, err := repo.Open(dbPath) + if err != nil { + t.Fatalf("repo.Open() error = %v", err) + } + defer repoStore.Close() + now := time.Now().UnixMilli() + if err := repoStore.UpsertConfig("github_proxy_enabled", "false", now); err != nil { + t.Fatalf("UpsertConfig() github_proxy_enabled error = %v", err) + } originalBase := systemUpgradeReleaseBaseURL - systemUpgradeReleaseBaseURL = server.URL + systemUpgradeReleaseBaseURL = "https://example.invalid" t.Cleanup(func() { systemUpgradeReleaseBaseURL = originalBase }) - h := &Handler{} - _, err := h.downloadReleaseAsset("2.1.9", "docker-compose-v4.yml") + originalGet := systemUpgradeHTTPGet + defer func() { systemUpgradeHTTPGet = originalGet }() + systemUpgradeHTTPGet = func(client *http.Client, url string) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(bytes.NewReader(bytes.Repeat([]byte("a"), maxSystemUpgradeComposeAssetBytes+1))), + }, nil + } + + h := &Handler{repo: repoStore} + _, err = h.downloadReleaseAsset("2.1.9", "docker-compose-v4.yml") if err == nil || !strings.Contains(err.Error(), "过大") { t.Fatalf("downloadReleaseAsset() error = %v, want oversized error", err) } @@ -272,7 +311,7 @@ func TestSystemUpgradeFailsFastBeforeMutatingFiles(t *testing.T) { fakeDockerDir := t.TempDir() fakeDockerPath := filepath.Join(fakeDockerDir, "docker") - fakeDockerScript := "#!/bin/sh\ncase \"$1\" in\n --version)\n echo 'Docker version 27.0.0'\n exit 0\n ;;\n compose)\n if [ \"$2\" = version ]; then\n echo 'Docker Compose version v2.33.0'\n exit 0\n fi\n exit 0\n ;;\n inspect)\n echo 'No such object: flux-panel-backend' >&2\n exit 1\n ;;\n *)\n exit 0\n ;;\n esac\n" + fakeDockerScript := "#!/bin/sh\ncase \"$1\" in\n --version)\n echo 'Docker version 27.0.0'\n exit 0\n ;;&\n compose)\n if [ \"$2\" = version ]; then\n echo 'Docker Compose version v2.33.0'\n exit 0\n fi\n exit 0\n ;;&\n inspect)\n echo 'No such object: flux-panel-backend' >&2\n exit 1\n ;;&\n *)\n exit 0\n ;;&\n esac\n" if err := os.WriteFile(fakeDockerPath, []byte(fakeDockerScript), 0o755); err != nil { t.Fatalf("WriteFile() fake docker error = %v", err) } @@ -281,7 +320,7 @@ func TestSystemUpgradeFailsFastBeforeMutatingFiles(t *testing.T) { t.Setenv(panelBackendContainerEnv, "flux-panel-backend") h := &Handler{} - req := httptest.NewRequest(http.MethodPost, "/api/v1/system/upgrade", strings.NewReader(`{"channel":"stable","version":"3.0.0"}`)) + req := httptest.NewRequest(http.MethodPost, "/api/v1/system/upgrade", strings.NewReader(`{"channel":"stable"}`)) rr := httptest.NewRecorder() h.systemUpgrade(rr, req) From b2454e86c93547daba025ae3931e4728e8515838 Mon Sep 17 00:00:00 2001 From: sagitchu Date: Fri, 15 May 2026 23:37:08 +0800 Subject: [PATCH 2/2] fix: allow monitor realtime websocket access --- go-backend/internal/ws/server.go | 183 ++++++++++++++++++-------- go-backend/internal/ws/server_test.go | 41 ++++++ 2 files changed, 168 insertions(+), 56 deletions(-) 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() +}