feat(java): add GlideCredentialProvider support to IamAuthConfig - #6671
feat(java): add GlideCredentialProvider support to IamAuthConfig#6671affonsov wants to merge 1 commit into
Conversation
![valkey-review-bot[bot]](https://avatars.githubusercontent.com/in/3633497?s=60&v=4){kind=small}
There was a problem hiding this comment.
Reviewed the core wiring, the JNI bridge, and the Java surface. The feature shape looks right — credentials_provider: None keeps the default AWS-chain path byte-identical, the single non-test IAMTokenManager::new caller and the only IamAuthenticationConfig struct literal are both updated, and GlideNativeBridge.createClient has exactly one Java caller, so nothing else breaks on the signature change.
Four points below: the provider is invoked synchronously from async token-generation paths that run on a runtime the Java binding configures with one worker thread; a null return from the provider isn't guarded before it reaches get_array_length; the Java exception is cleared without capturing its message even though a provider failure only ever surfaces as a log line; and nothing exercises the JNI marshalling path the PR adds.
| let creds = if let Some(callback) = &state.credentials_provider { | ||
| let (access_key_id, secret_access_key, session_token) = callback()?; |
There was a problem hiding this comment.
callback() is a blocking call made directly inside an async fn. On the Java binding that callback is a JNI upcall into arbitrary user code (java/src/iam_token_callback.rs:68-81 attaches the thread and calls getCredentials() synchronously), and the documented use cases — Vault, custom STS assume-role — are network round trips.
Every path that reaches generate_token_static from a task runs on the Tokio executor: the background refresh task (tokio::spawn at iam/mod.rs:301), the reconnect token refresh (glide-core/src/client/reconnecting_connection.rs:76), and user-triggered refreshIamToken (java/src/lib.rs:2148). The Java binding builds that runtime with worker_threads(DEFAULT_RUNTIME_WORKER_THREADS) where the constant is 1 (java/src/jni_client.rs:83, 140-141), so for the duration of the user's getCredentials() no other future on that runtime makes progress — for every client in the JVM, not just the one with IAM configured. With the default 250 ms request timeout (glide-core/src/client/mod.rs:55), a provider that takes a few hundred milliseconds is enough to time out in-flight commands. The reconnect path is the worst case, since it blocks while connections are already down.
The existing get_signing_identity branch below is fully async, so this is new exposure. Move the callback off the executor — tokio::task::spawn_blocking with a cloned Arc of the provider, awaiting the join handle (the Java runtime already reserves max_blocking_threads(worker_threads * 2) at jni_client.rs:142, and the current-thread runtime used by the other bindings has a blocking pool too). If you do that, switch iam_token_callback.rs from attach_current_thread_as_daemon() to the guard-based attach_current_thread(), otherwise pooled blocking threads will retire (thread_keep_alive(60s)) while still attached to the JVM.
| IamCallbackError::CallFailed(err) | ||
| })?; | ||
|
|
||
| let array_obj: JObjectArray = result.l().map_err(IamCallbackError::InvalidReturn)?.into(); |
There was a problem hiding this comment.
A null return from getCredentials() is not guarded. result.l() yields a null JObject, .into() reinterprets it as a JObjectArray without any check, and it goes straight into get_array_length on the next line. Nothing on the Java side prevents this — GlideCredentialProvider.getCredentials() returns a plain String[], and a provider that returns null on a cache miss or a failed lookup is an ordinary mistake.
The sibling implementation guards exactly this: java/src/address_resolver.rs:149-152 null-checks the result of result.l() before using it. This function already null-checks the array elements (lines 110, 141), so the array itself being null is the one case that slips through.
| let array_obj: JObjectArray = result.l().map_err(IamCallbackError::InvalidReturn)?.into(); | |
| let returned = result.l().map_err(IamCallbackError::InvalidReturn)?; | |
| if returned.is_null() { | |
| return Err(IamCallbackError::InvalidCredentials( | |
| "getCredentials() returned null".to_string(), | |
| )); | |
| } | |
| let array_obj: JObjectArray = returned.into(); |
| let result = result.map_err(|err| { | ||
| if let Ok(true) = env.exception_check() { | ||
| let _ = env.exception_describe(); | ||
| let _ = env.exception_clear(); | ||
| } | ||
| IamCallbackError::CallFailed(err) | ||
| })?; |
There was a problem hiding this comment.
The Java throwable is described to the JVM's stderr and then cleared, so nothing about it reaches the error that propagates: IamCallbackError::CallFailed(err) where err is jni's opaque JavaException, which becomes GlideIAMError::CredentialsError("Failed to call getCredentials() on Java callback: Java exception was thrown").
That matters more here than in the address-resolver equivalent, because a provider failure is never surfaced to the caller. After generate_token_with_backoff exhausts its 8 attempts, create_iam_token_manager logs and returns None (glide-core/src/client/mod.rs:1890-1894); client creation then continues, and since info.iam_config.is_some() && iam_token_manager.is_some() is false, get_valkey_connection_info falls into the non-IAM branch at client/mod.rs:236-247 with password: None. The user gets a client that was created "successfully" and fails on every command, and the only tie back to their provider is a log line that says "Java exception was thrown". The exception_describe() stack trace lands on the JVM's raw stderr rather than in glide's log stream.
Capture the throwable before clearing it — env.exception_occurred(), then toString() (or getMessage()) on the returned object, then exception_clear() — and carry that string in IamCallbackError::CallFailed so the class and message show up in the propagated CredentialsError. getCredentials() is declared throws Exception, so throwing is a documented path, not an edge case.
|
|
||
| assertNotNull(iamConfig.getCredentialsProvider()); | ||
| // Invoke the provider and verify the returned credentials | ||
| String[] creds = iamConfig.getCredentialsProvider().getCredentials(); |
There was a problem hiding this comment.
These tests invoke the lambda directly from Java, so they cover Lombok's generated getter and nothing else — the JNI marshalling this PR adds is untested. The method descriptor "()[Ljava/lang/String;" and the name lookup at java/src/iam_token_callback.rs:47, the array-length and element extraction, and the exception path are all validated only at runtime, and a mismatch there fails late and quietly: JavaIamTokenCallback::new returning None makes createClient return 0 (java/src/lib.rs:1242-1244), which ConnectionManager reports as ClosingException("Failed to create client - Connection refused"). The Rust tests pass a Rust closure, so they exercise the core plumbing but never the bridge.
The parallel callback feature is covered end to end — AddressResolver lambdas go through the real native bridge in java/integTest/src/test/java/glide/ConnectionTests.java:978-1087 — and IAM integ scaffolding already exists (TestUtilities.createTestIamConfig, used by StandaloneClientTests.createStandaloneClientWithIam). One integ test that builds a client with a credentialsProvider and asserts the provider was invoked would pin the descriptor and the array handling; token signing itself needs no real AWS account, as test_iam_token_manager_with_custom_callback demonstrates.
05ed584 to
4981e3f
Compare
Implements GitHub issue #5429. Users can now configure a custom IamCredentialsProvider on IamAuthConfig to supply AWS credentials (access key, secret key, optional session token) for IAM token signing, instead of relying on the default AWS SDK credential chain. This is useful in environments where credentials come from non-standard sources such as HashiCorp Vault, custom STS assume-role flows, etc. Signed-off-by: affonsov <67347924+affonsov@users.noreply.github.com>
4981e3f to
98da166
Compare
1 participant
Summary
Adds support for a custom credentials provider (
GlideCredentialProvider) toIamAuthConfigin the Java client. Users can now supply their own AWS credentials (e.g. from HashiCorp Vault, custom STS assume-role flows, or any non-standard source) for IAM token signing, instead of being locked into the default AWS SDK credential chain.Issue link
This Pull Request is linked to issue: IamAuthConfig does not support custom credentials provider
Closes #5429
Features / Behaviour Changes
credentialsProviderfield (typeGlideCredentialProvider) toIamAuthConfigGlideCredentialProviderfunctional interface — a simple@FunctionalInterfacereturningString[]{accessKeyId, secretAccessKey, sessionToken?}credentialsProvideris set, the Rust core calls back into Java to get credentials for each IAM token signing operationcredentialsProviderisnull(default), behaviour is completely unchanged — the default AWS SDK credential chain is usedImplementation
Java client
IamAuthConfig.java— newcredentialsProvider: GlideCredentialProviderfield (optional, Lombok@Builder)GlideCredentialProvider.java(new) —@FunctionalInterfacewithString[] getCredentials() throws Exceptionreturning[accessKeyId, secretAccessKey, sessionToken?]GlideNativeBridge.java— updatedcreateClientnative method to accept the provider as a 3rd parameterConnectionManager.java— passes the provider fromIamAuthConfigto the native bridge when setRust core (
glide-core)iam/mod.rs— addedcredentials_provider: Option<CredentialsProviderFn>toIamTokenState; whenSome, invokes the closure to get credentials instead of callingget_signing_identity()(the default AWS SDK chain); addedCredentialsProviderFntype aliasclient/types.rs— addedcredentials_providerfield toIamAuthenticationConfigclient/mod.rs— wires the provider throughcreate_iam_token_managerJNI bridge (
java/src/)iam_token_callback.rs(new) —JavaIamTokenCallbackholds a JNIGlobalRefto the Java provider object and callsgetCredentials()via JNI with proper thread attachment (attach_current_thread_as_daemon) for the async Rust refresh tasklib.rs— updatedJava_glide_internal_GlideNativeBridge_createClientto accept and wire the provider; logs a warning if a provider is supplied without IAM configUsage example:
Limitations
Testing
ServerCredentialsTestcovering: custom provider, null session token, backward-compat default, optional fieldglide-core/src/iam/mod.rscovering: provider happy path, error propagation, default-path regressioncargo clippy -D warningsclean,cargo fmtclean,spotlessCheckcleanChecklist