diff --git a/README.md b/README.md index 08857dca6..d1165f398 100644 --- a/README.md +++ b/README.md @@ -203,6 +203,48 @@ OpenAIClient client = OpenAIOkHttpClient.builder() .build(); ``` +#### X.509 client certificate authentication + +Applications with a client certificate can exchange that certificate directly for short-lived +OpenAI access tokens without providing a JWT or implementing `SubjectTokenProvider`: + +```java +import com.openai.client.OpenAIClient; +import com.openai.client.okhttp.OpenAIOkHttpClient; +import com.openai.client.okhttp.X509Transport; +import com.openai.client.okhttp.X509WorkloadIdentity; +import java.time.Duration; +import javax.net.ssl.X509ExtendedKeyManager; +import javax.net.ssl.X509TrustManager; + +X509ExtendedKeyManager keyManager = /* load your PKCS#12 key manager */; +X509TrustManager trustManager = /* load your trusted server roots */; + +X509Transport transport = X509Transport.builder() + .keyManager(keyManager) + .certificateAlias("client-certificate") + .trustManager(trustManager) + .build(); + +X509WorkloadIdentity identity = X509WorkloadIdentity.builder() + .identityProviderId("your-identity-provider-id") + .serviceAccountId("your-service-account-id") + .transport(transport) + .refreshBuffer(Duration.ofMinutes(10)) // Optional; defaults to 20 minutes. + .build(); + +OpenAIClient client = OpenAIOkHttpClient.builder() + .x509WorkloadIdentity(identity) + .build(); +``` + +`OpenAIOkHttpClientAsync.builder()` supports the same `x509WorkloadIdentity` option. Tokens are +obtained lazily, cached, and refreshed before expiration. Both the token exchange and API requests +use the configured fixed certificate alias, isolated direct mutual-TLS connections, and native +hostname verification; redirects and custom transport settings are not supported. Close the SDK +client to release both connection pools. For a complete compilable PKCS#12 setup, see +[`X509WorkloadIdentityExample`](openai-java-example/src/main/java/com/openai/example/X509WorkloadIdentityExample.java). + #### Kubernetes service account token provider ```java diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClient.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClient.kt index 6539ce5dd..10eb8e2a4 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClient.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClient.kt @@ -61,6 +61,8 @@ class OpenAIOkHttpClient private constructor() { private var sslSocketFactory: SSLSocketFactory? = null private var trustManager: X509TrustManager? = null private var hostnameVerifier: HostnameVerifier? = null + private var x509WorkloadIdentity: X509WorkloadIdentity? = null + private var x509BaseUrl: String? = null /** * The executor service to use for running HTTP requests. @@ -252,7 +254,10 @@ class OpenAIOkHttpClient private constructor() { * * Defaults to the production environment: `https://api.openai.com/v1`. */ - fun baseUrl(baseUrl: String?) = apply { clientOptions.baseUrl(baseUrl) } + fun baseUrl(baseUrl: String?) = apply { + clientOptions.baseUrl(baseUrl) + x509BaseUrl = baseUrl + } /** Alias for calling [Builder.baseUrl] with `baseUrl.orElse(null)`. */ fun baseUrl(baseUrl: Optional) = baseUrl(baseUrl.getOrNull()) @@ -352,6 +357,11 @@ class OpenAIOkHttpClient private constructor() { fun workloadIdentity(workloadIdentity: Optional) = workloadIdentity(workloadIdentity.getOrNull()) + /** Authenticates requests using a fixed X.509 client certificate instead of an API key. */ + fun x509WorkloadIdentity(x509WorkloadIdentity: X509WorkloadIdentity) = apply { + this.x509WorkloadIdentity = x509WorkloadIdentity + } + fun azureServiceVersion(azureServiceVersion: AzureOpenAIServiceVersion) = apply { clientOptions.azureServiceVersion(azureServiceVersion) } @@ -463,7 +473,12 @@ class OpenAIOkHttpClient private constructor() { * * @see ClientOptions.Builder.fromEnv */ - fun fromEnv() = apply { clientOptions.fromEnv() } + fun fromEnv() = apply { + clientOptions.fromEnv() + (System.getProperty("openai.baseUrl") ?: System.getenv("OPENAI_BASE_URL"))?.let { + x509BaseUrl = it + } + } /** * Returns an immutable instance of [OpenAIClient]. @@ -472,22 +487,37 @@ class OpenAIOkHttpClient private constructor() { */ fun build(): OpenAIClient = OpenAIClientImpl( - clientOptions - .httpClient( - OkHttpClient.builder() - .timeout(clientOptions.timeout()) - .followRedirects(followRedirects) - .proxy(proxy) - .proxyAuthenticator(proxyAuthenticator) - .maxIdleConnections(maxIdleConnections) - .keepAliveDuration(keepAliveDuration) - .dispatcherExecutorService(dispatcherExecutorService) - .sslSocketFactory(sslSocketFactory) - .trustManager(trustManager) - .hostnameVerifier(hostnameVerifier) - .build() - ) - .build() + x509WorkloadIdentity?.let { identity -> + require( + proxy == null && + proxyAuthenticator == null && + maxIdleConnections == null && + keepAliveDuration == null && + dispatcherExecutorService == null && + sslSocketFactory == null && + trustManager == null && + hostnameVerifier == null + ) { + "X.509 workload identity cannot be combined with custom transport settings" + } + x509ClientOptions(clientOptions, identity, x509BaseUrl) + } + ?: clientOptions + .httpClient( + OkHttpClient.builder() + .timeout(clientOptions.timeout()) + .followRedirects(followRedirects) + .proxy(proxy) + .proxyAuthenticator(proxyAuthenticator) + .maxIdleConnections(maxIdleConnections) + .keepAliveDuration(keepAliveDuration) + .dispatcherExecutorService(dispatcherExecutorService) + .sslSocketFactory(sslSocketFactory) + .trustManager(trustManager) + .hostnameVerifier(hostnameVerifier) + .build() + ) + .build() ) } } diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientAsync.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientAsync.kt index af1afa25f..26049634e 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientAsync.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientAsync.kt @@ -61,6 +61,8 @@ class OpenAIOkHttpClientAsync private constructor() { private var sslSocketFactory: SSLSocketFactory? = null private var trustManager: X509TrustManager? = null private var hostnameVerifier: HostnameVerifier? = null + private var x509WorkloadIdentity: X509WorkloadIdentity? = null + private var x509BaseUrl: String? = null /** * The executor service to use for running HTTP requests. @@ -252,7 +254,10 @@ class OpenAIOkHttpClientAsync private constructor() { * * Defaults to the production environment: `https://api.openai.com/v1`. */ - fun baseUrl(baseUrl: String?) = apply { clientOptions.baseUrl(baseUrl) } + fun baseUrl(baseUrl: String?) = apply { + clientOptions.baseUrl(baseUrl) + x509BaseUrl = baseUrl + } /** Alias for calling [Builder.baseUrl] with `baseUrl.orElse(null)`. */ fun baseUrl(baseUrl: Optional) = baseUrl(baseUrl.getOrNull()) @@ -352,6 +357,11 @@ class OpenAIOkHttpClientAsync private constructor() { fun workloadIdentity(workloadIdentity: Optional) = workloadIdentity(workloadIdentity.getOrNull()) + /** Authenticates requests using a fixed X.509 client certificate instead of an API key. */ + fun x509WorkloadIdentity(x509WorkloadIdentity: X509WorkloadIdentity) = apply { + this.x509WorkloadIdentity = x509WorkloadIdentity + } + fun azureServiceVersion(azureServiceVersion: AzureOpenAIServiceVersion) = apply { clientOptions.azureServiceVersion(azureServiceVersion) } @@ -463,7 +473,12 @@ class OpenAIOkHttpClientAsync private constructor() { * * @see ClientOptions.Builder.fromEnv */ - fun fromEnv() = apply { clientOptions.fromEnv() } + fun fromEnv() = apply { + clientOptions.fromEnv() + (System.getProperty("openai.baseUrl") ?: System.getenv("OPENAI_BASE_URL"))?.let { + x509BaseUrl = it + } + } /** * Returns an immutable instance of [OpenAIClientAsync]. @@ -472,22 +487,37 @@ class OpenAIOkHttpClientAsync private constructor() { */ fun build(): OpenAIClientAsync = OpenAIClientAsyncImpl( - clientOptions - .httpClient( - OkHttpClient.builder() - .timeout(clientOptions.timeout()) - .followRedirects(followRedirects) - .proxy(proxy) - .proxyAuthenticator(proxyAuthenticator) - .maxIdleConnections(maxIdleConnections) - .keepAliveDuration(keepAliveDuration) - .dispatcherExecutorService(dispatcherExecutorService) - .sslSocketFactory(sslSocketFactory) - .trustManager(trustManager) - .hostnameVerifier(hostnameVerifier) - .build() - ) - .build() + x509WorkloadIdentity?.let { identity -> + require( + proxy == null && + proxyAuthenticator == null && + maxIdleConnections == null && + keepAliveDuration == null && + dispatcherExecutorService == null && + sslSocketFactory == null && + trustManager == null && + hostnameVerifier == null + ) { + "X.509 workload identity cannot be combined with custom transport settings" + } + x509ClientOptions(clientOptions, identity, x509BaseUrl) + } + ?: clientOptions + .httpClient( + OkHttpClient.builder() + .timeout(clientOptions.timeout()) + .followRedirects(followRedirects) + .proxy(proxy) + .proxyAuthenticator(proxyAuthenticator) + .maxIdleConnections(maxIdleConnections) + .keepAliveDuration(keepAliveDuration) + .dispatcherExecutorService(dispatcherExecutorService) + .sslSocketFactory(sslSocketFactory) + .trustManager(trustManager) + .hostnameVerifier(hostnameVerifier) + .build() + ) + .build() ) } } diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509WorkloadIdentity.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509WorkloadIdentity.kt new file mode 100644 index 000000000..b7f0debf8 --- /dev/null +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509WorkloadIdentity.kt @@ -0,0 +1,84 @@ +package com.openai.client.okhttp + +import com.openai.core.Timeout +import com.openai.core.checkRequired +import java.net.Proxy +import java.time.Duration + +/** Configuration for workload identity authenticated with a fixed X.509 client certificate. */ +class X509WorkloadIdentity +private constructor( + /** Identity provider resource ID. */ + @get:JvmName("identityProviderId") val identityProviderId: String, + /** Service account ID associated with the certificate identity. */ + @get:JvmName("serviceAccountId") val serviceAccountId: String, + /** Caller-owned, fixed-alias mutual TLS configuration. */ + @get:JvmName("transport") val transport: X509Transport, + /** How early the SDK refreshes an access token. Defaults to 20 minutes. */ + @get:JvmName("refreshBuffer") val refreshBuffer: Duration, + private val exchangeTestProxy: Proxy? = null, + private val apiTestProxy: Proxy? = null, +) { + + companion object { + @JvmStatic fun builder() = Builder() + } + + /** A Java-compatible builder for certificate-based workload identity. */ + class Builder internal constructor() { + + private var identityProviderId: String? = null + private var serviceAccountId: String? = null + private var transport: X509Transport? = null + private var refreshBuffer: Duration = Duration.ofMinutes(20) + + /** Sets the identity provider resource ID. */ + fun identityProviderId(identityProviderId: String) = apply { + this.identityProviderId = identityProviderId + } + + /** Sets the service account ID associated with the certificate identity. */ + fun serviceAccountId(serviceAccountId: String) = apply { + this.serviceAccountId = serviceAccountId + } + + /** Sets the existing fixed-alias X.509 mutual TLS transport capability. */ + fun transport(transport: X509Transport) = apply { this.transport = transport } + + /** Sets how early a cached token should be refreshed. Defaults to 20 minutes. */ + fun refreshBuffer(refreshBuffer: Duration) = apply { + require(!refreshBuffer.isNegative) { "refreshBuffer must not be negative" } + this.refreshBuffer = refreshBuffer + } + + /** Returns immutable certificate-based workload identity configuration. */ + fun build(): X509WorkloadIdentity = + X509WorkloadIdentity( + checkRequired("identityProviderId", identityProviderId).also { + require(it.isNotBlank()) { "identityProviderId must not be blank" } + }, + checkRequired("serviceAccountId", serviceAccountId).also { + require(it.isNotBlank()) { "serviceAccountId must not be blank" } + }, + checkRequired("transport", transport), + refreshBuffer, + ) + } + + @JvmSynthetic + internal fun bind(timeout: Timeout): BoundX509Transport = + if (exchangeTestProxy == null || apiTestProxy == null) transport.bind(timeout) + else transport.bindForTest(timeout, exchangeTestProxy, apiTestProxy) + + /** Test-only loopback CONNECT seam; production bindings always connect directly. */ + @JvmSynthetic + internal fun withTestProxies(exchangeProxy: Proxy, apiProxy: Proxy): X509WorkloadIdentity = + X509WorkloadIdentity( + identityProviderId, + serviceAccountId, + transport, + refreshBuffer, + exchangeProxy, + apiProxy, + ) +} diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509WorkloadIdentityAuthenticator.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509WorkloadIdentityAuthenticator.kt new file mode 100644 index 000000000..e69ab0e5b --- /dev/null +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509WorkloadIdentityAuthenticator.kt @@ -0,0 +1,242 @@ +package com.openai.client.okhttp + +import com.openai.core.ClientOptions +import com.openai.core.RequestOptions +import com.openai.core.http.Headers +import com.openai.core.http.HttpClient +import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpRequestAuthenticator +import com.openai.core.http.HttpResponse +import com.openai.errors.OpenAIException +import com.openai.errors.OpenAIInvalidDataException +import java.time.DateTimeException +import java.time.Duration +import java.time.Instant +import java.util.concurrent.CompletableFuture +import java.util.concurrent.CompletionException +import okhttp3.HttpUrl.Companion.toHttpUrl + +/** Connects the existing X.509 exchange to the SDK's owned provider-authentication pipeline. */ +internal class X509WorkloadIdentityAuthenticator( + private val exchange: X509TokenExchange, + private val exchangeClient: OkHttpClient, + private val refreshBuffer: Duration, +) : HttpRequestAuthenticator { + + private val lock = Any() + private var cachedToken: String? = null + private var tokenExpiry: Instant? = null + private var refreshInFlight: CompletableFuture? = null + private var closed = false + + override fun authenticate(request: HttpRequest): HttpRequest = + try { + requireAuthorizedOrigin(request) + authenticated(request, token(async = false).join()) + } catch (failure: CompletionException) { + throw failure.cause ?: failure + } + + override fun authenticateAsync(request: HttpRequest): CompletableFuture = + try { + requireAuthorizedOrigin(request) + token(async = true).thenApply { value -> authenticated(request, value) } + } catch (failure: Throwable) { + CompletableFuture().apply { completeExceptionally(failure) } + } + + private fun requireAuthorizedOrigin(request: HttpRequest) { + val url = request.url().toHttpUrl() + if ( + url.scheme != "https" || + (url.host != "mtls.api.openai.com" && url.host != "mtls-eu.api.openai.com") || + url.port != 443 || + url.encodedUsername.isNotEmpty() || + url.encodedPassword.isNotEmpty() + ) { + throw OpenAIException("X.509 request destination is not authorized") + } + } + + private fun authenticated(request: HttpRequest, token: String): HttpRequest = + request.toBuilder().replaceHeaders("Authorization", "Bearer $token").build() + + fun invalidate(rejectedAuthorization: String?) { + synchronized(lock) { + val token = cachedToken + if (token != null && rejectedAuthorization == "Bearer $token") { + cachedToken = null + tokenExpiry = null + } + } + } + + private fun token(async: Boolean): CompletableFuture { + val refresh: CompletableFuture + synchronized(lock) { + check(!closed) { "X.509 workload identity authentication is closed" } + val token = cachedToken + val expiry = tokenExpiry + if ( + token != null && + expiry != null && + Duration.between(Instant.now(), expiry) > refreshBuffer + ) { + return CompletableFuture.completedFuture(token) + } + refreshInFlight?.let { + return it.thenApply { value -> value } + } + refresh = CompletableFuture() + refreshInFlight = refresh + } + + if (async) { + try { + exchange.executeAsync().whenComplete { accessToken, failure -> + completeRefresh(refresh, accessToken, failure) + } + } catch (failure: Throwable) { + completeRefresh(refresh, null, failure) + } + } else { + try { + completeRefresh(refresh, exchange.execute(), null) + } catch (failure: Throwable) { + completeRefresh(refresh, null, failure) + } + } + return refresh.thenApply { value -> value } + } + + private fun completeRefresh( + refresh: CompletableFuture, + accessToken: X509AccessToken?, + failure: Throwable?, + ) { + var error = + if (failure is CompletionException && failure.cause != null) failure.cause!! + else failure + val expiry = + try { + if (error == null && accessToken != null) Instant.now().plus(accessToken.expiresIn) + else null + } catch (invalidExpiry: DateTimeException) { + error = + OpenAIInvalidDataException( + "Invalid X.509 access token expiration", + invalidExpiry, + ) + null + } catch (invalidExpiry: ArithmeticException) { + error = + OpenAIInvalidDataException( + "Invalid X.509 access token expiration", + invalidExpiry, + ) + null + } + synchronized(lock) { + if (refreshInFlight === refresh) refreshInFlight = null + if (error == null && accessToken != null && expiry != null && !closed) { + cachedToken = accessToken.value + tokenExpiry = expiry + } + } + val finalError = error + when { + finalError != null -> refresh.completeExceptionally(finalError) + accessToken != null -> refresh.complete(accessToken.value) + else -> refresh.completeExceptionally(IllegalStateException("X.509 token unavailable")) + } + } + + override fun close() { + val pending = + synchronized(lock) { + if (closed) return + closed = true + cachedToken = null + tokenExpiry = null + refreshInFlight.also { refreshInFlight = null } + } + pending?.cancel(true) + exchange.use { exchangeClient.close() } + } +} + +/** Restores the existing workload-identity behavior when an API rejects a cached bearer. */ +private class X509RefreshingHttpClient( + private val delegate: OkHttpClient, + private val authenticator: X509WorkloadIdentityAuthenticator, +) : HttpClient { + + override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse = + checkResponse(request, delegate.execute(request, requestOptions)) + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = + delegate.executeAsync(request, requestOptions).thenApply { response -> + checkResponse(request, response) + } + + private fun checkResponse(request: HttpRequest, response: HttpResponse): HttpResponse { + if (response.statusCode() != 401) return response + authenticator.invalidate(request.headers.values("Authorization").singleOrNull()) + return object : HttpResponse by response { + override fun headers(): Headers = + response.headers().toBuilder().replace("X-Should-Retry", "true").build() + } + } + + override fun close() { + try { + authenticator.close() + } catch (authenticatorFailure: Throwable) { + try { + delegate.close() + } catch (delegateFailure: Throwable) { + if (delegateFailure !== authenticatorFailure) { + authenticatorFailure.addSuppressed(delegateFailure) + } + } + throw authenticatorFailure + } + delegate.close() + } +} + +internal fun x509ClientOptions( + clientOptions: ClientOptions.Builder, + identity: X509WorkloadIdentity, + baseUrl: String?, +): ClientOptions { + val bound = identity.bind(clientOptions.timeout()) + var authenticator: X509WorkloadIdentityAuthenticator? = null + try { + authenticator = + X509WorkloadIdentityAuthenticator( + X509TokenExchange( + identity.identityProviderId, + identity.serviceAccountId, + bound.exchangeClient, + ), + bound.exchangeClient, + identity.refreshBuffer, + ) + return clientOptions + .baseUrl(baseUrl ?: "https://mtls.api.openai.com/v1") + .httpClient(X509RefreshingHttpClient(bound.apiClient, authenticator)) + .httpRequestAuthenticator(authenticator) + .build() + } catch (failure: Throwable) { + try { + authenticator?.use { bound.apiClient.close() } ?: bound.close() + } catch (closeFailure: Throwable) { + if (closeFailure !== failure) failure.addSuppressed(closeFailure) + } + throw failure + } +} diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509WorkloadIdentityIntegrationTest.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509WorkloadIdentityIntegrationTest.kt new file mode 100644 index 000000000..aa57f533e --- /dev/null +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509WorkloadIdentityIntegrationTest.kt @@ -0,0 +1,569 @@ +package com.openai.client.okhttp + +import com.fasterxml.jackson.databind.ObjectMapper +import com.openai.errors.OpenAIException +import com.openai.errors.OpenAIInvalidDataException +import java.net.InetSocketAddress +import java.net.Proxy +import java.time.Duration +import java.util.concurrent.ExecutionException +import java.util.concurrent.TimeUnit +import okhttp3.mockwebserver.MockResponse +import okhttp3.mockwebserver.MockWebServer +import okhttp3.tls.HandshakeCertificates +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatThrownBy +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.parallel.ResourceLock +import org.junit.jupiter.api.parallel.Resources + +@ResourceLock(Resources.SYSTEM_PROPERTIES) +internal class X509WorkloadIdentityIntegrationTest { + + private val jsonMapper = ObjectMapper() + + @Test + fun synchronousPublicClientExchangesLazilyAndCachesBearerOverRealMutualTls() { + verifyPublicClient(async = false) + } + + @Test + fun asynchronousPublicClientExchangesLazilyAndCachesBearerOverRealMutualTls() { + verifyPublicClient(async = true) + } + + @Test + fun synchronousPublicClientPreservesTheExplicitEuMutualTlsEndpoint() { + verifyPublicClient(async = false, apiHost = EU_API_HOST) + } + + @Test + fun asynchronousPublicClientPreservesTheExplicitEuMutualTlsEndpoint() { + verifyPublicClient(async = true, apiHost = EU_API_HOST) + } + + @Test + fun synchronousPublicClientPreservesTheEnvironmentConfiguredEuMutualTlsEndpoint() { + verifyEnvironmentConfiguredPublicClient(async = false) + } + + @Test + fun asynchronousPublicClientPreservesTheEnvironmentConfiguredEuMutualTlsEndpoint() { + verifyEnvironmentConfiguredPublicClient(async = true) + } + + @Test + fun synchronousClonedPublicClientCannotSendBearerToAnUnauthorizedOrigin() { + verifyClonedPublicClientRejectsUnauthorizedOrigin(async = false) + } + + @Test + fun asynchronousClonedPublicClientCannotSendBearerToAnUnauthorizedOrigin() { + verifyClonedPublicClientRejectsUnauthorizedOrigin(async = true) + } + + @Test + fun synchronousPublicClientRefreshesTokensWithinTheConfiguredBuffer() { + verifyRefresh(async = false) + } + + @Test + fun asynchronousPublicClientRefreshesTokensWithinTheConfiguredBuffer() { + verifyRefresh(async = true) + } + + @Test + fun synchronousPublicClientRefreshesAndRetriesRejectedAccessTokens() { + verifyRejectedTokenRecovery(async = false) + } + + @Test + fun asynchronousPublicClientRefreshesAndRetriesRejectedAccessTokens() { + verifyRejectedTokenRecovery(async = true) + } + + @Test + fun synchronousPublicClientRejectsUnrepresentableTokenExpiration() { + verifyUnrepresentableTokenExpiration(async = false) + } + + @Test + fun asynchronousPublicClientRejectsUnrepresentableTokenExpirationWithoutHanging() { + verifyUnrepresentableTokenExpiration(async = true) + } + + @Test + fun closingSynchronousPublicClientClosesItsTokenExchangeAuthentication() { + verifyPublicClientClosesAuthentication(async = false) + } + + @Test + fun closingAsynchronousPublicClientClosesItsTokenExchangeAuthentication() { + verifyPublicClientClosesAuthentication(async = true) + } + + @Test + fun certificateAuthenticationCannotBeCombinedWithApiKeysOrCustomTransports() { + val identity = X509TestIdentity.create("configuration identity") + val configuration = configuration(identity, emptyList()) + + assertThatThrownBy { + OpenAIOkHttpClient.builder() + .apiKey("fake-api-key") + .x509WorkloadIdentity(configuration) + .build() + } + .isInstanceOf(IllegalStateException::class.java) + .hasMessageContaining("cannot be combined") + assertThatThrownBy { + OpenAIOkHttpClientAsync.builder() + .apiKey("fake-api-key") + .x509WorkloadIdentity(configuration) + .build() + } + .isInstanceOf(IllegalStateException::class.java) + .hasMessageContaining("cannot be combined") + + val proxy = Proxy(Proxy.Type.HTTP, InetSocketAddress("localhost", 12345)) + assertThatThrownBy { + OpenAIOkHttpClient.builder() + .proxy(proxy) + .x509WorkloadIdentity(configuration) + .build() + } + .isInstanceOf(IllegalArgumentException::class.java) + .hasMessageContaining("custom transport settings") + assertThatThrownBy { + OpenAIOkHttpClientAsync.builder() + .proxy(proxy) + .x509WorkloadIdentity(configuration) + .build() + } + .isInstanceOf(IllegalArgumentException::class.java) + .hasMessageContaining("custom transport settings") + } + + @Test + fun configurationRequiresIdsAndTransportAndRejectsNegativeRefreshBuffers() { + assertThatThrownBy { X509WorkloadIdentity.builder().build() } + .isInstanceOf(IllegalStateException::class.java) + .hasMessageContaining("identityProviderId") + assertThatThrownBy { X509WorkloadIdentity.builder().refreshBuffer(Duration.ofSeconds(-1)) } + .isInstanceOf(IllegalArgumentException::class.java) + .hasMessageContaining("refreshBuffer") + } + + private fun verifyEnvironmentConfiguredPublicClient(async: Boolean) { + val previousBaseUrl = System.setProperty("openai.baseUrl", "https://$EU_API_HOST/v1") + try { + verifyPublicClient(async, EU_API_HOST, fromEnv = true) + } finally { + if (previousBaseUrl == null) { + System.clearProperty("openai.baseUrl") + } else { + System.setProperty("openai.baseUrl", previousBaseUrl) + } + } + } + + private fun verifyPublicClient( + async: Boolean, + apiHost: String = API_HOST, + fromEnv: Boolean = false, + ) { + val identity = X509TestIdentity.create("public certificate identity") + X509TestPeer(AUTH_HOST, identity.root.certificate).use { authPeer -> + X509TestPeer(apiHost, identity.root.certificate).use { apiPeer -> + authPeer.enqueue(tokenResponse(ACCESS_TOKEN)) + apiPeer.enqueue(filesResponse()) + apiPeer.server.enqueue(filesResponse()) + val configuration = + configuration( + identity, + listOf(authPeer.serverRootCertificate, apiPeer.serverRootCertificate), + ) + .withTestProxies(authPeer.proxy, apiPeer.proxy) + + if (async) { + val client = + OpenAIOkHttpClientAsync.builder() + .apply { + if (fromEnv) { + baseUrl("https://$API_HOST/v1") + fromEnv() + apiKey(null as String?) + } else if (apiHost != API_HOST) { + baseUrl("https://$apiHost/v1") + } + } + .x509WorkloadIdentity(configuration) + .maxRetries(0) + .build() + try { + assertThat(authPeer.server.requestCount).isZero() + client.files().list().get(5, TimeUnit.SECONDS) + client.files().list().get(5, TimeUnit.SECONDS) + } finally { + client.close() + } + } else { + val client = + OpenAIOkHttpClient.builder() + .apply { + if (fromEnv) { + baseUrl("https://$API_HOST/v1") + fromEnv() + apiKey(null as String?) + } else if (apiHost != API_HOST) { + baseUrl("https://$apiHost/v1") + } + } + .x509WorkloadIdentity(configuration) + .maxRetries(0) + .build() + try { + assertThat(authPeer.server.requestCount).isZero() + client.files().list() + client.files().list() + } finally { + client.close() + } + } + + assertThat(authPeer.takeRequest().requestLine) + .isEqualTo("CONNECT $AUTH_HOST:443 HTTP/1.1") + val exchangeRequest = authPeer.takeRequest() + assertThat(exchangeRequest.path).isEqualTo("/oauth/token") + assertThat(exchangeRequest.getHeader("Authorization")).isNull() + assertThat(jsonMapper.readTree(exchangeRequest.body.readUtf8())) + .isEqualTo(jsonMapper.readTree(TOKEN_REQUEST)) + assertThat(requireNotNull(exchangeRequest.handshake).peerCertificates.first()) + .isEqualTo(identity.leaf.certificate) + assertThat(authPeer.server.requestCount).isEqualTo(2) + assertThat(apiPeer.takeRequest().requestLine) + .isEqualTo("CONNECT $apiHost:443 HTTP/1.1") + repeat(2) { + val apiRequest = apiPeer.takeRequest() + assertThat(apiRequest.path).isEqualTo("/v1/files") + assertThat(apiRequest.getHeader("Authorization")) + .isEqualTo("Bearer $ACCESS_TOKEN") + assertThat(requireNotNull(apiRequest.handshake).peerCertificates.first()) + .isEqualTo(identity.leaf.certificate) + } + assertThat(authPeer.requestedServerNames).containsExactly(AUTH_HOST) + assertThat(apiPeer.requestedServerNames).containsExactly(apiHost) + } + } + } + + private fun verifyPublicClientClosesAuthentication(async: Boolean) { + val identity = X509TestIdentity.create("closed certificate identity") + val configuration = configuration(identity, emptyList()) + + if (async) { + val client = + OpenAIOkHttpClientAsync.builder() + .x509WorkloadIdentity(configuration) + .maxRetries(0) + .build() + client.close() + assertThatThrownBy { client.files().list().get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasRootCauseInstanceOf(IllegalStateException::class.java) + .hasMessageContaining("authentication is closed") + } else { + val client = + OpenAIOkHttpClient.builder() + .x509WorkloadIdentity(configuration) + .maxRetries(0) + .build() + client.close() + assertThatThrownBy { client.files().list() } + .isInstanceOf(IllegalStateException::class.java) + .hasMessageContaining("authentication is closed") + } + } + + private fun verifyClonedPublicClientRejectsUnauthorizedOrigin(async: Boolean) { + val identity = X509TestIdentity.create("origin-bound certificate identity") + X509TestPeer(AUTH_HOST, identity.root.certificate).use { authPeer -> + X509TestPeer(API_HOST, identity.root.certificate).use { apiPeer -> + MockWebServer().use { unauthorizedPeer -> + authPeer.enqueue(tokenResponse(ACCESS_TOKEN)) + apiPeer.enqueue(filesResponse()) + unauthorizedPeer.enqueue(filesResponse()) + val configuration = + configuration( + identity, + listOf( + authPeer.serverRootCertificate, + apiPeer.serverRootCertificate, + ), + ) + .withTestProxies(authPeer.proxy, apiPeer.proxy) + + if (async) { + val client = + OpenAIOkHttpClientAsync.builder() + .x509WorkloadIdentity(configuration) + .maxRetries(0) + .build() + try { + client.files().list().get(5, TimeUnit.SECONDS) + val cloned = + client.withOptions { + it.baseUrl(unauthorizedPeer.url("/v1").toString()) + it.httpClient(OkHttpClient.builder().build()) + } + try { + assertThatThrownBy { + cloned.files().list().get(5, TimeUnit.SECONDS) + } + .isInstanceOf(ExecutionException::class.java) + .hasRootCauseInstanceOf(OpenAIException::class.java) + .hasMessageContaining("destination is not authorized") + } finally { + cloned.close() + } + } finally { + client.close() + } + } else { + val client = + OpenAIOkHttpClient.builder() + .x509WorkloadIdentity(configuration) + .maxRetries(0) + .build() + try { + client.files().list() + val cloned = + client.withOptions { + it.baseUrl(unauthorizedPeer.url("/v1").toString()) + it.httpClient(OkHttpClient.builder().build()) + } + try { + assertThatThrownBy { cloned.files().list() } + .isInstanceOf(OpenAIException::class.java) + .hasMessageContaining("destination is not authorized") + } finally { + cloned.close() + } + } finally { + client.close() + } + } + + assertThat(unauthorizedPeer.requestCount).isZero() + assertThat(authPeer.server.requestCount).isEqualTo(2) + } + } + } + } + + private fun verifyRefresh(async: Boolean) { + val identity = X509TestIdentity.create("refresh certificate identity") + X509TestPeer(AUTH_HOST, identity.root.certificate).use { authPeer -> + X509TestPeer(API_HOST, identity.root.certificate).use { apiPeer -> + authPeer.enqueue(tokenResponse("first-test-token", expiresInSeconds = 1)) + authPeer.server.enqueue(tokenResponse("second-test-token", expiresInSeconds = 3600)) + apiPeer.enqueue(filesResponse()) + apiPeer.server.enqueue(filesResponse()) + val configuration = + configuration( + identity, + listOf(authPeer.serverRootCertificate, apiPeer.serverRootCertificate), + Duration.ofSeconds(2), + ) + .withTestProxies(authPeer.proxy, apiPeer.proxy) + + if (async) { + val client = + OpenAIOkHttpClientAsync.builder() + .x509WorkloadIdentity(configuration) + .maxRetries(0) + .build() + try { + client.files().list().get(5, TimeUnit.SECONDS) + client.files().list().get(5, TimeUnit.SECONDS) + } finally { + client.close() + } + } else { + val client = + OpenAIOkHttpClient.builder() + .x509WorkloadIdentity(configuration) + .maxRetries(0) + .build() + try { + client.files().list() + client.files().list() + } finally { + client.close() + } + } + + authPeer.takeRequest() + assertThat(authPeer.takeRequest().path).isEqualTo("/oauth/token") + assertThat(authPeer.takeRequest().path).isEqualTo("/oauth/token") + apiPeer.takeRequest() + assertThat(apiPeer.takeRequest().getHeader("Authorization")) + .isEqualTo("Bearer first-test-token") + assertThat(apiPeer.takeRequest().getHeader("Authorization")) + .isEqualTo("Bearer second-test-token") + } + } + } + + private fun verifyUnrepresentableTokenExpiration(async: Boolean) { + val identity = X509TestIdentity.create("invalid expiration certificate identity") + X509TestPeer(AUTH_HOST, identity.root.certificate).use { authPeer -> + X509TestPeer(API_HOST, identity.root.certificate).use { apiPeer -> + authPeer.enqueue(tokenResponse("fake-invalid-expiration-token", Long.MAX_VALUE)) + val configuration = + configuration( + identity, + listOf(authPeer.serverRootCertificate, apiPeer.serverRootCertificate), + ) + .withTestProxies(authPeer.proxy, apiPeer.proxy) + + if (async) { + val client = + OpenAIOkHttpClientAsync.builder() + .x509WorkloadIdentity(configuration) + .maxRetries(0) + .build() + try { + assertThatThrownBy { client.files().list().get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasCauseInstanceOf(OpenAIInvalidDataException::class.java) + .hasMessageContaining("token expiration") + } finally { + client.close() + } + } else { + val client = + OpenAIOkHttpClient.builder() + .x509WorkloadIdentity(configuration) + .maxRetries(0) + .build() + try { + assertThatThrownBy { client.files().list() } + .isInstanceOf(OpenAIInvalidDataException::class.java) + .hasMessageContaining("token expiration") + } finally { + client.close() + } + } + + assertThat(apiPeer.server.requestCount).isZero() + } + } + } + + private fun verifyRejectedTokenRecovery(async: Boolean) { + val identity = X509TestIdentity.create("rejected token certificate identity") + X509TestPeer(AUTH_HOST, identity.root.certificate).use { authPeer -> + X509TestPeer(API_HOST, identity.root.certificate).use { apiPeer -> + authPeer.enqueue(tokenResponse("rejected-test-token")) + authPeer.server.enqueue(tokenResponse("refreshed-test-token")) + apiPeer.enqueue(MockResponse().setResponseCode(401)) + apiPeer.server.enqueue(filesResponse()) + val configuration = + configuration( + identity, + listOf(authPeer.serverRootCertificate, apiPeer.serverRootCertificate), + ) + .withTestProxies(authPeer.proxy, apiPeer.proxy) + + if (async) { + val client = + OpenAIOkHttpClientAsync.builder() + .x509WorkloadIdentity(configuration) + .maxRetries(1) + .build() + try { + client.files().list().get(5, TimeUnit.SECONDS) + } finally { + client.close() + } + } else { + val client = + OpenAIOkHttpClient.builder() + .x509WorkloadIdentity(configuration) + .maxRetries(1) + .build() + try { + client.files().list() + } finally { + client.close() + } + } + + authPeer.takeRequest() + assertThat(authPeer.takeRequest().path).isEqualTo("/oauth/token") + assertThat(authPeer.takeRequest().path).isEqualTo("/oauth/token") + apiPeer.takeRequest() + assertThat(apiPeer.takeRequest().getHeader("Authorization")) + .isEqualTo("Bearer rejected-test-token") + assertThat(apiPeer.takeRequest().getHeader("Authorization")) + .isEqualTo("Bearer refreshed-test-token") + } + } + } + + private fun configuration( + identity: X509TestIdentity, + serverRoots: Iterable, + refreshBuffer: Duration = Duration.ofMinutes(20), + ): X509WorkloadIdentity { + val trustManager = + HandshakeCertificates.Builder() + .apply { serverRoots.forEach { addTrustedCertificate(it) } } + .build() + .trustManager + val transport = + X509Transport.builder() + .keyManager(x509TestKeyManager(mapOf(CERTIFICATE_ALIAS to identity))) + .certificateAlias(CERTIFICATE_ALIAS) + .trustManager(trustManager) + .build() + return X509WorkloadIdentity.builder() + .identityProviderId(IDENTITY_PROVIDER_ID) + .serviceAccountId(SERVICE_ACCOUNT_ID) + .transport(transport) + .refreshBuffer(refreshBuffer) + .build() + } + + private fun tokenResponse(value: String, expiresInSeconds: Long = 3600): MockResponse = + MockResponse() + .setHeader("Content-Type", "application/json") + .setBody( + """{"access_token":"$value","token_type":"Bearer","issued_token_type":"urn:ietf:params:oauth:token-type:access_token","expires_in":$expiresInSeconds}""" + ) + + private fun filesResponse(): MockResponse = + MockResponse() + .setHeader("Content-Type", "application/json") + .setBody("""{"object":"list","data":[]}""") + + private companion object { + const val AUTH_HOST = "mtls.auth.openai.com" + const val API_HOST = "mtls.api.openai.com" + const val EU_API_HOST = "mtls-eu.api.openai.com" + const val CERTIFICATE_ALIAS = "fixed-test-alias" + const val IDENTITY_PROVIDER_ID = "idp_test" + const val SERVICE_ACCOUNT_ID = "svc_acct_test" + const val ACCESS_TOKEN = "fake-x509-access-token" + val TOKEN_REQUEST = + """ + { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "subject_token_type": "urn:openai:params:oauth:token-type:x509", + "identity_provider_id": "$IDENTITY_PROVIDER_ID", + "service_account_id": "$SERVICE_ACCOUNT_ID" + } + """ + .trimIndent() + } +} diff --git a/openai-java-example/src/main/java/com/openai/example/X509WorkloadIdentityExample.java b/openai-java-example/src/main/java/com/openai/example/X509WorkloadIdentityExample.java new file mode 100644 index 000000000..1b62e4ce2 --- /dev/null +++ b/openai-java-example/src/main/java/com/openai/example/X509WorkloadIdentityExample.java @@ -0,0 +1,78 @@ +package com.openai.example; + +import com.openai.client.OpenAIClient; +import com.openai.client.okhttp.OpenAIOkHttpClient; +import com.openai.client.okhttp.X509Transport; +import com.openai.client.okhttp.X509WorkloadIdentity; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.security.KeyStore; +import java.time.Duration; +import java.util.Arrays; +import javax.net.ssl.KeyManagerFactory; +import javax.net.ssl.TrustManagerFactory; +import javax.net.ssl.X509ExtendedKeyManager; +import javax.net.ssl.X509TrustManager; + +public final class X509WorkloadIdentityExample { + private X509WorkloadIdentityExample() {} + + public static void main(String[] args) throws Exception { + char[] password = requiredEnvironment("OPENAI_X509_KEYSTORE_PASSWORD").toCharArray(); + try { + KeyStore keyStore = KeyStore.getInstance("PKCS12"); + try (InputStream input = + Files.newInputStream(Paths.get(requiredEnvironment("OPENAI_X509_KEYSTORE_PATH")))) { + keyStore.load(input, password); + } + + KeyManagerFactory keyManagers = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm()); + keyManagers.init(keyStore, password); + X509ExtendedKeyManager keyManager = Arrays.stream(keyManagers.getKeyManagers()) + .filter(X509ExtendedKeyManager.class::isInstance) + .map(X509ExtendedKeyManager.class::cast) + .findFirst() + .orElseThrow(() -> new IllegalStateException("No X.509 key manager available")); + + TrustManagerFactory trustManagers = + TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm()); + trustManagers.init((KeyStore) null); + X509TrustManager trustManager = Arrays.stream(trustManagers.getTrustManagers()) + .filter(X509TrustManager.class::isInstance) + .map(X509TrustManager.class::cast) + .findFirst() + .orElseThrow(() -> new IllegalStateException("No X.509 trust manager available")); + + X509Transport transport = X509Transport.builder() + .keyManager(keyManager) + .certificateAlias(requiredEnvironment("OPENAI_X509_CERTIFICATE_ALIAS")) + .trustManager(trustManager) + .build(); + X509WorkloadIdentity identity = X509WorkloadIdentity.builder() + .identityProviderId(requiredEnvironment("OPENAI_IDENTITY_PROVIDER_ID")) + .serviceAccountId(requiredEnvironment("OPENAI_SERVICE_ACCOUNT_ID")) + .transport(transport) + .refreshBuffer(Duration.ofMinutes(10)) + .build(); + + OpenAIClient client = + OpenAIOkHttpClient.builder().x509WorkloadIdentity(identity).build(); + try { + client.files().list(); + } finally { + client.close(); + } + } finally { + Arrays.fill(password, '\0'); + } + } + + private static String requiredEnvironment(String name) { + String value = System.getenv(name); + if (value == null || value.isBlank()) { + throw new IllegalArgumentException("Missing required environment variable: " + name); + } + return value; + } +}