Implement support for kyber pre keys

[signal-cli] / lib / src / main / java / org / asamk / signal / manager / helper / PreKeyHelper.java

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 06dc31b91d0431ead7052a20dc99063be643c6bf..282b017aa435d57b4b6f22e7a5957da22c4bdfe0 100644 (file)
--- a/lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java
+++ b/lib/src/main/java/org/asamk/signal/manager/helper/PreKeyHelper.java
@@ -5,6 +5,7 @@ import org.asamk.signal.manager.internal.SignalDependencies;
 import org.asamk.signal.manager.storage.SignalAccount;
 import org.asamk.signal.manager.util.KeyUtils;
 import org.signal.libsignal.protocol.IdentityKeyPair;
+import org.signal.libsignal.protocol.state.KyberPreKeyRecord;
 import org.signal.libsignal.protocol.state.PreKeyRecord;
 import org.signal.libsignal.protocol.state.SignedPreKeyRecord;
 import org.slf4j.Logger;
@@ -39,7 +40,9 @@ public class PreKeyHelper {
         if (preKeyCounts.getEcCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) {
             refreshPreKeys(serviceIdType);
         }
-        // TODO kyber pre keys
+        if (preKeyCounts.getKyberCount() < ServiceConfig.PREKEY_MINIMUM_COUNT) {
+            refreshKyberPreKeys(serviceIdType);
+        }
     }
 
     public void refreshPreKeys() throws IOException {
@@ -47,7 +50,7 @@ public class PreKeyHelper {
         refreshPreKeys(ServiceIdType.PNI);
     }
 
-    public void refreshPreKeys(ServiceIdType serviceIdType) throws IOException {
+    private void refreshPreKeys(ServiceIdType serviceIdType) throws IOException {
         final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
         if (identityKeyPair == null) {
             return;
@@ -97,4 +100,61 @@ public class PreKeyHelper {
 
         return record;
     }
+
+    private void refreshKyberPreKeys(ServiceIdType serviceIdType) throws IOException {
+        final var identityKeyPair = account.getIdentityKeyPair(serviceIdType);
+        if (identityKeyPair == null) {
+            return;
+        }
+        final var accountId = account.getAccountId(serviceIdType);
+        if (accountId == null) {
+            return;
+        }
+        try {
+            refreshKyberPreKeys(serviceIdType, identityKeyPair);
+        } catch (Exception e) {
+            logger.warn("Failed to store new pre keys, resetting preKey id offset", e);
+            account.resetKyberPreKeyOffsets(serviceIdType);
+            refreshKyberPreKeys(serviceIdType, identityKeyPair);
+        }
+    }
+
+    private void refreshKyberPreKeys(
+            final ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
+    ) throws IOException {
+        final var oneTimePreKeys = generateKyberPreKeys(serviceIdType, identityKeyPair);
+        final var lastResortPreKeyRecord = generateLastResortKyberPreKey(serviceIdType, identityKeyPair);
+
+        final var preKeyUpload = new PreKeyUpload(serviceIdType,
+                identityKeyPair.getPublicKey(),
+                null,
+                null,
+                lastResortPreKeyRecord,
+                oneTimePreKeys);
+        dependencies.getAccountManager().setPreKeys(preKeyUpload);
+    }
+
+    private List<KyberPreKeyRecord> generateKyberPreKeys(
+            ServiceIdType serviceIdType, final IdentityKeyPair identityKeyPair
+    ) {
+        final var offset = account.getKyberPreKeyIdOffset(serviceIdType);
+
+        var records = KeyUtils.generateKyberPreKeyRecords(offset,
+                ServiceConfig.PREKEY_BATCH_SIZE,
+                identityKeyPair.getPrivateKey());
+        account.addKyberPreKeys(serviceIdType, records);
+
+        return records;
+    }
+
+    private KyberPreKeyRecord generateLastResortKyberPreKey(
+            ServiceIdType serviceIdType, IdentityKeyPair identityKeyPair
+    ) {
+        final var signedPreKeyId = account.getKyberPreKeyIdOffset(serviceIdType);
+
+        var record = KeyUtils.generateKyberPreKeyRecord(signedPreKeyId, identityKeyPair.getPrivateKey());
+        account.addLastResortKyberPreKey(serviceIdType, record);
+
+        return record;
+    }
 }