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                 final var keyPair = record.getKeyPair();
  97                 statement.setBytes(3, keyPair.getPublicKey().serialize());
  98                 statement.setBytes(4, keyPair.getPrivateKey().serialize());
  99                 statement.setBytes(5, record.getSignature());
 100                 statement.setLong(6, record.getTimestamp());
 101                 statement.executeUpdate();
 102             }
 103         } catch (SQLException e) {
 104             throw new RuntimeException("Failed update signed_pre_key store", e);
 105         }
 106     }
 107
 108     @Override
 109     public boolean containsSignedPreKey(int signedPreKeyId) {
 110         return getSignedPreKey(signedPreKeyId) != null;
 111     }
 112
 113     @Override
 114     public void removeSignedPreKey(int signedPreKeyId) {
 115         final var sql = (
 116                 """
 117                 DELETE FROM %s AS p
 118                 WHERE p.account_id_type = ? AND p.key_id = ?
 119                 """
 120         ).formatted(TABLE_SIGNED_PRE_KEY);
 121         try (final var connection = database.getConnection()) {
 122             try (final var statement = connection.prepareStatement(sql)) {
 123                 statement.setInt(1, accountIdType);
 124                 statement.setInt(2, signedPreKeyId);
 125                 statement.executeUpdate();
 126             }
 127         } catch (SQLException e) {
 128             throw new RuntimeException("Failed update signed_pre_key store", e);
 129         }
 130     }
 131
 132     public void removeAllSignedPreKeys() {
 133         final var sql = (
 134                 """
 135                 DELETE FROM %s AS p
 136                 WHERE p.account_id_type = ?
 137                 """
 138         ).formatted(TABLE_SIGNED_PRE_KEY);
 139         try (final var connection = database.getConnection()) {
 140             try (final var statement = connection.prepareStatement(sql)) {
 141                 statement.setInt(1, accountIdType);
 142                 statement.executeUpdate();
 143             }
 144         } catch (SQLException e) {
 145             throw new RuntimeException("Failed update signed_pre_key store", e);
 146         }
 147     }
 148
 149     public void removeOldSignedPreKeys(int activePreKeyId) {
 150         final var sql = (
 151                 """
 152                 DELETE FROM %s AS p
 153                 WHERE p._id IN (
 154                     SELECT p._id
 155                     FROM %s AS p
 156                     WHERE p.account_id_type = ?
 157                         AND p.key_id != ?
 158                         AND p.timestamp < ?
 159                     ORDER BY p.timestamp DESC
 160                     LIMIT -1 OFFSET 1
 161                 )
 162                 """
 163         ).formatted(TABLE_SIGNED_PRE_KEY, TABLE_SIGNED_PRE_KEY);
 164         try (final var connection = database.getConnection()) {
 165             try (final var statement = connection.prepareStatement(sql)) {
 166                 statement.setInt(1, accountIdType);
 167                 statement.setInt(2, activePreKeyId);
 168                 statement.setLong(3, System.currentTimeMillis() - PREKEY_ARCHIVE_AGE);
 169                 statement.executeUpdate();
 170             }
 171         } catch (SQLException e) {
 172             throw new RuntimeException("Failed update signed_pre_key store", e);
 173         }
 174     }
 175
 176     void addLegacySignedPreKeys(final Collection<SignedPreKeyRecord> signedPreKeys) {
 177         logger.debug("Migrating legacy signedPreKeys to database");
 178         long start = System.nanoTime();
 179         final var sql = (
 180                 """
 181                 INSERT INTO %s (account_id_type, key_id, public_key, private_key, signature, timestamp)
 182                 VALUES (?, ?, ?, ?, ?, ?)
 183                 """
 184         ).formatted(TABLE_SIGNED_PRE_KEY);
 185         try (final var connection = database.getConnection()) {
 186             connection.setAutoCommit(false);
 187             final var deleteSql = "DELETE FROM %s AS p WHERE p.account_id_type = ?".formatted(TABLE_SIGNED_PRE_KEY);
 188             try (final var statement = connection.prepareStatement(deleteSql)) {
 189                 statement.setInt(1, accountIdType);
 190                 statement.executeUpdate();
 191             }
 192             try (final var statement = connection.prepareStatement(sql)) {
 193                 for (final var record : signedPreKeys) {
 194                     statement.setInt(1, accountIdType);
 195                     statement.setInt(2, record.getId());
 196                     final var keyPair = record.getKeyPair();
 197                     statement.setBytes(3, keyPair.getPublicKey().serialize());
 198                     statement.setBytes(4, keyPair.getPrivateKey().serialize());
 199                     statement.setBytes(5, record.getSignature());
 200                     statement.setLong(6, record.getTimestamp());
 201                     statement.executeUpdate();
 202                 }
 203             }
 204             connection.commit();
 205         } catch (SQLException e) {
 206             throw new RuntimeException("Failed update signedPreKey store", e);
 207         }
 208         logger.debug("Complete signedPreKeys migration took {}ms", (System.nanoTime() - start) / 1000000);
 209     }
 210
 211     private SignedPreKeyRecord getSignedPreKey(int signedPreKeyId) {
 212         final var sql = (
 213                 """
 214                 SELECT p.key_id, p.public_key, p.private_key, p.signature, p.timestamp
 215                 FROM %s p
 216                 WHERE p.account_id_type = ? AND p.key_id = ?
 217                 """
 218         ).formatted(TABLE_SIGNED_PRE_KEY);
 219         try (final var connection = database.getConnection()) {
 220             try (final var statement = connection.prepareStatement(sql)) {
 221                 statement.setInt(1, accountIdType);
 222                 statement.setInt(2, signedPreKeyId);
 223                 return Utils.executeQueryForOptional(statement, this::getSignedPreKeyRecordFromResultSet).orElse(null);
 224             }
 225         } catch (SQLException e) {
 226             throw new RuntimeException("Failed read from signed_pre_key store", e);
 227         }
 228     }
 229
 230     private SignedPreKeyRecord getSignedPreKeyRecordFromResultSet(ResultSet resultSet) throws SQLException {
 231         try {
 232             final var keyId = resultSet.getInt("key_id");
 233             final var publicKey = Curve.decodePoint(resultSet.getBytes("public_key"), 0);
 234             final var privateKey = Curve.decodePrivatePoint(resultSet.getBytes("private_key"));
 235             final var signature = resultSet.getBytes("signature");
 236             final var timestamp = resultSet.getLong("timestamp");
 237             return new SignedPreKeyRecord(keyId, timestamp, new ECKeyPair(publicKey, privateKey), signature);
 238         } catch (InvalidKeyException e) {
 239             return null;
 240         }
 241     }
 242 }