From 8e0ecf4a9cea041e8014d2fa0e93ef3561af7745 Mon Sep 17 00:00:00 2001 From: yaswanth-pula-skyflow Date: Mon, 26 Oct 2026 20:27:28 +0530 Subject: [PATCH 1/4] SK-3253:Add gated retry on 401 status code for both skyvault & flowvault SDKs. --- .../java/com/skyflow/BaseVaultClient.java | 107 ++++++- .../main/java/com/skyflow/logs/InfoLogs.java | 3 + .../java/com/skyflow/logs/WarningLogs.java | 4 +- .../skyflow/serviceaccount/util/Token.java | 23 ++ .../com/skyflow/BaseVaultClientTests.java | 299 ++++++++++++++++++ .../serviceaccount/util/TokenTests.java | 42 +++ .../main/java/com/skyflow/VaultClient.java | 10 +- .../com/skyflow/AuthInterceptorTests.java | 101 ++++++ .../java/com/skyflow/utils/FakeChain.java | 4 + .../main/java/com/skyflow/VaultClient.java | 2 +- 10 files changed, 578 insertions(+), 17 deletions(-) diff --git a/common/src/main/java/com/skyflow/BaseVaultClient.java b/common/src/main/java/com/skyflow/BaseVaultClient.java index 3d5b0d32..0d3837e4 100644 --- a/common/src/main/java/com/skyflow/BaseVaultClient.java +++ b/common/src/main/java/com/skyflow/BaseVaultClient.java @@ -9,6 +9,7 @@ import com.skyflow.errors.SkyflowException; import com.skyflow.logs.ErrorLogs; import com.skyflow.logs.InfoLogs; +import com.skyflow.logs.WarningLogs; import com.skyflow.serviceaccount.util.Token; import com.skyflow.utils.BaseConstants; import com.skyflow.utils.BaseUtils; @@ -17,9 +18,12 @@ import io.github.cdimascio.dotenv.Dotenv; import io.github.cdimascio.dotenv.DotenvException; import okhttp3.ConnectionPool; +import okhttp3.Interceptor; import okhttp3.OkHttpClient; import okhttp3.Request; +import okhttp3.Response; +import java.net.HttpURLConnection; import java.util.List; import java.util.Map; import java.util.concurrent.TimeUnit; @@ -34,6 +38,13 @@ class BaseVaultClient { protected String token; protected String apiKey; + /** + * A 401 is treated as token expiry only if the token the SDK sent has an {@code exp} at or + * before now plus this window. By the time the 401 arrives the gateway has already judged the + * token expired, so the window only has to absorb clock skew between this host and the gateway. + */ + protected static final long BEARER_TOKEN_EXPIRY_WINDOW_SECONDS = 60; + protected BaseVaultClient(V vaultConfig, BaseCredentials credentials) { this.vaultConfig = vaultConfig; this.commonCredentials = credentials; @@ -43,18 +54,98 @@ protected V getVaultConfig() { return vaultConfig; } + /** + * The bearer token to send. Synchronized on the same monitor as {@link #setBearerToken(BaseCredentials)} + * and {@link #refreshBearerTokenAfterUnauthorized(String)}, so a request attempt on an OkHttp or + * batch thread sees the token those methods wrote, and waits for a regeneration in progress + * instead of sending the token it is replacing. + */ + protected synchronized String getToken() { + return token; + } + protected OkHttpClient buildSharedHttpClient(Supplier tokenSupplier) { return new OkHttpClient.Builder() .connectionPool(new ConnectionPool(10, 1, TimeUnit.MINUTES)) - .addInterceptor(chain -> { - Request requestWithAuth = chain.request().newBuilder() - .header("Authorization", "Bearer " + tokenSupplier.get()) - .build(); - return chain.proceed(requestWithAuth); - }) + .addInterceptor(authInterceptor(tokenSupplier)) .build(); } + /** + * The auth interceptor shared by every OkHttp client in the SDK. It attaches the current bearer + * token to each attempt, reading it fresh from {@code tokenSupplier} so a token refreshed between + * attempts is picked up. When the vault answers {@code 401} and + * {@link #refreshBearerTokenAfterUnauthorized(String)} can supply a replacement, the rejected + * response is closed and the same request is sent once more with the new token. Any other + * status, and a {@code 401} the refresh declines, is returned as is. + */ + protected Interceptor authInterceptor(Supplier tokenSupplier) { + return chain -> { + String sent = tokenSupplier.get(); + Response response = chain.proceed(withBearer(chain.request(), sent)); + if (response.code() != HttpURLConnection.HTTP_UNAUTHORIZED) { + return response; + } + String fresh = refreshBearerTokenAfterUnauthorized(sent); + if (fresh == null) { + return response; + } + response.close(); + return chain.proceed(withBearer(chain.request(), fresh)); + }; + } + + private static Request withBearer(Request request, String bearer) { + return request.newBuilder().header("Authorization", "Bearer " + bearer).build(); + } + + /** + * Called when the vault rejected {@code rejected} with {@code 401}. Returns the token to resend + * the request with, or {@code null} when the request must fail as it stands: + *
    + *
  • the effective credentials are a static token or an API key, which the SDK cannot replace;
    • + *
  • the rejected token is not at or past expiry (within {@link #BEARER_TOKEN_EXPIRY_WINDOW_SECONDS}), + * so the {@code 401} was not caused by expiry and a new token would not help;
    • + *
  • generating a replacement failed, or produced the very token that was rejected.
    • + *
+ * Synchronized with {@link #setBearerToken(BaseCredentials)} and {@link #getToken()} so concurrent + * {@code 401}s on one token share a single regeneration: a caller that finds the cached token + * already different from the one it sent reuses it without generating again. + */ + protected synchronized String refreshBearerTokenAfterUnauthorized(String rejected) { + BaseCredentials credentials = this.finalCredentials; + if (credentials == null || credentials.getApiKey() != null || credentials.getToken() != null) { + return null; + } + if (!Token.isExpiringWithin(rejected, BEARER_TOKEN_EXPIRY_WINDOW_SECONDS)) { + LogUtil.printInfoLog(InfoLogs.BEARER_TOKEN_REJECTED_NOT_DUE.getLog()); + return null; + } + String cached = this.token; + if (cached != null && !cached.equals(rejected)) { + LogUtil.printInfoLog(InfoLogs.BEARER_TOKEN_REJECTED_REUSING_REGENERATED.getLog()); + return cached; + } + LogUtil.printInfoLog(InfoLogs.BEARER_TOKEN_REJECTED_REGENERATING.getLog()); + try { + String fresh = generateBearerToken(credentials); + if (fresh == null || fresh.equals(rejected)) { + LogUtil.printWarningLog(WarningLogs.BEARER_TOKEN_REGENERATION_RETURNED_REJECTED_TOKEN.getLog()); + return null; + } + this.token = fresh; + return fresh; + } catch (SkyflowException e) { + LogUtil.printWarningLog(WarningLogs.BEARER_TOKEN_REGENERATION_AFTER_401_FAILED.getLog() + " " + e.getMessage()); + return null; + } + } + + /** Single point that turns credentials into a bearer token; tests substitute it. */ + protected String generateBearerToken(BaseCredentials credentials) throws SkyflowException { + return BaseUtils.generateBearerToken(credentials); + } + protected synchronized void prioritiseCredentials(BaseCredentials vaultSpecificCredentials) throws SkyflowException { try { BaseCredentials original = this.finalCredentials; @@ -95,10 +186,10 @@ protected synchronized void setBearerToken(BaseCredentials vaultSpecificCredenti LogUtil.printInfoLog(InfoLogs.USE_API_KEY.getLog()); token = this.finalCredentials.getApiKey(); } else if (token == null || token.trim().isEmpty()) { - token = BaseUtils.generateBearerToken(this.finalCredentials); + token = generateBearerToken(this.finalCredentials); } else if (Token.isExpired(token)) { LogUtil.printInfoLog(InfoLogs.BEARER_TOKEN_EXPIRED.getLog()); - token = BaseUtils.generateBearerToken(this.finalCredentials); + token = generateBearerToken(this.finalCredentials); } else { LogUtil.printInfoLog(InfoLogs.REUSE_BEARER_TOKEN.getLog()); } diff --git a/common/src/main/java/com/skyflow/logs/InfoLogs.java b/common/src/main/java/com/skyflow/logs/InfoLogs.java index b4045229..f30278b7 100644 --- a/common/src/main/java/com/skyflow/logs/InfoLogs.java +++ b/common/src/main/java/com/skyflow/logs/InfoLogs.java @@ -20,6 +20,9 @@ public enum InfoLogs { GET_SIGNED_DATA_TOKENS_TRIGGERED("getSignedDataTokens method triggered."), GET_SIGNED_DATA_TOKEN_SUCCESS("Signed data tokens generated."), REUSE_BEARER_TOKEN("Reusing bearer token."), + BEARER_TOKEN_REJECTED_REGENERATING("Bearer token rejected with 401 and is past expiry; regenerating once and resending the request."), + BEARER_TOKEN_REJECTED_REUSING_REGENERATED("Bearer token rejected with 401; reusing the token another request already regenerated."), + BEARER_TOKEN_REJECTED_NOT_DUE("Bearer token rejected with 401 but is not near expiry; not retrying."), USE_CLIENT_PROVIDED_BEARER_TOKEN("Using bearer token provided by client."), USE_API_KEY("Using api key."), REUSE_API_KEY("Reusing api key."), diff --git a/common/src/main/java/com/skyflow/logs/WarningLogs.java b/common/src/main/java/com/skyflow/logs/WarningLogs.java index 1eb1bcd8..9d001694 100644 --- a/common/src/main/java/com/skyflow/logs/WarningLogs.java +++ b/common/src/main/java/com/skyflow/logs/WarningLogs.java @@ -8,7 +8,9 @@ public enum WarningLogs { EMPTY_DELETE_TOKENS_RESPONSE("DeleteTokens response did not include any token results."), INCOMPLETE_DELETE_TOKENS_RESPONSE("DeleteTokens response did not account for all requested tokens."), EMPTY_TOKENIZE_RESPONSE("Tokenize response did not include any record results."), - INCOMPLETE_TOKENIZE_RESPONSE("Tokenize response did not account for all requested records.") + INCOMPLETE_TOKENIZE_RESPONSE("Tokenize response did not account for all requested records."), + BEARER_TOKEN_REGENERATION_AFTER_401_FAILED("Bearer token regeneration after 401 failed; returning the 401 response."), + BEARER_TOKEN_REGENERATION_RETURNED_REJECTED_TOKEN("Bearer token regeneration after 401 returned the rejected token; not resending.") ; private final String log; diff --git a/common/src/main/java/com/skyflow/serviceaccount/util/Token.java b/common/src/main/java/com/skyflow/serviceaccount/util/Token.java index 0359d944..e272df8b 100644 --- a/common/src/main/java/com/skyflow/serviceaccount/util/Token.java +++ b/common/src/main/java/com/skyflow/serviceaccount/util/Token.java @@ -40,6 +40,29 @@ public static boolean isExpired(String token) { return currentTime > expiryTime; } + /** + * Returns true when the token's {@code exp} claim is at or before now plus {@code windowSeconds}, + * that is, the token has expired or will within the window. Unlike {@link #isExpired(String)}, + * a token that is empty, cannot be decoded, or carries no {@code exp} is reported as + * not expiring: callers use this to decide whether a {@code 401} was caused by expiry, + * and a malformed token is not an expiry. + */ + public static boolean isExpiringWithin(String token, long windowSeconds) { + if (token == null || token.trim().isEmpty()) { + return false; + } + try { + JsonElement exp = decoded(token).get("exp"); + if (exp == null) { + return false; + } + long currentTime = new Date().getTime() / 1000; + return exp.getAsLong() <= currentTime + windowSeconds; + } catch (SkyflowException | RuntimeException e) { + return false; + } + } + static JsonObject decoded(String encodedToken) throws JsonSyntaxException, SkyflowException { String[] split = encodedToken.split("\\."); if (split.length < 3) { diff --git a/common/src/test/java/com/skyflow/BaseVaultClientTests.java b/common/src/test/java/com/skyflow/BaseVaultClientTests.java index 7a4c0da9..182b87a7 100644 --- a/common/src/test/java/com/skyflow/BaseVaultClientTests.java +++ b/common/src/test/java/com/skyflow/BaseVaultClientTests.java @@ -26,7 +26,15 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Base64; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Supplier; public class BaseVaultClientTests { @@ -301,4 +309,295 @@ public void testSetBearerToken_expiredToken_regeneratesToken() throws SkyflowExc Assert.assertEquals("x.eyJleHAiOjk5OTk5OTk5OTl9.y", client.token); Assert.assertNotEquals("", client.token); } + + // ── refresh-and-resend on 401 ──────────────────────────────────────────── + + private static String jwtExpiringIn(long seconds) { + Base64.Encoder enc = Base64.getUrlEncoder().withoutPadding(); + String header = enc.encodeToString("{\"alg\":\"none\"}".getBytes(StandardCharsets.UTF_8)); + String payload = enc.encodeToString(("{\"exp\":" + (System.currentTimeMillis() / 1000 + seconds) + "}") + .getBytes(StandardCharsets.UTF_8)); + return header + "." + payload + ".sig"; + } + + /** A client whose token generation is scripted, so no service-account file is needed. */ + private static final class ScriptedClient extends BaseVaultClient { + final AtomicInteger generations = new AtomicInteger(); + final java.util.List script = new ArrayList<>(); + SkyflowException failure; + + ScriptedClient(BaseCredentials credentials) { + super(new BaseVaultConfig(), credentials); + this.finalCredentials = credentials; + } + + @Override + protected String generateBearerToken(BaseCredentials credentials) throws SkyflowException { + int n = generations.getAndIncrement(); + if (failure != null) { + throw failure; + } + return script.get(Math.min(n, script.size() - 1)); + } + } + + private static BaseCredentials mintableCredentials() { + BaseCredentials credentials = new BaseCredentials(); + credentials.setCredentialsString("{\"clientID\":\"c\"}"); + return credentials; + } + + @Test + public void testRefreshAfter401_expiredSdkTokenIsRegeneratedOnceAndCached() { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = rejected; + client.script.add("fresh-token"); + + Assert.assertEquals("fresh-token", client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals("fresh-token", client.token); + Assert.assertEquals(1, client.generations.get()); + } + + @Test + public void testRefreshAfter401_tokenInsideWindowCountsAsExpired() { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(30); + client.token = rejected; + client.script.add("fresh-token"); + + Assert.assertEquals("fresh-token", client.refreshBearerTokenAfterUnauthorized(rejected)); + } + + @Test + public void testRefreshAfter401_tokenWithLifeLeftIsNotRetried() { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(600); + client.token = rejected; + client.script.add("fresh-token"); + + Assert.assertNull("a 401 on a live token is not expiry", client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals("must not spend a mint on a non-expiry 401", 0, client.generations.get()); + Assert.assertEquals(rejected, client.token); + } + + @Test + public void testRefreshAfter401_staticTokenIsNeverRegenerated() { + BaseCredentials credentials = new BaseCredentials(); + String rejected = jwtExpiringIn(-2); + credentials.setToken(rejected); + ScriptedClient client = new ScriptedClient(credentials); + client.token = rejected; + client.script.add("fresh-token"); + + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals(0, client.generations.get()); + } + + @Test + public void testRefreshAfter401_apiKeyIsNeverRegenerated() { + BaseCredentials credentials = new BaseCredentials(); + credentials.setApiKey("sky-test-key"); + ScriptedClient client = new ScriptedClient(credentials); + client.token = "sky-test-key"; + client.script.add("fresh-token"); + + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized("sky-test-key")); + Assert.assertEquals(0, client.generations.get()); + } + + @Test + public void testRefreshAfter401_noCredentialsResolvedYet() { + ScriptedClient client = new ScriptedClient(null); + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(jwtExpiringIn(-2))); + } + + @Test + public void testRefreshAfter401_generationFailureReturnsNullAndKeepsToken() { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = rejected; + client.failure = new SkyflowException("auth service down"); + + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals(rejected, client.token); + } + + @Test + public void testRefreshAfter401_sameTokenBackMeansStop() { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = rejected; + client.script.add(rejected); + + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals(1, client.generations.get()); + } + + @Test + public void testRefreshAfter401_reusesTokenAnotherRequestAlreadyRegenerated() { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = "already-fresh"; + client.script.add("never-used"); + + Assert.assertEquals("already-fresh", client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals(0, client.generations.get()); + } + + @Test + public void testRefreshAfter401_concurrentRejectionsShareOneGeneration() throws Exception { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = rejected; + client.script.add("fresh-token"); + int threads = 8; + ExecutorService pool = Executors.newFixedThreadPool(threads); + CountDownLatch start = new CountDownLatch(1); + java.util.List> results = new ArrayList<>(); + try { + for (int i = 0; i < threads; i++) { + results.add(pool.submit(() -> { + start.await(); + return client.refreshBearerTokenAfterUnauthorized(rejected); + })); + } + start.countDown(); + for (Future result : results) { + Assert.assertEquals("fresh-token", result.get(10, TimeUnit.SECONDS)); + } + } finally { + pool.shutdownNow(); + } + Assert.assertEquals("eight 401s on one token must cost one mint", 1, client.generations.get()); + } + + /** Fake chain that replays scripted status codes and records what was proceeded. */ + private static final class ScriptedChain implements Interceptor.Chain { + private final Request request = new Request.Builder().url("https://example.com/v1/x").build(); + private final int[] codes; + final java.util.List proceeded = new ArrayList<>(); + + ScriptedChain(int... codes) { + this.codes = codes; + } + + @Override + public Request request() { + return request; + } + + @Override + public Response proceed(Request request) { + proceeded.add(request); + int code = codes[Math.min(proceeded.size() - 1, codes.length - 1)]; + return new Response.Builder().request(request).protocol(Protocol.HTTP_1_1) + .code(code).message("status " + code) + .body(okhttp3.ResponseBody.create(null, new byte[0])) + .build(); + } + + @Override + public Connection connection() { + return null; + } + + @Override + public Call call() { + throw new UnsupportedOperationException(); + } + + @Override + public int connectTimeoutMillis() { + return 0; + } + + @Override + public Interceptor.Chain withConnectTimeout(int timeout, TimeUnit unit) { + return this; + } + + @Override + public int readTimeoutMillis() { + return 0; + } + + @Override + public Interceptor.Chain withReadTimeout(int timeout, TimeUnit unit) { + return this; + } + + @Override + public int writeTimeoutMillis() { + return 0; + } + + @Override + public Interceptor.Chain withWriteTimeout(int timeout, TimeUnit unit) { + return this; + } + } + + @Test + public void testAuthInterceptor_401OnExpiredTokenIsResentOnceWithFreshToken() throws IOException { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + client.token = jwtExpiringIn(-2); + client.script.add("fresh-token"); + ScriptedChain chain = new ScriptedChain(401, 200); + + Response response = client.authInterceptor(() -> client.token).intercept(chain); + + Assert.assertEquals(200, response.code()); + Assert.assertEquals(2, chain.proceeded.size()); + Assert.assertEquals("Bearer fresh-token", chain.proceeded.get(1).header("Authorization")); + Assert.assertEquals(chain.request().url(), chain.proceeded.get(1).url()); + } + + @Test + public void testAuthInterceptor_second401IsReturnedNotRetried() throws IOException { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + client.token = jwtExpiringIn(-2); + client.script.add("fresh-token"); + ScriptedChain chain = new ScriptedChain(401, 401); + + Assert.assertEquals(401, client.authInterceptor(() -> client.token).intercept(chain).code()); + Assert.assertEquals(2, chain.proceeded.size()); + Assert.assertEquals(1, client.generations.get()); + } + + @Test + public void testAuthInterceptor_401OnStaticTokenIsNotRetried() throws IOException { + BaseCredentials credentials = new BaseCredentials(); + credentials.setToken(jwtExpiringIn(-2)); + ScriptedClient client = new ScriptedClient(credentials); + client.token = credentials.getToken(); + ScriptedChain chain = new ScriptedChain(401, 200); + + Assert.assertEquals(401, client.authInterceptor(() -> client.token).intercept(chain).code()); + Assert.assertEquals(1, chain.proceeded.size()); + } + + @Test + public void testAuthInterceptor_401OnLiveTokenIsNotRetried() throws IOException { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + client.token = jwtExpiringIn(600); + client.script.add("fresh-token"); + ScriptedChain chain = new ScriptedChain(401, 200); + + Assert.assertEquals(401, client.authInterceptor(() -> client.token).intercept(chain).code()); + Assert.assertEquals(1, chain.proceeded.size()); + Assert.assertEquals(0, client.generations.get()); + } + + @Test + public void testAuthInterceptor_403IsNotRetried() throws IOException { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + client.token = jwtExpiringIn(-2); + client.script.add("fresh-token"); + ScriptedChain chain = new ScriptedChain(403, 200); + + Assert.assertEquals(403, client.authInterceptor(() -> client.token).intercept(chain).code()); + Assert.assertEquals(1, chain.proceeded.size()); + Assert.assertEquals(0, client.generations.get()); + } } diff --git a/common/src/test/java/com/skyflow/serviceaccount/util/TokenTests.java b/common/src/test/java/com/skyflow/serviceaccount/util/TokenTests.java index 00fc26b4..646f8f5c 100644 --- a/common/src/test/java/com/skyflow/serviceaccount/util/TokenTests.java +++ b/common/src/test/java/com/skyflow/serviceaccount/util/TokenTests.java @@ -21,6 +21,48 @@ public static void setup() { LogUtil.setupLogger(LogLevel.DEBUG); } + private static String unsignedJwt(String payloadJson) { + Base64.Encoder enc = Base64.getUrlEncoder().withoutPadding(); + String header = enc.encodeToString("{\"alg\":\"none\"}".getBytes(StandardCharsets.UTF_8)); + String payload = enc.encodeToString(payloadJson.getBytes(StandardCharsets.UTF_8)); + return header + "." + payload + ".sig"; + } + + private static String jwtExpiringIn(long seconds) { + return unsignedJwt("{\"exp\":" + (System.currentTimeMillis() / 1000 + seconds) + "}"); + } + + @Test + public void testIsExpiringWithin_expiredTokenIsDue() { + Assert.assertTrue(Token.isExpiringWithin(jwtExpiringIn(-5), 60)); + } + + @Test + public void testIsExpiringWithin_tokenInsideWindowIsDue() { + Assert.assertTrue(Token.isExpiringWithin(jwtExpiringIn(30), 60)); + } + + @Test + public void testIsExpiringWithin_tokenBeyondWindowIsNotDue() { + Assert.assertFalse(Token.isExpiringWithin(jwtExpiringIn(600), 60)); + } + + @Test + public void testIsExpiringWithin_zeroWindowMatchesExpiryOnly() { + Assert.assertTrue(Token.isExpiringWithin(jwtExpiringIn(-1), 0)); + Assert.assertFalse(Token.isExpiringWithin(jwtExpiringIn(5), 0)); + } + + @Test + public void testIsExpiringWithin_nullEmptyMalformedAndNoExpAreNotDue() { + Assert.assertFalse(Token.isExpiringWithin(null, 60)); + Assert.assertFalse(Token.isExpiringWithin("", 60)); + Assert.assertFalse(Token.isExpiringWithin("not-a-jwt", 60)); + Assert.assertFalse(Token.isExpiringWithin("a.b", 60)); + Assert.assertFalse(Token.isExpiringWithin(unsignedJwt("{\"sub\":\"x\"}"), 60)); + Assert.assertFalse(Token.isExpiringWithin(unsignedJwt("{\"exp\":\"soon\"}"), 60)); + } + @Test public void testNoTokenForIsExpiredToken() { try { diff --git a/flowvault/src/main/java/com/skyflow/VaultClient.java b/flowvault/src/main/java/com/skyflow/VaultClient.java index 93a677e0..78be997d 100644 --- a/flowvault/src/main/java/com/skyflow/VaultClient.java +++ b/flowvault/src/main/java/com/skyflow/VaultClient.java @@ -13,7 +13,6 @@ import okhttp3.ConnectionPool; import okhttp3.OkHttpClient; -import okhttp3.Request; public class VaultClient extends BaseVaultClient { private final ApiClientBuilder apiClientBuilder; @@ -166,12 +165,9 @@ protected void updateExecutorInHTTP() throws SkyflowException { // OUTER: retries. Must wrap the auth interceptor so each attempt re-reads the // (possibly refreshed) bearer token rather than replaying a stale one. .addInterceptor(new SkyflowRetryInterceptor(maxRetries, initialRetryDelayMillis, maxRetryDelayMillis)) - .addInterceptor(chain -> { // INNER: auth - Request requestWithAuth = chain.request().newBuilder() - .header("Authorization", "Bearer " + this.token) - .build(); - return chain.proceed(requestWithAuth); - }); + // INNER: auth, shared with skyvault. Attaches the token and, on a 401 for an + // expired SDK-generated token, regenerates once and resends once. + .addInterceptor(authInterceptor(this::getToken)); if (connectTimeout != null) { httpBuilder.connectTimeout(connectTimeout, TimeUnit.SECONDS); } diff --git a/flowvault/src/test/java/com/skyflow/AuthInterceptorTests.java b/flowvault/src/test/java/com/skyflow/AuthInterceptorTests.java index 631a1120..a53227f1 100644 --- a/flowvault/src/test/java/com/skyflow/AuthInterceptorTests.java +++ b/flowvault/src/test/java/com/skyflow/AuthInterceptorTests.java @@ -15,7 +15,9 @@ import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Paths; +import java.util.Base64; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; /** * The auth interceptor installed by {@link VaultClient#updateExecutorInHTTP()}. @@ -104,4 +106,103 @@ public void testAuthInterceptor_returnsTheChainResponseUntouched() throws Skyflo Assert.assertEquals(503, authInterceptorOf(client).intercept(chain).code()); Assert.assertEquals("auth must not retry - that is the outer interceptor's job", 1, chain.calls()); } + + // ── refresh-and-resend on 401, wired through the flowvault client ──────── + + private static String jwtExpiringIn(long seconds) { + Base64.Encoder enc = Base64.getUrlEncoder().withoutPadding(); + String header = enc.encodeToString("{\"alg\":\"none\"}".getBytes(StandardCharsets.UTF_8)); + String payload = enc.encodeToString(("{\"exp\":" + (System.currentTimeMillis() / 1000 + seconds) + "}") + .getBytes(StandardCharsets.UTF_8)); + return header + "." + payload + ".sig"; + } + + /** Flowvault client with mintable credentials and a scripted token source. */ + private static final class MintingClient extends VaultClient { + final AtomicInteger generations = new AtomicInteger(); + final String fresh; + + MintingClient(String fresh) throws SkyflowException { + super(configWithCredentialsString(), null); + this.fresh = fresh; + this.finalCredentials = getVaultConfig().getCredentials(); + } + + @Override + protected String generateBearerToken(com.skyflow.config.BaseCredentials credentials) { + generations.incrementAndGet(); + return fresh; + } + } + + private static VaultConfig configWithCredentialsString() { + Credentials credentials = new Credentials(); + credentials.setCredentialsString("{\"clientID\":\"c\"}"); + VaultConfig config = config(); + config.setCredentials(credentials); + return config; + } + + @Test + public void testAuthInterceptor_regeneratesOnceAndResendsOn401ForExpiredToken() throws SkyflowException, IOException { + MintingClient client = new MintingClient("fresh-token"); + client.token = jwtExpiringIn(-1); + FakeChain chain = new FakeChain(401, 200); + + Assert.assertEquals(200, authInterceptorOf(client).intercept(chain).code()); + Assert.assertEquals(2, chain.calls()); + Assert.assertEquals("Bearer fresh-token", chain.lastProceeded().header("Authorization")); + Assert.assertEquals("the cache must hold the new token for the batches that follow", "fresh-token", client.token); + Assert.assertEquals(1, client.generations.get()); + Assert.assertTrue("the rejected response must be closed before the resend", chain.bodies().get(0).isClosed()); + } + + @Test + public void testAuthInterceptor_401ForTokenWithLifeLeftIsReturnedWithoutRegenerating() throws SkyflowException, IOException { + MintingClient client = new MintingClient("fresh-token"); + String live = jwtExpiringIn(900); + client.token = live; + FakeChain chain = new FakeChain(401, 200); + + Assert.assertEquals(401, authInterceptorOf(client).intercept(chain).code()); + Assert.assertEquals(1, chain.calls()); + Assert.assertEquals(0, client.generations.get()); + Assert.assertEquals(live, client.token); + } + + @Test + public void testAuthInterceptor_401WithStaticTokenIsReturnedAsToday() throws SkyflowException, IOException { + Credentials credentials = new Credentials(); + credentials.setToken(jwtExpiringIn(-1)); + VaultConfig config = config(); + config.setCredentials(credentials); + VaultClient client = new VaultClient(config, null); + client.setBearerToken(); + FakeChain chain = new FakeChain(401, 200); + + Assert.assertEquals(401, authInterceptorOf(client).intercept(chain).code()); + Assert.assertEquals(1, chain.calls()); + } + + @Test + public void testAuthInterceptor_resendThatAlsoGets401IsNotRetriedAgain() throws SkyflowException, IOException { + MintingClient client = new MintingClient("fresh-token"); + client.token = jwtExpiringIn(-1); + FakeChain chain = new FakeChain(401, 401, 200); + + Assert.assertEquals(401, authInterceptorOf(client).intercept(chain).code()); + Assert.assertEquals(2, chain.calls()); + Assert.assertEquals(1, client.generations.get()); + } + + @Test + public void testAuthInterceptor_sitsInsideTheTransportRetryInterceptor() throws SkyflowException { + VaultClient client = new VaultClient(config(), null); + client.updateExecutorInHTTP(); + List interceptors = client.sharedHttpClient.interceptors(); + + Assert.assertEquals(2, interceptors.size()); + Assert.assertTrue("transport retries must wrap auth so each attempt re-reads the token", + interceptors.get(0) instanceof SkyflowRetryInterceptor); + } } diff --git a/flowvault/src/test/java/com/skyflow/utils/FakeChain.java b/flowvault/src/test/java/com/skyflow/utils/FakeChain.java index 97012f13..b803fe1c 100644 --- a/flowvault/src/test/java/com/skyflow/utils/FakeChain.java +++ b/flowvault/src/test/java/com/skyflow/utils/FakeChain.java @@ -132,5 +132,9 @@ public BufferedSource source() { public void close() { closed = true; } + + public boolean isClosed() { + return closed; + } } } diff --git a/skyvault/src/main/java/com/skyflow/VaultClient.java b/skyvault/src/main/java/com/skyflow/VaultClient.java index 509b70c7..3c97a312 100644 --- a/skyvault/src/main/java/com/skyflow/VaultClient.java +++ b/skyvault/src/main/java/com/skyflow/VaultClient.java @@ -832,7 +832,7 @@ private void updateVaultURL() { private void updateExecutorInHTTP() { if (sharedHttpClient == null) { - sharedHttpClient = buildSharedHttpClient(() -> this.token); + sharedHttpClient = buildSharedHttpClient(this::getToken); apiClientBuilder.httpClient(sharedHttpClient); } } From 1ed1f9e6ea73b1f4078dec2315151e79e6ebdb6d Mon Sep 17 00:00:00 2001 From: yaswanth-pula-skyflow Date: Tue, 27 Oct 2026 18:29:48 +0530 Subject: [PATCH 2/4] SK-3253:Fix failed token regeneration when concurent 401 failures. --- .../java/com/skyflow/BaseVaultClient.java | 28 ++++++- .../main/java/com/skyflow/logs/InfoLogs.java | 1 + .../com/skyflow/BaseVaultClientTests.java | 75 +++++++++++++++++++ 3 files changed, 100 insertions(+), 4 deletions(-) diff --git a/common/src/main/java/com/skyflow/BaseVaultClient.java b/common/src/main/java/com/skyflow/BaseVaultClient.java index 0d3837e4..c19ffa6f 100644 --- a/common/src/main/java/com/skyflow/BaseVaultClient.java +++ b/common/src/main/java/com/skyflow/BaseVaultClient.java @@ -37,6 +37,10 @@ class BaseVaultClient { protected BaseCredentials finalCredentials; protected String token; protected String apiKey; + // The token whose regeneration after a 401 last failed. Guarded by this. Further 401s on that + // same token return without generating again, so an auth outage costs one generation per + // token instead of one per rejected request. Cleared whenever a token is stored. + private String refreshFailedFor; /** * A 401 is treated as token expiry only if the token the SDK sent has an {@code exp} at or @@ -106,7 +110,10 @@ private static Request withBearer(Request request, String bearer) { *
  • the effective credentials are a static token or an API key, which the SDK cannot replace;
    • *
  • the rejected token is not at or past expiry (within {@link #BEARER_TOKEN_EXPIRY_WINDOW_SECONDS}), * so the {@code 401} was not caused by expiry and a new token would not help;
    • - *
  • generating a replacement failed, or produced the very token that was rejected.
    • + *
  • generating a replacement failed, or produced the very token that was rejected;
    • + *
  • a replacement for this same token was already attempted and failed, so the outcome is + * shared with every request still holding it. A later method call recovers through + * {@link #setBearerToken(BaseCredentials)}, which clears that record when it stores a token.
    • * * Synchronized with {@link #setBearerToken(BaseCredentials)} and {@link #getToken()} so concurrent * {@code 401}s on one token share a single regeneration: a caller that finds the cached token @@ -126,21 +133,33 @@ protected synchronized String refreshBearerTokenAfterUnauthorized(String rejecte LogUtil.printInfoLog(InfoLogs.BEARER_TOKEN_REJECTED_REUSING_REGENERATED.getLog()); return cached; } + if (rejected != null && rejected.equals(refreshFailedFor)) { + LogUtil.printInfoLog(InfoLogs.BEARER_TOKEN_REJECTED_REGENERATION_ALREADY_FAILED.getLog()); + return null; + } LogUtil.printInfoLog(InfoLogs.BEARER_TOKEN_REJECTED_REGENERATING.getLog()); try { String fresh = generateBearerToken(credentials); if (fresh == null || fresh.equals(rejected)) { LogUtil.printWarningLog(WarningLogs.BEARER_TOKEN_REGENERATION_RETURNED_REJECTED_TOKEN.getLog()); + refreshFailedFor = rejected; return null; } - this.token = fresh; + storeToken(fresh); return fresh; } catch (SkyflowException e) { LogUtil.printWarningLog(WarningLogs.BEARER_TOKEN_REGENERATION_AFTER_401_FAILED.getLog() + " " + e.getMessage()); + refreshFailedFor = rejected; return null; } } + /** Stores a token the SDK obtained and forgets any failed regeneration, so refresh can run again. */ + private void storeToken(String fresh) { + this.token = fresh; + this.refreshFailedFor = null; + } + /** Single point that turns credentials into a bearer token; tests substitute it. */ protected String generateBearerToken(BaseCredentials credentials) throws SkyflowException { return BaseUtils.generateBearerToken(credentials); @@ -169,6 +188,7 @@ protected synchronized void prioritiseCredentials(BaseCredentials vaultSpecificC if (original != null && !original.equals(this.finalCredentials)) { token = null; apiKey = null; + refreshFailedFor = null; } } catch (DotenvException e) { throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.EmptyCredentials.getMessage()); @@ -186,10 +206,10 @@ protected synchronized void setBearerToken(BaseCredentials vaultSpecificCredenti LogUtil.printInfoLog(InfoLogs.USE_API_KEY.getLog()); token = this.finalCredentials.getApiKey(); } else if (token == null || token.trim().isEmpty()) { - token = generateBearerToken(this.finalCredentials); + storeToken(generateBearerToken(this.finalCredentials)); } else if (Token.isExpired(token)) { LogUtil.printInfoLog(InfoLogs.BEARER_TOKEN_EXPIRED.getLog()); - token = generateBearerToken(this.finalCredentials); + storeToken(generateBearerToken(this.finalCredentials)); } else { LogUtil.printInfoLog(InfoLogs.REUSE_BEARER_TOKEN.getLog()); } diff --git a/common/src/main/java/com/skyflow/logs/InfoLogs.java b/common/src/main/java/com/skyflow/logs/InfoLogs.java index f30278b7..55732d1e 100644 --- a/common/src/main/java/com/skyflow/logs/InfoLogs.java +++ b/common/src/main/java/com/skyflow/logs/InfoLogs.java @@ -23,6 +23,7 @@ public enum InfoLogs { BEARER_TOKEN_REJECTED_REGENERATING("Bearer token rejected with 401 and is past expiry; regenerating once and resending the request."), BEARER_TOKEN_REJECTED_REUSING_REGENERATED("Bearer token rejected with 401; reusing the token another request already regenerated."), BEARER_TOKEN_REJECTED_NOT_DUE("Bearer token rejected with 401 but is not near expiry; not retrying."), + BEARER_TOKEN_REJECTED_REGENERATION_ALREADY_FAILED("Bearer token rejected with 401; regeneration for this token already failed, not retrying."), USE_CLIENT_PROVIDED_BEARER_TOKEN("Using bearer token provided by client."), USE_API_KEY("Using api key."), REUSE_API_KEY("Reusing api key."), diff --git a/common/src/test/java/com/skyflow/BaseVaultClientTests.java b/common/src/test/java/com/skyflow/BaseVaultClientTests.java index 182b87a7..6ae52796 100644 --- a/common/src/test/java/com/skyflow/BaseVaultClientTests.java +++ b/common/src/test/java/com/skyflow/BaseVaultClientTests.java @@ -434,6 +434,81 @@ public void testRefreshAfter401_sameTokenBackMeansStop() { Assert.assertEquals(1, client.generations.get()); } + @Test + public void testRefreshAfter401_failedGenerationIsNotRepeatedForTheSameToken() { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = rejected; + client.failure = new SkyflowException("auth service down"); + + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals("later 401s on the same token must not generate again", 1, client.generations.get()); + Assert.assertEquals(rejected, client.token); + } + + @Test + public void testRefreshAfter401_sameTokenBackIsNotRepeatedForTheSameToken() { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = rejected; + client.script.add(rejected); + + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals(1, client.generations.get()); + } + + @Test + public void testRefreshAfter401_recoversAfterStartOfMethodCheckStoresNewToken() throws Exception { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = rejected; + client.failure = new SkyflowException("auth service down"); + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + + // Auth comes back; the next method call's start-of-method check mints and stores a token. + client.failure = null; + String stored = jwtExpiringIn(-3); + client.script.add(stored); + client.setBearerToken(null); + Assert.assertEquals(stored, client.token); + + // That token expires too and is rejected: refresh must generate again, not stay stuck. + client.script.add("fresh-after-recovery"); + Assert.assertEquals("fresh-after-recovery", client.refreshBearerTokenAfterUnauthorized(stored)); + Assert.assertEquals(3, client.generations.get()); + } + + @Test + public void testRefreshAfter401_concurrentGenerationFailuresShareOneAttempt() throws Exception { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = rejected; + client.failure = new SkyflowException("auth service down"); + int threads = 8; + ExecutorService pool = Executors.newFixedThreadPool(threads); + CountDownLatch start = new CountDownLatch(1); + java.util.List> results = new ArrayList<>(); + try { + for (int i = 0; i < threads; i++) { + results.add(pool.submit(() -> { + start.await(); + return client.refreshBearerTokenAfterUnauthorized(rejected); + })); + } + start.countDown(); + for (Future result : results) { + Assert.assertNull(result.get(10, TimeUnit.SECONDS)); + } + } finally { + pool.shutdownNow(); + } + Assert.assertEquals("eight 401s during an auth outage must cost one failed mint, not eight", 1, client.generations.get()); + Assert.assertEquals(rejected, client.token); + } + @Test public void testRefreshAfter401_reusesTokenAnotherRequestAlreadyRegenerated() { ScriptedClient client = new ScriptedClient(mintableCredentials()); From f77ac9a651bbb287b39ea8138108f2752ef63bc5 Mon Sep 17 00:00:00 2001 From: yaswanth-pula-skyflow Date: Tue, 27 Oct 2026 22:12:11 +0530 Subject: [PATCH 3/4] SK-3253:Update unit tests to cover both skyvault & flowvault. --- .../com/skyflow/BaseVaultClientTests.java | 68 +++++++ .../com/skyflow/AuthInterceptorTests.java | 183 ++++++++++++++++++ .../java/com/skyflow/utils/FakeChain.java | 140 ++++++++++++++ .../java/com/skyflow/utils/UtilsTests.java | 17 ++ 4 files changed, 408 insertions(+) create mode 100644 skyvault/src/test/java/com/skyflow/AuthInterceptorTests.java create mode 100644 skyvault/src/test/java/com/skyflow/utils/FakeChain.java diff --git a/common/src/test/java/com/skyflow/BaseVaultClientTests.java b/common/src/test/java/com/skyflow/BaseVaultClientTests.java index 6ae52796..d00b1547 100644 --- a/common/src/test/java/com/skyflow/BaseVaultClientTests.java +++ b/common/src/test/java/com/skyflow/BaseVaultClientTests.java @@ -675,4 +675,72 @@ public void testAuthInterceptor_403IsNotRetried() throws IOException { Assert.assertEquals(1, chain.proceeded.size()); Assert.assertEquals(0, client.generations.get()); } + + // ── gaps found by line coverage of the fix ─────────────────────────────── + + @Test + public void testAuthInterceptor_readsTokenThroughGetTokenOnEveryAttempt() throws IOException { + // The SDKs hand the interceptor this::getToken, not a captured value: the resend must see the + // token the refresh just stored, and the synchronized getter is what makes that visible. + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = rejected; + client.script.add("fresh-token"); + ScriptedChain chain = new ScriptedChain(401, 200); + + Response response = client.authInterceptor(client::getToken).intercept(chain); + + Assert.assertEquals(200, response.code()); + Assert.assertEquals("Bearer " + rejected, chain.proceeded.get(0).header("Authorization")); + Assert.assertEquals("Bearer fresh-token", chain.proceeded.get(1).header("Authorization")); + Assert.assertEquals("fresh-token", client.getToken()); + } + + @Test + public void testRefreshAfter401_nullFromGenerationIsAFailureAndIsRemembered() { + // BearerToken.getBearerToken() returns null when the credentials match neither FILE nor STRING. + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = rejected; + client.script.add(null); + + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals("a null token must never be cached", rejected, client.token); + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals("the failed attempt is remembered for this token", 1, client.generations.get()); + } + + @Test + public void testRefreshAfter401_cacheClearedWhileRequestInFlightRegenerates() { + // prioritiseCredentials nulls the cache on a credential change. A 401 for the token that was + // already in flight must regenerate rather than treat the empty cache as "someone else refreshed". + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = null; + client.script.add("fresh-token"); + + Assert.assertEquals("fresh-token", client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals("fresh-token", client.token); + Assert.assertEquals(1, client.generations.get()); + } + + @Test + public void testRefreshAfter401_credentialChangeForgetsFailedRegeneration() throws SkyflowException { + ScriptedClient client = new ScriptedClient(mintableCredentials()); + String rejected = jwtExpiringIn(-2); + client.token = rejected; + client.failure = new SkyflowException("auth service down"); + Assert.assertNull(client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals(1, client.generations.get()); + + // New credentials: the failure recorded under the old ones must not block the first 401 on them. + client.prioritiseCredentials(mintableCredentials()); + client.failure = null; + client.token = rejected; // the expired token is still what the in-flight request sent + client.script.add("fresh-token"); + + Assert.assertEquals("fresh-token", client.refreshBearerTokenAfterUnauthorized(rejected)); + Assert.assertEquals("fresh-token", client.token); + Assert.assertEquals(2, client.generations.get()); + } } diff --git a/skyvault/src/test/java/com/skyflow/AuthInterceptorTests.java b/skyvault/src/test/java/com/skyflow/AuthInterceptorTests.java new file mode 100644 index 00000000..562a847d --- /dev/null +++ b/skyvault/src/test/java/com/skyflow/AuthInterceptorTests.java @@ -0,0 +1,183 @@ +package com.skyflow; + +import com.skyflow.config.BaseCredentials; +import com.skyflow.config.Credentials; +import com.skyflow.config.VaultConfig; +import com.skyflow.enums.Env; +import com.skyflow.errors.SkyflowException; +import com.skyflow.utils.FakeChain; +import okhttp3.Interceptor; +import org.junit.Assert; +import org.junit.Test; + +import java.io.IOException; +import java.lang.reflect.Method; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.util.Base64; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * The auth interceptor installed by skyvault's {@code VaultClient.updateExecutorInHTTP()}. + * + *

    Skyvault's shared OkHttp client carries this one interceptor; the generated {@code ClientOptions} + * adds no transport retry when a client is supplied. The refresh-and-resend logic lives in + * {@code BaseVaultClient} and is unit tested there; these tests prove the skyvault client wires it + * in with {@code this::getToken}, so a token refreshed after a {@code 401} reaches the resend. + * Mirrors flowvault's {@code AuthInterceptorTests}. + */ +public class AuthInterceptorTests { + + private static VaultConfig config() { + VaultConfig config = new VaultConfig(); + config.setVaultId("vault1"); + config.setClusterId("cluster1"); + config.setEnv(Env.DEV); + return config; + } + + private static VaultConfig configWithCredentialsString() { + Credentials credentials = new Credentials(); + credentials.setCredentialsString("{\"clientID\":\"c\"}"); + VaultConfig config = config(); + config.setCredentials(credentials); + return config; + } + + private static String jwtExpiringIn(long seconds) { + Base64.Encoder enc = Base64.getUrlEncoder().withoutPadding(); + String header = enc.encodeToString("{\"alg\":\"none\"}".getBytes(StandardCharsets.UTF_8)); + String payload = enc.encodeToString(("{\"exp\":" + (System.currentTimeMillis() / 1000 + seconds) + "}") + .getBytes(StandardCharsets.UTF_8)); + return header + "." + payload + ".sig"; + } + + /** Skyvault client with mintable credentials and a scripted token source. */ + private static final class MintingClient extends VaultClient { + final AtomicInteger generations = new AtomicInteger(); + final String fresh; + + MintingClient(String fresh) { + super(configWithCredentialsString(), null); + this.fresh = fresh; + this.finalCredentials = getVaultConfig().getCredentials(); + } + + @Override + protected String generateBearerToken(BaseCredentials credentials) { + generations.incrementAndGet(); + return fresh; + } + } + + /** updateExecutorInHTTP is private in skyvault; install the client the way setBearerToken does. */ + private static Interceptor authInterceptorOf(VaultClient client) throws Exception { + Method install = VaultClient.class.getDeclaredMethod("updateExecutorInHTTP"); + install.setAccessible(true); + install.invoke(client); + List interceptors = client.sharedHttpClient.interceptors(); + Assert.assertEquals("skyvault's shared client carries the auth interceptor alone", 1, interceptors.size()); + return interceptors.get(0); + } + + @Test + public void testAuthInterceptor_regeneratesOnceAndResendsOn401ForExpiredToken() throws Exception { + MintingClient client = new MintingClient("fresh-token"); + client.token = jwtExpiringIn(-1); + FakeChain chain = new FakeChain(401, 200); + + Assert.assertEquals(200, authInterceptorOf(client).intercept(chain).code()); + Assert.assertEquals(2, chain.calls()); + Assert.assertEquals("Bearer fresh-token", chain.lastProceeded().header("Authorization")); + Assert.assertEquals("the cache must hold the new token for the calls that follow", "fresh-token", client.token); + Assert.assertEquals(1, client.generations.get()); + Assert.assertTrue("the rejected response must be closed before the resend", chain.bodies().get(0).isClosed()); + } + + @Test + public void testAuthInterceptor_401ForTokenWithLifeLeftIsReturnedWithoutRegenerating() throws Exception { + MintingClient client = new MintingClient("fresh-token"); + String live = jwtExpiringIn(900); + client.token = live; + FakeChain chain = new FakeChain(401, 200); + + Assert.assertEquals(401, authInterceptorOf(client).intercept(chain).code()); + Assert.assertEquals(1, chain.calls()); + Assert.assertEquals(0, client.generations.get()); + Assert.assertEquals(live, client.token); + } + + @Test + public void testAuthInterceptor_401WithStaticTokenIsReturnedAsToday() throws Exception { + Credentials credentials = new Credentials(); + credentials.setToken(jwtExpiringIn(-1)); + VaultConfig config = config(); + config.setCredentials(credentials); + VaultClient client = new VaultClient(config, null); + client.setBearerToken(); + FakeChain chain = new FakeChain(401, 200); + + Assert.assertEquals(401, authInterceptorOf(client).intercept(chain).code()); + Assert.assertEquals(1, chain.calls()); + } + + @Test + public void testAuthInterceptor_401WithApiKeyIsReturnedAsToday() throws Exception { + // Loaded from dummy-non-secrets/ (Gitleaks-excluded) rather than a string literal, + // since an API-key-shaped literal would itself trip the scanner. + String apiKey = new String(Files.readAllBytes( + Paths.get("./src/test/resources/dummy-non-secrets/dummy-api-key.txt")), StandardCharsets.UTF_8).trim(); + Credentials credentials = new Credentials(); + credentials.setApiKey(apiKey); + VaultConfig config = config(); + config.setCredentials(credentials); + VaultClient client = new VaultClient(config, null); + client.setBearerToken(); + FakeChain chain = new FakeChain(401, 200); + + Assert.assertEquals(401, authInterceptorOf(client).intercept(chain).code()); + Assert.assertEquals(1, chain.calls()); + Assert.assertEquals("Bearer " + apiKey, chain.lastProceeded().header("Authorization")); + } + + @Test + public void testAuthInterceptor_resendThatAlsoGets401IsNotRetriedAgain() throws Exception { + MintingClient client = new MintingClient("fresh-token"); + client.token = jwtExpiringIn(-1); + FakeChain chain = new FakeChain(401, 401, 200); + + Assert.assertEquals(401, authInterceptorOf(client).intercept(chain).code()); + Assert.assertEquals(2, chain.calls()); + Assert.assertEquals(1, client.generations.get()); + } + + @Test + public void testAuthInterceptor_readsTheRefreshedTokenOnTheResend() throws Exception { + // setBearerToken() is what production code calls; it must install the same interceptor and + // the resend must carry the token that refresh stored, read back through getToken(). + MintingClient client = new MintingClient("fresh-token"); + client.setBearerToken(); + Assert.assertEquals("fresh-token", client.token); + client.token = jwtExpiringIn(-1); + FakeChain chain = new FakeChain(401, 200); + + Interceptor interceptor = client.sharedHttpClient.interceptors().get(0); + Assert.assertEquals(200, interceptor.intercept(chain).code()); + Assert.assertEquals("Bearer fresh-token", chain.lastProceeded().header("Authorization")); + } + + @Test + public void testAuthInterceptor_isTheOnlyInterceptorOnTheSharedClient() throws SkyflowException { + Credentials credentials = new Credentials(); + credentials.setToken(jwtExpiringIn(3600)); + VaultConfig config = config(); + config.setCredentials(credentials); + VaultClient client = new VaultClient(config, null); + client.setBearerToken(); + + Assert.assertEquals("no transport retry wraps auth in skyvault; a 401 reaches auth directly", + 1, client.sharedHttpClient.interceptors().size()); + } +} diff --git a/skyvault/src/test/java/com/skyflow/utils/FakeChain.java b/skyvault/src/test/java/com/skyflow/utils/FakeChain.java new file mode 100644 index 00000000..8945788a --- /dev/null +++ b/skyvault/src/test/java/com/skyflow/utils/FakeChain.java @@ -0,0 +1,140 @@ +package com.skyflow.utils; + +import okhttp3.Call; +import okhttp3.Connection; +import okhttp3.Interceptor; +import okhttp3.MediaType; +import okhttp3.Protocol; +import okhttp3.Request; +import okhttp3.Response; +import okhttp3.ResponseBody; +import okio.Buffer; +import okio.BufferedSource; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeUnit; + +/** + * Minimal Interceptor.Chain that replays a scripted list of status codes, so the retry loop can + * be driven without a network or a mock-server dependency. Records how many times it was called + * and which responses were closed. + */ +public final class FakeChain implements Interceptor.Chain { + + private final Request request = + new Request.Builder().url("https://cluster1.vault.skyflowapis.dev/v1/vaults/vault1/table1").build(); + private final int[] statusCodes; + private final List bodies = new ArrayList<>(); + private int calls; + private Request lastProceeded; + + public FakeChain(int... statusCodes) { + this.statusCodes = statusCodes; + } + + public int calls() { + return calls; + } + + /** Responses the interceptor superseded and should have closed. */ + public List bodies() { + return bodies; + } + + /** The request as it reached the next interceptor, i.e. after any rewriting. */ + public Request lastProceeded() { + return lastProceeded; + } + + @Override + public Request request() { + return request; + } + + @Override + public Response proceed(Request request) throws IOException { + this.lastProceeded = request; + // Past the end of the script, keep returning the last code. + int code = statusCodes[Math.min(calls, statusCodes.length - 1)]; + calls++; + TrackingBody body = new TrackingBody(); + bodies.add(body); + return new Response.Builder() + .request(request) + .protocol(Protocol.HTTP_1_1) + .code(code) + .message("status " + code) + .body(body) + .build(); + } + + @Override + public Connection connection() { + return null; + } + + @Override + public Call call() { + throw new UnsupportedOperationException(); + } + + @Override + public int connectTimeoutMillis() { + return 0; + } + + @Override + public Interceptor.Chain withConnectTimeout(int timeout, TimeUnit unit) { + return this; + } + + @Override + public int readTimeoutMillis() { + return 0; + } + + @Override + public Interceptor.Chain withReadTimeout(int timeout, TimeUnit unit) { + return this; + } + + @Override + public int writeTimeoutMillis() { + return 0; + } + + @Override + public Interceptor.Chain withWriteTimeout(int timeout, TimeUnit unit) { + return this; + } + + public static final class TrackingBody extends ResponseBody { + boolean closed; + + @Override + public MediaType contentType() { + return null; + } + + @Override + public long contentLength() { + return 0; + } + + @Override + public BufferedSource source() { + return new Buffer(); + } + + @Override + public void close() { + closed = true; + } + + public boolean isClosed() { + return closed; + } + } +} diff --git a/skyvault/src/test/java/com/skyflow/utils/UtilsTests.java b/skyvault/src/test/java/com/skyflow/utils/UtilsTests.java index d1e1c725..bba35780 100644 --- a/skyvault/src/test/java/com/skyflow/utils/UtilsTests.java +++ b/skyvault/src/test/java/com/skyflow/utils/UtilsTests.java @@ -76,6 +76,11 @@ public void testGetMetrics() { @Test public void testGetMetricsWithException() { + // The JDK reads os.name once, when its socket layer initialises; leaving it cleared breaks + // every later test in this JVM that opens a server socket. Restore the properties on exit. + String osName = System.getProperty("os.name"); + String osVersion = System.getProperty("os.version"); + String javaVersion = System.getProperty("java.version"); try { // Clearing System Properties explicitly to throw exception System.clearProperty("os.name"); @@ -89,6 +94,18 @@ public void testGetMetricsWithException() { Assert.assertTrue(metrics.get(Constants.SDK_METRIC_CLIENT_OS_DETAILS).getAsString().isEmpty()); } catch (Exception e) { Assert.fail(INVALID_EXCEPTION_THROWN); + } finally { + restore("os.name", osName); + restore("os.version", osVersion); + restore("java.version", javaVersion); + } + } + + private static void restore(String key, String value) { + if (value == null) { + System.clearProperty(key); + } else { + System.setProperty(key, value); } } } From d82f4f1d2c183f01aac5cb597b04bc7e799d6da2 Mon Sep 17 00:00:00 2001 From: yaswanth-pula-skyflow Date: Wed, 28 Oct 2026 09:46:31 +0530 Subject: [PATCH 4/4] SK-3252:Update null check. --- common/src/main/java/com/skyflow/BaseVaultClient.java | 3 ++- skyvault/src/test/java/com/skyflow/utils/UtilsTests.java | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/common/src/main/java/com/skyflow/BaseVaultClient.java b/common/src/main/java/com/skyflow/BaseVaultClient.java index c19ffa6f..9b1bc682 100644 --- a/common/src/main/java/com/skyflow/BaseVaultClient.java +++ b/common/src/main/java/com/skyflow/BaseVaultClient.java @@ -133,7 +133,8 @@ protected synchronized String refreshBearerTokenAfterUnauthorized(String rejecte LogUtil.printInfoLog(InfoLogs.BEARER_TOKEN_REJECTED_REUSING_REGENERATED.getLog()); return cached; } - if (rejected != null && rejected.equals(refreshFailedFor)) { + // rejected is non-null here: a null token fails the expiry gate above. + if (rejected.equals(refreshFailedFor)) { LogUtil.printInfoLog(InfoLogs.BEARER_TOKEN_REJECTED_REGENERATION_ALREADY_FAILED.getLog()); return null; } diff --git a/skyvault/src/test/java/com/skyflow/utils/UtilsTests.java b/skyvault/src/test/java/com/skyflow/utils/UtilsTests.java index bba35780..00b1ca9e 100644 --- a/skyvault/src/test/java/com/skyflow/utils/UtilsTests.java +++ b/skyvault/src/test/java/com/skyflow/utils/UtilsTests.java @@ -76,7 +76,7 @@ public void testGetMetrics() { @Test public void testGetMetricsWithException() { - // The JDK reads os.name once, when its socket layer initialises; leaving it cleared breaks + // The JDK reads os.name once, when its socket layer initializes; leaving it cleared breaks // every later test in this JVM that opens a server socket. Restore the properties on exit. String osName = System.getProperty("os.name"); String osVersion = System.getProperty("os.version");