package org.asamk.signal.manager.storage.prekeys; import org.asamk.signal.manager.storage.Database; import org.asamk.signal.manager.storage.Utils; import org.signal.libsignal.protocol.InvalidKeyException; import org.signal.libsignal.protocol.InvalidKeyIdException; import org.signal.libsignal.protocol.ecc.ECKeyPair; import org.signal.libsignal.protocol.ecc.ECPrivateKey; import org.signal.libsignal.protocol.ecc.ECPublicKey; import org.signal.libsignal.protocol.state.PreKeyRecord; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.whispersystems.signalservice.api.SignalServicePreKeyStore; import org.whispersystems.signalservice.api.push.ServiceIdType; import java.sql.Connection; import java.sql.ResultSet; import java.sql.SQLException; import java.util.Collection; public class PreKeyStore implements SignalServicePreKeyStore { private static final String TABLE_PRE_KEY = "pre_key"; private static final Logger logger = LoggerFactory.getLogger(PreKeyStore.class); private final Database database; private final int accountIdType; 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 pre_key ( _id INTEGER PRIMARY KEY, account_id_type INTEGER NOT NULL, key_id INTEGER NOT NULL, public_key BLOB NOT NULL, private_key BLOB NOT NULL, stale_timestamp INTEGER, UNIQUE(account_id_type, key_id) ) STRICT; """); } } public PreKeyStore(final Database database, final ServiceIdType serviceIdType) { this.database = database; this.accountIdType = Utils.getAccountIdType(serviceIdType); } @Override public PreKeyRecord loadPreKey(int preKeyId) throws InvalidKeyIdException { final var preKey = getPreKey(preKeyId); if (preKey == null) { throw new InvalidKeyIdException("No such pre key record: " + preKeyId); } return preKey; } @Override public void storePreKey(int preKeyId, PreKeyRecord record) { final var sql = ( """ INSERT INTO %s (account_id_type, key_id, public_key, private_key) VALUES (?, ?, ?, ?) """ ).formatted(TABLE_PRE_KEY); try (final var connection = database.getConnection()) { try (final var statement = connection.prepareStatement(sql)) { statement.setInt(1, accountIdType); statement.setInt(2, preKeyId); final var keyPair = record.getKeyPair(); statement.setBytes(3, keyPair.getPublicKey().serialize()); statement.setBytes(4, keyPair.getPrivateKey().serialize()); statement.executeUpdate(); } catch (InvalidKeyException ignored) { } } catch (SQLException e) { throw new RuntimeException("Failed update pre_key store", e); } } @Override public boolean containsPreKey(int preKeyId) { return getPreKey(preKeyId) != null; } @Override public void removePreKey(int preKeyId) { final var sql = ( """ DELETE FROM %s AS p WHERE p.account_id_type = ? AND p.key_id = ? """ ).formatted(TABLE_PRE_KEY); try (final var connection = database.getConnection()) { try (final var statement = connection.prepareStatement(sql)) { statement.setInt(1, accountIdType); statement.setInt(2, preKeyId); statement.executeUpdate(); } } catch (SQLException e) { throw new RuntimeException("Failed update pre_key store", e); } } public void removeAllPreKeys() { final var sql = ( """ DELETE FROM %s AS p WHERE p.account_id_type = ? """ ).formatted(TABLE_PRE_KEY); try (final var connection = database.getConnection()) { try (final var statement = connection.prepareStatement(sql)) { statement.setInt(1, accountIdType); statement.executeUpdate(); } } catch (SQLException e) { throw new RuntimeException("Failed update pre_key store", e); } } void addLegacyPreKeys(final Collection preKeys) { logger.debug("Migrating legacy preKeys to database"); long start = System.nanoTime(); final var sql = ( """ INSERT INTO %s (account_id_type, key_id, public_key, private_key) VALUES (?, ?, ?, ?) """ ).formatted(TABLE_PRE_KEY); try (final var connection = database.getConnection()) { connection.setAutoCommit(false); final var deleteSql = "DELETE FROM %s AS p WHERE p.account_id_type = ?".formatted(TABLE_PRE_KEY); try (final var statement = connection.prepareStatement(deleteSql)) { statement.setInt(1, accountIdType); statement.executeUpdate(); } try (final var statement = connection.prepareStatement(sql)) { for (final var record : preKeys) { statement.setInt(1, accountIdType); statement.setInt(2, record.getId()); final var keyPair = record.getKeyPair(); statement.setBytes(3, keyPair.getPublicKey().serialize()); statement.setBytes(4, keyPair.getPrivateKey().serialize()); statement.executeUpdate(); } } catch (InvalidKeyException ignored) { } connection.commit(); } catch (SQLException e) { throw new RuntimeException("Failed update preKey store", e); } logger.debug("Complete preKeys migration took {}ms", (System.nanoTime() - start) / 1000000); } private PreKeyRecord getPreKey(int preKeyId) { final var sql = ( """ SELECT p.key_id, p.public_key, p.private_key FROM %s p WHERE p.account_id_type = ? AND p.key_id = ? """ ).formatted(TABLE_PRE_KEY); try (final var connection = database.getConnection()) { try (final var statement = connection.prepareStatement(sql)) { statement.setInt(1, accountIdType); statement.setInt(2, preKeyId); return Utils.executeQueryForOptional(statement, this::getPreKeyRecordFromResultSet).orElse(null); } } catch (SQLException e) { throw new RuntimeException("Failed read from pre_key store", e); } } private PreKeyRecord getPreKeyRecordFromResultSet(ResultSet resultSet) throws SQLException { try { final var keyId = resultSet.getInt("key_id"); final var publicKey = new ECPublicKey(resultSet.getBytes("public_key")); final var privateKey = new ECPrivateKey(resultSet.getBytes("private_key")); return new PreKeyRecord(keyId, new ECKeyPair(publicKey, privateKey)); } catch (InvalidKeyException e) { return null; } } @Override public void deleteAllStaleOneTimeEcPreKeys(final long threshold, final int minCount) { final var sql = ( """ DELETE FROM %s AS p WHERE p.account_id_type = ?1 AND p.stale_timestamp < ?2 AND p._id NOT IN ( SELECT _id FROM %s AS p2 WHERE p2.account_id_type = ?1 ORDER BY CASE WHEN p2.stale_timestamp IS NULL THEN 1 ELSE 0 END DESC, p2.stale_timestamp DESC, p2._id DESC LIMIT ?3 ) """ ).formatted(TABLE_PRE_KEY, TABLE_PRE_KEY); try (final var connection = database.getConnection()) { try (final var statement = connection.prepareStatement(sql)) { statement.setInt(1, accountIdType); statement.setLong(2, threshold); statement.setInt(3, minCount); final var rowCount = statement.executeUpdate(); if (rowCount > 0) { logger.debug("Deleted {} stale one time pre keys", rowCount); } } } catch (SQLException e) { throw new RuntimeException("Failed update pre_key store", e); } } @Override public void markAllOneTimeEcPreKeysStaleIfNecessary(final long staleTime) { final var sql = ( """ UPDATE %s SET stale_timestamp = ? WHERE account_id_type = ? AND stale_timestamp IS NULL """ ).formatted(TABLE_PRE_KEY); try (final var connection = database.getConnection()) { try (final var statement = connection.prepareStatement(sql)) { statement.setLong(1, staleTime); statement.setInt(2, accountIdType); statement.executeUpdate(); } } catch (SQLException e) { throw new RuntimeException("Failed update pre_key store", e); } } }