routewatch/internal/routewatch/prefixhandler_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"

// 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) {
	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", "192.0.2.0/24"},
		}},
	})
	flush()

	for ip, want := range map[string]string{
		"2001:db8:1::1": "2001:db8:1::/48",
		"192.0.2.1":     "192.0.2.0/24",
	} {
		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", "192.0.2.0/24"},
	})
	flush()

	for _, ip := range []string{"2001:db8:1::1", "192.0.2.1"} {
		if info, err := db.GetIPInfo(ip); !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) {
	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) {
			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()}
			for _, batch := range tt.batches {
				for _, msg := range batch {
					h.HandleMessage(msg)
				}
				h.mu.Lock()
				h.flushBatchLocked()
				h.mu.Unlock()
			}

			ctx := context.Background()
			routes, err := db.GetPrefixDetailsContext(ctx, prefix)
			if err != nil && !errors.Is(err, database.ErrNoRoute) {
				t.Fatalf("GetPrefixDetailsContext: %v", err)
			}
			var gotPeers []string
			for _, route := range routes {
				gotPeers = append(gotPeers, route.PeerIP)
			}
			if !slices.Equal(gotPeers, tt.wantPeers) {
				t.Errorf("peers with a live route for %s = %v, want %v", prefix, gotPeers, tt.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(tt.wantPeers) || tableRoutesV4 != len(tt.wantPeers) {
				t.Errorf("IPv4 live routes: %d in memory, %d in the table, want %d",
					stats.IPv4Routes, tableRoutesV4, len(tt.wantPeers))
			}

			var wantDistribution []database.PrefixDistribution
			if len(tt.wantPeers) > 0 {
				wantDistribution = []database.PrefixDistribution{{MaskLength: prefixMaskLength, 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)
			}
		})
	}
}