Move session store to database

This commit is contained in:
AsamK 2022-06-09 16:54:48 +02:00
parent 65c9a2e185
commit 484daa4c69
4 changed files with 413 additions and 176 deletions

View file

@ -7,6 +7,7 @@ import org.asamk.signal.manager.storage.prekeys.PreKeyStore;
import org.asamk.signal.manager.storage.prekeys.SignedPreKeyStore; import org.asamk.signal.manager.storage.prekeys.SignedPreKeyStore;
import org.asamk.signal.manager.storage.recipients.RecipientStore; import org.asamk.signal.manager.storage.recipients.RecipientStore;
import org.asamk.signal.manager.storage.sendLog.MessageSendLogStore; import org.asamk.signal.manager.storage.sendLog.MessageSendLogStore;
import org.asamk.signal.manager.storage.sessions.SessionStore;
import org.asamk.signal.manager.storage.stickers.StickerStore; import org.asamk.signal.manager.storage.stickers.StickerStore;
import org.slf4j.Logger; import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
@ -18,7 +19,7 @@ import java.sql.SQLException;
public class AccountDatabase extends Database { public class AccountDatabase extends Database {
private final static Logger logger = LoggerFactory.getLogger(AccountDatabase.class); private final static Logger logger = LoggerFactory.getLogger(AccountDatabase.class);
private static final long DATABASE_VERSION = 5; private static final long DATABASE_VERSION = 6;
private AccountDatabase(final HikariDataSource dataSource) { private AccountDatabase(final HikariDataSource dataSource) {
super(logger, DATABASE_VERSION, dataSource); super(logger, DATABASE_VERSION, dataSource);
@ -36,6 +37,7 @@ public class AccountDatabase extends Database {
PreKeyStore.createSql(connection); PreKeyStore.createSql(connection);
SignedPreKeyStore.createSql(connection); SignedPreKeyStore.createSql(connection);
GroupStore.createSql(connection); GroupStore.createSql(connection);
SessionStore.createSql(connection);
} }
@Override @Override
@ -143,5 +145,20 @@ public class AccountDatabase extends Database {
"""); """);
} }
} }
if (oldVersion < 6) {
logger.debug("Updating database: Creating session tables");
try (final var statement = connection.createStatement()) {
statement.executeUpdate("""
CREATE TABLE session (
_id INTEGER PRIMARY KEY,
account_id_type INTEGER NOT NULL,
recipient_id INTEGER NOT NULL REFERENCES recipient (_id) ON DELETE CASCADE,
device_id INTEGER NOT NULL,
record BLOB NOT NULL,
UNIQUE(account_id_type, recipient_id, device_id)
);
""");
}
}
} }
} }

View file

@ -38,6 +38,7 @@ import org.asamk.signal.manager.storage.recipients.RecipientStore;
import org.asamk.signal.manager.storage.recipients.RecipientTrustedResolver; import org.asamk.signal.manager.storage.recipients.RecipientTrustedResolver;
import org.asamk.signal.manager.storage.sendLog.MessageSendLogStore; import org.asamk.signal.manager.storage.sendLog.MessageSendLogStore;
import org.asamk.signal.manager.storage.senderKeys.SenderKeyStore; import org.asamk.signal.manager.storage.senderKeys.SenderKeyStore;
import org.asamk.signal.manager.storage.sessions.LegacySessionStore;
import org.asamk.signal.manager.storage.sessions.SessionStore; import org.asamk.signal.manager.storage.sessions.SessionStore;
import org.asamk.signal.manager.storage.stickers.LegacyStickerStore; import org.asamk.signal.manager.storage.stickers.LegacyStickerStore;
import org.asamk.signal.manager.storage.stickers.StickerStore; import org.asamk.signal.manager.storage.stickers.StickerStore;
@ -634,6 +635,11 @@ public class SignalAccount implements Closeable {
LegacySignedPreKeyStore.migrate(legacyPniSignedPreKeysPath, getPniSignedPreKeyStore()); LegacySignedPreKeyStore.migrate(legacyPniSignedPreKeysPath, getPniSignedPreKeyStore());
migratedLegacyConfig = true; migratedLegacyConfig = true;
} }
final var legacySessionsPath = getSessionsPath(dataPath, accountPath);
if (legacySessionsPath.exists()) {
LegacySessionStore.migrate(legacySessionsPath, getRecipientResolver(), getSessionStore());
migratedLegacyConfig = true;
}
final var legacySignalProtocolStore = rootNode.hasNonNull("axolotlStore") final var legacySignalProtocolStore = rootNode.hasNonNull("axolotlStore")
? jsonProcessor.convertValue(Utils.getNotNullNode(rootNode, "axolotlStore"), ? jsonProcessor.convertValue(Utils.getNotNullNode(rootNode, "axolotlStore"),
LegacyJsonSignalProtocolStore.class) LegacyJsonSignalProtocolStore.class)
@ -1067,7 +1073,10 @@ public class SignalAccount implements Closeable {
public SessionStore getSessionStore() { public SessionStore getSessionStore() {
return getOrCreate(() -> sessionStore, return getOrCreate(() -> sessionStore,
() -> sessionStore = new SessionStore(getSessionsPath(dataPath, accountPath), getRecipientResolver())); () -> sessionStore = new SessionStore(getAccountDatabase(),
ServiceIdType.ACI,
getRecipientResolver(),
getRecipientIdCreator()));
} }
public IdentityKeyStore getIdentityKeyStore() { public IdentityKeyStore getIdentityKeyStore() {

View file

@ -0,0 +1,107 @@
package org.asamk.signal.manager.storage.sessions;
import org.asamk.signal.manager.api.Pair;
import org.asamk.signal.manager.storage.recipients.RecipientResolver;
import org.asamk.signal.manager.storage.sessions.SessionStore.Key;
import org.asamk.signal.manager.util.IOUtils;
import org.signal.libsignal.protocol.state.SessionRecord;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Objects;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
public class LegacySessionStore {
private final static Logger logger = LoggerFactory.getLogger(LegacySessionStore.class);
public static void migrate(
final File sessionsPath, final RecipientResolver resolver, final SessionStore sessionStore
) {
final var keys = getKeysLocked(sessionsPath, resolver);
final var sessions = keys.stream().map(key -> {
final var record = loadSessionLocked(key, sessionsPath);
if (record == null) {
return null;
}
return new Pair<>(key, record);
}).filter(Objects::nonNull).toList();
sessionStore.addLegacySessions(sessions);
deleteAllSessions(sessionsPath);
}
private static void deleteAllSessions(File sessionsPath) {
final var files = sessionsPath.listFiles();
if (files == null) {
return;
}
for (var file : files) {
try {
Files.delete(file.toPath());
} catch (IOException e) {
logger.error("Failed to delete session file {}: {}", file, e.getMessage());
}
}
try {
Files.delete(sessionsPath.toPath());
} catch (IOException e) {
logger.error("Failed to delete session directory {}: {}", sessionsPath, e.getMessage());
}
}
private static Collection<Key> getKeysLocked(File sessionsPath, final RecipientResolver resolver) {
final var files = sessionsPath.listFiles();
if (files == null) {
return List.of();
}
return parseFileNames(files, resolver);
}
static final Pattern sessionFileNamePattern = Pattern.compile("(\\d+)_(\\d+)");
private static List<Key> parseFileNames(final File[] files, final RecipientResolver resolver) {
return Arrays.stream(files)
.map(f -> sessionFileNamePattern.matcher(f.getName()))
.filter(Matcher::matches)
.map(matcher -> {
final var recipientId = resolver.resolveRecipient(Long.parseLong(matcher.group(1)));
if (recipientId == null) {
return null;
}
return new Key(recipientId, Integer.parseInt(matcher.group(2)));
})
.filter(Objects::nonNull)
.toList();
}
private static File getSessionFile(Key key, final File sessionsPath) {
try {
IOUtils.createPrivateDirectories(sessionsPath);
} catch (IOException e) {
throw new AssertionError("Failed to create sessions path", e);
}
return new File(sessionsPath, key.recipientId().id() + "_" + key.deviceId());
}
private static SessionRecord loadSessionLocked(final Key key, final File sessionsPath) {
final var file = getSessionFile(key, sessionsPath);
if (!file.exists()) {
return null;
}
try (var inputStream = new FileInputStream(file)) {
return new SessionRecord(inputStream.readAllBytes());
} catch (Exception e) {
logger.warn("Failed to load session, resetting session: {}", e.getMessage());
return null;
}
}
}

View file

@ -1,8 +1,12 @@
package org.asamk.signal.manager.storage.sessions; package org.asamk.signal.manager.storage.sessions;
import org.asamk.signal.manager.api.Pair;
import org.asamk.signal.manager.storage.Database;
import org.asamk.signal.manager.storage.Utils;
import org.asamk.signal.manager.storage.recipients.RecipientId; import org.asamk.signal.manager.storage.recipients.RecipientId;
import org.asamk.signal.manager.storage.recipients.RecipientIdCreator;
import org.asamk.signal.manager.storage.recipients.RecipientResolver; import org.asamk.signal.manager.storage.recipients.RecipientResolver;
import org.asamk.signal.manager.util.IOUtils; import org.signal.libsignal.protocol.InvalidMessageException;
import org.signal.libsignal.protocol.NoSessionException; import org.signal.libsignal.protocol.NoSessionException;
import org.signal.libsignal.protocol.SignalProtocolAddress; import org.signal.libsignal.protocol.SignalProtocolAddress;
import org.signal.libsignal.protocol.ecc.ECPublicKey; import org.signal.libsignal.protocol.ecc.ECPublicKey;
@ -11,50 +15,68 @@ import org.signal.libsignal.protocol.state.SessionRecord;
import org.slf4j.Logger; import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
import org.whispersystems.signalservice.api.SignalServiceSessionStore; import org.whispersystems.signalservice.api.SignalServiceSessionStore;
import org.whispersystems.signalservice.api.push.ServiceIdType;
import java.io.File; import java.sql.Connection;
import java.io.FileInputStream; import java.sql.ResultSet;
import java.io.FileOutputStream; import java.sql.SQLException;
import java.io.IOException; import java.util.ArrayList;
import java.nio.file.Files;
import java.util.Arrays;
import java.util.Collection; import java.util.Collection;
import java.util.HashMap; import java.util.HashMap;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
import java.util.Objects; import java.util.Objects;
import java.util.Set; import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors; import java.util.stream.Collectors;
public class SessionStore implements SignalServiceSessionStore { public class SessionStore implements SignalServiceSessionStore {
private static final String TABLE_SESSION = "session";
private final static Logger logger = LoggerFactory.getLogger(SessionStore.class); private final static Logger logger = LoggerFactory.getLogger(SessionStore.class);
private final Map<Key, SessionRecord> cachedSessions = new HashMap<>(); private final Map<Key, SessionRecord> cachedSessions = new HashMap<>();
private final File sessionsPath; private final Database database;
private final int accountIdType;
private final RecipientResolver resolver; private final RecipientResolver resolver;
private final RecipientIdCreator recipientIdCreator;
public static void createSql(Connection connection) throws SQLException {
// When modifying the CREATE statement here, also add a migration in AccountDatabase.java
try (final var statement = connection.createStatement()) {
statement.executeUpdate("""
CREATE TABLE session (
_id INTEGER PRIMARY KEY,
account_id_type INTEGER NOT NULL,
recipient_id INTEGER NOT NULL REFERENCES recipient (_id) ON DELETE CASCADE,
device_id INTEGER NOT NULL,
record BLOB NOT NULL,
UNIQUE(account_id_type, recipient_id, device_id)
);
""");
}
}
public SessionStore( public SessionStore(
final File sessionsPath, final RecipientResolver resolver final Database database,
final ServiceIdType serviceIdType,
final RecipientResolver resolver,
final RecipientIdCreator recipientIdCreator
) { ) {
this.sessionsPath = sessionsPath; this.database = database;
this.accountIdType = Utils.getAccountIdType(serviceIdType);
this.resolver = resolver; this.resolver = resolver;
this.recipientIdCreator = recipientIdCreator;
} }
@Override @Override
public SessionRecord loadSession(SignalProtocolAddress address) { public SessionRecord loadSession(SignalProtocolAddress address) {
final var key = getKey(address); final var key = getKey(address);
try (final var connection = database.getConnection()) {
synchronized (cachedSessions) { final var session = loadSession(connection, key);
final var session = loadSessionLocked(key); return Objects.requireNonNullElseGet(session, SessionRecord::new);
if (session == null) { } catch (SQLException e) {
return new SessionRecord(); throw new RuntimeException("Failed read from session store", e);
}
return session;
} }
} }
@ -62,8 +84,14 @@ public class SessionStore implements SignalServiceSessionStore {
public List<SessionRecord> loadExistingSessions(final List<SignalProtocolAddress> addresses) throws NoSessionException { public List<SessionRecord> loadExistingSessions(final List<SignalProtocolAddress> addresses) throws NoSessionException {
final var keys = addresses.stream().map(this::getKey).toList(); final var keys = addresses.stream().map(this::getKey).toList();
synchronized (cachedSessions) { try (final var connection = database.getConnection()) {
final var sessions = keys.stream().map(this::loadSessionLocked).filter(Objects::nonNull).toList(); final var sessions = new ArrayList<SessionRecord>();
for (final var key : keys) {
final var sessionRecord = loadSession(connection, key);
if (sessionRecord != null) {
sessions.add(sessionRecord);
}
}
if (sessions.size() != addresses.size()) { if (sessions.size() != addresses.size()) {
String message = "Mismatch! Asked for " String message = "Mismatch! Asked for "
@ -76,31 +104,44 @@ public class SessionStore implements SignalServiceSessionStore {
} }
return sessions; return sessions;
} catch (SQLException e) {
throw new RuntimeException("Failed read from session store", e);
} }
} }
@Override @Override
public List<Integer> getSubDeviceSessions(String name) { public List<Integer> getSubDeviceSessions(String name) {
final var recipientId = resolveRecipient(name); final var recipientId = resolver.resolveRecipient(name);
synchronized (cachedSessions) {
return getKeysLocked(recipientId).stream()
// get all sessions for recipient except primary device session // get all sessions for recipient except primary device session
.filter(key -> key.deviceId() != 1 && key.recipientId().equals(recipientId)) final var sql = (
.map(Key::deviceId) """
.toList(); SELECT s.device_id
FROM %s AS s
WHERE s.account_id_type = ? AND s.recipient_id = ? AND s.device_id != 1
"""
).formatted(TABLE_SESSION);
try (final var connection = database.getConnection()) {
try (final var statement = connection.prepareStatement(sql)) {
statement.setInt(1, accountIdType);
statement.setLong(2, recipientId.id());
return Utils.executeQueryForStream(statement, res -> res.getInt("device_id")).toList();
}
} catch (SQLException e) {
throw new RuntimeException("Failed read from session store", e);
} }
} }
public boolean isCurrentRatchetKey(RecipientId recipientId, int deviceId, ECPublicKey ratchetKey) { public boolean isCurrentRatchetKey(RecipientId recipientId, int deviceId, ECPublicKey ratchetKey) {
final var key = new Key(recipientId, deviceId); final var key = new Key(recipientId, deviceId);
synchronized (cachedSessions) { try (final var connection = database.getConnection()) {
final var session = loadSessionLocked(key); final var session = loadSession(connection, key);
if (session == null) { if (session == null) {
return false; return false;
} }
return session.currentRatchetKeyMatches(ratchetKey); return session.currentRatchetKeyMatches(ratchetKey);
} catch (SQLException e) {
throw new RuntimeException("Failed read from session store", e);
} }
} }
@ -108,8 +149,10 @@ public class SessionStore implements SignalServiceSessionStore {
public void storeSession(SignalProtocolAddress address, SessionRecord session) { public void storeSession(SignalProtocolAddress address, SessionRecord session) {
final var key = getKey(address); final var key = getKey(address);
synchronized (cachedSessions) { try (final var connection = database.getConnection()) {
storeSessionLocked(key, session); storeSession(connection, key, session);
} catch (SQLException e) {
throw new RuntimeException("Failed read from session store", e);
} }
} }
@ -117,9 +160,11 @@ public class SessionStore implements SignalServiceSessionStore {
public boolean containsSession(SignalProtocolAddress address) { public boolean containsSession(SignalProtocolAddress address) {
final var key = getKey(address); final var key = getKey(address);
synchronized (cachedSessions) { try (final var connection = database.getConnection()) {
final var session = loadSessionLocked(key); final var session = loadSession(connection, key);
return isActive(session); return isActive(session);
} catch (SQLException e) {
throw new RuntimeException("Failed read from session store", e);
} }
} }
@ -127,23 +172,24 @@ public class SessionStore implements SignalServiceSessionStore {
public void deleteSession(SignalProtocolAddress address) { public void deleteSession(SignalProtocolAddress address) {
final var key = getKey(address); final var key = getKey(address);
synchronized (cachedSessions) { try (final var connection = database.getConnection()) {
deleteSessionLocked(key); deleteSession(connection, key);
} catch (SQLException e) {
throw new RuntimeException("Failed update session store", e);
} }
} }
@Override @Override
public void deleteAllSessions(String name) { public void deleteAllSessions(String name) {
final var recipientId = resolveRecipient(name); final var recipientId = resolver.resolveRecipient(name);
deleteAllSessions(recipientId); deleteAllSessions(recipientId);
} }
public void deleteAllSessions(RecipientId recipientId) { public void deleteAllSessions(RecipientId recipientId) {
synchronized (cachedSessions) { try (final var connection = database.getConnection()) {
final var keys = getKeysLocked(recipientId); deleteAllSessions(connection, recipientId);
for (var key : keys) { } catch (SQLException e) {
deleteSessionLocked(key); throw new RuntimeException("Failed update session store", e);
}
} }
} }
@ -151,186 +197,244 @@ public class SessionStore implements SignalServiceSessionStore {
public void archiveSession(final SignalProtocolAddress address) { public void archiveSession(final SignalProtocolAddress address) {
final var key = getKey(address); final var key = getKey(address);
synchronized (cachedSessions) { try (final var connection = database.getConnection()) {
archiveSessionLocked(key); connection.setAutoCommit(false);
final var session = loadSession(connection, key);
if (session != null) {
session.archiveCurrentState();
storeSession(connection, key, session);
connection.commit();
}
} catch (SQLException e) {
throw new RuntimeException("Failed update session store", e);
} }
} }
@Override @Override
public Set<SignalProtocolAddress> getAllAddressesWithActiveSessions(final List<String> addressNames) { public Set<SignalProtocolAddress> getAllAddressesWithActiveSessions(final List<String> addressNames) {
final var recipientIdToNameMap = addressNames.stream() final var recipientIdToNameMap = addressNames.stream()
.collect(Collectors.toMap(this::resolveRecipient, name -> name)); .collect(Collectors.toMap(resolver::resolveRecipient, name -> name));
synchronized (cachedSessions) { final var recipientIdsCommaSeparated = recipientIdToNameMap.keySet()
return recipientIdToNameMap.keySet()
.stream() .stream()
.flatMap(recipientId -> getKeysLocked(recipientId).stream()) .map(recipientId -> String.valueOf(recipientId.id()))
.filter(key -> isActive(this.loadSessionLocked(key))) .collect(Collectors.joining(","));
.map(key -> new SignalProtocolAddress(recipientIdToNameMap.get(key.recipientId), key.deviceId())) final var sql = (
"""
SELECT s.recipient_id, s.device_id, s.record
FROM %s AS s
WHERE s.account_id_type = ? AND s.recipient_id IN (%s)
"""
).formatted(TABLE_SESSION, recipientIdsCommaSeparated);
try (final var connection = database.getConnection()) {
try (final var statement = connection.prepareStatement(sql)) {
statement.setInt(1, accountIdType);
return Utils.executeQueryForStream(statement,
res -> new Pair<>(getKeyFromResultSet(res), getSessionRecordFromResultSet(res)))
.filter(pair -> isActive(pair.second()))
.map(Pair::first)
.map(key -> new SignalProtocolAddress(recipientIdToNameMap.get(key.recipientId),
key.deviceId()))
.collect(Collectors.toSet()); .collect(Collectors.toSet());
} }
} catch (SQLException e) {
throw new RuntimeException("Failed read from session store", e);
}
} }
public void archiveAllSessions() { public void archiveAllSessions() {
synchronized (cachedSessions) { final var sql = (
final var keys = getKeysLocked(); """
for (var key : keys) { SELECT s.recipient_id, s.device_id, s.record
archiveSessionLocked(key); FROM %s AS s
WHERE s.account_id_type = ?
"""
).formatted(TABLE_SESSION);
try (final var connection = database.getConnection()) {
connection.setAutoCommit(false);
final List<Pair<Key, SessionRecord>> records;
try (final var statement = connection.prepareStatement(sql)) {
statement.setInt(1, accountIdType);
records = Utils.executeQueryForStream(statement,
res -> new Pair<>(getKeyFromResultSet(res), getSessionRecordFromResultSet(res))).toList();
} }
for (final var record : records) {
record.second().archiveCurrentState();
storeSession(connection, record.first(), record.second());
}
connection.commit();
} catch (SQLException e) {
throw new RuntimeException("Failed update session store", e);
} }
} }
public void archiveSessions(final RecipientId recipientId) { public void archiveSessions(final RecipientId recipientId) {
synchronized (cachedSessions) { final var sql = (
getKeysLocked().stream() """
.filter(key -> key.recipientId.equals(recipientId)) SELECT s.recipient_id, s.device_id, s.record
.forEach(this::archiveSessionLocked); FROM %s AS s
WHERE s.account_id_type = ? AND s.recipient_id = ?
"""
).formatted(TABLE_SESSION);
try (final var connection = database.getConnection()) {
connection.setAutoCommit(false);
final List<Pair<Key, SessionRecord>> records;
try (final var statement = connection.prepareStatement(sql)) {
statement.setInt(1, accountIdType);
statement.setLong(2, recipientId.id());
records = Utils.executeQueryForStream(statement,
res -> new Pair<>(getKeyFromResultSet(res), getSessionRecordFromResultSet(res))).toList();
}
for (final var record : records) {
record.second().archiveCurrentState();
storeSession(connection, record.first(), record.second());
}
connection.commit();
} catch (SQLException e) {
throw new RuntimeException("Failed update session store", e);
} }
} }
public void mergeRecipients(RecipientId recipientId, RecipientId toBeMergedRecipientId) { public void mergeRecipients(RecipientId recipientId, RecipientId toBeMergedRecipientId) {
try (final var connection = database.getConnection()) {
connection.setAutoCommit(false);
synchronized (cachedSessions) { synchronized (cachedSessions) {
final var keys = getKeysLocked(toBeMergedRecipientId); cachedSessions.clear();
final var otherHasSession = keys.size() > 0;
if (!otherHasSession) {
return;
} }
final var hasSession = getKeysLocked(recipientId).size() > 0; final var sql = """
if (hasSession) { UPDATE OR IGNORE %s
logger.debug("To be merged recipient had sessions, deleting."); SET recipient_id = ?
deleteAllSessions(toBeMergedRecipientId); WHERE account_id_type = ? AND recipient_id = ?
} else { """.formatted(TABLE_SESSION);
logger.debug("Only to be merged recipient had sessions, re-assigning to the new recipient."); try (final var statement = connection.prepareStatement(sql)) {
for (var key : keys) { statement.setLong(1, recipientId.id());
final var session = loadSessionLocked(key); statement.setInt(2, accountIdType);
deleteSessionLocked(key); statement.setLong(3, toBeMergedRecipientId.id());
if (session == null) { final var rows = statement.executeUpdate();
continue; if (rows > 0) {
} logger.debug("Reassigned {} sessions of to be merged recipient.", rows);
final var newKey = new Key(recipientId, key.deviceId());
storeSessionLocked(newKey, session);
} }
} }
// Delete all conflicting sessions now
deleteAllSessions(connection, toBeMergedRecipientId);
connection.commit();
} catch (SQLException e) {
throw new RuntimeException("Failed update session store", e);
} }
} }
/** void addLegacySessions(final Collection<Pair<Key, SessionRecord>> sessions) {
* @param identifier can be either a serialized uuid or a e164 phone number logger.debug("Migrating legacy sessions to database");
*/ long start = System.nanoTime();
private RecipientId resolveRecipient(String identifier) { try (final var connection = database.getConnection()) {
return resolver.resolveRecipient(identifier); connection.setAutoCommit(false);
for (final var pair : sessions) {
storeSession(connection, pair.first(), pair.second());
}
connection.commit();
} catch (SQLException e) {
throw new RuntimeException("Failed update session store", e);
}
logger.debug("Complete sessions migration took {}ms", (System.nanoTime() - start) / 1000000);
} }
private Key getKey(final SignalProtocolAddress address) { private Key getKey(final SignalProtocolAddress address) {
final var recipientId = resolveRecipient(address.getName()); final var recipientId = resolver.resolveRecipient(address.getName());
return new Key(recipientId, address.getDeviceId()); return new Key(recipientId, address.getDeviceId());
} }
private List<Key> getKeysLocked(RecipientId recipientId) { private SessionRecord loadSession(Connection connection, final Key key) throws SQLException {
final var files = sessionsPath.listFiles((_file, s) -> s.startsWith(recipientId.id() + "_")); synchronized (cachedSessions) {
if (files == null) {
return List.of();
}
return parseFileNames(files);
}
private Collection<Key> getKeysLocked() {
final var files = sessionsPath.listFiles();
if (files == null) {
return List.of();
}
return parseFileNames(files);
}
final Pattern sessionFileNamePattern = Pattern.compile("(\\d+)_(\\d+)");
private List<Key> parseFileNames(final File[] files) {
return Arrays.stream(files)
.map(f -> sessionFileNamePattern.matcher(f.getName()))
.filter(Matcher::matches)
.map(matcher -> {
final var recipientId = resolver.resolveRecipient(Long.parseLong(matcher.group(1)));
if (recipientId == null) {
return null;
}
return new Key(recipientId, Integer.parseInt(matcher.group(2)));
})
.filter(Objects::nonNull)
.toList();
}
private File getSessionFile(Key key) {
try {
IOUtils.createPrivateDirectories(sessionsPath);
} catch (IOException e) {
throw new AssertionError("Failed to create sessions path", e);
}
return new File(sessionsPath, key.recipientId().id() + "_" + key.deviceId());
}
private SessionRecord loadSessionLocked(final Key key) {
{
final var session = cachedSessions.get(key); final var session = cachedSessions.get(key);
if (session != null) { if (session != null) {
return session; return session;
} }
} }
final var sql = (
final var file = getSessionFile(key); """
if (!file.exists()) { SELECT s.record
return null; FROM %s AS s
WHERE s.account_id_type = ? AND s.recipient_id = ? AND s.device_id = ?
"""
).formatted(TABLE_SESSION);
try (final var statement = connection.prepareStatement(sql)) {
statement.setInt(1, accountIdType);
statement.setLong(2, key.recipientId().id());
statement.setInt(3, key.deviceId());
return Utils.executeQueryForOptional(statement, this::getSessionRecordFromResultSet).orElse(null);
} }
try (var inputStream = new FileInputStream(file)) { }
final var session = new SessionRecord(inputStream.readAllBytes());
cachedSessions.put(key, session); private Key getKeyFromResultSet(ResultSet resultSet) throws SQLException {
return session; final var recipientId = resultSet.getLong("recipient_id");
} catch (Exception e) { final var deviceId = resultSet.getInt("device_id");
return new Key(recipientIdCreator.create(recipientId), deviceId);
}
private SessionRecord getSessionRecordFromResultSet(ResultSet resultSet) throws SQLException {
try {
final var record = resultSet.getBytes("record");
return new SessionRecord(record);
} catch (InvalidMessageException e) {
logger.warn("Failed to load session, resetting session: {}", e.getMessage()); logger.warn("Failed to load session, resetting session: {}", e.getMessage());
return null; return null;
} }
} }
private void storeSessionLocked(final Key key, final SessionRecord session) { private void storeSession(
final Connection connection, final Key key, final SessionRecord session
) throws SQLException {
synchronized (cachedSessions) {
cachedSessions.put(key, session); cachedSessions.put(key, session);
}
final var file = getSessionFile(key); final var sql = """
try { INSERT OR REPLACE INTO %s (account_id_type, recipient_id, device_id, record)
try (var outputStream = new FileOutputStream(file)) { VALUES (?, ?, ?, ?)
outputStream.write(session.serialize()); """.formatted(TABLE_SESSION);
} try (final var statement = connection.prepareStatement(sql)) {
} catch (IOException e) { statement.setInt(1, accountIdType);
logger.warn("Failed to store session, trying to delete file and retry: {}", e.getMessage()); statement.setLong(2, key.recipientId().id());
try { statement.setInt(3, key.deviceId());
Files.delete(file.toPath()); statement.setBytes(4, session.serialize());
try (var outputStream = new FileOutputStream(file)) { statement.executeUpdate();
outputStream.write(session.serialize());
}
} catch (IOException e2) {
logger.error("Failed to store session file {}: {}", file, e2.getMessage());
}
} }
} }
private void archiveSessionLocked(final Key key) { private void deleteAllSessions(final Connection connection, final RecipientId recipientId) throws SQLException {
final var session = loadSessionLocked(key); synchronized (cachedSessions) {
if (session == null) { cachedSessions.clear();
return;
}
session.archiveCurrentState();
storeSessionLocked(key, session);
} }
private void deleteSessionLocked(final Key key) { final var sql = (
"""
DELETE FROM %s AS s
WHERE s.account_id_type = ? AND s.recipient_id = ?
"""
).formatted(TABLE_SESSION);
try (final var statement = connection.prepareStatement(sql)) {
statement.setInt(1, accountIdType);
statement.setLong(2, recipientId.id());
statement.executeUpdate();
}
}
private void deleteSession(Connection connection, final Key key) throws SQLException {
synchronized (cachedSessions) {
cachedSessions.remove(key); cachedSessions.remove(key);
final var file = getSessionFile(key);
if (!file.exists()) {
return;
} }
try {
Files.delete(file.toPath()); final var sql = (
} catch (IOException e) { """
logger.error("Failed to delete session file {}: {}", file, e.getMessage()); DELETE FROM %s AS s
WHERE s.account_id_type = ? AND s.recipient_id = ? AND s.device_id = ?
"""
).formatted(TABLE_SESSION);
try (final var statement = connection.prepareStatement(sql)) {
statement.setInt(1, accountIdType);
statement.setLong(2, key.recipientId().id());
statement.setInt(3, key.deviceId());
statement.executeUpdate();
} }
} }
@ -340,5 +444,5 @@ public class SessionStore implements SignalServiceSessionStore {
&& record.getSessionVersion() == CiphertextMessage.CURRENT_VERSION; && record.getSessionVersion() == CiphertextMessage.CURRENT_VERSION;
} }
private record Key(RecipientId recipientId, int deviceId) {} record Key(RecipientId recipientId, int deviceId) {}
} }