routewatch/internal/routewatch/prefixhandler_internal_test.go

package routewatch

import (
	"context"
	"errors"
	"slices"
	"testing"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/config"
	"git.eeqj.de/sneak/routewatch/internal/database"
	"git.eeqj.de/sneak/routewatch/internal/logger"
	"git.eeqj.de/sneak/routewatch/internal/ristypes"
)

const (
	testPeerIP   = "2001:db8:ffff::1"
	testV4Prefix = "192.0.2.0/24"
)

// TestPrefixHandlerStoresPrefixesTheIPLookupFinds runs announcements and
// withdrawals through the prefix handler into a real database and looks the
// addresses up. The prefixes are written the way the feed sends them: IPv6
// announcements compressed, IPv6 withdrawals uncompressed. The withdrawal must
// remove the route the announcement stored, and the IP lookup must find the
// stored prefix.
func TestPrefixHandlerStoresPrefixesTheIPLookupFinds(t *testing.T) {
	t.Parallel()

	db, err := database.New(&config.Config{StateDir: t.TempDir()}, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	// Built without NewPrefixHandler's flush timer, so the test flushes itself.
	h := &PrefixHandler{db: db, logger: logger.New()}
	flush := func() {
		h.mu.Lock()
		defer h.mu.Unlock()

		h.flushBatchLocked()
	}

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	h.HandleMessage(&ristypes.RISMessage{
		Peer:            testPeerIP,
		Path:            ristypes.ASPath{testASNA, testASNB},
		ParsedTimestamp: ts,
		Announcements: []ristypes.RISAnnouncement{{
			NextHop:  testPeerIP,
			Prefixes: []string{"2001:db8:1::/48", testV4Prefix},
		}},
	})
	flush()

	for ip, want := range map[string]string{
		"2001:db8:1::1": "2001:db8:1::/48",
		"192.0.2.1":     testV4Prefix,
	} {
		info, err := db.GetIPInfo(ip)
		if err != nil {
			t.Fatalf("GetIPInfo(%s) after announcement: %v", ip, err)
		}

		if info.Netblock != want || info.ASN != testASNB {
			t.Errorf("GetIPInfo(%s) = %s AS%d, want %s AS%d",
				ip, info.Netblock, info.ASN, want, testASNB)
		}
	}

	h.HandleMessage(&ristypes.RISMessage{
		Peer:            testPeerIP,
		ParsedTimestamp: ts.Add(time.Minute),
		Withdrawals:     []string{"2001:db8:1:0:0:0:0:0/48", testV4Prefix},
	})
	flush()

	for _, ip := range []string{"2001:db8:1::1", "192.0.2.1"} {
		info, err := db.GetIPInfo(ip)
		if !errors.Is(err, database.ErrNoRoute) {
			t.Errorf("GetIPInfo(%s) after withdrawal = %+v, %v; want ErrNoRoute",
				ip, info, err)
		}
	}
}

// TestPrefixHandlerAppliesEachPeersLastUpdateInABatch runs updates for one
// prefix through the prefix handler into a real database. All updates carry the
// same timestamp, as updates within one second of the feed do. Within a batch,
// the last update each peer sent for the prefix decides whether that peer has a
// live route for it, and updates from different peers are all applied. The
// in-memory route count and prefix distribution must match the route table.
func TestPrefixHandlerAppliesEachPeersLastUpdateInABatch(t *testing.T) {
	t.Parallel()

	const (
		prefix           = "198.51.100.0/24"
		prefixMaskLength = 24
		otherPeerIP      = "2001:db8:ffff::2"
	)

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	announce := func(peer string) *ristypes.RISMessage {
		return &ristypes.RISMessage{
			Peer:            peer,
			Path:            ristypes.ASPath{testASNA, testASNB},
			ParsedTimestamp: ts,
			Announcements: []ristypes.RISAnnouncement{
				{NextHop: peer, Prefixes: []string{prefix}},
			},
		}
	}
	withdraw := func(peer string) *ristypes.RISMessage {
		return &ristypes.RISMessage{
			Peer: peer, ParsedTimestamp: ts, Withdrawals: []string{prefix},
		}
	}
	announceAndWithdraw := func(peer string) *ristypes.RISMessage {
		msg := announce(peer)
		msg.Withdrawals = []string{prefix}

		return msg
	}

	tests := []struct {
		name string
		// Each inner list is handled and then flushed as one batch.
		batches [][]*ristypes.RISMessage
		// The peers left with a live route for the prefix, by peer IP.
		wantPeers []string
	}{
		{
			name:      "two peers announce the prefix",
			batches:   [][]*ristypes.RISMessage{{announce(testPeerIP), announce(otherPeerIP)}},
			wantPeers: []string{testPeerIP, otherPeerIP},
		},
		{
			name: "two peers withdraw the prefix",
			batches: [][]*ristypes.RISMessage{
				{announce(testPeerIP)},
				{announce(otherPeerIP)},
				{withdraw(testPeerIP), withdraw(otherPeerIP)},
			},
		},
		{
			name:    "a peer announces and then withdraws the prefix",
			batches: [][]*ristypes.RISMessage{{announce(testPeerIP), withdraw(testPeerIP)}},
		},
		{
			name: "a peer withdraws and then re-announces the prefix",
			batches: [][]*ristypes.RISMessage{
				{announce(testPeerIP)},
				{withdraw(testPeerIP), announce(testPeerIP)},
			},
			wantPeers: []string{testPeerIP},
		},
		{
			// RFC 4271 section 4.3: the announcement counts, not the withdrawal.
			name: "one message announces and withdraws the prefix",
			batches: [][]*ristypes.RISMessage{
				{announce(testPeerIP)},
				{announceAndWithdraw(testPeerIP)},
			},
			wantPeers: []string{testPeerIP},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			db := handleBatches(t, tt.batches)
			checkPrefixRoutes(t, db, prefix, prefixMaskLength, tt.wantPeers)
		})
	}
}

// handleBatches runs each batch of messages through a prefix handler on a new
// database, flushing after each batch, and returns the database. The database
// is closed when the test ends.
func handleBatches(
	t *testing.T, batches [][]*ristypes.RISMessage,
) *database.Database {
	t.Helper()

	db, err := database.New(&config.Config{StateDir: t.TempDir()}, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}

	t.Cleanup(func() { _ = db.Close() })

	// Built without NewPrefixHandler's flush timer, so the test flushes itself.
	h := &PrefixHandler{db: db, logger: logger.New()}

	for _, batch := range batches {
		for _, msg := range batch {
			h.HandleMessage(msg)
		}

		h.mu.Lock()
		h.flushBatchLocked()
		h.mu.Unlock()
	}

	return db
}

// checkPrefixRoutes checks that exactly wantPeers have a live route for the
// IPv4 prefix, and that the in-memory route count and prefix distribution
// match the route table.
func checkPrefixRoutes(
	t *testing.T, db *database.Database, prefix string, maskLength int,
	wantPeers []string,
) {
	t.Helper()

	ctx := context.Background()

	routes, _, err := db.GetPrefixDetailsContext(ctx, prefix)
	if err != nil && !errors.Is(err, database.ErrNoRoute) {
		t.Fatalf("GetPrefixDetailsContext: %v", err)
	}

	gotPeers := make([]string, 0, len(routes))
	for _, route := range routes {
		gotPeers = append(gotPeers, route.PeerIP)
	}

	if !slices.Equal(gotPeers, wantPeers) {
		t.Errorf("peers with a live route for %s = %v, want %v",
			prefix, gotPeers, wantPeers)
	}

	stats, err := db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}

	tableRoutesV4, _, err := db.GetLiveRouteCountsContext(ctx)
	if err != nil {
		t.Fatalf("GetLiveRouteCountsContext: %v", err)
	}

	if stats.IPv4Routes != len(wantPeers) || tableRoutesV4 != len(wantPeers) {
		t.Errorf("IPv4 live routes: %d in memory, %d in the table, want %d",
			stats.IPv4Routes, tableRoutesV4, len(wantPeers))
	}

	var wantDistribution []database.PrefixDistribution
	if len(wantPeers) > 0 {
		wantDistribution = []database.PrefixDistribution{
			{MaskLength: maskLength, Count: 1},
		}
	}

	queryDistributionV4, _, err := db.GetPrefixDistributionContext(ctx)
	if err != nil {
		t.Fatalf("GetPrefixDistributionContext: %v", err)
	}

	if !slices.Equal(stats.IPv4PrefixDistribution, wantDistribution) ||
		!slices.Equal(queryDistributionV4, wantDistribution) {
		t.Errorf("IPv4 prefix distribution: %v in memory, %v from the table, want %v",
			stats.IPv4PrefixDistribution, queryDistributionV4, wantDistribution)
	}
}