Add a gRPC service for working with devices

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

@@ -0,0 +1,212 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.grpc;
+
+import com.google.protobuf.ByteString;
+import io.grpc.Status;
+import java.util.Base64;
+import java.util.Objects;
+import javax.annotation.Nullable;
+import org.apache.commons.lang3.StringUtils;
+import org.signal.chat.device.ClearPushTokenRequest;
+import org.signal.chat.device.ClearPushTokenResponse;
+import org.signal.chat.device.GetDevicesRequest;
+import org.signal.chat.device.GetDevicesResponse;
+import org.signal.chat.device.ReactorDevicesGrpc;
+import org.signal.chat.device.RemoveDeviceRequest;
+import org.signal.chat.device.RemoveDeviceResponse;
+import org.signal.chat.device.SetCapabilitiesRequest;
+import org.signal.chat.device.SetCapabilitiesResponse;
+import org.signal.chat.device.SetDeviceNameRequest;
+import org.signal.chat.device.SetDeviceNameResponse;
+import org.signal.chat.device.SetPushTokenRequest;
+import org.signal.chat.device.SetPushTokenResponse;
+import org.whispersystems.textsecuregcm.auth.grpc.AuthenticatedDevice;
+import org.whispersystems.textsecuregcm.auth.grpc.AuthenticationUtil;
+import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.KeysManager;
+import org.whispersystems.textsecuregcm.storage.MessagesManager;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+public class DevicesGrpcService extends ReactorDevicesGrpc.DevicesImplBase {
+
+  private final AccountsManager accountsManager;
+  private final KeysManager keysManager;
+  private final MessagesManager messagesManager;
+
+  private static final int MAX_NAME_LENGTH = 256;
+
+  public DevicesGrpcService(final AccountsManager accountsManager,
+      final KeysManager keysManager,
+      final MessagesManager messagesManager) {
+
+    this.accountsManager = accountsManager;
+    this.keysManager = keysManager;
+    this.messagesManager = messagesManager;
+  }
+
+  @Override
+  public Mono<GetDevicesResponse> getDevices(final GetDevicesRequest request) {
+    final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice();
+
+    return Mono.fromFuture(() -> accountsManager.getByAccountIdentifierAsync(authenticatedDevice.accountIdentifier()))
+        .map(maybeAccount -> maybeAccount.orElseThrow(Status.UNAUTHENTICATED::asRuntimeException))
+        .flatMapMany(account -> Flux.fromIterable(account.getDevices()))
+        .reduce(GetDevicesResponse.newBuilder(), (builder, device) -> {
+          final GetDevicesResponse.LinkedDevice.Builder linkedDeviceBuilder = GetDevicesResponse.LinkedDevice.newBuilder();
+
+          if (StringUtils.isNotBlank(device.getName())) {
+            linkedDeviceBuilder.setName(ByteString.copyFrom(Base64.getDecoder().decode(device.getName())));
+          }
+
+          return builder.addDevices(linkedDeviceBuilder
+              .setId(device.getId())
+              .setCreated(device.getCreated())
+              .setLastSeen(device.getLastSeen())
+              .build());
+        })
+        .map(GetDevicesResponse.Builder::build);
+  }
+
+  @Override
+  public Mono<RemoveDeviceResponse> removeDevice(final RemoveDeviceRequest request) {
+    if (request.getId() == Device.MASTER_ID) {
+      throw Status.INVALID_ARGUMENT.withDescription("Cannot remove primary device").asRuntimeException();
+    }
+
+    final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedPrimaryDevice();
+
+    return Mono.fromFuture(() -> accountsManager.getByAccountIdentifierAsync(authenticatedDevice.accountIdentifier()))
+        .map(maybeAccount -> maybeAccount.orElseThrow(Status.UNAUTHENTICATED::asRuntimeException))
+        .flatMap(account -> Flux.merge(
+                Mono.fromFuture(() -> messagesManager.clear(account.getUuid(), request.getId())),
+                Mono.fromFuture(() -> keysManager.delete(account.getUuid(), request.getId())))
+            .then(Mono.fromFuture(() -> accountsManager.updateAsync(account, a -> a.removeDevice(request.getId()))))
+            // Some messages may have arrived while we were performing the other updates; make a best effort to clear
+            // those out, too
+            .then(Mono.fromFuture(() -> messagesManager.clear(account.getUuid(), request.getId()))))
+        .thenReturn(RemoveDeviceResponse.newBuilder().build());
+  }
+
+  @Override
+  public Mono<SetDeviceNameResponse> setDeviceName(final SetDeviceNameRequest request) {
+    final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice();
+
+    if (request.getName().isEmpty()) {
+      throw Status.INVALID_ARGUMENT.withDescription("Must specify a device name").asRuntimeException();
+    }
+
+    if (request.getName().size() > MAX_NAME_LENGTH) {
+      throw Status.INVALID_ARGUMENT.withDescription("Device name must be at most " + MAX_NAME_LENGTH + " bytes")
+          .asRuntimeException();
+    }
+
+    return Mono.fromFuture(() -> accountsManager.getByAccountIdentifierAsync(authenticatedDevice.accountIdentifier()))
+        .map(maybeAccount -> maybeAccount.orElseThrow(Status.UNAUTHENTICATED::asRuntimeException))
+        .flatMap(account -> Mono.fromFuture(() -> accountsManager.updateDeviceAsync(account, authenticatedDevice.deviceId(),
+            device -> device.setName(Base64.getEncoder().encodeToString(request.getName().toByteArray())))))
+        .thenReturn(SetDeviceNameResponse.newBuilder().build());
+  }
+
+  @Override
+  public Mono<SetPushTokenResponse> setPushToken(final SetPushTokenRequest request) {
+    final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice();
+
+    @Nullable final String apnsToken;
+    @Nullable final String apnsVoipToken;
+    @Nullable final String fcmToken;
+
+    switch (request.getTokenRequestCase()) {
+
+      case APNS_TOKEN_REQUEST -> {
+        final SetPushTokenRequest.ApnsTokenRequest apnsTokenRequest = request.getApnsTokenRequest();
+
+        if (StringUtils.isAllBlank(apnsTokenRequest.getApnsToken(), apnsTokenRequest.getApnsVoipToken())) {
+          throw Status.INVALID_ARGUMENT.withDescription("APNs tokens may not both be blank").asRuntimeException();
+        }
+
+        apnsToken = StringUtils.stripToNull(apnsTokenRequest.getApnsToken());
+        apnsVoipToken = StringUtils.stripToNull(apnsTokenRequest.getApnsVoipToken());
+        fcmToken = null;
+      }
+
+      case FCM_TOKEN_REQUEST -> {
+        final SetPushTokenRequest.FcmTokenRequest fcmTokenRequest = request.getFcmTokenRequest();
+
+        if (StringUtils.isBlank(fcmTokenRequest.getFcmToken())) {
+          throw Status.INVALID_ARGUMENT.withDescription("FCM token must not be blank").asRuntimeException();
+        }
+
+        apnsToken = null;
+        apnsVoipToken = null;
+        fcmToken = StringUtils.stripToNull(fcmTokenRequest.getFcmToken());
+      }
+
+      default -> throw Status.INVALID_ARGUMENT.withDescription("No tokens specified").asRuntimeException();
+    }
+
+    return Mono.fromFuture(() -> accountsManager.getByAccountIdentifierAsync(authenticatedDevice.accountIdentifier()))
+        .map(maybeAccount -> maybeAccount.orElseThrow(Status.UNAUTHENTICATED::asRuntimeException))
+        .flatMap(account -> {
+          final Device device = account.getDevice(authenticatedDevice.deviceId())
+              .orElseThrow(Status.UNAUTHENTICATED::asRuntimeException);
+
+          final boolean tokenUnchanged =
+              Objects.equals(device.getApnId(), apnsToken) &&
+                  Objects.equals(device.getVoipApnId(), apnsVoipToken) &&
+                  Objects.equals(device.getGcmId(), fcmToken);
+
+          return tokenUnchanged
+              ? Mono.empty()
+              : Mono.fromFuture(() -> accountsManager.updateDeviceAsync(account, authenticatedDevice.deviceId(), d -> {
+                d.setApnId(apnsToken);
+                d.setVoipApnId(apnsVoipToken);
+                d.setGcmId(fcmToken);
+                d.setFetchesMessages(false);
+              }));
+        })
+        .thenReturn(SetPushTokenResponse.newBuilder().build());
+  }
+
+  @Override
+  public Mono<ClearPushTokenResponse> clearPushToken(final ClearPushTokenRequest request) {
+    final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice();
+
+    return Mono.fromFuture(() -> accountsManager.getByAccountIdentifierAsync(authenticatedDevice.accountIdentifier()))
+        .map(maybeAccount -> maybeAccount.orElseThrow(Status.UNAUTHENTICATED::asRuntimeException))
+        .flatMap(account -> Mono.fromFuture(() -> accountsManager.updateDeviceAsync(account, authenticatedDevice.deviceId(), device -> {
+          if (StringUtils.isNotBlank(device.getApnId()) || StringUtils.isNotBlank(device.getVoipApnId())) {
+            device.setUserAgent(device.isMaster() ? "OWI" : "OWP");
+          } else if (StringUtils.isNotBlank(device.getGcmId())) {
+            device.setUserAgent("OWA");
+          }
+
+          device.setApnId(null);
+          device.setVoipApnId(null);
+          device.setGcmId(null);
+          device.setFetchesMessages(true);
+        })))
+        .thenReturn(ClearPushTokenResponse.newBuilder().build());
+  }
+
+  @Override
+  public Mono<SetCapabilitiesResponse> setCapabilities(final SetCapabilitiesRequest request) {
+    final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice();
+
+    return Mono.fromFuture(() -> accountsManager.getByAccountIdentifierAsync(authenticatedDevice.accountIdentifier()))
+        .map(maybeAccount -> maybeAccount.orElseThrow(Status.UNAUTHENTICATED::asRuntimeException))
+        .flatMap(account ->
+            Mono.fromFuture(() -> accountsManager.updateDeviceAsync(account, authenticatedDevice.deviceId(),
+                d -> d.setCapabilities(new Device.DeviceCapabilities(
+                    request.getStorage(),
+                    request.getTransfer(),
+                    request.getPni(),
+                    request.getPaymentActivation())))))
+        .thenReturn(SetCapabilitiesResponse.newBuilder().build());
+  }
+}

service/src/main/proto/org/signal/chat/device.proto

@@ -0,0 +1,149 @@
+syntax = "proto3";
+
+option java_multiple_files = true;
+
+package org.signal.chat.device;
+
+/**
+ * Provides methods for working with devices attached to a Signal account.
+ */
+service Devices {
+  /**
+   * Returns a list of devices associated with the caller's account.
+   */
+  rpc GetDevices(GetDevicesRequest) returns (GetDevicesResponse) {}
+
+  /**
+   * Removes a linked device from the caller's account. This call will fail with
+   * a status of `PERMISSION_DENIED` if not called from the primary device
+   * associated with an account. It will also fail with a status of
+   * `INVALID_ARGUMENT` if the targeted device is the primary device associated
+   * with the account.
+   */
+  rpc RemoveDevice(RemoveDeviceRequest) returns (RemoveDeviceResponse) {}
+
+  rpc SetDeviceName(SetDeviceNameRequest) returns (SetDeviceNameResponse) {}
+
+  /**
+   * Sets the token(s) the server should use to send new message notifications
+   * to the authenticated device.
+   */
+  rpc SetPushToken(SetPushTokenRequest) returns (SetPushTokenResponse) {}
+
+  /**
+   * Removes any push tokens associated with the authenticated device. After
+   * calling this method, the server will assume that the authenticated device
+   * will periodically poll for new messages.
+   */
+  rpc ClearPushToken(ClearPushTokenRequest) returns (ClearPushTokenResponse) {}
+
+  /**
+   * Declares that the authenticated device supports certain features.
+   */
+  rpc SetCapabilities(SetCapabilitiesRequest) returns (SetCapabilitiesResponse) {}
+}
+
+message GetDevicesRequest {}
+
+message GetDevicesResponse {
+  message LinkedDevice {
+    /**
+     * The identifier for the device within an account.
+     */
+    uint64 id = 1;
+
+    /**
+     * A sequence of bytes that encodes an encrypted human-readable name for
+     * this device.
+     */
+    bytes name = 2;
+
+    /**
+     * The time, in milliseconds since the epoch, at which this device was
+     * attached to its parent account.
+     */
+    uint64 created = 3;
+
+    /**
+     * The approximate time, in milliseconds since the epoch, at which this
+     * device last connected to the server.
+     */
+    uint64 last_seen = 4;
+  }
+
+  /**
+   * A list of devices linked to the authenticated account.
+   */
+  repeated LinkedDevice devices = 1;
+}
+
+message RemoveDeviceRequest {
+  /**
+   * The identifier for the device to remove from the authenticated account.
+   */
+  uint64 id = 1;
+}
+
+message SetDeviceNameRequest {
+  /**
+   * A sequence of bytes that encodes an encrypted human-readable name for this
+   * device.
+   */
+  bytes name = 1;
+}
+
+message SetDeviceNameResponse {}
+
+message RemoveDeviceResponse {}
+
+message SetPushTokenRequest {
+  message ApnsTokenRequest {
+    /**
+     * A "standard" APNs device token.
+     */
+    string apns_token = 1;
+
+    /**
+     * A VoIP APNs device token. If present, the server will prefer to send
+     * message notifications to the device using this token on a VOIP APNs
+     * topic.
+     */
+    string apns_voip_token = 2;
+  }
+
+  message FcmTokenRequest {
+    /**
+     * An FCM push token.
+     */
+    string fcm_token = 1;
+  }
+
+  oneof token_request {
+    /**
+     * If present, specifies the APNs device token(s) the server will use to
+     * send new message notifications to the authenticated device.
+     */
+    ApnsTokenRequest apns_token_request = 1;
+
+    /**
+     * If present, specifies the FCM push token the server will use to send new
+     * message notifications to the authenticated device.
+     */
+    FcmTokenRequest fcm_token_request = 2;
+  }
+}
+
+message SetPushTokenResponse {}
+
+message ClearPushTokenRequest {}
+
+message ClearPushTokenResponse {}
+
+message SetCapabilitiesRequest {
+  bool storage = 1;
+  bool transfer = 2;
+  bool pni = 3;
+  bool paymentActivation = 4;
+}
+
+message SetCapabilitiesResponse {}