From 25258db55d36af977444c5cc8bbf99dd6246aa74 Mon Sep 17 00:00:00 2001 From: AsamK Date: Sun, 18 Feb 2024 16:32:50 +0100 Subject: [PATCH] Restructure pre key refresh to be more robust --- .../signal/manager/helper/PreKeyHelper.java | 133 ++++++++++-------- 1 file changed, 73 insertions(+), 60 deletions(-) diff --git a/lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java b/lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java index 6dd579c6..f81623e4 100644 --- a/lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java +++ b/lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java @@ -51,6 +51,14 @@ public class PreKeyHelper { return; } + if (refreshPreKeysIfNecessary(serviceIdType, identityKeyPair)) { + refreshPreKeysIfNecessary(serviceIdType, identityKeyPair); + } + } + + private boolean refreshPreKeysIfNecessary( + final ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair + ) throws IOException { OneTimePreKeyCounts preKeyCounts; try { preKeyCounts = dependencies.getAccountManager().getPreKeyCounts(serviceIdType); @@ -59,65 +67,80 @@ public class PreKeyHelper { preKeyCounts = new OneTimePreKeyCounts(0, 0); } - SignedPreKeyRecord signedPreKeyRecord = null; List preKeyRecords = null; - KyberPreKeyRecord lastResortKyberPreKeyRecord = null; - List kyberPreKeyRecords = null; - - try { - if (preKeyCounts.getEcCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) { - logger.debug("Refreshing {} ec pre keys, because only {} of min {} pre keys remain", - serviceIdType, - preKeyCounts.getEcCount(), - ServiceConfig.PREKEY_MINIMUM_COUNT); - preKeyRecords = generatePreKeys(serviceIdType); - } - if (signedPreKeyNeedsRefresh(serviceIdType)) { - logger.debug("Refreshing {} signed pre key.", serviceIdType); - signedPreKeyRecord = generateSignedPreKey(serviceIdType, identityKeyPair); - } - } catch (Exception e) { - logger.warn("Failed to store new pre keys, resetting preKey id offset", e); - account.resetPreKeyOffsets(serviceIdType); + if (preKeyCounts.getEcCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) { + logger.debug("Refreshing {} ec pre keys, because only {} of min {} pre keys remain", + serviceIdType, + preKeyCounts.getEcCount(), + ServiceConfig.PREKEY_MINIMUM_COUNT); preKeyRecords = generatePreKeys(serviceIdType); + } + + SignedPreKeyRecord signedPreKeyRecord = null; + if (signedPreKeyNeedsRefresh(serviceIdType)) { + logger.debug("Refreshing {} signed pre key.", serviceIdType); signedPreKeyRecord = generateSignedPreKey(serviceIdType, identityKeyPair); } - try { - if (preKeyCounts.getKyberCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) { - logger.debug("Refreshing {} kyber pre keys, because only {} of min {} pre keys remain", - serviceIdType, - preKeyCounts.getKyberCount(), - ServiceConfig.PREKEY_MINIMUM_COUNT); - kyberPreKeyRecords = generateKyberPreKeys(serviceIdType, identityKeyPair); - } - if (lastResortKyberPreKeyNeedsRefresh(serviceIdType)) { - logger.debug("Refreshing {} last resort kyber pre key.", serviceIdType); - lastResortKyberPreKeyRecord = generateLastResortKyberPreKey(serviceIdType, identityKeyPair); - } - } catch (Exception e) { - logger.warn("Failed to store new kyber pre keys, resetting preKey id offset", e); - account.resetKyberPreKeyOffsets(serviceIdType); + List kyberPreKeyRecords = null; + if (preKeyCounts.getKyberCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) { + logger.debug("Refreshing {} kyber pre keys, because only {} of min {} pre keys remain", + serviceIdType, + preKeyCounts.getKyberCount(), + ServiceConfig.PREKEY_MINIMUM_COUNT); kyberPreKeyRecords = generateKyberPreKeys(serviceIdType, identityKeyPair); + } + + KyberPreKeyRecord lastResortKyberPreKeyRecord = null; + if (lastResortKyberPreKeyNeedsRefresh(serviceIdType)) { + logger.debug("Refreshing {} last resort kyber pre key.", serviceIdType); lastResortKyberPreKeyRecord = generateLastResortKyberPreKey(serviceIdType, identityKeyPair); } - if (signedPreKeyRecord != null - || preKeyRecords != null - || lastResortKyberPreKeyRecord != null - || kyberPreKeyRecords != null) { - final var preKeyUpload = new PreKeyUpload(serviceIdType, - signedPreKeyRecord, - preKeyRecords, - lastResortKyberPreKeyRecord, - kyberPreKeyRecords); + if (signedPreKeyRecord == null + && preKeyRecords == null + && lastResortKyberPreKeyRecord == null + && kyberPreKeyRecords == null) { + return false; + } + + final var preKeyUpload = new PreKeyUpload(serviceIdType, + signedPreKeyRecord, + preKeyRecords, + lastResortKyberPreKeyRecord, + kyberPreKeyRecords); + var needsReset = false; + try { + dependencies.getAccountManager().setPreKeys(preKeyUpload); try { - dependencies.getAccountManager().setPreKeys(preKeyUpload); - } catch (AuthorizationFailedException e) { - // This can happen when the primary device has changed phone number - logger.warn("Failed to updated pre keys: {}", e.getMessage()); + if (preKeyRecords != null) { + account.addPreKeys(serviceIdType, preKeyRecords); + } + if (signedPreKeyRecord != null) { + account.addSignedPreKey(serviceIdType, signedPreKeyRecord); + } + } catch (Exception e) { + logger.warn("Failed to store new pre keys, resetting preKey id offset", e); + account.resetPreKeyOffsets(serviceIdType); + needsReset = true; } + try { + if (kyberPreKeyRecords != null) { + account.addKyberPreKeys(serviceIdType, kyberPreKeyRecords); + } + if (lastResortKyberPreKeyRecord != null) { + account.addLastResortKyberPreKey(serviceIdType, lastResortKyberPreKeyRecord); + } + } catch (Exception e) { + logger.warn("Failed to store new kyber pre keys, resetting preKey id offset", e); + account.resetKyberPreKeyOffsets(serviceIdType); + needsReset = true; + } + } catch (AuthorizationFailedException e) { + // This can happen when the primary device has changed phone number + logger.warn("Failed to updated pre keys: {}", e.getMessage()); } + return needsReset; } public void cleanOldPreKeys() { @@ -135,7 +158,6 @@ public class PreKeyHelper { final var offset = accountData.getPreKeyMetadata().getNextPreKeyId(); var records = KeyUtils.generatePreKeyRecords(offset); - account.addPreKeys(serviceIdType, records); return records; } @@ -159,10 +181,7 @@ public class PreKeyHelper { final var accountData = account.getAccountData(serviceIdType); final var signedPreKeyId = accountData.getPreKeyMetadata().getNextSignedPreKeyId(); - var record = KeyUtils.generateSignedPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey()); - account.addSignedPreKey(serviceIdType, record); - - return record; + return KeyUtils.generateSignedPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey()); } private List generateKyberPreKeys( @@ -171,10 +190,7 @@ public class PreKeyHelper { final var accountData = account.getAccountData(serviceIdType); final var offset = accountData.getPreKeyMetadata().getNextKyberPreKeyId(); - var records = KeyUtils.generateKyberPreKeyRecords(offset, identityKeyPair.getPrivateKey()); - account.addKyberPreKeys(serviceIdType, records); - - return records; + return KeyUtils.generateKyberPreKeyRecords(offset, identityKeyPair.getPrivateKey()); } private boolean lastResortKyberPreKeyNeedsRefresh(ServiceIdType serviceIdType) { @@ -199,10 +215,7 @@ public class PreKeyHelper { final var accountData = account.getAccountData(serviceIdType); final var signedPreKeyId = accountData.getPreKeyMetadata().getNextKyberPreKeyId(); - var record = KeyUtils.generateKyberPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey()); - account.addLastResortKyberPreKey(serviceIdType, record); - - return record; + return KeyUtils.generateKyberPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey()); } private void cleanSignedPreKeys(ServiceIdType serviceIdType) { -- 2.50.1