Move session construction and KeyExchangeMessage into libaxolotl.

1) Add plain two-way key exchange support libaxolotl by moving
   all the KeyExchangeMessage code there.

2) Move the bulk of KeyExchangeProcessor code to libaxolotl
   for setting up sessions based on retrieved prekeys, received
   prekeybundles, or exchanged key exchange messages.

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

@@ -1,5 +1,5 @@
 /** 
- * Copyright (C) 2011 Whisper Systems
+ * 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
@@ -14,13 +14,9 @@
  * 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.textsecure.storage;
+package org.whispersystems.libaxolotl;
 
 public class InvalidKeyIdException extends Exception {
-
-  public InvalidKeyIdException() {
-  }
-
   public InvalidKeyIdException(String detailMessage) {
     super(detailMessage);
   }
@@ -28,9 +24,4 @@ public class InvalidKeyIdException extends Exception {
   public InvalidKeyIdException(Throwable throwable) {
     super(throwable);
   }
-
-  public InvalidKeyIdException(String detailMessage, Throwable throwable) {
-    super(detailMessage, throwable);
-  }
-
 }

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

@@ -0,0 +1,180 @@
+package org.whispersystems.libaxolotl;
+
+import android.util.Log;
+
+import org.whispersystems.libaxolotl.ecc.Curve;
+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.state.IdentityKeyStore;
+import org.whispersystems.libaxolotl.state.PreKey;
+import org.whispersystems.libaxolotl.state.PreKeyRecord;
+import org.whispersystems.libaxolotl.state.PreKeyStore;
+import org.whispersystems.libaxolotl.state.SessionRecord;
+import org.whispersystems.libaxolotl.state.SessionStore;
+import org.whispersystems.libaxolotl.util.Medium;
+
+public class SessionBuilder {
+
+  private static final String TAG = SessionBuilder.class.getSimpleName();
+
+  private final SessionStore     sessionStore;
+  private final PreKeyStore      preKeyStore;
+  private final IdentityKeyStore identityKeyStore;
+  private final long             recipientId;
+  private final int              deviceId;
+
+  public SessionBuilder(SessionStore sessionStore,
+                        PreKeyStore preKeyStore,
+                        IdentityKeyStore identityKeyStore,
+                        long recipientId, int deviceId)
+  {
+    this.sessionStore     = sessionStore;
+    this.preKeyStore      = preKeyStore;
+    this.identityKeyStore = identityKeyStore;
+    this.recipientId      = recipientId;
+    this.deviceId         = deviceId;
+  }
+
+  public void process(PreKeyWhisperMessage message)
+      throws InvalidKeyIdException, InvalidKeyException
+  {
+    int         preKeyId          = message.getPreKeyId();
+    ECPublicKey theirBaseKey      = message.getBaseKey();
+    ECPublicKey theirEphemeralKey = message.getWhisperMessage().getSenderEphemeral();
+    IdentityKey theirIdentityKey  = message.getIdentityKey();
+
+    Log.w(TAG, "Received pre-key with local key ID: " + preKeyId);
+
+    if (!preKeyStore.contains(preKeyId) &&
+        sessionStore.contains(recipientId, deviceId))
+    {
+      Log.w(TAG, "We've already processed the prekey part, letting bundled message fall through...");
+      return;
+    }
+
+    if (!preKeyStore.contains(preKeyId))
+      throw new InvalidKeyIdException("No such prekey: " + preKeyId);
+
+    SessionRecord   sessionRecord        = sessionStore.get(recipientId, deviceId);
+    PreKeyRecord    preKeyRecord         = preKeyStore.load(preKeyId);
+    ECKeyPair       ourBaseKey           = preKeyRecord.getKeyPair();
+    ECKeyPair       ourEphemeralKey      = ourBaseKey;
+    IdentityKeyPair ourIdentityKey       = identityKeyStore.getIdentityKeyPair();
+    boolean         simultaneousInitiate = sessionRecord.getSessionState().hasPendingPreKey();
+
+    if (!simultaneousInitiate) sessionRecord.reset();
+    else                       sessionRecord.archiveCurrentState();
+
+    RatchetingSession.initializeSession(sessionRecord.getSessionState(),
+                                        ourBaseKey, theirBaseKey,
+                                        ourEphemeralKey, theirEphemeralKey,
+                                        ourIdentityKey, theirIdentityKey);
+
+    sessionRecord.getSessionState().setLocalRegistrationId(identityKeyStore.getLocalRegistrationId());
+    sessionRecord.getSessionState().setRemoteRegistrationId(message.getRegistrationId());
+
+    if (simultaneousInitiate) sessionRecord.getSessionState().setNeedsRefresh(true);
+
+    sessionStore.put(recipientId, deviceId, sessionRecord);
+
+    if (preKeyId != Medium.MAX_VALUE) {
+      preKeyStore.remove(preKeyId);
+    }
+
+    identityKeyStore.saveIdentity(recipientId, theirIdentityKey);
+  }
+
+  public void process(PreKey preKey) throws InvalidKeyException {
+    SessionRecord   sessionRecord     = sessionStore.get(recipientId, deviceId);
+    ECKeyPair       ourBaseKey        = Curve.generateKeyPair(true);
+    ECKeyPair       ourEphemeralKey   = Curve.generateKeyPair(true);
+    ECPublicKey     theirBaseKey      = preKey.getPublicKey();
+    ECPublicKey     theirEphemeralKey = theirBaseKey;
+    IdentityKey     theirIdentityKey  = preKey.getIdentityKey();
+    IdentityKeyPair ourIdentityKey    = identityKeyStore.getIdentityKeyPair();
+
+    if (sessionRecord.getSessionState().getNeedsRefresh()) sessionRecord.archiveCurrentState();
+    else                                                   sessionRecord.reset();
+
+    RatchetingSession.initializeSession(sessionRecord.getSessionState(),
+                                        ourBaseKey, theirBaseKey, ourEphemeralKey,
+                                        theirEphemeralKey, ourIdentityKey, theirIdentityKey);
+
+    sessionRecord.getSessionState().setPendingPreKey(preKey.getKeyId(), ourBaseKey.getPublicKey());
+    sessionRecord.getSessionState().setLocalRegistrationId(identityKeyStore.getLocalRegistrationId());
+    sessionRecord.getSessionState().setRemoteRegistrationId(preKey.getRegistrationId());
+
+    sessionStore.put(recipientId, deviceId, sessionRecord);
+
+    identityKeyStore.saveIdentity(recipientId, preKey.getIdentityKey());
+  }
+
+  public KeyExchangeMessage process(KeyExchangeMessage message) throws InvalidKeyException {
+    KeyExchangeMessage responseMessage = null;
+    SessionRecord      sessionRecord   = sessionStore.get(recipientId, deviceId);
+
+    Log.w(TAG, "Received key exchange with sequence: " + message.getSequence());
+
+    if (message.isInitiate()) {
+      ECKeyPair       ourBaseKey, ourEphemeralKey;
+      IdentityKeyPair ourIdentityKey;
+
+      int flags = KeyExchangeMessage.RESPONSE_FLAG;
+
+      Log.w(TAG, "KeyExchange is an initiate.");
+
+      if (!sessionRecord.getSessionState().hasPendingKeyExchange()) {
+        Log.w(TAG, "We don't have a pending initiate...");
+        ourBaseKey      = Curve.generateKeyPair(true);
+        ourEphemeralKey = Curve.generateKeyPair(true);
+        ourIdentityKey  = identityKeyStore.getIdentityKeyPair();
+
+        sessionRecord.getSessionState().setPendingKeyExchange(message.getSequence(), ourBaseKey,
+                                                              ourEphemeralKey, ourIdentityKey);
+      } else {
+        Log.w(TAG, "We already have a pending initiate, responding as simultaneous initiate...");
+        ourBaseKey      = sessionRecord.getSessionState().getPendingKeyExchangeBaseKey();
+        ourEphemeralKey = sessionRecord.getSessionState().getPendingKeyExchangeEphemeralKey();
+        ourIdentityKey  = sessionRecord.getSessionState().getPendingKeyExchangeIdentityKey();
+        flags          |= KeyExchangeMessage.SIMULTAENOUS_INITIATE_FLAG;
+
+        sessionRecord.getSessionState().setPendingKeyExchange(message.getSequence(), ourBaseKey,
+                                                              ourEphemeralKey, ourIdentityKey);
+      }
+
+      responseMessage = new KeyExchangeMessage(message.getSequence(),
+                                               flags, ourBaseKey.getPublicKey(),
+                                               ourEphemeralKey.getPublicKey(),
+                                               ourIdentityKey.getPublicKey());
+    }
+
+    if (message.getSequence() != sessionRecord.getSessionState().getPendingKeyExchangeSequence()) {
+      Log.w("KeyExchangeProcessor", "No matching sequence for response. " +
+          "Is simultaneous initiate response: " + message.isResponseForSimultaneousInitiate());
+      return responseMessage;
+    }
+
+    ECKeyPair       ourBaseKey      = sessionRecord.getSessionState().getPendingKeyExchangeBaseKey();
+    ECKeyPair       ourEphemeralKey = sessionRecord.getSessionState().getPendingKeyExchangeEphemeralKey();
+    IdentityKeyPair ourIdentityKey  = sessionRecord.getSessionState().getPendingKeyExchangeIdentityKey();
+
+    sessionRecord.reset();
+
+    RatchetingSession.initializeSession(sessionRecord.getSessionState(),
+                                        ourBaseKey, message.getBaseKey(),
+                                        ourEphemeralKey, message.getEphemeralKey(),
+                                        ourIdentityKey, message.getIdentityKey());
+
+    sessionRecord.getSessionState().setSessionVersion(message.getVersion());
+    sessionStore.put(recipientId, deviceId, sessionRecord);
+
+    identityKeyStore.saveIdentity(recipientId, message.getIdentityKey());
+
+    return responseMessage;
+  }
+
+
+}

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

@@ -1,7 +1,7 @@
-package org.thoughtcrime.securesms.crypto.protocol;
+package org.whispersystems.libaxolotl.protocol;
+
 
 import com.google.protobuf.ByteString;
-import com.google.protobuf.InvalidProtocolBufferException;
 
 import org.whispersystems.libaxolotl.IdentityKey;
 import org.whispersystems.libaxolotl.InvalidKeyException;
@@ -10,11 +10,7 @@ import org.whispersystems.libaxolotl.InvalidVersionException;
 import org.whispersystems.libaxolotl.LegacyMessageException;
 import org.whispersystems.libaxolotl.ecc.Curve;
 import org.whispersystems.libaxolotl.ecc.ECPublicKey;
-import org.whispersystems.libaxolotl.protocol.CiphertextMessage;
+import org.whispersystems.libaxolotl.util.ByteUtil;
-import org.whispersystems.libaxolotl.protocol.WhisperProtos;
-import org.whispersystems.textsecure.util.Base64;
-import org.whispersystems.textsecure.util.Conversions;
-import org.whispersystems.textsecure.util.Util;
 
 import java.io.IOException;
 
@@ -46,26 +42,24 @@ public class KeyExchangeMessage {
     this.ephemeralKey     = ephemeralKey;
     this.identityKey      = identityKey;
 
-    byte[] version = {Conversions.intsToByteHighAndLow(this.version, this.supportedVersion)};
+    byte[] version = {ByteUtil.intsToByteHighAndLow(this.version, this.supportedVersion)};
     byte[] message = WhisperProtos.KeyExchangeMessage.newBuilder()
-                                  .setId((sequence << 5) | flags)
+                                                     .setId((sequence << 5) | flags)
-                                  .setBaseKey(ByteString.copyFrom(baseKey.serialize()))
+                                                     .setBaseKey(ByteString.copyFrom(baseKey.serialize()))
-                                  .setEphemeralKey(ByteString.copyFrom(ephemeralKey.serialize()))
+                                                     .setEphemeralKey(ByteString.copyFrom(ephemeralKey.serialize()))
-                                  .setIdentityKey(ByteString.copyFrom(identityKey.serialize()))
+                                                     .setIdentityKey(ByteString.copyFrom(identityKey.serialize()))
-                                  .build().toByteArray();
+                                                     .build().toByteArray();
 
-    this.serialized = Util.combine(version, message);
+    this.serialized = ByteUtil.combine(version, message);
   }
 
-  public KeyExchangeMessage(String serializedAndEncoded)
+  public KeyExchangeMessage(byte[] serialized)
       throws InvalidMessageException, InvalidVersionException, LegacyMessageException
   {
     try {
-      byte[]   serialized = Base64.decodeWithoutPadding(serializedAndEncoded);
+      byte[][] parts        = ByteUtil.split(serialized, 1, serialized.length - 1);
-      byte[][] parts      = Util.split(serialized, 1, serialized.length - 1);
+      this.version          = ByteUtil.highBitsToInt(parts[0][0]);
-
+      this.supportedVersion = ByteUtil.lowBitsToInt(parts[0][0]);
-      this.version          = Conversions.highBitsToInt(parts[0][0]);
-      this.supportedVersion = Conversions.lowBitsToInt(parts[0][0]);
 
       if (this.version <= CiphertextMessage.UNSUPPORTED_VERSION) {
         throw new LegacyMessageException("Unsupported legacy version: " + this.version);
@@ -138,7 +132,7 @@ public class KeyExchangeMessage {
     return sequence;
   }
 
-  public String serialize() {
+  public byte[] serialize() {
-    return Base64.encodeBytesWithoutPadding(serialized);
+    return serialized;
   }
 }

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

@@ -0,0 +1,12 @@
+package org.whispersystems.libaxolotl.state;
+
+import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.IdentityKeyPair;
+
+public interface IdentityKeyStore {
+
+  public IdentityKeyPair getIdentityKeyPair();
+  public int             getLocalRegistrationId();
+  public void            saveIdentity(long recipientId, IdentityKey identityKey);
+
+}

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

@@ -0,0 +1,12 @@
+package org.whispersystems.libaxolotl.state;
+
+import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.ecc.ECPublicKey;
+
+public interface PreKey {
+  public int         getDeviceId();
+  public int         getKeyId();
+  public ECPublicKey getPublicKey();
+  public IdentityKey getIdentityKey();
+  public int         getRegistrationId();
+}

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

@@ -0,0 +1,9 @@
+package org.whispersystems.libaxolotl.state;
+
+import org.whispersystems.libaxolotl.ecc.ECKeyPair;
+
+public interface PreKeyRecord {
+  public int       getId();
+  public ECKeyPair getKeyPair();
+  public byte[]    serialize();
+}

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

@@ -0,0 +1,12 @@
+package org.whispersystems.libaxolotl.state;
+
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
+
+public interface PreKeyStore {
+
+  public PreKeyRecord load(int preKeyId) throws InvalidKeyIdException;
+  public void         store(int preKeyId, PreKeyRecord record);
+  public boolean      contains(int preKeyId);
+  public void         remove(int preKeyId);
+
+}

libaxolotl/src/main/java/org/whispersystems/libaxolotl/util/Medium.java

@@ -1,4 +1,4 @@
-package org.whispersystems.textsecure.util;
+package org.whispersystems.libaxolotl.util;
 
 public class Medium {
   public static int MAX_VALUE = 0xFFFFFF;

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

@@ -24,9 +24,12 @@ import com.google.thoughtcrimegson.Gson;
 
 import org.whispersystems.libaxolotl.ecc.Curve25519;
 import org.whispersystems.libaxolotl.ecc.ECKeyPair;
-import org.whispersystems.textsecure.storage.InvalidKeyIdException;
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
-import org.whispersystems.textsecure.storage.PreKeyRecord;
+import org.whispersystems.libaxolotl.state.PreKeyRecord;
-import org.whispersystems.textsecure.util.Medium;
+import org.whispersystems.libaxolotl.state.PreKeyStore;
+import org.whispersystems.libaxolotl.util.Medium;
+import org.whispersystems.textsecure.storage.TextSecurePreKeyRecord;
+import org.whispersystems.textsecure.storage.TextSecurePreKeyStore;
 import org.whispersystems.textsecure.util.Util;
 
 import java.io.File;
@@ -43,15 +46,16 @@ public class PreKeyUtil {
   public static final int BATCH_SIZE = 100;
 
   public static List<PreKeyRecord> generatePreKeys(Context context, MasterSecret masterSecret) {
-    List<PreKeyRecord> records        = new LinkedList<PreKeyRecord>();
+    PreKeyStore        preKeyStore    = new TextSecurePreKeyStore(context, masterSecret);
+    List<PreKeyRecord> records        = new LinkedList<>();
     int                preKeyIdOffset = getNextPreKeyId(context);
 
     for (int i=0;i<BATCH_SIZE;i++) {
       int          preKeyId = (preKeyIdOffset + i) % Medium.MAX_VALUE;
       ECKeyPair    keyPair  = Curve25519.generateKeyPair(true);
-      PreKeyRecord record   = new PreKeyRecord(context, masterSecret, preKeyId, keyPair);
+      PreKeyRecord record   = new TextSecurePreKeyRecord(masterSecret, preKeyId, keyPair);
 
-      record.save();
+      preKeyStore.store(preKeyId, record);
       records.add(record);
     }
 
@@ -60,19 +64,21 @@ public class PreKeyUtil {
   }
 
   public static PreKeyRecord generateLastResortKey(Context context, MasterSecret masterSecret) {
-    if (PreKeyRecord.hasRecord(context, Medium.MAX_VALUE)) {
+    PreKeyStore preKeyStore = new TextSecurePreKeyStore(context, masterSecret);
+
+    if (preKeyStore.contains(Medium.MAX_VALUE)) {
       try {
-        return new PreKeyRecord(context, masterSecret, Medium.MAX_VALUE);
+        return preKeyStore.load(Medium.MAX_VALUE);
       } catch (InvalidKeyIdException e) {
         Log.w("PreKeyUtil", e);
-        PreKeyRecord.delete(context, Medium.MAX_VALUE);
+        preKeyStore.remove(Medium.MAX_VALUE);
       }
     }
 
     ECKeyPair    keyPair = Curve25519.generateKeyPair(true);
-    PreKeyRecord record  = new PreKeyRecord(context, masterSecret, Medium.MAX_VALUE, keyPair);
+    PreKeyRecord record  = new TextSecurePreKeyRecord(masterSecret, Medium.MAX_VALUE, keyPair);
 
-    record.save();
+    preKeyStore.store(Medium.MAX_VALUE, record);
 
     return record;
   }
@@ -140,7 +146,7 @@ public class PreKeyUtil {
   }
 
   private static File getPreKeysDirectory(Context context) {
-    File directory = new File(context.getFilesDir(), PreKeyRecord.PREKEY_DIRECTORY);
+    File directory = new File(context.getFilesDir(), TextSecurePreKeyStore.PREKEY_DIRECTORY);
 
     if (!directory.exists())
       directory.mkdirs();

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

@@ -14,12 +14,13 @@ 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 {
+public class PreKeyEntity implements PreKey {
 
   @Expose(serialize = false)
   private int         deviceId;

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

@@ -24,7 +24,7 @@ import com.google.thoughtcrimegson.JsonParseException;
 
 import org.apache.http.conn.ssl.StrictHostnameVerifier;
 import org.whispersystems.libaxolotl.IdentityKey;
-import org.whispersystems.textsecure.storage.PreKeyRecord;
+import org.whispersystems.libaxolotl.state.PreKeyRecord;
 import org.whispersystems.textsecure.util.Base64;
 import org.whispersystems.textsecure.util.BlacklistingTrustManager;
 import org.whispersystems.textsecure.util.Util;
@@ -129,7 +129,7 @@ public class PushServiceSocket {
                               List<PreKeyRecord> records)
       throws IOException
   {
-    List<PreKeyEntity> entities = new LinkedList<PreKeyEntity>();
+    List<PreKeyEntity> entities = new LinkedList<>();
 
     for (PreKeyRecord record : records) {
       PreKeyEntity entity = new PreKeyEntity(record.getId(),

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

@@ -1,130 +0,0 @@
-package org.whispersystems.textsecure.storage;
-
-import android.content.Context;
-import android.util.Log;
-
-import com.google.protobuf.ByteString;
-
-import org.whispersystems.libaxolotl.InvalidKeyException;
-import org.whispersystems.libaxolotl.InvalidMessageException;
-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.textsecure.crypto.MasterCipher;
-import org.whispersystems.textsecure.crypto.MasterSecret;
-
-import java.io.FileInputStream;
-import java.io.FileNotFoundException;
-import java.io.IOException;
-import java.io.RandomAccessFile;
-import java.nio.channels.FileChannel;
-
-public class PreKeyRecord extends Record {
-
-  private static final Object FILE_LOCK              = new Object();
-  private static final int    CURRENT_VERSION_MARKER = 1;
-
-  private final MasterSecret masterSecret;
-  private StorageProtos.PreKeyRecordStructure structure;
-
-  public PreKeyRecord(Context context, MasterSecret masterSecret, int id)
-      throws InvalidKeyIdException
-  {
-    super(context, PREKEY_DIRECTORY, id+"");
-
-    this.structure    = StorageProtos.PreKeyRecordStructure.newBuilder().setId(id).build();
-    this.masterSecret = masterSecret;
-
-    loadData();
-  }
-
-  public PreKeyRecord(Context context, MasterSecret masterSecret,
-                      int id, ECKeyPair keyPair)
-  {
-    super(context, PREKEY_DIRECTORY, id+"");
-    this.masterSecret = masterSecret;
-    this.structure    = StorageProtos.PreKeyRecordStructure.newBuilder()
-                                     .setId(id)
-                                     .setPublicKey(ByteString.copyFrom(keyPair.getPublicKey()
-                                                                              .serialize()))
-                                     .setPrivateKey(ByteString.copyFrom(keyPair.getPrivateKey()
-                                                                               .serialize()))
-                                     .build();
-  }
-
-  public int getId() {
-    return this.structure.getId();
-  }
-
-  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 static boolean hasRecord(Context context, long id) {
-    Log.w("PreKeyRecord", "Checking: " + id);
-    return Record.hasRecord(context, PREKEY_DIRECTORY, id+"");
-  }
-
-  public static void delete(Context context, long id) {
-    Record.delete(context, PREKEY_DIRECTORY, id+"");
-  }
-
-  public void save() {
-    synchronized (FILE_LOCK) {
-      try {
-        RandomAccessFile file = openRandomAccessFile();
-        FileChannel out       = file.getChannel();
-        out.position(0);
-
-        MasterCipher masterCipher = new MasterCipher(masterSecret);
-
-        writeInteger(CURRENT_VERSION_MARKER, out);
-        writeBlob(masterCipher.encryptBytes(structure.toByteArray()), out);
-
-        out.force(true);
-        out.truncate(out.position());
-        out.close();
-        file.close();
-      } catch (IOException ioe) {
-        Log.w("PreKeyRecord", ioe);
-      }
-    }
-  }
-
-  private void loadData() throws InvalidKeyIdException {
-    synchronized (FILE_LOCK) {
-      try {
-        MasterCipher    masterCipher  = new MasterCipher(masterSecret);
-        FileInputStream in            = this.openInputStream();
-        int             recordVersion = readInteger(in);
-
-        if (recordVersion != CURRENT_VERSION_MARKER) {
-          Log.w("PreKeyRecord", "Invalid version: " + recordVersion);
-          return;
-        }
-
-        this.structure =
-            StorageProtos.PreKeyRecordStructure.parseFrom(masterCipher.decryptBytes(readBlob(in)));
-
-        in.close();
-      } catch (FileNotFoundException e) {
-        Log.w("PreKeyRecord", e);
-        throw new InvalidKeyIdException(e);
-      } catch (IOException ioe) {
-        Log.w("PreKeyRecord", ioe);
-        throw new InvalidKeyIdException(ioe);
-      } catch (InvalidMessageException ime) {
-        Log.w("PreKeyRecord", ime);
-        throw new InvalidKeyIdException(ime);
-      }
-    }
-  }
-}

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

@@ -1,113 +0,0 @@
-/**
- * Copyright (C) 2011 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.textsecure.storage;
-
-import android.content.Context;
-
-import org.whispersystems.textsecure.util.Conversions;
-
-import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileNotFoundException;
-import java.io.IOException;
-import java.io.RandomAccessFile;
-import java.nio.ByteBuffer;
-import java.nio.channels.FileChannel;
-
-public abstract class Record {
-
-  public    static final String SESSIONS_DIRECTORY    = "sessions";
-  protected static final String SESSIONS_DIRECTORY_V2 = "sessions-v2";
-  public    static final String PREKEY_DIRECTORY      = "prekeys";
-
-  protected final String address;
-  protected final String directory;
-  protected final Context context;
-
-  public Record(Context context, String directory, String address) {
-    this.context   = context;
-    this.directory = directory;
-    this.address   = address;
-  }
-
-  public void delete() {
-    delete(this.context, this.directory, this.address);
-  }
-
-  public static void delete(Context context, String directory, String address) {
-    getAddressFile(context, directory, address).delete();
-  }
-
-  protected static  boolean hasRecord(Context context, String directory, String address) {
-    return getAddressFile(context, directory, address).exists();
-  }
-
-  protected RandomAccessFile openRandomAccessFile() throws FileNotFoundException {
-    return new RandomAccessFile(getAddressFile(), "rw");
-  }
-
-  protected FileInputStream openInputStream() throws FileNotFoundException {
-    return new FileInputStream(getAddressFile().getAbsolutePath());
-  }
-
-  private File getAddressFile() {
-    return getAddressFile(context, directory, address);
-  }
-
-  private static File getAddressFile(Context context, String directory, String address) {
-    File parent = getParentDirectory(context, directory);
-
-    return new File(parent, address);
-  }
-
-  protected static File getParentDirectory(Context context, String directory) {
-    File parent = new File(context.getFilesDir(), directory);
-
-    if (!parent.exists()) {
-      parent.mkdirs();
-    }
-
-    return parent;
-  }
-
-  protected byte[] readBlob(FileInputStream in) throws IOException {
-    int length       = readInteger(in);
-    byte[] blobBytes = new byte[length];
-
-    in.read(blobBytes, 0, blobBytes.length);
-    return blobBytes;
-  }
-
-  protected void writeBlob(byte[] blobBytes, FileChannel out) throws IOException {
-    writeInteger(blobBytes.length, out);
-    ByteBuffer buffer = ByteBuffer.wrap(blobBytes);
-    out.write(buffer);
-  }
-
-  protected int readInteger(FileInputStream in) throws IOException {
-    byte[] integer = new byte[4];
-    in.read(integer, 0, integer.length);
-    return Conversions.byteArrayToInt(integer);
-  }
-
-  protected void writeInteger(int value, FileChannel out) throws IOException {
-    byte[] valueBytes = Conversions.intToByteArray(value);
-    ByteBuffer buffer = ByteBuffer.wrap(valueBytes);
-    out.write(buffer);
-  }
-
-}

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

@@ -0,0 +1,117 @@
+package org.whispersystems.textsecure.storage;
+
+import android.util.Log;
+
+import com.google.protobuf.ByteString;
+
+import org.whispersystems.libaxolotl.InvalidKeyException;
+import org.whispersystems.libaxolotl.InvalidMessageException;
+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.state.PreKeyRecord;
+import org.whispersystems.textsecure.crypto.MasterCipher;
+import org.whispersystems.textsecure.crypto.MasterSecret;
+import org.whispersystems.textsecure.util.Conversions;
+
+import java.io.ByteArrayOutputStream;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+public class TextSecurePreKeyRecord implements PreKeyRecord {
+
+  private static final int    CURRENT_VERSION_MARKER = 1;
+
+  private final MasterSecret masterSecret;
+  private StorageProtos.PreKeyRecordStructure structure;
+
+  public TextSecurePreKeyRecord(MasterSecret masterSecret, int id, ECKeyPair keyPair) {
+    this.masterSecret = masterSecret;
+    this.structure    = StorageProtos.PreKeyRecordStructure.newBuilder()
+                                     .setId(id)
+                                     .setPublicKey(ByteString.copyFrom(keyPair.getPublicKey()
+                                                                              .serialize()))
+                                     .setPrivateKey(ByteString.copyFrom(keyPair.getPrivateKey()
+                                                                               .serialize()))
+                                     .build();
+  }
+
+  public TextSecurePreKeyRecord(MasterSecret masterSecret, FileInputStream in)
+      throws IOException, InvalidMessageException
+  {
+    this.masterSecret = masterSecret;
+
+    MasterCipher masterCipher  = new MasterCipher(masterSecret);
+    int          recordVersion = readInteger(in);
+
+    if (recordVersion != CURRENT_VERSION_MARKER) {
+      Log.w("PreKeyRecord", "Invalid version: " + recordVersion);
+      return;
+    }
+
+    this.structure =
+        StorageProtos.PreKeyRecordStructure.parseFrom(masterCipher.decryptBytes(readBlob(in)));
+
+    in.close();
+  }
+
+  @Override
+  public int getId() {
+    return this.structure.getId();
+  }
+
+  @Override
+  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[] serialize() {
+    try {
+      ByteArrayOutputStream out          = new ByteArrayOutputStream();
+      MasterCipher          masterCipher = new MasterCipher(masterSecret);
+
+      writeInteger(CURRENT_VERSION_MARKER, out);
+      writeBlob(masterCipher.encryptBytes(structure.toByteArray()), out);
+
+      return out.toByteArray();
+    } catch (IOException e) {
+      throw new AssertionError(e);
+    }
+  }
+
+
+  private byte[] readBlob(FileInputStream in) throws IOException {
+    int length       = readInteger(in);
+    byte[] blobBytes = new byte[length];
+
+    in.read(blobBytes, 0, blobBytes.length);
+    return blobBytes;
+  }
+
+  private void writeBlob(byte[] blobBytes, OutputStream out) throws IOException {
+    writeInteger(blobBytes.length, out);
+    out.write(blobBytes);
+  }
+
+  private int readInteger(FileInputStream in) throws IOException {
+    byte[] integer = new byte[4];
+    in.read(integer, 0, integer.length);
+    return Conversions.byteArrayToInt(integer);
+  }
+
+  private void writeInteger(int value, OutputStream out) throws IOException {
+    byte[] valueBytes = Conversions.intToByteArray(value);
+    out.write(valueBytes);
+  }
+
+}

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

@@ -0,0 +1,86 @@
+package org.whispersystems.textsecure.storage;
+
+import android.content.Context;
+import android.util.Log;
+
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
+import org.whispersystems.libaxolotl.InvalidMessageException;
+import org.whispersystems.libaxolotl.state.PreKeyRecord;
+import org.whispersystems.libaxolotl.state.PreKeyStore;
+import org.whispersystems.textsecure.crypto.MasterSecret;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.RandomAccessFile;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+
+public class TextSecurePreKeyStore implements PreKeyStore {
+
+  public  static final String PREKEY_DIRECTORY = "prekeys";
+  private static final String TAG              = TextSecurePreKeyStore.class.getSimpleName();
+
+  private final Context      context;
+  private final MasterSecret masterSecret;
+
+  public TextSecurePreKeyStore(Context context, MasterSecret masterSecret) {
+    this.context      = context;
+    this.masterSecret = masterSecret;
+  }
+
+  @Override
+  public PreKeyRecord load(int preKeyId) throws InvalidKeyIdException {
+    try {
+      FileInputStream fin = new FileInputStream(getPreKeyFile(preKeyId));
+      return new TextSecurePreKeyRecord(masterSecret, fin);
+    } catch (IOException | InvalidMessageException e) {
+      Log.w(TAG, e);
+      throw new InvalidKeyIdException(e);
+    }
+  }
+
+  @Override
+  public void store(int preKeyId, PreKeyRecord record) {
+    try {
+      RandomAccessFile recordFile = new RandomAccessFile(getPreKeyFile(preKeyId), "rw");
+      FileChannel      out        = recordFile.getChannel();
+
+      out.position(0);
+      out.write(ByteBuffer.wrap(record.serialize()));
+      out.truncate(out.position());
+
+      recordFile.close();
+    } catch (IOException e) {
+      throw new AssertionError(e);
+    }
+  }
+
+  @Override
+  public boolean contains(int preKeyId) {
+    File record = getPreKeyFile(preKeyId);
+    return record.exists();
+  }
+
+  @Override
+  public void remove(int preKeyId) {
+    File record = getPreKeyFile(preKeyId);
+    record.delete();
+  }
+
+  private File getPreKeyFile(int preKeyId) {
+    return new File(getPreKeyDirectory(), String.valueOf(preKeyId));
+  }
+
+  private File getPreKeyDirectory() {
+    File directory = new File(context.getFilesDir(), PREKEY_DIRECTORY);
+
+    if (!directory.exists()) {
+      if (!directory.mkdirs()) {
+        Log.w(TAG, "PreKey directory creation failed!");
+      }
+    }
+
+    return directory;
+  }
+}

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

@@ -106,7 +106,6 @@ public class TextSecureSessionStore implements SessionStore {
     return results;
   }
 
-
   private File getSessionFile(long recipientId, int deviceId) {
     return new File(getSessionDirectory(), getSessionName(recipientId, deviceId));
   }

src/org/thoughtcrime/securesms/ReceiveKeyActivity.java

@@ -32,7 +32,6 @@ import android.widget.TextView;
 
 import org.thoughtcrime.securesms.crypto.DecryptingQueue;
 import org.thoughtcrime.securesms.crypto.KeyExchangeProcessor;
-import org.thoughtcrime.securesms.crypto.protocol.KeyExchangeMessage;
 import org.thoughtcrime.securesms.database.DatabaseFactory;
 import org.thoughtcrime.securesms.recipients.Recipient;
 import org.thoughtcrime.securesms.service.SendReceiveService;
@@ -45,11 +44,12 @@ import org.whispersystems.libaxolotl.InvalidMessageException;
 import org.whispersystems.libaxolotl.InvalidVersionException;
 import org.whispersystems.libaxolotl.LegacyMessageException;
 import org.whispersystems.libaxolotl.protocol.CiphertextMessage;
+import org.whispersystems.libaxolotl.protocol.KeyExchangeMessage;
 import org.whispersystems.libaxolotl.protocol.PreKeyWhisperMessage;
 import org.whispersystems.textsecure.crypto.IdentityKeyParcelable;
 import org.whispersystems.textsecure.crypto.MasterSecret;
 import org.whispersystems.textsecure.push.IncomingPushMessage;
-import org.whispersystems.textsecure.storage.InvalidKeyIdException;
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
 import org.whispersystems.textsecure.storage.RecipientDevice;
 import org.whispersystems.textsecure.util.Base64;
 import org.whispersystems.textsecure.util.InvalidNumberException;
@@ -178,7 +178,7 @@ public class ReceiveKeyActivity extends Activity {
       } else if (getIntent().getBooleanExtra("is_identity_update", false)) {
         this.identityUpdateMessage = new IdentityKey(Base64.decodeWithoutPadding(messageBody), 0);
       } else {
-        this.keyExchangeMessage = new KeyExchangeMessage(messageBody);
+        this.keyExchangeMessage = new KeyExchangeMessage(Base64.decodeWithoutPadding(messageBody));
       }
     } catch (IOException e) {
       throw new AssertionError(e);
@@ -228,7 +228,7 @@ public class ReceiveKeyActivity extends Activity {
 
               DatabaseFactory.getEncryptingSmsDatabase(ReceiveKeyActivity.this)
                              .markAsProcessedKeyExchange(messageId);
-            } catch (InvalidMessageException e) {
+            } catch (InvalidKeyException e) {
               Log.w("ReceiveKeyActivity", e);
               DatabaseFactory.getEncryptingSmsDatabase(ReceiveKeyActivity.this)
                              .markAsCorruptKeyExchange(messageId);

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

@@ -22,7 +22,6 @@ import android.content.Intent;
 import android.database.Cursor;
 import android.util.Log;
 
-import org.thoughtcrime.securesms.crypto.protocol.KeyExchangeMessage;
 import org.thoughtcrime.securesms.database.DatabaseFactory;
 import org.thoughtcrime.securesms.database.EncryptingSmsDatabase;
 import org.thoughtcrime.securesms.database.MmsDatabase;
@@ -41,10 +40,12 @@ import org.thoughtcrime.securesms.service.SendReceiveService;
 import org.thoughtcrime.securesms.sms.SmsTransportDetails;
 import org.thoughtcrime.securesms.util.TextSecurePreferences;
 import org.whispersystems.libaxolotl.DuplicateMessageException;
+import org.whispersystems.libaxolotl.InvalidKeyException;
 import org.whispersystems.libaxolotl.InvalidMessageException;
 import org.whispersystems.libaxolotl.InvalidVersionException;
 import org.whispersystems.libaxolotl.LegacyMessageException;
 import org.whispersystems.libaxolotl.SessionCipher;
+import org.whispersystems.libaxolotl.protocol.KeyExchangeMessage;
 import org.whispersystems.libaxolotl.protocol.WhisperMessage;
 import org.whispersystems.libaxolotl.state.SessionStore;
 import org.whispersystems.textsecure.crypto.MasterSecret;
@@ -52,6 +53,7 @@ import org.whispersystems.textsecure.crypto.SessionCipherFactory;
 import org.whispersystems.textsecure.push.IncomingPushMessage;
 import org.whispersystems.textsecure.storage.RecipientDevice;
 import org.whispersystems.textsecure.storage.TextSecureSessionStore;
+import org.whispersystems.textsecure.util.Base64;
 import org.whispersystems.textsecure.util.Hex;
 import org.whispersystems.textsecure.util.Util;
 
@@ -428,7 +430,7 @@ public class DecryptingQueue {
           Recipient            recipient       = RecipientFactory.getRecipientsFromString(context, originator, false)
                                                                  .getPrimaryRecipient();
           RecipientDevice      recipientDevice = new RecipientDevice(recipient.getRecipientId(), deviceId);
-          KeyExchangeMessage   message         = new KeyExchangeMessage(plaintextBody);
+          KeyExchangeMessage   message         = new KeyExchangeMessage(Base64.decodeWithoutPadding(plaintextBody));
           KeyExchangeProcessor processor       = new KeyExchangeProcessor(context, masterSecret, recipientDevice);
 
           if (processor.isStale(message)) {
@@ -440,7 +442,7 @@ public class DecryptingQueue {
         } catch (InvalidVersionException e) {
           Log.w("DecryptingQueue", e);
           DatabaseFactory.getEncryptingSmsDatabase(context).markAsInvalidVersionKeyExchange(messageId);
-        } catch (InvalidMessageException | RecipientFormattingException e) {
+        } catch (InvalidMessageException | IOException | InvalidKeyException | RecipientFormattingException e) {
           Log.w("DecryptingQueue", e);
           DatabaseFactory.getEncryptingSmsDatabase(context).markAsCorruptKeyExchange(messageId);
         } catch (LegacyMessageException e) {

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

@@ -22,7 +22,6 @@ import android.content.Context;
 import android.content.DialogInterface;
 
 import org.thoughtcrime.securesms.R;
-import org.thoughtcrime.securesms.crypto.protocol.KeyExchangeMessage;
 import org.thoughtcrime.securesms.recipients.Recipient;
 import org.thoughtcrime.securesms.sms.MessageSender;
 import org.thoughtcrime.securesms.sms.OutgoingKeyExchangeMessage;
@@ -30,11 +29,13 @@ import org.thoughtcrime.securesms.util.Dialogs;
 import org.whispersystems.libaxolotl.IdentityKeyPair;
 import org.whispersystems.libaxolotl.ecc.Curve;
 import org.whispersystems.libaxolotl.ecc.ECKeyPair;
+import org.whispersystems.libaxolotl.protocol.KeyExchangeMessage;
 import org.whispersystems.libaxolotl.state.SessionRecord;
 import org.whispersystems.libaxolotl.state.SessionStore;
 import org.whispersystems.textsecure.crypto.MasterSecret;
 import org.whispersystems.textsecure.storage.RecipientDevice;
 import org.whispersystems.textsecure.storage.TextSecureSessionStore;
+import org.whispersystems.textsecure.util.Base64;
 
 import java.security.NoSuchAlgorithmException;
 import java.security.SecureRandom;
@@ -68,11 +69,11 @@ public class KeyExchangeInitiator {
     IdentityKeyPair identityKey  = IdentityKeyUtil.getIdentityKeyPair(context, masterSecret);
 
     KeyExchangeMessage message = new KeyExchangeMessage(sequence, flags,
-                                                            baseKey.getPublicKey(),
+                                                        baseKey.getPublicKey(),
-                                                            ephemeralKey.getPublicKey(),
+                                                        ephemeralKey.getPublicKey(),
-                                                            identityKey.getPublicKey());
+                                                        identityKey.getPublicKey());
 
-    OutgoingKeyExchangeMessage textMessage     = new OutgoingKeyExchangeMessage(recipient, message.serialize());
+    OutgoingKeyExchangeMessage textMessage     = new OutgoingKeyExchangeMessage(recipient, Base64.encodeBytesWithoutPadding(message.serialize()));
     SessionStore               sessionStore    = new TextSecureSessionStore(context, masterSecret);
     SessionRecord              sessionRecord   = sessionStore.get(recipient.getRecipientId(), RecipientDevice.DEFAULT_DEVICE_ID);
 

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

@@ -2,9 +2,7 @@ package org.thoughtcrime.securesms.crypto;
 
 import android.content.Context;
 import android.content.Intent;
-import android.util.Log;
 
-import org.thoughtcrime.securesms.crypto.protocol.KeyExchangeMessage;
 import org.thoughtcrime.securesms.database.DatabaseFactory;
 import org.thoughtcrime.securesms.recipients.Recipient;
 import org.thoughtcrime.securesms.recipients.RecipientFactory;
@@ -12,25 +10,22 @@ import org.thoughtcrime.securesms.service.KeyCachingService;
 import org.thoughtcrime.securesms.service.PreKeyService;
 import org.thoughtcrime.securesms.sms.MessageSender;
 import org.thoughtcrime.securesms.sms.OutgoingKeyExchangeMessage;
-import org.thoughtcrime.securesms.util.TextSecurePreferences;
 import org.whispersystems.libaxolotl.IdentityKey;
-import org.whispersystems.libaxolotl.IdentityKeyPair;
 import org.whispersystems.libaxolotl.InvalidKeyException;
-import org.whispersystems.libaxolotl.InvalidMessageException;
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
-import org.whispersystems.libaxolotl.ecc.Curve;
+import org.whispersystems.libaxolotl.SessionBuilder;
-import org.whispersystems.libaxolotl.ecc.ECKeyPair;
+import org.whispersystems.libaxolotl.protocol.KeyExchangeMessage;
-import org.whispersystems.libaxolotl.ecc.ECPublicKey;
 import org.whispersystems.libaxolotl.protocol.PreKeyWhisperMessage;
-import org.whispersystems.libaxolotl.ratchet.RatchetingSession;
+import org.whispersystems.libaxolotl.state.IdentityKeyStore;
+import org.whispersystems.libaxolotl.state.PreKeyStore;
 import org.whispersystems.libaxolotl.state.SessionRecord;
 import org.whispersystems.libaxolotl.state.SessionStore;
 import org.whispersystems.textsecure.crypto.MasterSecret;
 import org.whispersystems.textsecure.push.PreKeyEntity;
-import org.whispersystems.textsecure.storage.InvalidKeyIdException;
-import org.whispersystems.textsecure.storage.PreKeyRecord;
 import org.whispersystems.textsecure.storage.RecipientDevice;
+import org.whispersystems.textsecure.storage.TextSecurePreKeyStore;
 import org.whispersystems.textsecure.storage.TextSecureSessionStore;
-import org.whispersystems.textsecure.util.Medium;
+import org.whispersystems.textsecure.util.Base64;
 
 /**
  * This class processes key exchange interactions.
@@ -45,6 +40,7 @@ public class KeyExchangeProcessor {
   private Context         context;
   private RecipientDevice recipientDevice;
   private MasterSecret    masterSecret;
+  private SessionBuilder  sessionBuilder;
   private SessionStore    sessionStore;
 
   public KeyExchangeProcessor(Context context, MasterSecret masterSecret, RecipientDevice recipientDevice)
@@ -52,7 +48,14 @@ public class KeyExchangeProcessor {
     this.context         = context;
     this.recipientDevice = recipientDevice;
     this.masterSecret    = masterSecret;
-    this.sessionStore    = new TextSecureSessionStore(context, masterSecret);
+
+    IdentityKeyStore identityKeyStore = new TextSecureIdentityKeyStore(context, masterSecret);
+    PreKeyStore      preKeyStore      = new TextSecurePreKeyStore(context, masterSecret);
+
+    this.sessionStore   = new TextSecureSessionStore(context, masterSecret);
+    this.sessionBuilder = new SessionBuilder(sessionStore, preKeyStore, identityKeyStore,
+                                             recipientDevice.getRecipientId(),
+                                             recipientDevice.getDeviceId());
   }
 
   public boolean isTrusted(PreKeyWhisperMessage message) {
@@ -87,83 +90,14 @@ public class KeyExchangeProcessor {
   public void processKeyExchangeMessage(PreKeyWhisperMessage message)
       throws InvalidKeyIdException, InvalidKeyException
   {
-    int         preKeyId          = message.getPreKeyId();
+    sessionBuilder.process(message);
-    ECPublicKey theirBaseKey      = message.getBaseKey();
-    ECPublicKey theirEphemeralKey = message.getWhisperMessage().getSenderEphemeral();
-    IdentityKey theirIdentityKey  = message.getIdentityKey();
-
-    Log.w("KeyExchangeProcessor", "Received pre-key with local key ID: " + preKeyId);
-
-    if (!PreKeyRecord.hasRecord(context, preKeyId) &&
-        sessionStore.contains(recipientDevice.getRecipientId(), recipientDevice.getDeviceId()))
-    {
-      Log.w("KeyExchangeProcessor", "We've already processed the prekey part, letting bundled message fall through...");
-      return;
-    }
-
-    if (!PreKeyRecord.hasRecord(context, preKeyId))
-      throw new InvalidKeyIdException("No such prekey: " + preKeyId);
-
-    SessionRecord   sessionRecord        = sessionStore.get(recipientDevice.getRecipientId(),
-                                                            recipientDevice.getDeviceId());
-    PreKeyRecord    preKeyRecord         = new PreKeyRecord(context, masterSecret, preKeyId);
-    ECKeyPair       ourBaseKey           = preKeyRecord.getKeyPair();
-    ECKeyPair       ourEphemeralKey      = ourBaseKey;
-    IdentityKeyPair ourIdentityKey       = IdentityKeyUtil.getIdentityKeyPair(context, masterSecret);
-    boolean         simultaneousInitiate = sessionRecord.getSessionState().hasPendingPreKey();
-
-    if (!simultaneousInitiate) sessionRecord.reset();
-    else                       sessionRecord.archiveCurrentState();
-
-    RatchetingSession.initializeSession(sessionRecord.getSessionState(),
-                                        ourBaseKey, theirBaseKey,
-                                        ourEphemeralKey, theirEphemeralKey,
-                                        ourIdentityKey, theirIdentityKey);
-
-    sessionRecord.getSessionState().setLocalRegistrationId(TextSecurePreferences.getLocalRegistrationId(context));
-    sessionRecord.getSessionState().setRemoteRegistrationId(message.getRegistrationId());
-
-    if (simultaneousInitiate) sessionRecord.getSessionState().setNeedsRefresh(true);
-
-    sessionStore.put(recipientDevice.getRecipientId(), recipientDevice.getDeviceId(), sessionRecord);
-
-    if (preKeyId != Medium.MAX_VALUE) {
-      PreKeyRecord.delete(context, preKeyId);
-    }
-
     PreKeyService.initiateRefresh(context, masterSecret);
-
-    DatabaseFactory.getIdentityDatabase(context)
-                   .saveIdentity(masterSecret, recipientDevice.getRecipientId(), theirIdentityKey);
   }
 
   public void processKeyExchangeMessage(PreKeyEntity message, long threadId)
       throws InvalidKeyException
   {
-    SessionRecord   sessionRecord     = sessionStore.get(recipientDevice.getRecipientId(),
+    sessionBuilder.process(message);
-                                                         recipientDevice.getDeviceId());
-    ECKeyPair       ourBaseKey        = Curve.generateKeyPair(true);
-    ECKeyPair       ourEphemeralKey   = Curve.generateKeyPair(true);
-    ECPublicKey     theirBaseKey      = message.getPublicKey();
-    ECPublicKey     theirEphemeralKey = theirBaseKey;
-    IdentityKey     theirIdentityKey  = message.getIdentityKey();
-    IdentityKeyPair ourIdentityKey    = IdentityKeyUtil.getIdentityKeyPair(context, masterSecret);
-
-    if (sessionRecord.getSessionState().getNeedsRefresh()) sessionRecord.archiveCurrentState();
-    else                                                   sessionRecord.reset();
-
-    RatchetingSession.initializeSession(sessionRecord.getSessionState(),
-                                        ourBaseKey, theirBaseKey, ourEphemeralKey,
-                                        theirEphemeralKey, ourIdentityKey, theirIdentityKey);
-
-    sessionRecord.getSessionState().setPendingPreKey(message.getKeyId(), ourBaseKey.getPublicKey());
-    sessionRecord.getSessionState().setLocalRegistrationId(TextSecurePreferences.getLocalRegistrationId(context));
-    sessionRecord.getSessionState().setRemoteRegistrationId(message.getRegistrationId());
-
-    sessionStore.put(recipientDevice.getRecipientId(), recipientDevice.getDeviceId(), sessionRecord);
-
-    DatabaseFactory.getIdentityDatabase(context)
-                   .saveIdentity(masterSecret, recipientDevice.getRecipientId(), message.getIdentityKey());
 
     if (threadId != -1) {
       broadcastSecurityUpdateEvent(context, threadId);
@@ -171,84 +105,23 @@ public class KeyExchangeProcessor {
   }
 
   public void processKeyExchangeMessage(KeyExchangeMessage message, long threadId)
-      throws InvalidMessageException
+      throws InvalidKeyException
   {
-    try {
+    KeyExchangeMessage responseMessage = sessionBuilder.process(message);
-      SessionRecord        sessionRecord = sessionStore.get(recipientDevice.getRecipientId(),
+    Recipient            recipient     = RecipientFactory.getRecipientsForIds(context,
-                                                            recipientDevice.getDeviceId());
+                                                                              String.valueOf(recipientDevice.getRecipientId()),
-      Recipient            recipient     = RecipientFactory.getRecipientsForIds(context,
+                                                                              false)
-                                                                                String.valueOf(recipientDevice.getRecipientId()),
+                                                         .getPrimaryRecipient();
-                                                                                false)
-                                                           .getPrimaryRecipient();
 
-      Log.w("KeyExchangeProcessor", "Received key exchange with sequence: " + message.getSequence());
+    if (responseMessage != null) {
-
+      String                     serializedResponse = Base64.encodeBytesWithoutPadding(responseMessage.serialize());
-      if (message.isInitiate()) {
+      OutgoingKeyExchangeMessage textMessage        = new OutgoingKeyExchangeMessage(recipient, serializedResponse);
-        ECKeyPair       ourBaseKey, ourEphemeralKey;
+      MessageSender.send(context, masterSecret, textMessage, threadId, true);
-        IdentityKeyPair ourIdentityKey;
-
-        int flags = KeyExchangeMessage.RESPONSE_FLAG;
-
-        Log.w("KeyExchangeProcessor", "KeyExchange is an initiate.");
-
-        if (!sessionRecord.getSessionState().hasPendingKeyExchange()) {
-          Log.w("KeyExchangeProcessor", "We don't have a pending initiate...");
-          ourBaseKey      = Curve.generateKeyPair(true);
-          ourEphemeralKey = Curve.generateKeyPair(true);
-          ourIdentityKey  = IdentityKeyUtil.getIdentityKeyPair(context, masterSecret);
-
-          sessionRecord.getSessionState().setPendingKeyExchange(message.getSequence(), ourBaseKey,
-                                                                ourEphemeralKey, ourIdentityKey);
-        } else {
-          Log.w("KeyExchangeProcessor", "We alredy have a pending initiate, responding as simultaneous initiate...");
-          ourBaseKey      = sessionRecord.getSessionState().getPendingKeyExchangeBaseKey();
-          ourEphemeralKey = sessionRecord.getSessionState().getPendingKeyExchangeEphemeralKey();
-          ourIdentityKey  = sessionRecord.getSessionState().getPendingKeyExchangeIdentityKey();
-          flags          |= KeyExchangeMessage.SIMULTAENOUS_INITIATE_FLAG;
-
-          sessionRecord.getSessionState().setPendingKeyExchange(message.getSequence(), ourBaseKey,
-                                                                ourEphemeralKey, ourIdentityKey);
-        }
-
-        KeyExchangeMessage ourMessage = new KeyExchangeMessage(message.getSequence(),
-                                                                   flags, ourBaseKey.getPublicKey(),
-                                                                   ourEphemeralKey.getPublicKey(),
-                                                                   ourIdentityKey.getPublicKey());
-
-        OutgoingKeyExchangeMessage textMessage = new OutgoingKeyExchangeMessage(recipient,
-                                                                                ourMessage.serialize());
-        MessageSender.send(context, masterSecret, textMessage, threadId, true);
-      }
-
-      if (message.getSequence() != sessionRecord.getSessionState().getPendingKeyExchangeSequence()) {
-        Log.w("KeyExchangeProcessor", "No matching sequence for response. " +
-            "Is simultaneous initiate response: " + message.isResponseForSimultaneousInitiate());
-        return;
-      }
-
-      ECKeyPair       ourBaseKey      = sessionRecord.getSessionState().getPendingKeyExchangeBaseKey();
-      ECKeyPair       ourEphemeralKey = sessionRecord.getSessionState().getPendingKeyExchangeEphemeralKey();
-      IdentityKeyPair ourIdentityKey  = sessionRecord.getSessionState().getPendingKeyExchangeIdentityKey();
-
-      sessionRecord.reset();
-
-      RatchetingSession.initializeSession(sessionRecord.getSessionState(),
-                                          ourBaseKey, message.getBaseKey(),
-                                          ourEphemeralKey, message.getEphemeralKey(),
-                                          ourIdentityKey, message.getIdentityKey());
-
-      sessionRecord.getSessionState().setSessionVersion(message.getVersion());
-      sessionStore.put(recipientDevice.getRecipientId(), recipientDevice.getDeviceId(), sessionRecord);
-
-      DatabaseFactory.getIdentityDatabase(context)
-                     .saveIdentity(masterSecret, recipientDevice.getRecipientId(), message.getIdentityKey());
-
-      DecryptingQueue.scheduleRogueMessages(context, masterSecret, recipient);
-
-      broadcastSecurityUpdateEvent(context, threadId);
-    } catch (InvalidKeyException e) {
-      throw new InvalidMessageException(e);
     }
+
+    DecryptingQueue.scheduleRogueMessages(context, masterSecret, recipient);
+
+    broadcastSecurityUpdateEvent(context, threadId);
   }
 
   public static void broadcastSecurityUpdateEvent(Context context, long threadId) {

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

@@ -0,0 +1,36 @@
+package org.thoughtcrime.securesms.crypto;
+
+import android.content.Context;
+
+import org.thoughtcrime.securesms.database.DatabaseFactory;
+import org.thoughtcrime.securesms.util.TextSecurePreferences;
+import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.IdentityKeyPair;
+import org.whispersystems.libaxolotl.state.IdentityKeyStore;
+import org.whispersystems.textsecure.crypto.MasterSecret;
+
+public class TextSecureIdentityKeyStore implements IdentityKeyStore {
+
+  private final Context      context;
+  private final MasterSecret masterSecret;
+
+  public TextSecureIdentityKeyStore(Context context, MasterSecret masterSecret) {
+    this.context      = context;
+    this.masterSecret = masterSecret;
+  }
+
+  @Override
+  public IdentityKeyPair getIdentityKeyPair() {
+    return IdentityKeyUtil.getIdentityKeyPair(context, masterSecret);
+  }
+
+  @Override
+  public int getLocalRegistrationId() {
+    return TextSecurePreferences.getLocalRegistrationId(context);
+  }
+
+  @Override
+  public void saveIdentity(long recipientId, IdentityKey identityKey) {
+    DatabaseFactory.getIdentityDatabase(context).saveIdentity(masterSecret, recipientId, identityKey);
+  }
+}

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

@@ -10,10 +10,10 @@ 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.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.PreKeyRecord;
 
 import java.io.IOException;
 import java.util.List;

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

@@ -32,7 +32,7 @@ import org.whispersystems.libaxolotl.state.SessionStore;
 import org.whispersystems.textsecure.crypto.MasterSecret;
 import org.whispersystems.textsecure.push.IncomingPushMessage;
 import org.whispersystems.textsecure.push.PushMessageProtos.PushMessageContent;
-import org.whispersystems.textsecure.storage.InvalidKeyIdException;
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
 import org.whispersystems.textsecure.storage.RecipientDevice;
 import org.whispersystems.textsecure.storage.TextSecureSessionStore;
 import org.whispersystems.textsecure.util.Base64;

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

@@ -18,11 +18,11 @@ 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.state.PreKeyRecord;
 import org.whispersystems.textsecure.crypto.MasterSecret;
 import org.whispersystems.textsecure.crypto.PreKeyUtil;
 import org.whispersystems.textsecure.push.ExpectationFailedException;
 import org.whispersystems.textsecure.push.PushServiceSocket;
-import org.whispersystems.textsecure.storage.PreKeyRecord;
 import org.whispersystems.textsecure.util.Util;
 
 import java.io.IOException;

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

@@ -24,7 +24,6 @@ import android.util.Pair;
 import org.thoughtcrime.securesms.crypto.DecryptingQueue;
 import org.thoughtcrime.securesms.crypto.KeyExchangeProcessor;
 import org.thoughtcrime.securesms.crypto.MasterSecretUtil;
-import org.thoughtcrime.securesms.crypto.protocol.KeyExchangeMessage;
 import org.thoughtcrime.securesms.database.DatabaseFactory;
 import org.thoughtcrime.securesms.database.EncryptingSmsDatabase;
 import org.thoughtcrime.securesms.database.SmsDatabase;
@@ -44,11 +43,13 @@ import org.whispersystems.libaxolotl.InvalidKeyException;
 import org.whispersystems.libaxolotl.InvalidMessageException;
 import org.whispersystems.libaxolotl.InvalidVersionException;
 import org.whispersystems.libaxolotl.LegacyMessageException;
+import org.whispersystems.libaxolotl.protocol.KeyExchangeMessage;
 import org.whispersystems.libaxolotl.protocol.PreKeyWhisperMessage;
 import org.whispersystems.libaxolotl.protocol.WhisperMessage;
 import org.whispersystems.textsecure.crypto.MasterSecret;
-import org.whispersystems.textsecure.storage.InvalidKeyIdException;
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
 import org.whispersystems.textsecure.storage.RecipientDevice;
+import org.whispersystems.textsecure.util.Base64;
 
 import java.io.IOException;
 import java.util.List;
@@ -111,12 +112,12 @@ public class SmsReceiver {
     Log.w("SmsReceiver", "Processing prekey message...");
 
     try {
-      Recipient              recipient        = RecipientFactory.getRecipientsFromString(context, message.getSender(), false).getPrimaryRecipient();
+      Recipient            recipient        = RecipientFactory.getRecipientsFromString(context, message.getSender(), false).getPrimaryRecipient();
-      RecipientDevice        recipientDevice  = new RecipientDevice(recipient.getRecipientId(), message.getSenderDeviceId());
+      RecipientDevice      recipientDevice  = new RecipientDevice(recipient.getRecipientId(), message.getSenderDeviceId());
       KeyExchangeProcessor processor        = new KeyExchangeProcessor(context, masterSecret, recipientDevice);
-      SmsTransportDetails    transportDetails = new SmsTransportDetails();
+      SmsTransportDetails  transportDetails = new SmsTransportDetails();
-      byte[]                 rawMessage       = transportDetails.getDecodedMessage(message.getMessageBody().getBytes());
+      byte[]               rawMessage       = transportDetails.getDecodedMessage(message.getMessageBody().getBytes());
-      PreKeyWhisperMessage   preKeyExchange   = new PreKeyWhisperMessage(rawMessage);
+      PreKeyWhisperMessage preKeyExchange   = new PreKeyWhisperMessage(rawMessage);
 
       if (processor.isTrusted(preKeyExchange)) {
         processor.processKeyExchangeMessage(preKeyExchange);
@@ -163,7 +164,7 @@ public class SmsReceiver {
       try {
         Recipient            recipient       = RecipientFactory.getRecipientsFromString(context, message.getSender(), false).getPrimaryRecipient();
         RecipientDevice      recipientDevice = new RecipientDevice(recipient.getRecipientId(), message.getSenderDeviceId());
-        KeyExchangeMessage   exchangeMessage = new KeyExchangeMessage(message.getMessageBody());
+        KeyExchangeMessage   exchangeMessage = new KeyExchangeMessage(Base64.decodeWithoutPadding(message.getMessageBody()));
         KeyExchangeProcessor processor       = new KeyExchangeProcessor(context, masterSecret, recipientDevice);
 
         if (processor.isStale(exchangeMessage)) {
@@ -179,7 +180,7 @@ public class SmsReceiver {
       } catch (InvalidVersionException e) {
         Log.w("SmsReceiver", e);
         message.setInvalidVersion(true);
-      } catch (InvalidMessageException | RecipientFormattingException e) {
+      } catch (InvalidMessageException | InvalidKeyException | IOException | RecipientFormattingException e) {
         Log.w("SmsReceiver", e);
         message.setCorrupted(true);
       } catch (LegacyMessageException e) {