zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-01 21:37:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: fix SSRF (#17641)

Browse Source 
### Summary

Check the URL to prevent SSRF attack.

---------

Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-07-31 19:15:38 +08:00
committed by GitHub
parent f03a00c54c
commit b058229ab3
146 changed files with 833 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
internal/entity/models/cometapi_test.go
View File

@@ -77,6 +77,7 @@ func TestCometAPIFactoryRoute(t *testing.T) {
}
func TestCometAPIChatHappyPath(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newCometAPIServer(t, "/v1/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) {
if body["model"] != "gpt-5" {
@@ -115,6 +116,7 @@ func TestCometAPIChatHappyPath(t *testing.T) {
}
func TestCometAPIChatPropagatesConfig(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newCometAPIServer(t, "/v1/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) {
if body["max_tokens"] != float64(64) {
@@ -152,6 +154,7 @@ func TestCometAPIChatPropagatesConfig(t *testing.T) {
}
func TestCometAPIChatReturnsReasoningContent(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newCometAPIServer(t, "/v1/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
@@ -174,6 +177,7 @@ func TestCometAPIChatReturnsReasoningContent(t *testing.T) {
}
func TestCometAPIChatRequiresAPIKey(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
m := newCometAPIForTest("http://unused")
_, err := m.ChatWithMessages(ctx, "gpt-5", []Message{{Role: "user", Content: "x"}}, &APIConfig{}, nil, nil)
@@ -188,6 +192,7 @@ func TestCometAPIChatRequiresAPIKey(t *testing.T) {
}
func TestCometAPIChatRequiresModelName(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
m := newCometAPIForTest("http://unused")
apiKey := "test-key"
@@ -202,6 +207,7 @@ func TestCometAPIChatRequiresModelName(t *testing.T) {
}
func TestCometAPIChatRequiresMessages(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
m := newCometAPIForTest("http://unused")
apiKey := "test-key"
@@ -212,6 +218,7 @@ func TestCometAPIChatRequiresMessages(t *testing.T) {
}
func TestCometAPIChatRejectsHTTPError(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newCometAPIServer(t, "/v1/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) {
w.WriteHeader(http.StatusUnauthorized)
@@ -228,6 +235,7 @@ func TestCometAPIChatRejectsHTTPError(t *testing.T) {
}
func TestCometAPIChatFallsBackToDefaultOnEmptyRegion(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
// Empty *Region pointer must fall back to the "default" entry, not
// be treated as an explicit "" region (which would miss the lookup).
@@ -250,6 +258,7 @@ func TestCometAPIChatFallsBackToDefaultOnEmptyRegion(t *testing.T) {
}
func TestCometAPIListModelsFallsBackToDefaultOnEmptyRegion(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newCometAPIServer(t, "/api/models", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{"data": []map[string]interface{}{{"id": "x"}}})
@@ -265,6 +274,7 @@ func TestCometAPIListModelsFallsBackToDefaultOnEmptyRegion(t *testing.T) {
}
func TestCometAPIStreamRequiresSender(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
m := newCometAPIForTest("http://unused")
apiKey := "test-key"
@@ -277,6 +287,7 @@ func TestCometAPIStreamRequiresSender(t *testing.T) {
}
func TestCometAPIChatRejectsUnknownRegion(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
m := newCometAPIForTest("http://unused")
apiKey := "test-key"
@@ -289,6 +300,7 @@ func TestCometAPIChatRejectsUnknownRegion(t *testing.T) {
}
func TestCometAPIBaseURLNormalizesSlashes(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
tests := []struct {
name string
@@ -359,6 +371,7 @@ func TestCometAPIBaseURLNormalizesSlashes(t *testing.T) {
}
func TestCometAPIStreamHappyPath(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newCometAPIServer(t, "/v1/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) {
if body["stream"] != true {
@@ -407,6 +420,7 @@ func TestCometAPIStreamHappyPath(t *testing.T) {
}
func TestCometAPIStreamRejectsExplicitFalse(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
m := newCometAPIForTest("http://unused")
apiKey := "test-key"
@@ -424,6 +438,7 @@ func TestCometAPIStreamRejectsExplicitFalse(t *testing.T) {
}
func TestCometAPIStreamFailsWithoutTerminal(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
// Body closes before [DONE] or a finish_reason -> driver must complain
// instead of pretending the stream finished cleanly.
@@ -446,6 +461,7 @@ func TestCometAPIStreamFailsWithoutTerminal(t *testing.T) {
}
func TestCometAPIListModelsHappyPath(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newCometAPIServer(t, "/api/models", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
@@ -470,6 +486,7 @@ func TestCometAPIListModelsHappyPath(t *testing.T) {
}
func TestCometAPIListModelsAllowsNilAPIConfig(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newCometAPIServer(t, "/api/models", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{"data": []map[string]interface{}{{"id": "gpt-5"}}})
@@ -487,6 +504,7 @@ func TestCometAPIListModelsAllowsNilAPIConfig(t *testing.T) {
}
func TestCometAPICheckConnectionDelegatesToBalance(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
// 200 -> CheckConnection succeeds; 401 -> CheckConnection propagates.
okSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -518,6 +536,7 @@ func TestCometAPICheckConnectionDelegatesToBalance(t *testing.T) {
}
func TestCometAPIBalanceHappyPath(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/user/quota" {
@@ -551,6 +570,7 @@ func TestCometAPIBalanceHappyPath(t *testing.T) {
}
func TestCometAPIBalanceRequiresAPIKey(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
m := newCometAPIForTest("http://unused")
_, err := m.Balance(ctx, &APIConfig{})
@@ -560,6 +580,7 @@ func TestCometAPIBalanceRequiresAPIKey(t *testing.T) {
}
func TestCometAPIBalanceRequiresConfiguredURL(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
m := newCometAPIForTest("http://unused")
m.baseModel.URLSuffix.Balance = ""
@@ -571,6 +592,7 @@ func TestCometAPIBalanceRequiresConfiguredURL(t *testing.T) {
}
func TestCometAPIRerankReturnsNoSuchMethod(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
m := newCometAPIForTest("http://unused")
q := "gpt-5"
@@ -581,6 +603,7 @@ func TestCometAPIRerankReturnsNoSuchMethod(t *testing.T) {
}
func TestCometAPIEmbedHappyPath(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newCometAPIServer(t, "/v1/embeddings", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) {
if body["model"] != "text-embedding-3-small" {
@@ -619,6 +642,7 @@ func TestCometAPIEmbedHappyPath(t *testing.T) {
}
func TestCometAPIEmbedReordersByIndex(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
// Upstream returns the three vectors in shuffled order. The driver
// must reorder them so the slot at position i corresponds to input i.
@@ -648,6 +672,7 @@ func TestCometAPIEmbedReordersByIndex(t *testing.T) {
}
func TestCometAPIEmbedEmptyInputShortCircuits(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
// Empty input must NOT make an HTTP call; the test fails the request
// rather than the assertion if it does.
@@ -670,6 +695,7 @@ func TestCometAPIEmbedEmptyInputShortCircuits(t *testing.T) {
}
func TestCometAPIEmbedRequiresAPIKey(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
m := newCometAPIForTest("http://unused")
model := "text-embedding-3-small"
@@ -680,6 +706,7 @@ func TestCometAPIEmbedRequiresAPIKey(t *testing.T) {
}
func TestCometAPIEmbedRequiresModelName(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
m := newCometAPIForTest("http://unused")
apiKey := "test-key"
@@ -695,6 +722,7 @@ func TestCometAPIEmbedRequiresModelName(t *testing.T) {
}
func TestCometAPIEmbedRejectsDuplicateIndex(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
// A malformed upstream that repeats data[*].index would silently
// overwrite the earlier vector; the driver must fail loudly instead.
@@ -718,6 +746,7 @@ func TestCometAPIEmbedRejectsDuplicateIndex(t *testing.T) {
}
func TestCometAPIEmbedRejectsOutOfRangeIndex(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newCometAPIServer(t, "/v1/embeddings", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
@@ -738,6 +767,7 @@ func TestCometAPIEmbedRejectsOutOfRangeIndex(t *testing.T) {
}
func TestCometAPIEmbedRejectsMissingSlot(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
// Upstream returns only one of the two requested embeddings.
srv := newCometAPIServer(t, "/v1/embeddings", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
@@ -759,6 +789,7 @@ func TestCometAPIEmbedRejectsMissingSlot(t *testing.T) {
}
func TestCometAPIEmbedRejectsHTTPError(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newCometAPIServer(t, "/v1/embeddings", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
w.WriteHeader(http.StatusUnauthorized)