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
This commit is contained in:
parent
332159841a
commit
14f8b9c967
7 changed files with 499 additions and 141 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
46
pkg/agent/model_resolution.go
Normal file
46
pkg/agent/model_resolution.go
Normal file
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
@ -68,6 +71,7 @@ type translatedConversation struct {
|
|||
type OpenAIAPIChannel struct {
|
||||
*channels.BaseChannel
|
||||
config config.OpenAIAPIConfig
|
||||
allowedOrigins []string
|
||||
listenHost string
|
||||
models []config.ModelConfig
|
||||
messageBus *bus.MessageBus
|
||||
|
|
@ -87,16 +91,20 @@ 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,
|
||||
config: channelConfig,
|
||||
allowedOrigins: append([]string(nil), channelConfig.AllowOrigins...),
|
||||
listenHost: listenHost,
|
||||
models: append([]config.ModelConfig(nil), cfg.ModelList...),
|
||||
messageBus: messageBus,
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -478,6 +478,7 @@ 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"`
|
||||
AllowOrigins []string `json:"allow_origins,omitempty"`
|
||||
}
|
||||
|
||||
type IRCConfig struct {
|
||||
|
|
|
|||
|
|
@ -179,6 +179,7 @@ func DefaultConfig() *Config {
|
|||
Enabled: false,
|
||||
Port: 18794,
|
||||
APIKey: "",
|
||||
AllowOrigins: []string{"localhost"},
|
||||
},
|
||||
},
|
||||
Providers: ProvidersConfig{
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue