echolot/server/internal/session/grant_test.go

// SPDX-FileCopyrightText: 2026 Echolot contributors
// SPDX-License-Identifier: GPL-3.0-or-later

package session

import (
	"net/netip"
	"testing"
	"time"
)

func sessionWithSource(t *testing.T) *Session {
	t.Helper()
	m := NewManager(time.Minute)
	s, _, err := m.New("dev", "cred", netip.MustParseAddr("127.0.0.1"))
	if err != nil {
		t.Fatal(err)
	}
	s.NoteDataSource(netip.MustParseAddrPort("127.0.0.1:5000"))
	return s
}

// Without an observed data-plane source there is nowhere verified to send, so no grant may exist.
func TestNoGrantWithoutObservedSource(t *testing.T) {
	m := NewManager(time.Minute)
	s, _, _ := m.New("dev", "cred", netip.MustParseAddr("127.0.0.1"))
	if g := s.NewGrant("a1", 1000, 1000, DefaultGrantLimits); g != nil {
		t.Fatal("granted a send with no verified destination")
	}
}

func TestGrantIsBoundToObservedSourceAndClamped(t *testing.T) {
	s := sessionWithSource(t)
	g := s.NewGrant("a1", 1<<40, 1<<30, DefaultGrantLimits) // absurd request
	if g == nil {
		t.Fatal("expected a grant")
	}
	if g.Dest != "127.0.0.1:5000" {
		t.Fatalf("grant destination = %s, want the observed source", g.Dest)
	}
	if g.MaxBytes != DefaultGrantLimits.MaxBytes || g.MaxKbps != DefaultGrantLimits.MaxKbps {
		t.Fatalf("client request was not clamped to server limits: %d bytes / %d kbps",
			g.MaxBytes, g.MaxKbps)
	}
}

// The byte budget must actually stop sending — this is the anti-amplification guarantee.
func TestGrantStopsAtByteBudget(t *testing.T) {
	s := sessionWithSource(t)
	g := s.NewGrant("a1", 1000, 0, DefaultGrantLimits)
	sent := 0
	for i := 0; i < 100; i++ {
		if !g.Allow(100) {
			break
		}
		sent += 100
	}
	if sent != 1000 {
		t.Fatalf("sent %d bytes, want exactly the 1000-byte budget", sent)
	}
	if g.Allow(1) {
		t.Fatal("grant allowed a send after the budget was exhausted")
	}
	if g.Sent() != 1000 {
		t.Fatalf("Sent() = %d, want 1000", g.Sent())
	}
}

func TestExpiredGrantRefuses(t *testing.T) {
	s := sessionWithSource(t)
	g := s.NewGrant("a1", 10_000, 0, GrantLimits{MaxBytes: 10_000, MaxKbps: 1000, MaxHold: time.Millisecond})
	time.Sleep(5 * time.Millisecond)
	if g.Allow(10) {
		t.Fatal("expired grant still allowed a send")
	}
}

// The rate ceiling must throttle a burst that is well inside the byte budget.
func TestGrantEnforcesRate(t *testing.T) {
	s := sessionWithSource(t)
	// 8 kbps = 1000 bytes/s. A burst far beyond one second's worth must be refused.
	g := s.NewGrant("a1", 1<<20, 8, DefaultGrantLimits)
	time.Sleep(60 * time.Millisecond) // let the rate window open a little
	sent := 0
	for i := 0; i < 1000; i++ {
		if !g.Allow(100) {
			break
		}
		sent += 100
	}
	if sent > 5000 {
		t.Fatalf("rate limit let %d bytes through in ~60ms at 8kbps", sent)
	}
}

// The bug this pins: the rate check used to exempt the first 50 ms entirely, so a sender could
// dump an unbounded burst into that window and then be refused for as long as it took real time
// to catch up. Every short test passed; a sustained send died about fifty milliseconds in. A
// token bucket has no such cliff, and the property that matters is that a sender pacing *at* the
// allowed rate is never refused for long.
func TestSustainedSendAtTheAllowedRateIsNotCutOff(t *testing.T) {
	s := sessionWithSource(t)
	const kbps = 8000   // 1 MB/s
	const packet = 1200 // bytes
	g := s.NewGrant("a1", 8<<20, kbps, DefaultGrantLimits)

	// Pace at the allowed rate for a short run and count how much got through. A correct
	// limiter passes essentially all of it; the old one stopped almost immediately.
	perPacket := time.Duration(float64(packet) / (float64(kbps) * 125) * float64(time.Second))
	deadline := time.Now().Add(300 * time.Millisecond)
	sent, refusals := 0, 0
	for time.Now().Before(deadline) {
		if ok, why := g.TryAllow(packet); ok {
			sent += packet
		} else if why == RefusalRate {
			refusals++
		} else {
			t.Fatalf("unexpected terminal refusal %q after %d bytes", why, sent)
		}
		time.Sleep(perPacket)
	}

	// 300 ms at 1 MB/s is ~300 KB. Allow generous slack for scheduler granularity, but a run
	// that delivered only a few packets means the limiter cut it off.
	if sent < 100_000 {
		t.Fatalf("a sender pacing at the allowed rate got only %d bytes through in 300ms "+
			"(%d rate refusals) — the limiter is cutting off sustained sends", sent, refusals)
	}
}

// The other half: a rate refusal must be distinguishable from a spent budget, because one is
// transient and one is terminal, and a caller that cannot tell them apart either gives up early
// or spins forever.
func TestRefusalReasonsAreDistinguishable(t *testing.T) {
	s := sessionWithSource(t)

	// Budget: tiny ceiling, plenty of rate.
	g := s.NewGrant("a1", 1000, 100_000, DefaultGrantLimits)
	for i := 0; i < 20; i++ {
		g.TryAllow(100)
	}
	if ok, why := g.TryAllow(100); ok || why != RefusalBudget {
		t.Errorf("spent budget reported as ok=%v why=%q, want %q", ok, why, RefusalBudget)
	}

	// Rate: huge ceiling, minimal rate, so only the bucket can refuse.
	g2 := s.NewGrant("a2", 1<<20, 8, DefaultGrantLimits)
	sawRate := false
	for i := 0; i < 100; i++ {
		if ok, why := g2.TryAllow(1000); !ok && why == RefusalRate {
			sawRate = true
			break
		}
	}
	if !sawRate {
		t.Error("a sender far above the rate ceiling never got a rate refusal")
	}
}