Refactor recipients store

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

@@ -10,6 +10,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.SerializationFeature;
 
 import org.asamk.signal.manager.groups.GroupId;
+import org.asamk.signal.manager.storage.contacts.ContactInfo;
 import org.asamk.signal.manager.storage.contacts.JsonContactsStore;
 import org.asamk.signal.manager.storage.groups.GroupInfoV1;
 import org.asamk.signal.manager.storage.groups.JsonGroupStore;
@@ -17,6 +18,8 @@ import org.asamk.signal.manager.storage.messageCache.MessageCache;
 import org.asamk.signal.manager.storage.profiles.ProfileStore;
 import org.asamk.signal.manager.storage.protocol.JsonSignalProtocolStore;
 import org.asamk.signal.manager.storage.protocol.SignalServiceAddressResolver;
+import org.asamk.signal.manager.storage.recipients.LegacyRecipientStore;
+import org.asamk.signal.manager.storage.recipients.RecipientId;
 import org.asamk.signal.manager.storage.recipients.RecipientStore;
 import org.asamk.signal.manager.storage.stickers.StickerStore;
 import org.asamk.signal.manager.storage.threads.LegacyJsonThreadStore;
@@ -125,7 +128,8 @@ public class SignalAccount implements Closeable {
         account.signalProtocolStore = new JsonSignalProtocolStore(identityKey, registrationId);
         account.groupStore = new JsonGroupStore(getGroupCachePath(dataPath, username));
         account.contactStore = new JsonContactsStore();
-        account.recipientStore = new RecipientStore();
+        account.recipientStore = RecipientStore.load(getRecipientsStoreFile(dataPath, username),
+                account::mergeRecipients);
         account.profileStore = new ProfileStore();
         account.stickerStore = new StickerStore();
 
@@ -165,7 +169,8 @@ public class SignalAccount implements Closeable {
         account.signalProtocolStore = new JsonSignalProtocolStore(identityKey, registrationId);
         account.groupStore = new JsonGroupStore(getGroupCachePath(dataPath, username));
         account.contactStore = new JsonContactsStore();
-        account.recipientStore = new RecipientStore();
+        account.recipientStore = RecipientStore.load(getRecipientsStoreFile(dataPath, username),
+                account::mergeRecipients);
         account.profileStore = new ProfileStore();
         account.stickerStore = new StickerStore();
 
@@ -204,6 +209,10 @@ public class SignalAccount implements Closeable {
         getProfileStore().storeProfileKey(getSelfAddress(), getProfileKey());
     }
 
+    private void mergeRecipients(RecipientId recipientId, RecipientId toBeMergedRecipientId) {
+        // TODO
+    }
+
     public static File getFileName(File dataPath, String username) {
         return new File(dataPath, username);
     }
@@ -220,6 +229,10 @@ public class SignalAccount implements Closeable {
         return new File(getUserPath(dataPath, username), "group-cache");
     }
 
+    private static File getRecipientsStoreFile(File dataPath, String username) {
+        return new File(getUserPath(dataPath, username), "recipients-store");
+    }
+
     public static boolean userExists(File dataPath, String username) {
         if (username == null) {
             return false;
@@ -279,6 +292,16 @@ public class SignalAccount implements Closeable {
             }
         }
 
+        recipientStore = RecipientStore.load(getRecipientsStoreFile(dataPath, username), this::mergeRecipients);
+        var legacyRecipientStoreNode = rootNode.get("recipientStore");
+        if (legacyRecipientStoreNode != null) {
+            logger.debug("Migrating legacy recipient store.");
+            var legacyRecipientStore = jsonProcessor.convertValue(legacyRecipientStoreNode, LegacyRecipientStore.class);
+            if (legacyRecipientStore != null) {
+                recipientStore.resolveRecipients(legacyRecipientStore.getAddresses());
+            }
+        }
+
         signalProtocolStore = jsonProcessor.convertValue(Utils.getNotNullNode(rootNode, "axolotlStore"),
                 JsonSignalProtocolStore.class);
         registered = Utils.getNotNullNode(rootNode, "registered").asBoolean();
@@ -299,18 +322,29 @@ public class SignalAccount implements Closeable {
             contactStore = new JsonContactsStore();
         }
 
-        var recipientStoreNode = rootNode.get("recipientStore");
-        if (recipientStoreNode != null) {
-            recipientStore = jsonProcessor.convertValue(recipientStoreNode, RecipientStore.class);
+        var profileStoreNode = rootNode.get("profileStore");
+        if (profileStoreNode != null) {
+            profileStore = jsonProcessor.convertValue(profileStoreNode, ProfileStore.class);
+        }
+        if (profileStore == null) {
+            profileStore = new ProfileStore();
         }
-        if (recipientStore == null) {
-            recipientStore = new RecipientStore();
 
-            recipientStore.resolveServiceAddress(getSelfAddress());
+        var stickerStoreNode = rootNode.get("stickerStore");
+        if (stickerStoreNode != null) {
+            stickerStore = jsonProcessor.convertValue(stickerStoreNode, StickerStore.class);
+        }
+        if (stickerStore == null) {
+            stickerStore = new StickerStore();
+        }
 
-            for (var contact : contactStore.getContacts()) {
-                recipientStore.resolveServiceAddress(contact.getAddress());
-            }
+        if (recipientStore.isEmpty()) {
+            recipientStore.resolveRecipient(getSelfAddress());
+
+            recipientStore.resolveRecipients(contactStore.getContacts()
+                    .stream()
+                    .map(ContactInfo::getAddress)
+                    .collect(Collectors.toList()));
 
             for (var group : groupStore.getGroups()) {
                 if (group instanceof GroupInfoV1) {
@@ -330,22 +364,6 @@ public class SignalAccount implements Closeable {
             }
         }
 
-        var profileStoreNode = rootNode.get("profileStore");
-        if (profileStoreNode != null) {
-            profileStore = jsonProcessor.convertValue(profileStoreNode, ProfileStore.class);
-        }
-        if (profileStore == null) {
-            profileStore = new ProfileStore();
-        }
-
-        var stickerStoreNode = rootNode.get("stickerStore");
-        if (stickerStoreNode != null) {
-            stickerStore = jsonProcessor.convertValue(stickerStoreNode, StickerStore.class);
-        }
-        if (stickerStore == null) {
-            stickerStore = new StickerStore();
-        }
-
         messageCache = new MessageCache(getMessageCachePath(dataPath, username));
 
         var threadStoreNode = rootNode.get("threadStore");
@@ -396,7 +414,6 @@ public class SignalAccount implements Closeable {
                 .putPOJO("axolotlStore", signalProtocolStore)
                 .putPOJO("groupStore", groupStore)
                 .putPOJO("contactStore", contactStore)
-                .putPOJO("recipientStore", recipientStore)
                 .putPOJO("profileStore", profileStore)
                 .putPOJO("stickerStore", stickerStore);
         try {

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

@@ -0,0 +1,27 @@
+package org.asamk.signal.manager.storage;
+
+import com.fasterxml.jackson.annotation.JsonAutoDetect;
+import com.fasterxml.jackson.annotation.PropertyAccessor;
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializationFeature;
+
+public class Utils {
+
+    private Utils() {
+    }
+
+    public static ObjectMapper createStorageObjectMapper() {
+        final ObjectMapper jsonProcessor = new ObjectMapper();
+
+        jsonProcessor.setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.PUBLIC_ONLY);
+        jsonProcessor.enable(SerializationFeature.INDENT_OUTPUT); // for pretty print
+        jsonProcessor.disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES);
+        jsonProcessor.disable(JsonParser.Feature.AUTO_CLOSE_SOURCE);
+        jsonProcessor.disable(JsonGenerator.Feature.AUTO_CLOSE_TARGET);
+
+        return jsonProcessor;
+    }
+}

lib/src/main/java/org/asamk/signal/manager/storage/recipients/LegacyRecipientStore.java

@@ -0,0 +1,49 @@
+package org.asamk.signal.manager.storage.recipients;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.DeserializationContext;
+import com.fasterxml.jackson.databind.JsonDeserializer;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+
+import org.whispersystems.signalservice.api.push.SignalServiceAddress;
+import org.whispersystems.signalservice.api.util.UuidUtil;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+public class LegacyRecipientStore {
+
+    @JsonProperty("recipientStore")
+    @JsonDeserialize(using = RecipientStoreDeserializer.class)
+    private final List<SignalServiceAddress> addresses = new ArrayList<>();
+
+    public List<SignalServiceAddress> getAddresses() {
+        return addresses;
+    }
+
+    public static class RecipientStoreDeserializer extends JsonDeserializer<List<SignalServiceAddress>> {
+
+        @Override
+        public List<SignalServiceAddress> deserialize(
+                JsonParser jsonParser, DeserializationContext deserializationContext
+        ) throws IOException {
+            JsonNode node = jsonParser.getCodec().readTree(jsonParser);
+
+            var addresses = new ArrayList<SignalServiceAddress>();
+
+            if (node.isArray()) {
+                for (var recipient : node) {
+                    var recipientName = recipient.get("name").asText();
+                    var uuid = UuidUtil.parseOrThrow(recipient.get("uuid").asText());
+                    final var serviceAddress = new SignalServiceAddress(uuid, recipientName);
+                    addresses.add(serviceAddress);
+                }
+            }
+
+            return addresses;
+        }
+    }
+}

lib/src/main/java/org/asamk/signal/manager/storage/recipients/RecipientId.java

@@ -0,0 +1,29 @@
+package org.asamk.signal.manager.storage.recipients;
+
+public class RecipientId {
+
+    private final long id;
+
+    RecipientId(final long id) {
+        this.id = id;
+    }
+
+    long getId() {
+        return id;
+    }
+
+    @Override
+    public boolean equals(final Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        final RecipientId that = (RecipientId) o;
+
+        return id == that.id;
+    }
+
+    @Override
+    public int hashCode() {
+        return (int) (id ^ (id >>> 32));
+    }
+}

lib/src/main/java/org/asamk/signal/manager/storage/recipients/RecipientStore.java

@@ -1,87 +1,314 @@
 package org.asamk.signal.manager.storage.recipients;
 
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.fasterxml.jackson.core.JsonGenerator;
-import com.fasterxml.jackson.core.JsonParser;
-import com.fasterxml.jackson.databind.DeserializationContext;
-import com.fasterxml.jackson.databind.JsonDeserializer;
-import com.fasterxml.jackson.databind.JsonNode;
-import com.fasterxml.jackson.databind.JsonSerializer;
-import com.fasterxml.jackson.databind.SerializerProvider;
-import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
-import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+import com.fasterxml.jackson.databind.ObjectMapper;
 
+import org.asamk.signal.manager.storage.Utils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.libsignal.util.Pair;
 import org.whispersystems.signalservice.api.push.SignalServiceAddress;
 import org.whispersystems.signalservice.api.util.UuidUtil;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.FileOutputStream;
 import java.io.IOException;
-import java.util.HashSet;
-import java.util.Set;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.UUID;
+import java.util.stream.Collectors;
 
 public class RecipientStore {
 
-    @JsonProperty("recipientStore")
-    @JsonDeserialize(using = RecipientStoreDeserializer.class)
-    @JsonSerialize(using = RecipientStoreSerializer.class)
-    private final Set<SignalServiceAddress> addresses = new HashSet<>();
+    private final static Logger logger = LoggerFactory.getLogger(RecipientStore.class);
 
-    public SignalServiceAddress resolveServiceAddress(SignalServiceAddress serviceAddress) {
-        if (addresses.contains(serviceAddress)) {
-            // If the Set already contains the exact address with UUID and Number,
-            // we can just return it here.
-            return serviceAddress;
+    private final ObjectMapper objectMapper;
+    private final File file;
+    private final RecipientMergeHandler recipientMergeHandler;
+
+    private final Map<RecipientId, SignalServiceAddress> recipients;
+    private final Map<RecipientId, RecipientId> recipientsMerged = new HashMap<>();
+
+    private long lastId;
+
+    public static RecipientStore load(File file, RecipientMergeHandler recipientMergeHandler) throws IOException {
+        final var objectMapper = Utils.createStorageObjectMapper();
+        try (var inputStream = new FileInputStream(file)) {
+            var storage = objectMapper.readValue(inputStream, Storage.class);
+            return new RecipientStore(objectMapper,
+                    file,
+                    recipientMergeHandler,
+                    storage.recipients.stream()
+                            .collect(Collectors.toMap(r -> new RecipientId(r.id),
+                                    r -> new SignalServiceAddress(org.whispersystems.libsignal.util.guava.Optional.fromNullable(
+                                            r.uuid).transform(UuidUtil::parseOrThrow),
+                                            org.whispersystems.libsignal.util.guava.Optional.fromNullable(r.name)))),
+                    storage.lastId);
+        } catch (FileNotFoundException e) {
+            logger.debug("Creating new recipient store.");
+            return new RecipientStore(objectMapper, file, recipientMergeHandler, new HashMap<>(), 0);
         }
-
-        for (var address : addresses) {
-            if (address.matches(serviceAddress)) {
-                return address;
-            }
-        }
-
-        if (serviceAddress.getNumber().isPresent() && serviceAddress.getUuid().isPresent()) {
-            addresses.add(serviceAddress);
-        }
-
-        return serviceAddress;
     }
 
-    public static class RecipientStoreDeserializer extends JsonDeserializer<Set<SignalServiceAddress>> {
+    private RecipientStore(
+            final ObjectMapper objectMapper,
+            final File file,
+            final RecipientMergeHandler recipientMergeHandler,
+            final Map<RecipientId, SignalServiceAddress> recipients,
+            final long lastId
+    ) {
+        this.objectMapper = objectMapper;
+        this.file = file;
+        this.recipientMergeHandler = recipientMergeHandler;
+        this.recipients = recipients;
+        this.lastId = lastId;
+    }
 
-        @Override
-        public Set<SignalServiceAddress> deserialize(
-                JsonParser jsonParser, DeserializationContext deserializationContext
-        ) throws IOException {
-            JsonNode node = jsonParser.getCodec().readTree(jsonParser);
+    public SignalServiceAddress resolveServiceAddress(RecipientId recipientId) {
+        synchronized (recipients) {
+            while (recipientsMerged.containsKey(recipientId)) {
+                recipientId = recipientsMerged.get(recipientId);
+            }
+            return recipients.get(recipientId);
+        }
+    }
 
-            var addresses = new HashSet<SignalServiceAddress>();
+    @Deprecated
+    public SignalServiceAddress resolveServiceAddress(SignalServiceAddress address) {
+        return resolveServiceAddress(resolveRecipient(address, true));
+    }
 
-            if (node.isArray()) {
-                for (var recipient : node) {
-                    var recipientName = recipient.get("name").asText();
-                    var uuid = UuidUtil.parseOrThrow(recipient.get("uuid").asText());
-                    final var serviceAddress = new SignalServiceAddress(uuid, recipientName);
-                    addresses.add(serviceAddress);
+    public RecipientId resolveRecipient(UUID uuid) {
+        return resolveRecipient(new SignalServiceAddress(uuid, null), false);
+    }
+
+    public RecipientId resolveRecipient(String number) {
+        return resolveRecipient(new SignalServiceAddress(null, number), false);
+    }
+
+    public RecipientId resolveRecipient(SignalServiceAddress address) {
+        return resolveRecipient(address, true);
+    }
+
+    public List<RecipientId> resolveRecipients(List<SignalServiceAddress> addresses) {
+        final List<RecipientId> recipientIds;
+        final List<Pair<RecipientId, RecipientId>> toBeMerged = new ArrayList<>();
+        synchronized (recipients) {
+            recipientIds = addresses.stream().map(address -> {
+                final var pair = resolveRecipientLocked(address, true);
+                if (pair.second().isPresent()) {
+                    toBeMerged.add(new Pair<>(pair.first(), pair.second().get()));
                 }
+                return pair.first();
+            }).collect(Collectors.toList());
+        }
+        for (var pair : toBeMerged) {
+            recipientMergeHandler.mergeRecipients(pair.first(), pair.second());
+        }
+        return recipientIds;
+    }
+
+    public RecipientId resolveRecipientUntrusted(SignalServiceAddress address) {
+        return resolveRecipient(address, false);
+    }
+
+    /**
+     * @param isHighTrust true, if the number/uuid connection was obtained from a trusted source.
+     *                    Has no effect, if the address contains only a number or a uuid.
+     */
+    private RecipientId resolveRecipient(SignalServiceAddress address, boolean isHighTrust) {
+        final Pair<RecipientId, Optional<RecipientId>> pair;
+        synchronized (recipients) {
+            pair = resolveRecipientLocked(address, isHighTrust);
+            if (pair.second().isPresent()) {
+                recipientsMerged.put(pair.second().get(), pair.first());
+            }
+        }
+
+        if (pair.second().isPresent()) {
+            recipientMergeHandler.mergeRecipients(pair.first(), pair.second().get());
+        }
+        return pair.first();
+    }
+
+    private Pair<RecipientId, Optional<RecipientId>> resolveRecipientLocked(
+            SignalServiceAddress address, boolean isHighTrust
+    ) {
+        final var byNumber = !address.getNumber().isPresent()
+                ? Optional.<RecipientId>empty()
+                : findByName(address.getNumber().get());
+        final var byUuid = !address.getUuid().isPresent()
+                ? Optional.<RecipientId>empty()
+                : findByUuid(address.getUuid().get());
+
+        if (byNumber.isEmpty() && byUuid.isEmpty()) {
+            logger.debug("Got new recipient, both uuid and number are unknown");
+
+            if (isHighTrust || !address.getUuid().isPresent() || !address.getNumber().isPresent()) {
+                return new Pair<>(addNewRecipient(address), Optional.empty());
             }
 
-            return addresses;
+            return new Pair<>(addNewRecipient(new SignalServiceAddress(address.getUuid().get(), null)),
+                    Optional.empty());
+        }
+
+        if (!isHighTrust
+                || !address.getUuid().isPresent()
+                || !address.getNumber().isPresent()
+                || byNumber.equals(byUuid)) {
+            return new Pair<>(byUuid.orElseGet(byNumber::get), Optional.empty());
+        }
+
+        if (byNumber.isEmpty()) {
+            logger.debug("Got recipient existing with uuid, updating with high trust number");
+            recipients.put(byUuid.get(), address);
+            save();
+            return new Pair<>(byUuid.get(), Optional.empty());
+        }
+
+        if (byUuid.isEmpty()) {
+            logger.debug("Got recipient existing with number, updating with high trust uuid");
+            recipients.put(byNumber.get(), address);
+            save();
+            return new Pair<>(byNumber.get(), Optional.empty());
+        }
+
+        final var byNumberAddress = recipients.get(byNumber.get());
+        if (byNumberAddress.getUuid().isPresent()) {
+            logger.debug(
+                    "Got separate recipients for high trust number and uuid, recipient for number has different uuid, so stripping its number");
+
+            recipients.put(byNumber.get(), new SignalServiceAddress(byNumberAddress.getUuid().get(), null));
+            recipients.put(byUuid.get(), address);
+            save();
+            return new Pair<>(byUuid.get(), Optional.empty());
+        }
+
+        logger.debug("Got separate recipients for high trust number and uuid, need to merge them");
+        recipients.put(byUuid.get(), address);
+        recipients.remove(byNumber.get());
+        save();
+        return new Pair<>(byUuid.get(), byNumber);
+    }
+
+    private RecipientId addNewRecipient(final SignalServiceAddress serviceAddress) {
+        final var nextRecipientId = nextId();
+        recipients.put(nextRecipientId, serviceAddress);
+        save();
+        return nextRecipientId;
+    }
+
+    private Optional<RecipientId> findByName(final String number) {
+        return recipients.entrySet()
+                .stream()
+                .filter(entry -> entry.getValue().getNumber().isPresent() && number.equals(entry.getValue()
+                        .getNumber()
+                        .get()))
+                .findFirst()
+                .map(Map.Entry::getKey);
+    }
+
+    private Optional<RecipientId> findByUuid(final UUID uuid) {
+        return recipients.entrySet()
+                .stream()
+                .filter(entry -> entry.getValue().getUuid().isPresent() && uuid.equals(entry.getValue()
+                        .getUuid()
+                        .get()))
+                .findFirst()
+                .map(Map.Entry::getKey);
+    }
+
+    private RecipientId nextId() {
+        return new RecipientId(++this.lastId);
+    }
+
+    private void save() {
+        // Write to memory first to prevent corrupting the file in case of serialization errors
+        try (var inMemoryOutput = new ByteArrayOutputStream()) {
+            var storage = new Storage(recipients.entrySet()
+                    .stream()
+                    .map(pair -> new Storage.Recipient(pair.getKey().getId(),
+                            pair.getValue().getNumber().orNull(),
+                            pair.getValue().getUuid().transform(UUID::toString).orNull()))
+                    .collect(Collectors.toList()), lastId);
+            objectMapper.writeValue(inMemoryOutput, storage);
+
+            var input = new ByteArrayInputStream(inMemoryOutput.toByteArray());
+            try (var outputStream = new FileOutputStream(file)) {
+                input.transferTo(outputStream);
+            }
+        } catch (Exception e) {
+            logger.error("Error saving recipient store file: {}", e.getMessage());
         }
     }
 
-    public static class RecipientStoreSerializer extends JsonSerializer<Set<SignalServiceAddress>> {
-
-        @Override
-        public void serialize(
-                Set<SignalServiceAddress> addresses, JsonGenerator json, SerializerProvider serializerProvider
-        ) throws IOException {
-            json.writeStartArray();
-            for (var address : addresses) {
-                json.writeStartObject();
-                json.writeStringField("name", address.getNumber().get());
-                json.writeStringField("uuid", address.getUuid().get().toString());
-                json.writeEndObject();
-            }
-            json.writeEndArray();
+    public boolean isEmpty() {
+        synchronized (recipients) {
+            return recipients.isEmpty();
         }
     }
+
+    private static class Storage {
+
+        private List<Recipient> recipients;
+
+        private long lastId;
+
+        // For deserialization
+        private Storage() {
+        }
+
+        public Storage(final List<Recipient> recipients, final long lastId) {
+            this.recipients = recipients;
+            this.lastId = lastId;
+        }
+
+        public List<Recipient> getRecipients() {
+            return recipients;
+        }
+
+        public long getLastId() {
+            return lastId;
+        }
+
+        public static class Recipient {
+
+            private long id;
+            private String name;
+            private String uuid;
+
+            // For deserialization
+            private Recipient() {
+            }
+
+            public Recipient(final long id, final String name, final String uuid) {
+                this.id = id;
+                this.name = name;
+                this.uuid = uuid;
+            }
+
+            public long getId() {
+                return id;
+            }
+
+            public String getName() {
+                return name;
+            }
+
+            public String getUuid() {
+                return uuid;
+            }
+        }
+    }
+
+    public interface RecipientMergeHandler {
+
+        void mergeRecipients(RecipientId recipientId, RecipientId toBeMergedRecipientId);
+    }
 }