Always elide devices missing required PreKeys

2026-04-19 10:38:03 +01:00 · 2025-08-01 12:01:35 -05:00
parent e8a1854c5e
commit f6a87b2ec0
12 changed files with 302 additions and 473 deletions
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
@@ -34,7 +34,6 @@ import jakarta.ws.rs.QueryParam;
 import jakarta.ws.rs.WebApplicationException;
 import jakarta.ws.rs.core.MediaType;
 import jakarta.ws.rs.core.Response;
-import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
@@ -57,7 +56,6 @@ import org.whispersystems.textsecuregcm.auth.AuthenticatedDevice;
 import org.whispersystems.textsecuregcm.auth.GroupSendTokenHeader;
 import org.whispersystems.textsecuregcm.auth.OptionalAccess;
 import org.whispersystems.textsecuregcm.entities.CheckKeysRequest;
-import org.whispersystems.textsecuregcm.entities.ECPreKey;
 import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.PreKeyCount;
@@ -92,7 +90,6 @@ public class KeysController {
  private final ServerSecretParams serverSecretParams;
  private final Clock clock;

-  private static final String GET_KEYS_COUNTER_NAME = MetricsUtil.name(KeysController.class, "getKeys");
  private static final String STORE_KEYS_COUNTER_NAME = MetricsUtil.name(KeysController.class, "storeKeys");
  private static final String STORE_KEY_BUNDLE_SIZE_DISTRIBUTION_NAME =
      MetricsUtil.name(KeysController.class, "storeKeyBundleSize");
@@ -395,51 +392,14 @@ public class KeysController {

    final List<Device> devices = parseDeviceId(deviceId, target);

-    final List<PreKeyResponseItem> responseItems = Flux.fromIterable(devices)
-        .flatMap(device -> Mono.zip(
-            Mono.just(device),
-            Mono.fromFuture(keysManager.takeEC(targetIdentifier.uuid(), device.getId())),
-            Mono.fromFuture(keysManager.getEcSignedPreKey(targetIdentifier.uuid(), device.getId())),
-            Mono.fromFuture(keysManager.takePQ(targetIdentifier.uuid(), device.getId()))))
-        .flatMap(deviceAndPreKeys -> {
-          final Device device = deviceAndPreKeys.getT1();
-          final KEMSignedPreKey pqPreKey = deviceAndPreKeys.getT4().orElse(null);
-          final ECPreKey unsignedEcPreKey = deviceAndPreKeys.getT2().orElse(null);
-          final ECSignedPreKey signedEcPreKey = deviceAndPreKeys.getT3().orElse(null);
-          final int registrationId = device.getRegistrationId(targetIdentifier.identityType());
-
-          Metrics.counter(GET_KEYS_COUNTER_NAME, Tags.of(
-                  UserAgentTagUtil.getPlatformTag(userAgent),
-                  Tag.of(IDENTITY_TYPE_TAG_NAME, targetIdentifier.identityType().name()),
-                  Tag.of("oneTimeEcKeyAvailable", String.valueOf(unsignedEcPreKey != null)),
-                  Tag.of("signedEcKeyAvailable", String.valueOf(signedEcPreKey != null)),
-                  Tag.of("pqKeyAvailable", String.valueOf(pqPreKey != null))))
-              .increment();
-
-          if (pqPreKey == null) {
-            // The PQ prekey should never be null. This should only happen if the account or device has been
-            // removed.
-            return Mono.fromCompletionStage(() -> accounts.getByServiceIdentifierAsync(targetIdentifier))
-                .flatMap(maybeAccount -> maybeAccount
-                    .flatMap(rereadAccount -> rereadAccount.getDevice(device.getId()))
-                    .filter(rereadDevice ->
-                        registrationId == rereadDevice.getRegistrationId(targetIdentifier.identityType()))
-                    .map(rereadDevice -> {
-                      // The account and device still exist, and the device we originally read matches the current
-                      // registrationId, so the lastResort key should have existed
-                      log.error(
-                          "Target {}, Account {}, DeviceId {}, RegistrationId {} was missing a last resort prekey",
-                          targetIdentifier,
-                          target.getIdentifier(IdentityType.ACI),
-                          rereadDevice.getId(),
-                          rereadDevice.getRegistrationId(targetIdentifier.identityType()));
-                      return Mono.<PreKeyResponseItem>error(new IOException("Device missing last resort prekey"));
-                    })
-                    .orElse(Mono.empty()));
-          }
-          return Mono.just(new PreKeyResponseItem(
-              device.getId(), registrationId, signedEcPreKey, unsignedEcPreKey, pqPreKey));
-        })
+    final List<PreKeyResponseItem> responseItems = Flux.fromIterable(devices).flatMap(device -> Mono
+            .fromCompletionStage(keysManager.takeDevicePreKeys(device.getId(), targetIdentifier, userAgent))
+            .flatMap(Mono::justOrEmpty)
+            .map(devicePreKeys -> new PreKeyResponseItem(
+                device.getId(), device.getRegistrationId(targetIdentifier.identityType()),
+                devicePreKeys.ecSignedPreKey(),
+                devicePreKeys.ecPreKey().orElse(null),
+                devicePreKeys.kemSignedPreKey())))
        .collectList()
        .block();

--- a/service/src/main/java/org/whispersystems/textsecuregcm/entities/PreKeyResponseItem.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/entities/PreKeyResponseItem.java
@@ -19,7 +19,7 @@ public class PreKeyResponseItem {
  private int registrationId;

  @JsonProperty
-  @Schema(description="the signed elliptic-curve prekey for the device, if one has been set")
+  @Schema(description="the signed elliptic-curve prekey for the device")
  private ECSignedPreKey signedPreKey;

  @JsonProperty
@@ -28,7 +28,7 @@ public class PreKeyResponseItem {

  @JsonProperty
  @Schema(description="a signed post-quantum prekey for the device " +
-      "(a one-time prekey if any remain, otherwise the last-resort prekey if one has been set)")
+      "(a one-time prekey if any remain, otherwise the last-resort prekey)")
  private KEMSignedPreKey pqPreKey;

  public PreKeyResponseItem() {}
--- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysAnonymousGrpcService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysAnonymousGrpcService.java
@@ -56,7 +56,7 @@ public class KeysAnonymousGrpcService extends ReactorKeysAnonymousGrpc.KeysAnony
          groupSendTokenUtil.checkGroupSendToken(request.getGroupSendToken(), serviceIdentifier);

          yield lookUpAccount(serviceIdentifier, Status.NOT_FOUND)
-              .flatMap(targetAccount -> KeysGrpcHelper.getPreKeys(targetAccount, serviceIdentifier.identityType(), deviceId, keysManager));
+              .flatMap(targetAccount -> KeysGrpcHelper.getPreKeys(targetAccount, serviceIdentifier, deviceId, keysManager));
        } catch (final StatusException e) {
          yield Mono.error(e);
        }
@@ -66,7 +66,7 @@ public class KeysAnonymousGrpcService extends ReactorKeysAnonymousGrpc.KeysAnony
          lookUpAccount(serviceIdentifier, Status.UNAUTHENTICATED)
              .flatMap(targetAccount ->
                  UnidentifiedAccessUtil.checkUnidentifiedAccess(targetAccount, request.getUnidentifiedAccessKey().toByteArray())
-                  ? KeysGrpcHelper.getPreKeys(targetAccount, serviceIdentifier.identityType(), deviceId, keysManager)
+                  ? KeysGrpcHelper.getPreKeys(targetAccount, serviceIdentifier, deviceId, keysManager)
                  : Mono.error(Status.UNAUTHENTICATED.asException()));

      default -> Mono.error(Status.INVALID_ARGUMENT.asException());
--- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcHelper.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcHelper.java
@@ -11,10 +11,7 @@ import org.signal.chat.common.EcPreKey;
 import org.signal.chat.common.EcSignedPreKey;
 import org.signal.chat.common.KemSignedPreKey;
 import org.signal.chat.keys.GetPreKeysResponse;
-import org.whispersystems.textsecuregcm.entities.ECPreKey;
-import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
-import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
-import org.whispersystems.textsecuregcm.identity.IdentityType;
+import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.KeysManager;
@@ -28,7 +25,7 @@ class KeysGrpcHelper {
  static final byte ALL_DEVICES = 0;

  static Mono<GetPreKeysResponse> getPreKeys(final Account targetAccount,
-      final IdentityType identityType,
+      final ServiceIdentifier targetServiceIdentifier,
      final byte targetDeviceId,
      final KeysManager keysManager) {

@@ -36,42 +33,37 @@ class KeysGrpcHelper {
        ? Flux.fromIterable(targetAccount.getDevices())
        : Flux.from(Mono.justOrEmpty(targetAccount.getDevice(targetDeviceId)));

+    final String userAgent = RequestAttributesUtil.getUserAgent().orElse(null);
    return devices
-        .switchIfEmpty(Mono.error(Status.NOT_FOUND.asException()))
-        .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(device -> Mono
+            .fromFuture(keysManager.takeDevicePreKeys(device.getId(), targetServiceIdentifier, userAgent))
            .flatMap(Mono::justOrEmpty)
-            .reduce(GetPreKeysResponse.PreKeyBundle.newBuilder(), (builder, preKey) -> {
-              if (preKey instanceof ECPreKey ecPreKey) {
-                builder.setEcOneTimePreKey(EcPreKey.newBuilder()
+            .map(devicePreKeys -> {
+              final GetPreKeysResponse.PreKeyBundle.Builder builder = GetPreKeysResponse.PreKeyBundle.newBuilder()
+                  .setEcSignedPreKey(EcSignedPreKey.newBuilder()
+                      .setKeyId(devicePreKeys.ecSignedPreKey().keyId())
+                      .setPublicKey(ByteString.copyFrom(devicePreKeys.ecSignedPreKey().serializedPublicKey()))
+                      .setSignature(ByteString.copyFrom(devicePreKeys.ecSignedPreKey().signature()))
+                      .build())
+                  .setKemOneTimePreKey(KemSignedPreKey.newBuilder()
+                      .setKeyId(devicePreKeys.kemSignedPreKey().keyId())
+                      .setPublicKey(ByteString.copyFrom(devicePreKeys.kemSignedPreKey().serializedPublicKey()))
+                      .setSignature(ByteString.copyFrom(devicePreKeys.kemSignedPreKey().signature()))
+                      .build());
+              devicePreKeys.ecPreKey().ifPresent(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());
-              }
-
-              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())))
+                    .build()));
+              // Cast device IDs to `int` to match data types in the response object’s protobuf definition
+              return Tuples.of((int) device.getId(), builder.build());
+            }))
+        // If there were no devices with valid prekey bundles in the account, the account is gone
+        .switchIfEmpty(Mono.error(Status.NOT_FOUND.asException()))
        .collectMap(Tuple2::getT1, Tuple2::getT2)
        .map(preKeyBundles -> GetPreKeysResponse.newBuilder()
-            .setIdentityKey(ByteString.copyFrom(targetAccount.getIdentityKey(identityType).serialize()))
+            .setIdentityKey(ByteString
+                .copyFrom(targetAccount.getIdentityKey(targetServiceIdentifier.identityType())
+                .serialize()))
            .putAllPreKeys(preKeyBundles)
            .build());
  }
--- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcService.java
@@ -136,7 +136,7 @@ public class KeysGrpcService extends ReactorKeysGrpc.KeysImplBase {
            .flatMap(Mono::justOrEmpty))
        .switchIfEmpty(Mono.error(Status.NOT_FOUND.asException()))
        .flatMap(targetAccount ->
-            KeysGrpcHelper.getPreKeys(targetAccount, targetIdentifier.identityType(), deviceId, keysManager));
+            KeysGrpcHelper.getPreKeys(targetAccount, targetIdentifier, deviceId, keysManager));
  }

  @Override
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java
@@ -5,20 +5,31 @@

 package org.whispersystems.textsecuregcm.storage;

+import com.google.common.annotations.VisibleForTesting;
 import io.micrometer.core.instrument.Metrics;
 import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import io.micrometer.core.instrument.Tag;
+import io.micrometer.core.instrument.Tags;
+import org.whispersystems.textsecuregcm.controllers.KeysController;
 import org.whispersystems.textsecuregcm.entities.ECPreKey;
 import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
+import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
 import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
+import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
+import org.whispersystems.textsecuregcm.util.Futures;
+import org.whispersystems.textsecuregcm.util.Optionals;
 import reactor.core.publisher.Flux;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
+import javax.annotation.Nullable;

 public class KeysManager {
+  // KeysController for backwards compatibility
+  private static final String GET_KEYS_COUNTER_NAME = MetricsUtil.name(KeysController.class, "getKeys");

  private final SingleUseECPreKeyStore ecPreKeys;
  private final SingleUseKEMPreKeyStore pqPreKeys;
@@ -115,11 +126,13 @@ public class KeysManager {

  }

-  public CompletableFuture<Optional<ECPreKey>> takeEC(final UUID identifier, final byte deviceId) {
+  @VisibleForTesting
+  CompletableFuture<Optional<ECPreKey>> takeEC(final UUID identifier, final byte deviceId) {
    return ecPreKeys.take(identifier, deviceId);
  }

-  public CompletableFuture<Optional<KEMSignedPreKey>> takePQ(final UUID identifier, final byte deviceId) {
+  @VisibleForTesting
+  CompletableFuture<Optional<KEMSignedPreKey>> takePQ(final UUID identifier, final byte deviceId) {
    final boolean enrolledInPagedKeys = experimentEnrollmentManager.isEnrolled(identifier, PAGED_KEYS_EXPERIMENT_NAME);
    return tagTakePQ(pagedPqPreKeys.take(identifier, deviceId), PQSource.PAGE, enrolledInPagedKeys)
        .thenCompose(maybeSingleUsePreKey -> maybeSingleUsePreKey
@@ -209,4 +222,36 @@ public class KeysManager {
  public CompletableFuture<Void> pruneDeadPage(final UUID identifier, final byte deviceId, final UUID pageId) {
    return pagedPqPreKeys.deleteBundleFromS3(identifier, deviceId, pageId);
  }
+
+  public record DevicePreKeys(
+      ECSignedPreKey ecSignedPreKey,
+      Optional<ECPreKey> ecPreKey,
+      KEMSignedPreKey kemSignedPreKey) {}
+
+  public CompletableFuture<Optional<DevicePreKeys>> takeDevicePreKeys(
+      final byte deviceId,
+      final ServiceIdentifier serviceIdentifier,
+      final @Nullable String userAgent) {
+    final UUID uuid = serviceIdentifier.uuid();
+    return Futures.zipWith(
+            this.takeEC(uuid, deviceId),
+            this.getEcSignedPreKey(uuid, deviceId),
+            this.takePQ(uuid, deviceId),
+            (maybeUnsignedEcPreKey, maybeSignedEcPreKey, maybePqPreKey) -> {
+
+              Metrics.counter(GET_KEYS_COUNTER_NAME, Tags.of(
+                      UserAgentTagUtil.getPlatformTag(userAgent),
+                      Tag.of("identityType", serviceIdentifier.identityType().name()),
+                      Tag.of("oneTimeEcKeyAvailable", String.valueOf(maybeUnsignedEcPreKey.isPresent())),
+                      Tag.of("signedEcKeyAvailable", String.valueOf(maybeSignedEcPreKey.isPresent())),
+                      Tag.of("pqKeyAvailable", String.valueOf(maybePqPreKey.isPresent()))))
+                  .increment();
+
+              // The pq prekey and signed EC prekey should never be null for an existing account. This should only happen
+              // if the account or device has been removed and the read was split, so we can return empty in those cases.
+              return Optionals.zipWith(maybeSignedEcPreKey, maybePqPreKey, (signedEcPreKey, pqPreKey) ->
+                  new DevicePreKeys(signedEcPreKey, maybeUnsignedEcPreKey, pqPreKey));
+            })
+        .toCompletableFuture();
+  }
 }
--- a/service/src/main/java/org/whispersystems/textsecuregcm/util/Futures.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/util/Futures.java
@@ -0,0 +1,21 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.util;
+
+import java.util.concurrent.CompletionStage;
+import org.apache.commons.lang3.function.TriFunction;
+
+public class Futures {
+
+  public static <T, U, V, R> CompletionStage<R> zipWith(
+      CompletionStage<T> futureT,
+      CompletionStage<U> futureU,
+      CompletionStage<V> futureV,
+      TriFunction<T, U, V, R> fun) {
+
+    return futureT.thenCompose(t -> futureU.thenCombine(futureV, (u, v) -> fun.apply(t, u, v)));
+  }
+}