nmode's Git Repositories - signal-cli/blob - lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java

   1 package org.asamk.signal.manager.helper;
   2
   3 import org.asamk.signal.manager.config.ServiceConfig;
   4 import org.asamk.signal.manager.internal.SignalDependencies;
   5 import org.asamk.signal.manager.storage.SignalAccount;
   6 import org.asamk.signal.manager.util.KeyUtils;
   7 import org.signal.libsignal.protocol.IdentityKeyPair;
   8 import org.signal.libsignal.protocol.InvalidKeyIdException;
   9 import org.signal.libsignal.protocol.state.KyberPreKeyRecord;
  10 import org.signal.libsignal.protocol.state.PreKeyRecord;
  11 import org.signal.libsignal.protocol.state.SignedPreKeyRecord;
  12 import org.slf4j.Logger;
  13 import org.slf4j.LoggerFactory;
  14 import org.whispersystems.signalservice.api.account.PreKeyUpload;
  15 import org.whispersystems.signalservice.api.push.ServiceIdType;
  16 import org.whispersystems.signalservice.api.push.exceptions.AuthorizationFailedException;
  17 import org.whispersystems.signalservice.api.push.exceptions.NonSuccessfulResponseCodeException;
  18 import org.whispersystems.signalservice.internal.push.OneTimePreKeyCounts;
  19
  20 import java.io.IOException;
  21 import java.util.List;
  22
  23 import static org.asamk.signal.manager.config.ServiceConfig.PREKEY_STALE_AGE;
  24 import static org.asamk.signal.manager.config.ServiceConfig.SIGNED_PREKEY_ROTATE_AGE;
  25
  26 public class PreKeyHelper {
  27
  28     private static final Logger logger = LoggerFactory.getLogger(PreKeyHelper.class);
  29
  30     private final SignalAccount account;
  31     private final SignalDependencies dependencies;
  32
  33     public PreKeyHelper(
  34             final SignalAccount account, final SignalDependencies dependencies
  35     ) {
  36         this.account = account;
  37         this.dependencies = dependencies;
  38     }
  39
  40     public void refreshPreKeysIfNecessary() throws IOException {
  41         refreshPreKeysIfNecessary(ServiceIdType.ACI);
  42         refreshPreKeysIfNecessary(ServiceIdType.PNI);
  43     }
  44
  45     public void forceRefreshPreKeys() throws IOException {
  46         forceRefreshPreKeys(ServiceIdType.ACI);
  47         forceRefreshPreKeys(ServiceIdType.PNI);
  48     }
  49
  50     public void refreshPreKeysIfNecessary(ServiceIdType serviceIdType) throws IOException {
  51         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
  52         if (identityKeyPair == null) {
  53             return;
  54         }
  55         final var accountId = account.getAccountId(serviceIdType);
  56         if (accountId == null) {
  57             return;
  58         }
  59
  60         if (refreshPreKeysIfNecessary(serviceIdType, identityKeyPair)) {
  61             refreshPreKeysIfNecessary(serviceIdType, identityKeyPair);
  62         }
  63     }
  64
  65     public void forceRefreshPreKeys(ServiceIdType serviceIdType) throws IOException {
  66         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
  67         if (identityKeyPair == null) {
  68             return;
  69         }
  70         final var accountId = account.getAccountId(serviceIdType);
  71         if (accountId == null) {
  72             return;
  73         }
  74
  75         final var counts = new OneTimePreKeyCounts(0, 0);
  76         if (refreshPreKeysIfNecessary(serviceIdType, identityKeyPair, counts, true)) {
  77             refreshPreKeysIfNecessary(serviceIdType, identityKeyPair, counts, true);
  78         }
  79     }
  80
  81     private boolean refreshPreKeysIfNecessary(
  82             final ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
  83     ) throws IOException {
  84         OneTimePreKeyCounts preKeyCounts;
  85         try {
  86             preKeyCounts = dependencies.getAccountManager().getPreKeyCounts(serviceIdType);
  87         } catch (AuthorizationFailedException e) {
  88             logger.debug("Failed to get pre key count, ignoring: " + e.getClass().getSimpleName());
  89             preKeyCounts = new OneTimePreKeyCounts(0, 0);
  90         }
  91
  92         return refreshPreKeysIfNecessary(serviceIdType, identityKeyPair, preKeyCounts, false);
  93     }
  94
  95     private boolean refreshPreKeysIfNecessary(
  96             final ServiceIdType serviceIdType,
  97             final IdentityKeyPair identityKeyPair,
  98             final OneTimePreKeyCounts preKeyCounts,
  99             final boolean force
 100     ) throws IOException {
 101         List<PreKeyRecord> preKeyRecords = null;
 102         if (force || preKeyCounts.getEcCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) {
 103             logger.debug("Refreshing {} ec pre keys, because only {} of min {} pre keys remain",
 104                     serviceIdType,
 105                     preKeyCounts.getEcCount(),
 106                     ServiceConfig.PREKEY_MINIMUM_COUNT);
 107             preKeyRecords = generatePreKeys(serviceIdType);
 108         }
 109
 110         SignedPreKeyRecord signedPreKeyRecord = null;
 111         if (force || signedPreKeyNeedsRefresh(serviceIdType)) {
 112             logger.debug("Refreshing {} signed pre key.", serviceIdType);
 113             signedPreKeyRecord = generateSignedPreKey(serviceIdType, identityKeyPair);
 114         }
 115
 116         List<KyberPreKeyRecord> kyberPreKeyRecords = null;
 117         if (force || preKeyCounts.getKyberCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) {
 118             logger.debug("Refreshing {} kyber pre keys, because only {} of min {} pre keys remain",
 119                     serviceIdType,
 120                     preKeyCounts.getKyberCount(),
 121                     ServiceConfig.PREKEY_MINIMUM_COUNT);
 122             kyberPreKeyRecords = generateKyberPreKeys(serviceIdType, identityKeyPair);
 123         }
 124
 125         KyberPreKeyRecord lastResortKyberPreKeyRecord = null;
 126         if (force || lastResortKyberPreKeyNeedsRefresh(serviceIdType)) {
 127             logger.debug("Refreshing {} last resort kyber pre key.", serviceIdType);
 128             lastResortKyberPreKeyRecord = generateLastResortKyberPreKey(serviceIdType,
 129                     identityKeyPair,
 130                     kyberPreKeyRecords == null ? 0 : kyberPreKeyRecords.size());
 131         }
 132
 133         if (signedPreKeyRecord == null
 134                 && preKeyRecords == null
 135                 && lastResortKyberPreKeyRecord == null
 136                 && kyberPreKeyRecords == null) {
 137             return false;
 138         }
 139
 140         final var preKeyUpload = new PreKeyUpload(serviceIdType,
 141                 signedPreKeyRecord,
 142                 preKeyRecords,
 143                 lastResortKyberPreKeyRecord,
 144                 kyberPreKeyRecords);
 145         var needsReset = false;
 146         try {
 147             dependencies.getAccountManager().setPreKeys(preKeyUpload);
 148             try {
 149                 if (preKeyRecords != null) {
 150                     account.addPreKeys(serviceIdType, preKeyRecords);
 151                 }
 152                 if (signedPreKeyRecord != null) {
 153                     account.addSignedPreKey(serviceIdType, signedPreKeyRecord);
 154                 }
 155             } catch (Exception e) {
 156                 logger.warn("Failed to store new pre keys, resetting preKey id offset", e);
 157                 account.resetPreKeyOffsets(serviceIdType);
 158                 needsReset = true;
 159             }
 160             try {
 161                 if (kyberPreKeyRecords != null) {
 162                     account.addKyberPreKeys(serviceIdType, kyberPreKeyRecords);
 163                 }
 164                 if (lastResortKyberPreKeyRecord != null) {
 165                     account.addLastResortKyberPreKey(serviceIdType, lastResortKyberPreKeyRecord);
 166                 }
 167             } catch (Exception e) {
 168                 logger.warn("Failed to store new kyber pre keys, resetting preKey id offset", e);
 169                 account.resetKyberPreKeyOffsets(serviceIdType);
 170                 needsReset = true;
 171             }
 172         } catch (AuthorizationFailedException e) {
 173             // This can happen when the primary device has changed phone number
 174             logger.warn("Failed to updated pre keys: {}", e.getMessage());
 175         } catch (NonSuccessfulResponseCodeException e) {
 176             if (serviceIdType != ServiceIdType.PNI || e.getCode() != 422) {
 177                 throw e;
 178             }
 179             logger.warn("Failed to update PNI pre keys, ignoring.");
 180         }
 181         return needsReset;
 182     }
 183
 184     public void cleanOldPreKeys() {
 185         cleanOldPreKeys(ServiceIdType.ACI);
 186         cleanOldPreKeys(ServiceIdType.PNI);
 187     }
 188
 189     private void cleanOldPreKeys(final ServiceIdType serviceIdType) {
 190         cleanSignedPreKeys(serviceIdType);
 191         cleanOneTimePreKeys(serviceIdType);
 192     }
 193
 194     private List<PreKeyRecord> generatePreKeys(ServiceIdType serviceIdType) {
 195         final var accountData = account.getAccountData(serviceIdType);
 196         final var offset = accountData.getPreKeyMetadata().getNextPreKeyId();
 197
 198         return KeyUtils.generatePreKeyRecords(offset);
 199     }
 200
 201     private boolean signedPreKeyNeedsRefresh(ServiceIdType serviceIdType) {
 202         final var accountData = account.getAccountData(serviceIdType);
 203
 204         final var activeSignedPreKeyId = accountData.getPreKeyMetadata().getActiveSignedPreKeyId();
 205         if (activeSignedPreKeyId == -1) {
 206             return true;
 207         }
 208         try {
 209             final var signedPreKeyRecord = accountData.getSignedPreKeyStore().loadSignedPreKey(activeSignedPreKeyId);
 210             return signedPreKeyRecord.getTimestamp() < System.currentTimeMillis() - SIGNED_PREKEY_ROTATE_AGE;
 211         } catch (InvalidKeyIdException e) {
 212             return true;
 213         }
 214     }
 215
 216     private SignedPreKeyRecord generateSignedPreKey(ServiceIdType serviceIdType, IdentityKeyPair identityKeyPair) {
 217         final var accountData = account.getAccountData(serviceIdType);
 218         final var signedPreKeyId = accountData.getPreKeyMetadata().getNextSignedPreKeyId();
 219
 220         return KeyUtils.generateSignedPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey());
 221     }
 222
 223     private List<KyberPreKeyRecord> generateKyberPreKeys(
 224             ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
 225     ) {
 226         final var accountData = account.getAccountData(serviceIdType);
 227         final var offset = accountData.getPreKeyMetadata().getNextKyberPreKeyId();
 228
 229         return KeyUtils.generateKyberPreKeyRecords(offset, identityKeyPair.getPrivateKey());
 230     }
 231
 232     private boolean lastResortKyberPreKeyNeedsRefresh(ServiceIdType serviceIdType) {
 233         final var accountData = account.getAccountData(serviceIdType);
 234
 235         final var activeLastResortKyberPreKeyId = accountData.getPreKeyMetadata().getActiveLastResortKyberPreKeyId();
 236         if (activeLastResortKyberPreKeyId == -1) {
 237             return true;
 238         }
 239         try {
 240             final var kyberPreKeyRecord = accountData.getKyberPreKeyStore()
 241                     .loadKyberPreKey(activeLastResortKyberPreKeyId);
 242             return kyberPreKeyRecord.getTimestamp() < System.currentTimeMillis() - SIGNED_PREKEY_ROTATE_AGE;
 243         } catch (InvalidKeyIdException e) {
 244             return true;
 245         }
 246     }
 247
 248     private KyberPreKeyRecord generateLastResortKyberPreKey(
 249             ServiceIdType serviceIdType, IdentityKeyPair identityKeyPair, final int offset
 250     ) {
 251         final var accountData = account.getAccountData(serviceIdType);
 252         final var signedPreKeyId = accountData.getPreKeyMetadata().getNextKyberPreKeyId() + offset;
 253
 254         return KeyUtils.generateKyberPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey());
 255     }
 256
 257     private void cleanSignedPreKeys(ServiceIdType serviceIdType) {
 258         final var accountData = account.getAccountData(serviceIdType);
 259
 260         final var activeSignedPreKeyId = accountData.getPreKeyMetadata().getActiveSignedPreKeyId();
 261         accountData.getSignedPreKeyStore().removeOldSignedPreKeys(activeSignedPreKeyId);
 262
 263         final var activeLastResortKyberPreKeyId = accountData.getPreKeyMetadata().getActiveLastResortKyberPreKeyId();
 264         accountData.getKyberPreKeyStore().removeOldLastResortKyberPreKeys(activeLastResortKyberPreKeyId);
 265     }
 266
 267     private void cleanOneTimePreKeys(ServiceIdType serviceIdType) {
 268         long threshold = System.currentTimeMillis() - PREKEY_STALE_AGE;
 269         int minCount = 200;
 270
 271         final var accountData = account.getAccountData(serviceIdType);
 272         accountData.getPreKeyStore().deleteAllStaleOneTimeEcPreKeys(threshold, minCount);
 273         accountData.getKyberPreKeyStore().deleteAllStaleOneTimeKyberPreKeys(threshold, minCount);
 274     }
 275 }