Warp/server/internal/storage/remotekeys_test.go

package storage_test

import (
	"slices"
	"testing"

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

const carol = "@carol:remote.example"

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

			if _, ok, err := keys.RemoteKeys(ctx, carol); ok || err != nil {
				t.Fatalf("RemoteKeys before anything was cached = %v, %v", ok, err)
			}

			snapshot := storage.RemoteKeys{
				UserID: carol, StreamID: 5,
				Devices: []storage.RemoteDevice{
					{DeviceID: "C2", JSON: []byte(`{"d":2}`)},
					{DeviceID: "C1", DisplayName: "phone", JSON: []byte(`{"d":1}`)},
				},
				CrossSigning: map[string][]byte{
					storage.CrossSigningMaster:      []byte(`{"m":1}`),
					storage.CrossSigningSelfSigning: []byte(`{"s":1}`),
				},
			}
			if stored, err := keys.SetRemoteKeys(ctx, snapshot); !stored || err != nil {
				t.Fatalf("SetRemoteKeys = %v, %v", stored, err)
			}
			got, ok, err := keys.RemoteKeys(ctx, carol)
			if err != nil || !ok {
				t.Fatalf("RemoteKeys = %v, %v", ok, err)
			}
			if got.StreamID != 5 || got.Stale || len(got.Devices) != 2 ||
				got.Devices[0].DeviceID != "C1" || got.Devices[0].DisplayName != "phone" || string(got.Devices[1].JSON) != `{"d":2}` ||
				string(got.CrossSigning[storage.CrossSigningMaster]) != `{"m":1}` || len(got.CrossSigning) != 2 {
				t.Errorf("RemoteKeys = %+v", got)
			}

			// Caching moves nothing in the device list stream: that is the
			// keyserver's call, per update.
			if pos, _ := keys.DeviceListPosition(ctx, carol); pos != 0 {
				t.Errorf("caching moved the device list to %d", pos)
			}

			apply := func(u storage.RemoteDeviceUpdate) storage.RemoteUpdateResult {
				t.Helper()
				u.UserID = carol
				result, err := keys.UpdateRemoteDevice(ctx, u)
				if err != nil {
					t.Fatalf("UpdateRemoteDevice(%+v): %v", u, err)
				}
				return result
			}

			// Following on, a repeat, and a deletion.
			if r := apply(storage.RemoteDeviceUpdate{DeviceID: "C3", JSON: []byte(`{"d":3}`), StreamID: 6, PrevIDs: []int64{5}}); r != storage.RemoteUpdateApplied {
				t.Errorf("an update that follows on = %v", r)
			}
			if r := apply(storage.RemoteDeviceUpdate{DeviceID: "C3", JSON: []byte(`{"d":3}`), StreamID: 6, PrevIDs: []int64{5}}); r != storage.RemoteUpdateDuplicate {
				t.Errorf("the same update again = %v", r)
			}
			if r := apply(storage.RemoteDeviceUpdate{DeviceID: "C2", Deleted: true, StreamID: 7, PrevIDs: []int64{4, 6}}); r != storage.RemoteUpdateApplied {
				t.Errorf("a deletion naming the position held among others = %v", r)
			}
			got, _, _ = keys.RemoteKeys(ctx, carol)
			var ids []string
			for _, d := range got.Devices {
				ids = append(ids, d.DeviceID)
			}
			if got.StreamID != 7 || !slices.Equal(ids, []string{"C1", "C3"}) {
				t.Errorf("after the updates: position %d, devices %v", got.StreamID, ids)
			}

			// A gap marks the copy stale at the newest position heard of.
			if r := apply(storage.RemoteDeviceUpdate{DeviceID: "C4", JSON: []byte(`{}`), StreamID: 12, PrevIDs: []int64{11}}); r != storage.RemoteUpdateStale {
				t.Errorf("an update after a gap = %v", r)
			}
			got, _, _ = keys.RemoteKeys(ctx, carol)
			if !got.Stale || got.StreamID != 12 || len(got.Devices) != 2 {
				t.Errorf("after the gap: %+v", got)
			}
			if r := apply(storage.RemoteDeviceUpdate{DeviceID: "C4", JSON: []byte(`{}`), StreamID: 13, PrevIDs: []int64{12}}); r != storage.RemoteUpdateStale {
				t.Errorf("an update to a stale copy = %v", r)
			}
			if stored, _ := keys.SetRemoteCrossSigningKeys(ctx, carol, map[string][]byte{storage.CrossSigningMaster: []byte(`{}`)}); stored {
				t.Error("cross-signing keys were written into a stale copy")
			}

			// A snapshot older than the position held is not stored; one at
			// or past it replaces the copy whole.
			if stored, _ := keys.SetRemoteKeys(ctx, storage.RemoteKeys{UserID: carol, StreamID: 10}); stored {
				t.Error("a snapshot older than the copy was stored")
			}
			if stored, _ := keys.SetRemoteKeys(ctx, storage.RemoteKeys{
				UserID: carol, StreamID: 13, Devices: []storage.RemoteDevice{{DeviceID: "C4", JSON: []byte(`{"d":4}`)}},
			}); !stored {
				t.Error("a current snapshot was not stored")
			}
			got, _, _ = keys.RemoteKeys(ctx, carol)
			if got.Stale || got.StreamID != 13 || len(got.Devices) != 1 || len(got.CrossSigning) != 0 {
				t.Errorf("after the resync: %+v", got)
			}

			// Cross-signing keys are replaced as a set.
			if stored, err := keys.SetRemoteCrossSigningKeys(ctx, carol, map[string][]byte{
				storage.CrossSigningMaster: []byte(`{"m":2}`), storage.CrossSigningSelfSigning: nil,
			}); !stored || err != nil {
				t.Fatalf("SetRemoteCrossSigningKeys = %v, %v", stored, err)
			}
			got, _, _ = keys.RemoteKeys(ctx, carol)
			if string(got.CrossSigning[storage.CrossSigningMaster]) != `{"m":2}` || len(got.CrossSigning) != 1 {
				t.Errorf("cross-signing keys = %v", got.CrossSigning)
			}

			// An update about a user with nothing cached leaves a stale
			// marker, and a snapshot from before the update loses to it.
			const dave = "@dave:remote.example"
			if r, _ := keys.UpdateRemoteDevice(ctx, storage.RemoteDeviceUpdate{UserID: dave, DeviceID: "V1", StreamID: 3}); r != storage.RemoteUpdateNotCached {
				t.Errorf("an update about an uncached user = %v", r)
			}
			marker, ok, _ := keys.RemoteKeys(ctx, dave)
			if !ok || !marker.Stale || marker.StreamID != 3 {
				t.Errorf("the marker = %+v, %v", marker, ok)
			}
			if stored, _ := keys.SetRemoteKeys(ctx, storage.RemoteKeys{UserID: dave, StreamID: 2}); stored {
				t.Error("a snapshot from before the update was stored")
			}
			if stored, _ := keys.SetRemoteCrossSigningKeys(ctx, "@nobody:remote.example", map[string][]byte{}); stored {
				t.Error("cross-signing keys were cached for a user with no copy")
			}

			users, err := keys.RemoteKeyUsers(ctx)
			if err != nil || !slices.Equal(users, []string{carol, dave}) {
				t.Errorf("RemoteKeyUsers = %v, %v", users, err)
			}
			if err := keys.DeleteRemoteKeys(ctx, []string{carol, dave}); err != nil {
				t.Fatalf("DeleteRemoteKeys: %v", err)
			}
			if _, ok, _ := keys.RemoteKeys(ctx, carol); ok {
				t.Error("a dropped copy is still there")
			}
			if users, _ := keys.RemoteKeyUsers(ctx); len(users) != 0 {
				t.Errorf("RemoteKeyUsers after dropping everybody = %v", users)
			}
		})
	}
}

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

			// A user whose server had never moved their list is cached at
			// position 0, and the first update, with no prev_id, follows on.
			if _, err := keys.SetRemoteKeys(ctx, storage.RemoteKeys{UserID: carol}); err != nil {
				t.Fatal(err)
			}
			r, err := keys.UpdateRemoteDevice(ctx, storage.RemoteDeviceUpdate{UserID: carol, DeviceID: "C1", JSON: []byte(`{}`), StreamID: 1})
			if err != nil || r != storage.RemoteUpdateApplied {
				t.Errorf("the first update = %v, %v", r, err)
			}
			// A later one with no prev_id does not.
			r, _ = keys.UpdateRemoteDevice(ctx, storage.RemoteDeviceUpdate{UserID: carol, DeviceID: "C1", JSON: []byte(`{}`), StreamID: 2})
			if r != storage.RemoteUpdateStale {
				t.Errorf("a later update with no prev_id = %v", r)
			}
			// Nor does one that says a device changed without its keys.
			_, _ = keys.SetRemoteKeys(ctx, storage.RemoteKeys{UserID: carol, StreamID: 2})
			r, _ = keys.UpdateRemoteDevice(ctx, storage.RemoteDeviceUpdate{UserID: carol, DeviceID: "C1", StreamID: 3, PrevIDs: []int64{2}})
			if r != storage.RemoteUpdateStale {
				t.Errorf("an update with no keys = %v", r)
			}
		})
	}
}