Rewrite most of the device storage stuff

Gruntfile.js

@@ -81,6 +81,7 @@ module.exports = function(grunt) {
           'libaxolotl/crypto.js',
           'libaxolotl/protocol.js',
           'libaxolotl/protobufs.js',
+          'libaxolotl/session_storage.js',
         ],
         dest: 'libtextsecure/libaxolotl_concat.js',
       },

libaxolotl/protocol.js

@@ -69,25 +69,22 @@ window.axolotl.protocol = function() {
     }
 
     crypto_storage.saveSession = function(encodedNumber, session, registrationId) {
-        var device = axolotl.api.storage.sessions.get(encodedNumber);
-        if (device === undefined)
-            device = { sessions: {}, encodedNumber: encodedNumber };
-
-        if (registrationId !== undefined)
-            device.registrationId = registrationId;
+        var record = axolotl.api.storage.sessions.get(encodedNumber);
+        if (record === undefined) {
+            if (registrationId === undefined)
+                throw new Error("Tried to save a session for an existing device that didn't exist");
+            else
+                record = new axolotl.sessions.RecipientRecord(session.indexInfo.remoteIdentityKey, registrationId);
+        }
 
-        crypto_storage.saveSessionAndDevice(device, session);
-    }
+        var sessions = record._sessions;
 
-    crypto_storage.saveSessionAndDevice = function(device, session) {
-        if (device.sessions === undefined)
-            device.sessions = {};
-        var sessions = device.sessions;
+        if (record.identityKey === null)
+            record.identityKey = session.indexInfo.remoteIdentityKey;
+        if (getString(record.identityKey) !== getString(session.indexInfo.remoteIdentityKey))
+            throw new Error("Identity key changed at session save time");
 
         var doDeleteSession = false;
-        if (session.indexInfo.closed == -1 || device.identityKey === undefined)
-            device.identityKey = session.indexInfo.remoteIdentityKey;
-
         if (session.indexInfo.closed != -1) {
             doDeleteSession = (session.indexInfo.closed < (new Date().getTime() - MESSAGE_LOST_THRESHOLD_MS));
 
@@ -114,19 +111,17 @@ window.axolotl.protocol = function() {
         for (var key in sessions)
             if (sessions[key].indexInfo.closed == -1)
                 openSessionRemaining = true;
-        if (!openSessionRemaining)
-            try {
-                delete device['registrationId'];
-            } catch(_) {}
+        if (!openSessionRemaining) // Used as a flag to get new pre keys for the next session
+            record.registrationId = null;
 
-        axolotl.api.storage.sessions.put(device);
+        axolotl.api.storage.sessions.put(encodedNumber, record);
     }
 
     var getSessions = function(encodedNumber) {
-        var device = axolotl.api.storage.sessions.get(encodedNumber);
-        if (device === undefined || device.sessions === undefined)
+        var record = axolotl.api.storage.sessions.get(encodedNumber);
+        if (record === undefined)
             return undefined;
-        return device.sessions;
+        return record._sessions;
     }
 
     crypto_storage.getOpenSession = function(encodedNumber) {
@@ -164,17 +159,17 @@ window.axolotl.protocol = function() {
     }
 
     crypto_storage.getSessionOrIdentityKeyByBaseKey = function(encodedNumber, baseKey) {
-        var sessions = getSessions(encodedNumber);
-        var device = axolotl.api.storage.sessions.get(encodedNumber);
-        if (device === undefined)
+        var record = axolotl.api.storage.sessions.get(encodedNumber);
+        if (record === undefined)
             return undefined;
+        var sessions = record._sessions;
 
-        var preferredSession = device.sessions && device.sessions[getString(baseKey)];
+        var preferredSession = record._sessions[getString(baseKey)];
         if (preferredSession !== undefined)
             return preferredSession;
 
-        if (device.identityKey !== undefined)
-            return { indexInfo: { remoteIdentityKey: device.identityKey } };
+        if (record.identityKey !== undefined)
+            return { indexInfo: { remoteIdentityKey: record.identityKey } };
 
         throw new Error("Datastore inconsistency: device was stored without identity key");
     }
@@ -540,6 +535,7 @@ window.axolotl.protocol = function() {
     // return Promise(encoded [PreKey]WhisperMessage)
     self.encryptMessageFor = function(deviceObject, pushMessageContent) {
         var session = crypto_storage.getOpenSession(deviceObject.encodedNumber);
+        var hadSession = session !== undefined;
 
         var doEncryptPushMessageContent = function() {
             var msg = new axolotl.protobuf.WhisperMessage();
@@ -574,17 +570,9 @@ window.axolotl.protocol = function() {
                             result.set(new Uint8Array(encodedMsg), 1);
                             result.set(new Uint8Array(mac, 0, 8), encodedMsg.byteLength + 1);
 
-                            try {
-                                delete deviceObject['signedKey'];
-                                delete deviceObject['signedKeyId'];
-                                delete deviceObject['signedKeySignature'];
-                                delete deviceObject['preKey'];
-                                delete deviceObject['preKeyId'];
-                            } catch(_) {}
-
                             removeOldChains(session);
 
-                            crypto_storage.saveSessionAndDevice(deviceObject, session);
+                            crypto_storage.saveSession(deviceObject.encodedNumber, session, !hadSession ? deviceObject.registrationId : undefined);
                             return result;
                         });
                     });

libaxolotl/session_storage.js

@@ -0,0 +1,53 @@
+/* vim: ts=4:sw=4
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+;(function() {
+
+'use strict';
+window.axolotl = window.axolotl || {};
+
+var RecipientRecord = function(identityKey, registrationId) {
+    this._sessions = {};
+    this.identityKey = identityKey !== undefined ? getString(identityKey) : null;
+    this.registrationId = registrationId;
+
+    if (this.registrationId === undefined || typeof this.registrationId !== "number")
+        this.registrationId = null;
+};
+
+RecipientRecord.prototype.serialize = function() {
+    return textsecure.utils.jsonThing({sessions: this._sessions, registrationId: this.registrationId, identityKey: this.identityKey});
+}
+
+RecipientRecord.prototype.deserialize = function(serialized) {
+    var data = JSON.parse(serialized);
+    this._sessions = data.sessions;
+    if (this._sessions === undefined || this._sessions === null || typeof this._sessions !== "object" || Array.isArray(this._sessions))
+        throw new Error("Error deserializing RecipientRecord");
+    this.identityKey = data.identityKey;
+    this.registrationId = data.registrationId;
+    if (this.identityKey === undefined || this.registrationId === undefined)
+        throw new Error("Error deserializing RecipientRecord");
+}
+
+RecipientRecord.prototype.haveOpenSession = function() {
+    return this.registrationId !== null;
+}
+
+window.axolotl.sessions = {
+    RecipientRecord: RecipientRecord,
+};
+
+})();

libtextsecure/axolotl_wrapper.js

@@ -21,10 +21,32 @@
 
             sessions: {
                 get: function(identifier) {
-                    return textsecure.storage.devices.getDeviceObject(identifier);
+                    var device = textsecure.storage.devices.getDeviceObject(identifier, true);
+                    if (device === undefined)
+                        return undefined;
+                    var record = new axolotl.sessions.RecipientRecord();
+                    if (device.sessions !== undefined)
+                        record.deserialize(device.sessions);
+                    else
+                        // TODO: (even for numbers, not devices) We MUST return an identityKey if we have one available
+                        record.identityKey = device.identityKey;
+                    if (getString(device.identityKey) !== getString(record.identityKey))
+                        throw new Error("Got mismatched identity key on sessions load");
+                    return record;
                 },
-                put: function(object) {
-                    return textsecure.storage.devices.saveDeviceObject(object);
+                put: function(identifier, record) {
+                    var device = textsecure.storage.devices.getDeviceObject(identifier);
+                    if (device === undefined) {
+                        device = { encodedNumber: identifier,
+                                   //TODO: Remove this duplication (esp registrationId?)
+                                   identityKey: record.identityKey,
+                                   registrationId: record.registrationId
+                                 };
+                    }
+                    if (getString(device.identityKey) !== getString(record.identityKey))
+                        throw new Error("Tried to put session for device with changed identity key");
+                    device.sessions = record.serialize();
+                    return textsecure.storage.devices.saveDeviceObject(device);
                 }
             }
         },

libtextsecure/sendmessage.js

@@ -67,6 +67,7 @@ window.textsecure.messaging = function() {
             }
 
             return axolotl.protocol.encryptMessageFor(deviceObjectList[i], message).then(function(encryptedMsg) {
+                textsecure.storage.devices.removeTempKeysFromDevice(deviceObjectList[i].encodedNumber);
                 jsonData[i] = {
                     type: encryptedMsg.type,
                     destinationDeviceId: textsecure.utils.unencodeNumber(deviceObjectList[i].encodedNumber)[1],

libtextsecure/storage/devices.js

@@ -32,21 +32,42 @@
             return internalSaveDeviceObject(deviceObject, true);
         },
 
+        removeTempKeysFromDevice: function(encodedNumber) {
+            var deviceObject = textsecure.storage.devices.getDeviceObject(encodedNumber);
+            try {
+                delete deviceObject['signedKey'];
+                delete deviceObject['signedKeyId'];
+                delete deviceObject['signedKeySignature'];
+                delete deviceObject['preKey'];
+                delete deviceObject['preKeyId'];
+            } catch(_) {}
+            return internalSaveDeviceObject(deviceObject, false);
+        },
+
         getDeviceObjectsForNumber: function(number) {
             var map = textsecure.storage.getEncrypted("devices" + number);
             return map === undefined ? [] : map.devices;
         },
 
-        getDeviceObject: function(encodedNumber) {
-            var number = textsecure.utils.unencodeNumber(encodedNumber);
-            var devices = textsecure.storage.devices.getDeviceObjectsForNumber(number[0]);
-            if (devices === undefined)
+        getDeviceObject: function(encodedNumber, returnIdentityKey) {
+            var number = textsecure.utils.unencodeNumber(encodedNumber)[0];
+            var devices = textsecure.storage.devices.getDeviceObjectsForNumber(number);
+            if (devices.length == 0) {
+                if (returnIdentityKey) {
+                    var identityKey = textsecure.storage.devices.getIdentityKeyForNumber(number);
+                    if (identityKey !== undefined)
+                        return {identityKey: identityKey};
+                }
                 return undefined;
+            }
 
             for (var i in devices)
                 if (devices[i].encodedNumber == encodedNumber)
                     return devices[i];
 
+            if (returnIdentityKey)
+                return {identityKey: devices[0].identityKey};
+
             return undefined;
         },
 
@@ -96,6 +117,7 @@
                         map.devices[i].preKeyId = deviceObject.preKeyId;
                         map.devices[i].signedKey = deviceObject.signedKey;
                         map.devices[i].signedKeyId = deviceObject.signedKeyId;
+                        map.devices[i].signedKeySignature = deviceObject.signedKeySignature;
                         map.devices[i].registrationId = deviceObject.registrationId;
                     }
                     updated = true;