gravitational/teleport
1
0
Fork 0
mirror of https://github.com/gravitational/teleport.git synced 2026-10-12 02:15:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

limiter: Fix custom rate bucket clobbering (#64828)

Browse Source 
* limiter: Add failing test for custom rate bucket clobbering

Add an assertion to TestCustomRate that proves the bug: when a
default-rate RegisterRequest call interleaves with custom-rate calls
for the same client IP, the custom-rate bucket is clobbered because
both share a single TokenBucketSet keyed by IP alone.

The test fails on master at limiter_test.go:145. The next commit
fixes the bug.

* limiter: Use dedicated RateLimiter for account recovery

Replace the shared-bucket custom rate approach with a separate
RateLimiter instance for account recovery RPCs. This follows the
same pattern as gravitational/teleport#64559 and avoids the
shared-bucket clobber problem entirely: each concern gets its own
RateLimiter with its own FnCache, so default-rate requests cannot
reset stricter per-endpoint buckets.

The `getCustomRate` middleware passed per-endpoint rate sets through
a shared FnCache keyed only by client IP. When a default-rate
request hit the same IP, `TokenBucketSet.Update` overwrote the
cached bucket with the default rate parameters, resetting the
stricter custom-rate state.

Add a private `accountRecoveryLimiter` to the auth Middleware and
wire it into `rateLimitUnaryInterceptor`, which applies the default
limiter to all endpoints and the recovery limiter on top. Remove
`getCustomRate`, `CustomRateFunc`, `UnaryServerInterceptorWithCustomRate`,
and the `customRate` parameter from `RateLimiter.RegisterRequest` and
`RegisterRequestFromAddr`. Export `ClientIPFromContext` from the limiter
package for use by the middleware.

Keep a deprecated `Limiter.RegisterRequestWithCustomRate` shim for
backwards compatibility with enterprise callers (they all pass nil).

* limiter: Remove stale accountRecoveryEndpoints comment

Remove orphaned comment from the deleted `accountRecoveryEndpoints`
variable. The comment was left behind when the variable was replaced
by inline switch cases in `rateLimitUnaryInterceptor`.

* limiter: Guard nil accountRecoveryLimiter in interceptor

The accountRecoveryLimiter is nil when Middleware is constructed
without calling newAccountRecoveryLimiter, which happens in
initSecureGRPCServer for Kube-only traffic. The nil dereference is
not reachable on current code paths because recovery RPCs are never
routed to that server, but add a nil guard to make the interceptor
safe independently of how the server is wired.

* limiter: Remove RegisterRequestWithCustomRate

Remove the backwards-compatibility shim now that all enterprise
callers have been removed via gravitational/teleport.e#8359.

* limiter: Extract ClientIPFromAddr to lib/utils

Move the addr-to-IP parsing into `utils.ClientIPFromAddr` (next to
`ClientIPFromConn`) so `lib/auth` does not depend on `lib/limiter` for
IP extraction. Add an unexported `clientIPFromContext` in `lib/auth`
for the gRPC peer lookup.

Reduce the public API surface of `lib/limiter` that deals with gRPC
concepts, keeping it closer to its intended role as a general-purpose
rate limiter keyed by string tokens.
This commit is contained in:
Julia Ogris
2026-03-31 01:37:58 +00:00
committed by GitHub
parent 51f4d3bedd
commit 4df912cce9
10 changed files with 159 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lib/auth/grpcserver.go
+1 -1
View File
@@ -4889,7 +4889,7 @@ func (g *GRPCServer) CreateAuthenticateChallenge(ctx context.Context, req *authp
return nil, trace.BadParameter("unable to find peer")
}
if err := g.createAuthenticateChallengeLimiter.RegisterRequestFromAddr(peerInfo.Addr, nil); err != nil {
if err := g.createAuthenticateChallengeLimiter.RegisterRequestFromAddr(peerInfo.Addr); err != nil {
return nil, trace.LimitExceeded("rate limit exceeded")
}
}
lib/auth/grpcserver_test.go
+6
View File
@@ -4755,6 +4755,12 @@ func TestCustomRateLimiting(t *testing.T) {
return err
},
},
{
name: "RPC ChangePassword",
fn: func(ctx context.Context, clt *authclient.Client) error {
return clt.ChangePassword(ctx, &proto.ChangePasswordRequest{})
},
},
{
name: "RPC GetAccountRecoveryToken",
fn: func(ctx context.Context, clt *authclient.Client) error {
lib/auth/middleware.go
+57 -19
View File
@@ -173,6 +173,11 @@ func NewTLSServer(ctx context.Context, cfg TLSServerConfig) (*TLSServer, error)
return nil, trace.Wrap(err)
}
accountRecoveryLimiter, err := newAccountRecoveryLimiter()
if err != nil {
return nil, trace.Wrap(err)
}
// authMiddleware authenticates request assuming TLS client authentication
// adds authentication information to the context
// and passes it to the API server
@@ -183,6 +188,7 @@ func NewTLSServer(ctx context.Context, cfg TLSServerConfig) (*TLSServer, error)
Handler: apiServer,
},
Limiter: limiter,
accountRecoveryLimiter: accountRecoveryLimiter,
GRPCMetrics: grpcMetrics,
OldestSupportedVersion: oldestSupportedVersion,
AlertCreator: func(ctx context.Context, a types.ClusterAlert) error {
@@ -342,6 +348,9 @@ type Middleware struct {
authz.Middleware
// Limiter is a rate and connection limiter
Limiter *limiter.Limiter
// accountRecoveryLimiter rate-limits account recovery RPCs
// independently from the default limiter.
accountRecoveryLimiter *limiter.RateLimiter
// GRPCMetrics is the configured gRPC metrics for the interceptors
GRPCMetrics *grpcprom.ServerMetrics
@@ -358,26 +367,55 @@ type Middleware struct {
lastRejectedAlertTime atomic.Int64
}
func getCustomRate(endpoint string) *limiter.RateSet {
switch endpoint {
// Account recovery RPCs.
case
"/proto.AuthService/ChangeUserAuthentication",
"/proto.AuthService/ChangePassword",
"/proto.AuthService/GetAccountRecoveryToken",
"/proto.AuthService/StartAccountRecovery",
"/proto.AuthService/VerifyAccountRecovery":
rates := limiter.NewRateSet()
// This limit means: 1 request per minute with bursts up to 10 requests.
if err := rates.Add(time.Minute, 1, 10); err != nil {
logger.DebugContext(context.Background(), "Failed to define a custom rate for rpc method, using default rate",
"error", err,
"rpc_method", endpoint)
return nil
// newAccountRecoveryLimiter creates the dedicated limiter for
// account recovery RPCs (1 req/min, burst 10).
func newAccountRecoveryLimiter() (*limiter.RateLimiter, error) {
return limiter.NewRateLimiter(limiter.Config{
Rates: []limiter.Rate{{Period: time.Minute, Average: 1, Burst: 10}},
})
}
// clientIPFromContext extracts the client IP from the gRPC peer in
// the given context.
func clientIPFromContext(ctx context.Context) (string, error) {
peerInfo, ok := peer.FromContext(ctx)
if !ok {
return "", trace.AccessDenied("missing peer info")
}
return utils.ClientIPFromAddr(peerInfo.Addr)
}
// rateLimitUnaryInterceptor returns a gRPC unary interceptor that
// applies the default rate and connection limiter to all endpoints,
// plus the stricter account recovery limiter on recovery endpoints.
func (a *Middleware) rateLimitUnaryInterceptor() grpc.UnaryServerInterceptor {
return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
clientIP, err := clientIPFromContext(ctx)
if err != nil {
return nil, trace.Wrap(err)
}
return rates
release, err := a.Limiter.RegisterRequestAndConnection(clientIP)
if err != nil {
return nil, trace.LimitExceeded("rate limit exceeded")
}
defer release()
// Additional rate check for account recovery endpoints.
switch info.FullMethod {
case
"/proto.AuthService/ChangeUserAuthentication",
"/proto.AuthService/ChangePassword",
"/proto.AuthService/GetAccountRecoveryToken",
"/proto.AuthService/StartAccountRecovery",
"/proto.AuthService/VerifyAccountRecovery":
if a.accountRecoveryLimiter != nil {
err := a.accountRecoveryLimiter.RegisterRequest(clientIP)
if err != nil {
return nil, trace.LimitExceeded("rate limit exceeded")
}
}
}
return handler(ctx, req)
}
return nil
}
// ValidateClientVersion inspects the client version for the connection and terminates
@@ -573,7 +611,7 @@ func (a *Middleware) UnaryInterceptors() []grpc.UnaryServerInterceptor {
return append(is,
interceptors.GRPCServerUnaryErrorInterceptor,
metadata.UnaryServerInterceptor,
a.Limiter.UnaryServerInterceptorWithCustomRate(getCustomRate),
a.rateLimitUnaryInterceptor(),
a.withAuthenticatedUserUnaryInterceptor,
)
}
lib/limiter/internal/ratelimit/ratelimit.go
+3 -2
View File
@@ -134,7 +134,8 @@ func (tl *TokenLimiter) consumeRates(source string, amount int64) error {
return nil
}
// RateSet maintains a set of rates. It can contain only one rate per period at a time.
// RateSet maintains a set of rates. It can contain only one rate per
// period at a time.
type RateSet struct {
m map[time.Duration]*rate
}
@@ -146,7 +147,7 @@ type rate struct {
burst int64
}
// NewRateSet crates an empty rate set.
// NewRateSet creates an empty rate set.
func NewRateSet() *RateSet {
return &RateSet{m: make(map[time.Duration]*rate)}
}
lib/limiter/internal/ratelimit/tokenbucket.go
+1 -1
View File
@@ -137,7 +137,7 @@ type tokenBucket struct {
lastConsumed int64
}
// newTokenBucket crates a tokenBucket instance for the specified rate.
// newTokenBucket creates a tokenBucket instance for the specified rate.
func newTokenBucket(rate *rate, clock clockwork.Clock) *tokenBucket {
period := cmp.Or(rate.period, time.Nanosecond)
return &tokenBucket{
lib/limiter/limiter.go
+6 -44
View File
@@ -30,6 +30,7 @@ import (
"google.golang.org/grpc/peer"
"github.com/gravitational/teleport/lib/limiter/internal/ratelimit"
"github.com/gravitational/teleport/lib/utils"
)
// Limiter helps limiting connections and request rates
@@ -71,11 +72,7 @@ func (l *Limiter) GetNumConnection(token string) (int64, error) {
}
func (l *Limiter) RegisterRequest(token string) error {
return l.rateLimiter.RegisterRequest(token, nil)
}
func (l *Limiter) RegisterRequestWithCustomRate(token string, customRate *ratelimit.RateSet) error {
return l.rateLimiter.RegisterRequest(token, customRate)
return l.rateLimiter.RegisterRequest(token)
}
func (l *Limiter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -113,37 +110,23 @@ func (l *Limiter) RegisterRequestAndConnection(token string) (func(), error) {
type RateSet = ratelimit.RateSet
// NewRateSet crates an empty `RateSet` instance.
// NewRateSet creates an empty RateSet instance.
func NewRateSet() *RateSet { return ratelimit.NewRateSet() }
// UnaryServerInterceptor returns a gRPC unary interceptor which
// rate limits by client IP.
func (l *Limiter) UnaryServerInterceptor() grpc.UnaryServerInterceptor {
return l.UnaryServerInterceptorWithCustomRate(func(string) *RateSet {
return nil
})
}
// CustomRateFunc is a function type which returns a custom rate set
// for a given endpoint string.
type CustomRateFunc func(endpoint string) *RateSet
// UnaryServerInterceptorWithCustomRate returns a gRPC unary interceptor which
// rate limits by client IP. Accepts a CustomRateFunc to set custom rates for
// specific gRPC methods.
func (l *Limiter) UnaryServerInterceptorWithCustomRate(customRate CustomRateFunc) grpc.UnaryServerInterceptor {
return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
peerInfo, ok := peer.FromContext(ctx)
if !ok {
return nil, trace.AccessDenied("missing peer info")
}
// Limit requests per second and simultaneous connection by client IP.
clientIP, err := clientIPFromAddr(peerInfo.Addr)
clientIP, err := utils.ClientIPFromAddr(peerInfo.Addr)
if err != nil {
return nil, trace.Wrap(err)
}
if err := l.RegisterRequestWithCustomRate(clientIP, customRate(info.FullMethod)); err != nil {
if err := l.RegisterRequest(clientIP); err != nil {
return nil, trace.LimitExceeded("rate limit exceeded")
}
if err := l.connectionLimiter.AcquireConnection(clientIP); err != nil {
@@ -162,7 +145,7 @@ func (l *Limiter) StreamServerInterceptor(srv any, serverStream grpc.ServerStrea
return trace.AccessDenied("missing peer info")
}
// Limit requests per second and simultaneous connection by client IP.
clientIP, err := clientIPFromAddr(peerInfo.Addr)
clientIP, err := utils.ClientIPFromAddr(peerInfo.Addr)
if err != nil {
return trace.Wrap(err)
}
@@ -176,27 +159,6 @@ func (l *Limiter) StreamServerInterceptor(srv any, serverStream grpc.ServerStrea
return handler(srv, serverStream)
}
func clientIPFromAddr(addr net.Addr) (string, error) {
if addr == nil {
return "", trace.BadParameter("missing client IP")
}
s := addr.String()
// bufconn peers don't include host:port, so use a stable synthetic key
// for request/connection limiting in tests.
if s == "bufconn" && addr.Network() == "bufconn" {
return "bufconn", nil
}
clientIP, _, err := net.SplitHostPort(s)
if err == nil {
return clientIP, nil
}
return "", trace.BadParameter("missing client IP")
}
// WrapListener returns a [Listener] that wraps the provided listener
// with one that limits connections
func (l *Limiter) WrapListener(ln net.Listener) (*Listener, error) {
lib/limiter/limiter_test.go
+50 -81
View File
@@ -34,7 +34,6 @@ import (
"google.golang.org/grpc"
"google.golang.org/grpc/peer"
"github.com/gravitational/teleport/lib/limiter/internal/ratelimit"
"github.com/gravitational/teleport/lib/utils"
"github.com/gravitational/teleport/lib/utils/log/logtest"
)
@@ -110,41 +109,6 @@ func TestRateLimiter(t *testing.T) {
require.Error(t, err)
}
func TestCustomRate(t *testing.T) {
clock := clockwork.NewFakeClock()
limiter, err := NewLimiter(
Config{
Clock: clock,
Rates: []Rate{
// Default rate
{
Period: 10 * time.Millisecond,
Average: 10,
Burst: 20,
},
},
})
require.NoError(t, err)
customRate := ratelimit.NewRateSet()
err = customRate.Add(time.Minute, 1, 5)
require.NoError(t, err)
// Max out custom rate.
for range 5 {
require.NoError(t, limiter.RegisterRequestWithCustomRate("token1", customRate))
}
// Test rate limit exceeded with custom rate.
require.Error(t, limiter.RegisterRequestWithCustomRate("token1", customRate))
// Test default rate still works.
for range 20 {
require.NoError(t, limiter.RegisterRequest("token1"))
}
}
type mockAddr struct{}
func (a mockAddr) Network() string {
@@ -156,56 +120,24 @@ func (a mockAddr) String() string {
}
func TestLimiter_UnaryServerInterceptor(t *testing.T) {
ctx := peer.NewContext(t.Context(), &peer.Peer{Addr: mockAddr{}})
req := "request"
serverInfo := &grpc.UnaryServerInfo{FullMethod: "/method"}
handler := func(context.Context, any) (any, error) { return nil, nil }
limiter, err := NewLimiter(Config{
MaxConnections: 1,
Rates: []Rate{
{
Period: time.Minute,
Average: 1,
Burst: 1,
},
},
Rates: []Rate{{Period: time.Minute, Average: 1, Burst: 1}},
})
require.NoError(t, err)
ctx := peer.NewContext(context.Background(), &peer.Peer{Addr: mockAddr{}})
req := "request"
serverInfo := &grpc.UnaryServerInfo{
FullMethod: "/method",
}
handler := func(context.Context, any) (any, error) { return nil, nil }
interceptor := limiter.UnaryServerInterceptor()
unaryInterceptor := limiter.UnaryServerInterceptor()
// pass at least once
_, err = unaryInterceptor(ctx, req, serverInfo, handler)
_, err = interceptor(ctx, req, serverInfo, handler)
require.NoError(t, err)
// should eventually fail, not testing the limiter behavior here
for range 10 {
_, err = unaryInterceptor(ctx, req, serverInfo, handler)
if err != nil {
break
}
}
require.Error(t, err)
getCustomRate := func(endpoint string) *ratelimit.RateSet {
rates := ratelimit.NewRateSet()
err := rates.Add(2*time.Minute, 1, 2)
require.NoError(t, err)
return rates
}
unaryInterceptor = limiter.UnaryServerInterceptorWithCustomRate(getCustomRate)
// should pass at least once
_, err = unaryInterceptor(ctx, req, serverInfo, handler)
require.NoError(t, err)
// should eventually fail, not testing the limiter behavior here
for range 10 {
_, err = unaryInterceptor(ctx, req, serverInfo, handler)
_, err = interceptor(ctx, req, serverInfo, handler)
if err != nil {
break
}
@@ -235,7 +167,7 @@ func TestLimiter_StreamServerInterceptor(t *testing.T) {
})
require.NoError(t, err)
ctx := peer.NewContext(context.Background(), &peer.Peer{Addr: mockAddr{}})
ctx := peer.NewContext(t.Context(), &peer.Peer{Addr: mockAddr{}})
ss := mockServerStream{
ctx: ctx,
}
@@ -479,19 +411,19 @@ func TestRateLimiter_IsRateLimited(t *testing.T) {
// Consume some tokens but not all
for range 5 {
require.NoError(t, limiter.RegisterRequest("token1", nil))
require.NoError(t, limiter.RegisterRequest("token1"))
}
require.False(t, limiter.IsRateLimited("token1"))
// Consume the rest of the tokens
for range 4 {
require.NoError(t, limiter.RegisterRequest("token1", nil))
require.NoError(t, limiter.RegisterRequest("token1"))
}
require.False(t, limiter.IsRateLimited("token1"))
// Consume the last token
require.NoError(t, limiter.RegisterRequest("token1", nil))
require.NoError(t, limiter.RegisterRequest("token1"))
// Now token1 should be rate limited
require.True(t, limiter.IsRateLimited("token1"))
// token2 should not be rate limited
@@ -501,3 +433,40 @@ func TestRateLimiter_IsRateLimited(t *testing.T) {
// After time passes, token1 should not be rate limited anymore
require.False(t, limiter.IsRateLimited("token1"))
}
// TestIndependentLimiters verifies that two Limiter instances
// maintain independent token buckets for the same client IP.
func TestIndependentLimiters(t *testing.T) {
t.Parallel()
clock := clockwork.NewFakeClock()
defaultLimiter, err := NewLimiter(Config{
Clock: clock,
Rates: []Rate{{Period: 10 * time.Millisecond, Average: 10, Burst: 20}},
})
require.NoError(t, err)
recoveryLimiter, err := NewLimiter(Config{
Clock: clock,
Rates: []Rate{{Period: time.Minute, Average: 1, Burst: 5}},
})
require.NoError(t, err)
// Exhaust the recovery limiter (burst 5).
for range 5 {
require.NoError(t, recoveryLimiter.RegisterRequest("127.0.0.1"))
}
require.Error(t, recoveryLimiter.RegisterRequest("127.0.0.1"))
// Default limiter for the same IP must still work.
require.NoError(t, defaultLimiter.RegisterRequest("127.0.0.1"))
// Recovery limiter must still be exhausted.
require.Error(t, recoveryLimiter.RegisterRequest("127.0.0.1"))
// Default limiter has its own independent bucket.
for range 19 {
require.NoError(t, defaultLimiter.RegisterRequest("127.0.0.1"))
}
require.Error(t, defaultLimiter.RegisterRequest("127.0.0.1"))
}
lib/limiter/ratelimiter.go
+14 -15
View File
@@ -19,7 +19,6 @@
package limiter
import (
"cmp"
"context"
"net"
"net/http"
@@ -119,27 +118,26 @@ func (l *RateLimiter) IsRateLimited(token string) bool {
return bucket.IsRateLimited()
}
// RegisterRequest increases number of requests for the provided token
// Returns error if there are too many requests with the provided token.
func (l *RateLimiter) RegisterRequest(token string, customRate *ratelimit.RateSet) error {
// RegisterRequest increases number of requests for the provided token.
// It returns an error if there are too many requests with the provided
// token.
func (l *RateLimiter) RegisterRequest(token string) error {
l.mu.Lock()
defer l.mu.Unlock()
rate := cmp.Or(customRate, l.rates)
// We set the TTL as 10 times the rate period. E.g. if rate is 100 requests/second
// per client IP, the counters for this IP will expire after 10 seconds of inactivity.
ttl := rate.MaxPeriod()*10 + 1
ttl := l.rates.MaxPeriod()*10 + 1
bucketSet, err := utils.FnCacheGetWithTTL(context.TODO(), l.rateLimits, token, ttl,
func(ctx context.Context) (*ratelimit.TokenBucketSet, error) {
return ratelimit.NewTokenBucketSet(rate, l.clock), nil
return ratelimit.NewTokenBucketSet(l.rates, l.clock), nil
},
)
if err != nil {
return trace.Wrap(err)
}
bucketSet.Update(rate)
bucketSet.Update(l.rates)
delay, err := bucketSet.Consume(1)
if err != nil {
@@ -151,15 +149,16 @@ func (l *RateLimiter) RegisterRequest(token string, customRate *ratelimit.RateSe
return nil
}
// RegisterRequestFromAddr increases the number of requests coming for the given
// remote address, returning an error if the address is invalid or if there have
// been too many requests from that address recently.
func (l *RateLimiter) RegisterRequestFromAddr(addr net.Addr, customRate *ratelimit.RateSet) error {
token, err := clientIPFromAddr(addr)
// RegisterRequestFromAddr increases the number of requests coming
// for the given remote address, returning an error if the address is
// invalid or if there have been too many requests from that address
// recently.
func (l *RateLimiter) RegisterRequestFromAddr(addr net.Addr) error {
token, err := utils.ClientIPFromAddr(addr)
if err != nil {
return trace.Wrap(err)
}
return l.RegisterRequest(token, customRate)
return l.RegisterRequest(token)
}
// Add rate limiter to the handle
lib/utils/net.go
+19
View File
@@ -25,6 +25,25 @@ import (
"github.com/gravitational/trace"
)
// ClientIPFromAddr extracts the client IP from a network address.
// For bufconn test addresses it returns the literal string "bufconn".
func ClientIPFromAddr(addr net.Addr) (string, error) {
if addr == nil {
return "", trace.BadParameter("missing client IP")
}
s := addr.String()
// bufconn peers don't include host:port, so use a stable synthetic
// key for request/connection limiting in tests.
if s == "bufconn" && addr.Network() == "bufconn" {
return "bufconn", nil
}
clientIP, _, err := net.SplitHostPort(s)
if err != nil {
return "", trace.BadParameter("missing client IP")
}
return clientIP, nil
}
// ClientIPFromConn extracts host from provided remote address.
func ClientIPFromConn(conn net.Conn) (string, error) {
clientRemoteAddr := conn.RemoteAddr()
lib/web/apiserver.go
+2 -2
View File
@@ -5319,7 +5319,7 @@ func (h *Handler) conditionalLimiterFunc(fn httplib.HandlerFunc, limiter *limite
result, err := fn(w, r, p)
if err != nil {
if shouldLimit(err) {
if err := limiter.RegisterRequest(remoteAddr, nil); err != nil {
if err := limiter.RegisterRequest(remoteAddr); err != nil {
return nil, trace.Wrap(err)
}
}
@@ -5381,7 +5381,7 @@ func rateLimitRequest(r *http.Request, limiter *limiter.RateLimiter) error {
if err != nil {
return trace.Wrap(err)
}
return trace.Wrap(limiter.RegisterRequest(remote, nil /* customRate */))
return trace.Wrap(limiter.RegisterRequest(remote))
}
func (h *Handler) validateCookie(w http.ResponseWriter, r *http.Request) (*SessionContext, error) {