Reorganize session store load/store operations.

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

@@ -0,0 +1,54 @@
+package org.whispersystems.test;
+
+import org.whispersystems.libaxolotl.state.SessionRecord;
+import org.whispersystems.libaxolotl.state.SessionState;
+
+import java.util.LinkedList;
+import java.util.List;
+
+public class InMemorySessionRecord implements SessionRecord {
+
+  private SessionState       currentSessionState;
+  private List<SessionState> previousSessionStates;
+
+  public InMemorySessionRecord() {
+    currentSessionState   = new InMemorySessionState();
+    previousSessionStates = new LinkedList<>();
+  }
+
+  public InMemorySessionRecord(SessionRecord copy) {
+    currentSessionState   = new InMemorySessionState(copy.getSessionState());
+    previousSessionStates = new LinkedList<>();
+
+    for (SessionState previousState : copy.getPreviousSessionStates()) {
+      previousSessionStates.add(new InMemorySessionState(previousState));
+    }
+  }
+
+  @Override
+  public SessionState getSessionState() {
+    return currentSessionState;
+  }
+
+  @Override
+  public List<SessionState> getPreviousSessionStates() {
+    return previousSessionStates;
+  }
+
+  @Override
+  public void reset() {
+    this.currentSessionState   = new InMemorySessionState();
+    this.previousSessionStates = new LinkedList<>();
+  }
+
+  @Override
+  public void archiveCurrentState() {
+    this.previousSessionStates.add(currentSessionState);
+    this.currentSessionState = new InMemorySessionState();
+  }
+
+  @Override
+  public byte[] serialize() {
+    throw new AssertionError();
+  }
+}

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

@@ -3,7 +3,7 @@ package org.whispersystems.test;
 import org.whispersystems.libaxolotl.IdentityKey;
 import org.whispersystems.libaxolotl.IdentityKeyPair;
 import org.whispersystems.libaxolotl.InvalidKeyException;
-import org.whispersystems.libaxolotl.SessionState;
+import org.whispersystems.libaxolotl.state.SessionState;
 import org.whispersystems.libaxolotl.ecc.ECKeyPair;
 import org.whispersystems.libaxolotl.ecc.ECPublicKey;
 import org.whispersystems.libaxolotl.ratchet.ChainKey;

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

@@ -1,44 +1,63 @@
 package org.whispersystems.test;
 
-import org.whispersystems.libaxolotl.SessionState;
-import org.whispersystems.libaxolotl.SessionStore;
+import org.whispersystems.libaxolotl.state.SessionRecord;
+import org.whispersystems.libaxolotl.state.SessionStore;
+import org.whispersystems.libaxolotl.util.Pair;
 
+import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
 
 public class InMemorySessionStore implements SessionStore {
 
-  private SessionState currentSessionState;
-  private List<SessionState> previousSessionStates;
+  private Map<Pair<Long, Integer>, SessionRecord> sessions = new HashMap<>();
 
-  private SessionState       checkedOutSessionState;
-  private List<SessionState> checkedOutPreviousSessionStates;
+  public InMemorySessionStore() {}
 
-  public InMemorySessionStore(SessionState sessionState) {
-    this.currentSessionState             = sessionState;
-    this.previousSessionStates           = new LinkedList<>();
-    this.checkedOutPreviousSessionStates = new LinkedList<>();
+  @Override
+  public SessionRecord get(long recipientId, int deviceId) {
+    if (contains(recipientId, deviceId)) {
+      return new InMemorySessionRecord(sessions.get(new Pair<>(recipientId, deviceId)));
+    } else {
+      return new InMemorySessionRecord();
+    }
   }
 
   @Override
-  public SessionState getSessionState() {
-    checkedOutSessionState = new InMemorySessionState(currentSessionState);
-    return checkedOutSessionState;
-  }
+  public List<Integer> getSubDeviceSessions(long recipientId) {
+    List<Integer> deviceIds = new LinkedList<>();
 
-  @Override
-  public List<SessionState> getPreviousSessionStates() {
-    checkedOutPreviousSessionStates = new LinkedList<>();
-    for (SessionState state : previousSessionStates) {
-      checkedOutPreviousSessionStates.add(new InMemorySessionState(state));
+    for (Pair<Long, Integer> key : sessions.keySet()) {
+      if (key.first() == recipientId) {
+        deviceIds.add(key.second());
+      }
     }
 
-    return checkedOutPreviousSessionStates;
+    return deviceIds;
   }
 
   @Override
-  public void save() {
-    this.currentSessionState   = this.checkedOutSessionState;
-    this.previousSessionStates = this.checkedOutPreviousSessionStates;
+  public void put(long recipientId, int deviceId, SessionRecord record) {
+    sessions.put(new Pair<>(recipientId, deviceId), record);
+  }
+
+  @Override
+  public boolean contains(long recipientId, int deviceId) {
+    return sessions.containsKey(new Pair<>(recipientId, deviceId));
+  }
+
+  @Override
+  public void delete(long recipientId, int deviceId) {
+    sessions.remove(new Pair<>(recipientId, deviceId));
+  }
+
+  @Override
+  public void deleteAll(long recipientId) {
+    for (Pair<Long, Integer> key : sessions.keySet()) {
+      if (key.first() == recipientId) {
+        sessions.remove(key);
+      }
+    }
   }
 }

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

@@ -9,8 +9,9 @@ import org.whispersystems.libaxolotl.InvalidKeyException;
 import org.whispersystems.libaxolotl.InvalidMessageException;
 import org.whispersystems.libaxolotl.LegacyMessageException;
 import org.whispersystems.libaxolotl.SessionCipher;
-import org.whispersystems.libaxolotl.SessionState;
-import org.whispersystems.libaxolotl.SessionStore;
+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;
@@ -24,16 +25,19 @@ public class SessionCipherTest extends AndroidTestCase {
       throws InvalidKeyException, DuplicateMessageException,
       LegacyMessageException, InvalidMessageException
   {
-    SessionState aliceSessionState = new InMemorySessionState();
-    SessionState  bobSessionState   = new InMemorySessionState();
+    SessionRecord aliceSessionRecord = new InMemorySessionRecord();
+    SessionRecord bobSessionRecord   = new InMemorySessionRecord();
 
-    initializeSessions(aliceSessionState, bobSessionState);
+    initializeSessions(aliceSessionRecord.getSessionState(), bobSessionRecord.getSessionState());
 
-    SessionStore aliceSessionStore = new InMemorySessionStore(aliceSessionState);
-    SessionStore  bobSessionStore   = new InMemorySessionStore(bobSessionState);
+    SessionStore aliceSessionStore = new InMemorySessionStore();
+    SessionStore bobSessionStore   = new InMemorySessionStore();
 
-    SessionCipher aliceCipher       = new SessionCipher(aliceSessionStore);
-    SessionCipher bobCipher         = new SessionCipher(bobSessionStore);
+    aliceSessionStore.put(2L, 1, aliceSessionRecord);
+    bobSessionStore.put(3L, 1, bobSessionRecord);
+
+    SessionCipher     aliceCipher    = new SessionCipher(aliceSessionStore, 2L, 1);
+    SessionCipher     bobCipher      = new SessionCipher(bobSessionStore, 3L, 1);
 
     byte[]            alicePlaintext = "This is a plaintext message.".getBytes();
     CiphertextMessage message        = aliceCipher.encrypt(alicePlaintext);

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

@@ -6,7 +6,7 @@ import org.whispersystems.libaxolotl.IdentityKey;
 import org.whispersystems.libaxolotl.IdentityKeyPair;
 import org.whispersystems.test.InMemorySessionState;
 import org.whispersystems.libaxolotl.InvalidKeyException;
-import org.whispersystems.libaxolotl.SessionState;
+import org.whispersystems.libaxolotl.state.SessionState;
 import org.whispersystems.libaxolotl.ecc.Curve;
 import org.whispersystems.libaxolotl.ecc.ECKeyPair;
 import org.whispersystems.libaxolotl.ecc.ECPrivateKey;