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