Add current entitlements to whoami response

This commit is contained in:
Ravi Khadiwala
2024-12-11 11:51:38 -06:00
committed by ravi-signal
parent d5b39cd496
commit 33c0a27b85
6 changed files with 165 additions and 3 deletions

View File

@@ -42,6 +42,7 @@ import java.util.List;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.function.BiConsumer;
import java.util.stream.Stream;
import org.glassfish.jersey.server.ServerProperties;
import org.glassfish.jersey.test.grizzly.GrizzlyWebTestContainerFactory;
@@ -66,6 +67,7 @@ import org.whispersystems.textsecuregcm.entities.ApnRegistrationId;
import org.whispersystems.textsecuregcm.entities.ConfirmUsernameHashRequest;
import org.whispersystems.textsecuregcm.entities.DeviceName;
import org.whispersystems.textsecuregcm.entities.EncryptedUsername;
import org.whispersystems.textsecuregcm.entities.Entitlements;
import org.whispersystems.textsecuregcm.entities.GcmRegistrationId;
import org.whispersystems.textsecuregcm.entities.RegistrationLock;
import org.whispersystems.textsecuregcm.entities.ReserveUsernameHashRequest;
@@ -82,6 +84,7 @@ import org.whispersystems.textsecuregcm.mappers.JsonMappingExceptionMapper;
import org.whispersystems.textsecuregcm.mappers.NonNormalizedPhoneNumberExceptionMapper;
import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountBadge;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.DeviceCapability;
@@ -324,6 +327,15 @@ class AccountControllerTest {
@ParameterizedTest
@ValueSource(strings = {"/v1/accounts/whoami", "/v1/accounts/me"})
void testWhoAmI(final String path) {
final Instant expiration = Instant.now().plus(Duration.ofHours(1)).plusMillis(101);
final Instant truncatedExpiration = Instant.ofEpochSecond(expiration.getEpochSecond());
final AccountBadge badge1 = new AccountBadge("badge1", expiration, true);
final AccountBadge badge2 = new AccountBadge("badge2", expiration, true);
when(AuthHelper.VALID_ACCOUNT.getBackupVoucher())
.thenReturn(new Account.BackupVoucher(100, expiration));
when(AuthHelper.VALID_ACCOUNT.getBadges()).thenReturn(List.of(badge1, badge2));
try (final Response response = resources.getJerseyTest()
.target(path)
.request()
@@ -331,7 +343,19 @@ class AccountControllerTest {
.get()) {
assertThat(response.getStatus()).isEqualTo(200);
assertThat(response.readEntity(AccountIdentityResponse.class).uuid()).isEqualTo(AuthHelper.VALID_UUID);
final AccountIdentityResponse identityResponse = response.readEntity(AccountIdentityResponse.class);
assertThat(identityResponse.uuid()).isEqualTo(AuthHelper.VALID_UUID);
final BiConsumer<Entitlements.BadgeEntitlement, AccountBadge> compareBadge = (actual, expected) -> {
assertThat(actual.expiration()).isEqualTo(truncatedExpiration);
assertThat(actual.id()).isEqualTo(expected.id());
assertThat(actual.visible()).isEqualTo(expected.visible());
};
compareBadge.accept(identityResponse.entitlements().badges().getFirst(), badge1);
compareBadge.accept(identityResponse.entitlements().badges().getLast(), badge2);
assertThat(identityResponse.entitlements().backup().backupLevel()).isEqualTo(100);
assertThat(identityResponse.entitlements().backup().expiration()).isEqualTo(truncatedExpiration);
}
}

View File

@@ -0,0 +1,67 @@
/*
* Copyright 2024 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.controllers;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import java.time.Instant;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.whispersystems.textsecuregcm.entities.Entitlements;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountBadge;
import org.whispersystems.textsecuregcm.util.TestClock;
class AccountIdentityResponseBuilderTest {
@Test
void expiredBackupEntitlement() {
final Instant expiration = Instant.ofEpochSecond(101);
final Account account = mock(Account.class);
when(account.getBackupVoucher()).thenReturn(new Account.BackupVoucher(6, expiration));
Entitlements.BackupEntitlement backup = new AccountIdentityResponseBuilder(account)
.clock(TestClock.pinned(Instant.ofEpochSecond(101)))
.build().entitlements().backup();
assertThat(backup).isNull();
backup = new AccountIdentityResponseBuilder(account)
.clock(TestClock.pinned(Instant.ofEpochSecond(100)))
.build().entitlements().backup();
assertThat(backup).isNotNull();
assertThat(backup.expiration()).isEqualTo(expiration);
assertThat(backup.backupLevel()).isEqualTo(6);
}
@Test
void expiredBadgeEntitlement() {
final Account account = mock(Account.class);
when(account.getBadges()).thenReturn(List.of(
new AccountBadge("badge1", Instant.ofEpochSecond(10), false),
new AccountBadge("badge2", Instant.ofEpochSecond(11), true)));
// all should be expired
assertThat(new AccountIdentityResponseBuilder(account)
.clock(TestClock.pinned(Instant.ofEpochSecond(11)))
.build().entitlements().badges()).isEmpty();
// first badge should be expired
assertThat(new AccountIdentityResponseBuilder(account).clock(TestClock.pinned(Instant.ofEpochSecond(10))).build()
.entitlements()
.badges()
.stream().map(Entitlements.BadgeEntitlement::id).toList())
.containsExactly("badge2");
// no badges should be expired
assertThat(new AccountIdentityResponseBuilder(account).clock(TestClock.pinned(Instant.ofEpochSecond(9))).build()
.entitlements()
.badges()
.stream().map(Entitlements.BadgeEntitlement::id).toList())
.containsExactly("badge1", "badge2");
}
}

View File

@@ -152,6 +152,7 @@ public class AccountsHelper {
case "getIdentityKey" ->
when(updatedAccount.getIdentityKey(stubbing.getInvocation().getArgument(0))).thenAnswer(stubbing);
case "getBadges" -> when(updatedAccount.getBadges()).thenAnswer(stubbing);
case "getBackupVoucher" -> when(updatedAccount.getBackupVoucher()).thenAnswer(stubbing);
case "getLastSeen" -> when(updatedAccount.getLastSeen()).thenAnswer(stubbing);
case "hasLockedCredentials" -> when(updatedAccount.hasLockedCredentials()).thenAnswer(stubbing);
default -> throw new IllegalArgumentException("unsupported method: Account#" + stubbing.getInvocation().getMethod().getName());