diff --git a/config/config.example.json b/config/config.example.json index 77a8c0683..b2af3d86c 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -139,6 +139,13 @@ "webhook_path": "/webhook/wecom-app", "allow_from": [], "reply_timeout": 5 + }, + "websocket": { + "enabled": false, + "host": "127.0.0.1", + "port": 8080, + "token": "YOUR_TOKEN", + "allow_from": [] } }, "providers": { diff --git a/pkg/channels/websocket/channel.go b/pkg/channels/websocket/channel.go index d44de4cd3..1c6288b4b 100644 --- a/pkg/channels/websocket/channel.go +++ b/pkg/channels/websocket/channel.go @@ -8,6 +8,7 @@ import ( _ "embed" "encoding/json" "fmt" + "net" "net/http" "strings" "sync" @@ -57,6 +58,7 @@ type Channel struct { bus *bus.MessageBus running bool allowList []string + token string // Authentication token (optional) server *http.Server upgrader websocket.Upgrader clients sync.Map // map[string]*clientConn @@ -83,10 +85,16 @@ func NewChannel(cfg config.WebSocketConfig, messageBus *bus.MessageBus) (*Channe cfg.Host = "0.0.0.0" } + // Validate allow_from configuration at startup + if err := validateAllowList(cfg.AllowFrom); err != nil { + return nil, fmt.Errorf("invalid allow_from configuration: %w", err) + } + return &Channel{ config: cfg, bus: messageBus, allowList: cfg.AllowFrom, + token: strings.TrimSpace(cfg.Token), running: false, upgrader: websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { @@ -100,6 +108,55 @@ func NewChannel(cfg config.WebSocketConfig, messageBus *bus.MessageBus) (*Channe }, nil } +// validateAllowList validates the allow_from configuration at startup +// This prevents silent failures during runtime when invalid rules are encountered +func validateAllowList(allowList []string) error { + if len(allowList) == 0 { + // Empty list is valid (allows all) + return nil + } + + var invalidRules []string + validRuleCount := 0 + + for _, allowed := range allowList { + allowed = strings.TrimSpace(allowed) + if allowed == "" { + continue + } + + // Check if it's a CIDR notation + if strings.Contains(allowed, "/") { + if _, _, err := net.ParseCIDR(allowed); err != nil { + invalidRules = append(invalidRules, fmt.Sprintf("%s (CIDR parse error: %v)", allowed, err)) + continue + } + validRuleCount++ + } else { + // Validate as plain IP address + // Strip zone identifier for validation + ipStr := allowed + if idx := strings.IndexByte(ipStr, '%'); idx != -1 { + ipStr = ipStr[:idx] + } + if net.ParseIP(ipStr) == nil { + invalidRules = append(invalidRules, fmt.Sprintf("%s (invalid IP address)", allowed)) + continue + } + validRuleCount++ + } + } + + // If there are invalid rules, return error with details + if len(invalidRules) > 0 { + return fmt.Errorf("found %d invalid rule(s) in allow_from: %s", + len(invalidRules), strings.Join(invalidRules, "; ")) + } + + // All rules are valid + return nil +} + // Name returns the channel name func (c *Channel) Name() string { return "websocket" @@ -112,19 +169,93 @@ func (c *Channel) IsRunning() bool { return c.running } -// IsAllowed checks if a sender ID is allowed to use this channel -func (c *Channel) IsAllowed(senderID string) bool { +// IsAllowed checks if a client IP is allowed to use this channel +// The clientID should be in the format "ip:port" (from r.RemoteAddr) +// Supports: +// - Exact IP match: "192.168.1.5" +// - CIDR notation: "192.168.1.0/24", "10.0.0.0/8" +// - IPv6: "::1", "fe80::/10" +func (c *Channel) IsAllowed(clientID string) bool { if len(c.allowList) == 0 { return true } + + // Extract IP address from "ip:port" format + clientIPStr := c.extractIP(clientID) + if clientIPStr == "" { + logger.WarnCF("websocket", "Failed to extract IP from client ID", map[string]any{ + "client_id": clientID, + }) + return false + } + + // Parse client IP (strip zone identifier for IPv6 link-local addresses) + clientIPStr = c.stripZone(clientIPStr) + clientIP := net.ParseIP(clientIPStr) + if clientIP == nil { + logger.WarnCF("websocket", "Invalid client IP address", map[string]any{ + "client_ip": clientIPStr, + }) + return false + } + + // Check against allow list for _, allowed := range c.allowList { - if strings.EqualFold(allowed, senderID) { - return true + // Try CIDR notation first + if strings.Contains(allowed, "/") { + _, ipNet, err := net.ParseCIDR(allowed) + if err != nil { + // This should never happen as we validate at startup + // But keep for safety + logger.ErrorCF("websocket", "Invalid CIDR in allow list (should be caught at startup)", map[string]any{ + "cidr": allowed, + "error": err.Error(), + }) + continue + } + if ipNet.Contains(clientIP) { + return true + } + } else { + // Exact IP match (strip zone for comparison) + allowedStripped := c.stripZone(allowed) + if strings.EqualFold(allowedStripped, clientIPStr) { + return true + } } } return false } +// extractIP extracts the IP address from "ip:port" format +// Handles: +// - "192.168.1.5:8080" -> "192.168.1.5" +// - "[::1]:8080" -> "::1" +// - "[::1]" -> "::1" (brackets without port) +// - "::1" -> "::1" (no brackets, no port) +func (c *Channel) extractIP(addr string) string { + host, _, err := net.SplitHostPort(addr) + if err != nil { + // Handle IPv6 addresses with brackets but no port: [::1] -> ::1 + if strings.HasPrefix(addr, "[") && strings.HasSuffix(addr, "]") { + return addr[1 : len(addr)-1] + } + // Otherwise assume it's already just an IP (no port) + return addr + } + return host +} + +// stripZone removes the zone identifier from IPv6 addresses +// For example: "fe80::1%eth0" -> "fe80::1" +// This is necessary because net.ParseIP doesn't handle zone identifiers +func (c *Channel) stripZone(ipStr string) string { + if idx := strings.IndexByte(ipStr, '%'); idx != -1 { + return ipStr[:idx] + } + return ipStr +} + // setRunning sets the running state func (c *Channel) setRunning(running bool) { c.clientsMu.Lock() @@ -133,11 +264,8 @@ func (c *Channel) setRunning(running bool) { } // HandleMessage processes an incoming message and publishes it to the bus +// Note: Authorization is already checked during WebSocket connection establishment func (c *Channel) HandleMessage(senderID, chatID, content string, media []string, metadata map[string]string) { - if !c.IsAllowed(senderID) { - return - } - // Build session key: channel:chatID sessionKey := fmt.Sprintf("%s:%s", c.Name(), chatID) @@ -313,8 +441,74 @@ func min(a, b int) int { return b } +// validateToken validates the authentication token from the request +// Supports multiple token sources: +// 1. Query parameter: ?token=xxx +// 2. Authorization header: Bearer xxx +// 3. Sec-WebSocket-Protocol header: token (WebSocket standard approach) +func (c *Channel) validateToken(r *http.Request) bool { + if c.token == "" { + return true // No token configured, allow + } + + // Method 1: Check query parameter + if tokenParam := r.URL.Query().Get("token"); tokenParam != "" { + return tokenParam == c.token + } + + // Method 2: Check Authorization header (Bearer token) + authHeader := r.Header.Get("Authorization") + if authHeader != "" { + if strings.HasPrefix(authHeader, "Bearer ") { + token := strings.TrimPrefix(authHeader, "Bearer ") + return token == c.token + } + } + + // Method 3: Check Sec-WebSocket-Protocol header + // Client can send: Sec-WebSocket-Protocol: token + // This is a standard WebSocket approach for authentication + protocol := r.Header.Get("Sec-WebSocket-Protocol") + if protocol != "" { + // Support both "token" format and "bearer-" format + if protocol == c.token { + return true + } + if strings.HasPrefix(protocol, "bearer-") { + token := strings.TrimPrefix(protocol, "bearer-") + return token == c.token + } + } + + return false +} + // handleWebSocket handles WebSocket connection upgrades func (c *Channel) handleWebSocket(w http.ResponseWriter, r *http.Request) { + // Check token authentication if configured + if c.token != "" { + if !c.validateToken(r) { + logger.WarnCF("websocket", "Unauthorized connection attempt - invalid token", map[string]any{ + "remote_addr": r.RemoteAddr, + "user_agent": r.UserAgent(), + }) + http.Error(w, "Unauthorized: invalid or missing token", http.StatusUnauthorized) + return + } + } + + // Check IP authorization + clientID := r.RemoteAddr + if !c.IsAllowed(clientID) { + clientIP := c.extractIP(clientID) + logger.WarnCF("websocket", "Unauthorized connection attempt - IP not allowed", map[string]any{ + "client_ip": clientIP, + "client_id": clientID, + }) + http.Error(w, "Forbidden: IP not allowed", http.StatusForbidden) + return + } + conn, err := c.upgrader.Upgrade(w, r, nil) if err != nil { logger.ErrorCF("websocket", "Failed to upgrade connection", map[string]any{ @@ -323,13 +517,12 @@ func (c *Channel) handleWebSocket(w http.ResponseWriter, r *http.Request) { return } - // Generate client ID from remote address - clientID := r.RemoteAddr client := &clientConn{conn: conn} c.clients.Store(clientID, client) logger.InfoCF("websocket", "New client connected", map[string]any{ "client_id": clientID, + "client_ip": c.extractIP(clientID), }) // Handle client messages @@ -407,14 +600,6 @@ func (c *Channel) handleClient(clientID string, client *clientConn) { // Process chat messages if wsMsg.Type == "chat" { - // Check allowlist - if !c.IsAllowed(clientID) { - logger.WarnCF("websocket", "Unauthorized client", map[string]any{ - "client_id": clientID, - }) - continue - } - logger.DebugCF("websocket", "Received message", map[string]any{ "client_id": clientID, "content": wsMsg.Content, diff --git a/pkg/channels/websocket/chat.html b/pkg/channels/websocket/chat.html index 06e0c6104..507d4c76b 100644 --- a/pkg/channels/websocket/chat.html +++ b/pkg/channels/websocket/chat.html @@ -1004,6 +1004,8 @@ let reconnectCount = 0; let reconnectMessageElement = null; + let connectionAttempted = false; // Track if we've tried to connect + let hasEverConnected = false; // Track if we've successfully connected before let lastMessageTime = null; // Scroll to bottom helper @@ -1072,8 +1074,28 @@ } function connect() { + // Get token from multiple sources (priority order): + // 1. URL parameter: ?token=xxx + // 2. localStorage + // 3. Prompt user + let token = new URLSearchParams(window.location.search).get('token'); + + if (!token) { + token = localStorage.getItem('picoclaw-token'); + } + + // If still no token and this is first connection attempt, check if token might be required + // by attempting connection. Server will return 401 if token is needed. + const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; - const wsUrl = protocol + '//' + window.location.host + '/ws'; + let wsUrl = protocol + '//' + window.location.host + '/ws'; + + // Add token to URL if available + if (token) { + wsUrl += '?token=' + encodeURIComponent(token); + // Save token for future use + localStorage.setItem('picoclaw-token', token); + } // Close existing connection if any if (ws && ws.readyState !== WebSocket.CLOSED) { @@ -1119,9 +1141,109 @@ }; ws.onclose = (event) => { - console.log('Disconnected from PicoClaw', event.code, event.reason); + console.log('Disconnected from PicoClaw', event.code, event.reason, 'hasEverConnected:', hasEverConnected, 'connectionAttempted:', connectionAttempted, 'hasToken:', !!localStorage.getItem('picoclaw-token')); - // Increment reconnect count + // Scenario 1: First connection attempt failed without token + // This likely means server requires token (close code 1006 = abnormal closure from HTTP error) + if (!hasEverConnected && connectionAttempted && !localStorage.getItem('picoclaw-token') && event.code === 1006) { + const t = currentLang === 'zh' + ? { + authRequired: '连接失败', + authPrompt: '服务器可能需要访问令牌（Token）\n\n如果已配置 token，请输入；\n如果未配置 token，请留空并点击「确定」', + retry: '重试中...', + noAuth: '正在重试连接...' + } + : { + authRequired: 'Connection Failed', + authPrompt: 'Server may require access token\n\nIf token is configured, enter it;\nIf not configured, leave empty and click OK', + retry: 'Retrying...', + noAuth: 'Retrying connection...' + }; + + statusDiv.innerHTML = '
' + t.authRequired + ''; + statusDiv.className = 'status disconnected'; + messageInput.disabled = true; + sendButton.disabled = true; + + // Prompt for token (user can leave empty if server has no token) + setTimeout(() => { + const token = prompt(t.authPrompt); + // User clicked OK (with or without token) + if (token !== null) { + if (token.trim()) { + // User entered a token, save and retry + localStorage.setItem('picoclaw-token', token.trim()); + statusDiv.innerHTML = '
' + t.retry + ''; + addMessage('system', t.retry, 'system'); + setTimeout(connect, 500); + } else { + // User left empty and clicked OK - server might not require token + addMessage('system', t.noAuth, 'system'); + setTimeout(connect, 2000); + } + } else { + // User clicked Cancel - don't retry + const cancelMsg = currentLang === 'zh' ? '已取消连接' : 'Connection cancelled'; + addMessage('system', cancelMsg, 'system'); + addRetryButton(); + } + }, 100); + return; + } + + // Scenario 2: Connection failed with token present + // This means the token is invalid (close code 1006 from HTTP 401) + if (event.code === 1006 && localStorage.getItem('picoclaw-token') && !hasEverConnected) { + localStorage.removeItem('picoclaw-token'); + + const t = currentLang === 'zh' + ? { + authRequired: '认证失败', + authPrompt: 'Token 无效或已过期，请重新输入：', + authFailed: '认证失败，请检查令牌是否正确', + tokenRequired: '需要有效的访问令牌才能连接', + retry: '重试中...', + cancelled: '已取消连接' + } + : { + authRequired: 'Authentication Failed', + authPrompt: 'Token invalid or expired, please enter again:', + authFailed: 'Authentication failed, please check your token', + tokenRequired: 'Valid access token required to connect', + retry: 'Retrying...', + cancelled: 'Connection cancelled' + }; + + statusDiv.innerHTML = '
' + t.authRequired + ''; + statusDiv.className = 'status disconnected'; + messageInput.disabled = true; + sendButton.disabled = true; + + // Show auth failure message + addMessage('system', t.authFailed, 'system'); + + // Prompt for token with retry support + setTimeout(() => { + const token = prompt(t.authPrompt); + if (token && token.trim()) { + localStorage.setItem('picoclaw-token', token.trim()); + statusDiv.innerHTML = '
' + t.retry + ''; + addMessage('system', t.retry, 'system'); + setTimeout(connect, 500); + } else { + // User cancelled + addMessage('system', t.cancelled + ' - ' + t.tokenRequired, 'system'); + statusDiv.innerHTML = '
' + t.cancelled + ''; + messageInput.disabled = true; + sendButton.disabled = true; + addRetryButton(); + } + }, 100); + return; + } + + // Scenario 3: Normal disconnection after successful connection + // Increment reconnect count and try to reconnect reconnectCount++; const t = currentLang === 'zh' @@ -1151,7 +1273,9 @@ }; ws.onerror = (error) => { - console.error('WebSocket error:', error); + console.error('WebSocket error:', error, 'hasEverConnected:', hasEverConnected, 'hasToken:', !!localStorage.getItem('picoclaw-token')); + // Mark that connection has been attempted + connectionAttempted = true; }; } @@ -1219,6 +1343,70 @@ } } + // Add a retry button for reconnection with token + function addRetryButton() { + const t = currentLang === 'zh' + ? { retryButton: '点击重新连接', clickToRetry: '点击下方按钮重新输入 Token 并连接' } + : { retryButton: 'Retry Connection', clickToRetry: 'Click button below to re-enter token and connect' }; + + // Create a message with a retry button + const messageDiv = document.createElement('div'); + messageDiv.className = 'message system'; + messageDiv.dataset.timestamp = Date.now(); + + const contentDiv = document.createElement('div'); + contentDiv.className = 'message-content'; + contentDiv.style.cursor = 'pointer'; + contentDiv.style.userSelect = 'none'; + + const rawHtml = marked.parse( + t.clickToRetry + '\n\n**[' + t.retryButton + ']**' + ); + const cleanHtml = DOMPurify.sanitize(rawHtml); + contentDiv.innerHTML = cleanHtml; + + messageDiv.appendChild(contentDiv); + messagesDiv.insertBefore(messageDiv, typingIndicator); + + // Make it clickable + contentDiv.addEventListener('click', () => { + // Remove the retry button message + messageDiv.remove(); + // Trigger reconnection with token prompt + promptForToken(); + }); + + scrollToBottom(); + } + + // Prompt user for token and attempt reconnection + function promptForToken() { + const t = currentLang === 'zh' + ? { + authPrompt: '请输入访问令牌（Token）：', + retry: '重试中...', + cancelled: '已取消连接', + tokenRequired: '需要有效的访问令牌才能连接' + } + : { + authPrompt: 'Please enter access token:', + retry: 'Retrying...', + cancelled: 'Connection cancelled', + tokenRequired: 'Valid access token required to connect' + }; + + const token = prompt(t.authPrompt); + if (token && token.trim()) { + localStorage.setItem('picoclaw-token', token.trim()); + statusDiv.innerHTML = '
' + t.retry + ''; + addMessage('system', t.retry, 'system'); + setTimeout(connect, 500); + } else { + addMessage('system', t.cancelled + ' - ' + t.tokenRequired, 'system'); + addRetryButton(); + } + } + function sendMessage(content) { if (!content.trim() || !ws || ws.readyState !== WebSocket.OPEN) { return; diff --git a/pkg/config/config.go b/pkg/config/config.go index 0bb98abdf..f7053cd71 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -301,6 +301,7 @@ type WebSocketConfig struct { Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WEBSOCKET_ENABLED"` Host string `json:"host" env:"PICOCLAW_CHANNELS_WEBSOCKET_HOST"` Port int `json:"port" env:"PICOCLAW_CHANNELS_WEBSOCKET_PORT"` + Token string `json:"token" env:"PICOCLAW_CHANNELS_WEBSOCKET_TOKEN"` AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WEBSOCKET_ALLOW_FROM"` } diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index 5b906b0c3..65bb0f1de 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -115,8 +115,9 @@ func DefaultConfig() *Config { }, WebSocket: WebSocketConfig{ Enabled: false, - Host: "0.0.0.0", + Host: "127.0.0.1", Port: 8080, + Token: "YOUR_TOKEN", AllowFrom: FlexibleStringSlice{}, }, },