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