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 public class KyberPreKeyStore implements SignalServiceKyberPreKeyStore {
  19
  20     private static final String TABLE_KYBER_PRE_KEY = "kyber_pre_key";
  21     private final static Logger logger = LoggerFactory.getLogger(KyberPreKeyStore.class);
  22
  23     private final Database database;
  24     private final int accountIdType;
  25
  26     public static void createSql(Connection connection) throws SQLException {
  27         // When modifying the CREATE statement here, also add a migration in AccountDatabase.java
  28         try (final var statement = connection.createStatement()) {
  29             statement.executeUpdate("""
  30                                     CREATE TABLE kyber_pre_key (
  31                                       _id INTEGER PRIMARY KEY,
  32                                       account_id_type INTEGER NOT NULL,
  33                                       key_id INTEGER NOT NULL,
  34                                       serialized BLOB NOT NULL,
  35                                       is_last_resort INTEGER NOT NULL,
  36                                       UNIQUE(account_id_type, key_id)
  37                                     ) STRICT;
  38                                     """);
  39         }
  40     }
  41
  42     public KyberPreKeyStore(final Database database, final ServiceIdType serviceIdType) {
  43         this.database = database;
  44         this.accountIdType = Utils.getAccountIdType(serviceIdType);
  45     }
  46
  47     @Override
  48     public KyberPreKeyRecord loadKyberPreKey(final int keyId) throws InvalidKeyIdException {
  49         final var kyberPreKey = getPreKey(keyId);
  50         if (kyberPreKey == null) {
  51             throw new InvalidKeyIdException("No such kyber pre key record: " + keyId);
  52         }
  53         return kyberPreKey;
  54     }
  55
  56     @Override
  57     public List<KyberPreKeyRecord> loadKyberPreKeys() {
  58         final var sql = (
  59                 """
  60                 SELECT p.serialized
  61                 FROM %s p
  62                 WHERE p.account_id_type = ?
  63                 """
  64         ).formatted(TABLE_KYBER_PRE_KEY);
  65         try (final var connection = database.getConnection()) {
  66             try (final var statement = connection.prepareStatement(sql)) {
  67                 statement.setInt(1, accountIdType);
  68                 return Utils.executeQueryForStream(statement, this::getKyberPreKeyRecordFromResultSet).toList();
  69             }
  70         } catch (SQLException e) {
  71             throw new RuntimeException("Failed read from kyber_pre_key store", e);
  72         }
  73     }
  74
  75     @Override
  76     public List<KyberPreKeyRecord> loadLastResortKyberPreKeys() {
  77         final var sql = (
  78                 """
  79                 SELECT p.serialized
  80                 FROM %s p
  81                 WHERE p.account_id_type = ? AND p.is_last_resort = TRUE
  82                 """
  83         ).formatted(TABLE_KYBER_PRE_KEY);
  84         try (final var connection = database.getConnection()) {
  85             try (final var statement = connection.prepareStatement(sql)) {
  86                 statement.setInt(1, accountIdType);
  87                 return Utils.executeQueryForStream(statement, this::getKyberPreKeyRecordFromResultSet).toList();
  88             }
  89         } catch (SQLException e) {
  90             throw new RuntimeException("Failed read from kyber_pre_key store", e);
  91         }
  92     }
  93
  94     @Override
  95     public void storeLastResortKyberPreKey(final int keyId, final KyberPreKeyRecord record) {
  96         storeKyberPreKey(keyId, record, true);
  97     }
  98
  99     @Override
 100     public void storeKyberPreKey(final int keyId, final KyberPreKeyRecord record) {
 101         storeKyberPreKey(keyId, record, false);
 102     }
 103
 104     public void storeKyberPreKey(final int keyId, final KyberPreKeyRecord record, final boolean isLastResort) {
 105         final var sql = (
 106                 """
 107                 INSERT INTO %s (account_id_type, key_id, serialized, is_last_resort)
 108                 VALUES (?, ?, ?, ?)
 109                 """
 110         ).formatted(TABLE_KYBER_PRE_KEY);
 111         try (final var connection = database.getConnection()) {
 112             try (final var statement = connection.prepareStatement(sql)) {
 113                 statement.setInt(1, accountIdType);
 114                 statement.setInt(2, keyId);
 115                 statement.setBytes(3, record.serialize());
 116                 statement.setBoolean(4, isLastResort);
 117                 statement.executeUpdate();
 118             }
 119         } catch (SQLException e) {
 120             throw new RuntimeException("Failed update kyber_pre_key store", e);
 121         }
 122     }
 123
 124     @Override
 125     public boolean containsKyberPreKey(final int keyId) {
 126         return getPreKey(keyId) != null;
 127     }
 128
 129     @Override
 130     public void markKyberPreKeyUsed(final int keyId) {
 131         final var sql = (
 132                 """
 133                 DELETE FROM %s AS p
 134                 WHERE p.account_id_type = ? AND p.key_id = ? AND p.is_last_resort = FALSE
 135                 """
 136         ).formatted(TABLE_KYBER_PRE_KEY);
 137         try (final var connection = database.getConnection()) {
 138             try (final var statement = connection.prepareStatement(sql)) {
 139                 statement.setInt(1, accountIdType);
 140                 statement.setInt(2, keyId);
 141                 statement.executeUpdate();
 142             }
 143         } catch (SQLException e) {
 144             throw new RuntimeException("Failed update kyber_pre_key store", e);
 145         }
 146     }
 147
 148     @Override
 149     public void removeKyberPreKey(final int keyId) {
 150         final var sql = (
 151                 """
 152                 DELETE FROM %s AS p
 153                 WHERE p.account_id_type = ? AND p.key_id = ?
 154                 """
 155         ).formatted(TABLE_KYBER_PRE_KEY);
 156         try (final var connection = database.getConnection()) {
 157             try (final var statement = connection.prepareStatement(sql)) {
 158                 statement.setInt(1, accountIdType);
 159                 statement.setInt(2, keyId);
 160                 statement.executeUpdate();
 161             }
 162         } catch (SQLException e) {
 163             throw new RuntimeException("Failed update kyber_pre_key store", e);
 164         }
 165     }
 166
 167     public void removeAllKyberPreKeys() {
 168         final var sql = (
 169                 """
 170                 DELETE FROM %s AS p
 171                 WHERE p.account_id_type = ?
 172                 """
 173         ).formatted(TABLE_KYBER_PRE_KEY);
 174         try (final var connection = database.getConnection()) {
 175             try (final var statement = connection.prepareStatement(sql)) {
 176                 statement.setInt(1, accountIdType);
 177                 statement.executeUpdate();
 178             }
 179         } catch (SQLException e) {
 180             throw new RuntimeException("Failed update kyber_pre_key store", e);
 181         }
 182     }
 183
 184     private KyberPreKeyRecord getPreKey(int keyId) {
 185         final var sql = (
 186                 """
 187                 SELECT p.serialized
 188                 FROM %s p
 189                 WHERE p.account_id_type = ? AND p.key_id = ?
 190                 """
 191         ).formatted(TABLE_KYBER_PRE_KEY);
 192         try (final var connection = database.getConnection()) {
 193             try (final var statement = connection.prepareStatement(sql)) {
 194                 statement.setInt(1, accountIdType);
 195                 statement.setInt(2, keyId);
 196                 return Utils.executeQueryForOptional(statement, this::getKyberPreKeyRecordFromResultSet).orElse(null);
 197             }
 198         } catch (SQLException e) {
 199             throw new RuntimeException("Failed read from kyber_pre_key store", e);
 200         }
 201     }
 202
 203     private KyberPreKeyRecord getKyberPreKeyRecordFromResultSet(ResultSet resultSet) throws SQLException {
 204         try {
 205             final var serialized = resultSet.getBytes("serialized");
 206             return new KyberPreKeyRecord(serialized);
 207         } catch (InvalidMessageException e) {
 208             return null;
 209         }
 210     }
 211
 212     @Override
 213     public void deleteAllStaleOneTimeKyberPreKeys(final long threshold, final int minCount) {
 214         //TODO
 215     }
 216
 217     @Override
 218     public void markAllOneTimeKyberPreKeysStaleIfNecessary(final long staleTime) {
 219         //TODO
 220     }
 221 }