Implement support for kyber pre keys

[signal-cli] / lib / src / main / java / org / asamk / signal / manager / util / KeyUtils.java

diff --git a/lib/src/main/java/org/asamk/signal/manager/util/KeyUtils.java b/lib/src/main/java/org/asamk/signal/manager/util/KeyUtils.java
index d87868eedf754b7f3aa0eb4438ff68424019f99c..3ee5657ebda436abd6d4f3bbd81e7f5ff969419a 100644 (file)
--- a/lib/src/main/java/org/asamk/signal/manager/util/KeyUtils.java
+++ b/lib/src/main/java/org/asamk/signal/manager/util/KeyUtils.java
@@ -5,9 +5,11 @@ import org.signal.libsignal.protocol.IdentityKeyPair;
 import org.signal.libsignal.protocol.InvalidKeyException;
 import org.signal.libsignal.protocol.ecc.Curve;
 import org.signal.libsignal.protocol.ecc.ECPrivateKey;
+import org.signal.libsignal.protocol.kem.KEMKeyPair;
+import org.signal.libsignal.protocol.kem.KEMKeyType;
+import org.signal.libsignal.protocol.state.KyberPreKeyRecord;
 import org.signal.libsignal.protocol.state.PreKeyRecord;
 import org.signal.libsignal.protocol.state.SignedPreKeyRecord;
-import org.signal.libsignal.protocol.util.Medium;
 import org.signal.libsignal.zkgroup.InvalidInputException;
 import org.signal.libsignal.zkgroup.profiles.ProfileKey;
 import org.whispersystems.signalservice.api.kbs.MasterKey;
@@ -17,6 +19,8 @@ import java.util.ArrayList;
 import java.util.Base64;
 import java.util.List;
 
+import static org.asamk.signal.manager.config.ServiceConfig.PREKEY_MAXIMUM_ID;
+
 public class KeyUtils {
 
     private static final SecureRandom secureRandom = new SecureRandom();
@@ -46,7 +50,7 @@ public class KeyUtils {
     public static List<PreKeyRecord> generatePreKeyRecords(final int offset, final int batchSize) {
         var records = new ArrayList<PreKeyRecord>(batchSize);
         for (var i = 0; i < batchSize; i++) {
-            var preKeyId = (offset + i) % Medium.MAX_VALUE;
+            var preKeyId = (offset + i) % PREKEY_MAXIMUM_ID;
             var keyPair = Curve.generateKeyPair();
             var record = new PreKeyRecord(preKeyId, keyPair);
 
@@ -68,6 +72,24 @@ public class KeyUtils {
         return new SignedPreKeyRecord(signedPreKeyId, System.currentTimeMillis(), keyPair, signature);
     }
 
+    public static List<KyberPreKeyRecord> generateKyberPreKeyRecords(
+            final int offset, final int batchSize, final ECPrivateKey privateKey
+    ) {
+        var records = new ArrayList<KyberPreKeyRecord>(batchSize);
+        for (var i = 0; i < batchSize; i++) {
+            var preKeyId = (offset + i) % PREKEY_MAXIMUM_ID;
+            records.add(generateKyberPreKeyRecord(preKeyId, privateKey));
+        }
+        return records;
+    }
+
+    public static KyberPreKeyRecord generateKyberPreKeyRecord(final int preKeyId, final ECPrivateKey privateKey) {
+        KEMKeyPair keyPair = KEMKeyPair.generate(KEMKeyType.KYBER_1024);
+        byte[] signature = privateKey.calculateSignature(keyPair.getPublicKey().serialize());
+
+        return new KyberPreKeyRecord(preKeyId, System.currentTimeMillis(), keyPair, signature);
+    }
+
     public static ProfileKey createProfileKey() {
         try {
             return new ProfileKey(getSecretBytes(32));