diff --git a/internal/models/embedding/aliyun.go b/internal/models/embedding/aliyun.go index 5580fd2c..211878f7 100644 --- a/internal/models/embedding/aliyun.go +++ b/internal/models/embedding/aliyun.go @@ -112,15 +112,15 @@ func NewAliyunEmbedder(apiKey, baseURL, modelName string, timeout := 60 * time.Second - client := &http.Client{ - Timeout: timeout, + if err := validateEmbeddingBaseURL(baseURL); err != nil { + return nil, err } return &AliyunEmbedder{ apiKey: apiKey, baseURL: baseURL, modelName: modelName, - httpClient: client, + httpClient: newEmbeddingHTTPClient(timeout), truncatePromptTokens: truncatePromptTokens, EmbedderPooler: pooler, dimensions: dimensions, diff --git a/internal/models/embedding/azure_openai.go b/internal/models/embedding/azure_openai.go index c4a068ec..1ee59175 100644 --- a/internal/models/embedding/azure_openai.go +++ b/internal/models/embedding/azure_openai.go @@ -63,6 +63,10 @@ func NewAzureOpenAIEmbedder(apiKey, baseURL, modelName string, truncatePromptTokens = 511 } + if err := validateEmbeddingBaseURL(baseURL); err != nil { + return nil, err + } + return &AzureOpenAIEmbedder{ apiKey: apiKey, baseURL: baseURL, @@ -71,7 +75,7 @@ func NewAzureOpenAIEmbedder(apiKey, baseURL, modelName string, dimensions: dimensions, modelID: modelID, apiVersion: apiVersion, - httpClient: &http.Client{Timeout: 60 * time.Second}, + httpClient: newEmbeddingHTTPClient(60 * time.Second), maxRetries: 3, EmbedderPooler: pooler, }, nil diff --git a/internal/models/embedding/gemini.go b/internal/models/embedding/gemini.go index 589eeed0..1239d6fb 100644 --- a/internal/models/embedding/gemini.go +++ b/internal/models/embedding/gemini.go @@ -79,6 +79,11 @@ func NewGeminiEmbedder(apiKey, baseURL, modelName string, } timeout := 60 * time.Second + + if err := validateEmbeddingBaseURL(baseURL); err != nil { + return nil, err + } + return &GeminiEmbedder{ apiKey: apiKey, baseURL: baseURL, @@ -86,7 +91,7 @@ func NewGeminiEmbedder(apiKey, baseURL, modelName string, truncatePromptTokens: truncatePromptTokens, dimensions: dimensions, modelID: modelID, - httpClient: &http.Client{Timeout: timeout}, + httpClient: newEmbeddingHTTPClient(timeout), timeout: timeout, maxRetries: 3, EmbedderPooler: pooler, diff --git a/internal/models/embedding/gemini_test.go b/internal/models/embedding/gemini_test.go index 31babede..29562ab1 100644 --- a/internal/models/embedding/gemini_test.go +++ b/internal/models/embedding/gemini_test.go @@ -10,6 +10,8 @@ import ( ) func TestGeminiEmbedderBatchEmbedUsesNativeAPI(t *testing.T) { + t.Setenv("SSRF_WHITELIST", "127.0.0.1") + var gotPath string var gotAPIKey string var gotReq geminiBatchEmbedRequest @@ -66,6 +68,8 @@ func TestGeminiEmbedderBatchEmbedUsesNativeAPI(t *testing.T) { } func TestGeminiEmbedderBatchEmbedSendsOutputDimensionalityWhenOverrideEnabled(t *testing.T) { + t.Setenv("SSRF_WHITELIST", "127.0.0.1") + var gotReq geminiBatchEmbedRequest server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if err := json.NewDecoder(r.Body).Decode(&gotReq); err != nil { @@ -92,6 +96,8 @@ func TestGeminiEmbedderBatchEmbedSendsOutputDimensionalityWhenOverrideEnabled(t } func TestGeminiEmbedderReturnsAPIErrorBody(t *testing.T) { + t.Setenv("SSRF_WHITELIST", "127.0.0.1") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Error(w, `{"error":"not found"}`, http.StatusNotFound) })) diff --git a/internal/models/embedding/jina.go b/internal/models/embedding/jina.go index 770e6006..4c47da05 100644 --- a/internal/models/embedding/jina.go +++ b/internal/models/embedding/jina.go @@ -69,16 +69,15 @@ func NewJinaEmbedder(apiKey, baseURL, modelName string, timeout := 60 * time.Second - // Create HTTP client - client := &http.Client{ - Timeout: timeout, + if err := validateEmbeddingBaseURL(baseURL); err != nil { + return nil, err } return &JinaEmbedder{ apiKey: apiKey, baseURL: baseURL, modelName: modelName, - httpClient: client, + httpClient: newEmbeddingHTTPClient(timeout), EmbedderPooler: pooler, dimensions: dimensions, modelID: modelID, diff --git a/internal/models/embedding/nvidia.go b/internal/models/embedding/nvidia.go index 03f39448..453adb2a 100644 --- a/internal/models/embedding/nvidia.go +++ b/internal/models/embedding/nvidia.go @@ -70,16 +70,15 @@ func NewNvidiaEmbedder(apiKey, baseURL, modelName string, timeout := 60 * time.Second - // Create HTTP client - client := &http.Client{ - Timeout: timeout, + if err := validateEmbeddingBaseURL(baseURL); err != nil { + return nil, err } return &NvidiaEmbedder{ apiKey: apiKey, baseURL: baseURL, modelName: modelName, - httpClient: client, + httpClient: newEmbeddingHTTPClient(timeout), EmbedderPooler: pooler, dimensions: dimensions, modelID: modelID, diff --git a/internal/models/embedding/openai.go b/internal/models/embedding/openai.go index fce035ad..e55c3070 100644 --- a/internal/models/embedding/openai.go +++ b/internal/models/embedding/openai.go @@ -64,16 +64,15 @@ func NewOpenAIEmbedder(apiKey, baseURL, modelName string, timeout := 60 * time.Second - // Create HTTP client - client := &http.Client{ - Timeout: timeout, + if err := validateEmbeddingBaseURL(baseURL); err != nil { + return nil, err } return &OpenAIEmbedder{ apiKey: apiKey, baseURL: baseURL, modelName: modelName, - httpClient: client, + httpClient: newEmbeddingHTTPClient(timeout), truncatePromptTokens: truncatePromptTokens, EmbedderPooler: pooler, dimensions: dimensions, diff --git a/internal/models/embedding/openai_test.go b/internal/models/embedding/openai_test.go index 95ff4aee..c9650c16 100644 --- a/internal/models/embedding/openai_test.go +++ b/internal/models/embedding/openai_test.go @@ -46,6 +46,7 @@ func TestOpenAIEmbedderBatchEmbedOmitsDimensionsForFixedSizeModels(t *testing.T) func captureOpenAIEmbeddingRequest(t *testing.T, modelName string, dimensions int, supportsDimensionOverride bool) map[string]any { t.Helper() + t.Setenv("SSRF_WHITELIST", "127.0.0.1") requestBody := map[string]any{} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/internal/models/embedding/transport.go b/internal/models/embedding/transport.go new file mode 100644 index 00000000..1c3ae791 --- /dev/null +++ b/internal/models/embedding/transport.go @@ -0,0 +1,29 @@ +package embedding + +import ( + "fmt" + "net/http" + "time" + + secutils "github.com/Tencent/WeKnora/internal/utils" +) + +// validateEmbeddingBaseURL checks that a resolved embedding API base URL is safe +// for outbound requests. Empty URLs are allowed (callers apply provider defaults). +func validateEmbeddingBaseURL(baseURL string) error { + if baseURL == "" { + return nil + } + if err := secutils.ValidateURLForSSRF(baseURL); err != nil { + return fmt.Errorf("base URL SSRF check failed: %w", err) + } + return nil +} + +// newEmbeddingHTTPClient returns an HTTP client with connection-level SSRF +// protection and redirect validation, aligned with internal/models/chat/transport.go. +func newEmbeddingHTTPClient(timeout time.Duration) *http.Client { + cfg := secutils.DefaultSSRFSafeHTTPClientConfig() + cfg.Timeout = timeout + return secutils.NewSSRFSafeHTTPClient(cfg) +} diff --git a/internal/models/embedding/transport_test.go b/internal/models/embedding/transport_test.go new file mode 100644 index 00000000..af2db66b --- /dev/null +++ b/internal/models/embedding/transport_test.go @@ -0,0 +1,37 @@ +package embedding + +import ( + "strings" + "testing" +) + +func TestValidateEmbeddingBaseURL_RejectsLoopback(t *testing.T) { + err := validateEmbeddingBaseURL("http://169.254.169.254/latest/meta-data") + if err == nil { + t.Fatal("expected SSRF error for link-local metadata URL") + } + if !strings.Contains(err.Error(), "SSRF") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestValidateEmbeddingBaseURL_AllowsEmpty(t *testing.T) { + if err := validateEmbeddingBaseURL(""); err != nil { + t.Fatalf("empty base URL should be allowed: %v", err) + } +} + +func TestNewOpenAIEmbedder_RejectsPrivateBaseURL(t *testing.T) { + _, err := NewOpenAIEmbedder( + "test-key", + "http://169.254.169.254/latest/meta-data", + "text-embedding-3-small", + 511, + 256, + "model-id", + nil, + ) + if err == nil { + t.Fatal("expected SSRF rejection for link-local metadata URL") + } +} diff --git a/internal/models/embedding/volcengine.go b/internal/models/embedding/volcengine.go index 040d191c..9ef00a43 100644 --- a/internal/models/embedding/volcengine.go +++ b/internal/models/embedding/volcengine.go @@ -118,15 +118,15 @@ func NewVolcengineEmbedder(apiKey, baseURL, modelName string, timeout := 60 * time.Second - client := &http.Client{ - Timeout: timeout, + if err := validateEmbeddingBaseURL(baseURL); err != nil { + return nil, err } return &VolcengineEmbedder{ apiKey: apiKey, baseURL: baseURL, modelName: modelName, - httpClient: client, + httpClient: newEmbeddingHTTPClient(timeout), truncatePromptTokens: truncatePromptTokens, EmbedderPooler: pooler, dimensions: dimensions, diff --git a/internal/models/embedding/weknoracloud.go b/internal/models/embedding/weknoracloud.go index 421fb069..d2285056 100644 --- a/internal/models/embedding/weknoracloud.go +++ b/internal/models/embedding/weknoracloud.go @@ -10,6 +10,7 @@ import ( "strings" "time" + "github.com/Tencent/WeKnora/internal/models/provider" "github.com/Tencent/WeKnora/internal/models/utils" "github.com/google/uuid" ) @@ -42,16 +43,23 @@ func NewWeKnoraCloudEmbedder(config Config) (*WeKnoraCloudEmbedder, error) { if config.ExtraConfig != nil { remoteModelName = strings.TrimSpace(config.ExtraConfig["remote_model_name"]) } + baseURL := strings.TrimRight(config.BaseURL, "/") + if baseURL == "" { + baseURL = provider.WeKnoraCloudBaseURL + } + if err := validateEmbeddingBaseURL(baseURL); err != nil { + return nil, err + } return &WeKnoraCloudEmbedder{ modelName: config.ModelName, remoteModelName: remoteModelName, modelID: config.ModelID, appID: config.AppID, apiKey: config.AppSecret, - baseURL: strings.TrimRight(config.BaseURL, "/"), + baseURL: baseURL, dimensions: config.Dimensions, supportsDimensionOverride: config.SupportsDimensionOverride, - client: &http.Client{Timeout: 60 * time.Second}, + client: newEmbeddingHTTPClient(60 * time.Second), }, nil } diff --git a/internal/models/embedding/zhipu.go b/internal/models/embedding/zhipu.go index 4206cdcf..513487dc 100644 --- a/internal/models/embedding/zhipu.go +++ b/internal/models/embedding/zhipu.go @@ -65,16 +65,15 @@ func NewZhipuEmbedder(apiKey, baseURL, modelName string, timeout := 60 * time.Second - // Create HTTP client - client := &http.Client{ - Timeout: timeout, + if err := validateEmbeddingBaseURL(baseURL); err != nil { + return nil, err } return &ZhipuEmbedder{ apiKey: apiKey, baseURL: baseURL, modelName: modelName, - httpClient: client, + httpClient: newEmbeddingHTTPClient(timeout), truncatePromptTokens: truncatePromptTokens, EmbedderPooler: pooler, dimensions: dimensions,