Merge pull request #69 from matrix-org/kegan/keys-query-fails-after-room-key

Add TestToDeviceMessagesArentLostWhenKeysQueryFails

Author: kegsay

internal/api/client.go

@@ -65,6 +65,8 @@ type Client interface {
 	Logf(t ct.TestLike, format string, args ...interface{})
 	// The user for this client
 	UserID() string
+	// The current access token for this client
+	CurrentAccessToken(t ct.TestLike) string
 	Type() ClientTypeLang
 	Opts() ClientCreationOpts
 }
@@ -78,6 +80,13 @@ type LoggedClient struct {
 	Client
 }
 
+func (c *LoggedClient) CurrentAccessToken(t ct.TestLike) string {
+	t.Helper()
+	token := c.Client.CurrentAccessToken(t)
+	c.Logf(t, "%s CurrentAccessToken => %s", c.logPrefix(), token)
+	return token
+}
+
 func (c *LoggedClient) Login(t ct.TestLike, opts ClientCreationOpts) error {
 	t.Helper()
 	c.Logf(t, "%s Login %+v", c.logPrefix(), opts)

internal/api/js/js.go

@@ -251,6 +251,12 @@ func (c *JSClient) DeletePersistentStorage(t ct.TestLike) {
 	`, indexedDBName, indexedDBCryptoName))
 }
 
+func (c *JSClient) CurrentAccessToken(t ct.TestLike) string {
+	token := chrome.MustRunAsyncFn[string](t, c.browser.Ctx, `
+		return window.__client.getAccessToken();`)
+	return *token
+}
+
 func (c *JSClient) GetNotification(t ct.TestLike, roomID, eventID string) (*api.Notification, error) {
 	return nil, fmt.Errorf("not implemented yet") // TODO
 }

internal/api/rust/rust.go

@@ -188,6 +188,14 @@ func (c *RustClient) Login(t ct.TestLike, opts api.ClientCreationOpts) error {
 	return nil
 }
 
+func (c *RustClient) CurrentAccessToken(t ct.TestLike) string {
+	s, err := c.FFIClient.Session()
+	if err != nil {
+		ct.Fatalf(t, "error retrieving session: %s", err)
+	}
+	return s.AccessToken
+}
+
 func (c *RustClient) DeletePersistentStorage(t ct.TestLike) {
 	t.Helper()
 	if c.persistentStoragePath != "" {

internal/deploy/rpc_client.go

@@ -181,6 +181,15 @@ func (c *RPCClient) GetNotification(t ct.TestLike, roomID, eventID string) (*api
 	return &notification, err
 }
 
+func (c *RPCClient) CurrentAccessToken(t ct.TestLike) string {
+	var token string
+	err := c.client.Call("RPCServer.CurrentAccessToken", t.Name(), &token)
+	if err != nil {
+		ct.Fatalf(t, "RPCServer.CurrentAccessToken: %s", err)
+	}
+	return token
+}
+
 // Remove any persistent storage, if it was enabled.
 func (c *RPCClient) DeletePersistentStorage(t ct.TestLike) {
 	var void int

internal/deploy/rpc_server.go

@@ -100,6 +100,12 @@ func (s *RPCServer) DeletePersistentStorage(testName string, void *int) error {
 	return nil
 }
 
+func (s *RPCServer) CurrentAccessToken(testName string, token *string) error {
+	defer s.keepAlive()
+	*token = s.activeClient.CurrentAccessToken(&api.MockT{TestName: testName})
+	return nil
+}
+
 func (s *RPCServer) Login(opts api.ClientCreationOpts, void *int) error {
 	defer s.keepAlive()
 	return s.activeClient.Login(&api.MockT{}, opts)

tests/to_device_test.go

@@ -310,3 +310,75 @@ func TestToDeviceMessagesAreBatched(t *testing.T) {
 
 	})
 }
+
+// Regression test for https://github.com/element-hq/element-web/issues/24682
+//
+// When a to-device msg is received, the SDK may need to check that the device belongs
+// to the user in question. To do this, it needs an up-to-date device list. To get this,
+// it does a /keys/query request. If this request fails, the entire processing of the
+// to-device msg could fail, dropping the msg and the room key it contains.
+//
+// This test reproduces this by having an existing E2EE room between Alice and Bob, then:
+//   - Block /keys/query requests.
+//   - Alice logs in on a new device.
+//   - Alice sends a message on the new device.
+//   - Bob should get that message but may refuse to decrypt it as it cannot verify that the sender_key
+//     belongs to Alice.
+//   - Unblock /keys/query requests.
+//   - Bob should eventually retry and be able to decrypt the event.
+func TestToDeviceMessagesArentLostWhenKeysQueryFails(t *testing.T) {
+	ForEachClientType(t, func(t *testing.T, clientType api.ClientType) {
+		tc := CreateTestContext(t, clientType, clientType)
+		// get a normal E2EE room set up
+		roomID := tc.CreateNewEncryptedRoom(t, tc.Alice, EncRoomOptions.Invite([]string{tc.Bob.UserID}))
+		tc.Bob.MustJoinRoom(t, roomID, []string{clientType.HS})
+		tc.WithAliceAndBobSyncing(t, func(alice, bob api.Client) {
+			msg := "hello world"
+			msg2 := "new device message from alice"
+			alice.SendMessage(t, roomID, msg)
+			bob.WaitUntilEventInRoom(t, roomID, api.CheckEventHasBody(msg)).Waitf(t, 5*time.Second, "bob failed to see message from alice")
+			// Block /keys/query requests
+			waiter := helpers.NewWaiter()
+			callbackURL, closeCallbackServer := deploy.NewCallbackServer(t, tc.Deployment, func(cd deploy.CallbackData) {
+				t.Logf("%+v", cd)
+				waiter.Finish()
+			})
+			defer closeCallbackServer()
+			var eventID string
+			bobAccessToken := bob.CurrentAccessToken(t)
+			t.Logf("Bob's token => %s", bobAccessToken)
+			tc.Deployment.WithMITMOptions(t, map[string]interface{}{
+				"statuscode": map[string]interface{}{
+					"return_status": http.StatusGatewayTimeout,
+					"block_request": true,
+					"count":         3,
+					"filter":        "~u .*/keys/query.* ~hq " + bobAccessToken,
+				},
+				"callback": map[string]interface{}{
+					"callback_url": callbackURL,
+					"filter":       "~u .*/keys/query.*",
+				},
+			}, func() {
+				// Alice logs in on a new device.
+				csapiAlice2 := tc.MustRegisterNewDevice(t, tc.Alice, clientType.HS, "OTHER_DEVICE")
+				alice2 := tc.MustLoginClient(t, csapiAlice2, clientType)
+				defer alice2.Close(t)
+				alice2StopSyncing := alice2.MustStartSyncing(t)
+				defer alice2StopSyncing()
+				// we don't know how long it will take for the device list update to be processed, so wait 1s
+				time.Sleep(time.Second)
+
+				// Alice sends a message on the new device.
+				eventID = alice2.SendMessage(t, roomID, msg2)
+
+				waiter.Waitf(t, 3*time.Second, "did not see /keys/query")
+				time.Sleep(3 * time.Second) // let Bob retry /keys/query
+			})
+			// now we aren't blocking /keys/query anymore.
+			// Bob should be able to decrypt this message.
+			ev := bob.MustGetEvent(t, roomID, eventID)
+			must.Equal(t, ev.Text, msg2, "bob failed to decrypt "+eventID)
+		})
+
+	})
+}