5 files changed, 337 insertions, 0 deletions
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryAxolotlStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryAxolotlStore.java
new file mode 100644
index 00000000..4e3c818d
--- /dev/null
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryAxolotlStore.java
@@ -0,0 +1,120 @@
+package org.whispersystems.libaxolotl.state.impl;
+
+import org.whispersystems.libaxolotl.AxolotlAddress;
+import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.IdentityKeyPair;
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
+import org.whispersystems.libaxolotl.state.AxolotlStore;
+import org.whispersystems.libaxolotl.state.PreKeyRecord;
+import org.whispersystems.libaxolotl.state.SessionRecord;
+import org.whispersystems.libaxolotl.state.SignedPreKeyRecord;
+
+import java.util.List;
+
+public class InMemoryAxolotlStore implements AxolotlStore {
+
+  private final InMemoryPreKeyStore       preKeyStore       = new InMemoryPreKeyStore();
+  private final InMemorySessionStore      sessionStore      = new InMemorySessionStore();
+  private final InMemorySignedPreKeyStore signedPreKeyStore = new InMemorySignedPreKeyStore();
+
+  private final InMemoryIdentityKeyStore  identityKeyStore;
+
+  public InMemoryAxolotlStore(IdentityKeyPair identityKeyPair, int registrationId) {
+    this.identityKeyStore = new InMemoryIdentityKeyStore(identityKeyPair, registrationId);
+  }
+
+  @Override
+  public IdentityKeyPair getIdentityKeyPair() {
+    return identityKeyStore.getIdentityKeyPair();
+  }
+
+  @Override
+  public int getLocalRegistrationId() {
+    return identityKeyStore.getLocalRegistrationId();
+  }
+
+  @Override
+  public void saveIdentity(String name, IdentityKey identityKey) {
+    identityKeyStore.saveIdentity(name, identityKey);
+  }
+
+  @Override
+  public boolean isTrustedIdentity(String name, IdentityKey identityKey) {
+    return identityKeyStore.isTrustedIdentity(name, identityKey);
+  }
+
+  @Override
+  public PreKeyRecord loadPreKey(int preKeyId) throws InvalidKeyIdException {
+    return preKeyStore.loadPreKey(preKeyId);
+  }
+
+  @Override
+  public void storePreKey(int preKeyId, PreKeyRecord record) {
+    preKeyStore.storePreKey(preKeyId, record);
+  }
+
+  @Override
+  public boolean containsPreKey(int preKeyId) {
+    return preKeyStore.containsPreKey(preKeyId);
+  }
+
+  @Override
+  public void removePreKey(int preKeyId) {
+    preKeyStore.removePreKey(preKeyId);
+  }
+
+  @Override
+  public SessionRecord loadSession(AxolotlAddress address) {
+    return sessionStore.loadSession(address);
+  }
+
+  @Override
+  public List<Integer> getSubDeviceSessions(String name) {
+    return sessionStore.getSubDeviceSessions(name);
+  }
+
+  @Override
+  public void storeSession(AxolotlAddress address, SessionRecord record) {
+    sessionStore.storeSession(address, record);
+  }
+
+  @Override
+  public boolean containsSession(AxolotlAddress address) {
+    return sessionStore.containsSession(address);
+  }
+
+  @Override
+  public void deleteSession(AxolotlAddress address) {
+    sessionStore.deleteSession(address);
+  }
+
+  @Override
+  public void deleteAllSessions(String name) {
+    sessionStore.deleteAllSessions(name);
+  }
+
+  @Override
+  public SignedPreKeyRecord loadSignedPreKey(int signedPreKeyId) throws InvalidKeyIdException {
+    return signedPreKeyStore.loadSignedPreKey(signedPreKeyId);
+  }
+
+  @Override
+  public List<SignedPreKeyRecord> loadSignedPreKeys() {
+    return signedPreKeyStore.loadSignedPreKeys();
+  }
+
+  @Override
+  public void storeSignedPreKey(int signedPreKeyId, SignedPreKeyRecord record) {
+    signedPreKeyStore.storeSignedPreKey(signedPreKeyId, record);
+  }
+
+  @Override
+  public boolean containsSignedPreKey(int signedPreKeyId) {
+    return signedPreKeyStore.containsSignedPreKey(signedPreKeyId);
+  }
+
+  @Override
+  public void removeSignedPreKey(int signedPreKeyId) {
+    signedPreKeyStore.removeSignedPreKey(signedPreKeyId);
+  }
+}
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryIdentityKeyStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryIdentityKeyStore.java
new file mode 100644
index 00000000..b1b04e08
--- /dev/null
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryIdentityKeyStore.java
@@ -0,0 +1,46 @@
+package org.whispersystems.libaxolotl.state.impl;
+
+import org.whispersystems.libaxolotl.IdentityKey;
+import org.whispersystems.libaxolotl.IdentityKeyPair;
+import org.whispersystems.libaxolotl.ecc.Curve;
+import org.whispersystems.libaxolotl.ecc.ECKeyPair;
+import org.whispersystems.libaxolotl.state.IdentityKeyStore;
+
+import java.security.NoSuchAlgorithmException;
+import java.security.SecureRandom;
+import java.util.HashMap;
+import java.util.Map;
+
+public class InMemoryIdentityKeyStore implements IdentityKeyStore {
+
+  private final Map<String, IdentityKey> trustedKeys = new HashMap<>();
+
+  private final IdentityKeyPair identityKeyPair;
+  private final int             localRegistrationId;
+
+  public InMemoryIdentityKeyStore(IdentityKeyPair identityKeyPair, int localRegistrationId) {
+    this.identityKeyPair     = identityKeyPair;
+    this.localRegistrationId = localRegistrationId;
+  }
+
+  @Override
+  public IdentityKeyPair getIdentityKeyPair() {
+    return identityKeyPair;
+  }
+
+  @Override
+  public int getLocalRegistrationId() {
+    return localRegistrationId;
+  }
+
+  @Override
+  public void saveIdentity(String name, IdentityKey identityKey) {
+    trustedKeys.put(name, identityKey);
+  }
+
+  @Override
+  public boolean isTrustedIdentity(String name, IdentityKey identityKey) {
+    IdentityKey trusted = trustedKeys.get(name);
+    return (trusted == null || trusted.equals(identityKey));
+  }
+}
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryPreKeyStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryPreKeyStore.java
new file mode 100644
index 00000000..80651f21
--- /dev/null
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemoryPreKeyStore.java
@@ -0,0 +1,42 @@
+package org.whispersystems.libaxolotl.state.impl;
+
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
+import org.whispersystems.libaxolotl.state.PreKeyRecord;
+import org.whispersystems.libaxolotl.state.PreKeyStore;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+public class InMemoryPreKeyStore implements PreKeyStore {
+
+  private final Map<Integer, byte[]> store = new HashMap<>();
+
+  @Override
+  public PreKeyRecord loadPreKey(int preKeyId) throws InvalidKeyIdException {
+    try {
+      if (!store.containsKey(preKeyId)) {
+        throw new InvalidKeyIdException("No such prekeyrecord!");
+      }
+
+      return new PreKeyRecord(store.get(preKeyId));
+    } catch (IOException e) {
+      throw new AssertionError(e);
+    }
+  }
+
+  @Override
+  public void storePreKey(int preKeyId, PreKeyRecord record) {
+    store.put(preKeyId, record.serialize());
+  }
+
+  @Override
+  public boolean containsPreKey(int preKeyId) {
+    return store.containsKey(preKeyId);
+  }
+
+  @Override
+  public void removePreKey(int preKeyId) {
+    store.remove(preKeyId);
+  }
+}
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemorySessionStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemorySessionStore.java
new file mode 100644
index 00000000..f5270c7d
--- /dev/null
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemorySessionStore.java
@@ -0,0 +1,70 @@
+package org.whispersystems.libaxolotl.state.impl;
+
+import org.whispersystems.libaxolotl.AxolotlAddress;
+import org.whispersystems.libaxolotl.state.SessionRecord;
+import org.whispersystems.libaxolotl.state.SessionStore;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+
+public class InMemorySessionStore implements SessionStore {
+
+  private Map<AxolotlAddress, byte[]> sessions = new HashMap<>();
+
+  public InMemorySessionStore() {}
+
+  @Override
+  public synchronized SessionRecord loadSession(AxolotlAddress remoteAddress) {
+    try {
+      if (containsSession(remoteAddress)) {
+        return new SessionRecord(sessions.get(remoteAddress));
+      } else {
+        return new SessionRecord();
+      }
+    } catch (IOException e) {
+      throw new AssertionError(e);
+    }
+  }
+
+  @Override
+  public synchronized List<Integer> getSubDeviceSessions(String name) {
+    List<Integer> deviceIds = new LinkedList<>();
+
+    for (AxolotlAddress key : sessions.keySet()) {
+      if (key.getName().equals(name) &&
+          key.getDeviceId() != 1)
+      {
+        deviceIds.add(key.getDeviceId());
+      }
+    }
+
+    return deviceIds;
+  }
+
+  @Override
+  public synchronized void storeSession(AxolotlAddress address, SessionRecord record) {
+    sessions.put(address, record.serialize());
+  }
+
+  @Override
+  public synchronized boolean containsSession(AxolotlAddress address) {
+    return sessions.containsKey(address);
+  }
+
+  @Override
+  public synchronized void deleteSession(AxolotlAddress address) {
+    sessions.remove(address);
+  }
+
+  @Override
+  public synchronized void deleteAllSessions(String name) {
+    for (AxolotlAddress key : sessions.keySet()) {
+      if (key.getName().equals(name)) {
+        sessions.remove(key);
+      }
+    }
+  }
+}
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemorySignedPreKeyStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemorySignedPreKeyStore.java
new file mode 100644
index 00000000..bab83137
--- /dev/null
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/impl/InMemorySignedPreKeyStore.java
@@ -0,0 +1,59 @@
+package org.whispersystems.libaxolotl.state.impl;
+
+import org.whispersystems.libaxolotl.InvalidKeyIdException;
+import org.whispersystems.libaxolotl.state.SignedPreKeyRecord;
+import org.whispersystems.libaxolotl.state.SignedPreKeyStore;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+
+public class InMemorySignedPreKeyStore implements SignedPreKeyStore {
+
+  private final Map<Integer, byte[]> store = new HashMap<>();
+
+  @Override
+  public SignedPreKeyRecord loadSignedPreKey(int signedPreKeyId) throws InvalidKeyIdException {
+    try {
+      if (!store.containsKey(signedPreKeyId)) {
+        throw new InvalidKeyIdException("No such signedprekeyrecord! " + signedPreKeyId);
+      }
+
+      return new SignedPreKeyRecord(store.get(signedPreKeyId));
+    } catch (IOException e) {
+      throw new AssertionError(e);
+    }
+  }
+
+  @Override
+  public List<SignedPreKeyRecord> loadSignedPreKeys() {
+    try {
+      List<SignedPreKeyRecord> results = new LinkedList<>();
+
+      for (byte[] serialized : store.values()) {
+        results.add(new SignedPreKeyRecord(serialized));
+      }
+
+      return results;
+    } catch (IOException e) {
+      throw new AssertionError(e);
+    }
+  }
+
+  @Override
+  public void storeSignedPreKey(int signedPreKeyId, SignedPreKeyRecord record) {
+    store.put(signedPreKeyId, record.serialize());
+  }
+
+  @Override
+  public boolean containsSignedPreKey(int signedPreKeyId) {
+    return store.containsKey(signedPreKeyId);
+  }
+
+  @Override
+  public void removeSignedPreKey(int signedPreKeyId) {
+    store.remove(signedPreKeyId);
+  }
+}