Function StreamServerInterceptor

pkg/ratelimit/grpc.go:81–102

StreamServerInterceptor is a grpc.StreamServerInterceptor that rate limits new gRPC requests and messages sent by the client. If the X-Real-IP header is not set, it is assumed that the gRPC request originates from the cluster, and no rate limits are enforced.

(limiter Interface)

// StreamServerInterceptor is a grpc.StreamServerInterceptor that rate limits new gRPC requests and messages sent by the client.
// If the X-Real-IP header is not set, it is assumed that the gRPC request originates from the cluster, and no rate limits are enforced.
func StreamServerInterceptor(limiter Interface) grpc.StreamServerInterceptor {
	return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		ctx := stream.Context()

		if grpcIsClusterAuthCall(ctx) {
			return handler(ctx, stream)
		}

		acceptResource := grpcStreamAcceptResource(ctx, info.FullMethod)
		limit, result := limiter.RateLimit(acceptResource)
		stream.SetHeader(result.GRPCHeaders()) // nolint:errcheck
		if limit {
			return errRateLimitExceeded.WithAttributes("key", acceptResource.Key(), "rate", result.Limit)
		}

		return handler(srv, &rateLimitedServerStream{
			ServerStream: stream,
			limiter:      limiter,
			resource:     grpcStreamUpResource(ctx, info.FullMethod),
		})
	}
}