mirror of
https://github.com/status-im/consul.git
synced 2025-01-10 22:06:20 +00:00
d4c435856b
Adds automation for generating the map of `gRPC Method Name → Rate Limit Type` used by the middleware introduced in #15550, and will ensure we don't forget to add new endpoints. Engineers must annotate their RPCs in the proto file like so: ``` rpc Foo(FooRequest) returns (FooResponse) { option (consul.internal.ratelimit.spec) = { operation_type: READ, }; } ``` When they run `make proto` a protoc plugin `protoc-gen-consul-rate-limit` will be installed that writes rate-limit specs as a JSON array to a file called `.ratelimit.tmp` (one per protobuf package/directory). After running Buf, `make proto` will execute a post-process script that will ingest all of the `.ratelimit.tmp` files and generate a Go file containing the mappings in the `agent/grpc-middleware` package. In the enterprise repository, it will write an additional file with the enterprise-only endpoints. If an engineer forgets to add the annotation to a new RPC, the plugin will return an error like so: ``` RPC Foo is missing rate-limit specification, fix it with: import "proto-public/annotations/ratelimit/ratelimit.proto"; service Bar { rpc Foo(...) returns (...) { option (hashicorp.consul.internal.ratelimit.spec) = { operation_type: OPERATION_READ | OPERATION_WRITE | OPERATION_EXEMPT, }; } } ``` In the future, this annotation can be extended to support rate-limit category (e.g. KV vs Catalog) and to determine the retry policy.
121 lines
3.1 KiB
Go
121 lines
3.1 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/require"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
"google.golang.org/grpc/status"
|
|
|
|
pbacl "github.com/hashicorp/consul/proto-public/pbacl"
|
|
|
|
"github.com/hashicorp/go-hclog"
|
|
|
|
"github.com/hashicorp/consul/agent/consul/rate"
|
|
)
|
|
|
|
func TestServerRateLimiterMiddleware_Integration(t *testing.T) {
|
|
limiter := rate.NewMockRequestLimitsHandler(t)
|
|
|
|
logger := hclog.NewNullLogger()
|
|
server := grpc.NewServer(
|
|
grpc.InTapHandle(ServerRateLimiterMiddleware(limiter, NewPanicHandler(logger), logger)),
|
|
)
|
|
pbacl.RegisterACLServiceServer(server, mockACLServer{})
|
|
|
|
lis, err := net.Listen("tcp", "127.0.0.1:0")
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() {
|
|
if err := lis.Close(); err != nil {
|
|
t.Logf("failed to close listener: %v", err)
|
|
}
|
|
})
|
|
go server.Serve(lis)
|
|
t.Cleanup(server.Stop)
|
|
|
|
conn, err := grpc.Dial(
|
|
lis.Addr().String(),
|
|
grpc.WithTransportCredentials(insecure.NewCredentials()),
|
|
)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() {
|
|
if err := conn.Close(); err != nil {
|
|
t.Logf("failed to close client connection: %v", err)
|
|
}
|
|
})
|
|
client := pbacl.NewACLServiceClient(conn)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
t.Cleanup(cancel)
|
|
|
|
t.Run("ErrRetryElsewhere = ResourceExhausted", func(t *testing.T) {
|
|
limiter.On("Allow", mock.Anything).
|
|
Run(func(args mock.Arguments) {
|
|
op := args.Get(0).(rate.Operation)
|
|
require.Equal(t, "/hashicorp.consul.acl.ACLService/Login", op.Name)
|
|
|
|
addr := op.SourceAddr.(*net.TCPAddr)
|
|
require.True(t, addr.IP.IsLoopback())
|
|
}).
|
|
Return(rate.ErrRetryElsewhere).
|
|
Once()
|
|
|
|
_, err = client.Login(ctx, &pbacl.LoginRequest{})
|
|
require.Error(t, err)
|
|
require.Equal(t, codes.ResourceExhausted.String(), status.Code(err).String())
|
|
})
|
|
|
|
t.Run("ErrRetryLater = Unavailable", func(t *testing.T) {
|
|
limiter.On("Allow", mock.Anything).
|
|
Return(rate.ErrRetryLater).
|
|
Once()
|
|
|
|
_, err = client.Login(ctx, &pbacl.LoginRequest{})
|
|
require.Error(t, err)
|
|
require.Equal(t, codes.Unavailable.String(), status.Code(err).String())
|
|
})
|
|
|
|
t.Run("unexpected error", func(t *testing.T) {
|
|
limiter.On("Allow", mock.Anything).
|
|
Return(errors.New("uh oh")).
|
|
Once()
|
|
|
|
_, err = client.Login(ctx, &pbacl.LoginRequest{})
|
|
require.Error(t, err)
|
|
require.Equal(t, codes.Internal.String(), status.Code(err).String())
|
|
})
|
|
|
|
t.Run("operation allowed", func(t *testing.T) {
|
|
limiter.On("Allow", mock.Anything).
|
|
Return(nil).
|
|
Once()
|
|
|
|
_, err = client.Login(ctx, &pbacl.LoginRequest{})
|
|
require.NoError(t, err)
|
|
})
|
|
|
|
t.Run("Allow panics", func(t *testing.T) {
|
|
limiter.On("Allow", mock.Anything).
|
|
Panic("uh oh").
|
|
Once()
|
|
|
|
_, err = client.Login(ctx, &pbacl.LoginRequest{})
|
|
require.Error(t, err)
|
|
require.Equal(t, codes.Internal.String(), status.Code(err).String())
|
|
})
|
|
}
|
|
|
|
type mockACLServer struct {
|
|
pbacl.ACLServiceServer
|
|
}
|
|
|
|
func (mockACLServer) Login(context.Context, *pbacl.LoginRequest) (*pbacl.LoginResponse, error) {
|
|
return &pbacl.LoginResponse{}, nil
|
|
}
|