From ed9802ff2b56bc4e1d95aa83b5c5bf1597ed317b Mon Sep 17 00:00:00 2001 From: wizardchen Date: Fri, 12 Jun 2026 22:40:18 +0800 Subject: [PATCH] feat(embed): add backend embed channels with auth and rate limiting (#1657) Introduce embed_channels storage, publish-token and session-token auth, origin allowlist enforcement, shared Redis rate limiter, public embed API routes, webhooks with SSRF validation, and session listing by embed source. --- .../application/repository/embed_channel.go | 78 ++ internal/application/repository/session.go | 21 +- .../application/repository/session_test.go | 23 + internal/application/service/embed_channel.go | 402 +++++++++++ .../service/embed_channel_chunk_test.go | 123 ++++ .../embed_channel_public_config_test.go | 142 ++++ internal/application/service/embed_session.go | 143 ++++ .../application/service/embed_session_test.go | 85 +++ internal/application/service/embed_webhook.go | 111 +++ .../application/service/embed_webhook_test.go | 35 + internal/container/container.go | 3 + internal/handler/embed_channel.go | 681 ++++++++++++++++++ .../handler/embed_channel_exchange_test.go | 152 ++++ internal/handler/embed_flow_test.go | 305 ++++++++ internal/handler/embed_session_test.go | 369 ++++++++++ internal/im/ratelimit.go | 199 ----- internal/im/service.go | 18 +- internal/middleware/auth_public_ratelimit.go | 2 +- internal/middleware/embed_auth.go | 247 +++++++ internal/middleware/embed_auth_test.go | 408 +++++++++++ internal/ratelimit/limiter.go | 181 +++++ internal/ratelimit/limiter_test.go | 41 ++ internal/router/router.go | 185 ++++- internal/types/embed_channel.go | 150 ++++ internal/types/embed_channel_test.go | 40 + internal/types/interfaces/embed_channel.go | 38 + internal/types/session.go | 3 +- migrations/sqlite/000000_init.up.sql | 33 + .../versioned/000060_embed_channels.down.sql | 5 + .../versioned/000060_embed_channels.up.sql | 63 ++ 30 files changed, 4070 insertions(+), 216 deletions(-) create mode 100644 internal/application/repository/embed_channel.go create mode 100644 internal/application/service/embed_channel.go create mode 100644 internal/application/service/embed_channel_chunk_test.go create mode 100644 internal/application/service/embed_channel_public_config_test.go create mode 100644 internal/application/service/embed_session.go create mode 100644 internal/application/service/embed_session_test.go create mode 100644 internal/application/service/embed_webhook.go create mode 100644 internal/application/service/embed_webhook_test.go create mode 100644 internal/handler/embed_channel.go create mode 100644 internal/handler/embed_channel_exchange_test.go create mode 100644 internal/handler/embed_flow_test.go create mode 100644 internal/handler/embed_session_test.go delete mode 100644 internal/im/ratelimit.go create mode 100644 internal/middleware/embed_auth.go create mode 100644 internal/middleware/embed_auth_test.go create mode 100644 internal/ratelimit/limiter.go create mode 100644 internal/ratelimit/limiter_test.go create mode 100644 internal/types/embed_channel.go create mode 100644 internal/types/embed_channel_test.go create mode 100644 internal/types/interfaces/embed_channel.go create mode 100644 migrations/versioned/000060_embed_channels.down.sql create mode 100644 migrations/versioned/000060_embed_channels.up.sql diff --git a/internal/application/repository/embed_channel.go b/internal/application/repository/embed_channel.go new file mode 100644 index 00000000..60d700e8 --- /dev/null +++ b/internal/application/repository/embed_channel.go @@ -0,0 +1,78 @@ +package repository + +import ( + "context" + "errors" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "gorm.io/gorm" +) + +type embedChannelRepository struct { + db *gorm.DB +} + +func NewEmbedChannelRepository(db *gorm.DB) interfaces.EmbedChannelRepository { + return &embedChannelRepository{db: db} +} + +func (r *embedChannelRepository) Create(ctx context.Context, ch *types.EmbedChannel) error { + return r.db.WithContext(ctx).Create(ch).Error +} + +func (r *embedChannelRepository) GetByID(ctx context.Context, id string) (*types.EmbedChannel, error) { + var ch types.EmbedChannel + err := r.db.WithContext(ctx).Where("id = ?", id).First(&ch).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return &ch, nil +} + +func (r *embedChannelRepository) GetByPublishToken(ctx context.Context, token string) (*types.EmbedChannel, error) { + var ch types.EmbedChannel + err := r.db.WithContext(ctx).Where("publish_token = ?", token).First(&ch).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return &ch, nil +} + +func (r *embedChannelRepository) ListByAgent( + ctx context.Context, tenantID uint64, agentID string, +) ([]*types.EmbedChannel, error) { + var rows []*types.EmbedChannel + err := r.db.WithContext(ctx). + Where("tenant_id = ? AND agent_id = ?", tenantID, agentID). + Order("created_at DESC"). + Find(&rows).Error + return rows, err +} + +func (r *embedChannelRepository) ListByTenant( + ctx context.Context, tenantID uint64, +) ([]*types.EmbedChannel, error) { + var rows []*types.EmbedChannel + err := r.db.WithContext(ctx). + Where("tenant_id = ?", tenantID). + Order("created_at DESC"). + Find(&rows).Error + return rows, err +} + +func (r *embedChannelRepository) Update(ctx context.Context, ch *types.EmbedChannel) error { + return r.db.WithContext(ctx).Save(ch).Error +} + +func (r *embedChannelRepository) Delete(ctx context.Context, tenantID uint64, id string) error { + return r.db.WithContext(ctx). + Where("tenant_id = ? AND id = ?", tenantID, id). + Delete(&types.EmbedChannel{}).Error +} diff --git a/internal/application/repository/session.go b/internal/application/repository/session.go index 635de16d..250d13e8 100644 --- a/internal/application/repository/session.go +++ b/internal/application/repository/session.go @@ -148,13 +148,28 @@ func (r *sessionRepository) QueryPaged( joinClause := "LEFT JOIN im_channel_sessions ics ON ics.session_id = s.id" applySource := func(db *gorm.DB) *gorm.DB { - switch strings.ToLower(strings.TrimSpace(q.Source)) { + src := strings.TrimSpace(q.Source) + lower := strings.ToLower(src) + embedPrefix := types.EmbedSessionMarkerPrefix + switch lower { case "": return db case "web": - return db.Where("ics.id IS NULL") + // User web chats only — exclude embed-widget sessions (same IM-null row). + return db.Where( + "ics.id IS NULL AND (s.description = '' OR s.description NOT LIKE ?)", + embedPrefix+"%", + ) + case "embed": + return db.Where("ics.id IS NULL AND s.description LIKE ?", embedPrefix+"%") default: - return db.Where("ics.platform = ?", strings.ToLower(q.Source)) + if strings.HasPrefix(lower, "embed:") { + channelID := strings.TrimSpace(src[len("embed:"):]) + if channelID != "" { + return db.Where("ics.id IS NULL AND s.description = ?", embedPrefix+channelID) + } + } + return db.Where("ics.platform = ?", lower) } } applyAgent := func(db *gorm.DB) *gorm.DB { diff --git a/internal/application/repository/session_test.go b/internal/application/repository/session_test.go index f242f9a6..4c40e23b 100644 --- a/internal/application/repository/session_test.go +++ b/internal/application/repository/session_test.go @@ -222,3 +222,26 @@ func TestSessionRepositoryQueryPagedKeepsClearedIMSessionsOutOfWeb(t *testing.T) require.ElementsMatch(t, []string{active.ID, cleared.ID}, listItemIDsForTest(wecomItems), "wecom must include both the active and the cleared IM session") } + +func TestSessionRepositoryQueryPagedSplitsWebAndEmbedSessions(t *testing.T) { + repo, db := newSessionRepositoryForTest(t) + require.NoError(t, db.AutoMigrate(&testIMChannelSession{})) + ctx := context.Background() + + web := createSessionForTest(t, db, 1, "alice") + embed := createSessionForTest(t, db, 1, "alice") + require.NoError(t, db.Model(&types.Session{}).Where("id = ?", embed.ID). + Update("description", types.EmbedSessionMarkerPrefix+"ch-1").Error) + + webItems, _, err := repo.QueryPaged(ctx, &types.SessionListQuery{ + TenantID: 1, UserID: "alice", Source: "web", Page: 1, PageSize: 50, + }) + require.NoError(t, err) + require.Equal(t, []string{web.ID}, listItemIDsForTest(webItems)) + + embedItems, _, err := repo.QueryPaged(ctx, &types.SessionListQuery{ + TenantID: 1, UserID: "alice", Source: "embed:ch-1", Page: 1, PageSize: 50, + }) + require.NoError(t, err) + require.Equal(t, []string{embed.ID}, listItemIDsForTest(embedItems)) +} diff --git a/internal/application/service/embed_channel.go b/internal/application/service/embed_channel.go new file mode 100644 index 00000000..601ada01 --- /dev/null +++ b/internal/application/service/embed_channel.go @@ -0,0 +1,402 @@ +package service + +import ( + "context" + "crypto/rand" + "crypto/subtle" + "encoding/base64" + "errors" + "fmt" + "strings" + + apperrors "github.com/Tencent/WeKnora/internal/errors" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/redis/go-redis/v9" +) + +const embedTokenBytes = 32 + +var ( + ErrEmbedChannelNotFound = errors.New("embed channel not found") + ErrEmbedTokenInvalid = errors.New("embed publish token is invalid") + ErrEmbedChannelDisabled = errors.New("embed channel is disabled") + ErrEmbedChunkNotFound = errors.New("embed chunk not found") + ErrEmbedChunkForbidden = errors.New("embed chunk not accessible") +) + +type embedChannelService struct { + repo interfaces.EmbedChannelRepository + agentService interfaces.CustomAgentService + chunkService interfaces.ChunkService + redis *redis.Client +} + +func NewEmbedChannelService( + repo interfaces.EmbedChannelRepository, + agentService interfaces.CustomAgentService, + chunkService interfaces.ChunkService, + redisClient *redis.Client, +) interfaces.EmbedChannelService { + return &embedChannelService{ + repo: repo, + agentService: agentService, + chunkService: chunkService, + redis: redisClient, + } +} + +func generateEmbedPublishToken() (string, error) { + buf := make([]byte, embedTokenBytes) + if _, err := rand.Read(buf); err != nil { + return "", err + } + return "em_" + base64.RawURLEncoding.EncodeToString(buf), nil +} + +func (s *embedChannelService) Create( + ctx context.Context, tenantID uint64, agentID string, req *types.EmbedChannel, +) (*types.EmbedChannel, string, error) { + agentID = strings.TrimSpace(agentID) + if _, err := s.ensureAgentOwned(ctx, tenantID, agentID); err != nil { + return nil, "", err + } + token, err := generateEmbedPublishToken() + if err != nil { + return nil, "", err + } + originsJSON := req.AllowedOrigins + if len(originsJSON) == 0 { + originsJSON = []byte("[]") + } + ch := &types.EmbedChannel{ + TenantID: tenantID, + AgentID: agentID, + Name: strings.TrimSpace(req.Name), + Enabled: req.Enabled, + PublishToken: token, + AllowedOrigins: originsJSON, + WelcomeMessage: req.WelcomeMessage, + RateLimitPerMinute: req.RateLimitPerMinute, + RateLimitPerDay: req.RateLimitPerDay, + PrimaryColor: strings.TrimSpace(req.PrimaryColor), + PageTitle: strings.TrimSpace(req.PageTitle), + HeaderTitleMode: types.NormalizeEmbedHeaderTitleMode(req.HeaderTitleMode), + ShowSuggestedQuestions: req.ShowSuggestedQuestions, + WidgetPosition: types.NormalizeEmbedWidgetPosition(req.WidgetPosition), + AllowWebSearch: req.AllowWebSearch, + AllowMemory: req.AllowMemory, + AllowFileUpload: req.AllowFileUpload, + DefaultLocale: types.NormalizeEmbedDefaultLocale(req.DefaultLocale), + } + if ch.RateLimitPerMinute <= 0 { + ch.RateLimitPerMinute = 30 + } + if ch.RateLimitPerDay <= 0 { + ch.RateLimitPerDay = types.DefaultEmbedRateLimitPerDay + } + if err := s.repo.Create(ctx, ch); err != nil { + return nil, "", err + } + return ch, token, nil +} + +func (s *embedChannelService) ListByAgent( + ctx context.Context, tenantID uint64, agentID string, +) ([]*types.EmbedChannel, error) { + agentID = strings.TrimSpace(agentID) + if _, err := s.ensureAgentOwned(ctx, tenantID, agentID); err != nil { + return nil, err + } + return s.repo.ListByAgent(ctx, tenantID, agentID) +} + +func (s *embedChannelService) ListByTenant( + ctx context.Context, tenantID uint64, +) ([]*types.EmbedChannel, error) { + return s.repo.ListByTenant(ctx, tenantID) +} + +func (s *embedChannelService) Update( + ctx context.Context, tenantID uint64, id string, req *types.EmbedChannel, + enabled *bool, showSuggested *bool, allowWebSearch *bool, allowMemory *bool, allowFileUpload *bool, + defaultLocale *string, webhookURL *string, webhookSecret *string, +) (*types.EmbedChannel, error) { + ch, err := s.getOwned(ctx, tenantID, id) + if err != nil { + return nil, err + } + if req.Name != "" { + ch.Name = strings.TrimSpace(req.Name) + } + ch.WelcomeMessage = req.WelcomeMessage + ch.PrimaryColor = strings.TrimSpace(req.PrimaryColor) + ch.PageTitle = strings.TrimSpace(req.PageTitle) + ch.HeaderTitleMode = types.NormalizeEmbedHeaderTitleMode(req.HeaderTitleMode) + if showSuggested != nil { + ch.ShowSuggestedQuestions = *showSuggested + } + if allowWebSearch != nil { + ch.AllowWebSearch = *allowWebSearch + } + if allowMemory != nil { + ch.AllowMemory = *allowMemory + } + if allowFileUpload != nil { + ch.AllowFileUpload = *allowFileUpload + } + if defaultLocale != nil { + ch.DefaultLocale = types.NormalizeEmbedDefaultLocale(*defaultLocale) + } + if webhookURL != nil { + trimmed := strings.TrimSpace(*webhookURL) + if err := ValidateEmbedWebhookURL(trimmed); err != nil { + return nil, err + } + ch.WebhookURL = trimmed + } + if webhookSecret != nil { + ch.WebhookSecret = strings.TrimSpace(*webhookSecret) + } + if req.WidgetPosition != "" { + ch.WidgetPosition = types.NormalizeEmbedWidgetPosition(req.WidgetPosition) + } + if enabled != nil { + ch.Enabled = *enabled + } + if req.RateLimitPerMinute > 0 { + ch.RateLimitPerMinute = req.RateLimitPerMinute + } + if req.RateLimitPerDay > 0 { + ch.RateLimitPerDay = req.RateLimitPerDay + } + if req.AllowedOrigins != nil { + if len(req.AllowedOrigins) == 0 { + ch.AllowedOrigins = []byte("[]") + } else { + ch.AllowedOrigins = req.AllowedOrigins + } + } + if err := s.repo.Update(ctx, ch); err != nil { + return nil, err + } + return ch, nil +} + +func (s *embedChannelService) Delete(ctx context.Context, tenantID uint64, id string) error { + if _, err := s.getOwned(ctx, tenantID, id); err != nil { + return err + } + return s.repo.Delete(ctx, tenantID, id) +} + +func (s *embedChannelService) RotateToken( + ctx context.Context, tenantID uint64, id string, +) (*types.EmbedChannel, string, error) { + ch, err := s.getOwned(ctx, tenantID, id) + if err != nil { + return nil, "", err + } + token, err := generateEmbedPublishToken() + if err != nil { + return nil, "", err + } + ch.PublishToken = token + if err := s.repo.Update(ctx, ch); err != nil { + return nil, "", err + } + return ch, token, nil +} + +func (s *embedChannelService) LookupForEmbed( + ctx context.Context, channelID, token string, +) (*types.EmbedChannel, error) { + token = strings.TrimSpace(token) + if token == "" { + return nil, ErrEmbedTokenInvalid + } + ch, err := s.repo.GetByID(ctx, channelID) + if err != nil { + return nil, err + } + if ch == nil { + return nil, ErrEmbedTokenInvalid + } + if !ch.Enabled { + return nil, ErrEmbedChannelDisabled + } + if subtle.ConstantTimeCompare([]byte(ch.PublishToken), []byte(token)) != 1 { + return nil, ErrEmbedTokenInvalid + } + return ch, nil +} + +func (s *embedChannelService) PublicConfig(ctx context.Context, ch *types.EmbedChannel) types.EmbedChannelPublicConfig { + kbIDs := s.resolveKnowledgeBaseIDs(ctx, ch) + displayTitle, agentName, agentAvatar := s.resolveDisplayMeta(ctx, ch) + agentWebSearchEnabled := false + agentImageUploadEnabled := false + if agent, err := s.agentService.GetAgentByID(ctx, ch.AgentID); err == nil && agent != nil { + agentWebSearchEnabled = agent.Config.WebSearchEnabled + agentImageUploadEnabled = agent.Config.ImageUploadEnabled + } + return types.EmbedChannelPublicConfig{ + ChannelID: ch.ID, + Name: ch.Name, + DisplayTitle: displayTitle, + KnowledgeBaseIDs: kbIDs, + AgentID: ch.AgentID, + AgentName: agentName, + AgentAvatar: agentAvatar, + WelcomeMessage: ch.WelcomeMessage, + PrimaryColor: ch.PrimaryColor, + PageTitle: ch.PageTitle, + HeaderTitleMode: types.NormalizeEmbedHeaderTitleMode(ch.HeaderTitleMode), + ShowSuggestedQuestions: ch.ShowSuggestedQuestions, + AllowedOrigins: ch.AllowedOriginsList(), + WidgetPosition: types.NormalizeEmbedWidgetPosition(ch.WidgetPosition), + AllowWebSearch: ch.AllowWebSearch, + AllowMemory: ch.AllowMemory, + AllowFileUpload: ch.AllowFileUpload, + AgentWebSearchEnabled: agentWebSearchEnabled, + AgentImageUploadEnabled: agentImageUploadEnabled, + DefaultLocale: types.NormalizeEmbedDefaultLocale(ch.DefaultLocale), + } +} + +func (s *embedChannelService) EmbedChunk( + ctx context.Context, ch *types.EmbedChannel, chunkID string, +) (*types.Chunk, error) { + chunkID = strings.TrimSpace(chunkID) + if chunkID == "" { + return nil, ErrEmbedChunkNotFound + } + chunk, err := s.chunkService.GetChunkByIDOnly(ctx, chunkID) + if err != nil || chunk == nil { + return nil, ErrEmbedChunkNotFound + } + if !s.chunkAllowedForEmbed(ctx, ch, chunk) { + return nil, ErrEmbedChunkForbidden + } + return chunk, nil +} + +func (s *embedChannelService) chunkAllowedForEmbed( + ctx context.Context, ch *types.EmbedChannel, chunk *types.Chunk, +) bool { + if chunk == nil || chunk.KnowledgeBaseID == "" { + return false + } + // 显式重校验租户：GetChunkByIDOnly 无租户过滤，且 KBSelectionMode=="all"/默认 + // 分支会无条件放行，必须在此挡住跨租户 chunk id 撞库读取他人知识库正文。 + if ch == nil || chunk.TenantID != ch.TenantID { + return false + } + allowedKBs := s.resolveKnowledgeBaseIDs(ctx, ch) + if len(allowedKBs) > 0 { + for _, kbID := range allowedKBs { + if kbID == chunk.KnowledgeBaseID { + return true + } + } + return false + } + agent, err := s.agentService.GetAgentByID(ctx, ch.AgentID) + if err != nil || agent == nil { + return false + } + switch agent.Config.KBSelectionMode { + case "none": + return false + case "selected": + return false + default: + return true + } +} + +func (s *embedChannelService) SuggestedQuestions( + ctx context.Context, ch *types.EmbedChannel, limit int, +) ([]types.SuggestedQuestion, error) { + if ch == nil || !ch.ShowSuggestedQuestions { + return nil, nil + } + if limit <= 0 { + limit = 6 + } + kbIDs := s.resolveKnowledgeBaseIDs(ctx, ch) + return s.agentService.GetSuggestedQuestions(ctx, ch.AgentID, kbIDs, nil, limit) +} + +// EmbedDisplayTitle resolves the human-readable title for embed sessions and UI chrome. +func (s *embedChannelService) EmbedDisplayTitle(ctx context.Context, ch *types.EmbedChannel) string { + title, _, _ := s.resolveDisplayMeta(ctx, ch) + return title +} + +func (s *embedChannelService) resolveDisplayMeta( + ctx context.Context, ch *types.EmbedChannel, +) (displayTitle, agentName, agentAvatar string) { + if pageTitle := strings.TrimSpace(ch.PageTitle); pageTitle != "" { + displayTitle = pageTitle + } else if name := strings.TrimSpace(ch.Name); name != "" { + displayTitle = name + } + agent, err := s.agentService.GetAgentByID(ctx, ch.AgentID) + if err == nil && agent != nil { + agentName = strings.TrimSpace(agent.Name) + agentAvatar = strings.TrimSpace(agent.Avatar) + if displayTitle == "" && agentName != "" { + displayTitle = agentName + } + } + if displayTitle == "" { + displayTitle = "AI Assistant" + } + return displayTitle, agentName, agentAvatar +} + +func (s *embedChannelService) resolveKnowledgeBaseIDs(ctx context.Context, ch *types.EmbedChannel) []string { + agent, err := s.agentService.GetAgentByID(ctx, ch.AgentID) + if err == nil && agent != nil && agent.Config.KBSelectionMode == "selected" { + return append([]string(nil), agent.Config.KnowledgeBases...) + } + return nil +} + +func (s *embedChannelService) ensureAgentOwned(ctx context.Context, tenantID uint64, agentID string) (*types.CustomAgent, error) { + agentID = strings.TrimSpace(agentID) + if agentID == "" { + return nil, apperrors.NewBadRequestError("agent_id is required") + } + agent, err := s.agentService.GetAgentByID(ctx, agentID) + if err != nil { + return nil, err + } + if agent == nil || agent.TenantID != tenantID { + return nil, apperrors.NewNotFoundError("agent not found") + } + return agent, nil +} + +func (s *embedChannelService) GetOwnedChannel( + ctx context.Context, tenantID uint64, id string, +) (*types.EmbedChannel, error) { + return s.getOwned(ctx, tenantID, id) +} + +func (s *embedChannelService) getOwned(ctx context.Context, tenantID uint64, id string) (*types.EmbedChannel, error) { + ch, err := s.repo.GetByID(ctx, id) + if err != nil { + return nil, err + } + if ch == nil || ch.TenantID != tenantID { + return nil, ErrEmbedChannelNotFound + } + return ch, nil +} + +// EmbedSessionDescription returns the marker stored on embed-created sessions. +func EmbedSessionDescription(channelID string) string { + return fmt.Sprintf("%s%s", types.EmbedSessionMarkerPrefix, channelID) +} diff --git a/internal/application/service/embed_channel_chunk_test.go b/internal/application/service/embed_channel_chunk_test.go new file mode 100644 index 00000000..764a0b6a --- /dev/null +++ b/internal/application/service/embed_channel_chunk_test.go @@ -0,0 +1,123 @@ +package service + +import ( + "context" + "errors" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" +) + +type stubChunkForEmbed struct { + interfaces.ChunkService + chunks map[string]*types.Chunk +} + +func (s *stubChunkForEmbed) GetChunkByIDOnly(_ context.Context, id string) (*types.Chunk, error) { + chunk, ok := s.chunks[id] + if !ok { + return nil, ErrChunkNotFound + } + return chunk, nil +} + +func newEmbedChunkService(agent *types.CustomAgent, chunks map[string]*types.Chunk) *embedChannelService { + return &embedChannelService{ + agentService: &stubAgentForEmbed{agent: agent}, + chunkService: &stubChunkForEmbed{chunks: chunks}, + } +} + +func TestChunkAllowedForEmbedAllowedKB(t *testing.T) { + svc := newEmbedChunkService( + &types.CustomAgent{ + Config: types.CustomAgentConfig{ + KBSelectionMode: "selected", + KnowledgeBases: []string{"kb-allowed"}, + }, + }, + map[string]*types.Chunk{ + "chunk-1": {ID: "chunk-1", TenantID: 5, KnowledgeBaseID: "kb-allowed"}, + }, + ) + ch := &types.EmbedChannel{ID: "ch-1", TenantID: 5, AgentID: "agent-1"} + + if !svc.chunkAllowedForEmbed(context.Background(), ch, &types.Chunk{TenantID: 5, KnowledgeBaseID: "kb-allowed"}) { + t.Fatal("chunk in allowed KB should be permitted") + } + + chunk, err := svc.EmbedChunk(context.Background(), ch, "chunk-1") + if err != nil { + t.Fatalf("EmbedChunk() = %v", err) + } + if chunk == nil || chunk.ID != "chunk-1" { + t.Fatalf("unexpected chunk: %#v", chunk) + } +} + +func TestChunkAllowedForEmbedCrossTenantDenied(t *testing.T) { + svc := newEmbedChunkService( + &types.CustomAgent{ + Config: types.CustomAgentConfig{ + KBSelectionMode: "selected", + KnowledgeBases: []string{"kb-allowed"}, + }, + }, + map[string]*types.Chunk{ + "chunk-x": {ID: "chunk-x", TenantID: 999, KnowledgeBaseID: "kb-allowed"}, + }, + ) + ch := &types.EmbedChannel{ID: "ch-1", TenantID: 5, AgentID: "agent-1"} + chunk := &types.Chunk{ID: "chunk-x", TenantID: 999, KnowledgeBaseID: "kb-allowed"} + + if svc.chunkAllowedForEmbed(context.Background(), ch, chunk) { + t.Fatal("cross-tenant chunk must be denied before KB checks") + } + + _, err := svc.EmbedChunk(context.Background(), ch, "chunk-x") + if !errors.Is(err, ErrEmbedChunkForbidden) { + t.Fatalf("EmbedChunk() = %v, want ErrEmbedChunkForbidden", err) + } +} + +func TestChunkAllowedForEmbedWrongKBDenied(t *testing.T) { + svc := newEmbedChunkService( + &types.CustomAgent{ + Config: types.CustomAgentConfig{ + KBSelectionMode: "selected", + KnowledgeBases: []string{"kb-allowed"}, + }, + }, + map[string]*types.Chunk{ + "chunk-other": {ID: "chunk-other", TenantID: 5, KnowledgeBaseID: "kb-other"}, + }, + ) + ch := &types.EmbedChannel{ID: "ch-1", TenantID: 5, AgentID: "agent-1"} + chunk := &types.Chunk{ID: "chunk-other", TenantID: 5, KnowledgeBaseID: "kb-other"} + + if svc.chunkAllowedForEmbed(context.Background(), ch, chunk) { + t.Fatal("chunk outside allowed KB list must be denied") + } + + _, err := svc.EmbedChunk(context.Background(), ch, "chunk-other") + if !errors.Is(err, ErrEmbedChunkForbidden) { + t.Fatalf("EmbedChunk() = %v, want ErrEmbedChunkForbidden", err) + } +} + +func TestChunkAllowedForEmbedAllModePermitsInTenant(t *testing.T) { + svc := newEmbedChunkService( + &types.CustomAgent{ + Config: types.CustomAgentConfig{KBSelectionMode: "all"}, + }, + map[string]*types.Chunk{ + "chunk-all": {ID: "chunk-all", TenantID: 5, KnowledgeBaseID: "kb-any"}, + }, + ) + ch := &types.EmbedChannel{ID: "ch-1", TenantID: 5, AgentID: "agent-1"} + + if !svc.chunkAllowedForEmbed(context.Background(), ch, &types.Chunk{TenantID: 5, KnowledgeBaseID: "kb-any"}) { + t.Fatal("all-mode agent should permit in-tenant chunks") + } +} diff --git a/internal/application/service/embed_channel_public_config_test.go b/internal/application/service/embed_channel_public_config_test.go new file mode 100644 index 00000000..b44c74b4 --- /dev/null +++ b/internal/application/service/embed_channel_public_config_test.go @@ -0,0 +1,142 @@ +package service + +import ( + "context" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" +) + +type stubAgentForEmbed struct { + interfaces.CustomAgentService + agent *types.CustomAgent + err error +} + +func (s *stubAgentForEmbed) GetAgentByID(_ context.Context, _ string) (*types.CustomAgent, error) { + if s.err != nil { + return nil, s.err + } + return s.agent, nil +} + +func TestResolveKnowledgeBaseIDsFromAgent(t *testing.T) { + svc := &embedChannelService{ + agentService: &stubAgentForEmbed{ + agent: &types.CustomAgent{ + Config: types.CustomAgentConfig{ + KBSelectionMode: "selected", + KnowledgeBases: []string{"kb-a", "kb-b"}, + }, + }, + }, + } + ids := svc.resolveKnowledgeBaseIDs(context.Background(), &types.EmbedChannel{ + AgentID: "agent-1", + }) + if len(ids) != 2 || ids[0] != "kb-a" || ids[1] != "kb-b" { + t.Fatalf("unexpected kb ids: %#v", ids) + } +} + +func TestPublicConfigUsesAgentKBs(t *testing.T) { + svc := &embedChannelService{ + agentService: &stubAgentForEmbed{ + agent: &types.CustomAgent{ + Config: types.CustomAgentConfig{ + KBSelectionMode: "selected", + KnowledgeBases: []string{"kb-primary"}, + }, + }, + }, + } + cfg := svc.PublicConfig(context.Background(), &types.EmbedChannel{ + ID: "ch-1", + AgentID: "agent-1", + Name: "Support", + }) + if len(cfg.KnowledgeBaseIDs) != 1 || cfg.KnowledgeBaseIDs[0] != "kb-primary" { + t.Fatalf("unexpected public config: %#v", cfg) + } + if cfg.AgentID != "agent-1" || cfg.ChannelID != "ch-1" { + t.Fatalf("unexpected ids: %#v", cfg) + } + if cfg.DisplayTitle != "Support" { + t.Fatalf("display_title = %q, want channel name", cfg.DisplayTitle) + } +} + +func TestPublicConfigDisplayTitlePrefersPageTitle(t *testing.T) { + svc := &embedChannelService{ + agentService: &stubAgentForEmbed{ + agent: &types.CustomAgent{ + Name: "Agent Name", + }, + }, + } + cfg := svc.PublicConfig(context.Background(), &types.EmbedChannel{ + ID: "ch-3", + AgentID: "agent-3", + Name: "Internal Channel", + PageTitle: "Visitor Chat", + }) + if cfg.DisplayTitle != "Visitor Chat" { + t.Fatalf("display_title = %q, want page title", cfg.DisplayTitle) + } +} + +func TestPublicConfigIncludesCapabilityFlags(t *testing.T) { + svc := &embedChannelService{ + agentService: &stubAgentForEmbed{ + agent: &types.CustomAgent{ + Config: types.CustomAgentConfig{ + KBSelectionMode: "all", + WebSearchEnabled: true, + }, + }, + }, + } + cfg := svc.PublicConfig(context.Background(), &types.EmbedChannel{ + ID: "ch-cap", + AgentID: "agent-cap", + AllowWebSearch: true, + AllowMemory: true, + WidgetPosition: "top-left", + }) + if !cfg.AllowWebSearch { + t.Fatalf("allow_web_search = false, want true") + } + if !cfg.AllowMemory { + t.Fatalf("allow_memory = false, want true") + } + if !cfg.AgentWebSearchEnabled { + t.Fatalf("agent_web_search_enabled = false, want true") + } + if cfg.WidgetPosition != "top-left" { + t.Fatalf("widget_position = %q, want top-left", cfg.WidgetPosition) + } +} + +func TestPublicConfigDisplayTitleFallsBackToAgentName(t *testing.T) { + svc := &embedChannelService{ + agentService: &stubAgentForEmbed{ + agent: &types.CustomAgent{ + Name: "客服助手", + Config: types.CustomAgentConfig{ + KBSelectionMode: "all", + }, + }, + }, + } + cfg := svc.PublicConfig(context.Background(), &types.EmbedChannel{ + ID: "ch-2", + AgentID: "agent-2", + }) + if cfg.DisplayTitle != "客服助手" { + t.Fatalf("display_title = %q, want agent name", cfg.DisplayTitle) + } + if cfg.AgentName != "客服助手" { + t.Fatalf("agent_name = %q", cfg.AgentName) + } +} diff --git a/internal/application/service/embed_session.go b/internal/application/service/embed_session.go new file mode 100644 index 00000000..4459dde9 --- /dev/null +++ b/internal/application/service/embed_session.go @@ -0,0 +1,143 @@ +package service + +import ( + "context" + "crypto/hmac" + "crypto/rand" + "crypto/sha256" + "crypto/subtle" + "encoding/base64" + "errors" + "strings" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/redis/go-redis/v9" +) + +const ( + embedSessionTokenPrefix = "ems_" + embedSessionRedisPrefix = "embed:session:" + embedSessionTTL = 30 * time.Minute +) + +var ErrEmbedSessionUnavailable = errors.New("embed session tokens unavailable") + +func generateEmbedSessionToken() (string, error) { + buf := make([]byte, embedTokenBytes) + if _, err := rand.Read(buf); err != nil { + return "", err + } + return embedSessionTokenPrefix + base64.RawURLEncoding.EncodeToString(buf), nil +} + +// IssueSessionToken mints a short-lived session token bound to channelID. +func (s *embedChannelService) IssueSessionToken(ctx context.Context, channelID string) (string, int, error) { + if s.redis == nil { + return "", 0, ErrEmbedSessionUnavailable + } + channelID = strings.TrimSpace(channelID) + if channelID == "" { + return "", 0, ErrEmbedTokenInvalid + } + token, err := generateEmbedSessionToken() + if err != nil { + return "", 0, err + } + key := embedSessionRedisPrefix + token + if err := s.redis.Set(ctx, key, channelID, embedSessionTTL).Err(); err != nil { + return "", 0, err + } + return token, int(embedSessionTTL.Seconds()), nil +} + +// ResolveSessionToken returns the channel ID stored for a session token. +func (s *embedChannelService) ResolveSessionToken(ctx context.Context, token string) (string, error) { + token = strings.TrimSpace(token) + if !strings.HasPrefix(token, embedSessionTokenPrefix) { + return "", ErrEmbedTokenInvalid + } + if s.redis == nil { + return "", ErrEmbedSessionUnavailable + } + channelID, err := s.redis.Get(ctx, embedSessionRedisPrefix+token).Result() + if err == redis.Nil { + return "", ErrEmbedTokenInvalid + } + if err != nil { + return "", err + } + channelID = strings.TrimSpace(channelID) + if channelID == "" { + return "", ErrEmbedTokenInvalid + } + return channelID, nil +} + +// LookupEnabledChannel loads an embed channel and verifies it is enabled. +func (s *embedChannelService) LookupEnabledChannel(ctx context.Context, channelID string) (*types.EmbedChannel, error) { + channelID = strings.TrimSpace(channelID) + if channelID == "" { + return nil, ErrEmbedTokenInvalid + } + ch, err := s.repo.GetByID(ctx, channelID) + if err != nil { + return nil, err + } + if ch == nil { + return nil, ErrEmbedTokenInvalid + } + if !ch.Enabled { + return nil, ErrEmbedChannelDisabled + } + return ch, nil +} + +// IsEmbedSessionToken reports whether token is a session token (ems_ prefix). +func IsEmbedSessionToken(token string) bool { + return strings.HasPrefix(strings.TrimSpace(token), embedSessionTokenPrefix) +} + +// SignEmbedSessionHandle binds a chat session id to its embed channel with an +// HMAC keyed by the channel's (server-only) publish token. The handle is handed +// to the widget at session-creation time and must be presented on every history +// load / chat call. Because the session id travels in the request path (and can +// land in access logs), this signature — sent in a header, never logged — is the +// real authorization secret: a leaked session id is useless without it. Rotating +// the channel token invalidates outstanding handles, which is acceptable. +func SignEmbedSessionHandle(ch *types.EmbedChannel, sessionID string) string { + if ch == nil || strings.TrimSpace(sessionID) == "" { + return "" + } + mac := hmac.New(sha256.New, []byte(ch.PublishToken)) + mac.Write([]byte(ch.ID + "|" + sessionID)) + return base64.RawURLEncoding.EncodeToString(mac.Sum(nil)) +} + +// VerifyEmbedSessionHandle reports whether sig is a valid handle for sessionID +// on channel ch, using a constant-time comparison. +func VerifyEmbedSessionHandle(ch *types.EmbedChannel, sessionID, sig string) bool { + sig = strings.TrimSpace(sig) + if sig == "" { + return false + } + expected := SignEmbedSessionHandle(ch, sessionID) + if expected == "" { + return false + } + return subtle.ConstantTimeCompare([]byte(expected), []byte(sig)) == 1 +} + +// IssuePreviewSession mints a short-lived session token for management UI preview. +func (s *embedChannelService) IssuePreviewSession( + ctx context.Context, tenantID uint64, channelID string, +) (string, int, error) { + ch, err := s.getOwned(ctx, tenantID, channelID) + if err != nil { + return "", 0, err + } + if !ch.Enabled { + return "", 0, ErrEmbedChannelDisabled + } + return s.IssueSessionToken(ctx, ch.ID) +} diff --git a/internal/application/service/embed_session_test.go b/internal/application/service/embed_session_test.go new file mode 100644 index 00000000..b9f5721b --- /dev/null +++ b/internal/application/service/embed_session_test.go @@ -0,0 +1,85 @@ +package service + +import ( + "context" + "testing" + + "github.com/Tencent/WeKnora/internal/types" +) + +func TestEmbedSessionHandleSignVerify(t *testing.T) { + ch := &types.EmbedChannel{ID: "ch-1", PublishToken: "em_secret_token"} + const sessionID = "11111111-2222-3333-4444-555555555555" + + sig := SignEmbedSessionHandle(ch, sessionID) + if sig == "" { + t.Fatal("expected a non-empty signature") + } + if !VerifyEmbedSessionHandle(ch, sessionID, sig) { + t.Fatal("valid handle should verify") + } + + // Tampered session id must not verify with the same signature. + if VerifyEmbedSessionHandle(ch, "00000000-0000-0000-0000-000000000000", sig) { + t.Fatal("signature must be bound to the session id") + } + + // A different channel token (e.g. after rotation) invalidates the handle. + rotated := &types.EmbedChannel{ID: "ch-1", PublishToken: "em_rotated_token"} + if VerifyEmbedSessionHandle(rotated, sessionID, sig) { + t.Fatal("rotating the channel token must invalidate old handles") + } + + // A different channel id must not verify. + other := &types.EmbedChannel{ID: "ch-2", PublishToken: "em_secret_token"} + if VerifyEmbedSessionHandle(other, sessionID, sig) { + t.Fatal("signature must be bound to the channel id") + } + + // Empty / missing signature is rejected. + if VerifyEmbedSessionHandle(ch, sessionID, "") { + t.Fatal("empty signature must be rejected") + } + if SignEmbedSessionHandle(nil, sessionID) != "" { + t.Fatal("nil channel must yield empty signature") + } + if SignEmbedSessionHandle(ch, "") != "" { + t.Fatal("empty session id must yield empty signature") + } +} + +func TestIsEmbedSessionToken(t *testing.T) { + if !IsEmbedSessionToken("ems_abc123") { + t.Fatal("expected ems_ prefix to be session token") + } + if IsEmbedSessionToken("em_abc123") { + t.Fatal("publish token must not be treated as session token") + } + if IsEmbedSessionToken("") { + t.Fatal("empty token must not match") + } +} + +func TestIssueSessionTokenWithoutRedis(t *testing.T) { + svc := &embedChannelService{redis: nil} + _, _, err := svc.IssueSessionToken(context.Background(), "channel-1") + if err != ErrEmbedSessionUnavailable { + t.Fatalf("expected ErrEmbedSessionUnavailable, got %v", err) + } +} + +func TestResolveSessionTokenWithoutRedis(t *testing.T) { + svc := &embedChannelService{redis: nil} + _, err := svc.ResolveSessionToken(context.Background(), "ems_test") + if err != ErrEmbedSessionUnavailable { + t.Fatalf("expected ErrEmbedSessionUnavailable, got %v", err) + } +} + +func TestResolveSessionTokenRejectsPublishToken(t *testing.T) { + svc := &embedChannelService{redis: nil} + _, err := svc.ResolveSessionToken(context.Background(), "em_publish_only") + if err != ErrEmbedTokenInvalid { + t.Fatalf("expected ErrEmbedTokenInvalid for publish token, got %v", err) + } +} diff --git a/internal/application/service/embed_webhook.go b/internal/application/service/embed_webhook.go new file mode 100644 index 00000000..2f1b9d26 --- /dev/null +++ b/internal/application/service/embed_webhook.go @@ -0,0 +1,111 @@ +package service + +import ( + "bytes" + "context" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" + + "github.com/Tencent/WeKnora/internal/logger" + "github.com/Tencent/WeKnora/internal/types" + secutils "github.com/Tencent/WeKnora/internal/utils" +) + +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") + +// ValidateEmbedWebhookURL checks an optional outbound webhook URL. Empty is allowed. +func ValidateEmbedWebhookURL(raw string) error { + trimmed := strings.TrimSpace(raw) + if trimmed == "" { + return nil + } + parsed, err := url.Parse(trimmed) + if err != nil || parsed.Host == "" { + return fmt.Errorf("%w: webhook URL must be a valid http(s) URL", ErrEmbedWebhookURLInvalid) + } + switch parsed.Scheme { + case "http", "https": + default: + return fmt.Errorf("%w: webhook URL must use http or https", ErrEmbedWebhookURLInvalid) + } + if err := secutils.ValidateURLForSSRF(trimmed); err != nil { + if hint := secutils.FormatSSRFError("Webhook URL", trimmed, err); hint != "" { + return fmt.Errorf("%w: %s", ErrEmbedWebhookURLInvalid, hint) + } + return fmt.Errorf("%w: %v", ErrEmbedWebhookURLInvalid, err) + } + return nil +} + +// DispatchEmbedWebhook POSTs an event to the channel webhook URL (best-effort, async). +func DispatchEmbedWebhook(ch *types.EmbedChannel, eventType, sessionID string, payload map[string]any) { + if ch == nil { + return + } + url := strings.TrimSpace(ch.WebhookURL) + if url == "" { + return + } + if err := ValidateEmbedWebhookURL(url); err != nil { + logger.Warnf(context.Background(), "[embed_webhook] skip dispatch %s: %v", eventType, err) + return + } + body := map[string]any{ + "type": eventType, + "channel_id": ch.ID, + "session_id": sessionID, + "timestamp": time.Now().UTC().Format(time.RFC3339), + } + for k, v := range payload { + body[k] = v + } + raw, err := json.Marshal(body) + if err != nil { + return + } + secret := strings.TrimSpace(ch.WebhookSecret) + go func() { + ctx, cancel := context.WithTimeout(context.Background(), embedWebhookTimeout) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(raw)) + if err != nil { + return + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", "WeKnora-Embed-Webhook/1.0") + if secret != "" { + mac := hmac.New(sha256.New, []byte(secret)) + _, _ = mac.Write(raw) + req.Header.Set("X-WeKnora-Signature", "sha256="+hex.EncodeToString(mac.Sum(nil))) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + logger.Warnf(context.Background(), "[embed_webhook] dispatch %s failed: %v", eventType, err) + return + } + defer resp.Body.Close() + _, _ = io.Copy(io.Discard, resp.Body) + if resp.StatusCode >= 300 { + logger.Warnf(context.Background(), "[embed_webhook] dispatch %s HTTP %d", eventType, resp.StatusCode) + } + }() +} + +// SignEmbedWebhookBody returns the hex HMAC signature for tests. +func SignEmbedWebhookBody(secret string, raw []byte) string { + mac := hmac.New(sha256.New, []byte(secret)) + _, _ = mac.Write(raw) + return hex.EncodeToString(mac.Sum(nil)) +} diff --git a/internal/application/service/embed_webhook_test.go b/internal/application/service/embed_webhook_test.go new file mode 100644 index 00000000..f67732ab --- /dev/null +++ b/internal/application/service/embed_webhook_test.go @@ -0,0 +1,35 @@ +package service + +import "testing" + +func TestValidateEmbedWebhookURL(t *testing.T) { + withSSRFWhitelist(t, "*.example.com,127.0.0.1") + + if err := ValidateEmbedWebhookURL(""); err != nil { + t.Fatalf("empty URL should be allowed: %v", err) + } + if err := ValidateEmbedWebhookURL("https://hooks.example.com/weknora/events"); err != nil { + t.Fatalf("whitelisted public https URL should pass: %v", err) + } + if err := ValidateEmbedWebhookURL("ftp://hooks.example.com/x"); err == nil { + t.Fatal("expected non-http(s) scheme to fail") + } + if err := ValidateEmbedWebhookURL("http://127.0.0.1/webhook"); err != nil { + t.Fatalf("whitelisted loopback should pass: %v", err) + } + if err := ValidateEmbedWebhookURL("http://169.254.169.254/latest/meta-data/"); err == nil { + t.Fatal("expected link-local metadata URL to be blocked") + } +} + +func TestSignEmbedWebhookBody(t *testing.T) { + raw := []byte(`{"type":"message_sent","query":"hi"}`) + sig := SignEmbedWebhookBody("test-secret", raw) + if sig == "" || len(sig) != 64 { + t.Fatalf("unexpected signature: %q", sig) + } + sig2 := SignEmbedWebhookBody("test-secret", raw) + if sig != sig2 { + t.Fatal("signature not deterministic") + } +} diff --git a/internal/container/container.go b/internal/container/container.go index ebb49f3e..22727ee4 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -157,6 +157,7 @@ func BuildContainer(container *dig.Container) *dig.Container { must(container.Provide(repository.NewOrganizationRepository)) must(container.Provide(repository.NewKBShareRepository)) must(container.Provide(repository.NewAgentShareRepository)) + must(container.Provide(repository.NewEmbedChannelRepository)) must(container.Provide(repository.NewTenantDisabledSharedAgentRepository)) must(container.Provide(repository.NewUserResourceFavoriteRepository)) must(container.Provide(service.NewWebSearchStateService)) @@ -210,6 +211,7 @@ func BuildContainer(container *dig.Container) *dig.Container { must(container.Provide(service.NewWikiLogEntryService)) must(container.Provide(service.NewWikiIngestService, dig.Name("wikiIngest"))) must(container.Provide(service.NewWikiLintService)) + must(container.Provide(service.NewEmbedChannelService)) // Web search service (needed by AgentService) logger.Debugf(ctx, "[Container] Registering web search registry and providers...") @@ -344,6 +346,7 @@ func BuildContainer(container *dig.Container) *dig.Container { must(container.Provide(imPkg.NewService)) must(container.Invoke(registerIMAdapterFactories)) must(container.Provide(handler.NewIMHandler)) + must(container.Provide(handler.NewEmbedChannelHandler)) must(container.Provide(handler.NewWeKnoraCloudHandler)) logger.Debugf(ctx, "[Container] HTTP handlers registered") diff --git a/internal/handler/embed_channel.go b/internal/handler/embed_channel.go new file mode 100644 index 00000000..4714ddb0 --- /dev/null +++ b/internal/handler/embed_channel.go @@ -0,0 +1,681 @@ +package handler + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "strconv" + "strings" + + "github.com/Tencent/WeKnora/internal/application/service" + apperrors "github.com/Tencent/WeKnora/internal/errors" + "github.com/Tencent/WeKnora/internal/handler/session" + "github.com/Tencent/WeKnora/internal/logger" + "github.com/Tencent/WeKnora/internal/middleware" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + secutils "github.com/Tencent/WeKnora/internal/utils" + "github.com/gin-gonic/gin" + "github.com/redis/go-redis/v9" +) + +// EmbedChannelHandler manages web embed channel CRUD and public embed endpoints. +type EmbedChannelHandler struct { + embedSvc interfaces.EmbedChannelService + sessionService interfaces.SessionService + sessionHandler *session.Handler + messageHandler *MessageHandler + redis *redis.Client +} + +func NewEmbedChannelHandler( + embedSvc interfaces.EmbedChannelService, + sessionService interfaces.SessionService, + sessionHandler *session.Handler, + messageHandler *MessageHandler, + redisClient *redis.Client, +) *EmbedChannelHandler { + return &EmbedChannelHandler{ + embedSvc: embedSvc, + sessionService: sessionService, + sessionHandler: sessionHandler, + messageHandler: messageHandler, + redis: redisClient, + } +} + +type embedChannelRequest struct { + Name string `json:"name"` + Enabled *bool `json:"enabled"` + AllowedOrigins []string `json:"allowed_origins"` + WelcomeMessage string `json:"welcome_message"` + RateLimitPerMinute int `json:"rate_limit_per_minute"` + RateLimitPerDay int `json:"rate_limit_per_day"` + PrimaryColor string `json:"primary_color"` + PageTitle string `json:"page_title"` + HeaderTitleMode string `json:"header_title_mode"` + ShowSuggestedQuestions *bool `json:"show_suggested_questions"` + WidgetPosition string `json:"widget_position"` + AllowWebSearch *bool `json:"allow_web_search"` + AllowMemory *bool `json:"allow_memory"` + AllowFileUpload *bool `json:"allow_file_upload"` + DefaultLocale *string `json:"default_locale"` + WebhookURL *string `json:"webhook_url"` + WebhookSecret *string `json:"webhook_secret"` +} + +// isProductionMode reports whether the server runs in a hardened (release) mode. +func isProductionMode() bool { + return strings.EqualFold(strings.TrimSpace(os.Getenv("GIN_MODE")), "release") +} + +func stringOrEmpty(v *string) string { + if v == nil { + return "" + } + return *v +} + +// validateAllowedOrigins enforces that a public embed channel declares an +// explicit origin allowlist. An empty list means "allow any origin" in the +// auth middleware, which is unsafe for a publicly reachable widget, so it is +// rejected. In production a wildcard ("*") is also rejected; each entry must be +// a well-formed http(s) origin (optionally a "*." subdomain wildcard). +func validateAllowedOrigins(origins []string) error { + cleaned := make([]string, 0, len(origins)) + for _, o := range origins { + o = strings.TrimSpace(o) + if o == "" { + continue + } + cleaned = append(cleaned, o) + } + if len(cleaned) == 0 { + return fmt.Errorf("at least one allowed origin is required") + } + for _, o := range cleaned { + if o == "*" { + if isProductionMode() { + return fmt.Errorf("wildcard origin '*' is not allowed in production") + } + continue + } + host := o + if strings.HasPrefix(o, "*.") { + host = "https://" + strings.TrimPrefix(o, "*.") + } + u, err := url.Parse(host) + if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" { + return fmt.Errorf("invalid allowed origin: %q", o) + } + } + return nil +} + +func (h *EmbedChannelHandler) CreateEmbedChannel(c *gin.Context) { + agentID := secutils.SanitizeForLog(c.Param("id")) + tenantID := c.GetUint64(types.TenantIDContextKey.String()) + var req embedChannelRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := validateAllowedOrigins(req.AllowedOrigins); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + originsJSON, _ := json.Marshal(req.AllowedOrigins) + enabled := true + if req.Enabled != nil { + enabled = *req.Enabled + } + showSuggested := true + if req.ShowSuggestedQuestions != nil { + showSuggested = *req.ShowSuggestedQuestions + } + allowWebSearch := false + if req.AllowWebSearch != nil { + allowWebSearch = *req.AllowWebSearch + } + allowMemory := false + if req.AllowMemory != nil { + allowMemory = *req.AllowMemory + } + allowFileUpload := false + if req.AllowFileUpload != nil { + allowFileUpload = *req.AllowFileUpload + } + ch, token, err := h.embedSvc.Create(c.Request.Context(), tenantID, agentID, &types.EmbedChannel{ + Name: req.Name, + Enabled: enabled, + AllowedOrigins: originsJSON, + WelcomeMessage: req.WelcomeMessage, + RateLimitPerMinute: req.RateLimitPerMinute, + RateLimitPerDay: req.RateLimitPerDay, + PrimaryColor: req.PrimaryColor, + PageTitle: req.PageTitle, + HeaderTitleMode: req.HeaderTitleMode, + ShowSuggestedQuestions: showSuggested, + WidgetPosition: req.WidgetPosition, + AllowWebSearch: allowWebSearch, + AllowMemory: allowMemory, + AllowFileUpload: allowFileUpload, + DefaultLocale: types.NormalizeEmbedDefaultLocale(stringOrEmpty(req.DefaultLocale)), + }) + if err != nil { + writeEmbedMgmtError(c, err) + return + } + c.JSON(http.StatusCreated, gin.H{ + "success": true, + "data": embedChannelResponse(ch, token), + }) +} + +func (h *EmbedChannelHandler) ListEmbedChannels(c *gin.Context) { + agentID := secutils.SanitizeForLog(c.Param("id")) + tenantID := c.GetUint64(types.TenantIDContextKey.String()) + rows, err := h.embedSvc.ListByAgent(c.Request.Context(), tenantID, agentID) + if err != nil { + writeEmbedMgmtError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"success": true, "data": embedChannelsResponse(rows)}) +} + +// ListAllEmbedChannels lists every embed channel in the current tenant, across +// agents, for sidebar session grouping. Publish tokens are never included. +func (h *EmbedChannelHandler) ListAllEmbedChannels(c *gin.Context) { + tenantID := c.GetUint64(types.TenantIDContextKey.String()) + rows, err := h.embedSvc.ListByTenant(c.Request.Context(), tenantID) + if err != nil { + writeEmbedMgmtError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"success": true, "data": embedChannelsResponse(rows)}) +} + +func embedChannelsResponse(rows []*types.EmbedChannel) []gin.H { + data := make([]gin.H, 0, len(rows)) + for _, ch := range rows { + data = append(data, embedChannelResponse(ch, "")) + } + return data +} + +func (h *EmbedChannelHandler) UpdateEmbedChannel(c *gin.Context) { + channelID := secutils.SanitizeForLog(c.Param("channel_id")) + tenantID := c.GetUint64(types.TenantIDContextKey.String()) + var req embedChannelRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + // Only validate when the caller intends to change the allowlist. A nil slice + // means "leave unchanged"; a present slice must still be a valid allowlist. + if req.AllowedOrigins != nil { + if err := validateAllowedOrigins(req.AllowedOrigins); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + } + if req.WebhookURL != nil { + if err := service.ValidateEmbedWebhookURL(*req.WebhookURL); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + } + originsJSON, _ := json.Marshal(req.AllowedOrigins) + update := &types.EmbedChannel{ + Name: req.Name, + AllowedOrigins: originsJSON, + WelcomeMessage: req.WelcomeMessage, + RateLimitPerMinute: req.RateLimitPerMinute, + RateLimitPerDay: req.RateLimitPerDay, + PrimaryColor: req.PrimaryColor, + PageTitle: req.PageTitle, + HeaderTitleMode: req.HeaderTitleMode, + WidgetPosition: req.WidgetPosition, + } + ch, err := h.embedSvc.Update(c.Request.Context(), tenantID, channelID, update, req.Enabled, req.ShowSuggestedQuestions, req.AllowWebSearch, req.AllowMemory, req.AllowFileUpload, req.DefaultLocale, req.WebhookURL, req.WebhookSecret) + if err != nil { + writeEmbedMgmtError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"success": true, "data": embedChannelResponse(ch, "")}) +} + +func (h *EmbedChannelHandler) DeleteEmbedChannel(c *gin.Context) { + channelID := secutils.SanitizeForLog(c.Param("channel_id")) + tenantID := c.GetUint64(types.TenantIDContextKey.String()) + if err := h.embedSvc.Delete(c.Request.Context(), tenantID, channelID); err != nil { + writeEmbedMgmtError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +func (h *EmbedChannelHandler) RotateEmbedToken(c *gin.Context) { + channelID := secutils.SanitizeForLog(c.Param("channel_id")) + tenantID := c.GetUint64(types.TenantIDContextKey.String()) + ch, token, err := h.embedSvc.RotateToken(c.Request.Context(), tenantID, channelID) + if err != nil { + writeEmbedMgmtError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"success": true, "data": embedChannelResponse(ch, token)}) +} + +func (h *EmbedChannelHandler) IssuePreviewSession(c *gin.Context) { + channelID := secutils.SanitizeForLog(c.Param("channel_id")) + tenantID := c.GetUint64(types.TenantIDContextKey.String()) + sessionToken, expiresIn, err := h.embedSvc.IssuePreviewSession(c.Request.Context(), tenantID, channelID) + if err != nil { + if errors.Is(err, service.ErrEmbedChannelDisabled) { + c.JSON(http.StatusForbidden, gin.H{"error": "embed channel is disabled"}) + return + } + writeEmbedMgmtError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "data": gin.H{ + "session_token": sessionToken, + "expires_in": expiresIn, + }, + }) +} + +func (h *EmbedChannelHandler) ExchangeEmbedSession(c *gin.Context) { + ctx := c.Request.Context() + ch, ok := middleware.EmbedChannelFromContext(ctx) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return + } + // Only the long-lived publish token may mint session tokens. Accepting a + // session token here would let a holder renew it indefinitely without ever + // re-presenting the publish token. + if auth := strings.TrimSpace(c.GetHeader("Authorization")); !strings.HasPrefix(auth, "Embed ") || + service.IsEmbedSessionToken(strings.TrimPrefix(auth, "Embed ")) { + c.JSON(http.StatusForbidden, gin.H{"error": "publish token required"}) + return + } + sessionToken, expiresIn, err := h.embedSvc.IssueSessionToken(ctx, ch.ID) + if err != nil { + if errors.Is(err, service.ErrEmbedSessionUnavailable) { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "session tokens unavailable"}) + return + } + logger.ErrorWithFields(ctx, err, nil) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to issue session token"}) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "data": gin.H{ + "session_token": sessionToken, + "expires_in": expiresIn, + }, + }) +} + +func (h *EmbedChannelHandler) GetEmbedConfig(c *gin.Context) { + ch, ok := middleware.EmbedChannelFromContext(c.Request.Context()) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true, "data": h.embedSvc.PublicConfig(c.Request.Context(), ch)}) +} + +func (h *EmbedChannelHandler) GetEmbedChunk(c *gin.Context) { + ch, ok := middleware.EmbedChannelFromContext(c.Request.Context()) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return + } + chunkID := secutils.SanitizeForLog(c.Param("chunk_id")) + if chunkID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "chunk_id is required"}) + return + } + chunk, err := h.embedSvc.EmbedChunk(c.Request.Context(), ch, chunkID) + if err != nil { + switch { + case errors.Is(err, service.ErrEmbedChunkForbidden): + c.JSON(http.StatusForbidden, gin.H{"error": "chunk not accessible"}) + case errors.Is(err, service.ErrEmbedChunkNotFound): + c.JSON(http.StatusNotFound, gin.H{"error": "chunk not found"}) + default: + logger.Error(c.Request.Context(), "embed chunk lookup failed", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load chunk"}) + } + return + } + if chunk.Content != "" { + chunk.Content = secutils.SanitizeForDisplay(chunk.Content) + } + c.JSON(http.StatusOK, gin.H{"success": true, "data": chunk}) +} + +func (h *EmbedChannelHandler) GetEmbedSuggestedQuestions(c *gin.Context) { + ch, ok := middleware.EmbedChannelFromContext(c.Request.Context()) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return + } + if !ch.ShowSuggestedQuestions { + c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{"questions": []types.SuggestedQuestion{}}}) + return + } + limit := 6 + if raw := c.Query("limit"); raw != "" { + if n, err := strconv.Atoi(raw); err == nil && n > 0 { + limit = n + } + if limit > 12 { + limit = 12 + } + } + questions, err := h.embedSvc.SuggestedQuestions(c.Request.Context(), ch, limit) + if err != nil { + logger.Error(c.Request.Context(), "embed suggested questions failed", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load suggested questions"}) + return + } + if questions == nil { + questions = []types.SuggestedQuestion{} + } + c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{"questions": questions}}) +} + +func (h *EmbedChannelHandler) CreateEmbedSession(c *gin.Context) { + ctx := c.Request.Context() + ch, ok := middleware.EmbedChannelFromContext(ctx) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return + } + tenantID := c.GetUint64(types.TenantIDContextKey.String()) + // Leave Title empty so the first visitor message triggers the same async + // title generation as normal chat (see setupSSEStream in session/qa.go). + // Channel display name belongs on the embed page chrome, not every session row. + createdSession := &types.Session{ + TenantID: tenantID, + Title: "", + Description: service.EmbedSessionDescription(ch.ID), + } + created, err := h.sessionService.CreateSession(ctx, createdSession) + if err != nil { + logger.ErrorWithFields(ctx, err, nil) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create session"}) + return + } + // Hand back a signed handle bound to this session; the widget must echo it + // (X-Embed-Session header) on every subsequent load/chat call. + sig := service.SignEmbedSessionHandle(ch, created.ID) + c.JSON(http.StatusCreated, gin.H{"success": true, "data": gin.H{"id": created.ID, "sig": sig}}) +} + +func (h *EmbedChannelHandler) EmbedKnowledgeChat(c *gin.Context) { + h.delegateEmbedChat(c, false) +} + +func (h *EmbedChannelHandler) EmbedAgentChat(c *gin.Context) { + h.delegateEmbedChat(c, true) +} + +func (h *EmbedChannelHandler) EmbedLoadMessages(c *gin.Context) { + if err := h.ensureEmbedSession(c); err != nil { + return + } + h.messageHandler.LoadMessages(c) +} + +func (h *EmbedChannelHandler) EmbedStopSession(c *gin.Context) { + if err := h.ensureEmbedSession(c); err != nil { + return + } + h.sessionHandler.StopSession(c) +} + +type embedWebhookEventRequest struct { + Type string `json:"type"` + SessionID string `json:"session_id"` + Query string `json:"query"` + Content string `json:"content"` +} + +// EmbedRelayWebhookEvent forwards a visitor chat event to the channel webhook URL. +func (h *EmbedChannelHandler) EmbedRelayWebhookEvent(c *gin.Context) { + ch, ok := middleware.EmbedChannelFromContext(c.Request.Context()) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return + } + if err := h.ensureEmbedSession(c); err != nil { + return + } + var req embedWebhookEventRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"}) + return + } + eventType := strings.TrimSpace(req.Type) + switch eventType { + case "message_sent", "message_received": + default: + c.JSON(http.StatusBadRequest, gin.H{"error": "unsupported event type"}) + return + } + payload := map[string]any{} + if q := strings.TrimSpace(req.Query); q != "" { + payload["query"] = q + } + if content := strings.TrimSpace(req.Content); content != "" { + payload["content"] = content + } + sessionID := strings.TrimSpace(req.SessionID) + if sessionID == "" { + sessionID = secutils.SanitizeForLog(c.Param("session_id")) + } + service.DispatchEmbedWebhook(ch, eventType, sessionID, payload) + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +func (h *EmbedChannelHandler) delegateEmbedChat(c *gin.Context, agentMode bool) { + ch, ok := middleware.EmbedChannelFromContext(c.Request.Context()) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return + } + if err := h.ensureEmbedSession(c); err != nil { + return + } + patched, err := patchEmbedChatPayload(c.Request.Body, ch, agentMode) + if err != nil { + switch { + case errors.Is(err, errInvalidEmbedChatBody): + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"}) + case errors.Is(err, errInvalidEmbedChatJSON): + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json"}) + default: + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to prepare request"}) + } + return + } + c.Request.Body = io.NopCloser(bytes.NewReader(patched)) + c.Request.ContentLength = int64(len(patched)) + if agentMode && ch.AgentID != types.BuiltinQuickAnswerID { + h.sessionHandler.AgentQA(c) + return + } + h.sessionHandler.KnowledgeQA(c) +} + +func (h *EmbedChannelHandler) ensureEmbedSession(c *gin.Context) error { + ch, ok := middleware.EmbedChannelFromContext(c.Request.Context()) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return apperrors.NewUnauthorizedError("unauthorized") + } + sessionID := secutils.SanitizeForLog(c.Param("session_id")) + if sessionID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "session_id is required"}) + return apperrors.NewBadRequestError("session_id is required") + } + sess, err := h.sessionService.GetSession(c.Request.Context(), sessionID) + if err != nil || sess == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "session not found"}) + return apperrors.NewNotFoundError("session not found") + } + marker := service.EmbedSessionDescription(ch.ID) + if sess.TenantID != ch.TenantID || sess.Description != marker { + c.JSON(http.StatusForbidden, gin.H{"error": "session not allowed for this embed channel"}) + return apperrors.NewForbiddenError("session not allowed") + } + // Require the signed handle minted at creation. This is the per-visitor + // authorization secret: knowing the session id alone (e.g. from a leaked + // access log) is insufficient without the matching signature. + sig := c.GetHeader("X-Embed-Session") + if !service.VerifyEmbedSessionHandle(ch, sessionID, sig) { + c.JSON(http.StatusForbidden, gin.H{"error": "session signature invalid"}) + return apperrors.NewForbiddenError("session signature invalid") + } + return nil +} + +var ( + errInvalidEmbedChatBody = errors.New("invalid embed chat request body") + errInvalidEmbedChatJSON = errors.New("invalid embed chat json") +) + +// patchEmbedChatPayload merges embed-channel constraints into the client QA body. +func patchEmbedChatPayload(body io.Reader, ch *types.EmbedChannel, agentMode bool) ([]byte, error) { + raw, err := io.ReadAll(body) + if err != nil { + return nil, fmt.Errorf("%w: %v", errInvalidEmbedChatBody, err) + } + var payload map[string]any + if len(raw) > 0 { + if err := json.Unmarshal(raw, &payload); err != nil { + return nil, fmt.Errorf("%w: %v", errInvalidEmbedChatJSON, err) + } + } + if payload == nil { + payload = make(map[string]any) + } + payload["agent_id"] = ch.AgentID + payload["knowledge_base_ids"] = []string{} + clientWebSearch := false + if v, ok := payload["web_search_enabled"].(bool); ok { + clientWebSearch = v + } + // Channel allow_web_search only exposes the visitor toggle; the client must opt in. + payload["web_search_enabled"] = ch.AllowWebSearch && clientWebSearch + // Embed memory UI is disabled for now; always off regardless of channel flag. + payload["enable_memory"] = false + if !ch.AllowFileUpload { + delete(payload, "images") + delete(payload, "attachment_uploads") + } + payload["mcp_service_ids"] = []string{} + if agentMode { + payload["agent_enabled"] = true + } else { + payload["agent_enabled"] = false + } + patched, err := json.Marshal(payload) + if err != nil { + return nil, err + } + return patched, nil +} + +// GetEmbedChannelStats returns lightweight usage stats for an embed channel. +func (h *EmbedChannelHandler) GetEmbedChannelStats(c *gin.Context) { + channelID := strings.TrimSpace(c.Param("channel_id")) + tenantID := c.GetUint64(types.TenantIDContextKey.String()) + ctx := c.Request.Context() + + if _, err := h.embedSvc.GetOwnedChannel(ctx, tenantID, channelID); err != nil { + writeEmbedMgmtError(c, err) + return + } + + result, err := h.sessionService.ListSessions(ctx, &types.SessionListQuery{ + TenantID: tenantID, + Source: "embed:" + channelID, + Page: 1, + PageSize: 1, + }) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + total := int64(0) + if result != nil { + total = result.Total + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "data": gin.H{ + "session_count": total, + }, + }) +} + +func embedChannelResponse(ch *types.EmbedChannel, publishToken string) gin.H { + row := gin.H{ + "id": ch.ID, + "tenant_id": ch.TenantID, + "agent_id": ch.AgentID, + "name": ch.Name, + "enabled": ch.Enabled, + "allowed_origins": ch.AllowedOriginsList(), + "welcome_message": ch.WelcomeMessage, + "rate_limit_per_minute": ch.RateLimitPerMinute, + "rate_limit_per_day": ch.RateLimitPerDay, + "primary_color": ch.PrimaryColor, + "page_title": ch.PageTitle, + "header_title_mode": types.NormalizeEmbedHeaderTitleMode(ch.HeaderTitleMode), + "show_suggested_questions": ch.ShowSuggestedQuestions, + "widget_position": ch.WidgetPosition, + "allow_web_search": ch.AllowWebSearch, + "allow_memory": ch.AllowMemory, + "allow_file_upload": ch.AllowFileUpload, + "default_locale": ch.DefaultLocale, + "webhook_url": ch.WebhookURL, + "has_webhook_secret": ch.WebhookSecret != "", + "created_at": ch.CreatedAt, + "updated_at": ch.UpdatedAt, + } + if publishToken != "" { + row["publish_token"] = publishToken + } + return row +} + +func writeEmbedMgmtError(c *gin.Context, err error) { + switch { + case errors.Is(err, service.ErrEmbedChannelNotFound): + c.JSON(http.StatusNotFound, gin.H{"error": "embed channel not found"}) + case errors.Is(err, service.ErrEmbedWebhookURLInvalid): + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + default: + var appErr *apperrors.AppError + if errors.As(err, &appErr) && appErr.Code == apperrors.ErrNotFound { + c.JSON(http.StatusNotFound, gin.H{"error": appErr.Message}) + return + } + logger.Error(c.Request.Context(), "embed channel management failed", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "operation failed"}) + } +} diff --git a/internal/handler/embed_channel_exchange_test.go b/internal/handler/embed_channel_exchange_test.go new file mode 100644 index 00000000..ca7d28ff --- /dev/null +++ b/internal/handler/embed_channel_exchange_test.go @@ -0,0 +1,152 @@ +package handler + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Tencent/WeKnora/internal/application/service" + "github.com/Tencent/WeKnora/internal/middleware" + "github.com/Tencent/WeKnora/internal/types" + "github.com/gin-gonic/gin" +) + +type exchangeEmbedSvc struct { + sessionToken string + expiresIn int + err error +} + +func (f *exchangeEmbedSvc) Create(context.Context, uint64, string, *types.EmbedChannel) (*types.EmbedChannel, string, error) { + return nil, "", nil +} +func (f *exchangeEmbedSvc) ListByAgent(context.Context, uint64, string) ([]*types.EmbedChannel, error) { + return nil, nil +} +func (f *exchangeEmbedSvc) ListByTenant(context.Context, uint64) ([]*types.EmbedChannel, error) { + return nil, nil +} +func (f *exchangeEmbedSvc) Update(context.Context, uint64, string, *types.EmbedChannel, *bool, *bool, *bool, *bool, *bool, *string, *string, *string) (*types.EmbedChannel, error) { + return nil, nil +} +func (f *exchangeEmbedSvc) GetOwnedChannel(context.Context, uint64, string) (*types.EmbedChannel, error) { + return nil, service.ErrEmbedChannelNotFound +} +func (f *exchangeEmbedSvc) Delete(context.Context, uint64, string) error { return nil } +func (f *exchangeEmbedSvc) RotateToken(context.Context, uint64, string) (*types.EmbedChannel, string, error) { + return nil, "", nil +} +func (f *exchangeEmbedSvc) LookupForEmbed(context.Context, string, string) (*types.EmbedChannel, error) { + return nil, nil +} +func (f *exchangeEmbedSvc) LookupEnabledChannel(context.Context, string) (*types.EmbedChannel, error) { + return nil, nil +} +func (f *exchangeEmbedSvc) IssueSessionToken(context.Context, string) (string, int, error) { + if f.err != nil { + return "", 0, f.err + } + return f.sessionToken, f.expiresIn, nil +} +func (f *exchangeEmbedSvc) IssuePreviewSession(context.Context, uint64, string) (string, int, error) { + return f.IssueSessionToken(context.Background(), "") +} +func (f *exchangeEmbedSvc) ResolveSessionToken(context.Context, string) (string, error) { + return "", nil +} +func (f *exchangeEmbedSvc) PublicConfig(context.Context, *types.EmbedChannel) types.EmbedChannelPublicConfig { + return types.EmbedChannelPublicConfig{} +} +func (f *exchangeEmbedSvc) SuggestedQuestions(context.Context, *types.EmbedChannel, int) ([]types.SuggestedQuestion, error) { + return nil, nil +} +func (f *exchangeEmbedSvc) EmbedChunk(context.Context, *types.EmbedChannel, string) (*types.Chunk, error) { + return nil, nil +} +func (f *exchangeEmbedSvc) EmbedDisplayTitle(context.Context, *types.EmbedChannel) string { + return "AI Assistant" +} + +func TestExchangeEmbedSessionSuccess(t *testing.T) { + gin.SetMode(gin.TestMode) + h := &EmbedChannelHandler{embedSvc: &exchangeEmbedSvc{ + sessionToken: "ems_test_token", + expiresIn: 1800, + }} + + r := gin.New() + r.POST("/exchange", func(c *gin.Context) { + ch := &types.EmbedChannel{ID: "channel-1", Enabled: true} + ctx := context.WithValue(c.Request.Context(), middleware.EmbedChannelContextKey, ch) + c.Request = c.Request.WithContext(ctx) + h.ExchangeEmbedSession(c) + }) + + req := httptest.NewRequest(http.MethodPost, "/exchange", nil) + req.Header.Set("Authorization", "Embed em_publish_token") + w := httptest.NewRecorder() + r.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } + var resp struct { + Success bool `json:"success"` + Data struct { + SessionToken string `json:"session_token"` + ExpiresIn int `json:"expires_in"` + } `json:"data"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if !resp.Success || resp.Data.SessionToken != "ems_test_token" || resp.Data.ExpiresIn != 1800 { + t.Fatalf("unexpected response: %#v", resp) + } +} + +func TestExchangeEmbedSessionUnavailable(t *testing.T) { + gin.SetMode(gin.TestMode) + h := &EmbedChannelHandler{embedSvc: &exchangeEmbedSvc{err: service.ErrEmbedSessionUnavailable}} + + r := gin.New() + r.POST("/exchange", func(c *gin.Context) { + ch := &types.EmbedChannel{ID: "channel-1", Enabled: true} + ctx := context.WithValue(c.Request.Context(), middleware.EmbedChannelContextKey, ch) + c.Request = c.Request.WithContext(ctx) + h.ExchangeEmbedSession(c) + }) + + req := httptest.NewRequest(http.MethodPost, "/exchange", nil) + req.Header.Set("Authorization", "Embed em_publish_token") + w := httptest.NewRecorder() + r.ServeHTTP(w, req) + + if w.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want 503", w.Code) + } +} + +func TestExchangeEmbedSessionRejectsSessionToken(t *testing.T) { + gin.SetMode(gin.TestMode) + h := &EmbedChannelHandler{embedSvc: &exchangeEmbedSvc{sessionToken: "ems_new", expiresIn: 1800}} + + r := gin.New() + r.POST("/exchange", func(c *gin.Context) { + ch := &types.EmbedChannel{ID: "channel-1", Enabled: true} + ctx := context.WithValue(c.Request.Context(), middleware.EmbedChannelContextKey, ch) + c.Request = c.Request.WithContext(ctx) + h.ExchangeEmbedSession(c) + }) + + req := httptest.NewRequest(http.MethodPost, "/exchange", nil) + req.Header.Set("Authorization", "Embed ems_existing_session") + w := httptest.NewRecorder() + r.ServeHTTP(w, req) + + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want 403", w.Code) + } +} diff --git a/internal/handler/embed_flow_test.go b/internal/handler/embed_flow_test.go new file mode 100644 index 00000000..abf5b091 --- /dev/null +++ b/internal/handler/embed_flow_test.go @@ -0,0 +1,305 @@ +package handler + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Tencent/WeKnora/internal/application/service" + "github.com/Tencent/WeKnora/internal/middleware" + "github.com/Tencent/WeKnora/internal/types" + "github.com/gin-gonic/gin" +) + +type flowEmbedSvc struct { + sessionToken string + expiresIn int + issueErr error + channels map[string]*types.EmbedChannel +} + +func (f *flowEmbedSvc) Create(context.Context, uint64, string, *types.EmbedChannel) (*types.EmbedChannel, string, error) { + return nil, "", nil +} +func (f *flowEmbedSvc) ListByAgent(context.Context, uint64, string) ([]*types.EmbedChannel, error) { + return nil, nil +} +func (f *flowEmbedSvc) ListByTenant(context.Context, uint64) ([]*types.EmbedChannel, error) { + return nil, nil +} +func (f *flowEmbedSvc) Update(context.Context, uint64, string, *types.EmbedChannel, *bool, *bool, *bool, *bool, *bool, *string, *string, *string) (*types.EmbedChannel, error) { + return nil, nil +} +func (f *flowEmbedSvc) GetOwnedChannel(_ context.Context, tenantID uint64, id string) (*types.EmbedChannel, error) { + ch := f.channels[id] + if ch == nil || ch.TenantID != tenantID { + return nil, service.ErrEmbedChannelNotFound + } + return ch, nil +} +func (f *flowEmbedSvc) Delete(context.Context, uint64, string) error { return nil } +func (f *flowEmbedSvc) RotateToken(context.Context, uint64, string) (*types.EmbedChannel, string, error) { + return nil, "", nil +} +func (f *flowEmbedSvc) LookupForEmbed(_ context.Context, channelID, token string) (*types.EmbedChannel, error) { + ch := f.channels[channelID] + if ch == nil || ch.PublishToken != token { + return nil, service.ErrEmbedTokenInvalid + } + if !ch.Enabled { + return nil, service.ErrEmbedChannelDisabled + } + return ch, nil +} +func (f *flowEmbedSvc) LookupEnabledChannel(context.Context, string) (*types.EmbedChannel, error) { + return nil, nil +} +func (f *flowEmbedSvc) IssueSessionToken(context.Context, string) (string, int, error) { + if f.issueErr != nil { + return "", 0, f.issueErr + } + return f.sessionToken, f.expiresIn, nil +} +func (f *flowEmbedSvc) IssuePreviewSession(context.Context, uint64, string) (string, int, error) { + return f.IssueSessionToken(context.Background(), "") +} +func (f *flowEmbedSvc) ResolveSessionToken(context.Context, string) (string, error) { + return "", nil +} +func (f *flowEmbedSvc) PublicConfig(context.Context, *types.EmbedChannel) types.EmbedChannelPublicConfig { + return types.EmbedChannelPublicConfig{} +} +func (f *flowEmbedSvc) SuggestedQuestions(context.Context, *types.EmbedChannel, int) ([]types.SuggestedQuestion, error) { + return nil, nil +} +func (f *flowEmbedSvc) EmbedChunk(context.Context, *types.EmbedChannel, string) (*types.Chunk, error) { + return nil, nil +} +func (f *flowEmbedSvc) EmbedDisplayTitle(context.Context, *types.EmbedChannel) string { + return "AI Assistant" +} + +type flowTenantSvc struct { + tenant *types.Tenant +} + +func (f *flowTenantSvc) GetTenantByID(context.Context, uint64) (*types.Tenant, error) { + return f.tenant, nil +} +func (f *flowTenantSvc) CreateTenant(context.Context, *types.Tenant) (*types.Tenant, error) { + return nil, nil +} +func (f *flowTenantSvc) GetTenantsByIDs(context.Context, []uint64) (map[uint64]*types.Tenant, error) { + return nil, nil +} +func (f *flowTenantSvc) UpdateTenant(context.Context, *types.Tenant) (*types.Tenant, error) { + return nil, nil +} +func (f *flowTenantSvc) DeleteTenant(context.Context, uint64) error { return nil } +func (f *flowTenantSvc) ListTenants(context.Context) ([]*types.Tenant, error) { + return nil, nil +} +func (f *flowTenantSvc) UpdateAPIKey(context.Context, uint64) (string, error) { + return "", nil +} +func (f *flowTenantSvc) ExtractTenantIDFromAPIKey(string) (uint64, error) { return 0, nil } +func (f *flowTenantSvc) ListAllTenants(context.Context) ([]*types.Tenant, error) { + return nil, nil +} +func (f *flowTenantSvc) BulkSetStorageQuota(context.Context, int64) (int64, error) { + return 0, nil +} +func (f *flowTenantSvc) SearchTenants(context.Context, string, uint64, int, int) ([]*types.Tenant, int64, error) { + return nil, 0, nil +} +func (f *flowTenantSvc) GetTenantByIDForUser(context.Context, uint64, string) (*types.Tenant, error) { + return f.tenant, nil +} +func (f *flowTenantSvc) GetWeKnoraCloudCredentials(context.Context) *types.WeKnoraCloudCredentials { + return nil +} + +func TestEmbedExchangeFlowIntegration(t *testing.T) { + gin.SetMode(gin.TestMode) + + const ( + channelID = "ch-flow-1" + publishToken = "em_publish_valid" + ) + svc := &flowEmbedSvc{ + sessionToken: "ems_integration_token", + expiresIn: 1800, + channels: map[string]*types.EmbedChannel{ + channelID: { + ID: channelID, + TenantID: 7, + AgentID: "agent-flow-1", + Enabled: true, + PublishToken: publishToken, + AllowedOrigins: []byte(`["https://partner.example.com"]`), + RateLimitPerMinute: 0, + }, + }, + } + h := &EmbedChannelHandler{embedSvc: svc} + tenantSvc := &flowTenantSvc{tenant: &types.Tenant{ID: 7}} + + r := gin.New() + r.POST( + "/api/v1/embed/:channel_id/exchange", + middleware.EmbedAuth(svc, tenantSvc, nil), + h.ExchangeEmbedSession, + ) + + req := httptest.NewRequest(http.MethodPost, "/api/v1/embed/"+channelID+"/exchange", nil) + req.Header.Set("Authorization", "Embed "+publishToken) + req.Header.Set("Origin", "https://partner.example.com") + w := httptest.NewRecorder() + r.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } + var resp struct { + Success bool `json:"success"` + Data struct { + SessionToken string `json:"session_token"` + ExpiresIn int `json:"expires_in"` + } `json:"data"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if !resp.Success { + t.Fatalf("expected success, got %#v", resp) + } + if !strings.HasPrefix(resp.Data.SessionToken, "ems_") { + t.Fatalf("session_token = %q, want ems_ prefix", resp.Data.SessionToken) + } + if resp.Data.SessionToken != "ems_integration_token" || resp.Data.ExpiresIn != 1800 { + t.Fatalf("unexpected exchange payload: %#v", resp.Data) + } +} + +func TestPatchEmbedChatPayloadInjectsAgentID(t *testing.T) { + ch := &types.EmbedChannel{AgentID: "agent-embed-42"} + body := `{"query":"hello","agent_id":"client-override","web_search_enabled":true}` + + patched, err := patchEmbedChatPayload(strings.NewReader(body), ch, false) + if err != nil { + t.Fatal(err) + } + var payload map[string]any + if err := json.Unmarshal(patched, &payload); err != nil { + t.Fatal(err) + } + if payload["agent_id"] != "agent-embed-42" { + t.Fatalf("agent_id = %v, want channel agent", payload["agent_id"]) + } + if payload["query"] != "hello" { + t.Fatalf("query = %v, want preserved client field", payload["query"]) + } + if payload["web_search_enabled"] != false { + t.Fatalf("web_search_enabled = %v, want false", payload["web_search_enabled"]) + } + if payload["agent_enabled"] != false { + t.Fatalf("agent_enabled = %v, want false for knowledge mode", payload["agent_enabled"]) + } + kbIDs, ok := payload["knowledge_base_ids"].([]any) + if !ok || len(kbIDs) != 0 { + t.Fatalf("knowledge_base_ids = %v, want empty slice", payload["knowledge_base_ids"]) + } +} + +func TestPatchEmbedChatPayloadWebSearchRequiresClientOptIn(t *testing.T) { + ch := &types.EmbedChannel{AgentID: "agent-1", AllowWebSearch: true} + body := `{"query":"hello","web_search_enabled":false,"enable_memory":true}` + + patched, err := patchEmbedChatPayload(strings.NewReader(body), ch, false) + if err != nil { + t.Fatal(err) + } + var payload map[string]any + if err := json.Unmarshal(patched, &payload); err != nil { + t.Fatal(err) + } + if payload["web_search_enabled"] != false { + t.Fatalf("web_search_enabled = %v, want false when visitor did not opt in", payload["web_search_enabled"]) + } + if payload["enable_memory"] != false { + t.Fatalf("enable_memory = %v, want false (embed memory is disabled)", payload["enable_memory"]) + } + + bodyOn := `{"query":"hello","web_search_enabled":true}` + patchedOn, err := patchEmbedChatPayload(strings.NewReader(bodyOn), ch, false) + if err != nil { + t.Fatal(err) + } + var payloadOn map[string]any + if err := json.Unmarshal(patchedOn, &payloadOn); err != nil { + t.Fatal(err) + } + if payloadOn["web_search_enabled"] != true { + t.Fatalf("web_search_enabled = %v, want true when channel allows and visitor opted in", payloadOn["web_search_enabled"]) + } +} + +func TestPatchEmbedChatPayloadWebSearchBlockedWhenChannelDisabled(t *testing.T) { + ch := &types.EmbedChannel{AgentID: "agent-1", AllowWebSearch: false} + body := `{"query":"hello","web_search_enabled":true}` + + patched, err := patchEmbedChatPayload(strings.NewReader(body), ch, false) + if err != nil { + t.Fatal(err) + } + var payload map[string]any + if err := json.Unmarshal(patched, &payload); err != nil { + t.Fatal(err) + } + if payload["web_search_enabled"] != false { + t.Fatalf("web_search_enabled = %v, want false when channel disallows web search", payload["web_search_enabled"]) + } +} + +func TestPatchEmbedChatPayloadAgentMode(t *testing.T) { + ch := &types.EmbedChannel{AgentID: "agent-embed-99"} + patched, err := patchEmbedChatPayload(bytes.NewReader(nil), ch, true) + if err != nil { + t.Fatal(err) + } + var payload map[string]any + if err := json.Unmarshal(patched, &payload); err != nil { + t.Fatal(err) + } + if payload["agent_id"] != "agent-embed-99" { + t.Fatalf("agent_id = %v", payload["agent_id"]) + } + if payload["agent_enabled"] != true { + t.Fatalf("agent_enabled = %v, want true", payload["agent_enabled"]) + } +} + +func TestPatchEmbedChatPayloadInvalidJSON(t *testing.T) { + ch := &types.EmbedChannel{AgentID: "agent-1"} + _, err := patchEmbedChatPayload(strings.NewReader("{not-json"), ch, false) + if err == nil || !strings.Contains(err.Error(), "invalid embed chat json") { + t.Fatalf("expected invalid json error, got %v", err) + } +} + +func TestPatchEmbedChatPayloadInvalidBody(t *testing.T) { + ch := &types.EmbedChannel{AgentID: "agent-1"} + _, err := patchEmbedChatPayload(badReader{}, ch, false) + if err == nil || !strings.Contains(err.Error(), "invalid embed chat request body") { + t.Fatalf("expected invalid body error, got %v", err) + } +} + +type badReader struct{} + +func (badReader) Read([]byte) (int, error) { return 0, io.ErrUnexpectedEOF } diff --git a/internal/handler/embed_session_test.go b/internal/handler/embed_session_test.go new file mode 100644 index 00000000..951aa728 --- /dev/null +++ b/internal/handler/embed_session_test.go @@ -0,0 +1,369 @@ +package handler + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Tencent/WeKnora/internal/application/service" + "github.com/Tencent/WeKnora/internal/handler/session" + "github.com/Tencent/WeKnora/internal/middleware" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/gin-gonic/gin" +) + +const ( + testEmbedChannelID = "ch-session-test" + testEmbedTenantID = uint64(42) + testEmbedSessionID = "11111111-2222-3333-4444-555555555555" + testEmbedPublishTok = "em_publish_session_test" +) + +type stubSessionServiceForEmbed struct { + interfaces.SessionService + sessions map[string]*types.Session + created *types.Session +} + +func (s *stubSessionServiceForEmbed) GetSession(_ context.Context, id string) (*types.Session, error) { + sess, ok := s.sessions[id] + if !ok { + return nil, errors.New("session not found") + } + return sess, nil +} + +func (s *stubSessionServiceForEmbed) CreateSession(_ context.Context, session *types.Session) (*types.Session, error) { + created := *session + if created.ID == "" { + created.ID = testEmbedSessionID + } + s.created = &created + return &created, nil +} + +type sessionEmbedSvc struct { + interfaces.EmbedChannelService + embedChunk func(context.Context, *types.EmbedChannel, string) (*types.Chunk, error) +} + +func (s *sessionEmbedSvc) Create(context.Context, uint64, string, *types.EmbedChannel) (*types.EmbedChannel, string, error) { + return nil, "", nil +} +func (s *sessionEmbedSvc) ListByAgent(context.Context, uint64, string) ([]*types.EmbedChannel, error) { + return nil, nil +} +func (s *sessionEmbedSvc) ListByTenant(context.Context, uint64) ([]*types.EmbedChannel, error) { + return nil, nil +} +func (s *sessionEmbedSvc) Update(context.Context, uint64, string, *types.EmbedChannel, *bool, *bool, *bool, *bool, *bool, *string, *string, *string) (*types.EmbedChannel, error) { + return nil, nil +} +func (s *sessionEmbedSvc) GetOwnedChannel(context.Context, uint64, string) (*types.EmbedChannel, error) { + return nil, service.ErrEmbedChannelNotFound +} +func (s *sessionEmbedSvc) Delete(context.Context, uint64, string) error { return nil } +func (s *sessionEmbedSvc) RotateToken(context.Context, uint64, string) (*types.EmbedChannel, string, error) { + return nil, "", nil +} +func (s *sessionEmbedSvc) LookupForEmbed(context.Context, string, string) (*types.EmbedChannel, error) { + return nil, nil +} +func (s *sessionEmbedSvc) LookupEnabledChannel(context.Context, string) (*types.EmbedChannel, error) { + return nil, nil +} +func (s *sessionEmbedSvc) IssueSessionToken(context.Context, string) (string, int, error) { + return "", 0, nil +} +func (s *sessionEmbedSvc) IssuePreviewSession(context.Context, uint64, string) (string, int, error) { + return "", 0, nil +} +func (s *sessionEmbedSvc) ResolveSessionToken(context.Context, string) (string, error) { + return "", nil +} +func (s *sessionEmbedSvc) PublicConfig(context.Context, *types.EmbedChannel) types.EmbedChannelPublicConfig { + return types.EmbedChannelPublicConfig{} +} +func (s *sessionEmbedSvc) SuggestedQuestions(context.Context, *types.EmbedChannel, int) ([]types.SuggestedQuestion, error) { + return nil, nil +} +func (s *sessionEmbedSvc) EmbedChunk(ctx context.Context, ch *types.EmbedChannel, chunkID string) (*types.Chunk, error) { + if s.embedChunk != nil { + return s.embedChunk(ctx, ch, chunkID) + } + return nil, nil +} +func (s *sessionEmbedSvc) EmbedDisplayTitle(context.Context, *types.EmbedChannel) string { + return "" +} + +func testEmbedChannel() *types.EmbedChannel { + return &types.EmbedChannel{ + ID: testEmbedChannelID, + TenantID: testEmbedTenantID, + PublishToken: testEmbedPublishTok, + Enabled: true, + } +} + +func validEmbedSession(ch *types.EmbedChannel) *types.Session { + return &types.Session{ + ID: testEmbedSessionID, + TenantID: ch.TenantID, + Description: service.EmbedSessionDescription(ch.ID), + } +} + +func newEnsureEmbedSessionCtx(ch *types.EmbedChannel, sessionID, sig string) (*gin.Context, *httptest.ResponseRecorder) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/sessions/"+sessionID+"/messages", nil) + if sig != "" { + c.Request.Header.Set("X-Embed-Session", sig) + } + c.Params = gin.Params{{Key: "session_id", Value: sessionID}} + ctx := context.WithValue(c.Request.Context(), middleware.EmbedChannelContextKey, ch) + c.Request = c.Request.WithContext(ctx) + return c, w +} + +func TestEnsureEmbedSessionValid(t *testing.T) { + ch := testEmbedChannel() + sig := service.SignEmbedSessionHandle(ch, testEmbedSessionID) + h := &EmbedChannelHandler{ + sessionService: &stubSessionServiceForEmbed{ + sessions: map[string]*types.Session{ + testEmbedSessionID: validEmbedSession(ch), + }, + }, + } + c, w := newEnsureEmbedSessionCtx(ch, testEmbedSessionID, sig) + if err := h.ensureEmbedSession(c); err != nil { + t.Fatalf("ensureEmbedSession() = %v, want nil", err) + } + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want no response written on success", w.Code) + } +} + +func TestEnsureEmbedSessionWrongTenant(t *testing.T) { + ch := testEmbedChannel() + sig := service.SignEmbedSessionHandle(ch, testEmbedSessionID) + h := &EmbedChannelHandler{ + sessionService: &stubSessionServiceForEmbed{ + sessions: map[string]*types.Session{ + testEmbedSessionID: { + ID: testEmbedSessionID, + TenantID: 999, + Description: service.EmbedSessionDescription(ch.ID), + }, + }, + }, + } + c, w := newEnsureEmbedSessionCtx(ch, testEmbedSessionID, sig) + if err := h.ensureEmbedSession(c); err == nil { + t.Fatal("expected error for cross-tenant session") + } + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } +} + +func TestEnsureEmbedSessionWrongDescription(t *testing.T) { + ch := testEmbedChannel() + sig := service.SignEmbedSessionHandle(ch, testEmbedSessionID) + h := &EmbedChannelHandler{ + sessionService: &stubSessionServiceForEmbed{ + sessions: map[string]*types.Session{ + testEmbedSessionID: { + ID: testEmbedSessionID, + TenantID: ch.TenantID, + Description: "not-an-embed-session", + }, + }, + }, + } + c, w := newEnsureEmbedSessionCtx(ch, testEmbedSessionID, sig) + if err := h.ensureEmbedSession(c); err == nil { + t.Fatal("expected error for wrong description marker") + } + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } +} + +func TestEnsureEmbedSessionInvalidSig(t *testing.T) { + ch := testEmbedChannel() + h := &EmbedChannelHandler{ + sessionService: &stubSessionServiceForEmbed{ + sessions: map[string]*types.Session{ + testEmbedSessionID: validEmbedSession(ch), + }, + }, + } + c, w := newEnsureEmbedSessionCtx(ch, testEmbedSessionID, "invalid-signature") + if err := h.ensureEmbedSession(c); err == nil { + t.Fatal("expected error for invalid signature") + } + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } +} + +func TestEnsureEmbedSessionNotFound(t *testing.T) { + ch := testEmbedChannel() + h := &EmbedChannelHandler{sessionService: &stubSessionServiceForEmbed{sessions: map[string]*types.Session{}}} + c, w := newEnsureEmbedSessionCtx(ch, testEmbedSessionID, "anything") + if err := h.ensureEmbedSession(c); err == nil { + t.Fatal("expected error for missing session") + } + if w.Code != http.StatusNotFound { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } +} + +func TestCreateEmbedSessionSuccess(t *testing.T) { + gin.SetMode(gin.TestMode) + ch := testEmbedChannel() + stub := &stubSessionServiceForEmbed{sessions: map[string]*types.Session{}} + h := &EmbedChannelHandler{sessionService: stub} + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/sessions", nil) + c.Set(types.TenantIDContextKey.String(), ch.TenantID) + ctx := context.WithValue(c.Request.Context(), middleware.EmbedChannelContextKey, ch) + c.Request = c.Request.WithContext(ctx) + + h.CreateEmbedSession(c) + + if w.Code != http.StatusCreated { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } + var resp struct { + Success bool `json:"success"` + Data struct { + ID string `json:"id"` + Sig string `json:"sig"` + } `json:"data"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if !resp.Success || resp.Data.ID != testEmbedSessionID { + t.Fatalf("unexpected response: %#v", resp) + } + if stub.created == nil || stub.created.Description != service.EmbedSessionDescription(ch.ID) { + t.Fatalf("created session marker = %q", stub.created.Description) + } + if !service.VerifyEmbedSessionHandle(ch, resp.Data.ID, resp.Data.Sig) { + t.Fatal("returned sig must verify against created session") + } +} + +func TestGetEmbedChunkForbidden(t *testing.T) { + gin.SetMode(gin.TestMode) + ch := testEmbedChannel() + h := &EmbedChannelHandler{ + embedSvc: &sessionEmbedSvc{ + embedChunk: func(_ context.Context, _ *types.EmbedChannel, _ string) (*types.Chunk, error) { + return nil, service.ErrEmbedChunkForbidden + }, + }, + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/chunks/chunk-1", nil) + c.Params = gin.Params{{Key: "chunk_id", Value: "chunk-1"}} + ctx := context.WithValue(c.Request.Context(), middleware.EmbedChannelContextKey, ch) + c.Request = c.Request.WithContext(ctx) + + h.GetEmbedChunk(c) + + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } +} + +func TestGetEmbedChunkSuccess(t *testing.T) { + gin.SetMode(gin.TestMode) + ch := testEmbedChannel() + h := &EmbedChannelHandler{ + embedSvc: &sessionEmbedSvc{ + embedChunk: func(_ context.Context, _ *types.EmbedChannel, chunkID string) (*types.Chunk, error) { + return &types.Chunk{ID: chunkID, Content: "hello"}, nil + }, + }, + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/chunks/chunk-ok", nil) + c.Params = gin.Params{{Key: "chunk_id", Value: "chunk-ok"}} + ctx := context.WithValue(c.Request.Context(), middleware.EmbedChannelContextKey, ch) + c.Request = c.Request.WithContext(ctx) + + h.GetEmbedChunk(c) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } +} + +func newEmbedStopSessionCtx(ch *types.EmbedChannel, sessionID, sig, body string) (*gin.Context, *httptest.ResponseRecorder) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/sessions/"+sessionID+"/stop", strings.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + if sig != "" { + c.Request.Header.Set("X-Embed-Session", sig) + } + c.Params = gin.Params{{Key: "session_id", Value: sessionID}} + c.Set(types.TenantIDContextKey.String(), ch.TenantID) + ctx := context.WithValue(c.Request.Context(), middleware.EmbedChannelContextKey, ch) + c.Request = c.Request.WithContext(ctx) + return c, w +} + +func TestEmbedStopSessionInvalidSig(t *testing.T) { + ch := testEmbedChannel() + h := &EmbedChannelHandler{ + sessionService: &stubSessionServiceForEmbed{ + sessions: map[string]*types.Session{ + testEmbedSessionID: validEmbedSession(ch), + }, + }, + sessionHandler: &session.Handler{}, + } + c, w := newEmbedStopSessionCtx(ch, testEmbedSessionID, "bad-sig", `{"message_id":"msg-1"}`) + h.EmbedStopSession(c) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } +} + +func TestEmbedStopSessionMissingMessageID(t *testing.T) { + ch := testEmbedChannel() + sig := service.SignEmbedSessionHandle(ch, testEmbedSessionID) + h := &EmbedChannelHandler{ + sessionService: &stubSessionServiceForEmbed{ + sessions: map[string]*types.Session{ + testEmbedSessionID: validEmbedSession(ch), + }, + }, + sessionHandler: &session.Handler{}, + } + c, w := newEmbedStopSessionCtx(ch, testEmbedSessionID, sig, `{}`) + h.EmbedStopSession(c) + if w.Code != http.StatusBadRequest { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } +} diff --git a/internal/im/ratelimit.go b/internal/im/ratelimit.go deleted file mode 100644 index 6f1d891d..00000000 --- a/internal/im/ratelimit.go +++ /dev/null @@ -1,199 +0,0 @@ -package im - -import ( - "context" - "fmt" - "sync" - "time" - - "github.com/redis/go-redis/v9" -) - -const ( - // rateLimitWindow is the sliding window duration for rate limiting. - rateLimitWindow = 60 * time.Second - // rateLimitMaxRequests is the maximum number of requests allowed per window per key. - rateLimitMaxRequests = 10 - // rateLimitCleanupInterval is how often stale entries are purged. - rateLimitCleanupInterval = 1 * time.Minute -) - -// ────────────────────────────────────────────────────────────────────────────── -// distributedLimiter: Redis ZSET + local fallback -// ────────────────────────────────────────────────────────────────────────────── - -// rateLimitScript is an atomic Lua script that implements a sliding-window rate -// limiter on a Redis Sorted Set. It prunes expired entries, checks the count, -// and conditionally adds a new member — all in a single round-trip. -// -// KEYS[1] = the rate-limit key -// ARGV[1] = now (Unix milliseconds) -// ARGV[2] = window size (milliseconds) -// ARGV[3] = max allowed requests -// ARGV[4] = unique member value (e.g. now_ms as string) -// -// Returns 1 if the request is allowed, 0 if rate-limited. -var rateLimitScript = redis.NewScript(` -local key = KEYS[1] -local now = tonumber(ARGV[1]) -local window = tonumber(ARGV[2]) -local maxReq = tonumber(ARGV[3]) -local member = ARGV[4] - -redis.call('ZREMRANGEBYSCORE', key, 0, now - window) -local count = redis.call('ZCARD', key) -if count < maxReq then - redis.call('ZADD', key, now, member) - redis.call('PEXPIRE', key, window + 1000) - return 1 -end -return 0 -`) - -// distributedLimiter tries Redis first, falls back to a local sliding-window -// limiter when Redis is unavailable (nil client or transient error). -type distributedLimiter struct { - redisClient *redis.Client - local *slidingWindowLimiter - window time.Duration - maxRequests int - instanceID string // used to disambiguate ZSET members across instances -} - -func newDistributedLimiter(redisClient *redis.Client, window time.Duration, maxRequests int, instanceID string) *distributedLimiter { - return &distributedLimiter{ - redisClient: redisClient, - local: newSlidingWindowLimiter(window, maxRequests), - window: window, - maxRequests: maxRequests, - instanceID: instanceID, - } -} - -// Allow returns true if the request for the given key is within the rate limit. -func (d *distributedLimiter) Allow(key string) bool { - if d.redisClient != nil { - allowed, err := d.redisAllow(context.Background(), key) - if err == nil { - return allowed - } - // Redis failed — fall through to local limiter. - } - return d.local.Allow(key) -} - -func (d *distributedLimiter) redisAllow(ctx context.Context, key string) (bool, error) { - redisKey := RedisKeyRateLimit + key - nowMs := time.Now().UnixMilli() - windowMs := d.window.Milliseconds() - member := fmt.Sprintf("%s:%d", d.instanceID, nowMs) // instanceID prevents ZSET member collision across instances - - result, err := rateLimitScript.Run(ctx, d.redisClient, - []string{redisKey}, - nowMs, windowMs, d.maxRequests, member, - ).Int64() - if err != nil { - return false, err - } - return result == 1, nil -} - -// cleanupLoop delegates to the local limiter's cleanup for the fallback path. -func (d *distributedLimiter) cleanupLoop(stopCh <-chan struct{}) { - d.local.cleanupLoop(stopCh) -} - -// ────────────────────────────────────────────────────────────────────────────── -// slidingWindowLimiter: local in-memory fallback (original implementation) -// ────────────────────────────────────────────────────────────────────────────── - -// rateLimitEntry holds the request timestamps for a single key. -type rateLimitEntry struct { - mu sync.Mutex - timestamps []time.Time - deleted bool // marked true when removed from the map by cleanupLoop -} - -// slidingWindowLimiter implements per-key sliding window rate limiting. -type slidingWindowLimiter struct { - window time.Duration - maxRequests int - entries sync.Map // key -> *rateLimitEntry -} - -func newSlidingWindowLimiter(window time.Duration, maxRequests int) *slidingWindowLimiter { - return &slidingWindowLimiter{ - window: window, - maxRequests: maxRequests, - } -} - -// Allow checks if the request for the given key is within the rate limit. -// Returns true if allowed, false if rate limited. -func (l *slidingWindowLimiter) Allow(key string) bool { - now := time.Now() - cutoff := now.Add(-l.window) - - for { - val, _ := l.entries.LoadOrStore(key, &rateLimitEntry{}) - entry := val.(*rateLimitEntry) - - entry.mu.Lock() - // If the entry was concurrently deleted by cleanupLoop, retry with a fresh one. - if entry.deleted { - entry.mu.Unlock() - l.entries.Delete(key) // ensure stale entry is gone - continue - } - - // Remove expired timestamps - valid := entry.timestamps[:0] - for _, t := range entry.timestamps { - if t.After(cutoff) { - valid = append(valid, t) - } - } - entry.timestamps = valid - - if len(entry.timestamps) >= l.maxRequests { - entry.mu.Unlock() - return false - } - - entry.timestamps = append(entry.timestamps, now) - entry.mu.Unlock() - return true - } -} - -// cleanupLoop periodically removes stale entries from the limiter. -func (l *slidingWindowLimiter) cleanupLoop(stopCh <-chan struct{}) { - ticker := time.NewTicker(rateLimitCleanupInterval) - defer ticker.Stop() - - for { - select { - case <-ticker.C: - cutoff := time.Now().Add(-l.window) - l.entries.Range(func(key, val interface{}) bool { - entry := val.(*rateLimitEntry) - entry.mu.Lock() - allExpired := true - for _, t := range entry.timestamps { - if t.After(cutoff) { - allExpired = false - break - } - } - if allExpired { - entry.deleted = true - l.entries.Delete(key) - } - entry.mu.Unlock() - return true - }) - case <-stopCh: - return - } - } -} diff --git a/internal/im/service.go b/internal/im/service.go index 3f74f2af..ed03e454 100644 --- a/internal/im/service.go +++ b/internal/im/service.go @@ -20,6 +20,7 @@ import ( "github.com/Tencent/WeKnora/internal/event" "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/models/chat" + "github.com/Tencent/WeKnora/internal/ratelimit" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" "github.com/google/uuid" @@ -305,6 +306,9 @@ const ( RedisKeyQueueUser = "im:queue:user:" // + userKey — global per-user queue counter RedisKeyRateLimit = "im:ratelimit:" // + key — sliding-window rate limiting RedisKeyGlobalGate = "im:global:active" // global concurrent worker counter + + defaultRateLimitWindow = 60 * time.Second + defaultRateLimitMaxRequests = 10 ) // channelState holds runtime state for a running IM channel. @@ -372,7 +376,8 @@ type Service struct { // rateLimiter enforces per-user sliding window rate limiting. // Uses Redis ZSET when available, falls back to local sliding window. - rateLimiter *distributedLimiter + rateLimiter *ratelimit.Limiter + rateLimitMax int // inflight tracks in-progress QA requests, keyed by userKey // ("channelID:userID:chatID"). Allows /stop to abort a running request @@ -493,8 +498,8 @@ func resolveIMConfig(appCfg *config.Config) (workers, maxQueue, maxPerUser, glob workers = defaultWorkers maxQueue = defaultMaxQueueSize maxPerUser = defaultMaxPerUser - rlWindow = rateLimitWindow - rlMax = rateLimitMaxRequests + rlWindow = defaultRateLimitWindow + rlMax = defaultRateLimitMaxRequests if appCfg == nil || appCfg.IM == nil { return @@ -565,7 +570,8 @@ func NewService( cmdRegistry: registry, channels: make(map[string]*channelState), adapterFactories: make(map[string]AdapterFactory), - rateLimiter: newDistributedLimiter(redisClient, rlWindow, rlMax, instanceID), + rateLimiter: ratelimit.New(redisClient, RedisKeyRateLimit, rlWindow, instanceID), + rateLimitMax: rlMax, redis: redisClient, instanceID: instanceID, stopCh: make(chan struct{}), @@ -581,7 +587,7 @@ func NewService( if redisClient == nil { go s.dedupCleanupLoop() } - go s.rateLimiter.cleanupLoop(s.stopCh) + go s.rateLimiter.StartCleanup(s.stopCh) if redisClient != nil { globalInfo := "unlimited" @@ -1070,7 +1076,7 @@ func (s *Service) HandleMessage(ctx context.Context, msg *IncomingMessage, chann isCommand := s.cmdRegistry.IsRegistered(msg.Content) if !isCommand { rateLimitKey := makeUserKey(channelID, msg.UserID, msg.ChatID, threadID) - if !s.rateLimiter.Allow(rateLimitKey) { + if !s.rateLimiter.Allow(ctx, rateLimitKey, s.rateLimitMax) { logger.Warnf(ctx, "[IM] Rate limited: channel=%s user=%s chat=%s", channelID, msg.UserID, msg.ChatID) _ = adapter.SendReply(ctx, msg, &ReplyMessage{ Content: "您的消息发送过于频繁，请稍后再试。", diff --git a/internal/middleware/auth_public_ratelimit.go b/internal/middleware/auth_public_ratelimit.go index 17237462..6d032a28 100644 --- a/internal/middleware/auth_public_ratelimit.go +++ b/internal/middleware/auth_public_ratelimit.go @@ -24,7 +24,7 @@ import ( // Local in-memory only — fine for typical deployments since both // endpoints handle low absolute volumes; if/when WeKnora horizontally // scales the auth surface, swap to the Redis-backed limiter in -// internal/im/ratelimit.go (kept lowercase / package-private today). +// internal/ratelimit (shared with IM + embed surfaces). // publicAuthRateLimitWindow is the rolling window length per IP. const publicAuthRateLimitWindow = 60 * time.Second diff --git a/internal/middleware/embed_auth.go b/internal/middleware/embed_auth.go new file mode 100644 index 00000000..284a0c72 --- /dev/null +++ b/internal/middleware/embed_auth.go @@ -0,0 +1,247 @@ +package middleware + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/url" + "strings" + "sync" + "time" + + "github.com/Tencent/WeKnora/internal/application/service" + "github.com/Tencent/WeKnora/internal/logger" + "github.com/Tencent/WeKnora/internal/ratelimit" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/gin-gonic/gin" + "github.com/redis/go-redis/v9" +) + +const ( + embedRateLimitKeyPrefix = "embed:ratelimit:" + embedDailyRateLimitKeyPrefix = "embed:ratelimit:day:" + + // embedGlobalMinuteFactor derives a channel-wide per-minute cap from the + // per-IP cap. The publish token is publicly visible, so a single attacker + // can rotate IPs to defeat the per-IP limit; this bounds aggregate burst. + embedGlobalMinuteFactor = 20 + // embedGlobalMinuteFloor keeps the global per-minute cap usable even when + // the per-IP cap is tiny. + embedGlobalMinuteFloor = 120 +) + +// EmbedChannelContextKey stores the authenticated embed channel on the request context. +const EmbedChannelContextKey types.ContextKey = "EmbedChannel" + +var ( + embedLimiterOnce sync.Once + embedLimiter *ratelimit.Limiter + + embedDailyLimiterOnce sync.Once + embedDailyLimiter *ratelimit.Limiter +) + +func embedRateLimiter(redisClient *redis.Client) *ratelimit.Limiter { + embedLimiterOnce.Do(func() { + embedLimiter = ratelimit.New(redisClient, embedRateLimitKeyPrefix, time.Minute, "") + // Local-fallback eviction; Redis keys expire via PEXPIRE in the Lua script. + stopCh := make(chan struct{}) + go embedLimiter.StartCleanup(stopCh) + }) + return embedLimiter +} + +func embedDailyRateLimiter(redisClient *redis.Client) *ratelimit.Limiter { + embedDailyLimiterOnce.Do(func() { + embedDailyLimiter = ratelimit.New(redisClient, embedDailyRateLimitKeyPrefix, 24*time.Hour, "") + stopCh := make(chan struct{}) + go embedDailyLimiter.StartCleanup(stopCh) + }) + return embedDailyLimiter +} + +// embedGlobalPerMinute returns the channel-wide per-minute budget derived from +// the per-IP budget. +func embedGlobalPerMinute(perIP int) int { + budget := perIP * embedGlobalMinuteFactor + if budget < embedGlobalMinuteFloor { + budget = embedGlobalMinuteFloor + } + return budget +} + +// EmbedAuth validates publish tokens and injects a scoped tenant context for embed routes. +func EmbedAuth( + svc interfaces.EmbedChannelService, + tenantSvc interfaces.TenantService, + redisClient *redis.Client, +) gin.HandlerFunc { + limiter := embedRateLimiter(redisClient) + dailyLimiter := embedDailyRateLimiter(redisClient) + return func(c *gin.Context) { + channelID := strings.TrimSpace(c.Param("channel_id")) + if channelID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "channel_id is required"}) + c.Abort() + return + } + + token := extractEmbedToken(c) + if token == "" { + c.JSON(http.StatusUnauthorized, gin.H{"error": "embed publish token is required"}) + c.Abort() + return + } + + var ch *types.EmbedChannel + var err error + if service.IsEmbedSessionToken(token) { + resolvedID, resolveErr := svc.ResolveSessionToken(c.Request.Context(), token) + if resolveErr != nil || resolvedID != channelID { + c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid embed channel or token"}) + c.Abort() + return + } + ch, err = svc.LookupEnabledChannel(c.Request.Context(), channelID) + } else { + ch, err = svc.LookupForEmbed(c.Request.Context(), channelID, token) + } + if err != nil { + if errors.Is(err, service.ErrEmbedChannelDisabled) { + c.JSON(http.StatusForbidden, gin.H{"error": "embed channel is disabled"}) + c.Abort() + return + } + c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid embed channel or token"}) + c.Abort() + return + } + + origin := requestOrigin(c) + if !originAllowed(origin, ch.AllowedOriginsList()) { + logger.Warnf(c.Request.Context(), "[embed_auth] origin %q not allowed for channel %s", origin, channelID) + c.JSON(http.StatusForbidden, gin.H{"error": "origin not allowed"}) + c.Abort() + return + } + + // Per-IP per-minute cap. + rateKey := fmt.Sprintf("%s:%s", channelID, c.ClientIP()) + if !limiter.Allow(c.Request.Context(), rateKey, ch.RateLimitPerMinute) { + c.JSON(http.StatusTooManyRequests, gin.H{"error": "rate limit exceeded"}) + c.Abort() + return + } + // Channel-wide per-minute cap (bounds burst across rotating IPs since + // the publish token is publicly visible). + if !limiter.Allow(c.Request.Context(), channelID+":__global", embedGlobalPerMinute(ch.RateLimitPerMinute)) { + c.JSON(http.StatusTooManyRequests, gin.H{"error": "rate limit exceeded"}) + c.Abort() + return + } + // Channel-wide daily total cap (bounds sustained abuse). + if !dailyLimiter.Allow(c.Request.Context(), channelID, ch.RateLimitPerDay) { + c.JSON(http.StatusTooManyRequests, gin.H{"error": "daily request limit exceeded"}) + c.Abort() + return + } + + tenant, err := tenantSvc.GetTenantByID(c.Request.Context(), ch.TenantID) + if err != nil || tenant == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "tenant unavailable"}) + c.Abort() + return + } + + user := &types.User{ + ID: fmt.Sprintf("embed-%s", channelID), + Username: fmt.Sprintf("embed-%s", channelID), + Email: fmt.Sprintf("embed-%s@embed.local", channelID), + TenantID: ch.TenantID, + IsActive: true, + } + + c.Set(types.TenantIDContextKey.String(), ch.TenantID) + c.Set(types.TenantInfoContextKey.String(), tenant) + c.Set(types.UserContextKey.String(), user) + c.Set(types.UserIDContextKey.String(), user.ID) + c.Set(types.TenantRoleContextKey.String(), types.TenantRoleViewer) + c.Set(types.SystemAdminContextKey.String(), false) + c.Set(string(EmbedChannelContextKey), ch) + + ctx := c.Request.Context() + ctx = context.WithValue(ctx, types.TenantIDContextKey, ch.TenantID) + ctx = context.WithValue(ctx, types.TenantInfoContextKey, tenant) + ctx = context.WithValue(ctx, types.UserContextKey, user) + ctx = context.WithValue(ctx, types.UserIDContextKey, user.ID) + ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleViewer) + ctx = context.WithValue(ctx, types.SystemAdminContextKey, false) + ctx = context.WithValue(ctx, EmbedChannelContextKey, ch) + c.Request = c.Request.WithContext(ctx) + c.Next() + } +} + +func extractEmbedToken(c *gin.Context) string { + // Only accept the token via the Authorization header. A query-string token + // would be captured by proxy/access logs and browser history; the embed + // client always sends "Authorization: Embed ". + auth := c.GetHeader("Authorization") + if strings.HasPrefix(auth, "Embed ") { + return strings.TrimSpace(strings.TrimPrefix(auth, "Embed ")) + } + return "" +} + +func requestOrigin(c *gin.Context) string { + if o := strings.TrimSpace(c.GetHeader("Origin")); o != "" { + return o + } + ref := strings.TrimSpace(c.GetHeader("Referer")) + if ref == "" { + return "" + } + u, err := url.Parse(ref) + if err != nil { + return "" + } + if u.Scheme == "" || u.Host == "" { + return "" + } + return u.Scheme + "://" + u.Host +} + +func originAllowed(origin string, allowed []string) bool { + // Empty allowlist rejects all origins. Management create/update requires at + // least one origin; legacy rows with [] must be fixed before going live. + if len(allowed) == 0 { + return false + } + if origin == "" { + return false + } + for _, pattern := range allowed { + pattern = strings.TrimSpace(pattern) + if pattern == "" { + continue + } + if pattern == "*" || strings.EqualFold(pattern, origin) { + return true + } + if strings.HasPrefix(pattern, "*.") { + suffix := strings.TrimPrefix(pattern, "*") + if strings.HasSuffix(origin, suffix) { + return true + } + } + } + return false +} + +// EmbedChannelFromContext returns the authenticated embed channel, if any. +func EmbedChannelFromContext(ctx context.Context) (*types.EmbedChannel, bool) { + ch, ok := ctx.Value(EmbedChannelContextKey).(*types.EmbedChannel) + return ch, ok && ch != nil +} diff --git a/internal/middleware/embed_auth_test.go b/internal/middleware/embed_auth_test.go new file mode 100644 index 00000000..84296f60 --- /dev/null +++ b/internal/middleware/embed_auth_test.go @@ -0,0 +1,408 @@ +package middleware + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Tencent/WeKnora/internal/application/service" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/gin-gonic/gin" +) + +type fakeEmbedChannelService struct { + channels map[string]*types.EmbedChannel + sessions map[string]string +} + +func (f *fakeEmbedChannelService) Create( + ctx context.Context, tenantID uint64, agentID string, req *types.EmbedChannel, +) (*types.EmbedChannel, string, error) { + return nil, "", nil +} + +func (f *fakeEmbedChannelService) ListByAgent( + ctx context.Context, tenantID uint64, agentID string, +) ([]*types.EmbedChannel, error) { + return nil, nil +} + +func (f *fakeEmbedChannelService) ListByTenant( + ctx context.Context, tenantID uint64, +) ([]*types.EmbedChannel, error) { + return nil, nil +} + +func (f *fakeEmbedChannelService) Update( + ctx context.Context, tenantID uint64, id string, req *types.EmbedChannel, + enabled *bool, showSuggested *bool, allowWebSearch *bool, allowMemory *bool, allowFileUpload *bool, + defaultLocale *string, webhookURL *string, webhookSecret *string, +) (*types.EmbedChannel, error) { + return nil, nil +} + +func (f *fakeEmbedChannelService) GetOwnedChannel( + ctx context.Context, tenantID uint64, id string, +) (*types.EmbedChannel, error) { + ch := f.channels[id] + if ch == nil || ch.TenantID != tenantID { + return nil, service.ErrEmbedChannelNotFound + } + return ch, nil +} + +func (f *fakeEmbedChannelService) Delete(ctx context.Context, tenantID uint64, id string) error { + return nil +} + +func (f *fakeEmbedChannelService) RotateToken( + ctx context.Context, tenantID uint64, id string, +) (*types.EmbedChannel, string, error) { + return nil, "", nil +} + +func (f *fakeEmbedChannelService) LookupForEmbed( + ctx context.Context, channelID, token string, +) (*types.EmbedChannel, error) { + ch := f.channels[channelID] + if ch == nil || ch.PublishToken != token { + return nil, service.ErrEmbedTokenInvalid + } + if !ch.Enabled { + return nil, service.ErrEmbedChannelDisabled + } + return ch, nil +} + +func (f *fakeEmbedChannelService) LookupEnabledChannel( + ctx context.Context, channelID string, +) (*types.EmbedChannel, error) { + ch := f.channels[channelID] + if ch == nil { + return nil, service.ErrEmbedTokenInvalid + } + if !ch.Enabled { + return nil, service.ErrEmbedChannelDisabled + } + return ch, nil +} + +func (f *fakeEmbedChannelService) IssueSessionToken( + ctx context.Context, channelID string, +) (string, int, error) { + return "ems_testtoken", 1800, nil +} + +func (f *fakeEmbedChannelService) ResolveSessionToken(ctx context.Context, token string) (string, error) { + channelID, ok := f.sessions[token] + if !ok { + return "", service.ErrEmbedTokenInvalid + } + return channelID, nil +} + +func (f *fakeEmbedChannelService) PublicConfig( + ctx context.Context, ch *types.EmbedChannel, +) types.EmbedChannelPublicConfig { + return types.EmbedChannelPublicConfig{ChannelID: ch.ID} +} + +func (f *fakeEmbedChannelService) SuggestedQuestions( + ctx context.Context, ch *types.EmbedChannel, limit int, +) ([]types.SuggestedQuestion, error) { + return nil, nil +} + +func (f *fakeEmbedChannelService) EmbedChunk( + ctx context.Context, ch *types.EmbedChannel, chunkID string, +) (*types.Chunk, error) { + return nil, nil +} + +func (f *fakeEmbedChannelService) IssuePreviewSession( + ctx context.Context, tenantID uint64, channelID string, +) (string, int, error) { + return "", 0, nil +} + +func (f *fakeEmbedChannelService) EmbedDisplayTitle(ctx context.Context, ch *types.EmbedChannel) string { + return "" +} + +type fakeTenantService struct { + tenant *types.Tenant +} + +func (f *fakeTenantService) GetTenantByID(ctx context.Context, id uint64) (*types.Tenant, error) { + return f.tenant, nil +} + +func (f *fakeTenantService) CreateTenant(ctx context.Context, tenant *types.Tenant) (*types.Tenant, error) { + return nil, nil +} + +func (f *fakeTenantService) GetTenantsByIDs(ctx context.Context, ids []uint64) (map[uint64]*types.Tenant, error) { + return nil, nil +} + +func (f *fakeTenantService) UpdateTenant(ctx context.Context, tenant *types.Tenant) (*types.Tenant, error) { + return nil, nil +} + +func (f *fakeTenantService) DeleteTenant(ctx context.Context, id uint64) error { + return nil +} + +func (f *fakeTenantService) ListTenants(ctx context.Context) ([]*types.Tenant, error) { + return nil, nil +} + +func (f *fakeTenantService) UpdateAPIKey(ctx context.Context, id uint64) (string, error) { + return "", nil +} + +func (f *fakeTenantService) ExtractTenantIDFromAPIKey(apiKey string) (uint64, error) { + return 0, nil +} + +func (f *fakeTenantService) ListAllTenants(ctx context.Context) ([]*types.Tenant, error) { + return nil, nil +} + +func (f *fakeTenantService) BulkSetStorageQuota(ctx context.Context, quotaBytes int64) (int64, error) { + return 0, nil +} + +func (f *fakeTenantService) SearchTenants( + ctx context.Context, keyword string, tenantID uint64, page, pageSize int, +) ([]*types.Tenant, int64, error) { + return nil, 0, nil +} + +func (f *fakeTenantService) GetTenantByIDForUser( + ctx context.Context, tenantID uint64, userID string, +) (*types.Tenant, error) { + return f.tenant, nil +} + +func (f *fakeTenantService) GetWeKnoraCloudCredentials(ctx context.Context) *types.WeKnoraCloudCredentials { + return nil +} + +var ( + _ interfaces.EmbedChannelService = (*fakeEmbedChannelService)(nil) + _ interfaces.TenantService = (*fakeTenantService)(nil) +) + +func TestEmbedGlobalPerMinute(t *testing.T) { + tests := []struct { + perIP int + want int + }{ + {perIP: 0, want: 120}, + {perIP: 1, want: 120}, + {perIP: 6, want: 120}, + {perIP: 7, want: 140}, + {perIP: 10, want: 200}, + } + for _, tt := range tests { + name := fmt.Sprintf("perIP=%d", tt.perIP) + t.Run(name, func(t *testing.T) { + if got := embedGlobalPerMinute(tt.perIP); got != tt.want { + t.Fatalf("embedGlobalPerMinute(%d) = %d, want %d", tt.perIP, got, tt.want) + } + }) + } +} + +func TestOriginAllowed(t *testing.T) { + tests := []struct { + name string + origin string + allowed []string + want bool + }{ + {name: "empty allow list", origin: "https://evil.com", allowed: nil, want: false}, + {name: "exact match", origin: "https://app.example.com", allowed: []string{"https://app.example.com"}, want: true}, + {name: "wildcard star", origin: "https://any.example.com", allowed: []string{"*"}, want: true}, + {name: "subdomain suffix", origin: "https://app.example.com", allowed: []string{"*.example.com"}, want: true}, + {name: "missing origin", origin: "", allowed: []string{"https://app.example.com"}, want: false}, + {name: "not allowed", origin: "https://evil.com", allowed: []string{"https://app.example.com"}, want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := originAllowed(tt.origin, tt.allowed); got != tt.want { + t.Fatalf("originAllowed(%q, %v) = %v, want %v", tt.origin, tt.allowed, got, tt.want) + } + }) + } +} + +func TestExtractEmbedToken(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + header string + query string + want string + }{ + {name: "authorization header", header: "Embed em_publish", want: "em_publish"}, + {name: "query param rejected", query: "ems_session", want: ""}, + {name: "header preferred", header: "Embed em_header", query: "em_query", want: "em_header"}, + {name: "missing", want: ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + req := httptest.NewRequest(http.MethodGet, "/?token="+tt.query, nil) + if tt.header != "" { + req.Header.Set("Authorization", tt.header) + } + c.Request = req + if got := extractEmbedToken(c); got != tt.want { + t.Fatalf("extractEmbedToken() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestEmbedAuthSessionTokenPath(t *testing.T) { + gin.SetMode(gin.TestMode) + + const channelID = "ch-1" + svc := &fakeEmbedChannelService{ + channels: map[string]*types.EmbedChannel{ + channelID: { + ID: channelID, + TenantID: 42, + Enabled: true, + AllowedOrigins: []byte(`["https://app.example.com"]`), + RateLimitPerMinute: 0, + }, + }, + sessions: map[string]string{ + "ems_valid": channelID, + }, + } + tenantSvc := &fakeTenantService{tenant: &types.Tenant{ID: 42}} + + r := gin.New() + r.GET("/api/v1/embed/:channel_id/config", EmbedAuth(svc, tenantSvc, nil), func(c *gin.Context) { + ch, ok := EmbedChannelFromContext(c.Request.Context()) + if !ok || ch.ID != channelID { + c.JSON(http.StatusInternalServerError, gin.H{"error": "missing channel"}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true}) + }) + req := httptest.NewRequest(http.MethodGet, "/api/v1/embed/"+channelID+"/config", nil) + req.Header.Set("Authorization", "Embed ems_valid") + req.Header.Set("Origin", "https://app.example.com") + w := httptest.NewRecorder() + r.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } + var body map[string]any + if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil { + t.Fatalf("unmarshal response: %v", err) + } + if body["success"] != true { + t.Fatalf("expected success response, got %v", body) + } +} + +func TestEmbedAuthPublishTokenValid(t *testing.T) { + gin.SetMode(gin.TestMode) + + const ( + channelID = "ch-pub-1" + publishToken = "em_publish_ok" + ) + svc := &fakeEmbedChannelService{ + channels: map[string]*types.EmbedChannel{ + channelID: { + ID: channelID, + TenantID: 11, + Enabled: true, + PublishToken: publishToken, + AllowedOrigins: []byte(`["https://app.example.com"]`), + RateLimitPerMinute: 0, + }, + }, + } + tenantSvc := &fakeTenantService{tenant: &types.Tenant{ID: 11}} + + r := gin.New() + r.GET("/api/v1/embed/:channel_id/config", EmbedAuth(svc, tenantSvc, nil), func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"success": true}) + }) + + req := httptest.NewRequest(http.MethodGet, "/api/v1/embed/"+channelID+"/config", nil) + req.Header.Set("Authorization", "Embed "+publishToken) + req.Header.Set("Origin", "https://app.example.com") + w := httptest.NewRecorder() + r.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } +} + +func TestEmbedAuthPublishTokenInvalid(t *testing.T) { + gin.SetMode(gin.TestMode) + + const channelID = "ch-pub-1" + svc := &fakeEmbedChannelService{ + channels: map[string]*types.EmbedChannel{ + channelID: { + ID: channelID, + TenantID: 11, + Enabled: true, + PublishToken: "em_real_token", + AllowedOrigins: []byte(`["https://app.example.com"]`), + }, + }, + } + handler := EmbedAuth(svc, &fakeTenantService{tenant: &types.Tenant{ID: 11}}, nil) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/embed/"+channelID+"/config", nil) + c.Request.Header.Set("Authorization", "Embed em_wrong_token") + c.Request.Header.Set("Origin", "https://app.example.com") + c.Params = gin.Params{{Key: "channel_id", Value: channelID}} + handler(c) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d, body = %s", w.Code, http.StatusUnauthorized, w.Body.String()) + } +} + +func TestEmbedAuthSessionTokenMismatch(t *testing.T) { + gin.SetMode(gin.TestMode) + + const channelID = "ch-1" + svc := &fakeEmbedChannelService{ + sessions: map[string]string{ + "ems_other": "other-channel", + }, + } + handler := EmbedAuth(svc, &fakeTenantService{tenant: &types.Tenant{ID: 1}}, nil) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/embed/"+channelID+"/config", nil) + c.Request.Header.Set("Authorization", "Embed ems_other") + c.Params = gin.Params{{Key: "channel_id", Value: channelID}} + handler(c) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d, body = %s", w.Code, http.StatusUnauthorized, w.Body.String()) + } +} diff --git a/internal/ratelimit/limiter.go b/internal/ratelimit/limiter.go new file mode 100644 index 00000000..5f646b41 --- /dev/null +++ b/internal/ratelimit/limiter.go @@ -0,0 +1,181 @@ +// Package ratelimit provides a Redis-backed sliding-window rate limiter with a +// local in-memory fallback when Redis is unavailable (Lite / single-instance). +package ratelimit + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/google/uuid" + "github.com/redis/go-redis/v9" +) + +const localCleanupInterval = time.Minute + +// rateLimitScript atomically prunes expired ZSET members, checks the count, +// and conditionally records a new hit. +var rateLimitScript = redis.NewScript(` +local key = KEYS[1] +local now = tonumber(ARGV[1]) +local window = tonumber(ARGV[2]) +local maxReq = tonumber(ARGV[3]) +local member = ARGV[4] + +redis.call('ZREMRANGEBYSCORE', key, 0, now - window) +local count = redis.call('ZCARD', key) +if count < maxReq then + redis.call('ZADD', key, now, member) + redis.call('PEXPIRE', key, window + 1000) + return 1 +end +return 0 +`) + +// Limiter enforces per-key sliding-window limits. max is evaluated per Allow +// call so callers (e.g. embed channels) can vary budgets without rebuilding +// the limiter. +type Limiter struct { + redis *redis.Client + local *localLimiter + keyPrefix string + window time.Duration + instanceID string +} + +// New constructs a limiter. keyPrefix should include a trailing delimiter +// (e.g. "embed:ratelimit:"). When redis is nil, only the local fallback runs. +func New(redisClient *redis.Client, keyPrefix string, window time.Duration, instanceID string) *Limiter { + if window <= 0 { + window = time.Minute + } + if instanceID == "" { + instanceID = uuid.New().String() + } + return &Limiter{ + redis: redisClient, + local: newLocalLimiter(), + keyPrefix: keyPrefix, + window: window, + instanceID: instanceID, + } +} + +// Allow reports whether key is within budget for the current window. +func (l *Limiter) Allow(ctx context.Context, key string, max int) bool { + if max <= 0 { + return true + } + if l.redis != nil { + allowed, err := l.redisAllow(ctx, key, max) + if err == nil { + return allowed + } + } + return l.local.allow(key, l.window, max) +} + +func (l *Limiter) redisAllow(ctx context.Context, key string, max int) (bool, error) { + redisKey := l.keyPrefix + key + nowMs := time.Now().UnixMilli() + windowMs := l.window.Milliseconds() + member := fmt.Sprintf("%s:%d", l.instanceID, nowMs) + + result, err := rateLimitScript.Run(ctx, l.redis, + []string{redisKey}, + nowMs, windowMs, max, member, + ).Int64() + if err != nil { + return false, err + } + return result == 1, nil +} + +// StartCleanup runs periodic eviction for the local fallback map. No-op when +// only Redis is in use, but cheap to call either way. +func (l *Limiter) StartCleanup(stopCh <-chan struct{}) { + l.local.startCleanup(l.window, stopCh) +} + +type localEntry struct { + mu sync.Mutex + timestamps []time.Time + deleted bool +} + +type localLimiter struct { + entries sync.Map // key -> *localEntry +} + +func newLocalLimiter() *localLimiter { + return &localLimiter{} +} + +func (l *localLimiter) allow(key string, window time.Duration, max int) bool { + now := time.Now() + cutoff := now.Add(-window) + + for { + val, _ := l.entries.LoadOrStore(key, &localEntry{}) + entry := val.(*localEntry) + + entry.mu.Lock() + if entry.deleted { + entry.mu.Unlock() + l.entries.Delete(key) + continue + } + + valid := entry.timestamps[:0] + for _, t := range entry.timestamps { + if t.After(cutoff) { + valid = append(valid, t) + } + } + entry.timestamps = valid + + if len(entry.timestamps) >= max { + entry.mu.Unlock() + return false + } + entry.timestamps = append(entry.timestamps, now) + entry.mu.Unlock() + return true + } +} + +func (l *localLimiter) startCleanup(window time.Duration, stopCh <-chan struct{}) { + if window <= 0 { + window = time.Minute + } + ticker := time.NewTicker(localCleanupInterval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + // Evict against the limiter's own window, not a hardcoded minute, + // so non-minute limiters (e.g. a per-day cap) are not dropped early. + cutoff := time.Now().Add(-window) + l.entries.Range(func(key, val any) bool { + entry := val.(*localEntry) + entry.mu.Lock() + allExpired := true + for _, t := range entry.timestamps { + if t.After(cutoff) { + allExpired = false + break + } + } + if allExpired { + entry.deleted = true + l.entries.Delete(key) + } + entry.mu.Unlock() + return true + }) + case <-stopCh: + return + } + } +} diff --git a/internal/ratelimit/limiter_test.go b/internal/ratelimit/limiter_test.go new file mode 100644 index 00000000..c6779859 --- /dev/null +++ b/internal/ratelimit/limiter_test.go @@ -0,0 +1,41 @@ +package ratelimit + +import ( + "context" + "testing" + "time" +) + +func TestLocalLimiterAllowsWithinBudget(t *testing.T) { + l := New(nil, "test:", time.Minute, "inst") + ctx := context.Background() + const max = 3 + for i := 0; i < max; i++ { + if !l.Allow(ctx, "k1", max) { + t.Fatalf("request %d should be allowed", i+1) + } + } + if l.Allow(ctx, "k1", max) { + t.Fatal("request over budget should be denied") + } +} + +func TestLocalLimiterIndependentKeys(t *testing.T) { + l := New(nil, "test:", time.Minute, "inst") + ctx := context.Background() + if !l.Allow(ctx, "a", 1) { + t.Fatal("first key should be allowed") + } + if !l.Allow(ctx, "b", 1) { + t.Fatal("second key should be allowed") + } +} + +func TestAllowSkipsWhenMaxZero(t *testing.T) { + l := New(nil, "test:", time.Minute, "inst") + for i := 0; i < 100; i++ { + if !l.Allow(context.Background(), "k", 0) { + t.Fatal("max<=0 should always allow") + } + } +} diff --git a/internal/router/router.go b/internal/router/router.go index d4e16eb9..7b72dcd7 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -13,6 +13,7 @@ import ( filesvc "github.com/Tencent/WeKnora/internal/application/service/file" "github.com/gin-contrib/cors" "github.com/gin-gonic/gin" + "github.com/redis/go-redis/v9" swaggerFiles "github.com/swaggo/files" ginSwagger "github.com/swaggo/gin-swagger" "go.uber.org/dig" @@ -77,6 +78,9 @@ type RouterParams struct { SkillHandler *handler.SkillHandler OrganizationHandler *handler.OrganizationHandler IMHandler *handler.IMHandler + EmbedChannelHandler *handler.EmbedChannelHandler + EmbedChannelService interfaces.EmbedChannelService + RedisClient *redis.Client DataSourceHandler *handler.DataSourceHandler DataSourceCredentialsHandler *handler.DataSourceCredentialsHandler WeKnoraCloudHandler *handler.WeKnoraCloudHandler @@ -88,11 +92,21 @@ func NewRouter(params RouterParams) *gin.Engine { r := gin.New() r.ContextWithFallback = true + // Trusted proxies: gin defaults to trusting ALL proxies, which makes + // c.ClientIP() honor a client-supplied X-Forwarded-For. Public, unauthed + // embed endpoints rate-limit per (channel, ClientIP), so a spoofed XFF would + // trivially bypass the limiter. Restrict to the fronting proxy network so + // only the real client IP (appended by nginx) is returned. Configurable via + // WEKNORA_TRUSTED_PROXIES (comma-separated CIDRs/IPs). + if err := r.SetTrustedProxies(trustedProxies()); err != nil { + logger.Errorf(context.Background(), "[Router] failed to set trusted proxies: %v", err) + } + // CORS 中间件应放在最前面 r.Use(cors.New(cors.Config{ AllowOrigins: []string{"*"}, AllowMethods: []string{"GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"}, - AllowHeaders: []string{"Origin", "Content-Type", "Accept", "Authorization", "X-API-Key", "X-Request-ID"}, + AllowHeaders: []string{"Origin", "Content-Type", "Accept", "Authorization", "X-API-Key", "X-Request-ID", "X-Tenant-ID", "X-Embed-Session"}, ExposeHeaders: []string{"Content-Length", "Access-Control-Allow-Origin"}, AllowCredentials: true, MaxAge: 12 * time.Hour, @@ -121,6 +135,15 @@ func NewRouter(params RouterParams) *gin.Engine { )) } + // Embed page framing policy: emit a per-channel `frame-ancestors` CSP so the + // embed SPA page (/embed/:channelId) can only be iframed by the channel's + // allowed origins. This is the page-level counterpart to the API Origin + // allowlist enforced in EmbedAuth. Registered before the static handler so + // it runs for the embed HTML response. + if params.EmbedChannelService != nil { + r.Use(embedFrameAncestorsMiddleware(params.EmbedChannelService)) + } + // 前端静态文件（仅 Lite 版本内嵌前端） if handler.Edition == "lite" { serveFrontendStatic(r) @@ -129,6 +152,9 @@ func NewRouter(params RouterParams) *gin.Engine { // IM 回调路由（在认证中间件之前注册，使用各平台自身的签名验证） RegisterIMRoutes(r, params.IMHandler) + // Web embed 公开路由（使用 publish token 鉴权，不走全局 Auth） + RegisterEmbedPublicRoutes(r, params.EmbedChannelHandler, params.EmbedChannelService, params.TenantService, params.RedisClient, params.FileService) + // 认证中间件 r.Use(middleware.Auth(params.TenantService, params.UserService, params.TenantMemberService, params.Config)) @@ -199,6 +225,7 @@ func NewRouter(params RouterParams) *gin.Engine { RegisterSkillRoutes(v1, params.SkillHandler, rbacGuards) RegisterOrganizationRoutes(v1, params.OrganizationHandler, rbacGuards) RegisterIMChannelRoutes(v1, params.IMHandler, rbacGuards) + RegisterEmbedChannelRoutes(v1, params.EmbedChannelHandler, rbacGuards) RegisterDataSourceRoutes(v1, params.DataSourceHandler, params.DataSourceCredentialsHandler, rbacGuards) RegisterWeKnoraCloudRoutes(v1, params.WeKnoraCloudHandler, rbacGuards) RegisterWikiPageRoutes(v1, params.WikiPageHandler, rbacGuards) @@ -1088,6 +1115,58 @@ func RegisterOrganizationRoutes(r *gin.RouterGroup, orgHandler *handler.Organiza r.POST("/shared-agents/disabled", g.Admin(), orgHandler.SetSharedAgentDisabledByMe) } +// RegisterEmbedPublicRoutes registers anonymous embed endpoints secured by publish tokens. +func RegisterEmbedPublicRoutes( + r *gin.Engine, + embedHandler *handler.EmbedChannelHandler, + embedService interfaces.EmbedChannelService, + tenantService interfaces.TenantService, + redisClient *redis.Client, + fileService interfaces.FileService, +) { + if embedHandler == nil || embedService == nil { + return + } + embed := r.Group("/api/v1/embed/:channel_id", middleware.EmbedAuth(embedService, tenantService, redisClient)) + { + embed.POST("/exchange", embedHandler.ExchangeEmbedSession) + embed.GET("/config", embedHandler.GetEmbedConfig) + embed.GET("/suggested-questions", embedHandler.GetEmbedSuggestedQuestions) + embed.GET("/chunks/:chunk_id", embedHandler.GetEmbedChunk) + embed.POST("/sessions", embedHandler.CreateEmbedSession) + embed.POST("/knowledge-chat/:session_id", embedHandler.EmbedKnowledgeChat) + embed.POST("/agent-chat/:session_id", embedHandler.EmbedAgentChat) + embed.GET("/messages/:session_id/load", embedHandler.EmbedLoadMessages) + embed.POST("/sessions/:session_id/stop", embedHandler.EmbedStopSession) + embed.POST("/sessions/:session_id/events", embedHandler.EmbedRelayWebhookEvent) + // Serve images embedded in bot replies (e.g. chart exports). EmbedAuth + // injects the channel's tenant, and the handler enforces that the + // requested path belongs to that tenant. + embed.GET("/files", newFileServeHandler(fileService)) + } +} + +// RegisterEmbedChannelRoutes registers authenticated embed channel management routes. +func RegisterEmbedChannelRoutes(r *gin.RouterGroup, embedHandler *handler.EmbedChannelHandler, g *rbacGuards) { + if embedHandler == nil { + return + } + agentEmbed := r.Group("/agents/:id/embed-channels") + { + agentEmbed.POST("", g.Admin(), embedHandler.CreateEmbedChannel) + agentEmbed.GET("", g.Viewer(), embedHandler.ListEmbedChannels) + } + channels := r.Group("/embed-channels") + { + channels.GET("", g.Viewer(), embedHandler.ListAllEmbedChannels) + channels.PUT("/:channel_id", g.Admin(), embedHandler.UpdateEmbedChannel) + channels.DELETE("/:channel_id", g.Admin(), embedHandler.DeleteEmbedChannel) + channels.POST("/:channel_id/rotate-token", g.Admin(), embedHandler.RotateEmbedToken) + channels.POST("/:channel_id/preview-session", g.Viewer(), embedHandler.IssuePreviewSession) + channels.GET("/:channel_id/stats", g.Viewer(), embedHandler.GetEmbedChannelStats) + } +} + // RegisterIMRoutes registers IM callback routes. // These are registered BEFORE auth middleware since IM platforms use their own signature verification. func RegisterIMRoutes(r *gin.Engine, imHandler *handler.IMHandler) { @@ -1129,6 +1208,92 @@ func RegisterIMChannelRoutes(r *gin.RouterGroup, imHandler *handler.IMHandler, g } } +// trustedProxies returns the proxy CIDRs/IPs whose X-Forwarded-For headers +// gin should trust when resolving the client IP. Defaults to loopback and +// private ranges (covers the bundled nginx in a container network); override +// with WEKNORA_TRUSTED_PROXIES (comma-separated). An explicit empty value +// disables proxy trust entirely so ClientIP() returns the direct peer. +func trustedProxies() []string { + raw, ok := os.LookupEnv("WEKNORA_TRUSTED_PROXIES") + if !ok { + return []string{ + "127.0.0.0/8", + "::1/128", + "10.0.0.0/8", + "172.16.0.0/12", + "192.168.0.0/16", + "fc00::/7", + } + } + proxies := make([]string, 0) + for _, p := range strings.Split(raw, ",") { + if p = strings.TrimSpace(p); p != "" { + proxies = append(proxies, p) + } + } + return proxies +} + +// embedChannelIDFromPath extracts the channel id from an /embed/:channelID path. +func embedChannelIDFromPath(path string) string { + const prefix = "/embed/" + if !strings.HasPrefix(path, prefix) { + return "" + } + rest := strings.TrimPrefix(path, prefix) + if i := strings.IndexByte(rest, '/'); i >= 0 { + rest = rest[:i] + } + if i := strings.IndexByte(rest, '?'); i >= 0 { + rest = rest[:i] + } + return strings.TrimSpace(rest) +} + +// embedFrameAncestorsMiddleware sets a per-channel `frame-ancestors` CSP on the +// embed SPA page so it can only be framed by the channel's allowed origins. +// When the channel declares no origins (or "*"), no restriction is applied, +// matching the API allowlist semantics. Only GET/HEAD page loads are handled. +func embedFrameAncestorsMiddleware(svc interfaces.EmbedChannelService) gin.HandlerFunc { + return func(c *gin.Context) { + if c.Request.Method != http.MethodGet && c.Request.Method != http.MethodHead { + c.Next() + return + } + channelID := embedChannelIDFromPath(c.Request.URL.Path) + if channelID == "" { + c.Next() + return + } + ch, err := svc.LookupEnabledChannel(c.Request.Context(), channelID) + if err != nil || ch == nil { + c.Next() + return + } + origins := ch.AllowedOriginsList() + sources := make([]string, 0, len(origins)) + wildcard := false + for _, o := range origins { + o = strings.TrimSpace(o) + if o == "" { + continue + } + if o == "*" { + wildcard = true + break + } + sources = append(sources, o) + } + // No explicit origins or a wildcard => do not constrain framing here. + if wildcard || len(sources) == 0 { + c.Next() + return + } + c.Header("Content-Security-Policy", "frame-ancestors "+strings.Join(sources, " ")) + c.Next() + } +} + // serveFrontendStatic registers a middleware that serves the frontend SPA // from the ./web directory if it exists. Must be called BEFORE auth middleware // so static files are served without authentication. @@ -1191,7 +1356,12 @@ type getRouteRegistrar interface { GET(string, ...gin.HandlerFunc) gin.IRoutes } -func serveFiles(r getRouteRegistrar, globalFileService interfaces.FileService) { +// newFileServeHandler builds the file-proxy handler. It reads the tenant from +// the request context (set by whichever auth middleware precedes it), so the +// same handler backs both the authenticated /files route and the embed route +// (where EmbedAuth injects the channel's tenant). Tenant ownership of the +// requested path is enforced via ValidateStoragePathTenant either way. +func newFileServeHandler(globalFileService interfaces.FileService) gin.HandlerFunc { baseDir := os.Getenv("LOCAL_STORAGE_BASE_DIR") if baseDir == "" { baseDir = "/data/files" @@ -1203,9 +1373,7 @@ func serveFiles(r getRouteRegistrar, globalFileService interfaces.FileService) { } } - logger.Infof(context.Background(), "[Router] Serving files from /files (local base: %s)", absDir) - - r.GET("/files", func(c *gin.Context) { + return func(c *gin.Context) { filePath := strings.TrimSpace(c.Query("file_path")) if filePath == "" { c.JSON(http.StatusBadRequest, gin.H{"error": "missing required parameter: file_path"}) @@ -1292,7 +1460,12 @@ func serveFiles(r getRouteRegistrar, globalFileService interfaces.FileService) { if _, err := io.Copy(c.Writer, reader); err != nil { logger.Warnf(context.Background(), "[Router] /files write response failed: %v", err) } - }) + } +} + +func serveFiles(r getRouteRegistrar, globalFileService interfaces.FileService) { + logger.Infof(context.Background(), "[Router] Serving files from /files") + r.GET("/files", newFileServeHandler(globalFileService)) } // servePresignedFiles serves files via HMAC-signed URLs without requiring authentication. diff --git a/internal/types/embed_channel.go b/internal/types/embed_channel.go new file mode 100644 index 00000000..00f68f1a --- /dev/null +++ b/internal/types/embed_channel.go @@ -0,0 +1,150 @@ +package types + +import ( + "encoding/json" + "strings" + "time" + + "github.com/google/uuid" + "gorm.io/gorm" +) + +// EmbedChannel publishes an agent chat surface for external websites. +type EmbedChannel struct { + ID string `json:"id" gorm:"type:varchar(36);primaryKey"` + TenantID uint64 `json:"tenant_id" gorm:"not null;index:idx_embed_channels_tenant"` + AgentID string `json:"agent_id" gorm:"type:varchar(36);not null;index:idx_embed_channels_agent;default:'builtin-quick-answer'"` + Name string `json:"name" gorm:"type:varchar(255);not null;default:''"` + Enabled bool `json:"enabled" gorm:"not null;default:true"` + PublishToken string `json:"-" gorm:"type:varchar(64);not null;default:''"` + AllowedOrigins JSON `json:"allowed_origins" gorm:"type:jsonb;not null;default:'[]'"` + WelcomeMessage string `json:"welcome_message" gorm:"type:text;not null;default:''"` + RateLimitPerMinute int `json:"rate_limit_per_minute" gorm:"not null;default:30"` + RateLimitPerDay int `json:"rate_limit_per_day" gorm:"not null;default:10000"` + PrimaryColor string `json:"primary_color" gorm:"type:varchar(32);not null;default:''"` + PageTitle string `json:"page_title" gorm:"type:varchar(255);not null;default:''"` + HeaderTitleMode string `json:"header_title_mode" gorm:"type:varchar(32);not null;default:'channel'"` + ShowSuggestedQuestions bool `json:"show_suggested_questions" gorm:"not null;default:true"` + WidgetPosition string `json:"widget_position" gorm:"type:varchar(32);not null;default:'bottom-right'"` + AllowWebSearch bool `json:"allow_web_search" gorm:"not null;default:false"` + AllowMemory bool `json:"allow_memory" gorm:"not null;default:false"` + AllowFileUpload bool `json:"allow_file_upload" gorm:"not null;default:false"` + DefaultLocale string `json:"default_locale" gorm:"type:varchar(16);not null;default:''"` + WebhookURL string `json:"webhook_url" gorm:"type:varchar(512);not null;default:''"` + WebhookSecret string `json:"-" gorm:"type:varchar(128);not null;default:''"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + DeletedAt gorm.DeletedAt `json:"deleted_at" gorm:"index"` +} + +func (EmbedChannel) TableName() string { return "embed_channels" } + +func (ch *EmbedChannel) BeforeCreate(tx *gorm.DB) error { + if ch.ID == "" { + ch.ID = uuid.New().String() + } + if ch.AgentID == "" { + ch.AgentID = BuiltinQuickAnswerID + } + if ch.RateLimitPerMinute <= 0 { + ch.RateLimitPerMinute = 30 + } + if ch.RateLimitPerDay <= 0 { + ch.RateLimitPerDay = DefaultEmbedRateLimitPerDay + } + if ch.WidgetPosition == "" { + ch.WidgetPosition = DefaultEmbedWidgetPosition + } + if ch.HeaderTitleMode == "" { + ch.HeaderTitleMode = DefaultEmbedHeaderTitleMode + } + return nil +} + +// DefaultEmbedRateLimitPerDay caps total embed requests per channel per day +// (across all client IPs), bounding cost/abuse when the publicly visible +// publish token is copied and replayed from rotating IPs. +const DefaultEmbedRateLimitPerDay = 10000 + +const DefaultEmbedWidgetPosition = "bottom-right" +const DefaultEmbedHeaderTitleMode = "channel" +const EmbedHeaderTitleModeSession = "session" + +// NormalizeEmbedWidgetPosition returns a supported widget corner or the default. +func NormalizeEmbedWidgetPosition(position string) string { + switch strings.TrimSpace(position) { + case "bottom-left", "top-right", "top-left", "bottom-right": + return strings.TrimSpace(position) + default: + return DefaultEmbedWidgetPosition + } +} + +// NormalizeEmbedHeaderTitleMode returns a supported header title source. +func NormalizeEmbedHeaderTitleMode(mode string) string { + switch strings.TrimSpace(mode) { + case EmbedHeaderTitleModeSession: + return EmbedHeaderTitleModeSession + default: + return DefaultEmbedHeaderTitleMode + } +} + +// AllowedOriginsList decodes the JSON array of origin patterns. +func (ch *EmbedChannel) AllowedOriginsList() []string { + if len(ch.AllowedOrigins) == 0 { + return nil + } + var origins []string + if err := json.Unmarshal(ch.AllowedOrigins, &origins); err != nil { + return nil + } + return origins +} + +// EmbedChannelPublicConfig is returned to anonymous embed clients (no secrets). +type EmbedChannelPublicConfig struct { + ChannelID string `json:"channel_id"` + Name string `json:"name"` + DisplayTitle string `json:"display_title"` + KnowledgeBaseIDs []string `json:"knowledge_base_ids,omitempty"` + AgentID string `json:"agent_id"` + AgentName string `json:"agent_name,omitempty"` + AgentAvatar string `json:"agent_avatar,omitempty"` + WelcomeMessage string `json:"welcome_message"` + PrimaryColor string `json:"primary_color,omitempty"` + PageTitle string `json:"page_title,omitempty"` + HeaderTitleMode string `json:"header_title_mode,omitempty"` + ShowSuggestedQuestions bool `json:"show_suggested_questions"` + AllowedOrigins []string `json:"allowed_origins,omitempty"` + WidgetPosition string `json:"widget_position,omitempty"` + AllowWebSearch bool `json:"allow_web_search"` + AllowMemory bool `json:"allow_memory"` + AllowFileUpload bool `json:"allow_file_upload"` + // AgentWebSearchEnabled reflects whether the bound agent has web search configured. + AgentWebSearchEnabled bool `json:"agent_web_search_enabled"` + // AgentImageUploadEnabled reflects whether the bound agent supports image upload. + AgentImageUploadEnabled bool `json:"agent_image_upload_enabled"` + DefaultLocale string `json:"default_locale,omitempty"` +} + +// Supported embed UI locales. +var supportedEmbedLocales = map[string]struct{}{ + "zh-CN": {}, + "en-US": {}, + "ko-KR": {}, + "ru-RU": {}, +} + +// NormalizeEmbedDefaultLocale returns a supported locale tag or empty string +// (meaning follow browser / host widget locale). +func NormalizeEmbedDefaultLocale(locale string) string { + locale = strings.TrimSpace(locale) + if _, ok := supportedEmbedLocales[locale]; ok { + return locale + } + return "" +} + +// EmbedSessionMarkerPrefix tags sessions created through an embed channel. +const EmbedSessionMarkerPrefix = "embed_channel:" diff --git a/internal/types/embed_channel_test.go b/internal/types/embed_channel_test.go new file mode 100644 index 00000000..4a541c27 --- /dev/null +++ b/internal/types/embed_channel_test.go @@ -0,0 +1,40 @@ +package types + +import "testing" + +func TestNormalizeEmbedWidgetPosition(t *testing.T) { + cases := map[string]string{ + "bottom-right": DefaultEmbedWidgetPosition, + "bottom-left": "bottom-left", + "top-right": "top-right", + "top-left": "top-left", + "": DefaultEmbedWidgetPosition, + "center": DefaultEmbedWidgetPosition, + " bottom-left ": "bottom-left", + "BOTTOM-RIGHT": DefaultEmbedWidgetPosition, + "bottom-right-extra": DefaultEmbedWidgetPosition, + "middle": DefaultEmbedWidgetPosition, + "top-center": DefaultEmbedWidgetPosition, + " ": DefaultEmbedWidgetPosition, + } + for in, want := range cases { + if got := NormalizeEmbedWidgetPosition(in); got != want { + t.Fatalf("NormalizeEmbedWidgetPosition(%q) = %q, want %q", in, got, want) + } + } +} + +func TestNormalizeEmbedHeaderTitleMode(t *testing.T) { + cases := map[string]string{ + "": DefaultEmbedHeaderTitleMode, + "channel": DefaultEmbedHeaderTitleMode, + "session": EmbedHeaderTitleModeSession, + "SESSION": DefaultEmbedHeaderTitleMode, + "other": DefaultEmbedHeaderTitleMode, + } + for in, want := range cases { + if got := NormalizeEmbedHeaderTitleMode(in); got != want { + t.Fatalf("NormalizeEmbedHeaderTitleMode(%q) = %q, want %q", in, got, want) + } + } +} diff --git a/internal/types/interfaces/embed_channel.go b/internal/types/interfaces/embed_channel.go new file mode 100644 index 00000000..ee7de84d --- /dev/null +++ b/internal/types/interfaces/embed_channel.go @@ -0,0 +1,38 @@ +package interfaces + +import ( + "context" + + "github.com/Tencent/WeKnora/internal/types" +) + +// EmbedChannelRepository persists web embed channel rows. +type EmbedChannelRepository interface { + Create(ctx context.Context, ch *types.EmbedChannel) error + GetByID(ctx context.Context, id string) (*types.EmbedChannel, error) + GetByPublishToken(ctx context.Context, token string) (*types.EmbedChannel, error) + ListByAgent(ctx context.Context, tenantID uint64, agentID string) ([]*types.EmbedChannel, error) + ListByTenant(ctx context.Context, tenantID uint64) ([]*types.EmbedChannel, error) + Update(ctx context.Context, ch *types.EmbedChannel) error + Delete(ctx context.Context, tenantID uint64, id string) error +} + +// EmbedChannelService manages web embed channel lifecycle. +type EmbedChannelService interface { + Create(ctx context.Context, tenantID uint64, agentID string, req *types.EmbedChannel) (*types.EmbedChannel, string, error) + ListByAgent(ctx context.Context, tenantID uint64, agentID string) ([]*types.EmbedChannel, error) + ListByTenant(ctx context.Context, tenantID uint64) ([]*types.EmbedChannel, error) + Update(ctx context.Context, tenantID uint64, id string, req *types.EmbedChannel, enabled *bool, showSuggested *bool, allowWebSearch *bool, allowMemory *bool, allowFileUpload *bool, defaultLocale *string, webhookURL *string, webhookSecret *string) (*types.EmbedChannel, error) + Delete(ctx context.Context, tenantID uint64, id string) error + RotateToken(ctx context.Context, tenantID uint64, id string) (*types.EmbedChannel, string, error) + LookupForEmbed(ctx context.Context, channelID, token string) (*types.EmbedChannel, error) + LookupEnabledChannel(ctx context.Context, channelID string) (*types.EmbedChannel, error) + IssueSessionToken(ctx context.Context, channelID string) (sessionToken string, expiresIn int, err error) + IssuePreviewSession(ctx context.Context, tenantID uint64, channelID string) (sessionToken string, expiresIn int, err error) + ResolveSessionToken(ctx context.Context, token string) (channelID string, err error) + PublicConfig(ctx context.Context, ch *types.EmbedChannel) types.EmbedChannelPublicConfig + SuggestedQuestions(ctx context.Context, ch *types.EmbedChannel, limit int) ([]types.SuggestedQuestion, error) + EmbedChunk(ctx context.Context, ch *types.EmbedChannel, chunkID string) (*types.Chunk, error) + EmbedDisplayTitle(ctx context.Context, ch *types.EmbedChannel) string + GetOwnedChannel(ctx context.Context, tenantID uint64, id string) (*types.EmbedChannel, error) +} diff --git a/internal/types/session.go b/internal/types/session.go index 33a396ab..1dcb788c 100644 --- a/internal/types/session.go +++ b/internal/types/session.go @@ -130,7 +130,8 @@ func (s *Session) BeforeCreate(tx *gorm.DB) (err error) { // SessionListQuery bundles the parameters for listing sessions. // UserID empty means "tenant-wide" (used by API-key callers / legacy rows). // Keyword matches title ILIKE '%keyword%'. -// Source values: "web" (no IM mapping) or an IM platform name (e.g. "feishu", "wechat"). +// Source values: "web" (user chats, no IM/embed), "embed" / "embed:{channelID}", +// or an IM platform name (e.g. "feishu", "wechat"). // AgentID currently only filters sessions that have an IM channel mapping. type SessionListQuery struct { TenantID uint64 diff --git a/migrations/sqlite/000000_init.up.sql b/migrations/sqlite/000000_init.up.sql index f06b4f91..10cfadbe 100644 --- a/migrations/sqlite/000000_init.up.sql +++ b/migrations/sqlite/000000_init.up.sql @@ -598,6 +598,39 @@ CREATE UNIQUE INDEX IF NOT EXISTS idx_im_channels_bot_identity ON im_channels (bot_identity) WHERE deleted_at IS NULL AND bot_identity != ''; +CREATE TABLE IF NOT EXISTS embed_channels ( + id VARCHAR(36) PRIMARY KEY, + tenant_id INTEGER NOT NULL, + agent_id VARCHAR(36) NOT NULL DEFAULT 'builtin-quick-answer', + name VARCHAR(255) NOT NULL DEFAULT '', + enabled INTEGER NOT NULL DEFAULT 1, + publish_token VARCHAR(64) NOT NULL DEFAULT '', + allowed_origins TEXT NOT NULL DEFAULT '[]', + welcome_message TEXT NOT NULL DEFAULT '', + rate_limit_per_minute INTEGER NOT NULL DEFAULT 30, + rate_limit_per_day INTEGER NOT NULL DEFAULT 10000, + primary_color VARCHAR(32) NOT NULL DEFAULT '', + page_title VARCHAR(255) NOT NULL DEFAULT '', + header_title_mode VARCHAR(32) NOT NULL DEFAULT 'channel', + show_suggested_questions INTEGER NOT NULL DEFAULT 1, + widget_position VARCHAR(32) NOT NULL DEFAULT 'bottom-right', + allow_web_search INTEGER NOT NULL DEFAULT 0, + allow_memory INTEGER NOT NULL DEFAULT 0, + allow_file_upload INTEGER NOT NULL DEFAULT 0, + default_locale VARCHAR(16) NOT NULL DEFAULT '', + webhook_url VARCHAR(512) NOT NULL DEFAULT '', + webhook_secret VARCHAR(128) NOT NULL DEFAULT '', + created_at DATETIME DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME DEFAULT CURRENT_TIMESTAMP, + deleted_at DATETIME +); + +CREATE INDEX IF NOT EXISTS idx_embed_channels_tenant ON embed_channels (tenant_id); +CREATE INDEX IF NOT EXISTS idx_embed_channels_agent ON embed_channels (agent_id); +CREATE UNIQUE INDEX IF NOT EXISTS idx_embed_channels_publish_token + ON embed_channels (publish_token) + WHERE publish_token != '' AND deleted_at IS NULL; + CREATE TABLE IF NOT EXISTS data_sources ( id VARCHAR(36) NOT NULL PRIMARY KEY, tenant_id INTEGER NOT NULL, diff --git a/migrations/versioned/000060_embed_channels.down.sql b/migrations/versioned/000060_embed_channels.down.sql new file mode 100644 index 00000000..32cea1de --- /dev/null +++ b/migrations/versioned/000060_embed_channels.down.sql @@ -0,0 +1,5 @@ +DROP INDEX IF EXISTS idx_embed_channels_deleted; +DROP INDEX IF EXISTS idx_embed_channels_publish_token; +DROP INDEX IF EXISTS idx_embed_channels_agent; +DROP INDEX IF EXISTS idx_embed_channels_tenant; +DROP TABLE IF EXISTS embed_channels; diff --git a/migrations/versioned/000060_embed_channels.up.sql b/migrations/versioned/000060_embed_channels.up.sql new file mode 100644 index 00000000..3470d713 --- /dev/null +++ b/migrations/versioned/000060_embed_channels.up.sql @@ -0,0 +1,63 @@ +-- Migration: 000060_embed_channels +-- Web embed channels for publishing agent chat to external sites. +DO $$ BEGIN RAISE NOTICE '[Migration 000060] Creating embed_channels table'; END $$; + +CREATE TABLE IF NOT EXISTS embed_channels ( + id VARCHAR(36) PRIMARY KEY DEFAULT uuid_generate_v4(), + tenant_id BIGINT NOT NULL, + agent_id VARCHAR(36) NOT NULL DEFAULT 'builtin-quick-answer', + name VARCHAR(255) NOT NULL DEFAULT '', + enabled BOOLEAN NOT NULL DEFAULT true, + publish_token VARCHAR(64) NOT NULL DEFAULT '', + allowed_origins JSONB NOT NULL DEFAULT '[]', + welcome_message TEXT NOT NULL DEFAULT '', + rate_limit_per_minute INTEGER NOT NULL DEFAULT 30, + rate_limit_per_day INTEGER NOT NULL DEFAULT 10000, + primary_color VARCHAR(32) NOT NULL DEFAULT '', + page_title VARCHAR(255) NOT NULL DEFAULT '', + header_title_mode VARCHAR(32) NOT NULL DEFAULT 'channel', + show_suggested_questions BOOLEAN NOT NULL DEFAULT true, + widget_position VARCHAR(32) NOT NULL DEFAULT 'bottom-right', + allow_web_search BOOLEAN NOT NULL DEFAULT false, + allow_memory BOOLEAN NOT NULL DEFAULT false, + allow_file_upload BOOLEAN NOT NULL DEFAULT false, + default_locale VARCHAR(16) NOT NULL DEFAULT '', + webhook_url VARCHAR(512) NOT NULL DEFAULT '', + webhook_secret VARCHAR(128) NOT NULL DEFAULT '', + created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + deleted_at TIMESTAMP WITH TIME ZONE +); + +CREATE INDEX IF NOT EXISTS idx_embed_channels_tenant ON embed_channels (tenant_id); +CREATE INDEX IF NOT EXISTS idx_embed_channels_agent ON embed_channels (agent_id); +CREATE UNIQUE INDEX IF NOT EXISTS idx_embed_channels_publish_token + ON embed_channels (publish_token) + WHERE publish_token <> '' AND deleted_at IS NULL; +CREATE INDEX IF NOT EXISTS idx_embed_channels_deleted ON embed_channels (deleted_at) WHERE deleted_at IS NOT NULL; + +COMMENT ON TABLE embed_channels IS 'Web embed channels for publishing agent chat to external sites via iframe'; +COMMENT ON COLUMN embed_channels.publish_token IS 'Plaintext scoped token (em_ prefix); rotatable from management UI'; +COMMENT ON COLUMN embed_channels.allowed_origins IS 'JSON array of allowed HTTP(S) origins for embed API requests; empty rejects all (management UI requires at least one)'; +COMMENT ON COLUMN embed_channels.rate_limit_per_minute IS 'Per-IP per-minute request cap for the public embed endpoints'; +COMMENT ON COLUMN embed_channels.rate_limit_per_day IS + 'Channel-wide daily total request cap across all IPs; bounds cost/abuse since the publish token is publicly visible'; +COMMENT ON COLUMN embed_channels.primary_color IS 'CSS color for embed widget accent (e.g. #0052d9)'; +COMMENT ON COLUMN embed_channels.page_title IS 'Browser tab title for the embed page'; +COMMENT ON COLUMN embed_channels.header_title_mode IS + 'Embed header title source: channel (fixed page title) or session (auto-generated after first message)'; +COMMENT ON COLUMN embed_channels.show_suggested_questions IS + 'When true, embed chat shows suggested starter questions before the first visitor message'; +COMMENT ON COLUMN embed_channels.widget_position IS 'Floating widget corner: bottom-right, bottom-left, top-right, top-left'; +COMMENT ON COLUMN embed_channels.allow_web_search IS + 'When true, embed chat may show web search toggle; visitor must opt in per message'; +COMMENT ON COLUMN embed_channels.allow_memory IS + 'When true, embed chat may use agent memory; client cannot override when false'; +COMMENT ON COLUMN embed_channels.allow_file_upload IS + 'When true, embed chat may send images; client cannot override when false'; +COMMENT ON COLUMN embed_channels.default_locale IS + 'Default visitor UI locale (zh-CN, en-US, ko-KR, ru-RU); empty follows browser'; +COMMENT ON COLUMN embed_channels.webhook_url IS 'HTTPS endpoint for outbound message_sent / message_received events'; +COMMENT ON COLUMN embed_channels.webhook_secret IS 'Optional HMAC-SHA256 secret for X-WeKnora-Signature header'; + +DO $$ BEGIN RAISE NOTICE '[Migration 000060] embed_channels created'; END $$;