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