fix: Phase 1 - security & stability patches
- AuthMiddleware now requires auth on /v1/* routes (returns 401) - WebSocket origin check configurable via WSAllowedOrigin - Removed debug fmt.Printf leaks (config, ollama, server) - Registry access protected by sync.RWMutex (race condition fix) - Session cleanup goroutine runs every 15 min - RevokeSession returns error instead of silent no-op
This commit is contained in:
@@ -9,9 +9,9 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/go-resty/resty/v2"
|
||||
"gophergate/internal/config"
|
||||
"gophergate/internal/models"
|
||||
"github.com/go-resty/resty/v2"
|
||||
)
|
||||
|
||||
type OllamaProvider struct {
|
||||
@@ -26,7 +26,7 @@ func NewOllamaProvider(cfg config.OllamaConfig) *OllamaProvider {
|
||||
client.SetTimeout(15 * time.Minute)
|
||||
client.SetRetryCount(2)
|
||||
client.SetRetryWaitTime(1 * time.Second)
|
||||
|
||||
|
||||
return &OllamaProvider{
|
||||
client: client,
|
||||
config: cfg,
|
||||
@@ -46,9 +46,6 @@ func (p *OllamaProvider) ChatCompletion(ctx context.Context, req *models.Unified
|
||||
body := BuildOllamaBody(req, messagesJSON, false)
|
||||
url := fmt.Sprintf("%s/chat/completions", p.config.BaseURL)
|
||||
|
||||
// Log request for debugging
|
||||
fmt.Printf("[Ollama] Request to %s with model %s\n", url, req.Model)
|
||||
|
||||
resp, err := p.client.R().
|
||||
SetContext(ctx).
|
||||
SetBody(body).
|
||||
@@ -70,7 +67,6 @@ func (p *OllamaProvider) ChatCompletion(ctx context.Context, req *models.Unified
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
|
||||
fmt.Printf("[Ollama] Success response for model %s\n", req.Model)
|
||||
return ParseOllamaResponse(respJSON, req.Model)
|
||||
}
|
||||
|
||||
@@ -97,7 +93,7 @@ func (p *OllamaProvider) ChatCompletionStream(ctx context.Context, req *models.U
|
||||
}
|
||||
|
||||
ch := make(chan *models.ChatCompletionStreamResponse)
|
||||
|
||||
|
||||
go func() {
|
||||
defer close(ch)
|
||||
err := StreamOllama(resp.RawBody(), ch, req.Model)
|
||||
@@ -121,14 +117,14 @@ func BuildOllamaBody(request *models.UnifiedRequest, messagesJSON []interface{},
|
||||
|
||||
// Context window size (default 8k for all, 32k+ for modern large-context models)
|
||||
ctxSize := 8192
|
||||
if strings.Contains(modelLower, "llama") ||
|
||||
strings.Contains(modelLower, "gemma") ||
|
||||
strings.Contains(modelLower, "mistral") ||
|
||||
strings.Contains(modelLower, "mixtral") ||
|
||||
strings.Contains(modelLower, "qwen") ||
|
||||
strings.Contains(modelLower, "deepseek") ||
|
||||
strings.Contains(modelLower, "command-r") ||
|
||||
strings.Contains(modelLower, "phi") {
|
||||
if strings.Contains(modelLower, "llama") ||
|
||||
strings.Contains(modelLower, "gemma") ||
|
||||
strings.Contains(modelLower, "mistral") ||
|
||||
strings.Contains(modelLower, "mixtral") ||
|
||||
strings.Contains(modelLower, "qwen") ||
|
||||
strings.Contains(modelLower, "deepseek") ||
|
||||
strings.Contains(modelLower, "command-r") ||
|
||||
strings.Contains(modelLower, "phi") {
|
||||
ctxSize = 32768
|
||||
}
|
||||
options["num_ctx"] = ctxSize
|
||||
@@ -137,13 +133,13 @@ func BuildOllamaBody(request *models.UnifiedRequest, messagesJSON []interface{},
|
||||
body["temperature"] = *request.Temperature
|
||||
options["temperature"] = *request.Temperature
|
||||
}
|
||||
|
||||
|
||||
if request.MaxTokens != nil {
|
||||
body["max_tokens"] = *request.MaxTokens
|
||||
options["num_predict"] = *request.MaxTokens
|
||||
} else {
|
||||
// Default to 8192 for all Ollama models if not specified,
|
||||
// as Ollama's compatibility layer defaults to 128 if neither
|
||||
// Default to 8192 for all Ollama models if not specified,
|
||||
// as Ollama's compatibility layer defaults to 128 if neither
|
||||
// max_tokens nor num_predict are provided.
|
||||
body["max_tokens"] = 8192
|
||||
options["num_predict"] = 8192
|
||||
@@ -189,7 +185,7 @@ func ParseOllamaResponse(respJSON map[string]interface{}, model string) (*models
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
|
||||
var resp models.ChatCompletionResponse
|
||||
if err := json.Unmarshal(data, &resp); err != nil {
|
||||
return nil, err
|
||||
@@ -202,7 +198,7 @@ func ParseOllamaResponse(respJSON map[string]interface{}, model string) (*models
|
||||
resp.Usage = &usage
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
@@ -256,4 +252,4 @@ func StreamOllama(ctx io.ReadCloser, ch chan<- *models.ChatCompletionStreamRespo
|
||||
}
|
||||
}
|
||||
return scanner.Err()
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user