Warp/server/internal/storage/federation_test.go

package storage_test

import (
	"database/sql"
	"os"
	"slices"
	"testing"

	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// migratedFederationDB opens a driver with the room and federation tables
// empty.
func migratedFederationDB(t *testing.T, driver string) storage.Database {
	t.Helper()
	db := migratedRoomDB(t, driver)
	if driver == storage.DriverPostgres {
		pool, err := sql.Open("pgx", os.Getenv("WARP_TEST_POSTGRES_DSN"))
		if err != nil {
			t.Fatalf("open postgres to reset the federation tables: %v", err)
		}
		defer func() { _ = pool.Close() }()
		if _, err := pool.ExecContext(t.Context(), `TRUNCATE
			federation_server_keys, federation_queue,
			federation_destinations, federation_inbound_txns`); err != nil {
			t.Fatalf("reset the federation tables: %v", err)
		}
		if _, err := pool.ExecContext(t.Context(), `UPDATE federation_sequences SET value = 0`); err != nil {
			t.Fatalf("rewind the federation counters: %v", err)
		}
	}
	return db
}

func TestOutliersAreReadableButNotInTheTimeline(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			rooms := migratedFederationDB(t, driver).Rooms()
			if err := rooms.CreateRoom(ctx, newTestRoom(testRoomID)); err != nil {
				t.Fatalf("CreateRoom: %v", err)
			}

			outliers := []storage.StoredEvent{
				newTestEvent(testRoomID, "$create", "m.room.create", stateKey("")),
				newTestEvent(testRoomID, "$member", "m.room.member", stateKey("@bob:remote.example")),
			}
			if err := rooms.WriteOutliers(ctx, testRoomID, outliers); err != nil {
				t.Fatalf("WriteOutliers: %v", err)
			}
			// Writing them again is not an error: an event ID names its
			// contents, so the second copy is the first one.
			if err := rooms.WriteOutliers(ctx, testRoomID, outliers); err != nil {
				t.Fatalf("WriteOutliers twice: %v", err)
			}

			got, err := rooms.EventsByID(ctx, testRoomID, []string{"$create", "$member"})
			if err != nil {
				t.Fatalf("EventsByID: %v", err)
			}
			if len(got) != 2 {
				t.Fatalf("EventsByID found %d outliers, want 2", len(got))
			}
			for _, e := range got {
				if !e.Outlier {
					t.Errorf("%s came back as a timeline event", e.EventID)
				}
			}

			// A join that cites the outliers as its state.
			join := newTestEvent(testRoomID, "$join", "m.room.member", stateKey("@alice:warp.example"))
			if _, err := rooms.Write(ctx, storage.RoomWrite{
				Event: join,
				StateBefore: []storage.StateEntry{
					{Type: "m.room.create", EventID: "$create"},
					{Type: "m.room.member", StateKey: "@bob:remote.example", EventID: "$member", Membership: "join"},
				},
				StateDelta: []storage.StateEntry{
					{Type: "m.room.create", EventID: "$create"},
					{Type: "m.room.member", StateKey: "@bob:remote.example", EventID: "$member", Membership: "join"},
					{Type: "m.room.member", StateKey: "@alice:warp.example", EventID: "$join", Membership: "join"},
				},
				StateDeltaIsWholeState: true,
				Extremities:            []string{"$join"},
			}); err != nil {
				t.Fatalf("Write join: %v", err)
			}

			timeline, err := rooms.Timeline(ctx, storage.TimelineRequest{
				RoomID: testRoomID, Limit: 10, Direction: storage.DirectionBackward,
			})
			if err != nil {
				t.Fatalf("Timeline: %v", err)
			}
			if ids := eventIDsOf(timeline); !slices.Equal(ids, []string{"$join"}) {
				t.Errorf("timeline = %v, want only the join", ids)
			}

			// The current state reads through to the outliers: that is
			// the whole point of keeping them in the events table.
			state, err := rooms.CurrentStateEvents(ctx, testRoomID)
			if err != nil {
				t.Fatalf("CurrentStateEvents: %v", err)
			}
			if len(state) != 3 {
				t.Errorf("current state has %d events, want 3", len(state))
			}

			top, err := rooms.MaxStreamOrdering(ctx)
			if err != nil {
				t.Fatalf("MaxStreamOrdering: %v", err)
			}
			at, err := rooms.StateAtStream(ctx, testRoomID, top)
			if err != nil {
				t.Fatalf("StateAtStream: %v", err)
			}
			if len(at) != 3 {
				t.Errorf("state at the join has %d entries, want 3", len(at))
			}
		})
	}
}

func TestOutlierIsPromotedWhenItArrivesAgainAsAnEvent(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			rooms := migratedFederationDB(t, driver).Rooms()
			if err := rooms.CreateRoom(ctx, newTestRoom(testRoomID)); err != nil {
				t.Fatalf("CreateRoom: %v", err)
			}
			writeMessages(t, rooms, testRoomID, "$first")
			if err := rooms.WriteOutliers(ctx, testRoomID, []storage.StoredEvent{
				newTestEvent(testRoomID, "$late", "m.room.message", nil),
			}); err != nil {
				t.Fatalf("WriteOutliers: %v", err)
			}
			writeMessages(t, rooms, testRoomID, "$late")

			stored, err := rooms.Event(ctx, testRoomID, "$late")
			if err != nil {
				t.Fatalf("Event: %v", err)
			}
			if stored.Outlier {
				t.Error("the event is still an outlier after being written as one")
			}
			timeline, err := rooms.Timeline(ctx, storage.TimelineRequest{
				RoomID: testRoomID, Limit: 10, Direction: storage.DirectionForward,
			})
			if err != nil {
				t.Fatalf("Timeline: %v", err)
			}
			if ids := eventIDsOf(timeline); !slices.Equal(ids, []string{"$first", "$late"}) {
				t.Errorf("timeline = %v, want [$first $late]", ids)
			}
		})
	}
}

func TestInviteStateLastsUntilTheNextMembershipChange(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			rooms := migratedFederationDB(t, driver).Rooms()
			if err := rooms.CreateRoom(ctx, newTestRoom(testRoomID)); err != nil {
				t.Fatalf("CreateRoom: %v", err)
			}
			const invitee = "@alice:warp.example"
			stripped := []byte(`[{"type":"m.room.name","state_key":"","sender":"@bob:remote.example","content":{"name":"Hi"}}]`)

			invite := newTestEvent(testRoomID, "$invite", "m.room.member", stateKey(invitee))
			if _, err := rooms.Write(ctx, storage.RoomWrite{
				Event:       invite,
				Outlier:     true,
				StateDelta:  []storage.StateEntry{{Type: "m.room.member", StateKey: invitee, EventID: "$invite", Membership: "invite"}},
				InviteState: stripped,
			}); err != nil {
				t.Fatalf("Write invite: %v", err)
			}

			got, err := rooms.InviteState(ctx, testRoomID, invitee)
			if err != nil {
				t.Fatalf("InviteState: %v", err)
			}
			if string(got) != string(stripped) {
				t.Errorf("InviteState = %s, want %s", got, stripped)
			}
			memberships, err := rooms.MembershipsForUser(ctx, invitee, []string{"invite"})
			if err != nil {
				t.Fatalf("MembershipsForUser: %v", err)
			}
			if len(memberships) != 1 || memberships[0].StreamOrdering == 0 {
				t.Errorf("an outlier invite must still have a stream position for /sync, got %+v", memberships)
			}

			leave := newTestEvent(testRoomID, "$leave", "m.room.member", stateKey(invitee))
			if _, err := rooms.Write(ctx, storage.RoomWrite{
				Event:      leave,
				Outlier:    true,
				StateDelta: []storage.StateEntry{{Type: "m.room.member", StateKey: invitee, EventID: "$leave", Membership: "leave"}},
			}); err != nil {
				t.Fatalf("Write leave: %v", err)
			}
			if got, err := rooms.InviteState(ctx, testRoomID, invitee); err != nil || got != nil {
				t.Errorf("InviteState after a rejection = %s, %v; want nothing", got, err)
			}
		})
	}
}

func TestServerKeysAreUpsertedAndFiltered(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			fed := migratedFederationDB(t, driver).Federation()

			if err := fed.StoreServerKeys(ctx, []storage.ServerKey{
				{ServerName: "remote.example", KeyID: "ed25519:a", PublicKey: "AAA", ValidUntilTS: 10, FetchedAtMS: 1},
				{ServerName: "remote.example", KeyID: "ed25519:b", PublicKey: "BBB", ValidUntilTS: 20, ExpiredTS: 15, FetchedAtMS: 1},
				{ServerName: "other.example", KeyID: "ed25519:a", PublicKey: "CCC", ValidUntilTS: 30, FetchedAtMS: 1},
			}); err != nil {
				t.Fatalf("StoreServerKeys: %v", err)
			}
			if err := fed.StoreServerKeys(ctx, []storage.ServerKey{
				{ServerName: "remote.example", KeyID: "ed25519:a", PublicKey: "AAA", ValidUntilTS: 99, FetchedAtMS: 2},
			}); err != nil {
				t.Fatalf("StoreServerKeys again: %v", err)
			}

			all, err := fed.ServerKeys(ctx, "remote.example", nil)
			if err != nil {
				t.Fatalf("ServerKeys: %v", err)
			}
			if len(all) != 2 {
				t.Fatalf("ServerKeys found %d keys, want 2", len(all))
			}
			if all[0].KeyID != "ed25519:a" || all[0].ValidUntilTS != 99 {
				t.Errorf("the refetched key was not replaced: %+v", all[0])
			}
			if all[1].ExpiredTS != 15 {
				t.Errorf("an old key lost its expiry: %+v", all[1])
			}

			one, err := fed.ServerKeys(ctx, "remote.example", []string{"ed25519:b"})
			if err != nil {
				t.Fatalf("ServerKeys(b): %v", err)
			}
			if len(one) != 1 || one[0].PublicKey != "BBB" {
				t.Errorf("ServerKeys(b) = %+v", one)
			}
		})
	}
}

func TestQueueIsPerDestinationAndInOrder(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			fed := migratedFederationDB(t, driver).Federation()

			queued, err := fed.Enqueue(ctx, []storage.QueuedItem{
				{Destination: "a.example", Kind: storage.QueueKindPDU, EventID: "$1", Payload: []byte(`{"n":1}`)},
				{Destination: "b.example", Kind: storage.QueueKindPDU, EventID: "$1", Payload: []byte(`{"n":1}`)},
				{Destination: "a.example", Kind: storage.QueueKindEDU, Payload: []byte(`{"edu":true}`)},
				{Destination: "a.example", Kind: storage.QueueKindPDU, EventID: "$2", Payload: []byte(`{"n":2}`)},
			})
			if err != nil {
				t.Fatalf("Enqueue: %v", err)
			}
			for i := 1; i < len(queued); i++ {
				if queued[i].ID <= queued[i-1].ID {
					t.Fatalf("queue IDs are not increasing: %d then %d", queued[i-1].ID, queued[i].ID)
				}
			}

			pdus, err := fed.Queued(ctx, "a.example", storage.QueueKindPDU, 10)
			if err != nil {
				t.Fatalf("Queued: %v", err)
			}
			if len(pdus) != 2 || pdus[0].EventID != "$1" || pdus[1].EventID != "$2" {
				t.Fatalf("Queued(a, pdu) = %+v", pdus)
			}
			if string(pdus[1].Payload) != `{"n":2}` {
				t.Errorf("payload = %s", pdus[1].Payload)
			}

			destinations, err := fed.QueuedDestinations(ctx)
			if err != nil {
				t.Fatalf("QueuedDestinations: %v", err)
			}
			if !slices.Equal(destinations, []string{"a.example", "b.example"}) {
				t.Errorf("QueuedDestinations = %v", destinations)
			}

			if err := fed.Dequeue(ctx, []int64{pdus[0].ID}); err != nil {
				t.Fatalf("Dequeue: %v", err)
			}
			left, err := fed.Queued(ctx, "a.example", storage.QueueKindPDU, 10)
			if err != nil {
				t.Fatalf("Queued after Dequeue: %v", err)
			}
			if len(left) != 1 || left[0].EventID != "$2" {
				t.Errorf("after dequeue = %+v", left)
			}
		})
	}
}

func TestDestinationHealth(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			fed := migratedFederationDB(t, driver).Federation()

			d, err := fed.Destination(ctx, "never.example")
			if err != nil {
				t.Fatalf("Destination: %v", err)
			}
			if d.ServerName != "never.example" || d.Failures != 0 {
				t.Errorf("an unknown destination = %+v, want a healthy zero", d)
			}

			if err := fed.SetDestination(ctx, storage.Destination{
				ServerName: "down.example", Failures: 3, RetryAtMS: 500, LastError: "connection refused",
			}); err != nil {
				t.Fatalf("SetDestination: %v", err)
			}
			if err := fed.SetDestination(ctx, storage.Destination{
				ServerName: "down.example", Failures: 4, RetryAtMS: 900, LastError: "timeout",
			}); err != nil {
				t.Fatalf("SetDestination again: %v", err)
			}
			if _, err := fed.Enqueue(ctx, []storage.QueuedItem{
				{Destination: "busy.example", Kind: storage.QueueKindPDU, Payload: []byte(`{}`)},
				{Destination: "busy.example", Kind: storage.QueueKindEDU, Payload: []byte(`{}`)},
				{Destination: "down.example", Kind: storage.QueueKindPDU, Payload: []byte(`{}`)},
			}); err != nil {
				t.Fatalf("Enqueue: %v", err)
			}

			all, err := fed.Destinations(ctx)
			if err != nil {
				t.Fatalf("Destinations: %v", err)
			}
			if len(all) != 2 {
				t.Fatalf("Destinations = %+v, want busy and down", all)
			}
			busy, down := all[0], all[1]
			if busy.ServerName != "busy.example" || busy.QueuedPDUs != 1 || busy.QueuedEDUs != 1 || busy.Failures != 0 {
				t.Errorf("busy = %+v", busy)
			}
			if down.ServerName != "down.example" || down.Failures != 4 || down.LastError != "timeout" || down.QueuedPDUs != 1 {
				t.Errorf("down = %+v", down)
			}
		})
	}
}

func TestInboundTransactionsAreRememberedOnce(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			fed := migratedFederationDB(t, driver).Federation()

			if _, ok, err := fed.InboundTransaction(ctx, "remote.example", "t1"); err != nil || ok {
				t.Fatalf("InboundTransaction before = %v, %v", ok, err)
			}
			if err := fed.RecordInboundTransaction(ctx, "remote.example", "t1", []byte(`{"pdus":{}}`), 1); err != nil {
				t.Fatalf("RecordInboundTransaction: %v", err)
			}
			if err := fed.RecordInboundTransaction(ctx, "remote.example", "t1", []byte(`{"other":1}`), 2); err != nil {
				t.Fatalf("RecordInboundTransaction again: %v", err)
			}
			got, ok, err := fed.InboundTransaction(ctx, "remote.example", "t1")
			if err != nil || !ok || string(got) != `{"pdus":{}}` {
				t.Errorf("InboundTransaction = %s, %v, %v; want the first response", got, ok, err)
			}
		})
	}
}

func eventIDsOf(events []storage.StoredEvent) []string {
	out := make([]string, 0, len(events))
	for _, e := range events {
		out = append(out, e.EventID)
	}
	return out
}

func TestHasJoinedMemberOn(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			rooms := migratedFederationDB(t, driver).Rooms()
			if err := rooms.CreateRoom(ctx, newTestRoom(testRoomID)); err != nil {
				t.Fatal(err)
			}
			for _, m := range []struct{ user, membership string }{
				{"@a:one.example", "join"},
				{"@b:two.example", "leave"},
				{"@c:x_y.example", "join"},
			} {
				e := newTestEvent(testRoomID, "$"+m.user, "m.room.member", stateKey(m.user))
				if _, err := rooms.Write(ctx, storage.RoomWrite{
					Event:       e,
					StateDelta:  []storage.StateEntry{{Type: "m.room.member", StateKey: m.user, EventID: e.EventID, Membership: m.membership}},
					Extremities: []string{e.EventID},
				}); err != nil {
					t.Fatal(err)
				}
			}
			cases := map[string]bool{
				"one.example":  true,
				"two.example":  false, // left
				"x_y.example":  true,
				"xay.example":  false, // '_' is not a LIKE wildcard here
				"e.example":    false, // a suffix of a name is not the name
				"none.example": false,
			}
			for server, want := range cases {
				got, err := rooms.HasJoinedMemberOn(ctx, testRoomID, server)
				if err != nil {
					t.Fatal(err)
				}
				if got != want {
					t.Errorf("HasJoinedMemberOn(%s) = %v, want %v", server, got, want)
				}
			}
		})
	}
}