Warp/server/internal/storage/rooms_test.go

package storage_test

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

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

// migratedRoomDB opens a driver, migrates it and empties the roomserver
// tables, which is what every room-store test needs first.
func migratedRoomDB(t *testing.T, driver string) storage.Database {
	t.Helper()
	db := openTestDB(t, driver)
	if err := db.Migrate(t.Context()); err != nil {
		t.Fatalf("Migrate: %v", err)
	}
	resetRoomTables(t, driver)
	return db
}

// resetRoomTables empties the roomserver tables and rewinds the counters.
//
// SQLite gets a fresh file per test. The PostgreSQL tests share one
// database, so without this the second test to create "!room" would collide
// with the first — and only in CI, where PostgreSQL is available.
func resetRoomTables(t *testing.T, driver string) {
	t.Helper()
	if driver != storage.DriverPostgres {
		return
	}

	db, err := sql.Open("pgx", os.Getenv("WARP_TEST_POSTGRES_DSN"))
	if err != nil {
		t.Fatalf("open postgres to reset the room tables: %v", err)
	}
	defer func() { _ = db.Close() }()

	if _, err := db.ExecContext(t.Context(), `TRUNCATE
		roomserver_state_entries, roomserver_state_snapshots,
		roomserver_current_state, roomserver_extremities,
		roomserver_transactions, roomserver_aliases, roomserver_invite_state,
		roomserver_events, roomserver_rooms, roomserver_event_search`); err != nil {
		t.Fatalf("reset the room tables: %v", err)
	}
	if _, err := db.ExecContext(t.Context(),
		`UPDATE roomserver_sequences SET value = 0`); err != nil {
		t.Fatalf("rewind the room counters: %v", err)
	}
}

// sqlDB reaches the underlying pool, for the one assertion that is about
// the schema rather than the interface: that identical state is stored once.
func sqlDB(t *testing.T, db storage.Database) *sql.DB {
	t.Helper()
	pool, ok := db.(interface{ DB() *sql.DB })
	if !ok {
		t.Skipf("driver %s does not expose a database/sql pool", db.Driver())
	}
	return pool.DB()
}

// The room store never parses an event, so these tests feed it event JSON
// that is only plausible. What is under test is the bookkeeping: stream
// ordering, state snapshots, current state and the member counters.
const testRoomID = "!room:warp.example"

func newTestRoom(id string) storage.NewRoom {
	return storage.NewRoom{
		ID:          id,
		Version:     "11",
		Creator:     "@alice:warp.example",
		CreatedAtMS: 1_700_000_000_000,
	}
}

// stateKey returns a pointer to s, for StoredEvent.StateKey.
func stateKey(s string) *string { return &s }

func newTestEvent(roomID, eventID, eventType string, key *string) storage.StoredEvent {
	return storage.StoredEvent{
		RoomID:         roomID,
		EventID:        eventID,
		Type:           eventType,
		StateKey:       key,
		Sender:         "@alice:warp.example",
		Depth:          1,
		OriginServerTS: 1_700_000_000_000,
		JSON:           []byte(`{"type":"` + eventType + `"}`),
	}
}

func TestRoomCreateAndRead(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			rooms := migratedRoomDB(t, driver).Rooms()

			if _, err := rooms.Room(ctx, testRoomID); !errors.Is(err, storage.ErrNoSuchRoom) {
				t.Errorf("Room of an unknown room = %v, want ErrNoSuchRoom", err)
			}
			if err := rooms.CreateRoom(ctx, newTestRoom(testRoomID)); err != nil {
				t.Fatalf("CreateRoom: %v", err)
			}
			if err := rooms.CreateRoom(ctx, newTestRoom(testRoomID)); !errors.Is(err, storage.ErrRoomExists) {
				t.Errorf("CreateRoom twice = %v, want ErrRoomExists", err)
			}

			room, err := rooms.Room(ctx, testRoomID)
			if err != nil {
				t.Fatalf("Room: %v", err)
			}
			if room.Version != "11" || room.Creator != "@alice:warp.example" || room.Published {
				t.Errorf("Room returned %+v", room)
			}
		})
	}
}

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

			// The create event: no state before it, and it is the room's
			// only state and only extremity.
			create := newTestEvent(testRoomID, "$create", "m.room.create", stateKey(""))
			createState := []storage.StateEntry{
				{Type: "m.room.create", StateKey: "", EventID: "$create"},
			}
			ordering, err := rooms.Write(ctx, storage.RoomWrite{
				Event:       create,
				StateDelta:  createState,
				Extremities: []string{"$create"},
			})
			if err != nil {
				t.Fatalf("Write the create event: %v", err)
			}
			if ordering != 1 {
				t.Errorf("the first event got stream ordering %d, want 1", ordering)
			}

			// A join, authorised against the state above.
			join := newTestEvent(testRoomID, "$join", "m.room.member", stateKey("@alice:warp.example"))
			join.Depth = 2
			ordering, err = rooms.Write(ctx, storage.RoomWrite{
				Event:       join,
				StateBefore: createState,
				StateDelta: []storage.StateEntry{{
					Type: "m.room.member", StateKey: "@alice:warp.example",
					EventID: "$join", Membership: "join",
				}},
				Extremities: []string{"$join"},
			})
			if err != nil {
				t.Fatalf("Write the join: %v", err)
			}
			if ordering != 2 {
				t.Errorf("the second event got stream ordering %d, want 2", ordering)
			}

			// A message, which changes no state at all.
			message := newTestEvent(testRoomID, "$message", "m.room.message", nil)
			message.Depth = 3
			if _, err := rooms.Write(ctx, storage.RoomWrite{
				Event: message,
				StateBefore: append(createState, storage.StateEntry{
					Type: "m.room.member", StateKey: "@alice:warp.example", EventID: "$join",
				}),
				Extremities: []string{"$message"},
			}); err != nil {
				t.Fatalf("Write the message: %v", err)
			}

			state, err := rooms.CurrentState(ctx, testRoomID)
			if err != nil {
				t.Fatalf("CurrentState: %v", err)
			}
			if len(state) != 2 {
				t.Fatalf("current state has %d entries, want 2: %+v", len(state), state)
			}

			stored, err := rooms.Event(ctx, testRoomID, "$message")
			if err != nil {
				t.Fatalf("Event: %v", err)
			}
			if stored.IsState() {
				t.Error("a message event came back as a state event")
			}
			if stored.StreamOrdering != 3 {
				t.Errorf("stream ordering = %d, want 3", stored.StreamOrdering)
			}
			if _, err := rooms.Event(ctx, testRoomID, "$nope"); !errors.Is(err, storage.ErrNoSuchEvent) {
				t.Errorf("Event of an unknown ID = %v, want ErrNoSuchEvent", err)
			}

			// The state recorded against the message is the state it was
			// authorised against, not the state now.
			before, err := rooms.StateBefore(ctx, testRoomID, "$message")
			if err != nil {
				t.Fatalf("StateBefore: %v", err)
			}
			if len(before) != 2 {
				t.Errorf("the state before the message has %d entries, want 2", len(before))
			}

			extremities, err := rooms.Extremities(ctx, testRoomID)
			if err != nil {
				t.Fatalf("Extremities: %v", err)
			}
			if len(extremities) != 1 || extremities[0].EventID != "$message" {
				t.Errorf("extremities = %+v, want just $message", extremities)
			}

			// The create event is a state event with the empty state key,
			// which must not read back as "no state key".
			storedCreate, err := rooms.Event(ctx, testRoomID, "$create")
			if err != nil {
				t.Fatalf("Event of the create event: %v", err)
			}
			if !storedCreate.IsState() {
				t.Error("the create event came back without a state key")
			} else if *storedCreate.StateKey != "" {
				t.Errorf("the create event state key = %q, want empty", *storedCreate.StateKey)
			}
		})
	}
}

// TestRoomStateSnapshotsAreShared is the reason a per-event state snapshot
// is affordable: two events authorised against the same state must not
// store it twice.
func TestRoomStateSnapshotsAreShared(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			db := migratedRoomDB(t, driver)
			rooms := db.Rooms()
			if err := rooms.CreateRoom(ctx, newTestRoom(testRoomID)); err != nil {
				t.Fatalf("CreateRoom: %v", err)
			}

			state := []storage.StateEntry{
				{Type: "m.room.create", StateKey: "", EventID: "$create"},
			}
			for _, id := range []string{"$a", "$b", "$c"} {
				if _, err := rooms.Write(ctx, storage.RoomWrite{
					Event:       newTestEvent(testRoomID, id, "m.room.message", nil),
					StateBefore: state,
					Extremities: []string{id},
				}); err != nil {
					t.Fatalf("Write %s: %v", id, err)
				}
			}

			var snapshots int
			if err := sqlDB(t, db).QueryRowContext(ctx,
				`SELECT COUNT(*) FROM roomserver_state_snapshots`).Scan(&snapshots); err != nil {
				t.Fatalf("count snapshots: %v", err)
			}
			if snapshots != 1 {
				t.Errorf("three events sharing one state stored %d snapshots, want 1", snapshots)
			}
		})
	}
}

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

			write := func(eventID, userID, membership string) {
				t.Helper()
				e := newTestEvent(testRoomID, eventID, "m.room.member", stateKey(userID))
				if _, err := rooms.Write(ctx, storage.RoomWrite{
					Event: e,
					StateDelta: []storage.StateEntry{{
						Type: "m.room.member", StateKey: userID,
						EventID: eventID, Membership: membership,
					}},
					Extremities: []string{eventID},
				}); err != nil {
					t.Fatalf("Write %s: %v", eventID, err)
				}
			}
			write("$alice", "@alice:warp.example", "join")
			write("$bob", "@bob:warp.example", "invite")
			write("$carol", "@carol:warp.example", "join")

			// A user who was never in the room gets an empty membership
			// rather than an error.
			absent, err := rooms.Membership(ctx, testRoomID, "@dave:warp.example")
			if err != nil {
				t.Fatalf("Membership of a stranger: %v", err)
			}
			if absent.Membership != "" {
				t.Errorf("membership of a stranger = %q, want empty", absent.Membership)
			}

			alice, err := rooms.Membership(ctx, testRoomID, "@alice:warp.example")
			if err != nil {
				t.Fatalf("Membership: %v", err)
			}
			if alice.Membership != "join" || alice.EventID != "$alice" {
				t.Errorf("Membership returned %+v", alice)
			}

			joined, err := rooms.MembershipEntries(ctx, testRoomID, []string{"join"})
			if err != nil {
				t.Fatalf("MembershipEntries: %v", err)
			}
			if len(joined) != 2 {
				t.Errorf("%d joined members, want 2: %+v", len(joined), joined)
			}

			all, err := rooms.Members(ctx, testRoomID, nil)
			if err != nil {
				t.Fatalf("Members: %v", err)
			}
			if len(all) != 3 {
				t.Errorf("%d member events, want 3", len(all))
			}

			// Leaving replaces the entry rather than adding one, and the
			// member count follows.
			write("$aliceleft", "@alice:warp.example", "leave")
			summary, err := rooms.RoomSummary(ctx, testRoomID)
			if err != nil {
				t.Fatalf("RoomSummary: %v", err)
			}
			if summary.JoinedMembers != 1 {
				t.Errorf("joined members = %d, want 1", summary.JoinedMembers)
			}

			forAlice, err := rooms.RoomsForUser(ctx, "@alice:warp.example", []string{"join"})
			if err != nil {
				t.Fatalf("RoomsForUser: %v", err)
			}
			if len(forAlice) != 0 {
				t.Errorf("alice is still joined to %v after leaving", forAlice)
			}
			forCarol, err := rooms.RoomsForUser(ctx, "@carol:warp.example", []string{"join"})
			if err != nil {
				t.Fatalf("RoomsForUser: %v", err)
			}
			if len(forCarol) != 1 || forCarol[0] != testRoomID {
				t.Errorf("carol's joined rooms = %v, want [%s]", forCarol, testRoomID)
			}
		})
	}
}

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

			ref := storage.TransactionRef{Localpart: "alice", DeviceID: "DEV", TxnID: "txn1"}
			if _, err := rooms.TransactionEvent(ctx, ref); !errors.Is(err, storage.ErrNoSuchEvent) {
				t.Errorf("TransactionEvent of an unseen ID = %v, want ErrNoSuchEvent", err)
			}

			target := newTestEvent(testRoomID, "$target", "m.room.message", nil)
			if _, err := rooms.Write(ctx, storage.RoomWrite{
				Event:       target,
				Extremities: []string{"$target"},
				Transaction: &ref,
			}); err != nil {
				t.Fatalf("Write the target: %v", err)
			}

			got, err := rooms.TransactionEvent(ctx, ref)
			if err != nil {
				t.Fatalf("TransactionEvent: %v", err)
			}
			if got != "$target" {
				t.Errorf("TransactionEvent = %q, want $target", got)
			}

			redaction := newTestEvent(testRoomID, "$redaction", "m.room.redaction", nil)
			if _, err := rooms.Write(ctx, storage.RoomWrite{
				Event:        redaction,
				Extremities:  []string{"$redaction"},
				Redacts:      "$target",
				RedactedJSON: []byte(`{"type":"m.room.message","content":{}}`),
			}); err != nil {
				t.Fatalf("Write the redaction: %v", err)
			}

			stored, err := rooms.Event(ctx, testRoomID, "$target")
			if err != nil {
				t.Fatalf("Event: %v", err)
			}
			if stored.RedactedBy != "$redaction" {
				t.Errorf("redacted_by = %q, want $redaction", stored.RedactedBy)
			}
			if string(stored.JSON) != `{"type":"m.room.message","content":{}}` {
				t.Errorf("the stored event was not replaced with its redacted form: %s", stored.JSON)
			}
		})
	}
}

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

			const aliasName = "#general:warp.example"
			newAlias := storage.RoomAlias{
				Alias:       aliasName,
				RoomID:      testRoomID,
				Creator:     "@alice:warp.example",
				CreatedAtMS: 1,
			}
			if _, err := rooms.Alias(ctx, aliasName); !errors.Is(err, storage.ErrNoSuchAlias) {
				t.Errorf("Alias of an unknown alias = %v, want ErrNoSuchAlias", err)
			}
			if err := rooms.CreateAlias(ctx, newAlias); err != nil {
				t.Fatalf("CreateAlias: %v", err)
			}
			taken := newAlias
			taken.Creator = "@bob:warp.example"
			if err := rooms.CreateAlias(ctx, taken); !errors.Is(err, storage.ErrAliasExists) {
				t.Errorf("CreateAlias twice = %v, want ErrAliasExists", err)
			}

			stored, err := rooms.Alias(ctx, aliasName)
			if err != nil {
				t.Fatalf("Alias: %v", err)
			}
			if stored.RoomID != testRoomID || stored.Creator != "@alice:warp.example" {
				t.Errorf("Alias returned %+v", stored)
			}
			aliases, err := rooms.AliasesForRoom(ctx, testRoomID)
			if err != nil || len(aliases) != 1 || aliases[0] != aliasName {
				t.Fatalf("AliasesForRoom = %v, %v", aliases, err)
			}

			if err := rooms.DeleteAlias(ctx, aliasName); err != nil {
				t.Fatalf("DeleteAlias: %v", err)
			}
			if err := rooms.DeleteAlias(ctx, aliasName); !errors.Is(err, storage.ErrNoSuchAlias) {
				t.Errorf("DeleteAlias twice = %v, want ErrNoSuchAlias", err)
			}

			// The directory is empty until the room is published.
			published, total, err := rooms.PublishedRooms(ctx, storage.RoomFilter{})
			if err != nil {
				t.Fatalf("PublishedRooms: %v", err)
			}
			if total != 0 || len(published) != 0 {
				t.Errorf("an unpublished room is in the directory: %+v", published)
			}
			if err := rooms.SetPublished(ctx, testRoomID, true); err != nil {
				t.Fatalf("SetPublished: %v", err)
			}
			if err := rooms.SetPublished(ctx, "!nope:warp.example", true); !errors.Is(err, storage.ErrNoSuchRoom) {
				t.Errorf("SetPublished on an unknown room = %v, want ErrNoSuchRoom", err)
			}
			published, total, err = rooms.PublishedRooms(ctx, storage.RoomFilter{})
			if err != nil {
				t.Fatalf("PublishedRooms: %v", err)
			}
			if total != 1 || len(published) != 1 || published[0].ID != testRoomID {
				t.Errorf("PublishedRooms = %+v (total %d)", published, total)
			}
		})
	}
}

func TestRoomListingAndPurge(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			rooms := migratedRoomDB(t, driver).Rooms()

			for _, id := range []string{"!a:warp.example", "!b:warp.example"} {
				if err := rooms.CreateRoom(ctx, newTestRoom(id)); err != nil {
					t.Fatalf("CreateRoom %s: %v", id, err)
				}
				if _, err := rooms.Write(ctx, storage.RoomWrite{
					Event:       newTestEvent(id, "$create"+id, "m.room.create", stateKey("")),
					Extremities: []string{"$create" + id},
					Profile:     &storage.RoomProfile{Name: "Room " + id, JoinRule: "invite"},
				}); err != nil {
					t.Fatalf("Write %s: %v", id, err)
				}
			}

			listed, total, err := rooms.ListRooms(ctx, storage.RoomFilter{})
			if err != nil {
				t.Fatalf("ListRooms: %v", err)
			}
			if total != 2 || len(listed) != 2 {
				t.Fatalf("ListRooms returned %d of %d, want 2 of 2", len(listed), total)
			}
			if listed[0].Name == "" || listed[0].JoinRule != "invite" {
				t.Errorf("the denormalised profile was not stored: %+v", listed[0])
			}

			// The search covers the name as well as the ID.
			found, total, err := rooms.ListRooms(ctx, storage.RoomFilter{Search: "room !a"})
			if err != nil {
				t.Fatalf("ListRooms with a search: %v", err)
			}
			if total != 1 || len(found) != 1 || found[0].ID != "!a:warp.example" {
				t.Errorf("searching for a name returned %+v (total %d)", found, total)
			}

			stats, err := rooms.RoomStats(ctx)
			if err != nil {
				t.Fatalf("RoomStats: %v", err)
			}
			if stats.Rooms != 2 || stats.Events != 2 {
				t.Errorf("RoomStats = %+v, want 2 rooms and 2 events", stats)
			}

			if err := rooms.PurgeRoom(ctx, "!a:warp.example"); err != nil {
				t.Fatalf("PurgeRoom: %v", err)
			}
			if err := rooms.PurgeRoom(ctx, "!a:warp.example"); !errors.Is(err, storage.ErrNoSuchRoom) {
				t.Errorf("PurgeRoom twice = %v, want ErrNoSuchRoom", err)
			}
			if _, err := rooms.Room(ctx, "!a:warp.example"); !errors.Is(err, storage.ErrNoSuchRoom) {
				t.Errorf("the purged room is still there: %v", err)
			}
			if _, err := rooms.Event(ctx, "!b:warp.example", "$create!b:warp.example"); err != nil {
				t.Errorf("purging one room disturbed another: %v", err)
			}
		})
	}
}