From: AsamK Date: Sat, 17 Apr 2021 12:11:31 +0000 (+0200) Subject: Refactor pre key store X-Git-Tag: v0.8.2~46 X-Git-Url: https://git.nmode.ca/signal-cli/commitdiff_plain/ccc380f57528ccddf2aec0addf40bab60de9588f?ds=sidebyside Refactor pre key store --- diff --git a/lib/src/main/java/org/asamk/signal/manager/Manager.java b/lib/src/main/java/org/asamk/signal/manager/Manager.java index f672d0fe..c1189df3 100644 --- a/lib/src/main/java/org/asamk/signal/manager/Manager.java +++ b/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; } diff --git a/lib/src/main/java/org/asamk/signal/manager/storage/SignalAccount.java b/lib/src/main/java/org/asamk/signal/manager/storage/SignalAccount.java index 18526f2a..494a7e89 100644 --- a/lib/src/main/java/org/asamk/signal/manager/storage/SignalAccount.java +++ b/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 records) { + public void addPreKeys(List 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() { diff --git a/lib/src/main/java/org/asamk/signal/manager/storage/prekeys/PreKeyStore.java b/lib/src/main/java/org/asamk/signal/manager/storage/prekeys/PreKeyStore.java new file mode 100644 index 00000000..c8153c56 --- /dev/null +++ b/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)); + } +} diff --git a/lib/src/main/java/org/asamk/signal/manager/storage/protocol/JsonPreKeyStore.java b/lib/src/main/java/org/asamk/signal/manager/storage/protocol/JsonPreKeyStore.java deleted file mode 100644 index 9ff0d8ea..00000000 --- a/lib/src/main/java/org/asamk/signal/manager/storage/protocol/JsonPreKeyStore.java +++ /dev/null @@ -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 store = new HashMap<>(); - - public JsonPreKeyStore() { - - } - - private void addPreKeys(Map 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 { - - @Override - public JsonPreKeyStore deserialize( - JsonParser jsonParser, DeserializationContext deserializationContext - ) throws IOException { - JsonNode node = jsonParser.getCodec().readTree(jsonParser); - - var preKeyMap = new HashMap(); - 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 { - - @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(); - } - } -} diff --git a/lib/src/main/java/org/asamk/signal/manager/storage/protocol/JsonSignalProtocolStore.java b/lib/src/main/java/org/asamk/signal/manager/storage/protocol/JsonSignalProtocolStore.java index dc45a0da..29e5c031 100644 --- a/lib/src/main/java/org/asamk/signal/manager/storage/protocol/JsonSignalProtocolStore.java +++ b/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) - @JsonSerialize(using = JsonPreKeyStore.JsonPreKeyStoreSerializer.class) - private JsonPreKeyStore preKeyStore; + @JsonDeserialize(using = LegacyJsonPreKeyStore.JsonPreKeyStoreDeserializer.class) + private LegacyJsonPreKeyStore legacyPreKeyStore; @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) { - preKeyStore = new JsonPreKeyStore(); + public JsonSignalProtocolStore( + 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; } diff --git a/lib/src/main/java/org/asamk/signal/manager/storage/protocol/LegacyJsonPreKeyStore.java b/lib/src/main/java/org/asamk/signal/manager/storage/protocol/LegacyJsonPreKeyStore.java new file mode 100644 index 00000000..24101e10 --- /dev/null +++ b/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 preKeys; + + public LegacyJsonPreKeyStore(final Map preKeys) { + this.preKeys = preKeys; + } + + public Map getPreKeys() { + return preKeys; + } + + public static class JsonPreKeyStoreDeserializer extends JsonDeserializer { + + @Override + public LegacyJsonPreKeyStore deserialize( + JsonParser jsonParser, DeserializationContext deserializationContext + ) throws IOException { + JsonNode node = jsonParser.getCodec().readTree(jsonParser); + + var preKeyMap = new HashMap(); + 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); + } + } +}