Use more generic AxololAddress for identifying recipients.

6 files changed, 110 insertions, 85 deletions

diff --git a/java/src/main/java/org/whispersystems/libaxolotl/AxolotlAddress.java b/java/src/main/java/org/whispersystems/libaxolotl/AxolotlAddress.java
new file mode 100644
index 00000000..9d0476c0
--- /dev/null
+++ b/java/src/main/java/org/whispersystems/libaxolotl/AxolotlAddress.java
@@ -0,0 +1,39 @@
+package org.whispersystems.libaxolotl;
+
+public class AxolotlAddress {
+
+  private final String name;
+  private final int    deviceId;
+
+  public AxolotlAddress(String name, int deviceId) {
+    this.name     = name;
+    this.deviceId = deviceId;
+  }
+
+  public String getName() {
+    return name;
+  }
+
+  public int getDeviceId() {
+    return deviceId;
+  }
+
+  @Override
+  public String toString() {
+    return name + ":" + deviceId;
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    if (other == null)                       return false;
+    if (!(other instanceof  AxolotlAddress)) return false;
+
+    AxolotlAddress that = (AxolotlAddress)other;
+    return this.name.equals(that.name) && this.deviceId == that.deviceId;
+  }
+
+  @Override
+  public int hashCode() {
+    return this.name.hashCode() ^ this.deviceId;
+  }
+}
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/SessionBuilder.java b/java/src/main/java/org/whispersystems/libaxolotl/SessionBuilder.java
index 736d9ab1..826fcb44 100644
--- a/java/src/main/java/org/whispersystems/libaxolotl/SessionBuilder.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/SessionBuilder.java
@@ -49,8 +49,7 @@ public class SessionBuilder {
   private final PreKeyStore       preKeyStore;
   private final SignedPreKeyStore signedPreKeyStore;
   private final IdentityKeyStore  identityKeyStore;
-  private final long              recipientId;
-  private final int               deviceId;
+  private final AxolotlAddress    remoteAddress;
 
   /**
    * Constructs a SessionBuilder.
@@ -58,31 +57,28 @@ public class SessionBuilder {
    * @param sessionStore The {@link org.whispersystems.libaxolotl.state.SessionStore} to store the constructed session in.
    * @param preKeyStore The {@link  org.whispersystems.libaxolotl.state.PreKeyStore} where the client's local {@link org.whispersystems.libaxolotl.state.PreKeyRecord}s are stored.
    * @param identityKeyStore The {@link org.whispersystems.libaxolotl.state.IdentityKeyStore} containing the client's identity key information.
-   * @param recipientId The recipient ID of the remote user to build a session with.
-   * @param deviceId The device ID of the remote user's physical device.
+   * @param remoteAddress The address of the remote user to build a session with.
    */
   public SessionBuilder(SessionStore sessionStore,
                         PreKeyStore preKeyStore,
                         SignedPreKeyStore signedPreKeyStore,
                         IdentityKeyStore identityKeyStore,
-                        long recipientId, int deviceId)
+                        AxolotlAddress remoteAddress)
   {
     this.sessionStore      = sessionStore;
     this.preKeyStore       = preKeyStore;
     this.signedPreKeyStore = signedPreKeyStore;
     this.identityKeyStore  = identityKeyStore;
-    this.recipientId       = recipientId;
-    this.deviceId          = deviceId;
+    this.remoteAddress     = remoteAddress;
   }
 
   /**
    * Constructs a SessionBuilder
    * @param store The {@link org.whispersystems.libaxolotl.state.AxolotlStore} to store all state information in.
-   * @param recipientId The recipient ID of the remote user to build a session with.
-   * @param deviceId The device ID of the remote user's physical device.
+   * @param remoteAddress The address of the remote user to build a session with.
    */
-  public SessionBuilder(AxolotlStore store, long recipientId, int deviceId) {
-    this(store, store, store, store, recipientId, deviceId);
+  public SessionBuilder(AxolotlStore store, AxolotlAddress remoteAddress) {
+    this(store, store, store, store, remoteAddress);
   }
 
   /**
@@ -107,7 +103,7 @@ public class SessionBuilder {
 
     Optional<Integer> unsignedPreKeyId;
 
-    if (!identityKeyStore.isTrustedIdentity(recipientId, theirIdentityKey)) {
+    if (!identityKeyStore.isTrustedIdentity(remoteAddress.getName(), theirIdentityKey)) {
       throw new UntrustedIdentityException();
     }
 
@@ -117,7 +113,7 @@ public class SessionBuilder {
       default: throw new AssertionError("Unknown version: " + messageVersion);
     }
 
-    identityKeyStore.saveIdentity(recipientId, theirIdentityKey);
+    identityKeyStore.saveIdentity(remoteAddress.getName(), theirIdentityKey);
     return unsignedPreKeyId;
   }
 
@@ -169,7 +165,7 @@ public class SessionBuilder {
     }
 
     if (!preKeyStore.containsPreKey(message.getPreKeyId().get()) &&
-        sessionStore.containsSession(recipientId, deviceId))
+        sessionStore.containsSession(remoteAddress))
     {
       Log.w(TAG, "We've already processed the prekey part of this V2 session, letting bundled message fall through...");
       return Optional.absent();
@@ -214,7 +210,7 @@ public class SessionBuilder {
    */
   public void process(PreKeyBundle preKey) throws InvalidKeyException, UntrustedIdentityException {
     synchronized (SessionCipher.SESSION_LOCK) {
-      if (!identityKeyStore.isTrustedIdentity(recipientId, preKey.getIdentityKey())) {
+      if (!identityKeyStore.isTrustedIdentity(remoteAddress.getName(), preKey.getIdentityKey())) {
         throw new UntrustedIdentityException();
       }
 
@@ -231,7 +227,7 @@ public class SessionBuilder {
       }
 
       boolean               supportsV3           = preKey.getSignedPreKey() != null;
-      SessionRecord         sessionRecord        = sessionStore.loadSession(recipientId, deviceId);
+      SessionRecord         sessionRecord        = sessionStore.loadSession(remoteAddress);
       ECKeyPair             ourBaseKey           = Curve.generateKeyPair();
       ECPublicKey           theirSignedPreKey    = supportsV3 ? preKey.getSignedPreKey() : preKey.getPreKey();
       Optional<ECPublicKey> theirOneTimePreKey   = Optional.fromNullable(preKey.getPreKey());
@@ -258,8 +254,8 @@ public class SessionBuilder {
       sessionRecord.getSessionState().setRemoteRegistrationId(preKey.getRegistrationId());
       sessionRecord.getSessionState().setAliceBaseKey(ourBaseKey.getPublicKey().serialize());
 
-      sessionStore.storeSession(recipientId, deviceId, sessionRecord);
-      identityKeyStore.saveIdentity(recipientId, preKey.getIdentityKey());
+      sessionStore.storeSession(remoteAddress, sessionRecord);
+      identityKeyStore.saveIdentity(remoteAddress.getName(), preKey.getIdentityKey());
     }
   }
 
@@ -275,7 +271,7 @@ public class SessionBuilder {
       throws InvalidKeyException, UntrustedIdentityException, StaleKeyExchangeException
   {
     synchronized (SessionCipher.SESSION_LOCK) {
-      if (!identityKeyStore.isTrustedIdentity(recipientId, message.getIdentityKey())) {
+      if (!identityKeyStore.isTrustedIdentity(remoteAddress.getName(), message.getIdentityKey())) {
         throw new UntrustedIdentityException();
       }
 
@@ -290,7 +286,7 @@ public class SessionBuilder {
 
   private KeyExchangeMessage processInitiate(KeyExchangeMessage message) throws InvalidKeyException {
     int           flags         = KeyExchangeMessage.RESPONSE_FLAG;
-    SessionRecord sessionRecord = sessionStore.loadSession(recipientId, deviceId);
+    SessionRecord sessionRecord = sessionStore.loadSession(remoteAddress);
 
     if (message.getVersion() >= 3 &&
         !Curve.verifySignature(message.getIdentityKey().getPublicKey(),
@@ -325,8 +321,8 @@ public class SessionBuilder {
                                         Math.min(message.getMaxVersion(), CiphertextMessage.CURRENT_VERSION),
                                         parameters);
 
-    sessionStore.storeSession(recipientId, deviceId, sessionRecord);
-    identityKeyStore.saveIdentity(recipientId, message.getIdentityKey());
+    sessionStore.storeSession(remoteAddress, sessionRecord);
+    identityKeyStore.saveIdentity(remoteAddress.getName(), message.getIdentityKey());
 
     byte[] baseKeySignature = Curve.calculateSignature(parameters.getOurIdentityKey().getPrivateKey(),
                                                        parameters.getOurBaseKey().getPublicKey().serialize());
@@ -341,7 +337,7 @@ public class SessionBuilder {
   private void processResponse(KeyExchangeMessage message)
       throws StaleKeyExchangeException, InvalidKeyException
   {
-    SessionRecord sessionRecord                  = sessionStore.loadSession(recipientId, deviceId);
+    SessionRecord sessionRecord                  = sessionStore.loadSession(remoteAddress);
     SessionState  sessionState                   = sessionRecord.getSessionState();
     boolean       hasPendingKeyExchange          = sessionState.hasPendingKeyExchange();
     boolean       isSimultaneousInitiateResponse = message.isResponseForSimultaneousInitiate();
@@ -375,8 +371,8 @@ public class SessionBuilder {
       throw new InvalidKeyException("Base key signature doesn't match!");
     }
 
-    sessionStore.storeSession(recipientId, deviceId, sessionRecord);
-    identityKeyStore.saveIdentity(recipientId, message.getIdentityKey());
+    sessionStore.storeSession(remoteAddress, sessionRecord);
+    identityKeyStore.saveIdentity(remoteAddress.getName(), message.getIdentityKey());
 
   }
 
@@ -394,10 +390,10 @@ public class SessionBuilder {
         ECKeyPair       ratchetKey       = Curve.generateKeyPair();
         IdentityKeyPair identityKey      = identityKeyStore.getIdentityKeyPair();
         byte[]          baseKeySignature = Curve.calculateSignature(identityKey.getPrivateKey(), baseKey.getPublicKey().serialize());
-        SessionRecord   sessionRecord    = sessionStore.loadSession(recipientId, deviceId);
+        SessionRecord   sessionRecord    = sessionStore.loadSession(remoteAddress);
 
         sessionRecord.getSessionState().setPendingKeyExchange(sequence, baseKey, ratchetKey, identityKey);
-        sessionStore.storeSession(recipientId, deviceId, sessionRecord);
+        sessionStore.storeSession(remoteAddress, sessionRecord);
 
         return new KeyExchangeMessage(2, sequence, flags, baseKey.getPublicKey(), baseKeySignature,
                                       ratchetKey.getPublicKey(), identityKey.getPublicKey());
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/SessionCipher.java b/java/src/main/java/org/whispersystems/libaxolotl/SessionCipher.java
index 381dedb8..b9cc6ac1 100644
--- a/java/src/main/java/org/whispersystems/libaxolotl/SessionCipher.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/SessionCipher.java
@@ -67,8 +67,7 @@ public class SessionCipher {
   private final SessionStore   sessionStore;
   private final SessionBuilder sessionBuilder;
   private final PreKeyStore    preKeyStore;
-  private final long           recipientId;
-  private final int            deviceId;
+  private final AxolotlAddress remoteAddress;
 
   /**
    * Construct a SessionCipher for encrypt/decrypt operations on a session.
@@ -76,23 +75,21 @@ public class SessionCipher {
    * and stored using {@link SessionBuilder}.
    *
    * @param  sessionStore The {@link SessionStore} that contains a session for this recipient.
-   * @param  recipientId  The remote ID that messages will be encrypted to or decrypted from.
-   * @param  deviceId     The device corresponding to the recipientId.
+   * @param  remoteAddress  The remote address that messages will be encrypted to or decrypted from.
    */
   public SessionCipher(SessionStore sessionStore, PreKeyStore preKeyStore,
                        SignedPreKeyStore signedPreKeyStore, IdentityKeyStore identityKeyStore,
-                       long recipientId, int deviceId)
+                       AxolotlAddress remoteAddress)
   {
     this.sessionStore   = sessionStore;
-    this.recipientId    = recipientId;
-    this.deviceId       = deviceId;
     this.preKeyStore    = preKeyStore;
+    this.remoteAddress  = remoteAddress;
     this.sessionBuilder = new SessionBuilder(sessionStore, preKeyStore, signedPreKeyStore,
-                                             identityKeyStore, recipientId, deviceId);
+                                             identityKeyStore, remoteAddress);
   }
 
-  public SessionCipher(AxolotlStore store, long recipientId, int deviceId) {
-    this(store, store, store, store, recipientId, deviceId);
+  public SessionCipher(AxolotlStore store, AxolotlAddress remoteAddress) {
+    this(store, store, store, store, remoteAddress);
   }
 
   /**
@@ -103,7 +100,7 @@ public class SessionCipher {
    */
   public CiphertextMessage encrypt(byte[] paddedMessage) {
     synchronized (SESSION_LOCK) {
-      SessionRecord sessionRecord   = sessionStore.loadSession(recipientId, deviceId);
+      SessionRecord sessionRecord   = sessionStore.loadSession(remoteAddress);
       SessionState  sessionState    = sessionRecord.getSessionState();
       ChainKey      chainKey        = sessionState.getSenderChainKey();
       MessageKeys   messageKeys     = chainKey.getMessageKeys();
@@ -129,7 +126,7 @@ public class SessionCipher {
       }
 
       sessionState.setSenderChainKey(chainKey.getNextChainKey());
-      sessionStore.storeSession(recipientId, deviceId, sessionRecord);
+      sessionStore.storeSession(remoteAddress, sessionRecord);
       return ciphertextMessage;
     }
   }
@@ -182,13 +179,13 @@ public class SessionCipher {
              InvalidKeyIdException, InvalidKeyException, UntrustedIdentityException
   {
     synchronized (SESSION_LOCK) {
-      SessionRecord     sessionRecord    = sessionStore.loadSession(recipientId, deviceId);
+      SessionRecord     sessionRecord    = sessionStore.loadSession(remoteAddress);
       Optional<Integer> unsignedPreKeyId = sessionBuilder.process(sessionRecord, ciphertext);
       byte[]            plaintext        = decrypt(sessionRecord, ciphertext.getWhisperMessage());
 
       callback.handlePlaintext(plaintext);
 
-      sessionStore.storeSession(recipientId, deviceId, sessionRecord);
+      sessionStore.storeSession(remoteAddress, sessionRecord);
 
       if (unsignedPreKeyId.isPresent()) {
         preKeyStore.removePreKey(unsignedPreKeyId.get());
@@ -241,16 +238,16 @@ public class SessionCipher {
   {
     synchronized (SESSION_LOCK) {
 
-      if (!sessionStore.containsSession(recipientId, deviceId)) {
-        throw new NoSessionException("No session for: " + recipientId + ", " + deviceId);
+      if (!sessionStore.containsSession(remoteAddress)) {
+        throw new NoSessionException("No session for: " + remoteAddress);
       }
 
-      SessionRecord sessionRecord = sessionStore.loadSession(recipientId, deviceId);
+      SessionRecord sessionRecord = sessionStore.loadSession(remoteAddress);
       byte[]        plaintext     = decrypt(sessionRecord, ciphertext);
 
       callback.handlePlaintext(plaintext);
 
-      sessionStore.storeSession(recipientId, deviceId, sessionRecord);
+      sessionStore.storeSession(remoteAddress, sessionRecord);
 
       return plaintext;
     }
@@ -325,18 +322,18 @@ public class SessionCipher {
 
   public int getRemoteRegistrationId() {
     synchronized (SESSION_LOCK) {
-      SessionRecord record = sessionStore.loadSession(recipientId, deviceId);
+      SessionRecord record = sessionStore.loadSession(remoteAddress);
       return record.getSessionState().getRemoteRegistrationId();
     }
   }
 
   public int getSessionVersion() {
     synchronized (SESSION_LOCK) {
-      if (!sessionStore.containsSession(recipientId, deviceId)) {
-        throw new IllegalStateException(String.format("No session for (%d, %d)!", recipientId, deviceId));
+      if (!sessionStore.containsSession(remoteAddress)) {
+        throw new IllegalStateException(String.format("No session for (%s)!", remoteAddress));
       }
 
-      SessionRecord record = sessionStore.loadSession(recipientId, deviceId);
+      SessionRecord record = sessionStore.loadSession(remoteAddress);
       return record.getSessionState().getSessionVersion();
     }
   }
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/groups/SenderKeyName.java b/java/src/main/java/org/whispersystems/libaxolotl/groups/SenderKeyName.java
index ce4325ff..cf2e48f2 100644
--- a/java/src/main/java/org/whispersystems/libaxolotl/groups/SenderKeyName.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/groups/SenderKeyName.java
@@ -16,35 +16,31 @@
  */
 package org.whispersystems.libaxolotl.groups;
 
+import org.whispersystems.libaxolotl.AxolotlAddress;
+
 /**
  * A representation of a (groupId + senderId + deviceId) tuple.
  */
 public class SenderKeyName {
 
-  private final String groupId;
-  private final long   senderId;
-  private final int    deviceId;
+  private final String         groupId;
+  private final AxolotlAddress sender;
 
-  public SenderKeyName(String groupId, long senderId, int deviceId) {
+  public SenderKeyName(String groupId, AxolotlAddress sender) {
     this.groupId  = groupId;
-    this.senderId = senderId;
-    this.deviceId = deviceId;
+    this.sender   = sender;
   }
 
   public String getGroupId() {
     return groupId;
   }
 
-  public long getSenderId() {
-    return senderId;
-  }
-
-  public int getDeviceId() {
-    return deviceId;
+  public AxolotlAddress getSender() {
+    return sender;
   }
 
   public String serialize() {
-    return groupId + "::" + String.valueOf(senderId) + "::" + String.valueOf(deviceId);
+    return groupId + "::" + sender.getName() + "::" + String.valueOf(sender.getDeviceId());
   }
 
   @Override
@@ -56,13 +52,12 @@ public class SenderKeyName {
 
     return
         this.groupId.equals(that.groupId) &&
-        this.senderId == that.senderId &&
-        this.deviceId == that.deviceId;
+        this.sender.equals(that.sender);
   }
 
   @Override
   public int hashCode() {
-    return this.groupId.hashCode() ^ (int)this.senderId ^ this.deviceId;
+    return this.groupId.hashCode() ^ this.sender.hashCode();
   }
 
 }
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/state/IdentityKeyStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/IdentityKeyStore.java
index d2024f78..86e18adf 100644
--- a/java/src/main/java/org/whispersystems/libaxolotl/state/IdentityKeyStore.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/IdentityKeyStore.java
@@ -32,10 +32,10 @@ public interface IdentityKeyStore {
    * <p>
    * Store a remote client's identity key as trusted.
    *
-   * @param recipientId The recipient ID of the remote client.
+   * @param name        The name of the remote client.
    * @param identityKey The remote client's identity key.
    */
-  public void            saveIdentity(long recipientId, IdentityKey identityKey);
+  public void            saveIdentity(String name, IdentityKey identityKey);
 
 
   /**
@@ -48,10 +48,10 @@ public interface IdentityKeyStore {
    * store.  Only if it mismatches an entry in the local store is it considered
    * 'untrusted.'
    *
-   * @param recipientId The recipient ID of the remote client.
+   * @param name        The name of the remote client.
    * @param identityKey The identity key to verify.
    * @return true if trusted, false if untrusted.
    */
-  public boolean         isTrustedIdentity(long recipientId, IdentityKey identityKey);
+  public boolean         isTrustedIdentity(String name, IdentityKey identityKey);
 
 }
diff --git a/java/src/main/java/org/whispersystems/libaxolotl/state/SessionStore.java b/java/src/main/java/org/whispersystems/libaxolotl/state/SessionStore.java
index c5ad00b0..a2a20786 100644
--- a/java/src/main/java/org/whispersystems/libaxolotl/state/SessionStore.java
+++ b/java/src/main/java/org/whispersystems/libaxolotl/state/SessionStore.java
@@ -1,5 +1,7 @@
 package org.whispersystems.libaxolotl.state;
 
+import org.whispersystems.libaxolotl.AxolotlAddress;
+
 import java.util.List;
 
 /**
@@ -19,50 +21,46 @@ public interface SessionStore {
    * durable session state (what is returned by subsequent calls to this method) without the
    * store method being called here first.
    *
-   * @param recipientId The recipientID of the remote client.
-   * @param deviceId The deviceID of the remote client.
+   * @param address The name and device ID of the remote client.
    * @return a copy of the SessionRecord corresponding to the recipientId + deviceId tuple, or
    *         a new SessionRecord if one does not currently exist.
    */
-  public SessionRecord loadSession(long recipientId, int deviceId);
+  public SessionRecord loadSession(AxolotlAddress address);
 
   /**
    * Returns all known devices with active sessions for a recipient
    *
-   * @param recipientId the recipient ID.
+   * @param name the name of the client.
    * @return all known sub-devices with active sessions.
    */
-  public List<Integer> getSubDeviceSessions(long recipientId);
+  public List<Integer> getSubDeviceSessions(String name);
 
   /**
    * Commit to storage the {@link SessionRecord} for a given recipientId + deviceId tuple.
-   * @param recipientId the recipient ID of the remote client.
-   * @param deviceId the device ID of the remote client.
+   * @param address the address of the remote client.
    * @param record the current SessionRecord for the remote client.
    */
-  public void storeSession(long recipientId, int deviceId, SessionRecord record);
+  public void storeSession(AxolotlAddress address, SessionRecord record);
 
   /**
    * Determine whether there is a committed {@link SessionRecord} for a recipientId + deviceId tuple.
-   * @param recipientId the recipient ID of the remote client.
-   * @param deviceId the device ID of the remote client.
+   * @param address the address of the remote client.
    * @return true if a {@link SessionRecord} exists, false otherwise.
    */
-  public boolean containsSession(long recipientId, int deviceId);
+  public boolean containsSession(AxolotlAddress address);
 
   /**
    * Remove a {@link SessionRecord} for a recipientId + deviceId tuple.
    *
-   * @param recipientId the recipient ID of the remote client.
-   * @param deviceId the device ID of the remote client.
+   * @param address the address of the remote client.
    */
-  public void deleteSession(long recipientId, int deviceId);
+  public void deleteSession(AxolotlAddress address);
 
   /**
    * Remove the {@link SessionRecord}s corresponding to all devices of a recipientId.
    *
-   * @param recipientId the recipient ID of the remote client.
+   * @param name the name of the remote client.
    */
-  public void deleteAllSessions(long recipientId);
+  public void deleteAllSessions(String name);
 
 }