SK-3253:Add gated retry on 401 status code for flowvault SDK

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;
@@ -33,6 +37,17 @@ class BaseVaultClient<V extends BaseVaultConfig> {
     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
+     * 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;
@@ -43,18 +58,116 @@ 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<String> 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<String> 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:
+     * <ul>
+     *   <li>the effective credentials are a static token or an API key, which the SDK cannot replace;</li>
+     *   <li>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;</li>
+     *   <li>generating a replacement failed, or produced the very token that was rejected;</li>
+     *   <li>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.</li>
+     * </ul>
+     * 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;
+        }
+        // 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;
+        }
+        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;
+            }
+            storeToken(fresh);
+            return fresh;
+        } catch (SkyflowException e) {
+            LogUtil.printWarningLog(WarningLogs.BEARER_TOKEN_REGENERATION_AFTER_401_FAILED.getLog()
+                    + " httpCode=" + e.getHttpCode()
+                    + (e.getRequestId() != null ? " requestId=" + e.getRequestId() : ""));
+            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);
+    }
+
     protected synchronized void prioritiseCredentials(BaseCredentials vaultSpecificCredentials) throws SkyflowException {
         try {
             BaseCredentials original = this.finalCredentials;
@@ -78,6 +191,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());
@@ -95,10 +209,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);
+            storeToken(generateBearerToken(this.finalCredentials));
         } else if (Token.isExpired(token)) {
             LogUtil.printInfoLog(InfoLogs.BEARER_TOKEN_EXPIRED.getLog());
-            token = BaseUtils.generateBearerToken(this.finalCredentials);
+            storeToken(generateBearerToken(this.finalCredentials));
         } else {
             LogUtil.printInfoLog(InfoLogs.REUSE_BEARER_TOKEN.getLog());
         }

common/src/main/java/com/skyflow/logs/InfoLogs.java

@@ -20,6 +20,10 @@ 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."),
+    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."),

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;

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
+     * <em>not</em> 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) {