git » sdk » commit cf62ef3

diff --git a/doc/OMEMO.md b/doc/OMEMO.md
index edaa116..232867d 100644
--- a/doc/OMEMO.md
+++ b/doc/OMEMO.md
@@ -14,6 +14,5 @@ compile with the NO_OMEMO flag.
 - No API to determine cryptographic identity of message sender
 - Persistence: only IndexedDB backend is currently implemented
 - Encryption status reported by the API can be forged by sender
-- Consumed prekeys are not removed and replaced
 - Outgoing messages are not encrypted to the sending account's other devices
 - No support for group chats
diff --git a/snikket/OMEMO.hx b/snikket/OMEMO.hx
index bc6568a..5113b3f 100644
--- a/snikket/OMEMO.hx
+++ b/snikket/OMEMO.hx
@@ -87,6 +87,11 @@ class OMEMOBundle {
 	public function getRandomPreKey():PublicPreKey {
 		return prekeys[Std.random(prekeys.length-1)];
 	}
+
+	// Return a new bundle with an updated set of prekeys
+	public function withNewPreKeys(newPreKeys:Array<PublicPreKey>):OMEMOBundle {
+		return new OMEMOBundle(this.identity_key, this.device_id, newPreKeys, this.signed_prekey);
+	}
 }
 
 class OMEMOStore extends SignalProtocolStore {
@@ -142,6 +147,7 @@ class OMEMOStore extends SignalProtocolStore {
 	}
 
 	public function removePreKey(keyId:Int):Promise<Bool> {
+		trace("OMEMO: Removing prekey "+keyId);
 		persistence.removeOmemoPreKey(accountId, keyId);
 		// FIXME: Need to signal that we need to generate a replacement
 		// for the consumed prekey and republish our bundle
@@ -444,8 +450,8 @@ class OMEMO {
 			newBundle.prekeys = [
 				for(i in 0...prekeys.length) {
 					{
-						keyId: i+1,
-						pubKey: Base64.encode(Bytes.ofData(prekeys[i].pubKey)),
+						keyId: prekeys[i].keyId,
+						pubKey: Base64.encode(Bytes.ofData(prekeys[i].keyPair.pubKey)),
 					};
 				}
 			];
@@ -630,6 +636,8 @@ class OMEMO {
 	// Stuff that touches libsignal
 
 	// Build and store an OMEMO identity bundle
+	// This should only be called once, when setting up a
+	// new client!
 	private function buildBundle():Promise<Bool> {
 		final deviceId = KeyHelper.generateRegistrationId(); // FIXME: Check for collision
 		var identityKeyPair:IdentityKeyPair;
@@ -667,34 +675,75 @@ class OMEMO {
 		});
 	}
 
-	private function generatePreKeys(_:Bool):Promise<Array<PublicPreKey>> {
-		final amount = NUM_PREKEYS;
-		final keys = [];
+	private function storePreKeys(prekeys:Array<PreKey>):Promise<Array<PublicPreKey>> {
+		for(prekey in prekeys) {
+			// Store the full keypair
+			persistence.storeOmemoPreKey(client.accountId(), prekey.keyId, prekey.keyPair);
+		}
+		return Promise.resolve([
+			for(prekey in prekeys) {
+				// Emit the base64 public part for the application to publish
+				{
+					keyId: prekey.keyId,
+					pubKey: Base64.encode(Bytes.ofData(prekey.keyPair.pubKey))
+				};
+			}
+		]);
+	}
 
+	// Generate new prekeys
+	private function generatePreKeys(_:Bool):Promise<Array<PublicPreKey>> {
 		final generatedPreKeys:Promise<Array<PreKey>> = PromiseTools.all([
-			for(i in 1...(amount+1)) {
+			for(i in 1...(NUM_PREKEYS+1)) {
 				trace("Generating prekey "+Std.string(i));
 				KeyHelper.generatePreKey(i);
 			}
 		]);
 
-		final publicStoredPreKeys:Promise<Array<PublicPreKey>> = cast generatedPreKeys.then(cast function (prekeys:Array<PreKey>):Array<PublicPreKey> {
-			for(prekey in prekeys) {
-				// Store the full keypair
-				persistence.storeOmemoPreKey(client.accountId(), prekey.keyId, prekey.keyPair);
-			}
-			return [
-				for(prekey in prekeys) {
-					// Emit the base64 public part for the application to publish
-					{
-						keyId: prekey.keyId,
-						pubKey: Base64.encode(Bytes.ofData(prekey.keyPair.pubKey))
-					};
+		return generatedPreKeys.then(storePreKeys);
+	}
+
+	// Generate only prekeys which are "missing" (i.e. consumed)
+	// Returns an array of all prekeys, which can be used to update the
+	// published bundle.
+	private function generateMissingPreKeys():Promise<Array<PublicPreKey>> {
+		return new Promise((resolve, reject) -> {
+			persistence.getOmemoPreKeys(client.accountId(), function (prekeys:Array<PreKey>) {
+				// Generate an array of all keyIds we currently have in storage
+				final currentKeyIds:Array<Int> = prekeys.map(function (prekey) {
+					return prekey.keyId;
+				});
+
+				currentKeyIds.sort(function (a:Int, b:Int):Int {
+					if(a == b) {
+						return 0;
+					}
+					return a < b ? -1 : 1;
+				});
+
+				final generatedKeys:Array<Promise<PreKey>> = [];
+				var idx = 0;
+				for(keyId in 1...(NUM_PREKEYS+1)) {
+					if(currentKeyIds[idx] == keyId) {
+						// Key already present
+						idx++;
+					} else {
+						trace("Generating replacement prekey "+Std.string(keyId));
+						generatedKeys.push(KeyHelper.generatePreKey(keyId));
+					}
 				}
-			];
+
+				PromiseTools.all(generatedKeys).then(storePreKeys).then((storedPreKeys) -> {
+					resolve(prekeys.map(preKeyToPublicPreKey).concat(storedPreKeys));
+				});
+			});
 		});
+	}
 
-		return publicStoredPreKeys;
+	private function publishNewPreKeys(newPreKeys:Array<PublicPreKey>):Bool {
+		this.bundle = this.bundle.withNewPreKeys(newPreKeys);
+		publishBundle();
+		return true;
 	}
 
 	public function getDeviceId():Promise<Int> {
@@ -748,6 +797,17 @@ class OMEMO {
 				return cipher.decryptWhisperMessage(deviceKey.getRawKey());
 			}
 		});
+
+		if(deviceKey.prekey) {
+			promRawKeyWithTag.then((cipher) -> {
+				// Now it has been used, we need to replace the prekey that
+				// was used for this incoming message. libsignal has already
+				// removed it from the store, so we just need to regenerate
+				// any missing keys
+				generateMissingPreKeys().then(publishNewPreKeys);
+			});
+		}
+
 		final promPayload = promRawKeyWithTag.then((rawKeyWithTag) -> {
 			return decryptPayloadWithKey(payload.getRawPayload(), rawKeyWithTag, payload.getRawIv());
 		});
@@ -1074,4 +1134,11 @@ class OMEMO {
 			return header.toXml();
 		});
 	}
+
+	static private function preKeyToPublicPreKey(prekey:PreKey):PublicPreKey {
+		return {
+			keyId: prekey.keyId,
+			pubKey: Base64.encode(Bytes.ofData(prekey.keyPair.pubKey)),
+		};
+	}
 }
diff --git a/snikket/Persistence.hx b/snikket/Persistence.hx
index 65f5398..6801072 100644
--- a/snikket/Persistence.hx
+++ b/snikket/Persistence.hx
@@ -51,7 +51,7 @@ interface Persistence {
 	public function removeOmemoPreKey(identifier:String, keyId:Int):Void;
 	public function storeOmemoSignedPreKey(login:String, signedPreKey:SignedPreKey):Void;
 	public function getOmemoSignedPreKey(login:String, keyId:Int, callback: (SignedPreKey)->Void):Void;
-	public function getOmemoPreKeys(login:String, callback: (Array<PreKeyPair>)->Void):Void;
+	public function getOmemoPreKeys(login:String, callback: (Array<PreKey>)->Void):Void;
 	public function storeOmemoContactIdentityKey(account:String, address:String, identityKey:IdentityPublicKey):Void;
 	public function getOmemoContactIdentityKey(account:String, address:String, callback:(IdentityPublicKey)->Void):Void;
 	public function getOmemoSession(account:String, address:String, callback:(SignalSession)->Void):Void;
diff --git a/snikket/persistence/Dummy.hx b/snikket/persistence/Dummy.hx
index ebad151..aa9ddc9 100644
--- a/snikket/persistence/Dummy.hx
+++ b/snikket/persistence/Dummy.hx
@@ -174,7 +174,7 @@ class Dummy implements Persistence {
 	public function getOmemoSignedPreKey(login:String, keyId:Int, callback: (SignedPreKey)->Void):Void { }
 
 	@HaxeCBridge.noemit
-	public function getOmemoPreKeys(login:String, callback: (Array<PreKeyPair>)->Void):Void { }
+	public function getOmemoPreKeys(login:String, callback: (Array<PreKey>)->Void):Void { }
 
 	@HaxeCBridge.noemit
 	public function storeOmemoContactIdentityKey(account:String, address:String, identityKey:IdentityPublicKey):Void { }
diff --git a/snikket/persistence/IDB.js b/snikket/persistence/IDB.js
index d2e8f9d..9b27982 100644
--- a/snikket/persistence/IDB.js
+++ b/snikket/persistence/IDB.js
@@ -735,9 +735,14 @@ export default (dbname, media, tokenize, stemmer) => {
 			req.onsuccess = (event) => {
 				const cursor = event.target.result;
 				if(cursor) {
+					const splitDbKey = cursor.key.split(":");
+					const keyId = parseInt(splitDbKey[splitDbKey.length - 1], 10);
 					prekeys.push({
-						"privKey": base64ToArrayBuffer(cursor.value.privKey),
-						"pubKey": base64ToArrayBuffer(cursor.value.pubKey),
+						keyId: keyId,
+						keyPair: {
+							"privKey": base64ToArrayBuffer(cursor.value.privKey),
+							"pubKey": base64ToArrayBuffer(cursor.value.pubKey),
+						},
 					});
 					cursor.continue();
 				} else {
diff --git a/snikket/persistence/Sqlite.hx b/snikket/persistence/Sqlite.hx
index f1910db..6b548f9 100644
--- a/snikket/persistence/Sqlite.hx
+++ b/snikket/persistence/Sqlite.hx
@@ -865,7 +865,7 @@ class Sqlite implements Persistence implements KeyValueStore {
 	public function getOmemoSignedPreKey(login:String, keyId:Int, callback: (SignedPreKey)->Void):Void { }
 
 	@HaxeCBridge.noemit
-	public function getOmemoPreKeys(login:String, callback: (Array<PreKeyPair>)->Void):Void { }
+	public function getOmemoPreKeys(login:String, callback: (Array<PreKey>)->Void):Void { }
 
 	@HaxeCBridge.noemit
 	public function storeOmemoContactIdentityKey(account:String, address:String, identityKey:IdentityPublicKey):Void { }