MailUI4Agents/server/internal/handler/permission_request_test.go

package handler

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"

	"github.com/agentmail/gateway/internal/db"
	"github.com/agentmail/gateway/internal/middleware"
	"github.com/agentmail/gateway/internal/repo"
	"github.com/google/uuid"
)

// RequestPermission 的决策人解析必须在 HTTP handler 层有覆盖。
//
// 过去 `to` 是 Agent 可自由填写的用户名，只要对应人类账号存在就直接采用。
// 这允许 Agent 把危险操作的审批丢给任意一个与任务无关的人；repo 层的参与方
// helper 即使正确，handler 漏调一次仍会把漏洞带回来。
func TestRequestPermissionRestrictsExplicitDeciderToThread(t *testing.T) {
	setupPermissionHandlerDB(t)
	ctx := context.Background()

	mustPermissionUser(t, "owner")
	mustPermissionUser(t, "reviewer")
	mustPermissionUser(t, "outsider")
	if err := repo.CreateOrUpdateAgent(ctx, "pi", "secret", "test", nil); err != nil {
		t.Fatalf("注册 pi: %v", err)
	}

	ownerID := permissionUserID(t, "owner")
	sid, err := repo.CreateSession(ctx, nil, "pi", "授权路由", "/tmp")
	if err != nil {
		t.Fatalf("建会话: %v", err)
	}
	if err := repo.SetSessionOwner(ctx, sid, ownerID); err != nil {
		t.Fatalf("设置 owner: %v", err)
	}
	if _, err := repo.CreateMail(ctx, sid, nil, "reviewer", "", "pi", "/tmp",
		"任务", "请执行", nil); err != nil {
		t.Fatalf("建线索邮件: %v", err)
	}

	t.Run("无关人类被忽略并回落 owner", func(t *testing.T) {
		resp := requestPermission(t, sid, "outsider", "是否允许执行 bash？")
		if resp.Code != http.StatusOK {
			t.Fatalf("HTTP %d: %s", resp.Code, resp.Body.String())
		}
		var body map[string]any
		if err := json.Unmarshal(resp.Body.Bytes(), &body); err != nil {
			t.Fatalf("解析响应: %v", err)
		}
		if got := body["decider"]; got != "owner" {
			t.Fatalf("无关人类不应收到审批，decider=%v，期望 owner", got)
		}
		assertPermissionRecipient(t, body["mail_id"].(string), "owner")
	})

	t.Run("线索上的人类可被显式指定", func(t *testing.T) {
		resp := requestPermission(t, sid, "reviewer", "是否允许写文件？")
		if resp.Code != http.StatusOK {
			t.Fatalf("HTTP %d: %s", resp.Code, resp.Body.String())
		}
		var body map[string]any
		if err := json.Unmarshal(resp.Body.Bytes(), &body); err != nil {
			t.Fatalf("解析响应: %v", err)
		}
		if got := body["decider"]; got != "reviewer" {
			t.Fatalf("线索上的人类应保持显式优先级，decider=%v", got)
		}
		assertPermissionRecipient(t, body["mail_id"].(string), "reviewer")
	})

	t.Run("线索上的 Agent 仍不能成为决策人", func(t *testing.T) {
		resp := requestPermission(t, sid, "pi", "是否允许编辑？")
		if resp.Code != http.StatusOK {
			t.Fatalf("HTTP %d: %s", resp.Code, resp.Body.String())
		}
		var body map[string]any
		if err := json.Unmarshal(resp.Body.Bytes(), &body); err != nil {
			t.Fatalf("解析响应: %v", err)
		}
		if got := body["decider"]; got != "owner" {
			t.Fatalf("Agent 不能在 WebUI 决策，应回落 owner，实际 %v", got)
		}
	})
}

func TestRequestPermissionRejectsAllAgentThread(t *testing.T) {
	setupPermissionHandlerDB(t)
	ctx := context.Background()
	for _, name := range []string{"pi", "dsh"} {
		if err := repo.CreateOrUpdateAgent(ctx, name, "secret", "test", nil); err != nil {
			t.Fatalf("注册 %s: %v", name, err)
		}
	}
	sid, err := repo.CreateSession(ctx, nil, "pi", "纯 Agent 链", "/tmp")
	if err != nil {
		t.Fatalf("建会话: %v", err)
	}
	if _, err := repo.CreateMail(ctx, sid, nil, "dsh", "", "pi", "/tmp",
		"任务", "请执行", nil); err != nil {
		t.Fatalf("建线索邮件: %v", err)
	}

	resp := requestPermission(t, sid, "human", "是否允许执行 bash？")
	if resp.Code != http.StatusConflict {
		t.Fatalf("纯 Agent 链必须 fail closed，HTTP %d: %s", resp.Code, resp.Body.String())
	}
	if bytes.Contains(resp.Body.Bytes(), []byte(`"decider":"admin"`)) {
		t.Fatalf("不得退回无关管理员: %s", resp.Body.String())
	}
}

func setupPermissionHandlerDB(t *testing.T) {
	t.Helper()
	db.Close()
	path := filepath.Join(t.TempDir(), "permission-handler.db")
	if err := db.Connect(context.Background(), "sqlite://"+path); err != nil {
		t.Fatalf("连接测试库: %v", err)
	}
	if err := db.Migrate(context.Background()); err != nil {
		t.Fatalf("迁移测试库: %v", err)
	}
	t.Cleanup(db.Close)
}

func mustPermissionUser(t *testing.T, username string) {
	t.Helper()
	if _, err := db.DB.ExecContext(context.Background(),
		`INSERT INTO users (user_id, username, display_name, password_hash, role)
		 VALUES ($1, $2, $3, 'x', 'user')`, uuid.New(), username, username); err != nil {
		t.Fatalf("建用户 %s: %v", username, err)
	}
}

func permissionUserID(t *testing.T, username string) uuid.UUID {
	t.Helper()
	var id uuid.UUID
	if err := db.DB.QueryRowContext(context.Background(),
		`SELECT user_id FROM users WHERE username = $1`, username).Scan(&id); err != nil {
		t.Fatalf("查用户 %s: %v", username, err)
	}
	return id
}

func requestPermission(t *testing.T, sid uuid.UUID, to, question string) *httptest.ResponseRecorder {
	t.Helper()
	payload, _ := json.Marshal(map[string]any{
		"session_id": sid.String(),
		"to":         to,
		"question":   question,
		"options":    []string{"同意", "拒绝"},
	})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/permission/request", bytes.NewReader(payload))
	req.Header.Set("Content-Type", "application/json")
	ctx := context.WithValue(req.Context(), middleware.AgentNameKey, "pi")
	resp := httptest.NewRecorder()
	RequestPermission(resp, req.WithContext(ctx))
	return resp
}

func assertPermissionRecipient(t *testing.T, mailID, want string) {
	t.Helper()
	var got string
	if err := db.DB.QueryRowContext(context.Background(),
		`SELECT to_name FROM mails WHERE mail_id = $1`, mailID).Scan(&got); err != nil {
		t.Fatalf("查权限邮件 %s: %v", mailID, err)
	}
	if got != want {
		t.Fatalf("权限邮件收件人=%q，期望 %q", got, want)
	}
}