mirror of
https://github.com/AsamK/signal-cli
synced 2025-08-29 10:30:38 +00:00
Refactor pre key store
This commit is contained in:
AsamK
parent
f77519445c
commit
ccc380f575
6 changed files with 199 additions and 121 deletions
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
}
|
||||
preKeyIdOffset = (preKeyIdOffset + records.size()) % Medium.MAX_VALUE;
|
||||
preKeyStore.storePreKey(record.getId(), record);
|
||||
preKeyIdOffset = (preKeyIdOffset + 1) % 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() {
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
}
|
||||
}
|
||||
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue