xinmi/internal/router/router_files_test.go

package router

import (
	"context"
	"errors"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/gin-gonic/gin"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
)

var _ interfaces.FileService = (*stubFileService)(nil)

type stubFileService struct {
	getFile func(ctx context.Context, filePath string) (io.ReadCloser, error)
}

type stubResourceCatalog struct {
	resource *types.StoredResource
}

type authorizingFileCatalog struct {
	stubResourceCatalog
	authorize func(context.Context, string, interfaces.ResourceFileScope) error
}

func (s *authorizingFileCatalog) AuthorizeFileRead(ctx context.Context, value string, scope interfaces.ResourceFileScope) error {
	return s.authorize(ctx, value, scope)
}

func TestFileProxyAuthorizesBeforeOpeningStorage(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")
	const ref = "resource://AbCdEfGhIjKlMnOpQrStUv"
	opened := false
	catalog := &authorizingFileCatalog{
		stubResourceCatalog: stubResourceCatalog{resource: &types.StoredResource{TenantID: 42, PhysicalPath: "local://42/exports/a.png"}},
		authorize: func(ctx context.Context, value string, scope interfaces.ResourceFileScope) error {
			if scope.TenantID != 42 || value != ref {
				t.Fatal("wrong authorization scope")
			}
			return errors.New("private file owned by another user")
		},
	}
	engine := gin.New()
	serveFilesWithResources(engine, &stubFileService{getFile: func(context.Context, string) (io.ReadCloser, error) {
		opened = true
		return io.NopCloser(strings.NewReader("private")), nil
	}}, nil, catalog)
	req := httptest.NewRequest(http.MethodGet, "/files?file_path="+url.QueryEscape(ref), nil)
	req = req.WithContext(context.WithValue(req.Context(), types.TenantInfoContextKey, &types.Tenant{ID: 42}))
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)
	if recorder.Code != http.StatusForbidden || opened {
		t.Fatalf("status=%d storage_opened=%v", recorder.Code, opened)
	}
	catAllowed := false
	catalog.authorize = func(context.Context, string, interfaces.ResourceFileScope) error { catAllowed = true; return nil }
	recorder = httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)
	if recorder.Code != http.StatusOK || !catAllowed || recorder.Header().Get("Cache-Control") != "private, no-store" {
		t.Fatalf("authorized status=%d cache=%s", recorder.Code, recorder.Header().Get("Cache-Control"))
	}
}

type stubMessageFileLookup struct {
	get func(ctx context.Context, sessionID, messageID string) (*types.Message, error)
}

func (s *stubMessageFileLookup) GetMessage(
	ctx context.Context,
	sessionID, messageID string,
) (*types.Message, error) {
	return s.get(ctx, sessionID, messageID)
}

type stubSharedAgentFileLookup struct {
	get func(
		ctx context.Context,
		tenantID uint64,
		callerTenantRole types.TenantRole,
		agentID string,
		sourceTenantID ...uint64,
	) (*types.CustomAgent, error)
}

func (s *stubSharedAgentFileLookup) GetSharedAgentForTenant(
	ctx context.Context,
	tenantID uint64,
	callerTenantRole types.TenantRole,
	agentID string,
	sourceTenantID ...uint64,
) (*types.CustomAgent, error) {
	return s.get(ctx, tenantID, callerTenantRole, agentID, sourceTenantID...)
}

func (s *stubResourceCatalog) Register(
	context.Context,
	uint64,
	string,
	interfaces.ResourceRegistration,
) (string, error) {
	panic("unexpected Register")
}

func (s *stubResourceCatalog) Resolve(context.Context, string) (*types.StoredResource, error) {
	return s.resource, nil
}

func (s *stubResourceCatalog) ResolvePath(_ context.Context, value string) (string, *types.StoredResource, error) {
	if _, ok := types.ParseResourcePath(value); ok && s.resource != nil {
		return s.resource.PhysicalPath, s.resource, nil
	}
	return value, nil, nil
}

func (s *stubResourceCatalog) Bind(context.Context, string, string, string, string) error {
	panic("unexpected Bind")
}

func (s *stubResourceCatalog) Release(context.Context, string, string, string) (int64, error) {
	panic("unexpected Release")
}

func (s *stubResourceCatalog) MarkDeleted(context.Context, string) error {
	panic("unexpected MarkDeleted")
}

func (s *stubResourceCatalog) CreateAccessGrant(context.Context, string, time.Duration) (string, error) {
	panic("unexpected CreateAccessGrant")
}

func (s *stubResourceCatalog) ResolveAccessGrant(context.Context, string) (*types.StoredResource, error) {
	return s.resource, nil
}

func (s *stubFileService) CheckConnectivity(ctx context.Context) error {
	return nil
}

func (s *stubFileService) SaveFile(ctx context.Context, file *multipart.FileHeader, tenantID uint64, knowledgeID string) (string, error) {
	panic("unexpected call to SaveFile")
}

func (s *stubFileService) SaveBytes(ctx context.Context, data []byte, tenantID uint64, fileName string, temp bool) (string, error) {
	panic("unexpected call to SaveBytes")
}

func (s *stubFileService) GetFile(ctx context.Context, filePath string) (io.ReadCloser, error) {
	if s.getFile == nil {
		panic("unexpected call to GetFile")
	}
	return s.getFile(ctx, filePath)
}

func (s *stubFileService) GetFileURL(ctx context.Context, filePath string) (string, error) {
	panic("unexpected call to GetFileURL")
}

func (s *stubFileService) DeleteFile(ctx context.Context, filePath string) error {
	panic("unexpected call to DeleteFile")
}

func (s *stubFileService) CopyFile(ctx context.Context, srcPath string, tenantID uint64, knowledgeID string) (string, error) {
	panic("unexpected call to CopyFile")
}

func TestServeFilesFallsBackToGlobalFileService(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")

	engine := gin.New()
	var requestedPath string
	serveFiles(engine, &stubFileService{
		getFile: func(ctx context.Context, filePath string) (io.ReadCloser, error) {
			requestedPath = filePath
			return io.NopCloser(strings.NewReader("fallback-body")), nil
		},
	})

	filePath := "local://42/docs/example.txt"
	req := httptest.NewRequest(http.MethodGet, "/files?file_path="+url.QueryEscape(filePath), nil)
	req = req.WithContext(context.WithValue(req.Context(), types.TenantInfoContextKey, &types.Tenant{ID: 42}))

	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if got, want := recorder.Code, http.StatusOK; got != want {
		t.Fatalf("status = %d, want %d", got, want)
	}
	if requestedPath != filePath {
		t.Fatalf("requested path = %q, want %q", requestedPath, filePath)
	}
	if body := recorder.Body.String(); body != "fallback-body" {
		t.Fatalf("body = %q, want %q", body, "fallback-body")
	}
}

func TestServeFilesResolvesShortResourceReference(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")
	const ref = "resource://AbCdEfGhIjKlMnOpQrStUv"
	const physical = "local://42/exports/a.png"

	engine := gin.New()
	var requestedPath string
	serveFilesWithResources(engine, &stubFileService{getFile: func(_ context.Context, path string) (io.ReadCloser, error) {
		requestedPath = path
		return io.NopCloser(strings.NewReader("image")), nil
	}}, nil, &stubResourceCatalog{resource: &types.StoredResource{TenantID: 42, PhysicalPath: physical}})

	req := httptest.NewRequest(http.MethodGet, "/files?file_path="+url.QueryEscape(ref), nil)
	req = req.WithContext(context.WithValue(req.Context(), types.TenantInfoContextKey, &types.Tenant{ID: 42}))
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String())
	}
	if requestedPath != physical {
		t.Fatalf("requested path = %q, want %q", requestedPath, physical)
	}
}

func TestServeFilesRejectsCrossTenantResourceReference(t *testing.T) {
	gin.SetMode(gin.TestMode)
	const ref = "resource://AbCdEfGhIjKlMnOpQrStUv"
	engine := gin.New()
	serveFilesWithResources(engine, &stubFileService{getFile: func(context.Context, string) (io.ReadCloser, error) {
		t.Fatal("GetFile should not be called")
		return nil, nil
	}}, nil, &stubResourceCatalog{resource: &types.StoredResource{TenantID: 7, PhysicalPath: "local://7/exports/a.png"}})

	req := httptest.NewRequest(http.MethodGet, "/files?file_path="+url.QueryEscape(ref), nil)
	req = req.WithContext(context.WithValue(req.Context(), types.TenantInfoContextKey, &types.Tenant{ID: 42}))
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if recorder.Code != http.StatusForbidden {
		t.Fatalf("status = %d, want %d", recorder.Code, http.StatusForbidden)
	}
}

func TestResourceGrantServesShortPublicURL(t *testing.T) {
	gin.SetMode(gin.TestMode)
	physical := "local://42/exports/a.png"
	engine := gin.New()
	serveResourceGrants(
		engine,
		&stubResourceCatalog{resource: &types.StoredResource{
			ID:           "resource-1",
			TenantID:     42,
			PhysicalPath: physical,
			OriginalName: "a.png",
			MimeType:     "image/png",
		}},
		&stubTenantService{get: func(_ context.Context, id uint64) (*types.Tenant, error) {
			return &types.Tenant{ID: id}, nil
		}},
		&stubFileService{getFile: func(_ context.Context, path string) (io.ReadCloser, error) {
			if path != physical {
				t.Fatalf("path = %q, want %q", path, physical)
			}
			return io.NopCloser(strings.NewReader("image")), nil
		}},
		nil,
	)

	req := httptest.NewRequest(http.MethodGet, "/r/GrantTokenAbCdEfGhIjKlM", nil)
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)
	if recorder.Code != http.StatusOK || recorder.Body.String() != "image" {
		t.Fatalf("status=%d body=%q", recorder.Code, recorder.Body.String())
	}
	if got := recorder.Header().Get("X-Content-Type-Options"); got != "nosniff" {
		t.Fatalf("X-Content-Type-Options = %q", got)
	}
}

func TestServeFilesDoesNotFallbackWhenProviderDoesNotMatchGlobalStorage(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "minio")

	engine := gin.New()
	serveFiles(engine, &stubFileService{
		getFile: func(ctx context.Context, filePath string) (io.ReadCloser, error) {
			t.Fatalf("GetFile should not be called for mismatched provider, got %q", filePath)
			return nil, nil
		},
	})

	req := httptest.NewRequest(http.MethodGet, "/files?file_path="+url.QueryEscape("local://42/docs/example.txt"), nil)
	req = req.WithContext(context.WithValue(req.Context(), types.TenantInfoContextKey, &types.Tenant{ID: 42}))

	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if got, want := recorder.Code, http.StatusBadRequest; got != want {
		t.Fatalf("status = %d, want %d", got, want)
	}
}

func TestServeFilesRejectsCrossTenantPath(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")

	engine := gin.New()
	serveFiles(engine, &stubFileService{
		getFile: func(ctx context.Context, filePath string) (io.ReadCloser, error) {
			t.Fatalf("GetFile should not be called for cross-tenant path, got %q", filePath)
			return nil, nil
		},
	})

	req := httptest.NewRequest(http.MethodGet, "/files?file_path="+url.QueryEscape("local://7/knowledge/secret.pdf"), nil)
	req = req.WithContext(context.WithValue(req.Context(), types.TenantInfoContextKey, &types.Tenant{ID: 42}))

	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if got, want := recorder.Code, http.StatusForbidden; got != want {
		t.Fatalf("status = %d, want %d", got, want)
	}
}

func TestServeFilesRejectsPathWithoutTenantSegment(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")

	engine := gin.New()
	serveFiles(engine, &stubFileService{
		getFile: func(ctx context.Context, filePath string) (io.ReadCloser, error) {
			t.Fatalf("GetFile should not be called without tenant segment, got %q", filePath)
			return nil, nil
		},
	})

	req := httptest.NewRequest(http.MethodGet, "/files?file_path="+url.QueryEscape("local://docs/example.txt"), nil)
	req = req.WithContext(context.WithValue(req.Context(), types.TenantInfoContextKey, &types.Tenant{ID: 42}))

	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if got, want := recorder.Code, http.StatusForbidden; got != want {
		t.Fatalf("status = %d, want %d", got, want)
	}
}

// /files carries its own API-key guard (middleware.AllowFileServeAPIKey):
// full-access and tenant-wide retrieve keys may serve tenant-bounded paths,
// but KB-restricted keys (and keys lacking retrieve) are denied because a raw
// storage path cannot be bounded to a KB allow-list.
func TestServeFilesAPIKeyScopeMatrix(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")

	const filePath = "local://42/docs/example.txt"

	cases := []struct {
		name     string
		scope    types.TenantAPIKeyScope
		wantCode int
	}{
		{
			name:     "full access allowed",
			scope:    types.TenantAPIKeyScope{FullAccess: true},
			wantCode: http.StatusOK,
		},
		{
			name: "tenant-wide retrieve allowed",
			scope: types.TenantAPIKeyScope{
				Capabilities: types.StringArray{string(types.APIKeyCapabilityRetrieve)},
			},
			wantCode: http.StatusOK,
		},
		{
			name: "kb-restricted retrieve denied",
			scope: types.TenantAPIKeyScope{
				KnowledgeBaseIDs: types.StringArray{"kb-1"},
				Capabilities:     types.StringArray{string(types.APIKeyCapabilityRetrieve)},
			},
			wantCode: http.StatusForbidden,
		},
		{
			name: "non-retrieve capability denied",
			scope: types.TenantAPIKeyScope{
				Capabilities: types.StringArray{string(types.APIKeyCapabilityChat)},
			},
			wantCode: http.StatusForbidden,
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			engine := gin.New()
			serveFiles(engine, &stubFileService{
				getFile: func(_ context.Context, _ string) (io.ReadCloser, error) {
					return io.NopCloser(strings.NewReader("body")), nil
				},
			})

			req := httptest.NewRequest(http.MethodGet, "/files?file_path="+url.QueryEscape(filePath), nil)
			ctx := context.WithValue(req.Context(), types.TenantInfoContextKey, &types.Tenant{ID: 42})
			ctx = types.WithTenantAPIKeyScope(ctx, tc.scope)
			req = req.WithContext(ctx)

			recorder := httptest.NewRecorder()
			engine.ServeHTTP(recorder, req)

			if got := recorder.Code; got != tc.wantCode {
				t.Fatalf("status = %d, want %d body=%s", got, tc.wantCode, recorder.Body.String())
			}
		})
	}
}

// newKBScopedFilesTestEngine wires newKBScopedFileServeHandler behind a
// middleware that injects effectiveTenantID into the request context, mirroring
// what RequireKBAccess does after resolving an org-shared KB to its source
// tenant. This lets the handler be exercised without the full RBAC stack.
func newKBScopedFilesTestEngine(
	effectiveTenantID uint64,
	tenantSvc interfaces.TenantService,
	global interfaces.FileService,
) *gin.Engine {
	engine := gin.New()
	engine.GET("/knowledge-bases/:id/files",
		func(c *gin.Context) {
			ctx := context.WithValue(c.Request.Context(), types.TenantIDContextKey, effectiveTenantID)
			c.Request = c.Request.WithContext(ctx)
			c.Next()
		},
		newKBScopedFileServeHandler(tenantSvc, global),
	)
	return engine
}

// A tenant whose owner-tenant (10008) storage objects are requested by a
// borrowing tenant via a shared KB: the effective tenant in context is the
// owner, so the path validates and the file is served.
func TestKBScopedFilesServesOwnerTenantPath(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")

	const ownerTenantID = uint64(10008)
	var requestedPath string
	engine := newKBScopedFilesTestEngine(
		ownerTenantID,
		&stubTenantService{get: func(_ context.Context, id uint64) (*types.Tenant, error) {
			return &types.Tenant{ID: id}, nil
		}},
		&stubFileService{getFile: func(_ context.Context, filePath string) (io.ReadCloser, error) {
			requestedPath = filePath
			return io.NopCloser(strings.NewReader("shared-body")), nil
		}},
	)

	filePath := "local://10008/exports/img.jpg"
	req := httptest.NewRequest(http.MethodGet, "/knowledge-bases/kb-1/files?file_path="+url.QueryEscape(filePath), nil)
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if got, want := recorder.Code, http.StatusOK; got != want {
		t.Fatalf("status = %d, want %d body=%s", got, want, recorder.Body.String())
	}
	if requestedPath != filePath {
		t.Fatalf("requested path = %q, want %q", requestedPath, filePath)
	}
	if body := recorder.Body.String(); body != "shared-body" {
		t.Fatalf("body = %q, want %q", body, "shared-body")
	}
}

// The path must belong to the effective (owner) tenant. A path pointing at a
// different tenant than the resolved KB owner is still rejected, so the guard
// cannot be used to reach arbitrary tenants' files.
func TestKBScopedFilesRejectsPathNotOwnedByKBTenant(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")

	const ownerTenantID = uint64(10008)
	engine := newKBScopedFilesTestEngine(
		ownerTenantID,
		&stubTenantService{get: func(_ context.Context, id uint64) (*types.Tenant, error) {
			t.Fatalf("GetTenantByID should not be called for mismatched path, got %d", id)
			return nil, nil
		}},
		&stubFileService{getFile: func(_ context.Context, filePath string) (io.ReadCloser, error) {
			t.Fatalf("GetFile should not be called for mismatched path, got %q", filePath)
			return nil, nil
		}},
	)

	req := httptest.NewRequest(http.MethodGet,
		"/knowledge-bases/kb-1/files?file_path="+url.QueryEscape("local://9999/exports/other.jpg"), nil)
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if got, want := recorder.Code, http.StatusForbidden; got != want {
		t.Fatalf("status = %d, want %d", got, want)
	}
}

// KB-scoped proxy is for embedded exports/ images only; raw knowledge uploads
// must use /knowledge/:id/download even when the tenant matches.
func TestKBScopedFilesRejectsNonExportsPath(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")

	const ownerTenantID = uint64(10008)
	engine := newKBScopedFilesTestEngine(
		ownerTenantID,
		&stubTenantService{get: func(_ context.Context, id uint64) (*types.Tenant, error) {
			t.Fatalf("GetTenantByID should not be called for non-exports path, got %d", id)
			return nil, nil
		}},
		&stubFileService{getFile: func(_ context.Context, filePath string) (io.ReadCloser, error) {
			t.Fatalf("GetFile should not be called for non-exports path, got %q", filePath)
			return nil, nil
		}},
	)

	req := httptest.NewRequest(http.MethodGet,
		"/knowledge-bases/kb-1/files?file_path="+url.QueryEscape("local://10008/knowledge-id/123.pdf"), nil)
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if got, want := recorder.Code, http.StatusForbidden; got != want {
		t.Fatalf("status = %d, want %d", got, want)
	}
}

func TestKBScopedFilesRequiresFilePath(t *testing.T) {
	gin.SetMode(gin.TestMode)

	engine := newKBScopedFilesTestEngine(
		10008,
		&stubTenantService{},
		&stubFileService{},
	)

	req := httptest.NewRequest(http.MethodGet, "/knowledge-bases/kb-1/files", nil)
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if got, want := recorder.Code, http.StatusBadRequest; got != want {
		t.Fatalf("status = %d, want %d", got, want)
	}
}

func newMessageScopedFilesTestEngine(
	callerTenantID uint64,
	messageService messageFileLookup,
	agentShareService sharedAgentFileLookup,
	tenantService interfaces.TenantService,
	global interfaces.FileService,
	resourceCatalog interfaces.ResourceCatalog,
) *gin.Engine {
	engine := gin.New()
	engine.GET("/sessions/:id/messages/:message_id/files",
		func(c *gin.Context) {
			ctx := context.WithValue(c.Request.Context(), types.TenantIDContextKey, callerTenantID)
			c.Request = c.Request.WithContext(ctx)
			c.Next()
		},
		newMessageScopedFileServeHandler(
			messageService,
			agentShareService,
			tenantService,
			global,
			nil,
			resourceCatalog,
		),
	)
	return engine
}

func TestMessageScopedFilesServesSharedAgentResource(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")

	const (
		callerTenantID = uint64(42)
		ownerTenantID  = uint64(7)
		ref            = "resource://AbCdEfGhIjKlMnOpQrStUv"
		physical       = "local://7/exports/chart.png"
	)
	var requestedPath string
	engine := newMessageScopedFilesTestEngine(
		callerTenantID,
		&stubMessageFileLookup{get: func(_ context.Context, sessionID, messageID string) (*types.Message, error) {
			if sessionID != "session-1" || messageID != "message-1" {
				t.Fatalf("unexpected message scope %s/%s", sessionID, messageID)
			}
			return &types.Message{AgentID: "agent-1", AgentTenantID: ownerTenantID}, nil
		}},
		&stubSharedAgentFileLookup{get: func(
			_ context.Context,
			tenantID uint64,
			_ types.TenantRole,
			agentID string,
			sourceTenantID ...uint64,
		) (*types.CustomAgent, error) {
			if tenantID != callerTenantID || agentID != "agent-1" || len(sourceTenantID) != 1 || sourceTenantID[0] != ownerTenantID {
				t.Fatalf("unexpected shared-agent lookup tenant=%d agent=%s source=%v", tenantID, agentID, sourceTenantID)
			}
			return &types.CustomAgent{ID: agentID, TenantID: ownerTenantID}, nil
		}},
		&stubTenantService{get: func(_ context.Context, id uint64) (*types.Tenant, error) {
			return &types.Tenant{ID: id}, nil
		}},
		&stubFileService{getFile: func(_ context.Context, filePath string) (io.ReadCloser, error) {
			requestedPath = filePath
			return io.NopCloser(strings.NewReader("shared-agent-image")), nil
		}},
		&stubResourceCatalog{resource: &types.StoredResource{
			TenantID:     ownerTenantID,
			PhysicalPath: physical,
			MimeType:     "image/png",
		}},
	)

	req := httptest.NewRequest(http.MethodGet,
		"/sessions/session-1/messages/message-1/files?file_path="+url.QueryEscape(ref), nil)
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if recorder.Code != http.StatusOK || recorder.Body.String() != "shared-agent-image" {
		t.Fatalf("status=%d body=%q", recorder.Code, recorder.Body.String())
	}
	if requestedPath != physical {
		t.Fatalf("requested path = %q, want %q", requestedPath, physical)
	}
}

func TestMessageScopedFilesServesSameTenantResource(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")

	const (
		tenantID = uint64(42)
		ref      = "resource://AbCdEfGhIjKlMnOpQrStUv"
		physical = "local://42/exports/chart.png"
	)
	var requestedPath string
	engine := newMessageScopedFilesTestEngine(
		tenantID,
		&stubMessageFileLookup{get: func(_ context.Context, sessionID, messageID string) (*types.Message, error) {
			if sessionID != "session-1" || messageID != "message-1" {
				t.Fatalf("unexpected message scope %s/%s", sessionID, messageID)
			}
			return &types.Message{AgentTenantID: tenantID}, nil
		}},
		&stubSharedAgentFileLookup{get: func(
			context.Context, uint64, types.TenantRole, string, ...uint64,
		) (*types.CustomAgent, error) {
			t.Fatal("shared-agent lookup should not run for same-tenant resources")
			return nil, nil
		}},
		&stubTenantService{get: func(_ context.Context, id uint64) (*types.Tenant, error) {
			return &types.Tenant{ID: id}, nil
		}},
		&stubFileService{getFile: func(_ context.Context, filePath string) (io.ReadCloser, error) {
			requestedPath = filePath
			return io.NopCloser(strings.NewReader("same-tenant-image")), nil
		}},
		&stubResourceCatalog{resource: &types.StoredResource{
			TenantID:     tenantID,
			PhysicalPath: physical,
			MimeType:     "image/png",
		}},
	)

	req := httptest.NewRequest(http.MethodGet,
		"/sessions/session-1/messages/message-1/files?file_path="+url.QueryEscape(ref), nil)
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if recorder.Code != http.StatusOK || recorder.Body.String() != "same-tenant-image" {
		t.Fatalf("status=%d body=%q", recorder.Code, recorder.Body.String())
	}
	if requestedPath != physical {
		t.Fatalf("requested path = %q, want %q", requestedPath, physical)
	}
}

func TestMessageScopedFilesRequiresFilePath(t *testing.T) {
	gin.SetMode(gin.TestMode)

	engine := newMessageScopedFilesTestEngine(
		42,
		&stubMessageFileLookup{get: func(context.Context, string, string) (*types.Message, error) {
			return &types.Message{AgentTenantID: 42}, nil
		}},
		&stubSharedAgentFileLookup{},
		&stubTenantService{},
		&stubFileService{},
		&stubResourceCatalog{},
	)

	req := httptest.NewRequest(http.MethodGet, "/sessions/session-1/messages/message-1/files", nil)
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if got, want := recorder.Code, http.StatusBadRequest; got != want {
		t.Fatalf("status = %d, want %d", got, want)
	}
}

func TestMessageScopedFilesRejectsRevokedSharedAgent(t *testing.T) {
	gin.SetMode(gin.TestMode)
	const ref = "resource://AbCdEfGhIjKlMnOpQrStUv"

	engine := newMessageScopedFilesTestEngine(
		42,
		&stubMessageFileLookup{get: func(context.Context, string, string) (*types.Message, error) {
			return &types.Message{AgentID: "agent-1", AgentTenantID: 7}, nil
		}},
		&stubSharedAgentFileLookup{get: func(
			context.Context, uint64, types.TenantRole, string, ...uint64,
		) (*types.CustomAgent, error) {
			return nil, nil
		}},
		&stubTenantService{get: func(context.Context, uint64) (*types.Tenant, error) {
			t.Fatal("tenant lookup should not run after share revocation")
			return nil, nil
		}},
		&stubFileService{getFile: func(context.Context, string) (io.ReadCloser, error) {
			t.Fatal("GetFile should not run after share revocation")
			return nil, nil
		}},
		&stubResourceCatalog{resource: &types.StoredResource{TenantID: 7, PhysicalPath: "local://7/exports/chart.png"}},
	)

	req := httptest.NewRequest(http.MethodGet,
		"/sessions/session-1/messages/message-1/files?file_path="+url.QueryEscape(ref), nil)
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if recorder.Code != http.StatusForbidden {
		t.Fatalf("status=%d, want %d", recorder.Code, http.StatusForbidden)
	}
}

func TestMessageScopedFilesRejectsResourceOutsideMessageTenant(t *testing.T) {
	gin.SetMode(gin.TestMode)
	const ref = "resource://AbCdEfGhIjKlMnOpQrStUv"

	engine := newMessageScopedFilesTestEngine(
		42,
		&stubMessageFileLookup{get: func(context.Context, string, string) (*types.Message, error) {
			return &types.Message{AgentID: "agent-1", AgentTenantID: 8}, nil
		}},
		&stubSharedAgentFileLookup{get: func(
			context.Context, uint64, types.TenantRole, string, ...uint64,
		) (*types.CustomAgent, error) {
			t.Fatal("shared-agent lookup should not run for a mismatched resource tenant")
			return nil, nil
		}},
		&stubTenantService{},
		&stubFileService{},
		&stubResourceCatalog{resource: &types.StoredResource{TenantID: 7, PhysicalPath: "local://7/exports/chart.png"}},
	)

	req := httptest.NewRequest(http.MethodGet,
		"/sessions/session-1/messages/message-1/files?file_path="+url.QueryEscape(ref), nil)
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if recorder.Code != http.StatusForbidden {
		t.Fatalf("status=%d, want %d", recorder.Code, http.StatusForbidden)
	}
}

func TestServeFilesForcesActiveContentDownload(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("STORAGE_TYPE", "local")

	engine := gin.New()
	serveFiles(engine, &stubFileService{
		getFile: func(_ context.Context, _ string) (io.ReadCloser, error) {
			return io.NopCloser(strings.NewReader(`<svg onload="alert(1)"></svg>`)), nil
		},
	})

	filePath := "local://42/docs/payload.svg"
	req := httptest.NewRequest(http.MethodGet, "/files?file_path="+url.QueryEscape(filePath), nil)
	req = req.WithContext(context.WithValue(req.Context(), types.TenantInfoContextKey, &types.Tenant{ID: 42}))

	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)

	if got, want := recorder.Code, http.StatusOK; got != want {
		t.Fatalf("status = %d, want %d", got, want)
	}
	if got := recorder.Header().Get("Content-Type"); got != "application/octet-stream" {
		t.Fatalf("Content-Type = %q, want application/octet-stream", got)
	}
	if got := recorder.Header().Get("Content-Disposition"); got != "attachment" {
		t.Fatalf("Content-Disposition = %q, want attachment", got)
	}
	if got := recorder.Header().Get("X-Content-Type-Options"); got != "nosniff" {
		t.Fatalf("X-Content-Type-Options = %q, want nosniff", got)
	}
}