Add API endpoints for waiting for newly-linked devices

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

@@ -642,7 +642,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     ClientPublicKeysManager clientPublicKeysManager =
         new ClientPublicKeysManager(clientPublicKeys, accountLockManager, accountLockExecutor);
     AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
-        accountLockManager, keysManager, messagesManager, profilesManager,
+        pubsubClient, accountLockManager, keysManager, messagesManager, profilesManager,
         secureStorageClient, secureValueRecovery2Client,
         clientPresenceManager,
         registrationRecoveryPasswordsManager, clientPublicKeysManager, accountLockExecutor, clientPresenceExecutor,
@@ -764,6 +764,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     environment.lifecycle().manage(keyTransparencyServiceClient);
     environment.lifecycle().manage(clientReleaseManager);
     environment.lifecycle().manage(virtualThreadPinEventMonitor);
+    environment.lifecycle().manage(accountsManager);
 
     final RegistrationCaptchaManager registrationCaptchaManager = new RegistrationCaptchaManager(captchaChecker);
 

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

@@ -4,22 +4,36 @@
  */
 package org.whispersystems.textsecuregcm.controllers;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.net.HttpHeaders;
 import io.dropwizard.auth.Auth;
+import io.lettuce.core.RedisException;
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Tags;
+import io.micrometer.core.instrument.Timer;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.headers.Header;
+import io.swagger.v3.oas.annotations.media.Schema;
 import io.swagger.v3.oas.annotations.responses.ApiResponse;
-import io.swagger.v3.oas.annotations.tags.Tag;
-import java.util.LinkedList;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.EnumMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import javax.validation.Valid;
+import javax.validation.constraints.Max;
+import javax.validation.constraints.Min;
 import javax.validation.constraints.NotNull;
+import javax.validation.constraints.Size;
 import javax.ws.rs.Consumes;
 import javax.ws.rs.DELETE;
+import javax.ws.rs.DefaultValue;
 import javax.ws.rs.ForbiddenException;
 import javax.ws.rs.GET;
 import javax.ws.rs.HeaderParam;
@@ -27,10 +41,12 @@ import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
 import javax.ws.rs.Produces;
+import javax.ws.rs.QueryParam;
 import javax.ws.rs.WebApplicationException;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
+import io.swagger.v3.oas.annotations.tags.Tag;
 import org.glassfish.jersey.server.ContainerRequest;
 import org.whispersystems.textsecuregcm.auth.LinkedDeviceRefreshRequirementProvider;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedDevice;
@@ -47,6 +63,8 @@ import org.whispersystems.textsecuregcm.entities.ProvisioningMessage;
 import org.whispersystems.textsecuregcm.entities.SetPublicKeyRequest;
 import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
+import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
+import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.ClientPublicKeysManager;
@@ -54,7 +72,11 @@ import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.Device.DeviceCapabilities;
 import org.whispersystems.textsecuregcm.storage.DeviceSpec;
 import org.whispersystems.textsecuregcm.storage.LinkDeviceTokenAlreadyUsedException;
-import org.whispersystems.textsecuregcm.util.VerificationCode;
+import org.whispersystems.textsecuregcm.util.ExceptionUtils;
+import org.whispersystems.textsecuregcm.util.LinkDeviceToken;
+import org.whispersystems.textsecuregcm.util.ua.ClientPlatform;
+import org.whispersystems.textsecuregcm.util.ua.UnrecognizedUserAgentException;
+import org.whispersystems.textsecuregcm.util.ua.UserAgentUtil;
 import org.whispersystems.websocket.auth.Mutable;
 import org.whispersystems.websocket.auth.ReadOnly;
 
@@ -69,6 +91,21 @@ public class DeviceController {
   private final RateLimiters rateLimiters;
   private final Map<String, Integer> maxDeviceConfiguration;
 
+  private final EnumMap<ClientPlatform, AtomicInteger> linkedDeviceListenersByPlatform;
+  private final AtomicInteger linkedDeviceListenersForUnrecognizedPlatforms;
+
+  private static final String LINKED_DEVICE_LISTENER_GAUGE_NAME =
+      MetricsUtil.name(DeviceController.class, "linkedDeviceListeners");
+
+  private static final String WAIT_FOR_LINKED_DEVICE_TIMER_NAME =
+      MetricsUtil.name(DeviceController.class, "waitForLinkedDeviceDuration");
+
+  @VisibleForTesting
+  static final int MIN_TOKEN_IDENTIFIER_LENGTH = 32;
+
+  @VisibleForTesting
+  static final int MAX_TOKEN_IDENTIFIER_LENGTH = 64;
+
   public DeviceController(final AccountsManager accounts,
       final ClientPublicKeysManager clientPublicKeysManager,
       final RateLimiters rateLimiters,
@@ -78,19 +115,32 @@ public class DeviceController {
     this.clientPublicKeysManager = clientPublicKeysManager;
     this.rateLimiters = rateLimiters;
     this.maxDeviceConfiguration = maxDeviceConfiguration;
+
+    linkedDeviceListenersByPlatform = Arrays.stream(ClientPlatform.values())
+        .collect(Collectors.toMap(
+            Function.identity(),
+            clientPlatform -> buildGauge(clientPlatform.name().toLowerCase()),
+            (a, b) -> {
+              throw new AssertionError("Duplicate client platform enumeration key");
+            },
+            () -> new EnumMap<>(ClientPlatform.class)
+        ));
+
+    linkedDeviceListenersForUnrecognizedPlatforms = buildGauge("unknown");
+  }
+
+  private static AtomicInteger buildGauge(final String clientPlatformName) {
+    return Metrics.gauge(LINKED_DEVICE_LISTENER_GAUGE_NAME,
+        Tags.of(io.micrometer.core.instrument.Tag.of(UserAgentTagUtil.PLATFORM_TAG, clientPlatformName)),
+        new AtomicInteger(0));
   }
 
   @GET
   @Produces(MediaType.APPLICATION_JSON)
   public DeviceInfoList getDevices(@ReadOnly @Auth AuthenticatedDevice auth) {
-    List<DeviceInfo> devices = new LinkedList<>();
-
-    for (Device device : auth.getAccount().getDevices()) {
-      devices.add(new DeviceInfo(device.getId(), device.getName(),
-          device.getLastSeen(), device.getCreated()));
-    }
-
-    return new DeviceInfoList(devices);
+    return new DeviceInfoList(auth.getAccount().getDevices().stream()
+        .map(DeviceInfo::forDevice)
+        .toList());
   }
 
   @DELETE
@@ -138,7 +188,7 @@ public class DeviceController {
   @ApiResponse(responseCode = "429", description = "Too many attempts", headers = @Header(
       name = "Retry-After",
       description = "If present, an positive integer indicating the number of seconds before a subsequent attempt could succeed"))
-  public VerificationCode createDeviceToken(@ReadOnly @Auth AuthenticatedDevice auth)
+  public LinkDeviceToken createDeviceToken(@ReadOnly @Auth AuthenticatedDevice auth)
       throws RateLimitExceededException, DeviceLimitExceededException {
 
     final Account account = auth.getAccount();
@@ -159,7 +209,9 @@ public class DeviceController {
       throw new WebApplicationException(Response.Status.UNAUTHORIZED);
     }
 
-    return new VerificationCode(accounts.generateDeviceLinkingToken(account.getUuid()));
+    final String token = accounts.generateLinkDeviceToken(account.getUuid());
+
+    return new LinkDeviceToken(token, AccountsManager.getLinkDeviceTokenIdentifier(token));
   }
 
   @PUT
@@ -266,6 +318,83 @@ public class DeviceController {
     }
   }
 
+  @GET
+  @Path("/wait_for_linked_device/{tokenIdentifier}")
+  @Produces(MediaType.APPLICATION_JSON)
+  @Operation(summary = "Wait for a new device to be linked to an account",
+      description = """
+          Waits for a new device to be linked to an account and returns basic information about the new device when
+          available.
+          """)
+  @ApiResponse(responseCode = "200", description = "The specified was linked to an account")
+  @ApiResponse(responseCode = "204", description = "No device was linked to the account before the call completed")
+  @ApiResponse(responseCode = "400", description = "The given token identifier or timeout was invalid")
+  @ApiResponse(responseCode = "429", description = "Rate-limited; try again after the prescribed delay")
+  @Schema(description = "Basic information about the linked device", implementation = DeviceInfo.class)
+  public CompletableFuture<Response> waitForLinkedDevice(
+      @ReadOnly @Auth final AuthenticatedDevice authenticatedDevice,
+
+      @PathParam("tokenIdentifier")
+      @Schema(description = "A 'link device' token identifier provided by the 'create link device token' endpoint")
+      @Size(min = MIN_TOKEN_IDENTIFIER_LENGTH, max = MAX_TOKEN_IDENTIFIER_LENGTH)
+      final String tokenIdentifier,
+
+      @QueryParam("timeout")
+      @DefaultValue("30")
+      @Min(1)
+      @Max(3600)
+      @Schema(requiredMode = Schema.RequiredMode.NOT_REQUIRED,
+          minimum = "1",
+          maximum = "3600",
+          description = """
+                The amount of time (in seconds) to wait for a response. If the expected device is not linked within the
+                given amount of time, this endpoint will return a status of HTTP/204.
+              """) final int timeoutSeconds,
+
+      @HeaderParam(HttpHeaders.USER_AGENT) String userAgent) throws RateLimitExceededException {
+
+    rateLimiters.getWaitForLinkedDeviceLimiter().validate(authenticatedDevice.getAccount().getIdentifier(IdentityType.ACI));
+
+    final AtomicInteger linkedDeviceListenerCounter = getCounterForLinkedDeviceListeners(userAgent);
+    linkedDeviceListenerCounter.incrementAndGet();
+
+    final Timer.Sample sample = Timer.start();
+
+    try {
+      return accounts.waitForNewLinkedDevice(tokenIdentifier, Duration.ofSeconds(timeoutSeconds))
+          .thenApply(maybeDeviceInfo -> maybeDeviceInfo
+              .map(deviceInfo -> Response.status(Response.Status.OK).entity(deviceInfo).build())
+              .orElseGet(() -> Response.status(Response.Status.NO_CONTENT).build()))
+          .exceptionally(ExceptionUtils.exceptionallyHandler(IllegalArgumentException.class,
+              e -> Response.status(Response.Status.BAD_REQUEST).build()))
+          .whenComplete((response, throwable) -> {
+            linkedDeviceListenerCounter.decrementAndGet();
+
+            if (response != null) {
+              sample.stop(Timer.builder(WAIT_FOR_LINKED_DEVICE_TIMER_NAME)
+                  .publishPercentileHistogram(true)
+                  .tags(Tags.of(UserAgentTagUtil.getPlatformTag(userAgent),
+                      io.micrometer.core.instrument.Tag.of("deviceFound",
+                          String.valueOf(response.getStatus() == Response.Status.OK.getStatusCode()))))
+                  .register(Metrics.globalRegistry));
+            }
+          });
+    } catch (final RedisException e) {
+      // `waitForNewLinkedDevice` could fail synchronously if the Redis circuit breaker is open; prevent counter drift
+      // if that happens
+      linkedDeviceListenerCounter.decrementAndGet();
+      throw e;
+    }
+  }
+
+  private AtomicInteger getCounterForLinkedDeviceListeners(final String userAgent) {
+    try {
+      return linkedDeviceListenersByPlatform.get(UserAgentUtil.parseUserAgentString(userAgent).getPlatform());
+    } catch (final UnrecognizedUserAgentException ignored) {
+      return linkedDeviceListenersForUnrecognizedPlatforms;
+    }
+  }
+
   @PUT
   @Produces(MediaType.APPLICATION_JSON)
   @Path("/unauthenticated_delivery")

service/src/main/java/org/whispersystems/textsecuregcm/entities/DeviceInfo.java

@@ -7,6 +7,7 @@ package org.whispersystems.textsecuregcm.entities;
 
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.util.ByteArrayBase64WithPaddingAdapter;
 
 public record DeviceInfo(long id,
@@ -17,4 +18,8 @@ public record DeviceInfo(long id,
 
                          long lastSeen,
                          long created) {
+
+  public static DeviceInfo forDevice(final Device device) {
+    return new DeviceInfo(device.getId(), device.getName(), device.getLastSeen(), device.getCreated());
+  }
 }

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

@@ -50,6 +50,7 @@ public class RateLimiters extends BaseRateLimiters<RateLimiters.For> {
     EXTERNAL_SERVICE_CREDENTIALS("externalServiceCredentials", true, new RateLimiterConfig(100, Duration.ofMinutes(15))),
     KEY_TRANSPARENCY_SEARCH_PER_IP("keyTransparencySearch", true, new RateLimiterConfig(100, Duration.ofSeconds(15))),
     KEY_TRANSPARENCY_MONITOR_PER_IP("keyTransparencyMonitor", true, new RateLimiterConfig(100, Duration.ofSeconds(15))),
+    WAIT_FOR_LINKED_DEVICE("waitForLinkedDevice", true, new RateLimiterConfig(10, Duration.ofSeconds(30))),
     ;
 
     private final String id;
@@ -205,4 +206,8 @@ public class RateLimiters extends BaseRateLimiters<RateLimiters.For> {
   public RateLimiter getStoriesLimiter() {
     return forDescriptor(For.STORIES);
   }
+
+  public RateLimiter getWaitForLinkedDeviceLimiter() {
+    return forDescriptor(For.WAIT_FOR_LINKED_DEVICE);
+  }
 }

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

@@ -12,8 +12,11 @@ import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectWriter;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import io.dropwizard.lifecycle.Managed;
 import io.lettuce.core.RedisException;
+import io.lettuce.core.SetArgs;
 import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
+import io.lettuce.core.pubsub.RedisPubSubAdapter;
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Tag;
 import io.micrometer.core.instrument.Tags;
@@ -42,7 +45,9 @@ import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
@@ -61,13 +66,16 @@ import org.whispersystems.textsecuregcm.auth.SaltedTokenHash;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.controllers.MismatchedDevicesException;
 import org.whispersystems.textsecuregcm.entities.AccountAttributes;
+import org.whispersystems.textsecuregcm.entities.DeviceInfo;
 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.metrics.UserAgentTagUtil;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantPubSubConnection;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisClusterClient;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisClient;
 import org.whispersystems.textsecuregcm.redis.RedisOperation;
 import org.whispersystems.textsecuregcm.securestorage.SecureStorageClient;
 import org.whispersystems.textsecuregcm.securevaluerecovery.SecureValueRecovery2Client;
@@ -82,14 +90,13 @@ import reactor.core.scheduler.Scheduler;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
 import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException;
 
-public class AccountsManager {
+public class AccountsManager extends RedisPubSubAdapter<String, String> implements Managed {
 
   private static final Timer createTimer = Metrics.timer(name(AccountsManager.class, "create"));
   private static final Timer updateTimer = Metrics.timer(name(AccountsManager.class, "update"));
   private static final Timer getByNumberTimer = Metrics.timer(name(AccountsManager.class, "getByNumber"));
   private static final Timer getByUsernameHashTimer = Metrics.timer(name(AccountsManager.class, "getByUsernameHash"));
-  private static final Timer getByUsernameLinkHandleTimer = Metrics.timer(
-      name(AccountsManager.class, "getByUsernameLinkHandle"));
+  private static final Timer getByUsernameLinkHandleTimer = Metrics.timer(name(AccountsManager.class, "getByUsernameLinkHandle"));
   private static final Timer getByUuidTimer = Metrics.timer(name(AccountsManager.class, "getByUuid"));
   private static final Timer deleteTimer = Metrics.timer(name(AccountsManager.class, "delete"));
 
@@ -108,6 +115,7 @@ public class AccountsManager {
   private final Accounts accounts;
   private final PhoneNumberIdentifiers phoneNumberIdentifiers;
   private final FaultTolerantRedisClusterClient cacheCluster;
+  private final FaultTolerantRedisClient pubSubRedisSingleton;
   private final AccountLockManager accountLockManager;
   private final KeysManager keysManager;
   private final MessagesManager messagesManager;
@@ -124,6 +132,16 @@ public class AccountsManager {
 
   private final Key verificationTokenKey;
 
+  private final FaultTolerantPubSubConnection<String, String> pubSubConnection;
+
+  private final Map<String, CompletableFuture<Optional<DeviceInfo>>> waitForDeviceFuturesByTokenIdentifier =
+      new ConcurrentHashMap<>();
+
+  private static final int SHA256_HASH_LENGTH = getSha256MessageDigest().getDigestLength();
+  private static final Duration RECENTLY_ADDED_DEVICE_TTL = Duration.ofHours(1);
+  private static final String LINKED_DEVICE_PREFIX = "linked_device::";
+  private static final String LINKED_DEVICE_KEYSPACE_PATTERN = "__keyspace@0__:" + LINKED_DEVICE_PREFIX + "*";
+
   private static final ObjectWriter ACCOUNT_REDIS_JSON_WRITER = SystemMapper.jsonMapper()
       .writer(SystemMapper.excludingField(Account.class, List.of("uuid")));
 
@@ -158,6 +176,7 @@ public class AccountsManager {
   public AccountsManager(final Accounts accounts,
       final PhoneNumberIdentifiers phoneNumberIdentifiers,
       final FaultTolerantRedisClusterClient cacheCluster,
+      final FaultTolerantRedisClient pubSubRedisSingleton,
       final AccountLockManager accountLockManager,
       final KeysManager keysManager,
       final MessagesManager messagesManager,
@@ -175,6 +194,7 @@ public class AccountsManager {
     this.accounts = accounts;
     this.phoneNumberIdentifiers = phoneNumberIdentifiers;
     this.cacheCluster = cacheCluster;
+    this.pubSubRedisSingleton = pubSubRedisSingleton;
     this.accountLockManager = accountLockManager;
     this.keysManager = keysManager;
     this.messagesManager = messagesManager;
@@ -197,6 +217,20 @@ public class AccountsManager {
     } catch (final InvalidKeyException e) {
       throw new IllegalArgumentException(e);
     }
+
+    this.pubSubConnection = pubSubRedisSingleton.createPubSubConnection();
+  }
+
+  @Override
+  public void start() {
+    pubSubConnection.usePubSubConnection(connection -> connection.addListener(this));
+    pubSubConnection.usePubSubConnection(connection -> connection.sync().psubscribe(LINKED_DEVICE_KEYSPACE_PATTERN));
+  }
+
+  @Override
+  public void stop() {
+    pubSubConnection.usePubSubConnection(connection -> connection.sync().punsubscribe());
+    pubSubConnection.usePubSubConnection(connection -> connection.removeListener(this));
   }
 
   public Account create(final String number,
@@ -363,6 +397,26 @@ public class AccountsManager {
           }
 
           return CompletableFuture.failedFuture(throwable);
+        })
+        .whenComplete((updatedAccountAndDevice, throwable) -> {
+          if (updatedAccountAndDevice != null) {
+            final String key = getLinkedDeviceKey(getLinkDeviceTokenIdentifier(linkDeviceToken));
+            final String deviceInfoJson;
+
+            try {
+              deviceInfoJson = SystemMapper.jsonMapper().writeValueAsString(DeviceInfo.forDevice(updatedAccountAndDevice.second()));
+            } catch (final JsonProcessingException e) {
+              throw new UncheckedIOException(e);
+            }
+
+            pubSubRedisSingleton.withConnection(connection ->
+                connection.async().set(key, deviceInfoJson, SetArgs.Builder.ex(RECENTLY_ADDED_DEVICE_TTL)))
+                .whenComplete((ignored, pubSubThrowable) -> {
+                  if (pubSubThrowable != null) {
+                    logger.warn("Failed to record recently-created device", pubSubThrowable);
+                  }
+                });
+          }
         });
   }
 
@@ -386,7 +440,7 @@ public class AccountsManager {
     }
   }
 
-  public String generateDeviceLinkingToken(final UUID aci) {
+  public String generateLinkDeviceToken(final UUID aci) {
     final String claims = aci + "." + clock.instant().toEpochMilli();
     final byte[] signature = getInitializedMac().doFinal(claims.getBytes(StandardCharsets.UTF_8));
 
@@ -394,7 +448,7 @@ public class AccountsManager {
   }
 
   @VisibleForTesting
-  static String generateDeviceLinkingToken(final UUID aci, final Key linkDeviceTokenKey, final Clock clock)
+  static String generateLinkDeviceToken(final UUID aci, final Key linkDeviceTokenKey, final Clock clock)
       throws InvalidKeyException {
 
     final String claims = aci + "." + clock.instant().toEpochMilli();
@@ -403,6 +457,11 @@ public class AccountsManager {
     return claims + ":" + Base64.getUrlEncoder().encodeToString(signature);
   }
 
+  public static String getLinkDeviceTokenIdentifier(final String linkDeviceToken) {
+    return Base64.getUrlEncoder().withoutPadding().encodeToString(
+        getSha256MessageDigest().digest(linkDeviceToken.getBytes(StandardCharsets.UTF_8)));
+  }
+
   /**
    * Checks that a device-linking token is valid and returns the account identifier from the token if so, or empty if
    * the token was invalid
@@ -1340,4 +1399,75 @@ public class AccountsManager {
         .whenComplete((ignoredResult, ignoredException) -> sample.stop(redisDeleteTimer))
         .thenRun(Util.NOOP);
   }
+
+  public CompletableFuture<Optional<DeviceInfo>> waitForNewLinkedDevice(final String linkDeviceTokenIdentifier, final Duration timeout) {
+    // Unbeknownst to callers but beknownst to us, the "link device token identifier" is the base64/url-encoded SHA256
+    // hash of a device-linking token. Before we use the string anywhere, make sure it's the right "shape" for a hash.
+    if (Base64.getUrlDecoder().decode(linkDeviceTokenIdentifier).length != SHA256_HASH_LENGTH) {
+      return CompletableFuture.failedFuture(new IllegalArgumentException("Invalid token identifier"));
+    }
+
+    final CompletableFuture<Optional<DeviceInfo>> waitForDeviceFuture = new CompletableFuture<>();
+
+    waitForDeviceFuture
+        .completeOnTimeout(Optional.empty(), TimeUnit.MILLISECONDS.convert(timeout), TimeUnit.MILLISECONDS)
+        .whenComplete((maybeDevice, throwable) -> waitForDeviceFuturesByTokenIdentifier.compute(linkDeviceTokenIdentifier,
+            (ignored, existingFuture) -> {
+              // Only remove the future from the map if it's THIS future, and not one that later displaced this one
+              return existingFuture == waitForDeviceFuture ? null : existingFuture;
+            }));
+
+    {
+      final CompletableFuture<Optional<DeviceInfo>> displacedFuture =
+          waitForDeviceFuturesByTokenIdentifier.put(linkDeviceTokenIdentifier, waitForDeviceFuture);
+
+      if (displacedFuture != null) {
+        displacedFuture.complete(Optional.empty());
+      }
+    }
+
+    // The device may already have been linked by the time the caller started watching for it; perform an immediate
+    // check to see if the device is already there.
+    pubSubRedisSingleton.withConnection(connection -> connection.async().get(getLinkedDeviceKey(linkDeviceTokenIdentifier)))
+        .thenAccept(response -> {
+          if (StringUtils.isNotBlank(response)) {
+            handleDeviceAdded(waitForDeviceFuture, response);
+          }
+        });
+
+    return waitForDeviceFuture;
+  }
+
+  private static String getLinkedDeviceKey(final String linkDeviceTokenIdentifier) {
+    return LINKED_DEVICE_PREFIX + linkDeviceTokenIdentifier;
+  }
+
+  @Override
+  public void message(final String pattern, final String channel, final String message) {
+    if (LINKED_DEVICE_KEYSPACE_PATTERN.equals(pattern) && "set".equalsIgnoreCase(message)) {
+      // The `- 1` here compensates for the '*' in the pattern
+      final String tokenIdentifier = channel.substring(LINKED_DEVICE_KEYSPACE_PATTERN.length() - 1);
+
+      Optional.ofNullable(waitForDeviceFuturesByTokenIdentifier.remove(tokenIdentifier))
+          .ifPresent(future -> pubSubRedisSingleton.withConnection(connection -> connection.async().get(getLinkedDeviceKey(tokenIdentifier)))
+              .thenAccept(deviceInfoJson -> handleDeviceAdded(future, deviceInfoJson)));
+    }
+  }
+
+  private void handleDeviceAdded(final CompletableFuture<Optional<DeviceInfo>> future, final String deviceInfoJson) {
+    try {
+      future.complete(Optional.of(SystemMapper.jsonMapper().readValue(deviceInfoJson, DeviceInfo.class)));
+    } catch (final JsonProcessingException e) {
+      logger.error("Could not parse device json", e);
+      future.completeExceptionally(e);
+    }
+  }
+
+  private static MessageDigest getSha256MessageDigest() {
+    try {
+      return MessageDigest.getInstance("SHA-256");
+    } catch (final NoSuchAlgorithmException e) {
+      throw new AssertionError("Every implementation of the Java platform is required to support the SHA-256 MessageDigest algorithm", e);
+    }
+  }
 }

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

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.textsecuregcm.util;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+import io.swagger.v3.oas.annotations.media.Schema;
+
+public record LinkDeviceToken(
+    @Schema(description = """
+        An opaque token to send to a new linked device that authorizes the new device to link itself to the account that
+        requested this token.
+        """)
+    @JsonProperty("verificationCode") String token,
+
+    @Schema(description = """
+        An opaque identifier for the generated token that the caller may use to watch for a new device to complete the
+        linking process.
+        """)
+    String tokenIdentifier) {
+}

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

@@ -1,8 +0,0 @@
-/*
- * Copyright 2013-2020 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-package org.whispersystems.textsecuregcm.util;
-
-public record VerificationCode(String verificationCode) {
-}

service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java

@@ -40,6 +40,7 @@ import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 import org.whispersystems.textsecuregcm.push.FcmSender;
 import org.whispersystems.textsecuregcm.push.PushNotificationManager;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisClusterClient;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisClient;
 import org.whispersystems.textsecuregcm.securestorage.SecureStorageClient;
 import org.whispersystems.textsecuregcm.securevaluerecovery.SecureValueRecovery2Client;
 import org.whispersystems.textsecuregcm.storage.AccountLockManager;
@@ -112,6 +113,8 @@ record CommandDependencies(
         .build("main_cache", redisClientResourcesBuilder);
     FaultTolerantRedisClusterClient pushSchedulerCluster = configuration.getPushSchedulerCluster()
         .build("push_scheduler", redisClientResourcesBuilder);
+    FaultTolerantRedisClient pubsubClient =
+        configuration.getRedisPubSubConfiguration().build("pubsub", redisClientResourcesBuilder.build());
 
     ScheduledExecutorService recurringJobExecutor = environment.lifecycle()
         .scheduledExecutorService(name(name, "recurringJob-%d")).threads(2).build();
@@ -225,7 +228,7 @@ record CommandDependencies(
     ClientPublicKeysManager clientPublicKeysManager =
         new ClientPublicKeysManager(clientPublicKeys, accountLockManager, accountLockExecutor);
     AccountsManager accountsManager = new AccountsManager(accounts, phoneNumberIdentifiers, cacheCluster,
-        accountLockManager, keys, messagesManager, profilesManager,
+        pubsubClient, accountLockManager, keys, messagesManager, profilesManager,
         secureStorageClient, secureValueRecovery2Client, clientPresenceManager,
         registrationRecoveryPasswordsManager, clientPublicKeysManager, accountLockExecutor, clientPresenceExecutor,
         clock, configuration.getLinkDeviceSecretConfiguration().secret().value(), dynamicConfigurationManager);