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

   1 package org.asamk.signal.manager.storage.senderKeys;
   2
   3 import org.asamk.signal.manager.helper.RecipientAddressResolver;
   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.RecipientIdCreator;
   8 import org.asamk.signal.manager.storage.recipients.RecipientResolver;
   9 import org.signal.libsignal.protocol.SignalProtocolAddress;
  10 import org.slf4j.Logger;
  11 import org.slf4j.LoggerFactory;
  12 import org.whispersystems.signalservice.api.push.DistributionId;
  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.Map;
  20 import java.util.Set;
  21 import java.util.stream.Collectors;
  22
  23 public class SenderKeySharedStore {
  24
  25     private final static Logger logger = LoggerFactory.getLogger(SenderKeySharedStore.class);
  26     private final static String TABLE_SENDER_KEY_SHARED = "sender_key_shared";
  27
  28     private final Database database;
  29     private final RecipientIdCreator recipientIdCreator;
  30     private final RecipientResolver resolver;
  31     private final RecipientAddressResolver addressResolver;
  32
  33     public static void createSql(Connection connection) throws SQLException {
  34         // When modifying the CREATE statement here, also add a migration in AccountDatabase.java
  35         try (final var statement = connection.createStatement()) {
  36             statement.executeUpdate("""
  37                                     CREATE TABLE sender_key_shared (
  38                                       _id INTEGER PRIMARY KEY,
  39                                       recipient_id INTEGER NOT NULL REFERENCES recipient (_id) ON DELETE CASCADE,
  40                                       device_id INTEGER NOT NULL,
  41                                       distribution_id BLOB NOT NULL,
  42                                       timestamp INTEGER NOT NULL,
  43                                       UNIQUE(recipient_id, device_id, distribution_id)
  44                                     );
  45                                     """);
  46         }
  47     }
  48
  49     SenderKeySharedStore(
  50             final Database database,
  51             final RecipientIdCreator recipientIdCreator,
  52             final RecipientAddressResolver addressResolver,
  53             final RecipientResolver resolver
  54     ) {
  55         this.database = database;
  56         this.recipientIdCreator = recipientIdCreator;
  57         this.addressResolver = addressResolver;
  58         this.resolver = resolver;
  59     }
  60
  61     public Set<SignalProtocolAddress> getSenderKeySharedWith(final DistributionId distributionId) {
  62         try (final var connection = database.getConnection()) {
  63             final var sql = (
  64                     """
  65                     SELECT s.recipient_id, s.device_id
  66                     FROM %s AS s
  67                     WHERE s.distribution_id = ?
  68                     """
  69             ).formatted(TABLE_SENDER_KEY_SHARED);
  70             try (final var statement = connection.prepareStatement(sql)) {
  71                 statement.setBytes(1, UuidUtil.toByteArray(distributionId.asUuid()));
  72                 return Utils.executeQueryForStream(statement, this::getSenderKeySharedEntryFromResultSet)
  73                         .map(k -> new SignalProtocolAddress(addressResolver.resolveRecipientAddress(k.recipientId())
  74                                 .getIdentifier(), k.deviceId()))
  75                         .collect(Collectors.toSet());
  76             }
  77         } catch (SQLException e) {
  78             throw new RuntimeException("Failed read from shared sender key store", e);
  79         }
  80     }
  81
  82     public void markSenderKeySharedWith(
  83             final DistributionId distributionId, final Collection<SignalProtocolAddress> addresses
  84     ) {
  85         final var newEntries = addresses.stream()
  86                 .map(a -> new SenderKeySharedEntry(resolver.resolveRecipient(a.getName()), a.getDeviceId()))
  87                 .collect(Collectors.toSet());
  88
  89         try (final var connection = database.getConnection()) {
  90             connection.setAutoCommit(false);
  91             markSenderKeysSharedWith(connection, distributionId, newEntries);
  92             connection.commit();
  93         } catch (SQLException e) {
  94             throw new RuntimeException("Failed update shared sender key store", e);
  95         }
  96     }
  97
  98     public void clearSenderKeySharedWith(final Collection<SignalProtocolAddress> addresses) {
  99         final var entriesToDelete = addresses.stream()
 100                 .map(a -> new SenderKeySharedEntry(resolver.resolveRecipient(a.getName()), a.getDeviceId()))
 101                 .collect(Collectors.toSet());
 102
 103         try (final var connection = database.getConnection()) {
 104             connection.setAutoCommit(false);
 105             final var sql = (
 106                     """
 107                     DELETE FROM %s AS s
 108                     WHERE recipient_id = ? AND device_id = ?
 109                     """
 110             ).formatted(TABLE_SENDER_KEY_SHARED);
 111             try (final var statement = connection.prepareStatement(sql)) {
 112                 for (final var entry : entriesToDelete) {
 113                     statement.setLong(1, entry.recipientId().id());
 114                     statement.setInt(2, entry.deviceId());
 115                     statement.executeUpdate();
 116                 }
 117             }
 118             connection.commit();
 119         } catch (SQLException e) {
 120             throw new RuntimeException("Failed update shared sender key store", e);
 121         }
 122     }
 123
 124     public void deleteAll() {
 125         try (final var connection = database.getConnection()) {
 126             final var sql = (
 127                     """
 128                     DELETE FROM %s AS s
 129                     """
 130             ).formatted(TABLE_SENDER_KEY_SHARED);
 131             try (final var statement = connection.prepareStatement(sql)) {
 132                 statement.executeUpdate();
 133             }
 134         } catch (SQLException e) {
 135             throw new RuntimeException("Failed update shared sender key store", e);
 136         }
 137     }
 138
 139     public void deleteAllFor(final RecipientId recipientId) {
 140         try (final var connection = database.getConnection()) {
 141             final var sql = (
 142                     """
 143                     DELETE FROM %s AS s
 144                     WHERE recipient_id = ?
 145                     """
 146             ).formatted(TABLE_SENDER_KEY_SHARED);
 147             try (final var statement = connection.prepareStatement(sql)) {
 148                 statement.setLong(1, recipientId.id());
 149                 statement.executeUpdate();
 150             }
 151         } catch (SQLException e) {
 152             throw new RuntimeException("Failed update shared sender key store", e);
 153         }
 154     }
 155
 156     public void deleteSharedWith(
 157             final RecipientId recipientId, final int deviceId, final DistributionId distributionId
 158     ) {
 159         try (final var connection = database.getConnection()) {
 160             final var sql = (
 161                     """
 162                     DELETE FROM %s AS s
 163                     WHERE recipient_id = ? AND device_id = ? AND distribution_id = ?
 164                     """
 165             ).formatted(TABLE_SENDER_KEY_SHARED);
 166             try (final var statement = connection.prepareStatement(sql)) {
 167                 statement.setLong(1, recipientId.id());
 168                 statement.setInt(2, deviceId);
 169                 statement.setBytes(3, UuidUtil.toByteArray(distributionId.asUuid()));
 170                 statement.executeUpdate();
 171             }
 172         } catch (SQLException e) {
 173             throw new RuntimeException("Failed update shared sender key store", e);
 174         }
 175     }
 176
 177     public void deleteAllFor(final DistributionId distributionId) {
 178         try (final var connection = database.getConnection()) {
 179             final var sql = (
 180                     """
 181                     DELETE FROM %s AS s
 182                     WHERE distribution_id = ?
 183                     """
 184             ).formatted(TABLE_SENDER_KEY_SHARED);
 185             try (final var statement = connection.prepareStatement(sql)) {
 186                 statement.setBytes(1, UuidUtil.toByteArray(distributionId.asUuid()));
 187                 statement.executeUpdate();
 188             }
 189         } catch (SQLException e) {
 190             throw new RuntimeException("Failed update shared sender key store", e);
 191         }
 192     }
 193
 194     public void mergeRecipients(RecipientId recipientId, RecipientId toBeMergedRecipientId) {
 195         try (final var connection = database.getConnection()) {
 196             final var sql = (
 197                     """
 198                     UPDATE OR REPLACE %s
 199                     SET recipient_id = ?
 200                     WHERE recipient_id = ?
 201                     """
 202             ).formatted(TABLE_SENDER_KEY_SHARED);
 203             try (final var statement = connection.prepareStatement(sql)) {
 204                 statement.setLong(1, recipientId.id());
 205                 statement.setLong(2, toBeMergedRecipientId.id());
 206                 statement.executeUpdate();
 207             }
 208         } catch (SQLException e) {
 209             throw new RuntimeException("Failed update shared sender key store", e);
 210         }
 211     }
 212
 213     void addLegacySenderKeysShared(final Map<DistributionId, Set<SenderKeySharedEntry>> sharedSenderKeys) {
 214         logger.debug("Migrating legacy sender keys shared to database");
 215         long start = System.nanoTime();
 216         try (final var connection = database.getConnection()) {
 217             connection.setAutoCommit(false);
 218             for (final var entry : sharedSenderKeys.entrySet()) {
 219                 markSenderKeysSharedWith(connection, entry.getKey(), entry.getValue());
 220             }
 221             connection.commit();
 222         } catch (SQLException e) {
 223             throw new RuntimeException("Failed update shared sender key store", e);
 224         }
 225         logger.debug("Complete sender keys shared migration took {}ms", (System.nanoTime() - start) / 1000000);
 226     }
 227
 228     private void markSenderKeysSharedWith(
 229             final Connection connection, final DistributionId distributionId, final Set<SenderKeySharedEntry> newEntries
 230     ) throws SQLException {
 231         final var sql = (
 232                 """
 233                 INSERT OR REPLACE INTO %s (recipient_id, device_id, distribution_id, timestamp)
 234                 VALUES (?, ?, ?, ?)
 235                 """
 236         ).formatted(TABLE_SENDER_KEY_SHARED);
 237         try (final var statement = connection.prepareStatement(sql)) {
 238             for (final var entry : newEntries) {
 239                 statement.setLong(1, entry.recipientId().id());
 240                 statement.setInt(2, entry.deviceId());
 241                 statement.setBytes(3, UuidUtil.toByteArray(distributionId.asUuid()));
 242                 statement.setLong(4, System.currentTimeMillis());
 243                 statement.executeUpdate();
 244             }
 245         }
 246     }
 247
 248     private SenderKeySharedEntry getSenderKeySharedEntryFromResultSet(ResultSet resultSet) throws SQLException {
 249         final var recipientId = resultSet.getLong("recipient_id");
 250         final var deviceId = resultSet.getInt("device_id");
 251         return new SenderKeySharedEntry(recipientIdCreator.create(recipientId), deviceId);
 252     }
 253
 254     record SenderKeySharedEntry(RecipientId recipientId, int deviceId) {}
 255 }