Compare commits

...

7 Commits

Author SHA1 Message Date
sagit 2d0c993c90 fix: allow admin access to sensitive configs (#511) 2026-05-17 23:09:31 +08:00
sagitchu 8b64542c94 fix(handler): allow admin access to sensitive configs 2026-05-17 23:07:37 +08:00
sagit 032b0f0cfd fix: serialize websocket writes in realtime server (#510) 2026-05-17 21:52:54 +08:00
sagitchu 7008717a49 fix(ws): serialize websocket writes in realtime server 2026-05-17 21:48:37 +08:00
sagit 8552a70355 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
2026-05-15 23:54:27 +08:00
sagitchu b2454e86c9 fix: allow monitor realtime websocket access 2026-05-15 23:37:08 +08:00
sagitchu 312c9a9c5c fix: harden panel self-upgrade 2026-05-15 23:16:14 +08:00
6 changed files with 482 additions and 79 deletions
@@ -67,7 +67,22 @@ func TestConfigGetRejectsSensitiveKeysWithoutAuth(t *testing.T) {
assertHandlerCodeMsg(t, resp, 403, "禁止访问敏感配置")
}
func TestConfigUpdateRejectsSensitiveKeys(t *testing.T) {
func TestConfigGetAllowsSensitiveKeysForAdmin(t *testing.T) {
router, r := setupConfigAccessTestRouter(t)
adminToken := mustGenerateConfigAccessToken(t, 1, "admin_user", 0)
seedConfigValue(t, r, "jwt_secret", "jwt-secret")
req := httptest.NewRequest(http.MethodPost, "/api/v1/config/get", bytes.NewBufferString(`{"name":"jwt_secret"}`))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", adminToken)
resp := httptest.NewRecorder()
router.ServeHTTP(resp, req)
assertHandlerConfigValue(t, resp, "jwt_secret", "jwt-secret")
}
func TestConfigUpdateAllowsSensitiveKeysForAdmin(t *testing.T) {
router, _ := setupConfigAccessTestRouter(t)
adminToken := mustGenerateConfigAccessToken(t, 1, "admin_user", 0)
@@ -78,10 +93,10 @@ func TestConfigUpdateRejectsSensitiveKeys(t *testing.T) {
router.ServeHTTP(resp, req)
assertHandlerCodeMsg(t, resp, 403, "禁止访问敏感配置")
assertHandlerCode(t, resp, 0)
}
func TestConfigUpdateSingleRejectsSensitiveKeys(t *testing.T) {
func TestConfigUpdateSingleAllowsSensitiveKeysForAdmin(t *testing.T) {
router, _ := setupConfigAccessTestRouter(t)
adminToken := mustGenerateConfigAccessToken(t, 1, "admin_user", 0)
@@ -92,7 +107,95 @@ func TestConfigUpdateSingleRejectsSensitiveKeys(t *testing.T) {
router.ServeHTTP(resp, req)
assertHandlerCodeMsg(t, resp, 403, "禁止访问敏感配置")
assertHandlerCode(t, resp, 0)
}
func TestConfigUpdateAllowsCloudflareSecretKeyWrite(t *testing.T) {
router, r := setupConfigAccessTestRouter(t)
adminToken := mustGenerateConfigAccessToken(t, 1, "admin_user", 0)
req := httptest.NewRequest(http.MethodPost, "/api/v1/config/update", bytes.NewBufferString(`{"cloudflare_secret_key":"turnstile-secret"}`))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", adminToken)
resp := httptest.NewRecorder()
router.ServeHTTP(resp, req)
assertHandlerCode(t, resp, 0)
cfg, err := r.GetConfigByName("cloudflare_secret_key")
if err != nil {
t.Fatalf("get config: %v", err)
}
if cfg == nil || cfg.Value != "turnstile-secret" {
t.Fatalf("expected cloudflare_secret_key to be updated, got %#v", cfg)
}
}
func TestConfigUpdateSingleAllowsCloudflareSecretKeyWrite(t *testing.T) {
router, r := setupConfigAccessTestRouter(t)
adminToken := mustGenerateConfigAccessToken(t, 1, "admin_user", 0)
req := httptest.NewRequest(http.MethodPost, "/api/v1/config/update-single", bytes.NewBufferString(`{"name":"cloudflare_secret_key","value":"turnstile-secret"}`))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", adminToken)
resp := httptest.NewRecorder()
router.ServeHTTP(resp, req)
assertHandlerCode(t, resp, 0)
cfg, err := r.GetConfigByName("cloudflare_secret_key")
if err != nil {
t.Fatalf("get config: %v", err)
}
if cfg == nil || cfg.Value != "turnstile-secret" {
t.Fatalf("expected cloudflare_secret_key to be updated, got %#v", cfg)
}
}
func TestConfigUpdateAllowsLicenseKeyWrite(t *testing.T) {
router, r := setupConfigAccessTestRouter(t)
adminToken := mustGenerateConfigAccessToken(t, 1, "admin_user", 0)
req := httptest.NewRequest(http.MethodPost, "/api/v1/config/update", bytes.NewBufferString(`{"license_key":"license-secret"}`))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", adminToken)
resp := httptest.NewRecorder()
router.ServeHTTP(resp, req)
assertHandlerCode(t, resp, 0)
cfg, err := r.GetConfigByName("license_key")
if err != nil {
t.Fatalf("get config: %v", err)
}
if cfg == nil || cfg.Value != "license-secret" {
t.Fatalf("expected license_key to be updated, got %#v", cfg)
}
}
func TestConfigUpdateSingleAllowsLicenseKeyWrite(t *testing.T) {
router, r := setupConfigAccessTestRouter(t)
adminToken := mustGenerateConfigAccessToken(t, 1, "admin_user", 0)
req := httptest.NewRequest(http.MethodPost, "/api/v1/config/update-single", bytes.NewBufferString(`{"name":"license_key","value":"license-secret"}`))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", adminToken)
resp := httptest.NewRecorder()
router.ServeHTTP(resp, req)
assertHandlerCode(t, resp, 0)
cfg, err := r.GetConfigByName("license_key")
if err != nil {
t.Fatalf("get config: %v", err)
}
if cfg == nil || cfg.Value != "license-secret" {
t.Fatalf("expected license_key to be updated, got %#v", cfg)
}
}
func setupConfigAccessTestRouter(t *testing.T) (http.Handler, *repo.Repository) {
+11 -4
View File
@@ -393,8 +393,7 @@ func (h *Handler) getConfigByName(w http.ResponseWriter, r *http.Request) {
return
}
configName := strings.ToLower(strings.TrimSpace(req.Name))
switch configName {
case "license_key", "cloudflare_secret_key", "jwt_secret":
if repo.IsSensitiveConfigKey(configName) && !isAdminRequest(r) {
response.WriteJSON(w, response.Err(403, "禁止访问敏感配置"))
return
}
@@ -989,7 +988,7 @@ func (h *Handler) updateConfigs(w http.ResponseWriter, r *http.Request) {
if key == "" {
continue
}
if repo.IsSensitiveConfigKey(key) {
if repo.IsSensitiveConfigKey(key) && !isAdminRequest(r) {
response.WriteJSON(w, response.Err(403, "禁止访问敏感配置"))
return
}
@@ -1030,7 +1029,7 @@ func (h *Handler) updateSingleConfig(w http.ResponseWriter, r *http.Request) {
response.WriteJSON(w, response.ErrDefault("配置名称不能为空"))
return
}
if repo.IsSensitiveConfigKey(name) {
if repo.IsSensitiveConfigKey(name) && !isAdminRequest(r) {
response.WriteJSON(w, response.Err(403, "禁止访问敏感配置"))
return
}
@@ -1060,6 +1059,14 @@ func (h *Handler) updateSingleConfig(w http.ResponseWriter, r *http.Request) {
response.WriteJSON(w, response.OKEmpty())
}
func isAdminRequest(r *http.Request) bool {
if r == nil {
return false
}
claims, ok := r.Context().Value(middleware.ClaimsContextKey).(auth.Claims)
return ok && claims.RoleID == 0
}
func normalizeAndValidateConfigValue(key, value string) (string, error) {
switch strings.TrimSpace(key) {
case "app_logo", "app_favicon":
@@ -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)
+118 -34
View File
@@ -48,6 +48,12 @@ type adminSession struct {
conn *connWrap
}
type monitorSession struct {
userID int64
claims auth.Claims
conn *connWrap
}
type commandResponse struct {
Type string `json:"type"`
Success bool `json:"success"`
@@ -83,10 +89,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 +137,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 +180,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
}
@@ -222,6 +234,41 @@ func (s *Server) handleAdmin(w http.ResponseWriter, r *http.Request, userID int6
}
}
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{})
session := &monitorSession{userID: userID, claims: claims, conn: cw}
go startKeepalive(cw, 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 {
@@ -270,7 +317,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 != nil && current.conn.conn == conn {
delete(s.nodes, nodeID)
needOfflineBroadcast = true
}
@@ -450,11 +497,7 @@ 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()
err = writeWSMessage(ns.conn, websocket.TextMessage, messageData)
if err != nil {
cleanup()
return CommandResult{}, err
@@ -570,38 +613,47 @@ 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 {
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()
err := writeWSMessage(c.conn, websocket.TextMessage, []byte(message))
if err != nil {
log.Printf("websocket broadcast failed: %v", err)
}
}
for _, c := range monitors {
if c == nil || c.conn == nil || c.conn.conn == nil {
continue
}
err := writeWSMessage(c.conn, websocket.TextMessage, []byte(message))
if err != nil {
log.Printf("websocket broadcast failed: %v", err)
}
@@ -655,6 +707,42 @@ func (s *Server) validateAdminSession(userID int64, claims auth.Claims) bool {
return true
}
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 writeWSMessage(cw *connWrap, messageType int, payload []byte) error {
if cw == nil || cw.conn == nil {
return errors.New("websocket connection not initialized")
}
cw.mu.Lock()
defer cw.mu.Unlock()
_ = cw.conn.SetWriteDeadline(time.Now().Add(wsWriteWait))
err := cw.conn.WriteMessage(messageType, payload)
_ = cw.conn.SetWriteDeadline(time.Time{})
return err
}
func startKeepalive(cw *connWrap, done <-chan struct{}, validate func() bool) {
if cw == nil || cw.conn == nil {
return
@@ -671,11 +759,7 @@ func startKeepalive(cw *connWrap, done <-chan struct{}, validate func() bool) {
_ = cw.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()
err := writeWSMessage(cw, websocket.PingMessage, nil)
if err != nil {
_ = cw.conn.Close()
return
+105
View File
@@ -4,9 +4,14 @@ import (
"net/http"
"net/http/httptest"
"net/url"
"strings"
"sync"
"testing"
"go-backend/internal/auth"
"go-backend/internal/store/repo"
"github.com/gorilla/websocket"
)
func TestServeHTTPRejectsDisabledAdminToken(t *testing.T) {
@@ -108,3 +113,103 @@ 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()
}
func TestConnWrapSerializesConcurrentWrites(t *testing.T) {
serverConn, clientConn := websocketTestPipe(t)
defer serverConn.Close()
defer clientConn.Close()
cw := &connWrap{conn: serverConn}
readerDone := make(chan struct{})
go func() {
defer close(readerDone)
for i := 0; i < 64; i++ {
if _, _, err := clientConn.ReadMessage(); err != nil {
return
}
}
}()
var wg sync.WaitGroup
for i := 0; i < 64; i++ {
wg.Add(1)
go func() {
defer wg.Done()
if err := writeWSMessage(cw, websocket.TextMessage, []byte("x")); err != nil {
t.Errorf("writeWSMessage() error = %v", err)
}
}()
}
wg.Wait()
_ = clientConn.Close()
<-readerDone
}
func websocketTestPipe(t *testing.T) (*websocket.Conn, *websocket.Conn) {
t.Helper()
upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
serverConnCh := make(chan *websocket.Conn, 1)
serverErrCh := make(chan error, 1)
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
serverErrCh <- err
return
}
serverConnCh <- conn
}))
t.Cleanup(ts.Close)
clientConn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(ts.URL, "http"), nil)
if err != nil {
t.Fatalf("dial websocket: %v", err)
}
select {
case err := <-serverErrCh:
t.Fatalf("upgrade websocket: %v", err)
case serverConn := <-serverConnCh:
return serverConn, clientConn
}
return nil, nil
}