From 14f8b9c967818dc1a84ada53a6c470f7ac296c68 Mon Sep 17 00:00:00 2001 From: j4ckzh0u Date: Tue, 24 Mar 2026 15:17:20 +0800 Subject: [PATCH] fix: address security and code quality issues in OpenAI channel - Make CORS policy configurable with AllowOrigins config - Default to localhost only for security - Refactor resolveFromModelCandidates to reduce complexity - Improve error handling consistency with proper logging - Fix translateConversation boundary case handling - Add comprehensive tests for CORS and message validation Fixes review comments from @yinwm and @alexhoshina --- pkg/agent/instance.go | 48 +--- pkg/agent/loop.go | 40 +-- pkg/agent/model_resolution.go | 46 ++++ pkg/channels/openai_api/openai_api.go | 271 ++++++++++++++++----- pkg/channels/openai_api/openai_api_test.go | 221 +++++++++++++++++ pkg/config/config.go | 7 +- pkg/config/defaults.go | 7 +- 7 files changed, 499 insertions(+), 141 deletions(-) create mode 100644 pkg/agent/model_resolution.go diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 1c3635322..54e5ec87c 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -154,47 +154,9 @@ func NewAgentInstance( Primary: model, Fallbacks: fallbacks, } - resolveFromModelList := func(raw string) (string, bool) { - ensureProtocol := func(model string) string { - model = strings.TrimSpace(model) - if model == "" { - return "" - } - if strings.Contains(model, "/") { - return model - } - return "openai/" + model - } - - raw = strings.TrimSpace(raw) - if raw == "" { - return "", false - } - - if cfg != nil { - if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" { - return ensureProtocol(mc.Model), true - } - - for i := range cfg.ModelList { - fullModel := strings.TrimSpace(cfg.ModelList[i].Model) - if fullModel == "" { - continue - } - if fullModel == raw { - return ensureProtocol(fullModel), true - } - _, modelID := providers.ExtractProtocol(fullModel) - if modelID == raw { - return ensureProtocol(fullModel), true - } - } - } - - return "", false - } - - candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList) + candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, func(raw string) (string, bool) { + return resolveFromModelList(cfg, raw) + }) // Model routing setup: pre-resolve light model candidates at creation time // to avoid repeated model_list lookups on every incoming message. @@ -202,7 +164,9 @@ func NewAgentInstance( var lightCandidates []providers.FallbackCandidate if rc := defaults.Routing; rc != nil && rc.Enabled && rc.LightModel != "" { lightModelCfg := providers.ModelConfig{Primary: rc.LightModel} - resolved := providers.ResolveCandidatesWithLookup(lightModelCfg, defaults.Provider, resolveFromModelList) + resolved := providers.ResolveCandidatesWithLookup(lightModelCfg, defaults.Provider, func(raw string) (string, bool) { + return resolveFromModelList(cfg, raw) + }) if len(resolved) > 0 { router = routing.New(routing.RouterConfig{ LightModel: rc.LightModel, diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index ee5fc30bf..944198ccc 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -1514,47 +1514,13 @@ func (al *AgentLoop) resolveRequestedModelCandidates( } cfg := al.GetConfig() - resolveFromModelList := func(raw string) (string, bool) { - ensureProtocol := func(model string) string { - if model == "" { - return "" - } - if strings.Contains(model, "/") { - return model - } - return "openai/" + model - } - - raw = strings.TrimSpace(raw) - if raw == "" || cfg == nil { - return "", false - } - - if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" { - return ensureProtocol(mc.Model), true - } - - for i := range cfg.ModelList { - fullModel := strings.TrimSpace(cfg.ModelList[i].Model) - if fullModel == "" { - continue - } - if fullModel == raw { - return ensureProtocol(fullModel), true - } - _, modelID := providers.ExtractProtocol(fullModel) - if modelID == raw { - return ensureProtocol(fullModel), true - } - } - - return "", false - } candidates := providers.ResolveCandidatesWithLookup( providers.ModelConfig{Primary: requestedModel}, al.cfg.Agents.Defaults.Provider, - resolveFromModelList, + func(raw string) (string, bool) { + return resolveFromModelList(cfg, raw) + }, ) if len(candidates) == 0 { return nil, "", fmt.Errorf("requested model %q not found in model_list", requestedModel) diff --git a/pkg/agent/model_resolution.go b/pkg/agent/model_resolution.go new file mode 100644 index 000000000..cf9d3c524 --- /dev/null +++ b/pkg/agent/model_resolution.go @@ -0,0 +1,46 @@ +package agent + +import ( + "strings" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers" +) + +func ensureProtocol(model string) string { + model = strings.TrimSpace(model) + if model == "" { + return "" + } + if strings.Contains(model, "/") { + return model + } + return "openai/" + model +} + +func resolveFromModelList(cfg *config.Config, raw string) (string, bool) { + raw = strings.TrimSpace(raw) + if raw == "" || cfg == nil { + return "", false + } + + if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" { + return ensureProtocol(mc.Model), true + } + + for i := range cfg.ModelList { + fullModel := strings.TrimSpace(cfg.ModelList[i].Model) + if fullModel == "" { + continue + } + if fullModel == raw { + return ensureProtocol(fullModel), true + } + _, modelID := providers.ExtractProtocol(fullModel) + if modelID == raw { + return ensureProtocol(fullModel), true + } + } + + return "", false +} diff --git a/pkg/channels/openai_api/openai_api.go b/pkg/channels/openai_api/openai_api.go index 4d7604e5c..1c2382b85 100644 --- a/pkg/channels/openai_api/openai_api.go +++ b/pkg/channels/openai_api/openai_api.go @@ -7,6 +7,7 @@ import ( "fmt" "net" "net/http" + "net/url" "sort" "strconv" "strings" @@ -28,6 +29,8 @@ const ( responseWaitTimeout = 5 * time.Minute ) +var defaultAllowedOrigins = []string{"localhost"} + type responseTask struct { ctx context.Context cancel context.CancelFunc @@ -67,16 +70,17 @@ type translatedConversation struct { type OpenAIAPIChannel struct { *channels.BaseChannel - config config.OpenAIAPIConfig - listenHost string - models []config.ModelConfig - messageBus *bus.MessageBus - server *http.Server - listener net.Listener - ctx context.Context - cancel context.CancelFunc - taskMu sync.RWMutex - tasks map[string]*responseTask + config config.OpenAIAPIConfig + allowedOrigins []string + listenHost string + models []config.ModelConfig + messageBus *bus.MessageBus + server *http.Server + listener net.Listener + ctx context.Context + cancel context.CancelFunc + taskMu sync.RWMutex + tasks map[string]*responseTask } func NewOpenAIAPIChannel(cfg *config.Config, messageBus *bus.MessageBus) (*OpenAIAPIChannel, error) { @@ -87,20 +91,24 @@ func NewOpenAIAPIChannel(cfg *config.Config, messageBus *bus.MessageBus) (*OpenA return nil, fmt.Errorf("openai_api api_key is required") } + channelConfig := cfg.Channels.OpenAIAPI + channelConfig.AllowOrigins = normalizeAllowedOrigins(channelConfig.AllowOrigins) + listenHost := strings.TrimSpace(cfg.Gateway.Host) if listenHost == "" { listenHost = "127.0.0.1" } - base := channels.NewBaseChannel("openai_api", cfg.Channels.OpenAIAPI, messageBus, nil) + base := channels.NewBaseChannel("openai_api", channelConfig, messageBus, nil) return &OpenAIAPIChannel{ - BaseChannel: base, - config: cfg.Channels.OpenAIAPI, - listenHost: listenHost, - models: append([]config.ModelConfig(nil), cfg.ModelList...), - messageBus: messageBus, - tasks: make(map[string]*responseTask), + BaseChannel: base, + config: channelConfig, + allowedOrigins: append([]string(nil), channelConfig.AllowOrigins...), + listenHost: listenHost, + models: append([]config.ModelConfig(nil), cfg.ModelList...), + messageBus: messageBus, + tasks: make(map[string]*responseTask), }, nil } @@ -197,19 +205,30 @@ func (c *OpenAIAPIChannel) Send(ctx context.Context, msg bus.OutboundMessage) er } func (c *OpenAIAPIChannel) handleOptions(w http.ResponseWriter, r *http.Request) { - setCORSHeaders(w) + if !c.applyCORSHeaders(w, r) { + return + } w.WriteHeader(http.StatusNoContent) } func (c *OpenAIAPIChannel) handleHealth(w http.ResponseWriter, r *http.Request) { + if !c.applyCORSHeaders(w, r) { + return + } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(map[string]any{"status": "ok"}) + if err := writeJSONResponse(w, http.StatusOK, map[string]any{"status": "ok"}); err != nil { + return + } } func (c *OpenAIAPIChannel) handleModels(w http.ResponseWriter, r *http.Request) { - setCORSHeaders(w) + if !c.applyCORSHeaders(w, r) { + return + } if !c.authenticate(r) { - writeOpenAIError(w, http.StatusUnauthorized, "Invalid API key", "invalid_request_error", "invalid_api_key") + if err := writeOpenAIError(w, http.StatusUnauthorized, "Invalid API key", "invalid_request_error", "invalid_api_key"); err != nil { + return + } return } @@ -241,42 +260,64 @@ func (c *OpenAIAPIChannel) handleModels(w http.ResponseWriter, r *http.Request) }) w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(map[string]any{ + if err := writeJSONResponse(w, http.StatusOK, map[string]any{ "object": "list", "data": items, - }) + }); err != nil { + return + } } func (c *OpenAIAPIChannel) handleChatCompletions(w http.ResponseWriter, r *http.Request) { - setCORSHeaders(w) + if !c.applyCORSHeaders(w, r) { + return + } if !c.authenticate(r) { - writeOpenAIError(w, http.StatusUnauthorized, "Invalid API key", "invalid_request_error", "invalid_api_key") + if err := writeOpenAIError(w, http.StatusUnauthorized, "Invalid API key", "invalid_request_error", "invalid_api_key"); err != nil { + return + } return } var req chatCompletionRequest decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, maxRequestBodySize)) if err := decoder.Decode(&req); err != nil { - writeOpenAIError(w, http.StatusBadRequest, "Invalid JSON request body", "invalid_request_error", "invalid_json") + logger.WarnCF("openai_api", "Invalid chat completion request body", map[string]any{ + "error": err.Error(), + }) + if err := writeOpenAIError(w, http.StatusBadRequest, "Invalid JSON request body", "invalid_request_error", "invalid_json"); err != nil { + return + } return } if strings.TrimSpace(req.Model) == "" { - writeOpenAIError(w, http.StatusBadRequest, "model is required", "invalid_request_error", "missing_model") + if err := writeOpenAIError(w, http.StatusBadRequest, "model is required", "invalid_request_error", "missing_model"); err != nil { + return + } return } if !c.supportsModel(req.Model) { - writeOpenAIError(w, http.StatusBadRequest, fmt.Sprintf("model %q is not configured", req.Model), "invalid_request_error", "model_not_found") + if err := writeOpenAIError(w, http.StatusBadRequest, fmt.Sprintf("model %q is not configured", req.Model), "invalid_request_error", "model_not_found"); err != nil { + return + } return } if len(req.Messages) == 0 { - writeOpenAIError(w, http.StatusBadRequest, "messages must not be empty", "invalid_request_error", "missing_messages") + if err := writeOpenAIError(w, http.StatusBadRequest, "messages must not be empty", "invalid_request_error", "missing_messages"); err != nil { + return + } return } translated, err := translateConversation(req.Messages) if err != nil { - writeOpenAIError(w, http.StatusBadRequest, err.Error(), "invalid_request_error", "invalid_messages") + logger.WarnCF("openai_api", "Invalid chat completion message sequence", map[string]any{ + "error": err.Error(), + }) + if err := writeOpenAIError(w, http.StatusBadRequest, err.Error(), "invalid_request_error", "invalid_messages"); err != nil { + return + } return } @@ -302,7 +343,12 @@ func (c *OpenAIAPIChannel) handleChatCompletions(w http.ResponseWriter, r *http. if len(translated.InjectedHistory) > 0 { rawHistory, err := json.Marshal(translated.InjectedHistory) if err != nil { - writeOpenAIError(w, http.StatusInternalServerError, "Failed to encode conversation history", "server_error", "history_encode_failed") + logger.ErrorCF("openai_api", "Failed to encode injected history", map[string]any{ + "error": err.Error(), + }) + if err := writeOpenAIError(w, http.StatusInternalServerError, "Failed to encode conversation history", "server_error", "history_encode_failed"); err != nil { + return + } return } metadata["injected_history"] = string(rawHistory) @@ -327,13 +373,23 @@ func (c *OpenAIAPIChannel) handleChatCompletions(w http.ResponseWriter, r *http. Peer: bus.Peer{Kind: "direct", ID: senderID}, Metadata: metadata, }); err != nil { - writeOpenAIError(w, http.StatusBadGateway, fmt.Sprintf("Failed to submit request: %v", err), "server_error", "publish_failed") + logger.ErrorCF("openai_api", "Failed to publish inbound OpenAI API request", map[string]any{ + "error": err.Error(), + }) + if err := writeOpenAIError(w, http.StatusBadGateway, fmt.Sprintf("Failed to submit request: %v", err), "server_error", "publish_failed"); err != nil { + return + } return } firstChunk, err := waitForFirstChunk(reqCtx, task) if err != nil { - writeOpenAIError(w, http.StatusGatewayTimeout, "Timed out waiting for assistant response", "server_error", "response_timeout") + logger.ErrorCF("openai_api", "Timed out waiting for first assistant chunk", map[string]any{ + "error": err.Error(), + }) + if err := writeOpenAIError(w, http.StatusGatewayTimeout, "Timed out waiting for assistant response", "server_error", "response_timeout"); err != nil { + return + } return } @@ -343,7 +399,9 @@ func (c *OpenAIAPIChannel) handleChatCompletions(w http.ResponseWriter, r *http. if req.Stream { flusher, ok := w.(http.Flusher) if !ok { - writeOpenAIError(w, http.StatusInternalServerError, "Streaming is not supported by this server", "server_error", "stream_not_supported") + if err := writeOpenAIError(w, http.StatusInternalServerError, "Streaming is not supported by this server", "server_error", "stream_not_supported"); err != nil { + return + } return } @@ -352,28 +410,49 @@ func (c *OpenAIAPIChannel) handleChatCompletions(w http.ResponseWriter, r *http. w.Header().Set("Connection", "keep-alive") if err := writeChatCompletionChunk(w, completionID, createdAt, req.Model, firstChunk, true, false); err != nil { + logger.ErrorCF("openai_api", "Failed to write first streaming chunk", map[string]any{ + "error": err.Error(), + }) return } flusher.Flush() if err := streamRemainingChunks(reqCtx, task, w, flusher, completionID, createdAt, req.Model); err != nil { + logger.ErrorCF("openai_api", "Failed to stream assistant chunks", map[string]any{ + "error": err.Error(), + }) return } - _ = writeChatCompletionChunk(w, completionID, createdAt, req.Model, "", false, true) - _, _ = fmt.Fprint(w, "data: [DONE]\n\n") + if err := writeChatCompletionChunk(w, completionID, createdAt, req.Model, "", false, true); err != nil { + logger.ErrorCF("openai_api", "Failed to write final streaming chunk", map[string]any{ + "error": err.Error(), + }) + return + } + if _, err := fmt.Fprint(w, "data: [DONE]\n\n"); err != nil { + logger.ErrorCF("openai_api", "Failed to write stream terminator", map[string]any{ + "error": err.Error(), + }) + return + } flusher.Flush() return } chunks, err := collectRemainingChunks(reqCtx, task, []string{firstChunk}) if err != nil { - writeOpenAIError(w, http.StatusGatewayTimeout, "Timed out waiting for assistant response", "server_error", "response_timeout") + logger.ErrorCF("openai_api", "Timed out collecting assistant response chunks", map[string]any{ + "error": err.Error(), + }) + if err := writeOpenAIError(w, http.StatusGatewayTimeout, "Timed out waiting for assistant response", "server_error", "response_timeout"); err != nil { + return + } return } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(map[string]any{ + if err := writeJSONResponse(w, http.StatusOK, map[string]any{ "id": completionID, "object": "chat.completion", "created": createdAt, @@ -388,7 +467,9 @@ func (c *OpenAIAPIChannel) handleChatCompletions(w http.ResponseWriter, r *http. "finish_reason": "stop", }, }, - }) + }); err != nil { + return + } } func translateConversation(messages []chatCompletionMessage) (translatedConversation, error) { @@ -432,21 +513,17 @@ func translateConversation(messages []chatCompletionMessage) (translatedConversa out.ExtraSystemPrompt = strings.Join(systemPrompts, "\n\n") last := nonSystem[len(nonSystem)-1] - if last.Role == "user" && strings.TrimSpace(last.Content) != "" { - out.CurrentMessage = last.Content - out.InjectedHistory = append([]providers.Message(nil), nonSystem[:len(nonSystem)-1]...) - return out, nil - } - - out.InjectedHistory = append([]providers.Message(nil), nonSystem...) switch last.Role { case "assistant": - out.CurrentMessage = "Continue the conversation with the next assistant response." + return translatedConversation{}, fmt.Errorf("last message must be a user message, got assistant") case "tool": - out.CurrentMessage = "Continue the conversation after the tool result above." - default: - out.CurrentMessage = "Continue the conversation based on the previous messages." + return translatedConversation{}, fmt.Errorf("last message must be a user message, got tool") } + if strings.TrimSpace(last.Content) == "" { + return translatedConversation{}, fmt.Errorf("last user message must not be empty") + } + out.CurrentMessage = last.Content + out.InjectedHistory = append([]providers.Message(nil), nonSystem[:len(nonSystem)-1]...) return out, nil } @@ -675,11 +752,9 @@ func writeChatCompletionChunk( return err } -func writeOpenAIError(w http.ResponseWriter, status int, message, errorType, code string) { - setCORSHeaders(w) +func writeOpenAIError(w http.ResponseWriter, status int, message, errorType, code string) error { w.Header().Set("Content-Type", "application/json") - w.WriteHeader(status) - _ = json.NewEncoder(w).Encode(map[string]any{ + return writeJSONResponse(w, status, map[string]any{ "error": map[string]any{ "message": message, "type": errorType, @@ -688,10 +763,94 @@ func writeOpenAIError(w http.ResponseWriter, status int, message, errorType, cod }) } -func setCORSHeaders(w http.ResponseWriter) { - w.Header().Set("Access-Control-Allow-Origin", "*") +func writeJSONResponse(w http.ResponseWriter, status int, payload any) error { + w.WriteHeader(status) + if err := json.NewEncoder(w).Encode(payload); err != nil { + logger.ErrorCF("openai_api", "Failed to write JSON response", map[string]any{ + "error": err.Error(), + }) + return err + } + return nil +} + +func normalizeAllowedOrigins(origins []string) []string { + normalized := make([]string, 0, len(origins)) + for _, origin := range origins { + origin = strings.TrimSpace(origin) + if origin == "" { + continue + } + normalized = append(normalized, origin) + } + if len(normalized) == 0 { + return append([]string(nil), defaultAllowedOrigins...) + } + return normalized +} + +func (c *OpenAIAPIChannel) applyCORSHeaders(w http.ResponseWriter, r *http.Request) bool { w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type") w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS") + origin := strings.TrimSpace(r.Header.Get("Origin")) + if origin == "" { + return true + } + if !originAllowed(c.allowedOrigins, origin) { + logger.WarnCF("openai_api", "Rejected request from disallowed origin", map[string]any{ + "origin": origin, + }) + if err := writeOpenAIError(w, http.StatusForbidden, "Origin is not allowed", "invalid_request_error", "origin_not_allowed"); err != nil { + return false + } + return false + } + w.Header().Set("Access-Control-Allow-Origin", origin) + addVaryHeader(w.Header(), "Origin") + return true +} + +func originAllowed(allowedOrigins []string, origin string) bool { + parsedOrigin, err := url.Parse(origin) + if err != nil || parsedOrigin.Scheme == "" || parsedOrigin.Host == "" { + return false + } + + requestHost := strings.ToLower(parsedOrigin.Hostname()) + for _, allowed := range allowedOrigins { + allowed = strings.TrimSpace(allowed) + if allowed == "" { + continue + } + + if strings.Contains(allowed, "://") { + parsedAllowed, err := url.Parse(allowed) + if err != nil || parsedAllowed.Scheme == "" || parsedAllowed.Host == "" { + continue + } + if strings.EqualFold(parsedAllowed.Scheme, parsedOrigin.Scheme) && strings.EqualFold(parsedAllowed.Host, parsedOrigin.Host) { + return true + } + continue + } + + if strings.EqualFold(allowed, requestHost) { + return true + } + } + + return false +} + +func addVaryHeader(headers http.Header, value string) { + for _, existing := range headers.Values("Vary") { + for _, part := range strings.Split(existing, ",") { + if strings.EqualFold(strings.TrimSpace(part), value) { + return + } + } + } + headers.Add("Vary", value) } func (c *OpenAIAPIChannel) authenticate(r *http.Request) bool { diff --git a/pkg/channels/openai_api/openai_api_test.go b/pkg/channels/openai_api/openai_api_test.go index 122dfe85b..6a60a2b71 100644 --- a/pkg/channels/openai_api/openai_api_test.go +++ b/pkg/channels/openai_api/openai_api_test.go @@ -3,22 +3,33 @@ package openai_api import ( "context" "encoding/json" + "errors" "net/http" "net/http/httptest" + "reflect" "strings" "testing" "time" "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers" ) func newTestChannel(t *testing.T) (*OpenAIAPIChannel, *bus.MessageBus) { t.Helper() + return newTestChannelWithConfig(t, nil) +} + +func newTestChannelWithConfig(t *testing.T, mutate func(*config.Config)) (*OpenAIAPIChannel, *bus.MessageBus) { + t.Helper() cfg := config.DefaultConfig() cfg.Channels.OpenAIAPI.APIKey = "test-key" cfg.Channels.OpenAIAPI.Port = 0 + if mutate != nil { + mutate(cfg) + } messageBus := bus.NewMessageBus() channel, err := NewOpenAIAPIChannel(cfg, messageBus) @@ -30,6 +41,23 @@ func newTestChannel(t *testing.T) (*OpenAIAPIChannel, *bus.MessageBus) { return channel, messageBus } +type failingResponseWriter struct { + headers http.Header +} + +func (w *failingResponseWriter) Header() http.Header { + if w.headers == nil { + w.headers = make(http.Header) + } + return w.headers +} + +func (w *failingResponseWriter) WriteHeader(statusCode int) {} + +func (w *failingResponseWriter) Write(p []byte) (int, error) { + return 0, errors.New("write failed") +} + func TestHandleChatCompletions_NonStreaming(t *testing.T) { channel, messageBus := newTestChannel(t) @@ -164,3 +192,196 @@ func TestHandleModels_RequiresAuth(t *testing.T) { t.Fatalf("status = %d, want %d", rec.Code, http.StatusUnauthorized) } } + +func TestHandleModels_AllowsConfiguredOrigin(t *testing.T) { + channel, _ := newTestChannelWithConfig(t, func(cfg *config.Config) { + cfg.Channels.OpenAIAPI.AllowOrigins = []string{"https://console.example.com"} + }) + + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + req.Header.Set("Authorization", "Bearer test-key") + req.Header.Set("Origin", "https://console.example.com") + rec := httptest.NewRecorder() + + channel.handleModels(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if got := rec.Header().Get("Access-Control-Allow-Origin"); got != "https://console.example.com" { + t.Fatalf("Access-Control-Allow-Origin = %q, want %q", got, "https://console.example.com") + } +} + +func TestHandleModels_DefaultCORSAllowsLocalhost(t *testing.T) { + channel, _ := newTestChannelWithConfig(t, func(cfg *config.Config) { + cfg.Channels.OpenAIAPI.AllowOrigins = nil + }) + + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + req.Header.Set("Authorization", "Bearer test-key") + req.Header.Set("Origin", "http://localhost:3000") + rec := httptest.NewRecorder() + + channel.handleModels(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if got := rec.Header().Get("Access-Control-Allow-Origin"); got != "http://localhost:3000" { + t.Fatalf("Access-Control-Allow-Origin = %q, want %q", got, "http://localhost:3000") + } +} + +func TestHandleModels_RejectsDisallowedOrigin(t *testing.T) { + channel, _ := newTestChannelWithConfig(t, func(cfg *config.Config) { + cfg.Channels.OpenAIAPI.AllowOrigins = []string{"https://console.example.com"} + }) + + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + req.Header.Set("Authorization", "Bearer test-key") + req.Header.Set("Origin", "https://evil.example.com") + rec := httptest.NewRecorder() + + channel.handleModels(rec, req) + + if rec.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusForbidden, rec.Body.String()) + } + if got := rec.Header().Get("Access-Control-Allow-Origin"); got != "" { + t.Fatalf("Access-Control-Allow-Origin = %q, want empty", got) + } + if !strings.Contains(rec.Body.String(), "origin_not_allowed") { + t.Fatalf("response body missing origin_not_allowed: %s", rec.Body.String()) + } +} + +func TestHandleChatCompletions_RejectsInvalidLastMessage(t *testing.T) { + channel, _ := newTestChannel(t) + + testCases := []struct { + name string + body string + }{ + { + name: "assistant last", + body: `{ + "model": "gpt-5.4", + "messages": [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"} + ] + }`, + }, + { + name: "tool last", + body: `{ + "model": "gpt-5.4", + "messages": [ + {"role": "user", "content": "hello"}, + {"role": "tool", "content": "tool result", "tool_call_id": "call_1"} + ] + }`, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(tc.body)) + req.Header.Set("Authorization", "Bearer test-key") + rec := httptest.NewRecorder() + + channel.handleChatCompletions(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusBadRequest, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "invalid_messages") { + t.Fatalf("response body missing invalid_messages: %s", rec.Body.String()) + } + }) + } +} + +func TestTranslateConversation_RejectsInvalidLastMessage(t *testing.T) { + testCases := []struct { + name string + messages []chatCompletionMessage + wantErr string + }{ + { + name: "assistant last", + messages: []chatCompletionMessage{ + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "hi"}, + }, + wantErr: "got assistant", + }, + { + name: "tool last", + messages: []chatCompletionMessage{ + {Role: "user", Content: "hello"}, + {Role: "tool", Content: "tool result", ToolCallID: "call_1"}, + }, + wantErr: "got tool", + }, + { + name: "empty user last", + messages: []chatCompletionMessage{ + {Role: "user", Content: "hello"}, + {Role: "user", Content: " "}, + }, + wantErr: "must not be empty", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + _, err := translateConversation(tc.messages) + if err == nil { + t.Fatal("translateConversation() error = nil, want non-nil") + } + if !strings.Contains(err.Error(), tc.wantErr) { + t.Fatalf("translateConversation() error = %q, want substring %q", err.Error(), tc.wantErr) + } + }) + } +} + +func TestTranslateConversation_ExtractsHistoryAndCurrentUserMessage(t *testing.T) { + translated, err := translateConversation([]chatCompletionMessage{ + {Role: "system", Content: "be concise"}, + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "hi"}, + {Role: "user", Content: "what next?"}, + }) + if err != nil { + t.Fatalf("translateConversation() error = %v", err) + } + + if translated.CurrentMessage != "what next?" { + t.Fatalf("CurrentMessage = %q, want %q", translated.CurrentMessage, "what next?") + } + if translated.ExtraSystemPrompt != "be concise" { + t.Fatalf("ExtraSystemPrompt = %q, want %q", translated.ExtraSystemPrompt, "be concise") + } + wantHistory := []providers.Message{ + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "hi"}, + } + if len(translated.InjectedHistory) != len(wantHistory) { + t.Fatalf("len(InjectedHistory) = %d, want %d", len(translated.InjectedHistory), len(wantHistory)) + } + for i := range wantHistory { + if !reflect.DeepEqual(translated.InjectedHistory[i], wantHistory[i]) { + t.Fatalf("InjectedHistory[%d] = %+v, want %+v", i, translated.InjectedHistory[i], wantHistory[i]) + } + } +} + +func TestWriteOpenAIError_ReturnsWriteError(t *testing.T) { + err := writeOpenAIError(&failingResponseWriter{}, http.StatusBadRequest, "bad request", "invalid_request_error", "bad_request") + if err == nil { + t.Fatal("writeOpenAIError() error = nil, want non-nil") + } +} diff --git a/pkg/config/config.go b/pkg/config/config.go index f1cafb638..65e16aef0 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -475,9 +475,10 @@ type PicoConfig struct { } type OpenAIAPIConfig struct { - Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_OPENAI_API_ENABLED"` - Port int `json:"port" env:"PICOCLAW_CHANNELS_OPENAI_API_PORT"` - APIKey string `json:"api_key" env:"PICOCLAW_CHANNELS_OPENAI_API_API_KEY"` + Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_OPENAI_API_ENABLED"` + Port int `json:"port" env:"PICOCLAW_CHANNELS_OPENAI_API_PORT"` + APIKey string `json:"api_key" env:"PICOCLAW_CHANNELS_OPENAI_API_API_KEY"` + AllowOrigins []string `json:"allow_origins,omitempty"` } type IRCConfig struct { diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index bdd036555..72fd79ba2 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -176,9 +176,10 @@ func DefaultConfig() *Config { AllowFrom: FlexibleStringSlice{}, }, OpenAIAPI: OpenAIAPIConfig{ - Enabled: false, - Port: 18794, - APIKey: "", + Enabled: false, + Port: 18794, + APIKey: "", + AllowOrigins: []string{"localhost"}, }, }, Providers: ProvidersConfig{