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(final SignalAccount account, final SignalDependencies dependencies) {
  34         this.account = account;
  35         this.dependencies = dependencies;
  36     }
  37
  38     public void refreshPreKeysIfNecessary() throws IOException {
  39         refreshPreKeysIfNecessary(ServiceIdType.ACI);
  40         refreshPreKeysIfNecessary(ServiceIdType.PNI);
  41     }
  42
  43     public void forceRefreshPreKeys() throws IOException {
  44         forceRefreshPreKeys(ServiceIdType.ACI);
  45         forceRefreshPreKeys(ServiceIdType.PNI);
  46     }
  47
  48     public void refreshPreKeysIfNecessary(ServiceIdType serviceIdType) throws IOException {
  49         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
  50         if (identityKeyPair == null) {
  51             return;
  52         }
  53         final var accountId = account.getAccountId(serviceIdType);
  54         if (accountId == null) {
  55             return;
  56         }
  57
  58         if (refreshPreKeysIfNecessary(serviceIdType, identityKeyPair)) {
  59             refreshPreKeysIfNecessary(serviceIdType, identityKeyPair);
  60         }
  61     }
  62
  63     public void forceRefreshPreKeys(ServiceIdType serviceIdType) throws IOException {
  64         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
  65         if (identityKeyPair == null) {
  66             return;
  67         }
  68         final var accountId = account.getAccountId(serviceIdType);
  69         if (accountId == null) {
  70             return;
  71         }
  72
  73         final var counts = new OneTimePreKeyCounts(0, 0);
  74         if (refreshPreKeysIfNecessary(serviceIdType, identityKeyPair, counts, true)) {
  75             refreshPreKeysIfNecessary(serviceIdType, identityKeyPair, counts, true);
  76         }
  77     }
  78
  79     private boolean refreshPreKeysIfNecessary(
  80             final ServiceIdType serviceIdType,
  81             final IdentityKeyPair identityKeyPair
  82     ) throws IOException {
  83         OneTimePreKeyCounts preKeyCounts;
  84         try {
  85             preKeyCounts = dependencies.getAccountManager().getPreKeyCounts(serviceIdType);
  86         } catch (AuthorizationFailedException e) {
  87             logger.debug("Failed to get pre key count, ignoring: " + e.getClass().getSimpleName());
  88             preKeyCounts = new OneTimePreKeyCounts(0, 0);
  89         }
  90
  91         return refreshPreKeysIfNecessary(serviceIdType, identityKeyPair, preKeyCounts, false);
  92     }
  93
  94     private boolean refreshPreKeysIfNecessary(
  95             final ServiceIdType serviceIdType,
  96             final IdentityKeyPair identityKeyPair,
  97             final OneTimePreKeyCounts preKeyCounts,
  98             final boolean force
  99     ) throws IOException {
 100         List<PreKeyRecord> preKeyRecords = null;
 101         if (force || preKeyCounts.getEcCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) {
 102             logger.debug("Refreshing {} ec pre keys, because only {} of min {} pre keys remain",
 103                     serviceIdType,
 104                     preKeyCounts.getEcCount(),
 105                     ServiceConfig.PREKEY_MINIMUM_COUNT);
 106             preKeyRecords = generatePreKeys(serviceIdType);
 107         }
 108
 109         SignedPreKeyRecord signedPreKeyRecord = null;
 110         if (force || signedPreKeyNeedsRefresh(serviceIdType)) {
 111             logger.debug("Refreshing {} signed pre key.", serviceIdType);
 112             signedPreKeyRecord = generateSignedPreKey(serviceIdType, identityKeyPair);
 113         }
 114
 115         List<KyberPreKeyRecord> kyberPreKeyRecords = null;
 116         if (force || preKeyCounts.getKyberCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) {
 117             logger.debug("Refreshing {} kyber pre keys, because only {} of min {} pre keys remain",
 118                     serviceIdType,
 119                     preKeyCounts.getKyberCount(),
 120                     ServiceConfig.PREKEY_MINIMUM_COUNT);
 121             kyberPreKeyRecords = generateKyberPreKeys(serviceIdType, identityKeyPair);
 122         }
 123
 124         KyberPreKeyRecord lastResortKyberPreKeyRecord = null;
 125         if (force || lastResortKyberPreKeyNeedsRefresh(serviceIdType)) {
 126             logger.debug("Refreshing {} last resort kyber pre key.", serviceIdType);
 127             lastResortKyberPreKeyRecord = generateLastResortKyberPreKey(serviceIdType,
 128                     identityKeyPair,
 129                     kyberPreKeyRecords == null ? 0 : kyberPreKeyRecords.size());
 130         }
 131
 132         if (signedPreKeyRecord == null
 133                 && preKeyRecords == null
 134                 && lastResortKyberPreKeyRecord == null
 135                 && kyberPreKeyRecords == null) {
 136             return false;
 137         }
 138
 139         final var preKeyUpload = new PreKeyUpload(serviceIdType,
 140                 signedPreKeyRecord,
 141                 preKeyRecords,
 142                 lastResortKyberPreKeyRecord,
 143                 kyberPreKeyRecords);
 144         var needsReset = false;
 145         try {
 146             dependencies.getAccountManager().setPreKeys(preKeyUpload);
 147             try {
 148                 if (preKeyRecords != null) {
 149                     account.addPreKeys(serviceIdType, preKeyRecords);
 150                 }
 151                 if (signedPreKeyRecord != null) {
 152                     account.addSignedPreKey(serviceIdType, signedPreKeyRecord);
 153                 }
 154             } catch (Exception e) {
 155                 logger.warn("Failed to store new pre keys, resetting preKey id offset", e);
 156                 account.resetPreKeyOffsets(serviceIdType);
 157                 needsReset = true;
 158             }
 159             try {
 160                 if (kyberPreKeyRecords != null) {
 161                     account.addKyberPreKeys(serviceIdType, kyberPreKeyRecords);
 162                 }
 163                 if (lastResortKyberPreKeyRecord != null) {
 164                     account.addLastResortKyberPreKey(serviceIdType, lastResortKyberPreKeyRecord);
 165                 }
 166             } catch (Exception e) {
 167                 logger.warn("Failed to store new kyber pre keys, resetting preKey id offset", e);
 168                 account.resetKyberPreKeyOffsets(serviceIdType);
 169                 needsReset = true;
 170             }
 171         } catch (AuthorizationFailedException e) {
 172             // This can happen when the primary device has changed phone number
 173             logger.warn("Failed to updated pre keys: {}", e.getMessage());
 174         } catch (NonSuccessfulResponseCodeException e) {
 175             if (serviceIdType != ServiceIdType.PNI || e.code != 422) {
 176                 throw e;
 177             }
 178             logger.warn("Failed to set PNI pre keys, ignoring for now. Account needs to be reregistered to fix this.");
 179         }
 180         return needsReset;
 181     }
 182
 183     public void cleanOldPreKeys() {
 184         cleanOldPreKeys(ServiceIdType.ACI);
 185         cleanOldPreKeys(ServiceIdType.PNI);
 186     }
 187
 188     private void cleanOldPreKeys(final ServiceIdType serviceIdType) {
 189         cleanSignedPreKeys(serviceIdType);
 190         cleanOneTimePreKeys(serviceIdType);
 191     }
 192
 193     private List<PreKeyRecord> generatePreKeys(ServiceIdType serviceIdType) {
 194         final var accountData = account.getAccountData(serviceIdType);
 195         final var offset = accountData.getPreKeyMetadata().getNextPreKeyId();
 196
 197         return KeyUtils.generatePreKeyRecords(offset);
 198     }
 199
 200     private boolean signedPreKeyNeedsRefresh(ServiceIdType serviceIdType) {
 201         final var accountData = account.getAccountData(serviceIdType);
 202
 203         final var activeSignedPreKeyId = accountData.getPreKeyMetadata().getActiveSignedPreKeyId();
 204         if (activeSignedPreKeyId == -1) {
 205             return true;
 206         }
 207         try {
 208             final var signedPreKeyRecord = accountData.getSignedPreKeyStore().loadSignedPreKey(activeSignedPreKeyId);
 209             return signedPreKeyRecord.getTimestamp() < System.currentTimeMillis() - SIGNED_PREKEY_ROTATE_AGE;
 210         } catch (InvalidKeyIdException e) {
 211             return true;
 212         }
 213     }
 214
 215     private SignedPreKeyRecord generateSignedPreKey(ServiceIdType serviceIdType, IdentityKeyPair identityKeyPair) {
 216         final var accountData = account.getAccountData(serviceIdType);
 217         final var signedPreKeyId = accountData.getPreKeyMetadata().getNextSignedPreKeyId();
 218
 219         return KeyUtils.generateSignedPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey());
 220     }
 221
 222     private List<KyberPreKeyRecord> generateKyberPreKeys(
 223             ServiceIdType serviceIdType,
 224             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,
 250             IdentityKeyPair identityKeyPair,
 251             final int offset
 252     ) {
 253         final var accountData = account.getAccountData(serviceIdType);
 254         final var signedPreKeyId = accountData.getPreKeyMetadata().getNextKyberPreKeyId() + offset;
 255
 256         return KeyUtils.generateKyberPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey());
 257     }
 258
 259     private void cleanSignedPreKeys(ServiceIdType serviceIdType) {
 260         final var accountData = account.getAccountData(serviceIdType);
 261
 262         final var activeSignedPreKeyId = accountData.getPreKeyMetadata().getActiveSignedPreKeyId();
 263         accountData.getSignedPreKeyStore().removeOldSignedPreKeys(activeSignedPreKeyId);
 264
 265         final var activeLastResortKyberPreKeyId = accountData.getPreKeyMetadata().getActiveLastResortKyberPreKeyId();
 266         accountData.getKyberPreKeyStore().removeOldLastResortKyberPreKeys(activeLastResortKyberPreKeyId);
 267     }
 268
 269     private void cleanOneTimePreKeys(ServiceIdType serviceIdType) {
 270         long threshold = System.currentTimeMillis() - PREKEY_STALE_AGE;
 271         int minCount = 200;
 272
 273         final var accountData = account.getAccountData(serviceIdType);
 274         accountData.getPreKeyStore().deleteAllStaleOneTimeEcPreKeys(threshold, minCount);
 275         accountData.getKyberPreKeyStore().deleteAllStaleOneTimeKyberPreKeys(threshold, minCount);
 276     }
 277 }