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 forceRefreshPreKeys() throws IOException {
  45         forceRefreshPreKeys(ServiceIdType.ACI);
  46         forceRefreshPreKeys(ServiceIdType.PNI);
  47     }
  48
  49     public void refreshPreKeysIfNecessary(ServiceIdType serviceIdType) throws IOException {
  50         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
  51         if (identityKeyPair == null) {
  52             return;
  53         }
  54         final var accountId = account.getAccountId(serviceIdType);
  55         if (accountId == null) {
  56             return;
  57         }
  58
  59         if (refreshPreKeysIfNecessary(serviceIdType, identityKeyPair)) {
  60             refreshPreKeysIfNecessary(serviceIdType, identityKeyPair);
  61         }
  62     }
  63
  64     public void forceRefreshPreKeys(ServiceIdType serviceIdType) throws IOException {
  65         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
  66         if (identityKeyPair == null) {
  67             return;
  68         }
  69         final var accountId = account.getAccountId(serviceIdType);
  70         if (accountId == null) {
  71             return;
  72         }
  73
  74         final var counts = new OneTimePreKeyCounts(0, 0);
  75         if (refreshPreKeysIfNecessary(serviceIdType, identityKeyPair, counts, true)) {
  76             refreshPreKeysIfNecessary(serviceIdType, identityKeyPair, counts, true);
  77         }
  78     }
  79
  80     private boolean refreshPreKeysIfNecessary(
  81             final ServiceIdType serviceIdType, 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         }
 175         return needsReset;
 176     }
 177
 178     public void cleanOldPreKeys() {
 179         cleanOldPreKeys(ServiceIdType.ACI);
 180         cleanOldPreKeys(ServiceIdType.PNI);
 181     }
 182
 183     private void cleanOldPreKeys(final ServiceIdType serviceIdType) {
 184         cleanSignedPreKeys(serviceIdType);
 185         cleanOneTimePreKeys(serviceIdType);
 186     }
 187
 188     private List<PreKeyRecord> generatePreKeys(ServiceIdType serviceIdType) {
 189         final var accountData = account.getAccountData(serviceIdType);
 190         final var offset = accountData.getPreKeyMetadata().getNextPreKeyId();
 191
 192         return KeyUtils.generatePreKeyRecords(offset);
 193     }
 194
 195     private boolean signedPreKeyNeedsRefresh(ServiceIdType serviceIdType) {
 196         final var accountData = account.getAccountData(serviceIdType);
 197
 198         final var activeSignedPreKeyId = accountData.getPreKeyMetadata().getActiveSignedPreKeyId();
 199         if (activeSignedPreKeyId == -1) {
 200             return true;
 201         }
 202         try {
 203             final var signedPreKeyRecord = accountData.getSignedPreKeyStore().loadSignedPreKey(activeSignedPreKeyId);
 204             return signedPreKeyRecord.getTimestamp() < System.currentTimeMillis() - SIGNED_PREKEY_ROTATE_AGE;
 205         } catch (InvalidKeyIdException e) {
 206             return true;
 207         }
 208     }
 209
 210     private SignedPreKeyRecord generateSignedPreKey(ServiceIdType serviceIdType, IdentityKeyPair identityKeyPair) {
 211         final var accountData = account.getAccountData(serviceIdType);
 212         final var signedPreKeyId = accountData.getPreKeyMetadata().getNextSignedPreKeyId();
 213
 214         return KeyUtils.generateSignedPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey());
 215     }
 216
 217     private List<KyberPreKeyRecord> generateKyberPreKeys(
 218             ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
 219     ) {
 220         final var accountData = account.getAccountData(serviceIdType);
 221         final var offset = accountData.getPreKeyMetadata().getNextKyberPreKeyId();
 222
 223         return KeyUtils.generateKyberPreKeyRecords(offset, identityKeyPair.getPrivateKey());
 224     }
 225
 226     private boolean lastResortKyberPreKeyNeedsRefresh(ServiceIdType serviceIdType) {
 227         final var accountData = account.getAccountData(serviceIdType);
 228
 229         final var activeLastResortKyberPreKeyId = accountData.getPreKeyMetadata().getActiveLastResortKyberPreKeyId();
 230         if (activeLastResortKyberPreKeyId == -1) {
 231             return true;
 232         }
 233         try {
 234             final var kyberPreKeyRecord = accountData.getKyberPreKeyStore()
 235                     .loadKyberPreKey(activeLastResortKyberPreKeyId);
 236             return kyberPreKeyRecord.getTimestamp() < System.currentTimeMillis() - SIGNED_PREKEY_ROTATE_AGE;
 237         } catch (InvalidKeyIdException e) {
 238             return true;
 239         }
 240     }
 241
 242     private KyberPreKeyRecord generateLastResortKyberPreKey(
 243             ServiceIdType serviceIdType, IdentityKeyPair identityKeyPair, final int offset
 244     ) {
 245         final var accountData = account.getAccountData(serviceIdType);
 246         final var signedPreKeyId = accountData.getPreKeyMetadata().getNextKyberPreKeyId() + offset;
 247
 248         return KeyUtils.generateKyberPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey());
 249     }
 250
 251     private void cleanSignedPreKeys(ServiceIdType serviceIdType) {
 252         final var accountData = account.getAccountData(serviceIdType);
 253
 254         final var activeSignedPreKeyId = accountData.getPreKeyMetadata().getActiveSignedPreKeyId();
 255         accountData.getSignedPreKeyStore().removeOldSignedPreKeys(activeSignedPreKeyId);
 256
 257         final var activeLastResortKyberPreKeyId = accountData.getPreKeyMetadata().getActiveLastResortKyberPreKeyId();
 258         accountData.getKyberPreKeyStore().removeOldLastResortKyberPreKeys(activeLastResortKyberPreKeyId);
 259     }
 260
 261     private void cleanOneTimePreKeys(ServiceIdType serviceIdType) {
 262         long threshold = System.currentTimeMillis() - PREKEY_STALE_AGE;
 263         int minCount = 200;
 264
 265         final var accountData = account.getAccountData(serviceIdType);
 266         accountData.getPreKeyStore().deleteAllStaleOneTimeEcPreKeys(threshold, minCount);
 267         accountData.getKyberPreKeyStore().deleteAllStaleOneTimeKyberPreKeys(threshold, minCount);
 268     }
 269 }