diff --git a/lib/src/omemo/events.dart b/lib/src/omemo/events.dart index b41a35a..89359c1 100644 --- a/lib/src/omemo/events.dart +++ b/lib/src/omemo/events.dart @@ -12,6 +12,14 @@ class RatchetModifiedEvent extends OmemoEvent { final OmemoDoubleRatchet ratchet; } +/// Triggered when a ratchet has been removed and should be removed from storage. +class RatchetRemovedEvent extends OmemoEvent { + + RatchetRemovedEvent(this.jid, this.deviceId); + final String jid; + final int deviceId; +} + /// Triggered when the device map has been modified class DeviceMapModifiedEvent extends OmemoEvent { diff --git a/lib/src/omemo/sessionmanager.dart b/lib/src/omemo/sessionmanager.dart index 5df282d..b995347 100644 --- a/lib/src/omemo/sessionmanager.dart +++ b/lib/src/omemo/sessionmanager.dart @@ -3,7 +3,8 @@ import 'dart:convert'; import 'package:collection/collection.dart'; import 'package:cryptography/cryptography.dart'; import 'package:hex/hex.dart'; -import 'package:meta/meta.dart'; import 'package:omemo_dart/src/crypto.dart'; +import 'package:meta/meta.dart'; +import 'package:omemo_dart/src/crypto.dart'; import 'package:omemo_dart/src/double_ratchet/double_ratchet.dart'; import 'package:omemo_dart/src/errors.dart'; import 'package:omemo_dart/src/helpers.dart'; @@ -93,6 +94,7 @@ class OmemoSessionManager { /// A stream that receives events regarding the session Stream get eventStream => _eventStreamController.stream; + /// Returns our own device. Future getDevice() async { Device? dev; await _deviceLock.synchronized(() async { @@ -403,11 +405,31 @@ class OmemoSessionManager { Map>? map; await _lock.synchronized(() async { - map = _deviceMap; + map = _deviceMap; }); return map!; } + + /// Removes the ratchet identified by [jid] and [deviceId] from the session manager. + /// Also triggers events for commiting the new device map to storage and removing + /// the old ratchet. + Future removeRatchet(String jid, int deviceId) async { + await _lock.synchronized(() async { + // Remove the ratchet + _ratchetMap.remove(RatchetMapKey(jid, deviceId)); + // Commit it + _eventStreamController.add(RatchetRemovedEvent(jid, deviceId)); + + // Remove the device from jid + _deviceMap[jid]!.remove(deviceId); + if (_deviceMap[jid]!.isEmpty) { + _deviceMap.remove(jid); + } + // Commit it + _eventStreamController.add(DeviceMapModifiedEvent(_deviceMap)); + }); + } @visibleForTesting OmemoDoubleRatchet getRatchet(String jid, int deviceId) => _ratchetMap[RatchetMapKey(jid, deviceId)]!; diff --git a/test/omemo_test.dart b/test/omemo_test.dart index e39f95e..27f072c 100644 --- a/test/omemo_test.dart +++ b/test/omemo_test.dart @@ -422,4 +422,81 @@ void main() { expect(messageText, aliceReceivedMessage); } }); + + group('Test removing a ratchet', () { + test('Test removing a ratchet when the user has multiple', () async { + const aliceJid = 'alice@server.local'; + const bobJid = 'bob@some.server.local'; + final aliceSession = await OmemoSessionManager.generateNewIdentity( + aliceJid, + AlwaysTrustingTrustManager(), + opkAmount: 1, + ); + final bobSession1 = await OmemoSessionManager.generateNewIdentity( + bobJid, + AlwaysTrustingTrustManager(), + opkAmount: 1, + ); + final bobSession2 = await OmemoSessionManager.generateNewIdentity( + bobJid, + AlwaysTrustingTrustManager(), + opkAmount: 1, + ); + + // Alice sends a message to those two Bobs + final aliceMessage = await aliceSession.encryptToJid( + bobJid, + 'Hallo Welt', + newSessions: [ + await (await bobSession1.getDevice()).toBundle(), + await (await bobSession2.getDevice()).toBundle(), + ], + ); + + // One of those two sessions is broken, so Alice removes the session2 ratchet + final id1 = (await bobSession1.getDevice()).id; + final id2 = (await bobSession2.getDevice()).id; + await aliceSession.removeRatchet(bobJid, id1); + + final map = await aliceSession.getRatchetMap(); + expect(map.containsKey(RatchetMapKey(bobJid, id1)), false); + expect(map.containsKey(RatchetMapKey(bobJid, id2)), true); + final deviceMap = await aliceSession.getDeviceMap(); + expect(deviceMap.containsKey(bobJid), true); + expect(deviceMap[bobJid], [id2]); + }); + + test('Test removing a ratchet when the user has only one', () async { + const aliceJid = 'alice@server.local'; + const bobJid = 'bob@some.server.local'; + final aliceSession = await OmemoSessionManager.generateNewIdentity( + aliceJid, + AlwaysTrustingTrustManager(), + opkAmount: 1, + ); + final bobSession = await OmemoSessionManager.generateNewIdentity( + bobJid, + AlwaysTrustingTrustManager(), + opkAmount: 1, + ); + + // Alice sends a message to those two Bobs + final aliceMessage = await aliceSession.encryptToJid( + bobJid, + 'Hallo Welt', + newSessions: [ + await (await bobSession.getDevice()).toBundle(), + ], + ); + + // One of those two sessions is broken, so Alice removes the session2 ratchet + final id = (await bobSession.getDevice()).id; + await aliceSession.removeRatchet(bobJid, id); + + final map = await aliceSession.getRatchetMap(); + expect(map.containsKey(RatchetMapKey(bobJid, id)), false); + final deviceMap = await aliceSession.getDeviceMap(); + expect(deviceMap.containsKey(bobJid), false); + }); + }); }