Automatically create new Olm sessions if decryption fails

crypto/decryptmegolm.go

@@ -41,7 +41,6 @@ func (mach *OlmMachine) DecryptMegolmEvent(evt *event.Event) (*event.Event, erro
 	if err != nil {
 		return nil, fmt.Errorf("failed to get group session: %w", err)
 	} else if sess == nil {
-		mach.checkIfWedged(evt)
 		return nil, fmt.Errorf("%w (ID %s)", NoSessionFound, content.SessionID)
 	}
 	plaintext, messageIndex, err := sess.Internal.Decrypt(content.MegolmCiphertext)

crypto/decryptolm.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2020 Tulir Asokan
+// Copyright (c) 2021 Tulir Asokan
 //
 // This Source Code Form is subject to the terms of the Mozilla Public
 // License, v. 2.0. If a copy of the MPL was not distributed with this
@@ -106,7 +106,7 @@ func (mach *OlmMachine) tryDecryptOlmCiphertext(sender id.UserID, deviceID id.De
 	if err != nil {
 		if err == DecryptionFailedWithMatchingSession {
 			mach.Log.Warn("Found matching session yet decryption failed for sender %s with key %s", sender, senderKey)
-			mach.markDeviceForUnwedging(sender, senderKey)
+			go mach.unwedgeDevice(sender, deviceID, senderKey)
 		}
 		return nil, fmt.Errorf("failed to decrypt olm event: %w", err)
 	}
@@ -121,14 +121,14 @@ func (mach *OlmMachine) tryDecryptOlmCiphertext(sender id.UserID, deviceID id.De
 	// New sessions can only be created if it's a prekey message, we can't decrypt the message
 	// if it isn't one at this point in time anymore, so return early.
 	if olmType != id.OlmMsgTypePreKey {
-		mach.markDeviceForUnwedging(sender, senderKey)
+		go mach.unwedgeDevice(sender, deviceID, senderKey)
 		return nil, DecryptionFailedForNormalMessage
 	}
 
 	mach.Log.Trace("Trying to create inbound session for %s/%s", sender, deviceID)
 	session, err := mach.createInboundSession(senderKey, ciphertext)
 	if err != nil {
-		mach.markDeviceForUnwedging(sender, senderKey)
+		go mach.unwedgeDevice(sender, deviceID, senderKey)
 		return nil, fmt.Errorf("failed to create new session from prekey message: %w", err)
 	}
 	mach.Log.Debug("Created inbound olm session %s for %s/%s (sender key: %s)", session.ID(), sender, deviceID, senderKey)
@@ -187,3 +187,16 @@ func (mach *OlmMachine) createInboundSession(senderKey id.SenderKey, ciphertext
 	}
 	return session, nil
 }
+
+func (mach *OlmMachine) unwedgeDevice(sender id.UserID, deviceID id.DeviceID, senderKey id.SenderKey) {
+	mach.Log.Debug("Creating new Olm session with %s/%s...", sender, deviceID)
+	mach.devicesToUnwedge.Store(senderKey, true)
+	err := mach.SendEncryptedToDevice(&DeviceIdentity{
+		UserID:      sender,
+		DeviceID:    deviceID,
+		IdentityKey: senderKey,
+	}, event.ToDeviceDummy, event.Content{})
+	if err != nil {
+		mach.Log.Error("Failed to send dummy event to unwedge session with %s/%s: %v", sender, deviceID, err)
+	}
+}

crypto/encryptolm.go

@@ -54,6 +54,9 @@ func (mach *OlmMachine) createOutboundSessions(input map[id.UserID]map[id.Device
 		for deviceID, identity := range devices {
 			if !mach.CryptoStore.HasSession(identity.IdentityKey) {
 				request[userID][deviceID] = id.KeyAlgorithmSignedCurve25519
+			} else if _, shouldUnwedge := mach.devicesToUnwedge.LoadAndDelete(identity.IdentityKey); shouldUnwedge {
+				mach.Log.Trace("Adding %s/%s to claim key request for unwedging", userID, deviceID)
+				request[userID][deviceID] = id.KeyAlgorithmSignedCurve25519
 			}
 		}
 		if len(request[userID]) == 0 {
@@ -77,7 +80,7 @@ func (mach *OlmMachine) createOutboundSessions(input map[id.UserID]map[id.Device
 			for keyID, oneTimeKey = range oneTimeKeys {
 				break
 			}
-			keyAlg, _ := keyID.Parse()
+			keyAlg, keyIndex := keyID.Parse()
 			if keyAlg != id.KeyAlgorithmSignedCurve25519 {
 				mach.Log.Warn("Unexpected key ID algorithm in one-time key response for %s of %s: %s", deviceID, userID, keyID)
 				continue
@@ -94,6 +97,8 @@ func (mach *OlmMachine) createOutboundSessions(input map[id.UserID]map[id.Device
 				err = mach.CryptoStore.AddSession(identity.IdentityKey, wrapped)
 				if err != nil {
 					mach.Log.Error("Failed to store created session for %s of %s: %v", deviceID, userID, err)
+				} else {
+					mach.Log.Debug("Created new Olm session with %s/%s (OTK ID: %d)", userID, deviceID, keyIndex)
 				}
 			}
 		}

crypto/machine.go

@@ -55,6 +55,8 @@ type OlmMachine struct {
 	keyWaiters     map[id.SessionID]chan struct{}
 	keyWaitersLock sync.Mutex
 
+	devicesToUnwedge *sync.Map
+
 	olmLock sync.Mutex
 
 	CrossSigningKeys    *CrossSigningKeysCache
@@ -93,6 +95,8 @@ func NewOlmMachine(client *mautrix.Client, log Logger, cryptoStore Store, stateS
 		keyVerificationTransactionState: &sync.Map{},
 
 		keyWaiters: make(map[id.SessionID]chan struct{}),
+
+		devicesToUnwedge: &sync.Map{},
 	}
 	mach.AllowKeyShare = mach.defaultAllowKeyShare
 	return mach
@@ -277,7 +281,8 @@ func (mach *OlmMachine) HandleToDeviceEvent(evt *event.Event) {
 					close(ch.(chan struct{}))
 				}
 			}
-			// TODO handle m.dummy encrypted to-device event?
+		case *event.DummyEventContent:
+			mach.Log.Debug("Received encrypted dummy event from %s/%s", decryptedEvt.Sender, decryptedEvt.SenderDevice)
 		default:
 			mach.Log.Debug("Unhandled encrypted to-device event of type %s from %s/%s", decryptedEvt.Type.String(), decryptedEvt.Sender, decryptedEvt.SenderDevice)
 		}
@@ -327,7 +332,6 @@ func (mach *OlmMachine) GetOrFetchDevice(userID id.UserID, deviceID id.DeviceID)
 
 // SendEncryptedToDevice sends an Olm-encrypted event to the given user device.
 func (mach *OlmMachine) SendEncryptedToDevice(device *DeviceIdentity, evtType event.Type, content event.Content) error {
-	// create outbound sessions if missing
 	if err := mach.createOutboundSessions(map[id.UserID]map[id.DeviceID]*DeviceIdentity{
 		device.UserID: {
 			device.DeviceID: device,
@@ -339,7 +343,6 @@ func (mach *OlmMachine) SendEncryptedToDevice(device *DeviceIdentity, evtType ev
 	mach.olmLock.Lock()
 	defer mach.olmLock.Unlock()
 
-	// get Olm session
 	olmSess, err := mach.CryptoStore.GetLatestSession(device.IdentityKey)
 	if err != nil {
 		return err

crypto/unwedge.go

@@ -1,21 +0,0 @@
-// Copyright (c) 2020 Tulir Asokan
-//
-// This Source Code Form is subject to the terms of the Mozilla Public
-// License, v. 2.0. If a copy of the MPL was not distributed with this
-// file, You can obtain one at http://mozilla.org/MPL/2.0/.
-
-package crypto
-
-import (
-	"maunium.net/go/mautrix/event"
-	"maunium.net/go/mautrix/id"
-)
-
-func (mach *OlmMachine) markDeviceForUnwedging(sender id.UserID, senderKey id.SenderKey) {
-	// TODO implement
-}
-
-func (mach *OlmMachine) checkIfWedged(evt *event.Event) {
-	//content := evt.Content.AsEncrypted()
-	// TODO implement
-}

event/content.go

@@ -60,6 +60,7 @@ var TypeMap = map[Type]reflect.Type{
 	ToDeviceRoomKeyRequest:   reflect.TypeOf(RoomKeyRequestEventContent{}),
 	ToDeviceEncrypted:        reflect.TypeOf(EncryptedEventContent{}),
 	ToDeviceRoomKeyWithheld:  reflect.TypeOf(RoomKeyWithheldEventContent{}),
+	ToDeviceDummy:            reflect.TypeOf(DummyEventContent{}),
 
 	ToDeviceVerificationStart:   reflect.TypeOf(VerificationStartEventContent{}),
 	ToDeviceVerificationAccept:  reflect.TypeOf(VerificationAcceptEventContent{}),

event/encryption.go

@@ -139,3 +139,5 @@ type RoomKeyWithheldEventContent struct {
 	Code      RoomKeyWithheldCode `json:"code"`
 	Reason    string              `json:"reason,omitempty"`
 }
+
+type DummyEventContent struct{}

event/type.go

@@ -226,6 +226,7 @@ var (
 	ToDeviceForwardedRoomKey    = Type{"m.forwarded_room_key", ToDeviceEventType}
 	ToDeviceEncrypted           = Type{"m.room.encrypted", ToDeviceEventType}
 	ToDeviceRoomKeyWithheld     = Type{"m.room_key.withheld", ToDeviceEventType}
+	ToDeviceDummy               = Type{"m.dummy", ToDeviceEventType}
 	ToDeviceVerificationRequest = Type{"m.key.verification.request", ToDeviceEventType}
 	ToDeviceVerificationStart   = Type{"m.key.verification.start", ToDeviceEventType}
 	ToDeviceVerificationAccept  = Type{"m.key.verification.accept", ToDeviceEventType}