Ratelimit interceptor

README.md

@@ -71,6 +71,7 @@ myServer := grpc.NewServer(
 #### Server
    * [`grpc_validator`](validator/) - codegen inbound message validation from `.proto` options
    * [`grpc_recovery`](recovery/) - turn panics into gRPC errors
+   * [`ratelimit`](ratelimit/) - grpc rate limiting by your own limiter
 
 
 ## Status

ratelimit/doc.go

@@ -0,0 +1,12 @@
+/*
+`ratelimit` a generic server-side ratelimit middleware for gRPC.
+
+Server Side Ratelimit Middleware
+
+It allows to do grpc rate limit by your own rate limiter (e.g. token bucket, leaky bucket, etc.)
+
+`ratelimit/tokenbucket`provides an implementation based on token bucket `github.com/juju/ratelimit`.
+
+Please see examples for simple examples of use.
+*/
+package ratelimit

ratelimit/examples_test.go

@@ -0,0 +1,26 @@
+package ratelimit_test
+
+import (
+	"time"
+
+	"github.com/ceshihao/ratelimiter/tokenbucket"
+	"github.com/grpc-ecosystem/go-grpc-middleware"
+	"github.com/grpc-ecosystem/go-grpc-middleware/ratelimit"
+	"google.golang.org/grpc"
+)
+
+// Simple example of server initialization code.
+func Example() {
+	// Create unary/stream rateLimiters, based on token bucket here.
+	// You can implement your own ratelimiter for the interface.
+	unaryRateLimiter := tokenbucket.NewTokenBucketRateLimiter(1*time.Second, 10, 10, 10*time.Second)
+	streamRateLimiter := tokenbucket.NewTokenBucketRateLimiter(1*time.Second, 5, 5, 5*time.Second)
+	_ = grpc.NewServer(
+		grpc_middleware.WithUnaryServerChain(
+			ratelimit.UnaryServerInterceptor(unaryRateLimiter),
+		),
+		grpc_middleware.WithStreamServerChain(
+			ratelimit.StreamServerInterceptor(streamRateLimiter),
+		),
+	)
+}

ratelimit/ratelimit.go

@@ -0,0 +1,35 @@
+package ratelimit
+
+import (
+	"golang.org/x/net/context"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+// Limiter defines the interface to perform request rate limiting.
+// If Limit function return true, the request will be rejected.
+// Otherwise, the request will pass.
+type Limiter interface {
+	Limit() bool
+}
+
+// UnaryServerInterceptor returns a new unary server interceptors that performs request rate limiting.
+func UnaryServerInterceptor(limiter Limiter) grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+		if limiter.Limit() {
+			return nil, status.Errorf(codes.ResourceExhausted, "%s is rejected by grpc_ratelimit middleare, please retry later.", info.FullMethod)
+		}
+		return handler(ctx, req)
+	}
+}
+
+// StreamServerInterceptor returns a new stream server interceptor that performs rate limiting on the request.
+func StreamServerInterceptor(limiter Limiter) grpc.StreamServerInterceptor {
+	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+		if limiter.Limit() {
+			return status.Errorf(codes.ResourceExhausted, "%s is rejected by grpc_ratelimit middleare, please retry later.", info.FullMethod)
+		}
+		return handler(srv, stream)
+	}
+}

ratelimit/ratelimit_test.go

@@ -0,0 +1,75 @@
+package ratelimit
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"google.golang.org/grpc"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const errMsgFake = "fake error"
+
+type mockPassLimiter struct{}
+
+func (*mockPassLimiter) Limit() bool {
+	return false
+}
+
+func TestUnaryServerInterceptor_RateLimitPass(t *testing.T) {
+	interceptor := UnaryServerInterceptor(&mockPassLimiter{})
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return nil, errors.New(errMsgFake)
+	}
+	info := &grpc.UnaryServerInfo{
+		FullMethod: "FakeMethod",
+	}
+	req, err := interceptor(nil, nil, info, handler)
+	assert.Nil(t, req)
+	assert.EqualError(t, err, errMsgFake)
+}
+
+type mockFailLimiter struct{}
+
+func (*mockFailLimiter) Limit() bool {
+	return true
+}
+
+func TestUnaryServerInterceptor_RateLimitFail(t *testing.T) {
+	interceptor := UnaryServerInterceptor(&mockFailLimiter{})
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return nil, errors.New(errMsgFake)
+	}
+	info := &grpc.UnaryServerInfo{
+		FullMethod: "FakeMethod",
+	}
+	req, err := interceptor(nil, nil, info, handler)
+	assert.Nil(t, req)
+	assert.EqualError(t, err, "rpc error: code = ResourceExhausted desc = FakeMethod is rejected by grpc_ratelimit middleare, please retry later.")
+}
+
+func TestStreamServerInterceptor_RateLimitPass(t *testing.T) {
+	interceptor := StreamServerInterceptor(&mockPassLimiter{})
+	handler := func(srv interface{}, stream grpc.ServerStream) error {
+		return errors.New(errMsgFake)
+	}
+	info := &grpc.StreamServerInfo{
+		FullMethod: "FakeMethod",
+	}
+	err := interceptor(nil, nil, info, handler)
+	assert.EqualError(t, err, errMsgFake)
+}
+
+func TestStreamServerInterceptor_RateLimitFail(t *testing.T) {
+	interceptor := StreamServerInterceptor(&mockFailLimiter{})
+	handler := func(srv interface{}, stream grpc.ServerStream) error {
+		return errors.New(errMsgFake)
+	}
+	info := &grpc.StreamServerInfo{
+		FullMethod: "FakeMethod",
+	}
+	err := interceptor(nil, nil, info, handler)
+	assert.EqualError(t, err, "rpc error: code = ResourceExhausted desc = FakeMethod is rejected by grpc_ratelimit middleare, please retry later.")
+}