proxy-pool/internal/gateway/server/protection.go

package server

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"sync"

	"proxy-pool/internal/gateway/policy"
	"proxy-pool/internal/platform/httpsecurity"
)

type HTTPError struct {
	StatusCode int
	Header     http.Header
	Cause      error
}

func (err *HTTPError) Error() string {
	if err.Cause != nil {
		return err.Cause.Error()
	}
	return http.StatusText(err.StatusCode)
}

func (err *HTTPError) Unwrap() error { return err.Cause }

type ClientIPResolver = httpsecurity.ClientIPResolver

var NewClientIPResolver = httpsecurity.NewClientIPResolver

type AccessGuard struct {
	resolver     *ClientIPResolver
	allow        policy.CIDRMatcher
	unrestricted bool
}

func NewAccessGuard(resolver *ClientIPResolver, allowCIDRs []string) (*AccessGuard, error) {
	if resolver == nil {
		return nil, errors.New("create access guard: client IP resolver is required")
	}
	allow, err := policy.NewCIDRMatcher(allowCIDRs)
	if err != nil {
		return nil, fmt.Errorf("create access guard: %w", err)
	}
	return &AccessGuard{resolver: resolver, allow: allow, unrestricted: len(allowCIDRs) == 0}, nil
}

func (guard *AccessGuard) Check(_ context.Context, request *http.Request) error {
	address, err := guard.resolver.Resolve(request)
	if err != nil {
		return &HTTPError{StatusCode: http.StatusBadRequest, Cause: err}
	}
	if guard.unrestricted || guard.allow.Match(address) {
		return nil
	}
	return &HTTPError{StatusCode: http.StatusForbidden, Cause: fmt.Errorf("client address %s is not allowed", address)}
}

type Admitter interface {
	Admit(context.Context, string) error
}

type AdmissionGuard struct {
	resolver *ClientIPResolver
	admitter Admitter
}

func NewAdmissionGuard(resolver *ClientIPResolver, admitter Admitter) *AdmissionGuard {
	return &AdmissionGuard{resolver: resolver, admitter: admitter}
}

func (guard *AdmissionGuard) Check(ctx context.Context, request *http.Request) error {
	if guard == nil || guard.resolver == nil || guard.admitter == nil {
		return &HTTPError{StatusCode: http.StatusInternalServerError, Cause: errors.New("gateway admission is not configured")}
	}
	address, err := guard.resolver.Resolve(request)
	if err != nil {
		return &HTTPError{StatusCode: http.StatusBadRequest, Cause: err}
	}
	if err := guard.admitter.Admit(ctx, address.String()); err != nil {
		return &HTTPError{StatusCode: http.StatusTooManyRequests, Cause: err}
	}
	return nil
}

// CredentialAdmissionGuard enforces the request rate embedded in the
// authenticated credential policy. It executes after authentication has
// attached an immutable Client identity to the request.
type CredentialAdmissionGuard struct {
	limiters map[int]Admitter
}

type credentialReservationContextKey struct{}

type clientConcurrency struct {
	max int
	mu  sync.Mutex
	use map[string]int
}

func (limiter *clientConcurrency) acquire(clientID string) (func(), bool) {
	limiter.mu.Lock()
	defer limiter.mu.Unlock()
	if limiter.use[clientID] >= limiter.max {
		return nil, false
	}
	limiter.use[clientID]++
	var once sync.Once
	return func() {
		once.Do(func() {
			limiter.mu.Lock()
			defer limiter.mu.Unlock()
			if limiter.use[clientID] <= 1 {
				delete(limiter.use, clientID)
				return
			}
			limiter.use[clientID]--
		})
	}, true
}

type CredentialConcurrencyGuard struct {
	limiters map[int]*clientConcurrency
}

func NewCredentialConcurrencyGuard(limits map[int]struct{}) *CredentialConcurrencyGuard {
	if len(limits) == 0 {
		return nil
	}
	limiters := make(map[int]*clientConcurrency, len(limits))
	for limit := range limits {
		if limit > 0 {
			limiters[limit] = &clientConcurrency{max: limit, use: make(map[string]int)}
		}
	}
	if len(limiters) == 0 {
		return nil
	}
	return &CredentialConcurrencyGuard{limiters: limiters}
}

func (guard *CredentialConcurrencyGuard) Check(_ context.Context, request *http.Request) error {
	if guard == nil {
		return nil
	}
	identity, authenticated := httpsecurity.IdentityFromRequest(request)
	if !authenticated || identity.ClientPolicy.MaxConcurrentConnections == 0 {
		return nil
	}
	limiter := guard.limiters[identity.ClientPolicy.MaxConcurrentConnections]
	if limiter == nil || identity.ClientID == "" {
		return &HTTPError{StatusCode: http.StatusInternalServerError, Cause: errors.New("gateway credential concurrency is not configured")}
	}
	release, acquired := limiter.acquire(identity.ClientID)
	if !acquired {
		return &HTTPError{StatusCode: http.StatusTooManyRequests, Cause: errors.New("gateway credential concurrency limit exceeded")}
	}
	*request = *request.WithContext(context.WithValue(request.Context(), credentialReservationContextKey{}, release))
	return nil
}

func releaseCredentialReservation(request *http.Request) {
	if request == nil {
		return
	}
	release, found := request.Context().Value(credentialReservationContextKey{}).(func())
	if found && release != nil {
		release()
	}
}

func NewCredentialAdmissionGuard(limiters map[int]Admitter) *CredentialAdmissionGuard {
	if len(limiters) == 0 {
		return nil
	}
	cloned := make(map[int]Admitter, len(limiters))
	for limit, limiter := range limiters {
		if limit > 0 && limiter != nil {
			cloned[limit] = limiter
		}
	}
	if len(cloned) == 0 {
		return nil
	}
	return &CredentialAdmissionGuard{limiters: cloned}
}

func (guard *CredentialAdmissionGuard) Check(ctx context.Context, request *http.Request) error {
	if guard == nil {
		return nil
	}
	identity, authenticated := httpsecurity.IdentityFromRequest(request)
	if !authenticated || identity.ClientPolicy.RequestsPerMinute == 0 {
		return nil
	}
	limiter, exists := guard.limiters[identity.ClientPolicy.RequestsPerMinute]
	if !exists || limiter == nil || identity.ClientID == "" {
		return &HTTPError{StatusCode: http.StatusInternalServerError, Cause: errors.New("gateway credential admission is not configured")}
	}
	if err := limiter.Admit(ctx, identity.ClientID); err != nil {
		return &HTTPError{StatusCode: http.StatusTooManyRequests, Cause: err}
	}
	return nil
}

func chainGuards(guards ...Guard) Guard {
	filtered := make([]Guard, 0, len(guards))
	for _, guard := range guards {
		if guard != nil {
			filtered = append(filtered, guard)
		}
	}
	if len(filtered) == 0 {
		return nil
	}
	if len(filtered) == 1 {
		return filtered[0]
	}
	return GuardFunc(func(ctx context.Context, request *http.Request) error {
		for _, guard := range filtered {
			if err := guard.Check(ctx, request); err != nil {
				return err
			}
		}
		return nil
	})
}