routewatch/internal/routewatch/prefixhandler_test.go

package routewatch

import (
	"context"
	"errors"
	"fmt"
	"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()}

	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"},
		}},
	})
	h.flushBatch()

	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"},
	})
	h.flushBatch()

	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.flushBatch()
			}

			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)
			}

			var wantDistribution []database.PrefixDistribution
			if len(tt.wantPeers) > 0 {
				wantDistribution = []database.PrefixDistribution{{MaskLength: prefixMaskLength, Count: 1}}
			}
			assertCounts(t, db, ipv4Version, len(tt.wantPeers), wantDistribution)
		})
	}
}

// TestPrefixHandlerKeepsOneRoutePerPrefixAndPeer runs updates from one peer
// through the prefix handler into a real database, flushing each message as its
// own batch, once with IPv4 prefixes and once with IPv6 prefixes. A peer has at
// most one route for a prefix: a later announcement replaces it whatever its
// origin AS, and a withdrawal removes it whatever AS the message's AS path ends
// in. The in-memory route count and prefix distribution must match the route
// table.
func TestPrefixHandlerKeepsOneRoutePerPrefixAndPeer(t *testing.T) {
	families := []struct {
		ipVersion           int
		prefix, otherPrefix string
		// The mask length of both prefixes.
		maskLength int
	}{
		{ipVersion: ipv4Version, prefix: "198.51.100.0/24", otherPrefix: "203.0.113.0/24", maskLength: 24},
		{ipVersion: ipv6Version, prefix: "2001:db8:1::/48", otherPrefix: "2001:db8:2::/48", maskLength: 48},
	}

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	announce := func(path ristypes.ASPath, announced string) *ristypes.RISMessage {
		return &ristypes.RISMessage{
			Peer:            testPeerIP,
			Path:            path,
			ParsedTimestamp: ts,
			Announcements:   []ristypes.RISAnnouncement{{NextHop: testPeerIP, Prefixes: []string{announced}}},
		}
	}

	for _, family := range families {
		prefix, otherPrefix := family.prefix, family.otherPrefix

		// One message that announces otherPrefix with origin AS testASNB and
		// withdraws prefix.
		announceOtherAndWithdraw := announce(ristypes.ASPath{testASNA, testASNB}, otherPrefix)
		announceOtherAndWithdraw.Withdrawals = []string{prefix}

		tests := []struct {
			name     string
			messages []*ristypes.RISMessage
			// The origin AS of the peer's live route for each prefix. A prefix
			// left out must have no live route.
			wantOrigins map[string]int
		}{
			{
				name: "the peer announces the prefix with a new origin AS",
				messages: []*ristypes.RISMessage{
					announce(ristypes.ASPath{testASNA}, prefix),
					announce(ristypes.ASPath{testASNA, testASNB}, prefix),
				},
				wantOrigins: map[string]int{prefix: testASNB},
			},
			{
				name: "one message announces a prefix with another origin AS and withdraws the prefix",
				messages: []*ristypes.RISMessage{
					announce(ristypes.ASPath{testASNA}, prefix),
					announceOtherAndWithdraw,
				},
				wantOrigins: map[string]int{otherPrefix: testASNB},
			},
		}

		for _, tt := range tests {
			t.Run(fmt.Sprintf("IPv%d %s", family.ipVersion, 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 _, msg := range tt.messages {
					h.HandleMessage(msg)
					h.flushBatch()
				}

				for _, p := range []string{prefix, otherPrefix} {
					routes, _, err := db.GetPrefixDetailsContext(context.Background(), p)
					if err != nil && !errors.Is(err, database.ErrNoRoute) {
						t.Fatalf("GetPrefixDetailsContext(%s): %v", p, err)
					}
					var gotOrigins []int
					for _, route := range routes {
						gotOrigins = append(gotOrigins, route.OriginASN)
					}
					var wantOrigins []int
					if origin, ok := tt.wantOrigins[p]; ok {
						wantOrigins = []int{origin}
					}
					if !slices.Equal(gotOrigins, wantOrigins) {
						t.Errorf("origin ASes of the live routes for %s = %v, want %v", p, gotOrigins, wantOrigins)
					}
				}

				wantDistribution := []database.PrefixDistribution{
					{MaskLength: family.maskLength, Count: len(tt.wantOrigins)},
				}
				assertCounts(t, db, family.ipVersion, len(tt.wantOrigins), wantDistribution)
			})
		}
	}
}

// assertCounts checks that the live route count and prefix distribution of one
// IP version are wantRoutes and wantDistribution, both as held in memory and as
// read from the route table.
func assertCounts(t *testing.T, db *database.Database, ipVersion, wantRoutes int,
	wantDistribution []database.PrefixDistribution) {
	t.Helper()

	ctx := context.Background()
	stats, err := db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	tableRoutesV4, tableRoutesV6, err := db.GetLiveRouteCountsContext(ctx)
	if err != nil {
		t.Fatalf("GetLiveRouteCountsContext: %v", err)
	}
	queryDistributionV4, queryDistributionV6, err := db.GetPrefixDistributionContext(ctx)
	if err != nil {
		t.Fatalf("GetPrefixDistributionContext: %v", err)
	}

	memoryRoutes, tableRoutes := stats.IPv4Routes, tableRoutesV4
	memoryDistribution, queryDistribution := stats.IPv4PrefixDistribution, queryDistributionV4
	if ipVersion == ipv6Version {
		memoryRoutes, tableRoutes = stats.IPv6Routes, tableRoutesV6
		memoryDistribution, queryDistribution = stats.IPv6PrefixDistribution, queryDistributionV6
	}

	if memoryRoutes != wantRoutes || tableRoutes != wantRoutes {
		t.Errorf("IPv%d live routes: %d in memory, %d in the table, want %d",
			ipVersion, memoryRoutes, tableRoutes, wantRoutes)
	}
	if !slices.Equal(memoryDistribution, wantDistribution) ||
		!slices.Equal(queryDistribution, wantDistribution) {
		t.Errorf("IPv%d prefix distribution: %v in memory, %v from the table, want %v",
			ipVersion, memoryDistribution, queryDistribution, wantDistribution)
	}
}