From b61d441b5ddb69331962d334f2984bc295107d2e Mon Sep 17 00:00:00 2001
From: Kegan Dougal <7190048+kegsay@users.noreply.github.com>
Date: Fri, 23 Feb 2024 14:05:16 +0000
Subject: [PATCH] Add TestRoomKeyIsCycledOnMemberLeaving

---
 TEST_HITLIST.md         |  6 ++--
 tests/room_keys_test.go | 67 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/TEST_HITLIST.md b/TEST_HITLIST.md
index 9d31ece..c3df3d1 100644
--- a/TEST_HITLIST.md
+++ b/TEST_HITLIST.md
@@ -93,15 +93,15 @@ This refers to cases where the client has some state and wishes to synchronise i
 - [ ] If a server is terminated mid-way through sending a to-device message over federation, it retries on startup.
 
 ### Room Keys
-- [ ] The room key is cycled when a user leaves a room.
-- [ ] The room key is cycled when one of a user's devices logs out.
+- [x] The room key is cycled when a user leaves a room.
+- [x] The room key is cycled when one of a user's devices logs out.
 - [ ] The room key is cycled when one of a user's devices is blacklisted.
 - [ ] The room key is cycled when history visibility changes to something more restrictive TODO: define precisely.
 - [ ] The room key is cycled when the encryption algorithm changes.
 - [ ] The room key is cycled when `rotation_period_msgs` is met (default: 100).
 - [ ] The room key is cycled when `rotation_period_ms` is exceeded (default: 1 week).
 - [ ] The room key is not cycled when one of a user's devices logs in.
-- [ ] The room key is not cycled when the client restarts.
+- [x] The room key is not cycled when the client restarts.
 - [ ] The room key is not cycled when users change their display name or avatar.
 
 ### Adversarial Attacks
diff --git a/tests/room_keys_test.go b/tests/room_keys_test.go
index a313923..b9345c4 100644
--- a/tests/room_keys_test.go
+++ b/tests/room_keys_test.go
@@ -99,6 +99,73 @@ func TestRoomKeyIsCycledOnDeviceLogout(t *testing.T) {
 	})
 }
 
+func TestRoomKeyIsCycledOnMemberLeaving(t *testing.T) {
+	ClientTypeMatrix(t, func(t *testing.T, clientTypeA, clientTypeB api.ClientType) {
+		tc := CreateTestContext(t, clientTypeA, clientTypeB, clientTypeB)
+		roomID := tc.CreateNewEncryptedRoom(t, tc.Alice, "trusted_private_chat", []string{tc.Bob.UserID, tc.Charlie.UserID})
+		tc.Bob.MustJoinRoom(t, roomID, []string{clientTypeA.HS})
+		tc.Charlie.MustJoinRoom(t, roomID, []string{clientTypeA.HS})
+
+		// Alice, Bob and Charlie are in a room.
+		alice := tc.MustLoginClient(t, tc.Alice, clientTypeA)
+		defer alice.Close(t)
+		bob := tc.MustLoginClient(t, tc.Bob, clientTypeB)
+		defer bob.Close(t)
+		charlie := tc.MustLoginClient(t, tc.Charlie, clientTypeB)
+		defer charlie.Close(t)
+		aliceStopSyncing := alice.MustStartSyncing(t)
+		defer aliceStopSyncing()
+		bobStopSyncing := bob.MustStartSyncing(t)
+		defer bobStopSyncing()
+		charlieStopSyncing := charlie.MustStartSyncing(t)
+		defer charlieStopSyncing()
+
+		// check the room works
+		wantMsgBody := "Test Message"
+		waiter := bob.WaitUntilEventInRoom(t, roomID, api.CheckEventHasBody(wantMsgBody))
+		waiter2 := charlie.WaitUntilEventInRoom(t, roomID, api.CheckEventHasBody(wantMsgBody))
+		alice.SendMessage(t, roomID, wantMsgBody)
+		waiter.Wait(t, 5*time.Second)
+		waiter2.Wait(t, 5*time.Second)
+
+		// we're going to sniff calls to /sendToDevice to ensure we see the new room key being sent.
+		ch := make(chan deploy.CallbackData, 10)
+		callbackURL, close := sniffToDeviceEvent(t, ch)
+		defer close()
+
+		// we don't know when the new room key will be sent, it could be sent as soon as the device list update
+		// is sent, or it could be delayed until message send. We want to handle both cases so we start sniffing
+		// traffic now.
+		tc.Deployment.WithMITMOptions(t, map[string]interface{}{
+			"callback": map[string]interface{}{
+				"callback_url": callbackURL,
+				"filter":       "~u .*\\/sendToDevice.*",
+			},
+		}, func() {
+			// now Charlie is going to leave the room, causing her user ID to appear in device_lists.left
+			// which should trigger a new room key to be sent (on message send)
+			tc.Charlie.MustDo(t, "POST", []string{"_matrix", "client", "v3", "logout"}, client.WithJSONBody(t, map[string]any{}))
+
+			// we don't know how long it will take for the device list update to be processed, so wait 1s
+			time.Sleep(time.Second)
+
+			// now send another message from Alice, who should negotiate a new room key
+			wantMsgBody = "Another Test Message"
+			waiter = bob.WaitUntilEventInRoom(t, roomID, api.CheckEventHasBody(wantMsgBody))
+			alice.SendMessage(t, roomID, wantMsgBody)
+			waiter.Wait(t, 5*time.Second)
+		})
+
+		// we should have seen a /sendToDevice call by now. If we didn't, this implies we didn't cycle
+		// the room key.
+		select {
+		case <-ch:
+		default:
+			ct.Fatalf(t, "did not see /sendToDevice when logging out and sending a new message")
+		}
+	})
+}
+
 // Test that the m.room_key is NOT cycled when the client is restarted, but there is no change in devices
 // in the room. This is important to ensure that we don't cycle m.room_keys too frequently, which increases
 // the chances of seeing undecryptable events.