Merge pull request router-for-me#161 from router-for-me/aistudio

Add websocket provider

Author: luispater

config.example.yaml

@@ -43,6 +43,9 @@ quota-exceeded:
   switch-project: true # Whether to automatically switch to another project when a quota is exceeded
   switch-preview-model: true # Whether to automatically switch to a preview model when a quota is exceeded
 
+# When true, enable authentication for the WebSocket API (/v1/ws).
+ws-auth: false
+
 # API keys for official Generative Language API
 #generative-language-api-key:
 #  - "AIzaSy...01"

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/gin-gonic/gin v1.10.1
 	github.com/go-git/go-git/v6 v6.0.0-20251009132922-75a182125145
 	github.com/google/uuid v1.6.0
+	github.com/gorilla/websocket v1.5.3
 	github.com/jackc/pgx/v5 v5.7.6
 	github.com/joho/godotenv v1.5.1
 	github.com/klauspost/compress v1.17.4

go.sum

@@ -66,6 +66,8 @@ github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
 github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
 github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
@@ -80,8 +82,6 @@ github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnr
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
 github.com/kevinburke/ssh_config v1.4.0 h1:6xxtP5bZ2E4NF5tuQulISpTO2z8XbtH8cg1PWkxoFkQ=
 github.com/kevinburke/ssh_config v1.4.0/go.mod h1:q2RIzfka+BXARoNexmF9gkxEX7DmvbW9P4hIVx2Kg4M=
-github.com/klauspost/compress v1.17.3 h1:qkRjuerhUU1EmXLYGkSH6EZL+vPSxIrYjLNAK4slzwA=
-github.com/klauspost/compress v1.17.3/go.mod h1:/dCuZOvVtNoHsyb+cuJD3itjs3NbnF6KH9zAO4BDxPM=
 github.com/klauspost/compress v1.17.4 h1:Ej5ixsIri7BrIjBkRZLTo6ghwrEtHFk7ijlczPW4fZ4=
 github.com/klauspost/compress v1.17.4/go.mod h1:/dCuZOvVtNoHsyb+cuJD3itjs3NbnF6KH9zAO4BDxPM=
 github.com/klauspost/cpuid/v2 v2.0.1/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=

internal/access/config_access/provider.go

@@ -57,10 +57,12 @@ func (p *provider) Authenticate(_ context.Context, r *http.Request) (*sdkaccess.
 	authHeaderGoogle := r.Header.Get("X-Goog-Api-Key")
 	authHeaderAnthropic := r.Header.Get("X-Api-Key")
 	queryKey := ""
+	queryAuthToken := ""
 	if r.URL != nil {
 		queryKey = r.URL.Query().Get("key")
+		queryAuthToken = r.URL.Query().Get("auth_token")
 	}
-	if authHeader == "" && authHeaderGoogle == "" && authHeaderAnthropic == "" && queryKey == "" {
+	if authHeader == "" && authHeaderGoogle == "" && authHeaderAnthropic == "" && queryKey == "" && queryAuthToken == "" {
 		return nil, sdkaccess.ErrNoCredentials
 	}
 
@@ -74,6 +76,7 @@ func (p *provider) Authenticate(_ context.Context, r *http.Request) (*sdkaccess.
 		{authHeaderGoogle, "x-goog-api-key"},
 		{authHeaderAnthropic, "x-api-key"},
 		{queryKey, "query-key"},
+		{queryAuthToken, "query-auth-token"},
 	}
 
 	for _, candidate := range candidates {

internal/api/middleware/request_logging.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
 )
 
 // RequestLoggingMiddleware creates a Gin middleware that logs HTTP requests and responses.
@@ -63,13 +64,11 @@ func RequestLoggingMiddleware(logger logging.RequestLogger) gin.HandlerFunc {
 // It captures the URL, method, headers, and body. The request body is read and then
 // restored so that it can be processed by subsequent handlers.
 func captureRequestInfo(c *gin.Context) (*RequestInfo, error) {
-	// Capture URL
-	url := c.Request.URL.String()
-	if c.Request.URL.Path != "" {
-		url = c.Request.URL.Path
-		if c.Request.URL.RawQuery != "" {
-			url += "?" + c.Request.URL.RawQuery
-		}
+	// Capture URL with sensitive query parameters masked
+	maskedQuery := util.MaskSensitiveQuery(c.Request.URL.RawQuery)
+	url := c.Request.URL.Path
+	if maskedQuery != "" {
+		url += "?" + maskedQuery
 	}
 
 	// Capture method

internal/api/server.go

@@ -13,6 +13,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -138,6 +139,12 @@ type Server struct {
 	// currentPath is the absolute path to the current working directory.
 	currentPath string
 
+	// wsRoutes tracks registered websocket upgrade paths.
+	wsRouteMu     sync.Mutex
+	wsRoutes      map[string]struct{}
+	wsAuthChanged func(bool, bool)
+	wsAuthEnabled atomic.Bool
+
 	// management handler
 	mgmt *managementHandlers.Handler
 
@@ -228,7 +235,9 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
 		configFilePath:      configFilePath,
 		currentPath:         wd,
 		envManagementSecret: envManagementSecret,
+		wsRoutes:            make(map[string]struct{}),
 	}
+	s.wsAuthEnabled.Store(cfg.WebsocketAuth)
 	// Save initial YAML snapshot
 	s.oldConfigYaml, _ = yaml.Marshal(cfg)
 	s.applyAccessConfig(nil, cfg)
@@ -371,6 +380,43 @@ func (s *Server) setupRoutes() {
 	// Management routes are registered lazily by registerManagementRoutes when a secret is configured.
 }
 
+// AttachWebsocketRoute registers a websocket upgrade handler on the primary Gin engine.
+// The handler is served as-is without additional middleware beyond the standard stack already configured.
+func (s *Server) AttachWebsocketRoute(path string, handler http.Handler) {
+	if s == nil || s.engine == nil || handler == nil {
+		return
+	}
+	trimmed := strings.TrimSpace(path)
+	if trimmed == "" {
+		trimmed = "/v1/ws"
+	}
+	if !strings.HasPrefix(trimmed, "/") {
+		trimmed = "/" + trimmed
+	}
+	s.wsRouteMu.Lock()
+	if _, exists := s.wsRoutes[trimmed]; exists {
+		s.wsRouteMu.Unlock()
+		return
+	}
+	s.wsRoutes[trimmed] = struct{}{}
+	s.wsRouteMu.Unlock()
+
+	authMiddleware := AuthMiddleware(s.accessManager)
+	conditionalAuth := func(c *gin.Context) {
+		if !s.wsAuthEnabled.Load() {
+			c.Next()
+			return
+		}
+		authMiddleware(c)
+	}
+	finalHandler := func(c *gin.Context) {
+		handler.ServeHTTP(c.Writer, c.Request)
+		c.Abort()
+	}
+
+	s.engine.GET(trimmed, conditionalAuth, finalHandler)
+}
+
 func (s *Server) registerManagementRoutes() {
 	if s == nil || s.engine == nil || s.mgmt == nil {
 		return
@@ -770,6 +816,10 @@ func (s *Server) UpdateClients(cfg *config.Config) {
 
 	s.applyAccessConfig(oldCfg, cfg)
 	s.cfg = cfg
+	s.wsAuthEnabled.Store(cfg.WebsocketAuth)
+	if oldCfg != nil && s.wsAuthChanged != nil && oldCfg.WebsocketAuth != cfg.WebsocketAuth {
+		s.wsAuthChanged(oldCfg.WebsocketAuth, cfg.WebsocketAuth)
+	}
 	managementasset.SetCurrentConfig(cfg)
 	// Save YAML snapshot for next comparison
 	s.oldConfigYaml, _ = yaml.Marshal(cfg)
@@ -810,6 +860,13 @@ func (s *Server) UpdateClients(cfg *config.Config) {
 	)
 }
 
+func (s *Server) SetWebsocketAuthChangeHandler(fn func(bool, bool)) {
+	if s == nil {
+		return
+	}
+	s.wsAuthChanged = fn
+}
+
 // (management handlers moved to internal/api/handlers/management)
 
 // AuthMiddleware returns a Gin middleware handler that authenticates requests

internal/config/config.go

@@ -40,6 +40,9 @@ type Config struct {
 	// QuotaExceeded defines the behavior when a quota is exceeded.
 	QuotaExceeded QuotaExceeded `yaml:"quota-exceeded" json:"quota-exceeded"`
 
+	// WebsocketAuth enables or disables authentication for the WebSocket API.
+	WebsocketAuth bool `yaml:"ws-auth" json:"ws-auth"`
+
 	// GlAPIKey is the API key for the generative language API.
 	GlAPIKey []string `yaml:"generative-language-api-key" json:"generative-language-api-key"`
 

internal/logging/gin_logger.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -23,7 +24,7 @@ func GinLogrusLogger() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		start := time.Now()
 		path := c.Request.URL.Path
-		raw := c.Request.URL.RawQuery
+		raw := util.MaskSensitiveQuery(c.Request.URL.RawQuery)
 
 		c.Next()