fix(embedding): add SSRF validation and safe HTTP client for embedders

Align embedding outbound requests with chat module defenses by validating base URLs at construction time and using SSRFSafeDialContext with redirect checks.
2026-06-26 22:37:30 +08:00
parent 500b223acf
commit 6b60d292f8
13 changed files with 112 additions and 26 deletions
@@ -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,
@@ -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
@@ -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,
@@ -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)
 	}))
@@ -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,
@@ -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,
@@ -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,
@@ -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) {
@@ -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)
+}
@@ -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")
+	}
+}
@@ -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,
@@ -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
 }

@@ -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,