From 3c91c66ee24bbe1dd270752585c601b99a539bcc Mon Sep 17 00:00:00 2001
From: Valere <valeref@matrix.org>
Date: Tue, 14 Dec 2021 16:40:22 +0100
Subject: [PATCH] Unpublished fallback key bindings + forget

---
 xcode/OLMKit/OLMAccount.h       | 15 ++++++++++++-
 xcode/OLMKit/OLMAccount.m       | 27 ++++++++++++++++++++++++
 xcode/OLMKitTests/OLMKitTests.m | 37 ++++++++++++++++++++++++++++-----
 3 files changed, 73 insertions(+), 6 deletions(-)

diff --git a/xcode/OLMKit/OLMAccount.h b/xcode/OLMKit/OLMAccount.h
index 5e5e5b6..49cab5f 100644
--- a/xcode/OLMKit/OLMAccount.h
+++ b/xcode/OLMKit/OLMAccount.h
@@ -35,14 +35,27 @@
 /** Public parts of the unpublished one time keys for the account */
 - (NSDictionary*) oneTimeKeys;
 
-/** Public part of the unpublished fallback key for the account */
+/**
+ * Deprecated use unPublishedFallbackKey
+ */
 - (NSDictionary*) fallbackKey;
 
+/**
+ Public part of the unpublished fallback key for the account, if present and unublished.
+ */
+- (NSDictionary*) unpublishedFallbackKey;
+
 - (BOOL) removeOneTimeKeysForSession:(OLMSession*)session;
 
 /** Marks the current set of one time keys as being published. */
 - (void) markOneTimeKeysAsPublished;
 
+/** Forget about the old fallback key.
+ * This should be called once you are reasonably certain that you will not
+ * receive any more messages that use the old fallback key
+ */
+- (void) forgetFallbackKey;
+
 /** The largest number of one time keys this account can store. */
 - (NSUInteger) maxOneTimeKeys;
 
diff --git a/xcode/OLMKit/OLMAccount.m b/xcode/OLMKit/OLMAccount.m
index 0aa8beb..76a47cd 100644
--- a/xcode/OLMKit/OLMAccount.m
+++ b/xcode/OLMKit/OLMAccount.m
@@ -179,6 +179,33 @@
     return keyDictionary;
 }
 
+- (NSDictionary *) unpublishedFallbackKey {
+    size_t fallbackKeyLength = olm_account_unpublished_fallback_key_length(_account);
+    uint8_t *fallbackKeyBytes = malloc(fallbackKeyLength);
+    if (!fallbackKeyBytes) {
+        return nil;
+    }
+    
+    size_t result = olm_account_unpublished_fallback_key(_account, fallbackKeyBytes, fallbackKeyLength);
+    if (result == olm_error()) {
+        const char *error = olm_account_last_error(_account);
+        NSLog(@"error getting unpublished fallback key: %s", error);
+        free(fallbackKeyBytes);
+        return nil;
+    }
+    NSData *fallbackKeyData = [NSData dataWithBytesNoCopy:fallbackKeyBytes length:fallbackKeyLength freeWhenDone:YES];
+    NSError *error = nil;
+    NSDictionary *keyDictionary = [NSJSONSerialization JSONObjectWithData:fallbackKeyData options:0 error:&error];
+    if (error) {
+        NSLog(@"Could not decode JSON for unpublished fallback: %@", error.localizedDescription);
+    }
+    return keyDictionary;
+}
+
+- (void) forgetFallbackKey {
+    olm_account_forget_old_fallback_key(self.account);
+}
+
 - (void) generateFallbackKey {
     size_t randomLength = olm_account_generate_fallback_key_random_length(_account);
     NSMutableData *random = [OLMUtility randomBytesOfLength:randomLength];
diff --git a/xcode/OLMKitTests/OLMKitTests.m b/xcode/OLMKitTests/OLMKitTests.m
index 2cf8a0d..1d75a81 100644
--- a/xcode/OLMKitTests/OLMKitTests.m
+++ b/xcode/OLMKitTests/OLMKitTests.m
@@ -36,7 +36,34 @@ limitations under the License.
     OLMAccount *bob = [[OLMAccount alloc] initNewAccount];
     [bob generateFallbackKey];
     
-    [self _testAliceAndBob:bob withBobKeys:bob.fallbackKey];
+    [self _testAliceAndBob:bob withBobKeys:bob.unpublishedFallbackKey];
+
+}
+
+- (void)testMarkAsPublishedFallbackKey {
+    OLMAccount *bob = [[OLMAccount alloc] initNewAccount];
+    [bob generateFallbackKey];
+    
+    
+    NSDictionary *unpublished = bob.unpublishedFallbackKey;
+    __block NSString *bobKeyValue = nil;
+    [unpublished[@"curve25519"] enumerateKeysAndObjectsUsingBlock:^(id  _Nonnull key, id  _Nonnull obj, BOOL * _Nonnull stop) {
+        bobKeyValue = obj;
+    }];
+    
+    
+    XCTAssertNotNil(bobKeyValue);
+    
+    [bob markOneTimeKeysAsPublished];
+    
+    NSDictionary *unpublishedAfter = bob.unpublishedFallbackKey;
+    
+    __block NSString *bobKeyValueAfter = nil;
+    [unpublishedAfter[@"curve25519"] enumerateKeysAndObjectsUsingBlock:^(id  _Nonnull key, id  _Nonnull obj, BOOL * _Nonnull stop) {
+        bobKeyValueAfter = obj;
+    }];
+    
+    XCTAssertNil(bobKeyValueAfter);
 }
 
 - (void)_testAliceAndBob:(OLMAccount *)bob withBobKeys:(NSDictionary *)bobKeys {
@@ -89,7 +116,7 @@ limitations under the License.
     OLMAccount *bob = [[OLMAccount alloc] initNewAccount];
     [bob generateFallbackKey];
     
-    [self _testBackAndForthWithBob:bob andBobKeys:bob.fallbackKey];
+    [self _testBackAndForthWithBob:bob andBobKeys:bob.unpublishedFallbackKey];
 }
 
 - (void)_testBackAndForthWithBob:(OLMAccount *)bob andBobKeys:(NSDictionary *)bobKeys {
@@ -140,7 +167,7 @@ limitations under the License.
     [bob generateFallbackKey];
     NSDictionary *bobIdKeys = bob.identityKeys;
     NSDictionary *bobOneTimeKeys = bob.oneTimeKeys;
-    NSDictionary *bobFallbackKey = bob.fallbackKey;
+    NSDictionary *bobFallbackKey = bob.unpublishedFallbackKey;
     
     NSError *error;
     NSData *bobData = [NSKeyedArchiver archivedDataWithRootObject:bob requiringSecureCoding:NO error:&error];
@@ -151,7 +178,7 @@ limitations under the License.
     
     NSDictionary *bobIdKeys2 = bob2.identityKeys;
     NSDictionary *bobOneTimeKeys2 = bob2.oneTimeKeys;
-    NSDictionary *bobFallbackKey2 = bob2.fallbackKey;
+    NSDictionary *bobFallbackKey2 = bob2.unpublishedFallbackKey;
     
     XCTAssertEqualObjects(bobIdKeys, bobIdKeys2);
     XCTAssertEqualObjects(bobOneTimeKeys, bobOneTimeKeys2);
@@ -169,7 +196,7 @@ limitations under the License.
     OLMAccount *bob = [[OLMAccount alloc] initNewAccount];
     [bob generateFallbackKey];
     
-    [self _testSessionSerializationWithBob:bob bobKeys:bob.fallbackKey];
+    [self _testSessionSerializationWithBob:bob bobKeys:bob.unpublishedFallbackKey];
 }
 
 - (void)_testSessionSerializationWithBob:(OLMAccount *)bob bobKeys:(NSDictionary *)bobKeys {