Make identity token fetcher more async friendly

After the identity token expires a subsequent call would do a blocking
operation to retrieve the new token. Since we're making use of an async
gRPC client, this tends to block a thread we don't want to be blocking
on.

Instead, switch to periodically refreshing the token on a dedicated
thread.
This commit is contained in:
Ravi Khadiwala
2024-01-23 16:55:57 -06:00
committed by ravi-signal
parent 498ace0488
commit 1428ca73de
4 changed files with 160 additions and 29 deletions

View File

@@ -0,0 +1,61 @@
/*
* Copyright 2024 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.registration;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.auth.oauth2.IdToken;
import com.google.auth.oauth2.ImpersonatedCredentials;
import io.github.resilience4j.core.IntervalFunction;
import io.github.resilience4j.retry.RetryConfig;
import io.grpc.CallCredentials;
import java.io.IOException;
import java.time.Duration;
import java.util.concurrent.Executors;
import org.junit.jupiter.api.Test;
public class IdentityTokenCallCredentialsTest {
@Test
public void retryErrors() throws IOException {
final ImpersonatedCredentials impersonatedCredentials = mock(ImpersonatedCredentials.class);
when(impersonatedCredentials.getSourceCredentials()).thenReturn(mock(GoogleCredentials.class));
final IdentityTokenCallCredentials creds = new IdentityTokenCallCredentials(
RetryConfig.custom()
.retryOnException(throwable -> true)
.maxAttempts(Integer.MAX_VALUE)
.intervalFunction(IntervalFunction.ofExponentialRandomBackoff(
Duration.ofMillis(100), 1.5, Duration.ofSeconds(5)))
.build(),
impersonatedCredentials,
"test",
Executors.newSingleThreadScheduledExecutor());
final IdToken idToken = mock(IdToken.class);
when(idToken.getTokenValue()).thenReturn("testtoken");
// throw exception first two calls, then succeed
when(impersonatedCredentials.idTokenWithAudience(anyString(), any()))
.thenThrow(new IOException("uh oh 1"))
.thenThrow(new IOException("uh oh 2"))
.thenReturn(idToken)
.thenThrow(new IOException("uh oh 3"));
creds.refreshIdentityToken();
CallCredentials.MetadataApplier metadataApplier = mock(CallCredentials.MetadataApplier.class);
creds.applyRequestMetadata(null, null, metadataApplier);
verify(metadataApplier, times(1))
.apply(argThat(metadata -> "Bearer testtoken".equals(metadata.get(IdentityTokenCallCredentials.AUTHORIZATION_METADATA_KEY))));
}
}