From da467c9297e641e2ab9a1889252589503791b7e9 Mon Sep 17 00:00:00 2001
From: NanamiAdmin <yzt114@yeah.net>
Date: Wed, 23 Sep 2026 11:20:52 +0800
Subject: [PATCH] feat(auth): implement WebSocket authentication for admin
 access to logs

---
 README.md                   |  11 ++-
 frontend/README.md          |   2 +-
 frontend/src/views/Logs.vue |  15 +++-
 router.go                   |   8 +-
 utils/auth.go               | 171 ++++++++++++++++++++++++++----------
 5 files changed, 152 insertions(+), 55 deletions(-)

diff --git a/README.md b/README.md
index 9ee2dbf..1287bf3 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@ Nukumizu connects to a Komari Dashboard instance, keeps an in-memory view of eve
 - **Network proxy** — a global proxy URL can be enabled per controller (`networkUseProxy`) for HTTP, WebSocket, and even SMTP (HTTP CONNECT tunnel).
 - **Customizable message templates** — every bot/notification message is rendered from a template in `config.json`.
 - **Storage** — SQLite (pure-Go driver) for `user.db` and `log.db`; safe on network shares (WAL disabled).
-- **Dashboard API** — token-authenticated REST API plus a live log-streaming WebSocket.
+- **Dashboard API** — token-authenticated REST API plus an admin-only live log-streaming WebSocket.
 - **Web console** — a Vue 3 admin UI for browsing nodes, editing `config.json`, managing bot trust, and tailing logs. The built bundle is embedded in the binary, so a single executable serves both the API and the console.
 
 ## How it works
@@ -45,12 +45,12 @@ nukumizu-backend/
 │   └── user.go                   # user.db (SQLite) user store
 ├── utils/
 │   ├── auth.go                   # Token management, Auth middleware, JSON responses
-│   └── middleware.go             # Rate limit, CORS, XSS/security headers
+│   └── middleware.go             # Rate limit, CORS, XSS headers, WebSocket auth
 ├── postLog/                      # Logging subsystem
 │   ├── postLog.go                # Leveled logger (stdout + broadcast)
 │   ├── database.go               # log.db (SQLite, one table per run)
 │   ├── logBroadcaster.go         # Fan-out to WebSocket clients
-│   └── logSocketHandler.go       # /api/system/getLogs WebSocket handler
+│   └── logSocketHandler.go       # /api/system/getLogs WebSocket handler (admin only)
 ├── web/
 │   ├── embed.go                  # Embeds the built console (web/dist) in the binary
 │   ├── dist/                     # Vite build output — generated, gitignored
@@ -287,6 +287,8 @@ Requests are authenticated with HTTP headers:
 
 Tokens idle for more than 1 hour are expired (cleaned every 10 minutes); any authenticated call refreshes the timer. Permission levels: `None`, `bot`, `admin`. Endpoints requiring `bot` accept both `bot` and `admin` tokens. Currently registration/login always issue `admin`-level tokens.
 
+Browser WebSocket handshakes cannot carry custom headers, so `/api/system/getLogs` also accepts the same credentials as `?token=` and `?timestamp=` query parameters (headers still take precedence when both are present). Only `admin` tokens are accepted; the timestamp is checked once at handshake time, so an accepted connection stays open past its tolerance window. Because the query string can leak into proxy and access logs, a token-carrying WebSocket URL should be treated as a secret.
+
 ### Endpoints
 
 | Endpoint | Method | Permission | Description |
@@ -300,13 +302,14 @@ Tokens idle for more than 1 hour are expired (cleaned every 10 minutes); any aut
 | `/api/settings/get` | GET | admin | `?type=global\|bot_user_config\|bot_node_config` | Returns `data: {config}`, where `config` is the selected config file's content (same layout as the JSON file). |
 | `/api/settings/set` | POST | admin | `?type=<same types>` + JSON body of partial updates, e.g. `{"system":{"debugMode":true}}` | Deep-merges the body into the selected config file, persists it, and reloads it in memory. Only the given keys change; arrays replace. |
 | `/health` | GET | None | Health check. Returns `data: {status, database}`. |
-| `/api/system/getLogs` | WebSocket | None | Streams logs. Sends the last 100 buffered entries, then live `{level, content, timestamp}` events. |
+| `/api/system/getLogs` | WebSocket | admin | Streams logs. Sends the last 100 buffered entries, then live `{level, content, timestamp}` events. Credentials via `X-Token`/`X-Timestamp` headers or `?token=`/`?timestamp=` query parameters; a failed check answers with the JSON error and no upgrade. |
 
 Middleware applied to the whole server:
 
 - **Rate limit** — token bucket, 100 requests/minute per client IP.
 - **CORS** — `Access-Control-Allow-Origin: *`, allows `Content-Type`, `X-Token`, `X-Timestamp`, `Authorization`.
 - **Security headers** — `X-XSS-Protection`, `X-Content-Type-Options: nosniff`, `X-Frame-Options: DENY`, `Referrer-Policy`, a restrictive CSP.
+- **WebSocket auth** — `utils.WebSocketAuthMiddleware` is attached to `/api/system/getLogs` (route-level, not global): it authenticates the upgrade request and requires an `admin` token before the connection is handed to the log handler.
 
 ## Bots
 
diff --git a/frontend/README.md b/frontend/README.md
index 91aa240..513277e 100644
--- a/frontend/README.md
+++ b/frontend/README.md
@@ -58,5 +58,5 @@ frontend/
 
 ## Caveats
 
-- The backend's `/api/system/getLogs` websocket is currently unauthenticated — anyone who can reach the port can read logs. Consider gating it in a future backend change.
+- The log websocket needs an `admin` token, and a browser cannot set headers on a WebSocket handshake, so the token rides in the query string (`/api/system/getLogs?token=…&timestamp=…`). That URL is a credential: it can end up in proxy and access logs, so don't paste it into third-party tools. The view reconnects with a fresh token from `localStorage` on every attempt.
 - Registering more than one user is intentionally impossible; the backend only accepts the very first registration.
diff --git a/frontend/src/views/Logs.vue b/frontend/src/views/Logs.vue
index 1dfca51..5405622 100644
--- a/frontend/src/views/Logs.vue
+++ b/frontend/src/views/Logs.vue
@@ -1,5 +1,6 @@
 <script setup>
 import { computed, onBeforeUnmount, onMounted, reactive, ref } from 'vue';
+import { getToken } from '../utils/auth.js';
 import { LOG_LEVELS } from '../utils/fmt.js';
 
 const MAX_LOGS = 1200;
@@ -38,7 +39,14 @@ const statusText = computed(() => {
 
 function wsUrl() {
     const proto = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
-    return `${proto}//${window.location.host}/api/system/getLogs`;
+    // The backend only upgrades the request for an admin token. A browser cannot
+    // set headers on a WebSocket handshake, so the credentials travel in the
+    // query string — the URL itself is therefore a secret.
+    const params = new URLSearchParams({
+        token: getToken(),
+        timestamp: String(Math.floor(Date.now() / 1000))
+    });
+    return `${proto}//${window.location.host}/api/system/getLogs?${params}`;
 }
 
 function connect() {
@@ -49,6 +57,11 @@ function connect() {
     if (ws) {
         try { ws.close(); } catch { /* ignore */ }
     }
+    // Signed out: the handshake would be rejected, so don't spin on reconnects.
+    if (!getToken()) {
+        status.value = 'closed';
+        return;
+    }
     status.value = 'connecting';
 
     try {
diff --git a/router.go b/router.go
index 8603550..4195779 100644
--- a/router.go
+++ b/router.go
@@ -6,6 +6,7 @@ import (
 
 	"nukumizu-backend/handler"
 	"nukumizu-backend/postLog"
+	"nukumizu-backend/utils"
 	"nukumizu-backend/web"
 )
 
@@ -32,11 +33,14 @@ func SetupRouter() *http.ServeMux {
 	// Health check endpoint.
 	mux.HandleFunc("/health", handler.HealthHandler)
 
-	// WebSocket log streaming endpoint.
+	// WebSocket log streaming endpoint (admin only). The middleware authenticates
+	// the upgrade request, so an anonymous or non-admin client is rejected before
+	// any log entry leaves the server.
 	logBroadcaster := postLog.GetLogBroadcaster()
 	if logBroadcaster != nil {
 		logSocketHandler := postLog.NewLogSocketHandler(logBroadcaster)
-		mux.HandleFunc("/api/system/getLogs", logSocketHandler.Handle)
+		adminOnly := utils.AuthWS("admin")
+		mux.Handle("/api/system/getLogs", adminOnly(http.HandlerFunc(logSocketHandler.Handle)))
 	}
 
 	// Static file serving for the web frontend.
diff --git a/utils/auth.go b/utils/auth.go
index 91e4bfe..011a8d0 100644
--- a/utils/auth.go
+++ b/utils/auth.go
@@ -100,62 +100,58 @@ func GetUserLevelFromRequest(r *http.Request) string {
 	return tokenInfo.Level
 }
 
-// Auth is the central authentication and authorization function.
-// It validates the request method, X-Timestamp header (30min tolerance),
-// X-Token header, and permission level. Returns true if the request is authorized.
+// checkTimestamp validates a Unix timestamp in seconds against the server clock
+// (30 minute tolerance per agent.md). The check is skipped entirely in debug
+// mode. It returns 0 when the timestamp is acceptable, otherwise the HTTP status
+// and message to reject the request with.
+func checkTimestamp(timestamp string) (int, string) {
+	if config.IsDebugMode() {
+		return 0, ""
+	}
+
+	if timestamp == "" {
+		return http.StatusUnauthorized, "missing timestamp"
+	}
+
+	ts, err := strconv.ParseInt(timestamp, 10, 64)
+	if err != nil {
+		return http.StatusUnauthorized, "invalid timestamp"
+	}
+
+	now := time.Now().Unix()
+	diff := now - ts
+	if diff < 0 {
+		diff = -diff
+	}
+	if diff > 1800 {
+		return http.StatusUnauthorized, "request expired"
+	}
+
+	return 0, ""
+}
+
+// checkPermission validates a session token against the required permission
+// level and refreshes the token's idle timer on success.
 //
 // Permission levels: "None" (public, no token required), "bot", "admin".
 // When level is "bot", both "bot" and "admin" tokens are accepted.
 // When level is "admin", only "admin" tokens are accepted.
-func Auth(w http.ResponseWriter, r *http.Request, targetMethod string, targetLevel string) bool {
-	// Validate HTTP method.
-	if r.Method != targetMethod {
-		SendErrorResponse(w, http.StatusMethodNotAllowed, "method not allowed")
-		return false
-	}
-
-	// Validate X-Timestamp.
-	timestamp := r.Header.Get("X-Timestamp")
-	if !config.IsDebugMode() {
-		if timestamp == "" {
-			SendErrorResponse(w, http.StatusUnauthorized, "missing timestamp")
-			return false
-		}
-
-		ts, err := strconv.ParseInt(timestamp, 10, 64)
-		if err != nil {
-			SendErrorResponse(w, http.StatusUnauthorized, "invalid timestamp")
-			return false
-		}
-
-		now := time.Now().Unix()
-		diff := now - ts
-		if diff < 0 {
-			diff = -diff
-		}
-		// 30 minute tolerance per agent.md.
-		if diff > 1800 {
-			SendErrorResponse(w, http.StatusUnauthorized, "request expired")
-			return false
-		}
-	}
-
+//
+// It returns 0 when the token is authorized, otherwise the HTTP status and
+// message to reject the request with.
+func checkPermission(token string, targetLevel string) (int, string) {
 	// Public endpoints require no token.
 	if targetLevel == "None" {
-		return true
+		return 0, ""
 	}
 
-	// Validate X-Token.
-	token := r.Header.Get("X-Token")
 	if token == "" {
-		SendErrorResponse(w, http.StatusUnauthorized, "missing token")
-		return false
+		return http.StatusUnauthorized, "missing token"
 	}
 
 	tokenInfo, exists := GetTokenInfo(token)
 	if !exists {
-		SendErrorResponse(w, http.StatusUnauthorized, "invalid token")
-		return false
+		return http.StatusUnauthorized, "invalid token"
 	}
 
 	// Check permission level.
@@ -164,21 +160,102 @@ func Auth(w http.ResponseWriter, r *http.Request, targetMethod string, targetLev
 	switch targetLevel {
 	case "admin":
 		if tokenInfo.Level != "admin" {
-			SendErrorResponse(w, http.StatusForbidden, "permission denied")
-			return false
+			return http.StatusForbidden, "permission denied"
 		}
 	case "bot":
 		if tokenInfo.Level != "bot" && tokenInfo.Level != "admin" {
-			SendErrorResponse(w, http.StatusForbidden, "permission denied")
-			return false
+			return http.StatusForbidden, "permission denied"
 		}
 	}
 
 	// Refresh token last access time.
 	RefreshToken(token)
+	return 0, ""
+}
+
+// Auth is the central authentication and authorization function.
+// It validates the request method, X-Timestamp header (30min tolerance),
+// X-Token header, and permission level. Returns true if the request is authorized.
+//
+// Permission levels: "None" (public, no token required), "bot", "admin".
+// When level is "bot", both "bot" and "admin" tokens are accepted.
+// When level is "admin", only "admin" tokens are accepted.
+//
+// WebSocket upgrades cannot carry custom headers from a browser; those endpoints
+// use WebSocketAuthMiddleware instead, which also accepts the credentials as
+// query parameters.
+func Auth(w http.ResponseWriter, r *http.Request, targetMethod string, targetLevel string) bool {
+	// Validate HTTP method.
+	if r.Method != targetMethod {
+		SendErrorResponse(w, http.StatusMethodNotAllowed, "method not allowed")
+		return false
+	}
+
+	// Validate X-Timestamp.
+	if status, message := checkTimestamp(r.Header.Get("X-Timestamp")); status != 0 {
+		SendErrorResponse(w, status, message)
+		return false
+	}
+
+	// Validate X-Token and its permission level.
+	if status, message := checkPermission(r.Header.Get("X-Token"), targetLevel); status != 0 {
+		SendErrorResponse(w, status, message)
+		return false
+	}
+
 	return true
 }
 
+
+// AuthWS gates a WebSocket endpoint behind the given permission
+// level ("bot" or "admin"), authenticating the upgrade request before the
+// connection is handed to the handler. Unauthorized requests are answered with
+// the standard JSON error response and are never upgraded.
+//
+// A browser cannot set custom headers on a WebSocket handshake, so the session
+// token and timestamp are read from the X-Token / X-Timestamp headers when
+// present and otherwise from the "token" and "timestamp" query parameters:
+//
+//	ws://host/api/system/getLogs?token=<token>&timestamp=<unix seconds>
+//
+// The timestamp is only checked at handshake time, so a long-lived connection
+// stays open past its tolerance window. Because a query string commonly ends up
+// in proxy and access logs, a token-carrying URL should be treated as a secret.
+func AuthWS(targetLevel string) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// An upgrade request is always a GET.
+			if r.Method != http.MethodGet {
+				SendErrorResponse(w, http.StatusMethodNotAllowed, "method not allowed")
+				return
+			}
+
+			// Headers win over query parameters so programmatic clients can keep
+			// the credentials out of the URL.
+			token := r.Header.Get("X-Token")
+			timestamp := r.Header.Get("X-Timestamp")
+			query := r.URL.Query()
+			if token == "" {
+				token = query.Get("token")
+			}
+			if timestamp == "" {
+				timestamp = query.Get("timestamp")
+			}
+
+			if status, message := checkTimestamp(timestamp); status != 0 {
+				SendErrorResponse(w, status, message)
+				return
+			}
+			if status, message := checkPermission(token, targetLevel); status != 0 {
+				SendErrorResponse(w, status, message)
+				return
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}
+
 // CleanExpiredTokens removes tokens that have been idle for over 1 hour.
 func CleanExpiredTokens() {
 	tokenStoreLock.Lock()