Warp/server/internal/storage/keys_test.go

package storage_test

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

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

// migratedKeysDB opens a driver with the keyserver tables empty. The
// device list counter is left where it is: every assertion below is about
// positions relative to one read before the write.
func migratedKeysDB(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)
	}
	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 keyserver tables: %v", err)
		}
		defer func() { _ = pool.Close() }()
		if _, err := pool.ExecContext(t.Context(), `TRUNCATE keyserver_device_keys, keyserver_one_time_keys,
			keyserver_fallback_keys, keyserver_cross_signing_keys, keyserver_signatures,
			keyserver_device_lists, keyserver_backup_versions, keyserver_backup_keys`); err != nil {
			t.Fatalf("reset the keyserver tables: %v", err)
		}
	}
	return db
}

const alice = "@alice:warp.example"

func TestDeviceKeysAndTheDeviceListStream(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			db := migratedKeysDB(t, driver)
			keys := db.Keys()

			before, err := db.Sync().StreamPositions(ctx)
			if err != nil {
				t.Fatalf("StreamPositions: %v", err)
			}

			first, err := keys.SetDeviceKeys(ctx, storage.DeviceKeys{UserID: alice, DeviceID: "D1", JSON: []byte(`{"k":1}`)})
			if err != nil {
				t.Fatalf("SetDeviceKeys: %v", err)
			}
			if first.StreamPos != before.DeviceList+1 || first.PrevPos != 0 || first.DeviceID != "D1" {
				t.Errorf("first update = %+v, want position %d with no previous", first, before.DeviceList+1)
			}

			// The same keys again are not a change.
			again, err := keys.SetDeviceKeys(ctx, storage.DeviceKeys{UserID: alice, DeviceID: "D1", JSON: []byte(`{"k":1}`)})
			if err != nil {
				t.Fatalf("SetDeviceKeys again: %v", err)
			}
			if again != (storage.DeviceListUpdate{}) {
				t.Errorf("re-uploading the same keys moved the stream: %+v", again)
			}

			second, err := keys.SetDeviceKeys(ctx, storage.DeviceKeys{UserID: alice, DeviceID: "D2", JSON: []byte(`{"k":2}`)})
			if err != nil {
				t.Fatalf("SetDeviceKeys D2: %v", err)
			}
			if second.PrevPos != first.StreamPos {
				t.Errorf("second update's previous = %d, want %d", second.PrevPos, first.StreamPos)
			}

			after, _ := db.Sync().StreamPositions(ctx)
			if after.DeviceList != second.StreamPos {
				t.Errorf("StreamPositions.DeviceList = %d, want %d", after.DeviceList, second.StreamPos)
			}

			all, err := keys.DeviceKeys(ctx, alice, nil)
			if err != nil || len(all) != 2 || all[0].DeviceID != "D1" || string(all[1].JSON) != `{"k":2}` {
				t.Fatalf("DeviceKeys = %+v, %v", all, err)
			}
			one, _ := keys.DeviceKeys(ctx, alice, []string{"D2", "nope"})
			if len(one) != 1 || one[0].DeviceID != "D2" {
				t.Errorf("DeviceKeys(D2) = %+v", one)
			}

			changes, err := keys.DeviceListChanges(ctx, []string{alice, "@bob:warp.example"}, first.StreamPos, 0)
			if err != nil || len(changes) != 1 || changes[0].StreamPos != second.StreamPos {
				t.Errorf("DeviceListChanges after the first = %+v, %v", changes, err)
			}
			if changes, _ := keys.DeviceListChanges(ctx, []string{alice}, second.StreamPos, 0); len(changes) != 0 {
				t.Errorf("DeviceListChanges after the last = %+v, want none", changes)
			}

			// Deleting two devices moves the stream once for each, so a
			// remote server sees two updates rather than one and a repeat.
			if err := keys.AddOneTimeKeys(ctx, alice, "D1", []storage.OneTimeKey{{KeyID: "curve25519:a", JSON: []byte(`"x"`)}}); err != nil {
				t.Fatal(err)
			}
			deleted, err := keys.DeleteDeviceKeys(ctx, alice, nil)
			if err != nil || len(deleted) != 2 {
				t.Fatalf("DeleteDeviceKeys = %+v, %v", deleted, err)
			}
			if deleted[0].StreamPos == deleted[1].StreamPos || deleted[1].PrevPos != deleted[0].StreamPos {
				t.Errorf("deletions share or skip positions: %+v", deleted)
			}
			if left, _ := keys.DeviceKeys(ctx, alice, nil); len(left) != 0 {
				t.Errorf("keys survived DeleteDeviceKeys: %+v", left)
			}
			if counts, _ := keys.OneTimeKeyCounts(ctx, alice, "D1"); len(counts) != 0 {
				t.Errorf("one-time keys survived DeleteDeviceKeys: %v", counts)
			}
			if pos, _ := keys.DeviceListPosition(ctx, alice); pos != deleted[1].StreamPos {
				t.Errorf("DeviceListPosition = %d, want %d", pos, deleted[1].StreamPos)
			}

			// Deleting devices that had no keys moves nothing.
			if none, err := keys.DeleteDeviceKeys(ctx, alice, []string{"D9"}); err != nil || len(none) != 0 {
				t.Errorf("DeleteDeviceKeys of a keyless device = %+v, %v", none, err)
			}

			top, err := keys.MarkDeviceListsChanged(ctx, []string{"@carol:remote.example", "@carol:remote.example"})
			if err != nil || top != deleted[1].StreamPos+1 {
				t.Errorf("MarkDeviceListsChanged = %d, %v; want one move to %d", top, err, deleted[1].StreamPos+1)
			}
		})
	}
}

func TestOneTimeAndFallbackKeys(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			keys := migratedKeysDB(t, driver).Keys()

			otks := []storage.OneTimeKey{
				{KeyID: "signed_curve25519:AAAA", JSON: []byte(`{"key":"a"}`)},
				{KeyID: "signed_curve25519:BBBB", JSON: []byte(`{"key":"b"}`)},
				{KeyID: "curve25519:CCCC", JSON: []byte(`"c"`)},
			}
			if err := keys.AddOneTimeKeys(ctx, alice, "D1", otks); err != nil {
				t.Fatalf("AddOneTimeKeys: %v", err)
			}
			// The same keys again are fine; a different key under a used
			// ID is not, and nothing of that batch is kept.
			if err := keys.AddOneTimeKeys(ctx, alice, "D1", otks[:1]); err != nil {
				t.Errorf("re-uploading the same key: %v", err)
			}
			err := keys.AddOneTimeKeys(ctx, alice, "D1", []storage.OneTimeKey{
				{KeyID: "signed_curve25519:DDDD", JSON: []byte(`{"key":"d"}`)},
				{KeyID: "signed_curve25519:AAAA", JSON: []byte(`{"key":"other"}`)},
			})
			if !errors.Is(err, storage.ErrKeyConflict) {
				t.Errorf("conflicting key = %v, want ErrKeyConflict", err)
			}

			counts, err := keys.OneTimeKeyCounts(ctx, alice, "D1")
			if err != nil || counts["signed_curve25519"] != 2 || counts["curve25519"] != 1 {
				t.Fatalf("OneTimeKeyCounts = %v, %v", counts, err)
			}

			fallback := storage.OneTimeKey{KeyID: "signed_curve25519:FFFF", JSON: []byte(`{"key":"f","fallback":true}`)}
			if err := keys.SetFallbackKeys(ctx, alice, "D1", []storage.OneTimeKey{fallback}); err != nil {
				t.Fatalf("SetFallbackKeys: %v", err)
			}
			if types, _ := keys.UnusedFallbackKeyTypes(ctx, alice, "D1"); !slices.Equal(types, []string{"signed_curve25519"}) {
				t.Errorf("UnusedFallbackKeyTypes = %v", types)
			}

			// Two one-time keys, oldest ID first, then the fallback key,
			// as often as it is asked for.
			var claimed []string
			for range 4 {
				k, ok, err := keys.ClaimKey(ctx, alice, "D1", "signed_curve25519")
				if err != nil || !ok {
					t.Fatalf("ClaimKey = %+v, %v, %v", k, ok, err)
				}
				claimed = append(claimed, k.KeyID)
				if k.Fallback != (k.KeyID == fallback.KeyID) {
					t.Errorf("claimed %s with Fallback=%v", k.KeyID, k.Fallback)
				}
			}
			want := []string{"signed_curve25519:AAAA", "signed_curve25519:BBBB", fallback.KeyID, fallback.KeyID}
			if !slices.Equal(claimed, want) {
				t.Errorf("claimed %v, want %v", claimed, want)
			}
			if types, _ := keys.UnusedFallbackKeyTypes(ctx, alice, "D1"); len(types) != 0 {
				t.Errorf("a claimed fallback key is still reported unused: %v", types)
			}

			// Re-uploading the same fallback key keeps it used; a new one
			// is fresh.
			_ = keys.SetFallbackKeys(ctx, alice, "D1", []storage.OneTimeKey{fallback})
			if types, _ := keys.UnusedFallbackKeyTypes(ctx, alice, "D1"); len(types) != 0 {
				t.Errorf("re-uploading a used fallback key made it unused: %v", types)
			}
			_ = keys.SetFallbackKeys(ctx, alice, "D1", []storage.OneTimeKey{{KeyID: "signed_curve25519:GGGG", JSON: []byte(`{}`)}})
			if types, _ := keys.UnusedFallbackKeyTypes(ctx, alice, "D1"); len(types) != 1 {
				t.Errorf("a new fallback key is not reported unused: %v", types)
			}

			if _, ok, err := keys.ClaimKey(ctx, alice, "D2", "signed_curve25519"); ok || err != nil {
				t.Errorf("claiming from a device with no keys = %v, %v", ok, err)
			}
		})
	}
}

func TestCrossSigningKeysAndSignatures(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			keys := migratedKeysDB(t, driver).Keys()

			u1, err := keys.SetCrossSigningKeys(ctx, alice, map[string][]byte{
				storage.CrossSigningMaster:      []byte(`{"m":1}`),
				storage.CrossSigningSelfSigning: []byte(`{"s":1}`),
			})
			if err != nil || u1.StreamPos == 0 {
				t.Fatalf("SetCrossSigningKeys = %+v, %v", u1, err)
			}
			_, err = keys.SetCrossSigningKeys(ctx, alice, map[string][]byte{
				storage.CrossSigningMaster:      []byte(`{"m":2}`),
				storage.CrossSigningSelfSigning: nil,
			})
			if err != nil {
				t.Fatal(err)
			}
			got, err := keys.CrossSigningKeys(ctx, alice)
			if err != nil || len(got) != 1 || string(got[storage.CrossSigningMaster]) != `{"m":2}` {
				t.Errorf("CrossSigningKeys = %q, %v", got, err)
			}

			updates, err := keys.AddSignatures(ctx, []storage.KeySignature{
				{TargetUserID: alice, TargetKeyID: "D1", SignerUserID: alice, SignerKeyID: "ed25519:ssk", Signature: "s1"},
				{TargetUserID: alice, TargetKeyID: "D1", SignerUserID: alice, SignerKeyID: "ed25519:ssk", Signature: "s2"},
				{TargetUserID: "@bob:warp.example", TargetKeyID: "mk", SignerUserID: alice, SignerKeyID: "ed25519:usk", Signature: "u1"},
			})
			if err != nil || len(updates) != 2 {
				t.Fatalf("AddSignatures = %+v, %v; want one update per signed key", updates, err)
			}
			sigs, err := keys.Signatures(ctx, alice)
			if err != nil || len(sigs) != 1 || sigs[0].Signature != "s2" {
				t.Errorf("Signatures = %+v, %v; want the later signature only", sigs, err)
			}
		})
	}
}

func TestKeyBackup(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			keys := migratedKeysDB(t, driver).Keys()

			if _, err := keys.Backup(ctx, alice, ""); !errors.Is(err, storage.ErrNoSuchBackup) {
				t.Fatalf("Backup with none = %v, want ErrNoSuchBackup", err)
			}
			v1, err := keys.CreateBackup(ctx, alice, "m.megolm_backup.v1.curve25519-aes-sha2", []byte(`{"public_key":"a"}`))
			if err != nil || v1 != "1" {
				t.Fatalf("CreateBackup = %q, %v", v1, err)
			}

			k := storage.BackupKey{RoomID: "!r:x", SessionID: "s1", FirstMessageIndex: 5, ForwardedCount: 1, SessionData: []byte(`{"d":1}`)}
			v, err := keys.PutBackupKeys(ctx, alice, v1, []storage.BackupKey{k})
			if err != nil || v.Count != 1 || v.ETag != "1" {
				t.Fatalf("PutBackupKeys = %+v, %v", v, err)
			}

			// A worse key does not replace the one there and does not
			// move the etag; a better one does.
			worse := k
			worse.FirstMessageIndex, worse.SessionData = 9, []byte(`{"d":"worse"}`)
			if v, _ = keys.PutBackupKeys(ctx, alice, v1, []storage.BackupKey{worse}); v.ETag != "1" {
				t.Errorf("a worse key moved the etag to %s", v.ETag)
			}
			verified := worse
			verified.IsVerified, verified.SessionData = true, []byte(`{"d":"verified"}`)
			if v, _ = keys.PutBackupKeys(ctx, alice, v1, []storage.BackupKey{verified}); v.ETag != "2" {
				t.Errorf("a verified key did not replace an unverified one: etag %s", v.ETag)
			}
			got, err := keys.BackupKeys(ctx, alice, v1, "!r:x", "s1")
			if err != nil || len(got) != 1 || string(got[0].SessionData) != `{"d":"verified"}` {
				t.Fatalf("BackupKeys = %+v, %v", got, err)
			}

			_, _ = keys.PutBackupKeys(ctx, alice, v1, []storage.BackupKey{{RoomID: "!r:x", SessionID: "s2", SessionData: []byte(`{}`)}})
			if all, _ := keys.BackupKeys(ctx, alice, v1, "", ""); len(all) != 2 {
				t.Errorf("BackupKeys(all) = %d keys, want 2", len(all))
			}
			v, err = keys.DeleteBackupKeys(ctx, alice, v1, "!r:x", "s2")
			if err != nil || v.Count != 1 {
				t.Errorf("DeleteBackupKeys = %+v, %v", v, err)
			}

			if err := keys.UpdateBackup(ctx, alice, v1, []byte(`{"public_key":"b"}`)); err != nil {
				t.Fatalf("UpdateBackup: %v", err)
			}
			current, _ := keys.Backup(ctx, alice, "")
			if current.Version != v1 || string(current.AuthData) != `{"public_key":"b"}` {
				t.Errorf("Backup = %+v", current)
			}

			// Deleting a version deletes its keys, and the next version
			// does not reuse its number.
			if err := keys.DeleteBackup(ctx, alice, v1); err != nil {
				t.Fatalf("DeleteBackup: %v", err)
			}
			if _, err := keys.Backup(ctx, alice, v1); !errors.Is(err, storage.ErrNoSuchBackup) {
				t.Errorf("Backup of a deleted version = %v", err)
			}
			if _, err := keys.BackupKeys(ctx, alice, v1, "", ""); !errors.Is(err, storage.ErrNoSuchBackup) {
				t.Errorf("BackupKeys of a deleted version = %v", err)
			}
			if v2, _ := keys.CreateBackup(ctx, alice, "alg", []byte(`{}`)); v2 != "2" {
				t.Errorf("the next version is %q, want 2", v2)
			}
			if _, err := keys.Backup(ctx, alice, "not-a-number"); !errors.Is(err, storage.ErrNoSuchBackup) {
				t.Errorf("Backup of a malformed version = %v", err)
			}
		})
	}
}

func TestUsersSharingRoomsWithNobody(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			db := migratedKeysDB(t, driver)
			users, err := db.Rooms().UsersSharingRooms(t.Context(), "@nobody:warp.example")
			if err != nil || len(users) != 0 {
				t.Errorf("UsersSharingRooms = %v, %v", users, err)
			}
			changes, err := db.Rooms().MemberChanges(t.Context(), []string{"!none:warp.example"}, 0, 0)
			if err != nil || len(changes) != 0 {
				t.Errorf("MemberChanges = %v, %v", changes, err)
			}
		})
	}
}