From fa258db6a3b4ad07779feed91cd73859c4aea93a Mon Sep 17 00:00:00 2001 From: Jin Hai Date: Sat, 19 Sep 2026 23:32:26 +0800 Subject: [PATCH] Go: reuse a global HTTP client and tune connection pool to prevent port exhaustion (#19888) Signed-off-by: Jin Hai --- internal/agent/component/agent_mcp_test.go | 14 +- internal/agent/component/invoke.go | 9 +- internal/agent/component/invoke_test.go | 16 +- internal/agent/tool/exesql_ssrf_test.go | 10 +- internal/agent/tool/mcp_test.go | 7 +- internal/agent/tool/ssrf.go | 8 +- internal/common/http.go | 583 ++++++++++++++++++ internal/entity/models/base_model.go | 7 +- .../entity/models/model_test_helpers_test.go | 8 +- internal/entity/models/ollama.go | 2 +- internal/handler/chat_audio_test.go | 7 +- .../component/ssrf_test_helper_test.go | 8 +- .../parser/parser/ssrf_test_helper_test.go | 9 +- internal/service/agent_dbcheck.go | 2 +- internal/service/file/file.go | 3 +- internal/service/file/file_test.go | 13 +- internal/service/mcp.go | 2 +- internal/syncer/connector/azure_blob.go | 7 +- internal/syncer/connector/bitbucket_test.go | 9 +- internal/syncer/connector/rest_api_test.go | 9 +- .../syncer/connector/s3_compatible_test.go | 9 +- internal/syncer/connector/salesforce_test.go | 15 +- internal/syncer/connector/sitemap.go | 3 +- internal/syncer/connector/ssrf.go | 11 +- internal/syncer/connector/ssrf_test.go | 14 +- internal/syncer/connector/teams_test.go | 9 +- internal/utility/mcp_call.go | 3 +- internal/utility/mcp_client.go | 5 +- internal/utility/mcp_client_test.go | 7 +- internal/utility/ssrf.go | 270 -------- internal/utility/ssrf_test.go | 10 +- internal/utility/upload.go | 3 +- internal/utility/upload_test.go | 13 +- 33 files changed, 709 insertions(+), 396 deletions(-) diff --git a/internal/agent/component/agent_mcp_test.go b/internal/agent/component/agent_mcp_test.go index 4cd7ff741c..d1715e4c47 100644 --- a/internal/agent/component/agent_mcp_test.go +++ b/internal/agent/component/agent_mcp_test.go @@ -3,13 +3,11 @@ package component import ( "encoding/json" "fmt" - "github.com/cloudwego/eino/compose" - "github.com/cloudwego/eino/flow/agent/react" - "github.com/cloudwego/eino/schema" "io" "net/http" "net/http/httptest" "ragflow/internal/agent/runtime" + "ragflow/internal/common" "ragflow/internal/entity" "ragflow/internal/entity/models" "ragflow/internal/utility" @@ -17,6 +15,10 @@ import ( "sync/atomic" "testing" "time" + + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/flow/agent/react" + "github.com/cloudwego/eino/schema" ) func TestAgentMCPRequiresRuntimeServer(t *testing.T) { @@ -33,9 +35,9 @@ func TestAgentMCPModelDispatch(t *testing.T) { t.Fatal(err) } pushComponentDB(t, db) - lookup := utility.LookupHost - utility.LookupHost = func(string) ([]string, error) { return []string{"8.8.8.8"}, nil } - t.Cleanup(func() { utility.LookupHost = lookup }) + lookup := common.LookupHost + common.LookupHost = func(string) ([]string, error) { return []string{"8.8.8.8"}, nil } + t.Cleanup(func() { common.LookupHost = lookup }) var dispatched atomic.Bool mcp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method == "DELETE" { diff --git a/internal/agent/component/invoke.go b/internal/agent/component/invoke.go index e2591cd988..b0ed168380 100644 --- a/internal/agent/component/invoke.go +++ b/internal/agent/component/invoke.go @@ -24,7 +24,7 @@ // traceparent headers. // // SSRF guard (PR #15426): every outbound URL is validated against -// the shared utility.AssertURLSafe before any network I/O. Both the +// the shared common.AssertURLSafe before any network I/O. Both the // target URL and an optional proxy URL are checked — the proxy // vector matters because the Go transport hands the request to the // proxy host, which would otherwise re-resolve the original host @@ -44,6 +44,7 @@ import ( "net" "net/http" "net/url" + "ragflow/internal/common" "strings" "time" @@ -88,7 +89,7 @@ func (i *InvokeComponent) Name() string { return i.name } // param // // SSRF flow (PR #15426): -// 1. Validate the target URL via utility.AssertURLSafe (loopback / +// 1. Validate the target URL via common.AssertURLSafe (loopback / // link-local / RFC1918 / metadata / unresolvable are rejected). // 2. Validate the optional proxy URL the same way (the proxy // re-resolves the target host; an unsafe proxy would defeat @@ -136,7 +137,7 @@ func (i *InvokeComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[st // Step 1: SSRF guard for the target URL. The validated // hostname + resolved public IP are reused for DNS pinning. - host, pinnedIP, err := utility.AssertURLSafe(rawURL) + host, pinnedIP, err := common.AssertURLSafe(rawURL) if err != nil { return invokeSSRFError("url", rawURL, err), nil } @@ -175,7 +176,7 @@ func (i *InvokeComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[st return invokeSSRFError("url", rawURL, fmt.Errorf("Invoke: proxy mode requires a literal-IP target URL (hostnames are unsafe because the proxy re-resolves them)")), nil } - ph, pip, perr := utility.AssertURLSafe(proxyStr) + ph, pip, perr := common.AssertURLSafe(proxyStr) if perr != nil { return invokeSSRFError("proxy", proxyStr, perr), nil } diff --git a/internal/agent/component/invoke_test.go b/internal/agent/component/invoke_test.go index 228ea6d1b6..3735a11715 100644 --- a/internal/agent/component/invoke_test.go +++ b/internal/agent/component/invoke_test.go @@ -24,6 +24,7 @@ import ( "net/http" "net/http/httptest" "net/url" + "ragflow/internal/common" "reflect" "strings" "testing" @@ -31,13 +32,12 @@ import ( "go.uber.org/zap" "go.uber.org/zap/zaptest/observer" - "ragflow/internal/utility" ) // Test setup: enable the test-only SSRF bypass for tests in this // file (the happy-path httptest server lives on 127.0.0.1, which // the production guard rejects). The bypass is now a process- -// memory boolean (utility.AllowAnyHostForTest) instead of an +// memory boolean (common.AllowAnyHostForTest) instead of an // env var — the previous form (ALLOW_ANY_HOST env) was a live // runtime toggle any operator could flip to disable the guard // globally. PR review round 6, Major #3. @@ -47,9 +47,9 @@ import ( // tests below rely on the bypass being on. func setupAllowAnyHost(t *testing.T, enabled bool) { t.Helper() - prev := utility.AllowAnyHostForTest - utility.AllowAnyHostForTest = enabled - t.Cleanup(func() { utility.AllowAnyHostForTest = prev }) + prev := common.AllowAnyHostForTest + common.AllowAnyHostForTest = enabled + t.Cleanup(func() { common.AllowAnyHostForTest = prev }) } // TestInvoke_GET exercises the happy path: a GET request to a stub @@ -441,14 +441,14 @@ func TestInvoke_ProxyDNSPin(t *testing.T) { // "rebinds" to a different answer afterward. // We achieve "rebinding" by stubbing the DNS lookup // to return a different IP on a second call. - originalLookup := utility.LookupHost - utility.LookupHost = func(host string) ([]string, error) { + originalLookup := common.LookupHost + common.LookupHost = func(host string) ([]string, error) { // Always return the already-running fake proxy. If the // Invoke transport re-resolves proxy.test.invalid instead // of using the pinned IP, the request will never hit it. return []string{pinnedProxyIP}, nil } - t.Cleanup(func() { utility.LookupHost = originalLookup }) + t.Cleanup(func() { common.LookupHost = originalLookup }) c, _ := NewInvokeComponent(nil) ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second) diff --git a/internal/agent/tool/exesql_ssrf_test.go b/internal/agent/tool/exesql_ssrf_test.go index 66968912ee..bfd8d8693c 100644 --- a/internal/agent/tool/exesql_ssrf_test.go +++ b/internal/agent/tool/exesql_ssrf_test.go @@ -24,7 +24,7 @@ import ( "sync" "testing" - "ragflow/internal/utility" + "ragflow/internal/common" ) // recordingDialer is an exesqlDialer that records the host it was @@ -131,18 +131,18 @@ func TestExeSQL_SSRF_RejectsEmptyHost(t *testing.T) { // TestExeSQL_SSRF_PinsToValidatedIP ensures that when a DNS name // resolves to a public IP, InvokableRun dials the validated IP, not // the original hostname — closing the TOCTOU window for DNS -// rebinding. The resolver is stubbed via utility.LookupHost so the +// rebinding. The resolver is stubbed via common.LookupHost so the // test does not depend on real DNS. func TestExeSQL_SSRF_PinsToValidatedIP(t *testing.T) { // Stub the resolver so example.test -> 1.2.3.4 (public, stable). - origLookup := utility.LookupHost - utility.LookupHost = func(host string) ([]string, error) { + origLookup := common.LookupHost + common.LookupHost = func(host string) ([]string, error) { if host == "example.test" { return []string{"1.2.3.4"}, nil } return origLookup(host) } - t.Cleanup(func() { utility.LookupHost = origLookup }) + t.Cleanup(func() { common.LookupHost = origLookup }) // failFast:true makes the test dialer return an error instead of // opening a real *sql.DB, so the InvokableRun path stops after diff --git a/internal/agent/tool/mcp_test.go b/internal/agent/tool/mcp_test.go index a67188a161..4c7e04b7e0 100644 --- a/internal/agent/tool/mcp_test.go +++ b/internal/agent/tool/mcp_test.go @@ -21,6 +21,7 @@ import ( "io" "net/http" "net/http/httptest" + "ragflow/internal/common" "slices" "strings" "testing" @@ -432,9 +433,9 @@ func TestMCPToolAdapter_InvokableRunIsError(t *testing.T) { // utility/mcp_client_test.go's allowLoopbackForTests helper. func mcpLoopbackOverride(t *testing.T) func() { t.Helper() - orig := mcpclient.LookupHost - mcpclient.LookupHost = func(_ string) ([]string, error) { + orig := common.LookupHost + common.LookupHost = func(_ string) ([]string, error) { return []string{"8.8.8.8"}, nil } - return func() { mcpclient.LookupHost = orig } + return func() { common.LookupHost = orig } } diff --git a/internal/agent/tool/ssrf.go b/internal/agent/tool/ssrf.go index a40e7aee19..e85bed4eee 100644 --- a/internal/agent/tool/ssrf.go +++ b/internal/agent/tool/ssrf.go @@ -23,8 +23,6 @@ import ( "net/url" "ragflow/internal/common" "strings" - - "ragflow/internal/utility" ) // ErrSSRFBlocked is returned when a tool is asked to fetch a URL whose @@ -187,7 +185,7 @@ func ValidateDBHost(host string) (string, error) { if ip := net.ParseIP(host); ip != nil { return ip.String(), nil } - addrs, lerr := utility.LookupHost(host) + addrs, lerr := common.LookupHost(host) if lerr != nil { return "", fmt.Errorf("ssrf: resolve %s: %w", host, lerr) } @@ -215,10 +213,10 @@ func ValidateDBHost(host string) (string, error) { return ip.String(), nil } - // Resolve via utility.LookupHost so tests can stub DNS without + // Resolve via common.LookupHost so tests can stub DNS without // touching real network — matches the stubbing pattern used by // the utility package (see internal/utility/ssrf.go LookupHost). - addrs, lerr := utility.LookupHost(host) + addrs, lerr := common.LookupHost(host) if lerr != nil { return "", fmt.Errorf("ssrf: resolve %s: %w", host, lerr) } diff --git a/internal/common/http.go b/internal/common/http.go index 190c7cea50..b5ddd3cb66 100644 --- a/internal/common/http.go +++ b/internal/common/http.go @@ -17,13 +17,596 @@ package common import ( + "bytes" + "encoding/json" "fmt" + "io" + "net" "net/http" + "net/url" + "slices" + "sort" "strconv" + "strings" + "sync" + "time" "github.com/gin-gonic/gin" ) +const ( + redactedLogValue = "[REDACTED]" + maxLoggedVectorFloats = 3 +) + +var ssrfHttpClient *http.Client +var schemeSafeHttpClient *http.Client + +func GetSSRFHTTPClient() *http.Client { + if ssrfHttpClient == nil { + var t *http.Transport + if dt, ok := http.DefaultTransport.(*http.Transport); ok { + t = dt.Clone() + } else { + t = &http.Transport{Proxy: http.ProxyFromEnvironment} + } + t.MaxIdleConns = 100 + t.MaxIdleConnsPerHost = 10 + t.IdleConnTimeout = 90 * time.Second + t.DisableCompression = false + t.ResponseHeaderTimeout = 20 * time.Minute + t.TLSHandshakeTimeout = 30 * time.Second + + var rt http.RoundTripper = t + rt = &strictSSRFTransport{base: rt} + rt = newProviderLoggingTransport(rt) + ssrfHttpClient = &http.Client{Transport: rt} + } + return ssrfHttpClient +} + +func GetSchemeSafeHTTPClient() *http.Client { + if schemeSafeHttpClient == nil { + var t *http.Transport + if dt, ok := http.DefaultTransport.(*http.Transport); ok { + t = dt.Clone() + } else { + t = &http.Transport{Proxy: http.ProxyFromEnvironment} + } + t.MaxIdleConns = 5000 + t.MaxIdleConnsPerHost = 500 + t.IdleConnTimeout = 90 * time.Second + t.DisableCompression = false + t.ResponseHeaderTimeout = 60 * time.Second + t.TLSHandshakeTimeout = 30 * time.Second + + var rt http.RoundTripper = t + rt = &schemeSafeTransport{base: rt} + rt = newProviderLoggingTransport(rt) + schemeSafeHttpClient = &http.Client{Transport: rt} + } + return schemeSafeHttpClient +} + +func newProviderLoggingTransport(base http.RoundTripper) http.RoundTripper { + if base == nil { + base = http.DefaultTransport + } + if !IsLLMDebugEnabled() { + return base + } + return &providerLoggingTransport{base: base, now: time.Now} +} + +type providerLoggingTransport struct { + base http.RoundTripper + now func() time.Time +} + +func (t *providerLoggingTransport) RoundTrip(req *http.Request) (*http.Response, error) { + payload, err := readAndRestoreRequestBody(req) + if err != nil { + return nil, err + } + providerURL := redactProviderURL(req.URL) + logPayload := redactProviderBody(payload) + + startedAt := t.now() + resp, err := t.base.RoundTrip(req) + if err != nil { + logProviderCall(providerURL, logPayload, 0, "", t.now().Sub(startedAt), 0, err) + return nil, err + } + if resp.Body == nil { + logProviderCall(providerURL, logPayload, resp.StatusCode, "", t.now().Sub(startedAt), 0, nil) + return resp, nil + } + + resp.Body = &providerResponseBody{ + ReadCloser: resp.Body, + startedAt: startedAt, + now: t.now, + log: func(body []byte, took, firstToken time.Duration) { + logProviderCall(providerURL, logPayload, resp.StatusCode, redactProviderBody(body), took, firstToken, nil) + }, + } + return resp, nil +} + +// providerResponseBody captures bytes while callers consume them, preserving +// streaming delivery instead of eagerly reading the entire provider response. +type providerResponseBody struct { + io.ReadCloser + body bytes.Buffer + startedAt time.Time + now func() time.Time + firstToken time.Duration + firstOnce sync.Once + logOnce sync.Once + log func([]byte, time.Duration, time.Duration) +} + +func (b *providerResponseBody) Read(p []byte) (int, error) { + n, err := b.ReadCloser.Read(p) + if n > 0 { + b.firstOnce.Do(func() { + b.firstToken = b.now().Sub(b.startedAt) + }) + _, _ = b.body.Write(p[:n]) + } + if err == io.EOF { + b.writeLogOnce() + } + return n, err +} + +func (b *providerResponseBody) Close() error { + err := b.ReadCloser.Close() + b.writeLogOnce() + return err +} + +func (b *providerResponseBody) writeLogOnce() { + b.logOnce.Do(func() { + b.log(b.body.Bytes(), b.now().Sub(b.startedAt), b.firstToken) + }) +} + +func readAndRestoreRequestBody(req *http.Request) ([]byte, error) { + if req.Body == nil { + return nil, nil + } + body, err := io.ReadAll(req.Body) + if err != nil { + return nil, fmt.Errorf("failed to read provider request body for logging: %w", err) + } + _ = req.Body.Close() + req.Body = io.NopCloser(bytes.NewReader(body)) + req.ContentLength = int64(len(body)) + return body, nil +} + +func redactProviderURL(requestURL *url.URL) string { + if requestURL == nil { + return "" + } + redacted := *requestURL + if redacted.User != nil { + redacted.User = url.User(redacted.User.Username()) + } + query := redacted.Query() + for key := range query { + if isSensitiveLogKey(key) { + query.Set(key, redactedLogValue) + } + } + redacted.RawQuery = query.Encode() + return redacted.String() +} + +func redactProviderBody(body []byte) string { + if len(body) == 0 { + return "" + } + + var value any + if err := json.Unmarshal(body, &value); err == nil { + redactProviderValue(value) + if redacted, err := json.Marshal(value); err == nil { + return string(redacted) + } + } + + lines := strings.Split(string(body), "\n") + redactedAny := false + for i, line := range lines { + prefix, data, ok := strings.Cut(line, "data:") + if !ok || strings.TrimSpace(data) == "[DONE]" { + continue + } + var event any + if err := json.Unmarshal([]byte(strings.TrimSpace(data)), &event); err != nil { + continue + } + redactProviderValue(event) + redacted, err := json.Marshal(event) + if err != nil { + continue + } + lines[i] = prefix + "data: " + string(redacted) + redactedAny = true + } + if redactedAny { + return strings.Join(lines, "\n") + } + return string(body) +} + +func redactProviderValue(value any) { + switch value := value.(type) { + case map[string]any: + for key, child := range value { + if isSensitiveLogKey(key) { + value[key] = redactedLogValue + continue + } + if isVectorLogKey(key) { + child = truncateLoggedVectors(child) + value[key] = child + } + redactProviderValue(child) + } + case []any: + for _, child := range value { + redactProviderValue(child) + } + } +} + +func isVectorLogKey(key string) bool { + switch strings.ToLower(key) { + case "embedding", "embeddings", "vector", "vectors": + return true + default: + return false + } +} + +func truncateLoggedVectors(value any) any { + switch value := value.(type) { + case []any: + if isNumericVector(value) { + if len(value) > maxLoggedVectorFloats { + return value[:maxLoggedVectorFloats] + } + return value + } + for i, child := range value { + value[i] = truncateLoggedVectors(child) + } + case map[string]any: + for key, child := range value { + value[key] = truncateLoggedVectors(child) + } + } + return value +} + +func isNumericVector(value []any) bool { + if len(value) == 0 { + return false + } + for _, item := range value { + if _, ok := item.(float64); !ok { + return false + } + } + return true +} + +func isSensitiveLogKey(key string) bool { + normalized := strings.NewReplacer("-", "", "_", "", ".", "").Replace(strings.ToLower(key)) + switch normalized { + case "apikey", "authorization", "accesstoken", "refreshtoken", "password", "secret", "token", "key": + return true + default: + return false + } +} + +func logProviderCall(providerURL, payload string, statusCode int, responseBody string, took, firstToken time.Duration, err error) { + request := fmt.Sprintf("url=%s payload=%s", providerURL, payload) + response := fmt.Sprintf("response_code=%d took=%s first-token=%s response_body=%s", statusCode, took, firstToken, responseBody) + if err != nil { + response += " error=" + err.Error() + } + LogRequestResponseInfo(request, response, err == nil && statusCode >= 200 && statusCode < 300) +} + +// schemeSafeTransport wraps an http.RoundTripper so every outgoing request is +// validated by the lenient SSRF guard (http/https scheme + non-empty host). +// Private and loopback hosts are permitted. Used only by local-inference +// drivers that may target a user's own network. +type schemeSafeTransport struct{ base http.RoundTripper } + +func (t *schemeSafeTransport) RoundTrip(req *http.Request) (*http.Response, error) { + if err := AssertURLSchemeSafe(req.URL.String()); err != nil { + return nil, err + } + return t.base.RoundTrip(req) +} + +// strictSSRFTransport wraps an http.RoundTripper so every outgoing request is +// validated by the strict SSRF guard (scheme + host + globally routable IP). +// This is the default for cloud-hosted model drivers and closes the +// go/request-forgery data flow: the user-controllable BaseURL cannot be made to +// point at private hosts, loopback, link-local, or cloud metadata endpoints. +type strictSSRFTransport struct{ base http.RoundTripper } + +func (t *strictSSRFTransport) RoundTrip(req *http.Request) (*http.Response, error) { + if _, _, err := AssertURLSafe(req.URL.String()); err != nil { + return nil, err + } + return t.base.RoundTrip(req) +} + +var AllowedURLSchemes = []string{"http", "https"} + +// LookupHost is the indirection used to resolve hostnames. Tests override it. +var LookupHost = net.LookupHost + +// AllowAnyHostForTest is a test-only override that skips the +// public-IP routability check in AssertURLSafe and AssertHostSafe. +// Scheme, host, and DNS resolution checks are unchanged: hostnames +// still resolve, unresolvable hostnames still fail, and callers +// still pin connections to the returned resolved address. Production +// code MUST leave this at its zero value (false). Tests that need to +// talk to a local httptest server flip it on and reset it in +// t.Cleanup. +// +// The previous form (env-var ALLOW_ANY_HOST) was a live runtime +// toggle that any operator could flip to disable the SSRF guard +// globally — including the DNS pinning that the Invoke component +// relies on. PR review round 6, Major #3: this variable lives in +// process memory only, so it cannot be enabled by an env var or +// a deployment mistake. The explicit "_ForTest" suffix is the +// signal that production code must never touch it. +var AllowAnyHostForTest = false + +// allowAnyHost reads the test-only override. Kept as a private +// helper so the call sites don't all have to know about the +// exported variable name. +func allowAnyHost() bool { + return AllowAnyHostForTest +} + +// AssertURLSchemeSafe is a lenient SSRF guard for drivers that may legitimately +// target private networks or loopback addresses (e.g. self-hosted Ollama, vLLM, +// Xinference). It only rejects dangerous schemes and empty hosts; it does not +// resolve DNS and does not require public routability. Use this ONLY for +// local-inference model drivers — cloud-hosted drivers must use AssertURLSafe. +var AssertURLSchemeSafe = func(rawURL string) error { + parsed, err := url.Parse(strings.TrimSpace(rawURL)) + if err != nil { + return fmt.Errorf("invalid url") + } + + scheme := strings.ToLower(parsed.Scheme) + if !slices.Contains(AllowedURLSchemes, scheme) { + sorted := append([]string(nil), AllowedURLSchemes...) + sort.Strings(sorted) + return fmt.Errorf("disallowed URL scheme: '%s'. Only %v are allowed", scheme, sorted) + } + + if parsed.Hostname() == "" { + return fmt.Errorf("URL is missing a host") + } + + return nil +} + +var AssertURLSafe = func(rawURL string) (hostname, resolvedIP string, err error) { + parsed, err := url.Parse(strings.TrimSpace(rawURL)) + if err != nil { + return "", "", fmt.Errorf("invalid url") + } + + scheme := strings.ToLower(parsed.Scheme) + if !slices.Contains(AllowedURLSchemes, scheme) { + sorted := append([]string(nil), AllowedURLSchemes...) + sort.Strings(sorted) + return "", "", fmt.Errorf("disallowed URL scheme: '%s'. Only %v are allowed", scheme, sorted) + } + + hostname = parsed.Hostname() + if hostname == "" { + return "", "", fmt.Errorf("URL is missing a host") + } + + allowAny := allowAnyHost() + addresses, err := LookupHost(hostname) + if err != nil { + return "", "", fmt.Errorf("could not resolve hostname '%s': %w", hostname, err) + } + if len(addresses) == 0 { + return "", "", fmt.Errorf("hostname '%s' resolved to no addresses", hostname) + } + + for _, addr := range addresses { + ip := net.ParseIP(addr) + if ip == nil { + return "", "", fmt.Errorf("could not parse resolved address '%s' for hostname '%s'", addr, hostname) + } + if !allowAny && !isGlobalIP(effectiveIP(ip)) { + return "", "", fmt.Errorf("URL resolves to a non-public address (%s), which is not allowed", ip.String()) + } + if resolvedIP == "" { + resolvedIP = ip.String() + } + } + return hostname, resolvedIP, nil +} + +// AssertHostSafe validates a bare host (a hostname or a literal IP, with no +// scheme or port) and returns the first resolved public IP. It is the +// host-type counterpart of AssertURLSafe: every resolved address must be +// globally routable (private, loopback, link-local, metadata, multicast and +// reserved ranges are rejected). Callers dial the returned IP directly so DNS +// cannot rebind the connection to an internal address between validation and +// the TCP connect. +// +// Used by host-based data sources (IMAP/MySQL/PostgreSQL) and by the ExeSQL / +// test_db_connection host guards, mirroring common/ssrf_guard.py: +// assert_host_is_safe. +var AssertHostSafe = func(host string) (resolvedIP string, err error) { + host = strings.TrimSpace(host) + if host == "" { + return "", fmt.Errorf("host is missing") + } + + allowAny := allowAnyHost() + if ip := net.ParseIP(host); ip != nil { + if !allowAny && !isGlobalIP(effectiveIP(ip)) { + return "", fmt.Errorf("host is not a public address (%s), which is not allowed", ip.String()) + } + return ip.String(), nil + } + + addresses, err := LookupHost(host) + if err != nil { + return "", fmt.Errorf("could not resolve hostname '%s': %w", host, err) + } + if len(addresses) == 0 { + return "", fmt.Errorf("hostname '%s' resolved to no addresses", host) + } + + for _, addr := range addresses { + ip := net.ParseIP(addr) + if ip == nil { + return "", fmt.Errorf("could not parse resolved address '%s' for hostname '%s'", addr, host) + } + if !allowAny && !isGlobalIP(effectiveIP(ip)) { + return "", fmt.Errorf("hostname '%s' resolves to a non-public address (%s), which is not allowed", host, ip.String()) + } + if resolvedIP == "" { + resolvedIP = ip.String() + } + } + return resolvedIP, nil +} + +// effectiveIP unwraps IPv4-mapped IPv6 addresses (e.g. ::ffff:127.0.0.1) so +// the routability check sees the IPv4 form. Without this, an attacker could +// bypass the guard with an IPv4-mapped IPv6 representation of a private host. +func effectiveIP(ip net.IP) net.IP { + if v4 := ip.To4(); v4 != nil { + return v4 + } + return ip +} + +// isGlobalIP mirrors Python's ipaddress.IPv*Address.is_global: an address is +// global if it is none of {unspecified, loopback, multicast, link-local, +// private (including CGNAT and IPv6 ULA), benchmarking, documentation, +// reserved}. +func isGlobalIP(ip net.IP) bool { + if ip == nil || ip.IsUnspecified() || ip.IsLoopback() || ip.IsMulticast() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsInterfaceLocalMulticast() || ip.IsPrivate() { + return false + } + if v4 := ip.To4(); v4 != nil { + // 0.0.0.0/8 — "this network"; 0.x.y.z routes to localhost on Linux. + if v4[0] == 0 { + return false + } + // CGNAT 100.64.0.0/10 — not flagged by IsPrivate in older Go versions. + if v4[0] == 100 && v4[1]&0xC0 == 64 { + return false + } + // 192.0.0.0/24 reserved for IETF protocol assignments. + if v4[0] == 192 && v4[1] == 0 && v4[2] == 0 { + return false + } + // 192.0.2.0/24, 198.51.100.0/24, 203.0.113.0/24 documentation (TEST-NET-1/2/3). + if v4[0] == 192 && v4[1] == 0 && v4[2] == 2 { + return false + } + if v4[0] == 198 && v4[1] == 51 && v4[2] == 100 { + return false + } + if v4[0] == 203 && v4[1] == 0 && v4[2] == 113 { + return false + } + // 198.18.0.0/15 benchmarking. + if v4[0] == 198 && (v4[1] == 18 || v4[1] == 19) { + return false + } + // 240.0.0.0/4 reserved (excluding 255.255.255.255 which IsUnspecified misses). + if v4[0] >= 240 { + return false + } + } else if v6 := ip.To16(); v6 != nil { + // 2001:db8::/32 documentation prefix. + if v6[0] == 0x20 && v6[1] == 0x01 && v6[2] == 0x0d && v6[3] == 0xb8 { + return false + } + // 100::/64 discard-only address block. + if v6[0] == 0x01 && v6[1] == 0x00 && allZero(v6[2:8]) { + return false + } + // IPv6 transition addresses (6to4, NAT64, Teredo, IPv4-compatible) embed + // an arbitrary IPv4 address that none of the checks above look at. Unwrap + // and re-check it so 2002:7f00:1::1 is treated as 127.0.0.1. + for _, inner := range embeddedIPv4(v6) { + if !isGlobalIP(inner) { + return false + } + } + } + return true +} + +// embeddedIPv4 returns the IPv4 addresses carried inside an IPv6 transition +// address, or nil when it carries none. Teredo yields two: the relay server and +// the (obfuscated) client. +func embeddedIPv4(v6 net.IP) []net.IP { + switch { + // 6to4 — RFC 3056, 2002::/16, IPv4 in bytes 2-6. + case v6[0] == 0x20 && v6[1] == 0x02: + return []net.IP{net.IPv4(v6[2], v6[3], v6[4], v6[5])} + + // NAT64 well-known prefix — RFC 6052, 64:ff9b::/96, IPv4 in the low 32 bits. + case v6[0] == 0x00 && v6[1] == 0x64 && v6[2] == 0xff && v6[3] == 0x9b && allZero(v6[4:12]): + return []net.IP{net.IPv4(v6[12], v6[13], v6[14], v6[15])} + + // NAT64 local-use prefix — RFC 8215, 64:ff9b:1::/48. The embedded IPv4 + // position depends on the operator's prefix length, so block the range. + case v6[0] == 0x00 && v6[1] == 0x64 && v6[2] == 0xff && v6[3] == 0x9b && v6[4] == 0x00 && v6[5] == 0x01: + return []net.IP{net.IPv4zero} + + // Teredo — RFC 4380, 2001::/32. Server IPv4 in bytes 4-8, client IPv4 in + // bytes 12-16 obfuscated by XOR with 0xff. + case v6[0] == 0x20 && v6[1] == 0x01 && v6[2] == 0x00 && v6[3] == 0x00: + return []net.IP{ + net.IPv4(v6[4], v6[5], v6[6], v6[7]), + net.IPv4(v6[12]^0xff, v6[13]^0xff, v6[14]^0xff, v6[15]^0xff), + } + + // IPv4-compatible — deprecated ::a.b.c.d, not unwrapped by net.IP.To4. + case allZero(v6[0:12]): + return []net.IP{net.IPv4(v6[12], v6[13], v6[14], v6[15])} + } + return nil +} + +func allZero(b []byte) bool { + for _, x := range b { + if x != 0 { + return false + } + } + return true +} + type response struct { Code ErrorCode `json:"code"` Data interface{} `json:"data"` diff --git a/internal/entity/models/base_model.go b/internal/entity/models/base_model.go index 20d000363c..ef56635bd6 100644 --- a/internal/entity/models/base_model.go +++ b/internal/entity/models/base_model.go @@ -27,7 +27,6 @@ import ( "net/url" "ragflow/internal/common" "ragflow/internal/engine/clickhouse" - "ragflow/internal/utility" "sort" "strings" "sync" @@ -471,7 +470,7 @@ func ParseListModel(modelList ModelList) []ListModelResponse { // // allowPrivate selects the guard strictness: // - false (cloud-hosted drivers): every request is validated with -// utility.AssertURLSafe — scheme + host must be present and every resolved +// common.AssertURLSafe — scheme + host must be present and every resolved // IP must be globally routable (private/loopback/link-local/metadata are // rejected). This is the default and closes the go/request-forgery sink. // - true (local-inference drivers): requests are validated with @@ -743,7 +742,7 @@ func logProviderCall(providerURL, payload string, statusCode int, responseBody s type schemeSafeTransport struct{ base http.RoundTripper } func (t *schemeSafeTransport) RoundTrip(req *http.Request) (*http.Response, error) { - if err := utility.AssertURLSchemeSafe(req.URL.String()); err != nil { + if err := common.AssertURLSchemeSafe(req.URL.String()); err != nil { return nil, err } return t.base.RoundTrip(req) @@ -757,7 +756,7 @@ func (t *schemeSafeTransport) RoundTrip(req *http.Request) (*http.Response, erro type strictSSRFTransport struct{ base http.RoundTripper } func (t *strictSSRFTransport) RoundTrip(req *http.Request) (*http.Response, error) { - if _, _, err := utility.AssertURLSafe(req.URL.String()); err != nil { + if _, _, err := common.AssertURLSafe(req.URL.String()); err != nil { return nil, err } return t.base.RoundTrip(req) diff --git a/internal/entity/models/model_test_helpers_test.go b/internal/entity/models/model_test_helpers_test.go index be97ba4d2e..6096841caa 100644 --- a/internal/entity/models/model_test_helpers_test.go +++ b/internal/entity/models/model_test_helpers_test.go @@ -4,7 +4,7 @@ import ( "strings" "testing" - "ragflow/internal/utility" + "ragflow/internal/common" ) // withSSRFBypass enables the test-only SSRF bypass for the duration of a test @@ -14,9 +14,9 @@ import ( // its own dedicated coverage in internal/utility/ssrf_test.go. func withSSRFBypass(t *testing.T) { t.Helper() - prev := utility.AllowAnyHostForTest - utility.AllowAnyHostForTest = true - t.Cleanup(func() { utility.AllowAnyHostForTest = prev }) + prev := common.AllowAnyHostForTest + common.AllowAnyHostForTest = true + t.Cleanup(func() { common.AllowAnyHostForTest = prev }) } func requireNoSuchMethod(t *testing.T, name string, err error) { diff --git a/internal/entity/models/ollama.go b/internal/entity/models/ollama.go index 5abcf4a468..bf40b9d3ea 100644 --- a/internal/entity/models/ollama.go +++ b/internal/entity/models/ollama.go @@ -40,7 +40,7 @@ func NewOllamaModel(baseURL map[string]string, urlSuffix URLSuffix) *OllamaModel BaseURL: baseURL, URLSuffix: urlSuffix, AllowEmptyAPIKey: true, - httpClient: NewDriverHTTPClient(true), + httpClient: common.GetSchemeSafeHTTPClient(), }, } } diff --git a/internal/handler/chat_audio_test.go b/internal/handler/chat_audio_test.go index c8efc8924a..ecd82f7d29 100644 --- a/internal/handler/chat_audio_test.go +++ b/internal/handler/chat_audio_test.go @@ -15,7 +15,6 @@ import ( "ragflow/internal/dao" "ragflow/internal/entity" "ragflow/internal/service" - "ragflow/internal/utility" ) // The chat TTS response may carry model/driver context but never the provider's @@ -28,9 +27,9 @@ func TestChatAudioSpeechFailureOmitsProviderError(t *testing.T) { })) defer server.Close() - prevAllow := utility.AllowAnyHostForTest - utility.AllowAnyHostForTest = true - t.Cleanup(func() { utility.AllowAnyHostForTest = prevAllow }) + prevAllow := common.AllowAnyHostForTest + common.AllowAnyHostForTest = true + t.Cleanup(func() { common.AllowAnyHostForTest = prevAllow }) db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{TranslateError: true}) if err != nil { diff --git a/internal/ingestion/component/ssrf_test_helper_test.go b/internal/ingestion/component/ssrf_test_helper_test.go index 9bb883efe8..d714d61152 100644 --- a/internal/ingestion/component/ssrf_test_helper_test.go +++ b/internal/ingestion/component/ssrf_test_helper_test.go @@ -3,7 +3,7 @@ package component import ( "testing" - "ragflow/internal/utility" + "ragflow/internal/common" ) // withSSRFBypass enables the test-only SSRF bypass for the duration of a test @@ -13,7 +13,7 @@ import ( // coverage in internal/utility/ssrf_test.go. func withSSRFBypass(t *testing.T) { t.Helper() - prev := utility.AllowAnyHostForTest - utility.AllowAnyHostForTest = true - t.Cleanup(func() { utility.AllowAnyHostForTest = prev }) + prev := common.AllowAnyHostForTest + common.AllowAnyHostForTest = true + t.Cleanup(func() { common.AllowAnyHostForTest = prev }) } diff --git a/internal/parser/parser/ssrf_test_helper_test.go b/internal/parser/parser/ssrf_test_helper_test.go index a0b76f4604..eea94a5306 100644 --- a/internal/parser/parser/ssrf_test_helper_test.go +++ b/internal/parser/parser/ssrf_test_helper_test.go @@ -1,9 +1,8 @@ package parser import ( + "ragflow/internal/common" "testing" - - "ragflow/internal/utility" ) // withSSRFBypass enables the test-only SSRF bypass for the duration of a test @@ -13,7 +12,7 @@ import ( // internal/utility/ssrf_test.go. func withSSRFBypass(t *testing.T) { t.Helper() - prev := utility.AllowAnyHostForTest - utility.AllowAnyHostForTest = true - t.Cleanup(func() { utility.AllowAnyHostForTest = prev }) + prev := common.AllowAnyHostForTest + common.AllowAnyHostForTest = true + t.Cleanup(func() { common.AllowAnyHostForTest = prev }) } diff --git a/internal/service/agent_dbcheck.go b/internal/service/agent_dbcheck.go index 3f56a66677..3a71460e19 100644 --- a/internal/service/agent_dbcheck.go +++ b/internal/service/agent_dbcheck.go @@ -44,7 +44,7 @@ type TestDBConnectionRequest struct { Password string `json:"password"` } -// AllowAnyHostForTest mirrors utility.AllowAnyHostForTest: a test-only +// AllowAnyHostForTest mirrors common.AllowAnyHostForTest: a test-only // override that disables the SSRF guard in AssertHostIsSafe. Production // code MUST leave this at false. Tests that need to talk to a local // httptest server or stub-resolved DB host flip it on and reset it diff --git a/internal/service/file/file.go b/internal/service/file/file.go index 3088835069..2bf316119e 100644 --- a/internal/service/file/file.go +++ b/internal/service/file/file.go @@ -19,6 +19,7 @@ package file import ( "context" "errors" + "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/entity" "ragflow/internal/utility" @@ -30,7 +31,7 @@ var ErrNoAuthorization = errors.New("no authorization") var ( // assertURLSafe and pinnedHTTPClient are aliased from utility so tests // can override them for mock injection. - assertURLSafe = utility.AssertURLSafe + assertURLSafe = common.AssertURLSafe pinnedHTTPClient = utility.PinnedHTTPClient ) diff --git a/internal/service/file/file_test.go b/internal/service/file/file_test.go index a2b3fa66de..c79385f384 100644 --- a/internal/service/file/file_test.go +++ b/internal/service/file/file_test.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "net/http/httptest" + "ragflow/internal/common" "strings" "testing" "time" @@ -366,16 +367,16 @@ func TestFileService_UploadFromURL_PDFAddsExtensionAndStoresToDownloads(t *testi })) defer server.Close() - origAssert := utility.AssertURLSafe + origAssert := common.AssertURLSafe origPinned := utility.PinnedHTTPClient - utility.AssertURLSafe = func(rawURL string) (string, string, error) { + common.AssertURLSafe = func(rawURL string) (string, string, error) { return "127.0.0.1", "127.0.0.1", nil } utility.PinnedHTTPClient = func(hostname, resolvedIP string, timeout time.Duration) *http.Client { return server.Client() } t.Cleanup(func() { - utility.AssertURLSafe = origAssert + common.AssertURLSafe = origAssert utility.PinnedHTTPClient = origPinned }) @@ -413,16 +414,16 @@ func TestFileService_UploadFromURL_HTMLNormalizesReadableContent(t *testing.T) { })) defer server.Close() - origAssert := utility.AssertURLSafe + origAssert := common.AssertURLSafe origPinned := utility.PinnedHTTPClient - utility.AssertURLSafe = func(rawURL string) (string, string, error) { + common.AssertURLSafe = func(rawURL string) (string, string, error) { return "127.0.0.1", "127.0.0.1", nil } utility.PinnedHTTPClient = func(hostname, resolvedIP string, timeout time.Duration) *http.Client { return server.Client() } t.Cleanup(func() { - utility.AssertURLSafe = origAssert + common.AssertURLSafe = origAssert utility.PinnedHTTPClient = origPinned }) diff --git a/internal/service/mcp.go b/internal/service/mcp.go index b2d2c28c33..b5f65a3e5d 100644 --- a/internal/service/mcp.go +++ b/internal/service/mcp.go @@ -716,7 +716,7 @@ func (s *MCPService) TestServer(ctx context.Context, mcpID string, req *TestServ // ErrMCPInvalidURL data errors instead of being swallowed inside the // generic FetchTools error and re-classified by the handler as a 500. // FetchTools repeats the check internally; the second call is cheap. - if _, _, err := utility.AssertURLSafe(req.URL); err != nil { + if _, _, err := common.AssertURLSafe(req.URL); err != nil { return nil, fmt.Errorf("%w: %w", ErrMCPInvalidURL, err) } diff --git a/internal/syncer/connector/azure_blob.go b/internal/syncer/connector/azure_blob.go index 964fa343f8..b1124e71ce 100644 --- a/internal/syncer/connector/azure_blob.go +++ b/internal/syncer/connector/azure_blob.go @@ -23,6 +23,7 @@ import ( "fmt" "io" "net/http" + "ragflow/internal/common" "sort" "strings" "time" @@ -30,8 +31,6 @@ import ( "github.com/Azure/azure-sdk-for-go/sdk/azcore" "github.com/Azure/azure-sdk-for-go/sdk/storage/azblob" "github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/container" - - "ragflow/internal/utility" ) const ( @@ -42,9 +41,9 @@ const ( ) // azureAssertURLSafe is the SSRF guard used for caller-supplied endpoint URLs. -// It is an indirection over utility.AssertURLSafe so unit tests can substitute +// It is an indirection over common.AssertURLSafe so unit tests can substitute // a stub without touching the shared global guard. -var azureAssertURLSafe = utility.AssertURLSafe +var azureAssertURLSafe = common.AssertURLSafe // azureBlobTextExtensions mirrors the extension set used by the OneDrive // family of file sources in the blob connector. diff --git a/internal/syncer/connector/bitbucket_test.go b/internal/syncer/connector/bitbucket_test.go index 071528e402..90adc3c96b 100644 --- a/internal/syncer/connector/bitbucket_test.go +++ b/internal/syncer/connector/bitbucket_test.go @@ -8,12 +8,11 @@ import ( "net/http" "net/http/httptest" "net/url" + "ragflow/internal/common" "strings" "sync/atomic" "testing" "time" - - "ragflow/internal/utility" ) func TestNewBitbucketConnectorParsesConfig(t *testing.T) { @@ -480,18 +479,18 @@ func TestBitbucketConnectorValidateMissingConfig(t *testing.T) { } func TestBitbucketGetJSONRetriesOnTransientStatuses(t *testing.T) { - originalAllowAny := utility.AllowAnyHostForTest + originalAllowAny := common.AllowAnyHostForTest originalTries := bitbucketRetryTries originalDelay := bitbucketRetryBaseDelay originalBackoff := bitbucketRetryBackoff originalMaxDelay := bitbucketRetryMaxDelay - utility.AllowAnyHostForTest = true + common.AllowAnyHostForTest = true bitbucketRetryTries = 3 bitbucketRetryBaseDelay = time.Millisecond bitbucketRetryBackoff = 1 bitbucketRetryMaxDelay = 10 * time.Millisecond t.Cleanup(func() { - utility.AllowAnyHostForTest = originalAllowAny + common.AllowAnyHostForTest = originalAllowAny bitbucketRetryTries = originalTries bitbucketRetryBaseDelay = originalDelay bitbucketRetryBackoff = originalBackoff diff --git a/internal/syncer/connector/rest_api_test.go b/internal/syncer/connector/rest_api_test.go index a1d37ac08e..0218db06c5 100644 --- a/internal/syncer/connector/rest_api_test.go +++ b/internal/syncer/connector/rest_api_test.go @@ -24,13 +24,12 @@ import ( "net" "net/http" "net/http/httptest" + "ragflow/internal/common" "strings" "sync" "sync/atomic" "testing" "time" - - "ragflow/internal/utility" ) // restAPITestPublicIP is the fixed public address used to satisfy the @@ -51,7 +50,7 @@ func restAPITestLookupHost(host string) ([]string, error) { func withRestAPITestHooks(t *testing.T) { t.Helper() origLoopback := connectorAllowLoopbackForTest - origLookup := utility.LookupHost + origLookup := common.LookupHost origTries := restAPIRetryTries origBaseDelay := restAPIRetryBaseDelay origMaxDelay := restAPIRetryMaxDelay @@ -60,7 +59,7 @@ func withRestAPITestHooks(t *testing.T) { orig429Waits := restAPI429MaxWaits orig429Wait := restAPI429DefaultWait connectorAllowLoopbackForTest = true - utility.LookupHost = restAPITestLookupHost + common.LookupHost = restAPITestLookupHost restAPIRetryTries = 3 restAPIRetryBaseDelay = time.Millisecond restAPIRetryMaxDelay = 10 * time.Millisecond @@ -70,7 +69,7 @@ func withRestAPITestHooks(t *testing.T) { restAPI429DefaultWait = time.Millisecond t.Cleanup(func() { connectorAllowLoopbackForTest = origLoopback - utility.LookupHost = origLookup + common.LookupHost = origLookup restAPIRetryTries = origTries restAPIRetryBaseDelay = origBaseDelay restAPIRetryMaxDelay = origMaxDelay diff --git a/internal/syncer/connector/s3_compatible_test.go b/internal/syncer/connector/s3_compatible_test.go index 6cd720df87..3ce4c94baf 100644 --- a/internal/syncer/connector/s3_compatible_test.go +++ b/internal/syncer/connector/s3_compatible_test.go @@ -5,10 +5,9 @@ import ( "context" "errors" "io" + "ragflow/internal/common" "strings" "testing" - - "ragflow/internal/utility" ) func TestS3CompatibleConnectorOpenSyncUsesFingerprintAndFetch(t *testing.T) { @@ -426,14 +425,14 @@ func TestS3CompatibleHelpers(t *testing.T) { // depend on DNS. func stubS3EndpointResolver(t *testing.T) { t.Helper() - orig := utility.LookupHost - utility.LookupHost = func(host string) ([]string, error) { + orig := common.LookupHost + common.LookupHost = func(host string) ([]string, error) { if host == "s3.example.com" { return []string{"93.184.216.34"}, nil } return orig(host) } - t.Cleanup(func() { utility.LookupHost = orig }) + t.Cleanup(func() { common.LookupHost = orig }) } func newTestS3CompatibleConnector(t *testing.T, objects []s3Object) *S3CompatibleConnector { diff --git a/internal/syncer/connector/salesforce_test.go b/internal/syncer/connector/salesforce_test.go index 413afe85c4..cc5d38405d 100644 --- a/internal/syncer/connector/salesforce_test.go +++ b/internal/syncer/connector/salesforce_test.go @@ -23,11 +23,10 @@ import ( "io" "net/http" "net/url" + "ragflow/internal/common" "strings" "testing" "time" - - "ragflow/internal/utility" ) func TestNewSalesforceConnectorDefaults(t *testing.T) { @@ -1020,11 +1019,11 @@ func TestSalesforceHostAllowed(t *testing.T) { } func TestRequestAccessTokenRejectsNonSalesforceHost(t *testing.T) { - origLookup := utility.LookupHost - utility.LookupHost = func(host string) ([]string, error) { + origLookup := common.LookupHost + common.LookupHost = func(host string) ([]string, error) { return []string{"93.184.216.34"}, nil } - t.Cleanup(func() { utility.LookupHost = origLookup }) + t.Cleanup(func() { common.LookupHost = origLookup }) connector := &SalesforceConnector{ instanceURL: "https://evil.example.com", @@ -1039,11 +1038,11 @@ func TestRequestAccessTokenRejectsNonSalesforceHost(t *testing.T) { } func TestRequestAccessTokenRequiresHTTPS(t *testing.T) { - origLookup := utility.LookupHost - utility.LookupHost = func(host string) ([]string, error) { + origLookup := common.LookupHost + common.LookupHost = func(host string) ([]string, error) { return []string{"93.184.216.34"}, nil } - t.Cleanup(func() { utility.LookupHost = origLookup }) + t.Cleanup(func() { common.LookupHost = origLookup }) connector := &SalesforceConnector{ instanceURL: "http://acme.my.salesforce.com", diff --git a/internal/syncer/connector/sitemap.go b/internal/syncer/connector/sitemap.go index 2e25edf1e8..cc83e13744 100644 --- a/internal/syncer/connector/sitemap.go +++ b/internal/syncer/connector/sitemap.go @@ -27,6 +27,7 @@ import ( "log/slog" "net/http" "net/url" + "ragflow/internal/common" "regexp" "strings" "time" @@ -50,7 +51,7 @@ const ( type sitemapFetchFunc func(ctx context.Context, rawURL string) ([]byte, string, error) // sitemapAssertURLSafe is the SSRF guard, indirected so tests can stub DNS resolution. -var sitemapAssertURLSafe = utility.AssertURLSafe +var sitemapAssertURLSafe = common.AssertURLSafe // SitemapConnector ingests the web pages listed in a sitemap.xml. // diff --git a/internal/syncer/connector/ssrf.go b/internal/syncer/connector/ssrf.go index c07498797b..27ba85bb86 100644 --- a/internal/syncer/connector/ssrf.go +++ b/internal/syncer/connector/ssrf.go @@ -27,23 +27,22 @@ import ( "net" "net/http" "net/url" + "ragflow/internal/common" "strings" "time" - - "ragflow/internal/utility" ) // connectorAssertURLSafe is the SSRF guard shared by all connectors. It is an -// indirection over utility.AssertURLSafe so unit tests can substitute a stub +// indirection over common.AssertURLSafe so unit tests can substitute a stub // without touching the shared utility guard (same pattern as sitemap.go and // azure_blob.go). -var connectorAssertURLSafe = utility.AssertURLSafe +var connectorAssertURLSafe = common.AssertURLSafe // connectorAssertHostSafe is the host-type SSRF guard shared by the host-based // connectors (IMAP/MySQL/PostgreSQL). It is an indirection over // utility.AssertHostSafe so unit tests can substitute a stub without touching // the shared utility guard (same pattern as connectorAssertURLSafe). -var connectorAssertHostSafe = utility.AssertHostSafe +var connectorAssertHostSafe = common.AssertHostSafe // assertConnectorHostSafe validates a bare host (hostname or literal IP, no // scheme or port) with the shared strict SSRF guard and returns the first @@ -132,7 +131,7 @@ func loopbackTestAllow(hostname string) (string, net.IP, bool) { if lower != "localhost" && !strings.HasSuffix(lower, ".localhost") { return "", nil, false } - addrs, err := utility.LookupHost(hostname) + addrs, err := common.LookupHost(hostname) if err != nil { return "", nil, false } diff --git a/internal/syncer/connector/ssrf_test.go b/internal/syncer/connector/ssrf_test.go index 1dae568656..47ee859d09 100644 --- a/internal/syncer/connector/ssrf_test.go +++ b/internal/syncer/connector/ssrf_test.go @@ -28,7 +28,7 @@ import ( "testing" "time" - "ragflow/internal/utility" + "ragflow/internal/common" ) // withConnectorLoopbackTestHook enables the loopback test seam for the duration @@ -147,9 +147,9 @@ func TestJiraConnectorRejectsInternalBaseURL(t *testing.T) { func TestConnectorRequestPinsRedirectHopAgainstDNSRebinding(t *testing.T) { withConnectorLoopbackTestHook(t) - origLookup := utility.LookupHost + origLookup := common.LookupHost lookups := 0 - utility.LookupHost = func(host string) ([]string, error) { + common.LookupHost = func(host string) ([]string, error) { if host == "localhost" { lookups++ if lookups == 1 { @@ -161,7 +161,7 @@ func TestConnectorRequestPinsRedirectHopAgainstDNSRebinding(t *testing.T) { } return origLookup(host) } - t.Cleanup(func() { utility.LookupHost = origLookup }) + t.Cleanup(func() { common.LookupHost = origLookup }) var targetHit atomic.Bool target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -250,11 +250,11 @@ func TestConnectorRequestRejectsCrossOriginBodyRedirect(t *testing.T) { } func TestAssertConnectorURLSafeHTTPS(t *testing.T) { - orig := utility.LookupHost - utility.LookupHost = func(host string) ([]string, error) { + orig := common.LookupHost + common.LookupHost = func(host string) ([]string, error) { return []string{"93.184.216.34"}, nil } - t.Cleanup(func() { utility.LookupHost = orig }) + t.Cleanup(func() { common.LookupHost = orig }) connectorAllowLoopbackForTest = false t.Cleanup(func() { connectorAllowLoopbackForTest = false }) diff --git a/internal/syncer/connector/teams_test.go b/internal/syncer/connector/teams_test.go index e80b20c587..d9f4adecef 100644 --- a/internal/syncer/connector/teams_test.go +++ b/internal/syncer/connector/teams_test.go @@ -22,11 +22,10 @@ import ( "io" "net/http" "net/http/httptest" + "ragflow/internal/common" "strings" "testing" "time" - - "ragflow/internal/utility" ) func TestNewTeamsConnectorDefaults(t *testing.T) { @@ -114,9 +113,9 @@ func TestTeamsConnectorValidateQueriesTeams(t *testing.T) { } func TestTeamsGetJSONReadsBodyBeforeCancel(t *testing.T) { - previousAllowAnyHost := utility.AllowAnyHostForTest - utility.AllowAnyHostForTest = true - t.Cleanup(func() { utility.AllowAnyHostForTest = previousAllowAnyHost }) + previousAllowAnyHost := common.AllowAnyHostForTest + common.AllowAnyHostForTest = true + t.Cleanup(func() { common.AllowAnyHostForTest = previousAllowAnyHost }) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") diff --git a/internal/utility/mcp_call.go b/internal/utility/mcp_call.go index 68e0ad9b0f..9dce50b50b 100644 --- a/internal/utility/mcp_call.go +++ b/internal/utility/mcp_call.go @@ -22,6 +22,7 @@ import ( "errors" "fmt" "net/http" + "ragflow/internal/common" "time" ) @@ -72,7 +73,7 @@ func CallTool(ctx context.Context, opts CallOptions) (*CallResult, error) { if opts.Timeout <= 0 { opts.Timeout = 10 * time.Second } - hostname, resolvedIP, err := AssertURLSafe(opts.URL) + hostname, resolvedIP, err := common.AssertURLSafe(opts.URL) if err != nil { return nil, err } diff --git a/internal/utility/mcp_client.go b/internal/utility/mcp_client.go index 9b26677b4c..7fb49fbc08 100644 --- a/internal/utility/mcp_client.go +++ b/internal/utility/mcp_client.go @@ -40,6 +40,7 @@ import ( "io" "net/http" "net/url" + "ragflow/internal/common" "strings" "sync" "time" @@ -92,7 +93,7 @@ func FetchTools(ctx context.Context, opts FetchOptions) ([]Tool, error) { opts.Timeout = 10 * time.Second } - hostname, resolvedIP, err := AssertURLSafe(opts.URL) + hostname, resolvedIP, err := common.AssertURLSafe(opts.URL) if err != nil { return nil, err } @@ -414,7 +415,7 @@ func requestSSE(ctx context.Context, endpoint string, headers map[string]string, // differs from the original SSE host — swap in a fresh pinned // client so the dial-time IP override still applies. postClient := client - if postHost, postIP, vErr := AssertURLSafe(postURL); vErr != nil { + if postHost, postIP, vErr := common.AssertURLSafe(postURL); vErr != nil { return nil, vErr } else if u, perr := url.Parse(postURL); perr == nil && u.Hostname() != "" { if u.Hostname() != originalHost(endpoint) { diff --git a/internal/utility/mcp_client_test.go b/internal/utility/mcp_client_test.go index 8293998b4c..3b7adfeaa9 100644 --- a/internal/utility/mcp_client_test.go +++ b/internal/utility/mcp_client_test.go @@ -22,6 +22,7 @@ import ( "io" "net/http" "net/http/httptest" + "ragflow/internal/common" "strings" "sync" "sync/atomic" @@ -33,13 +34,13 @@ import ( // targets used by httptest are accepted by AssertURLSafe. func allowLoopbackForTests(t *testing.T) func() { t.Helper() - orig := LookupHost - LookupHost = func(host string) ([]string, error) { + orig := common.LookupHost + common.LookupHost = func(host string) ([]string, error) { // Return a public IPv4 so the guard sees the host as global; the // httptest server is on loopback but we connect via raw URL. return []string{"8.8.8.8"}, nil } - return func() { LookupHost = orig } + return func() { common.LookupHost = orig } } func TestStreamableSessionCleanupBudget(t *testing.T) { diff --git a/internal/utility/ssrf.go b/internal/utility/ssrf.go index e20a36d9f0..6fa1a6455c 100644 --- a/internal/utility/ssrf.go +++ b/internal/utility/ssrf.go @@ -18,281 +18,11 @@ package utility import ( "context" - "fmt" "net" "net/http" - "net/url" - "slices" - "sort" - "strings" "time" ) -// AllowedURLSchemes are the schemes accepted by AssertURLSafe. -var AllowedURLSchemes = []string{"http", "https"} - -// LookupHost is the indirection used to resolve hostnames. Tests override it. -var LookupHost = net.LookupHost - -// AllowAnyHostForTest is a test-only override that skips the -// public-IP routability check in AssertURLSafe and AssertHostSafe. -// Scheme, host, and DNS resolution checks are unchanged: hostnames -// still resolve, unresolvable hostnames still fail, and callers -// still pin connections to the returned resolved address. Production -// code MUST leave this at its zero value (false). Tests that need to -// talk to a local httptest server flip it on and reset it in -// t.Cleanup. -// -// The previous form (env-var ALLOW_ANY_HOST) was a live runtime -// toggle that any operator could flip to disable the SSRF guard -// globally — including the DNS pinning that the Invoke component -// relies on. PR review round 6, Major #3: this variable lives in -// process memory only, so it cannot be enabled by an env var or -// a deployment mistake. The explicit "_ForTest" suffix is the -// signal that production code must never touch it. -var AllowAnyHostForTest = false - -// allowAnyHost reads the test-only override. Kept as a private -// helper so the call sites don't all have to know about the -// exported variable name. -func allowAnyHost() bool { - return AllowAnyHostForTest -} - -// AssertURLSafe parses rawURL and rejects it if the scheme is disallowed, -// the host is missing, or any resolved IP is not globally routable -// (private, loopback, link-local, multicast, reserved). Returns the hostname -// and the first validated public IP so callers can DNS-pin the address and -// prevent rebinding between validation and the actual TCP connection. -// -// Mirrors common/ssrf_guard.py:assert_url_is_safe. -var AssertURLSafe = func(rawURL string) (hostname, resolvedIP string, err error) { - parsed, err := url.Parse(strings.TrimSpace(rawURL)) - if err != nil { - return "", "", fmt.Errorf("invalid url") - } - - scheme := strings.ToLower(parsed.Scheme) - if !slices.Contains(AllowedURLSchemes, scheme) { - sorted := append([]string(nil), AllowedURLSchemes...) - sort.Strings(sorted) - return "", "", fmt.Errorf("disallowed URL scheme: '%s'. Only %v are allowed", scheme, sorted) - } - - hostname = parsed.Hostname() - if hostname == "" { - return "", "", fmt.Errorf("URL is missing a host") - } - - allowAny := allowAnyHost() - addresses, err := LookupHost(hostname) - if err != nil { - return "", "", fmt.Errorf("could not resolve hostname '%s': %w", hostname, err) - } - if len(addresses) == 0 { - return "", "", fmt.Errorf("hostname '%s' resolved to no addresses", hostname) - } - - for _, addr := range addresses { - ip := net.ParseIP(addr) - if ip == nil { - return "", "", fmt.Errorf("could not parse resolved address '%s' for hostname '%s'", addr, hostname) - } - if !allowAny && !isGlobalIP(effectiveIP(ip)) { - return "", "", fmt.Errorf("URL resolves to a non-public address (%s), which is not allowed", ip.String()) - } - if resolvedIP == "" { - resolvedIP = ip.String() - } - } - return hostname, resolvedIP, nil -} - -// AssertHostSafe validates a bare host (a hostname or a literal IP, with no -// scheme or port) and returns the first resolved public IP. It is the -// host-type counterpart of AssertURLSafe: every resolved address must be -// globally routable (private, loopback, link-local, metadata, multicast and -// reserved ranges are rejected). Callers dial the returned IP directly so DNS -// cannot rebind the connection to an internal address between validation and -// the TCP connect. -// -// Used by host-based data sources (IMAP/MySQL/PostgreSQL) and by the ExeSQL / -// test_db_connection host guards, mirroring common/ssrf_guard.py: -// assert_host_is_safe. -var AssertHostSafe = func(host string) (resolvedIP string, err error) { - host = strings.TrimSpace(host) - if host == "" { - return "", fmt.Errorf("host is missing") - } - - allowAny := allowAnyHost() - if ip := net.ParseIP(host); ip != nil { - if !allowAny && !isGlobalIP(effectiveIP(ip)) { - return "", fmt.Errorf("host is not a public address (%s), which is not allowed", ip.String()) - } - return ip.String(), nil - } - - addresses, err := LookupHost(host) - if err != nil { - return "", fmt.Errorf("could not resolve hostname '%s': %w", host, err) - } - if len(addresses) == 0 { - return "", fmt.Errorf("hostname '%s' resolved to no addresses", host) - } - - for _, addr := range addresses { - ip := net.ParseIP(addr) - if ip == nil { - return "", fmt.Errorf("could not parse resolved address '%s' for hostname '%s'", addr, host) - } - if !allowAny && !isGlobalIP(effectiveIP(ip)) { - return "", fmt.Errorf("hostname '%s' resolves to a non-public address (%s), which is not allowed", host, ip.String()) - } - if resolvedIP == "" { - resolvedIP = ip.String() - } - } - return resolvedIP, nil -} - -// effectiveIP unwraps IPv4-mapped IPv6 addresses (e.g. ::ffff:127.0.0.1) so -// the routability check sees the IPv4 form. Without this, an attacker could -// bypass the guard with an IPv4-mapped IPv6 representation of a private host. -func effectiveIP(ip net.IP) net.IP { - if v4 := ip.To4(); v4 != nil { - return v4 - } - return ip -} - -// isGlobalIP mirrors Python's ipaddress.IPv*Address.is_global: an address is -// global if it is none of {unspecified, loopback, multicast, link-local, -// private (including CGNAT and IPv6 ULA), benchmarking, documentation, -// reserved}. -func isGlobalIP(ip net.IP) bool { - if ip == nil || ip.IsUnspecified() || ip.IsLoopback() || ip.IsMulticast() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsInterfaceLocalMulticast() || ip.IsPrivate() { - return false - } - if v4 := ip.To4(); v4 != nil { - // 0.0.0.0/8 — "this network"; 0.x.y.z routes to localhost on Linux. - if v4[0] == 0 { - return false - } - // CGNAT 100.64.0.0/10 — not flagged by IsPrivate in older Go versions. - if v4[0] == 100 && v4[1]&0xC0 == 64 { - return false - } - // 192.0.0.0/24 reserved for IETF protocol assignments. - if v4[0] == 192 && v4[1] == 0 && v4[2] == 0 { - return false - } - // 192.0.2.0/24, 198.51.100.0/24, 203.0.113.0/24 documentation (TEST-NET-1/2/3). - if v4[0] == 192 && v4[1] == 0 && v4[2] == 2 { - return false - } - if v4[0] == 198 && v4[1] == 51 && v4[2] == 100 { - return false - } - if v4[0] == 203 && v4[1] == 0 && v4[2] == 113 { - return false - } - // 198.18.0.0/15 benchmarking. - if v4[0] == 198 && (v4[1] == 18 || v4[1] == 19) { - return false - } - // 240.0.0.0/4 reserved (excluding 255.255.255.255 which IsUnspecified misses). - if v4[0] >= 240 { - return false - } - } else if v6 := ip.To16(); v6 != nil { - // 2001:db8::/32 documentation prefix. - if v6[0] == 0x20 && v6[1] == 0x01 && v6[2] == 0x0d && v6[3] == 0xb8 { - return false - } - // 100::/64 discard-only address block. - if v6[0] == 0x01 && v6[1] == 0x00 && allZero(v6[2:8]) { - return false - } - // IPv6 transition addresses (6to4, NAT64, Teredo, IPv4-compatible) embed - // an arbitrary IPv4 address that none of the checks above look at. Unwrap - // and re-check it so 2002:7f00:1::1 is treated as 127.0.0.1. - for _, inner := range embeddedIPv4(v6) { - if !isGlobalIP(inner) { - return false - } - } - } - return true -} - -// embeddedIPv4 returns the IPv4 addresses carried inside an IPv6 transition -// address, or nil when it carries none. Teredo yields two: the relay server and -// the (obfuscated) client. -func embeddedIPv4(v6 net.IP) []net.IP { - switch { - // 6to4 — RFC 3056, 2002::/16, IPv4 in bytes 2-6. - case v6[0] == 0x20 && v6[1] == 0x02: - return []net.IP{net.IPv4(v6[2], v6[3], v6[4], v6[5])} - - // NAT64 well-known prefix — RFC 6052, 64:ff9b::/96, IPv4 in the low 32 bits. - case v6[0] == 0x00 && v6[1] == 0x64 && v6[2] == 0xff && v6[3] == 0x9b && allZero(v6[4:12]): - return []net.IP{net.IPv4(v6[12], v6[13], v6[14], v6[15])} - - // NAT64 local-use prefix — RFC 8215, 64:ff9b:1::/48. The embedded IPv4 - // position depends on the operator's prefix length, so block the range. - case v6[0] == 0x00 && v6[1] == 0x64 && v6[2] == 0xff && v6[3] == 0x9b && v6[4] == 0x00 && v6[5] == 0x01: - return []net.IP{net.IPv4zero} - - // Teredo — RFC 4380, 2001::/32. Server IPv4 in bytes 4-8, client IPv4 in - // bytes 12-16 obfuscated by XOR with 0xff. - case v6[0] == 0x20 && v6[1] == 0x01 && v6[2] == 0x00 && v6[3] == 0x00: - return []net.IP{ - net.IPv4(v6[4], v6[5], v6[6], v6[7]), - net.IPv4(v6[12]^0xff, v6[13]^0xff, v6[14]^0xff, v6[15]^0xff), - } - - // IPv4-compatible — deprecated ::a.b.c.d, not unwrapped by net.IP.To4. - case allZero(v6[0:12]): - return []net.IP{net.IPv4(v6[12], v6[13], v6[14], v6[15])} - } - return nil -} - -func allZero(b []byte) bool { - for _, x := range b { - if x != 0 { - return false - } - } - return true -} - -// AssertURLSchemeSafe is a lenient SSRF guard for drivers that may legitimately -// target private networks or loopback addresses (e.g. self-hosted Ollama, vLLM, -// Xinference). It only rejects dangerous schemes and empty hosts; it does not -// resolve DNS and does not require public routability. Use this ONLY for -// local-inference model drivers — cloud-hosted drivers must use AssertURLSafe. -var AssertURLSchemeSafe = func(rawURL string) error { - parsed, err := url.Parse(strings.TrimSpace(rawURL)) - if err != nil { - return fmt.Errorf("invalid url") - } - - scheme := strings.ToLower(parsed.Scheme) - if !slices.Contains(AllowedURLSchemes, scheme) { - sorted := append([]string(nil), AllowedURLSchemes...) - sort.Strings(sorted) - return fmt.Errorf("disallowed URL scheme: '%s'. Only %v are allowed", scheme, sorted) - } - - if parsed.Hostname() == "" { - return fmt.Errorf("URL is missing a host") - } - - return nil -} - // PinnedHTTPClient returns an HTTP client whose Transport rewrites every // outbound dial for hostname:port to resolvedIP:port, closing the TOCTOU // window between AssertURLSafe and the actual TCP connection. Pins are diff --git a/internal/utility/ssrf_test.go b/internal/utility/ssrf_test.go index d4d0c0bc6e..655dc464e9 100644 --- a/internal/utility/ssrf_test.go +++ b/internal/utility/ssrf_test.go @@ -17,13 +17,14 @@ package utility import ( + "ragflow/internal/common" "strings" "testing" ) func TestAssertURLSafe(t *testing.T) { - orig := LookupHost - defer func() { LookupHost = orig }() + orig := common.LookupHost + defer func() { common.LookupHost = orig }() type want struct { errSubstr string @@ -165,15 +166,14 @@ func TestAssertURLSafe(t *testing.T) { } for _, tc := range cases { - tc := tc t.Run(tc.name, func(t *testing.T) { - LookupHost = func(host string) ([]string, error) { + common.LookupHost = func(host string) ([]string, error) { if tc.err != "" { return nil, &mockErr{tc.err} } return tc.ips, nil } - host, ip, err := AssertURLSafe(tc.url) + host, ip, err := common.AssertURLSafe(tc.url) if tc.want.errSubstr != "" { if err == nil || !strings.Contains(err.Error(), tc.want.errSubstr) { t.Fatalf("expected error containing %q, got %v", tc.want.errSubstr, err) diff --git a/internal/utility/upload.go b/internal/utility/upload.go index ea72cd1473..07394b4617 100644 --- a/internal/utility/upload.go +++ b/internal/utility/upload.go @@ -23,6 +23,7 @@ import ( "io" "net/http" "net/url" + "ragflow/internal/common" "regexp" "strings" "time" @@ -49,7 +50,7 @@ func FetchRemoteFileSafelyWithTimeout(ctx context.Context, rawURL string, maxSiz } currentURL := rawURL for redirects := 0; redirects < 10; redirects++ { - hostname, resolvedIP, err := AssertURLSafe(currentURL) + hostname, resolvedIP, err := common.AssertURLSafe(currentURL) if err != nil { return nil, nil, "", err } diff --git a/internal/utility/upload_test.go b/internal/utility/upload_test.go index fee42df9e3..96914df965 100644 --- a/internal/utility/upload_test.go +++ b/internal/utility/upload_test.go @@ -20,6 +20,7 @@ import ( "bytes" "net/http" "net/http/httptest" + "ragflow/internal/common" "strings" "testing" "time" @@ -42,16 +43,16 @@ func TestFetchRemoteFileSafely_PDFAddsExtension(t *testing.T) { })) defer server.Close() - origAssert := AssertURLSafe + origAssert := common.AssertURLSafe origPinned := PinnedHTTPClient - AssertURLSafe = func(rawURL string) (string, string, error) { + common.AssertURLSafe = func(rawURL string) (string, string, error) { return "127.0.0.1", "127.0.0.1", nil } PinnedHTTPClient = func(hostname, resolvedIP string, timeout time.Duration) *http.Client { return server.Client() } t.Cleanup(func() { - AssertURLSafe = origAssert + common.AssertURLSafe = origAssert PinnedHTTPClient = origPinned }) @@ -75,16 +76,16 @@ func TestFetchRemoteFileSafely_ReturnsContentAndHeaders(t *testing.T) { })) defer server.Close() - origAssert := AssertURLSafe + origAssert := common.AssertURLSafe origPinned := PinnedHTTPClient - AssertURLSafe = func(rawURL string) (string, string, error) { + common.AssertURLSafe = func(rawURL string) (string, string, error) { return "127.0.0.1", "127.0.0.1", nil } PinnedHTTPClient = func(hostname, resolvedIP string, timeout time.Duration) *http.Client { return server.Client() } t.Cleanup(func() { - AssertURLSafe = origAssert + common.AssertURLSafe = origAssert PinnedHTTPClient = origPinned })