/*
 * Copyright 2023 Signal Messenger, LLC
 * SPDX-License-Identifier: AGPL-3.0-only
 */

package org.whispersystems.textsecuregcm.grpc;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyByte;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;
import static org.whispersystems.textsecuregcm.grpc.GrpcTestUtils.assertStatusException;

import com.google.protobuf.ByteString;
import com.google.protobuf.Empty;
import io.grpc.Status;
import io.grpc.stub.StreamObserver;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.time.Duration;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junitpioneer.jupiter.cartesian.CartesianTest;
import org.mockito.Mock;
import org.signal.chat.common.EcPreKey;
import org.signal.chat.common.EcSignedPreKey;
import org.signal.chat.common.KemSignedPreKey;
import org.signal.chat.common.ServiceIdentifier;
import org.signal.chat.keys.AccountPreKeyBundles;
import org.signal.chat.keys.CheckIdentityKeyRequest;
import org.signal.chat.keys.CheckIdentityKeyResponse;
import org.signal.chat.keys.DevicePreKeyBundle;
import org.signal.chat.keys.GetPreKeysAnonymousRequest;
import org.signal.chat.keys.GetPreKeysAnonymousResponse;
import org.signal.chat.keys.GetPreKeysRequest;
import org.signal.chat.keys.KeysAnonymousGrpc;
import org.signal.libsignal.protocol.IdentityKey;
import org.signal.libsignal.protocol.InvalidKeyException;
import org.signal.libsignal.protocol.ecc.ECKeyPair;
import org.signal.libsignal.zkgroup.ServerSecretParams;
import org.whispersystems.textsecuregcm.auth.UnidentifiedAccessUtil;
import org.whispersystems.textsecuregcm.entities.ECPreKey;
import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
import org.whispersystems.textsecuregcm.identity.IdentityType;
import org.whispersystems.textsecuregcm.identity.PniServiceIdentifier;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.KeyIdUtil;
import org.whispersystems.textsecuregcm.storage.KeysManager;
import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
import org.whispersystems.textsecuregcm.tests.util.DevicesHelper;
import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
import org.whispersystems.textsecuregcm.util.TestClock;
import org.whispersystems.textsecuregcm.util.TestRandomUtil;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
import org.whispersystems.textsecuregcm.util.Util;

class KeysAnonymousGrpcServiceTest extends SimpleBaseGrpcTest<KeysAnonymousGrpcService, KeysAnonymousGrpc.KeysAnonymousBlockingStub> {

  private static final ServerSecretParams SERVER_SECRET_PARAMS = ServerSecretParams.generate();
  private static final TestClock CLOCK = TestClock.now();

  @Mock
  private AccountsManager accountsManager;

  @Mock
  private KeysManager keysManager;

  @Override
  protected KeysAnonymousGrpcService createServiceBeforeEachTest() {
    return new KeysAnonymousGrpcService(accountsManager, keysManager, SERVER_SECRET_PARAMS, CLOCK);
  }

  @Test
  void getPreKeysUnidentifiedAccessKey() {
    final Account targetAccount = mock(Account.class);

    final Device targetDevice = DevicesHelper.createDevice(Device.PRIMARY_ID);
    when(targetAccount.getDevice(Device.PRIMARY_ID)).thenReturn(Optional.of(targetDevice));

    final ECKeyPair identityKeyPair = ECKeyPair.generate();
    final IdentityKey identityKey = new IdentityKey(identityKeyPair.getPublicKey());
    final UUID uuid = UUID.randomUUID();
    final AciServiceIdentifier identifier = new AciServiceIdentifier(uuid);
    final byte[] unidentifiedAccessKey = TestRandomUtil.nextBytes(UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH);

    when(targetAccount.getUnidentifiedAccessKey()).thenReturn(Optional.of(unidentifiedAccessKey));
    when(targetAccount.getIdentifier(IdentityType.ACI)).thenReturn(uuid);
    when(targetAccount.getIdentityKey(IdentityType.ACI)).thenReturn(identityKey);
    when(accountsManager.getByServiceIdentifier(identifier))
        .thenReturn(Optional.of(targetAccount));

    final ECPreKey ecPreKey = new ECPreKey(1, ECKeyPair.generate().getPublicKey());
    final ECSignedPreKey ecSignedPreKey = KeysHelper.signedECPreKey(2, identityKeyPair);
    final KEMSignedPreKey kemSignedPreKey = KeysHelper.signedKEMPreKey(3, identityKeyPair);
    final KeysManager.DevicePreKeys devicePreKeys =
        new KeysManager.DevicePreKeys(ecSignedPreKey, Optional.of(ecPreKey), kemSignedPreKey);

    when(keysManager.takeDevicePreKeys(eq(Device.PRIMARY_ID), eq(identifier), any()))
        .thenReturn(CompletableFuture.completedFuture(Optional.of(devicePreKeys)));

    final GetPreKeysAnonymousResponse response = unauthenticatedServiceStub().getPreKeys(GetPreKeysAnonymousRequest.newBuilder()
        .setUnidentifiedAccessKey(ByteString.copyFrom(unidentifiedAccessKey))
        .setRequest(GetPreKeysRequest.newBuilder()
            .setTargetIdentifier(ServiceIdentifierUtil.toGrpcServiceIdentifier(identifier))
            .setDeviceId(Device.PRIMARY_ID))
        .build());

    final GetPreKeysAnonymousResponse expectedResponse = GetPreKeysAnonymousResponse.newBuilder()
        .setPreKeys(AccountPreKeyBundles.newBuilder()
            .setIdentityKey(ByteString.copyFrom(identityKey.serialize()))
            .putDevicePreKeys(Device.PRIMARY_ID, DevicePreKeyBundle.newBuilder()
                .setEcOneTimePreKey(toGrpcEcPreKey(ecPreKey))
                .setEcSignedPreKey(toGrpcEcSignedPreKey(ecSignedPreKey))
                .setKemOneTimePreKey(toGrpcKemSignedPreKey(kemSignedPreKey))
                .build()))
        .build();

    assertEquals(expectedResponse, response);
  }

  @Test
  void getPreKeysGroupSendEndorsement() throws Exception {
    final Account targetAccount = mock(Account.class);

    final Device targetDevice = DevicesHelper.createDevice(Device.PRIMARY_ID);
    when(targetAccount.getDevice(Device.PRIMARY_ID)).thenReturn(Optional.of(targetDevice));

    final ECKeyPair identityKeyPair = ECKeyPair.generate();
    final IdentityKey identityKey = new IdentityKey(identityKeyPair.getPublicKey());
    final UUID uuid = UUID.randomUUID();
    final AciServiceIdentifier identifier = new AciServiceIdentifier(uuid);
    final byte[] unidentifiedAccessKey = TestRandomUtil.nextBytes(UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH);

    when(targetAccount.getUnidentifiedAccessKey()).thenReturn(Optional.of(unidentifiedAccessKey));
    when(targetAccount.getIdentifier(IdentityType.ACI)).thenReturn(uuid);
    when(targetAccount.getIdentityKey(IdentityType.ACI)).thenReturn(identityKey);
    when(accountsManager.getByServiceIdentifier(identifier))
        .thenReturn(Optional.of(targetAccount));

    final ECPreKey ecPreKey = new ECPreKey(1, ECKeyPair.generate().getPublicKey());
    final ECSignedPreKey ecSignedPreKey = KeysHelper.signedECPreKey(2, identityKeyPair);
    final KEMSignedPreKey kemSignedPreKey = KeysHelper.signedKEMPreKey(3, identityKeyPair);
    final KeysManager.DevicePreKeys devicePreKeys =
        new KeysManager.DevicePreKeys(ecSignedPreKey, Optional.of(ecPreKey), kemSignedPreKey);

    when(keysManager.takeDevicePreKeys(eq(Device.PRIMARY_ID), eq(identifier), any()))
        .thenReturn(CompletableFuture.completedFuture(Optional.of(devicePreKeys)));

    // Expirations must be on day boundaries or libsignal will refuse to create or verify the token
    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 byte[] token = AuthHelper.validGroupSendToken(SERVER_SECRET_PARAMS, List.of(identifier), expiration);

    final GetPreKeysAnonymousResponse response = unauthenticatedServiceStub().getPreKeys(GetPreKeysAnonymousRequest.newBuilder()
        .setGroupSendToken(ByteString.copyFrom(token))
        .setRequest(GetPreKeysRequest.newBuilder()
            .setTargetIdentifier(ServiceIdentifierUtil.toGrpcServiceIdentifier(identifier))
            .setDeviceId(Device.PRIMARY_ID))
        .build());

    final GetPreKeysAnonymousResponse expectedResponse = GetPreKeysAnonymousResponse.newBuilder()
        .setPreKeys(AccountPreKeyBundles.newBuilder()
            .setIdentityKey(ByteString.copyFrom(identityKey.serialize()))
            .putDevicePreKeys(Device.PRIMARY_ID, DevicePreKeyBundle.newBuilder()
                .setEcOneTimePreKey(toGrpcEcPreKey(ecPreKey))
                .setEcSignedPreKey(toGrpcEcSignedPreKey(ecSignedPreKey))
                .setKemOneTimePreKey(toGrpcKemSignedPreKey(kemSignedPreKey))
                .build()))
        .build();

    assertEquals(expectedResponse, response);
  }

  @CartesianTest
  void getPreKeysUnrestricted(@CartesianTest.Values(booleans = {true, false}) boolean includeUak) {
    final Account targetAccount = mock(Account.class);

    final Device targetDevice = DevicesHelper.createDevice(Device.PRIMARY_ID);
    when(targetAccount.getDevice(Device.PRIMARY_ID)).thenReturn(Optional.of(targetDevice));

    final ECKeyPair identityKeyPair = ECKeyPair.generate();
    final IdentityKey identityKey = new IdentityKey(identityKeyPair.getPublicKey());
    final UUID uuid = UUID.randomUUID();
    final AciServiceIdentifier identifier = new AciServiceIdentifier(uuid);
    final byte[] unidentifiedAccessKey = TestRandomUtil.nextBytes(UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH);

    when(targetAccount.isUnrestrictedUnidentifiedAccess()).thenReturn(true);
    when(targetAccount.getUnidentifiedAccessKey()).thenReturn(Optional.of(unidentifiedAccessKey));
    when(targetAccount.getIdentifier(IdentityType.ACI)).thenReturn(uuid);
    when(targetAccount.getIdentityKey(IdentityType.ACI)).thenReturn(identityKey);
    when(accountsManager.getByServiceIdentifier(identifier))
        .thenReturn(Optional.of(targetAccount));

    final ECPreKey ecPreKey = new ECPreKey(1, ECKeyPair.generate().getPublicKey());
    final ECSignedPreKey ecSignedPreKey = KeysHelper.signedECPreKey(2, identityKeyPair);
    final KEMSignedPreKey kemSignedPreKey = KeysHelper.signedKEMPreKey(3, identityKeyPair);
    final KeysManager.DevicePreKeys devicePreKeys =
        new KeysManager.DevicePreKeys(ecSignedPreKey, Optional.of(ecPreKey), kemSignedPreKey);

    when(keysManager.takeDevicePreKeys(eq(Device.PRIMARY_ID), eq(identifier), any()))
        .thenReturn(CompletableFuture.completedFuture(Optional.of(devicePreKeys)));
    final GetPreKeysAnonymousRequest.Builder request = GetPreKeysAnonymousRequest.newBuilder()
        .setRequest(GetPreKeysRequest.newBuilder()
            .setTargetIdentifier(ServiceIdentifierUtil.toGrpcServiceIdentifier(identifier))
            .setDeviceId(Device.PRIMARY_ID));

    if (includeUak) {
      request.setUnidentifiedAccessKey(ByteString.copyFrom(TestRandomUtil.nextBytes(16)));
    } else {
      request.setUnrestrictedAccess(Empty.getDefaultInstance());
    }

    final GetPreKeysAnonymousResponse response = unauthenticatedServiceStub().getPreKeys(request.build());
    final GetPreKeysAnonymousResponse expectedResponse = GetPreKeysAnonymousResponse.newBuilder()
        .setPreKeys(AccountPreKeyBundles.newBuilder()
            .setIdentityKey(ByteString.copyFrom(identityKey.serialize()))
            .putDevicePreKeys(Device.PRIMARY_ID, DevicePreKeyBundle.newBuilder()
                .setEcOneTimePreKey(toGrpcEcPreKey(ecPreKey))
                .setEcSignedPreKey(toGrpcEcSignedPreKey(ecSignedPreKey))
                .setKemOneTimePreKey(toGrpcKemSignedPreKey(kemSignedPreKey))
                .build()))
        .build();

    assertEquals(expectedResponse, response);
  }


  @Test
  void getPreKeysNoAuth() {
    assertGetKeysFailure(Status.INVALID_ARGUMENT, GetPreKeysAnonymousRequest.newBuilder()
        .setRequest(GetPreKeysRequest.newBuilder()
            .setTargetIdentifier(ServiceIdentifierUtil.toGrpcServiceIdentifier(new AciServiceIdentifier(UUID.randomUUID())))
            .setDeviceId(Device.PRIMARY_ID))
        .build());

    verifyNoInteractions(accountsManager);
    verifyNoInteractions(keysManager);
  }

  @Test
  void getPreKeysIncorrectUnidentifiedAccessKey() {
    final Account targetAccount = mock(Account.class);

    final UUID uuid = UUID.randomUUID();
    final AciServiceIdentifier identifier = new AciServiceIdentifier(uuid);
    final byte[] unidentifiedAccessKey = TestRandomUtil.nextBytes(UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH);

    when(targetAccount.getUnidentifiedAccessKey()).thenReturn(Optional.of(unidentifiedAccessKey));
    when(accountsManager.getByServiceIdentifier(identifier))
        .thenReturn(Optional.of(targetAccount));

    final GetPreKeysAnonymousResponse response = unauthenticatedServiceStub().getPreKeys(
        GetPreKeysAnonymousRequest.newBuilder()
            .setUnidentifiedAccessKey(UUIDUtil.toByteString(UUID.randomUUID()))
            .setRequest(GetPreKeysRequest.newBuilder()
                .setTargetIdentifier(ServiceIdentifierUtil.toGrpcServiceIdentifier(identifier))
                .setDeviceId(Device.PRIMARY_ID))
            .build());

    assertTrue(response.hasFailedUnidentifiedAuthorization());
    verifyNoInteractions(keysManager);
  }

  @Test
  void getPreKeysExpiredGroupSendEndorsement() throws Exception {
    final UUID uuid = UUID.randomUUID();
    final AciServiceIdentifier identifier = new AciServiceIdentifier(uuid);

    // Expirations must be on day boundaries or libsignal will refuse to create or verify the token
    final Instant expiration = Instant.now().truncatedTo(ChronoUnit.DAYS);
    CLOCK.pin(expiration.plus(Duration.ofHours(1))); // set time so our token is already expired

    final byte[] token = AuthHelper.validGroupSendToken(SERVER_SECRET_PARAMS, List.of(identifier), expiration);

    final GetPreKeysAnonymousResponse preKeysResponse =
        unauthenticatedServiceStub().getPreKeys(GetPreKeysAnonymousRequest.newBuilder()
            .setGroupSendToken(ByteString.copyFrom(token))
            .setRequest(GetPreKeysRequest.newBuilder()
                .setTargetIdentifier(ServiceIdentifierUtil.toGrpcServiceIdentifier(identifier))
                .setDeviceId(Device.PRIMARY_ID))
            .build());
    assertTrue(preKeysResponse.hasFailedUnidentifiedAuthorization());

    verifyNoInteractions(accountsManager);
    verifyNoInteractions(keysManager);
  }

  @Test
  void getPreKeysIncorrectGroupSendEndorsement() throws Exception {
    final AciServiceIdentifier authorizedIdentifier = new AciServiceIdentifier(UUID.randomUUID());
    final AciServiceIdentifier targetIdentifier = new AciServiceIdentifier(UUID.randomUUID());

    // Expirations must be on day boundaries or libsignal will refuse to create or verify the token
    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 byte[] token = AuthHelper.validGroupSendToken(SERVER_SECRET_PARAMS, List.of(authorizedIdentifier), expiration);

    final GetPreKeysAnonymousResponse response = unauthenticatedServiceStub().getPreKeys(
        GetPreKeysAnonymousRequest.newBuilder()
            .setGroupSendToken(ByteString.copyFrom(token))
            .setRequest(GetPreKeysRequest.newBuilder()
                .setTargetIdentifier(ServiceIdentifierUtil.toGrpcServiceIdentifier(targetIdentifier))
                .setDeviceId(Device.PRIMARY_ID))
            .build());
    assertTrue(response.hasFailedUnidentifiedAuthorization());
    verifyNoInteractions(accountsManager);
    verifyNoInteractions(keysManager);
  }

  @Test
  void getPreKeysAccountNotFoundUnidentifiedAccessKey() {
    final AciServiceIdentifier nonexistentAci = new AciServiceIdentifier(UUID.randomUUID());
    when(accountsManager.getByServiceIdentifier(nonexistentAci))
        .thenReturn(Optional.empty());

    final GetPreKeysAnonymousResponse preKeysResponse =
        unauthenticatedServiceStub().getPreKeys(GetPreKeysAnonymousRequest.newBuilder()
            .setUnidentifiedAccessKey(UUIDUtil.toByteString(UUID.randomUUID()))
            .setRequest(GetPreKeysRequest.newBuilder()
                .setTargetIdentifier(ServiceIdentifierUtil.toGrpcServiceIdentifier(nonexistentAci)))
            .build());
    assertTrue(preKeysResponse.hasFailedUnidentifiedAuthorization());
    verifyNoInteractions(keysManager);
  }

  @Test
  void getPreKeysAccountNotFoundGroupSendEndorsement() throws Exception {
    final AciServiceIdentifier nonexistentAci = new AciServiceIdentifier(UUID.randomUUID());

    // Expirations must be on day boundaries or libsignal will refuse to create or verify the token
    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 byte[] token = AuthHelper.validGroupSendToken(SERVER_SECRET_PARAMS, List.of(nonexistentAci), expiration);

    when(accountsManager.getByServiceIdentifier(nonexistentAci))
        .thenReturn(Optional.empty());

    final GetPreKeysAnonymousResponse preKeysResponse =
        unauthenticatedServiceStub().getPreKeys(GetPreKeysAnonymousRequest.newBuilder()
            .setGroupSendToken(ByteString.copyFrom(token))
            .setRequest(GetPreKeysRequest.newBuilder()
                .setTargetIdentifier(ServiceIdentifierUtil.toGrpcServiceIdentifier(nonexistentAci)))
        .build());
    assertTrue(preKeysResponse.hasTargetNotFound());
    verifyNoInteractions(keysManager);
  }

  @Test
  void getPreKeysDeviceNotFound() {
    final UUID accountIdentifier = UUID.randomUUID();

    final byte[] unidentifiedAccessKey = TestRandomUtil.nextBytes(UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH);

    final Account targetAccount = mock(Account.class);
    when(targetAccount.getUuid()).thenReturn(accountIdentifier);
    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(targetAccount.getUnidentifiedAccessKey()).thenReturn(Optional.of(unidentifiedAccessKey));

    when(accountsManager.getByServiceIdentifier(new AciServiceIdentifier(accountIdentifier)))
        .thenReturn(Optional.of(targetAccount));

    final GetPreKeysAnonymousResponse response = unauthenticatedServiceStub().getPreKeys(
        GetPreKeysAnonymousRequest.newBuilder()
            .setUnidentifiedAccessKey(ByteString.copyFrom(unidentifiedAccessKey))
            .setRequest(GetPreKeysRequest.newBuilder()
                .setTargetIdentifier(ServiceIdentifier.newBuilder()
                    .setIdentityType(org.signal.chat.common.IdentityType.IDENTITY_TYPE_ACI)
                    .setUuid(UUIDUtil.toByteString(accountIdentifier)))
                .setDeviceId(Device.PRIMARY_ID))
            .build());

    assertTrue(response.hasFailedUnidentifiedAuthorization());
  }

  @Test
  void checkIdentityKeys() throws InterruptedException {
    final KeysAnonymousGrpc.KeysAnonymousStub keysAnonymousStub =
        KeysAnonymousGrpc.newStub(SimpleBaseGrpcTest.GRPC_SERVER_EXTENSION_UNAUTHENTICATED.getChannel());

    when(accountsManager.getByServiceIdentifierAsync(any()))
        .thenReturn(CompletableFuture.completedFuture(Optional.empty()));

    final Account mismatchedAciFingerprintAccount = mock(Account.class);
    final UUID mismatchedAciFingerprintAccountIdentifier = UUID.randomUUID();
    final IdentityKey mismatchedAciFingerprintAccountIdentityKey = new IdentityKey(ECKeyPair.generate().getPublicKey());

    final Account matchingAciFingerprintAccount = mock(Account.class);
    final UUID matchingAciFingerprintAccountIdentifier = UUID.randomUUID();
    final IdentityKey matchingAciFingerprintAccountIdentityKey = new IdentityKey(ECKeyPair.generate().getPublicKey());

    final Account mismatchedPniFingerprintAccount = mock(Account.class);
    final UUID mismatchedPniFingerprintAccountIdentifier = UUID.randomUUID();
    final IdentityKey mismatchedPniFingerpringAccountIdentityKey = new IdentityKey(ECKeyPair.generate().getPublicKey());

    when(mismatchedAciFingerprintAccount.getIdentityKey(IdentityType.ACI)).thenReturn(mismatchedAciFingerprintAccountIdentityKey);
    when(accountsManager.getByServiceIdentifierAsync(new AciServiceIdentifier(mismatchedAciFingerprintAccountIdentifier)))
        .thenReturn(CompletableFuture.completedFuture(Optional.of(mismatchedAciFingerprintAccount)));

    when(matchingAciFingerprintAccount.getIdentityKey(IdentityType.ACI)).thenReturn(matchingAciFingerprintAccountIdentityKey);
    when(accountsManager.getByServiceIdentifierAsync(new AciServiceIdentifier(matchingAciFingerprintAccountIdentifier)))
        .thenReturn(CompletableFuture.completedFuture(Optional.of(matchingAciFingerprintAccount)));

    when(mismatchedPniFingerprintAccount.getIdentityKey(IdentityType.PNI)).thenReturn(mismatchedPniFingerpringAccountIdentityKey);
    when(accountsManager.getByServiceIdentifierAsync(new PniServiceIdentifier(mismatchedPniFingerprintAccountIdentifier)))
        .thenReturn(CompletableFuture.completedFuture(Optional.of(mismatchedPniFingerprintAccount)));

    final Map<UUID, IdentityKey> expectedResponses = Map.of(
        mismatchedAciFingerprintAccountIdentifier, mismatchedAciFingerprintAccountIdentityKey,
        mismatchedPniFingerprintAccountIdentifier, mismatchedPniFingerpringAccountIdentityKey);

    final Map<UUID, IdentityKey> responses = new ConcurrentHashMap<>();
    final CountDownLatch completedLatch = new CountDownLatch(1);
    final AtomicReference<Throwable> error = new AtomicReference<>();

    final StreamObserver<CheckIdentityKeyRequest> requestStreamObserver =
        keysAnonymousStub.checkIdentityKeys(new StreamObserver<>() {
          @Override
          public void onNext(final CheckIdentityKeyResponse checkIdentityKeyResponse) {
            try {
              responses.put(
                  ServiceIdentifierUtil.fromGrpcServiceIdentifier(checkIdentityKeyResponse.getTargetIdentifier()).uuid(),
                  new IdentityKey(checkIdentityKeyResponse.getIdentityKey().toByteArray()));
            } catch (final InvalidKeyException e) {
              throw new RuntimeException(e);
            }
          }

          @Override
          public void onError(final Throwable throwable) {
            error.set(throwable);
            completedLatch.countDown();
          }

          @Override
          public void onCompleted() {
            completedLatch.countDown();
          }
        });

    requestStreamObserver.onNext(buildCheckIdentityKeyRequest(org.signal.chat.common.IdentityType.IDENTITY_TYPE_ACI, mismatchedAciFingerprintAccountIdentifier,
        new IdentityKey(ECKeyPair.generate().getPublicKey())));

    requestStreamObserver.onNext(buildCheckIdentityKeyRequest(org.signal.chat.common.IdentityType.IDENTITY_TYPE_ACI, matchingAciFingerprintAccountIdentifier,
        matchingAciFingerprintAccountIdentityKey));

    requestStreamObserver.onNext(buildCheckIdentityKeyRequest(org.signal.chat.common.IdentityType.IDENTITY_TYPE_PNI, UUID.randomUUID(),
        new IdentityKey(ECKeyPair.generate().getPublicKey())));

    requestStreamObserver.onNext(buildCheckIdentityKeyRequest(org.signal.chat.common.IdentityType.IDENTITY_TYPE_PNI, mismatchedPniFingerprintAccountIdentifier,
        new IdentityKey(ECKeyPair.generate().getPublicKey())));

    requestStreamObserver.onCompleted();

    if (!completedLatch.await(5, TimeUnit.SECONDS)) {
      fail("Timed out waiting for countdown latch");
    }

    assertNull(error.get());
    assertEquals(expectedResponses, responses);
  }

  private static CheckIdentityKeyRequest buildCheckIdentityKeyRequest(final org.signal.chat.common.IdentityType identityType,
      final UUID uuid, final IdentityKey identityKey) {
    return CheckIdentityKeyRequest.newBuilder()
        .setTargetIdentifier(ServiceIdentifier.newBuilder()
            .setIdentityType(identityType)
            .setUuid(ByteString.copyFrom(UUIDUtil.toBytes(uuid))))
        .setFingerprint(ByteString.copyFrom(getFingerprint(identityKey)))
        .build();
  }

  private static byte[] getFingerprint(final IdentityKey publicKey) {
    try {
      return Util.truncate(MessageDigest.getInstance("SHA-256").digest(publicKey.serialize()), 4);
    } catch (final NoSuchAlgorithmException e) {
      throw new AssertionError("All Java implementations must support SHA-256 MessageDigest algorithm", e);
    }
  }

  private void assertGetKeysFailure(Status code, GetPreKeysAnonymousRequest request) {
    assertStatusException(code, () -> unauthenticatedServiceStub().getPreKeys(request));
  }

  private static EcPreKey toGrpcEcPreKey(final ECPreKey preKey) {
    return EcPreKey.newBuilder()
        .setKeyId(KeyIdUtil.toUnsignedInt(preKey.keyId()))
        .setPublicKey(ByteString.copyFrom(preKey.publicKey().serialize()))
        .build();
  }

  private static EcSignedPreKey toGrpcEcSignedPreKey(final ECSignedPreKey preKey) {
    return EcSignedPreKey.newBuilder()
        .setKeyId(KeyIdUtil.toUnsignedInt(preKey.keyId()))
        .setPublicKey(ByteString.copyFrom(preKey.publicKey().serialize()))
        .setSignature(ByteString.copyFrom(preKey.signature()))
        .build();
  }

  private static KemSignedPreKey toGrpcKemSignedPreKey(final KEMSignedPreKey preKey) {
    return KemSignedPreKey.newBuilder()
        .setKeyId(KeyIdUtil.toUnsignedInt(preKey.keyId()))
        .setPublicKey(ByteString.copyFrom(preKey.publicKey().serialize()))
        .setSignature(ByteString.copyFrom(preKey.signature()))
        .build();
  }

}
