Use simple-grpc in KeysGrpcService and KeysAnonymousGrpcService

This commit is contained in:
Jon Chambers
2026-02-27 16:49:15 -05:00
committed by GitHub
parent c63cbc304f
commit 065809e34f
5 changed files with 216 additions and 244 deletions

View File

@@ -98,8 +98,8 @@ class KeysAnonymousGrpcServiceTest extends SimpleBaseGrpcTest<KeysAnonymousGrpcS
when(targetAccount.getUnidentifiedAccessKey()).thenReturn(Optional.of(unidentifiedAccessKey));
when(targetAccount.getIdentifier(IdentityType.ACI)).thenReturn(uuid);
when(targetAccount.getIdentityKey(IdentityType.ACI)).thenReturn(identityKey);
when(accountsManager.getByServiceIdentifierAsync(identifier))
.thenReturn(CompletableFuture.completedFuture(Optional.of(targetAccount)));
when(accountsManager.getByServiceIdentifier(identifier))
.thenReturn(Optional.of(targetAccount));
final ECPreKey ecPreKey = new ECPreKey(1, ECKeyPair.generate().getPublicKey());
final ECSignedPreKey ecSignedPreKey = KeysHelper.signedECPreKey(2, identityKeyPair);
@@ -146,8 +146,8 @@ class KeysAnonymousGrpcServiceTest extends SimpleBaseGrpcTest<KeysAnonymousGrpcS
when(targetAccount.getUnidentifiedAccessKey()).thenReturn(Optional.of(unidentifiedAccessKey));
when(targetAccount.getIdentifier(IdentityType.ACI)).thenReturn(uuid);
when(targetAccount.getIdentityKey(IdentityType.ACI)).thenReturn(identityKey);
when(accountsManager.getByServiceIdentifierAsync(identifier))
.thenReturn(CompletableFuture.completedFuture(Optional.of(targetAccount)));
when(accountsManager.getByServiceIdentifier(identifier))
.thenReturn(Optional.of(targetAccount));
final ECPreKey ecPreKey = new ECPreKey(1, ECKeyPair.generate().getPublicKey());
final ECSignedPreKey ecSignedPreKey = KeysHelper.signedECPreKey(2, identityKeyPair);
@@ -204,8 +204,8 @@ class KeysAnonymousGrpcServiceTest extends SimpleBaseGrpcTest<KeysAnonymousGrpcS
final byte[] unidentifiedAccessKey = TestRandomUtil.nextBytes(UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH);
when(targetAccount.getUnidentifiedAccessKey()).thenReturn(Optional.of(unidentifiedAccessKey));
when(accountsManager.getByServiceIdentifierAsync(identifier))
.thenReturn(CompletableFuture.completedFuture(Optional.of(targetAccount)));
when(accountsManager.getByServiceIdentifier(identifier))
.thenReturn(Optional.of(targetAccount));
final GetPreKeysAnonymousResponse response = unauthenticatedServiceStub().getPreKeys(
GetPreKeysAnonymousRequest.newBuilder()
@@ -252,7 +252,6 @@ class KeysAnonymousGrpcServiceTest extends SimpleBaseGrpcTest<KeysAnonymousGrpcS
final Instant expiration = Instant.now().truncatedTo(ChronoUnit.DAYS);
CLOCK.pin(expiration.minus(Duration.ofHours(1))); // set time so the credential isn't expired yet
final AciServiceIdentifier wrongAci = new AciServiceIdentifier(UUID.randomUUID());
final byte[] token = AuthHelper.validGroupSendToken(SERVER_SECRET_PARAMS, List.of(authorizedIdentifier), expiration);
final GetPreKeysAnonymousResponse response = unauthenticatedServiceStub().getPreKeys(
@@ -270,8 +269,8 @@ class KeysAnonymousGrpcServiceTest extends SimpleBaseGrpcTest<KeysAnonymousGrpcS
@Test
void getPreKeysAccountNotFoundUnidentifiedAccessKey() {
final AciServiceIdentifier nonexistentAci = new AciServiceIdentifier(UUID.randomUUID());
when(accountsManager.getByServiceIdentifierAsync(nonexistentAci))
.thenReturn(CompletableFuture.completedFuture(Optional.empty()));
when(accountsManager.getByServiceIdentifier(nonexistentAci))
.thenReturn(Optional.empty());
final GetPreKeysAnonymousResponse preKeysResponse =
unauthenticatedServiceStub().getPreKeys(GetPreKeysAnonymousRequest.newBuilder()
@@ -293,8 +292,8 @@ class KeysAnonymousGrpcServiceTest extends SimpleBaseGrpcTest<KeysAnonymousGrpcS
final byte[] token = AuthHelper.validGroupSendToken(SERVER_SECRET_PARAMS, List.of(nonexistentAci), expiration);
when(accountsManager.getByServiceIdentifierAsync(nonexistentAci))
.thenReturn(CompletableFuture.completedFuture(Optional.empty()));
when(accountsManager.getByServiceIdentifier(nonexistentAci))
.thenReturn(Optional.empty());
final GetPreKeysAnonymousResponse preKeysResponse =
unauthenticatedServiceStub().getPreKeys(GetPreKeysAnonymousRequest.newBuilder()
@@ -319,8 +318,8 @@ class KeysAnonymousGrpcServiceTest extends SimpleBaseGrpcTest<KeysAnonymousGrpcS
when(targetAccount.getDevice(anyByte())).thenReturn(Optional.empty());
when(targetAccount.getUnidentifiedAccessKey()).thenReturn(Optional.of(unidentifiedAccessKey));
when(accountsManager.getByServiceIdentifierAsync(new AciServiceIdentifier(accountIdentifier)))
.thenReturn(CompletableFuture.completedFuture(Optional.of(targetAccount)));
when(accountsManager.getByServiceIdentifier(new AciServiceIdentifier(accountIdentifier)))
.thenReturn(Optional.of(targetAccount));
final GetPreKeysAnonymousResponse response = unauthenticatedServiceStub().getPreKeys(
GetPreKeysAnonymousRequest.newBuilder()

View File

@@ -12,6 +12,7 @@ import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyByte;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
@@ -31,7 +32,6 @@ import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.stream.Stream;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
@@ -115,8 +115,8 @@ class KeysGrpcServiceTest extends SimpleBaseGrpcTest<KeysGrpcService, KeysGrpc.K
when(accountsManager.getByAccountIdentifier(AUTHENTICATED_ACI)).thenReturn(Optional.of(authenticatedAccount));
when(accountsManager.getByPhoneNumberIdentifier(AUTHENTICATED_PNI)).thenReturn(Optional.of(authenticatedAccount));
when(accountsManager.getByAccountIdentifierAsync(AUTHENTICATED_ACI)).thenReturn(CompletableFuture.completedFuture(Optional.of(authenticatedAccount)));
when(accountsManager.getByPhoneNumberIdentifierAsync(AUTHENTICATED_PNI)).thenReturn(CompletableFuture.completedFuture(Optional.of(authenticatedAccount)));
when(accountsManager.getByAccountIdentifier(AUTHENTICATED_ACI)).thenReturn(Optional.of(authenticatedAccount));
when(accountsManager.getByPhoneNumberIdentifier(AUTHENTICATED_PNI)).thenReturn(Optional.of(authenticatedAccount));
return new KeysGrpcService(accountsManager, keysManager, rateLimiters);
}
@@ -296,16 +296,6 @@ class KeysGrpcServiceTest extends SimpleBaseGrpcTest<KeysGrpcService, KeysGrpc.K
@ParameterizedTest
@EnumSource(value = org.signal.chat.common.IdentityType.class, names = {"IDENTITY_TYPE_ACI", "IDENTITY_TYPE_PNI"})
void setSignedPreKey(final org.signal.chat.common.IdentityType identityType) {
when(accountsManager.updateDeviceAsync(any(), anyByte(), any())).thenAnswer(invocation -> {
final Account account = invocation.getArgument(0);
final byte deviceId = invocation.getArgument(1);
final Consumer<Device> deviceUpdater = invocation.getArgument(2);
account.getDevice(deviceId).ifPresent(deviceUpdater);
return CompletableFuture.completedFuture(account);
});
when(keysManager.storeEcSignedPreKeys(any(), anyByte(), any())).thenReturn(CompletableFuture.completedFuture(null));
final ECKeyPair identityKeyPair = switch (IdentityTypeUtil.fromGrpcIdentityType(identityType)) {
@@ -476,8 +466,8 @@ class KeysGrpcServiceTest extends SimpleBaseGrpcTest<KeysGrpcService, KeysGrpc.K
when(targetAccount.getUuid()).thenReturn(UUID.randomUUID());
when(targetAccount.getIdentifier(identityType)).thenReturn(identifier);
when(targetAccount.getIdentityKey(identityType)).thenReturn(identityKey);
when(accountsManager.getByServiceIdentifierAsync(serviceIdentifier))
.thenReturn(CompletableFuture.completedFuture(Optional.of(targetAccount)));
when(accountsManager.getByServiceIdentifier(serviceIdentifier))
.thenReturn(Optional.of(targetAccount));
final Map<Byte, KeysManager.DevicePreKeys> devicePreKeysMap = new HashMap<>();
@@ -573,8 +563,8 @@ class KeysGrpcServiceTest extends SimpleBaseGrpcTest<KeysGrpcService, KeysGrpc.K
@Test
void getPreKeysAccountNotFound() {
when(accountsManager.getByServiceIdentifierAsync(any()))
.thenReturn(CompletableFuture.completedFuture(Optional.empty()));
when(accountsManager.getByServiceIdentifier(any()))
.thenReturn(Optional.empty());
final GetPreKeysResponse response = authenticatedServiceStub().getPreKeys(GetPreKeysRequest.newBuilder()
.setTargetIdentifier(ServiceIdentifier.newBuilder()
@@ -595,8 +585,8 @@ class KeysGrpcServiceTest extends SimpleBaseGrpcTest<KeysGrpcService, KeysGrpc.K
when(targetAccount.getDevices()).thenReturn(Collections.emptyList());
when(targetAccount.getDevice(anyByte())).thenReturn(Optional.empty());
when(accountsManager.getByServiceIdentifierAsync(new AciServiceIdentifier(accountIdentifier)))
.thenReturn(CompletableFuture.completedFuture(Optional.of(targetAccount)));
when(accountsManager.getByServiceIdentifier(new AciServiceIdentifier(accountIdentifier)))
.thenReturn(Optional.of(targetAccount));
final GetPreKeysResponse response = authenticatedServiceStub().getPreKeys(GetPreKeysRequest.newBuilder()
.setTargetIdentifier(ServiceIdentifier.newBuilder()
@@ -609,19 +599,20 @@ class KeysGrpcServiceTest extends SimpleBaseGrpcTest<KeysGrpcService, KeysGrpc.K
}
@Test
void getPreKeysRateLimited() {
void getPreKeysRateLimited() throws RateLimitExceededException {
final Account targetAccount = mock(Account.class);
when(targetAccount.getUuid()).thenReturn(UUID.randomUUID());
when(targetAccount.getIdentityKey(IdentityType.ACI)).thenReturn(new IdentityKey(ECKeyPair.generate().getPublicKey()));
when(targetAccount.getDevices()).thenReturn(Collections.emptyList());
when(targetAccount.getDevice(anyByte())).thenReturn(Optional.empty());
when(accountsManager.getByServiceIdentifierAsync(any()))
.thenReturn(CompletableFuture.completedFuture(Optional.of(targetAccount)));
when(accountsManager.getByServiceIdentifier(any()))
.thenReturn(Optional.of(targetAccount));
final Duration retryAfterDuration = Duration.ofMinutes(7);
when(preKeysRateLimiter.validateReactive(anyString()))
.thenReturn(Mono.error(new RateLimitExceededException(retryAfterDuration)));
doThrow(new RateLimitExceededException(retryAfterDuration))
.when(preKeysRateLimiter).validate(anyString());
assertRateLimitExceeded(retryAfterDuration, () -> authenticatedServiceStub().getPreKeys(GetPreKeysRequest.newBuilder()
.setTargetIdentifier(ServiceIdentifier.newBuilder()