Add API endpoints for waiting for account restoration requests

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

@@ -29,6 +29,7 @@ import javax.annotation.Nullable;
 import javax.validation.Valid;
 import javax.validation.constraints.Max;
 import javax.validation.constraints.Min;
+import javax.validation.constraints.NotBlank;
 import javax.validation.constraints.NotNull;
 import javax.validation.constraints.Size;
 import javax.ws.rs.Consumes;
@@ -56,6 +57,7 @@ import org.whispersystems.textsecuregcm.entities.AccountAttributes;
 import org.whispersystems.textsecuregcm.entities.DeviceActivationRequest;
 import org.whispersystems.textsecuregcm.entities.DeviceInfo;
 import org.whispersystems.textsecuregcm.entities.DeviceInfoList;
+import org.whispersystems.textsecuregcm.entities.RestoreAccountRequest;
 import org.whispersystems.textsecuregcm.entities.LinkDeviceResponse;
 import org.whispersystems.textsecuregcm.entities.LinkDeviceRequest;
 import org.whispersystems.textsecuregcm.entities.PreKeySignatureValidator;
@@ -64,6 +66,7 @@ import org.whispersystems.textsecuregcm.entities.RemoteAttachment;
 import org.whispersystems.textsecuregcm.entities.SetPublicKeyRequest;
 import org.whispersystems.textsecuregcm.entities.TransferArchiveUploadedRequest;
 import org.whispersystems.textsecuregcm.identity.IdentityType;
+import org.whispersystems.textsecuregcm.limits.RateLimitedByIp;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
@@ -437,6 +440,70 @@ public class DeviceController {
     return isDowngrade;
   }
 
+  @PUT
+  @Consumes(MediaType.APPLICATION_JSON)
+  @Produces(MediaType.APPLICATION_JSON)
+  @Path("/restore_account/{token}")
+  @Operation(
+      summary = "Signals that a new device is requesting restoration of account data by some method",
+      description = """
+          Signals that a new device is requesting restoration of account data by some method. Devices waiting via the
+          "wait for 'restore account' request" endpoint will be notified that the request has been issued.
+          """)
+  @ApiResponse(responseCode = "204", description = "Success")
+  @ApiResponse(responseCode = "422", description = "The request object could not be parsed or was otherwise invalid")
+  @ApiResponse(responseCode = "429", description = "Rate-limited; try again after the prescribed delay")
+  @RateLimitedByIp(RateLimiters.For.RECORD_DEVICE_TRANSFER_REQUEST)
+  public CompletionStage<Void> recordRestoreAccountRequest(
+      @PathParam("token")
+      @NotBlank
+      @Size(max = 64)
+      @Schema(description = "A randomly-generated token identifying the request for device-to-device transfer.",
+          requiredMode = Schema.RequiredMode.REQUIRED,
+          maximum = "64") final String token,
+
+      @Valid
+      final RestoreAccountRequest restoreAccountRequest) {
+
+    return accounts.recordRestoreAccountRequest(token, restoreAccountRequest);
+  }
+
+  @GET
+  @Produces(MediaType.APPLICATION_JSON)
+  @Path("/restore_account/{token}")
+  @Operation(summary = "Wait for 'restore account' request")
+  @ApiResponse(responseCode = "200", description = "A 'restore account' request was received for the given token",
+      content = @Content(schema = @Schema(implementation = RestoreAccountRequest.class)))
+  @ApiResponse(responseCode = "204", description = "No 'restore account' request for the given token was received before the call completed; clients may repeat the call to continue waiting")
+  @ApiResponse(responseCode = "400", description = "The given token or timeout was invalid")
+  @ApiResponse(responseCode = "429", description = "Rate-limited; try again after the prescribed delay")
+  @RateLimitedByIp(RateLimiters.For.WAIT_FOR_DEVICE_TRANSFER_REQUEST)
+  public CompletionStage<Response> waitForDeviceTransferRequest(
+      @PathParam("token")
+      @NotBlank
+      @Size(max = 64)
+      @Schema(description = "A randomly-generated token identifying the request for device-to-device transfer.",
+          requiredMode = Schema.RequiredMode.REQUIRED,
+          maximum = "64") final String token,
+
+      @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 a transfer archive for the authenticated
+                device is not available within the given amount of time, this endpoint will return a status of HTTP/204.
+              """) final int timeoutSeconds) {
+
+    return accounts.waitForRestoreAccountRequest(token, Duration.ofSeconds(timeoutSeconds))
+        .thenApply(maybeRequestReceived -> maybeRequestReceived
+            .map(restoreAccountRequest -> Response.status(Response.Status.OK).entity(restoreAccountRequest).build())
+            .orElseGet(() -> Response.status(Response.Status.NO_CONTENT).build()));
+  }
+
   @PUT
   @Consumes(MediaType.APPLICATION_JSON)
   @Produces(MediaType.APPLICATION_JSON)

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.entities;
+
+import io.swagger.v3.oas.annotations.media.Schema;
+import javax.validation.constraints.NotNull;
+
+@Schema(description = """
+    Represents a request from a new device to restore account data by some method.
+    """)
+public record RestoreAccountRequest(
+    @NotNull
+    @Schema(description = "The method by which the new device has requested account data restoration")
+    Method method) {
+
+  public enum Method {
+    @Schema(description = "Restore account data from a remote message history backup")
+    REMOTE_BACKUP,
+
+    @Schema(description = "Restore account data from a local backup archive")
+    LOCAL_BACKUP,
+
+    @Schema(description = "Restore account data via direct device-to-device transfer")
+    DEVICE_TRANSFER,
+
+    @Schema(description = "Do not restore account data")
+    DECLINE,
+  }
+}

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

@@ -53,6 +53,8 @@ public class RateLimiters extends BaseRateLimiters<RateLimiters.For> {
     WAIT_FOR_LINKED_DEVICE("waitForLinkedDevice", true, new RateLimiterConfig(10, Duration.ofSeconds(30))),
     UPLOAD_TRANSFER_ARCHIVE("uploadTransferArchive", true, new RateLimiterConfig(10, Duration.ofMinutes(1))),
     WAIT_FOR_TRANSFER_ARCHIVE("waitForTransferArchive", true, new RateLimiterConfig(10, Duration.ofSeconds(30))),
+    RECORD_DEVICE_TRANSFER_REQUEST("recordDeviceTransferRequest", true, new RateLimiterConfig(10, Duration.ofMillis(100))),
+    WAIT_FOR_DEVICE_TRANSFER_REQUEST("waitForDeviceTransferRequest", true, new RateLimiterConfig(10, Duration.ofMillis(100))),
     ;
 
     private final String id;

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

@@ -68,6 +68,7 @@ import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfigurati
 import org.whispersystems.textsecuregcm.controllers.MismatchedDevicesException;
 import org.whispersystems.textsecuregcm.entities.AccountAttributes;
 import org.whispersystems.textsecuregcm.entities.DeviceInfo;
+import org.whispersystems.textsecuregcm.entities.RestoreAccountRequest;
 import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
 import org.whispersystems.textsecuregcm.entities.RemoteAttachment;
@@ -142,6 +143,9 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
   private final Map<TimestampedDeviceIdentifier, CompletableFuture<Optional<RemoteAttachment>>> waitForTransferArchiveFuturesByDeviceIdentifier =
       new ConcurrentHashMap<>();
 
+  private final Map<String, CompletableFuture<Optional<RestoreAccountRequest>>> waitForRestoreAccountRequestFuturesByToken =
+      new ConcurrentHashMap<>();
+
   private static final int SHA256_HASH_LENGTH = getSha256MessageDigest().getDigestLength();
 
   private static final Duration RECENTLY_ADDED_DEVICE_TTL = Duration.ofHours(1);
@@ -152,6 +156,10 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
   private static final String TRANSFER_ARCHIVE_PREFIX = "transfer_archive::";
   private static final String TRANSFER_ARCHIVE_KEYSPACE_PATTERN = "__keyspace@0__:" + TRANSFER_ARCHIVE_PREFIX + "*";
 
+  private static final Duration RESTORE_ACCOUNT_REQUEST_TTL = Duration.ofHours(1);
+  private static final String RESTORE_ACCOUNT_REQUEST_PREFIX = "restore_account::";
+  private static final String RESTORE_ACCOUNT_REQUEST_KEYSPACE_PATTERN = "__keyspace@0__:" + RESTORE_ACCOUNT_REQUEST_PREFIX + "*";
+
   private static final ObjectWriter ACCOUNT_REDIS_JSON_WRITER = SystemMapper.jsonMapper()
       .writer(SystemMapper.excludingField(Account.class, List.of("uuid")));
 
@@ -238,7 +246,8 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
   public void start() {
     pubSubConnection.usePubSubConnection(connection -> {
       connection.addListener(this);
-      connection.sync().psubscribe(LINKED_DEVICE_KEYSPACE_PATTERN, TRANSFER_ARCHIVE_KEYSPACE_PATTERN);
+      connection.sync().psubscribe(LINKED_DEVICE_KEYSPACE_PATTERN, TRANSFER_ARCHIVE_KEYSPACE_PATTERN,
+          RESTORE_ACCOUNT_REQUEST_KEYSPACE_PATTERN);
     });
   }
 
@@ -1496,6 +1505,44 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
         ":" + destinationDeviceCreationTimestamp.toEpochMilli();
   }
 
+  public CompletableFuture<Optional<RestoreAccountRequest>> waitForRestoreAccountRequest(final String token, final Duration timeout) {
+    return waitForPubSubKey(waitForRestoreAccountRequestFuturesByToken,
+        token,
+        getRestoreAccountRequestKey(token),
+        timeout,
+        this::handleRestoreAccountRequest);
+  }
+
+  public CompletableFuture<Void> recordRestoreAccountRequest(final String token, final RestoreAccountRequest restoreAccountRequest) {
+    final String key = getRestoreAccountRequestKey(token);
+
+    final String requestJson;
+
+    try {
+      requestJson = SystemMapper.jsonMapper().writeValueAsString(restoreAccountRequest);
+    } catch (final JsonProcessingException e) {
+      throw new UncheckedIOException(e);
+    }
+
+    return pubSubRedisClient.withConnection(connection ->
+            connection.async().set(key, requestJson, SetArgs.Builder.ex(RESTORE_ACCOUNT_REQUEST_TTL)))
+        .thenRun(Util.NOOP)
+        .toCompletableFuture();
+  }
+
+  private void handleRestoreAccountRequest(final CompletableFuture<Optional<RestoreAccountRequest>> future, final String transferRequestJson) {
+    try {
+      future.complete(Optional.of(SystemMapper.jsonMapper().readValue(transferRequestJson, RestoreAccountRequest.class)));
+    } catch (final JsonProcessingException e) {
+      logger.error("Could not parse device transfer request JSON", e);
+      future.completeExceptionally(e);
+    }
+  }
+
+  private static String getRestoreAccountRequestKey(final String token) {
+    return RESTORE_ACCOUNT_REQUEST_PREFIX + token;
+  }
+
   private <K, T> CompletableFuture<Optional<T>> waitForPubSubKey(final Map<K, CompletableFuture<Optional<T>>> futureMap,
       final K mapKey,
       final String redisKey,
@@ -1564,6 +1611,14 @@ public class AccountsManager extends RedisPubSubAdapter<String, String> implemen
       } catch (final IllegalArgumentException e) {
         logger.error("Could not parse timestamped device identifier", e);
       }
+    } else if (RESTORE_ACCOUNT_REQUEST_KEYSPACE_PATTERN.equalsIgnoreCase(pattern) && "set".equalsIgnoreCase(message)) {
+      // The `- 1` here compensates for the '*' in the pattern
+      final String token = channel.substring(RESTORE_ACCOUNT_REQUEST_KEYSPACE_PATTERN.length() - 1);
+
+      Optional.ofNullable(waitForRestoreAccountRequestFuturesByToken.remove(token))
+          .ifPresent(future -> pubSubRedisClient.withConnection(connection -> connection.async().get(
+                  getRestoreAccountRequestKey(token)))
+              .thenAccept(requestJson -> handleRestoreAccountRequest(future, requestJson)));
     }
   }