routewatch/internal/routewatch/prefixhandler_internal_test.go

package routewatch

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net/netip"
	"slices"
	"strings"
	"sync"
	"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"

// An IPv4 and an IPv6 prefix the tests announce.
const (
	testV4Prefix = "192.0.2.0/24"
	testV6Prefix = "2001:db8:1::/48"
)

// secondUpdateAfter is how long after the first update the tests of first_seen
// and last_seen send the second, in the same batch.
const secondUpdateAfter = 5 * time.Second

// 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 := newTestDatabase(t)

	// 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{Sequence: []int{testASNA, testASNB}},
		ParsedTimestamp: ts,
		Announcements: []ristypes.RISAnnouncement{{
			NextHop:  testPeerIP,
			Prefixes: []string{testV6Prefix, testV4Prefix},
		}},
	})
	h.flushBatch()

	for ip, want := range map[string]string{
		"2001:db8:1::1": testV6Prefix,
		"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},
	})
	h.flushBatch()

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

// TestPrefixHandlerStoresANewPrefixFirstSeenAtItsEarliestUpdate announces a
// prefix the database does not hold yet twice in one batch, the second time 5
// seconds after the first, through the prefix handler into a real database. The
// prefix is stored first seen at the time of the first announcement, whether
// the two come from two peers or from one, and the IP lookup answers with that
// time.
func TestPrefixHandlerStoresANewPrefixFirstSeenAtItsEarliestUpdate(
	t *testing.T,
) {
	t.Parallel()

	const otherPeerIP = "2001:db8:ffff::2"

	tests := []struct {
		name                  string
		firstPeer, secondPeer string
	}{
		{name: "two peers", firstPeer: testPeerIP, secondPeer: otherPeerIP},
		{name: "one peer", firstPeer: testPeerIP, secondPeer: testPeerIP},
	}

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

			first := peerAnnouncement(tt.firstPeer, testV4Prefix)
			second := peerAnnouncement(tt.secondPeer, testV4Prefix)
			second.ParsedTimestamp = first.ParsedTimestamp.Add(secondUpdateAfter)
			db := handleBatches(t, [][]*ristypes.RISMessage{{first, second}})

			info, err := db.GetIPInfo("192.0.2.1")
			if err != nil {
				t.Fatalf("GetIPInfo(192.0.2.1): %v", err)
			}

			if !info.FirstSeen.Equal(first.ParsedTimestamp) {
				t.Errorf("first_seen of %s = %s, want %s", testV4Prefix,
					info.FirstSeen, first.ParsedTimestamp)
			}
		})
	}
}

// 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{Sequence: []int{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)
			checkLivePeers(t, db, prefix, prefixMaskLength, tt.wantPeers)
		})
	}
}

// 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) {
	t.Parallel()

	families := []struct {
		ipVersion           int
		prefix, otherPrefix string
		// The mask length of both prefixes.
		maskLength int
	}{
		{
			ipVersion: ipv4Version, maskLength: 24,
			prefix: "198.51.100.0/24", otherPrefix: "203.0.113.0/24",
		},
		{
			ipVersion: ipv6Version, maskLength: 48,
			prefix: testV6Prefix, otherPrefix: "2001:db8:2::/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{Sequence: []int{testASNA, testASNB}}, otherPrefix)
		announceOtherAndWithdraw.Withdrawals = []string{prefix}

		tests := []struct {
			name string
			// Each message is handled and then flushed as its own batch.
			batches [][]*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",
				batches: [][]*ristypes.RISMessage{
					{announce(ristypes.ASPath{Sequence: []int{testASNA}}, prefix)},
					{announce(ristypes.ASPath{Sequence: []int{testASNA, testASNB}}, prefix)},
				},
				wantOrigins: map[string]int{prefix: testASNB},
			},
			{
				name: "one message announces a prefix with another origin AS " +
					"and withdraws the prefix",
				batches: [][]*ristypes.RISMessage{
					{announce(ristypes.ASPath{Sequence: []int{testASNA}}, prefix)},
					{announceOtherAndWithdraw},
				},
				wantOrigins: map[string]int{otherPrefix: testASNB},
			},
		}

		for _, tt := range tests {
			name := fmt.Sprintf("IPv%d %s", family.ipVersion, tt.name)
			t.Run(name, func(t *testing.T) {
				t.Parallel()

				db := handleBatches(t, tt.batches)
				checkLiveOrigins(t, db, tt.wantOrigins, prefix, otherPrefix)

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

// TestPrefixHandlerStoresNoOriginForAPathEndingInASet announces a prefix with
// AS paths that end in an AS set through the prefix handler into a real
// database. Any member of a set of several may be the origin, so the route is
// stored with origin 0 (none); a set of one member counts as that AS. The
// stored AS path lists the set's members after the rest of the path.
func TestPrefixHandlerStoresNoOriginForAPathEndingInASet(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name       string
		path       []any
		wantOrigin int
		wantPath   []int
	}{
		{
			name:     "a set of several members",
			path:     []any{testASNA, testASNB, []int{testASNC, testASND}},
			wantPath: []int{testASNA, testASNB, testASNC, testASND},
		},
		{
			name:       "a set of one member",
			path:       []any{testASNA, testASNB, []int{testASNC}},
			wantOrigin: testASNC,
			wantPath:   []int{testASNA, testASNB, testASNC},
		},
	}

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

			msg := peerAnnouncement(testPeerIP, testV4Prefix)
			msg.Path = decodePath(t, tt.path...)
			db := handleBatches(t, [][]*ristypes.RISMessage{{msg}})

			routes, _, err := db.GetPrefixDetailsContext(
				context.Background(), testV4Prefix)
			if err != nil {
				t.Fatalf("GetPrefixDetailsContext(%s): %v", testV4Prefix, err)
			}

			if len(routes) != 1 || routes[0].OriginASN != tt.wantOrigin ||
				!slices.Equal(routes[0].ASPath, tt.wantPath) {
				t.Errorf("live routes for %s = %+v, "+
					"want one with origin AS%d and AS path %v",
					testV4Prefix, routes, tt.wantOrigin, tt.wantPath)
			}
		})
	}
}

// TestPrefixHandlerDoesNotStoreAnAnnouncementWithAnEmptyPath announces a
// prefix with an empty AS path, which gives no origin and does not end in an
// AS set, so no live route is stored.
func TestPrefixHandlerDoesNotStoreAnAnnouncementWithAnEmptyPath(t *testing.T) {
	t.Parallel()

	msg := peerAnnouncement(testPeerIP, testV4Prefix)
	msg.Path = ristypes.ASPath{}
	db := handleBatches(t, [][]*ristypes.RISMessage{{msg}})

	assertPeers(t, db, map[string][]string{testV4Prefix: nil})
}

// TestPrefixHandlerRemovesTheRouteOnAnAnnouncementWithoutAnOrigin runs, through
// the prefix handler into a real database, a peer's announcement of a prefix
// with origin AS testASNB and then the peer's announcement of the prefix with
// an AS path that gives no origin: one whose last AS is 0, or an empty one. As
// any later announcement from the peer does, the second replaces the peer's
// route, here with none, whether it comes in a later batch or follows a
// withdrawal in one batch. The in-memory route count and prefix distribution
// must match the route table.
func TestPrefixHandlerRemovesTheRouteOnAnAnnouncementWithoutAnOrigin(
	t *testing.T,
) {
	t.Parallel()

	const prefixMaskLength = 24

	noOrigin := func(path ristypes.ASPath) *ristypes.RISMessage {
		msg := peerAnnouncement(testPeerIP, testV4Prefix)
		msg.Path = path

		return msg
	}
	endsInZero := ristypes.ASPath{Sequence: []int{testASNA, 0}}
	withdrawal := &ristypes.RISMessage{
		Type:            ristypes.MessageTypeUpdate,
		Peer:            testPeerIP,
		ParsedTimestamp: time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC),
		Withdrawals:     []string{testV4Prefix},
	}

	tests := []struct {
		name string
		// Handled and flushed as one batch after the announcement with origin
		// AS testASNB.
		batch []*ristypes.RISMessage
	}{
		{
			name:  "a path whose last AS is 0",
			batch: []*ristypes.RISMessage{noOrigin(endsInZero)},
		},
		{
			name:  "an empty path",
			batch: []*ristypes.RISMessage{noOrigin(ristypes.ASPath{})},
		},
		{
			name:  "a withdrawal and then a path whose last AS is 0",
			batch: []*ristypes.RISMessage{withdrawal, noOrigin(endsInZero)},
		},
	}

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

			db := handleBatches(t, [][]*ristypes.RISMessage{
				{peerAnnouncement(testPeerIP, testV4Prefix)},
				tt.batch,
			})
			checkLivePeers(t, db, testV4Prefix, prefixMaskLength, nil)
		})
	}
}

// TestPrefixHandlerDoesNotStoreAPrefixThatDoesNotParse runs an announcement of
// "not-a-prefix" and a withdrawal of "also/bad" through the prefix handler into
// a real database. Neither is stored in the prefix tables: the prefix count in
// the stats stays 0, both as kept in memory and as counted from the tables when
// the database is opened again.
func TestPrefixHandlerDoesNotStoreAPrefixThatDoesNotParse(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

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

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

	withdrawal := &ristypes.RISMessage{
		Type:            ristypes.MessageTypeUpdate,
		Peer:            testPeerIP,
		ParsedTimestamp: time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC),
		Withdrawals:     []string{"also/bad"},
	}
	h.HandleMessage(peerAnnouncement(testPeerIP, "not-a-prefix"))
	h.HandleMessage(withdrawal)
	h.flushBatch()

	checkNoPrefixes := func(when string) {
		t.Helper()

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

		if stats.Prefixes != 0 {
			t.Errorf("prefix count %s = %d, want 0", when, stats.Prefixes)
		}
	}

	checkNoPrefixes("after the flush")

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

	db, err = database.New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to open the database again: %v", err)
	}

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

	checkNoPrefixes("after opening the database again")
}

// TestStateMessagesGoToThePrefixHandlerNotThePeerHandler checks that STATE
// messages go to the prefix handler, which acts on a down, and not to the peer
// handler, which would record a peer whose session never comes up as active.
func TestStateMessagesGoToThePrefixHandlerNotThePeerHandler(t *testing.T) {
	t.Parallel()

	if !(&PrefixHandler{}).WantsMessage(ristypes.MessageTypeState) {
		t.Error("the prefix handler does not want STATE messages")
	}

	if (&PeerHandler{}).WantsMessage(ristypes.MessageTypeState) {
		t.Error("the peer handler wants STATE messages")
	}

	if !(&PeerHandler{}).WantsMessage(ristypes.MessageTypeUpdate) {
		t.Error("the peer handler does not want UPDATE messages")
	}
}

// TestPrefixHandlerDeletesAPeersRoutesWhenItsSessionGoesDown runs
// announcements from two peers through the prefix handler into a real
// database, then a batch in which one of the peers announces a prefix and its
// session goes down. The down deletes every live route of that peer, IPv4 and
// IPv6, the one announced just before it included, and keeps the other peer's
// routes for the same prefixes. An announcement the peer sends after the down,
// in the same batch, is kept. The in-memory route counts and prefix
// distribution must match the route tables.
func TestPrefixHandlerDeletesAPeersRoutesWhenItsSessionGoesDown(t *testing.T) {
	t.Parallel()

	const (
		otherPeerIP = "2001:db8:ffff::2"
		// Both peers announce the shared prefixes, only testPeerIP the others.
		sharedV4   = "198.51.100.0/24"
		sharedV6   = "2001:db8:1::/48"
		onlyV4     = "203.0.113.0/24"
		onlyV6     = "2001:db8:2::/48"
		beforeDown = "192.0.2.0/24"
		afterDown  = "2001:db8:3::/48"
		// The mask lengths of the IPv4 and of the IPv6 prefixes.
		maskLengthV4 = 24
		maskLengthV6 = 48
	)

	tests := []struct {
		name string
		// Handled and flushed as one batch after both peers' announcements.
		batch []*ristypes.RISMessage
		// Whether testPeerIP is left with a live route for afterDown.
		wantAfterDown bool
	}{
		{
			name: "the down ends the batch",
			batch: []*ristypes.RISMessage{
				peerAnnouncement(testPeerIP, beforeDown),
				sessionDown(),
			},
		},
		{
			name: "the peer announces a prefix after the down",
			batch: []*ristypes.RISMessage{
				peerAnnouncement(testPeerIP, beforeDown),
				sessionDown(),
				peerAnnouncement(testPeerIP, afterDown),
			},
			wantAfterDown: true,
		},
	}

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

			db := handleBatches(t, [][]*ristypes.RISMessage{
				{
					peerAnnouncement(testPeerIP,
						sharedV4, sharedV6, onlyV4, onlyV6),
					peerAnnouncement(otherPeerIP, sharedV4, sharedV6),
				},
				tt.batch,
			})

			// The peers left with a live route for each prefix.
			wantPeers := map[string][]string{
				sharedV4:   {otherPeerIP},
				sharedV6:   {otherPeerIP},
				onlyV4:     nil,
				onlyV6:     nil,
				beforeDown: nil,
				afterDown:  nil,
			}
			wantRoutesV6 := 1

			if tt.wantAfterDown {
				wantPeers[afterDown] = []string{testPeerIP}
				wantRoutesV6++
			}

			assertPeers(t, db, wantPeers)
			assertCounts(t, db, ipv4Version, 1, []database.PrefixDistribution{
				{MaskLength: maskLengthV4, Count: 1},
			})
			assertCounts(t, db, ipv6Version, wantRoutesV6,
				[]database.PrefixDistribution{
					{MaskLength: maskLengthV6, Count: wantRoutesV6},
				})
		})
	}
}

// TestPrefixHandlerKeepsWhatAPeerAnnouncesAfterTheFlushOfItsDown runs, through
// the prefix handler into a real database, one flush after another: a peer's
// announcement, its session going down, a new announcement from the peer, and
// another peer's announcement. The down deletes the peer's routes on the flush
// that handles it and on no later flush, so the route the peer announced after
// it is kept.
func TestPrefixHandlerKeepsWhatAPeerAnnouncesAfterTheFlushOfItsDown(t *testing.T) {
	t.Parallel()

	const (
		otherPeerIP = "2001:db8:ffff::2"
		beforeDown  = "192.0.2.0/24"
		afterDown   = "198.51.100.0/24"
		otherPeers  = "203.0.113.0/24"
		maskLength  = 24
	)

	// Each message is handled and then flushed as its own batch.
	db := handleBatches(t, [][]*ristypes.RISMessage{
		{peerAnnouncement(testPeerIP, beforeDown)},
		{sessionDown()},
		{peerAnnouncement(testPeerIP, afterDown)},
		{peerAnnouncement(otherPeerIP, otherPeers)},
	})

	assertPeers(t, db, map[string][]string{
		beforeDown: nil,
		afterDown:  {testPeerIP},
		otherPeers: {otherPeerIP},
	})
	assertCounts(t, db, ipv4Version, 2, []database.PrefixDistribution{
		{MaskLength: maskLength, Count: 2},
	})
}

// firstWriteHeldStore is a real database whose first live route batch write
// waits until the test closes release. writeStarted is closed when it starts.
type firstWriteHeldStore struct {
	*database.Database

	writeStarted chan struct{}
	release      chan struct{}
	once         sync.Once
}

// UpsertLiveRouteBatch holds the first call until release is closed, then
// writes routes to the database.
func (s *firstWriteHeldStore) UpsertLiveRouteBatch(
	routes []*database.LiveRoute,
) error {
	s.once.Do(func() {
		close(s.writeStarted)
		<-s.release
	})

	return s.Database.UpsertLiveRouteBatch(routes)
}

// TestPrefixHandlerHandlesADownWhileAFlushWrites holds the write of a flush
// that stores a peer's route and meanwhile hands the prefix handler the peer's
// session going down, then a new announcement from the peer. The next flush
// must delete the route the held write stored and keep the one announced after
// the down.
func TestPrefixHandlerHandlesADownWhileAFlushWrites(t *testing.T) {
	t.Parallel()

	const (
		beforeDown = "192.0.2.0/24"
		afterDown  = "198.51.100.0/24"
		maskLength = 24
	)

	db := newTestDatabase(t)
	store := &firstWriteHeldStore{
		Database:     db,
		writeStarted: make(chan struct{}),
		release:      make(chan struct{}),
	}

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

	flushed := make(chan struct{})

	go func() {
		h.flushBatch()
		close(flushed)
	}()

	select {
	case <-store.writeStarted:
	case <-time.After(handleMessageTimeout):
		t.Fatal("the flush did not start its write")
	}

	if !returnsInTime(func() {
		h.HandleMessage(sessionDown())
		h.HandleMessage(peerAnnouncement(testPeerIP, afterDown))
	}) {
		t.Fatal("HandleMessage did not return while a database write was held")
	}

	close(store.release)

	select {
	case <-flushed:
	case <-time.After(handleMessageTimeout):
		t.Fatal("the flush did not finish after its write was let go")
	}

	h.flushBatch()

	assertPeers(t, db, map[string][]string{
		beforeDown: nil,
		afterDown:  {testPeerIP},
	})
	assertCounts(t, db, ipv4Version, 1, []database.PrefixDistribution{
		{MaskLength: maskLength, Count: 1},
	})
}

// errTestPeerDelete is the error a test store's failing delete of a peer's
// routes returns.
var errTestPeerDelete = errors.New("test: the delete of a peer's routes failed")

// firstPeerDeleteFailsStore is a real database whose first delete of a peer's
// live routes fails, deleting nothing. deletes counts the deletes asked for.
type firstPeerDeleteFailsStore struct {
	*database.Database

	// whileFailing, if not nil, runs during the delete that fails.
	whileFailing func()
	deletes      int
}

// DeleteLiveRoutesOfPeer fails the first call and deletes the peer's routes on
// every later one.
func (s *firstPeerDeleteFailsStore) DeleteLiveRoutesOfPeer(
	peerIP string,
) (int, error) {
	s.deletes++
	if s.deletes > 1 {
		return s.Database.DeleteLiveRoutesOfPeer(peerIP)
	}

	if s.whileFailing != nil {
		s.whileFailing()
	}

	return 0, errTestPeerDelete
}

// TestPrefixHandlerRetriesAFailedDeleteOfAPeersRoutes runs, through the prefix
// handler into a database whose first delete of a peer's routes fails, one flush
// after another: a peer's announcement, then its session going down and a new
// announcement from the peer, then nothing. The failed delete must be retried
// on the next flush, and the announcement the peer sent after the down written
// only after that retry, so it is kept. In the other cases the peer's session
// goes down again, and the peer announces another prefix, while the delete
// fails or after the flush whose delete failed: the announcement sent between
// the two downs must not be written, and the next flush must delete the peer's
// routes once, not once for the retry and once for the down.
func TestPrefixHandlerRetriesAFailedDeleteOfAPeersRoutes(t *testing.T) {
	t.Parallel()

	const (
		beforeDown      = "192.0.2.0/24"
		afterDown       = "198.51.100.0/24"
		afterSecondDown = "203.0.113.0/24"
		maskLength      = 24
		// The delete that fails and its retry, which is also the delete for
		// a second down.
		wantDeletes = 2
	)

	tests := []struct {
		name string
		// Handled while the delete fails, if not nil.
		whileFailing []*ristypes.RISMessage
		// Handled after the flush whose delete fails, if not nil.
		afterFailing []*ristypes.RISMessage
		// The peers left with a live route for each prefix.
		wantPeers map[string][]string
	}{
		{
			name: "the delete fails once",
			wantPeers: map[string][]string{
				beforeDown: nil, afterDown: {testPeerIP}, afterSecondDown: nil,
			},
		},
		{
			name: "the session goes down again while the delete fails",
			whileFailing: []*ristypes.RISMessage{
				sessionDown(), peerAnnouncement(testPeerIP, afterSecondDown),
			},
			wantPeers: map[string][]string{
				beforeDown: nil, afterDown: nil, afterSecondDown: {testPeerIP},
			},
		},
		{
			name: "the session goes down again after the delete failed",
			afterFailing: []*ristypes.RISMessage{
				sessionDown(), peerAnnouncement(testPeerIP, afterSecondDown),
			},
			wantPeers: map[string][]string{
				beforeDown: nil, afterDown: nil, afterSecondDown: {testPeerIP},
			},
		},
	}

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

			db := newTestDatabase(t)
			store := &firstPeerDeleteFailsStore{Database: db}

			// Built without NewPrefixHandler's flush timer, so the test
			// flushes itself.
			h := &PrefixHandler{db: store, logger: logger.New()}
			store.whileFailing = func() {
				for _, msg := range tt.whileFailing {
					h.HandleMessage(msg)
				}
			}

			h.HandleMessage(peerAnnouncement(testPeerIP, beforeDown))
			h.flushBatch()
			h.HandleMessage(sessionDown())
			h.HandleMessage(peerAnnouncement(testPeerIP, afterDown))
			h.flushBatch() // the delete fails

			for _, msg := range tt.afterFailing {
				h.HandleMessage(msg)
			}

			h.flushBatch()

			if store.deletes != wantDeletes {
				t.Errorf("the store was asked for %d deletes of the peer's "+
					"routes, want %d", store.deletes, wantDeletes)
			}

			assertPeers(t, db, tt.wantPeers)
			assertCounts(t, db, ipv4Version, 1, []database.PrefixDistribution{
				{MaskLength: maskLength, Count: 1},
			})
		})
	}
}

// peerDeleteAlwaysFailsStore is a real database whose every delete of a peer's
// live routes fails, deleting nothing.
type peerDeleteAlwaysFailsStore struct {
	*database.Database
}

// DeleteLiveRoutesOfPeer fails.
func (peerDeleteAlwaysFailsStore) DeleteLiveRoutesOfPeer(string) (int, error) {
	return 0, errTestPeerDelete
}

// TestPrefixHandlerKeepsAtMostABatchOfUpdatesWhileADeleteKeepsFailing runs,
// through the prefix handler into a database whose every delete of a peer's
// routes fails, a peer's session going down and then three flushes. Before each
// flush the peer announces prefixBatchSize new prefixes and another peer
// announces one. After each flush the handler must hold only the peer's last
// prefixBatchSize updates, not every update since the down, and the other
// peer's route must be written. After the last flush, one more update must not
// wake the flush goroutine: held updates do not count toward a full batch.
func TestPrefixHandlerKeepsAtMostABatchOfUpdatesWhileADeleteKeepsFailing(
	t *testing.T,
) {
	t.Parallel()

	const otherPeerIP = "2001:db8:ffff::2"

	// One prefix the other peer announces before each flush.
	otherPeersPrefixes := []string{testV4Prefix, testV6Prefix, "2001:db8:2::/48"}

	db := newTestDatabase(t)

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

	// The peer announces a /32 for each address from 10.0.0.0 on, each
	// address once.
	addr := netip.MustParseAddr("10.0.0.0")

	for flush, otherPeersPrefix := range otherPeersPrefixes {
		prefixes := make([]string, prefixBatchSize)
		for i := range prefixes {
			prefixes[i] = netip.PrefixFrom(addr, addr.BitLen()).String()
			addr = addr.Next()
		}

		h.HandleMessage(peerAnnouncement(testPeerIP, prefixes...))
		h.HandleMessage(peerAnnouncement(otherPeerIP, otherPeersPrefix))
		h.flushBatch()

		kept := make([]string, 0, len(h.held))
		for _, update := range h.held {
			kept = append(kept, update.prefix)
		}

		if !slices.Equal(kept, prefixes) {
			t.Fatalf("after flush %d the handler holds %d updates, not just "+
				"the peer's last %d", flush+1, len(kept), len(prefixes))
		}

		assertPeers(t, db, map[string][]string{otherPeersPrefix: {otherPeerIP}})
	}

	// A new channel, so that a wake-up in it comes from the update below.
	h.batchFullCh = make(chan struct{}, 1)
	h.HandleMessage(peerAnnouncement(otherPeerIP, testV4Prefix))

	if len(h.batchFullCh) > 0 {
		t.Error("one update after a flush that held a batch of updates " +
			"woke the flush goroutine")
	}
}

// TestPrefixHandlerStopLogsAFailedDeleteWithoutTheDeadline stops a prefix
// handler, with no deadline, after a peer's session went down and the peer
// announced a prefix, into a database whose every delete of a peer's routes
// fails. Stop must log the peer down and the announcement as not written, and
// must not say that the stop deadline passed.
func TestPrefixHandlerStopLogsAFailedDeleteWithoutTheDeadline(t *testing.T) {
	t.Parallel()

	var logs bytes.Buffer

	h := NewPrefixHandler(peerDeleteAlwaysFailsStore{Database: newTestDatabase(t)},
		newErrorLogger(&logs))
	h.HandleMessage(sessionDown())
	h.HandleMessage(peerAnnouncement(testPeerIP, testV4Prefix))
	h.Stop(t.Context())

	logged := lastLogged(t, &logs)
	if strings.Contains(logged.Msg, "deadline") {
		t.Errorf("Stop logged %q although its deadline had not passed",
			logged.Msg)
	}

	if logged.Updates != 1 || logged.PeersDown != 1 {
		t.Errorf("Stop logged %d updates and %d peer downs not written, "+
			"want 1 and 1", logged.Updates, logged.PeersDown)
	}
}

// TestPrefixHandlerStopLogsNothingAfterADeleteThatSucceeds stops a prefix
// handler, with no deadline, after a peer's session went down, into a database
// whose delete of the peer's routes succeeds. Stop must log nothing: a peer
// whose routes are deleted no longer counts as a peer down left undone.
func TestPrefixHandlerStopLogsNothingAfterADeleteThatSucceeds(t *testing.T) {
	t.Parallel()

	var logs bytes.Buffer

	h := NewPrefixHandler(newTestDatabase(t), newErrorLogger(&logs))
	h.HandleMessage(sessionDown())
	h.Stop(t.Context())

	if logs.Len() != 0 {
		t.Errorf("Stop logged %q after the routes of the peer whose session "+
			"went down were deleted", logs.String())
	}
}

// newTestDatabase opens a database in a new temporary directory. It is closed
// when the test ends.
func newTestDatabase(t *testing.T) *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() })

	return db
}

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

	db := newTestDatabase(t)

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

	return db
}

// peerAnnouncement returns an UPDATE message in which peer announces prefixes
// with origin AS testASNB.
func peerAnnouncement(peer string, prefixes ...string) *ristypes.RISMessage {
	return &ristypes.RISMessage{
		Type:            ristypes.MessageTypeUpdate,
		Peer:            peer,
		Path:            ristypes.ASPath{Sequence: []int{testASNA, testASNB}},
		ParsedTimestamp: time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC),
		Announcements: []ristypes.RISAnnouncement{
			{NextHop: peer, Prefixes: prefixes},
		},
	}
}

// sessionDown returns the STATE message that says the BGP session of the peer
// testPeerIP went down.
func sessionDown() *ristypes.RISMessage {
	return &ristypes.RISMessage{
		Type:  ristypes.MessageTypeState,
		Peer:  testPeerIP,
		State: ristypes.StateDown,
	}
}

// checkLivePeers 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 checkLivePeers(
	t *testing.T, db *database.Database, prefix string, maskLength int,
	wantPeers []string,
) {
	t.Helper()

	assertPeers(t, db, map[string][]string{prefix: wantPeers})

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

	assertCounts(t, db, ipv4Version, len(wantPeers), wantDistribution)
}

// assertPeers checks, for each prefix in wantPeers, that the peers with a live
// route for it are the ones listed, in peer IP order; nil means none.
func assertPeers(
	t *testing.T, db *database.Database, wantPeers map[string][]string,
) {
	t.Helper()

	for prefix, want := range wantPeers {
		routes, _, err := db.GetPrefixDetailsContext(context.Background(), prefix)
		if err != nil && !errors.Is(err, database.ErrNoRoute) {
			t.Fatalf("GetPrefixDetailsContext(%s): %v", prefix, err)
		}

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

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

// checkLiveOrigins checks that each of prefixes has one live route, with the
// origin AS wantOrigins gives for it, or none when wantOrigins leaves it out.
func checkLiveOrigins(
	t *testing.T, db *database.Database, wantOrigins map[string]int,
	prefixes ...string,
) {
	t.Helper()

	for _, prefix := range prefixes {
		routes, _, err := db.GetPrefixDetailsContext(context.Background(), prefix)
		if err != nil && !errors.Is(err, database.ErrNoRoute) {
			t.Fatalf("GetPrefixDetailsContext(%s): %v", prefix, err)
		}

		gotOrigins := make([]int, 0, len(routes))
		for _, route := range routes {
			gotOrigins = append(gotOrigins, route.OriginASN)
		}

		var want []int
		if origin, ok := wantOrigins[prefix]; ok {
			want = []int{origin}
		}

		if !slices.Equal(gotOrigins, want) {
			t.Errorf("origin ASes of the live routes for %s = %v, want %v",
				prefix, gotOrigins, want)
		}
	}
}

// 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 := stats.IPv4PrefixDistribution
	queryDistribution := queryDistributionV4

	if ipVersion == ipv6Version {
		memoryRoutes, tableRoutes = stats.IPv6Routes, tableRoutesV6
		memoryDistribution = stats.IPv6PrefixDistribution
		queryDistribution = 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)
	}
}