keyserver: cache remote users' device and cross-signing keys

POST /keys/query for a remote user who shares a room with a local user is
now answered from a cache, fetched whole with GET /user/devices/{userId}
on a miss. Incoming m.device_list_update EDUs are applied when their
prev_id matches the cached stream_id; anything else marks the copy stale
and it is fetched again in the background, and a stale copy is never
served. m.signing_key_update replaces the cached cross-signing keys. A
user's copy is dropped once no local user shares a room with them, by a
keyserver listener on the roomserver's emit stage.

Migration 0009 adds keyserver_remote_users, keyserver_remote_device_keys
and keyserver_remote_cross_signing_keys for sqlite and postgres.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

server/internal/setup/federation.go

@@ -9,6 +9,8 @@ import (
 	"sync"
 
 	"github.com/dariuszjaniak/warp/server/internal/federationapi"
+	"github.com/dariuszjaniak/warp/server/internal/keyserver"
+	"github.com/dariuszjaniak/warp/server/internal/matrix/event"
 	"github.com/dariuszjaniak/warp/server/internal/roomserver"
 	"github.com/dariuszjaniak/warp/server/internal/storage"
 )
@@ -39,6 +41,27 @@ func (f *fanout) EventPersisted(ctx context.Context, n roomserver.Notification)
 	}
 }
 
+// keyMemberships is the keyserver's listener on the emit stage: a
+// membership ending is when a remote user may stop sharing a room with
+// anybody here, and the keyserver drops its cached copy of their keys once
+// they have. MembershipLeft does not block, as the emit stage requires.
+type keyMemberships struct{ keys *keyserver.API }
+
+// EventPersisted implements roomserver.Notifier.
+func (l keyMemberships) EventPersisted(_ context.Context, n roomserver.Notification) {
+	if n.Event == nil || n.Event.Type() != event.TypeMember || !n.StateChanged {
+		return
+	}
+	userID, ok := n.Event.StateKey()
+	if !ok {
+		return
+	}
+	switch membership, _ := n.Event.Membership(); membership {
+	case event.MembershipLeave, event.MembershipBan:
+		l.keys.MembershipLeft(userID)
+	}
+}
+
 // federationStatus adapts the module to the admin API.
 type federationStatus struct {
 	module *federationapi.Module

server/internal/setup/phase7_test.go

@@ -8,6 +8,8 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"strings"
+	"sync/atomic"
 	"testing"
 
 	"github.com/dariuszjaniak/warp/server/internal/matrix/signing"
@@ -280,6 +282,19 @@ func TestKeysOverFederation(t *testing.T) {
 	a, b := nodes["hs-a.example"].srv, nodes["hs-b.example"].srv
 	const alice, bob = "@alice:hs-a.example", "@bob:hs-b.example"
 
+	// Count the key requests that reach A, which is how the test tells a
+	// query B answered from its cache from one it passed on.
+	var keyRequests atomic.Int64
+	inner := *nodes["hs-a.example"].handler.Load()
+	counting := http.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if strings.HasPrefix(r.URL.Path, "/_matrix/federation/v1/user/devices/") ||
+			r.URL.Path == "/_matrix/federation/v1/user/keys/query" {
+			keyRequests.Add(1)
+		}
+		inner.ServeHTTP(w, r)
+	}))
+	nodes["hs-a.example"].handler.Store(&counting)
+
 	aliceToken := registerThrough(t, a, "alice")
 	bobToken := registerThrough(t, b, "bob")
 	aliceDevice := deviceIDOf(t, a, aliceToken)
@@ -297,12 +312,43 @@ func TestKeysOverFederation(t *testing.T) {
 		return contains(got.DeviceLists.Changed, alice)
 	})
 
-	// Bob's query and claim are answered by Alice's server.
-	query := mustOK(t, post(t, b, "/_matrix/client/v3/keys/query", map[string]any{
-		"device_keys": map[string]any{alice: []string{}},
-	}, bobToken), "remote keys/query")
-	if devices, _ := query["device_keys"].(map[string]any)[alice].(map[string]any); devices[aliceDevice] == nil {
-		t.Fatalf("remote keys/query = %v", query)
+	// Bob's first query is answered by Alice's server, and cached: the
+	// second does not leave B.
+	queryAlice := func(what string) map[string]any {
+		t.Helper()
+		query := mustOK(t, post(t, b, "/_matrix/client/v3/keys/query", map[string]any{
+			"device_keys": map[string]any{alice: []string{}},
+		}, bobToken), what)
+		devices, _ := query["device_keys"].(map[string]any)[alice].(map[string]any)
+		return devices
+	}
+	if devices := queryAlice("remote keys/query"); devices[aliceDevice] == nil {
+		t.Fatalf("remote keys/query = %v", devices)
+	}
+	if n := keyRequests.Load(); n != 1 {
+		t.Errorf("the first query made %d key requests to A, want 1", n)
+	}
+	if devices := queryAlice("second remote keys/query"); devices[aliceDevice] == nil {
+		t.Fatalf("second remote keys/query = %v", devices)
+	}
+	if n := keyRequests.Load(); n != 1 {
+		t.Errorf("the second query reached A: %d key requests, want 1", n)
+	}
+
+	// A new device of Alice's reaches B as an m.device_list_update that
+	// follows on from the cached copy, which B applies: Bob sees the new
+	// device, and A is still not asked.
+	login := mustOK(t, post(t, a, "/_matrix/client/v3/login", map[string]any{
+		"type": "m.login.password", "identifier": map[string]string{"type": "m.id.user", "user": "alice"},
+		"password": "sUp3rs3kr1t", "device_id": "PHONE",
+	}, ""), "login")
+	uploadDeviceKeys(t, a, login["access_token"].(string), alice, "PHONE", 1)
+	eventually(t, "alice's new device in bob's query", func() bool {
+		devices := queryAlice("keys/query after the update")
+		return devices["PHONE"] != nil && devices[aliceDevice] != nil
+	})
+	if n := keyRequests.Load(); n != 1 {
+		t.Errorf("the update was not applied from the EDU: %d key requests to A, want 1", n)
 	}
 	claim := mustOK(t, post(t, b, "/_matrix/client/v3/keys/claim", map[string]any{
 		"one_time_keys": map[string]any{alice: map[string]string{aliceDevice: "signed_curve25519"}},
@@ -320,9 +366,15 @@ func TestKeysOverFederation(t *testing.T) {
 	if err := a.Settings().SetFederationEnabled(t.Context(), false); err != nil {
 		t.Fatal(err)
 	}
-	query = mustOK(t, post(t, b, "/_matrix/client/v3/keys/query", map[string]any{
-		"device_keys": map[string]any{alice: []string{}}, "timeout": 5000,
-	}, bobToken), "keys/query with A's federation off")
+	// The cached copy is still B's to serve.
+	if devices := queryAlice("keys/query with A's federation off"); devices["PHONE"] == nil {
+		t.Errorf("with A's federation off, B no longer serves its cached copy: %v", devices)
+	}
+	// A user B has nothing cached for is asked of A, and A is reported as
+	// a failure instead of hanging.
+	query := mustOK(t, post(t, b, "/_matrix/client/v3/keys/query", map[string]any{
+		"device_keys": map[string]any{"@stranger:hs-a.example": []string{}}, "timeout": 5000,
+	}, bobToken), "keys/query for a stranger with A's federation off")
 	if failures, _ := query["failures"].(map[string]any); failures["hs-a.example"] == nil {
 		t.Errorf("failures = %v, want hs-a.example", query["failures"])
 	}

server/internal/setup/server.go

@@ -295,6 +295,7 @@ func New(ctx context.Context, cfg *config.Config, opts Options) (*Server, error)
 			return nil, err
 		}
 		s.users.SetDeviceObserver(s.keys)
+		emit.add(keyMemberships{keys: s.keys})
 
 		s.sync, err = syncapi.New(syncapi.Options{
 			ServerName:      cfg.Global.ServerName,
@@ -576,6 +577,11 @@ func (s *Server) Close() error {
 	if s.fedModule != nil {
 		s.fedModule.Stop()
 	}
+	if s.keys != nil {
+		// Background fetches of remote keys write to the database, which
+		// is about to close.
+		s.keys.Close()
+	}
 	if s.sync != nil {
 		// The typing tracker holds a timer per typing user; a server
 		// shutting down should not wait for them.