From 43d58fbad28e6baaf7265daed92c82bfe45e3fed Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Tue, 11 Apr 2023 12:55:23 +0300
Subject: [PATCH] Redact megolm sessions when device is deleted

---
 bridge/crypto.go     |  1 +
 crypto/devicelist.go | 21 +++++++++++++++++++++
 crypto/machine.go    |  1 +
 3 files changed, 23 insertions(+)

diff --git a/bridge/crypto.go b/bridge/crypto.go
index 8d2171a4..ab26c384 100644
--- a/bridge/crypto.go
+++ b/bridge/crypto.go
@@ -99,6 +99,7 @@ func (helper *CryptoHelper) Init() error {
 	helper.mach.RatchetKeysOnDecrypt = encryptionConfig.DeleteKeys.RatchetOnDecrypt
 	helper.mach.DeleteFullyUsedKeysOnDecrypt = encryptionConfig.DeleteKeys.DeleteFullyUsedOnDecrypt
 	helper.mach.DeletePreviousKeysOnReceive = encryptionConfig.DeleteKeys.DeletePrevOnNewSession
+	helper.mach.DeleteKeysOnDeviceDelete = encryptionConfig.DeleteKeys.DeleteOnDeviceDelete
 
 	helper.client.Syncer = &cryptoSyncer{helper.mach}
 	helper.client.Store = helper.store
diff --git a/crypto/devicelist.go b/crypto/devicelist.go
index fc204d1f..5bc43c1b 100644
--- a/crypto/devicelist.go
+++ b/crypto/devicelist.go
@@ -128,6 +128,7 @@ func (mach *OlmMachine) fetchKeys(ctx context.Context, users []id.UserID, sinceT
 			log.Warn().Err(err).Msg("Failed to get existing devices for user")
 			existingDevices = make(map[id.DeviceID]*id.Device)
 		}
+
 		log.Debug().
 			Int("new_device_count", len(devices)).
 			Int("old_device_count", len(existingDevices)).
@@ -158,6 +159,26 @@ func (mach *OlmMachine) fetchKeys(ctx context.Context, users []id.UserID, sinceT
 
 		changed = changed || len(newDevices) != len(existingDevices)
 		if changed {
+			if mach.DeleteKeysOnDeviceDelete {
+				for deviceID := range newDevices {
+					delete(existingDevices, deviceID)
+				}
+				for _, device := range existingDevices {
+					log := log.With().
+						Str("device_id", device.DeviceID.String()).
+						Str("identity_key", device.IdentityKey.String()).
+						Str("signing_key", device.SigningKey.String()).
+						Logger()
+					sessionIDs, err := mach.CryptoStore.RedactGroupSessions("", device.IdentityKey, "device removed")
+					if err != nil {
+						log.Err(err).Msg("Failed to redact megolm sessions from deleted device")
+					} else {
+						log.Info().
+							Strs("session_ids", stringifyArray(sessionIDs)).
+							Msg("Redacted megolm sessions from deleted device")
+					}
+				}
+			}
 			mach.OnDevicesChanged(userID)
 		}
 	}
diff --git a/crypto/machine.go b/crypto/machine.go
index e1682084..f4970de6 100644
--- a/crypto/machine.go
+++ b/crypto/machine.go
@@ -71,6 +71,7 @@ type OlmMachine struct {
 	DeletePreviousKeysOnReceive  bool
 	RatchetKeysOnDecrypt         bool
 	DeleteFullyUsedKeysOnDecrypt bool
+	DeleteKeysOnDeviceDelete     bool
 }
 
 // StateStore is used by OlmMachine to get room state information that's needed for encryption.