perms-system-server/internal/server/grpc_rate_limit_audit_test.go

package server

import (
	"context"
	"database/sql"
	"net"
	"testing"
	"time"

	authHelper "perms-system-server/internal/logic/auth"
	userModel "perms-system-server/internal/model/user"
	"perms-system-server/internal/svc"
	"perms-system-server/internal/testutil"
	"perms-system-server/pb"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/zeromicro/go-zero/core/limit"
	"github.com/zeromicro/go-zero/core/stores/redis"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/status"
)

// ---------------------------------------------------------------------------
// 覆盖目标：
//   H-2：gRPC RefreshToken / VerifyToken 必须受 IP 维度限流保护；
//   M-7：extractClientIP 必须剥离端口，同一 IP 的不同 TCP 端口共享同一限流桶；
//   H-1：gRPC RefreshToken 里走的是 IncrementTokenVersionIfMatch，
//        一次成功后旧 refreshToken 立刻失效；重放必须返回 Unauthenticated。
// ---------------------------------------------------------------------------

// withPeerIP 往 ctx 注入指定 "host:port" 的 peer，模拟 gRPC 上游的 PeerAddr。
func withPeerIP(ctx context.Context, hostPort string) context.Context {
	addr, err := net.ResolveTCPAddr("tcp", hostPort)
	if err != nil {
		panic(err)
	}
	return peer.NewContext(ctx, &peer.Peer{Addr: addr})
}

// TC-0828: H-2 —— GrpcRefreshLimiter 在配额用尽后对同 IP 新请求返回 ResourceExhausted。
func TestGrpcRefreshToken_RateLimit_OverIP(t *testing.T) {
	ctx := context.Background()
	svcCtx := svc.NewServiceContext(testutil.GetTestConfig())
	cfg := testutil.GetTestConfig()
	rds := redis.MustNewRedis(cfg.CacheRedis.Nodes[0].RedisConf)

	// quota=1 的定制 limiter，让第 2 次必然 429/ResourceExhausted。
	svcCtx.GrpcRefreshLimiter = limit.NewPeriodLimit(
		60, 1, rds, cfg.CacheRedis.KeyPrefix+":rl:grpc:refresh:ut:"+testutil.UniqueId())
	svcCtx.TokenOpLimiter = nil

	srv := NewPermServer(svcCtx)

	// 第 1 次：故意用个无效 token，让 limiter 放行、业务层兜底返回 Unauthenticated。
	// 这里只关心 limiter 是否"吃掉 1 个配额"。
	ctx1 := withPeerIP(ctx, "10.1.2.3:11111")
	_, err1 := srv.RefreshToken(ctx1, &pb.RefreshTokenReq{RefreshToken: "invalid"})
	require.Error(t, err1)
	st1, _ := status.FromError(err1)
	assert.Equal(t, codes.Unauthenticated, st1.Code(),
		"首次放行，业务层应返回 Unauthenticated（token 无效），不应是 ResourceExhausted")

	// 第 2 次：同 IP 但端口不同（模拟新 TCP 连接），必须被同一限流桶拦住。
	ctx2 := withPeerIP(ctx, "10.1.2.3:22222")
	_, err2 := srv.RefreshToken(ctx2, &pb.RefreshTokenReq{RefreshToken: "anything"})
	require.Error(t, err2)
	st2, _ := status.FromError(err2)
	assert.Equal(t, codes.ResourceExhausted, st2.Code(),
		"H-2 + M-7：同 IP 第 2 次刷新必须 429；端口变化不得绕过限流（extractClientIP 剥端口）")
	assert.Contains(t, st2.Message(), "过于频繁")
}

// TC-0829: H-2 —— GrpcVerifyLimiter 在配额用尽后对同 IP 新请求返回 ResourceExhausted。
// VerifyToken 契约是"非法 token 返回 Valid=false 而不是 error"，因此限流是唯一能让接口返回 gRPC error 的路径。
func TestGrpcVerifyToken_RateLimit_OverIP(t *testing.T) {
	ctx := context.Background()
	svcCtx := svc.NewServiceContext(testutil.GetTestConfig())
	cfg := testutil.GetTestConfig()
	rds := redis.MustNewRedis(cfg.CacheRedis.Nodes[0].RedisConf)
	svcCtx.GrpcVerifyLimiter = limit.NewPeriodLimit(
		60, 1, rds, cfg.CacheRedis.KeyPrefix+":rl:grpc:verify:ut:"+testutil.UniqueId())
	srv := NewPermServer(svcCtx)

	ctx1 := withPeerIP(ctx, "10.9.8.7:30001")
	resp1, err1 := srv.VerifyToken(ctx1, &pb.VerifyTokenReq{AccessToken: "invalid"})
	require.NoError(t, err1, "首次放行：VerifyToken 对非法 token 只返回 Valid=false，不 error")
	require.NotNil(t, resp1)
	assert.False(t, resp1.Valid)

	// 同 IP 不同端口 → 必须被限流拦住。
	ctx2 := withPeerIP(ctx, "10.9.8.7:30002")
	_, err2 := srv.VerifyToken(ctx2, &pb.VerifyTokenReq{AccessToken: "whatever"})
	require.Error(t, err2)
	st2, _ := status.FromError(err2)
	assert.Equal(t, codes.ResourceExhausted, st2.Code(),
		"H-2：gRPC VerifyToken 必须受 IP 级限流保护，防止下游被当 token oracle 爆破")
}

// TC-0830: M-7 —— extractClientIP 对 "host:port" 必须剥成 host；
// 缺失 peer 时返回 error，由上层决定降级到 unknown 桶。
func TestExtractClientIP_StripsPort(t *testing.T) {
	addr, err := net.ResolveTCPAddr("tcp", "192.168.0.1:54321")
	require.NoError(t, err)
	ctx := peer.NewContext(context.Background(), &peer.Peer{Addr: addr})

	ip, err := extractClientIP(ctx)
	require.NoError(t, err)
	assert.Equal(t, "192.168.0.1", ip,
		"M-7：gRPC peer.Addr 必须剥成纯 host；保留端口会导致限流形同虚设")

	// 无 peer 的 context
	_, err2 := extractClientIP(context.Background())
	assert.Error(t, err2, "无 peer 时必须返回 error，让上层选择 fail-close 或降级到 unknown 桶")
}

// TC-0831: H-1 + M-7 —— gRPC RefreshToken 成功一次后，旧 refreshToken 立刻失效；
// 换用同 IP 重放旧 token 必须返回 Unauthenticated（"登录状态已失效"），
// 而不是因端口变化绕过限流或因 CAS 失败被伪装成 500。
func TestGrpcRefreshToken_CASInvalidatesOldToken(t *testing.T) {
	ctx := context.Background()
	svcCtx := svc.NewServiceContext(testutil.GetTestConfig())
	conn := testutil.GetTestSqlConn()
	cfg := testutil.GetTestConfig()
	rds := redis.MustNewRedis(cfg.CacheRedis.Nodes[0].RedisConf)

	// 放开限流以聚焦 CAS 正确性（quota 大）。
	svcCtx.GrpcRefreshLimiter = limit.NewPeriodLimit(
		60, 100, rds, cfg.CacheRedis.KeyPrefix+":rl:grpc:refresh:cas:"+testutil.UniqueId())
	svcCtx.TokenOpLimiter = nil

	now := time.Now().Unix()
	uid := testutil.UniqueId()
	uRes, err := svcCtx.SysUserModel.Insert(ctx, &userModel.SysUser{
		Username: uid, Password: testutil.HashPassword("pass123"), Nickname: "n",
		Avatar: sql.NullString{}, IsSuperAdmin: 2, MustChangePassword: 2,
		Status: 1, CreateTime: now, UpdateTime: now,
	})
	require.NoError(t, err)
	userId, _ := uRes.LastInsertId()
	t.Cleanup(func() { testutil.CleanTable(ctx, conn, "`sys_user`", userId) })

	rt, err := authHelper.GenerateRefreshToken(cfg.Auth.RefreshSecret, cfg.Auth.RefreshExpire, userId, "", 0)
	require.NoError(t, err)

	srv := NewPermServer(svcCtx)

	// 第一次成功刷新。
	ctx1 := withPeerIP(ctx, "172.16.0.1:11001")
	resp, err := srv.RefreshToken(ctx1, &pb.RefreshTokenReq{RefreshToken: rt})
	require.NoError(t, err)
	require.NotEmpty(t, resp.RefreshToken)

	// 用同一个旧 rt 重放，应当 Unauthenticated；
	// 注意：旧 token 里 tokenVersion=0，DB 已被 CAS 推到 1，所以 "claims.TokenVersion != ud.TokenVersion" 这一步就会拦住。
	// 端口换掉以确保不是限流在帮我们挡。
	ctx2 := withPeerIP(ctx, "172.16.0.1:11002")
	_, err = srv.RefreshToken(ctx2, &pb.RefreshTokenReq{RefreshToken: rt})
	require.Error(t, err, "H-1：旧 refreshToken 成功刷新一次后必须失效")
	st, _ := status.FromError(err)
	assert.Equal(t, codes.Unauthenticated, st.Code(),
		"旧 token 重放必须返回 Unauthenticated，不能是 Internal/ResourceExhausted")
	assert.Contains(t, st.Message(), "登录状态已失效")
}