Add a gRPC service for working with pre-keys

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -24,6 +24,7 @@ import io.dropwizard.auth.basic.BasicCredentials;
 import io.dropwizard.setup.Bootstrap;
 import io.dropwizard.setup.Environment;
 import io.grpc.ServerBuilder;
+import io.grpc.ServerInterceptors;
 import io.lettuce.core.metrics.MicrometerCommandLatencyRecorder;
 import io.lettuce.core.metrics.MicrometerOptions;
 import io.lettuce.core.resource.ClientResources;
@@ -64,6 +65,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AccountAuthenticator;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
+import org.whispersystems.textsecuregcm.auth.BaseAccountAuthenticator;
 import org.whispersystems.textsecuregcm.auth.CertificateGenerator;
 import org.whispersystems.textsecuregcm.auth.DisabledPermittedAccountAuthenticator;
 import org.whispersystems.textsecuregcm.auth.DisabledPermittedAuthenticatedAccount;
@@ -72,6 +74,7 @@ import org.whispersystems.textsecuregcm.auth.PhoneVerificationTokenManager;
 import org.whispersystems.textsecuregcm.auth.RegistrationLockVerificationManager;
 import org.whispersystems.textsecuregcm.auth.TurnTokenGenerator;
 import org.whispersystems.textsecuregcm.auth.WebsocketRefreshApplicationEventListener;
+import org.whispersystems.textsecuregcm.auth.grpc.BasicCredentialAuthenticationInterceptor;
 import org.whispersystems.textsecuregcm.badges.ConfiguredProfileBadgeConverter;
 import org.whispersystems.textsecuregcm.badges.ResourceBundleLevelTranslator;
 import org.whispersystems.textsecuregcm.captcha.CaptchaChecker;
@@ -115,6 +118,8 @@ import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
 import org.whispersystems.textsecuregcm.filters.RemoteDeprecationFilter;
 import org.whispersystems.textsecuregcm.filters.RequestStatisticsFilter;
 import org.whispersystems.textsecuregcm.filters.TimestampResponseFilter;
+import org.whispersystems.textsecuregcm.grpc.KeysGrpcService;
+import org.whispersystems.textsecuregcm.grpc.KeysAnonymousGrpcService;
 import org.whispersystems.textsecuregcm.limits.PushChallengeManager;
 import org.whispersystems.textsecuregcm.limits.RateLimitChallengeManager;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
@@ -401,6 +406,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
                 .build(),
             MetricsUtil.name(getClass(), "messageDeliveryExecutor"), MetricsUtil.PREFIX),
         "messageDelivery");
+
     // TODO: generally speaking this is a DynamoDB I/O executor for the accounts table; we should eventually have a general executor for speaking to the accounts table, but most of the server is still synchronous so this isn't widely useful yet
     ExecutorService batchIdentityCheckExecutor = environment.lifecycle().executorService(name(getClass(), "batchIdentityCheck-%d")).minThreads(32).maxThreads(32).build();
     ExecutorService multiRecipientMessageExecutor = environment.lifecycle()
@@ -606,8 +612,14 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     AuthFilter<BasicCredentials, DisabledPermittedAuthenticatedAccount> disabledPermittedAccountAuthFilter = new BasicCredentialAuthFilter.Builder<DisabledPermittedAuthenticatedAccount>().setAuthenticator(
         disabledPermittedAccountAuthenticator).buildAuthFilter();
 
+    final BasicCredentialAuthenticationInterceptor basicCredentialAuthenticationInterceptor =
+        new BasicCredentialAuthenticationInterceptor(new BaseAccountAuthenticator(accountsManager));
+
     final ServerBuilder<?> grpcServer = ServerBuilder.forPort(config.getGrpcPort())
-        .intercept(new MetricCollectingServerInterceptor(Metrics.globalRegistry)); /* TODO: specialize metrics with user-agent platform */
+        // TODO: specialize metrics with user-agent platform
+        .intercept(new MetricCollectingServerInterceptor(Metrics.globalRegistry))
+        .addService(ServerInterceptors.intercept(new KeysGrpcService(accountsManager, keys, rateLimiters), basicCredentialAuthenticationInterceptor))
+        .addService(new KeysAnonymousGrpcService(accountsManager, keys));
 
     RemoteDeprecationFilter remoteDeprecationFilter = new RemoteDeprecationFilter(dynamicConfigurationManager);
     environment.servlets()

service/src/main/java/org/whispersystems/textsecuregcm/auth/UnidentifiedAccessUtil.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.auth;
+
+import org.whispersystems.textsecuregcm.storage.Account;
+import java.security.MessageDigest;
+
+public class UnidentifiedAccessUtil {
+
+  private UnidentifiedAccessUtil() {
+  }
+
+  /**
+   * Checks whether an action (e.g. sending a message or retrieving pre-keys) may be taken on the target account by an
+   * actor presenting the given unidentified access key.
+   *
+   * @param targetAccount the account on which an actor wishes to take an action
+   * @param unidentifiedAccessKey the unidentified access key presented by the actor
+   *
+   * @return {@code true} if an actor presenting the given unidentified access key has permission to take an action on
+   * the target account or {@code false} otherwise
+   */
+  public static boolean checkUnidentifiedAccess(final Account targetAccount, final byte[] unidentifiedAccessKey) {
+    return targetAccount.isUnrestrictedUnidentifiedAccess()
+        || targetAccount.getUnidentifiedAccessKey()
+        .map(targetUnidentifiedAccessKey -> MessageDigest.isEqual(targetUnidentifiedAccessKey, unidentifiedAccessKey))
+        .orElse(false);
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/auth/grpc/AuthenticatedDevice.java

@@ -0,0 +1,11 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.auth.grpc;
+
+import java.util.UUID;
+
+public record AuthenticatedDevice(UUID accountIdentifier, long deviceId) {
+}

service/src/main/java/org/whispersystems/textsecuregcm/auth/grpc/AuthenticationUtil.java

@@ -6,8 +6,12 @@
 package org.whispersystems.textsecuregcm.auth.grpc;
 
 import io.grpc.Context;
-import java.util.Optional;
+import io.grpc.Status;
 import java.util.UUID;
+import javax.annotation.Nullable;
+import reactor.core.publisher.Mono;
+import reactor.util.function.Tuple2;
+import reactor.util.function.Tuples;
 
 /**
  * Provides utility methods for working with authentication in the context of gRPC calls.
@@ -17,11 +21,23 @@ public class AuthenticationUtil {
   static final Context.Key<UUID> CONTEXT_AUTHENTICATED_ACCOUNT_IDENTIFIER_KEY = Context.key("authenticated-aci");
   static final Context.Key<Long> CONTEXT_AUTHENTICATED_DEVICE_IDENTIFIER_KEY = Context.key("authenticated-device-id");
 
-  public static Optional<UUID> getAuthenticatedAccountIdentifier() {
-    return Optional.ofNullable(CONTEXT_AUTHENTICATED_ACCOUNT_IDENTIFIER_KEY.get());
-  }
+  /**
+   * Returns the account/device authenticated in the current gRPC context or throws an "unauthenticated" exception if
+   * no authenticated account/device is available.
+   *
+   * @return the account/device authenticated in the current gRPC context
+   *
+   * @throws io.grpc.StatusRuntimeException with a status of {@code UNAUTHENTICATED} if no authenticated account/device
+   * could be retrieved from the current gRPC context
+   */
+  public static AuthenticatedDevice requireAuthenticatedDevice() {
+    @Nullable final UUID accountIdentifier = CONTEXT_AUTHENTICATED_ACCOUNT_IDENTIFIER_KEY.get();
+    @Nullable final Long deviceId = CONTEXT_AUTHENTICATED_DEVICE_IDENTIFIER_KEY.get();
 
-  public static Optional<Long> getAuthenticatedDeviceIdentifier() {
-    return Optional.ofNullable(CONTEXT_AUTHENTICATED_DEVICE_IDENTIFIER_KEY.get());
+    if (accountIdentifier != null && deviceId != null) {
+      return new AuthenticatedDevice(accountIdentifier, deviceId);
+    }
+
+    throw Status.UNAUTHENTICATED.asRuntimeException();
   }
 }

service/src/main/java/org/whispersystems/textsecuregcm/auth/grpc/BasicCredentialAuthenticationInterceptor.java

@@ -24,9 +24,8 @@ import org.whispersystems.textsecuregcm.auth.BaseAccountAuthenticator;
  * Callers supply credentials by providing a username (UUID and optional device ID) and password pair in the
  * {@code x-signal-basic-auth-credentials} call header.
  * <p/>
- * Downstream services can retrieve the identity of the authenticated caller using
- * {@link AuthenticationUtil#getAuthenticatedAccountIdentifier()} and
- * {@link AuthenticationUtil#getAuthenticatedDeviceIdentifier()}.
+ * Downstream services can retrieve the identity of the authenticated caller using methods in
+ * {@link AuthenticationUtil}.
  * <p/>
  * Note that this authentication, while fully functional, is intended only for development and testing purposes and is
  * intended to be replaced with a more robust and efficient strategy before widespread client adoption.

service/src/main/java/org/whispersystems/textsecuregcm/auth/grpc/NotAuthenticatedException.java

@@ -1,13 +0,0 @@
-/*
- * Copyright 2023 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.whispersystems.textsecuregcm.auth.grpc;
-
-/**
- * Indicates that a caller tried to get information about the authenticated gRPC caller, but no caller has been
- * authenticated.
- */
-public class NotAuthenticatedException extends Exception {
-}

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

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.grpc;
+
+import io.grpc.Status;
+
+public enum IdentityType {
+  ACI,
+  PNI;
+
+  public static IdentityType fromGrpcIdentityType(final org.signal.chat.common.IdentityType grpcIdentityType) {
+    return switch (grpcIdentityType) {
+      case IDENTITY_TYPE_ACI -> ACI;
+      case IDENTITY_TYPE_PNI -> PNI;
+      case IDENTITY_TYPE_UNSPECIFIED, UNRECOGNIZED -> throw Status.INVALID_ARGUMENT.asRuntimeException();
+    };
+  }
+}

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

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.grpc;
+
+import io.grpc.Status;
+import org.signal.chat.keys.GetPreKeysAnonymousRequest;
+import org.signal.chat.keys.GetPreKeysResponse;
+import org.signal.chat.keys.ReactorKeysAnonymousGrpc;
+import org.whispersystems.textsecuregcm.auth.UnidentifiedAccessUtil;
+import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.KeysManager;
+import reactor.core.publisher.Mono;
+
+public class KeysAnonymousGrpcService extends ReactorKeysAnonymousGrpc.KeysAnonymousImplBase {
+
+  private final AccountsManager accountsManager;
+  private final KeysManager keysManager;
+
+  public KeysAnonymousGrpcService(final AccountsManager accountsManager, final KeysManager keysManager) {
+    this.accountsManager = accountsManager;
+    this.keysManager = keysManager;
+  }
+
+  @Override
+  public Mono<GetPreKeysResponse> getPreKeys(final GetPreKeysAnonymousRequest request) {
+    return KeysGrpcHelper.findAccount(request.getTargetIdentifier(), accountsManager)
+        .switchIfEmpty(Mono.error(Status.UNAUTHENTICATED.asException()))
+        .flatMap(targetAccount -> {
+          final IdentityType identityType =
+              IdentityType.fromGrpcIdentityType(request.getTargetIdentifier().getIdentityType());
+
+          return UnidentifiedAccessUtil.checkUnidentifiedAccess(targetAccount, request.getUnidentifiedAccessKey().toByteArray())
+              ? KeysGrpcHelper.getPreKeys(targetAccount, identityType, request.getDeviceId(), keysManager)
+              : Mono.error(Status.UNAUTHENTICATED.asException());
+        });
+  }
+}

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

@@ -0,0 +1,107 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.grpc;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.protobuf.ByteString;
+import io.grpc.Status;
+import java.util.UUID;
+import org.signal.chat.common.EcPreKey;
+import org.signal.chat.common.EcSignedPreKey;
+import org.signal.chat.common.KemSignedPreKey;
+import org.signal.chat.common.ServiceIdentifier;
+import org.signal.chat.keys.GetPreKeysResponse;
+import org.signal.libsignal.protocol.IdentityKey;
+import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.KeysManager;
+import org.whispersystems.textsecuregcm.util.UUIDUtil;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.util.function.Tuple2;
+import reactor.util.function.Tuples;
+
+class KeysGrpcHelper {
+
+  @VisibleForTesting
+  static final long ALL_DEVICES = 0;
+
+  static Mono<Account> findAccount(final ServiceIdentifier targetIdentifier, final AccountsManager accountsManager) {
+
+    return Mono.just(IdentityType.fromGrpcIdentityType(targetIdentifier.getIdentityType()))
+        .flatMap(identityType -> {
+          final UUID uuid = UUIDUtil.fromByteString(targetIdentifier.getUuid());
+
+          return Mono.fromFuture(switch (identityType) {
+            case ACI -> accountsManager.getByAccountIdentifierAsync(uuid);
+            case PNI -> accountsManager.getByPhoneNumberIdentifierAsync(uuid);
+          });
+        })
+        .flatMap(Mono::justOrEmpty)
+        .onErrorMap(IllegalArgumentException.class, throwable -> Status.INVALID_ARGUMENT.asException());
+  }
+
+  static Tuple2<UUID, IdentityKey> getIdentifierAndIdentityKey(final Account account, final IdentityType identityType) {
+    final UUID identifier = switch (identityType) {
+      case ACI -> account.getUuid();
+      case PNI -> account.getPhoneNumberIdentifier();
+    };
+
+    final IdentityKey identityKey = switch (identityType) {
+      case ACI -> account.getIdentityKey();
+      case PNI -> account.getPhoneNumberIdentityKey();
+    };
+
+    return Tuples.of(identifier, identityKey);
+  }
+
+  static Mono<GetPreKeysResponse> getPreKeys(final Account targetAccount, final IdentityType identityType, final long targetDeviceId, final KeysManager keysManager) {
+    final Tuple2<UUID, IdentityKey> identifierAndIdentityKey = getIdentifierAndIdentityKey(targetAccount, identityType);
+
+    final Flux<Device> devices = targetDeviceId == ALL_DEVICES
+        ? Flux.fromIterable(targetAccount.getDevices())
+        : Flux.from(Mono.justOrEmpty(targetAccount.getDevice(targetDeviceId)));
+
+    return devices
+        .filter(Device::isEnabled)
+        .switchIfEmpty(Mono.error(Status.NOT_FOUND.asException()))
+        .flatMap(device -> Mono.zip(Mono.fromFuture(keysManager.takeEC(identifierAndIdentityKey.getT1(), device.getId())),
+                Mono.fromFuture(keysManager.takePQ(identifierAndIdentityKey.getT1(), device.getId())))
+            .map(oneTimePreKeys -> {
+              final ECSignedPreKey ecSignedPreKey = switch (identityType) {
+                case ACI -> device.getSignedPreKey();
+                case PNI -> device.getPhoneNumberIdentitySignedPreKey();
+              };
+
+              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());
+
+              oneTimePreKeys.getT1().ifPresent(ecPreKey -> preKeyBundleBuilder.setEcOneTimePreKey(EcPreKey.newBuilder()
+                  .setKeyId(ecPreKey.keyId())
+                  .setPublicKey(ByteString.copyFrom(ecPreKey.serializedPublicKey()))
+                  .build()));
+
+              oneTimePreKeys.getT2().ifPresent(kemSignedPreKey -> preKeyBundleBuilder.setKemOneTimePreKey(KemSignedPreKey.newBuilder()
+                  .setKeyId(kemSignedPreKey.keyId())
+                  .setPublicKey(ByteString.copyFrom(kemSignedPreKey.serializedPublicKey()))
+                  .setSignature(ByteString.copyFrom(kemSignedPreKey.signature()))
+                  .build()));
+
+              return Tuples.of(device.getId(), preKeyBundleBuilder.build());
+            }))
+        .collectMap(Tuple2::getT1, Tuple2::getT2)
+        .map(preKeyBundles -> GetPreKeysResponse.newBuilder()
+            .setIdentityKey(ByteString.copyFrom(identifierAndIdentityKey.getT2().serialize()))
+            .putAllPreKeys(preKeyBundles)
+            .build());
+  }
+}

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

@@ -0,0 +1,307 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.grpc;
+
+import static org.whispersystems.textsecuregcm.grpc.IdentityType.ACI;
+
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.BiFunction;
+import java.util.function.Consumer;
+import org.signal.chat.common.EcPreKey;
+import org.signal.chat.common.EcSignedPreKey;
+import org.signal.chat.common.KemSignedPreKey;
+import org.signal.chat.keys.GetPreKeyCountRequest;
+import org.signal.chat.keys.GetPreKeyCountResponse;
+import org.signal.chat.keys.GetPreKeysRequest;
+import org.signal.chat.keys.GetPreKeysResponse;
+import org.signal.chat.keys.ReactorKeysGrpc;
+import org.signal.chat.keys.SetEcSignedPreKeyRequest;
+import org.signal.chat.keys.SetKemLastResortPreKeyRequest;
+import org.signal.chat.keys.SetOneTimeEcPreKeysRequest;
+import org.signal.chat.keys.SetOneTimeKemSignedPreKeysRequest;
+import org.signal.chat.keys.SetPreKeyResponse;
+import org.signal.libsignal.protocol.IdentityKey;
+import org.signal.libsignal.protocol.InvalidKeyException;
+import org.signal.libsignal.protocol.ecc.ECPublicKey;
+import org.signal.libsignal.protocol.kem.KEMPublicKey;
+import org.whispersystems.textsecuregcm.auth.grpc.AuthenticatedDevice;
+import org.whispersystems.textsecuregcm.auth.grpc.AuthenticationUtil;
+import org.whispersystems.textsecuregcm.entities.ECPreKey;
+import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
+import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
+import org.whispersystems.textsecuregcm.limits.RateLimiters;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.KeysManager;
+import org.whispersystems.textsecuregcm.util.UUIDUtil;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.util.function.Tuple2;
+import reactor.util.function.Tuples;
+
+public class KeysGrpcService extends ReactorKeysGrpc.KeysImplBase {
+
+  private final AccountsManager accountsManager;
+  private final KeysManager keysManager;
+  private final RateLimiters rateLimiters;
+
+  private static final StatusRuntimeException INVALID_PUBLIC_KEY_EXCEPTION = Status.fromCode(Status.Code.INVALID_ARGUMENT)
+      .withDescription("Invalid public key")
+      .asRuntimeException();
+
+  private static final StatusRuntimeException INVALID_SIGNATURE_EXCEPTION = Status.fromCode(Status.Code.INVALID_ARGUMENT)
+      .withDescription("Invalid signature")
+      .asRuntimeException();
+
+  private enum PreKeyType {
+    EC,
+    KEM
+  }
+
+  public KeysGrpcService(final AccountsManager accountsManager,
+      final KeysManager keysManager,
+      final RateLimiters rateLimiters) {
+
+    this.accountsManager = accountsManager;
+    this.keysManager = keysManager;
+    this.rateLimiters = rateLimiters;
+  }
+
+  @Override
+  protected Throwable onErrorMap(final Throwable throwable) {
+    return RateLimitUtil.mapRateLimitExceededException(throwable);
+  }
+
+  @Override
+  public Mono<GetPreKeyCountResponse> getPreKeyCount(final GetPreKeyCountRequest request) {
+    return Mono.fromSupplier(AuthenticationUtil::requireAuthenticatedDevice)
+        .flatMap(authenticatedDevice -> Mono.fromFuture(accountsManager.getByAccountIdentifierAsync(authenticatedDevice.accountIdentifier()))
+            .map(maybeAccount -> maybeAccount
+                .map(account -> Tuples.of(account, authenticatedDevice.deviceId()))
+                .orElseThrow(Status.UNAUTHENTICATED::asRuntimeException)))
+        .flatMapMany(accountAndDeviceId -> Flux.just(
+            Tuples.of(ACI, accountAndDeviceId.getT1().getUuid(), accountAndDeviceId.getT2()),
+            Tuples.of(IdentityType.PNI, accountAndDeviceId.getT1().getPhoneNumberIdentifier(), accountAndDeviceId.getT2())
+        ))
+        .flatMap(identityTypeUuidAndDeviceId -> Flux.merge(
+            Mono.fromFuture(keysManager.getEcCount(identityTypeUuidAndDeviceId.getT2(), identityTypeUuidAndDeviceId.getT3()))
+                .map(ecKeyCount -> Tuples.of(identityTypeUuidAndDeviceId.getT1(), PreKeyType.EC, ecKeyCount)),
+
+            Mono.fromFuture(keysManager.getPqCount(identityTypeUuidAndDeviceId.getT2(), identityTypeUuidAndDeviceId.getT3()))
+                .map(ecKeyCount -> Tuples.of(identityTypeUuidAndDeviceId.getT1(), PreKeyType.KEM, ecKeyCount))
+        ))
+        .reduce(GetPreKeyCountResponse.newBuilder(), (builder, tuple) -> {
+          final IdentityType identityType = tuple.getT1();
+          final PreKeyType preKeyType = tuple.getT2();
+          final int count = tuple.getT3();
+
+          switch (identityType) {
+            case ACI -> {
+              switch (preKeyType) {
+                case EC -> builder.setAciEcPreKeyCount(count);
+                case KEM -> builder.setAciKemPreKeyCount(count);
+              }
+            }
+            case PNI -> {
+              switch (preKeyType) {
+                case EC -> builder.setPniEcPreKeyCount(count);
+                case KEM -> builder.setPniKemPreKeyCount(count);
+              }
+            }
+          }
+
+          return builder;
+        })
+        .map(GetPreKeyCountResponse.Builder::build);
+  }
+
+  @Override
+  public Mono<GetPreKeysResponse> getPreKeys(final GetPreKeysRequest request) {
+    final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice();
+
+    final String rateLimitKey;
+    {
+      final UUID targetUuid;
+
+      try {
+        targetUuid = UUIDUtil.fromByteString(request.getTargetIdentifier().getUuid());
+      } catch (final IllegalArgumentException e) {
+        throw Status.INVALID_ARGUMENT.asRuntimeException();
+      }
+
+      rateLimitKey = authenticatedDevice.accountIdentifier() + "." +
+          authenticatedDevice.deviceId() + "__" +
+          targetUuid + "." +
+          request.getDeviceId();
+    }
+
+    return rateLimiters.getPreKeysLimiter().validateReactive(rateLimitKey)
+        .then(KeysGrpcHelper.findAccount(request.getTargetIdentifier(), accountsManager))
+        .switchIfEmpty(Mono.error(Status.NOT_FOUND.asException()))
+        .flatMap(targetAccount -> {
+          final IdentityType identityType =
+              IdentityType.fromGrpcIdentityType(request.getTargetIdentifier().getIdentityType());
+
+          return KeysGrpcHelper.getPreKeys(targetAccount, identityType, request.getDeviceId(), keysManager);
+        });
+  }
+
+  @Override
+  public Mono<SetPreKeyResponse> setOneTimeEcPreKeys(final SetOneTimeEcPreKeysRequest request) {
+    return Mono.fromSupplier(AuthenticationUtil::requireAuthenticatedDevice)
+        .flatMap(authenticatedDevice -> storeOneTimePreKeys(authenticatedDevice.accountIdentifier(),
+            request.getPreKeysList(),
+            IdentityType.fromGrpcIdentityType(request.getIdentityType()),
+            (requestPreKey, ignored) -> checkEcPreKey(requestPreKey),
+            (identifier, preKeys) -> keysManager.storeEcOneTimePreKeys(identifier, authenticatedDevice.deviceId(), preKeys)));
+  }
+
+  @Override
+  public Mono<SetPreKeyResponse> setOneTimeKemSignedPreKeys(final SetOneTimeKemSignedPreKeysRequest request) {
+    return Mono.fromSupplier(AuthenticationUtil::requireAuthenticatedDevice)
+        .flatMap(authenticatedDevice -> storeOneTimePreKeys(authenticatedDevice.accountIdentifier(),
+            request.getPreKeysList(),
+            IdentityType.fromGrpcIdentityType(request.getIdentityType()),
+            KeysGrpcService::checkKemSignedPreKey,
+            (identifier, preKeys) -> keysManager.storeKemOneTimePreKeys(identifier, authenticatedDevice.deviceId(), preKeys)));
+  }
+
+  private <K, R> Mono<SetPreKeyResponse> storeOneTimePreKeys(final UUID authenticatedAccountUuid,
+      final List<R> requestPreKeys,
+      final IdentityType identityType,
+      final BiFunction<R, IdentityKey, K> extractPreKeyFunction,
+      final BiFunction<UUID, List<K>, CompletableFuture<Void>> storeKeysFunction) {
+
+    return Mono.fromFuture(accountsManager.getByAccountIdentifierAsync(authenticatedAccountUuid))
+        .map(maybeAccount -> maybeAccount.orElseThrow(Status.UNAUTHENTICATED::asRuntimeException))
+        .map(account -> {
+          final Tuple2<UUID, IdentityKey> identifierAndIdentityKey =
+              KeysGrpcHelper.getIdentifierAndIdentityKey(account, identityType);
+
+          final List<K> preKeys = requestPreKeys.stream()
+              .map(requestPreKey -> extractPreKeyFunction.apply(requestPreKey, identifierAndIdentityKey.getT2()))
+              .toList();
+
+          if (preKeys.isEmpty()) {
+            throw Status.INVALID_ARGUMENT.asRuntimeException();
+          }
+
+          return Tuples.of(identifierAndIdentityKey.getT1(), preKeys);
+        })
+        .flatMap(identifierAndPreKeys -> Mono.fromFuture(storeKeysFunction.apply(identifierAndPreKeys.getT1(), identifierAndPreKeys.getT2())))
+        .thenReturn(SetPreKeyResponse.newBuilder().build());
+  }
+
+  @Override
+  public Mono<SetPreKeyResponse> setEcSignedPreKey(final SetEcSignedPreKeyRequest request) {
+    return Mono.fromSupplier(AuthenticationUtil::requireAuthenticatedDevice)
+        .flatMap(authenticatedDevice -> storeRepeatedUseKey(authenticatedDevice.accountIdentifier(),
+            request.getIdentityType(),
+            request.getSignedPreKey(),
+            KeysGrpcService::checkEcSignedPreKey,
+            (account, signedPreKey) -> {
+              final Consumer<Device> deviceUpdater = switch (IdentityType.fromGrpcIdentityType(request.getIdentityType())) {
+                case ACI -> device -> device.setSignedPreKey(signedPreKey);
+                case PNI -> device -> device.setPhoneNumberIdentitySignedPreKey(signedPreKey);
+              };
+
+              final UUID identifier = switch (IdentityType.fromGrpcIdentityType(request.getIdentityType())) {
+                case ACI -> account.getUuid();
+                case PNI -> account.getPhoneNumberIdentifier();
+              };
+
+              return Flux.merge(
+                      Mono.fromFuture(keysManager.storeEcSignedPreKeys(identifier, Map.of(authenticatedDevice.deviceId(), signedPreKey))),
+                      Mono.fromFuture(accountsManager.updateDeviceAsync(account, authenticatedDevice.deviceId(), deviceUpdater)))
+                  .then();
+            }));
+  }
+
+  @Override
+  public Mono<SetPreKeyResponse> setKemLastResortPreKey(final SetKemLastResortPreKeyRequest request) {
+    return Mono.fromSupplier(AuthenticationUtil::requireAuthenticatedDevice)
+        .flatMap(authenticatedDevice -> storeRepeatedUseKey(authenticatedDevice.accountIdentifier(),
+            request.getIdentityType(),
+            request.getSignedPreKey(),
+            KeysGrpcService::checkKemSignedPreKey,
+            (account, lastResortKey) -> {
+              final UUID identifier = switch (IdentityType.fromGrpcIdentityType(request.getIdentityType())) {
+                case ACI -> account.getUuid();
+                case PNI -> account.getPhoneNumberIdentifier();
+              };
+
+              return Mono.fromFuture(keysManager.storePqLastResort(identifier, Map.of(authenticatedDevice.deviceId(), lastResortKey)));
+            }));
+  }
+
+  private <K, R> Mono<SetPreKeyResponse> storeRepeatedUseKey(final UUID authenticatedAccountUuid,
+      final org.signal.chat.common.IdentityType identityType,
+      final R storeKeyRequest,
+      final BiFunction<R, IdentityKey, K> extractKeyFunction,
+      final BiFunction<Account, K, Mono<?>> storeKeyFunction) {
+
+    return Mono.fromFuture(accountsManager.getByAccountIdentifierAsync(authenticatedAccountUuid))
+        .map(maybeAccount -> maybeAccount.orElseThrow(Status.UNAUTHENTICATED::asRuntimeException))
+        .map(account -> {
+          final IdentityKey identityKey = switch (IdentityType.fromGrpcIdentityType(identityType)) {
+            case ACI -> account.getIdentityKey();
+            case PNI -> account.getPhoneNumberIdentityKey();
+          };
+
+          final K key = extractKeyFunction.apply(storeKeyRequest, identityKey);
+
+          return Tuples.of(account, key);
+        })
+        .flatMap(accountAndKey -> storeKeyFunction.apply(accountAndKey.getT1(), accountAndKey.getT2()))
+        .thenReturn(SetPreKeyResponse.newBuilder().build());
+  }
+
+  private static ECPreKey checkEcPreKey(final EcPreKey preKey) {
+    try {
+      return new ECPreKey(preKey.getKeyId(), new ECPublicKey(preKey.getPublicKey().toByteArray()));
+    } catch (final InvalidKeyException e) {
+      throw INVALID_PUBLIC_KEY_EXCEPTION;
+    }
+  }
+
+  private static ECSignedPreKey checkEcSignedPreKey(final EcSignedPreKey preKey, final IdentityKey identityKey) {
+    try {
+      final ECSignedPreKey ecSignedPreKey = new ECSignedPreKey(preKey.getKeyId(),
+          new ECPublicKey(preKey.getPublicKey().toByteArray()),
+          preKey.getSignature().toByteArray());
+
+      if (ecSignedPreKey.signatureValid(identityKey)) {
+        return ecSignedPreKey;
+      } else {
+        throw INVALID_SIGNATURE_EXCEPTION;
+      }
+    } catch (final InvalidKeyException e) {
+      throw INVALID_PUBLIC_KEY_EXCEPTION;
+    }
+  }
+
+  private static KEMSignedPreKey checkKemSignedPreKey(final KemSignedPreKey preKey, final IdentityKey identityKey) {
+    try {
+      final KEMSignedPreKey kemSignedPreKey = new KEMSignedPreKey(preKey.getKeyId(),
+          new KEMPublicKey(preKey.getPublicKey().toByteArray()),
+          preKey.getSignature().toByteArray());
+
+      if (kemSignedPreKey.signatureValid(identityKey)) {
+        return kemSignedPreKey;
+      } else {
+        throw INVALID_SIGNATURE_EXCEPTION;
+      }
+    } catch (final InvalidKeyException e) {
+      throw INVALID_PUBLIC_KEY_EXCEPTION;
+    }
+  }
+}

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

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.grpc;
+
+import io.grpc.Metadata;
+import io.grpc.Status;
+import io.grpc.StatusException;
+import java.time.Duration;
+import javax.annotation.Nullable;
+import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
+
+public class RateLimitUtil {
+
+  public static final Metadata.Key<Duration> RETRY_AFTER_DURATION_KEY =
+      Metadata.Key.of("retry-after", new Metadata.AsciiMarshaller<>() {
+        @Override
+        public String toAsciiString(final Duration value) {
+          return value.toString();
+        }
+
+        @Override
+        public Duration parseAsciiString(final String serialized) {
+          return Duration.parse(serialized);
+        }
+      });
+
+  public static Throwable mapRateLimitExceededException(final Throwable throwable) {
+    if (throwable instanceof RateLimitExceededException rateLimitExceededException) {
+      @Nullable final Metadata trailers = rateLimitExceededException.getRetryDuration()
+          .map(duration -> {
+            final Metadata metadata = new Metadata();
+            metadata.put(RETRY_AFTER_DURATION_KEY, duration);
+
+            return metadata;
+          }).orElse(null);
+
+      return new StatusException(Status.RESOURCE_EXHAUSTED, trailers);
+    }
+
+    return throwable;
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiter.java

@@ -8,6 +8,7 @@ package org.whispersystems.textsecuregcm.limits;
 import java.util.UUID;
 import java.util.concurrent.CompletionStage;
 import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
+import reactor.core.publisher.Mono;
 
 public interface RateLimiter {
 
@@ -53,6 +54,10 @@ public interface RateLimiter {
     return validateAsync(srcAccountUuid.toString() + "__" + dstAccountUuid.toString());
   }
 
+  default Mono<Void> validateReactive(final String key) {
+    return Mono.fromFuture(validateAsync(key).toCompletableFuture());
+  }
+
   default boolean hasAvailablePermits(final UUID accountUuid, final int permits) {
     return hasAvailablePermits(accountUuid.toString(), permits);
   }

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

@@ -90,6 +90,14 @@ public class KeysManager {
     return pqLastResortKeys.store(identifier, keys);
   }
 
+  public CompletableFuture<Void> storeEcOneTimePreKeys(final UUID identifier, final long deviceId, final List<ECPreKey> preKeys) {
+    return ecPreKeys.store(identifier, deviceId, preKeys);
+  }
+
+  public CompletableFuture<Void> storeKemOneTimePreKeys(final UUID identifier, final long deviceId, final List<KEMSignedPreKey> preKeys) {
+    return pqPreKeys.store(identifier, deviceId, preKeys);
+  }
+
   public CompletableFuture<Optional<ECPreKey>> takeEC(final UUID identifier, final long deviceId) {
     return ecPreKeys.take(identifier, deviceId);
   }

service/src/main/java/org/whispersystems/textsecuregcm/util/NoStackTraceRuntimeException.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.util;
+
+/**
+ * An abstract base class for runtime exceptions that do not include a stack trace. Stackless exceptions are generally
+ * intended for internal error-handling cases where the error will never be logged or otherwise reported.
+ */
+public abstract class NoStackTraceRuntimeException extends RuntimeException {
+
+  public NoStackTraceRuntimeException() {
+    super(null, null, true, false);
+  }
+
+  public NoStackTraceRuntimeException(final String message) {
+    super(message, null, true, false);
+  }
+
+  public NoStackTraceRuntimeException(final String message, final Throwable cause) {
+    super(message, cause, true, false);
+  }
+
+  public NoStackTraceRuntimeException(final Throwable cause) {
+    super(null, cause, true, false);
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/util/UUIDUtil.java

@@ -5,6 +5,7 @@
 
 package org.whispersystems.textsecuregcm.util;
 
+import com.google.protobuf.ByteString;
 import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
 import java.util.Optional;
@@ -27,6 +28,14 @@ public final class UUIDUtil {
     return byteBuffer.flip();
   }
 
+  public static ByteString toByteString(final UUID uuid) {
+    return ByteString.copyFrom(toByteBuffer(uuid));
+  }
+
+  public static UUID fromByteString(final ByteString byteString) {
+    return fromBytes(byteString.toByteArray());
+  }
+
   public static UUID fromBytes(final byte[] bytes) {
     return fromByteBuffer(ByteBuffer.wrap(bytes));
   }