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

   1 package org.asamk.signal.manager.storage.prekeys;
   2
   3 import org.asamk.signal.manager.storage.Database;
   4 import org.asamk.signal.manager.storage.Utils;
   5 import org.signal.libsignal.protocol.InvalidKeyIdException;
   6 import org.signal.libsignal.protocol.InvalidMessageException;
   7 import org.signal.libsignal.protocol.state.KyberPreKeyRecord;
   8 import org.slf4j.Logger;
   9 import org.slf4j.LoggerFactory;
  10 import org.whispersystems.signalservice.api.SignalServiceKyberPreKeyStore;
  11 import org.whispersystems.signalservice.api.push.ServiceIdType;
  12
  13 import java.sql.Connection;
  14 import java.sql.ResultSet;
  15 import java.sql.SQLException;
  16 import java.util.List;
  17
  18 import static org.asamk.signal.manager.config.ServiceConfig.PREKEY_ARCHIVE_AGE;
  19
  20 public class KyberPreKeyStore implements SignalServiceKyberPreKeyStore {
  21
  22     private static final String TABLE_KYBER_PRE_KEY = "kyber_pre_key";
  23     private final static Logger logger = LoggerFactory.getLogger(KyberPreKeyStore.class);
  24
  25     private final Database database;
  26     private final int accountIdType;
  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 kyber_pre_key (
  33                                       _id INTEGER PRIMARY KEY,
  34                                       account_id_type INTEGER NOT NULL,
  35                                       key_id INTEGER NOT NULL,
  36                                       serialized BLOB NOT NULL,
  37                                       is_last_resort INTEGER NOT NULL,
  38                                       stale_timestamp INTEGER,
  39                                       timestamp INTEGER DEFAULT 0,
  40                                       UNIQUE(account_id_type, key_id)
  41                                     ) STRICT;
  42                                     """);
  43         }
  44     }
  45
  46     public KyberPreKeyStore(final Database database, final ServiceIdType serviceIdType) {
  47         this.database = database;
  48         this.accountIdType = Utils.getAccountIdType(serviceIdType);
  49     }
  50
  51     @Override
  52     public KyberPreKeyRecord loadKyberPreKey(final int keyId) throws InvalidKeyIdException {
  53         final var kyberPreKey = getPreKey(keyId);
  54         if (kyberPreKey == null) {
  55             throw new InvalidKeyIdException("No such kyber pre key record: " + keyId);
  56         }
  57         return kyberPreKey;
  58     }
  59
  60     @Override
  61     public List<KyberPreKeyRecord> loadKyberPreKeys() {
  62         final var sql = (
  63                 """
  64                 SELECT p.serialized
  65                 FROM %s p
  66                 WHERE p.account_id_type = ?
  67                 """
  68         ).formatted(TABLE_KYBER_PRE_KEY);
  69         try (final var connection = database.getConnection()) {
  70             try (final var statement = connection.prepareStatement(sql)) {
  71                 statement.setInt(1, accountIdType);
  72                 return Utils.executeQueryForStream(statement, this::getKyberPreKeyRecordFromResultSet).toList();
  73             }
  74         } catch (SQLException e) {
  75             throw new RuntimeException("Failed read from kyber_pre_key store", e);
  76         }
  77     }
  78
  79     @Override
  80     public List<KyberPreKeyRecord> loadLastResortKyberPreKeys() {
  81         final var sql = (
  82                 """
  83                 SELECT p.serialized
  84                 FROM %s p
  85                 WHERE p.account_id_type = ? AND p.is_last_resort = TRUE
  86                 """
  87         ).formatted(TABLE_KYBER_PRE_KEY);
  88         try (final var connection = database.getConnection()) {
  89             try (final var statement = connection.prepareStatement(sql)) {
  90                 statement.setInt(1, accountIdType);
  91                 return Utils.executeQueryForStream(statement, this::getKyberPreKeyRecordFromResultSet).toList();
  92             }
  93         } catch (SQLException e) {
  94             throw new RuntimeException("Failed read from kyber_pre_key store", e);
  95         }
  96     }
  97
  98     @Override
  99     public void storeLastResortKyberPreKey(final int keyId, final KyberPreKeyRecord record) {
 100         storeKyberPreKey(keyId, record, true);
 101     }
 102
 103     @Override
 104     public void storeKyberPreKey(final int keyId, final KyberPreKeyRecord record) {
 105         storeKyberPreKey(keyId, record, false);
 106     }
 107
 108     public void storeKyberPreKey(final int keyId, final KyberPreKeyRecord record, final boolean isLastResort) {
 109         final var sql = (
 110                 """
 111                 INSERT INTO %s (account_id_type, key_id, serialized, is_last_resort, timestamp)
 112                 VALUES (?, ?, ?, ?, ?)
 113                 """
 114         ).formatted(TABLE_KYBER_PRE_KEY);
 115         try (final var connection = database.getConnection()) {
 116             try (final var statement = connection.prepareStatement(sql)) {
 117                 statement.setInt(1, accountIdType);
 118                 statement.setInt(2, keyId);
 119                 statement.setBytes(3, record.serialize());
 120                 statement.setBoolean(4, isLastResort);
 121                 statement.setLong(5, record.getTimestamp());
 122                 statement.executeUpdate();
 123             }
 124         } catch (SQLException e) {
 125             throw new RuntimeException("Failed update kyber_pre_key store", e);
 126         }
 127     }
 128
 129     @Override
 130     public boolean containsKyberPreKey(final int keyId) {
 131         return getPreKey(keyId) != null;
 132     }
 133
 134     @Override
 135     public void markKyberPreKeyUsed(final int keyId) {
 136         final var sql = (
 137                 """
 138                 DELETE FROM %s AS p
 139                 WHERE p.account_id_type = ? AND p.key_id = ? AND p.is_last_resort = FALSE
 140                 """
 141         ).formatted(TABLE_KYBER_PRE_KEY);
 142         try (final var connection = database.getConnection()) {
 143             try (final var statement = connection.prepareStatement(sql)) {
 144                 statement.setInt(1, accountIdType);
 145                 statement.setInt(2, keyId);
 146                 statement.executeUpdate();
 147             }
 148         } catch (SQLException e) {
 149             throw new RuntimeException("Failed update kyber_pre_key store", e);
 150         }
 151     }
 152
 153     @Override
 154     public void removeKyberPreKey(final int keyId) {
 155         final var sql = (
 156                 """
 157                 DELETE FROM %s AS p
 158                 WHERE p.account_id_type = ? AND p.key_id = ?
 159                 """
 160         ).formatted(TABLE_KYBER_PRE_KEY);
 161         try (final var connection = database.getConnection()) {
 162             try (final var statement = connection.prepareStatement(sql)) {
 163                 statement.setInt(1, accountIdType);
 164                 statement.setInt(2, keyId);
 165                 statement.executeUpdate();
 166             }
 167         } catch (SQLException e) {
 168             throw new RuntimeException("Failed update kyber_pre_key store", e);
 169         }
 170     }
 171
 172     public void removeAllKyberPreKeys() {
 173         final var sql = (
 174                 """
 175                 DELETE FROM %s AS p
 176                 WHERE p.account_id_type = ?
 177                 """
 178         ).formatted(TABLE_KYBER_PRE_KEY);
 179         try (final var connection = database.getConnection()) {
 180             try (final var statement = connection.prepareStatement(sql)) {
 181                 statement.setInt(1, accountIdType);
 182                 statement.executeUpdate();
 183             }
 184         } catch (SQLException e) {
 185             throw new RuntimeException("Failed update kyber_pre_key store", e);
 186         }
 187     }
 188
 189     private KyberPreKeyRecord getPreKey(int keyId) {
 190         final var sql = (
 191                 """
 192                 SELECT p.serialized
 193                 FROM %s p
 194                 WHERE p.account_id_type = ? AND p.key_id = ?
 195                 """
 196         ).formatted(TABLE_KYBER_PRE_KEY);
 197         try (final var connection = database.getConnection()) {
 198             try (final var statement = connection.prepareStatement(sql)) {
 199                 statement.setInt(1, accountIdType);
 200                 statement.setInt(2, keyId);
 201                 return Utils.executeQueryForOptional(statement, this::getKyberPreKeyRecordFromResultSet).orElse(null);
 202             }
 203         } catch (SQLException e) {
 204             throw new RuntimeException("Failed read from kyber_pre_key store", e);
 205         }
 206     }
 207
 208     private KyberPreKeyRecord getKyberPreKeyRecordFromResultSet(ResultSet resultSet) throws SQLException {
 209         try {
 210             final var serialized = resultSet.getBytes("serialized");
 211             return new KyberPreKeyRecord(serialized);
 212         } catch (InvalidMessageException e) {
 213             return null;
 214         }
 215     }
 216
 217     public void removeOldLastResortKyberPreKeys(int activeLastResortKyberPreKeyId) {
 218         final var sql = (
 219                 """
 220                 DELETE FROM %s AS p
 221                 WHERE p._id IN (
 222                     SELECT p._id
 223                     FROM %s AS p
 224                     WHERE p.account_id_type = ?
 225                         AND p.is_last_resort = TRUE
 226                         AND p.key_id != ?
 227                         AND p.timestamp < ?
 228                     ORDER BY p.timestamp DESC
 229                     LIMIT -1 OFFSET 1
 230                 )
 231                 """
 232         ).formatted(TABLE_KYBER_PRE_KEY, TABLE_KYBER_PRE_KEY);
 233         try (final var connection = database.getConnection()) {
 234             try (final var statement = connection.prepareStatement(sql)) {
 235                 statement.setInt(1, accountIdType);
 236                 statement.setInt(2, activeLastResortKyberPreKeyId);
 237                 statement.setLong(3, System.currentTimeMillis() - PREKEY_ARCHIVE_AGE);
 238                 statement.executeUpdate();
 239             }
 240         } catch (SQLException e) {
 241             throw new RuntimeException("Failed update kyber_pre_key store", e);
 242         }
 243     }
 244
 245     @Override
 246     public void deleteAllStaleOneTimeKyberPreKeys(final long threshold, final int minCount) {
 247         final var sql = (
 248                 """
 249                 DELETE FROM %s AS p
 250                 WHERE p.account_id_type = ?1
 251                     AND p.stale_timestamp < ?2
 252                     AND p.is_last_resort = FALSE
 253                     AND p._id NOT IN (
 254                         SELECT _id
 255                         FROM %s p2
 256                         WHERE p2.account_id_type = ?1
 257                         ORDER BY
 258                           CASE WHEN p2.stale_timestamp IS NULL THEN 1 ELSE 0 END DESC,
 259                           p2.stale_timestamp DESC,
 260                           p2._id DESC
 261                         LIMIT ?3
 262                     )
 263                 """
 264         ).formatted(TABLE_KYBER_PRE_KEY, TABLE_KYBER_PRE_KEY);
 265         try (final var connection = database.getConnection()) {
 266             try (final var statement = connection.prepareStatement(sql)) {
 267                 statement.setInt(1, accountIdType);
 268                 statement.setLong(2, threshold);
 269                 statement.setInt(3, minCount);
 270                 final var rowCount = statement.executeUpdate();
 271                 if (rowCount > 0) {
 272                     logger.debug("Deleted {} stale one time kyber pre keys", rowCount);
 273                 }
 274             }
 275         } catch (SQLException e) {
 276             throw new RuntimeException("Failed update kyber_pre_key store", e);
 277         }
 278     }
 279
 280     @Override
 281     public void markAllOneTimeKyberPreKeysStaleIfNecessary(final long staleTime) {
 282         final var sql = (
 283                 """
 284                 UPDATE %s
 285                 SET stale_timestamp = ?
 286                 WHERE account_id_type = ? AND stale_timestamp IS NULL AND is_last_resort = FALSE
 287                 """
 288         ).formatted(TABLE_KYBER_PRE_KEY);
 289         try (final var connection = database.getConnection()) {
 290             try (final var statement = connection.prepareStatement(sql)) {
 291                 statement.setLong(1, staleTime);
 292                 statement.setInt(2, accountIdType);
 293                 statement.executeUpdate();
 294             }
 295         } catch (SQLException e) {
 296             throw new RuntimeException("Failed update kyber_pre_key store", e);
 297         }
 298     }
 299 }