ai-gateway-go/internal/gateway/admission.go

package gateway

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"time"

	"aigateway.local/core/internal/apikey"
	"github.com/redis/go-redis/v9"
)

var ErrAdmissionUnavailable = errors.New("admission control unavailable")

type AdmissionReason string

const (
	AdmissionAllowed      AdmissionReason = ""
	AdmissionRateLimit    AdmissionReason = "rate_limit"
	AdmissionMonthlyQuota AdmissionReason = "monthly_quota"
)

type AdmissionDecision struct {
	Allowed    bool
	Reason     AdmissionReason
	Limit      int64
	Remaining  int64
	ResetAt    time.Time
	RetryAfter time.Duration
}

type AdmissionController interface {
	Allow(context.Context, apikey.Principal, time.Time) (AdmissionDecision, error)
}

type RedisAdmissionController struct {
	client *redis.Client
	script *redis.Script
}

func NewRedisAdmissionController(client *redis.Client) *RedisAdmissionController {
	return &RedisAdmissionController{client: client, script: redis.NewScript(admissionScript)}
}

func (c *RedisAdmissionController) Allow(ctx context.Context, principal apikey.Principal, now time.Time) (AdmissionDecision, error) {
	if principal.RequestsPerMinute == 0 && principal.MonthlyRequestQuota == 0 {
		return AdmissionDecision{Allowed: true}, nil
	}
	if c == nil || c.client == nil || principal.APIKeyID == "" {
		return AdmissionDecision{}, ErrAdmissionUnavailable
	}
	now = now.UTC()
	minuteReset := now.Truncate(time.Minute).Add(time.Minute)
	monthReset := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC)
	minuteKey := fmt.Sprintf("gateway:limit:api-key:%s:minute:%d", principal.APIKeyID, now.Unix()/60)
	monthKey := fmt.Sprintf("gateway:limit:api-key:%s:month:%s", principal.APIKeyID, now.Format("200601"))
	result, err := c.script.Run(ctx, c.client, []string{minuteKey, monthKey},
		principal.RequestsPerMinute, principal.MonthlyRequestQuota,
		int64(minuteReset.Sub(now).Seconds())+2, int64(monthReset.Sub(now).Seconds())+86400,
	).Slice()
	if err != nil {
		return AdmissionDecision{}, fmt.Errorf("%w: %v", ErrAdmissionUnavailable, err)
	}
	if len(result) != 3 {
		return AdmissionDecision{}, ErrAdmissionUnavailable
	}
	code, err := redisInteger(result[0])
	if err != nil {
		return AdmissionDecision{}, ErrAdmissionUnavailable
	}
	minuteCount, err := redisInteger(result[1])
	if err != nil {
		return AdmissionDecision{}, ErrAdmissionUnavailable
	}
	monthCount, err := redisInteger(result[2])
	if err != nil {
		return AdmissionDecision{}, ErrAdmissionUnavailable
	}
	switch code {
	case 0:
		decision := AdmissionDecision{Allowed: true}
		if principal.RequestsPerMinute > 0 {
			decision.Limit = int64(principal.RequestsPerMinute)
			decision.Remaining = max(decision.Limit-minuteCount, 0)
			decision.ResetAt = minuteReset
		}
		return decision, nil
	case 1:
		return AdmissionDecision{
			Allowed: false, Reason: AdmissionRateLimit, Limit: int64(principal.RequestsPerMinute),
			Remaining: 0, ResetAt: minuteReset, RetryAfter: minuteReset.Sub(now),
		}, nil
	case 2:
		return AdmissionDecision{
			Allowed: false, Reason: AdmissionMonthlyQuota, Limit: principal.MonthlyRequestQuota,
			Remaining: max(principal.MonthlyRequestQuota-monthCount, 0), ResetAt: monthReset, RetryAfter: monthReset.Sub(now),
		}, nil
	default:
		return AdmissionDecision{}, ErrAdmissionUnavailable
	}
}

func redisInteger(value any) (int64, error) {
	switch number := value.(type) {
	case int64:
		return number, nil
	case string:
		return strconv.ParseInt(number, 10, 64)
	case []byte:
		return strconv.ParseInt(string(number), 10, 64)
	default:
		return 0, fmt.Errorf("unexpected redis integer %T", value)
	}
}

const admissionScript = `
local rpm = tonumber(ARGV[1])
local monthly_quota = tonumber(ARGV[2])
local minute_count = tonumber(redis.call('GET', KEYS[1]) or '0')
local month_count = tonumber(redis.call('GET', KEYS[2]) or '0')

if monthly_quota > 0 and month_count >= monthly_quota then
  return {2, minute_count, month_count}
end
if rpm > 0 and minute_count >= rpm then
  return {1, minute_count, month_count}
end

if rpm > 0 then
  minute_count = redis.call('INCR', KEYS[1])
  if minute_count == 1 then redis.call('EXPIRE', KEYS[1], tonumber(ARGV[3])) end
end
if monthly_quota > 0 then
  month_count = redis.call('INCR', KEYS[2])
  if month_count == 1 then redis.call('EXPIRE', KEYS[2], tonumber(ARGV[4])) end
end
return {0, minute_count, month_count}
`