370 lines
10 KiB
Go
370 lines
10 KiB
Go
package repo
|
||
|
||
import (
|
||
"context"
|
||
"testing"
|
||
|
||
"github.com/agentmail/gateway/internal/models"
|
||
"github.com/google/uuid"
|
||
)
|
||
|
||
// ─── InheritedMode：继承、收紧、不存在的父会话 ───
|
||
|
||
// parent nil → 返回 requested 的规范化值
|
||
func TestInheritedMode_NilParent(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
got := InheritedMode(ctx, nil, "plan")
|
||
if got != models.ModePlan {
|
||
t.Errorf("nil parent + plan: got %q, want plan", got)
|
||
}
|
||
|
||
got = InheritedMode(ctx, nil, "workspace")
|
||
if got != models.ModeWorkspace {
|
||
t.Errorf("nil parent + workspace: got %q, want workspace", got)
|
||
}
|
||
|
||
got = InheritedMode(ctx, nil, "full")
|
||
if got != models.ModeFull {
|
||
t.Errorf("nil parent + full: got %q, want full", got)
|
||
}
|
||
}
|
||
|
||
// parent plan → 子会话只能 plan（不能提权）
|
||
func TestInheritedMode_ParentPlan_Tightens(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
seedAgent(t, "pi", 20)
|
||
|
||
id := createTestSession(t, ctx, "pi", "/ws")
|
||
perm, err := SetSessionPermissionMode(ctx, id, models.ModePlan)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if perm.Mode != models.ModePlan {
|
||
t.Fatal("expected plan")
|
||
}
|
||
|
||
// 请求 workspace（更宽松）→ 应被收紧为 plan
|
||
got := InheritedMode(ctx, &id, models.ModeWorkspace)
|
||
if got != models.ModePlan {
|
||
t.Errorf("plan parent + workspace request: got %q, want plan", got)
|
||
}
|
||
|
||
// 请求 full → 同样收紧
|
||
got = InheritedMode(ctx, &id, models.ModeFull)
|
||
if got != models.ModePlan {
|
||
t.Errorf("plan parent + full request: got %q, want plan", got)
|
||
}
|
||
|
||
// 请求 plan → 保持 plan
|
||
got = InheritedMode(ctx, &id, models.ModePlan)
|
||
if got != models.ModePlan {
|
||
t.Errorf("plan parent + plan request: got %q, want plan", got)
|
||
}
|
||
}
|
||
|
||
// parent workspace → 子会话 workspace 或更严
|
||
func TestInheritedMode_ParentWorkspace(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
seedAgent(t, "pi", 20)
|
||
|
||
id := createTestSession(t, ctx, "pi", "/ws")
|
||
_, err := SetSessionPermissionMode(ctx, id, models.ModeWorkspace)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
// 请求 full → 收紧为 workspace（子不能比父更松）
|
||
got := InheritedMode(ctx, &id, models.ModeFull)
|
||
if got != models.ModeWorkspace {
|
||
t.Errorf("workspace parent + full: got %q, want workspace", got)
|
||
}
|
||
|
||
// 请求 plan → 保留 plan（比父更严，允许）
|
||
got = InheritedMode(ctx, &id, models.ModePlan)
|
||
if got != models.ModePlan {
|
||
t.Errorf("workspace parent + plan: got %q, want plan", got)
|
||
}
|
||
}
|
||
|
||
// parent full → 子会话可请求任意档位
|
||
func TestInheritedMode_ParentFull(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
seedAgent(t, "pi", 20)
|
||
|
||
id := createTestSession(t, ctx, "pi", "/ws")
|
||
_, err := SetSessionPermissionMode(ctx, id, models.ModeFull)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
got := InheritedMode(ctx, &id, models.ModeWorkspace)
|
||
if got != models.ModeWorkspace {
|
||
t.Errorf("full parent + workspace: got %q, want workspace", got)
|
||
}
|
||
got = InheritedMode(ctx, &id, models.ModePlan)
|
||
if got != models.ModePlan {
|
||
t.Errorf("full parent + plan: got %q, want plan", got)
|
||
}
|
||
}
|
||
|
||
// 脏值 fallback：非法的 requested 值在 InheritedMode 里被规范化为默认档
|
||
func TestInheritedMode_InvalidRequested(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
got := InheritedMode(ctx, nil, "elephant")
|
||
if got != models.DefaultPermissionMode {
|
||
t.Errorf("invalid requested: got %q, want %q", got, models.DefaultPermissionMode)
|
||
}
|
||
}
|
||
|
||
// parent 不存在时（查不到）回落：ModeAtMost(default, req)
|
||
func TestInheritedMode_InvalidParent(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
fakeID := uuid.New()
|
||
got := InheritedMode(ctx, &fakeID, "full")
|
||
want := models.ModeAtMost(models.DefaultPermissionMode, "full")
|
||
if got != want {
|
||
t.Errorf("invalid parent + full: got %q, want %q (modeAtMost(default, full))", got, want)
|
||
}
|
||
}
|
||
|
||
// ─── SetSessionPermissionMode roundtrip + dirty value normalization ───
|
||
|
||
func TestSetSessionPermissionMode_Roundtrip(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
seedAgent(t, "dsh", 20)
|
||
|
||
id := createTestSession(t, ctx, "dsh", "/ws")
|
||
perm, err := SetSessionPermissionMode(ctx, id, models.ModeFull)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if perm.Mode != models.ModeFull || perm.Enforcement != "advisory" {
|
||
t.Errorf("full mode: got mode=%q enforcement=%q", perm.Mode, perm.Enforcement)
|
||
}
|
||
|
||
// 读出来一致
|
||
got := SessionPermissionMode(ctx, id)
|
||
if got != models.ModeFull {
|
||
t.Errorf("read back: got %q, want full", got)
|
||
}
|
||
}
|
||
|
||
func TestSetSessionPermissionMode_DirtyValue_FailClosed(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
seedAgent(t, "pi", 20)
|
||
|
||
id := createTestSession(t, ctx, "pi", "/ws")
|
||
perm, err := SetSessionPermissionMode(ctx, id, "BOGUS")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if perm.Mode != models.DefaultPermissionMode {
|
||
t.Errorf("dirty value: got %q, want %q (fail-closed to default)", perm.Mode, models.DefaultPermissionMode)
|
||
}
|
||
}
|
||
|
||
// ─── 脏值归一化测试（覆盖 NormalizePermissionMode 本身） ───
|
||
|
||
func TestNormalizePermissionMode_Inputs(t *testing.T) {
|
||
tests := []struct {
|
||
input string
|
||
want string
|
||
}{
|
||
{"plan", "plan"},
|
||
{"workspace", "workspace"},
|
||
{"full", "full"},
|
||
// 大小写/空白不归一：NormalizePermissionMode 只接受精确匹配的合法档位，
|
||
// 其余一律 fail-closed 到默认档（workspace）—— 不 trim 不 lowercase，
|
||
// 避免「我以为给了 plan 实际拿到别的」这种隐式转换造成的安全错觉。
|
||
{"Plan", models.DefaultPermissionMode},
|
||
{" PLAN ", models.DefaultPermissionMode},
|
||
{"", models.DefaultPermissionMode},
|
||
{"bogus", models.DefaultPermissionMode},
|
||
{"F ULL", models.DefaultPermissionMode},
|
||
}
|
||
for _, tt := range tests {
|
||
got := models.NormalizePermissionMode(tt.input)
|
||
if got != tt.want {
|
||
t.Errorf("NormalizePermissionMode(%q) = %q, want %q", tt.input, got, tt.want)
|
||
}
|
||
}
|
||
}
|
||
|
||
// ─── 负向对照：plan 档不能提权 ───
|
||
|
||
func TestInheritedMode_PlanCannotEscalate(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
seedAgent(t, "pi", 20)
|
||
|
||
parent := createTestSession(t, ctx, "pi", "/ws")
|
||
_, _ = SetSessionPermissionMode(ctx, parent, models.ModePlan)
|
||
|
||
child := InheritedMode(ctx, &parent, models.ModeFull)
|
||
if child != models.ModePlan {
|
||
t.Errorf("SECURITY FAIL: plan session escalated to %q via InheritedMode", child)
|
||
}
|
||
}
|
||
|
||
// ─── 三层继承链 ───
|
||
|
||
func TestInheritedMode_ThreeLevelChain(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
seedAgent(t, "pi", 20)
|
||
|
||
root := createTestSession(t, ctx, "pi", "/ws")
|
||
_, _ = SetSessionPermissionMode(ctx, root, models.ModeFull)
|
||
|
||
child := InheritedMode(ctx, &root, models.ModeWorkspace) // workspace < full → workspace
|
||
childID := createTestSession(t, ctx, "pi", "/ws")
|
||
_, _ = SetSessionPermissionMode(ctx, childID, child)
|
||
|
||
grandchild := InheritedMode(ctx, &childID, models.ModeFull) // full vs workspace → workspace
|
||
if grandchild != models.ModeWorkspace {
|
||
t.Errorf("grandchild: got %q, want workspace", grandchild)
|
||
}
|
||
|
||
// plan → workspace → plan chain
|
||
planChild := InheritedMode(ctx, &root, models.ModePlan) // plan < full → plan
|
||
planChildID := createTestSession(t, ctx, "pi", "/ws")
|
||
_, _ = SetSessionPermissionMode(ctx, planChildID, planChild)
|
||
|
||
grandchild2 := InheritedMode(ctx, &planChildID, models.ModeFull) // full vs plan → plan
|
||
if grandchild2 != models.ModePlan {
|
||
t.Errorf("plan chain grandchild: got %q, want plan", grandchild2)
|
||
}
|
||
}
|
||
|
||
// ─── calendar_events permission_mode roundtrip ───
|
||
|
||
func TestCalendarEventPermissionMode_Roundtrip(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
seedAgent(t, "pi", 20)
|
||
|
||
// CreateCalendarEvent 应规范化档位
|
||
e := &models.CalendarEvent{
|
||
Title: "测试日程",
|
||
AgentName: "pi",
|
||
ToAddress: "pi@/home/program/agentmail",
|
||
PermissionMode: "full",
|
||
Status: "active",
|
||
CreatedBy: "jianf",
|
||
}
|
||
created, err := CreateCalendarEvent(ctx, e)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if created.PermissionMode != "full" {
|
||
t.Errorf("created event permission_mode: got %q, want full", created.PermissionMode)
|
||
}
|
||
|
||
// 读回来一致
|
||
got, err := GetCalendarEvent(ctx, created.EventID)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if got.PermissionMode != "full" {
|
||
t.Errorf("read back: got %q, want full", got.PermissionMode)
|
||
}
|
||
}
|
||
|
||
func TestCalendarEventPermissionMode_DirtyValue_Normalized(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
seedAgent(t, "pi", 20)
|
||
|
||
e := &models.CalendarEvent{
|
||
Title: "脏值日程",
|
||
AgentName: "pi",
|
||
ToAddress: "pi@/home/program/agentmail",
|
||
PermissionMode: "INVALID",
|
||
Status: "active",
|
||
CreatedBy: "jianf",
|
||
}
|
||
created, err := CreateCalendarEvent(ctx, e)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if created.PermissionMode != models.DefaultPermissionMode {
|
||
t.Errorf("dirty value: got %q, want %q", created.PermissionMode, models.DefaultPermissionMode)
|
||
}
|
||
}
|
||
|
||
func TestCalendarEventPermissionMode_UpdateRoundtrip(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
seedAgent(t, "pi", 20)
|
||
|
||
e := &models.CalendarEvent{
|
||
Title: "更新日程",
|
||
AgentName: "pi",
|
||
ToAddress: "pi@/home/program/agentmail",
|
||
PermissionMode: "workspace",
|
||
Status: "active",
|
||
CreatedBy: "jianf",
|
||
}
|
||
created, err := CreateCalendarEvent(ctx, e)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
created.PermissionMode = "plan"
|
||
if err := UpdateCalendarEvent(ctx, created.EventID, created); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
got, err := GetCalendarEvent(ctx, created.EventID)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if got.PermissionMode != "plan" {
|
||
t.Errorf("after update: got %q, want plan", got.PermissionMode)
|
||
}
|
||
}
|
||
|
||
// ─── adopt 接管时会话档位必须写入 ───
|
||
|
||
func TestAdoptPlatformSession_WritesDefaultMode(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
seedAgent(t, "pi", 20)
|
||
|
||
id, err := AdoptPlatformSession(ctx, "pi", "plat-123", "my-proj", "/ws", "接管测试")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
// 接管会话应显式写入默认档位（不是靠 DB 默认值）
|
||
mode := SessionPermissionMode(ctx, id)
|
||
if mode != models.DefaultPermissionMode {
|
||
t.Errorf("adopt session mode: got %q, want %q", mode, models.DefaultPermissionMode)
|
||
}
|
||
}
|
||
|
||
// ─── helpers ───
|
||
|
||
func createTestSession(t *testing.T, ctx context.Context, agent, workspace string) uuid.UUID {
|
||
t.Helper()
|
||
id, err := CreateSession(ctx, nil, agent, "test subject", workspace)
|
||
if err != nil {
|
||
t.Fatalf("create session: %v", err)
|
||
}
|
||
return id
|
||
}
|
||
|
||
// init 确保每个 test 函数执行前 DB 足够干净
|
||
func init() {
|
||
// 空 —— setupTestDB 在每个测试函数内调用
|
||
}
|