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

This reverts commit 3f9edfe597.

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

@@ -38,8 +38,6 @@ import javax.ws.rs.WebApplicationException;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
 import org.signal.libsignal.protocol.IdentityKey;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.Anonymous;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
 import org.whispersystems.textsecuregcm.auth.OptionalAccess;
@@ -51,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;
@@ -74,14 +71,11 @@ 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 String GET_KEYS_COUNTER_NAME = MetricsUtil.name(KeysController.class, "getKeys");
 
   private static final CompletableFuture<?>[] EMPTY_FUTURE_ARRAY = new CompletableFuture[0];
 
-  private static final Logger logger = LoggerFactory.getLogger(KeysController.class);
-
   public KeysController(RateLimiters rateLimiters, KeysManager keys, AccountsManager accounts) {
     this.rateLimiters = rateLimiters;
     this.keys = keys;
@@ -247,35 +241,13 @@ public class KeysController {
         .increment();
 
     final List<PreKeyResponseItem> responseItems = Flux.fromIterable(parseDeviceId(deviceId, target))
-        .flatMap(device -> {
-          final ECSignedPreKey ecSignedPreKey = device.getSignedPreKey(targetIdentifier.identityType());
-          final CompletableFuture<Optional<ECSignedPreKey>> signedPreKeyFuture =
-              keys.getEcSignedPreKey(targetIdentifier.uuid(), device.getId());
-
-          compareSignedEcPreKeysExperiment.compareFutureResult(Optional.ofNullable(ecSignedPreKey), signedPreKeyFuture);
-
-          signedPreKeyFuture.whenComplete((maybeSignedPreKey, throwable) -> {
-            if (throwable == null) {
-              if (!Optional.ofNullable(ecSignedPreKey).equals(maybeSignedPreKey)) {
-                logger.warn("Signed pre-keys do not match for {}, device {}. From device: {}; from table: {}",
-                    targetIdentifier,
-                    deviceId,
-                    Optional.ofNullable(ecSignedPreKey).map(ECSignedPreKey::keyId),
-                    maybeSignedPreKey.map(ECSignedPreKey::keyId));
-              }
-            } else {
-              logger.error("Failed to get signed pre-key for {}, device {}", targetIdentifier, deviceId, throwable);
-            }
-          });
-
-          return Mono.zip(
-              Mono.just(device),
-              Mono.just(Optional.ofNullable(ecSignedPreKey)),
-              Mono.fromFuture(() -> keys.takeEC(targetIdentifier.uuid(), device.getId())),
-              Mono.fromFuture(() -> returnPqKey ? keys.takePQ(targetIdentifier.uuid(), device.getId())
-                  : CompletableFuture.<Optional<KEMSignedPreKey>>completedFuture(Optional.empty()))
-          );
-        }).filter(keys -> keys.getT2().isPresent() || keys.getT3().isPresent() || keys.getT4().isPresent())
+        .flatMap(device -> Mono.zip(
+            Mono.just(device),
+            Mono.fromFuture(() -> keys.getEcSignedPreKey(targetIdentifier.uuid(), device.getId())),
+            Mono.fromFuture(() -> keys.takeEC(targetIdentifier.uuid(), device.getId())),
+            Mono.fromFuture(() -> returnPqKey ? keys.takePQ(targetIdentifier.uuid(), device.getId())
+                : CompletableFuture.<Optional<KEMSignedPreKey>>completedFuture(Optional.empty()))
+        )).filter(keys -> keys.getT2().isPresent() || keys.getT3().isPresent() || keys.getT4().isPresent())
         .map(deviceAndKeys -> {
           final Device device = deviceAndKeys.getT1();
           final int registrationId = switch (targetIdentifier.identityType()) {
@@ -296,6 +268,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

@@ -11,6 +11,7 @@ import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import java.time.Duration;
 import java.util.List;
 import java.util.OptionalInt;
+import java.util.UUID;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import javax.annotation.Nullable;
@@ -246,6 +247,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;