From 93c78b6e400d9b4e26dbeddee5eecdb03f96fd08 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Tue, 20 Jun 2023 09:59:40 -0400
Subject: [PATCH] Introduce `RepeatedUseECSignedPreKeyStore`

---
 .../RepeatedUseECSignedPreKeyStore.java       | 46 +++++++++++++++++++
 .../RepeatedUseECSignedPreKeyStoreTest.java   | 42 +++++++++++++++++
 2 files changed, 88 insertions(+)
 create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/RepeatedUseECSignedPreKeyStore.java
 create mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/storage/RepeatedUseECSignedPreKeyStoreTest.java

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/RepeatedUseECSignedPreKeyStore.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/RepeatedUseECSignedPreKeyStore.java
new file mode 100644
index 000000000..7d1f2a119
--- /dev/null
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/RepeatedUseECSignedPreKeyStore.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import java.util.Map;
+import java.util.UUID;
+import org.signal.libsignal.protocol.InvalidKeyException;
+import org.signal.libsignal.protocol.ecc.ECPublicKey;
+import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
+import org.whispersystems.textsecuregcm.util.AttributeValues;
+import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
+import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+
+public class RepeatedUseECSignedPreKeyStore extends RepeatedUseSignedPreKeyStore<ECSignedPreKey> {
+
+  public RepeatedUseECSignedPreKeyStore(final DynamoDbAsyncClient dynamoDbAsyncClient, final String tableName) {
+    super(dynamoDbAsyncClient, tableName);
+  }
+
+  @Override
+  protected Map<String, AttributeValue> getItemFromPreKey(final UUID accountUuid, final long deviceId, final ECSignedPreKey signedPreKey) {
+
+    return Map.of(
+        KEY_ACCOUNT_UUID, getPartitionKey(accountUuid),
+        KEY_DEVICE_ID, getSortKey(deviceId),
+        ATTR_KEY_ID, AttributeValues.fromLong(signedPreKey.keyId()),
+        ATTR_PUBLIC_KEY, AttributeValues.fromByteArray(signedPreKey.serializedPublicKey()),
+        ATTR_SIGNATURE, AttributeValues.fromByteArray(signedPreKey.signature()));
+  }
+
+  @Override
+  protected ECSignedPreKey getPreKeyFromItem(final Map<String, AttributeValue> item) {
+    try {
+      return new ECSignedPreKey(
+          Long.parseLong(item.get(ATTR_KEY_ID).n()),
+          new ECPublicKey(item.get(ATTR_PUBLIC_KEY).b().asByteArray()),
+          item.get(ATTR_SIGNATURE).b().asByteArray());
+    } catch (final InvalidKeyException e) {
+      // This should never happen since we're serializing keys directly from `ECPublicKey` instances on the way in
+      throw new IllegalArgumentException(e);
+    }
+  }
+}
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/RepeatedUseECSignedPreKeyStoreTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/RepeatedUseECSignedPreKeyStoreTest.java
new file mode 100644
index 000000000..8eaea0dd8
--- /dev/null
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/RepeatedUseECSignedPreKeyStoreTest.java
@@ -0,0 +1,42 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.signal.libsignal.protocol.ecc.Curve;
+import org.signal.libsignal.protocol.ecc.ECKeyPair;
+import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
+import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
+
+class RepeatedUseECSignedPreKeyStoreTest extends RepeatedUseSignedPreKeyStoreTest<ECSignedPreKey> {
+
+  private RepeatedUseECSignedPreKeyStore keyStore;
+
+  private int currentKeyId = 1;
+
+  @RegisterExtension
+  static final DynamoDbExtension DYNAMO_DB_EXTENSION =
+      new DynamoDbExtension(DynamoDbExtensionSchema.Tables.REPEATED_USE_SIGNED_PRE_KEYS);
+
+  private static final ECKeyPair IDENTITY_KEY_PAIR = Curve.generateKeyPair();
+
+  @BeforeEach
+  void setUp() {
+    keyStore = new RepeatedUseECSignedPreKeyStore(DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
+        DynamoDbExtensionSchema.Tables.REPEATED_USE_SIGNED_PRE_KEYS.tableName());
+  }
+
+  @Override
+  protected RepeatedUseSignedPreKeyStore<ECSignedPreKey> getKeyStore() {
+    return keyStore;
+  }
+
+  @Override
+  protected ECSignedPreKey generateSignedPreKey() {
+    return KeysHelper.signedECPreKey(currentKeyId++, IDENTITY_KEY_PAIR);
+  }
+}