Refactor identity key store

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

@@ -1049,7 +1049,7 @@ class ManagerImpl implements Manager {
         IdentityInfo identity;
         try {
             identity = account.getIdentityKeyStore()
-                    .getIdentity(context.getRecipientHelper().resolveRecipient(recipient));
+                    .getIdentityInfo(context.getRecipientHelper().resolveRecipient(recipient));
         } catch (UnregisteredRecipientException e) {
             identity = null;
         }

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

@@ -16,7 +16,6 @@ import org.whispersystems.signalservice.api.push.SignalServiceAddress;
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Date;
 import java.util.function.Function;
 
 import static org.asamk.signal.manager.config.ServiceConfig.capabilities;
@@ -85,7 +84,7 @@ public class IdentityHelper {
     private boolean trustIdentity(
             RecipientId recipientId, Function<IdentityKey, Boolean> verifier, TrustLevel trustLevel
     ) {
-        var identity = account.getIdentityKeyStore().getIdentity(recipientId);
+        var identity = account.getIdentityKeyStore().getIdentityInfo(recipientId);
         if (identity == null) {
             return false;
         }
@@ -110,7 +109,7 @@ public class IdentityHelper {
     ) {
         final var identityKey = identityFailure.getIdentityKey();
         if (identityKey != null) {
-            account.getIdentityKeyStore().saveIdentity(recipientId, identityKey, new Date());
+            account.getIdentityKeyStore().saveIdentity(recipientId, identityKey);
         } else {
             // Retrieve profile to get the current identity key from the server
             context.getProfileHelper().refreshRecipientProfile(recipientId);

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

@@ -35,7 +35,6 @@ import java.io.OutputStream;
 import java.nio.file.Files;
 import java.util.Base64;
 import java.util.Collection;
-import java.util.Date;
 import java.util.List;
 import java.util.Locale;
 import java.util.Objects;
@@ -354,7 +353,7 @@ public final class ProfileHelper {
             try {
                 logger.trace("Storing identity");
                 final var identityKey = new IdentityKey(Base64.getDecoder().decode(encryptedProfile.getIdentityKey()));
-                account.getIdentityKeyStore().saveIdentity(recipientId, identityKey, new Date());
+                account.getIdentityKeyStore().saveIdentity(recipientId, identityKey);
             } catch (InvalidKeyException ignored) {
                 logger.warn("Got invalid identity key in profile for {}",
                         context.getRecipientHelper().resolveSignalServiceAddress(recipientId).getIdentifier());

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

@@ -477,7 +477,7 @@ public class SendHelper {
                 continue;
             }
 
-            final var identity = account.getIdentityKeyStore().getIdentity(recipientId);
+            final var identity = account.getIdentityKeyStore().getIdentityInfo(recipientId);
             if (identity == null || !identity.getTrustLevel().isTrusted()) {
                 continue;
             }

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

@@ -24,7 +24,6 @@ import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Date;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -127,7 +126,7 @@ public class StorageHelper {
             try {
                 logger.trace("Storing identity key {}", recipientId);
                 final var identityKey = new IdentityKey(contactRecord.getIdentityKey().get());
-                account.getIdentityKeyStore().saveIdentity(recipientId, identityKey, new Date());
+                account.getIdentityKeyStore().saveIdentity(recipientId, identityKey);
 
                 final var trustLevel = TrustLevel.fromIdentityState(contactRecord.getIdentityState());
                 if (trustLevel != null) {

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

@@ -128,7 +128,7 @@ public class SyncHelper {
                     final var contact = contactPair.second();
                     final var address = context.getRecipientHelper().resolveSignalServiceAddress(recipientId);
 
-                    var currentIdentity = account.getIdentityKeyStore().getIdentity(recipientId);
+                    var currentIdentity = account.getIdentityKeyStore().getIdentityInfo(recipientId);
                     VerifiedMessage verifiedMessage = null;
                     if (currentIdentity != null) {
                         verifiedMessage = new VerifiedMessage(address,

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

@@ -15,6 +15,7 @@ import org.asamk.signal.manager.storage.groups.GroupInfoV1;
 import org.asamk.signal.manager.storage.groups.GroupInfoV2;
 import org.asamk.signal.manager.storage.groups.GroupStore;
 import org.asamk.signal.manager.storage.identities.IdentityKeyStore;
+import org.asamk.signal.manager.storage.identities.SignalIdentityKeyStore;
 import org.asamk.signal.manager.storage.identities.TrustNewIdentity;
 import org.asamk.signal.manager.storage.messageCache.MessageCache;
 import org.asamk.signal.manager.storage.prekeys.PreKeyStore;
@@ -81,7 +82,6 @@ import java.security.SecureRandom;
 import java.sql.SQLException;
 import java.util.Base64;
 import java.util.Comparator;
-import java.util.Date;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
@@ -137,6 +137,7 @@ public class SignalAccount implements Closeable {
     private SignedPreKeyStore pniSignedPreKeyStore;
     private SessionStore sessionStore;
     private IdentityKeyStore identityKeyStore;
+    private SignalIdentityKeyStore aciIdentityKeyStore;
     private SenderKeyStore senderKeyStore;
     private GroupStore groupStore;
     private GroupStore.Storage groupStoreStorage;
@@ -1016,7 +1017,7 @@ public class SignalAccount implements Closeable {
                 () -> signalProtocolStore = new SignalProtocolStore(getAciPreKeyStore(),
                         getAciSignedPreKeyStore(),
                         getSessionStore(),
-                        getIdentityKeyStore(),
+                        getAciIdentityKeyStore(),
                         getSenderKeyStore(),
                         this::isMultiDevice));
     }
@@ -1050,11 +1051,17 @@ public class SignalAccount implements Closeable {
         return getOrCreate(() -> identityKeyStore,
                 () -> identityKeyStore = new IdentityKeyStore(getIdentitiesPath(dataPath, accountPath),
                         getRecipientResolver(),
-                        aciIdentityKeyPair,
-                        localRegistrationId,
                         trustNewIdentity));
     }
 
+    public SignalIdentityKeyStore getAciIdentityKeyStore() {
+        return getOrCreate(() -> aciIdentityKeyStore,
+                () -> aciIdentityKeyStore = new SignalIdentityKeyStore(getRecipientResolver(),
+                        () -> aciIdentityKeyPair,
+                        localRegistrationId,
+                        getIdentityKeyStore()));
+    }
+
     public GroupStore getGroupStore() {
         return groupStore;
     }
@@ -1390,7 +1397,7 @@ public class SignalAccount implements Closeable {
         getSenderKeyStore().deleteAll();
         final var recipientId = getRecipientTrustedResolver().resolveSelfRecipientTrusted(getSelfRecipientAddress());
         final var publicKey = getAciIdentityKeyPair().getPublicKey();
-        getIdentityKeyStore().saveIdentity(recipientId, publicKey, new Date());
+        getIdentityKeyStore().saveIdentity(recipientId, publicKey);
         getIdentityKeyStore().setIdentityTrustLevel(recipientId, publicKey, TrustLevel.TRUSTED_VERIFIED);
     }
 

lib/src/main/java/org/asamk/signal/manager/storage/identities/IdentityKeyStore.java

@@ -7,9 +7,8 @@ import org.asamk.signal.manager.storage.recipients.RecipientId;
 import org.asamk.signal.manager.storage.recipients.RecipientResolver;
 import org.asamk.signal.manager.util.IOUtils;
 import org.signal.libsignal.protocol.IdentityKey;
-import org.signal.libsignal.protocol.IdentityKeyPair;
 import org.signal.libsignal.protocol.InvalidKeyException;
-import org.signal.libsignal.protocol.SignalProtocolAddress;
+import org.signal.libsignal.protocol.state.IdentityKeyStore.Direction;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -32,7 +31,7 @@ import java.util.regex.Pattern;
 import io.reactivex.rxjava3.subjects.PublishSubject;
 import io.reactivex.rxjava3.subjects.Subject;
 
-public class IdentityKeyStore implements org.signal.libsignal.protocol.state.IdentityKeyStore {
+public class IdentityKeyStore {
 
     private final static Logger logger = LoggerFactory.getLogger(IdentityKeyStore.class);
     private final ObjectMapper objectMapper = org.asamk.signal.manager.storage.Utils.createStorageObjectMapper();
@@ -42,24 +41,16 @@ public class IdentityKeyStore implements org.signal.libsignal.protocol.state.Ide
     private final File identitiesPath;
 
     private final RecipientResolver resolver;
-    private final IdentityKeyPair identityKeyPair;
-    private final int localRegistrationId;
     private final TrustNewIdentity trustNewIdentity;
     private final PublishSubject<RecipientId> identityChanges = PublishSubject.create();
 
     private boolean isRetryingDecryption = false;
 
     public IdentityKeyStore(
-            final File identitiesPath,
+            final File identitiesPath, final RecipientResolver resolver, final TrustNewIdentity trustNewIdentity
-            final RecipientResolver resolver,
-            final IdentityKeyPair identityKeyPair,
-            final int localRegistrationId,
-            final TrustNewIdentity trustNewIdentity
     ) {
         this.identitiesPath = identitiesPath;
         this.resolver = resolver;
-        this.identityKeyPair = identityKeyPair;
-        this.localRegistrationId = localRegistrationId;
         this.trustNewIdentity = trustNewIdentity;
     }
 
@@ -67,21 +58,8 @@ public class IdentityKeyStore implements org.signal.libsignal.protocol.state.Ide
         return identityChanges;
     }
 
-    @Override
+    public boolean saveIdentity(final RecipientId recipientId, final IdentityKey identityKey) {
-    public IdentityKeyPair getIdentityKeyPair() {
+        return saveIdentity(recipientId, identityKey, null);
-        return identityKeyPair;
-    }
-
-    @Override
-    public int getLocalRegistrationId() {
-        return localRegistrationId;
-    }
-
-    @Override
-    public boolean saveIdentity(SignalProtocolAddress address, IdentityKey identityKey) {
-        final var recipientId = resolveRecipient(address.getName());
-
-        return saveIdentity(recipientId, identityKey, new Date());
     }
 
     public boolean saveIdentity(final RecipientId recipientId, final IdentityKey identityKey, Date added) {
@@ -100,7 +78,10 @@ public class IdentityKeyStore implements org.signal.libsignal.protocol.state.Ide
                     trustNewIdentity == TrustNewIdentity.ON_FIRST_USE && identityInfo == null
             ) ? TrustLevel.TRUSTED_UNVERIFIED : TrustLevel.UNTRUSTED;
             logger.debug("Storing new identity for recipient {} with trust {}", recipientId, trustLevel);
-            final var newIdentityInfo = new IdentityInfo(recipientId, identityKey, trustLevel, added);
+            final var newIdentityInfo = new IdentityInfo(recipientId,
+                    identityKey,
+                    trustLevel,
+                    added == null ? new Date() : added);
             storeIdentityLocked(recipientId, newIdentityInfo);
             identityChanges.onNext(recipientId);
             return true;
@@ -137,26 +118,23 @@ public class IdentityKeyStore implements org.signal.libsignal.protocol.state.Ide
         }
     }
 
-    @Override
+    public boolean isTrustedIdentity(RecipientId recipientId, IdentityKey identityKey, Direction direction) {
-    public boolean isTrustedIdentity(SignalProtocolAddress address, IdentityKey identityKey, Direction direction) {
         if (trustNewIdentity == TrustNewIdentity.ALWAYS) {
             return true;
         }
 
-        var recipientId = resolveRecipient(address.getName());
-
         synchronized (cachedIdentities) {
             // TODO implement possibility for different handling of incoming/outgoing trust decisions
             var identityInfo = loadIdentityLocked(recipientId);
             if (identityInfo == null) {
                 logger.debug("Initial identity found for {}, saving.", recipientId);
-                saveIdentity(address, identityKey);
+                saveIdentity(recipientId, identityKey);
                 identityInfo = loadIdentityLocked(recipientId);
             } else if (!identityInfo.getIdentityKey().equals(identityKey)) {
                 // Identity found, but different
                 if (direction == Direction.SENDING) {
                     logger.debug("Changed identity found for {}, saving.", recipientId);
-                    saveIdentity(address, identityKey);
+                    saveIdentity(recipientId, identityKey);
                     identityInfo = loadIdentityLocked(recipientId);
                 } else {
                     logger.trace("Trusting identity for {} for {}: {}", recipientId, direction, false);
@@ -170,17 +148,14 @@ public class IdentityKeyStore implements org.signal.libsignal.protocol.state.Ide
         }
     }
 
-    @Override
+    public IdentityKey getIdentity(RecipientId recipientId) {
-    public IdentityKey getIdentity(SignalProtocolAddress address) {
-        var recipientId = resolveRecipient(address.getName());
-
         synchronized (cachedIdentities) {
             var identity = loadIdentityLocked(recipientId);
             return identity == null ? null : identity.getIdentityKey();
         }
     }
 
-    public IdentityInfo getIdentity(RecipientId recipientId) {
+    public IdentityInfo getIdentityInfo(RecipientId recipientId) {
         synchronized (cachedIdentities) {
             return loadIdentityLocked(recipientId);
         }
@@ -214,13 +189,6 @@ public class IdentityKeyStore implements org.signal.libsignal.protocol.state.Ide
         }
     }
 
-    /**
-     * @param identifier can be either a serialized uuid or a e164 phone number
-     */
-    private RecipientId resolveRecipient(String identifier) {
-        return resolver.resolveRecipient(identifier);
-    }
-
     private File getIdentityFile(final RecipientId recipientId) {
         try {
             IOUtils.createPrivateDirectories(identitiesPath);

lib/src/main/java/org/asamk/signal/manager/storage/identities/SignalIdentityKeyStore.java

@@ -0,0 +1,66 @@
+package org.asamk.signal.manager.storage.identities;
+
+import org.asamk.signal.manager.storage.recipients.RecipientId;
+import org.asamk.signal.manager.storage.recipients.RecipientResolver;
+import org.signal.libsignal.protocol.IdentityKey;
+import org.signal.libsignal.protocol.IdentityKeyPair;
+import org.signal.libsignal.protocol.SignalProtocolAddress;
+
+import java.util.function.Supplier;
+
+public class SignalIdentityKeyStore implements org.signal.libsignal.protocol.state.IdentityKeyStore {
+
+    private final RecipientResolver resolver;
+    private final Supplier<IdentityKeyPair> identityKeyPairSupplier;
+    private final int localRegistrationId;
+    private final IdentityKeyStore identityKeyStore;
+
+    public SignalIdentityKeyStore(
+            final RecipientResolver resolver,
+            final Supplier<IdentityKeyPair> identityKeyPairSupplier,
+            final int localRegistrationId,
+            final IdentityKeyStore identityKeyStore
+    ) {
+        this.resolver = resolver;
+        this.identityKeyPairSupplier = identityKeyPairSupplier;
+        this.localRegistrationId = localRegistrationId;
+        this.identityKeyStore = identityKeyStore;
+    }
+
+    @Override
+    public IdentityKeyPair getIdentityKeyPair() {
+        return identityKeyPairSupplier.get();
+    }
+
+    @Override
+    public int getLocalRegistrationId() {
+        return localRegistrationId;
+    }
+
+    @Override
+    public boolean saveIdentity(SignalProtocolAddress address, IdentityKey identityKey) {
+        final var recipientId = resolveRecipient(address.getName());
+
+        return identityKeyStore.saveIdentity(recipientId, identityKey);
+    }
+
+    @Override
+    public boolean isTrustedIdentity(SignalProtocolAddress address, IdentityKey identityKey, Direction direction) {
+        var recipientId = resolveRecipient(address.getName());
+
+        return identityKeyStore.isTrustedIdentity(recipientId, identityKey, direction);
+    }
+
+    @Override
+    public IdentityKey getIdentity(SignalProtocolAddress address) {
+        var recipientId = resolveRecipient(address.getName());
+        return identityKeyStore.getIdentity(recipientId);
+    }
+
+    /**
+     * @param identifier can be either a serialized uuid or an e164 phone number
+     */
+    private RecipientId resolveRecipient(String identifier) {
+        return resolver.resolveRecipient(identifier);
+    }
+}