Separate session store for PNI.

app/src/main/java/org/thoughtcrime/securesms/components/settings/conversation/InternalConversationSettingsFragment.kt

@@ -154,10 +154,16 @@ class InternalConversationSettingsFragment : DSLSettingsFragment(
               .setNegativeButton(android.R.string.cancel) { d, _ -> d.dismiss() }
               .setPositiveButton(android.R.string.ok) { _, _ ->
                 if (recipient.hasAci()) {
-                  SignalDatabase.sessions.deleteAllFor(recipient.requireAci().toString())
+                  SignalDatabase.sessions.deleteAllFor(accountId = Recipient.self().requireAci(), addressName = recipient.requireAci().toString())
+                  SignalDatabase.sessions.deleteAllFor(accountId = Recipient.self().requirePni(), addressName = recipient.requireAci().toString())
                 }
                 if (recipient.hasE164()) {
-                  SignalDatabase.sessions.deleteAllFor(recipient.requireE164())
+                  SignalDatabase.sessions.deleteAllFor(accountId = Recipient.self().requireAci(), addressName = recipient.requireE164())
+                  SignalDatabase.sessions.deleteAllFor(accountId = Recipient.self().requirePni(), addressName = recipient.requireE164())
+                }
+                if (recipient.hasPni()) {
+                  SignalDatabase.sessions.deleteAllFor(accountId = Recipient.self().requireAci(), addressName = recipient.requirePni().toString())
+                  SignalDatabase.sessions.deleteAllFor(accountId = Recipient.self().requirePni(), addressName = recipient.requirePni().toString())
                 }
               }
               .show()

app/src/main/java/org/thoughtcrime/securesms/contacts/sync/DirectoryHelper.java

@@ -23,12 +23,12 @@ import org.thoughtcrime.securesms.BuildConfig;
 import org.thoughtcrime.securesms.R;
 import org.thoughtcrime.securesms.contacts.ContactAccessor;
 import org.thoughtcrime.securesms.contacts.ContactsDatabase;
-import org.thoughtcrime.securesms.crypto.SessionUtil;
 import org.thoughtcrime.securesms.database.MessageDatabase.InsertResult;
 import org.thoughtcrime.securesms.database.RecipientDatabase;
 import org.thoughtcrime.securesms.database.RecipientDatabase.BulkOperationsHandle;
 import org.thoughtcrime.securesms.database.RecipientDatabase.RegisteredState;
 import org.thoughtcrime.securesms.database.SignalDatabase;
+import org.whispersystems.libsignal.SignalProtocolAddress;
 import org.whispersystems.signalservice.api.push.ACI;
 import org.thoughtcrime.securesms.dependencies.ApplicationDependencies;
 import org.thoughtcrime.securesms.jobs.MultiDeviceContactUpdateJob;
@@ -54,6 +54,7 @@ import org.whispersystems.libsignal.util.Pair;
 import org.whispersystems.libsignal.util.guava.Optional;
 import org.whispersystems.signalservice.api.profiles.ProfileAndCredential;
 import org.whispersystems.signalservice.api.profiles.SignalServiceProfile;
+import org.whispersystems.signalservice.api.push.SignalServiceAddress;
 import org.whispersystems.signalservice.api.services.ProfileService;
 import org.whispersystems.signalservice.api.util.UuidUtil;
 import org.whispersystems.signalservice.internal.ServiceResponse;
@@ -177,7 +178,7 @@ public class DirectoryHelper {
       } else {
         Log.w(TAG, "Registered number set had a null ACI!");
       }
-    } else if (recipient.hasAci() && recipient.isRegistered() && hasCommunicatedWith(context, recipient)) {
+    } else if (recipient.hasAci() && recipient.isRegistered() && hasCommunicatedWith(recipient)) {
       if (ApplicationDependencies.getSignalServiceAccountManager().isIdentifierRegistered(recipient.requireAci())) {
         recipientDatabase.markRegistered(recipient.getId(), recipient.requireAci());
       } else {
@@ -464,9 +465,9 @@ public class DirectoryHelper {
 
     for (RecipientId newUser: newUsers) {
       Recipient recipient = Recipient.resolved(newUser);
-      if (!SessionUtil.hasSession(recipient.getId()) &&
+      if (!recipient.isSelf() &&
-          !recipient.isSelf()                        &&
+          recipient.hasAUserSetDisplayName(context) &&
-          recipient.hasAUserSetDisplayName(context))
+          !hasSession(recipient.getId()))
       {
         IncomingJoinedMessage  message      = new IncomingJoinedMessage(recipient.getId());
         Optional<InsertResult> insertResult = SignalDatabase.sms().insertMessageInbox(message);
@@ -483,6 +484,13 @@ public class DirectoryHelper {
     }
   }
 
+  public static boolean hasSession(@NonNull RecipientId id) {
+    SignalProtocolAddress protocolAddress = new SignalProtocolAddress(Recipient.resolved(id).requireServiceId(), SignalServiceAddress.DEFAULT_DEVICE_ID);
+
+    return ApplicationDependencies.getProtocolStore().aci().containsSession(protocolAddress) ||
+           ApplicationDependencies.getProtocolStore().pni().containsSession(protocolAddress);
+  }
+
   private static Set<String> sanitizeNumbers(@NonNull Set<String> numbers) {
     return Stream.of(numbers).filter(number -> {
       try {
@@ -504,7 +512,7 @@ public class DirectoryHelper {
                                              .map(Recipient::resolved)
                                              .filter(Recipient::isRegistered)
                                              .filter(Recipient::hasAci)
-                                             .filter(r -> hasCommunicatedWith(context, r))
+                                             .filter(DirectoryHelper::hasCommunicatedWith)
                                              .toList();
 
     ProfileService profileService = new ProfileService(ApplicationDependencies.getGroupsV2Operations().getProfileOperations(),
@@ -537,10 +545,13 @@ public class DirectoryHelper {
                      .blockingGet();
   }
 
-  private static boolean hasCommunicatedWith(@NonNull Context context, @NonNull Recipient recipient) {
+  private static boolean hasCommunicatedWith(@NonNull Recipient recipient) {
+    ACI localAci = Recipient.self().requireAci();
+
     return SignalDatabase.threads().hasThread(recipient.getId()) ||
-           (recipient.hasAci() && SignalDatabase.sessions().hasSessionFor(recipient.requireAci().toString())) ||
+           (recipient.hasAci() && SignalDatabase.sessions().hasSessionFor(localAci, recipient.requireAci().toString())) ||
-           (recipient.hasE164() && SignalDatabase.sessions().hasSessionFor(recipient.requireE164()));
+           (recipient.hasPni() && SignalDatabase.sessions().hasSessionFor(localAci, recipient.requirePni().toString())) ||
+           (recipient.hasE164() && SignalDatabase.sessions().hasSessionFor(localAci, recipient.requireE164()));
   }
 
   static class DirectoryResult {

app/src/main/java/org/thoughtcrime/securesms/conversation/ui/error/SafetyNumberChangeRepository.java

@@ -12,7 +12,6 @@ import com.annimon.stream.Stream;
 
 import org.signal.core.util.concurrent.SignalExecutors;
 import org.signal.core.util.logging.Log;
-import org.thoughtcrime.securesms.crypto.SessionUtil;
 import org.thoughtcrime.securesms.crypto.ReentrantSessionLock;
 import org.thoughtcrime.securesms.crypto.storage.SignalIdentityKeyStore;
 import org.thoughtcrime.securesms.database.IdentityDatabase;
@@ -119,7 +118,8 @@ final class SafetyNumberChangeRepository {
 
   @WorkerThread
   private TrustAndVerifyResult trustOrVerifyChangedRecipientsAndResendInternal(@NonNull List<ChangedRecipient> changedRecipients,
-                                                                               @NonNull MessageRecord messageRecord) {
+                                                                               @NonNull MessageRecord messageRecord)
+  {
     if (changedRecipients.isEmpty()) {
       Log.d(TAG, "No changed recipients to process, will still process message record");
     }
@@ -134,8 +134,8 @@ final class SafetyNumberChangeRepository {
         Log.d(TAG, "Saving identity result: " + result);
         if (result == SignalIdentityKeyStore.SaveResult.NO_CHANGE) {
           Log.i(TAG, "Archiving sessions explicitly as they appear to be out of sync.");
-          SessionUtil.archiveSession(changedRecipient.getRecipient().getId(), SignalServiceAddress.DEFAULT_DEVICE_ID);
+          ApplicationDependencies.getProtocolStore().aci().sessions().archiveSession(changedRecipient.getRecipient().getId(), SignalServiceAddress.DEFAULT_DEVICE_ID);
-          SessionUtil.archiveSiblingSessions(mismatchAddress);
+          ApplicationDependencies.getProtocolStore().aci().sessions().archiveSiblingSessions(mismatchAddress);
           SignalDatabase.senderKeyShared().deleteAllFor(changedRecipient.getRecipient().getId());
         }
       }

app/src/main/java/org/thoughtcrime/securesms/crypto/SessionUtil.java

@@ -1,39 +0,0 @@
-package org.thoughtcrime.securesms.crypto;
-
-import androidx.annotation.NonNull;
-
-import org.thoughtcrime.securesms.dependencies.ApplicationDependencies;
-import org.thoughtcrime.securesms.recipients.Recipient;
-import org.thoughtcrime.securesms.recipients.RecipientId;
-import org.whispersystems.libsignal.SignalProtocolAddress;
-import org.whispersystems.libsignal.ecc.ECPublicKey;
-import org.whispersystems.libsignal.state.SessionRecord;
-import org.whispersystems.signalservice.api.push.SignalServiceAddress;
-
-public class SessionUtil {
-
-  public static boolean hasSession(@NonNull RecipientId id) {
-    SignalProtocolAddress axolotlAddress = new SignalProtocolAddress(Recipient.resolved(id).requireServiceId(), SignalServiceAddress.DEFAULT_DEVICE_ID);
-
-    return ApplicationDependencies.getProtocolStore().aci().containsSession(axolotlAddress);
-  }
-
-  public static void archiveSiblingSessions(SignalProtocolAddress address) {
-    ApplicationDependencies.getProtocolStore().aci().sessions().archiveSiblingSessions(address);
-  }
-
-  public static void archiveAllSessions() {
-    ApplicationDependencies.getProtocolStore().aci().sessions().archiveAllSessions();
-  }
-
-  public static void archiveSession(RecipientId recipientId, int deviceId) {
-    ApplicationDependencies.getProtocolStore().aci().sessions().archiveSession(recipientId, deviceId);
-  }
-
-  public static boolean ratchetKeyMatches(@NonNull Recipient recipient, int deviceId, @NonNull ECPublicKey ratchetKey) {
-    SignalProtocolAddress address = new SignalProtocolAddress(recipient.resolve().requireServiceId(), deviceId);
-    SessionRecord         session = ApplicationDependencies.getProtocolStore().aci().loadSession(address);
-
-    return session.currentRatchetKeyMatches(ratchetKey);
-  }
-}

app/src/main/java/org/thoughtcrime/securesms/crypto/storage/SignalBaseIdentityKeyStore.java

@@ -6,8 +6,6 @@ import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import org.signal.core.util.logging.Log;
-import org.thoughtcrime.securesms.crypto.IdentityKeyUtil;
-import org.thoughtcrime.securesms.crypto.SessionUtil;
 import org.thoughtcrime.securesms.crypto.storage.SignalIdentityKeyStore.SaveResult;
 import org.thoughtcrime.securesms.database.IdentityDatabase;
 import org.thoughtcrime.securesms.database.IdentityDatabase.VerifiedStatus;
@@ -15,17 +13,16 @@ import org.thoughtcrime.securesms.database.SignalDatabase;
 import org.thoughtcrime.securesms.database.identity.IdentityRecordList;
 import org.thoughtcrime.securesms.database.model.IdentityRecord;
 import org.thoughtcrime.securesms.database.model.IdentityStoreRecord;
+import org.thoughtcrime.securesms.dependencies.ApplicationDependencies;
 import org.thoughtcrime.securesms.keyvalue.SignalStore;
 import org.thoughtcrime.securesms.recipients.Recipient;
 import org.thoughtcrime.securesms.recipients.RecipientId;
 import org.thoughtcrime.securesms.util.IdentityUtil;
 import org.thoughtcrime.securesms.util.LRUCache;
 import org.whispersystems.libsignal.IdentityKey;
-import org.whispersystems.libsignal.IdentityKeyPair;
 import org.whispersystems.libsignal.SignalProtocolAddress;
 import org.whispersystems.libsignal.state.IdentityKeyStore;
 import org.whispersystems.libsignal.util.guava.Optional;
-import org.whispersystems.signalservice.api.push.AccountIdentifier;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -90,7 +87,7 @@ public class SignalBaseIdentityKeyStore {
 
         cache.save(address.getName(), recipientId, identityKey, verifiedStatus, false, System.currentTimeMillis(), nonBlockingApproval);
         IdentityUtil.markIdentityUpdate(context, recipientId);
-        SessionUtil.archiveSiblingSessions(address);
+        ApplicationDependencies.getProtocolStore().aci().sessions().archiveSiblingSessions(address);
         SignalDatabase.senderKeyShared().deleteAllFor(recipientId);
         return SaveResult.UPDATE;
       }

app/src/main/java/org/thoughtcrime/securesms/crypto/storage/TextSecureSessionStore.java

@@ -15,6 +15,7 @@ import org.whispersystems.libsignal.SignalProtocolAddress;
 import org.whispersystems.libsignal.protocol.CiphertextMessage;
 import org.whispersystems.libsignal.state.SessionRecord;
 import org.whispersystems.signalservice.api.SignalServiceSessionStore;
+import org.whispersystems.signalservice.api.push.AccountIdentifier;
 
 import java.util.List;
 import java.util.Objects;
@@ -27,16 +28,16 @@ public class TextSecureSessionStore implements SignalServiceSessionStore {
 
   private static final Object LOCK = new Object();
 
-  @NonNull private final Context context;
+  private final AccountIdentifier accountId;
 
-  public TextSecureSessionStore(@NonNull Context context) {
+  public TextSecureSessionStore(@NonNull AccountIdentifier accountId) {
-    this.context = context;
+    this.accountId = accountId;
   }
 
   @Override
   public SessionRecord loadSession(@NonNull SignalProtocolAddress address) {
     synchronized (LOCK) {
-      SessionRecord sessionRecord = SignalDatabase.sessions().load(address);
+      SessionRecord sessionRecord = SignalDatabase.sessions().load(accountId, address);
 
       if (sessionRecord == null) {
         Log.w(TAG, "No existing session information found for " + address);
@@ -50,7 +51,7 @@ public class TextSecureSessionStore implements SignalServiceSessionStore {
   @Override
   public List<SessionRecord> loadExistingSessions(List<SignalProtocolAddress> addresses) throws NoSessionException {
     synchronized (LOCK) {
-      List<SessionRecord> sessionRecords = SignalDatabase.sessions().load(addresses);
+      List<SessionRecord> sessionRecords = SignalDatabase.sessions().load(accountId, addresses);
 
       if (sessionRecords.size() != addresses.size()) {
         String message = "Mismatch! Asked for " + addresses.size() + " sessions, but only found " + sessionRecords.size() + "!";
@@ -69,14 +70,14 @@ public class TextSecureSessionStore implements SignalServiceSessionStore {
   @Override
   public void storeSession(@NonNull SignalProtocolAddress address, @NonNull SessionRecord record) {
     synchronized (LOCK) {
-      SignalDatabase.sessions().store(address, record);
+      SignalDatabase.sessions().store(accountId, address, record);
     }
   }
 
   @Override
   public boolean containsSession(SignalProtocolAddress address) {
     synchronized (LOCK) {
-      SessionRecord sessionRecord = SignalDatabase.sessions().load(address);
+      SessionRecord sessionRecord = SignalDatabase.sessions().load(accountId, address);
 
       return sessionRecord != null &&
              sessionRecord.hasSenderChain() &&
@@ -88,7 +89,7 @@ public class TextSecureSessionStore implements SignalServiceSessionStore {
   public void deleteSession(SignalProtocolAddress address) {
     synchronized (LOCK) {
       Log.w(TAG, "Deleting session for " + address);
-      SignalDatabase.sessions().delete(address);
+      SignalDatabase.sessions().delete(accountId, address);
     }
   }
 
@@ -96,14 +97,14 @@ public class TextSecureSessionStore implements SignalServiceSessionStore {
   public void deleteAllSessions(String name) {
     synchronized (LOCK) {
       Log.w(TAG, "Deleting all sessions for " + name);
-      SignalDatabase.sessions().deleteAllFor(name);
+      SignalDatabase.sessions().deleteAllFor(accountId, name);
     }
   }
 
   @Override
   public List<Integer> getSubDeviceSessions(String name) {
     synchronized (LOCK) {
-      return SignalDatabase.sessions().getSubDevices(name);
+      return SignalDatabase.sessions().getSubDevices(accountId, name);
     }
   }
 
@@ -111,7 +112,7 @@ public class TextSecureSessionStore implements SignalServiceSessionStore {
   public Set<SignalProtocolAddress> getAllAddressesWithActiveSessions(List<String> addressNames) {
     synchronized (LOCK) {
       return SignalDatabase.sessions()
-                           .getAllFor(addressNames)
+                           .getAllFor(accountId, addressNames)
                            .stream()
                            .filter(row -> isActive(row.getRecord()))
                            .map(row -> new SignalProtocolAddress(row.getAddress(), row.getDeviceId()))
@@ -122,10 +123,10 @@ public class TextSecureSessionStore implements SignalServiceSessionStore {
   @Override
   public void archiveSession(SignalProtocolAddress address) {
     synchronized (LOCK) {
-      SessionRecord session = SignalDatabase.sessions().load(address);
+      SessionRecord session = SignalDatabase.sessions().load(accountId, address);
       if (session != null) {
         session.archiveCurrentState();
-        SignalDatabase.sessions().store(address, session);
+        SignalDatabase.sessions().store(accountId, address, session);
       }
     }
   }
@@ -146,7 +147,7 @@ public class TextSecureSessionStore implements SignalServiceSessionStore {
 
   public void archiveSiblingSessions(@NonNull SignalProtocolAddress address) {
     synchronized (LOCK) {
-      List<SessionDatabase.SessionRow> sessions = SignalDatabase.sessions().getAllFor(address.getName());
+      List<SessionDatabase.SessionRow> sessions = SignalDatabase.sessions().getAllFor(accountId, address.getName());
 
       for (SessionDatabase.SessionRow row : sessions) {
         if (row.getDeviceId() != address.getDeviceId()) {
@@ -159,7 +160,7 @@ public class TextSecureSessionStore implements SignalServiceSessionStore {
 
   public void archiveAllSessions() {
     synchronized (LOCK) {
-      List<SessionDatabase.SessionRow> sessions = SignalDatabase.sessions().getAll();
+      List<SessionDatabase.SessionRow> sessions = SignalDatabase.sessions().getAll(accountId);
 
       for (SessionDatabase.SessionRow row : sessions) {
         row.getRecord().archiveCurrentState();
@@ -173,8 +174,4 @@ public class TextSecureSessionStore implements SignalServiceSessionStore {
            record.hasSenderChain() &&
            record.getSessionVersion() == CiphertextMessage.CURRENT_VERSION;
   }
-
-  private static boolean isValidRegistrationId(int registrationId) {
-    return (registrationId & 0x3fff) == registrationId;
-  }
 }

app/src/main/java/org/thoughtcrime/securesms/database/OneTimePreKeyDatabase.kt

@@ -12,7 +12,7 @@ import org.whispersystems.libsignal.state.PreKeyRecord
 import org.whispersystems.signalservice.api.push.AccountIdentifier
 import java.io.IOException
 
-class OneTimePreKeyDatabase internal constructor(context: Context?, databaseHelper: SignalDatabase?) : Database(context, databaseHelper) {
+class OneTimePreKeyDatabase(context: Context, databaseHelper: SignalDatabase) : Database(context, databaseHelper) {
   companion object {
     private val TAG = Log.tag(OneTimePreKeyDatabase::class.java)
 

app/src/main/java/org/thoughtcrime/securesms/database/RecipientDatabase.kt

@@ -2604,19 +2604,20 @@ open class RecipientDatabase(context: Context, databaseHelper: SignalDatabase) :
     }
 
     // Sessions
+    val localAci: ACI = Recipient.self().aci.get()
     val sessionDatabase = sessions
-    val hasE164Session = sessionDatabase.getAllFor(e164Record.e164).size > 0
+    val hasE164Session = sessionDatabase.getAllFor(localAci, e164Record.e164).isNotEmpty()
-    val hasAciSession = sessionDatabase.getAllFor(aciRecord.aci.toString()).size > 0
+    val hasAciSession = sessionDatabase.getAllFor(localAci, aciRecord.aci.toString()).isNotEmpty()
 
     if (hasE164Session && hasAciSession) {
       Log.w(TAG, "Had a session for both users. Deleting the E164.", true)
-      sessionDatabase.deleteAllFor(e164Record.e164)
+      sessionDatabase.deleteAllFor(localAci, e164Record.e164)
     } else if (hasE164Session && !hasAciSession) {
       Log.w(TAG, "Had a session for E164, but not ACI. Re-assigning to the ACI.", true)
       val values = ContentValues().apply {
         put(SessionDatabase.ADDRESS, aciRecord.aci.toString())
       }
-      db.update(SessionDatabase.TABLE_NAME, values, SessionDatabase.ADDRESS + " = ?", SqlUtil.buildArgs(e164Record.e164))
+      db.update(SessionDatabase.TABLE_NAME, values, "${SessionDatabase.ACCOUNT_ID} = ? AND ${SessionDatabase.ADDRESS} = ?", SqlUtil.buildArgs(localAci, e164Record.e164))
     } else if (!hasE164Session && hasAciSession) {
       Log.w(TAG, "Had a session for ACI, but not E164. No action necessary.", true)
     } else {

app/src/main/java/org/thoughtcrime/securesms/database/SessionDatabase.java

@@ -1,241 +0,0 @@
-package org.thoughtcrime.securesms.database;
-
-
-import android.content.Context;
-import android.database.Cursor;
-
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-
-import net.zetetic.database.sqlcipher.SQLiteStatement;
-
-import org.signal.core.util.logging.Log;
-import org.thoughtcrime.securesms.util.CursorUtil;
-import org.thoughtcrime.securesms.util.SqlUtil;
-import org.whispersystems.libsignal.SignalProtocolAddress;
-import org.whispersystems.libsignal.state.SessionRecord;
-import org.whispersystems.signalservice.api.push.SignalServiceAddress;
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.LinkedHashMap;
-import java.util.LinkedList;
-import java.util.List;
-
-public class SessionDatabase extends Database {
-
-  private static final String TAG = Log.tag(SessionDatabase.class);
-
-  public static final String TABLE_NAME = "sessions";
-
-  private static final String ID      = "_id";
-  public static final  String ADDRESS = "address";
-  public static final  String DEVICE  = "device";
-  public static final  String RECORD  = "record";
-
-  public static final String CREATE_TABLE = "CREATE TABLE " + TABLE_NAME + "(" + ID      + " INTEGER PRIMARY KEY AUTOINCREMENT, " +
-                                                                                 ADDRESS + " TEXT NOT NULL, " +
-                                                                                 DEVICE  + " INTEGER NOT NULL, " +
-                                                                                 RECORD  + " BLOB NOT NULL, " +
-                                                                                 "UNIQUE(" + ADDRESS + "," + DEVICE + "));";
-
-  SessionDatabase(Context context, SignalDatabase databaseHelper) {
-    super(context, databaseHelper);
-  }
-
-  public void store(@NonNull SignalProtocolAddress address, @NonNull SessionRecord record) {
-    if (address.getName().charAt(0) == '+') {
-      throw new IllegalArgumentException("Cannot insert an e164 into this table!");
-    }
-
-    SQLiteDatabase db = databaseHelper.getSignalWritableDatabase();
-
-    try (SQLiteStatement statement = db.compileStatement("INSERT INTO " + TABLE_NAME + " (" + ADDRESS + ", " + DEVICE + ", " + RECORD + ") VALUES (?, ?, ?) " +
-                                                         "ON CONFLICT (" + ADDRESS + ", " + DEVICE + ") DO UPDATE SET " + RECORD + " = excluded." + RECORD))
-    {
-      statement.bindString(1, address.getName());
-      statement.bindLong(2, address.getDeviceId());
-      statement.bindBlob(3, record.serialize());
-      statement.execute();
-    }
-  }
-
-  public @Nullable SessionRecord load(@NonNull SignalProtocolAddress address) {
-    SQLiteDatabase database   = databaseHelper.getSignalReadableDatabase();
-    String[]       projection = new String[] { RECORD };
-    String         selection  = ADDRESS + " = ? AND " + DEVICE + " = ?";
-    String[]       args       = SqlUtil.buildArgs(address.getName(), address.getDeviceId());
-
-    try (Cursor cursor = database.query(TABLE_NAME, projection, selection, args, null, null, null)) {
-      if (cursor.moveToFirst()) {
-        try {
-          return new SessionRecord(cursor.getBlob(cursor.getColumnIndexOrThrow(RECORD)));
-        } catch (IOException e) {
-          Log.w(TAG, e);
-        }
-      }
-    }
-
-    return null;
-  }
-
-  public @NonNull List<SessionRecord> load(@NonNull List<SignalProtocolAddress> addresses) {
-    SQLiteDatabase database = databaseHelper.getSignalReadableDatabase();
-    String         query    = ADDRESS + " = ? AND " + DEVICE + " = ?";
-    List<String[]> args     = new ArrayList<>(addresses.size());
-
-    HashMap<SignalProtocolAddress, SessionRecord> sessions = new LinkedHashMap<>(addresses.size());
-
-    for (SignalProtocolAddress address : addresses) {
-      args.add(SqlUtil.buildArgs(address.getName(), address.getDeviceId()));
-      sessions.put(address, null);
-    }
-
-    String[] projection = new String[] { ADDRESS, DEVICE, RECORD };
-
-    for (SqlUtil.Query combinedQuery : SqlUtil.buildCustomCollectionQuery(query, args)) {
-      try (Cursor cursor = database.query(TABLE_NAME, projection, combinedQuery.getWhere(), combinedQuery.getWhereArgs(), null, null, null)) {
-        while (cursor.moveToNext()) {
-          String address = CursorUtil.requireString(cursor, ADDRESS);
-          int    device  = CursorUtil.requireInt(cursor, DEVICE);
-
-          try {
-            SessionRecord record = new SessionRecord(cursor.getBlob(cursor.getColumnIndexOrThrow(RECORD)));
-            sessions.put(new SignalProtocolAddress(address, device), record);
-          } catch (IOException e) {
-            Log.w(TAG, e);
-          }
-        }
-      }
-    }
-
-    return new ArrayList<>(sessions.values());
-  }
-
-  public @NonNull List<SessionRow> getAllFor(@NonNull String addressName) {
-    SQLiteDatabase   database = databaseHelper.getSignalReadableDatabase();
-    List<SessionRow> results  = new LinkedList<>();
-
-    try (Cursor cursor = database.query(TABLE_NAME, null, ADDRESS + " = ?", SqlUtil.buildArgs(addressName), null, null, null)) {
-      while (cursor.moveToNext()) {
-        try {
-          results.add(new SessionRow(CursorUtil.requireString(cursor, ADDRESS),
-                                     CursorUtil.requireInt(cursor, DEVICE),
-                                     new SessionRecord(CursorUtil.requireBlob(cursor, RECORD))));
-        } catch (IOException e) {
-          Log.w(TAG, e);
-        }
-      }
-    }
-
-    return results;
-  }
-
-  public @NonNull List<SessionRow> getAllFor(@NonNull List<String> addressNames) {
-    SQLiteDatabase   database = databaseHelper.getSignalReadableDatabase();
-    SqlUtil.Query    query    = SqlUtil.buildCollectionQuery(ADDRESS, addressNames);
-    List<SessionRow> results  = new LinkedList<>();
-
-    try (Cursor cursor = database.query(TABLE_NAME, null, query.getWhere(), query.getWhereArgs(), null, null, null)) {
-      while (cursor.moveToNext()) {
-        try {
-          results.add(new SessionRow(CursorUtil.requireString(cursor, ADDRESS),
-                                     CursorUtil.requireInt(cursor, DEVICE),
-                                     new SessionRecord(CursorUtil.requireBlob(cursor, RECORD))));
-        } catch (IOException e) {
-          Log.w(TAG, e);
-        }
-      }
-    }
-
-    return results;
-  }
-
-  public @NonNull List<SessionRow> getAll() {
-    SQLiteDatabase   database = databaseHelper.getSignalReadableDatabase();
-    List<SessionRow> results  = new LinkedList<>();
-
-    try (Cursor cursor = database.query(TABLE_NAME, null, null, null, null, null, null)) {
-      while (cursor.moveToNext()) {
-        try {
-          results.add(new SessionRow(CursorUtil.requireString(cursor, ADDRESS),
-                                     CursorUtil.requireInt(cursor, DEVICE),
-                                     new SessionRecord(cursor.getBlob(cursor.getColumnIndexOrThrow(RECORD)))));
-        } catch (IOException e) {
-          Log.w(TAG, e);
-        }
-      }
-    }
-
-    return results;
-  }
-
-  public @NonNull List<Integer> getSubDevices(@NonNull String addressName) {
-    SQLiteDatabase database   = databaseHelper.getSignalReadableDatabase();
-    List<Integer>  results    = new LinkedList<>();
-    String[]       projection = new String[] { DEVICE };
-    String         selection  = ADDRESS + " = ?";
-    String[]       args       = SqlUtil.buildArgs(addressName);
-
-    try (Cursor cursor = database.query(TABLE_NAME, projection, selection, args, null, null, null)) {
-      while (cursor.moveToNext()) {
-        int device = cursor.getInt(cursor.getColumnIndexOrThrow(DEVICE));
-
-        if (device != SignalServiceAddress.DEFAULT_DEVICE_ID) {
-          results.add(device);
-        }
-      }
-    }
-
-    return results;
-  }
-
-  public void delete(@NonNull SignalProtocolAddress address) {
-    SQLiteDatabase database = databaseHelper.getSignalWritableDatabase();
-    String         selection  = ADDRESS + " = ? AND " + DEVICE + " = ?";
-    String[]       args       = SqlUtil.buildArgs(address.getName(), address.getDeviceId());
-
-
-    database.delete(TABLE_NAME, selection, args);
-  }
-
-  public void deleteAllFor(@NonNull String addressName) {
-    SQLiteDatabase database = databaseHelper.getSignalWritableDatabase();
-    database.delete(TABLE_NAME, ADDRESS + " = ?", SqlUtil.buildArgs(addressName));
-  }
-
-  public boolean hasSessionFor(@NonNull String addressName) {
-    SQLiteDatabase database = databaseHelper.getSignalReadableDatabase();
-    String         query    = ADDRESS + " = ?";
-    String[]       args     = SqlUtil.buildArgs(addressName);
-
-    try (Cursor cursor = database.query(TABLE_NAME, new String[] { "1" }, query, args, null, null, null, "1")) {
-      return cursor.moveToFirst();
-    }
-  }
-
-  public static final class SessionRow {
-    private final String        address;
-    private final int           deviceId;
-    private final SessionRecord record;
-
-    public SessionRow(@NonNull String address, int deviceId, SessionRecord record) {
-      this.address  = address;
-      this.deviceId = deviceId;
-      this.record   = record;
-    }
-
-    public @NonNull String getAddress() {
-      return address;
-    }
-
-    public int getDeviceId() {
-      return deviceId;
-    }
-
-    public SessionRecord getRecord() {
-      return record;
-    }
-  }
-}

app/src/main/java/org/thoughtcrime/securesms/database/SessionDatabase.kt

@@ -0,0 +1,198 @@
+package org.thoughtcrime.securesms.database
+
+import android.content.Context
+import org.signal.core.util.logging.Log
+import org.thoughtcrime.securesms.util.CursorUtil
+import org.thoughtcrime.securesms.util.SqlUtil
+import org.whispersystems.libsignal.SignalProtocolAddress
+import org.whispersystems.libsignal.state.SessionRecord
+import org.whispersystems.signalservice.api.push.AccountIdentifier
+import org.whispersystems.signalservice.api.push.SignalServiceAddress
+import java.io.IOException
+import java.util.ArrayList
+import java.util.HashMap
+import java.util.LinkedHashMap
+import java.util.LinkedList
+
+class SessionDatabase(context: Context, databaseHelper: SignalDatabase) : Database(context, databaseHelper) {
+  companion object {
+    private val TAG = Log.tag(SessionDatabase::class.java)
+
+    const val TABLE_NAME = "sessions"
+    const val ID = "_id"
+    const val ACCOUNT_ID = "account_id"
+    const val ADDRESS = "address"
+    const val DEVICE = "device"
+    const val RECORD = "record"
+    const val CREATE_TABLE = """
+      CREATE TABLE $TABLE_NAME (
+        $ID INTEGER PRIMARY KEY AUTOINCREMENT,
+        $ACCOUNT_ID TEXT NOT NULL,
+        $ADDRESS TEXT NOT NULL,
+        $DEVICE INTEGER NOT NULL,
+        $RECORD BLOB NOT NULL,
+        UNIQUE($ACCOUNT_ID, $ADDRESS, $DEVICE)
+    """
+  }
+
+  fun store(accountId: AccountIdentifier, address: SignalProtocolAddress, record: SessionRecord) {
+    require(address.name[0] != '+') { "Cannot insert an e164 into this table!" }
+
+    writableDatabase.compileStatement("INSERT INTO $TABLE_NAME ($ACCOUNT_ID, $ADDRESS, $DEVICE, $RECORD) VALUES (?, ?, ?, ?) ON CONFLICT ($ACCOUNT_ID, $ADDRESS, $DEVICE) DO UPDATE SET $RECORD = excluded.$RECORD").use { statement ->
+      statement.apply {
+        bindString(1, accountId.toString())
+        bindString(2, address.name)
+        bindLong(3, address.deviceId.toLong())
+        bindBlob(4, record.serialize())
+        execute()
+      }
+    }
+  }
+
+  fun load(accountId: AccountIdentifier, address: SignalProtocolAddress): SessionRecord? {
+    val projection = arrayOf(RECORD)
+    val selection = "$ACCOUNT_ID = ? AND $ADDRESS = ? AND $DEVICE = ?"
+    val args = SqlUtil.buildArgs(accountId, address.name, address.deviceId)
+
+    readableDatabase.query(TABLE_NAME, projection, selection, args, null, null, null).use { cursor ->
+      if (cursor.moveToFirst()) {
+        try {
+          return SessionRecord(cursor.requireNonNullBlob(RECORD))
+        } catch (e: IOException) {
+          Log.w(TAG, e)
+        }
+      }
+    }
+
+    return null
+  }
+
+  fun load(accountId: AccountIdentifier, addresses: List<SignalProtocolAddress>): List<SessionRecord?> {
+    val projection = arrayOf(ADDRESS, DEVICE, RECORD)
+    val query = "$ACCOUNT_ID = ? AND $ADDRESS = ? AND $DEVICE = ?"
+    val args: MutableList<Array<String>> = ArrayList(addresses.size)
+    val sessions: HashMap<SignalProtocolAddress, SessionRecord?> = LinkedHashMap(addresses.size)
+
+    for (address in addresses) {
+      args.add(SqlUtil.buildArgs(accountId, address.name, address.deviceId))
+      sessions[address] = null
+    }
+
+    for (combinedQuery in SqlUtil.buildCustomCollectionQuery(query, args)) {
+      readableDatabase.query(TABLE_NAME, projection, combinedQuery.where, combinedQuery.whereArgs, null, null, null).use { cursor ->
+        while (cursor.moveToNext()) {
+          val address = cursor.requireNonNullString(ADDRESS)
+          val device = cursor.requireInt(DEVICE)
+          try {
+            val record = SessionRecord(cursor.requireNonNullBlob(RECORD))
+            sessions[SignalProtocolAddress(address, device)] = record
+          } catch (e: IOException) {
+            Log.w(TAG, e)
+          }
+        }
+      }
+    }
+
+    return sessions.values.toList()
+  }
+
+  fun getAllFor(accountId: AccountIdentifier, addressName: String): List<SessionRow> {
+    val results: MutableList<SessionRow> = mutableListOf()
+
+    readableDatabase.query(TABLE_NAME, null, "$ACCOUNT_ID = ? AND $ADDRESS = ?", SqlUtil.buildArgs(accountId, addressName), null, null, null).use { cursor ->
+      while (cursor.moveToNext()) {
+        try {
+          results.add(
+            SessionRow(
+              CursorUtil.requireString(cursor, ADDRESS),
+              CursorUtil.requireInt(cursor, DEVICE),
+              SessionRecord(CursorUtil.requireBlob(cursor, RECORD))
+            )
+          )
+        } catch (e: IOException) {
+          Log.w(TAG, e)
+        }
+      }
+    }
+    return results
+  }
+
+  fun getAllFor(accountId: AccountIdentifier, addressNames: List<String?>): List<SessionRow> {
+    val query: SqlUtil.Query = SqlUtil.buildCollectionQuery(ADDRESS, addressNames)
+    val results: MutableList<SessionRow> = LinkedList()
+
+    val queryString = "$ACCOUNT_ID = ? AND (${query.where})"
+    val queryArgs: Array<String> = arrayOf(accountId.toString()) + query.whereArgs
+
+    readableDatabase.query(TABLE_NAME, null, queryString, queryArgs, null, null, null).use { cursor ->
+      while (cursor.moveToNext()) {
+        try {
+          results.add(
+            SessionRow(
+              address = CursorUtil.requireString(cursor, ADDRESS),
+              deviceId = CursorUtil.requireInt(cursor, DEVICE),
+              record = SessionRecord(cursor.requireNonNullBlob(RECORD))
+            )
+          )
+        } catch (e: IOException) {
+          Log.w(TAG, e)
+        }
+      }
+    }
+    return results
+  }
+
+  fun getAll(accountId: AccountIdentifier): List<SessionRow> {
+    val results: MutableList<SessionRow> = mutableListOf()
+
+    readableDatabase.query(TABLE_NAME, null, "$ACCOUNT_ID = ?", SqlUtil.buildArgs(accountId), null, null, null).use { cursor ->
+      while (cursor.moveToNext()) {
+        try {
+          results.add(
+            SessionRow(
+              address = cursor.requireNonNullString(ADDRESS),
+              deviceId = cursor.requireInt(DEVICE),
+              record = SessionRecord(cursor.requireNonNullBlob(RECORD))
+            )
+          )
+        } catch (e: IOException) {
+          Log.w(TAG, e)
+        }
+      }
+    }
+    return results
+  }
+
+  fun getSubDevices(accountId: AccountIdentifier, addressName: String): List<Int> {
+    val projection = arrayOf(DEVICE)
+    val selection = "$ACCOUNT_ID = ? AND $ADDRESS = ? AND $DEVICE != ?"
+    val args = SqlUtil.buildArgs(accountId, addressName, SignalServiceAddress.DEFAULT_DEVICE_ID)
+
+    val results: MutableList<Int> = mutableListOf()
+
+    readableDatabase.query(TABLE_NAME, projection, selection, args, null, null, null).use { cursor ->
+      while (cursor.moveToNext()) {
+        results.add(cursor.requireInt(DEVICE))
+      }
+    }
+    return results
+  }
+
+  fun delete(accountId: AccountIdentifier, address: SignalProtocolAddress) {
+    writableDatabase.delete(TABLE_NAME, "$ACCOUNT_ID = ? AND $ADDRESS = ? AND $DEVICE = ?", SqlUtil.buildArgs(accountId, address.name, address.deviceId))
+  }
+
+  fun deleteAllFor(accountId: AccountIdentifier, addressName: String) {
+    writableDatabase.delete(TABLE_NAME, "$ACCOUNT_ID = ? AND $ADDRESS = ?", SqlUtil.buildArgs(accountId, addressName))
+  }
+
+  fun hasSessionFor(accountId: AccountIdentifier, addressName: String): Boolean {
+    val query = "$ACCOUNT_ID = ? AND $ADDRESS = ?"
+    val args = SqlUtil.buildArgs(accountId, addressName)
+    readableDatabase.query(TABLE_NAME, arrayOf("1"), query, args, null, null, null, "1").use { cursor ->
+      return cursor.moveToFirst()
+    }
+  }
+
+  class SessionRow(val address: String, val deviceId: Int, val record: SessionRecord)
+}

app/src/main/java/org/thoughtcrime/securesms/database/SignedPreKeyDatabase.kt

@@ -13,7 +13,7 @@ import org.whispersystems.signalservice.api.push.AccountIdentifier
 import java.io.IOException
 import java.util.LinkedList
 
-class SignedPreKeyDatabase internal constructor(context: Context?, databaseHelper: SignalDatabase?) : Database(context, databaseHelper) {
+class SignedPreKeyDatabase(context: Context, databaseHelper: SignalDatabase) : Database(context, databaseHelper) {
   companion object {
     private val TAG = Log.tag(SignedPreKeyDatabase::class.java)
 

app/src/main/java/org/thoughtcrime/securesms/database/helpers/SignalDatabaseMigrations.kt

@@ -2363,6 +2363,39 @@ object SignalDatabaseMigrations {
 
       db.execSQL("DROP TABLE signed_prekeys")
       db.execSQL("ALTER TABLE signed_prekeys_tmp RENAME TO signed_prekeys")
+
+      // Sessions
+      db.execSQL(
+        """
+        CREATE TABLE sessions_tmp (
+          _id INTEGER PRIMARY KEY AUTOINCREMENT,
+          account_id TEXT NOT NULL,
+          address TEXT NOT NULL,
+          device INTEGER NOT NULL,
+          record BLOB NOT NULL,
+          UNIQUE(account_id, address, device)
+        )
+        """.trimIndent()
+      )
+
+      if (localAci != null) {
+        db.execSQL(
+          """
+          INSERT INTO sessions_tmp (account_id, address, device, record)
+          SELECT
+            '$localAci' AS account_id,
+            sessions.address,
+            sessions.device,
+            sessions.record
+          FROM sessions
+          """.trimIndent()
+        )
+      } else {
+        Log.w(TAG, "No local ACI set. Not migrating any existing sessions.")
+      }
+
+      db.execSQL("DROP TABLE sessions")
+      db.execSQL("ALTER TABLE sessions_tmp RENAME TO sessions")
     }
   }
 

app/src/main/java/org/thoughtcrime/securesms/dependencies/ApplicationDependencyProvider.java

@@ -294,13 +294,13 @@ public class ApplicationDependencyProvider implements ApplicationDependencies.Pr
     SignalServiceAccountDataStoreImpl aciStore = new SignalServiceAccountDataStoreImpl(context,
                                                                                        new TextSecurePreKeyStore(localAci),
                                                                                        new SignalIdentityKeyStore(baseIdentityStore, () -> SignalStore.account().getAciIdentityKey()),
-                                                                                       new TextSecureSessionStore(context),
+                                                                                       new TextSecureSessionStore(localAci),
                                                                                        new SignalSenderKeyStore(context));
 
     SignalServiceAccountDataStoreImpl pniStore = new SignalServiceAccountDataStoreImpl(context,
                                                                                        new TextSecurePreKeyStore(localPni),
                                                                                        new SignalIdentityKeyStore(baseIdentityStore, () -> SignalStore.account().getPniIdentityKey()),
-                                                                                       new TextSecureSessionStore(context),
+                                                                                       new TextSecureSessionStore(localPni),
                                                                                        new SignalSenderKeyStore(context));
     return new SignalServiceDataStoreImpl(context, aciStore, pniStore);
   }

app/src/main/java/org/thoughtcrime/securesms/jobs/AutomaticSessionResetJob.java

@@ -3,7 +3,6 @@ package org.thoughtcrime.securesms.jobs;
 import androidx.annotation.NonNull;
 
 import org.signal.core.util.logging.Log;
-import org.thoughtcrime.securesms.crypto.SessionUtil;
 import org.thoughtcrime.securesms.crypto.UnidentifiedAccessUtil;
 import org.thoughtcrime.securesms.database.MessageDatabase;
 import org.thoughtcrime.securesms.database.SignalDatabase;
@@ -84,7 +83,7 @@ public class AutomaticSessionResetJob extends BaseJob {
 
   @Override
   protected void onRun() throws Exception {
-    SessionUtil.archiveSession(recipientId, deviceId);
+    ApplicationDependencies.getProtocolStore().aci().sessions().archiveSession(recipientId, deviceId);
     SignalDatabase.senderKeyShared().deleteAllFor(recipientId);
     insertLocalMessage();
 

app/src/main/java/org/thoughtcrime/securesms/messages/MessageContentProcessor.java

@@ -25,7 +25,6 @@ import org.thoughtcrime.securesms.contactshare.Contact;
 import org.thoughtcrime.securesms.contactshare.ContactModelMapper;
 import org.thoughtcrime.securesms.crypto.ProfileKeyUtil;
 import org.thoughtcrime.securesms.crypto.SecurityEvent;
-import org.thoughtcrime.securesms.crypto.SessionUtil;
 import org.thoughtcrime.securesms.database.AttachmentDatabase;
 import org.thoughtcrime.securesms.database.GroupDatabase;
 import org.thoughtcrime.securesms.database.GroupDatabase.GroupRecord;
@@ -121,7 +120,9 @@ import org.thoughtcrime.securesms.util.RemoteDeleteUtil;
 import org.thoughtcrime.securesms.util.TextSecurePreferences;
 import org.thoughtcrime.securesms.util.Util;
 import org.whispersystems.libsignal.SignalProtocolAddress;
+import org.whispersystems.libsignal.ecc.ECPublicKey;
 import org.whispersystems.libsignal.protocol.DecryptionErrorMessage;
+import org.whispersystems.libsignal.state.SessionRecord;
 import org.whispersystems.libsignal.util.guava.Optional;
 import org.whispersystems.signalservice.api.messages.SignalServiceAttachment;
 import org.whispersystems.signalservice.api.messages.SignalServiceAttachmentPointer;
@@ -1977,11 +1978,13 @@ public final class MessageContentProcessor {
   private void handleIndividualRetryReceipt(@NonNull Recipient requester, @Nullable MessageLogEntry messageLogEntry, @NonNull SignalServiceContent content, @NonNull DecryptionErrorMessage decryptionErrorMessage) {
     boolean archivedSession = false;
 
+    // TODO [pnp] Ignore retry receipts that have a PNI destinationUuid
+
     if (decryptionErrorMessage.getRatchetKey().isPresent() &&
-        SessionUtil.ratchetKeyMatches(requester, content.getSenderDevice(), decryptionErrorMessage.getRatchetKey().get()))
+        ratchetKeyMatches(requester, content.getSenderDevice(), decryptionErrorMessage.getRatchetKey().get()))
     {
       warn(content.getTimestamp(), "[RetryReceipt-I] Ratchet key matches. Archiving the session.");
-      SessionUtil.archiveSession(requester.getId(), content.getSenderDevice());
+      ApplicationDependencies.getProtocolStore().aci().sessions().archiveSession(requester.getId(), content.getSenderDevice());
       archivedSession = true;
     }
 
@@ -2015,6 +2018,13 @@ public final class MessageContentProcessor {
     }
   }
 
+  public static boolean ratchetKeyMatches(@NonNull Recipient recipient, int deviceId, @NonNull ECPublicKey ratchetKey) {
+    SignalProtocolAddress address = new SignalProtocolAddress(recipient.resolve().requireServiceId(), deviceId);
+    SessionRecord         session = ApplicationDependencies.getProtocolStore().aci().loadSession(address);
+
+    return session.currentRatchetKeyMatches(ratchetKey);
+  }
+
   private static boolean isInvalidMessage(@NonNull SignalServiceDataMessage message) {
     if (message.isViewOnce()) {
       List<SignalServiceAttachment> attachments = message.getAttachments().or(Collections.emptyList());

app/src/main/java/org/thoughtcrime/securesms/recipients/Recipient.java

@@ -656,6 +656,15 @@ public class Recipient {
     return resolved;
   }
 
+  public @NonNull PNI requirePni() {
+    PNI resolved = resolving ? resolve().pni : pni;
+
+    if (resolved == null) {
+      throw new MissingAddressError(id);
+    }
+
+    return resolved;
+  }
 
   public @NonNull String requireE164() {
     String resolved = resolving ? resolve().e164 : e164;
@@ -701,6 +710,10 @@ public class Recipient {
     return getAci().isPresent();
   }
 
+  public boolean hasPni() {
+    return getPni().isPresent();
+  }
+
   public boolean isAciOnly() {
     return hasAci() && !hasSmsAddress();
   }

app/src/main/java/org/thoughtcrime/securesms/registration/RegistrationRepository.java

@@ -1,7 +1,6 @@
 package org.thoughtcrime.securesms.registration;
 
 import android.app.Application;
-import android.content.Context;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -12,10 +11,8 @@ import org.signal.zkgroup.profiles.ProfileKey;
 import org.thoughtcrime.securesms.crypto.PreKeyUtil;
 import org.thoughtcrime.securesms.crypto.ProfileKeyUtil;
 import org.thoughtcrime.securesms.crypto.SenderKeyUtil;
-import org.thoughtcrime.securesms.crypto.SessionUtil;
 import org.thoughtcrime.securesms.crypto.storage.PreKeyMetadataStore;
 import org.thoughtcrime.securesms.crypto.storage.SignalServiceAccountDataStoreImpl;
-import org.thoughtcrime.securesms.crypto.storage.SignalServiceDataStoreImpl;
 import org.thoughtcrime.securesms.database.IdentityDatabase;
 import org.thoughtcrime.securesms.database.RecipientDatabase;
 import org.thoughtcrime.securesms.database.SignalDatabase;
@@ -137,7 +134,8 @@ public final class RegistrationRepository {
     SignalStore.account().setAci(aci);
     SignalStore.account().setPni(pni);
 
-    SessionUtil.archiveAllSessions();
+    ApplicationDependencies.getProtocolStore().aci().sessions().archiveAllSessions();
+    ApplicationDependencies.getProtocolStore().pni().sessions().archiveAllSessions();
     SenderKeyUtil.clearAllState(context);
 
     SignalServiceAccountManager       accountManager   = AccountManagerFactory.createAuthenticated(context, aci, registrationData.getE164(), SignalServiceAddress.DEFAULT_DEVICE_ID, registrationData.getPassword());