Move "remove device" logic into AccountsManager

This commit is contained in:
Jon Chambers
2023-11-29 12:07:57 -05:00
committed by Jon Chambers
parent 4f42c10d60
commit 37e3bcfc3e
9 changed files with 88 additions and 58 deletions

View File

@@ -135,8 +135,7 @@ public class DeviceController {
@Produces(MediaType.APPLICATION_JSON)
@Path("/{device_id}")
@ChangesDeviceEnabledState
public void removeDevice(@Auth AuthenticatedAccount auth, @PathParam("device_id") byte deviceId) {
Account account = auth.getAccount();
public CompletableFuture<Response> removeDevice(@Auth AuthenticatedAccount auth, @PathParam("device_id") byte deviceId) {
if (auth.getAuthenticatedDevice().getId() != Device.PRIMARY_ID) {
throw new WebApplicationException(Response.Status.UNAUTHORIZED);
}
@@ -145,14 +144,7 @@ public class DeviceController {
throw new ForbiddenException();
}
final CompletableFuture<Void> deleteKeysFuture = keys.delete(account.getUuid(), deviceId);
messages.clear(account.getUuid(), deviceId).join();
account = accounts.update(account, a -> a.removeDevice(deviceId));
// ensure any messages that came in after the first clear() are also removed
messages.clear(account.getUuid(), deviceId).join();
deleteKeysFuture.join();
return accounts.removeDevice(auth.getAccount(), deviceId).thenApply(Util.ASYNC_EMPTY_RESPONSE);
}
@GET

View File

@@ -78,19 +78,14 @@ public class DevicesGrpcService extends ReactorDevicesGrpc.DevicesImplBase {
if (request.getId() == Device.PRIMARY_ID) {
throw Status.INVALID_ARGUMENT.withDescription("Cannot remove primary device").asRuntimeException();
}
final byte deviceId = DeviceIdUtil.validate(request.getId());
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(), deviceId)),
Mono.fromFuture(() -> keysManager.delete(account.getUuid(), deviceId)))
.then(Mono.fromFuture(() -> accountsManager.updateAsync(account, a -> a.removeDevice(deviceId))))
// 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(), deviceId))))
.flatMap(account -> Mono.fromFuture(accountsManager.removeDevice(account, deviceId)))
.thenReturn(RemoveDeviceResponse.newBuilder().build());
}

View File

@@ -296,6 +296,25 @@ public class AccountsManager {
}
}
public CompletableFuture<Account> removeDevice(final Account account, final byte deviceId) {
if (deviceId == Device.PRIMARY_ID) {
throw new IllegalArgumentException("Cannot remove primary device");
}
return CompletableFuture.allOf(
keysManager.delete(account.getUuid(), deviceId),
messagesManager.clear(account.getUuid(), deviceId))
.thenCompose(ignored -> updateAsync(account, (Consumer<Account>) a -> a.removeDevice(deviceId)))
// ensure any messages that came in after the first clear() are also removed
.thenCompose(updatedAccount -> messagesManager.clear(account.getUuid(), deviceId)
.thenApply(ignored -> updatedAccount))
.whenComplete((ignored, throwable) -> {
if (throwable == null) {
clientPresenceManager.disconnectPresence(account.getUuid(), deviceId);
}
});
}
public Account changeNumber(final Account account,
final String targetNumber,
@Nullable final IdentityKey pniIdentityKey,

View File

@@ -16,11 +16,6 @@ import com.codahale.metrics.Timer;
import com.google.common.annotations.VisibleForTesting;
import io.dropwizard.lifecycle.Managed;
import io.micrometer.core.instrument.Counter;
import reactor.core.publisher.Flux;
import reactor.util.function.Tuple2;
import reactor.util.function.Tuples;
import software.amazon.awssdk.services.dynamodb.model.ItemCollectionSizeLimitExceededException;
import java.time.Duration;
import java.time.Instant;
import java.util.Comparator;
@@ -29,12 +24,7 @@ import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
@@ -42,6 +32,10 @@ import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
import org.whispersystems.textsecuregcm.util.Constants;
import org.whispersystems.textsecuregcm.util.Util;
import reactor.core.publisher.Flux;
import reactor.util.function.Tuple2;
import reactor.util.function.Tuples;
import software.amazon.awssdk.services.dynamodb.model.ItemCollectionSizeLimitExceededException;
public class MessagePersister implements Managed {
@@ -276,7 +270,7 @@ public class MessagePersister implements Managed {
messagesManager.clear(account.getUuid(), deviceToDelete.getId()))
.orTimeout((UNLINK_TIMEOUT.toSeconds() * 3) / 4, TimeUnit.SECONDS)
.join();
accountsManager.update(updatedAccount, a -> a.removeDevice(deviceToDelete.getId()));
accountsManager.removeDevice(updatedAccount, deviceToDelete.getId()).join();
} finally {
messagesCache.unlockAccountForMessagePersisterCleanup(account.getUuid());
if (deviceToDelete.getId() != destinationDeviceId) { // no point in persisting a device we just purged

View File

@@ -71,18 +71,7 @@ public class UnlinkDeviceCommand extends EnvironmentCommand<WhisperServerConfigu
for (byte deviceId : deviceIds) {
/** see {@link org.whispersystems.textsecuregcm.controllers.DeviceController#removeDevice} */
System.out.format("Removing device %s::%d\n", aci, deviceId);
account = deps.accountsManager().update(account, a -> a.removeDevice(deviceId));
System.out.format("Removing keys for device %s::%d\n", aci, deviceId);
deps.keysManager().delete(account.getUuid(), deviceId).join();
System.out.format("Clearing additional messages for %s::%d\n", aci, deviceId);
deps.messagesManager().clear(account.getUuid(), deviceId).join();
System.out.format("Clearing presence state for %s::%d\n", aci, deviceId);
deps.clientPresenceManager().disconnectPresence(aci, deviceId);
System.out.format("Device %s::%d successfully removed\n", aci, deviceId);
deps.accountsManager().removeDevice(account, deviceId).join();
}
} finally {
commandStopListener.stop();