From a9f76c2f63d6aab0fa61d1b115e3a65fbf153d8b Mon Sep 17 00:00:00 2001
From: Kegan Dougal <7190048+kegsay@users.noreply.github.com>
Date: Thu, 16 May 2024 16:21:56 +0100
Subject: [PATCH] Add regression test to ensure we batch /sendToDevice

Fixes https://github.com/matrix-org/complement-crypto/issues/34
Regression test for https://github.com/element-hq/element-web/issues/24680
---
 tests/main_test.go      |  7 -----
 tests/to_device_test.go | 69 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 7 deletions(-)

diff --git a/tests/main_test.go b/tests/main_test.go
index 60b5acb..ff1fdf8 100644
--- a/tests/main_test.go
+++ b/tests/main_test.go
@@ -109,13 +109,6 @@ func MustCreateClient(t *testing.T, clientType api.ClientType, cfg api.ClientCre
 	return c
 }
 
-// WithDoLogin is an option which can be provided to MustCreateClient which will automatically login, else fail the test.
-func WithDoLogin(t *testing.T) func(api.Client, *api.ClientCreationOpts) {
-	return func(c api.Client, opts *api.ClientCreationOpts) {
-		must.NotError(t, "failed to login", c.Login(t, *opts))
-	}
-}
-
 // WithPersistentStorage is an option which can be provided to MustCreateClient which will configure clients to use persistent storage,
 // e.g IndexedDB or sqlite3 files.
 func WithPersistentStorage() func(*api.ClientCreationOpts) {
diff --git a/tests/to_device_test.go b/tests/to_device_test.go
index 3497f9d..457a51a 100644
--- a/tests/to_device_test.go
+++ b/tests/to_device_test.go
@@ -241,3 +241,72 @@ func testUnprocessedToDeviceMessagesArentLostOnRestartJS(t *testing.T, tc *TestC
 		must.Equal(t, ev.Text, "Kick to make a new room key!", "event text mismatch")
 	})
 }
+
+// Regression test for https://github.com/element-hq/element-web/issues/24680
+//
+// It's important that room keys are sent out ASAP, else the encrypted event may arrive
+// before the keys, causing a temporary unable-to-decrypt error. Clients SHOULD be batching
+// to-device messages, but old implementations batched too low (20 messages per request).
+// This test asserts we batch at least 100 per request.
+//
+// It does this by creating an E2EE room with 100 E2EE users, and forces a key rotation
+// by sending a message with rotation_period_msgs=1. It does not ensure that the room key
+// is correctly sent to all 100 users as that would entail having 100 users running at
+// the same time (think 100 browsers = expensive). Instead, we sequentially spin up 100
+// clients and then close them before doing the test, and assert we send 100 events.
+//
+// In the future, it may be difficult to run this test for 1 user with 100 devices due to
+// HS limits on the number of devices and forced cross-signing.
+func TestToDeviceMessagesAreBatched(t *testing.T) {
+	ForEachClientType(t, func(t *testing.T, clientType api.ClientType) {
+		tc := CreateTestContext(t, clientType)
+		roomID := tc.CreateNewEncryptedRoom(t, tc.Alice, EncRoomOptions.RotationPeriodMsgs(1), EncRoomOptions.PresetPublicChat())
+		// create 100 users
+		for i := 0; i < 100; i++ {
+			cli := tc.Deployment.Register(t, clientType.HS, helpers.RegistrationOpts{
+				LocalpartSuffix: fmt.Sprintf("bob-%d", i),
+				Password:        "complement-crypto-password",
+			})
+			cli.MustJoinRoom(t, roomID, []string{clientType.HS})
+			// this blocks until it has uploaded OTKs/device keys
+			clientUnderTest := tc.MustLoginClient(t, cli, tc.AliceClientType)
+			clientUnderTest.Close(t)
+		}
+		waiter := helpers.NewWaiter()
+		tc.WithAliceSyncing(t, func(alice api.Client) {
+			// intercept /sendToDevice and check we are sending 100 messages per request
+			tc.Deployment.WithSniffedEndpoint(t, "/sendToDevice", func(cd deploy.CallbackData) {
+				if cd.Method != "PUT" {
+					return
+				}
+				// format is:
+				/*
+					{
+					  "messages": {
+					    "@alice:example.com": {
+					      "TLLBEANAAG": {
+					        "example_content_key": "value"
+					      }
+					    }
+					  }
+					}
+				*/
+				usersMap := gjson.GetBytes(cd.RequestBody, "messages")
+				if !usersMap.Exists() {
+					t.Logf("intercepted PUT /sendToDevice but no messages existed")
+					return
+				}
+				if len(usersMap.Map()) != 100 {
+					t.Errorf("PUT /sendToDevice did not batch messages, got %d want 100", len(usersMap.Map()))
+					t.Logf(usersMap.Raw)
+				}
+				waiter.Finish()
+			}, func() {
+				alice.SendMessage(t, roomID, "this should cause to-device msgs to be sent")
+				time.Sleep(time.Second)
+				waiter.Waitf(t, 5*time.Second, "did not see /sendToDevice")
+			})
+		})
+
+	})
+}