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 static final Logger logger = LoggerFactory.getLogger(SenderKeySharedStore.class);
  23     private static final 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                                       address TEXT NOT NULL,
  34                                       device_id INTEGER NOT NULL,
  35                                       distribution_id BLOB NOT NULL,
  36                                       timestamp INTEGER NOT NULL,
  37                                       UNIQUE(address, 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.address, 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 -> new SignalProtocolAddress(k.address, 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(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                 .map(a -> new SenderKeySharedEntry(a.getName(), a.getDeviceId()))
  86                 .collect(Collectors.toSet());
  87
  88         try (final var connection = database.getConnection()) {
  89             connection.setAutoCommit(false);
  90             final var sql = (
  91                     """
  92                     DELETE FROM %s AS s
  93                     WHERE address = ? AND device_id = ?
  94                     """
  95             ).formatted(TABLE_SENDER_KEY_SHARED);
  96             try (final var statement = connection.prepareStatement(sql)) {
  97                 for (final var entry : entriesToDelete) {
  98                     statement.setString(1, entry.address());
  99                     statement.setInt(2, entry.deviceId());
 100                     statement.executeUpdate();
 101                 }
 102             }
 103             connection.commit();
 104         } catch (SQLException e) {
 105             throw new RuntimeException("Failed update shared sender key store", e);
 106         }
 107     }
 108
 109     public void deleteAll() {
 110         try (final var connection = database.getConnection()) {
 111             final var sql = (
 112                     """
 113                     DELETE FROM %s AS s
 114                     """
 115             ).formatted(TABLE_SENDER_KEY_SHARED);
 116             try (final var statement = connection.prepareStatement(sql)) {
 117                 statement.executeUpdate();
 118             }
 119         } catch (SQLException e) {
 120             throw new RuntimeException("Failed update shared sender key store", e);
 121         }
 122     }
 123
 124     public void deleteAllFor(final ServiceId serviceId) {
 125         try (final var connection = database.getConnection()) {
 126             final var sql = (
 127                     """
 128                     DELETE FROM %s AS s
 129                     WHERE address = ?
 130                     """
 131             ).formatted(TABLE_SENDER_KEY_SHARED);
 132             try (final var statement = connection.prepareStatement(sql)) {
 133                 statement.setString(1, serviceId.toString());
 134                 statement.executeUpdate();
 135             }
 136         } catch (SQLException e) {
 137             throw new RuntimeException("Failed update shared sender key store", e);
 138         }
 139     }
 140
 141     public void deleteSharedWith(
 142             final ServiceId serviceId, final int deviceId, final DistributionId distributionId
 143     ) {
 144         try (final var connection = database.getConnection()) {
 145             final var sql = (
 146                     """
 147                     DELETE FROM %s AS s
 148                     WHERE address = ? AND device_id = ? AND distribution_id = ?
 149                     """
 150             ).formatted(TABLE_SENDER_KEY_SHARED);
 151             try (final var statement = connection.prepareStatement(sql)) {
 152                 statement.setString(1, serviceId.toString());
 153                 statement.setInt(2, deviceId);
 154                 statement.setBytes(3, UuidUtil.toByteArray(distributionId.asUuid()));
 155                 statement.executeUpdate();
 156             }
 157         } catch (SQLException e) {
 158             throw new RuntimeException("Failed update shared sender key store", e);
 159         }
 160     }
 161
 162     public void deleteAllFor(final DistributionId distributionId) {
 163         try (final var connection = database.getConnection()) {
 164             final var sql = (
 165                     """
 166                     DELETE FROM %s AS s
 167                     WHERE distribution_id = ?
 168                     """
 169             ).formatted(TABLE_SENDER_KEY_SHARED);
 170             try (final var statement = connection.prepareStatement(sql)) {
 171                 statement.setBytes(1, UuidUtil.toByteArray(distributionId.asUuid()));
 172                 statement.executeUpdate();
 173             }
 174         } catch (SQLException e) {
 175             throw new RuntimeException("Failed update shared sender key store", e);
 176         }
 177     }
 178
 179     void addLegacySenderKeysShared(final Map<DistributionId, Set<SenderKeySharedEntry>> sharedSenderKeys) {
 180         logger.debug("Migrating legacy sender keys shared to database");
 181         long start = System.nanoTime();
 182         try (final var connection = database.getConnection()) {
 183             connection.setAutoCommit(false);
 184             for (final var entry : sharedSenderKeys.entrySet()) {
 185                 markSenderKeysSharedWith(connection, entry.getKey(), entry.getValue());
 186             }
 187             connection.commit();
 188         } catch (SQLException e) {
 189             throw new RuntimeException("Failed update shared sender key store", e);
 190         }
 191         logger.debug("Complete sender keys shared migration took {}ms", (System.nanoTime() - start) / 1000000);
 192     }
 193
 194     private void markSenderKeysSharedWith(
 195             final Connection connection, final DistributionId distributionId, final Set<SenderKeySharedEntry> newEntries
 196     ) throws SQLException {
 197         final var sql = (
 198                 """
 199                 INSERT OR REPLACE INTO %s (address, device_id, distribution_id, timestamp)
 200                 VALUES (?, ?, ?, ?)
 201                 """
 202         ).formatted(TABLE_SENDER_KEY_SHARED);
 203         try (final var statement = connection.prepareStatement(sql)) {
 204             for (final var entry : newEntries) {
 205                 statement.setString(1, entry.toString());
 206                 statement.setInt(2, entry.deviceId());
 207                 statement.setBytes(3, UuidUtil.toByteArray(distributionId.asUuid()));
 208                 statement.setLong(4, System.currentTimeMillis());
 209                 statement.executeUpdate();
 210             }
 211         }
 212     }
 213
 214     private SenderKeySharedEntry getSenderKeySharedEntryFromResultSet(ResultSet resultSet) throws SQLException {
 215         final var address = resultSet.getString("address");
 216         final var deviceId = resultSet.getInt("device_id");
 217         return new SenderKeySharedEntry(address, deviceId);
 218     }
 219
 220     record SenderKeySharedEntry(String address, int deviceId) {}
 221 }