nmode's Git Repositories - signal-cli/blob - lib/src/main/java/org/asamk/signal/manager/storage/senderKeys/SenderKeyRecordStore.java

   1 package org.asamk.signal.manager.storage.senderKeys;
   2
   3 import org.asamk.signal.manager.api.Pair;
   4 import org.asamk.signal.manager.storage.Database;
   5 import org.asamk.signal.manager.storage.Utils;
   6 import org.signal.libsignal.protocol.InvalidMessageException;
   7 import org.signal.libsignal.protocol.SignalProtocolAddress;
   8 import org.signal.libsignal.protocol.groups.state.SenderKeyRecord;
   9 import org.signal.libsignal.protocol.groups.state.SenderKeyStore;
  10 import org.slf4j.Logger;
  11 import org.slf4j.LoggerFactory;
  12 import org.whispersystems.signalservice.api.push.ServiceId;
  13 import org.whispersystems.signalservice.api.util.UuidUtil;
  14
  15 import java.sql.Connection;
  16 import java.sql.ResultSet;
  17 import java.sql.SQLException;
  18 import java.util.Collection;
  19 import java.util.UUID;
  20
  21 public class SenderKeyRecordStore implements SenderKeyStore {
  22
  23     private final static Logger logger = LoggerFactory.getLogger(SenderKeyRecordStore.class);
  24     private final static String TABLE_SENDER_KEY = "sender_key";
  25
  26     private final Database database;
  27
  28     public static void createSql(Connection connection) throws SQLException {
  29         // When modifying the CREATE statement here, also add a migration in AccountDatabase.java
  30         try (final var statement = connection.createStatement()) {
  31             statement.executeUpdate("""
  32                                     CREATE TABLE sender_key (
  33                                       _id INTEGER PRIMARY KEY,
  34                                       uuid BLOB NOT NULL,
  35                                       device_id INTEGER NOT NULL,
  36                                       distribution_id BLOB NOT NULL,
  37                                       record BLOB NOT NULL,
  38                                       created_timestamp INTEGER NOT NULL,
  39                                       UNIQUE(uuid, device_id, distribution_id)
  40                                     ) STRICT;
  41                                     """);
  42         }
  43     }
  44
  45     SenderKeyRecordStore(final Database database) {
  46         this.database = database;
  47     }
  48
  49     @Override
  50     public SenderKeyRecord loadSenderKey(final SignalProtocolAddress address, final UUID distributionId) {
  51         final var key = getKey(address, distributionId);
  52
  53         try (final var connection = database.getConnection()) {
  54             return loadSenderKey(connection, key);
  55         } catch (SQLException e) {
  56             throw new RuntimeException("Failed read from sender key store", e);
  57         }
  58     }
  59
  60     @Override
  61     public void storeSenderKey(
  62             final SignalProtocolAddress address, final UUID distributionId, final SenderKeyRecord record
  63     ) {
  64         final var key = getKey(address, distributionId);
  65
  66         try (final var connection = database.getConnection()) {
  67             storeSenderKey(connection, key, record);
  68         } catch (SQLException e) {
  69             throw new RuntimeException("Failed update sender key store", e);
  70         }
  71     }
  72
  73     long getCreateTimeForKey(final ServiceId selfServiceId, final int selfDeviceId, final UUID distributionId) {
  74         final var sql = (
  75                 """
  76                 SELECT s.created_timestamp
  77                 FROM %s AS s
  78                 WHERE s.uuid = ? AND s.device_id = ? AND s.distribution_id = ?
  79                 """
  80         ).formatted(TABLE_SENDER_KEY);
  81         try (final var connection = database.getConnection()) {
  82             try (final var statement = connection.prepareStatement(sql)) {
  83                 statement.setBytes(1, selfServiceId.toByteArray());
  84                 statement.setInt(2, selfDeviceId);
  85                 statement.setBytes(3, UuidUtil.toByteArray(distributionId));
  86                 return Utils.executeQueryForOptional(statement, res -> res.getLong("created_timestamp")).orElse(-1L);
  87             }
  88         } catch (SQLException e) {
  89             throw new RuntimeException("Failed read from sender key store", e);
  90         }
  91     }
  92
  93     void deleteSenderKey(final ServiceId serviceId, final UUID distributionId) {
  94         final var sql = (
  95                 """
  96                 DELETE FROM %s AS s
  97                 WHERE s.uuid = ? AND s.distribution_id = ?
  98                 """
  99         ).formatted(TABLE_SENDER_KEY);
 100         try (final var connection = database.getConnection()) {
 101             try (final var statement = connection.prepareStatement(sql)) {
 102                 statement.setBytes(1, serviceId.toByteArray());
 103                 statement.setBytes(2, UuidUtil.toByteArray(distributionId));
 104                 statement.executeUpdate();
 105             }
 106         } catch (SQLException e) {
 107             throw new RuntimeException("Failed update sender key store", e);
 108         }
 109     }
 110
 111     void deleteAll() {
 112         final var sql = """
 113                         DELETE FROM %s AS s
 114                         """.formatted(TABLE_SENDER_KEY);
 115         try (final var connection = database.getConnection()) {
 116             try (final var statement = connection.prepareStatement(sql)) {
 117                 statement.executeUpdate();
 118             }
 119         } catch (SQLException e) {
 120             throw new RuntimeException("Failed update sender key store", e);
 121         }
 122     }
 123
 124     void deleteAllFor(final ServiceId serviceId) {
 125         try (final var connection = database.getConnection()) {
 126             deleteAllFor(connection, serviceId);
 127         } catch (SQLException e) {
 128             throw new RuntimeException("Failed update sender key store", e);
 129         }
 130     }
 131
 132     void addLegacySenderKeys(final Collection<Pair<Key, SenderKeyRecord>> senderKeys) {
 133         logger.debug("Migrating legacy sender keys to database");
 134         long start = System.nanoTime();
 135         try (final var connection = database.getConnection()) {
 136             connection.setAutoCommit(false);
 137             for (final var pair : senderKeys) {
 138                 storeSenderKey(connection, pair.first(), pair.second());
 139             }
 140             connection.commit();
 141         } catch (SQLException e) {
 142             throw new RuntimeException("Failed update sender keys store", e);
 143         }
 144         logger.debug("Complete sender keys migration took {}ms", (System.nanoTime() - start) / 1000000);
 145     }
 146
 147     private Key getKey(final SignalProtocolAddress address, final UUID distributionId) {
 148         final var serviceId = ServiceId.parseOrThrow(address.getName());
 149         return new Key(serviceId, address.getDeviceId(), distributionId);
 150     }
 151
 152     private SenderKeyRecord loadSenderKey(final Connection connection, final Key key) throws SQLException {
 153         final var sql = (
 154                 """
 155                 SELECT s.record
 156                 FROM %s AS s
 157                 WHERE s.uuid = ? AND s.device_id = ? AND s.distribution_id = ?
 158                 """
 159         ).formatted(TABLE_SENDER_KEY);
 160         try (final var statement = connection.prepareStatement(sql)) {
 161             statement.setBytes(1, key.serviceId().toByteArray());
 162             statement.setInt(2, key.deviceId());
 163             statement.setBytes(3, UuidUtil.toByteArray(key.distributionId()));
 164             return Utils.executeQueryForOptional(statement, this::getSenderKeyRecordFromResultSet).orElse(null);
 165         }
 166     }
 167
 168     private void storeSenderKey(
 169             final Connection connection, final Key key, final SenderKeyRecord senderKeyRecord
 170     ) throws SQLException {
 171         final var sqlUpdate = """
 172                               UPDATE %s
 173                               SET record = ?
 174                               WHERE uuid = ? AND device_id = ? and distribution_id = ?
 175                               """.formatted(TABLE_SENDER_KEY);
 176         try (final var statement = connection.prepareStatement(sqlUpdate)) {
 177             statement.setBytes(1, senderKeyRecord.serialize());
 178             statement.setBytes(2, key.serviceId().toByteArray());
 179             statement.setLong(3, key.deviceId());
 180             statement.setBytes(4, UuidUtil.toByteArray(key.distributionId()));
 181             final var rows = statement.executeUpdate();
 182             if (rows > 0) {
 183                 return;
 184             }
 185         }
 186
 187         // Record doesn't exist yet, creating a new one
 188         final var sqlInsert = (
 189                 """
 190                 INSERT OR REPLACE INTO %s (uuid, device_id, distribution_id, record, created_timestamp)
 191                 VALUES (?, ?, ?, ?, ?)
 192                 """
 193         ).formatted(TABLE_SENDER_KEY);
 194         try (final var statement = connection.prepareStatement(sqlInsert)) {
 195             statement.setBytes(1, key.serviceId().toByteArray());
 196             statement.setInt(2, key.deviceId());
 197             statement.setBytes(3, UuidUtil.toByteArray(key.distributionId()));
 198             statement.setBytes(4, senderKeyRecord.serialize());
 199             statement.setLong(5, System.currentTimeMillis());
 200             statement.executeUpdate();
 201         }
 202     }
 203
 204     private void deleteAllFor(final Connection connection, final ServiceId serviceId) throws SQLException {
 205         final var sql = (
 206                 """
 207                 DELETE FROM %s AS s
 208                 WHERE s.uuid = ?
 209                 """
 210         ).formatted(TABLE_SENDER_KEY);
 211         try (final var statement = connection.prepareStatement(sql)) {
 212             statement.setBytes(1, serviceId.toByteArray());
 213             statement.executeUpdate();
 214         }
 215     }
 216
 217     private SenderKeyRecord getSenderKeyRecordFromResultSet(ResultSet resultSet) throws SQLException {
 218         try {
 219             final var record = resultSet.getBytes("record");
 220
 221             return new SenderKeyRecord(record);
 222         } catch (InvalidMessageException e) {
 223             logger.warn("Failed to load sender key, resetting: {}", e.getMessage());
 224             return null;
 225         }
 226     }
 227
 228     record Key(ServiceId serviceId, int deviceId, UUID distributionId) {}
 229 }