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
+}