FranklinWH-Toolbox/service/utils/agentserver/guard.go

package agentserver

import (
	"context"
	"errors"
)

// OwnershipStore 会话归属与子资源索引的存储（Mongo 实现见 services/agentServices.go）
type OwnershipStore interface {
	// SessionOwner 返回未删除会话的所属用户；不存在或已删除时 found=false
	SessionOwner(ctx context.Context, sessionId string) (owner string, found bool, err error)
	// ResourceSession 返回子资源所属的会话
	ResourceSession(ctx context.Context, resourceId string) (sessionId string, found bool, err error)
	// RecentSessionIds 用户最近活跃的未删除会话（反查交互、文件时遍历）
	RecentSessionIds(ctx context.Context, userId string, limit int) ([]string, error)
}

// ResourceSink 记录子资源索引（异步）
type ResourceSink interface {
	Record(rs []Resource)
}

// Upstream 归属反查用到的 AgentServer 接口（*Client 实现）
type Upstream interface {
	GetRun(ctx context.Context, runId string) (*RunDto, error)
	ListFiles(ctx context.Context, sessionId string) ([]FileDto, error)
	PendingInteractions(ctx context.Context, sessionId string) ([]InteractionRefDto, error)
}

// fallbackSessions 反查交互、文件时最多遍历的最近会话数
const fallbackSessions = 10

// Guard 代理请求的归属校验（technical-design §3.5 第 2 点）
type Guard struct {
	Store    OwnershipStore
	Upstream Upstream
	Sink     ResourceSink
}

// Check 校验 route 指向的资源属于 userId。
// 返回资源所属的会话 ID（OwnerNone 时为空）；ok=false 表示不属于本人或不存在（调用方统一返回 404）。
// err 只在存储或上游出错时返回。
func (g *Guard) Check(ctx context.Context, userId string, r Route) (sessionId string, ok bool, err error) {
	switch r.Owner {
	case OwnerNone:
		return "", true, nil
	case OwnerSession, OwnerQuerySession:
		owned, err := g.ownsSession(ctx, userId, r.Id)
		if err != nil || !owned {
			return "", false, err
		}
		return r.Id, true, nil
	case OwnerRun, OwnerInteraction, OwnerFile:
		return g.checkResource(ctx, userId, r.Owner, r.Id)
	}
	return "", false, nil
}

func (g *Guard) ownsSession(ctx context.Context, userId, sessionId string) (bool, error) {
	owner, found, err := g.Store.SessionOwner(ctx, sessionId)
	if err != nil {
		return false, err
	}
	return found && owner == userId, nil
}

func (g *Guard) checkResource(ctx context.Context, userId string, kind OwnerKind, id string) (string, bool, error) {
	sid, found, err := g.Store.ResourceSession(ctx, id)
	if err != nil {
		return "", false, err
	}
	if found {
		owned, err := g.ownsSession(ctx, userId, sid)
		if err != nil || !owned {
			return "", false, err
		}
		return sid, true, nil
	}

	// 索引中没有：向 AgentServer 反查，确认后补写索引
	sid, err = g.lookup(ctx, userId, kind, id)
	if err != nil || sid == "" {
		return "", false, err
	}
	if g.Sink != nil {
		g.Sink.Record([]Resource{{Id: id, SessionId: sid}})
	}
	return sid, true, nil
}

// lookup 反查子资源所属会话，返回空表示不属于本人
func (g *Guard) lookup(ctx context.Context, userId string, kind OwnerKind, id string) (string, error) {
	if g.Upstream == nil {
		return "", nil
	}
	if kind == OwnerRun {
		run, err := g.Upstream.GetRun(ctx, id)
		if errors.Is(err, ErrNotFound) {
			return "", nil
		}
		if err != nil {
			return "", err
		}
		owned, err := g.ownsSession(ctx, userId, run.SessionId)
		if err != nil || !owned {
			return "", err
		}
		return run.SessionId, nil
	}

	sessions, err := g.Store.RecentSessionIds(ctx, userId, fallbackSessions)
	if err != nil {
		return "", err
	}
	for _, sid := range sessions {
		hit, err := g.sessionHas(ctx, kind, sid, id)
		if errors.Is(err, ErrNotFound) {
			continue
		}
		if err != nil {
			return "", err
		}
		if hit {
			return sid, nil
		}
	}
	return "", nil
}

func (g *Guard) sessionHas(ctx context.Context, kind OwnerKind, sessionId, id string) (bool, error) {
	switch kind {
	case OwnerFile:
		files, err := g.Upstream.ListFiles(ctx, sessionId)
		if err != nil {
			return false, err
		}
		for _, f := range files {
			if f.Id == id {
				return true, nil
			}
		}
	case OwnerInteraction:
		items, err := g.Upstream.PendingInteractions(ctx, sessionId)
		if err != nil {
			return false, err
		}
		for _, it := range items {
			if it.InteractionId == id {
				return true, nil
			}
		}
	}
	return false, nil
}