From 6ac759c8ff4c1041b6f4edbdc323cc1e12324deb Mon Sep 17 00:00:00 2001
From: Toni Spets <toni.spets@beeper.com>
Date: Wed, 17 Jan 2024 09:26:13 +0200
Subject: [PATCH] Only skip fetching keys during Megolm decryption if disabled

Blanket disabling caused a lot of side effects which were hard to
deal with without major refactoring.

This should probably be an argument to DecryptMegolm instead of a
flag.
---
 crypto/decryptmegolm.go |  6 +++++-
 crypto/machine.go       | 14 +++++---------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/crypto/decryptmegolm.go b/crypto/decryptmegolm.go
index 540f99ca..abe01871 100644
--- a/crypto/decryptmegolm.go
+++ b/crypto/decryptmegolm.go
@@ -72,7 +72,11 @@ func (mach *OlmMachine) DecryptMegolmEvent(ctx context.Context, evt *event.Event
 	if sess.SigningKey == ownSigningKey && sess.SenderKey == ownIdentityKey && len(sess.ForwardingChains) == 0 {
 		trustLevel = id.TrustStateVerified
 	} else {
-		device, err = mach.GetOrFetchDeviceByKey(ctx, evt.Sender, sess.SenderKey)
+		if mach.DisableDecryptKeyFetching {
+			device, err = mach.CryptoStore.FindDeviceByKey(ctx, evt.Sender, sess.SenderKey)
+		} else {
+			device, err = mach.GetOrFetchDeviceByKey(ctx, evt.Sender, sess.SenderKey)
+		}
 		if err != nil {
 			// We don't want to throw these errors as the message can still be decrypted.
 			log.Debug().Err(err).Msg("Failed to get device to verify session")
diff --git a/crypto/machine.go b/crypto/machine.go
index b7c41ab0..77d99a8f 100644
--- a/crypto/machine.go
+++ b/crypto/machine.go
@@ -33,8 +33,8 @@ type OlmMachine struct {
 
 	PlaintextMentions bool
 
-	// Never ask the server for keys automatically as a side effect.
-	DisableKeyFetching bool
+	// Never ask the server for keys automatically as a side effect during Megolm decryption.
+	DisableDecryptKeyFetching bool
 
 	SendKeysMinTrust  id.TrustState
 	ShareKeysMinTrust id.TrustState
@@ -227,11 +227,7 @@ func (mach *OlmMachine) HandleDeviceLists(ctx context.Context, dl *mautrix.Devic
 			Str("trace_id", traceID).
 			Interface("changes", dl.Changed).
 			Msg("Device list changes in /sync")
-		if mach.DisableKeyFetching {
-			mach.CryptoStore.MarkTrackedUsersOutdated(ctx, dl.Changed)
-		} else {
-			mach.FetchKeys(ctx, dl.Changed, false)
-		}
+		mach.FetchKeys(ctx, dl.Changed, false)
 		mach.Log.Debug().Str("trace_id", traceID).Msg("Finished handling device list changes")
 	}
 }
@@ -420,7 +416,7 @@ func (mach *OlmMachine) GetOrFetchDevice(ctx context.Context, userID id.UserID,
 	device, err := mach.CryptoStore.GetDevice(ctx, userID, deviceID)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get sender device from store: %w", err)
-	} else if device != nil || mach.DisableKeyFetching {
+	} else if device != nil {
 		return device, nil
 	}
 	if usersToDevices, err := mach.FetchKeys(ctx, []id.UserID{userID}, true); err != nil {
@@ -439,7 +435,7 @@ func (mach *OlmMachine) GetOrFetchDevice(ctx context.Context, userID id.UserID,
 // the given identity key.
 func (mach *OlmMachine) GetOrFetchDeviceByKey(ctx context.Context, userID id.UserID, identityKey id.IdentityKey) (*id.Device, error) {
 	deviceIdentity, err := mach.CryptoStore.FindDeviceByKey(ctx, userID, identityKey)
-	if err != nil || deviceIdentity != nil || mach.DisableKeyFetching {
+	if err != nil || deviceIdentity != nil {
 		return deviceIdentity, err
 	}
 	mach.machOrContextLog(ctx).Debug().