mirror of
https://github.com/Tencent/WeKnora.git
synced 2026-11-04 23:32:11 +08:00
fix(embedding): reuse SSRF-safe HTTP transport
This commit is contained in:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user