Treat the stand-alone signed pre-keys table as the source of truth for signed pre-keys

service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java

@@ -49,7 +49,6 @@ import org.whispersystems.textsecuregcm.entities.PreKeyResponseItem;
 import org.whispersystems.textsecuregcm.entities.PreKeySignatureValidator;
 import org.whispersystems.textsecuregcm.entities.SetKeysRequest;
 import org.whispersystems.textsecuregcm.entities.SignedPreKey;
-import org.whispersystems.textsecuregcm.experiment.Experiment;
 import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
@@ -67,7 +66,6 @@ public class KeysController {
   private final RateLimiters rateLimiters;
   private final KeysManager keys;
   private final AccountsManager accounts;
-  private final Experiment compareSignedEcPreKeysExperiment = new Experiment("compareSignedEcPreKeys");
 
   private static final CompletableFuture<?>[] EMPTY_FUTURE_ARRAY = new CompletableFuture[0];
 
@@ -237,37 +235,33 @@ public class KeysController {
     final List<PreKeyResponseItem> responseItems = new ArrayList<>(devices.size());
 
     final List<CompletableFuture<Void>> tasks = devices.stream().map(device -> {
+          final CompletableFuture<Optional<ECPreKey>> unsignedEcPreKeyFuture =
+              keys.takeEC(targetIdentifier.uuid(), device.getId());
 
-          ECSignedPreKey signedECPreKey = device.getSignedPreKey(targetIdentifier.identityType());
+          final CompletableFuture<Optional<ECSignedPreKey>> signedEcPreKeyFuture =
+              keys.getEcSignedPreKey(targetIdentifier.uuid(), device.getId());
 
-          final CompletableFuture<Optional<ECPreKey>> unsignedEcPreKeyFuture = keys.takeEC(targetIdentifier.uuid(),
-              device.getId());
           final CompletableFuture<Optional<KEMSignedPreKey>> pqPreKeyFuture = returnPqKey
               ? keys.takePQ(targetIdentifier.uuid(), device.getId())
               : CompletableFuture.completedFuture(Optional.empty());
 
-          return pqPreKeyFuture.thenCombine(unsignedEcPreKeyFuture,
-              (maybePqPreKey, maybeUnsignedEcPreKey) -> {
+          return CompletableFuture.allOf(unsignedEcPreKeyFuture, signedEcPreKeyFuture, pqPreKeyFuture)
+              .thenAccept(ignored -> {
+                final KEMSignedPreKey pqPreKey = pqPreKeyFuture.join().orElse(null);
+                final ECPreKey unsignedEcPreKey = unsignedEcPreKeyFuture.join().orElse(null);
+                final ECSignedPreKey signedEcPreKey = signedEcPreKeyFuture.join().orElse(null);
 
-                KEMSignedPreKey pqPreKey = pqPreKeyFuture.join().orElse(null);
-                ECPreKey unsignedECPreKey = unsignedEcPreKeyFuture.join().orElse(null);
-
-                compareSignedEcPreKeysExperiment.compareFutureResult(Optional.ofNullable(signedECPreKey),
-                    keys.getEcSignedPreKey(targetIdentifier.uuid(), device.getId()));
-
-                if (signedECPreKey != null || unsignedECPreKey != null || pqPreKey != null) {
+                if (signedEcPreKey != null || unsignedEcPreKey != null || pqPreKey != null) {
                   final int registrationId = switch (targetIdentifier.identityType()) {
                     case ACI -> device.getRegistrationId();
                     case PNI -> device.getPhoneNumberIdentityRegistrationId().orElse(device.getRegistrationId());
                   };
 
                   responseItems.add(
-                      new PreKeyResponseItem(device.getId(), registrationId, signedECPreKey, unsignedECPreKey,
+                      new PreKeyResponseItem(device.getId(), registrationId, signedEcPreKey, unsignedEcPreKey,
                           pqPreKey));
                 }
-
-                return null;
-              }).thenRun(Util.NOOP);
+              });
         })
         .toList();
 
@@ -278,6 +272,7 @@ public class KeysController {
     if (responseItems.isEmpty()) {
       throw new WebApplicationException(Response.Status.NOT_FOUND);
     }
+
     return new PreKeyResponse(identityKey, responseItems);
   }
 

service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcHelper.java

@@ -41,41 +41,37 @@ class KeysGrpcHelper {
     return devices
         .filter(Device::isEnabled)
         .switchIfEmpty(Mono.error(Status.NOT_FOUND.asException()))
-        .flatMap(device -> {
-          final ECSignedPreKey ecSignedPreKey = device.getSignedPreKey(identityType);
+        .flatMap(device -> Flux.merge(
+                Mono.fromFuture(() -> keysManager.takeEC(targetAccount.getIdentifier(identityType), device.getId())),
+                Mono.fromFuture(() -> keysManager.getEcSignedPreKey(targetAccount.getIdentifier(identityType), device.getId())),
+                Mono.fromFuture(() -> keysManager.takePQ(targetAccount.getIdentifier(identityType), device.getId())))
+            .flatMap(Mono::justOrEmpty)
+            .reduce(GetPreKeysResponse.PreKeyBundle.newBuilder(), (builder, preKey) -> {
+              if (preKey instanceof ECPreKey ecPreKey) {
+                builder.setEcOneTimePreKey(EcPreKey.newBuilder()
+                    .setKeyId(ecPreKey.keyId())
+                    .setPublicKey(ByteString.copyFrom(ecPreKey.serializedPublicKey()))
+                    .build());
+              } else if (preKey instanceof ECSignedPreKey ecSignedPreKey) {
+                builder.setEcSignedPreKey(EcSignedPreKey.newBuilder()
+                    .setKeyId(ecSignedPreKey.keyId())
+                    .setPublicKey(ByteString.copyFrom(ecSignedPreKey.serializedPublicKey()))
+                    .setSignature(ByteString.copyFrom(ecSignedPreKey.signature()))
+                    .build());
+              } else if (preKey instanceof KEMSignedPreKey kemSignedPreKey) {
+                builder.setKemOneTimePreKey(KemSignedPreKey.newBuilder()
+                    .setKeyId(kemSignedPreKey.keyId())
+                    .setPublicKey(ByteString.copyFrom(kemSignedPreKey.serializedPublicKey()))
+                    .setSignature(ByteString.copyFrom(kemSignedPreKey.signature()))
+                    .build());
+              } else {
+                throw new AssertionError("Unexpected pre-key type: " + preKey.getClass());
+              }
 
-          final GetPreKeysResponse.PreKeyBundle.Builder preKeyBundleBuilder = GetPreKeysResponse.PreKeyBundle.newBuilder()
-              .setEcSignedPreKey(EcSignedPreKey.newBuilder()
-                  .setKeyId(ecSignedPreKey.keyId())
-                  .setPublicKey(ByteString.copyFrom(ecSignedPreKey.serializedPublicKey()))
-                  .setSignature(ByteString.copyFrom(ecSignedPreKey.signature()))
-                  .build());
-
-          return Flux.merge(
-                  Mono.fromFuture(() -> keysManager.takeEC(targetAccount.getIdentifier(identityType), device.getId())),
-                  Mono.fromFuture(() -> keysManager.takePQ(targetAccount.getIdentifier(identityType), device.getId())))
-              .flatMap(Mono::justOrEmpty)
-              .reduce(preKeyBundleBuilder, (builder, preKey) -> {
-                if (preKey instanceof ECPreKey ecPreKey) {
-                  builder.setEcOneTimePreKey(EcPreKey.newBuilder()
-                      .setKeyId(ecPreKey.keyId())
-                      .setPublicKey(ByteString.copyFrom(ecPreKey.serializedPublicKey()))
-                      .build());
-                } else if (preKey instanceof KEMSignedPreKey kemSignedPreKey) {
-                  preKeyBundleBuilder.setKemOneTimePreKey(KemSignedPreKey.newBuilder()
-                      .setKeyId(kemSignedPreKey.keyId())
-                      .setPublicKey(ByteString.copyFrom(kemSignedPreKey.serializedPublicKey()))
-                      .setSignature(ByteString.copyFrom(kemSignedPreKey.signature()))
-                      .build());
-                } else {
-                  throw new AssertionError("Unexpected pre-key type: " + preKey.getClass());
-                }
-
-                return builder;
-              })
-              // Cast device IDs to `int` to match data types in the response object’s protobuf definition
-              .map(builder -> Tuples.of((int) device.getId(), builder.build()));
+              return builder;
             })
+            // Cast device IDs to `int` to match data types in the response object’s protobuf definition
+            .map(builder -> Tuples.of((int) device.getId(), builder.build())))
         .collectMap(Tuple2::getT1, Tuple2::getT2)
         .map(preKeyBundles -> GetPreKeysResponse.newBuilder()
             .setIdentityKey(ByteString.copyFrom(targetAccount.getIdentityKey(identityType).serialize()))

service/src/main/java/org/whispersystems/textsecuregcm/storage/Device.java

@@ -10,6 +10,7 @@ import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import java.util.List;
 import java.util.OptionalInt;
+import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -238,6 +239,10 @@ public class Device {
     this.phoneNumberIdentityRegistrationId = phoneNumberIdentityRegistrationId;
   }
 
+  /**
+   * @deprecated Please retrieve signed pre-keys via {@link KeysManager#getEcSignedPreKey(UUID, byte)} instead
+   */
+  @Deprecated
   public ECSignedPreKey getSignedPreKey(final IdentityType identityType) {
     return switch (identityType) {
       case ACI -> signedPreKey;