nmode's Git Repositories - signal-cli/blob - lib/src/main/java/org/asamk/signal/manager/storage/prekeys/PreKeyStore.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.InvalidKeyException;
   6 import org.signal.libsignal.protocol.InvalidKeyIdException;
   7 import org.signal.libsignal.protocol.ecc.Curve;
   8 import org.signal.libsignal.protocol.ecc.ECKeyPair;
   9 import org.signal.libsignal.protocol.state.PreKeyRecord;
  10 import org.slf4j.Logger;
  11 import org.slf4j.LoggerFactory;
  12 import org.whispersystems.signalservice.api.push.ServiceIdType;
  13
  14 import java.sql.Connection;
  15 import java.sql.ResultSet;
  16 import java.sql.SQLException;
  17 import java.util.Collection;
  18
  19 public class PreKeyStore implements org.signal.libsignal.protocol.state.PreKeyStore {
  20
  21     private static final String TABLE_PRE_KEY = "pre_key";
  22     private final static Logger logger = LoggerFactory.getLogger(PreKeyStore.class);
  23
  24     private final Database database;
  25     private final int accountIdType;
  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 pre_key (
  32                                       _id INTEGER PRIMARY KEY,
  33                                       account_id_type INTEGER NOT NULL,
  34                                       key_id INTEGER NOT NULL,
  35                                       public_key BLOB NOT NULL,
  36                                       private_key BLOB NOT NULL,
  37                                       UNIQUE(account_id_type, key_id)
  38                                     ) STRICT;
  39                                     """);
  40         }
  41     }
  42
  43     public PreKeyStore(final Database database, final ServiceIdType serviceIdType) {
  44         this.database = database;
  45         this.accountIdType = Utils.getAccountIdType(serviceIdType);
  46     }
  47
  48     @Override
  49     public PreKeyRecord loadPreKey(int preKeyId) throws InvalidKeyIdException {
  50         final var preKey = getPreKey(preKeyId);
  51         if (preKey == null) {
  52             throw new InvalidKeyIdException("No such signed pre key record!");
  53         }
  54         return preKey;
  55     }
  56
  57     @Override
  58     public void storePreKey(int preKeyId, PreKeyRecord record) {
  59         final var sql = (
  60                 """
  61                 INSERT INTO %s (account_id_type, key_id, public_key, private_key)
  62                 VALUES (?, ?, ?, ?)
  63                 """
  64         ).formatted(TABLE_PRE_KEY);
  65         try (final var connection = database.getConnection()) {
  66             try (final var statement = connection.prepareStatement(sql)) {
  67                 statement.setInt(1, accountIdType);
  68                 statement.setLong(2, preKeyId);
  69                 final var keyPair = record.getKeyPair();
  70                 statement.setBytes(3, keyPair.getPublicKey().serialize());
  71                 statement.setBytes(4, keyPair.getPrivateKey().serialize());
  72                 statement.executeUpdate();
  73             } catch (InvalidKeyException ignored) {
  74             }
  75         } catch (SQLException e) {
  76             throw new RuntimeException("Failed update pre_key store", e);
  77         }
  78     }
  79
  80     @Override
  81     public boolean containsPreKey(int preKeyId) {
  82         return getPreKey(preKeyId) != null;
  83     }
  84
  85     @Override
  86     public void removePreKey(int preKeyId) {
  87         final var sql = (
  88                 """
  89                 DELETE FROM %s AS p
  90                 WHERE p.account_id_type = ? AND p.key_id = ?
  91                 """
  92         ).formatted(TABLE_PRE_KEY);
  93         try (final var connection = database.getConnection()) {
  94             try (final var statement = connection.prepareStatement(sql)) {
  95                 statement.setInt(1, accountIdType);
  96                 statement.setLong(2, preKeyId);
  97                 statement.executeUpdate();
  98             }
  99         } catch (SQLException e) {
 100             throw new RuntimeException("Failed update pre_key store", e);
 101         }
 102     }
 103
 104     public void removeAllPreKeys() {
 105         final var sql = (
 106                 """
 107                 DELETE FROM %s AS p
 108                 WHERE p.account_id_type = ?
 109                 """
 110         ).formatted(TABLE_PRE_KEY);
 111         try (final var connection = database.getConnection()) {
 112             try (final var statement = connection.prepareStatement(sql)) {
 113                 statement.setInt(1, accountIdType);
 114                 statement.executeUpdate();
 115             }
 116         } catch (SQLException e) {
 117             throw new RuntimeException("Failed update pre_key store", e);
 118         }
 119     }
 120
 121     void addLegacyPreKeys(final Collection<PreKeyRecord> preKeys) {
 122         logger.debug("Migrating legacy preKeys to database");
 123         long start = System.nanoTime();
 124         final var sql = (
 125                 """
 126                 INSERT INTO %s (account_id_type, key_id, public_key, private_key)
 127                 VALUES (?, ?, ?, ?)
 128                 """
 129         ).formatted(TABLE_PRE_KEY);
 130         try (final var connection = database.getConnection()) {
 131             connection.setAutoCommit(false);
 132             final var deleteSql = "DELETE FROM %s AS p WHERE p.account_id_type = ?".formatted(TABLE_PRE_KEY);
 133             try (final var statement = connection.prepareStatement(deleteSql)) {
 134                 statement.setInt(1, accountIdType);
 135                 statement.executeUpdate();
 136             }
 137             try (final var statement = connection.prepareStatement(sql)) {
 138                 for (final var record : preKeys) {
 139                     statement.setInt(1, accountIdType);
 140                     statement.setLong(2, record.getId());
 141                     final var keyPair = record.getKeyPair();
 142                     statement.setBytes(3, keyPair.getPublicKey().serialize());
 143                     statement.setBytes(4, keyPair.getPrivateKey().serialize());
 144                     statement.executeUpdate();
 145                 }
 146             } catch (InvalidKeyException ignored) {
 147             }
 148             connection.commit();
 149         } catch (SQLException e) {
 150             throw new RuntimeException("Failed update preKey store", e);
 151         }
 152         logger.debug("Complete preKeys migration took {}ms", (System.nanoTime() - start) / 1000000);
 153     }
 154
 155     private PreKeyRecord getPreKey(int preKeyId) {
 156         final var sql = (
 157                 """
 158                 SELECT p.key_id, p.public_key, p.private_key
 159                 FROM %s p
 160                 WHERE p.account_id_type = ? AND p.key_id = ?
 161                 """
 162         ).formatted(TABLE_PRE_KEY);
 163         try (final var connection = database.getConnection()) {
 164             try (final var statement = connection.prepareStatement(sql)) {
 165                 statement.setInt(1, accountIdType);
 166                 statement.setLong(2, preKeyId);
 167                 return Utils.executeQueryForOptional(statement, this::getPreKeyRecordFromResultSet).orElse(null);
 168             }
 169         } catch (SQLException e) {
 170             throw new RuntimeException("Failed read from pre_key store", e);
 171         }
 172     }
 173
 174     private PreKeyRecord getPreKeyRecordFromResultSet(ResultSet resultSet) throws SQLException {
 175         try {
 176             final var keyId = resultSet.getInt("key_id");
 177             final var publicKey = Curve.decodePoint(resultSet.getBytes("public_key"), 0);
 178             final var privateKey = Curve.decodePrivatePoint(resultSet.getBytes("private_key"));
 179             return new PreKeyRecord(keyId, new ECKeyPair(publicKey, privateKey));
 180         } catch (InvalidKeyException e) {
 181             return null;
 182         }
 183     }
 184 }