fix: harden panel self-upgrade (#506)

This fixes panel self-upgrade and monitor realtime access for permitted
non-admin users.

Changes:
- Reuse GitHub proxy config for system upgrade release/API fetches.
- Make panel self-upgrade helper fail fast on compose command errors.
- Allow users with explicit monitor permission to connect to the
realtime websocket and receive broadcasts, not just admins.

Validation:
- cd go-backend && GOCACHE=/private/tmp/flvx-gocache go test
./internal/http/handler -count=1
- cd go-backend && go test ./internal/ws -count=1
This commit is contained in:
sagit
2026-05-15 23:54:27 +08:00
committed by GitHub
4 changed files with 309 additions and 93 deletions
@@ -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()))
@@ -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)
+127 -56
View File
@@ -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
}
}
+41
View File
@@ -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()
}