nmode's Git Repositories - signal-cli/blob - lib/src/main/java/org/asamk/signal/manager/storage/prekeys/SignedPreKeyStore.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.SignedPreKeyRecord;
  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 import java.util.List;
  19 import java.util.Objects;
  20
  21 public class SignedPreKeyStore implements org.signal.libsignal.protocol.state.SignedPreKeyStore {
  22
  23     private static final String TABLE_SIGNED_PRE_KEY = "signed_pre_key";
  24     private final static Logger logger = LoggerFactory.getLogger(SignedPreKeyStore.class);
  25
  26     private final Database database;
  27     private final int accountIdType;
  28
  29     public static void createSql(Connection connection) throws SQLException {
  30         // When modifying the CREATE statement here, also add a migration in AccountDatabase.java
  31         try (final var statement = connection.createStatement()) {
  32             statement.executeUpdate("""
  33                                     CREATE TABLE signed_pre_key (
  34                                       _id INTEGER PRIMARY KEY,
  35                                       account_id_type INTEGER NOT NULL,
  36                                       key_id INTEGER NOT NULL,
  37                                       public_key BLOB NOT NULL,
  38                                       private_key BLOB NOT NULL,
  39                                       signature BLOB NOT NULL,
  40                                       timestamp INTEGER DEFAULT 0,
  41                                       UNIQUE(account_id_type, key_id)
  42                                     ) STRICT;
  43                                     """);
  44         }
  45     }
  46
  47     public SignedPreKeyStore(final Database database, final ServiceIdType serviceIdType) {
  48         this.database = database;
  49         this.accountIdType = Utils.getAccountIdType(serviceIdType);
  50     }
  51
  52     @Override
  53     public SignedPreKeyRecord loadSignedPreKey(int signedPreKeyId) throws InvalidKeyIdException {
  54         final SignedPreKeyRecord signedPreKeyRecord = getSignedPreKey(signedPreKeyId);
  55         if (signedPreKeyRecord == null) {
  56             throw new InvalidKeyIdException("No such signed pre key record!");
  57         }
  58         return signedPreKeyRecord;
  59     }
  60
  61     @Override
  62     public List<SignedPreKeyRecord> loadSignedPreKeys() {
  63         final var sql = (
  64                 """
  65                 SELECT p.key_id, p.public_key, p.private_key, p.signature, p.timestamp
  66                 FROM %s p
  67                 WHERE p.account_id_type = ?
  68                 """
  69         ).formatted(TABLE_SIGNED_PRE_KEY);
  70         try (final var connection = database.getConnection()) {
  71             try (final var statement = connection.prepareStatement(sql)) {
  72                 statement.setInt(1, accountIdType);
  73                 return Utils.executeQueryForStream(statement, this::getSignedPreKeyRecordFromResultSet)
  74                         .filter(Objects::nonNull)
  75                         .toList();
  76             }
  77         } catch (SQLException e) {
  78             throw new RuntimeException("Failed read from signed_pre_key store", e);
  79         }
  80     }
  81
  82     @Override
  83     public void storeSignedPreKey(int signedPreKeyId, SignedPreKeyRecord record) {
  84         final var sql = (
  85                 """
  86                 INSERT INTO %s (account_id_type, key_id, public_key, private_key, signature, timestamp)
  87                 VALUES (?, ?, ?, ?, ?, ?)
  88                 """
  89         ).formatted(TABLE_SIGNED_PRE_KEY);
  90         try (final var connection = database.getConnection()) {
  91             try (final var statement = connection.prepareStatement(sql)) {
  92                 statement.setInt(1, accountIdType);
  93                 statement.setLong(2, signedPreKeyId);
  94                 final var keyPair = record.getKeyPair();
  95                 statement.setBytes(3, keyPair.getPublicKey().serialize());
  96                 statement.setBytes(4, keyPair.getPrivateKey().serialize());
  97                 statement.setBytes(5, record.getSignature());
  98                 statement.setLong(6, record.getTimestamp());
  99                 statement.executeUpdate();
 100             }
 101         } catch (SQLException e) {
 102             throw new RuntimeException("Failed update signed_pre_key store", e);
 103         }
 104     }
 105
 106     @Override
 107     public boolean containsSignedPreKey(int signedPreKeyId) {
 108         return getSignedPreKey(signedPreKeyId) != null;
 109     }
 110
 111     @Override
 112     public void removeSignedPreKey(int signedPreKeyId) {
 113         final var sql = (
 114                 """
 115                 DELETE FROM %s AS p
 116                 WHERE p.account_id_type = ? AND p.key_id = ?
 117                 """
 118         ).formatted(TABLE_SIGNED_PRE_KEY);
 119         try (final var connection = database.getConnection()) {
 120             try (final var statement = connection.prepareStatement(sql)) {
 121                 statement.setInt(1, accountIdType);
 122                 statement.setLong(2, signedPreKeyId);
 123                 statement.executeUpdate();
 124             }
 125         } catch (SQLException e) {
 126             throw new RuntimeException("Failed update signed_pre_key store", e);
 127         }
 128     }
 129
 130     public void removeAllSignedPreKeys() {
 131         final var sql = (
 132                 """
 133                 DELETE FROM %s AS p
 134                 WHERE p.account_id_type = ?
 135                 """
 136         ).formatted(TABLE_SIGNED_PRE_KEY);
 137         try (final var connection = database.getConnection()) {
 138             try (final var statement = connection.prepareStatement(sql)) {
 139                 statement.setInt(1, accountIdType);
 140                 statement.executeUpdate();
 141             }
 142         } catch (SQLException e) {
 143             throw new RuntimeException("Failed update signed_pre_key store", e);
 144         }
 145     }
 146
 147     void addLegacySignedPreKeys(final Collection<SignedPreKeyRecord> signedPreKeys) {
 148         logger.debug("Migrating legacy signedPreKeys to database");
 149         long start = System.nanoTime();
 150         final var sql = (
 151                 """
 152                 INSERT INTO %s (account_id_type, key_id, public_key, private_key, signature, timestamp)
 153                 VALUES (?, ?, ?, ?, ?, ?)
 154                 """
 155         ).formatted(TABLE_SIGNED_PRE_KEY);
 156         try (final var connection = database.getConnection()) {
 157             connection.setAutoCommit(false);
 158             final var deleteSql = "DELETE FROM %s AS p WHERE p.account_id_type = ?".formatted(TABLE_SIGNED_PRE_KEY);
 159             try (final var statement = connection.prepareStatement(deleteSql)) {
 160                 statement.setInt(1, accountIdType);
 161                 statement.executeUpdate();
 162             }
 163             try (final var statement = connection.prepareStatement(sql)) {
 164                 for (final var record : signedPreKeys) {
 165                     statement.setInt(1, accountIdType);
 166                     statement.setLong(2, record.getId());
 167                     final var keyPair = record.getKeyPair();
 168                     statement.setBytes(3, keyPair.getPublicKey().serialize());
 169                     statement.setBytes(4, keyPair.getPrivateKey().serialize());
 170                     statement.setBytes(5, record.getSignature());
 171                     statement.setLong(6, record.getTimestamp());
 172                     statement.executeUpdate();
 173                 }
 174             }
 175             connection.commit();
 176         } catch (SQLException e) {
 177             throw new RuntimeException("Failed update signedPreKey store", e);
 178         }
 179         logger.debug("Complete signedPreKeys migration took {}ms", (System.nanoTime() - start) / 1000000);
 180     }
 181
 182     private SignedPreKeyRecord getSignedPreKey(int signedPreKeyId) {
 183         final var sql = (
 184                 """
 185                 SELECT p.key_id, p.public_key, p.private_key, p.signature, p.timestamp
 186                 FROM %s p
 187                 WHERE p.account_id_type = ? AND p.key_id = ?
 188                 """
 189         ).formatted(TABLE_SIGNED_PRE_KEY);
 190         try (final var connection = database.getConnection()) {
 191             try (final var statement = connection.prepareStatement(sql)) {
 192                 statement.setInt(1, accountIdType);
 193                 statement.setLong(2, signedPreKeyId);
 194                 return Utils.executeQueryForOptional(statement, this::getSignedPreKeyRecordFromResultSet).orElse(null);
 195             }
 196         } catch (SQLException e) {
 197             throw new RuntimeException("Failed read from signed_pre_key store", e);
 198         }
 199     }
 200
 201     private SignedPreKeyRecord getSignedPreKeyRecordFromResultSet(ResultSet resultSet) throws SQLException {
 202         try {
 203             final var keyId = resultSet.getInt("key_id");
 204             final var publicKey = Curve.decodePoint(resultSet.getBytes("public_key"), 0);
 205             final var privateKey = Curve.decodePrivatePoint(resultSet.getBytes("private_key"));
 206             final var signature = resultSet.getBytes("signature");
 207             final var timestamp = resultSet.getLong("timestamp");
 208             return new SignedPreKeyRecord(keyId, timestamp, new ECKeyPair(publicKey, privateKey), signature);
 209         } catch (InvalidKeyException e) {
 210             return null;
 211         }
 212     }
 213 }