fix(embedding): reuse SSRF-safe HTTP transport

This commit is contained in:
zhangzherui
2026-07-23 11:51:54 +08:00
committed by lyingbug
parent 8ca9c117bb
commit ccf02e1a2a
2 changed files with 35 additions and 1 deletions
+11 -1
View File
@@ -8,6 +8,14 @@ import (
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// sharedEmbeddingHTTPTransport keeps a single SSRF-safe connection pool for
// all embedding clients. Embedders are recreated as model configuration changes,
// but their outbound connections can be safely reused across client instances.
var sharedEmbeddingHTTPTransport = func() http.RoundTripper {
cfg := secutils.DefaultSSRFSafeHTTPClientConfig()
return secutils.NewSSRFSafeHTTPClient(cfg).Transport
}()
// 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 {
@@ -25,5 +33,7 @@ func validateEmbeddingBaseURL(baseURL string) error {
func newEmbeddingHTTPClient(timeout time.Duration) *http.Client {
cfg := secutils.DefaultSSRFSafeHTTPClientConfig()
cfg.Timeout = timeout
return secutils.NewSSRFSafeHTTPClient(cfg)
client := secutils.NewSSRFSafeHTTPClient(cfg)
client.Transport = sharedEmbeddingHTTPTransport
return client
}
@@ -3,8 +3,32 @@ package embedding
import (
"strings"
"testing"
"time"
)
func TestNewEmbeddingHTTPClient_ReusesTransport(t *testing.T) {
firstTimeout := 15 * time.Second
secondTimeout := 45 * time.Second
first := newEmbeddingHTTPClient(firstTimeout)
second := newEmbeddingHTTPClient(secondTimeout)
if first == second {
t.Fatal("expected distinct HTTP clients")
}
if first.Transport != second.Transport {
t.Fatal("expected embedding HTTP clients to share a transport")
}
if first.Transport != sharedEmbeddingHTTPTransport {
t.Fatal("expected embedding HTTP client to use the shared transport")
}
if first.Timeout != firstTimeout {
t.Fatalf("unexpected first client timeout: got %v, want %v", first.Timeout, firstTimeout)
}
if second.Timeout != secondTimeout {
t.Fatalf("unexpected second client timeout: got %v, want %v", second.Timeout, secondTimeout)
}
}
func TestValidateEmbeddingBaseURL_RejectsLoopback(t *testing.T) {
err := validateEmbeddingBaseURL("http://169.254.169.254/latest/meta-data")
if err == nil {