Refactor pre key store

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

@@ -482,7 +482,6 @@ public class Manager implements Closeable {
 
         var records = KeyUtils.generatePreKeyRecords(offset, ServiceConfig.PREKEY_BATCH_SIZE);
         account.addPreKeys(records);
-        account.save();
 
         return records;
     }
@@ -492,7 +491,6 @@ public class Manager implements Closeable {
 
         var record = KeyUtils.generateSignedPreKeyRecord(identityKeyPair, signedPreKeyId);
         account.addSignedPreKey(record);
-        account.save();
 
         return record;
     }

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

@@ -15,6 +15,7 @@ import org.asamk.signal.manager.storage.contacts.JsonContactsStore;
 import org.asamk.signal.manager.storage.groups.GroupInfoV1;
 import org.asamk.signal.manager.storage.groups.JsonGroupStore;
 import org.asamk.signal.manager.storage.messageCache.MessageCache;
+import org.asamk.signal.manager.storage.prekeys.PreKeyStore;
 import org.asamk.signal.manager.storage.profiles.ProfileStore;
 import org.asamk.signal.manager.storage.protocol.JsonSignalProtocolStore;
 import org.asamk.signal.manager.storage.protocol.SignalServiceAddressResolver;
@@ -54,7 +55,7 @@ import java.nio.channels.ClosedChannelException;
 import java.nio.channels.FileChannel;
 import java.nio.channels.FileLock;
 import java.util.Base64;
-import java.util.Collection;
+import java.util.List;
 import java.util.UUID;
 import java.util.stream.Collectors;
 
@@ -80,6 +81,7 @@ public class SignalAccount implements Closeable {
     private boolean registered = false;
 
     private JsonSignalProtocolStore signalProtocolStore;
+    private PreKeyStore preKeyStore;
     private SessionStore sessionStore;
     private JsonGroupStore groupStore;
     private JsonContactsStore contactStore;
@@ -133,9 +135,13 @@ public class SignalAccount implements Closeable {
         account.contactStore = new JsonContactsStore();
         account.recipientStore = RecipientStore.load(getRecipientsStoreFile(dataPath, username),
                 account::mergeRecipients);
+        account.preKeyStore = new PreKeyStore(getPreKeysPath(dataPath, username));
         account.sessionStore = new SessionStore(getSessionsPath(dataPath, username),
                 account.recipientStore::resolveRecipient);
-        account.signalProtocolStore = new JsonSignalProtocolStore(identityKey, registrationId, account.sessionStore);
+        account.signalProtocolStore = new JsonSignalProtocolStore(identityKey,
+                registrationId,
+                account.preKeyStore,
+                account.sessionStore);
         account.profileStore = new ProfileStore();
         account.stickerStore = new StickerStore();
 
@@ -176,9 +182,13 @@ public class SignalAccount implements Closeable {
         account.contactStore = new JsonContactsStore();
         account.recipientStore = RecipientStore.load(getRecipientsStoreFile(dataPath, username),
                 account::mergeRecipients);
+        account.preKeyStore = new PreKeyStore(getPreKeysPath(dataPath, username));
         account.sessionStore = new SessionStore(getSessionsPath(dataPath, username),
                 account.recipientStore::resolveRecipient);
-        account.signalProtocolStore = new JsonSignalProtocolStore(identityKey, registrationId, account.sessionStore);
+        account.signalProtocolStore = new JsonSignalProtocolStore(identityKey,
+                registrationId,
+                account.preKeyStore,
+                account.sessionStore);
         account.profileStore = new ProfileStore();
         account.stickerStore = new StickerStore();
 
@@ -237,6 +247,10 @@ public class SignalAccount implements Closeable {
         return new File(getUserPath(dataPath, username), "group-cache");
     }
 
+    private static File getPreKeysPath(File dataPath, String username) {
+        return new File(getUserPath(dataPath, username), "pre-keys");
+    }
+
     private static File getSessionsPath(File dataPath, String username) {
         return new File(getUserPath(dataPath, username), "sessions");
     }
@@ -316,6 +330,18 @@ public class SignalAccount implements Closeable {
 
         signalProtocolStore = jsonProcessor.convertValue(Utils.getNotNullNode(rootNode, "axolotlStore"),
                 JsonSignalProtocolStore.class);
+        preKeyStore = new PreKeyStore(getPreKeysPath(dataPath, username));
+        if (signalProtocolStore.getLegacyPreKeyStore() != null) {
+            logger.debug("Migrating legacy pre key store.");
+            for (var entry : signalProtocolStore.getLegacyPreKeyStore().getPreKeys().entrySet()) {
+                try {
+                    preKeyStore.storePreKey(entry.getKey(), new PreKeyRecord(entry.getValue()));
+                } catch (IOException e) {
+                    logger.warn("Failed to migrate pre key, ignoring", e);
+                }
+            }
+        }
+        signalProtocolStore.setPreKeyStore(preKeyStore);
         sessionStore = new SessionStore(getSessionsPath(dataPath, username), recipientStore::resolveRecipient);
         if (signalProtocolStore.getLegacySessionStore() != null) {
             logger.debug("Migrating legacy session store.");
@@ -469,16 +495,26 @@ public class SignalAccount implements Closeable {
         signalProtocolStore.setResolver(resolver);
     }
 
-    public void addPreKeys(Collection<PreKeyRecord> records) {
+    public void addPreKeys(List<PreKeyRecord> records) {
         for (var record : records) {
-            signalProtocolStore.storePreKey(record.getId(), record);
+            if (preKeyIdOffset != record.getId()) {
+                logger.error("Invalid pre key id {}, expected {}", record.getId(), preKeyIdOffset);
+                throw new AssertionError("Invalid pre key id");
+            }
+            preKeyStore.storePreKey(record.getId(), record);
+            preKeyIdOffset = (preKeyIdOffset + 1) % Medium.MAX_VALUE;
         }
-        preKeyIdOffset = (preKeyIdOffset + records.size()) % Medium.MAX_VALUE;
+        save();
     }
 
     public void addSignedPreKey(SignedPreKeyRecord record) {
+        if (nextSignedPreKeyId != record.getId()) {
+            logger.error("Invalid signed pre key id {}, expected {}", record.getId(), nextSignedPreKeyId);
+            throw new AssertionError("Invalid signed pre key id");
+        }
         signalProtocolStore.storeSignedPreKey(record.getId(), record);
         nextSignedPreKeyId = (nextSignedPreKeyId + 1) % Medium.MAX_VALUE;
+        save();
     }
 
     public JsonSignalProtocolStore getSignalProtocolStore() {

lib/src/main/java/org/asamk/signal/manager/storage/prekeys/PreKeyStore.java

@@ -0,0 +1,89 @@
+package org.asamk.signal.manager.storage.prekeys;
+
+import org.asamk.signal.manager.util.IOUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.libsignal.InvalidKeyIdException;
+import org.whispersystems.libsignal.state.PreKeyRecord;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.nio.file.Files;
+
+public class PreKeyStore implements org.whispersystems.libsignal.state.PreKeyStore {
+
+    private final static Logger logger = LoggerFactory.getLogger(PreKeyStore.class);
+
+    private final File preKeysPath;
+
+    public PreKeyStore(final File preKeysPath) {
+        this.preKeysPath = preKeysPath;
+    }
+
+    @Override
+    public PreKeyRecord loadPreKey(int preKeyId) throws InvalidKeyIdException {
+        final var file = getPreKeyFile(preKeyId);
+
+        if (!file.exists()) {
+            throw new InvalidKeyIdException("No such pre key record!");
+        }
+        try (var inputStream = new FileInputStream(file)) {
+            return new PreKeyRecord(inputStream.readAllBytes());
+        } catch (IOException e) {
+            logger.error("Failed to load pre key: {}", e.getMessage());
+            throw new AssertionError(e);
+        }
+    }
+
+    @Override
+    public void storePreKey(int preKeyId, PreKeyRecord record) {
+        final var file = getPreKeyFile(preKeyId);
+        try {
+            try (var outputStream = new FileOutputStream(file)) {
+                outputStream.write(record.serialize());
+            }
+        } catch (IOException e) {
+            logger.warn("Failed to store pre key, trying to delete file and retry: {}", e.getMessage());
+            try {
+                Files.delete(file.toPath());
+                try (var outputStream = new FileOutputStream(file)) {
+                    outputStream.write(record.serialize());
+                }
+            } catch (IOException e2) {
+                logger.error("Failed to store pre key file {}: {}", file, e2.getMessage());
+            }
+        }
+    }
+
+    @Override
+    public boolean containsPreKey(int preKeyId) {
+        final var file = getPreKeyFile(preKeyId);
+
+        return file.exists();
+    }
+
+    @Override
+    public void removePreKey(int preKeyId) {
+        final var file = getPreKeyFile(preKeyId);
+
+        if (!file.exists()) {
+            return;
+        }
+        try {
+            Files.delete(file.toPath());
+        } catch (IOException e) {
+            logger.error("Failed to delete pre key file {}: {}", file, e.getMessage());
+        }
+    }
+
+    private File getPreKeyFile(int preKeyId) {
+        try {
+            IOUtils.createPrivateDirectories(preKeysPath);
+        } catch (IOException e) {
+            throw new AssertionError("Failed to create pre keys path", e);
+        }
+        return new File(preKeysPath, String.valueOf(preKeyId));
+    }
+}

lib/src/main/java/org/asamk/signal/manager/storage/protocol/JsonPreKeyStore.java

@@ -1,104 +0,0 @@
-package org.asamk.signal.manager.storage.protocol;
-
-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 org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import org.whispersystems.libsignal.InvalidKeyIdException;
-import org.whispersystems.libsignal.state.PreKeyRecord;
-import org.whispersystems.libsignal.state.PreKeyStore;
-
-import java.io.IOException;
-import java.util.Base64;
-import java.util.HashMap;
-import java.util.Map;
-
-class JsonPreKeyStore implements PreKeyStore {
-
-    private final static Logger logger = LoggerFactory.getLogger(JsonPreKeyStore.class);
-
-    private final Map<Integer, byte[]> store = new HashMap<>();
-
-    public JsonPreKeyStore() {
-
-    }
-
-    private void addPreKeys(Map<Integer, byte[]> preKeys) {
-        store.putAll(preKeys);
-    }
-
-    @Override
-    public PreKeyRecord loadPreKey(int preKeyId) throws InvalidKeyIdException {
-        try {
-            if (!store.containsKey(preKeyId)) {
-                throw new InvalidKeyIdException("No such prekeyrecord!");
-            }
-
-            return new PreKeyRecord(store.get(preKeyId));
-        } catch (IOException e) {
-            throw new AssertionError(e);
-        }
-    }
-
-    @Override
-    public void storePreKey(int preKeyId, PreKeyRecord record) {
-        store.put(preKeyId, record.serialize());
-    }
-
-    @Override
-    public boolean containsPreKey(int preKeyId) {
-        return store.containsKey(preKeyId);
-    }
-
-    @Override
-    public void removePreKey(int preKeyId) {
-        store.remove(preKeyId);
-    }
-
-    public static class JsonPreKeyStoreDeserializer extends JsonDeserializer<JsonPreKeyStore> {
-
-        @Override
-        public JsonPreKeyStore deserialize(
-                JsonParser jsonParser, DeserializationContext deserializationContext
-        ) throws IOException {
-            JsonNode node = jsonParser.getCodec().readTree(jsonParser);
-
-            var preKeyMap = new HashMap<Integer, byte[]>();
-            if (node.isArray()) {
-                for (var preKey : node) {
-                    final var preKeyId = preKey.get("id").asInt();
-                    final var preKeyRecord = Base64.getDecoder().decode(preKey.get("record").asText());
-                    preKeyMap.put(preKeyId, preKeyRecord);
-                }
-            }
-
-            var keyStore = new JsonPreKeyStore();
-            keyStore.addPreKeys(preKeyMap);
-
-            return keyStore;
-        }
-    }
-
-    public static class JsonPreKeyStoreSerializer extends JsonSerializer<JsonPreKeyStore> {
-
-        @Override
-        public void serialize(
-                JsonPreKeyStore jsonPreKeyStore, JsonGenerator json, SerializerProvider serializerProvider
-        ) throws IOException {
-            json.writeStartArray();
-            for (var preKey : jsonPreKeyStore.store.entrySet()) {
-                json.writeStartObject();
-                json.writeNumberField("id", preKey.getKey());
-                json.writeStringField("record", Base64.getEncoder().encodeToString(preKey.getValue()));
-                json.writeEndObject();
-            }
-            json.writeEndArray();
-        }
-    }
-}

lib/src/main/java/org/asamk/signal/manager/storage/protocol/JsonSignalProtocolStore.java

@@ -6,26 +6,26 @@ import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 
 import org.asamk.signal.manager.TrustLevel;
-import org.asamk.signal.manager.storage.sessions.SessionStore;
 import org.whispersystems.libsignal.IdentityKey;
 import org.whispersystems.libsignal.IdentityKeyPair;
 import org.whispersystems.libsignal.InvalidKeyIdException;
 import org.whispersystems.libsignal.SignalProtocolAddress;
 import org.whispersystems.libsignal.state.PreKeyRecord;
+import org.whispersystems.libsignal.state.PreKeyStore;
 import org.whispersystems.libsignal.state.SessionRecord;
 import org.whispersystems.libsignal.state.SignedPreKeyRecord;
 import org.whispersystems.signalservice.api.SignalServiceProtocolStore;
+import org.whispersystems.signalservice.api.SignalServiceSessionStore;
 import org.whispersystems.signalservice.api.push.SignalServiceAddress;
 
 import java.util.List;
 
-@JsonIgnoreProperties(value = "sessionStore", allowSetters = true)
+@JsonIgnoreProperties(value = {"sessionStore", "preKeys"}, allowSetters = true)
 public class JsonSignalProtocolStore implements SignalServiceProtocolStore {
 
     @JsonProperty("preKeys")
-    @JsonDeserialize(using = JsonPreKeyStore.JsonPreKeyStoreDeserializer.class)
+    @JsonDeserialize(using = LegacyJsonPreKeyStore.JsonPreKeyStoreDeserializer.class)
-    @JsonSerialize(using = JsonPreKeyStore.JsonPreKeyStoreSerializer.class)
+    private LegacyJsonPreKeyStore legacyPreKeyStore;
-    private JsonPreKeyStore preKeyStore;
 
     @JsonProperty("sessionStore")
     @JsonDeserialize(using = LegacyJsonSessionStore.JsonSessionStoreDeserializer.class)
@@ -41,13 +41,19 @@ public class JsonSignalProtocolStore implements SignalServiceProtocolStore {
     @JsonSerialize(using = JsonIdentityKeyStore.JsonIdentityKeyStoreSerializer.class)
     private JsonIdentityKeyStore identityKeyStore;
 
-    private SessionStore sessionStore;
+    private PreKeyStore preKeyStore;
+    private SignalServiceSessionStore sessionStore;
 
     public JsonSignalProtocolStore() {
     }
 
-    public JsonSignalProtocolStore(IdentityKeyPair identityKeyPair, int registrationId, SessionStore sessionStore) {
+    public JsonSignalProtocolStore(
-        preKeyStore = new JsonPreKeyStore();
+            IdentityKeyPair identityKeyPair,
+            int registrationId,
+            PreKeyStore preKeyStore,
+            SignalServiceSessionStore sessionStore
+    ) {
+        this.preKeyStore = preKeyStore;
         this.sessionStore = sessionStore;
         signedPreKeyStore = new JsonSignedPreKeyStore();
         this.identityKeyStore = new JsonIdentityKeyStore(identityKeyPair, registrationId);
@@ -57,10 +63,18 @@ public class JsonSignalProtocolStore implements SignalServiceProtocolStore {
         identityKeyStore.setResolver(resolver);
     }
 
-    public void setSessionStore(final SessionStore sessionStore) {
+    public void setPreKeyStore(final PreKeyStore preKeyStore) {
+        this.preKeyStore = preKeyStore;
+    }
+
+    public void setSessionStore(final SignalServiceSessionStore sessionStore) {
         this.sessionStore = sessionStore;
     }
 
+    public LegacyJsonPreKeyStore getLegacyPreKeyStore() {
+        return legacyPreKeyStore;
+    }
+
     public LegacyJsonSessionStore getLegacySessionStore() {
         return legacySessionStore;
     }

lib/src/main/java/org/asamk/signal/manager/storage/protocol/LegacyJsonPreKeyStore.java

@@ -0,0 +1,45 @@
+package org.asamk.signal.manager.storage.protocol;
+
+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 java.io.IOException;
+import java.util.Base64;
+import java.util.HashMap;
+import java.util.Map;
+
+public class LegacyJsonPreKeyStore {
+
+    private final Map<Integer, byte[]> preKeys;
+
+    public LegacyJsonPreKeyStore(final Map<Integer, byte[]> preKeys) {
+        this.preKeys = preKeys;
+    }
+
+    public Map<Integer, byte[]> getPreKeys() {
+        return preKeys;
+    }
+
+    public static class JsonPreKeyStoreDeserializer extends JsonDeserializer<LegacyJsonPreKeyStore> {
+
+        @Override
+        public LegacyJsonPreKeyStore deserialize(
+                JsonParser jsonParser, DeserializationContext deserializationContext
+        ) throws IOException {
+            JsonNode node = jsonParser.getCodec().readTree(jsonParser);
+
+            var preKeyMap = new HashMap<Integer, byte[]>();
+            if (node.isArray()) {
+                for (var preKey : node) {
+                    final var preKeyId = preKey.get("id").asInt();
+                    final var preKeyRecord = Base64.getDecoder().decode(preKey.get("record").asText());
+                    preKeyMap.put(preKeyId, preKeyRecord);
+                }
+            }
+
+            return new LegacyJsonPreKeyStore(preKeyMap);
+        }
+    }
+}