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