Add first cut of protocol v3 support.

1) Use the new /v2/keys API for storing/retrieving prekey bundles.

2) For sessions built with PreKeyBundle and PreKeyWhisperMessage,
   use a v3 ratcheting session when available.

build.gradle

@@ -51,7 +51,7 @@ dependencyVerification {
         'com.astuetz:pagerslidingtabstrip:f1641396732c7132a7abb837e482e5ee2b0ebb8d10813fc52bbaec2c15c184c2',
         'org.w3c:smil:085dc40f2bb249651578bfa07499fd08b16ad0886dbe2c4078586a408da62f9b',
         'org.apache.httpcomponents:httpclient-android:6f56466a9bd0d42934b90bfbfe9977a8b654c058bf44a12bdc2877c4e1f033f1',
-        'com.google.protobuf:protobuf-java:ad9769a22989e688a46af4d3accc348cc501ced22118033230542bc916e33f0b',
+        'com.google.protobuf:protobuf-java:e0c1c64575c005601725e7c6a02cebf9e1285e888f756b2a1d73ffa8d725cc74',
         'com.madgag:sc-light-jdk15on:931f39d351429fb96c2f749e7ecb1a256a8ebbf5edca7995c9cc085b94d1841d',
         'com.googlecode.libphonenumber:libphonenumber:eba17eae81dd622ea89a00a3a8c025b2f25d342e0d9644c5b62e16f15687c3ab',
         'org.whispersystems:gson:08f4f7498455d1539c9233e5aac18e9b1805815ef29221572996508eb512fe51',

libaxolotl/build.gradle

@@ -4,7 +4,7 @@ buildscript {
     }
 
     dependencies {
-        classpath 'com.android.tools.build:gradle:0.9.+'
+        classpath 'com.android.tools.build:gradle:0.11.+'
     }
 }
 
@@ -15,7 +15,7 @@ repositories {
 }
 
 dependencies {
-    compile 'com.google.protobuf:protobuf-java:2.4.1'
+    compile 'com.google.protobuf:protobuf-java:2.5.0'
 }
 
 android {

libaxolotl/protobuf/LocalStorageProtocol.proto

@@ -35,8 +35,9 @@ message SessionStructure {
     }
 
     message PendingPreKey {
-        optional uint32 preKeyId = 1;
+        optional uint32 preKeyId    = 1;
-        optional bytes  baseKey  = 2;
+        optional int32  deviceKeyId = 3;
+        optional bytes  baseKey     = 2;
     }
 
     optional uint32 sessionVersion      = 1;
@@ -56,6 +57,7 @@ message SessionStructure {
     optional uint32 localRegistrationId  = 11;
 
     optional bool needsRefresh = 12;
+    optional bytes aliceBaseKey = 13;
 }
 
 message RecordStructure {
@@ -69,6 +71,14 @@ message PreKeyRecordStructure {
     optional bytes  privateKey = 3;
 }
 
+message DeviceKeyRecordStructure {
+    optional uint32  id         = 1;
+    optional bytes   publicKey  = 2;
+    optional bytes   privateKey = 3;
+    optional bytes   signature  = 4;
+    optional fixed64 timestamp  = 5;
+}
+
 message IdentityKeyPairStructure {
     optional bytes publicKey  = 1;
     optional bytes privateKey = 2;

libaxolotl/protobuf/WhisperTextProtocol.proto

@@ -12,10 +12,11 @@ message WhisperMessage {
 
 message PreKeyWhisperMessage {
   optional uint32 registrationId = 5;
-  optional uint32 preKeyId    = 1;
+  optional uint32 preKeyId       = 1;
-  optional bytes  baseKey     = 2;
+  optional uint32 deviceKeyId    = 6;
-  optional bytes  identityKey = 3;
+  optional bytes  baseKey        = 2;
-  optional bytes  message     = 4; // WhisperMessage
+  optional bytes  identityKey    = 3;
+  optional bytes  message        = 4; // WhisperMessage
 }
 
 message KeyExchangeMessage {

libaxolotl/src/androidTest/java/org/whispersystems/test/InMemoryDeviceKeyStore.java

@@ -0,0 +1,58 @@
+package org.whispersystems.test;
+
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
+import org.whispersystems.libaxolotl.state.DeviceKeyRecord;
+import org.whispersystems.libaxolotl.state.DeviceKeyStore;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+
+public class InMemoryDeviceKeyStore implements DeviceKeyStore {
+
+  private final Map<Integer, byte[]> store = new HashMap<>();
+
+  @Override
+  public DeviceKeyRecord loadDeviceKey(int deviceKeyId) throws InvalidKeyIdException {
+    try {
+      if (!store.containsKey(deviceKeyId)) {
+        throw new InvalidKeyIdException("No such devicekeyrecord!");
+      }
+
+      return new DeviceKeyRecord(store.get(deviceKeyId));
+    } catch (IOException e) {
+      throw new AssertionError(e);
+    }
+  }
+
+  public List<DeviceKeyRecord> loadDeviceKeys() {
+    try {
+      List<DeviceKeyRecord> results = new LinkedList<>();
+
+      for (byte[] serialized : store.values()) {
+        results.add(new DeviceKeyRecord(serialized));
+      }
+
+      return results;
+    } catch (IOException e) {
+      throw new AssertionError(e);
+    }
+  }
+
+  @Override
+  public void storeDeviceKey(int deviceKeyId, DeviceKeyRecord record) {
+    store.put(deviceKeyId, record.serialize());
+  }
+
+  @Override
+  public boolean containsDeviceKey(int deviceKeyId) {
+    return store.containsKey(deviceKeyId);
+  }
+
+  @Override
+  public void removeDeviceKey(int deviceKeyId) {
+    store.remove(deviceKeyId);
+  }
+}

libaxolotl/src/androidTest/java/org/whispersystems/test/InMemoryPreKeyStore.java

@@ -13,7 +13,7 @@ public class InMemoryPreKeyStore implements PreKeyStore {
   private final Map<Integer, byte[]> store = new HashMap<>();
 
   @Override
-  public PreKeyRecord load(int preKeyId) throws InvalidKeyIdException {
+  public PreKeyRecord loadPreKey(int preKeyId) throws InvalidKeyIdException {
     try {
       if (!store.containsKey(preKeyId)) {
         throw new InvalidKeyIdException("No such prekeyrecord!");
@@ -26,17 +26,17 @@ public class InMemoryPreKeyStore implements PreKeyStore {
   }
 
   @Override
-  public void store(int preKeyId, PreKeyRecord record) {
+  public void storePreKey(int preKeyId, PreKeyRecord record) {
     store.put(preKeyId, record.serialize());
   }
 
   @Override
-  public boolean contains(int preKeyId) {
+  public boolean containsPreKey(int preKeyId) {
     return store.containsKey(preKeyId);
   }
 
   @Override
-  public void remove(int preKeyId) {
+  public void removePreKey(int preKeyId) {
     store.remove(preKeyId);
   }
 }

libaxolotl/src/androidTest/java/org/whispersystems/test/InMemorySessionStore.java

@@ -17,9 +17,9 @@ public class InMemorySessionStore implements SessionStore {
   public InMemorySessionStore() {}
 
   @Override
-  public synchronized SessionRecord load(long recipientId, int deviceId) {
+  public synchronized SessionRecord loadSession(long recipientId, int deviceId) {
     try {
-      if (contains(recipientId, deviceId)) {
+      if (containsSession(recipientId, deviceId)) {
         return new SessionRecord(sessions.get(new Pair<>(recipientId, deviceId)));
       } else {
         return new SessionRecord();
@@ -43,22 +43,22 @@ public class InMemorySessionStore implements SessionStore {
   }
 
   @Override
-  public synchronized void store(long recipientId, int deviceId, SessionRecord record) {
+  public synchronized void storeSession(long recipientId, int deviceId, SessionRecord record) {
     sessions.put(new Pair<>(recipientId, deviceId), record.serialize());
   }
 
   @Override
-  public synchronized boolean contains(long recipientId, int deviceId) {
+  public synchronized boolean containsSession(long recipientId, int deviceId) {
     return sessions.containsKey(new Pair<>(recipientId, deviceId));
   }
 
   @Override
-  public synchronized void delete(long recipientId, int deviceId) {
+  public synchronized void deleteSession(long recipientId, int deviceId) {
     sessions.remove(new Pair<>(recipientId, deviceId));
   }
 
   @Override
-  public synchronized void deleteAll(long recipientId) {
+  public synchronized void deleteAllSessions(long recipientId) {
     for (Pair<Long, Integer> key : sessions.keySet()) {
       if (key.first() == recipientId) {
         sessions.remove(key);

libaxolotl/src/androidTest/java/org/whispersystems/test/SessionBuilderTest.java

@@ -3,7 +3,6 @@ package org.whispersystems.test;
 import android.test.AndroidTestCase;
 
 import org.whispersystems.libaxolotl.DuplicateMessageException;
-import org.whispersystems.libaxolotl.IdentityKey;
 import org.whispersystems.libaxolotl.InvalidKeyException;
 import org.whispersystems.libaxolotl.InvalidKeyIdException;
 import org.whispersystems.libaxolotl.InvalidMessageException;
@@ -15,12 +14,13 @@ import org.whispersystems.libaxolotl.StaleKeyExchangeException;
 import org.whispersystems.libaxolotl.UntrustedIdentityException;
 import org.whispersystems.libaxolotl.ecc.Curve;
 import org.whispersystems.libaxolotl.ecc.ECKeyPair;
-import org.whispersystems.libaxolotl.ecc.ECPublicKey;
 import org.whispersystems.libaxolotl.protocol.CiphertextMessage;
 import org.whispersystems.libaxolotl.protocol.KeyExchangeMessage;
 import org.whispersystems.libaxolotl.protocol.PreKeyWhisperMessage;
+import org.whispersystems.libaxolotl.state.DeviceKeyRecord;
+import org.whispersystems.libaxolotl.state.DeviceKeyStore;
 import org.whispersystems.libaxolotl.state.IdentityKeyStore;
-import org.whispersystems.libaxolotl.state.PreKey;
+import org.whispersystems.libaxolotl.state.PreKeyBundle;
 import org.whispersystems.libaxolotl.state.PreKeyRecord;
 import org.whispersystems.libaxolotl.state.PreKeyStore;
 import org.whispersystems.libaxolotl.state.SessionStore;
@@ -34,30 +34,37 @@ public class SessionBuilderTest extends AndroidTestCase {
   private static final long ALICE_RECIPIENT_ID = 5L;
   private static final long BOB_RECIPIENT_ID   = 2L;
 
-  public void testBasicPreKey()
+  public void testBasicPreKeyV2()
       throws InvalidKeyException, InvalidVersionException, InvalidMessageException, InvalidKeyIdException, DuplicateMessageException, LegacyMessageException, UntrustedIdentityException {
     SessionStore     aliceSessionStore     = new InMemorySessionStore();
+    DeviceKeyStore   aliceDeviceKeyStore   = new InMemoryDeviceKeyStore();
     PreKeyStore      alicePreKeyStore      = new InMemoryPreKeyStore();
     IdentityKeyStore aliceIdentityKeyStore = new InMemoryIdentityKeyStore();
     SessionBuilder   aliceSessionBuilder   = new SessionBuilder(aliceSessionStore, alicePreKeyStore,
+                                                                aliceDeviceKeyStore,
                                                                 aliceIdentityKeyStore,
                                                                 BOB_RECIPIENT_ID, 1);
 
     SessionStore     bobSessionStore     = new InMemorySessionStore();
     PreKeyStore      bobPreKeyStore      = new InMemoryPreKeyStore();
+    DeviceKeyStore   bobDeviceKeyStore   = new InMemoryDeviceKeyStore();
     IdentityKeyStore bobIdentityKeyStore = new InMemoryIdentityKeyStore();
     SessionBuilder   bobSessionBuilder   = new SessionBuilder(bobSessionStore, bobPreKeyStore,
+                                                              bobDeviceKeyStore,
                                                               bobIdentityKeyStore,
                                                               ALICE_RECIPIENT_ID, 1);
 
-    InMemoryPreKey bobPreKey = new InMemoryPreKey(31337, Curve.generateKeyPair(true),
+    ECKeyPair  bobPreKeyPair = Curve.generateKeyPair(true);
-                                                  bobIdentityKeyStore.getIdentityKeyPair().getPublicKey(),
+    PreKeyBundle bobPreKey = new PreKeyBundle(bobIdentityKeyStore.getLocalRegistrationId(), 1,
-                                                  bobIdentityKeyStore.getLocalRegistrationId());
+                                              31337, bobPreKeyPair.getPublicKey(),
+                                              0, null, null,
+                                              bobIdentityKeyStore.getIdentityKeyPair().getPublicKey());
 
     aliceSessionBuilder.process(bobPreKey);
 
-    assertTrue(aliceSessionStore.contains(BOB_RECIPIENT_ID, 1));
+    assertTrue(aliceSessionStore.containsSession(BOB_RECIPIENT_ID, 1));
-    assertTrue(!aliceSessionStore.load(BOB_RECIPIENT_ID, 1).getSessionState().getNeedsRefresh());
+    assertTrue(!aliceSessionStore.loadSession(BOB_RECIPIENT_ID, 1).getSessionState().getNeedsRefresh());
+    assertTrue(aliceSessionStore.loadSession(BOB_RECIPIENT_ID, 1).getSessionState().getSessionVersion() == 2);
 
     String            originalMessage    = "L'homme est condamné à être libre";
     SessionCipher     aliceSessionCipher = new SessionCipher(aliceSessionStore, BOB_RECIPIENT_ID, 1);
@@ -66,10 +73,11 @@ public class SessionBuilderTest extends AndroidTestCase {
     assertTrue(outgoingMessage.getType() == CiphertextMessage.PREKEY_TYPE);
 
     PreKeyWhisperMessage incomingMessage = new PreKeyWhisperMessage(outgoingMessage.serialize());
-    bobPreKeyStore.store(31337, bobPreKey);
+    bobPreKeyStore.storePreKey(31337, new PreKeyRecord(bobPreKey.getPreKeyId(), bobPreKeyPair));
     bobSessionBuilder.process(incomingMessage);
 
-    assertTrue(bobSessionStore.contains(ALICE_RECIPIENT_ID, 1));
+    assertTrue(bobSessionStore.containsSession(ALICE_RECIPIENT_ID, 1));
+    assertTrue(bobSessionStore.loadSession(ALICE_RECIPIENT_ID, 1).getSessionState().getSessionVersion() == 2);
 
     SessionCipher bobSessionCipher = new SessionCipher(bobSessionStore, ALICE_RECIPIENT_ID, 1);
     byte[]        plaintext        = bobSessionCipher.decrypt(incomingMessage.getWhisperMessage().serialize());
@@ -87,15 +95,17 @@ public class SessionBuilderTest extends AndroidTestCase {
     aliceSessionStore     = new InMemorySessionStore();
     aliceIdentityKeyStore = new InMemoryIdentityKeyStore();
     aliceSessionBuilder   = new SessionBuilder(aliceSessionStore, alicePreKeyStore,
+                                               aliceDeviceKeyStore,
                                                aliceIdentityKeyStore,
                                                BOB_RECIPIENT_ID, 1);
     aliceSessionCipher = new SessionCipher(aliceSessionStore, BOB_RECIPIENT_ID, 1);
 
-    bobPreKey = new InMemoryPreKey(31338, Curve.generateKeyPair(true),
+    bobPreKeyPair = Curve.generateKeyPair(true);
-                                   bobIdentityKeyStore.getIdentityKeyPair().getPublicKey(),
+    bobPreKey = new PreKeyBundle(bobIdentityKeyStore.getLocalRegistrationId(),
-                                   bobIdentityKeyStore.getLocalRegistrationId());
+                                 1, 31338, bobPreKeyPair.getPublicKey(),
+                                 0, null, null, bobIdentityKeyStore.getIdentityKeyPair().getPublicKey());
 
-    bobPreKeyStore.store(31338, bobPreKey);
+    bobPreKeyStore.storePreKey(31338, new PreKeyRecord(bobPreKey.getPreKeyId(), bobPreKeyPair));
     aliceSessionBuilder.process(bobPreKey);
 
     outgoingMessage = aliceSessionCipher.encrypt(originalMessage.getBytes());
@@ -111,9 +121,10 @@ public class SessionBuilderTest extends AndroidTestCase {
     plaintext = bobSessionCipher.decrypt(new PreKeyWhisperMessage(outgoingMessage.serialize()).getWhisperMessage().serialize());
     assertTrue(new String(plaintext).equals(originalMessage));
 
-    bobPreKey = new InMemoryPreKey(31337, Curve.generateKeyPair(true),
+    bobPreKey = new PreKeyBundle(bobIdentityKeyStore.getLocalRegistrationId(), 1,
-                                   aliceIdentityKeyStore.getIdentityKeyPair().getPublicKey(),
+                                 31337, Curve.generateKeyPair(true).getPublicKey(),
-                                   bobIdentityKeyStore.getLocalRegistrationId());
+                                 0, null, null,
+                                 aliceIdentityKeyStore.getIdentityKeyPair().getPublicKey());
 
     try {
       aliceSessionBuilder.process(bobPreKey);
@@ -123,18 +134,315 @@ public class SessionBuilderTest extends AndroidTestCase {
     }
   }
 
-  public void testBasicKeyExchange() throws InvalidKeyException, LegacyMessageException, InvalidMessageException, DuplicateMessageException, UntrustedIdentityException, StaleKeyExchangeException {
+  public void testBasicPreKeyV3()
+      throws InvalidKeyException, InvalidVersionException, InvalidMessageException, InvalidKeyIdException, DuplicateMessageException, LegacyMessageException, UntrustedIdentityException {
     SessionStore     aliceSessionStore     = new InMemorySessionStore();
+    DeviceKeyStore   aliceDeviceKeyStore   = new InMemoryDeviceKeyStore();
     PreKeyStore      alicePreKeyStore      = new InMemoryPreKeyStore();
     IdentityKeyStore aliceIdentityKeyStore = new InMemoryIdentityKeyStore();
     SessionBuilder   aliceSessionBuilder   = new SessionBuilder(aliceSessionStore, alicePreKeyStore,
+                                                                aliceDeviceKeyStore,
                                                                 aliceIdentityKeyStore,
                                                                 BOB_RECIPIENT_ID, 1);
 
     SessionStore     bobSessionStore     = new InMemorySessionStore();
     PreKeyStore      bobPreKeyStore      = new InMemoryPreKeyStore();
+    DeviceKeyStore   bobDeviceKeyStore   = new InMemoryDeviceKeyStore();
     IdentityKeyStore bobIdentityKeyStore = new InMemoryIdentityKeyStore();
     SessionBuilder   bobSessionBuilder   = new SessionBuilder(bobSessionStore, bobPreKeyStore,
+                                                              bobDeviceKeyStore,
+                                                              bobIdentityKeyStore,
+                                                              ALICE_RECIPIENT_ID, 1);
+
+    ECKeyPair bobPreKeyPair         = Curve.generateKeyPair(true);
+    ECKeyPair bobDeviceKeyPair      = Curve.generateKeyPair(true);
+    byte[]    bobDeviceKeySignature = Curve.calculateSignature(bobIdentityKeyStore.getIdentityKeyPair().getPrivateKey(),
+                                                               bobDeviceKeyPair.getPublicKey().serialize());
+
+    PreKeyBundle bobPreKey = new PreKeyBundle(bobIdentityKeyStore.getLocalRegistrationId(), 1,
+                                              31337, bobPreKeyPair.getPublicKey(),
+                                              22, bobDeviceKeyPair.getPublicKey(), bobDeviceKeySignature,
+                                              bobIdentityKeyStore.getIdentityKeyPair().getPublicKey());
+
+    aliceSessionBuilder.process(bobPreKey);
+
+    assertTrue(aliceSessionStore.containsSession(BOB_RECIPIENT_ID, 1));
+    assertTrue(!aliceSessionStore.loadSession(BOB_RECIPIENT_ID, 1).getSessionState().getNeedsRefresh());
+    assertTrue(aliceSessionStore.loadSession(BOB_RECIPIENT_ID, 1).getSessionState().getSessionVersion() == 3);
+
+    String            originalMessage    = "L'homme est condamné à être libre";
+    SessionCipher     aliceSessionCipher = new SessionCipher(aliceSessionStore, BOB_RECIPIENT_ID, 1);
+    CiphertextMessage outgoingMessage    = aliceSessionCipher.encrypt(originalMessage.getBytes());
+
+    assertTrue(outgoingMessage.getType() == CiphertextMessage.PREKEY_TYPE);
+
+    PreKeyWhisperMessage incomingMessage = new PreKeyWhisperMessage(outgoingMessage.serialize());
+    bobPreKeyStore.storePreKey(31337, new PreKeyRecord(bobPreKey.getPreKeyId(), bobPreKeyPair));
+    bobDeviceKeyStore.storeDeviceKey(22, new DeviceKeyRecord(22, System.currentTimeMillis(), bobDeviceKeyPair, bobDeviceKeySignature));
+    bobSessionBuilder.process(incomingMessage);
+
+    assertTrue(bobSessionStore.containsSession(ALICE_RECIPIENT_ID, 1));
+    assertTrue(bobSessionStore.loadSession(ALICE_RECIPIENT_ID, 1).getSessionState().getSessionVersion() == 3);
+    assertTrue(bobSessionStore.loadSession(ALICE_RECIPIENT_ID, 1).getSessionState().getAliceBaseKey() != null);
+
+    SessionCipher bobSessionCipher = new SessionCipher(bobSessionStore, ALICE_RECIPIENT_ID, 1);
+    byte[]        plaintext        = bobSessionCipher.decrypt(incomingMessage.getWhisperMessage().serialize());
+
+    assertTrue(originalMessage.equals(new String(plaintext)));
+
+    CiphertextMessage bobOutgoingMessage = bobSessionCipher.encrypt(originalMessage.getBytes());
+    assertTrue(bobOutgoingMessage.getType() == CiphertextMessage.WHISPER_TYPE);
+
+    byte[] alicePlaintext = aliceSessionCipher.decrypt(bobOutgoingMessage.serialize());
+    assertTrue(new String(alicePlaintext).equals(originalMessage));
+
+    runInteraction(aliceSessionStore, bobSessionStore);
+
+    aliceSessionStore     = new InMemorySessionStore();
+    aliceIdentityKeyStore = new InMemoryIdentityKeyStore();
+    aliceSessionBuilder   = new SessionBuilder(aliceSessionStore, alicePreKeyStore,
+                                               aliceDeviceKeyStore,
+                                               aliceIdentityKeyStore,
+                                               BOB_RECIPIENT_ID, 1);
+    aliceSessionCipher = new SessionCipher(aliceSessionStore, BOB_RECIPIENT_ID, 1);
+
+    bobPreKeyPair    = Curve.generateKeyPair(true);
+    bobDeviceKeyPair = Curve.generateKeyPair(true);
+    bobDeviceKeySignature = Curve.calculateSignature(bobIdentityKeyStore.getIdentityKeyPair().getPrivateKey(), bobDeviceKeyPair.getPublicKey().serialize());
+    bobPreKey = new PreKeyBundle(bobIdentityKeyStore.getLocalRegistrationId(),
+                                 1, 31338, bobPreKeyPair.getPublicKey(),
+                                 23, bobDeviceKeyPair.getPublicKey(), bobDeviceKeySignature,
+                                 bobIdentityKeyStore.getIdentityKeyPair().getPublicKey());
+
+    bobPreKeyStore.storePreKey(31338, new PreKeyRecord(bobPreKey.getPreKeyId(), bobPreKeyPair));
+    bobDeviceKeyStore.storeDeviceKey(23, new DeviceKeyRecord(23, System.currentTimeMillis(), bobDeviceKeyPair, bobDeviceKeySignature));
+    aliceSessionBuilder.process(bobPreKey);
+
+    outgoingMessage = aliceSessionCipher.encrypt(originalMessage.getBytes());
+
+    try {
+      bobSessionBuilder.process(new PreKeyWhisperMessage(outgoingMessage.serialize()));
+      throw new AssertionError("shouldn't be trusted!");
+    } catch (UntrustedIdentityException uie) {
+      bobIdentityKeyStore.saveIdentity(ALICE_RECIPIENT_ID, new PreKeyWhisperMessage(outgoingMessage.serialize()).getIdentityKey());
+      bobSessionBuilder.process(new PreKeyWhisperMessage(outgoingMessage.serialize()));
+    }
+
+    plaintext = bobSessionCipher.decrypt(new PreKeyWhisperMessage(outgoingMessage.serialize()).getWhisperMessage().serialize());
+    assertTrue(new String(plaintext).equals(originalMessage));
+
+    bobPreKey = new PreKeyBundle(bobIdentityKeyStore.getLocalRegistrationId(), 1,
+                                 31337, Curve.generateKeyPair(true).getPublicKey(),
+                                 23, bobDeviceKeyPair.getPublicKey(), bobDeviceKeySignature,
+                                 aliceIdentityKeyStore.getIdentityKeyPair().getPublicKey());
+
+    try {
+      aliceSessionBuilder.process(bobPreKey);
+      throw new AssertionError("shoulnd't be trusted!");
+    } catch (UntrustedIdentityException uie) {
+      // good
+    }
+  }
+
+  public void testBadDeviceKeySignature() throws InvalidKeyException, UntrustedIdentityException {
+    SessionStore     aliceSessionStore     = new InMemorySessionStore();
+    DeviceKeyStore   aliceDeviceKeyStore   = new InMemoryDeviceKeyStore();
+    PreKeyStore      alicePreKeyStore      = new InMemoryPreKeyStore();
+    IdentityKeyStore aliceIdentityKeyStore = new InMemoryIdentityKeyStore();
+    SessionBuilder   aliceSessionBuilder   = new SessionBuilder(aliceSessionStore, alicePreKeyStore,
+                                                                aliceDeviceKeyStore,
+                                                                aliceIdentityKeyStore,
+                                                                BOB_RECIPIENT_ID, 1);
+
+    IdentityKeyStore bobIdentityKeyStore = new InMemoryIdentityKeyStore();
+
+    ECKeyPair bobPreKeyPair         = Curve.generateKeyPair(true);
+    ECKeyPair bobDeviceKeyPair      = Curve.generateKeyPair(true);
+    byte[]    bobDeviceKeySignature = Curve.calculateSignature(bobIdentityKeyStore.getIdentityKeyPair().getPrivateKey(),
+                                                               bobDeviceKeyPair.getPublicKey().serialize());
+
+
+    for (int i=0;i<bobDeviceKeySignature.length * 8;i++) {
+      byte[] modifiedSignature = new byte[bobDeviceKeySignature.length];
+      System.arraycopy(bobDeviceKeySignature, 0, modifiedSignature, 0, modifiedSignature.length);
+
+      modifiedSignature[i/8] ^= (0x01 << (i % 8));
+
+      PreKeyBundle bobPreKey = new PreKeyBundle(bobIdentityKeyStore.getLocalRegistrationId(), 1,
+                                                31337, bobPreKeyPair.getPublicKey(),
+                                                22, bobDeviceKeyPair.getPublicKey(), modifiedSignature,
+                                                bobIdentityKeyStore.getIdentityKeyPair().getPublicKey());
+
+      try {
+        aliceSessionBuilder.process(bobPreKey);
+        throw new AssertionError("Accepted modified device key signature!");
+      } catch (InvalidKeyException ike) {
+        // good
+      }
+    }
+
+    PreKeyBundle bobPreKey = new PreKeyBundle(bobIdentityKeyStore.getLocalRegistrationId(), 1,
+                                              31337, bobPreKeyPair.getPublicKey(),
+                                              22, bobDeviceKeyPair.getPublicKey(), bobDeviceKeySignature,
+                                              bobIdentityKeyStore.getIdentityKeyPair().getPublicKey());
+
+    aliceSessionBuilder.process(bobPreKey);
+  }
+
+  public void testRepeatBundleMessageV2() throws InvalidKeyException, UntrustedIdentityException, InvalidVersionException, InvalidMessageException, InvalidKeyIdException, DuplicateMessageException, LegacyMessageException {
+    SessionStore     aliceSessionStore     = new InMemorySessionStore();
+    DeviceKeyStore   aliceDeviceKeyStore   = new InMemoryDeviceKeyStore();
+    PreKeyStore      alicePreKeyStore      = new InMemoryPreKeyStore();
+    IdentityKeyStore aliceIdentityKeyStore = new InMemoryIdentityKeyStore();
+    SessionBuilder   aliceSessionBuilder   = new SessionBuilder(aliceSessionStore, alicePreKeyStore,
+                                                                aliceDeviceKeyStore,
+                                                                aliceIdentityKeyStore,
+                                                                BOB_RECIPIENT_ID, 1);
+
+    SessionStore     bobSessionStore     = new InMemorySessionStore();
+    PreKeyStore      bobPreKeyStore      = new InMemoryPreKeyStore();
+    DeviceKeyStore   bobDeviceKeyStore   = new InMemoryDeviceKeyStore();
+    IdentityKeyStore bobIdentityKeyStore = new InMemoryIdentityKeyStore();
+    SessionBuilder   bobSessionBuilder   = new SessionBuilder(bobSessionStore, bobPreKeyStore,
+                                                              bobDeviceKeyStore,
+                                                              bobIdentityKeyStore,
+                                                              ALICE_RECIPIENT_ID, 1);
+
+    ECKeyPair bobPreKeyPair         = Curve.generateKeyPair(true);
+    ECKeyPair bobDeviceKeyPair      = Curve.generateKeyPair(true);
+    byte[]    bobDeviceKeySignature = Curve.calculateSignature(bobIdentityKeyStore.getIdentityKeyPair().getPrivateKey(),
+                                                               bobDeviceKeyPair.getPublicKey().serialize());
+
+    PreKeyBundle bobPreKey = new PreKeyBundle(bobIdentityKeyStore.getLocalRegistrationId(), 1,
+                                              31337, bobPreKeyPair.getPublicKey(),
+                                              0, null, null,
+                                              bobIdentityKeyStore.getIdentityKeyPair().getPublicKey());
+
+    bobPreKeyStore.storePreKey(31337, new PreKeyRecord(bobPreKey.getPreKeyId(), bobPreKeyPair));
+    bobDeviceKeyStore.storeDeviceKey(22, new DeviceKeyRecord(22, System.currentTimeMillis(), bobDeviceKeyPair, bobDeviceKeySignature));
+
+    aliceSessionBuilder.process(bobPreKey);
+
+    String            originalMessage    = "L'homme est condamné à être libre";
+    SessionCipher     aliceSessionCipher = new SessionCipher(aliceSessionStore, BOB_RECIPIENT_ID, 1);
+    CiphertextMessage outgoingMessageOne = aliceSessionCipher.encrypt(originalMessage.getBytes());
+    CiphertextMessage outgoingMessageTwo = aliceSessionCipher.encrypt(originalMessage.getBytes());
+
+    assertTrue(outgoingMessageOne.getType() == CiphertextMessage.PREKEY_TYPE);
+
+    PreKeyWhisperMessage incomingMessage = new PreKeyWhisperMessage(outgoingMessageOne.serialize());
+    bobSessionBuilder.process(incomingMessage);
+
+    SessionCipher bobSessionCipher = new SessionCipher(bobSessionStore, ALICE_RECIPIENT_ID, 1);
+
+    byte[]        plaintext        = bobSessionCipher.decrypt(incomingMessage.getWhisperMessage().serialize());
+    assertTrue(originalMessage.equals(new String(plaintext)));
+
+    CiphertextMessage bobOutgoingMessage = bobSessionCipher.encrypt(originalMessage.getBytes());
+
+    byte[] alicePlaintext = aliceSessionCipher.decrypt(bobOutgoingMessage.serialize());
+    assertTrue(originalMessage.equals(new String(alicePlaintext)));
+
+    // The test
+
+    PreKeyWhisperMessage incomingMessageTwo = new PreKeyWhisperMessage(outgoingMessageTwo.serialize());
+    bobSessionBuilder.process(incomingMessageTwo);
+
+    plaintext = bobSessionCipher.decrypt(incomingMessageTwo.getWhisperMessage().serialize());
+    assertTrue(originalMessage.equals(new String(plaintext)));
+
+    bobOutgoingMessage = bobSessionCipher.encrypt(originalMessage.getBytes());
+    alicePlaintext = aliceSessionCipher.decrypt(bobOutgoingMessage.serialize());
+    assertTrue(originalMessage.equals(new String(alicePlaintext)));
+
+  }
+
+
+  public void testRepeatBundleMessageV3() throws InvalidKeyException, UntrustedIdentityException, InvalidVersionException, InvalidMessageException, InvalidKeyIdException, DuplicateMessageException, LegacyMessageException {
+    SessionStore     aliceSessionStore     = new InMemorySessionStore();
+    DeviceKeyStore   aliceDeviceKeyStore   = new InMemoryDeviceKeyStore();
+    PreKeyStore      alicePreKeyStore      = new InMemoryPreKeyStore();
+    IdentityKeyStore aliceIdentityKeyStore = new InMemoryIdentityKeyStore();
+    SessionBuilder   aliceSessionBuilder   = new SessionBuilder(aliceSessionStore, alicePreKeyStore,
+                                                                aliceDeviceKeyStore,
+                                                                aliceIdentityKeyStore,
+                                                                BOB_RECIPIENT_ID, 1);
+
+    SessionStore     bobSessionStore     = new InMemorySessionStore();
+    PreKeyStore      bobPreKeyStore      = new InMemoryPreKeyStore();
+    DeviceKeyStore   bobDeviceKeyStore   = new InMemoryDeviceKeyStore();
+    IdentityKeyStore bobIdentityKeyStore = new InMemoryIdentityKeyStore();
+    SessionBuilder   bobSessionBuilder   = new SessionBuilder(bobSessionStore, bobPreKeyStore,
+                                                              bobDeviceKeyStore,
+                                                              bobIdentityKeyStore,
+                                                              ALICE_RECIPIENT_ID, 1);
+
+    ECKeyPair bobPreKeyPair         = Curve.generateKeyPair(true);
+    ECKeyPair bobDeviceKeyPair      = Curve.generateKeyPair(true);
+    byte[]    bobDeviceKeySignature = Curve.calculateSignature(bobIdentityKeyStore.getIdentityKeyPair().getPrivateKey(),
+                                                               bobDeviceKeyPair.getPublicKey().serialize());
+
+    PreKeyBundle bobPreKey = new PreKeyBundle(bobIdentityKeyStore.getLocalRegistrationId(), 1,
+                                              31337, bobPreKeyPair.getPublicKey(),
+                                              22, bobDeviceKeyPair.getPublicKey(), bobDeviceKeySignature,
+                                              bobIdentityKeyStore.getIdentityKeyPair().getPublicKey());
+
+    bobPreKeyStore.storePreKey(31337, new PreKeyRecord(bobPreKey.getPreKeyId(), bobPreKeyPair));
+    bobDeviceKeyStore.storeDeviceKey(22, new DeviceKeyRecord(22, System.currentTimeMillis(), bobDeviceKeyPair, bobDeviceKeySignature));
+
+    aliceSessionBuilder.process(bobPreKey);
+
+    String            originalMessage    = "L'homme est condamné à être libre";
+    SessionCipher     aliceSessionCipher = new SessionCipher(aliceSessionStore, BOB_RECIPIENT_ID, 1);
+    CiphertextMessage outgoingMessageOne = aliceSessionCipher.encrypt(originalMessage.getBytes());
+    CiphertextMessage outgoingMessageTwo = aliceSessionCipher.encrypt(originalMessage.getBytes());
+
+    assertTrue(outgoingMessageOne.getType() == CiphertextMessage.PREKEY_TYPE);
+
+    PreKeyWhisperMessage incomingMessage = new PreKeyWhisperMessage(outgoingMessageOne.serialize());
+    bobSessionBuilder.process(incomingMessage);
+
+    SessionCipher bobSessionCipher = new SessionCipher(bobSessionStore, ALICE_RECIPIENT_ID, 1);
+
+    byte[]        plaintext        = bobSessionCipher.decrypt(incomingMessage.getWhisperMessage().serialize());
+    assertTrue(originalMessage.equals(new String(plaintext)));
+
+    CiphertextMessage bobOutgoingMessage = bobSessionCipher.encrypt(originalMessage.getBytes());
+
+    byte[] alicePlaintext = aliceSessionCipher.decrypt(bobOutgoingMessage.serialize());
+    assertTrue(originalMessage.equals(new String(alicePlaintext)));
+
+    // The test
+
+    PreKeyWhisperMessage incomingMessageTwo = new PreKeyWhisperMessage(outgoingMessageTwo.serialize());
+    bobSessionBuilder.process(incomingMessageTwo);
+
+    plaintext = bobSessionCipher.decrypt(incomingMessageTwo.getWhisperMessage().serialize());
+    assertTrue(originalMessage.equals(new String(plaintext)));
+
+    bobOutgoingMessage = bobSessionCipher.encrypt(originalMessage.getBytes());
+    alicePlaintext = aliceSessionCipher.decrypt(bobOutgoingMessage.serialize());
+    assertTrue(originalMessage.equals(new String(alicePlaintext)));
+
+  }
+
+
+  public void testBasicKeyExchange() throws InvalidKeyException, LegacyMessageException, InvalidMessageException, DuplicateMessageException, UntrustedIdentityException, StaleKeyExchangeException {
+    SessionStore     aliceSessionStore     = new InMemorySessionStore();
+    PreKeyStore      alicePreKeyStore      = new InMemoryPreKeyStore();
+    DeviceKeyStore   aliceDeviceKeyStore   = new InMemoryDeviceKeyStore();
+    IdentityKeyStore aliceIdentityKeyStore = new InMemoryIdentityKeyStore();
+    SessionBuilder   aliceSessionBuilder   = new SessionBuilder(aliceSessionStore, alicePreKeyStore,
+                                                                aliceDeviceKeyStore,
+                                                                aliceIdentityKeyStore,
+                                                                BOB_RECIPIENT_ID, 1);
+
+    SessionStore     bobSessionStore     = new InMemorySessionStore();
+    PreKeyStore      bobPreKeyStore      = new InMemoryPreKeyStore();
+    DeviceKeyStore   bobDeviceKeyStore   = new InMemoryDeviceKeyStore();
+    IdentityKeyStore bobIdentityKeyStore = new InMemoryIdentityKeyStore();
+    SessionBuilder   bobSessionBuilder   = new SessionBuilder(bobSessionStore, bobPreKeyStore,
+                                                              bobDeviceKeyStore,
                                                               bobIdentityKeyStore,
                                                               ALICE_RECIPIENT_ID, 1);
 
@@ -147,14 +455,15 @@ public class SessionBuilderTest extends AndroidTestCase {
     KeyExchangeMessage response = aliceSessionBuilder.process(bobKeyExchangeMessage);
 
     assertTrue(response == null);
-    assertTrue(aliceSessionStore.contains(BOB_RECIPIENT_ID, 1));
+    assertTrue(aliceSessionStore.containsSession(BOB_RECIPIENT_ID, 1));
-    assertTrue(bobSessionStore.contains(ALICE_RECIPIENT_ID, 1));
+    assertTrue(bobSessionStore.containsSession(ALICE_RECIPIENT_ID, 1));
 
     runInteraction(aliceSessionStore, bobSessionStore);
 
     aliceSessionStore       = new InMemorySessionStore();
     aliceIdentityKeyStore   = new InMemoryIdentityKeyStore();
     aliceSessionBuilder     = new SessionBuilder(aliceSessionStore, alicePreKeyStore,
+                                                 aliceDeviceKeyStore,
                                                  aliceIdentityKeyStore, BOB_RECIPIENT_ID, 1);
     aliceKeyExchangeMessage = aliceSessionBuilder.process();
 
@@ -175,15 +484,19 @@ public class SessionBuilderTest extends AndroidTestCase {
       throws InvalidKeyException, DuplicateMessageException, LegacyMessageException, InvalidMessageException, UntrustedIdentityException, StaleKeyExchangeException {
     SessionStore     aliceSessionStore     = new InMemorySessionStore();
     PreKeyStore      alicePreKeyStore      = new InMemoryPreKeyStore();
+    DeviceKeyStore   aliceDeviceKeyStore   = new InMemoryDeviceKeyStore();
     IdentityKeyStore aliceIdentityKeyStore = new InMemoryIdentityKeyStore();
     SessionBuilder   aliceSessionBuilder   = new SessionBuilder(aliceSessionStore, alicePreKeyStore,
+                                                                aliceDeviceKeyStore,
                                                                 aliceIdentityKeyStore,
                                                                 BOB_RECIPIENT_ID, 1);
 
     SessionStore     bobSessionStore     = new InMemorySessionStore();
     PreKeyStore      bobPreKeyStore      = new InMemoryPreKeyStore();
+    DeviceKeyStore   bobDeviceKeyStore   = new InMemoryDeviceKeyStore();
     IdentityKeyStore bobIdentityKeyStore = new InMemoryIdentityKeyStore();
     SessionBuilder   bobSessionBuilder   = new SessionBuilder(bobSessionStore, bobPreKeyStore,
+                                                              bobDeviceKeyStore,
                                                               bobIdentityKeyStore,
                                                               ALICE_RECIPIENT_ID, 1);
 
@@ -284,41 +597,5 @@ public class SessionBuilderTest extends AndroidTestCase {
     }
   }
 
-  private class InMemoryPreKey extends PreKeyRecord implements PreKey {
-
-    private final IdentityKey identityKey;
-    private final int         registrationId;
-
-    public InMemoryPreKey(int keyId, ECKeyPair keyPair, IdentityKey identityKey, int registrationId) {
-      super(keyId, keyPair);
-      this.identityKey    = identityKey;
-      this.registrationId = registrationId;
-    }
-
-    @Override
-    public int getDeviceId() {
-      return 1;
-    }
-
-    @Override
-    public int getKeyId() {
-      return getId();
-    }
-
-    @Override
-    public ECPublicKey getPublicKey() {
-      return getKeyPair().getPublicKey();
-    }
-
-    @Override
-    public IdentityKey getIdentityKey() {
-      return identityKey;
-    }
-
-    @Override
-    public int getRegistrationId() {
-      return registrationId;
-    }
-  }
 
 }

libaxolotl/src/androidTest/java/org/whispersystems/test/SessionCipherTest.java

@@ -9,32 +9,53 @@ import org.whispersystems.libaxolotl.InvalidKeyException;
 import org.whispersystems.libaxolotl.InvalidMessageException;
 import org.whispersystems.libaxolotl.LegacyMessageException;
 import org.whispersystems.libaxolotl.SessionCipher;
-import org.whispersystems.libaxolotl.state.SessionRecord;
-import org.whispersystems.libaxolotl.state.SessionState;
-import org.whispersystems.libaxolotl.state.SessionStore;
 import org.whispersystems.libaxolotl.ecc.Curve;
 import org.whispersystems.libaxolotl.ecc.ECKeyPair;
 import org.whispersystems.libaxolotl.protocol.CiphertextMessage;
-import org.whispersystems.libaxolotl.ratchet.RatchetingSession;
+import org.whispersystems.libaxolotl.ratchet.RatchetingSessionV2;
+import org.whispersystems.libaxolotl.ratchet.RatchetingSessionV3;
+import org.whispersystems.libaxolotl.state.SessionRecord;
+import org.whispersystems.libaxolotl.state.SessionState;
+import org.whispersystems.libaxolotl.state.SessionStore;
 
+import java.security.NoSuchAlgorithmException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Random;
 
 public class SessionCipherTest extends AndroidTestCase {
 
-  public void testBasicSession()
+  public void testBasicSessionV2()
       throws InvalidKeyException, DuplicateMessageException,
-      LegacyMessageException, InvalidMessageException
+      LegacyMessageException, InvalidMessageException, NoSuchAlgorithmException
   {
     SessionRecord aliceSessionRecord = new SessionRecord();
     SessionRecord bobSessionRecord   = new SessionRecord();
 
-    initializeSessions(aliceSessionRecord.getSessionState(), bobSessionRecord.getSessionState());
+    initializeSessionsV2(aliceSessionRecord.getSessionState(), bobSessionRecord.getSessionState());
+    runInteraction(aliceSessionRecord, bobSessionRecord);
+  }
 
+  public void testBasicSessionV3()
+      throws InvalidKeyException, DuplicateMessageException,
+      LegacyMessageException, InvalidMessageException, NoSuchAlgorithmException
+  {
+    SessionRecord aliceSessionRecord = new SessionRecord();
+    SessionRecord bobSessionRecord   = new SessionRecord();
+
+    initializeSessionsV3(aliceSessionRecord.getSessionState(), bobSessionRecord.getSessionState());
+    runInteraction(aliceSessionRecord, bobSessionRecord);
+  }
+
+  private void runInteraction(SessionRecord aliceSessionRecord, SessionRecord bobSessionRecord)
+      throws DuplicateMessageException, LegacyMessageException, InvalidMessageException, NoSuchAlgorithmException {
     SessionStore aliceSessionStore = new InMemorySessionStore();
     SessionStore bobSessionStore   = new InMemorySessionStore();
 
-    aliceSessionStore.store(2L, 1, aliceSessionRecord);
+    aliceSessionStore.storeSession(2L, 1, aliceSessionRecord);
-    bobSessionStore.store(3L, 1, bobSessionRecord);
+    bobSessionStore.storeSession(3L, 1, bobSessionRecord);
 
     SessionCipher     aliceCipher    = new SessionCipher(aliceSessionStore, 2L, 1);
     SessionCipher     bobCipher      = new SessionCipher(bobSessionStore, 3L, 1);
@@ -50,10 +71,56 @@ public class SessionCipherTest extends AndroidTestCase {
     byte[]            receivedReply = aliceCipher.decrypt(reply.serialize());
 
     assertTrue(Arrays.equals(bobReply, receivedReply));
+
+    List<CiphertextMessage> aliceCiphertextMessages = new ArrayList<>();
+    List<byte[]>            alicePlaintextMessages  = new ArrayList<>();
+
+    for (int i=0;i<50;i++) {
+      alicePlaintextMessages.add(("смерть за смерть " + i).getBytes());
+      aliceCiphertextMessages.add(aliceCipher.encrypt(("смерть за смерть " + i).getBytes()));
+    }
+
+    long seed = System.currentTimeMillis();
+
+    Collections.shuffle(aliceCiphertextMessages, new Random(seed));
+    Collections.shuffle(alicePlaintextMessages, new Random(seed));
+
+    for (int i=0;i<aliceCiphertextMessages.size() / 2;i++) {
+      byte[] receivedPlaintext = bobCipher.decrypt(aliceCiphertextMessages.get(i).serialize());
+      assertTrue(Arrays.equals(receivedPlaintext, alicePlaintextMessages.get(i)));
+    }
+
+    List<CiphertextMessage> bobCiphertextMessages = new ArrayList<>();
+    List<byte[]>            bobPlaintextMessages  = new ArrayList<>();
+
+    for (int i=0;i<20;i++) {
+      bobPlaintextMessages.add(("смерть за смерть " + i).getBytes());
+      bobCiphertextMessages.add(bobCipher.encrypt(("смерть за смерть " + i).getBytes()));
+    }
+
+    seed = System.currentTimeMillis();
+
+    Collections.shuffle(bobCiphertextMessages, new Random(seed));
+    Collections.shuffle(bobPlaintextMessages, new Random(seed));
+
+    for (int i=0;i<bobCiphertextMessages.size() / 2;i++) {
+      byte[] receivedPlaintext = aliceCipher.decrypt(bobCiphertextMessages.get(i).serialize());
+      assertTrue(Arrays.equals(receivedPlaintext, bobPlaintextMessages.get(i)));
+    }
+
+    for (int i=aliceCiphertextMessages.size()/2;i<aliceCiphertextMessages.size();i++) {
+      byte[] receivedPlaintext = bobCipher.decrypt(aliceCiphertextMessages.get(i).serialize());
+      assertTrue(Arrays.equals(receivedPlaintext, alicePlaintextMessages.get(i)));
+    }
+
+    for (int i=bobCiphertextMessages.size() / 2;i<bobCiphertextMessages.size();i++) {
+      byte[] receivedPlaintext = aliceCipher.decrypt(bobCiphertextMessages.get(i).serialize());
+      assertTrue(Arrays.equals(receivedPlaintext, bobPlaintextMessages.get(i)));
+    }
   }
 
 
-  private void initializeSessions(SessionState aliceSessionState, SessionState bobSessionState)
+  private void initializeSessionsV2(SessionState aliceSessionState, SessionState bobSessionState)
       throws InvalidKeyException
   {
     ECKeyPair       aliceIdentityKeyPair = Curve.generateKeyPair(false);
@@ -69,12 +136,44 @@ public class SessionCipherTest extends AndroidTestCase {
     ECKeyPair       bobEphemeralKey      = bobBaseKey;
 
 
-    RatchetingSession.initializeSession(aliceSessionState, aliceBaseKey, bobBaseKey.getPublicKey(),
+    RatchetingSessionV2.initializeSession(aliceSessionState, aliceBaseKey, bobBaseKey.getPublicKey(),
-                                        aliceEphemeralKey, bobEphemeralKey.getPublicKey(),
+                                          aliceEphemeralKey, bobEphemeralKey.getPublicKey(),
-                                        aliceIdentityKey, bobIdentityKey.getPublicKey());
+                                          aliceIdentityKey, bobIdentityKey.getPublicKey());
+
+    RatchetingSessionV2.initializeSession(bobSessionState, bobBaseKey, aliceBaseKey.getPublicKey(),
+                                          bobEphemeralKey, aliceEphemeralKey.getPublicKey(),
+                                          bobIdentityKey, aliceIdentityKey.getPublicKey());
+  }
+
+  private void initializeSessionsV3(SessionState aliceSessionState, SessionState bobSessionState)
+      throws InvalidKeyException
+  {
+    ECKeyPair       aliceIdentityKeyPair = Curve.generateKeyPair(false);
+    IdentityKeyPair aliceIdentityKey     = new IdentityKeyPair(new IdentityKey(aliceIdentityKeyPair.getPublicKey()),
+                                                               aliceIdentityKeyPair.getPrivateKey());
+    ECKeyPair       aliceBaseKey         = Curve.generateKeyPair(true);
+    ECKeyPair       aliceEphemeralKey    = Curve.generateKeyPair(true);
+
+    ECKeyPair       alicePreKey          = aliceBaseKey;
+
+    ECKeyPair       bobIdentityKeyPair   = Curve.generateKeyPair(false);
+    IdentityKeyPair bobIdentityKey       = new IdentityKeyPair(new IdentityKey(bobIdentityKeyPair.getPublicKey()),
+                                                               bobIdentityKeyPair.getPrivateKey());
+    ECKeyPair       bobBaseKey           = Curve.generateKeyPair(true);
+    ECKeyPair       bobEphemeralKey      = bobBaseKey;
+
+    ECKeyPair       bobPreKey            = Curve.generateKeyPair(true);
+
+
+    RatchetingSessionV3.initializeSession(aliceSessionState, aliceBaseKey, bobBaseKey.getPublicKey(),
+                                          aliceEphemeralKey, bobEphemeralKey.getPublicKey(),
+                                          alicePreKey, bobPreKey.getPublicKey(),
+                                          aliceIdentityKey, bobIdentityKey.getPublicKey());
+
+    RatchetingSessionV3.initializeSession(bobSessionState, bobBaseKey, aliceBaseKey.getPublicKey(),
+                                          bobEphemeralKey, aliceEphemeralKey.getPublicKey(),
+                                          bobPreKey, alicePreKey.getPublicKey(),
+                                          bobIdentityKey, aliceIdentityKey.getPublicKey());
 
-    RatchetingSession.initializeSession(bobSessionState, bobBaseKey, aliceBaseKey.getPublicKey(),
-                                        bobEphemeralKey, aliceEphemeralKey.getPublicKey(),
-                                        bobIdentityKey, aliceIdentityKey.getPublicKey());
   }
 }

libaxolotl/src/androidTest/java/org/whispersystems/test/ratchet/RatchetingSessionV2Test.java

@@ -10,11 +10,11 @@ import org.whispersystems.libaxolotl.ecc.Curve;
 import org.whispersystems.libaxolotl.ecc.ECKeyPair;
 import org.whispersystems.libaxolotl.ecc.ECPrivateKey;
 import org.whispersystems.libaxolotl.ecc.ECPublicKey;
-import org.whispersystems.libaxolotl.ratchet.RatchetingSession;
+import org.whispersystems.libaxolotl.ratchet.RatchetingSessionV2;
 
 import java.util.Arrays;
 
-public class RatchetingSessionTest extends AndroidTestCase {
+public class RatchetingSessionV2Test extends AndroidTestCase {
 
   public void testRatchetingSessionAsBob() throws InvalidKeyException {
     byte[] bobPublic             = {(byte) 0x05, (byte) 0x2c, (byte) 0xb4, (byte) 0x97,
@@ -107,9 +107,9 @@ public class RatchetingSessionTest extends AndroidTestCase {
 
     SessionState session = new SessionState();
 
-    RatchetingSession.initializeSession(session, bobBaseKey, aliceBasePublicKey,
+    RatchetingSessionV2.initializeSession(session, bobBaseKey, aliceBasePublicKey,
-                                        bobEphemeralKey, aliceEphemeralPublicKey,
+                                          bobEphemeralKey, aliceEphemeralPublicKey,
-                                        bobIdentityKey, aliceIdentityPublicKey);
+                                          bobIdentityKey, aliceIdentityPublicKey);
 
     assertTrue(session.getLocalIdentityKey().equals(bobIdentityKey.getPublicKey()));
     assertTrue(session.getRemoteIdentityKey().equals(aliceIdentityPublicKey));
@@ -204,9 +204,9 @@ public class RatchetingSessionTest extends AndroidTestCase {
 
     SessionState session = new SessionState();
 
-    RatchetingSession.initializeSession(session, aliceBaseKey, bobBasePublicKey,
+    RatchetingSessionV2.initializeSession(session, aliceBaseKey, bobBasePublicKey,
-                                        aliceEphemeralKey, bobEphemeralPublicKey,
+                                          aliceEphemeralKey, bobEphemeralPublicKey,
-                                        aliceIdentityKey, bobIdentityKey);
+                                          aliceIdentityKey, bobIdentityKey);
 
     assertTrue(session.getLocalIdentityKey().equals(aliceIdentityKey.getPublicKey()));
     assertTrue(session.getRemoteIdentityKey().equals(bobIdentityKey));

libaxolotl/src/main/java/org/whispersystems/libaxolotl/SessionBuilder.java

@@ -7,9 +7,12 @@ import org.whispersystems.libaxolotl.ecc.ECKeyPair;
 import org.whispersystems.libaxolotl.ecc.ECPublicKey;
 import org.whispersystems.libaxolotl.protocol.KeyExchangeMessage;
 import org.whispersystems.libaxolotl.protocol.PreKeyWhisperMessage;
-import org.whispersystems.libaxolotl.ratchet.RatchetingSession;
+import org.whispersystems.libaxolotl.ratchet.RatchetingSessionV2;
+import org.whispersystems.libaxolotl.ratchet.RatchetingSessionV3;
+import org.whispersystems.libaxolotl.state.DeviceKeyRecord;
+import org.whispersystems.libaxolotl.state.DeviceKeyStore;
 import org.whispersystems.libaxolotl.state.IdentityKeyStore;
-import org.whispersystems.libaxolotl.state.PreKey;
+import org.whispersystems.libaxolotl.state.PreKeyBundle;
 import org.whispersystems.libaxolotl.state.PreKeyRecord;
 import org.whispersystems.libaxolotl.state.PreKeyStore;
 import org.whispersystems.libaxolotl.state.SessionRecord;
@@ -18,6 +21,8 @@ import org.whispersystems.libaxolotl.state.SessionStore;
 import org.whispersystems.libaxolotl.util.KeyHelper;
 import org.whispersystems.libaxolotl.util.Medium;
 
+import java.util.Arrays;
+
 /**
  * SessionBuilder is responsible for setting up encrypted sessions.
  * Once a session has been established, {@link org.whispersystems.libaxolotl.SessionCipher}
@@ -25,7 +30,7 @@ import org.whispersystems.libaxolotl.util.Medium;
  * <p>
  * Sessions are built from one of three different possible vectors:
  * <ol>
- *   <li>A {@link org.whispersystems.libaxolotl.state.PreKey} retrieved from a server.</li>
+ *   <li>A {@link org.whispersystems.libaxolotl.state.PreKeyBundle} retrieved from a server.</li>
  *   <li>A {@link org.whispersystems.libaxolotl.protocol.PreKeyWhisperMessage} received from a client.</li>
  *   <li>A {@link org.whispersystems.libaxolotl.protocol.KeyExchangeMessage} sent to or received from a client.</li>
  * </ol>
@@ -41,6 +46,7 @@ public class SessionBuilder {
 
   private final SessionStore     sessionStore;
   private final PreKeyStore      preKeyStore;
+  private final DeviceKeyStore   deviceKeyStore;
   private final IdentityKeyStore identityKeyStore;
   private final long             recipientId;
   private final int              deviceId;
@@ -56,11 +62,13 @@ public class SessionBuilder {
    */
   public SessionBuilder(SessionStore sessionStore,
                         PreKeyStore preKeyStore,
+                        DeviceKeyStore deviceKeyStore,
                         IdentityKeyStore identityKeyStore,
                         long recipientId, int deviceId)
   {
     this.sessionStore     = sessionStore;
     this.preKeyStore      = preKeyStore;
+    this.deviceKeyStore   = deviceKeyStore;
     this.identityKeyStore = identityKeyStore;
     this.recipientId      = recipientId;
     this.deviceId         = deviceId;
@@ -83,27 +91,93 @@ public class SessionBuilder {
   public void process(PreKeyWhisperMessage message)
       throws InvalidKeyIdException, InvalidKeyException, UntrustedIdentityException
   {
-    int         preKeyId          = message.getPreKeyId();
+    int         messageVersion   = message.getMessageVersion();
-    ECPublicKey theirBaseKey      = message.getBaseKey();
+    IdentityKey theirIdentityKey = message.getIdentityKey();
-    ECPublicKey theirEphemeralKey = message.getWhisperMessage().getSenderEphemeral();
-    IdentityKey theirIdentityKey  = message.getIdentityKey();
 
     if (!identityKeyStore.isTrustedIdentity(recipientId, theirIdentityKey)) {
       throw new UntrustedIdentityException();
     }
 
-    if (!preKeyStore.contains(preKeyId) &&
+    if      (messageVersion == 2) processV2(message);
-        sessionStore.contains(recipientId, deviceId))
+    else if (messageVersion == 3) processV3(message);
-    {
+    else                          throw new AssertionError("Unknown version: " + messageVersion);
-      Log.w(TAG, "We've already processed the prekey part, letting bundled message fall through...");
+
+    identityKeyStore.saveIdentity(recipientId, theirIdentityKey);
+  }
+
+  private void processV3(PreKeyWhisperMessage message)
+      throws UntrustedIdentityException, InvalidKeyIdException, InvalidKeyException
+  {
+    SessionRecord sessionRecord     = sessionStore.loadSession(recipientId, deviceId);
+    int           preKeyId          = message.getPreKeyId();
+    int           deviceKeyId       = message.getDeviceKeyId();
+    ECPublicKey   theirBaseKey      = message.getBaseKey();
+    ECPublicKey   theirEphemeralKey = message.getWhisperMessage().getSenderEphemeral();
+    IdentityKey   theirIdentityKey  = message.getIdentityKey();
+
+    if (sessionRecord.hasSessionState(message.getMessageVersion(), theirBaseKey.serialize())) {
+      Log.w(TAG, "We've already setup a session for this V3 message, letting bundled message fall through...");
       return;
     }
 
-    if (!preKeyStore.contains(preKeyId))
+    if (preKeyId >=0 && !preKeyStore.containsPreKey(preKeyId))
       throw new InvalidKeyIdException("No such prekey: " + preKeyId);
 
-    SessionRecord   sessionRecord        = sessionStore.load(recipientId, deviceId);
+    if (!deviceKeyStore.containsDeviceKey(deviceKeyId))
-    PreKeyRecord    preKeyRecord         = preKeyStore.load(preKeyId);
+      throw new InvalidKeyIdException("No such device key: " + deviceKeyId);
+
+    PreKeyRecord    preKeyRecord         = preKeyId >= 0 ? preKeyStore.loadPreKey(preKeyId) : null;
+    DeviceKeyRecord deviceKeyRecord      = deviceKeyStore.loadDeviceKey(deviceKeyId);
+    ECKeyPair       ourPreKey            = preKeyRecord != null ? preKeyRecord.getKeyPair() : null;
+    ECPublicKey     theirPreKey          = theirBaseKey;
+    ECKeyPair       ourBaseKey           = deviceKeyRecord.getKeyPair();
+    ECKeyPair       ourEphemeralKey      = ourBaseKey;
+    IdentityKeyPair ourIdentityKey       = identityKeyStore.getIdentityKeyPair();
+    boolean         simultaneousInitiate = sessionRecord.getSessionState().hasPendingPreKey();
+
+    if (!simultaneousInitiate) sessionRecord.reset();
+    else                       sessionRecord.archiveCurrentState();
+
+    RatchetingSessionV3.initializeSession(sessionRecord.getSessionState(),
+                                          ourBaseKey, theirBaseKey,
+                                          ourEphemeralKey, theirEphemeralKey,
+                                          ourPreKey, theirPreKey,
+                                          ourIdentityKey, theirIdentityKey);
+
+    sessionRecord.getSessionState().setLocalRegistrationId(identityKeyStore.getLocalRegistrationId());
+    sessionRecord.getSessionState().setRemoteRegistrationId(message.getRegistrationId());
+    sessionRecord.getSessionState().setAliceBaseKey(theirBaseKey.serialize());
+    sessionRecord.getSessionState().setSessionVersion(message.getMessageVersion());
+
+    if (simultaneousInitiate) sessionRecord.getSessionState().setNeedsRefresh(true);
+
+    sessionStore.storeSession(recipientId, deviceId, sessionRecord);
+
+    if (preKeyId >= 0 && preKeyId != Medium.MAX_VALUE) {
+      preKeyStore.removePreKey(preKeyId);
+    }
+  }
+
+  private void processV2(PreKeyWhisperMessage message)
+      throws UntrustedIdentityException, InvalidKeyIdException, InvalidKeyException
+  {
+    int           preKeyId          = message.getPreKeyId();
+    ECPublicKey   theirBaseKey      = message.getBaseKey();
+    ECPublicKey   theirEphemeralKey = message.getWhisperMessage().getSenderEphemeral();
+    IdentityKey   theirIdentityKey  = message.getIdentityKey();
+
+    if (!preKeyStore.containsPreKey(preKeyId) &&
+        sessionStore.containsSession(recipientId, deviceId))
+    {
+      Log.w(TAG, "We've already processed the prekey part of this V2 session, letting bundled message fall through...");
+      return;
+    }
+
+    if (!preKeyStore.containsPreKey(preKeyId))
+      throw new InvalidKeyIdException("No such prekey: " + preKeyId);
+
+    SessionRecord   sessionRecord        = sessionStore.loadSession(recipientId, deviceId);
+    PreKeyRecord    preKeyRecord         = preKeyStore.loadPreKey(preKeyId);
     ECKeyPair       ourBaseKey           = preKeyRecord.getKeyPair();
     ECKeyPair       ourEphemeralKey      = ourBaseKey;
     IdentityKeyPair ourIdentityKey       = identityKeyStore.getIdentityKeyPair();
@@ -112,63 +186,77 @@ public class SessionBuilder {
     if (!simultaneousInitiate) sessionRecord.reset();
     else                       sessionRecord.archiveCurrentState();
 
-    RatchetingSession.initializeSession(sessionRecord.getSessionState(),
+    RatchetingSessionV2.initializeSession(sessionRecord.getSessionState(),
-                                        ourBaseKey, theirBaseKey,
+                                          ourBaseKey, theirBaseKey,
-                                        ourEphemeralKey, theirEphemeralKey,
+                                          ourEphemeralKey, theirEphemeralKey,
-                                        ourIdentityKey, theirIdentityKey);
+                                          ourIdentityKey, theirIdentityKey);
 
     sessionRecord.getSessionState().setLocalRegistrationId(identityKeyStore.getLocalRegistrationId());
     sessionRecord.getSessionState().setRemoteRegistrationId(message.getRegistrationId());
+    sessionRecord.getSessionState().setSessionVersion(message.getMessageVersion());
 
     if (simultaneousInitiate) sessionRecord.getSessionState().setNeedsRefresh(true);
 
-    sessionStore.store(recipientId, deviceId, sessionRecord);
+    sessionStore.storeSession(recipientId, deviceId, sessionRecord);
 
     if (preKeyId != Medium.MAX_VALUE) {
-      preKeyStore.remove(preKeyId);
+      preKeyStore.removePreKey(preKeyId);
     }
-
-    identityKeyStore.saveIdentity(recipientId, theirIdentityKey);
   }
 
   /**
-   * Build a new session from a {@link org.whispersystems.libaxolotl.state.PreKey} retrieved from
+   * Build a new session from a {@link org.whispersystems.libaxolotl.state.PreKeyBundle} retrieved from
    * a server.
    *
    * @param preKey A PreKey for the destination recipient, retrieved from a server.
-   * @throws InvalidKeyException when the {@link org.whispersystems.libaxolotl.state.PreKey} is
+   * @throws InvalidKeyException when the {@link org.whispersystems.libaxolotl.state.PreKeyBundle} is
    *                             badly formatted.
    * @throws org.whispersystems.libaxolotl.UntrustedIdentityException when the sender's
    *                                                                  {@link IdentityKey} is not
    *                                                                  trusted.
    */
-  public void process(PreKey preKey) throws InvalidKeyException, UntrustedIdentityException {
+  public void process(PreKeyBundle preKey) throws InvalidKeyException, UntrustedIdentityException {
-
     if (!identityKeyStore.isTrustedIdentity(recipientId, preKey.getIdentityKey())) {
       throw new UntrustedIdentityException();
     }
 
-    SessionRecord   sessionRecord     = sessionStore.load(recipientId, deviceId);
+    if (preKey.getDeviceKey() != null &&
+        !Curve.verifySignature(preKey.getIdentityKey().getPublicKey(),
+                               preKey.getDeviceKey().serialize(),
+                               preKey.getDeviceKeySignature()))
+    {
+      throw new InvalidKeyException("Invalid signature on device key!");
+    }
+
+    SessionRecord   sessionRecord     = sessionStore.loadSession(recipientId, deviceId);
+    IdentityKeyPair ourIdentityKey    = identityKeyStore.getIdentityKeyPair();
     ECKeyPair       ourBaseKey        = Curve.generateKeyPair(true);
     ECKeyPair       ourEphemeralKey   = Curve.generateKeyPair(true);
-    ECPublicKey     theirBaseKey      = preKey.getPublicKey();
+    ECKeyPair       ourPreKey         = ourBaseKey;
-    ECPublicKey     theirEphemeralKey = theirBaseKey;
+
     IdentityKey     theirIdentityKey  = preKey.getIdentityKey();
-    IdentityKeyPair ourIdentityKey    = identityKeyStore.getIdentityKeyPair();
+    ECPublicKey     theirPreKey       = preKey.getPreKey();
+    ECPublicKey     theirBaseKey      = preKey.getDeviceKey() == null ? preKey.getPreKey() : preKey.getDeviceKey();
+    ECPublicKey     theirEphemeralKey = theirBaseKey;
 
     if (sessionRecord.getSessionState().getNeedsRefresh()) sessionRecord.archiveCurrentState();
     else                                                   sessionRecord.reset();
 
-    RatchetingSession.initializeSession(sessionRecord.getSessionState(),
+    if (preKey.getDeviceKey() == null) {
-                                        ourBaseKey, theirBaseKey, ourEphemeralKey,
+      RatchetingSessionV2.initializeSession(sessionRecord.getSessionState(),
-                                        theirEphemeralKey, ourIdentityKey, theirIdentityKey);
+                                            ourBaseKey, theirBaseKey, ourEphemeralKey,
+                                            theirEphemeralKey, ourIdentityKey, theirIdentityKey);
+    } else {
+      RatchetingSessionV3.initializeSession(sessionRecord.getSessionState(),
+                                            ourBaseKey, theirBaseKey, ourEphemeralKey, theirEphemeralKey,
+                                            ourPreKey, theirPreKey, ourIdentityKey, theirIdentityKey);
+    }
 
-    sessionRecord.getSessionState().setPendingPreKey(preKey.getKeyId(), ourBaseKey.getPublicKey());
+    sessionRecord.getSessionState().setPendingPreKey(preKey.getPreKeyId(), preKey.getDeviceKeyId(), ourBaseKey.getPublicKey());
     sessionRecord.getSessionState().setLocalRegistrationId(identityKeyStore.getLocalRegistrationId());
     sessionRecord.getSessionState().setRemoteRegistrationId(preKey.getRegistrationId());
 
-    sessionStore.store(recipientId, deviceId, sessionRecord);
+    sessionStore.storeSession(recipientId, deviceId, sessionRecord);
-
     identityKeyStore.saveIdentity(recipientId, preKey.getIdentityKey());
   }
 
@@ -189,7 +277,7 @@ public class SessionBuilder {
     }
 
     KeyExchangeMessage responseMessage = null;
-    SessionRecord      sessionRecord   = sessionStore.load(recipientId, deviceId);
+    SessionRecord      sessionRecord   = sessionStore.loadSession(recipientId, deviceId);
 
     Log.w(TAG, "Received key exchange with sequence: " + message.getSequence());
 
@@ -222,13 +310,13 @@ public class SessionBuilder {
 
     sessionRecord.reset();
 
-    RatchetingSession.initializeSession(sessionRecord.getSessionState(),
+    RatchetingSessionV2.initializeSession(sessionRecord.getSessionState(),
-                                        ourBaseKey, message.getBaseKey(),
+                                          ourBaseKey, message.getBaseKey(),
-                                        ourEphemeralKey, message.getEphemeralKey(),
+                                          ourEphemeralKey, message.getEphemeralKey(),
-                                        ourIdentityKey, message.getIdentityKey());
+                                          ourIdentityKey, message.getIdentityKey());
 
     sessionRecord.getSessionState().setSessionVersion(message.getVersion());
-    sessionStore.store(recipientId, deviceId, sessionRecord);
+    sessionStore.storeSession(recipientId, deviceId, sessionRecord);
 
     identityKeyStore.saveIdentity(recipientId, message.getIdentityKey());
 
@@ -279,10 +367,10 @@ public class SessionBuilder {
     ECKeyPair       baseKey       = Curve.generateKeyPair(true);
     ECKeyPair       ephemeralKey  = Curve.generateKeyPair(true);
     IdentityKeyPair identityKey   = identityKeyStore.getIdentityKeyPair();
-    SessionRecord   sessionRecord = sessionStore.load(recipientId, deviceId);
+    SessionRecord   sessionRecord = sessionStore.loadSession(recipientId, deviceId);
 
     sessionRecord.getSessionState().setPendingKeyExchange(sequence, baseKey, ephemeralKey, identityKey);
-    sessionStore.store(recipientId, deviceId, sessionRecord);
+    sessionStore.storeSession(recipientId, deviceId, sessionRecord);
 
     return new KeyExchangeMessage(sequence, flags,
                                   baseKey.getPublicKey(),

libaxolotl/src/main/java/org/whispersystems/libaxolotl/SessionCipher.java

@@ -83,30 +83,34 @@ public class SessionCipher {
    */
   public CiphertextMessage encrypt(byte[] paddedMessage) {
     synchronized (SESSION_LOCK) {
-      SessionRecord sessionRecord   = sessionStore.load(recipientId, deviceId);
+      SessionRecord sessionRecord   = sessionStore.loadSession(recipientId, deviceId);
       SessionState  sessionState    = sessionRecord.getSessionState();
       ChainKey      chainKey        = sessionState.getSenderChainKey();
       MessageKeys   messageKeys     = chainKey.getMessageKeys();
       ECPublicKey   senderEphemeral = sessionState.getSenderEphemeral();
       int           previousCounter = sessionState.getPreviousCounter();
+      int           sessionVersion  = sessionState.getSessionVersion();
 
       byte[]            ciphertextBody    = getCiphertext(messageKeys, paddedMessage);
-      CiphertextMessage ciphertextMessage = new WhisperMessage(messageKeys.getMacKey(),
+      CiphertextMessage ciphertextMessage = new WhisperMessage(sessionVersion, messageKeys.getMacKey(),
                                                                senderEphemeral, chainKey.getIndex(),
                                                                previousCounter, ciphertextBody);
 
       if (sessionState.hasPendingPreKey()) {
-        Pair<Integer, ECPublicKey> pendingPreKey       = sessionState.getPendingPreKey();
+        int         pendingPreKeyId     = sessionState.getPendingPreKeyId();
-        int                        localRegistrationId = sessionState.getLocalRegistrationId();
+        int         pendingDeviceKeyId  = sessionState.getPendingDeviceKeyId();
+        ECPublicKey pendingBaseKey      = sessionState.getPendingBaseKey();
+        int         localRegistrationId = sessionState.getLocalRegistrationId();
 
-        ciphertextMessage = new PreKeyWhisperMessage(localRegistrationId, pendingPreKey.first(),
+        ciphertextMessage = new PreKeyWhisperMessage(sessionVersion,
-                                                     pendingPreKey.second(),
+                                                     localRegistrationId, pendingPreKeyId,
+                                                     pendingDeviceKeyId, pendingBaseKey,
                                                      sessionState.getLocalIdentityKey(),
                                                      (WhisperMessage) ciphertextMessage);
       }
 
       sessionState.setSenderChainKey(chainKey.getNextChainKey());
-      sessionStore.store(recipientId, deviceId, sessionRecord);
+      sessionStore.storeSession(recipientId, deviceId, sessionRecord);
       return ciphertextMessage;
     }
   }
@@ -126,14 +130,14 @@ public class SessionCipher {
       throws InvalidMessageException, DuplicateMessageException, LegacyMessageException
   {
     synchronized (SESSION_LOCK) {
-      SessionRecord      sessionRecord  = sessionStore.load(recipientId, deviceId);
+      SessionRecord      sessionRecord  = sessionStore.loadSession(recipientId, deviceId);
       SessionState       sessionState   = sessionRecord.getSessionState();
       List<SessionState> previousStates = sessionRecord.getPreviousSessionStates();
       List<Exception>    exceptions     = new LinkedList<>();
 
       try {
         byte[] plaintext = decrypt(sessionState, ciphertext);
-        sessionStore.store(recipientId, deviceId, sessionRecord);
+        sessionStore.storeSession(recipientId, deviceId, sessionRecord);
 
         return plaintext;
       } catch (InvalidMessageException e) {
@@ -143,7 +147,7 @@ public class SessionCipher {
       for (SessionState previousState : previousStates) {
         try {
           byte[] plaintext = decrypt(previousState, ciphertext);
-          sessionStore.store(recipientId, deviceId, sessionRecord);
+          sessionStore.storeSession(recipientId, deviceId, sessionRecord);
 
           return plaintext;
         } catch (InvalidMessageException e) {
@@ -181,7 +185,7 @@ public class SessionCipher {
 
   public int getRemoteRegistrationId() {
     synchronized (SESSION_LOCK) {
-      SessionRecord record = sessionStore.load(recipientId, deviceId);
+      SessionRecord record = sessionStore.loadSession(recipientId, deviceId);
       return record.getSessionState().getRemoteRegistrationId();
     }
   }

libaxolotl/src/main/java/org/whispersystems/libaxolotl/ecc/Curve.java

@@ -55,4 +55,24 @@ public class Curve {
       throw new InvalidKeyException("Unknown type: " + publicKey.getType());
     }
   }
+
+  public static boolean verifySignature(ECPublicKey signingKey, byte[] message, byte[] signature)
+      throws InvalidKeyException
+  {
+    if (signingKey.getType() == DJB_TYPE) {
+      return Curve25519.verifySignature(signingKey, message, signature);
+    } else {
+      throw new InvalidKeyException("Unknown type: " + signingKey.getType());
+    }
+  }
+
+  public static byte[] calculateSignature(ECPrivateKey signingKey, byte[] message)
+      throws InvalidKeyException
+  {
+    if (signingKey.getType() == DJB_TYPE) {
+      return Curve25519.calculateSignature(signingKey, message);
+    } else {
+      throw new InvalidKeyException("Unknown type: " + signingKey.getType());
+    }
+  }
 }

libaxolotl/src/main/java/org/whispersystems/libaxolotl/protocol/CiphertextMessage.java

@@ -19,7 +19,7 @@ package org.whispersystems.libaxolotl.protocol;
 public interface CiphertextMessage {
 
   public static final int UNSUPPORTED_VERSION = 1;
-  public static final int CURRENT_VERSION     = 2;
+  public static final int CURRENT_VERSION     = 3;
 
   public static final int WHISPER_TYPE = 2;
   public static final int PREKEY_TYPE  = 3;

libaxolotl/src/main/java/org/whispersystems/libaxolotl/protocol/PreKeyWhisperMessage.java

@@ -34,6 +34,7 @@ public class PreKeyWhisperMessage implements CiphertextMessage {
   private final int            version;
   private final int            registrationId;
   private final int            preKeyId;
+  private final int            deviceKeyId;
   private final ECPublicKey    baseKey;
   private final IdentityKey    identityKey;
   private final WhisperMessage message;
@@ -53,9 +54,10 @@ public class PreKeyWhisperMessage implements CiphertextMessage {
           = WhisperProtos.PreKeyWhisperMessage.parseFrom(ByteString.copyFrom(serialized, 1,
                                                                              serialized.length-1));
 
-      if (!preKeyWhisperMessage.hasPreKeyId()       ||
+      if ((version == 2 && !preKeyWhisperMessage.hasPreKeyId())    ||
-          !preKeyWhisperMessage.hasBaseKey()        ||
+          (version == 3 && !preKeyWhisperMessage.hasDeviceKeyId()) ||
-          !preKeyWhisperMessage.hasIdentityKey()    ||
+          !preKeyWhisperMessage.hasBaseKey()                       ||
+          !preKeyWhisperMessage.hasIdentityKey()                   ||
           !preKeyWhisperMessage.hasMessage())
       {
         throw new InvalidMessageException("Incomplete message.");
@@ -63,32 +65,31 @@ public class PreKeyWhisperMessage implements CiphertextMessage {
 
       this.serialized     = serialized;
       this.registrationId = preKeyWhisperMessage.getRegistrationId();
-      this.preKeyId       = preKeyWhisperMessage.getPreKeyId();
+      this.preKeyId       = preKeyWhisperMessage.hasPreKeyId() ? preKeyWhisperMessage.getPreKeyId() : -1;
+      this.deviceKeyId    = preKeyWhisperMessage.hasDeviceKeyId() ? preKeyWhisperMessage.getDeviceKeyId() : -1;
       this.baseKey        = Curve.decodePoint(preKeyWhisperMessage.getBaseKey().toByteArray(), 0);
       this.identityKey    = new IdentityKey(Curve.decodePoint(preKeyWhisperMessage.getIdentityKey().toByteArray(), 0));
       this.message        = new WhisperMessage(preKeyWhisperMessage.getMessage().toByteArray());
-    } catch (InvalidProtocolBufferException e) {
+    } catch (InvalidProtocolBufferException | InvalidKeyException | LegacyMessageException e) {
-      throw new InvalidMessageException(e);
-    } catch (InvalidKeyException e) {
-      throw new InvalidMessageException(e);
-    } catch (LegacyMessageException e) {
       throw new InvalidMessageException(e);
     }
   }
 
-  public PreKeyWhisperMessage(int registrationId, int preKeyId, ECPublicKey baseKey,
+  public PreKeyWhisperMessage(int messageVersion, int registrationId, int preKeyId, int deviceKeyId,
-                              IdentityKey identityKey, WhisperMessage message)
+                              ECPublicKey baseKey, IdentityKey identityKey, WhisperMessage message)
   {
-    this.version        = CiphertextMessage.CURRENT_VERSION;
+    this.version        = messageVersion;
     this.registrationId = registrationId;
     this.preKeyId       = preKeyId;
+    this.deviceKeyId    = deviceKeyId;
     this.baseKey        = baseKey;
     this.identityKey    = identityKey;
     this.message        = message;
 
-    byte[] versionBytes = {ByteUtil.intsToByteHighAndLow(CURRENT_VERSION, this.version)};
+    byte[] versionBytes = {ByteUtil.intsToByteHighAndLow(this.version, CURRENT_VERSION)};
     byte[] messageBytes = WhisperProtos.PreKeyWhisperMessage.newBuilder()
                                        .setPreKeyId(preKeyId)
+                                       .setDeviceKeyId(deviceKeyId)
                                        .setBaseKey(ByteString.copyFrom(baseKey.serialize()))
                                        .setIdentityKey(ByteString.copyFrom(identityKey.serialize()))
                                        .setMessage(ByteString.copyFrom(message.serialize()))
@@ -98,6 +99,10 @@ public class PreKeyWhisperMessage implements CiphertextMessage {
     this.serialized = ByteUtil.combine(versionBytes, messageBytes);
   }
 
+  public int getMessageVersion() {
+    return version;
+  }
+
   public IdentityKey getIdentityKey() {
     return identityKey;
   }
@@ -110,6 +115,10 @@ public class PreKeyWhisperMessage implements CiphertextMessage {
     return preKeyId;
   }
 
+  public int getDeviceKeyId() {
+    return deviceKeyId;
+  }
+
   public ECPublicKey getBaseKey() {
     return baseKey;
   }

libaxolotl/src/main/java/org/whispersystems/libaxolotl/protocol/WhisperMessage.java

@@ -37,6 +37,7 @@ public class WhisperMessage implements CiphertextMessage {
 
   private static final int MAC_LENGTH = 8;
 
+  private final int         messageVersion;
   private final ECPublicKey senderEphemeral;
   private final int         counter;
   private final int         previousCounter;
@@ -54,7 +55,7 @@ public class WhisperMessage implements CiphertextMessage {
         throw new LegacyMessageException("Legacy message: " + ByteUtil.highBitsToInt(version));
       }
 
-      if (ByteUtil.highBitsToInt(version) != CURRENT_VERSION) {
+      if (ByteUtil.highBitsToInt(version) > CURRENT_VERSION) {
         throw new InvalidMessageException("Unknown version: " + ByteUtil.highBitsToInt(version));
       }
 
@@ -69,22 +70,19 @@ public class WhisperMessage implements CiphertextMessage {
 
       this.serialized      = serialized;
       this.senderEphemeral = Curve.decodePoint(whisperMessage.getEphemeralKey().toByteArray(), 0);
+      this.messageVersion  = ByteUtil.highBitsToInt(version);
       this.counter         = whisperMessage.getCounter();
       this.previousCounter = whisperMessage.getPreviousCounter();
       this.ciphertext      = whisperMessage.getCiphertext().toByteArray();
-    } catch (InvalidProtocolBufferException e) {
+    } catch (InvalidProtocolBufferException | InvalidKeyException | ParseException e) {
-      throw new InvalidMessageException(e);
-    } catch (InvalidKeyException e) {
-      throw new InvalidMessageException(e);
-    } catch (ParseException e) {
       throw new InvalidMessageException(e);
     }
   }
 
-  public WhisperMessage(SecretKeySpec macKey, ECPublicKey senderEphemeral,
+  public WhisperMessage(int messageVersion, SecretKeySpec macKey, ECPublicKey senderEphemeral,
                         int counter, int previousCounter, byte[] ciphertext)
   {
-    byte[] version = {ByteUtil.intsToByteHighAndLow(CURRENT_VERSION, CURRENT_VERSION)};
+    byte[] version = {ByteUtil.intsToByteHighAndLow(messageVersion, CURRENT_VERSION)};
     byte[] message = WhisperProtos.WhisperMessage.newBuilder()
                                    .setEphemeralKey(ByteString.copyFrom(senderEphemeral.serialize()))
                                    .setCounter(counter)
@@ -98,12 +96,17 @@ public class WhisperMessage implements CiphertextMessage {
     this.counter         = counter;
     this.previousCounter = previousCounter;
     this.ciphertext      = ciphertext;
+    this.messageVersion  = messageVersion;
   }
 
   public ECPublicKey getSenderEphemeral()  {
     return senderEphemeral;
   }
 
+  public int getMessageVersion() {
+    return messageVersion;
+  }
+
   public int getCounter() {
     return counter;
   }
@@ -131,9 +134,7 @@ public class WhisperMessage implements CiphertextMessage {
 
       byte[] fullMac = mac.doFinal(serialized);
       return ByteUtil.trim(fullMac, MAC_LENGTH);
-    } catch (NoSuchAlgorithmException e) {
+    } catch (NoSuchAlgorithmException | java.security.InvalidKeyException e) {
-      throw new AssertionError(e);
-    } catch (java.security.InvalidKeyException e) {
       throw new AssertionError(e);
     }
   }

libaxolotl/src/main/java/org/whispersystems/libaxolotl/ratchet/RatchetingSessionV2.java

@@ -30,7 +30,7 @@ import org.whispersystems.libaxolotl.util.Pair;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 
-public class RatchetingSession {
+public class RatchetingSessionV2 {
 
   public static void initializeSession(SessionState sessionState,
                                        ECKeyPair ourBaseKey,
@@ -48,6 +48,8 @@ public class RatchetingSession {
       initializeSessionAsBob(sessionState, ourBaseKey, theirBaseKey,
                              ourEphemeralKey, ourIdentityKey, theirIdentityKey);
     }
+
+    sessionState.setSessionVersion(2);
   }
 
   private static void initializeSessionAsAlice(SessionState sessionState,

libaxolotl/src/main/java/org/whispersystems/libaxolotl/ratchet/RatchetingSessionV3.java

@@ -0,0 +1,155 @@
+/**
+ * Copyright (C) 2014 Open Whisper Systems
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU 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 General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.whispersystems.libaxolotl.ratchet;
+
+import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.IdentityKeyPair;
+import org.whispersystems.libaxolotl.InvalidKeyException;
+import org.whispersystems.libaxolotl.state.SessionState;
+import org.whispersystems.libaxolotl.ecc.Curve;
+import org.whispersystems.libaxolotl.ecc.ECKeyPair;
+import org.whispersystems.libaxolotl.ecc.ECPublicKey;
+import org.whispersystems.libaxolotl.kdf.DerivedSecrets;
+import org.whispersystems.libaxolotl.kdf.HKDF;
+import org.whispersystems.libaxolotl.util.Pair;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.Arrays;
+
+public class RatchetingSessionV3 {
+
+  public static void initializeSession(SessionState    sessionState,
+                                       ECKeyPair       ourBaseKey,
+                                       ECPublicKey     theirBaseKey,
+                                       ECKeyPair       ourEphemeralKey,
+                                       ECPublicKey     theirEphemeralKey,
+                                       ECKeyPair       ourPreKey,
+                                       ECPublicKey     theirPreKey,
+                                       IdentityKeyPair ourIdentityKey,
+                                       IdentityKey     theirIdentityKey)
+      throws InvalidKeyException
+  {
+    if (isAlice(ourBaseKey.getPublicKey(), theirBaseKey, ourEphemeralKey.getPublicKey(), theirEphemeralKey)) {
+      initializeSessionAsAlice(sessionState, ourBaseKey, theirBaseKey, theirEphemeralKey,
+                               ourPreKey, theirPreKey, ourIdentityKey, theirIdentityKey);
+    } else {
+      initializeSessionAsBob(sessionState, ourBaseKey, theirBaseKey, ourEphemeralKey,
+                             ourPreKey, theirPreKey, ourIdentityKey, theirIdentityKey);
+    }
+
+    sessionState.setSessionVersion(3);
+  }
+
+  private static void initializeSessionAsAlice(SessionState sessionState,
+                                               ECKeyPair ourBaseKey, ECPublicKey theirBaseKey,
+                                               ECPublicKey theirEphemeralKey,
+                                               ECKeyPair ourPreKey, ECPublicKey theirPreKey,
+                                               IdentityKeyPair ourIdentityKey,
+                                               IdentityKey theirIdentityKey)
+      throws InvalidKeyException
+  {
+    sessionState.setRemoteIdentityKey(theirIdentityKey);
+    sessionState.setLocalIdentityKey(ourIdentityKey.getPublicKey());
+
+    ECKeyPair               sendingKey     = Curve.generateKeyPair(true);
+    Pair<RootKey, ChainKey> receivingChain = calculate4DHE(true, ourBaseKey, theirBaseKey,
+                                                           ourPreKey, theirPreKey,
+                                                           ourIdentityKey, theirIdentityKey);
+    Pair<RootKey, ChainKey> sendingChain   = receivingChain.first().createChain(theirEphemeralKey, sendingKey);
+
+    sessionState.addReceiverChain(theirEphemeralKey, receivingChain.second());
+    sessionState.setSenderChain(sendingKey, sendingChain.second());
+    sessionState.setRootKey(sendingChain.first());
+  }
+
+  private static void initializeSessionAsBob(SessionState sessionState,
+                                             ECKeyPair ourBaseKey, ECPublicKey theirBaseKey,
+                                             ECKeyPair ourEphemeralKey,
+                                             ECKeyPair ourPreKey, ECPublicKey theirPreKey,
+                                             IdentityKeyPair ourIdentityKey,
+                                             IdentityKey theirIdentityKey)
+      throws InvalidKeyException
+  {
+    sessionState.setRemoteIdentityKey(theirIdentityKey);
+    sessionState.setLocalIdentityKey(ourIdentityKey.getPublicKey());
+
+    Pair<RootKey, ChainKey> sendingChain = calculate4DHE(false, ourBaseKey, theirBaseKey,
+                                                         ourPreKey, theirPreKey,
+                                                         ourIdentityKey, theirIdentityKey);
+
+    sessionState.setSenderChain(ourEphemeralKey, sendingChain.second());
+    sessionState.setRootKey(sendingChain.first());
+  }
+
+  private static Pair<RootKey, ChainKey> calculate4DHE(boolean isAlice,
+                                                       ECKeyPair ourEphemeral, ECPublicKey theirEphemeral,
+                                                       ECKeyPair ourPreKey, ECPublicKey theirPreKey,
+                                                       IdentityKeyPair ourIdentity, IdentityKey theirIdentity)
+      throws InvalidKeyException
+  {
+    try {
+      byte[]                discontinuity = new byte[32];
+      ByteArrayOutputStream secrets       = new ByteArrayOutputStream();
+
+      Arrays.fill(discontinuity, (byte)0xFF);
+      secrets.write(discontinuity);
+
+      if (isAlice) {
+        secrets.write(Curve.calculateAgreement(theirEphemeral, ourIdentity.getPrivateKey()));
+        secrets.write(Curve.calculateAgreement(theirIdentity.getPublicKey(), ourEphemeral.getPrivateKey()));
+      } else {
+        secrets.write(Curve.calculateAgreement(theirIdentity.getPublicKey(), ourEphemeral.getPrivateKey()));
+        secrets.write(Curve.calculateAgreement(theirEphemeral, ourIdentity.getPrivateKey()));
+      }
+
+      secrets.write(Curve.calculateAgreement(theirEphemeral, ourEphemeral.getPrivateKey()));
+
+      if (ourPreKey != null && theirPreKey != null) {
+        secrets.write(Curve.calculateAgreement(theirPreKey, ourPreKey.getPrivateKey()));
+      }
+
+      DerivedSecrets derivedSecrets = new HKDF().deriveSecrets(secrets.toByteArray(),
+                                                               "WhisperText".getBytes());
+
+      return new Pair<>(new RootKey(derivedSecrets.getCipherKey().getEncoded()),
+                        new ChainKey(derivedSecrets.getMacKey().getEncoded(), 0));
+    } catch (IOException e) {
+      throw new AssertionError(e);
+    }
+  }
+
+  private static boolean isAlice(ECPublicKey ourBaseKey, ECPublicKey theirBaseKey,
+                                 ECPublicKey ourEphemeralKey, ECPublicKey theirEphemeralKey)
+  {
+    if (ourEphemeralKey.equals(ourBaseKey)) {
+      return false;
+    }
+
+    if (theirEphemeralKey.equals(theirBaseKey)) {
+      return true;
+    }
+
+    return isLowEnd(ourBaseKey, theirBaseKey);
+  }
+
+  private static boolean isLowEnd(ECPublicKey ourKey, ECPublicKey theirKey) {
+    return ourKey.compareTo(theirKey) < 0;
+  }
+
+
+}

libaxolotl/src/main/java/org/whispersystems/libaxolotl/state/DeviceKeyRecord.java

@@ -0,0 +1,61 @@
+package org.whispersystems.libaxolotl.state;
+
+import com.google.protobuf.ByteString;
+
+import org.whispersystems.libaxolotl.InvalidKeyException;
+import org.whispersystems.libaxolotl.ecc.Curve;
+import org.whispersystems.libaxolotl.ecc.ECKeyPair;
+import org.whispersystems.libaxolotl.ecc.ECPrivateKey;
+import org.whispersystems.libaxolotl.ecc.ECPublicKey;
+
+import java.io.IOException;
+
+import static org.whispersystems.libaxolotl.state.StorageProtos.DeviceKeyRecordStructure;
+
+public class DeviceKeyRecord {
+
+  private DeviceKeyRecordStructure structure;
+
+  public DeviceKeyRecord(int id, long timestamp, ECKeyPair keyPair, byte[] signature) {
+    this.structure = DeviceKeyRecordStructure.newBuilder()
+                                             .setId(id)
+                                             .setPublicKey(ByteString.copyFrom(keyPair.getPublicKey()
+                                                                                      .serialize()))
+                                             .setPrivateKey(ByteString.copyFrom(keyPair.getPrivateKey()
+                                                                                       .serialize()))
+                                             .setSignature(ByteString.copyFrom(signature))
+                                             .setTimestamp(timestamp)
+                                             .build();
+  }
+
+  public DeviceKeyRecord(byte[] serialized) throws IOException {
+    this.structure = DeviceKeyRecordStructure.parseFrom(serialized);
+  }
+
+  public int getId() {
+    return this.structure.getId();
+  }
+
+  public long getTimestamp() {
+    return this.structure.getTimestamp();
+  }
+
+  public ECKeyPair getKeyPair() {
+    try {
+      ECPublicKey publicKey = Curve.decodePoint(this.structure.getPublicKey().toByteArray(), 0);
+      ECPrivateKey privateKey = Curve.decodePrivatePoint(this.structure.getPrivateKey().toByteArray());
+
+      return new ECKeyPair(publicKey, privateKey);
+    } catch (InvalidKeyException e) {
+      throw new AssertionError(e);
+    }
+  }
+
+  public byte[] getSignature() {
+    return this.structure.getSignature().toByteArray();
+  }
+
+  public byte[] serialize() {
+    return this.structure.toByteArray();
+  }
+}

libaxolotl/src/main/java/org/whispersystems/libaxolotl/state/DeviceKeyStore.java

@@ -0,0 +1,47 @@
+package org.whispersystems.libaxolotl.state;
+
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
+
+import java.util.List;
+
+public interface DeviceKeyStore {
+
+
+  /**
+   * Load a local DeviceKeyRecord.
+   *
+   * @param deviceKeyId the ID of the local DeviceKeyRecord.
+   * @return the corresponding DeviceKeyRecord.
+   * @throws InvalidKeyIdException when there is no corresponding DeviceKeyRecord.
+   */
+  public DeviceKeyRecord loadDeviceKey(int deviceKeyId) throws InvalidKeyIdException;
+
+  /**
+   * Load all local DeviceKeyRecords.
+   *
+   * @return All stored DeviceKeyRecords.
+   */
+  public List<DeviceKeyRecord> loadDeviceKeys();
+
+  /**
+   * Store a local DeviceKeyRecord.
+   *
+   * @param deviceKeyId the ID of the DeviceKeyRecord to store.
+   * @param record the DeviceKeyRecord.
+   */
+  public void         storeDeviceKey(int deviceKeyId, DeviceKeyRecord record);
+
+  /**
+   * @param deviceKeyId A DeviceKeyRecord ID.
+   * @return true if the store has a record for the deviceKeyId, otherwise false.
+   */
+  public boolean      containsDeviceKey(int deviceKeyId);
+
+  /**
+   * Delete a DeviceKeyRecord from local storage.
+   *
+   * @param deviceKeyId The ID of the PreKeyRecord to remove.
+   */
+  public void         removeDeviceKey(int deviceKeyId);
+
+}

libaxolotl/src/main/java/org/whispersystems/libaxolotl/state/PreKey.java

@@ -1,36 +0,0 @@
-package org.whispersystems.libaxolotl.state;
-
-import org.whispersystems.libaxolotl.IdentityKey;
-import org.whispersystems.libaxolotl.ecc.ECPublicKey;
-
-/**
- * An interface that describes a remote PreKey.
- *
- * @author Moxie Marlinspike
- */
-public interface PreKey {
-  /**
-   * @return the device ID this PreKey belongs to.
-   */
-  public int         getDeviceId();
-
-  /**
-   * @return the unique key ID for this PreKey.
-   */
-  public int         getKeyId();
-
-  /**
-   * @return the public key for this PreKey.
-   */
-  public ECPublicKey getPublicKey();
-
-  /**
-   * @return the {@link org.whispersystems.libaxolotl.IdentityKey} of this PreKeys owner.
-   */
-  public IdentityKey getIdentityKey();
-
-  /**
-   * @return the registration ID associated with this PreKey.
-   */
-  public int         getRegistrationId();
-}

libaxolotl/src/main/java/org/whispersystems/libaxolotl/state/PreKeyBundle.java

@@ -0,0 +1,96 @@
+package org.whispersystems.libaxolotl.state;
+
+import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.ecc.ECPublicKey;
+
+/**
+ * A class that contains a remote PreKey and collection
+ * of associated items.
+ *
+ * @author Moxie Marlinspike
+ */
+public class PreKeyBundle {
+
+  private int         registrationId;
+
+  private int        deviceId;
+
+  private int         preKeyId;
+  private ECPublicKey preKeyPublic;
+
+  private int         deviceKeyId;
+  private ECPublicKey deviceKeyPublic;
+  private byte[]      deviceKeySignature;
+
+  private IdentityKey identityKey;
+
+  public PreKeyBundle(int registrationId, int deviceId, int preKeyId, ECPublicKey preKeyPublic,
+                      int deviceKeyId, ECPublicKey deviceKeyPublic, byte[] deviceKeySignature,
+                      IdentityKey identityKey)
+  {
+    this.registrationId     = registrationId;
+    this.deviceId           = deviceId;
+    this.preKeyId           = preKeyId;
+    this.preKeyPublic       = preKeyPublic;
+    this.deviceKeyId        = deviceKeyId;
+    this.deviceKeyPublic    = deviceKeyPublic;
+    this.deviceKeySignature = deviceKeySignature;
+    this.identityKey        = identityKey;
+  }
+
+  /**
+   * @return the device ID this PreKey belongs to.
+   */
+  public int getDeviceId() {
+    return deviceId;
+  }
+
+  /**
+   * @return the unique key ID for this PreKey.
+   */
+  public int getPreKeyId() {
+    return preKeyId;
+  }
+
+  /**
+   * @return the public key for this PreKey.
+   */
+  public ECPublicKey getPreKey() {
+    return preKeyPublic;
+  }
+
+  /**
+   * @return the unique key ID for this DeviceKey.
+   */
+  public int getDeviceKeyId() {
+    return deviceKeyId;
+  }
+
+  /**
+   * @return the device key for this PreKey.
+   */
+  public ECPublicKey getDeviceKey() {
+    return deviceKeyPublic;
+  }
+
+  /**
+   * @return the signature over the device key.
+   */
+  public byte[] getDeviceKeySignature() {
+    return deviceKeySignature;
+  }
+
+  /**
+   * @return the {@link org.whispersystems.libaxolotl.IdentityKey} of this PreKeys owner.
+   */
+  public IdentityKey getIdentityKey() {
+    return identityKey;
+  }
+
+  /**
+   * @return the registration ID associated with this PreKey.
+   */
+  public int getRegistrationId() {
+    return registrationId;
+  }
+}

libaxolotl/src/main/java/org/whispersystems/libaxolotl/state/PreKeyStore.java

@@ -16,7 +16,7 @@ public interface PreKeyStore {
    * @return the corresponding PreKeyRecord.
    * @throws InvalidKeyIdException when there is no corresponding PreKeyRecord.
    */
-  public PreKeyRecord load(int preKeyId) throws InvalidKeyIdException;
+  public PreKeyRecord loadPreKey(int preKeyId) throws InvalidKeyIdException;
 
   /**
    * Store a local PreKeyRecord.
@@ -24,19 +24,19 @@ public interface PreKeyStore {
    * @param preKeyId the ID of the PreKeyRecord to store.
    * @param record the PreKeyRecord.
    */
-  public void         store(int preKeyId, PreKeyRecord record);
+  public void         storePreKey(int preKeyId, PreKeyRecord record);
 
   /**
    * @param preKeyId A PreKeyRecord ID.
    * @return true if the store has a record for the preKeyId, otherwise false.
    */
-  public boolean      contains(int preKeyId);
+  public boolean      containsPreKey(int preKeyId);
 
   /**
    * Delete a PreKeyRecord from local storage.
    *
    * @param preKeyId The ID of the PreKeyRecord to remove.
    */
-  public void         remove(int preKeyId);
+  public void         removePreKey(int preKeyId);
 
 }

libaxolotl/src/main/java/org/whispersystems/libaxolotl/state/SessionRecord.java

@@ -1,6 +1,7 @@
 package org.whispersystems.libaxolotl.state;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.LinkedList;
 import java.util.List;
 
@@ -32,6 +33,24 @@ public class SessionRecord {
     }
   }
 
+  public boolean hasSessionState(int version, byte[] aliceBaseKey) {
+    if (sessionState.getSessionVersion() == version &&
+        Arrays.equals(aliceBaseKey, sessionState.getAliceBaseKey()))
+    {
+      return true;
+    }
+
+    for (SessionState state : previousStates) {
+      if (state.getSessionVersion() == version &&
+          Arrays.equals(aliceBaseKey, state.getAliceBaseKey()))
+      {
+        return true;
+      }
+    }
+
+    return false;
+  }
+
   public SessionState getSessionState() {
     return sessionState;
   }

libaxolotl/src/main/java/org/whispersystems/libaxolotl/state/SessionState.java

@@ -60,6 +60,16 @@ public class SessionState {
     return sessionStructure;
   }
 
+  public byte[] getAliceBaseKey() {
+    return this.sessionStructure.getAliceBaseKey().toByteArray();
+  }
+
+  public void setAliceBaseKey(byte[] aliceBaseKey) {
+    this.sessionStructure = this.sessionStructure.toBuilder()
+                                                 .setAliceBaseKey(ByteString.copyFrom(aliceBaseKey))
+                                                 .build();
+  }
+
   public void setNeedsRefresh(boolean needsRefresh) {
     this.sessionStructure = this.sessionStructure.toBuilder()
                                                  .setNeedsRefresh(needsRefresh)
@@ -77,7 +87,10 @@ public class SessionState {
   }
 
   public int getSessionVersion() {
-    return this.sessionStructure.getSessionVersion();
+    int sessionVersion = this.sessionStructure.getSessionVersion();
+
+    if (sessionVersion == 0) return 2;
+    else                     return sessionVersion;
   }
 
   public void setRemoteIdentityKey(IdentityKey identityKey) {
@@ -395,9 +408,10 @@ public class SessionState {
     return sessionStructure.hasPendingKeyExchange();
   }
 
-  public void setPendingPreKey(int preKeyId, ECPublicKey baseKey) {
+  public void setPendingPreKey(int preKeyId, int deviceKeyId, ECPublicKey baseKey) {
     PendingPreKey pending = PendingPreKey.newBuilder()
                                          .setPreKeyId(preKeyId)
+                                         .setDeviceKeyId(deviceKeyId)
                                          .setBaseKey(ByteString.copyFrom(baseKey.serialize()))
                                          .build();
 
@@ -410,12 +424,17 @@ public class SessionState {
     return this.sessionStructure.hasPendingPreKey();
   }
 
-  public Pair<Integer, ECPublicKey> getPendingPreKey() {
+  public int getPendingPreKeyId() {
+    return sessionStructure.getPendingPreKey().getPreKeyId();
+  }
+
+  public int getPendingDeviceKeyId() {
+    return sessionStructure.getPendingPreKey().getDeviceKeyId();
+  }
+
+  public ECPublicKey getPendingBaseKey() {
     try {
-      return new Pair<Integer, ECPublicKey>(sessionStructure.getPendingPreKey().getPreKeyId(),
+      return Curve.decodePoint(sessionStructure.getPendingPreKey().getBaseKey().toByteArray(), 0);
-                                            Curve.decodePoint(sessionStructure.getPendingPreKey()
-                                                                              .getBaseKey()
-                                                                              .toByteArray(), 0));
     } catch (InvalidKeyException e) {
       throw new AssertionError(e);
     }

libaxolotl/src/main/java/org/whispersystems/libaxolotl/state/SessionStore.java

@@ -24,7 +24,7 @@ public interface SessionStore {
    * @return a copy of the SessionRecord corresponding to the recipientId + deviceId tuple, or
    *         a new SessionRecord if one does not currently exist.
    */
-  public SessionRecord load(long recipientId, int deviceId);
+  public SessionRecord loadSession(long recipientId, int deviceId);
 
   /**
    * Returns all known devices with active sessions for a recipient
@@ -40,7 +40,7 @@ public interface SessionStore {
    * @param deviceId the device ID of the remote client.
    * @param record the current SessionRecord for the remote client.
    */
-  public void store(long recipientId, int deviceId, SessionRecord record);
+  public void storeSession(long recipientId, int deviceId, SessionRecord record);
 
   /**
    * Determine whether there is a committed {@link SessionRecord} for a recipientId + deviceId tuple.
@@ -48,7 +48,7 @@ public interface SessionStore {
    * @param deviceId the device ID of the remote client.
    * @return true if a {@link SessionRecord} exists, false otherwise.
    */
-  public boolean contains(long recipientId, int deviceId);
+  public boolean containsSession(long recipientId, int deviceId);
 
   /**
    * Remove a {@link SessionRecord} for a recipientId + deviceId tuple.
@@ -56,13 +56,13 @@ public interface SessionStore {
    * @param recipientId the recipient ID of the remote client.
    * @param deviceId the device ID of the remote client.
    */
-  public void delete(long recipientId, int deviceId);
+  public void deleteSession(long recipientId, int deviceId);
 
   /**
    * Remove the {@link SessionRecord}s corresponding to all devices of a recipientId.
    *
    * @param recipientId the recipient ID of the remote client.
    */
-  public void deleteAll(long recipientId);
+  public void deleteAllSessions(long recipientId);
 
 }

library/src/org/whispersystems/textsecure/crypto/PreKeyUtil.java

@@ -22,9 +22,14 @@ import android.util.Log;
 
 import com.google.thoughtcrimegson.Gson;
 
+import org.whispersystems.libaxolotl.IdentityKeyPair;
+import org.whispersystems.libaxolotl.InvalidKeyException;
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
+import org.whispersystems.libaxolotl.ecc.Curve;
 import org.whispersystems.libaxolotl.ecc.Curve25519;
 import org.whispersystems.libaxolotl.ecc.ECKeyPair;
-import org.whispersystems.libaxolotl.InvalidKeyIdException;
+import org.whispersystems.libaxolotl.state.DeviceKeyRecord;
+import org.whispersystems.libaxolotl.state.DeviceKeyStore;
 import org.whispersystems.libaxolotl.state.PreKeyRecord;
 import org.whispersystems.libaxolotl.state.PreKeyStore;
 import org.whispersystems.libaxolotl.util.Medium;
@@ -36,7 +41,6 @@ import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStreamReader;
-import java.util.Comparator;
 import java.util.LinkedList;
 import java.util.List;
 
@@ -54,7 +58,7 @@ public class PreKeyUtil {
       ECKeyPair    keyPair  = Curve25519.generateKeyPair(true);
       PreKeyRecord record   = new PreKeyRecord(preKeyId, keyPair);
 
-      preKeyStore.store(preKeyId, record);
+      preKeyStore.storePreKey(preKeyId, record);
       records.add(record);
     }
 
@@ -62,59 +66,45 @@ public class PreKeyUtil {
     return records;
   }
 
+  public static DeviceKeyRecord generateDeviceKey(Context context, MasterSecret masterSecret,
+                                                  IdentityKeyPair identityKeyPair)
+  {
+    try {
+      DeviceKeyStore  deviceKeyStore = new TextSecurePreKeyStore(context, masterSecret);
+      int             deviceKeyId    = getNextDeviceKeyId(context);
+      ECKeyPair       keyPair        = Curve25519.generateKeyPair(true);
+      byte[]          signature      = Curve.calculateSignature(identityKeyPair.getPrivateKey(), keyPair.getPublicKey().serialize());
+      DeviceKeyRecord record         = new DeviceKeyRecord(deviceKeyId, System.currentTimeMillis(), keyPair, signature);
+
+      deviceKeyStore.storeDeviceKey(deviceKeyId, record);
+      setNextDeviceKeyId(context, (deviceKeyId + 1) % Medium.MAX_VALUE);
+
+      return record;
+    } catch (InvalidKeyException e) {
+      throw new AssertionError(e);
+    }
+  }
+
   public static PreKeyRecord generateLastResortKey(Context context, MasterSecret masterSecret) {
     PreKeyStore preKeyStore = new TextSecurePreKeyStore(context, masterSecret);
 
-    if (preKeyStore.contains(Medium.MAX_VALUE)) {
+    if (preKeyStore.containsPreKey(Medium.MAX_VALUE)) {
       try {
-        return preKeyStore.load(Medium.MAX_VALUE);
+        return preKeyStore.loadPreKey(Medium.MAX_VALUE);
       } catch (InvalidKeyIdException e) {
         Log.w("PreKeyUtil", e);
-        preKeyStore.remove(Medium.MAX_VALUE);
+        preKeyStore.removePreKey(Medium.MAX_VALUE);
       }
     }
 
     ECKeyPair    keyPair = Curve25519.generateKeyPair(true);
     PreKeyRecord record  = new PreKeyRecord(Medium.MAX_VALUE, keyPair);
 
-    preKeyStore.store(Medium.MAX_VALUE, record);
+    preKeyStore.storePreKey(Medium.MAX_VALUE, record);
 
     return record;
   }
 
-//  public static List<PreKeyRecord> getPreKeys(Context context, MasterSecret masterSecret) {
-//    List<PreKeyRecord> records      = new LinkedList<PreKeyRecord>();
-//    File               directory    = getPreKeysDirectory(context);
-//    String[]           keyRecordIds = directory.list();
-//
-//    Arrays.sort(keyRecordIds, new PreKeyRecordIdComparator());
-//
-//    for (String keyRecordId : keyRecordIds) {
-//      try {
-//        if (!keyRecordId.equals(PreKeyIndex.FILE_NAME) && Integer.parseInt(keyRecordId) != Medium.MAX_VALUE) {
-//          records.add(new PreKeyRecord(context, masterSecret, Integer.parseInt(keyRecordId)));
-//        }
-//      } catch (InvalidKeyIdException e) {
-//        Log.w("PreKeyUtil", e);
-//        new File(getPreKeysDirectory(context), keyRecordId).delete();
-//      } catch (NumberFormatException nfe) {
-//        Log.w("PreKeyUtil", nfe);
-//        new File(getPreKeysDirectory(context), keyRecordId).delete();
-//      }
-//    }
-//
-//    return records;
-//  }
-//
-//  public static void clearPreKeys(Context context) {
-//    File     directory  = getPreKeysDirectory(context);
-//    String[] keyRecords = directory.list();
-//
-//    for (String keyRecord : keyRecords) {
-//      new File(directory, keyRecord).delete();
-//    }
-//  }
-
   private static void setNextPreKeyId(Context context, int id) {
     try {
       File             nextFile = new File(getPreKeysDirectory(context), PreKeyIndex.FILE_NAME);
@@ -126,6 +116,17 @@ public class PreKeyUtil {
     }
   }
 
+  private static void setNextDeviceKeyId(Context context, int id) {
+    try {
+      File             nextFile = new File(getDeviceKeysDirectory(context), DeviceKeyIndex.FILE_NAME);
+      FileOutputStream fout     = new FileOutputStream(nextFile);
+      fout.write(new Gson().toJson(new DeviceKeyIndex(id)).getBytes());
+      fout.close();
+    } catch (IOException e) {
+      Log.w("PreKeyUtil", e);
+    }
+  }
+
   private static int getNextPreKeyId(Context context) {
     try {
       File nextFile = new File(getPreKeysDirectory(context), PreKeyIndex.FILE_NAME);
@@ -144,8 +145,34 @@ public class PreKeyUtil {
     }
   }
 
+  private static int getNextDeviceKeyId(Context context) {
+    try {
+      File nextFile = new File(getDeviceKeysDirectory(context), DeviceKeyIndex.FILE_NAME);
+
+      if (!nextFile.exists()) {
+        return Util.getSecureRandom().nextInt(Medium.MAX_VALUE);
+      } else {
+        InputStreamReader reader = new InputStreamReader(new FileInputStream(nextFile));
+        DeviceKeyIndex    index  = new Gson().fromJson(reader, DeviceKeyIndex.class);
+        reader.close();
+        return index.nextDeviceKeyId;
+      }
+    } catch (IOException e) {
+      Log.w("PreKeyUtil", e);
+      return Util.getSecureRandom().nextInt(Medium.MAX_VALUE);
+    }
+  }
+
   private static File getPreKeysDirectory(Context context) {
-    File directory = new File(context.getFilesDir(), TextSecurePreKeyStore.PREKEY_DIRECTORY);
+    return getKeysDirectory(context, TextSecurePreKeyStore.PREKEY_DIRECTORY);
+  }
+
+  private static File getDeviceKeysDirectory(Context context) {
+    return getKeysDirectory(context, TextSecurePreKeyStore.DEVICE_KEY_DIRECTORY);
+  }
+
+  private static File getKeysDirectory(Context context, String name) {
+    File directory = new File(context.getFilesDir(), name);
 
     if (!directory.exists())
       directory.mkdirs();
@@ -153,25 +180,6 @@ public class PreKeyUtil {
     return directory;
   }
 
-  private static class PreKeyRecordIdComparator implements Comparator<String> {
-    @Override
-    public int compare(String lhs, String rhs) {
-      if      (lhs.equals(PreKeyIndex.FILE_NAME)) return -1;
-      else if (rhs.equals(PreKeyIndex.FILE_NAME)) return 1;
-
-      try {
-        long lhsLong = Long.parseLong(lhs);
-        long rhsLong = Long.parseLong(rhs);
-
-        if      (lhsLong < rhsLong) return -1;
-        else if (lhsLong > rhsLong) return 1;
-        else                        return 0;
-      } catch (NumberFormatException e) {
-        return 0;
-      }
-    }
-  }
-
   private static class PreKeyIndex {
     public static final String FILE_NAME = "index.dat";
 
@@ -184,4 +192,17 @@ public class PreKeyUtil {
     }
   }
 
+  private static class DeviceKeyIndex {
+    public static final String FILE_NAME = "index.dat";
+
+    private int nextDeviceKeyId;
+
+    public DeviceKeyIndex() {}
+
+    public DeviceKeyIndex(int nextDeviceKeyId) {
+      this.nextDeviceKeyId = nextDeviceKeyId;
+    }
+  }
+
+
 }

library/src/org/whispersystems/textsecure/crypto/SessionCipherFactory.java

@@ -32,7 +32,7 @@ public class SessionCipherFactory {
   {
     SessionStore sessionStore = new TextSecureSessionStore(context, masterSecret);
 
-    if (sessionStore.contains(recipient.getRecipientId(), recipient.getDeviceId())) {
+    if (sessionStore.containsSession(recipient.getRecipientId(), recipient.getDeviceId())) {
       return new SessionCipher(sessionStore, recipient.getRecipientId(), recipient.getDeviceId());
     } else {
       throw new AssertionError("Attempt to initialize cipher for non-existing session.");

library/src/org/whispersystems/textsecure/push/DeviceKeyEntity.java

@@ -0,0 +1,61 @@
+package org.whispersystems.textsecure.push;
+
+import com.google.thoughtcrimegson.GsonBuilder;
+import com.google.thoughtcrimegson.JsonDeserializationContext;
+import com.google.thoughtcrimegson.JsonDeserializer;
+import com.google.thoughtcrimegson.JsonElement;
+import com.google.thoughtcrimegson.JsonParseException;
+import com.google.thoughtcrimegson.JsonPrimitive;
+import com.google.thoughtcrimegson.JsonSerializationContext;
+import com.google.thoughtcrimegson.JsonSerializer;
+
+import org.whispersystems.libaxolotl.ecc.ECPublicKey;
+import org.whispersystems.textsecure.util.Base64;
+
+import java.io.IOException;
+import java.lang.reflect.Type;
+
+public class DeviceKeyEntity extends PreKeyEntity {
+
+  private byte[] signature;
+
+  public DeviceKeyEntity() {}
+
+  public DeviceKeyEntity(int keyId, ECPublicKey publicKey, byte[] signature) {
+    super(keyId, publicKey);
+    this.signature = signature;
+  }
+
+  public byte[] getSignature() {
+    return signature;
+  }
+
+  public static GsonBuilder forBuilder(GsonBuilder builder) {
+    return PreKeyEntity.forBuilder(builder)
+                       .registerTypeAdapter(byte[].class, new ByteArrayJsonAdapter());
+
+  }
+
+  private static class ByteArrayJsonAdapter
+      implements JsonSerializer<byte[]>, JsonDeserializer<byte[]>
+  {
+    @Override
+    public JsonElement serialize(byte[] signature, Type type,
+                                 JsonSerializationContext jsonSerializationContext)
+    {
+      return new JsonPrimitive(Base64.encodeBytesWithoutPadding(signature));
+    }
+
+    @Override
+    public byte[] deserialize(JsonElement jsonElement, Type type,
+                              JsonDeserializationContext jsonDeserializationContext)
+        throws JsonParseException
+    {
+      try {
+        return Base64.decodeWithoutPadding(jsonElement.getAsJsonPrimitive().getAsString());
+      } catch (IOException e) {
+        throw new JsonParseException(e);
+      }
+    }
+  }
+}

library/src/org/whispersystems/textsecure/push/PreKeyEntity.java

@@ -8,38 +8,26 @@ import com.google.thoughtcrimegson.JsonParseException;
 import com.google.thoughtcrimegson.JsonPrimitive;
 import com.google.thoughtcrimegson.JsonSerializationContext;
 import com.google.thoughtcrimegson.JsonSerializer;
-import com.google.thoughtcrimegson.annotations.Expose;
 
 import org.whispersystems.libaxolotl.IdentityKey;
 import org.whispersystems.libaxolotl.InvalidKeyException;
 import org.whispersystems.libaxolotl.ecc.Curve;
 import org.whispersystems.libaxolotl.ecc.ECPublicKey;
-import org.whispersystems.libaxolotl.state.PreKey;
 import org.whispersystems.textsecure.util.Base64;
 
 import java.io.IOException;
 import java.lang.reflect.Type;
 
-public class PreKeyEntity implements PreKey {
+public class PreKeyEntity {
-
-  @Expose(serialize = false)
-  private int         deviceId;
-
-  @Expose(serialize = false)
-  private int         registrationId;
 
   private int         keyId;
   private ECPublicKey publicKey;
-  private IdentityKey identityKey;
 
-  public PreKeyEntity(int keyId, ECPublicKey publicKey, IdentityKey identityKey) {
+  public PreKeyEntity() {}
-    this.keyId       = keyId;
-    this.publicKey   = publicKey;
-    this.identityKey = identityKey;
-  }
 
-  public int getDeviceId() {
+  public PreKeyEntity(int keyId, ECPublicKey publicKey) {
-    return deviceId;
+    this.keyId     = keyId;
+    this.publicKey = publicKey;
   }
 
   public int getKeyId() {
@@ -50,28 +38,8 @@ public class PreKeyEntity implements PreKey {
     return publicKey;
   }
 
-  public IdentityKey getIdentityKey() {
+  public static GsonBuilder forBuilder(GsonBuilder builder) {
-    return identityKey;
+    return builder.registerTypeAdapter(ECPublicKey.class, new ECPublicKeyJsonAdapter());
-  }
-
-  public int getRegistrationId() {
-    return registrationId;
-  }
-
-  public static String toJson(PreKeyEntity entity) {
-    return getBuilder().create().toJson(entity);
-  }
-
-  public static PreKeyEntity fromJson(String encoded) {
-    return getBuilder().create().fromJson(encoded, PreKeyEntity.class);
-  }
-
-  public static GsonBuilder getBuilder() {
-    GsonBuilder builder = new GsonBuilder();
-    builder.registerTypeAdapter(ECPublicKey.class, new ECPublicKeyJsonAdapter());
-    builder.registerTypeAdapter(IdentityKey.class, new IdentityKeyJsonAdapter());
-
-    return builder;
   }
 
 
@@ -92,34 +60,7 @@ public class PreKeyEntity implements PreKey {
     {
       try {
         return Curve.decodePoint(Base64.decodeWithoutPadding(jsonElement.getAsJsonPrimitive().getAsString()), 0);
-      } catch (InvalidKeyException e) {
+      } catch (InvalidKeyException | IOException e) {
-        throw new JsonParseException(e);
-      } catch (IOException e) {
-        throw new JsonParseException(e);
-      }
-    }
-  }
-
-  private static class IdentityKeyJsonAdapter
-      implements JsonSerializer<IdentityKey>, JsonDeserializer<IdentityKey>
-  {
-    @Override
-    public JsonElement serialize(IdentityKey identityKey, Type type,
-                                 JsonSerializationContext jsonSerializationContext)
-    {
-      return new JsonPrimitive(Base64.encodeBytesWithoutPadding(identityKey.serialize()));
-    }
-
-    @Override
-    public IdentityKey deserialize(JsonElement jsonElement, Type type,
-                                   JsonDeserializationContext jsonDeserializationContext)
-        throws JsonParseException
-    {
-      try {
-        return new IdentityKey(Base64.decodeWithoutPadding(jsonElement.getAsJsonPrimitive().getAsString()), 0);
-      } catch (InvalidKeyException e) {
-        throw new JsonParseException(e);
-      } catch (IOException e) {
         throw new JsonParseException(e);
       }
     }

library/src/org/whispersystems/textsecure/push/PreKeyList.java

@@ -1,30 +0,0 @@
-package org.whispersystems.textsecure.push;
-
-import java.util.List;
-
-public class PreKeyList {
-
-  private PreKeyEntity lastResortKey;
-  private List<PreKeyEntity> keys;
-
-  public PreKeyList(PreKeyEntity lastResortKey, List<PreKeyEntity> keys) {
-    this.keys          = keys;
-    this.lastResortKey = lastResortKey;
-  }
-
-  public List<PreKeyEntity> getKeys() {
-    return keys;
-  }
-
-  public static String toJson(PreKeyList entity) {
-    return PreKeyEntity.getBuilder().create().toJson(entity);
-  }
-
-  public static PreKeyList fromJson(String serialized) {
-    return PreKeyEntity.getBuilder().create().fromJson(serialized, PreKeyList.class);
-  }
-
-  public PreKeyEntity getLastResortKey() {
-    return lastResortKey;
-  }
-}

library/src/org/whispersystems/textsecure/push/PreKeyResponse.java

@@ -0,0 +1,65 @@
+package org.whispersystems.textsecure.push;
+
+import com.google.thoughtcrimegson.GsonBuilder;
+import com.google.thoughtcrimegson.JsonDeserializationContext;
+import com.google.thoughtcrimegson.JsonDeserializer;
+import com.google.thoughtcrimegson.JsonElement;
+import com.google.thoughtcrimegson.JsonParseException;
+import com.google.thoughtcrimegson.JsonPrimitive;
+import com.google.thoughtcrimegson.JsonSerializationContext;
+import com.google.thoughtcrimegson.JsonSerializer;
+
+import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.InvalidKeyException;
+import org.whispersystems.textsecure.util.Base64;
+
+import java.io.IOException;
+import java.lang.reflect.Type;
+import java.util.List;
+
+public class PreKeyResponse {
+
+  private IdentityKey              identityKey;
+  private List<PreKeyResponseItem> devices;
+
+  public IdentityKey getIdentityKey() {
+    return identityKey;
+  }
+
+  public List<PreKeyResponseItem> getDevices() {
+    return devices;
+  }
+
+  public static PreKeyResponse fromJson(String serialized) {
+    GsonBuilder builder = new GsonBuilder();
+    return PreKeyResponseItem.forBuilder(builder)
+                             .registerTypeAdapter(IdentityKey.class, new IdentityKeyJsonAdapter())
+                             .create().fromJson(serialized, PreKeyResponse.class);
+  }
+
+  public static class IdentityKeyJsonAdapter
+      implements JsonSerializer<IdentityKey>, JsonDeserializer<IdentityKey>
+  {
+    @Override
+    public JsonElement serialize(IdentityKey identityKey, Type type,
+                                 JsonSerializationContext jsonSerializationContext)
+    {
+      return new JsonPrimitive(Base64.encodeBytesWithoutPadding(identityKey.serialize()));
+    }
+
+    @Override
+    public IdentityKey deserialize(JsonElement jsonElement, Type type,
+                                   JsonDeserializationContext jsonDeserializationContext)
+        throws JsonParseException
+    {
+      try {
+        return new IdentityKey(Base64.decodeWithoutPadding(jsonElement.getAsJsonPrimitive().getAsString()), 0);
+      } catch (InvalidKeyException | IOException e) {
+        throw new JsonParseException(e);
+      }
+    }
+  }
+
+
+
+}

library/src/org/whispersystems/textsecure/push/PreKeyResponseItem.java

@@ -0,0 +1,31 @@
+package org.whispersystems.textsecure.push;
+
+import com.google.thoughtcrimegson.GsonBuilder;
+
+public class PreKeyResponseItem {
+
+  private int             deviceId;
+  private int             registrationId;
+  private DeviceKeyEntity deviceKey;
+  private PreKeyEntity    preKey;
+
+  public int getDeviceId() {
+    return deviceId;
+  }
+
+  public int getRegistrationId() {
+    return registrationId;
+  }
+
+  public DeviceKeyEntity getDeviceKey() {
+    return deviceKey;
+  }
+
+  public PreKeyEntity getPreKey() {
+    return preKey;
+  }
+
+  public static GsonBuilder forBuilder(GsonBuilder builder) {
+    return DeviceKeyEntity.forBuilder(builder);
+  }
+}

library/src/org/whispersystems/textsecure/push/PreKeyState.java

@@ -0,0 +1,32 @@
+package org.whispersystems.textsecure.push;
+
+import com.google.thoughtcrimegson.GsonBuilder;
+
+import org.whispersystems.libaxolotl.IdentityKey;
+
+import java.util.List;
+
+public class PreKeyState {
+
+  private IdentityKey        identityKey;
+  private List<PreKeyEntity> preKeys;
+  private PreKeyEntity       lastResortKey;
+  private DeviceKeyEntity    deviceKey;
+
+
+  public PreKeyState(List<PreKeyEntity> preKeys, PreKeyEntity lastResortKey,
+                     DeviceKeyEntity deviceKey, IdentityKey identityKey)
+  {
+    this.preKeys       = preKeys;
+    this.lastResortKey = lastResortKey;
+    this.deviceKey     = deviceKey;
+    this.identityKey   = identityKey;
+  }
+
+  public static String toJson(PreKeyState state) {
+    GsonBuilder builder = new GsonBuilder();
+    return DeviceKeyEntity.forBuilder(builder)
+                          .registerTypeAdapter(IdentityKey.class, new PreKeyResponse.IdentityKeyJsonAdapter())
+                          .create().toJson(state);
+  }
+}

library/src/org/whispersystems/textsecure/push/PushServiceSocket.java

@@ -24,6 +24,9 @@ import com.google.thoughtcrimegson.JsonParseException;
 
 import org.apache.http.conn.ssl.StrictHostnameVerifier;
 import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.ecc.ECPublicKey;
+import org.whispersystems.libaxolotl.state.DeviceKeyRecord;
+import org.whispersystems.libaxolotl.state.PreKeyBundle;
 import org.whispersystems.libaxolotl.state.PreKeyRecord;
 import org.whispersystems.textsecure.util.Base64;
 import org.whispersystems.textsecure.util.BlacklistingTrustManager;
@@ -64,9 +67,9 @@ public class PushServiceSocket {
   private static final String CREATE_ACCOUNT_VOICE_PATH = "/v1/accounts/voice/code/%s";
   private static final String VERIFY_ACCOUNT_PATH       = "/v1/accounts/code/%s";
   private static final String REGISTER_GCM_PATH         = "/v1/accounts/gcm/";
-  private static final String PREKEY_METADATA_PATH      = "/v1/keys/";
+  private static final String PREKEY_METADATA_PATH      = "/v2/keys/";
-  private static final String PREKEY_PATH               = "/v1/keys/%s";
+  private static final String PREKEY_PATH               = "/v2/keys/%s";
-  private static final String PREKEY_DEVICE_PATH        = "/v1/keys/%s/%s";
+  private static final String PREKEY_DEVICE_PATH        = "/v2/keys/%s/%s";
 
   private static final String DIRECTORY_TOKENS_PATH     = "/v1/directory/tokens";
   private static final String DIRECTORY_VERIFY_PATH     = "/v1/directory/%s";
@@ -126,6 +129,7 @@ public class PushServiceSocket {
 
   public void registerPreKeys(IdentityKey identityKey,
                               PreKeyRecord lastResortKey,
+                              DeviceKeyRecord deviceKey,
                               List<PreKeyRecord> records)
       throws IOException
   {
@@ -133,19 +137,20 @@ public class PushServiceSocket {
 
     for (PreKeyRecord record : records) {
       PreKeyEntity entity = new PreKeyEntity(record.getId(),
-                                             record.getKeyPair().getPublicKey(),
+                                             record.getKeyPair().getPublicKey());
-                                             identityKey);
 
       entities.add(entity);
     }
 
     PreKeyEntity lastResortEntity = new PreKeyEntity(lastResortKey.getId(),
-                                                     lastResortKey.getKeyPair().getPublicKey(),
+                                                     lastResortKey.getKeyPair().getPublicKey());
-                                                     identityKey);
 
+    DeviceKeyEntity deviceKeyEntity = new DeviceKeyEntity(deviceKey.getId(),
+                                                          deviceKey.getKeyPair().getPublicKey(),
+                                                          deviceKey.getSignature());
 
     makeRequest(String.format(PREKEY_PATH, ""), "PUT",
-                PreKeyList.toJson(new PreKeyList(lastResortEntity, entities)));
+                PreKeyState.toJson(new PreKeyState(entities, lastResortEntity, deviceKeyEntity, identityKey)));
   }
 
   public int getAvailablePreKeys() throws IOException {
@@ -155,7 +160,7 @@ public class PushServiceSocket {
     return preKeyStatus.getCount();
   }
 
-  public List<PreKeyEntity> getPreKeys(PushAddress destination) throws IOException {
+  public List<PreKeyBundle> getPreKeys(PushAddress destination) throws IOException {
     try {
       String deviceId = String.valueOf(destination.getDeviceId());
 
@@ -168,10 +173,34 @@ public class PushServiceSocket {
         path = path + "?relay=" + destination.getRelay();
       }
 
-      String responseText = makeRequest(path, "GET", null);
+      String             responseText = makeRequest(path, "GET", null);
-      PreKeyList response = PreKeyList.fromJson(responseText);
+      PreKeyResponse     response     = PreKeyResponse.fromJson(responseText);
+      List<PreKeyBundle> bundles      = new LinkedList<>();
 
-      return response.getKeys();
+      for (PreKeyResponseItem device : response.getDevices()) {
+        ECPublicKey preKey             = null;
+        ECPublicKey deviceKey          = null;
+        byte[]      deviceKeySignature = null;
+        int         preKeyId           = -1;
+        int         deviceKeyId        = -1;
+
+        if (device.getDeviceKey() != null) {
+          deviceKey          = device.getDeviceKey().getPublicKey();
+          deviceKeyId        = device.getDeviceKey().getKeyId();
+          deviceKeySignature = device.getDeviceKey().getSignature();
+        }
+
+        if (device.getPreKey() != null) {
+          preKeyId = device.getPreKey().getKeyId();
+          preKey   = device.getPreKey().getPublicKey();
+        }
+
+        bundles.add(new PreKeyBundle(device.getRegistrationId(), device.getDeviceId(), preKeyId,
+                                     preKey, deviceKeyId, deviceKey, deviceKeySignature,
+                                     response.getIdentityKey()));
+      }
+
+      return bundles;
     } catch (JsonParseException e) {
       throw new IOException(e);
     } catch (NotFoundException nfe) {
@@ -179,7 +208,7 @@ public class PushServiceSocket {
     }
   }
 
-  public PreKeyEntity getPreKey(PushAddress destination) throws IOException {
+  public PreKeyBundle getPreKey(PushAddress destination) throws IOException {
     try {
       String path = String.format(PREKEY_DEVICE_PATH, destination.getNumber(),
                                   String.valueOf(destination.getDeviceId()));
@@ -188,13 +217,32 @@ public class PushServiceSocket {
         path = path + "?relay=" + destination.getRelay();
       }
 
-      String     responseText = makeRequest(path, "GET", null);
+      String         responseText = makeRequest(path, "GET", null);
-      PreKeyList response     = PreKeyList.fromJson(responseText);
+      PreKeyResponse response     = PreKeyResponse.fromJson(responseText);
 
-      if (response.getKeys() == null || response.getKeys().size() < 1)
+      if (response.getDevices() == null || response.getDevices().size() < 1)
         throw new IOException("Empty prekey list");
 
-      return response.getKeys().get(0);
+      PreKeyResponseItem device             = response.getDevices().get(0);
+      ECPublicKey        preKey             = null;
+      ECPublicKey        deviceKey          = null;
+      byte[]             deviceKeySignature = null;
+      int                preKeyId           = -1;
+      int                deviceKeyId        = -1;
+
+      if (device.getPreKey() != null) {
+        preKeyId = device.getPreKey().getKeyId();
+        preKey   = device.getPreKey().getPublicKey();
+      }
+
+      if (device.getDeviceKey() != null) {
+        deviceKeyId        = device.getDeviceKey().getKeyId();
+        deviceKey          = device.getDeviceKey().getPublicKey();
+        deviceKeySignature = device.getDeviceKey().getSignature();
+      }
+
+      return new PreKeyBundle(device.getRegistrationId(), device.getDeviceId(), preKeyId, preKey,
+                              deviceKeyId, deviceKey, deviceKeySignature, response.getIdentityKey());
     } catch (JsonParseException e) {
       throw new IOException(e);
     } catch (NotFoundException nfe) {

library/src/org/whispersystems/textsecure/storage/SessionUtil.java

@@ -25,8 +25,8 @@ public class SessionUtil {
     SessionStore sessionStore = new TextSecureSessionStore(context, masterSecret);
 
     return
-        sessionStore.contains(recipientId, deviceId) &&
+        sessionStore.containsSession(recipientId, deviceId) &&
-        !sessionStore.load(recipientId, deviceId).getSessionState().getNeedsRefresh();
+        !sessionStore.loadSession(recipientId, deviceId).getSessionState().getNeedsRefresh();
   }
 
 }

library/src/org/whispersystems/textsecure/storage/TextSecurePreKeyStore.java

@@ -5,6 +5,8 @@ import android.util.Log;
 
 import org.whispersystems.libaxolotl.InvalidKeyIdException;
 import org.whispersystems.libaxolotl.InvalidMessageException;
+import org.whispersystems.libaxolotl.state.DeviceKeyRecord;
+import org.whispersystems.libaxolotl.state.DeviceKeyStore;
 import org.whispersystems.libaxolotl.state.PreKeyRecord;
 import org.whispersystems.libaxolotl.state.PreKeyStore;
 import org.whispersystems.textsecure.crypto.MasterCipher;
@@ -17,10 +19,15 @@ import java.io.IOException;
 import java.io.RandomAccessFile;
 import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
+import java.util.LinkedList;
+import java.util.List;
+
+public class TextSecurePreKeyStore implements PreKeyStore, DeviceKeyStore {
+
+  public  static final String PREKEY_DIRECTORY     = "prekeys";
+  public  static final String DEVICE_KEY_DIRECTORY = "device_keys";
 
-public class TextSecurePreKeyStore implements PreKeyStore {
 
-  public  static final String PREKEY_DIRECTORY       = "prekeys";
   private static final int    CURRENT_VERSION_MARKER = 1;
   private static final Object FILE_LOCK              = new Object();
   private static final String TAG                    = TextSecurePreKeyStore.class.getSimpleName();
@@ -34,20 +41,10 @@ public class TextSecurePreKeyStore implements PreKeyStore {
   }
 
   @Override
-  public PreKeyRecord load(int preKeyId) throws InvalidKeyIdException {
+  public PreKeyRecord loadPreKey(int preKeyId) throws InvalidKeyIdException {
     synchronized (FILE_LOCK) {
       try {
-        MasterCipher    masterCipher  = new MasterCipher(masterSecret);
+        return new PreKeyRecord(loadSerializedRecord(getPreKeyFile(preKeyId)));
-        FileInputStream fin           = new FileInputStream(getPreKeyFile(preKeyId));
-        int             recordVersion = readInteger(fin);
-
-        if (recordVersion != CURRENT_VERSION_MARKER) {
-          throw new AssertionError("Invalid version: " + recordVersion);
-        }
-
-        byte[] serializedRecord = masterCipher.decryptBytes(readBlob(fin));
-        return new PreKeyRecord(serializedRecord);
-
       } catch (IOException | InvalidMessageException e) {
         Log.w(TAG, e);
         throw new InvalidKeyIdException(e);
@@ -56,19 +53,40 @@ public class TextSecurePreKeyStore implements PreKeyStore {
   }
 
   @Override
-  public void store(int preKeyId, PreKeyRecord record) {
+  public DeviceKeyRecord loadDeviceKey(int deviceKeyId) throws InvalidKeyIdException {
     synchronized (FILE_LOCK) {
       try {
-        MasterCipher     masterCipher = new MasterCipher(masterSecret);
+        return new DeviceKeyRecord(loadSerializedRecord(getDeviceKeyFile(deviceKeyId)));
-        RandomAccessFile recordFile   = new RandomAccessFile(getPreKeyFile(preKeyId), "rw");
+      } catch (IOException | InvalidMessageException e) {
-        FileChannel      out          = recordFile.getChannel();
+        Log.w(TAG, e);
+        throw new InvalidKeyIdException(e);
+      }
+    }
+  }
 
-        out.position(0);
+  @Override
-        writeInteger(CURRENT_VERSION_MARKER, out);
+  public List<DeviceKeyRecord> loadDeviceKeys() {
-        writeBlob(masterCipher.encryptBytes(record.serialize()), out);
+    synchronized (FILE_LOCK) {
-        out.truncate(out.position());
+      File                  directory = getDeviceKeyDirectory();
+      List<DeviceKeyRecord> results   = new LinkedList<>();
 
-        recordFile.close();
+      for (File deviceKeyFile : directory.listFiles()) {
+        try {
+          results.add(new DeviceKeyRecord(loadSerializedRecord(deviceKeyFile)));
+        } catch (IOException | InvalidMessageException e) {
+          Log.w(TAG, e);
+        }
+      }
+
+      return results;
+    }
+  }
+
+  @Override
+  public void storePreKey(int preKeyId, PreKeyRecord record) {
+    synchronized (FILE_LOCK) {
+      try {
+        storeSerializedRecord(getPreKeyFile(preKeyId), record.serialize());
       } catch (IOException e) {
         throw new AssertionError(e);
       }
@@ -76,23 +94,85 @@ public class TextSecurePreKeyStore implements PreKeyStore {
   }
 
   @Override
-  public boolean contains(int preKeyId) {
+  public void storeDeviceKey(int deviceKeyId, DeviceKeyRecord record) {
+    synchronized (FILE_LOCK) {
+      try {
+        storeSerializedRecord(getDeviceKeyFile(deviceKeyId), record.serialize());
+      } catch (IOException e) {
+        throw new AssertionError(e);
+      }
+    }
+  }
+
+  @Override
+  public boolean containsPreKey(int preKeyId) {
     File record = getPreKeyFile(preKeyId);
     return record.exists();
   }
 
   @Override
-  public void remove(int preKeyId) {
+  public boolean containsDeviceKey(int deviceKeyId) {
+    File record = getDeviceKeyFile(deviceKeyId);
+    return record.exists();
+  }
+
+
+  @Override
+  public void removePreKey(int preKeyId) {
     File record = getPreKeyFile(preKeyId);
     record.delete();
   }
 
+  @Override
+  public void removeDeviceKey(int deviceKeyId) {
+    File record = getDeviceKeyFile(deviceKeyId);
+    record.delete();
+  }
+
+  private byte[] loadSerializedRecord(File recordFile)
+      throws IOException, InvalidMessageException
+  {
+    MasterCipher    masterCipher  = new MasterCipher(masterSecret);
+    FileInputStream fin           = new FileInputStream(recordFile);
+    int             recordVersion = readInteger(fin);
+
+    if (recordVersion != CURRENT_VERSION_MARKER) {
+      throw new AssertionError("Invalid version: " + recordVersion);
+    }
+
+    return masterCipher.decryptBytes(readBlob(fin));
+  }
+
+  private void storeSerializedRecord(File file, byte[] serialized) throws IOException {
+    MasterCipher     masterCipher = new MasterCipher(masterSecret);
+    RandomAccessFile recordFile   = new RandomAccessFile(file, "rw");
+    FileChannel      out          = recordFile.getChannel();
+
+    out.position(0);
+    writeInteger(CURRENT_VERSION_MARKER, out);
+    writeBlob(masterCipher.encryptBytes(serialized), out);
+    out.truncate(out.position());
+    recordFile.close();
+  }
+
   private File getPreKeyFile(int preKeyId) {
     return new File(getPreKeyDirectory(), String.valueOf(preKeyId));
   }
 
+  private File getDeviceKeyFile(int deviceKeyId) {
+    return new File(getDeviceKeyDirectory(), String.valueOf(deviceKeyId));
+  }
+
   private File getPreKeyDirectory() {
-    File directory = new File(context.getFilesDir(), PREKEY_DIRECTORY);
+    return getRecordsDirectory(PREKEY_DIRECTORY);
+  }
+
+  private File getDeviceKeyDirectory() {
+    return getRecordsDirectory(DEVICE_KEY_DIRECTORY);
+  }
+
+  private File getRecordsDirectory(String directoryName) {
+    File directory = new File(context.getFilesDir(), directoryName);
 
     if (!directory.exists()) {
       if (!directory.mkdirs()) {
@@ -127,4 +207,6 @@ public class TextSecurePreKeyStore implements PreKeyStore {
     out.write(ByteBuffer.wrap(valueBytes));
   }
 
+
+
 }

library/src/org/whispersystems/textsecure/storage/TextSecureSessionStore.java

@@ -41,7 +41,7 @@ public class TextSecureSessionStore implements SessionStore {
   }
 
   @Override
-  public SessionRecord load(long recipientId, int deviceId) {
+  public SessionRecord loadSession(long recipientId, int deviceId) {
     synchronized (FILE_LOCK) {
       try {
         MasterCipher    cipher = new MasterCipher(masterSecret);
@@ -73,7 +73,7 @@ public class TextSecureSessionStore implements SessionStore {
   }
 
   @Override
-  public void store(long recipientId, int deviceId, SessionRecord record) {
+  public void storeSession(long recipientId, int deviceId, SessionRecord record) {
     synchronized (FILE_LOCK) {
       try {
         MasterCipher     masterCipher = new MasterCipher(masterSecret);
@@ -93,24 +93,24 @@ public class TextSecureSessionStore implements SessionStore {
   }
 
   @Override
-  public boolean contains(long recipientId, int deviceId) {
+  public boolean containsSession(long recipientId, int deviceId) {
     return getSessionFile(recipientId, deviceId).exists() &&
-        load(recipientId, deviceId).getSessionState().hasSenderChain();
+        loadSession(recipientId, deviceId).getSessionState().hasSenderChain();
   }
 
   @Override
-  public void delete(long recipientId, int deviceId) {
+  public void deleteSession(long recipientId, int deviceId) {
     getSessionFile(recipientId, deviceId).delete();
   }
 
   @Override
-  public void deleteAll(long recipientId) {
+  public void deleteAllSessions(long recipientId) {
     List<Integer> devices = getSubDeviceSessions(recipientId);
 
-    delete(recipientId, RecipientDevice.DEFAULT_DEVICE_ID);
+    deleteSession(recipientId, RecipientDevice.DEFAULT_DEVICE_ID);
 
     for (int device : devices) {
-      delete(recipientId, device);
+      deleteSession(recipientId, device);
     }
   }
 

src/org/thoughtcrime/securesms/AutoInitiateActivity.java

@@ -118,6 +118,6 @@ public class AutoInitiateActivity extends Activity {
                                              Recipient recipient)
   {
     SessionStore sessionStore = new TextSecureSessionStore(context, masterSecret);
-    return sessionStore.contains(recipient.getRecipientId(), RecipientDevice.DEFAULT_DEVICE_ID);
+    return sessionStore.containsSession(recipient.getRecipientId(), RecipientDevice.DEFAULT_DEVICE_ID);
   }
 }

src/org/thoughtcrime/securesms/ConversationActivity.java

@@ -319,8 +319,8 @@ public class ConversationActivity extends PassphraseRequiredSherlockFragmentActi
       SessionStore sessionStore      = new TextSecureSessionStore(this, masterSecret);
       Recipient  primaryRecipient    = getRecipients() == null ? null : getRecipients().getPrimaryRecipient();
       boolean    isPushDestination   = DirectoryHelper.isPushDestination(this, getRecipients());
-      boolean    isSecureDestination = isSingleConversation() && sessionStore.contains(primaryRecipient.getRecipientId(),
+      boolean    isSecureDestination = isSingleConversation() && sessionStore.containsSession(primaryRecipient.getRecipientId(),
-                                                                                       RecipientDevice.DEFAULT_DEVICE_ID);
+                                                                                              RecipientDevice.DEFAULT_DEVICE_ID);
 
       getMenuInflater().inflate(R.menu.conversation_button_context, menu);
 
@@ -693,8 +693,8 @@ public class ConversationActivity extends PassphraseRequiredSherlockFragmentActi
     SessionStore sessionStore      = new TextSecureSessionStore(this, masterSecret);
     Recipient  primaryRecipient    = getRecipients() == null ? null : getRecipients().getPrimaryRecipient();
     boolean    isPushDestination   = DirectoryHelper.isPushDestination(this, getRecipients());
-    boolean    isSecureDestination = isSingleConversation() && sessionStore.contains(primaryRecipient.getRecipientId(),
+    boolean    isSecureDestination = isSingleConversation() && sessionStore.containsSession(primaryRecipient.getRecipientId(),
-                                                                                     RecipientDevice.DEFAULT_DEVICE_ID);
+                                                                                            RecipientDevice.DEFAULT_DEVICE_ID);
 
     if (isPushDestination || isSecureDestination) {
       this.isEncryptedConversation = true;

src/org/thoughtcrime/securesms/VerifyIdentityActivity.java

@@ -183,8 +183,8 @@ public class VerifyIdentityActivity extends KeyScanningActivity {
 
   private IdentityKey getRemoteIdentityKey(MasterSecret masterSecret, Recipient recipient) {
     SessionStore  sessionStore = new TextSecureSessionStore(this, masterSecret);
-    SessionRecord record       = sessionStore.load(recipient.getRecipientId(),
+    SessionRecord record       = sessionStore.loadSession(recipient.getRecipientId(),
-                                                   RecipientDevice.DEFAULT_DEVICE_ID);
+                                                          RecipientDevice.DEFAULT_DEVICE_ID);
 
     if (record == null) {
       return null;

src/org/thoughtcrime/securesms/crypto/DecryptingQueue.java

@@ -205,7 +205,7 @@ public class DecryptingQueue {
         Recipient       recipient       = recipients.getPrimaryRecipient();
         RecipientDevice recipientDevice = new RecipientDevice(recipient.getRecipientId(), message.getSourceDevice());
 
-        if (!sessionStore.contains(recipientDevice.getRecipientId(), recipientDevice.getDeviceId())) {
+        if (!sessionStore.containsSession(recipientDevice.getRecipientId(), recipientDevice.getDeviceId())) {
           sendResult(PushReceiver.RESULT_NO_SESSION);
           return;
         }
@@ -280,7 +280,7 @@ public class DecryptingQueue {
           return;
         }
 
-        if (!sessionStore.contains(recipientDevice.getRecipientId(), recipientDevice.getDeviceId())) {
+        if (!sessionStore.containsSession(recipientDevice.getRecipientId(), recipientDevice.getDeviceId())) {
           Log.w("DecryptingQueue", "No such recipient session for MMS...");
           database.markAsNoSession(messageId, threadId);
           return;
@@ -371,7 +371,7 @@ public class DecryptingQueue {
         SmsTransportDetails transportDetails  = new SmsTransportDetails();
         byte[]              decodedCiphertext = transportDetails.getDecodedMessage(body.getBytes());
 
-        if (!sessionStore.contains(recipientDevice.getRecipientId(), recipientDevice.getDeviceId())) {
+        if (!sessionStore.containsSession(recipientDevice.getRecipientId(), recipientDevice.getDeviceId())) {
           if (WhisperMessage.isLegacy(decodedCiphertext)) database.markAsLegacyVersion(messageId);
           else                                            database.markAsNoSession(messageId);
           return;
@@ -384,9 +384,9 @@ public class DecryptingQueue {
 
         if (isEndSession &&
             "TERMINATE".equals(plaintextBody) &&
-            sessionStore.contains(recipientDevice.getRecipientId(), recipientDevice.getDeviceId()))
+            sessionStore.containsSession(recipientDevice.getRecipientId(), recipientDevice.getDeviceId()))
         {
-          sessionStore.delete(recipientDevice.getRecipientId(), recipientDevice.getDeviceId());
+          sessionStore.deleteSession(recipientDevice.getRecipientId(), recipientDevice.getDeviceId());
         }
       } catch (InvalidMessageException | IOException | RecipientFormattingException e) {
         Log.w("DecryptionQueue", e);

src/org/thoughtcrime/securesms/crypto/KeyExchangeInitiator.java

@@ -28,6 +28,7 @@ import org.thoughtcrime.securesms.sms.OutgoingKeyExchangeMessage;
 import org.thoughtcrime.securesms.util.Dialogs;
 import org.whispersystems.libaxolotl.SessionBuilder;
 import org.whispersystems.libaxolotl.protocol.KeyExchangeMessage;
+import org.whispersystems.libaxolotl.state.DeviceKeyStore;
 import org.whispersystems.libaxolotl.state.IdentityKeyStore;
 import org.whispersystems.libaxolotl.state.PreKeyStore;
 import org.whispersystems.libaxolotl.state.SessionRecord;
@@ -62,10 +63,11 @@ public class KeyExchangeInitiator {
   private static void initiateKeyExchange(Context context, MasterSecret masterSecret, Recipient recipient) {
     SessionStore     sessionStore     = new TextSecureSessionStore(context, masterSecret);
     PreKeyStore      preKeyStore      = new TextSecurePreKeyStore(context, masterSecret);
+    DeviceKeyStore   deviceKeyStore   = new TextSecurePreKeyStore(context, masterSecret);
     IdentityKeyStore identityKeyStore = new TextSecureIdentityKeyStore(context, masterSecret);
 
-    SessionBuilder   sessionBuilder   = new SessionBuilder(sessionStore, preKeyStore, identityKeyStore,
+    SessionBuilder   sessionBuilder   = new SessionBuilder(sessionStore, preKeyStore, deviceKeyStore,
-                                                           recipient.getRecipientId(),
+                                                           identityKeyStore, recipient.getRecipientId(),
                                                            RecipientDevice.DEFAULT_DEVICE_ID);
 
     KeyExchangeMessage         keyExchangeMessage = sessionBuilder.process();
@@ -79,7 +81,7 @@ public class KeyExchangeInitiator {
                                              Recipient recipient)
   {
     SessionStore  sessionStore  = new TextSecureSessionStore(context, masterSecret);
-    SessionRecord sessionRecord = sessionStore.load(recipient.getRecipientId(), RecipientDevice.DEFAULT_DEVICE_ID);
+    SessionRecord sessionRecord = sessionStore.loadSession(recipient.getRecipientId(), RecipientDevice.DEFAULT_DEVICE_ID);
 
     return sessionRecord.getSessionState().hasPendingPreKey();
   }

src/org/thoughtcrime/securesms/crypto/KeyExchangeProcessor.java

@@ -15,7 +15,9 @@ import org.whispersystems.libaxolotl.StaleKeyExchangeException;
 import org.whispersystems.libaxolotl.UntrustedIdentityException;
 import org.whispersystems.libaxolotl.protocol.KeyExchangeMessage;
 import org.whispersystems.libaxolotl.protocol.PreKeyWhisperMessage;
+import org.whispersystems.libaxolotl.state.DeviceKeyStore;
 import org.whispersystems.libaxolotl.state.IdentityKeyStore;
+import org.whispersystems.libaxolotl.state.PreKeyBundle;
 import org.whispersystems.libaxolotl.state.PreKeyStore;
 import org.whispersystems.libaxolotl.state.SessionStore;
 import org.whispersystems.textsecure.crypto.MasterSecret;
@@ -48,10 +50,11 @@ public class KeyExchangeProcessor {
 
     IdentityKeyStore identityKeyStore = new TextSecureIdentityKeyStore(context, masterSecret);
     PreKeyStore      preKeyStore      = new TextSecurePreKeyStore(context, masterSecret);
+    DeviceKeyStore   deviceKeyStore   = new TextSecurePreKeyStore(context, masterSecret);
     SessionStore     sessionStore     = new TextSecureSessionStore(context, masterSecret);
 
-    this.sessionBuilder = new SessionBuilder(sessionStore, preKeyStore, identityKeyStore,
+    this.sessionBuilder = new SessionBuilder(sessionStore, preKeyStore, deviceKeyStore,
-                                             recipientDevice.getRecipientId(),
+                                             identityKeyStore, recipientDevice.getRecipientId(),
                                              recipientDevice.getDeviceId());
   }
 
@@ -62,10 +65,10 @@ public class KeyExchangeProcessor {
     PreKeyService.initiateRefresh(context, masterSecret);
   }
 
-  public void processKeyExchangeMessage(PreKeyEntity message, long threadId)
+  public void processKeyExchangeMessage(PreKeyBundle bundle, long threadId)
       throws InvalidKeyException, UntrustedIdentityException
   {
-    sessionBuilder.process(message);
+    sessionBuilder.process(bundle);
 
     if (threadId != -1) {
       broadcastSecurityUpdateEvent(context, threadId);

src/org/thoughtcrime/securesms/service/PreKeyService.java

@@ -10,12 +10,19 @@ import org.thoughtcrime.securesms.crypto.IdentityKeyUtil;
 import org.thoughtcrime.securesms.push.PushServiceSocketFactory;
 import org.thoughtcrime.securesms.util.TextSecurePreferences;
 import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.IdentityKeyPair;
+import org.whispersystems.libaxolotl.state.DeviceKeyRecord;
+import org.whispersystems.libaxolotl.state.DeviceKeyStore;
 import org.whispersystems.libaxolotl.state.PreKeyRecord;
 import org.whispersystems.textsecure.crypto.MasterSecret;
 import org.whispersystems.textsecure.crypto.PreKeyUtil;
 import org.whispersystems.textsecure.push.PushServiceSocket;
+import org.whispersystems.textsecure.storage.TextSecurePreKeyStore;
 
 import java.io.IOException;
+import java.util.Arrays;
+import java.util.Comparator;
+import java.util.Iterator;
 import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
@@ -62,6 +69,8 @@ public class PreKeyService extends Service {
     }
 
     public void run() {
+      DeviceKeyRecord deviceKeyRecord = null;
+
       try {
         if (!TextSecurePreferences.isPushRegistered(context)) return;
 
@@ -75,13 +84,60 @@ public class PreKeyService extends Service {
 
         List<PreKeyRecord> preKeyRecords       = PreKeyUtil.generatePreKeys(context, masterSecret);
         PreKeyRecord       lastResortKeyRecord = PreKeyUtil.generateLastResortKey(context, masterSecret);
-        IdentityKey        identityKey         = IdentityKeyUtil.getIdentityKey(context);
+        IdentityKeyPair    identityKey         = IdentityKeyUtil.getIdentityKeyPair(context, masterSecret);
+
+        deviceKeyRecord = PreKeyUtil.generateDeviceKey(context, masterSecret, identityKey);
 
         Log.w(TAG, "Registering new prekeys...");
 
-        socket.registerPreKeys(identityKey, lastResortKeyRecord, preKeyRecords);
+        socket.registerPreKeys(identityKey.getPublicKey(), lastResortKeyRecord,
+                               deviceKeyRecord, preKeyRecords);
+
+        removeOldDeviceKeysIfNecessary(deviceKeyRecord);
       } catch (IOException e) {
         Log.w(TAG, e);
+        if (deviceKeyRecord != null) {
+          Log.w(TAG, "Remote store failed, removing generated device key: " + deviceKeyRecord.getId());
+          new TextSecurePreKeyStore(context, masterSecret).removeDeviceKey(deviceKeyRecord.getId());
+        }
+      }
+    }
+
+    private void removeOldDeviceKeysIfNecessary(DeviceKeyRecord currentDeviceKey) {
+      DeviceKeyStore            deviceKeyStore = new TextSecurePreKeyStore(context, masterSecret);
+      List<DeviceKeyRecord>     records        = deviceKeyStore.loadDeviceKeys();
+      Iterator<DeviceKeyRecord> iterator       = records.iterator();
+
+      while (iterator.hasNext()) {
+        if (iterator.next().getId() == currentDeviceKey.getId()) {
+          iterator.remove();
+        }
+      }
+
+      DeviceKeyRecord[] recordsArray = (DeviceKeyRecord[])records.toArray();
+      Arrays.sort(recordsArray, new Comparator<DeviceKeyRecord>() {
+        @Override
+        public int compare(DeviceKeyRecord lhs, DeviceKeyRecord rhs) {
+          if      (lhs.getTimestamp() < rhs.getTimestamp()) return -1;
+          else if (lhs.getTimestamp() > rhs.getTimestamp()) return 1;
+          else                                              return 0;
+        }
+      });
+
+      Log.w(TAG, "Existing device key record count: " + recordsArray.length);
+
+      if (recordsArray.length > 3) {
+        long              oldTimestamp = System.currentTimeMillis() - (14 * 24 * 60 * 60 * 1000);
+        DeviceKeyRecord[] oldRecords   = Arrays.copyOf(recordsArray, recordsArray.length - 1);
+
+        for (DeviceKeyRecord oldRecord : oldRecords) {
+          Log.w(TAG, "Old device key record timestamp: " + oldRecord.getTimestamp());
+
+          if (oldRecord.getTimestamp() <= oldTimestamp) {
+            Log.w(TAG, "Remove device key record: " + oldRecord.getId());
+            deviceKeyStore.removeDeviceKey(oldRecord.getId());
+          }
+        }
       }
     }
   }

src/org/thoughtcrime/securesms/service/PushReceiver.java

@@ -185,7 +185,7 @@ public class PushReceiver {
       database.updateMessageBody(masterSecret, messageAndThreadId.first, messageContent.getBody());
 
       SessionStore sessionStore = new TextSecureSessionStore(context, masterSecret);
-      sessionStore.deleteAll(recipient.getRecipientId());
+      sessionStore.deleteAllSessions(recipient.getRecipientId());
 
       KeyExchangeProcessor.broadcastSecurityUpdateEvent(context, messageAndThreadId.second);
     } catch (RecipientFormattingException e) {

src/org/thoughtcrime/securesms/service/RegistrationService.java

@@ -17,7 +17,8 @@ import org.thoughtcrime.securesms.crypto.IdentityKeyUtil;
 import org.thoughtcrime.securesms.push.PushServiceSocketFactory;
 import org.thoughtcrime.securesms.util.DirectoryHelper;
 import org.thoughtcrime.securesms.util.TextSecurePreferences;
-import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.IdentityKeyPair;
+import org.whispersystems.libaxolotl.state.DeviceKeyRecord;
 import org.whispersystems.libaxolotl.state.PreKeyRecord;
 import org.whispersystems.libaxolotl.util.KeyHelper;
 import org.whispersystems.textsecure.crypto.MasterSecret;
@@ -227,10 +228,11 @@ public class RegistrationService extends Service {
       throws IOException
   {
     setState(new RegistrationState(RegistrationState.STATE_GENERATING_KEYS, number));
-    IdentityKey        identityKey = IdentityKeyUtil.getIdentityKey(this);
+    IdentityKeyPair    identityKey = IdentityKeyUtil.getIdentityKeyPair(this, masterSecret);
     List<PreKeyRecord> records     = PreKeyUtil.generatePreKeys(this, masterSecret);
     PreKeyRecord       lastResort  = PreKeyUtil.generateLastResortKey(this, masterSecret);
-    socket.registerPreKeys(identityKey, lastResort, records);
+    DeviceKeyRecord    deviceKey   = PreKeyUtil.generateDeviceKey(this, masterSecret, identityKey);
+    socket.registerPreKeys(identityKey.getPublicKey(), lastResort, deviceKey, records);
 
     setState(new RegistrationState(RegistrationState.STATE_GCM_REGISTERING, number));
 

src/org/thoughtcrime/securesms/service/SmsSender.java

@@ -140,7 +140,7 @@ public class SmsSender {
       if (record != null && record.isEndSession()) {
         Log.w("SmsSender", "Ending session...");
         SessionStore sessionStore = new TextSecureSessionStore(context, masterSecret);
-        sessionStore.deleteAll(record.getIndividualRecipient().getRecipientId());
+        sessionStore.deleteAllSessions(record.getIndividualRecipient().getRecipientId());
         KeyExchangeProcessor.broadcastSecurityUpdateEvent(context, record.getThreadId());
       }
 

src/org/thoughtcrime/securesms/transport/PushTransport.java

@@ -37,6 +37,7 @@ import org.thoughtcrime.securesms.util.Util;
 import org.whispersystems.libaxolotl.InvalidKeyException;
 import org.whispersystems.libaxolotl.SessionCipher;
 import org.whispersystems.libaxolotl.protocol.CiphertextMessage;
+import org.whispersystems.libaxolotl.state.PreKeyBundle;
 import org.whispersystems.libaxolotl.state.SessionStore;
 import org.whispersystems.textsecure.crypto.AttachmentCipher;
 import org.whispersystems.textsecure.crypto.MasterSecret;
@@ -45,7 +46,6 @@ import org.whispersystems.textsecure.push.MismatchedDevices;
 import org.whispersystems.textsecure.push.MismatchedDevicesException;
 import org.whispersystems.textsecure.push.OutgoingPushMessage;
 import org.whispersystems.textsecure.push.OutgoingPushMessageList;
-import org.whispersystems.textsecure.push.PreKeyEntity;
 import org.whispersystems.textsecure.push.PushAddress;
 import org.whispersystems.textsecure.push.PushAttachmentData;
 import org.whispersystems.textsecure.push.PushAttachmentPointer;
@@ -95,7 +95,7 @@ public class PushTransport extends BaseTransport {
 
       if (message.isEndSession()) {
         SessionStore sessionStore = new TextSecureSessionStore(context, masterSecret);
-        sessionStore.deleteAll(recipient.getRecipientId());
+        sessionStore.deleteAllSessions(recipient.getRecipientId());
         KeyExchangeProcessor.broadcastSecurityUpdateEvent(context, threadId);
       }
 
@@ -205,12 +205,12 @@ public class PushTransport extends BaseTransport {
       long         recipientId  = recipient.getRecipientId();
 
       for (int extraDeviceId : mismatchedDevices.getExtraDevices()) {
-        sessionStore.delete(recipientId, extraDeviceId);
+        sessionStore.deleteSession(recipientId, extraDeviceId);
       }
 
       for (int missingDeviceId : mismatchedDevices.getMissingDevices()) {
-        PushAddress            address   = PushAddress.create(context, recipientId, e164number, missingDeviceId);
+        PushAddress          address   = PushAddress.create(context, recipientId, e164number, missingDeviceId);
-        PreKeyEntity           preKey    = socket.getPreKey(address);
+        PreKeyBundle         preKey    = socket.getPreKey(address);
         KeyExchangeProcessor processor = new KeyExchangeProcessor(context, masterSecret, address);
 
         try {
@@ -230,7 +230,7 @@ public class PushTransport extends BaseTransport {
     long         recipientId  = recipient.getRecipientId();
 
     for (int staleDeviceId : staleDevices.getStaleDevices()) {
-      sessionStore.delete(recipientId, staleDeviceId);
+      sessionStore.deleteSession(recipientId, staleDeviceId);
     }
   }
 
@@ -327,10 +327,10 @@ public class PushTransport extends BaseTransport {
   {
     if (!SessionUtil.hasEncryptCapableSession(context, masterSecret, pushAddress)) {
       try {
-        List<PreKeyEntity> preKeys = socket.getPreKeys(pushAddress);
+        List<PreKeyBundle> preKeys = socket.getPreKeys(pushAddress);
 
-        for (PreKeyEntity preKey : preKeys) {
+        for (PreKeyBundle preKey : preKeys) {
-          PushAddress            device    = PushAddress.create(context, pushAddress.getRecipientId(), pushAddress.getNumber(), preKey.getDeviceId());
+          PushAddress          device    = PushAddress.create(context, pushAddress.getRecipientId(), pushAddress.getNumber(), preKey.getDeviceId());
           KeyExchangeProcessor processor = new KeyExchangeProcessor(context, masterSecret, device);
 
           try {