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