internal/server/api/ratelimit_test.go

Ref: Size: 5.5 KiB History
package api

import (
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestLimiterRefillAndCap pins the token-bucket math against a fake clock:
// a drained bucket refills one token per enrollRefillEvery, capped at burst.
func TestLimiterRefillAndCap(t *testing.T) {
	now := time.Unix(1_750_000_000, 0)
	l := newIPLimiter(func() time.Time { return now })

	for i := range enrollBurst {
		assert.True(t, l.allow("10.0.0.1"), "burst request %d", i)
	}
	assert.False(t, l.allow("10.0.0.1"), "burst exhausted")

	now = now.Add(enrollRefillEvery)
	assert.True(t, l.allow("10.0.0.1"), "one token refilled")
	assert.False(t, l.allow("10.0.0.1"), "only one refilled")

	// A long idle period refills to the cap, not beyond.
	now = now.Add(24 * time.Hour)
	for i := range enrollBurst {
		assert.True(t, l.allow("10.0.0.1"), "cap request %d", i)
	}
	assert.False(t, l.allow("10.0.0.1"), "cap enforced")
}

// TestLimiterPruneAndFailOpen pins the capacity behavior: at the entry cap,
// idle buckets are pruned to admit new IPs; when nothing is prunable the
// limiter fails OPEN (it is a brake, not the auth boundary).
func TestLimiterPruneAndFailOpen(t *testing.T) {
	now := time.Unix(1_750_000_000, 0)
	l := newIPLimiter(func() time.Time { return now })

	for i := range maxLimiterEntries {
		l.allow(fmt.Sprintf("10.%d.%d.%d", i>>16&0xff, i>>8&0xff, i&0xff))
	}
	assert.Len(t, l.buckets, maxLimiterEntries)

	// Nothing is idle yet: a new IP must still be allowed (fail open) and
	// must not grow the map.
	assert.True(t, l.allow("192.0.2.1"), "fail open when full and nothing prunable")
	assert.Len(t, l.buckets, maxLimiterEntries, "fail-open must not grow the map")

	// After everything goes idle past a full refill, pruning admits new IPs.
	now = now.Add(time.Duration(enrollBurst)*enrollRefillEvery + time.Second)
	assert.True(t, l.allow("192.0.2.2"))
	assert.Len(t, l.buckets, 1, "idle buckets pruned, new bucket tracked")
}

// The three tests below pin the limiter's numbers against literals. The rest of
// this file is written in terms of enrollBurst, enrollRefillEvery and
// maxLimiterEntries — correct as construction, but it means the whole file
// passes with the refill at 30ms and the entry cap at 10, which is a limiter
// that limits nothing.

func TestEnrollBurstIsFiveRequests(t *testing.T) {
	now := time.Unix(1_750_000_000, 0)
	l := newIPLimiter(func() time.Time { return now })

	allowed := 0
	for range 100 {
		if !l.allow("10.0.0.1") {
			break
		}
		allowed++
	}

	assert.Equal(t, 5, allowed,
		"POST /api/v1/enroll is the plane's only unauthenticated endpoint, and this burst is how many enrollment tokens one "+
			"address may guess before it has to wait. Enrollment is operator-paced — one paste per host — so five covers a "+
			"fat-fingered retry; raise it and token-guessing gets cheaper by exactly that factor, drop it to zero and no "+
			"host can ever enroll")
}

func TestEnrollRefillTakesThirtySeconds(t *testing.T) {
	now := time.Unix(1_750_000_000, 0)
	l := newIPLimiter(func() time.Time { return now })
	for range enrollBurst {
		require.True(t, l.allow("10.0.0.1"))
	}
	require.False(t, l.allow("10.0.0.1"), "burst exhausted")

	now = now.Add(29 * time.Second)
	assert.False(t, l.allow("10.0.0.1"),
		"a token must take a full 30 seconds to drip back. The burst is only a brake if the wait after it costs an attacker "+
			"real time: at a refill of milliseconds the bucket is always full and the limiter is decoration")

	now = now.Add(time.Second)
	assert.True(t, l.allow("10.0.0.1"),
		"and no longer than 30 seconds: an operator enrolling a rack should not be made to wait minutes between hosts")
}

func TestLimiterTracksTenThousandAddresses(t *testing.T) {
	now := time.Unix(1_750_000_000, 0)
	l := newIPLimiter(func() time.Time { return now })

	for i := range 10_000 {
		l.allow(fmt.Sprintf("10.%d.%d.%d", i>>16&0xff, i>>8&0xff, i&0xff))
	}

	assert.Len(t, l.buckets, 10_000,
		"the bucket map must hold ten thousand live addresses before it starts pruning. The cap is there so a spoofed-source "+
			"flood cannot grow the map without bound, but it is also how many distinct clients the limiter can brake at once: "+
			"set it low and the eleventh address evicts a real one, so a flood of fresh sources prunes the buckets holding "+
			"the actual attacker and the requests that matter fail open")
}

// TestBucketKeyGroupsIPv6BySlash64 pins the bucket-key normalization: IPv4
// (and v4-mapped-v6) keep per-address buckets; IPv6 collapses to the /64 so a
// single host's billions of addresses share one bucket; garbage stays raw.
func TestBucketKeyGroupsIPv6BySlash64(t *testing.T) {
	assert.Equal(t, "203.0.113.9", bucketKey("203.0.113.9"))
	assert.Equal(t, "203.0.113.9", bucketKey("::ffff:203.0.113.9"), "v4-mapped stays per-address")
	a := bucketKey("2001:db8:1:2:aaaa:bbbb:cccc:dddd")
	b := bucketKey("2001:db8:1:2:1111:2222:3333:4444")
	c := bucketKey("2001:db8:1:3::1")
	assert.Equal(t, a, b, "same /64 must share a bucket")
	assert.NotEqual(t, a, c, "different /64 must not")
	assert.Equal(t, "not-an-ip", bucketKey("not-an-ip"))
}

// TestLimiterSharesBucketAcrossSameSlash64 pins end behavior: burst spent
// from one v6 address exhausts the whole /64.
func TestLimiterSharesBucketAcrossSameSlash64(t *testing.T) {
	now := time.Unix(1_750_000_000, 0)
	l := newIPLimiter(func() time.Time { return now })
	for range enrollBurst {
		assert.True(t, l.allow(bucketKey("2001:db8::1")))
	}
	assert.False(t, l.allow(bucketKey("2001:db8::2")), "same /64: bucket shared")
	assert.True(t, l.allow(bucketKey("2001:db9::1")), "different /64: fresh bucket")
}