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