From 1605676509c5ba260d7e30cc9b9806726e7c2cf7 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Wed, 12 Jul 2023 12:41:08 -0400
Subject: [PATCH] Store signed EC pre-keys in a dedicated table when setting
 signed pre-keys individually

---
 .../textsecuregcm/controllers/KeysController.java             | 3 +++
 .../textsecuregcm/tests/controllers/KeysControllerTest.java   | 4 ++++
 2 files changed, 7 insertions(+)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
index 1bd650ff3..897d122ed 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
@@ -22,6 +22,7 @@ import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.UUID;
@@ -296,6 +297,8 @@ public class KeysController {
         d.setSignedPreKey(signedPreKey);
       }
     });
+
+    keys.storeEcSignedPreKeys(getIdentifier(auth.getAccount(), identityType), Map.of(device.getId(), signedPreKey)).join();
   }
 
   private static boolean usePhoneNumberIdentity(final Optional<String> identityType) {
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java
index a3133e1cb..6095fa0d9 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java
@@ -30,6 +30,7 @@ import io.dropwizard.testing.junit5.ResourceExtension;
 import java.time.Duration;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.UUID;
@@ -224,6 +225,7 @@ class KeysControllerTest {
 
     when(KEYS.store(any(), anyLong(), any(), any(), any(), any())).thenReturn(CompletableFuture.completedFuture(null));
     when(KEYS.getEcSignedPreKey(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(Optional.empty()));
+    when(KEYS.storeEcSignedPreKeys(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
 
     when(KEYS.takeEC(EXISTS_UUID, 1)).thenReturn(CompletableFuture.completedFuture(Optional.of(SAMPLE_KEY)));
     when(KEYS.takePQ(EXISTS_UUID, 1)).thenReturn(CompletableFuture.completedFuture(Optional.of(SAMPLE_PQ_KEY)));
@@ -281,6 +283,7 @@ class KeysControllerTest {
     verify(AuthHelper.VALID_DEVICE).setSignedPreKey(eq(test));
     verify(AuthHelper.VALID_DEVICE, never()).setPhoneNumberIdentitySignedPreKey(any());
     verify(accounts).updateDevice(eq(AuthHelper.VALID_ACCOUNT), anyLong(), any());
+    verify(KEYS).storeEcSignedPreKeys(AuthHelper.VALID_UUID, Map.of(Device.MASTER_ID, test));
   }
 
   @Test
@@ -299,6 +302,7 @@ class KeysControllerTest {
     verify(AuthHelper.VALID_DEVICE).setPhoneNumberIdentitySignedPreKey(eq(replacementKey));
     verify(AuthHelper.VALID_DEVICE, never()).setSignedPreKey(any());
     verify(accounts).updateDevice(eq(AuthHelper.VALID_ACCOUNT), anyLong(), any());
+    verify(KEYS).storeEcSignedPreKeys(AuthHelper.VALID_PNI, Map.of(Device.MASTER_ID, replacementKey));
   }
 
   @Test