diff --git a/internal/application/service/embed_webhook.go b/internal/application/service/embed_webhook.go index 2f1b9d26e..85a8b3e8f 100644 --- a/internal/application/service/embed_webhook.go +++ b/internal/application/service/embed_webhook.go @@ -25,6 +25,13 @@ const embedWebhookTimeout = 5 * time.Second // ErrEmbedWebhookURLInvalid is returned when a webhook URL fails format or SSRF checks. var ErrEmbedWebhookURLInvalid = errors.New("invalid embed webhook URL") +func newEmbedWebhookHTTPClient() *http.Client { + return secutils.NewSSRFSafeHTTPClient(secutils.SSRFSafeHTTPClientConfig{ + Timeout: embedWebhookTimeout, + MaxRedirects: 5, + }) +} + // ValidateEmbedWebhookURL checks an optional outbound webhook URL. Empty is allowed. func ValidateEmbedWebhookURL(raw string) error { trimmed := strings.TrimSpace(raw) @@ -90,7 +97,7 @@ func DispatchEmbedWebhook(ch *types.EmbedChannel, eventType, sessionID string, p _, _ = mac.Write(raw) req.Header.Set("X-WeKnora-Signature", "sha256="+hex.EncodeToString(mac.Sum(nil))) } - resp, err := http.DefaultClient.Do(req) + resp, err := newEmbedWebhookHTTPClient().Do(req) if err != nil { logger.Warnf(context.Background(), "[embed_webhook] dispatch %s failed: %v", eventType, err) return diff --git a/internal/application/service/embed_webhook_test.go b/internal/application/service/embed_webhook_test.go index f67732ab5..7719433b8 100644 --- a/internal/application/service/embed_webhook_test.go +++ b/internal/application/service/embed_webhook_test.go @@ -1,6 +1,13 @@ package service -import "testing" +import ( + stderrors "errors" + "net/http" + "net/http/httptest" + "testing" + + secutils "github.com/Tencent/WeKnora/internal/utils" +) func TestValidateEmbedWebhookURL(t *testing.T) { withSSRFWhitelist(t, "*.example.com,127.0.0.1") @@ -33,3 +40,27 @@ func TestSignEmbedWebhookBody(t *testing.T) { t.Fatal("signature not deterministic") } } + +func TestEmbedWebhookHTTPClientBlocksRedirectToInternalURL(t *testing.T) { + withSSRFWhitelist(t, "127.0.0.1") + + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, "http://169.254.169.254/latest/meta-data/", http.StatusFound) + })) + defer origin.Close() + + req, err := http.NewRequest(http.MethodPost, origin.URL, nil) + if err != nil { + t.Fatalf("NewRequest: %v", err) + } + resp, err := newEmbedWebhookHTTPClient().Do(req) + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + if err == nil { + t.Fatal("expected redirect to link-local metadata endpoint to be blocked") + } + if !stderrors.Is(err, secutils.ErrSSRFRedirectBlocked) { + t.Fatalf("error = %v, want ErrSSRFRedirectBlocked", err) + } +} diff --git a/internal/application/service/user.go b/internal/application/service/user.go index 4dbf7d329..670d85e8c 100644 --- a/internal/application/service/user.go +++ b/internal/application/service/user.go @@ -1051,6 +1051,39 @@ type oidcTokenResponse struct { TokenType string `json:"token_type"` } +func newOIDCHTTPClient() *http.Client { + cfg := secutils.DefaultSSRFSafeHTTPClientConfig() + cfg.Timeout = 30 * time.Second + return secutils.NewSSRFSafeHTTPClient(cfg) +} + +func validateOIDCEndpoint(label, endpoint string, required bool) error { + endpoint = strings.TrimSpace(endpoint) + if endpoint == "" { + if required { + return fmt.Errorf("OIDC %s endpoint is required", label) + } + return nil + } + if err := secutils.ValidateURLForSSRF(endpoint); err != nil { + return fmt.Errorf("OIDC %s endpoint failed SSRF validation: %w", label, err) + } + return nil +} + +func validateOIDCEndpoints(cfg *config.OIDCAuthConfig) error { + if err := validateOIDCEndpoint("authorization", cfg.AuthorizationEndpoint, true); err != nil { + return err + } + if err := validateOIDCEndpoint("token", cfg.TokenEndpoint, true); err != nil { + return err + } + if err := validateOIDCEndpoint("userinfo", cfg.UserInfoEndpoint, false); err != nil { + return err + } + return nil +} + func (s *userService) getOIDCConfig(ctx context.Context) (*config.OIDCAuthConfig, error) { if s.config == nil || s.config.OIDCAuth == nil || !s.config.OIDCAuth.Enable { return nil, errors.New("OIDC login is disabled") @@ -1067,18 +1100,21 @@ func (s *userService) getOIDCConfig(ctx context.Context) (*config.OIDCAuthConfig func (s *userService) populateOIDCEndpoints(ctx context.Context, cfg *config.OIDCAuthConfig) error { if strings.TrimSpace(cfg.AuthorizationEndpoint) != "" && strings.TrimSpace(cfg.TokenEndpoint) != "" { - return nil + return validateOIDCEndpoints(cfg) } if strings.TrimSpace(cfg.DiscoveryURL) == "" { return errors.New("OIDC discovery_url or explicit endpoints are required") } + if err := validateOIDCEndpoint("discovery", cfg.DiscoveryURL, true); err != nil { + return err + } req, err := http.NewRequestWithContext(ctx, http.MethodGet, cfg.DiscoveryURL, nil) if err != nil { return fmt.Errorf("failed to create OIDC discovery request: %w", err) } - resp, err := http.DefaultClient.Do(req) + resp, err := newOIDCHTTPClient().Do(req) if err != nil { return fmt.Errorf("failed to load OIDC discovery document: %w", err) } @@ -1104,10 +1140,14 @@ func (s *userService) populateOIDCEndpoints(ctx context.Context, cfg *config.OID if cfg.AuthorizationEndpoint == "" || cfg.TokenEndpoint == "" { return errors.New("OIDC discovery document missing required endpoints") } - return nil + return validateOIDCEndpoints(cfg) } func (s *userService) exchangeOIDCCode(ctx context.Context, cfg *config.OIDCAuthConfig, code, redirectURI string) (*oidcTokenResponse, error) { + if err := validateOIDCEndpoint("token", cfg.TokenEndpoint, true); err != nil { + return nil, err + } + form := url.Values{} form.Set("grant_type", "authorization_code") form.Set("code", code) @@ -1122,14 +1162,14 @@ func (s *userService) exchangeOIDCCode(ctx context.Context, cfg *config.OIDCAuth req.Header.Set("Content-Type", "application/x-www-form-urlencoded") req.Header.Set("Accept", "application/json") - resp, err := http.DefaultClient.Do(req) + resp, err := newOIDCHTTPClient().Do(req) if err != nil { return nil, fmt.Errorf("failed to exchange OIDC code: %w", err) } defer resp.Body.Close() if resp.StatusCode < 200 || resp.StatusCode >= 300 { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) - return nil, fmt.Errorf("OIDC token exchange failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(body))) + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 2048)) + return nil, fmt.Errorf("OIDC token exchange failed: status=%d", resp.StatusCode) } var tokenResp oidcTokenResponse @@ -1186,6 +1226,10 @@ func (s *userService) resolveOIDCUserInfo(ctx context.Context, cfg *config.OIDCA } func (s *userService) fetchOIDCUserInfo(ctx context.Context, endpoint, accessToken string) (map[string]interface{}, error) { + if err := validateOIDCEndpoint("userinfo", endpoint, true); err != nil { + return nil, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) if err != nil { return nil, err @@ -1193,14 +1237,14 @@ func (s *userService) fetchOIDCUserInfo(ctx context.Context, endpoint, accessTok req.Header.Set("Authorization", "Bearer "+accessToken) req.Header.Set("Accept", "application/json") - resp, err := http.DefaultClient.Do(req) + resp, err := newOIDCHTTPClient().Do(req) if err != nil { return nil, err } defer resp.Body.Close() if resp.StatusCode < 200 || resp.StatusCode >= 300 { - body, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) - return nil, fmt.Errorf("userinfo request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(body))) + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 2048)) + return nil, fmt.Errorf("userinfo request failed: status=%d", resp.StatusCode) } var claims map[string]interface{} diff --git a/internal/application/service/user_oidc_security_test.go b/internal/application/service/user_oidc_security_test.go new file mode 100644 index 000000000..1ed336d1b --- /dev/null +++ b/internal/application/service/user_oidc_security_test.go @@ -0,0 +1,87 @@ +package service + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Tencent/WeKnora/internal/config" + secutils "github.com/Tencent/WeKnora/internal/utils" +) + +func withOIDCSSRFWhitelist(t *testing.T, raw string) { + t.Helper() + t.Setenv("SSRF_WHITELIST", raw) + secutils.ResetSSRFWhitelistForTest() + t.Cleanup(secutils.ResetSSRFWhitelistForTest) +} + +func TestOIDCDiscoveryRejectsInternalDiscoveredEndpoint(t *testing.T) { + withOIDCSSRFWhitelist(t, "127.0.0.1") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(oidcDiscoveryDocument{ + AuthorizationEndpoint: serverURL(r, "/authorize"), + TokenEndpoint: "http://169.254.169.254/latest/meta-data/", + }) + })) + defer server.Close() + + svc := &userService{} + cfg := &config.OIDCAuthConfig{DiscoveryURL: server.URL} + err := svc.populateOIDCEndpoints(context.Background(), cfg) + if err == nil || !strings.Contains(err.Error(), "OIDC token endpoint failed SSRF validation") { + t.Fatalf("populateOIDCEndpoints error = %v, want token SSRF validation failure", err) + } +} + +func TestOIDCTokenExchangeBlocksRedirectToInternalURL(t *testing.T) { + withOIDCSSRFWhitelist(t, "127.0.0.1") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, "http://169.254.169.254/latest/meta-data/", http.StatusFound) + })) + defer server.Close() + + svc := &userService{} + cfg := &config.OIDCAuthConfig{TokenEndpoint: server.URL, ClientID: "client-id", ClientSecret: "client-secret"} + _, err := svc.exchangeOIDCCode(context.Background(), cfg, "code", "https://app.example/callback") + if err == nil || !strings.Contains(err.Error(), secutils.ErrSSRFRedirectBlocked.Error()) { + t.Fatalf("exchangeOIDCCode error = %v, want redirect SSRF block", err) + } +} + +func TestOIDCErrorsDoNotEchoSecrets(t *testing.T) { + withOIDCSSRFWhitelist(t, "127.0.0.1") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte("client-secret bearer-token")) + })) + defer server.Close() + + svc := &userService{} + cfg := &config.OIDCAuthConfig{TokenEndpoint: server.URL, ClientID: "client-id", ClientSecret: "client-secret"} + _, err := svc.exchangeOIDCCode(context.Background(), cfg, "code", "https://app.example/callback") + if err == nil { + t.Fatalf("exchangeOIDCCode returned nil error") + } + if strings.Contains(err.Error(), "client-secret") { + t.Fatalf("exchangeOIDCCode error leaked secret: %v", err) + } + + _, err = svc.fetchOIDCUserInfo(context.Background(), server.URL, "bearer-token") + if err == nil { + t.Fatalf("fetchOIDCUserInfo returned nil error") + } + if strings.Contains(err.Error(), "bearer-token") { + t.Fatalf("fetchOIDCUserInfo error leaked bearer token: %v", err) + } +} + +func serverURL(r *http.Request, path string) string { + return "http://" + r.Host + path +} diff --git a/internal/handler/knowledge.go b/internal/handler/knowledge.go index 591918283..fc840dac9 100644 --- a/internal/handler/knowledge.go +++ b/internal/handler/knowledge.go @@ -1267,49 +1267,8 @@ func (h *KnowledgeHandler) DownloadKnowledgeFile(c *gin.Context) { // mimeTypeByExt returns the MIME type for a given file extension. func mimeTypeByExt(filename string) string { - ext := strings.ToLower(filename) - if idx := strings.LastIndex(ext, "."); idx >= 0 { - ext = ext[idx:] - } else { - ext = "" - } - m := map[string]string{ - ".pdf": "application/pdf", - ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", - ".doc": "application/msword", - ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation", - ".ppt": "application/vnd.ms-powerpoint", - ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", - ".xls": "application/vnd.ms-excel", - ".csv": "text/csv", - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".png": "image/png", - ".gif": "image/gif", - ".bmp": "image/bmp", - ".webp": "image/webp", - ".svg": "image/svg+xml", - ".tiff": "image/tiff", - ".txt": "text/plain; charset=utf-8", - ".md": "text/markdown; charset=utf-8", - ".markdown": "text/markdown; charset=utf-8", - ".json": "application/json; charset=utf-8", - ".xml": "application/xml; charset=utf-8", - ".html": "text/html; charset=utf-8", - ".css": "text/css; charset=utf-8", - ".js": "text/javascript; charset=utf-8", - ".ts": "text/typescript; charset=utf-8", - ".py": "text/x-python; charset=utf-8", - ".go": "text/x-go; charset=utf-8", - ".java": "text/x-java; charset=utf-8", - ".yaml": "text/yaml; charset=utf-8", - ".yml": "text/yaml; charset=utf-8", - ".sh": "text/x-shellscript; charset=utf-8", - } - if ct, ok := m[ext]; ok { - return ct - } - return "application/octet-stream" + ct, _ := secutils.SafeContentTypeByFilename(filename) + return ct } // PreviewKnowledgeFile godoc @@ -1347,9 +1306,14 @@ func (h *KnowledgeHandler) PreviewKnowledgeFile(c *gin.Context) { } defer file.Close() - contentType := mimeTypeByExt(filename) + contentType, inline := secutils.SafeContentTypeByFilename(filename) c.Header("Content-Type", contentType) - c.Header("Content-Disposition", mime.FormatMediaType("inline", map[string]string{"filename": filename})) + c.Header("X-Content-Type-Options", "nosniff") + disposition := "inline" + if !inline { + disposition = "attachment" + } + c.Header("Content-Disposition", mime.FormatMediaType(disposition, map[string]string{"filename": filename})) c.Header("Cache-Control", "private, max-age=3600") c.Stream(func(w io.Writer) bool { diff --git a/internal/handler/knowledge_preview_security_test.go b/internal/handler/knowledge_preview_security_test.go new file mode 100644 index 000000000..972e818c5 --- /dev/null +++ b/internal/handler/knowledge_preview_security_test.go @@ -0,0 +1,67 @@ +package handler + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Tencent/WeKnora/internal/middleware" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/gin-gonic/gin" +) + +type previewKnowledgeServiceStub struct { + interfaces.KnowledgeService + filename string +} + +type closeNotifyRecorder struct { + *httptest.ResponseRecorder +} + +func (r *closeNotifyRecorder) CloseNotify() <-chan bool { + ch := make(chan bool, 1) + return ch +} + +func (s *previewKnowledgeServiceStub) GetKnowledgeByIDOnly(context.Context, string) (*types.Knowledge, error) { + return &types.Knowledge{ID: "k1", TenantID: 42, KnowledgeBaseID: "kb1"}, nil +} + +func (s *previewKnowledgeServiceStub) GetKnowledgeFile(context.Context, string) (io.ReadCloser, string, error) { + return io.NopCloser(strings.NewReader(``)), s.filename, nil +} + +func TestPreviewKnowledgeFileForcesActiveContentDownload(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(middleware.ErrorHandler()) + router.Use(func(c *gin.Context) { + c.Set(types.TenantIDContextKey.String(), uint64(42)) + c.Next() + }) + + h := &KnowledgeHandler{kgService: &previewKnowledgeServiceStub{filename: "payload.html"}} + router.GET("/knowledge/:id/preview", h.PreviewKnowledgeFile) + + req := httptest.NewRequest(http.MethodGet, "/knowledge/k1/preview", nil) + w := &closeNotifyRecorder{ResponseRecorder: httptest.NewRecorder()} + router.ServeHTTP(w, req) + + if got, want := w.Code, http.StatusOK; got != want { + t.Fatalf("status = %d, want %d; body=%s", got, want, w.Body.String()) + } + if got := w.Header().Get("Content-Type"); got != "application/octet-stream" { + t.Fatalf("Content-Type = %q, want application/octet-stream", got) + } + if got := w.Header().Get("X-Content-Type-Options"); got != "nosniff" { + t.Fatalf("X-Content-Type-Options = %q, want nosniff", got) + } + if got := w.Header().Get("Content-Disposition"); !strings.HasPrefix(got, "attachment") { + t.Fatalf("Content-Disposition = %q, want attachment", got) + } +} diff --git a/internal/im/mattermost/adapter_test.go b/internal/im/mattermost/adapter_test.go index dc551cda3..d2ec460a3 100644 --- a/internal/im/mattermost/adapter_test.go +++ b/internal/im/mattermost/adapter_test.go @@ -1,16 +1,29 @@ package mattermost import ( + "context" "encoding/json" + stderrors "errors" + "net/http" + "net/http/httptest" "testing" + + secutils "github.com/Tencent/WeKnora/internal/utils" ) +func withMattermostSSRFWhitelist(t *testing.T, whitelist string) { + t.Helper() + t.Setenv("SSRF_WHITELIST", whitelist) + secutils.ResetSSRFWhitelistForTest() + t.Cleanup(secutils.ResetSSRFWhitelistForTest) +} + func TestParseOutgoingBody_ThreadRoot(t *testing.T) { tests := []struct { - name string - payload outgoingPayload + name string + payload outgoingPayload postReplyToMain bool - wantThreadRoot string + wantThreadRoot string }{ { name: "threaded reply has RootID", @@ -60,6 +73,35 @@ func TestParseOutgoingBody_ThreadRoot(t *testing.T) { } } +func TestNewClientRejectsPrivateSiteURL(t *testing.T) { + withMattermostSSRFWhitelist(t, "") + + if _, err := NewClient("http://127.0.0.1:8065", "bot-token"); err == nil { + t.Fatal("expected private Mattermost site_url to be rejected") + } +} + +func TestMattermostClientBlocksRedirectToInternalURL(t *testing.T) { + withMattermostSSRFWhitelist(t, "127.0.0.1") + + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, "http://169.254.169.254/latest/meta-data/", http.StatusFound) + })) + defer origin.Close() + + client, err := NewClient(origin.URL, "bot-token") + if err != nil { + t.Fatalf("NewClient: %v", err) + } + _, err = client.CreatePost(context.Background(), "channel", "", "hello") + if err == nil { + t.Fatal("expected redirect to link-local metadata endpoint to be blocked") + } + if !stderrors.Is(err, secutils.ErrSSRFRedirectBlocked) { + t.Fatalf("error = %v, want ErrSSRFRedirectBlocked", err) + } +} + func TestParseOutgoingBody_JSON(t *testing.T) { body := `{"token":"tok","user_id":"u1","user_name":"alice","channel_id":"ch1","post_id":"p1","text":"hello","root_id":"r1"}` diff --git a/internal/im/mattermost/client.go b/internal/im/mattermost/client.go index 501c3d749..a2727363e 100644 --- a/internal/im/mattermost/client.go +++ b/internal/im/mattermost/client.go @@ -7,8 +7,11 @@ import ( "fmt" "io" "net/http" + "net/url" "strings" "time" + + secutils "github.com/Tencent/WeKnora/internal/utils" ) // Client calls Mattermost REST API v4. @@ -25,14 +28,28 @@ func NewClient(siteURL, botToken string) (*Client, error) { if siteURL == "" { return nil, fmt.Errorf("site_url is required") } + parsed, err := url.Parse(siteURL) + if err != nil || parsed.Host == "" { + return nil, fmt.Errorf("invalid site_url: must be a valid http(s) URL") + } + if parsed.Scheme != "http" && parsed.Scheme != "https" { + return nil, fmt.Errorf("invalid site_url: must use http or https") + } + if err := secutils.ValidateURLForSSRF(siteURL); err != nil { + return nil, fmt.Errorf( + "invalid site_url: %w (for private Mattermost deployments, add the hostname to SSRF_WHITELIST)", + err, + ) + } if strings.TrimSpace(botToken) == "" { return nil, fmt.Errorf("bot_token is required") } return &Client{ baseURL: siteURL + "/api/v4", - httpClient: &http.Client{ - Timeout: 60 * time.Second, - }, + httpClient: secutils.NewSSRFSafeHTTPClient(secutils.SSRFSafeHTTPClientConfig{ + Timeout: 60 * time.Second, + MaxRedirects: 5, + }), token: strings.TrimSpace(botToken), }, nil } diff --git a/internal/im/qqbot/adapter_test.go b/internal/im/qqbot/adapter_test.go index 288973294..f2902d23d 100644 --- a/internal/im/qqbot/adapter_test.go +++ b/internal/im/qqbot/adapter_test.go @@ -5,8 +5,16 @@ import ( "testing" "github.com/Tencent/WeKnora/internal/im" + secutils "github.com/Tencent/WeKnora/internal/utils" ) +func withQQBotSSRFWhitelist(t *testing.T, whitelist string) { + t.Helper() + t.Setenv("SSRF_WHITELIST", whitelist) + secutils.ResetSSRFWhitelistForTest() + t.Cleanup(secutils.ResetSSRFWhitelistForTest) +} + func TestParseGatewayPayloadC2CMessage(t *testing.T) { event := messageEvent{ ID: "msg-1", @@ -76,3 +84,27 @@ func TestParseExpiresIn(t *testing.T) { }) } } + +func TestNewClientRejectsPrivateAPIBaseURL(t *testing.T) { + withQQBotSSRFWhitelist(t, "") + + if _, err := NewClient("app", "secret", "http://127.0.0.1:8080", ""); err == nil { + t.Fatal("expected private qqbot api_base_url to be rejected") + } +} + +func TestNewClientRejectsInsecureGatewayURL(t *testing.T) { + withQQBotSSRFWhitelist(t, "127.0.0.1") + + if _, err := NewClient("app", "secret", "http://127.0.0.1:8080", "ws://127.0.0.1/gateway"); err == nil { + t.Fatal("expected non-wss qqbot gateway_url to be rejected") + } +} + +func TestNewClientAllowsWhitelistedPrivateDeployment(t *testing.T) { + withQQBotSSRFWhitelist(t, "127.0.0.1") + + if _, err := NewClient("app", "secret", "http://127.0.0.1:8080", "wss://127.0.0.1/gateway"); err != nil { + t.Fatalf("expected whitelisted private qqbot endpoints to pass: %v", err) + } +} diff --git a/internal/im/qqbot/client.go b/internal/im/qqbot/client.go index 49ebc3ed7..ce8c41b58 100644 --- a/internal/im/qqbot/client.go +++ b/internal/im/qqbot/client.go @@ -6,9 +6,12 @@ import ( "encoding/json" "fmt" "net/http" + "net/url" "strings" "sync" "time" + + secutils "github.com/Tencent/WeKnora/internal/utils" ) type Client struct { @@ -35,12 +38,23 @@ func NewClient(appID, clientSecret, apiBaseURL, gatewayURL string) (*Client, err if apiBaseURL == "" { apiBaseURL = defaultAPIBaseURL } + apiBaseURL = strings.TrimRight(strings.TrimSpace(apiBaseURL), "/") + if err := validateHTTPAPIBaseURL(apiBaseURL); err != nil { + return nil, err + } + gatewayURL = strings.TrimSpace(gatewayURL) + if err := validateGatewayURL(gatewayURL); err != nil { + return nil, err + } return &Client{ appID: appID, clientSecret: clientSecret, - apiBaseURL: strings.TrimRight(apiBaseURL, "/"), - gatewayURL: strings.TrimSpace(gatewayURL), - httpClient: &http.Client{Timeout: 15 * time.Second}, + apiBaseURL: apiBaseURL, + gatewayURL: gatewayURL, + httpClient: secutils.NewSSRFSafeHTTPClient(secutils.SSRFSafeHTTPClientConfig{ + Timeout: 15 * time.Second, + MaxRedirects: 5, + }), }, nil } @@ -55,9 +69,48 @@ func (c *Client) GatewayURL(ctx context.Context) (string, error) { if result.URL == "" { return "", fmt.Errorf("empty qqbot gateway url") } + if err := validateGatewayURL(result.URL); err != nil { + return "", fmt.Errorf("invalid qqbot gateway url: %w", err) + } return result.URL, nil } +func validateHTTPAPIBaseURL(raw string) error { + u, err := url.Parse(raw) + if err != nil || u.Host == "" { + return fmt.Errorf("invalid qqbot api_base_url: must be a valid http(s) URL") + } + if u.Scheme != "http" && u.Scheme != "https" { + return fmt.Errorf("invalid qqbot api_base_url: must use http or https") + } + if err := secutils.ValidateURLForSSRF(raw); err != nil { + return fmt.Errorf("invalid qqbot api_base_url: %w (for private deployments, add the hostname to SSRF_WHITELIST)", err) + } + return nil +} + +func validateGatewayURL(raw string) error { + if strings.TrimSpace(raw) == "" { + return nil + } + u, err := url.Parse(raw) + if err != nil || u.Host == "" { + return fmt.Errorf("gateway_url must be a valid wss URL") + } + if u.Scheme != "wss" { + return fmt.Errorf("gateway_url must use wss") + } + checkURL := *u + checkURL.Scheme = "https" + if err := secutils.ValidateURLForSSRF(checkURL.String()); err != nil { + return fmt.Errorf( + "gateway_url failed SSRF validation: %w (for private deployments, add the hostname to SSRF_WHITELIST)", + err, + ) + } + return nil +} + func (c *Client) SendC2CMessage(ctx context.Context, openID, content, msgID string) error { path := fmt.Sprintf("/v2/users/%s/messages", openID) return c.sendText(ctx, path, content, msgID) diff --git a/internal/models/chat/ollama.go b/internal/models/chat/ollama.go index 989a06fad..45f397d69 100644 --- a/internal/models/chat/ollama.go +++ b/internal/models/chat/ollama.go @@ -5,7 +5,6 @@ import ( "encoding/json" "fmt" "io" - "net/http" "strconv" "strings" "time" @@ -13,6 +12,7 @@ import ( "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/models/utils/ollama" "github.com/Tencent/WeKnora/internal/types" + secutils "github.com/Tencent/WeKnora/internal/utils" ollamaapi "github.com/ollama/ollama/api" ) @@ -63,7 +63,13 @@ func resolveImageForOllama(imageURL string) ollamaapi.ImageData { return data } if strings.HasPrefix(imageURL, "http://") || strings.HasPrefix(imageURL, "https://") { - client := &http.Client{Timeout: 30 * time.Second} + if err := secutils.ValidateURLForSSRF(imageURL); err != nil { + return nil + } + client := secutils.NewSSRFSafeHTTPClient(secutils.SSRFSafeHTTPClientConfig{ + Timeout: 30 * time.Second, + MaxRedirects: 5, + }) resp, err := client.Get(imageURL) if err != nil { return nil diff --git a/internal/models/chat/ollama_test.go b/internal/models/chat/ollama_test.go new file mode 100644 index 000000000..837702699 --- /dev/null +++ b/internal/models/chat/ollama_test.go @@ -0,0 +1,34 @@ +package chat + +import ( + "net/http" + "net/http/httptest" + "testing" + + secutils "github.com/Tencent/WeKnora/internal/utils" +) + +func TestResolveImageForOllamaRejectsInternalURL(t *testing.T) { + t.Setenv("SSRF_WHITELIST", "") + secutils.ResetSSRFWhitelistForTest() + t.Cleanup(secutils.ResetSSRFWhitelistForTest) + + if data := resolveImageForOllama("http://169.254.169.254/latest/meta-data/"); data != nil { + t.Fatalf("resolveImageForOllama returned data for blocked internal URL") + } +} + +func TestResolveImageForOllamaBlocksRedirectToInternalURL(t *testing.T) { + t.Setenv("SSRF_WHITELIST", "127.0.0.1") + secutils.ResetSSRFWhitelistForTest() + t.Cleanup(secutils.ResetSSRFWhitelistForTest) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, "http://169.254.169.254/latest/meta-data/", http.StatusFound) + })) + defer server.Close() + + if data := resolveImageForOllama(server.URL); data != nil { + t.Fatalf("resolveImageForOllama returned data after redirect to blocked internal URL") + } +} diff --git a/internal/router/router.go b/internal/router/router.go index 095429462..ce7c8df0c 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -1465,28 +1465,12 @@ func newFileServeHandler(globalFileService interfaces.FileService) gin.HandlerFu } defer reader.Close() - ext := filepath.Ext(filePath) - contentType := "application/octet-stream" - switch strings.ToLower(ext) { - case ".png": - contentType = "image/png" - case ".jpg", ".jpeg": - contentType = "image/jpeg" - case ".gif": - contentType = "image/gif" - case ".webp": - contentType = "image/webp" - case ".bmp": - contentType = "image/bmp" - case ".svg": - contentType = "image/svg+xml" - case ".pdf": - contentType = "application/pdf" - case ".csv": - contentType = "text/csv; charset=utf-8" - } - + contentType, inline := secutils.SafeContentTypeByFilename(filePath) c.Header("Content-Type", contentType) + c.Header("X-Content-Type-Options", "nosniff") + if !inline { + c.Header("Content-Disposition", "attachment") + } c.Header("Cache-Control", "public, max-age=86400") c.Status(http.StatusOK) if _, err := io.Copy(c.Writer, reader); err != nil { @@ -1587,24 +1571,7 @@ func presignedFileHandler(tenantService interfaces.TenantService, absDir string) return } - ext := filepath.Ext(filePath) - contentType := "application/octet-stream" - switch strings.ToLower(ext) { - case ".png": - contentType = "image/png" - case ".jpg", ".jpeg": - contentType = "image/jpeg" - case ".gif": - contentType = "image/gif" - case ".webp": - contentType = "image/webp" - case ".bmp": - contentType = "image/bmp" - case ".svg": - contentType = "image/svg+xml" - case ".pdf": - contentType = "application/pdf" - } + contentType, inline := secutils.SafeContentTypeByFilename(filePath) // HEAD short-circuits the body read. We still need to confirm the // object exists, but we use a 0-byte content length and skip io.Copy. @@ -1621,6 +1588,10 @@ func presignedFileHandler(tenantService interfaces.TenantService, absDir string) defer reader.Close() c.Header("Content-Type", contentType) + c.Header("X-Content-Type-Options", "nosniff") + if !inline { + c.Header("Content-Disposition", "attachment") + } c.Header("Cache-Control", "public, max-age=86400") if c.Request.Method == http.MethodHead { c.Status(http.StatusOK) diff --git a/internal/router/router_files_test.go b/internal/router/router_files_test.go index 8923f5f2f..f0eec566a 100644 --- a/internal/router/router_files_test.go +++ b/internal/router/router_files_test.go @@ -152,3 +152,35 @@ func TestServeFilesRejectsPathWithoutTenantSegment(t *testing.T) { t.Fatalf("status = %d, want %d", got, want) } } + +func TestServeFilesForcesActiveContentDownload(t *testing.T) { + gin.SetMode(gin.TestMode) + t.Setenv("STORAGE_TYPE", "local") + + engine := gin.New() + serveFiles(engine, &stubFileService{ + getFile: func(_ context.Context, _ string) (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader(``)), nil + }, + }) + + filePath := "local://42/docs/payload.svg" + req := httptest.NewRequest(http.MethodGet, "/files?file_path="+url.QueryEscape(filePath), nil) + req = req.WithContext(context.WithValue(req.Context(), types.TenantInfoContextKey, &types.Tenant{ID: 42})) + + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, req) + + if got, want := recorder.Code, http.StatusOK; got != want { + t.Fatalf("status = %d, want %d", got, want) + } + if got := recorder.Header().Get("Content-Type"); got != "application/octet-stream" { + t.Fatalf("Content-Type = %q, want application/octet-stream", got) + } + if got := recorder.Header().Get("Content-Disposition"); got != "attachment" { + t.Fatalf("Content-Disposition = %q, want attachment", got) + } + if got := recorder.Header().Get("X-Content-Type-Options"); got != "nosniff" { + t.Fatalf("X-Content-Type-Options = %q, want nosniff", got) + } +} diff --git a/internal/router/router_presigned_test.go b/internal/router/router_presigned_test.go index 186e2ea7a..6c4cc8573 100644 --- a/internal/router/router_presigned_test.go +++ b/internal/router/router_presigned_test.go @@ -170,6 +170,28 @@ func TestPresignedFile_GET_ReturnsContent(t *testing.T) { } } +func TestPresignedFile_ForcesActiveContentDownload(t *testing.T) { + engine, baseDir, signURL := setupPresignedTestServer(t) + storagePath := writeTestFile(t, baseDir, "1/payload.svg", ``) + + req := httptest.NewRequest(http.MethodGet, signURL(storagePath, 1, time.Hour), nil) + w := httptest.NewRecorder() + engine.ServeHTTP(w, req) + + if got, want := w.Code, http.StatusOK; got != want { + t.Fatalf("status = %d, want %d", got, want) + } + if got := w.Header().Get("Content-Type"); got != "application/octet-stream" { + t.Fatalf("Content-Type = %q, want application/octet-stream", got) + } + if got := w.Header().Get("Content-Disposition"); got != "attachment" { + t.Fatalf("Content-Disposition = %q, want attachment", got) + } + if got := w.Header().Get("X-Content-Type-Options"); got != "nosniff" { + t.Fatalf("X-Content-Type-Options = %q, want nosniff", got) + } +} + func TestPresignedFile_InvalidSig_403(t *testing.T) { engine, _, _ := setupPresignedTestServer(t) diff --git a/internal/utils/fileutil.go b/internal/utils/fileutil.go index 341ddcc0d..6d53ee825 100644 --- a/internal/utils/fileutil.go +++ b/internal/utils/fileutil.go @@ -2,8 +2,24 @@ package utils import "strings" +// IsActiveBrowserContentExt reports whether a file extension can execute script +// or otherwise participate as active same-origin browser content when served +// inline. User-uploaded files with these extensions must be downloaded instead +// of previewed inline. +func IsActiveBrowserContentExt(ext string) bool { + switch strings.ToLower(ext) { + case ".svg", ".svgz", ".html", ".htm", ".xhtml", ".xml", ".js", ".mjs", ".css": + return true + default: + return false + } +} + // GetContentTypeByExt returns the content type based on file extension func GetContentTypeByExt(ext string) string { + if IsActiveBrowserContentExt(ext) { + return "application/octet-stream" + } switch strings.ToLower(ext) { case ".csv": return "text/csv; charset=utf-8" @@ -27,12 +43,8 @@ func GetContentTypeByExt(ext string) string { return "image/gif" case ".webp": return "image/webp" - case ".svg": - return "image/svg+xml" case ".txt": return "text/plain; charset=utf-8" - case ".html", ".htm": - return "text/html; charset=utf-8" case ".mp4": return "video/mp4" case ".mp3": @@ -51,3 +63,18 @@ func GetContentTypeByExt(ext string) string { return "application/octet-stream" } } + +// SafeContentTypeByFilename returns a browser-safe content type for serving +// user-controlled files and whether that file may be displayed inline. +func SafeContentTypeByFilename(filename string) (string, bool) { + ext := filename + if idx := strings.LastIndex(ext, "."); idx >= 0 { + ext = ext[idx:] + } else { + ext = "" + } + if IsActiveBrowserContentExt(ext) { + return "application/octet-stream", false + } + return GetContentTypeByExt(ext), true +}