Implement SignalServiceAccountDataStore for PNI

lib/src/main/java/org/asamk/signal/manager/ManagerImpl.java

@@ -167,7 +167,8 @@ class ManagerImpl implements Manager {
         });
         disposable.add(account.getIdentityKeyStore().getIdentityChanges().subscribe(recipientId -> {
             logger.trace("Archiving old sessions for {}", recipientId);
-            account.getSessionStore().archiveSessions(recipientId);
+            account.getAciSessionStore().archiveSessions(recipientId);
+            account.getPniSessionStore().archiveSessions(recipientId);
             account.getSenderKeyStore().deleteSharedWith(recipientId);
             final var profile = account.getProfileStore().getProfile(recipientId);
             if (profile != null) {
@@ -688,7 +689,7 @@ class ManagerImpl implements Manager {
                 } catch (UnregisteredRecipientException e) {
                     continue;
                 }
-                account.getSessionStore().deleteAllSessions(recipientId);
+                account.getAciSessionStore().deleteAllSessions(recipientId);
             }
         }
     }

lib/src/main/java/org/asamk/signal/manager/actions/RenewSessionAction.java

@@ -13,7 +13,7 @@ public class RenewSessionAction implements HandleAction {
 
     @Override
     public void execute(Context context) throws Throwable {
-        context.getAccount().getSessionStore().archiveSessions(recipientId);
+        context.getAccount().getAciSessionStore().archiveSessions(recipientId);
         if (!recipientId.equals(context.getAccount().getSelfRecipientId())) {
             context.getSendHelper().sendNullMessage(recipientId);
         }

lib/src/main/java/org/asamk/signal/manager/actions/SendRetryMessageRequestAction.java

@@ -29,7 +29,7 @@ public class SendRetryMessageRequestAction implements HandleAction {
 
     @Override
     public void execute(Context context) throws Throwable {
-        context.getAccount().getSessionStore().archiveSessions(recipientId);
+        context.getAccount().getAciSessionStore().archiveSessions(recipientId);
 
         int senderDevice = protocolException.getSenderDevice();
         Optional<GroupId> groupId = protocolException.getGroupId().isPresent() ? Optional.of(GroupId.unknownVersion(

lib/src/main/java/org/asamk/signal/manager/helper/IncomingMessageHandler.java

@@ -214,6 +214,7 @@ public final class IncomingMessageHandler {
         final var senderPair = getSender(envelope, content);
         final var sender = senderPair.first();
         final var senderDeviceId = senderPair.second();
+        final var destination = getDestination(envelope);
 
         if (content.getReceiptMessage().isPresent()) {
             final var message = content.getReceiptMessage().get();
@@ -274,7 +275,7 @@ public final class IncomingMessageHandler {
             actions.addAll(handleSignalServiceDataMessage(message,
                     false,
                     sender,
-                    account.getSelfRecipientId(),
+                    destination,
                     receiveConfig.ignoreAttachments()));
         }
 
@@ -305,13 +306,14 @@ public final class IncomingMessageHandler {
         }
 
         if (message.getRatchetKey().isPresent()) {
-            if (account.getSessionStore().isCurrentRatchetKey(sender, senderDeviceId, message.getRatchetKey().get())) {
+            if (account.getAciSessionStore()
+                    .isCurrentRatchetKey(sender, senderDeviceId, message.getRatchetKey().get())) {
                 if (logEntries.isEmpty()) {
                     logger.debug("Renewing the session with sender");
                     actions.add(new RenewSessionAction(sender));
                 } else {
                     logger.trace("Archiving the session with sender, a resend message has already been queued");
-                    context.getAccount().getSessionStore().archiveSessions(sender);
+                    context.getAccount().getAciSessionStore().archiveSessions(sender);
                 }
             }
             return;
@@ -641,7 +643,7 @@ public final class IncomingMessageHandler {
 
         final var conversationPartnerAddress = isSync ? destination : source;
         if (conversationPartnerAddress != null && message.isEndSession()) {
-            account.getSessionStore().deleteAllSessions(conversationPartnerAddress);
+            account.getAciSessionStore().deleteAllSessions(conversationPartnerAddress);
         }
         if (message.isExpirationUpdate() || message.getBody().isPresent()) {
             if (message.getGroupContext().isPresent()) {
@@ -776,4 +778,15 @@ public final class IncomingMessageHandler {
                     content.getSenderDevice());
         }
     }
+
+    private RecipientId getDestination(SignalServiceEnvelope envelope) {
+        if (!envelope.hasDestinationUuid()) {
+            return account.getSelfRecipientId();
+        }
+        final var addressOptional = SignalServiceAddress.fromRaw(envelope.getDestinationUuid(), null);
+        if (addressOptional.isEmpty()) {
+            return account.getSelfRecipientId();
+        }
+        return context.getRecipientHelper().resolveRecipient(addressOptional.get());
+    }
 }

lib/src/main/java/org/asamk/signal/manager/storage/SignalAccount.java

@@ -137,14 +137,17 @@ public class SignalAccount implements Closeable {
 
     private boolean registered = false;
 
-    private SignalProtocolStore signalProtocolStore;
+    private SignalProtocolStore aciSignalProtocolStore;
+    private SignalProtocolStore pniSignalProtocolStore;
     private PreKeyStore aciPreKeyStore;
     private SignedPreKeyStore aciSignedPreKeyStore;
     private PreKeyStore pniPreKeyStore;
     private SignedPreKeyStore pniSignedPreKeyStore;
-    private SessionStore sessionStore;
+    private SessionStore aciSessionStore;
+    private SessionStore pniSessionStore;
     private IdentityKeyStore identityKeyStore;
     private SignalIdentityKeyStore aciIdentityKeyStore;
+    private SignalIdentityKeyStore pniIdentityKeyStore;
     private SenderKeyStore senderKeyStore;
     private GroupStore groupStore;
     private RecipientStore recipientStore;
@@ -275,7 +278,8 @@ public class SignalAccount implements Closeable {
                 profileKey);
         signalAccount.getRecipientTrustedResolver()
                 .resolveSelfRecipientTrusted(signalAccount.getSelfRecipientAddress());
-        signalAccount.getSessionStore().archiveAllSessions();
+        signalAccount.getAciSessionStore().archiveAllSessions();
+        signalAccount.getPniSessionStore().archiveAllSessions();
         signalAccount.getSenderKeyStore().deleteAll();
         signalAccount.clearAllPreKeys();
         return signalAccount;
@@ -393,7 +397,8 @@ public class SignalAccount implements Closeable {
     }
 
     private void mergeRecipients(RecipientId recipientId, RecipientId toBeMergedRecipientId) {
-        getSessionStore().mergeRecipients(recipientId, toBeMergedRecipientId);
+        getAciSessionStore().mergeRecipients(recipientId, toBeMergedRecipientId);
+        getPniSessionStore().mergeRecipients(recipientId, toBeMergedRecipientId);
         getIdentityKeyStore().mergeRecipients(recipientId, toBeMergedRecipientId);
         getMessageCache().mergeRecipients(recipientId, toBeMergedRecipientId);
         getGroupStore().mergeRecipients(recipientId, toBeMergedRecipientId);
@@ -401,7 +406,8 @@ public class SignalAccount implements Closeable {
     }
 
     public void removeRecipient(final RecipientId recipientId) {
-        getSessionStore().deleteAllSessions(recipientId);
+        getAciSessionStore().deleteAllSessions(recipientId);
+        getPniSessionStore().deleteAllSessions(recipientId);
         getIdentityKeyStore().deleteIdentity(recipientId);
         getMessageCache().deleteMessages(recipientId);
         getSenderKeyStore().deleteAll(recipientId);
@@ -637,7 +643,7 @@ public class SignalAccount implements Closeable {
         }
         final var legacySessionsPath = getSessionsPath(dataPath, accountPath);
         if (legacySessionsPath.exists()) {
-            LegacySessionStore.migrate(legacySessionsPath, getRecipientResolver(), getSessionStore());
+            LegacySessionStore.migrate(legacySessionsPath, getRecipientResolver(), getAciSessionStore());
             migratedLegacyConfig = true;
         }
         final var legacySignalProtocolStore = rootNode.hasNonNull("axolotlStore")
@@ -734,7 +740,7 @@ public class SignalAccount implements Closeable {
             logger.debug("Migrating legacy session store.");
             for (var session : legacySignalProtocolStore.getLegacySessionStore().getSessions()) {
                 try {
-                    getSessionStore().storeSession(new SignalProtocolAddress(session.address.getIdentifier(),
+                    getAciSessionStore().storeSession(new SignalProtocolAddress(session.address.getIdentifier(),
                             session.deviceId), new SessionRecord(session.sessionRecord));
                 } catch (Exception e) {
                     logger.warn("Failed to migrate session, ignoring", e);
@@ -1016,9 +1022,9 @@ public class SignalAccount implements Closeable {
             @Override
             public SignalServiceAccountDataStore get(final ServiceId accountIdentifier) {
                 if (accountIdentifier.equals(aci)) {
-                    return getAciSignalServiceAccountDataStore();
+                    return aci();
                 } else if (accountIdentifier.equals(pni)) {
-                    throw new AssertionError("PNI not to be used yet!");
+                    return pni();
                 } else {
                     throw new IllegalArgumentException("No matching store found for " + accountIdentifier);
                 }
@@ -1031,7 +1037,7 @@ public class SignalAccount implements Closeable {
 
             @Override
             public SignalServiceAccountDataStore pni() {
-                throw new AssertionError("PNI not to be used yet!");
+                return getPniSignalServiceAccountDataStore();
             }
 
             @Override
@@ -1042,15 +1048,25 @@ public class SignalAccount implements Closeable {
     }
 
     private SignalServiceAccountDataStore getAciSignalServiceAccountDataStore() {
-        return getOrCreate(() -> signalProtocolStore,
+        return getOrCreate(() -> aciSignalProtocolStore,
-                () -> signalProtocolStore = new SignalProtocolStore(getAciPreKeyStore(),
+                () -> aciSignalProtocolStore = new SignalProtocolStore(getAciPreKeyStore(),
                         getAciSignedPreKeyStore(),
-                        getSessionStore(),
+                        getAciSessionStore(),
                         getAciIdentityKeyStore(),
                         getSenderKeyStore(),
                         this::isMultiDevice));
     }
 
+    private SignalServiceAccountDataStore getPniSignalServiceAccountDataStore() {
+        return getOrCreate(() -> pniSignalProtocolStore,
+                () -> pniSignalProtocolStore = new SignalProtocolStore(getPniPreKeyStore(),
+                        getPniSignedPreKeyStore(),
+                        getPniSessionStore(),
+                        getPniIdentityKeyStore(),
+                        getSenderKeyStore(),
+                        this::isMultiDevice));
+    }
+
     private PreKeyStore getAciPreKeyStore() {
         return getOrCreate(() -> aciPreKeyStore,
                 () -> aciPreKeyStore = new PreKeyStore(getAccountDatabase(), ServiceIdType.ACI));
@@ -1071,14 +1087,22 @@ public class SignalAccount implements Closeable {
                 () -> pniSignedPreKeyStore = new SignedPreKeyStore(getAccountDatabase(), ServiceIdType.PNI));
     }
 
-    public SessionStore getSessionStore() {
+    public SessionStore getAciSessionStore() {
-        return getOrCreate(() -> sessionStore,
+        return getOrCreate(() -> aciSessionStore,
-                () -> sessionStore = new SessionStore(getAccountDatabase(),
+                () -> aciSessionStore = new SessionStore(getAccountDatabase(),
                         ServiceIdType.ACI,
                         getRecipientResolver(),
                         getRecipientIdCreator()));
     }
 
+    public SessionStore getPniSessionStore() {
+        return getOrCreate(() -> pniSessionStore,
+                () -> pniSessionStore = new SessionStore(getAccountDatabase(),
+                        ServiceIdType.PNI,
+                        getRecipientResolver(),
+                        getRecipientIdCreator()));
+    }
+
     public IdentityKeyStore getIdentityKeyStore() {
         return getOrCreate(() -> identityKeyStore,
                 () -> identityKeyStore = new IdentityKeyStore(getIdentitiesPath(dataPath, accountPath),
@@ -1094,6 +1118,14 @@ public class SignalAccount implements Closeable {
                         getIdentityKeyStore()));
     }
 
+    public SignalIdentityKeyStore getPniIdentityKeyStore() {
+        return getOrCreate(() -> pniIdentityKeyStore,
+                () -> pniIdentityKeyStore = new SignalIdentityKeyStore(getRecipientResolver(),
+                        () -> pniIdentityKeyPair,
+                        localRegistrationId,
+                        getIdentityKeyStore()));
+    }
+
     public GroupStore getGroupStore() {
         return getOrCreate(() -> groupStore,
                 () -> groupStore = new GroupStore(getAccountDatabase(),
@@ -1500,7 +1532,8 @@ public class SignalAccount implements Closeable {
         save();
 
         clearAllPreKeys();
-        getSessionStore().archiveAllSessions();
+        getAciSessionStore().archiveAllSessions();
+        getPniSessionStore().archiveAllSessions();
         getSenderKeyStore().deleteAll();
         final var recipientId = getRecipientTrustedResolver().resolveSelfRecipientTrusted(getSelfRecipientAddress());
         final var publicKey = getAciIdentityKeyPair().getPublicKey();