MailUI4Agents/server/internal/repo/disabled_auth_test.go

package repo

/*
停用必须拦住**鉴权**，不只是拦住「别人给它发信」（2026-10-03）。

# 这条缺口是怎么被发现的

要做公网 MCP 接入验证，注册了一个一次性探针身份
（`mcp-wan-probe`），收尾时 `UPDATE agents SET status='disabled'`。
本以为凭证就此失效，实测：

    POST /api/v1/mcp（带该凭证）        → 200
    tools/call send_mail                 → 已发送（真的发出去了）

`SetAgentDisabled` 这个 API 存在、也返回成功，但**只挡住了投递方向**
（`EnsureAgentDeliverable` 里查 status），鉴权两条路径都不查：

    VerifyAgent(ctx, name, secret)      选了 status 却从不判断 ← 缺口
    VerifyAgentKey(ctx, token)          拿到名字直接返回       ← 同一个缺口，更严重

缺口形状与这轮反复出现的那类一致：**接口存在、返回成功、但只做了一半**。
比「没有这个接口」更危险 —— 运维会以为停用已经生效。

`key_token` 那条尤其要紧：它是 `middleware/auth.go` 注释里标「推荐」的
路径，官方建议用它，于是它成了绕过停用的最短路。

# 判据要点

停用后必须**拿不到 Agent 身份**（AgentAuth 直接 401），而不是拿到身份后
在某个业务分支里拒绝 —— 后者会让「列表类」API 仍然泄露该身份的存在。
返回 `sql.ErrNoRows`（而非自定义错误）是为了与「凭证不存在」不可区分。
*/

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"testing"

	"github.com/agentmail/gateway/internal/db"
	"github.com/agentmail/gateway/internal/models"
)

func setupDisabledAgent(t *testing.T, name string, status string) {
	t.Helper()
	db.Close()
	if err := db.Connect(context.Background(), "sqlite://"+t.TempDir()+"/disabled.db"); err != nil {
		t.Fatalf("连接测试库: %v", err)
	}
	if err := db.Migrate(context.Background()); err != nil {
		t.Fatalf("迁移测试库: %v", err)
	}
	t.Cleanup(db.Close)
	if _, err := db.DB.ExecContext(context.Background(),
		`INSERT INTO agents (agent_name, secret, platform, status, default_rounds)
		 VALUES ($1, 'sec', 'test', $2, 50)`, name, status); err != nil {
		t.Fatalf("建 Agent: %v", err)
	}
}

// ★ 主判据：停用 ⇒ 两条鉴权路径都必须失败。

func TestDisabledAgentFailsAuth(t *testing.T) {
	setupDisabledAgent(t, "stopped-agent", "disabled")

	// 路径①name/secret
	if _, err := VerifyAgent(context.Background(), "stopped-agent", "sec"); err == nil {
		t.Error("★ VerifyAgent 接受了已停用的身份 ⇒ 停用形同虚设，" +
			"它照样能用凭证调任何 API（含 send_mail）")
	} else if !errors.Is(err, sql.ErrNoRows) {
		t.Errorf("应返回 sql.ErrNoRows（与凭证不存在不可区分），实际 %v", err)
	}

	// 路径② Bearer key_token —— 注释里标「推荐」的那条
	token := seedKeyFor(t, "stopped-agent")
	name, err := VerifyAgentKey(context.Background(), token)
	if err == nil && name != "" {
		t.Error("★ VerifyAgentKey 接受了已停用身份的令牌 ⇒ Bearer 这条" +
			"（推荐路径）能绕过停用")
	}
}

// ★ 反向对照：status='online'（或任何非 disabled 值）必须照常通过。
//   没有这一格，一个「永远返回 ErrNoRows」的 VerifyAgent 也是全绿的，
//   而那会让**所有** agent 都登不上 —— 一个比原缺口更大的事故。

func TestNonDisabledAgentStillAuthenticates(t *testing.T) {
	for _, st := range []string{"online", "active", ""} {
		setupDisabledAgent(t, "live-agent", st)
		a, err := VerifyAgent(context.Background(), "live-agent", "sec")
		if err != nil {
			t.Errorf("★ status=%q 的在跑身份被拒了 —— 判据只能拦 disabled，%v", st, err)
			continue
		}
		if a.Name != "live-agent" {
			t.Errorf("返回了错误身份: %q", a.Name)
		}
		token := seedKeyFor(t, "live-agent")
		name, err := VerifyAgentKey(context.Background(), token)
		if err != nil || name != "live-agent" {
			t.Errorf("★ status=%q 的在跑身份过不了 Bearer 路径: name=%q err=%v", st, name, err)
		}
	}
}

// ★ 停用不能泄露「这个身份存在」：错误必须与凭证错配时完全一致。
//   若停用返回自定义错误，攻击者可用它枚举出「哪些名字是真实 Agent」。

func TestDisabledIndistinguishableFromWrongSecret(t *testing.T) {
	setupDisabledAgent(t, "stopped-agent", "disabled")

	_, errDisabled := VerifyAgent(context.Background(), "stopped-agent", "sec")
	_, errWrongSecret := VerifyAgent(context.Background(), "stopped-agent", "wrong-secret")
	_, errNoSuchAgent := VerifyAgent(context.Background(), "nobody", "sec")

	for _, c := range []struct {
		name string
		err  error
	}{
		{"凭证错配", errWrongSecret},
		{"身份不存在", errNoSuchAgent},
	} {
		if !errors.Is(errDisabled, c.err) {
			t.Errorf("★ 停用(%v) 必须与「%s」(%v) 返回同一个错误，"+
				"否则可用于枚举真实 Agent 名", errDisabled, c.name, c.err)
		}
	}
}

// ★ 停用后不得解析出工作区：那会让调用方以为身份还「活着」。
//   （实现上把 status 检查放在 wsJSON 解析之前，这里钉住那个顺序。）

func TestDisabledAgentDoesNotExposeWorkspaces(t *testing.T) {
	setupDisabledAgent(t, "stopped-agent", "disabled")
	dbMustExec(t, `UPDATE agents SET workspaces = '[{"path":"/secret/project"}]' WHERE agent_name='stopped-agent'`)

	a, err := VerifyAgent(context.Background(), "stopped-agent", "sec")
	if err == nil && a != nil && len(a.Workspaces) > 0 {
		t.Errorf("★ 停用身份仍被解析出工作区 %+v ⇒ 调用方可能以为它还活着",
			a.Workspaces)
	}
}

// seedKeyFor 给某 Agent 建一张永久 key_token，返回 token。
// 用真实的表结构与 key_type 取值（models.KeyPermanent），别自造。
func seedKeyFor(t *testing.T, agentName string) string {
	t.Helper()
	token := "tok-" + agentName + "-0123456789abcdef"
	if _, err := db.DB.ExecContext(context.Background(),
		`INSERT INTO agent_keys (key_token, agent_name, key_type, label)
		 VALUES ($1, $2, $3, 'test')`,
		token, agentName, models.KeyPermanent); err != nil {
		t.Fatalf("建 key_token: %v", err)
	}
	return token
}

// dbMustExec 让判据里的 SQL 失败时直接报出来，而不是静默跳过
// （踩过：断言语句自己没跑，判据永远绿）。
func dbMustExec(t *testing.T, sqlText string) {
	t.Helper()
	if _, err := db.DB.ExecContext(context.Background(), sqlText); err != nil {
		t.Fatalf("执行 %s: %v", sqlText, err)
	}
}

var _ = fmt.Sprintf