grpc/interceptors/ban/options.go

package ban

import (
	"context"
	"net"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/peer"
)

const (
	Unauthorized    = "Unauthorized"
	Unauthenticated = "Unauthenticated"
)

var (
	defaultOptions = options{
		cap:                 1024,
		reaperInterval:      10 * time.Minute,
		rules:               defaultRules,
		actorFunc:           DefaultActorFunc,
		defaultCallback:     nil,
		defaultJailDuration: 10 * time.Second,
	}

	defaultRules = []Rule{
		{
			Name:        Unauthorized,
			Message:     "Too many unauthorized requests",
			Code:        codes.PermissionDenied,
			StrikeLimit: 3,
		},
		{
			Name:        Unauthenticated,
			Message:     "Too many unauthenticated requests",
			Code:        codes.Unauthenticated,
			StrikeLimit: 3,
		},
	}
)

func DefaultActorFunc(ctx context.Context) (string, bool, error) {
	p, ok := peer.FromContext(ctx)
	if !ok {
		return "", false, nil
	}
	if host, _, err := net.SplitHostPort(p.Addr.String()); err == nil {
		return host, true, nil
	}
	return p.Addr.String(), true, nil
}

type Option func(*options)

func WithCapacity(cap int32) Option {
	return func(o *options) {
		o.cap = cap
	}
}

func WithRules(rules ...Rule) Option {
	return func(o *options) {
		o.rules = rules
	}
}

func WithReaperInterval(interval time.Duration) Option {
	return func(o *options) {
		o.reaperInterval = interval
	}
}

func WithActorFunc(f func(context.Context) (name string, found bool, err error)) Option {
	return func(o *options) {
		o.actorFunc = f
	}
}

func WithDefaultCallback(f ActionCallback) Option {
	return func(o *options) {
		o.defaultCallback = f
	}
}

func WithDefaultJailDuration(expire time.Duration) Option {
	return func(o *options) {
		o.defaultJailDuration = expire
	}
}

type options struct {
	cap                 int32
	rules               []Rule
	reaperInterval      time.Duration
	actorFunc           func(ctx context.Context) (name string, found bool, err error)
	defaultCallback     ActionCallback
	defaultJailDuration time.Duration
}