Skip to content

Commit 1661827

Browse files
lyingbuglyingbug
committed
fix(security): harden model client redirects
1 parent 663d13e commit 1661827

15 files changed

Lines changed: 226 additions & 20 deletions

‎internal/models/asr/openai.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"context"
66
"encoding/json"
77
"fmt"
8-
"net/http"
98
"path/filepath"
109
"strings"
1110
"time"
@@ -30,11 +29,15 @@ type OpenAIASR struct {
3029

3130
// NewOpenAIASR creates an OpenAI-compatible ASR instance.
3231
func NewOpenAIASR(config *Config) (*OpenAIASR, error) {
32+
if err := validateASRBaseURL(config.BaseURL); err != nil {
33+
return nil, err
34+
}
35+
3336
apiCfg := openai.DefaultConfig(config.APIKey)
3437
if config.BaseURL != "" {
3538
apiCfg.BaseURL = config.BaseURL
3639
}
37-
httpClient := &http.Client{Timeout: asrDefaultTimeout}
40+
httpClient := newASRHTTPClient(asrDefaultTimeout)
3841

3942
// 注入用户自定义 HTTP header（类似 OpenAI Python SDK 的 extra_headers）
4043
if len(config.CustomHeaders) > 0 {

‎internal/models/asr/transport.go‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
package asr
2+
3+
import (
4+
"fmt"
5+
"net/http"
6+
"time"
7+
8+
secutils "github.com/Tencent/WeKnora/internal/utils"
9+
)
10+
11+
func validateASRBaseURL(baseURL string) error {
12+
if baseURL == "" {
13+
return nil
14+
}
15+
if err := secutils.ValidateURLForSSRF(baseURL); err != nil {
16+
return fmt.Errorf("base URL SSRF check failed: %w", err)
17+
}
18+
return nil
19+
}
20+
21+
func newASRHTTPClient(timeout time.Duration) *http.Client {
22+
cfg := secutils.DefaultSSRFSafeHTTPClientConfig()
23+
cfg.Timeout = timeout
24+
return secutils.NewSSRFSafeHTTPClient(cfg)
25+
}

‎internal/models/asr/transport_security_test.go‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
package asr
2+
3+
import (
4+
"testing"
5+
6+
secutils "github.com/Tencent/WeKnora/internal/utils"
7+
)
8+
9+
func withASRSSRFWhitelist(t *testing.T, raw string) {
10+
t.Helper()
11+
t.Setenv("SSRF_WHITELIST", raw)
12+
secutils.ResetSSRFWhitelistForTest()
13+
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
14+
}
15+
16+
func TestOpenAIASRRejectsInternalBaseURL(t *testing.T) {
17+
withASRSSRFWhitelist(t, "")
18+
19+
_, err := NewOpenAIASR(&Config{
20+
BaseURL: "http://169.254.169.254/latest/meta-data/",
21+
ModelName: "asr-test",
22+
})
23+
if err == nil {
24+
t.Fatalf("NewOpenAIASR returned nil error for blocked internal BaseURL")
25+
}
26+
}

‎internal/models/rerank/aliyun_reranker.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,13 +81,16 @@ func NewAliyunReranker(config *RerankerConfig) (*AliyunReranker, error) {
8181
if url := config.BaseURL; url != "" {
8282
baseURL = url
8383
}
84+
if err := validateRerankBaseURL(baseURL); err != nil {
85+
return nil, err
86+
}
8487

8588
return &AliyunReranker{
8689
modelName: config.ModelName,
8790
modelID: config.ModelID,
8891
apiKey: apiKey,
8992
baseURL: baseURL,
90-
client: &http.Client{},
93+
client: newRerankHTTPClient(0),
9194
}, nil
9295
}
9396

‎internal/models/rerank/jina_reranker.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,13 +54,16 @@ func NewJinaReranker(config *RerankerConfig) (*JinaReranker, error) {
5454
if url := config.BaseURL; url != "" {
5555
baseURL = url
5656
}
57+
if err := validateRerankBaseURL(baseURL); err != nil {
58+
return nil, err
59+
}
5760

5861
return &JinaReranker{
5962
modelName: config.ModelName,
6063
modelID: config.ModelID,
6164
apiKey: apiKey,
6265
baseURL: baseURL,
63-
client: &http.Client{},
66+
client: newRerankHTTPClient(0),
6467
}, nil
6568
}
6669

‎internal/models/rerank/nvidia_reranker.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ type NvidiaReranker struct {
2727
func (r *NvidiaReranker) SetCustomHeaders(headers map[string]string) {
2828
r.customHeaders = headers
2929
}
30+
3031
type NvidiaRerankDocument struct {
3132
Text string `json:"text"`
3233
}
@@ -57,17 +58,16 @@ func NewNvidiaReranker(config *RerankerConfig) (*NvidiaReranker, error) {
5758
if url := config.BaseURL; url != "" {
5859
baseURL = url
5960
}
61+
if err := validateRerankBaseURL(baseURL); err != nil {
62+
return nil, err
63+
}
6064

6165
return &NvidiaReranker{
6266
modelName: config.ModelName,
6367
modelID: config.ModelID,
6468
apiKey: apiKey,
6569
baseURL: baseURL,
66-
client: &http.Client{
67-
Transport: &http.Transport{
68-
Proxy: http.ProxyFromEnvironment,
69-
},
70-
},
70+
client: newRerankHTTPClient(0),
7171
}, nil
7272
}
7373

‎internal/models/rerank/remote_api.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,13 +56,16 @@ func NewOpenAIReranker(config *RerankerConfig) (*OpenAIReranker, error) {
5656
if url := config.BaseURL; url != "" {
5757
baseURL = url
5858
}
59+
if err := validateRerankBaseURL(baseURL); err != nil {
60+
return nil, err
61+
}
5962

6063
return &OpenAIReranker{
6164
modelName: config.ModelName,
6265
modelID: config.ModelID,
6366
apiKey: apiKey,
6467
baseURL: baseURL,
65-
client: &http.Client{},
68+
client: newRerankHTTPClient(0),
6669
}, nil
6770
}
6871

‎internal/models/rerank/transport.go‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
package rerank
2+
3+
import (
4+
"fmt"
5+
"net/http"
6+
"time"
7+
8+
secutils "github.com/Tencent/WeKnora/internal/utils"
9+
)
10+
11+
func validateRerankBaseURL(baseURL string) error {
12+
if baseURL == "" {
13+
return nil
14+
}
15+
if err := secutils.ValidateURLForSSRF(baseURL); err != nil {
16+
return fmt.Errorf("base URL SSRF check failed: %w", err)
17+
}
18+
return nil
19+
}
20+
21+
func newRerankHTTPClient(timeout time.Duration) *http.Client {
22+
cfg := secutils.DefaultSSRFSafeHTTPClientConfig()
23+
cfg.Timeout = timeout
24+
return secutils.NewSSRFSafeHTTPClient(cfg)
25+
}

‎internal/models/rerank/transport_security_test.go‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
package rerank
2+
3+
import (
4+
stderrors "errors"
5+
"net/http"
6+
"net/http/httptest"
7+
"testing"
8+
9+
secutils "github.com/Tencent/WeKnora/internal/utils"
10+
)
11+
12+
func withRerankSSRFWhitelist(t *testing.T, raw string) {
13+
t.Helper()
14+
t.Setenv("SSRF_WHITELIST", raw)
15+
secutils.ResetSSRFWhitelistForTest()
16+
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
17+
}
18+
19+
func TestOpenAIRerankerRejectsInternalBaseURL(t *testing.T) {
20+
withRerankSSRFWhitelist(t, "")
21+
22+
_, err := NewOpenAIReranker(&RerankerConfig{
23+
BaseURL: "http://169.254.169.254/latest/meta-data/",
24+
ModelName: "rerank-test",
25+
})
26+
if err == nil {
27+
t.Fatalf("NewOpenAIReranker returned nil error for blocked internal BaseURL")
28+
}
29+
}
30+
31+
func TestOpenAIRerankerBlocksRedirectToInternalURL(t *testing.T) {
32+
withRerankSSRFWhitelist(t, "127.0.0.1")
33+
34+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
35+
http.Redirect(w, r, "http://169.254.169.254/latest/meta-data/", http.StatusFound)
36+
}))
37+
defer server.Close()
38+
39+
reranker, err := NewOpenAIReranker(&RerankerConfig{
40+
BaseURL: server.URL,
41+
ModelName: "rerank-test",
42+
APIKey: "sk-test",
43+
})
44+
if err != nil {
45+
t.Fatalf("NewOpenAIReranker: %v", err)
46+
}
47+
48+
_, err = reranker.Rerank(t.Context(), "query", []string{"doc"})
49+
if !stderrors.Is(err, secutils.ErrSSRFRedirectBlocked) {
50+
t.Fatalf("Rerank error = %v, want ErrSSRFRedirectBlocked", err)
51+
}
52+
}

‎internal/models/rerank/weknoracloud.go‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,10 @@ func NewWeKnoraCloudReranker(config *RerankerConfig) (*WeKnoraCloudReranker, err
3535
if config.AppSecret == "" {
3636
return nil, fmt.Errorf("WeKnoraCloud reranker: AppSecret is required")
3737
}
38+
baseURL := strings.TrimRight(config.BaseURL, "/")
39+
if err := validateRerankBaseURL(baseURL); err != nil {
40+
return nil, err
41+
}
3842
remoteModelName := ""
3943
if config.ExtraConfig != nil {
4044
remoteModelName = strings.TrimSpace(config.ExtraConfig["remote_model_name"])
@@ -45,8 +49,8 @@ func NewWeKnoraCloudReranker(config *RerankerConfig) (*WeKnoraCloudReranker, err
4549
modelID: config.ModelID,
4650
appID: config.AppID,
4751
apiKey: config.AppSecret,
48-
baseURL: strings.TrimRight(config.BaseURL, "/"),
49-
client: &http.Client{Timeout: 60 * time.Second},
52+
baseURL: baseURL,
53+
client: newRerankHTTPClient(60 * time.Second),
5054
}, nil
5155
}
5256

0 commit comments

Comments
 (0)