mirror of
https://github.com/siderolabs/discovery-service.git
synced 2025-03-14 09:55:08 +00:00
c55142668f
Update removing multiple old middlewares, rework the way data is passed through the context, logging fields, etc. Fix minimum keepalive interval enforcement. Signed-off-by: Andrey Smirnov <andrey.smirnov@siderolabs.com>
57 lines
1.4 KiB
Go
57 lines
1.4 KiB
Go
// Copyright (c) 2024 Sidero Labs, Inc.
|
|
//
|
|
// Use of this software is governed by the Business Source License
|
|
// included in the LICENSE file.
|
|
|
|
package server
|
|
|
|
import (
|
|
"context"
|
|
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/status"
|
|
|
|
"github.com/siderolabs/discovery-service/internal/limiter"
|
|
)
|
|
|
|
func pause(ctx context.Context, limiter *limiter.IPRateLimiter) error {
|
|
iPAddr := PeerAddress(ctx)
|
|
if !IsZero(iPAddr) {
|
|
limit := limiter.Get(iPAddr)
|
|
|
|
err := limit.Wait(ctx)
|
|
if err != nil {
|
|
return status.Error(codes.ResourceExhausted, "rate limit exceeds request timeout")
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// RateLimitUnaryServerInterceptor limits Unary PRCs from an IPAdress.
|
|
func RateLimitUnaryServerInterceptor(limiter *limiter.IPRateLimiter) grpc.UnaryServerInterceptor {
|
|
return func(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
|
|
err = pause(ctx, limiter)
|
|
if err != nil {
|
|
return resp, err
|
|
}
|
|
|
|
return handler(ctx, req)
|
|
}
|
|
}
|
|
|
|
// RateLimitStreamServerInterceptor limits Stream PRCs from an IPAdress.
|
|
func RateLimitStreamServerInterceptor(limiter *limiter.IPRateLimiter) grpc.StreamServerInterceptor {
|
|
return func(srv interface{}, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
|
|
ctx := ss.Context()
|
|
|
|
err := pause(ctx, limiter)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return handler(srv, ss)
|
|
}
|
|
}
|