coredns-rfc2136/ratelimit_test.go

package rfc2136

import (
	"testing"
	"time"
)

func TestRateLimiter_FirstCallAllowed(t *testing.T) {
	rl := newRateLimiter(5, time.Minute)
	now := time.Now()
	if !rl.allow("key-a", now) {
		t.Errorf("first call for new key must be allowed")
	}
}

func TestRateLimiter_BurstExhausts(t *testing.T) {
	rl := newRateLimiter(3, time.Minute)
	now := time.Now()
	// First 3 calls succeed.
	for i := 0; i < 3; i++ {
		if !rl.allow("key-a", now) {
			t.Fatalf("call %d should be allowed (burst=3)", i+1)
		}
	}
	// 4th immediately after burst should be denied (no time elapsed
	// for refill).
	if rl.allow("key-a", now) {
		t.Errorf("4th call exceeded burst; should be denied")
	}
}

func TestRateLimiter_RefillsOverTime(t *testing.T) {
	// burst=2, period=1s → refill rate is 2 tokens/sec.
	rl := newRateLimiter(2, time.Second)
	t0 := time.Now()
	if !rl.allow("k", t0) {
		t.Fatal("call 1")
	}
	if !rl.allow("k", t0) {
		t.Fatal("call 2")
	}
	if rl.allow("k", t0) {
		t.Fatal("call 3 should be denied; bucket empty")
	}
	// Advance time by 500ms — should refill ~1 token.
	if !rl.allow("k", t0.Add(500*time.Millisecond)) {
		t.Errorf("expected refill after 500ms")
	}
}

func TestRateLimiter_PerKeyIsolation(t *testing.T) {
	rl := newRateLimiter(2, time.Minute)
	now := time.Now()
	// Exhaust key-a.
	rl.allow("key-a", now)
	rl.allow("key-a", now)
	if rl.allow("key-a", now) {
		t.Fatal("key-a still has tokens; setup wrong")
	}
	// key-b is independent — must still be allowed.
	if !rl.allow("key-b", now) {
		t.Errorf("key-b was rate-limited despite no prior use")
	}
}

// TestRateLimiter_DoesNotOverflow guards against refill math
// accumulating beyond burst (which would let an attacker burst more
// after a long idle period than the configured cap).
func TestRateLimiter_DoesNotOverflow(t *testing.T) {
	rl := newRateLimiter(5, time.Second)
	t0 := time.Now()
	rl.allow("k", t0) // create bucket
	// Advance time 1 hour. Refill should cap at burst=5.
	tFuture := t0.Add(time.Hour)
	for i := 0; i < 5; i++ {
		if !rl.allow("k", tFuture) {
			t.Fatalf("post-idle call %d should be allowed (cap=5)", i+1)
		}
	}
	if rl.allow("k", tFuture) {
		t.Errorf("post-idle call 6 should be denied; cap exceeded")
	}
}