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

package server

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

	"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
}

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
	})
}