From 5892dc71fa548cb4491b10274566f2e2b47a392e Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Thu, 15 Aug 2024 15:16:28 -0400
Subject: [PATCH] Allow linked devices to unlink themselves via the gRPC API

---
 .../grpc/DevicesGrpcService.java              | 22 +++++-------
 .../grpc/DevicesGrpcServiceTest.java          | 35 ++++++++++---------
 2 files changed, 27 insertions(+), 30 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/DevicesGrpcService.java b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/DevicesGrpcService.java
index 9954e36a3..5f687465e 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/DevicesGrpcService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/DevicesGrpcService.java
@@ -7,7 +7,6 @@
 
 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;
@@ -28,26 +27,17 @@
 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) {
-
+  public DevicesGrpcService(final AccountsManager accountsManager) {
     this.accountsManager = accountsManager;
-    this.keysManager = keysManager;
-    this.messagesManager = messagesManager;
   }
 
   @Override
@@ -79,9 +69,15 @@ public Mono<RemoveDeviceResponse> removeDevice(final RemoveDeviceRequest request
       throw Status.INVALID_ARGUMENT.withDescription("Cannot remove primary device").asRuntimeException();
     }
 
-    final byte deviceId = DeviceIdUtil.validate(request.getId());
+    final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice();
 
-    final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedPrimaryDevice();
+    if (authenticatedDevice.deviceId() != Device.PRIMARY_ID && request.getId() != authenticatedDevice.deviceId()) {
+      throw Status.PERMISSION_DENIED
+          .withDescription("Linked devices cannot remove devices other than themselves")
+          .asRuntimeException();
+    }
+
+    final byte deviceId = DeviceIdUtil.validate(request.getId());
 
     return Mono.fromFuture(() -> accountsManager.getByAccountIdentifierAsync(authenticatedDevice.accountIdentifier()))
         .map(maybeAccount -> maybeAccount.orElseThrow(Status.UNAUTHENTICATED::asRuntimeException))
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/DevicesGrpcServiceTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/DevicesGrpcServiceTest.java
index 277112725..a5968aa38 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/DevicesGrpcServiceTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/DevicesGrpcServiceTest.java
@@ -50,8 +50,6 @@
 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.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.util.TestRandomUtil;
 
 class DevicesGrpcServiceTest extends SimpleBaseGrpcTest<DevicesGrpcService, DevicesGrpc.DevicesBlockingStub> {
@@ -59,16 +57,9 @@ class DevicesGrpcServiceTest extends SimpleBaseGrpcTest<DevicesGrpcService, Devi
   @Mock
   private AccountsManager accountsManager;
 
-  @Mock
-  private KeysManager keysManager;
-
-  @Mock
-  private MessagesManager messagesManager;
-
   @Mock
   private Account authenticatedAccount;
 
-
   @Override
   protected DevicesGrpcService createServiceBeforeEachTest() {
     when(authenticatedAccount.getUuid()).thenReturn(AUTHENTICATED_ACI);
@@ -76,6 +67,9 @@ protected DevicesGrpcService createServiceBeforeEachTest() {
     when(accountsManager.getByAccountIdentifierAsync(AUTHENTICATED_ACI))
         .thenReturn(CompletableFuture.completedFuture(Optional.of(authenticatedAccount)));
 
+    when(accountsManager.removeDevice(any(), anyByte()))
+        .thenReturn(CompletableFuture.completedFuture(authenticatedAccount));
+
     when(accountsManager.updateAsync(any(), any()))
         .thenAnswer(invocation -> {
           final Account account = invocation.getArgument(0);
@@ -97,10 +91,7 @@ protected DevicesGrpcService createServiceBeforeEachTest() {
           return CompletableFuture.completedFuture(account);
         });
 
-    when(keysManager.deleteSingleUsePreKeys(any(), anyByte())).thenReturn(CompletableFuture.completedFuture(null));
-    when(messagesManager.clear(any(), anyByte())).thenReturn(CompletableFuture.completedFuture(null));
-
-    return new DevicesGrpcService(accountsManager, keysManager, messagesManager);
+    return new DevicesGrpcService(accountsManager);
   }
 
   @Test
@@ -146,9 +137,6 @@ void getDevices() {
   void removeDevice() {
     final byte deviceId = 17;
 
-    when(accountsManager.removeDevice(any(), anyByte()))
-        .thenReturn(CompletableFuture.completedFuture(authenticatedAccount));
-
     final RemoveDeviceResponse ignored = authenticatedServiceStub().removeDevice(RemoveDeviceRequest.newBuilder()
         .setId(deviceId)
         .build());
@@ -166,7 +154,20 @@ void removeDevicePrimary() {
   }
 
   @Test
-  void removeDeviceNonPrimaryAuthenticated() {
+  void removeDeviceNonPrimaryMatchAuthenticated() {
+    final byte deviceId = Device.PRIMARY_ID + 1;
+
+    mockAuthenticationInterceptor().setAuthenticatedDevice(AUTHENTICATED_ACI, deviceId);
+
+    final RemoveDeviceResponse ignored = authenticatedServiceStub().removeDevice(RemoveDeviceRequest.newBuilder()
+        .setId(deviceId)
+        .build());
+
+    verify(accountsManager).removeDevice(authenticatedAccount, deviceId);
+  }
+
+  @Test
+  void removeDeviceNonPrimaryMismatchAuthenticated() {
     mockAuthenticationInterceptor().setAuthenticatedDevice(AUTHENTICATED_ACI, (byte) (Device.PRIMARY_ID + 1));
     assertStatusException(Status.PERMISSION_DENIED, () -> authenticatedServiceStub().removeDevice(RemoveDeviceRequest.newBuilder()
         .setId(17)