142 lines
3.5 KiB
Go
142 lines
3.5 KiB
Go
package redisadmission
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
_ "embed"
|
|
"encoding/binary"
|
|
"encoding/hex"
|
|
"errors"
|
|
"reflect"
|
|
"regexp"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
|
|
"proxy-pool/internal/platform/admission"
|
|
)
|
|
|
|
const MaximumLuaInteger int64 = 1<<53 - 1
|
|
|
|
var (
|
|
ErrInvalidOptions = errors.New("invalid Redis admission options")
|
|
namespacePattern = regexp.MustCompile(`^[A-Za-z0-9._-]+$`)
|
|
redisAdmissionKey = "pp:{admission}:"
|
|
redisAdmissionTail = ":window"
|
|
)
|
|
|
|
//go:embed scripts/fixed_window.lua
|
|
var fixedWindowSource string
|
|
|
|
var fixedWindowScript = redis.NewScript(fixedWindowSource)
|
|
|
|
type Options struct {
|
|
Namespace string
|
|
Window time.Duration
|
|
Global int64
|
|
PerKey int64
|
|
}
|
|
|
|
type Adapter struct {
|
|
client redis.Scripter
|
|
key string
|
|
windowMillis int64
|
|
global int64
|
|
perKey int64
|
|
}
|
|
|
|
var _ admission.Admitter = (*Adapter)(nil)
|
|
|
|
func New(client redis.Scripter, options Options) (*Adapter, error) {
|
|
if nilInterface(client) || options.Namespace != strings.TrimSpace(options.Namespace) ||
|
|
!namespacePattern.MatchString(options.Namespace) || options.Window <= 0 ||
|
|
options.Window%time.Millisecond != 0 || options.Global < 0 || options.PerKey < 0 ||
|
|
(options.Global == 0 && options.PerKey == 0) || options.Global > MaximumLuaInteger ||
|
|
options.PerKey > MaximumLuaInteger {
|
|
return nil, ErrInvalidOptions
|
|
}
|
|
windowMillis := options.Window.Milliseconds()
|
|
if windowMillis <= 0 || windowMillis > MaximumLuaInteger {
|
|
return nil, ErrInvalidOptions
|
|
}
|
|
return &Adapter{
|
|
client: client,
|
|
key: redisAdmissionKey + options.Namespace + redisAdmissionTail,
|
|
windowMillis: windowMillis,
|
|
global: options.Global,
|
|
perKey: options.PerKey,
|
|
}, nil
|
|
}
|
|
|
|
func (adapter *Adapter) Admit(ctx context.Context, identity string) error {
|
|
if ctx == nil || identity == "" {
|
|
return admission.ErrInvalidIdentity
|
|
}
|
|
if adapter == nil || nilInterface(adapter.client) || adapter.key == "" || adapter.windowMillis <= 0 ||
|
|
adapter.windowMillis > MaximumLuaInteger || adapter.global < 0 || adapter.perKey < 0 ||
|
|
(adapter.global == 0 && adapter.perKey == 0) || adapter.global > MaximumLuaInteger ||
|
|
adapter.perKey > MaximumLuaInteger {
|
|
return admission.ErrUnavailable
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
|
|
result, err := fixedWindowScript.Run(ctx, adapter.client, []string{adapter.key},
|
|
strconv.FormatInt(adapter.windowMillis, 10),
|
|
strconv.FormatInt(adapter.global, 10),
|
|
strconv.FormatInt(adapter.perKey, 10),
|
|
hashedIdentityField(identity),
|
|
).Result()
|
|
if err != nil {
|
|
if ctxErr := ctx.Err(); ctxErr != nil {
|
|
return ctxErr
|
|
}
|
|
return admission.ErrUnavailable
|
|
}
|
|
|
|
var status string
|
|
switch value := result.(type) {
|
|
case string:
|
|
status = value
|
|
case []byte:
|
|
status = string(value)
|
|
default:
|
|
return admission.ErrUnavailable
|
|
}
|
|
switch status {
|
|
case "ok":
|
|
return nil
|
|
case "global":
|
|
return admission.ErrGlobalLimit
|
|
case "per_key":
|
|
return admission.ErrPerKeyLimit
|
|
default:
|
|
return admission.ErrUnavailable
|
|
}
|
|
}
|
|
|
|
func hashedIdentityField(identity string) string {
|
|
var size [8]byte
|
|
binary.BigEndian.PutUint64(size[:], uint64(len(identity)))
|
|
digest := sha256.New()
|
|
_, _ = digest.Write(size[:])
|
|
_, _ = digest.Write([]byte(identity))
|
|
return "client:" + hex.EncodeToString(digest.Sum(nil))
|
|
}
|
|
|
|
func nilInterface(value any) bool {
|
|
if value == nil {
|
|
return true
|
|
}
|
|
reflected := reflect.ValueOf(value)
|
|
switch reflected.Kind() {
|
|
case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice:
|
|
return reflected.IsNil()
|
|
default:
|
|
return false
|
|
}
|
|
}
|