nmode's Git Repositories - signal-cli/blob - lib/src/main/java/org/asamk/signal/manager/storage/identities/IdentityKeyStore.java

   1 package org.asamk.signal.manager.storage.identities;
   2
   3 import org.asamk.signal.manager.api.TrustLevel;
   4 import org.asamk.signal.manager.storage.Database;
   5 import org.asamk.signal.manager.storage.Utils;
   6 import org.asamk.signal.manager.storage.recipients.RecipientId;
   7 import org.asamk.signal.manager.storage.recipients.RecipientIdCreator;
   8 import org.signal.libsignal.protocol.IdentityKey;
   9 import org.signal.libsignal.protocol.InvalidKeyException;
  10 import org.signal.libsignal.protocol.state.IdentityKeyStore.Direction;
  11 import org.slf4j.Logger;
  12 import org.slf4j.LoggerFactory;
  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
  20 import io.reactivex.rxjava3.core.Observable;
  21 import io.reactivex.rxjava3.subjects.PublishSubject;
  22
  23 public class IdentityKeyStore {
  24
  25     private final static Logger logger = LoggerFactory.getLogger(IdentityKeyStore.class);
  26     private static final String TABLE_IDENTITY = "identity";
  27     private final Database database;
  28     private final RecipientIdCreator recipientIdCreator;
  29     private final TrustNewIdentity trustNewIdentity;
  30     private final PublishSubject<RecipientId> identityChanges = PublishSubject.create();
  31
  32     private boolean isRetryingDecryption = false;
  33
  34     public static void createSql(Connection connection) throws SQLException {
  35         // When modifying the CREATE statement here, also add a migration in AccountDatabase.java
  36         try (final var statement = connection.createStatement()) {
  37             statement.executeUpdate("""
  38                                     CREATE TABLE identity (
  39                                       _id INTEGER PRIMARY KEY,
  40                                       recipient_id INTEGER UNIQUE NOT NULL REFERENCES recipient (_id) ON DELETE CASCADE,
  41                                       identity_key BLOB NOT NULL,
  42                                       added_timestamp INTEGER NOT NULL,
  43                                       trust_level INTEGER NOT NULL
  44                                     );
  45                                     """);
  46         }
  47     }
  48
  49     public IdentityKeyStore(
  50             final Database database,
  51             final RecipientIdCreator recipientIdCreator,
  52             final TrustNewIdentity trustNewIdentity
  53     ) {
  54         this.database = database;
  55         this.recipientIdCreator = recipientIdCreator;
  56         this.trustNewIdentity = trustNewIdentity;
  57     }
  58
  59     public Observable<RecipientId> getIdentityChanges() {
  60         return identityChanges;
  61     }
  62
  63     public boolean saveIdentity(final RecipientId recipientId, final IdentityKey identityKey) {
  64         if (isRetryingDecryption) {
  65             return false;
  66         }
  67         try (final var connection = database.getConnection()) {
  68             final var identityInfo = loadIdentity(connection, recipientId);
  69             if (identityInfo != null && identityInfo.getIdentityKey().equals(identityKey)) {
  70                 // Identity already exists, not updating the trust level
  71                 logger.trace("Not storing new identity for recipient {}, identity already stored", recipientId);
  72                 return false;
  73             }
  74
  75             saveNewIdentity(connection, recipientId, identityKey, identityInfo == null);
  76             return true;
  77         } catch (SQLException e) {
  78             throw new RuntimeException("Failed update identity store", e);
  79         }
  80     }
  81
  82     public void setRetryingDecryption(final boolean retryingDecryption) {
  83         isRetryingDecryption = retryingDecryption;
  84     }
  85
  86     public boolean setIdentityTrustLevel(RecipientId recipientId, IdentityKey identityKey, TrustLevel trustLevel) {
  87         try (final var connection = database.getConnection()) {
  88             final var identityInfo = loadIdentity(connection, recipientId);
  89             if (identityInfo == null) {
  90                 logger.debug("Not updating trust level for recipient {}, identity not found", recipientId);
  91                 return false;
  92             }
  93             if (!identityInfo.getIdentityKey().equals(identityKey)) {
  94                 logger.debug("Not updating trust level for recipient {}, different identity found", recipientId);
  95                 return false;
  96             }
  97             if (identityInfo.getTrustLevel() == trustLevel) {
  98                 logger.trace("Not updating trust level for recipient {}, trust level already matches", recipientId);
  99                 return false;
 100             }
 101
 102             logger.debug("Updating trust level for recipient {} with trust {}", recipientId, trustLevel);
 103             final var newIdentityInfo = new IdentityInfo(recipientId,
 104                     identityKey,
 105                     trustLevel,
 106                     identityInfo.getDateAddedTimestamp());
 107             storeIdentity(connection, newIdentityInfo);
 108             return true;
 109         } catch (SQLException e) {
 110             throw new RuntimeException("Failed update identity store", e);
 111         }
 112     }
 113
 114     public boolean isTrustedIdentity(RecipientId recipientId, IdentityKey identityKey, Direction direction) {
 115         if (trustNewIdentity == TrustNewIdentity.ALWAYS) {
 116             return true;
 117         }
 118
 119         try (final var connection = database.getConnection()) {
 120             // TODO implement possibility for different handling of incoming/outgoing trust decisions
 121             var identityInfo = loadIdentity(connection, recipientId);
 122             if (identityInfo == null) {
 123                 logger.debug("Initial identity found for {}, saving.", recipientId);
 124                 saveNewIdentity(connection, recipientId, identityKey, true);
 125                 identityInfo = loadIdentity(connection, recipientId);
 126             } else if (!identityInfo.getIdentityKey().equals(identityKey)) {
 127                 // Identity found, but different
 128                 if (direction == Direction.SENDING) {
 129                     logger.debug("Changed identity found for {}, saving.", recipientId);
 130                     saveNewIdentity(connection, recipientId, identityKey, false);
 131                     identityInfo = loadIdentity(connection, recipientId);
 132                 } else {
 133                     logger.trace("Trusting identity for {} for {}: {}", recipientId, direction, false);
 134                     return false;
 135                 }
 136             }
 137
 138             final var isTrusted = identityInfo != null && identityInfo.isTrusted();
 139             logger.trace("Trusting identity for {} for {}: {}", recipientId, direction, isTrusted);
 140             return isTrusted;
 141         } catch (SQLException e) {
 142             throw new RuntimeException("Failed read from identity store", e);
 143         }
 144     }
 145
 146     public IdentityInfo getIdentityInfo(RecipientId recipientId) {
 147         try (final var connection = database.getConnection()) {
 148             return loadIdentity(connection, recipientId);
 149         } catch (SQLException e) {
 150             throw new RuntimeException("Failed read from identity store", e);
 151         }
 152     }
 153
 154     public List<IdentityInfo> getIdentities() {
 155         try (final var connection = database.getConnection()) {
 156             final var sql = (
 157                     """
 158                     SELECT i.recipient_id, i.identity_key, i.added_timestamp, i.trust_level
 159                     FROM %s AS i
 160                     """
 161             ).formatted(TABLE_IDENTITY);
 162             try (final var statement = connection.prepareStatement(sql)) {
 163                 return Utils.executeQueryForStream(statement, this::getIdentityInfoFromResultSet).toList();
 164             }
 165         } catch (SQLException e) {
 166             throw new RuntimeException("Failed read from identity store", e);
 167         }
 168     }
 169
 170     public void mergeRecipients(final RecipientId recipientId, final RecipientId toBeMergedRecipientId) {
 171         try (final var connection = database.getConnection()) {
 172             connection.setAutoCommit(false);
 173             final var sql = (
 174                     """
 175                     UPDATE OR IGNORE %s
 176                     SET recipient_id = ?
 177                     WHERE recipient_id = ?
 178                     """
 179             ).formatted(TABLE_IDENTITY);
 180             try (final var statement = connection.prepareStatement(sql)) {
 181                 statement.setLong(1, recipientId.id());
 182                 statement.setLong(2, toBeMergedRecipientId.id());
 183                 statement.executeUpdate();
 184             }
 185
 186             deleteIdentity(connection, toBeMergedRecipientId);
 187             connection.commit();
 188         } catch (SQLException e) {
 189             throw new RuntimeException("Failed update identity store", e);
 190         }
 191     }
 192
 193     public void deleteIdentity(final RecipientId recipientId) {
 194         try (final var connection = database.getConnection()) {
 195             deleteIdentity(connection, recipientId);
 196         } catch (SQLException e) {
 197             throw new RuntimeException("Failed update identity store", e);
 198         }
 199     }
 200
 201     void addLegacyIdentities(final Collection<IdentityInfo> identities) {
 202         logger.debug("Migrating legacy identities to database");
 203         long start = System.nanoTime();
 204         try (final var connection = database.getConnection()) {
 205             connection.setAutoCommit(false);
 206             for (final var identityInfo : identities) {
 207                 storeIdentity(connection, identityInfo);
 208             }
 209             connection.commit();
 210         } catch (SQLException e) {
 211             throw new RuntimeException("Failed update identity store", e);
 212         }
 213         logger.debug("Complete identities migration took {}ms", (System.nanoTime() - start) / 1000000);
 214     }
 215
 216     private IdentityInfo loadIdentity(
 217             final Connection connection, final RecipientId recipientId
 218     ) throws SQLException {
 219         final var sql = (
 220                 """
 221                 SELECT i.recipient_id, i.identity_key, i.added_timestamp, i.trust_level
 222                 FROM %s AS i
 223                 WHERE i.recipient_id = ?
 224                 """
 225         ).formatted(TABLE_IDENTITY);
 226         try (final var statement = connection.prepareStatement(sql)) {
 227             statement.setLong(1, recipientId.id());
 228             return Utils.executeQueryForOptional(statement, this::getIdentityInfoFromResultSet).orElse(null);
 229         }
 230     }
 231
 232     private void saveNewIdentity(
 233             final Connection connection,
 234             final RecipientId recipientId,
 235             final IdentityKey identityKey,
 236             final boolean firstIdentity
 237     ) throws SQLException {
 238         final var trustLevel = trustNewIdentity == TrustNewIdentity.ALWAYS || (
 239                 trustNewIdentity == TrustNewIdentity.ON_FIRST_USE && firstIdentity
 240         ) ? TrustLevel.TRUSTED_UNVERIFIED : TrustLevel.UNTRUSTED;
 241         logger.debug("Storing new identity for recipient {} with trust {}", recipientId, trustLevel);
 242         final var newIdentityInfo = new IdentityInfo(recipientId, identityKey, trustLevel, System.currentTimeMillis());
 243         storeIdentity(connection, newIdentityInfo);
 244         identityChanges.onNext(recipientId);
 245     }
 246
 247     private void storeIdentity(final Connection connection, final IdentityInfo identityInfo) throws SQLException {
 248         logger.trace("Storing identity info for {}, trust: {}, added: {}",
 249                 identityInfo.getRecipientId(),
 250                 identityInfo.getTrustLevel(),
 251                 identityInfo.getDateAddedTimestamp());
 252         final var sql = (
 253                 """
 254                 INSERT OR REPLACE INTO %s (recipient_id, identity_key, added_timestamp, trust_level)
 255                 VALUES (?, ?, ?, ?)
 256                 """
 257         ).formatted(TABLE_IDENTITY);
 258         try (final var statement = connection.prepareStatement(sql)) {
 259             statement.setLong(1, identityInfo.getRecipientId().id());
 260             statement.setBytes(2, identityInfo.getIdentityKey().serialize());
 261             statement.setLong(3, identityInfo.getDateAddedTimestamp());
 262             statement.setInt(4, identityInfo.getTrustLevel().ordinal());
 263             statement.executeUpdate();
 264         }
 265     }
 266
 267     private void deleteIdentity(final Connection connection, final RecipientId recipientId) throws SQLException {
 268         final var sql = (
 269                 """
 270                 DELETE FROM %s AS i
 271                 WHERE i.recipient_id = ?
 272                 """
 273         ).formatted(TABLE_IDENTITY);
 274         try (final var statement = connection.prepareStatement(sql)) {
 275             statement.setLong(1, recipientId.id());
 276             statement.executeUpdate();
 277         }
 278     }
 279
 280     private IdentityInfo getIdentityInfoFromResultSet(ResultSet resultSet) throws SQLException {
 281         try {
 282             final var recipientId = recipientIdCreator.create(resultSet.getLong("recipient_id"));
 283             final var id = new IdentityKey(resultSet.getBytes("identity_key"));
 284             final var trustLevel = TrustLevel.fromInt(resultSet.getInt("trust_level"));
 285             final var added = resultSet.getLong("added_timestamp");
 286
 287             return new IdentityInfo(recipientId, id, trustLevel, added);
 288         } catch (InvalidKeyException e) {
 289             logger.warn("Failed to load identity key, resetting: {}", e.getMessage());
 290             return null;
 291         }
 292     }
 293 }