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