diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt index 5a034f7bd..13ee78016 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt @@ -64,7 +64,10 @@ internal constructor(@JvmSynthetic internal val okHttpClient: okhttp3.OkHttpClie call.enqueue( object : Callback { override fun onResponse(call: Call, response: Response) { - future.complete(response.toHttpResponse()) + val httpResponse = response.toHttpResponse() + if (!future.complete(httpResponse)) { + httpResponse.close() + } } override fun onFailure(call: Call, e: IOException) { diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509TokenExchange.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509TokenExchange.kt new file mode 100644 index 000000000..73818bdcf --- /dev/null +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509TokenExchange.kt @@ -0,0 +1,708 @@ +package com.openai.client.okhttp + +import com.fasterxml.jackson.core.JsonParser +import com.fasterxml.jackson.core.JsonProcessingException +import com.fasterxml.jackson.core.JsonToken +import com.fasterxml.jackson.core.StreamReadFeature +import com.openai.core.http.Headers +import com.openai.core.http.HttpClient +import com.openai.core.http.HttpMethod +import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpRequestBody +import com.openai.core.http.HttpResponse +import com.openai.core.jsonMapper +import com.openai.errors.OpenAIInvalidDataException +import com.openai.errors.OpenAIIoException +import com.openai.errors.UnexpectedStatusCodeException +import com.openai.models.ErrorObject +import java.io.IOException +import java.io.InputStreamReader +import java.io.OutputStream +import java.io.Reader +import java.nio.charset.CharacterCodingException +import java.nio.charset.CodingErrorAction +import java.nio.charset.StandardCharsets +import java.time.Duration +import java.util.concurrent.ArrayBlockingQueue +import java.util.concurrent.CompletableFuture +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.ExecutionException +import java.util.concurrent.ThreadFactory +import java.util.concurrent.ThreadPoolExecutor +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.atomic.AtomicInteger +import java.util.concurrent.atomic.AtomicReference + +/** One validated access token. Its string representation never includes the credential. */ +internal class X509AccessToken(val value: String, val expiresIn: Duration) { + override fun toString(): String = "X509AccessToken{value=, expiresIn=$expiresIn}" +} + +/** Executes the fixed X.509 workload-identity token exchange without caching or retries. */ +internal class X509TokenExchange( + private val identityProviderId: String, + private val serviceAccountId: String, + private val httpClient: HttpClient, +) : AutoCloseable { + private val jsonMapper = jsonMapper() + private val responseReader = + jsonMapper.reader().with(StreamReadFeature.STRICT_DUPLICATE_DETECTION) + private val responseExecutor = + ThreadPoolExecutor( + MAX_RESPONSE_THREADS, + MAX_RESPONSE_THREADS, + RESPONSE_THREAD_IDLE_SECONDS, + TimeUnit.SECONDS, + ArrayBlockingQueue(MAX_QUEUED_EXCHANGES), + ResponseThreadFactory, + ) + .apply { allowCoreThreadTimeOut(true) } + private val lifecycleLock = Any() + private val operations = ConcurrentHashMap.newKeySet() + private val closed = AtomicBoolean() + + init { + require(identityProviderId.isNotBlank()) { "identityProviderId must not be blank" } + require(serviceAccountId.isNotBlank()) { "serviceAccountId must not be blank" } + } + + fun execute(): X509AccessToken { + checkOpen() + return httpClient.execute(request()).use(::parse) + } + + fun executeAsync(): CompletableFuture { + val operation = + synchronized(lifecycleLock) { + checkOpen() + AsyncOperation().also { operation -> + operations.add(operation) + try { + responseExecutor.execute(operation) + } catch (_: java.util.concurrent.RejectedExecutionException) { + operation.result.completeExceptionally( + OpenAIIoException("X.509 token exchange processing unavailable") + ) + } + } + } + return operation.result + } + + override fun close() { + val cancellations = + synchronized(lifecycleLock) { + if (closed.compareAndSet(false, true)) { + operations + .toTypedArray() + .filter { operation -> operation.prepareCancellation(true) } + .also { responseExecutor.shutdown() } + } else { + emptyList() + } + } + cancellations.forEach { operation -> operation.publishCancellation(true) } + } + + private fun checkOpen() { + check(!closed.get()) { "X.509 token exchange is closed" } + } + + private inner class AsyncOperation : Runnable { + private val terminal = AtomicBoolean() + private val cancellationRequested = AtomicBoolean() + private val enrollmentLock = Any() + val responseFuture = AtomicReference?>() + private val responseLeaseFuture = CompletableFuture() + private val activeResponse = AtomicReference() + val result = OperationResult(this) + + override fun run() { + try { + if (!initiateRequest()) return + val lease = + try { + responseLeaseFuture.get() + } catch (error: ExecutionException) { + throw error.cause ?: error + } + val token = lease.use { if (terminal.get()) null else parse(lease.response) } + if (token != null) result.complete(token) + } catch (error: Throwable) { + if (!closed.get() && !terminal.get()) result.completeExceptionally(error) + } finally { + activeResponse.getAndSet(null)?.close() + } + } + + private fun initiateRequest(): Boolean = + synchronized(enrollmentLock) { + if (terminal.get() || closed.get()) return@synchronized false + + val future = httpClient.executeAsync(request()) + future.whenComplete(::acceptResponse) + responseFuture.set(future) + if (terminal.get() || closed.get()) future.cancel(true) + true + } + + private fun acceptResponse(response: HttpResponse?, error: Throwable?) { + if (error != null) { + responseLeaseFuture.completeExceptionally(error) + return + } + if (response == null) { + responseLeaseFuture.completeExceptionally( + IllegalStateException("X.509 token exchange completed without a response") + ) + return + } + + val lease = ResponseLease(response) + activeResponse.set(lease) + if (terminal.get() && activeResponse.compareAndSet(lease, null)) lease.close() + responseLeaseFuture.complete(lease) + } + + private fun cancelResources(mayInterruptIfRunning: Boolean) { + val completedFuture = + synchronized(enrollmentLock) { + val future = responseFuture.get() + val canceled = future?.cancel(mayInterruptIfRunning) == true + activeResponse.getAndSet(null)?.close() + future?.takeIf { !canceled && it.isDone } + } + if (completedFuture != null) { + responseLeaseFuture.handle { _, _ -> Unit }.join() + activeResponse.getAndSet(null)?.close() + } + } + + fun cancel(mayInterruptIfRunning: Boolean): Boolean { + if (!prepareCancellation(mayInterruptIfRunning)) return false + return publishCancellation(mayInterruptIfRunning) + } + + fun prepareCancellation(mayInterruptIfRunning: Boolean): Boolean { + val firstCompletion = terminal.compareAndSet(false, true) + if (firstCompletion) cancellationRequested.set(true) + if (cancellationRequested.get()) cancelResources(mayInterruptIfRunning) + if (!firstCompletion) return false + + responseExecutor.remove(this) + operations.remove(this) + return true + } + + fun publishCancellation(mayInterruptIfRunning: Boolean): Boolean = + result.publishCancellation(mayInterruptIfRunning) + + fun complete(value: X509AccessToken): Boolean { + if (!terminal.compareAndSet(false, true)) return false + operations.remove(this) + return result.publishValue(value) + } + + fun completeExceptionally(error: Throwable): Boolean { + if (!terminal.compareAndSet(false, true)) return false + operations.remove(this) + return result.publishException(error) + } + } + + private inner class OperationResult(private val operation: AsyncOperation) : + CompletableFuture() { + override fun cancel(mayInterruptIfRunning: Boolean): Boolean = + operation.cancel(mayInterruptIfRunning) + + override fun complete(value: X509AccessToken): Boolean = operation.complete(value) + + override fun completeExceptionally(error: Throwable): Boolean = + operation.completeExceptionally(error) + + fun publishCancellation(mayInterruptIfRunning: Boolean): Boolean = + super.cancel(mayInterruptIfRunning) + + fun publishValue(value: X509AccessToken): Boolean = super.complete(value) + + fun publishException(error: Throwable): Boolean = super.completeExceptionally(error) + } + + private class ResponseLease(val response: HttpResponse) : AutoCloseable { + private val closed = AtomicBoolean() + + override fun close() { + if (closed.compareAndSet(false, true)) response.close() + } + } + + private fun request(): HttpRequest { + val bytes = + jsonMapper.writeValueAsBytes( + linkedMapOf( + "grant_type" to TOKEN_EXCHANGE_GRANT_TYPE, + "subject_token_type" to X509_TOKEN_TYPE, + "identity_provider_id" to identityProviderId, + "service_account_id" to serviceAccountId, + ) + ) + val body = + object : HttpRequestBody { + override fun writeTo(outputStream: OutputStream) = outputStream.write(bytes) + + override fun contentType(): String = "application/json" + + override fun contentLength(): Long = bytes.size.toLong() + + override fun repeatable(): Boolean = true + + override fun close() {} + } + return HttpRequest.builder() + .method(HttpMethod.POST) + .baseUrl(TOKEN_EXCHANGE_URL) + .body(body) + .build() + } + + private fun parse(response: HttpResponse): X509AccessToken { + val statusCode = response.statusCode() + if (statusCode != 200) { + val builder = + UnexpectedStatusCodeException.builder() + .statusCode(statusCode) + .headers(safeDiagnosticHeaders(response.headers())) + readOAuthError(response)?.let(builder::error) + throw builder.build() + } + + return try { + responseParser(response, SUCCESS_STRING_LIMITS).use(::parseSuccessResponse) + } catch (error: JsonProcessingException) { + if (isInvalidResponseFailure(error)) throw invalidResponse() + transportFailure(error)?.let { throw readFailure(it) } + throw invalidResponse() + } catch (_: X509ResponseConstraintException) { + throw invalidResponse() + } catch (_: CharacterCodingException) { + throw invalidResponse() + } catch (error: IOException) { + throw readFailure(error) + } + } + + private fun parseSuccessResponse(parser: JsonParser): X509AccessToken { + if (parser.nextToken() != JsonToken.START_OBJECT) throw invalidResponse() + + var accessToken: String? = null + var tokenType: String? = null + var issuedTokenType: String? = null + var expiresIn: Long? = null + while (parser.nextToken() != JsonToken.END_OBJECT) { + if (parser.currentToken() != JsonToken.FIELD_NAME) throw invalidResponse() + val field = parser.currentName() + val valueToken = parser.nextToken() ?: throw invalidResponse() + when (field) { + "access_token" -> { + if (valueToken != JsonToken.VALUE_STRING) { + throw invalidResponse("access_token") + } + accessToken = parser.text.takeIf(String::isNotBlank) + } + "token_type" -> { + if (valueToken != JsonToken.VALUE_STRING) throw invalidResponse("token_type") + tokenType = parser.text + } + "issued_token_type" -> { + if (valueToken != JsonToken.VALUE_STRING) { + throw invalidResponse("issued_token_type") + } + issuedTokenType = parser.text + } + "expires_in" -> { + if (valueToken != JsonToken.VALUE_NUMBER_INT) { + throw invalidResponse("expires_in") + } + expiresIn = parser.longValue.takeIf { it > 0 } + } + else -> parser.skipChildren() + } + } + if (parser.nextToken() != null) throw invalidResponse() + + val validatedAccessToken = + accessToken?.takeIf(BEARER_TOKEN_PATTERN::matches) + ?: throw invalidResponse("access_token") + if (tokenType?.equals("Bearer", ignoreCase = true) != true) { + throw invalidResponse("token_type") + } + if (issuedTokenType != ACCESS_TOKEN_TYPE) { + throw invalidResponse("issued_token_type") + } + return X509AccessToken( + validatedAccessToken, + Duration.ofSeconds(expiresIn ?: throw invalidResponse("expires_in")), + ) + } + + private fun readOAuthError(response: HttpResponse): ErrorObject? = + try { + responseParser(response, OAUTH_ERROR_STRING_LIMITS).use(::parseOAuthError) + } catch (error: JsonProcessingException) { + if (!isInvalidResponseFailure(error)) { + transportFailure(error)?.let { throw readFailure(it) } + } + null + } catch (_: X509ResponseConstraintException) { + null + } catch (_: CharacterCodingException) { + null + } catch (error: IOException) { + throw readFailure(error) + } catch (_: RuntimeException) { + null + } + + private fun parseOAuthError(parser: JsonParser): ErrorObject? { + if (parser.nextToken() != JsonToken.START_OBJECT) return null + + var errorCode: String? = null + while (parser.nextToken() != JsonToken.END_OBJECT) { + if (parser.currentToken() != JsonToken.FIELD_NAME) return null + val field = parser.currentName() + val valueToken = parser.nextToken() ?: return null + when (field) { + "error" -> { + if (valueToken != JsonToken.VALUE_STRING) return null + errorCode = parser.text + } + "error_description" -> if (valueToken != JsonToken.VALUE_STRING) return null + else -> parser.skipChildren() + } + } + if (parser.nextToken() != null) return null + + val safeCode = errorCode?.takeIf(SAFE_OAUTH_ERROR_CODES::contains) + val message = safeCode ?: return null + return ErrorObject.builder() + .code(safeCode) + .message(message) + .param(null) + .type("oauth_error") + .build() + } + + private fun responseParser(response: HttpResponse, stringLimits: Map): JsonParser { + val decoder = + StandardCharsets.UTF_8.newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT) + val reader = BoundedFieldReader(InputStreamReader(response.body(), decoder), stringLimits) + return responseReader.createParser(reader) + } + + private fun invalidResponse(field: String? = null): OpenAIInvalidDataException = + OpenAIInvalidDataException( + if (field == null) "Invalid X.509 token exchange response" + else "Invalid X.509 token exchange response field: $field" + ) + + private fun readFailure(cause: IOException): OpenAIIoException = + OpenAIIoException("Failed to read X.509 token exchange response", cause) + + private fun transportFailure(error: JsonProcessingException): IOException? { + var cause = error.cause + while (cause != null && cause !== error) { + if (cause is IOException && cause !is JsonProcessingException) return cause + cause = cause.cause + } + return null + } + + private fun isInvalidResponseFailure(error: Throwable): Boolean { + var cause: Throwable? = error + while (cause != null) { + if (cause is X509ResponseConstraintException || cause is CharacterCodingException) { + return true + } + cause = cause.cause + } + return false + } + + private fun safeDiagnosticHeaders(headers: Headers): Headers = + Headers.builder() + .apply { + SAFE_DIAGNOSTIC_HEADERS.forEach { name -> + put( + name, + headers.values(name).mapNotNull { value -> + safeDiagnosticHeaderValue(name, value) + }, + ) + } + } + .build() + + private fun safeDiagnosticHeaderValue(name: String, value: String): String? { + if (name in OPAQUE_DIAGNOSTIC_HEADERS) { + return if (value.isEmpty()) null else "" + } + if (value.length > MAX_DIAGNOSTIC_HEADER_CHARS) return null + + return when (name) { + "Content-Length", + "Retry-After", + "Retry-After-Ms" -> value.takeIf { it.isNotEmpty() && it.all { it in '0'..'9' } } + "Content-Type" -> + value.trim().lowercase().takeIf(SAFE_DIAGNOSTIC_CONTENT_TYPES::contains) + "X-Should-Retry" -> value.lowercase().takeIf { it == "true" || it == "false" } + else -> null + } + } + + private companion object { + const val TOKEN_EXCHANGE_URL = "https://mtls.auth.openai.com/oauth/token" + const val TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" + const val X509_TOKEN_TYPE = "urn:openai:params:oauth:token-type:x509" + const val ACCESS_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token" + const val MAX_RESPONSE_THREADS = 4 + const val MAX_QUEUED_EXCHANGES = MAX_RESPONSE_THREADS + const val RESPONSE_THREAD_IDLE_SECONDS = 30L + const val MAX_DIAGNOSTIC_HEADER_CHARS = 64 + const val MAX_TOKEN_TYPE_CHARS = 32 + const val MAX_ISSUED_TOKEN_TYPE_CHARS = 128 + const val MAX_OAUTH_ERROR_CODE_CHARS = 128 + const val MAX_OAUTH_ERROR_DESCRIPTION_CHARS = 1024 + val SUCCESS_STRING_LIMITS = + mapOf( + "token_type" to MAX_TOKEN_TYPE_CHARS, + "issued_token_type" to MAX_ISSUED_TOKEN_TYPE_CHARS, + ) + val OAUTH_ERROR_STRING_LIMITS = + mapOf( + "error" to MAX_OAUTH_ERROR_CODE_CHARS, + "error_description" to MAX_OAUTH_ERROR_DESCRIPTION_CHARS, + ) + val BEARER_TOKEN_PATTERN = Regex("[A-Za-z0-9._~+/-]+=*") + val SAFE_OAUTH_ERROR_CODES = + setOf( + "invalid_client", + "invalid_grant", + "invalid_request", + "invalid_scope", + "invalid_target", + "unauthorized_client", + "unsupported_grant_type", + ) + val SAFE_DIAGNOSTIC_HEADERS = + setOf( + "Content-Length", + "Content-Type", + "OpenAI-Request-ID", + "Request-ID", + "Retry-After", + "Retry-After-Ms", + "Traceparent", + "Tracestate", + "X-Request-ID", + "X-Should-Retry", + ) + val OPAQUE_DIAGNOSTIC_HEADERS = + setOf("OpenAI-Request-ID", "Request-ID", "Traceparent", "Tracestate", "X-Request-ID") + val SAFE_DIAGNOSTIC_CONTENT_TYPES = + setOf( + "application/json", + "application/json; charset=utf-8", + "application/problem+json", + "application/problem+json; charset=utf-8", + ) + } +} + +private class X509ResponseConstraintException : IOException("Invalid X.509 issuer response") + +/** + * Stops recognized top-level strings before Jackson materializes an oversized value. Jackson still + * owns JSON grammar, duplicate detection, unknown-value skipping, and trailing-token validation. + */ +private class BoundedFieldReader( + private val delegate: Reader, + private val stringLimits: Map, +) : Reader() { + private enum class StringKind { + FIELD_NAME, + FIELD_VALUE, + OTHER, + } + + private val maxFieldNameChars = stringLimits.keys.maxOf(String::length) + private var depth = 0 + private var previousSignificant: Char? = null + private var stringKind: StringKind? = null + private var fieldName = StringBuilder(maxFieldNameChars) + private var fieldNameTooLong = false + private var pendingFieldName: String? = null + private var valueLimit: Int? = null + private var valueChars = 0 + private var escaped = false + private var unicodeEscapeDigits = 0 + private var unicodeEscapeValue = 0 + private var firstCharacter = true + + override fun read(characters: CharArray, offset: Int, length: Int): Int { + if (length == 0) return 0 + + while (true) { + val read = delegate.read(characters, offset, length) + if (read <= 0) return read + + var monitoredOffset = offset + var monitoredLength = read + if (firstCharacter) { + firstCharacter = false + if (characters[monitoredOffset] == UTF8_BOM) { + monitoredOffset++ + monitoredLength-- + if (monitoredLength == 0) continue + characters.copyInto( + characters, + destinationOffset = offset, + startIndex = monitoredOffset, + endIndex = monitoredOffset + monitoredLength, + ) + monitoredOffset = offset + } + } + + for (index in monitoredOffset until monitoredOffset + monitoredLength) { + inspect(characters[index]) + } + return monitoredLength + } + } + + override fun close() = delegate.close() + + private fun inspect(character: Char) { + if (stringKind != null) { + inspectString(character) + return + } + + if (character == '"') { + startString() + return + } + if (!character.isWhitespace()) { + when (character) { + '{', + '[' -> depth++ + '}', + ']' -> depth-- + } + previousSignificant = character + } + } + + private fun startString() { + stringKind = + when { + depth == 1 && (previousSignificant == '{' || previousSignificant == ',') -> { + fieldName = StringBuilder(maxFieldNameChars) + fieldNameTooLong = false + StringKind.FIELD_NAME + } + depth == 1 && previousSignificant == ':' -> { + valueLimit = stringLimits[pendingFieldName] + valueChars = 0 + StringKind.FIELD_VALUE + } + else -> StringKind.OTHER + } + escaped = false + unicodeEscapeDigits = 0 + unicodeEscapeValue = 0 + } + + private fun inspectString(character: Char) { + if (unicodeEscapeDigits > 0) { + val digit = Character.digit(character, 16) + if (digit < 0) { + fieldNameTooLong = true + } else { + unicodeEscapeValue = (unicodeEscapeValue shl 4) or digit + } + unicodeEscapeDigits-- + if (unicodeEscapeDigits == 0) acceptDecoded(unicodeEscapeValue.toChar()) + return + } + if (escaped) { + escaped = false + if (character == 'u') { + unicodeEscapeDigits = 4 + unicodeEscapeValue = 0 + } else { + acceptDecoded( + when (character) { + '"', + '\\', + '/' -> character + 'b' -> '\b' + 'f' -> '\u000C' + 'n' -> '\n' + 'r' -> '\r' + 't' -> '\t' + else -> character + } + ) + } + return + } + when (character) { + '\\' -> escaped = true + '"' -> finishString() + else -> acceptDecoded(character) + } + } + + private fun acceptDecoded(character: Char) { + when (stringKind) { + StringKind.FIELD_NAME -> { + if (fieldName.length < maxFieldNameChars) fieldName.append(character) + else fieldNameTooLong = true + } + StringKind.FIELD_VALUE -> { + valueChars++ + if (valueLimit?.let { valueChars > it } == true) { + throw X509ResponseConstraintException() + } + } + StringKind.OTHER, + null -> {} + } + } + + private fun finishString() { + if (stringKind == StringKind.FIELD_NAME) { + pendingFieldName = fieldName.takeUnless { fieldNameTooLong }?.toString() + } + stringKind = null + valueLimit = null + previousSignificant = '"' + } + + private companion object { + const val UTF8_BOM = '\uFEFF' + } +} + +private object ResponseThreadFactory : ThreadFactory { + private val threadNumber = AtomicInteger() + + override fun newThread(runnable: Runnable): Thread = + Thread(runnable, "openai-x509-response-${threadNumber.incrementAndGet()}").apply { + isDaemon = true + } +} diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt index b202c83c6..7f87a5137 100644 --- a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt @@ -5,6 +5,17 @@ import com.github.tomakehurst.wiremock.junit5.WireMockRuntimeInfo import com.github.tomakehurst.wiremock.junit5.WireMockTest import com.openai.core.http.HttpMethod import com.openai.core.http.HttpRequest +import java.util.concurrent.CountDownLatch +import java.util.concurrent.TimeUnit +import okhttp3.Call +import okhttp3.EventListener +import okhttp3.Interceptor +import okhttp3.MediaType +import okhttp3.Protocol +import okhttp3.Response +import okhttp3.ResponseBody +import okio.Buffer +import okio.BufferedSource import org.assertj.core.api.Assertions.assertThat import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test @@ -41,4 +52,67 @@ internal class OkHttpClientTest { // Should have cancelled the underlying call assertThat(call.isCanceled()).isTrue() } + + @Test + fun executeAsync_whenResponseLosesCancellationRace_closesDroppedResponse() { + val handoffStarted = CountDownLatch(1) + val releaseHandoff = CountDownLatch(1) + val responseBody = TrackingResponseBody() + val interceptor = Interceptor { chain -> + Response.Builder() + .request(chain.request()) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(responseBody) + .build() + } + val handoffListener = + object : EventListener() { + override fun callEnd(call: Call) { + handoffStarted.countDown() + check(releaseHandoff.await(5, TimeUnit.SECONDS)) + } + } + val client = + OkHttpClient( + okhttp3.OkHttpClient.Builder() + .addInterceptor(interceptor) + .eventListener(handoffListener) + .build() + ) + + client.use { + val responseFuture = + client.executeAsync( + HttpRequest.builder() + .method(HttpMethod.POST) + .baseUrl("https://example.test") + .build() + ) + assertThat(handoffStarted.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(responseFuture.cancel(true)).isTrue() + + releaseHandoff.countDown() + + assertThat(responseBody.closed.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(responseFuture.isCancelled).isTrue() + } + } +} + +private class TrackingResponseBody : ResponseBody() { + val closed = CountDownLatch(1) + private val source = Buffer() + + override fun contentType(): MediaType? = null + + override fun contentLength(): Long = 0 + + override fun source(): BufferedSource = source + + override fun close() { + super.close() + closed.countDown() + } } diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509TokenExchangeTest.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509TokenExchangeTest.kt new file mode 100644 index 000000000..1c150d508 --- /dev/null +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509TokenExchangeTest.kt @@ -0,0 +1,1087 @@ +package com.openai.client.okhttp + +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.node.ObjectNode +import com.openai.core.RequestOptions +import com.openai.core.Timeout +import com.openai.core.http.Headers +import com.openai.core.http.HttpClient +import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpResponse +import com.openai.errors.OpenAIInvalidDataException +import com.openai.errors.OpenAIIoException +import com.openai.errors.UnexpectedStatusCodeException +import java.io.ByteArrayInputStream +import java.io.IOException +import java.io.InputStream +import java.net.SocketTimeoutException +import java.security.cert.X509Certificate +import java.time.Duration +import java.util.concurrent.CompletableFuture +import java.util.concurrent.ConcurrentLinkedQueue +import java.util.concurrent.CountDownLatch +import java.util.concurrent.ExecutionException +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicInteger +import java.util.function.BiConsumer +import okhttp3.mockwebserver.MockResponse +import okhttp3.tls.HandshakeCertificates +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test + +internal class X509TokenExchangeTest { + + @Test + fun exchangesSynchronouslyOverPinnedMtlsUsingTheExactProtocol() { + exchangesOverPinnedMtlsUsingTheExactProtocol(async = false) + } + + @Test + fun exchangesAsynchronouslyOverPinnedMtlsUsingTheExactProtocol() { + exchangesOverPinnedMtlsUsingTheExactProtocol(async = true) + } + + private fun exchangesOverPinnedMtlsUsingTheExactProtocol(async: Boolean) { + val clientIdentity = X509TestIdentity.create("exchange client") + X509TestPeer(AUTH_HOST, clientIdentity.root.certificate).use { authPeer -> + authPeer.enqueue( + MockResponse().setHeader("Content-Type", "application/json").setBody(TOKEN_RESPONSE) + ) + val transport = transport(clientIdentity, listOf(authPeer.serverRootCertificate)) + + val token = + transport.bindForTest(Timeout.default(), authPeer.proxy, authPeer.proxy).use { bound + -> + X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, bound.exchangeClient).use { + exchange -> + if (async) exchange.executeAsync().get(5, TimeUnit.SECONDS) + else exchange.execute() + } + } + + assertThat(token.value).isEqualTo(ACCESS_TOKEN) + assertThat(token.expiresIn).isEqualTo(Duration.ofHours(1)) + assertThat(token.toString()).doesNotContain(ACCESS_TOKEN) + + val connect = authPeer.takeRequest() + val request = authPeer.takeRequest() + assertThat(connect.requestLine).isEqualTo("CONNECT $AUTH_HOST:443 HTTP/1.1") + assertThat(request.method).isEqualTo("POST") + assertThat(request.path).isEqualTo("/oauth/token") + assertThat(request.getHeader("Authorization")).isNull() + assertThat(request.getHeader("Cookie")).isNull() + assertThat(request.getHeader("Content-Type")).startsWith("application/json") + assertThat(ObjectMapper().readTree(request.body.readUtf8())) + .isEqualTo(ObjectMapper().readTree(TOKEN_REQUEST)) + val presentedChain = requireNotNull(request.handshake).peerCertificates + assertThat(presentedChain.first()).isEqualTo(clientIdentity.leaf.certificate) + assertThat(presentedChain).contains(clientIdentity.root.certificate) + assertThat(authPeer.requestedServerNames).containsExactly(AUTH_HOST) + } + } + + @Test + fun rejectsMalformedOrSemanticallyInvalidSuccessResponsesWithoutLeakingCredentials() { + val invalidBodies = + mapOf( + "malformed" to "not-json", + "trailing garbage" to "${validResponse()} trailing", + "second root" to "${validResponse()} {}", + "duplicate access_token" to + validResponseWithDuplicate( + "access_token", + "\"secret-duplicate-token-must-not-leak\"", + ), + "duplicate expires_in" to validResponseWithDuplicate("expires_in", "60"), + "missing access_token" to validResponseWithout("access_token"), + "missing token_type" to validResponseWithout("token_type"), + "null token_type" to validResponseWithNull("token_type"), + "missing issued_token_type" to validResponseWithout("issued_token_type"), + "null issued_token_type" to validResponseWithNull("issued_token_type"), + "invalid access_token" to validResponse(accessToken = "secret token with spaces"), + "invalid token_type" to + validResponse(accessToken = "secret-token-must-not-leak", tokenType = "MAC"), + "invalid issued_token_type" to validResponse(issuedTokenType = "refresh_token"), + "missing expires_in" to validResponseWithout("expires_in"), + "zero expires_in" to validResponse(expiresIn = "0"), + "fractional expires_in" to validResponse(expiresIn = "1.5"), + "overflow expires_in" to validResponse(expiresIn = "9223372036854775808"), + ) + + invalidBodies.forEach { (description, body) -> + listOf(false, true).forEach { async -> + val response = TestResponse(200, body) + + assertThat(exchangeFailure(response, async)) + .describedAs("$description (async=$async)") + .isInstanceOf(OpenAIInvalidDataException::class.java) + .hasMessageNotContaining("secret") + .hasMessageNotContaining(body) + assertThat(response.closed).describedAs(description).isTrue() + } + } + } + + @Test + fun acceptsLargeForwardCompatibleResponsesWithoutAnArbitraryLimit() { + val response = + TestResponse( + 200, + """ + { + "forward_compatible_field": { + "nested": [{"value": "${"x".repeat(2 * 1024 * 1024)}"}] + }, + "access_token": "$ACCESS_TOKEN", + "issued_token_type": "$ACCESS_TOKEN_TYPE", + "token_type": "Bearer", + "expires_in": 3600 + } + """ + .trimIndent(), + ) + + val token = + X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, SingleResponseClient(response)).use { + it.execute() + } + + assertThat(token.value).isEqualTo(ACCESS_TOKEN) + assertThat(response.closed).isTrue() + } + + @Test + fun acceptsAccessTokensAboveMetadataDiagnosticLimits() { + val largeAccessToken = "a".repeat(2 * 1024) + + listOf(false, true).forEach { async -> + val response = TestResponse(200, validResponse(accessToken = largeAccessToken)) + + val token = + X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, SingleResponseClient(response)).use { + exchange -> + if (async) exchange.executeAsync().get(5, TimeUnit.SECONDS) + else exchange.execute() + } + + assertThat(token.value).isEqualTo(largeAccessToken) + assertThat(response.closed).isTrue() + } + } + + @Test + fun rejectsOversizedRecognizedSuccessStringsBeforeReadingTheCompleteValue() { + mapOf( + "token_type" to validResponse(tokenType = OVERSIZED_RECOGNIZED_VALUE), + "issued_token_type" to validResponse(issuedTokenType = OVERSIZED_RECOGNIZED_VALUE), + "token\\u005ftype" to + validResponse(tokenType = OVERSIZED_RECOGNIZED_VALUE) + .replaceFirst("token_type", "token\\u005ftype"), + ) + .forEach { (field, body) -> + listOf(false, true).forEach { async -> + val response = TestResponse(200, body) + + assertThat(exchangeFailure(response, async)) + .describedAs("$field (async=$async)") + .isInstanceOf(OpenAIInvalidDataException::class.java) + .hasMessage("Invalid X.509 token exchange response") + assertThat(response.bytesRead.get()).isLessThan(body.toByteArray().size) + assertThat(response.closed).isTrue() + } + } + } + + @Test + fun rejectsNestedRecognizedSuccessFieldsBeforeReadingLargeNestedStrings() { + listOf("access_token", "token_type", "issued_token_type", "expires_in").forEach { field -> + val body = + validResponseWithRawValue(field, """{"$field":"$OVERSIZED_RECOGNIZED_VALUE"}""") + listOf(false, true).forEach { async -> + val response = TestResponse(200, body) + + assertThat(exchangeFailure(response, async)) + .describedAs("$field (async=$async)") + .isInstanceOf(OpenAIInvalidDataException::class.java) + .hasMessage("Invalid X.509 token exchange response field: $field") + assertThat(response.bytesRead.get()).isLessThan(body.toByteArray().size) + assertThat(response.closed).isTrue() + } + } + } + + @Test + fun acceptsPositiveTokenLifetimesAboveOneHour() { + listOf(false, true).forEach { async -> + val response = TestResponse(200, validResponse(expiresIn = "86400")) + + val token = + X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, SingleResponseClient(response)).use { + exchange -> + if (async) exchange.executeAsync().get(5, TimeUnit.SECONDS) + else exchange.execute() + } + + assertThat(token.expiresIn).isEqualTo(Duration.ofDays(1)) + assertThat(response.closed).isTrue() + } + } + + @Test + fun failureStatusPreservesSafeOAuthCodeAndHeadersWithoutFreeFormDescription() { + listOf(false, true).forEach { async -> + val response = + TestResponse( + 503, + """ + { + "error": "invalid_grant", + "error_description": "Certificate is not authorized; Authorization: Bearer secret-body-token; Cookie: session=secret-body-cookie; access_token=secret-body-api-key" + } + """ + .trimIndent(), + Headers.builder() + .put("Set-Cookie", "session=secret-cookie") + .put("Authorization", "Bearer secret-token") + .put("X-Api-Key", "secret-api-key") + .put("X-Request-ID", "req_safe") + .put("Retry-After", "1") + .put("Traceparent", "aaaaaaaa.bbbbbbbb.cccccccc") + .put("Tracestate", "vendor=sk-test-not-a-real-credential-00000000") + .build(), + ) + + val failure = exchangeFailure(response, async) + + assertThat(failure).isInstanceOf(UnexpectedStatusCodeException::class.java) + val statusError = failure as UnexpectedStatusCodeException + assertThat(statusError.statusCode()).isEqualTo(503) + assertThat(statusError.headers().values("Set-Cookie")).isEmpty() + assertThat(statusError.headers().values("Authorization")).isEmpty() + assertThat(statusError.headers().values("X-Api-Key")).isEmpty() + assertThat(statusError.headers().values("X-Request-ID")).containsExactly("") + assertThat(statusError.headers().values("Retry-After")).containsExactly("1") + assertThat(statusError.headers().values("Traceparent")).containsExactly("") + assertThat(statusError.headers().values("Tracestate")).containsExactly("") + assertThat(statusError.code()).contains("invalid_grant") + assertThat(statusError.message).contains("invalid_grant") + assertThat(statusError.toString()) + .doesNotContain( + "Certificate is not authorized", + "secret-cookie", + "secret-token", + "secret-api-key", + "secret-body-token", + "secret-body-cookie", + "secret-body-api-key", + "aaaaaaaa.bbbbbbbb.cccccccc", + "sk-test-not-a-real-credential-00000000", + ) + assertThat(statusError.headers().toString()) + .doesNotContain( + "aaaaaaaa.bbbbbbbb.cccccccc", + "sk-test-not-a-real-credential-00000000", + ) + assertThat(statusError.body().toString()) + .doesNotContain("secret-body-token", "secret-body-cookie", "secret-body-api-key") + assertThat(response.bodyRead).isTrue() + assertThat(response.closed).isTrue() + } + } + + @Test + fun doesNotExposeShortOpaqueSecretsFromDiagnosticHeaders() { + listOf(false, true).forEach { async -> + val response = + TestResponse( + 400, + """{"error":"invalid_request"}""", + Headers.builder() + .put("X-Request-ID", "short-request-secret") + .put("Tracestate", "short-trace-secret") + .put("Retry-After", "short-retry-secret") + .put("Content-Type", "short-content-secret") + .build(), + ) + + val failure = exchangeFailure(response, async) + + assertThat(failure).isInstanceOf(UnexpectedStatusCodeException::class.java) + val statusError = failure as UnexpectedStatusCodeException + assertThat(statusError.headers().values("X-Request-ID")).containsExactly("") + assertThat(statusError.headers().values("Tracestate")).containsExactly("") + assertThat(statusError.headers().values("Retry-After")).isEmpty() + assertThat(statusError.headers().values("Content-Type")).isEmpty() + assertThat(statusError.headers().toString()) + .doesNotContain( + "short-request-secret", + "short-trace-secret", + "short-retry-secret", + "short-content-secret", + ) + assertThat(failure.toString()) + .doesNotContain( + "short-request-secret", + "short-trace-secret", + "short-retry-secret", + "short-content-secret", + ) + assertThat(response.closed).isTrue() + } + } + + @Test + fun rejectsNon200SuccessWithoutRetainingCredentialFields() { + listOf(false, true).forEach { async -> + val response = + TestResponse(201, validResponse(accessToken = "secret-token-must-not-leak")) + + val failure = exchangeFailure(response, async) + + assertThat(failure).isInstanceOf(UnexpectedStatusCodeException::class.java) + assertThat((failure as UnexpectedStatusCodeException).statusCode()).isEqualTo(201) + assertThat(failure.toString()).doesNotContain("secret-token-must-not-leak") + assertThat(response.bodyRead).isTrue() + assertThat(response.closed).isTrue() + } + } + + @Test + fun doesNotExposeFreeFormOAuthErrorDescriptions() { + mapOf( + "short-unlabeled-secret" to "short-unlabeled-secret", + "Authorization: Digest username=alice, response=secret-digest-value" to + "secret-digest-value", + "Cookie: first=secret-first-cookie; second=secret-second-cookie" to + "secret-second-cookie", + "client_secret=\"part-one,part-two\"" to "part-two", + "password=\"part-one;part-two\"" to "part-two", + "api_key=\"part-one part-two\"" to "part-two", + "Bearer \"part-one,part-two\"" to "part-two", + ) + .forEach { (diagnostic, secret) -> + listOf(false, true).forEach { async -> + val response = + TestResponse( + 400, + ObjectMapper() + .writeValueAsString( + mapOf( + "error" to "invalid_grant", + "error_description" to "Safe prefix; $diagnostic", + ) + ), + ) + + val failure = exchangeFailure(response, async) + + assertThat(failure).isInstanceOf(UnexpectedStatusCodeException::class.java) + val statusError = failure as UnexpectedStatusCodeException + assertThat(statusError.code()).contains("invalid_grant") + assertThat(failure.message).doesNotContain("Safe prefix", diagnostic, secret) + assertThat(failure.toString()).doesNotContain(secret) + assertThat(statusError.body().toString()).doesNotContain(diagnostic, secret) + assertThat(response.closed).isTrue() + } + } + } + + @Test + fun rejectsCredentialLikeOAuthErrorCodesFromDiagnostics() { + listOf( + "short-unlabeled-secret", + "sk-test-not-a-real-credential-00000000", + "aaaaaaaa.bbbbbbbb.cccccccc", + ) + .forEach { credential -> + listOf(false, true).forEach { async -> + val response = + TestResponse( + 400, + """{"error":"$credential","error_description":"Issuer rejected certificate"}""", + ) + + val failure = exchangeFailure(response, async) + + assertThat(failure).isInstanceOf(UnexpectedStatusCodeException::class.java) + val statusError = failure as UnexpectedStatusCodeException + assertThat(statusError.code()).isEmpty() + assertThat(statusError.message).doesNotContain("Issuer rejected certificate") + assertThat(statusError.message).doesNotContain(credential) + assertThat(statusError.body().toString()).doesNotContain(credential) + assertThat(statusError.toString()).doesNotContain(credential) + assertThat(response.closed).isTrue() + } + } + } + + @Test + fun ignoresOversizedOAuthDiagnosticsBeforeReadingTheCompleteValue() { + mapOf( + "error" to """{"error":"$OVERSIZED_RECOGNIZED_VALUE","error_description":"safe"}""", + "error_description" to + """{"error":"invalid_grant","error_description":"$OVERSIZED_RECOGNIZED_VALUE"}""", + "error\\u005fdescription" to + "{\"error\":\"invalid_grant\",\"error\\u005fdescription\":" + + "\"$OVERSIZED_RECOGNIZED_VALUE\"}", + ) + .forEach { (field, body) -> + listOf(false, true).forEach { async -> + val response = TestResponse(503, body) + + val failure = exchangeFailure(response, async) + + assertThat(failure) + .describedAs("$field (async=$async)") + .isInstanceOf(UnexpectedStatusCodeException::class.java) + val statusError = failure as UnexpectedStatusCodeException + assertThat(statusError.statusCode()).isEqualTo(503) + assertThat(statusError.code()).isEmpty() + assertThat(response.bytesRead.get()).isLessThan(body.toByteArray().size) + assertThat(response.closed).isTrue() + } + } + } + + @Test + fun ignoresNestedOAuthDiagnosticsBeforeReadingLargeNestedStrings() { + mapOf( + "error" to + """{"error":{"error":"$OVERSIZED_RECOGNIZED_VALUE"},"error_description":"safe"}""", + "error_description" to + """{"error":"invalid_grant","error_description":{"error_description":"$OVERSIZED_RECOGNIZED_VALUE"}}""", + ) + .forEach { (field, body) -> + listOf(false, true).forEach { async -> + val response = TestResponse(503, body) + + val failure = exchangeFailure(response, async) + + assertThat(failure) + .describedAs("$field (async=$async)") + .isInstanceOf(UnexpectedStatusCodeException::class.java) + val statusError = failure as UnexpectedStatusCodeException + assertThat(statusError.statusCode()).isEqualTo(503) + assertThat(statusError.code()).isEmpty() + assertThat(response.bytesRead.get()).isLessThan(body.toByteArray().size) + assertThat(response.closed).isTrue() + } + } + } + + @Test + fun preservesIssuerBodyTimeoutAndIoFailuresAsSanitizedRetryableIo() { + listOf( + SocketTimeoutException("issuer body stalled"), + IOException("issuer body disconnected"), + ) + .forEach { cause -> + listOf(200, 503).forEach { statusCode -> + listOf(false, true).forEach { async -> + val response = FailingBodyResponse(cause, statusCode) + + val failure = exchangeFailure(response, async) + + assertThat(failure).isInstanceOf(OpenAIIoException::class.java) + assertThat(failure) + .hasMessage("Failed to read X.509 token exchange response") + assertThat(failure.cause).isSameAs(cause) + assertThat(response.closed).isTrue() + } + } + } + } + + @Test + fun completedAsyncResponseParsesOnOwnedIoThreadAfterReturningCancellationHandle() { + val response = BlockingResponse() + val exchange = X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, SingleResponseClient(response)) + + val result = exchange.executeAsync() + assertThat(response.bodyStarted.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(response.bodyThreadName).startsWith("openai-x509-response-") + assertThat(result.cancel(true)).isTrue() + + assertThat(response.closed.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(response.closeCount).hasValue(1) + exchange.close() + } + + @Test + fun cancelingBeforeResponseDeliveryCancelsTheUnderlyingCall() { + val responseFuture = CompletableFuture() + val client = DeferredResponseClient(responseFuture) + val exchange = X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, client) + + val result = exchange.executeAsync() + assertThat(client.started.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(result.cancel(true)).isTrue() + + assertThat(responseFuture.isCancelled).isTrue() + assertThat(client.closed).isFalse() + exchange.close() + assertThat(client.closed).isFalse() + } + + @Test + fun cancellationWaitsForInFlightRequestEnrollmentBeforePublishing() { + val responseFuture = CompletableFuture() + val client = BlockingStartClient(responseFuture) + val exchange = X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, client) + val result = exchange.executeAsync() + val cancellationPublished = CountDownLatch(1) + val cancellationFinished = CountDownLatch(1) + result.whenComplete { _, _ -> cancellationPublished.countDown() } + + assertThat(client.started.await(5, TimeUnit.SECONDS)).isTrue() + Thread { + result.cancel(true) + cancellationFinished.countDown() + } + .start() + + try { + assertThat(cancellationFinished.await(200, TimeUnit.MILLISECONDS)).isFalse() + assertThat(cancellationPublished.count).isEqualTo(1) + + client.release.countDown() + + assertThat(cancellationFinished.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(cancellationPublished.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(result.isCancelled).isTrue() + assertThat(responseFuture.isCancelled).isTrue() + } finally { + client.release.countDown() + exchange.close() + } + } + + @Test + fun cancellationWaitsForCompletedResponseHandoffBeforePublishing() { + val response = BlockingResponse() + val responseFuture = BlockingHandoffFuture(response) + val client = DeferredResponseClient(responseFuture) + val exchange = X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, client) + val result = exchange.executeAsync() + val cancellationPublished = CountDownLatch(1) + val cancellationFinished = CountDownLatch(1) + result.whenComplete { _, _ -> cancellationPublished.countDown() } + + assertThat(responseFuture.handoffStarted.await(5, TimeUnit.SECONDS)).isTrue() + Thread { + result.cancel(true) + cancellationFinished.countDown() + } + .start() + + try { + assertThat(cancellationFinished.await(200, TimeUnit.MILLISECONDS)).isFalse() + assertThat(cancellationPublished.count).isEqualTo(1) + + responseFuture.releaseHandoff.countDown() + + assertThat(cancellationFinished.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(cancellationPublished.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(response.closed.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(response.bodyStarted.count).isEqualTo(1) + assertThat(response.closeCount).hasValue(1) + } finally { + responseFuture.releaseHandoff.countDown() + exchange.close() + } + } + + @Test + fun cancelingClosesALateResponseWhenUnderlyingCancellationLosesTheRace() { + val responseFuture = NonCancellableFuture() + val client = DeferredResponseClient(responseFuture) + val exchange = X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, client) + val result = exchange.executeAsync() + val response = BlockingResponse() + + assertThat(client.started.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(result.cancel(true)).isTrue() + assertThat(responseFuture.complete(response)).isTrue() + + assertThat(result.isCancelled).isTrue() + assertThat(response.bodyStarted.count).isEqualTo(1) + assertThat(response.closed.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(response.closeCount).hasValue(1) + exchange.close() + } + + @Test + fun cancellationClosesActiveResponseBeforeRunningCallerCallbacks() { + val response = BlockingResponse() + val exchange = X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, SingleResponseClient(response)) + val result = exchange.executeAsync() + val callbackStarted = CountDownLatch(1) + val releaseCallback = CountDownLatch(1) + val cancellationFinished = CountDownLatch(1) + + assertThat(response.bodyStarted.await(5, TimeUnit.SECONDS)).isTrue() + result.whenComplete { _, _ -> + callbackStarted.countDown() + releaseCallback.await() + } + Thread { + result.cancel(true) + cancellationFinished.countDown() + } + .start() + + try { + assertThat(callbackStarted.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(response.closed.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(response.closeCount).hasValue(1) + } finally { + releaseCallback.countDown() + assertThat(cancellationFinished.await(5, TimeUnit.SECONDS)).isTrue() + exchange.close() + } + } + + @Test + fun closeClosesALateResponseWhenUnderlyingCancellationLosesTheRace() { + val responseFuture = NonCancellableFuture() + val client = DeferredResponseClient(responseFuture) + val exchange = X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, client) + val result = exchange.executeAsync() + val response = BlockingResponse() + + assertThat(client.started.await(5, TimeUnit.SECONDS)).isTrue() + exchange.close() + assertThat(responseFuture.complete(response)).isTrue() + + assertThat(result.isCancelled).isTrue() + assertThat(response.bodyStarted.count).isEqualTo(1) + assertThat(response.closed.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(response.closeCount).hasValue(1) + } + + @Test + fun closeCleansAllAdmittedOperationsBeforeRunningCallerCallbacks() { + val responses = List(4) { BlockingResponse() } + val client = SequenceResponseClient(responses) + val exchange = X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, client) + val activeResults = responses.map { exchange.executeAsync() } + val queuedResult = exchange.executeAsync() + val results = activeResults + queuedResult + val callbackStarted = CountDownLatch(1) + val releaseCallbacks = CountDownLatch(1) + val closeFinished = CountDownLatch(1) + + responses.forEach { response -> + assertThat(response.bodyStarted.await(5, TimeUnit.SECONDS)).isTrue() + } + assertThat(client.executeCount).hasValue(4) + results.forEach { result -> + result.whenComplete { _, _ -> + callbackStarted.countDown() + releaseCallbacks.await() + } + } + Thread { + exchange.close() + closeFinished.countDown() + } + .start() + + try { + assertThat(callbackStarted.await(5, TimeUnit.SECONDS)).isTrue() + responses.forEach { response -> + assertThat(response.closed.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(response.closeCount).hasValue(1) + } + assertThat(client.executeCount).hasValue(4) + } finally { + releaseCallbacks.countDown() + assertThat(closeFinished.await(5, TimeUnit.SECONDS)).isTrue() + } + } + + @Test + fun asyncExecutorQueuesWholeExchangesBeforeIssuingAdditionalRequests() { + val release = CountDownLatch(1) + val activeResponses = List(4) { GatedResponse(validResponse(), release) } + val queuedResponse = TestResponse(200, validResponse()) + val responses = activeResponses + queuedResponse + val client = SequenceResponseClient(responses) + val exchange = X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, client) + try { + val results = responses.map { exchange.executeAsync() } + + activeResponses.forEach { response -> + assertThat(response.bodyStarted.await(5, TimeUnit.SECONDS)).isTrue() + } + assertThat(client.executeCount).hasValue(4) + assertThat(queuedResponse.bodyRead).isFalse() + assertThat(results.last().isDone).isFalse() + + release.countDown() + + assertThat(results.map { it.get(5, TimeUnit.SECONDS).value }) + .containsExactlyElementsOf(List(5) { ACCESS_TOKEN }) + assertThat(activeResponses).allMatch { it.closed } + assertThat(queuedResponse.closed).isTrue() + } finally { + release.countDown() + exchange.close() + } + assertThat(client.closed).isFalse() + } + + @Test + fun asyncExecutorBoundsAndRemovesCanceledQueuedExchanges() { + val release = CountDownLatch(1) + val activeResponses = List(4) { GatedResponse(validResponse(), release) } + val admittedResponses = List(4) { TestResponse(200, validResponse()) } + val client = SequenceResponseClient(activeResponses + admittedResponses) + val exchange = X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, client) + try { + val activeResults = List(4) { exchange.executeAsync() } + val queuedResults = List(4) { exchange.executeAsync() } + + activeResponses.forEach { response -> + assertThat(response.bodyStarted.await(5, TimeUnit.SECONDS)).isTrue() + } + assertThat(client.executeCount).hasValue(4) + + val overflow = exchange.executeAsync() + val overflowFailure = + runCatching { overflow.get(5, TimeUnit.SECONDS) }.exceptionOrNull() + ?: error("Expected bounded X.509 executor rejection") + assertThat(overflowFailure).isInstanceOf(ExecutionException::class.java) + assertThat(overflowFailure.cause).isInstanceOf(OpenAIIoException::class.java) + assertThat(overflowFailure.cause) + .hasMessage("X.509 token exchange processing unavailable") + + assertThat(queuedResults.first().cancel(true)).isTrue() + val replacement = exchange.executeAsync() + assertThat(replacement.isDone).isFalse() + assertThat(client.executeCount).hasValue(4) + + release.countDown() + + val successful = activeResults + queuedResults.drop(1) + replacement + assertThat(successful.map { it.get(5, TimeUnit.SECONDS).value }) + .containsExactlyElementsOf(List(8) { ACCESS_TOKEN }) + assertThat(queuedResults.first().isCancelled).isTrue() + assertThat(client.executeCount).hasValue(8) + assertThat(activeResponses).allMatch { it.closed } + assertThat(admittedResponses).allMatch { it.closed } + } finally { + release.countDown() + exchange.close() + } + } + + private fun exchangeFailure(response: HttpResponse, async: Boolean): Throwable { + val result = + runCatching { + X509TokenExchange(IDP_ID, SERVICE_ACCOUNT_ID, SingleResponseClient(response)) + .use { exchange -> + if (async) exchange.executeAsync().get(5, TimeUnit.SECONDS) + else exchange.execute() + } + } + .exceptionOrNull() ?: error("Expected X.509 exchange to fail") + return if (result is ExecutionException) result.cause ?: result else result + } + + private fun transport( + clientIdentity: X509TestIdentity, + trustedServerRoots: Iterable, + ): X509Transport { + val trustManager = + HandshakeCertificates.Builder() + .apply { + trustedServerRoots.forEach { certificate -> addTrustedCertificate(certificate) } + } + .build() + .trustManager + return X509Transport.builder() + .keyManager(x509TestKeyManager(mapOf(CERTIFICATE_ALIAS to clientIdentity))) + .certificateAlias(CERTIFICATE_ALIAS) + .trustManager(trustManager) + .build() + } + + private fun validResponseWithout(field: String): String = + (ObjectMapper().readTree(validResponse()) as ObjectNode).apply { remove(field) }.toString() + + private fun validResponseWithNull(field: String): String = + (ObjectMapper().readTree(validResponse()) as ObjectNode).apply { putNull(field) }.toString() + + private fun validResponseWithDuplicate(field: String, firstValue: String): String = + validResponse().replaceFirst("\"$field\":", "\"$field\": $firstValue,\n \"$field\":") + + private fun validResponseWithRawValue(field: String, value: String): String = + validResponse() + .replaceFirst(Regex("\"$field\"\\s*:\\s*(?:\"[^\"]*\"|[0-9]+)"), "\"$field\": $value") + + private fun validResponse( + accessToken: String = ACCESS_TOKEN, + tokenType: String = "Bearer", + issuedTokenType: String = ACCESS_TOKEN_TYPE, + expiresIn: String = "3600", + ): String = + """ + { + "access_token": "$accessToken", + "issued_token_type": "$issuedTokenType", + "token_type": "$tokenType", + "expires_in": $expiresIn + } + """ + .trimIndent() + + private companion object { + const val AUTH_HOST = "mtls.auth.openai.com" + const val CERTIFICATE_ALIAS = "x509" + const val IDP_ID = "idp_test" + const val SERVICE_ACCOUNT_ID = "svc_acct_test" + const val ACCESS_TOKEN = "test-x509-access-token" + const val ACCESS_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token" + val OVERSIZED_RECOGNIZED_VALUE = "x".repeat(256 * 1024) + 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": "$IDP_ID", + "service_account_id": "$SERVICE_ACCOUNT_ID" + } + """ + .trimIndent() + val TOKEN_RESPONSE = + """ + { + "access_token": "$ACCESS_TOKEN", + "issued_token_type": "$ACCESS_TOKEN_TYPE", + "token_type": "Bearer", + "expires_in": 3600, + "forward_compatible_field": true + } + """ + .trimIndent() + } +} + +private open class SingleResponseClient(private val response: HttpResponse) : HttpClient { + var closed = false + private set + + override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse = + response + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = CompletableFuture.completedFuture(response) + + override fun close() { + closed = true + } +} + +private class DeferredResponseClient(private val responseFuture: CompletableFuture) : + HttpClient { + val started = CountDownLatch(1) + var closed = false + private set + + override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse = + error("Unexpected synchronous call") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = responseFuture.also { started.countDown() } + + override fun close() { + closed = true + } +} + +private class BlockingStartClient(private val responseFuture: CompletableFuture) : + HttpClient { + val started = CountDownLatch(1) + val release = CountDownLatch(1) + + override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse = + error("Unexpected synchronous call") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture { + started.countDown() + release.await() + return responseFuture + } + + override fun close() {} +} + +private class NonCancellableFuture : CompletableFuture() { + override fun cancel(mayInterruptIfRunning: Boolean): Boolean = false +} + +private class BlockingHandoffFuture(value: T) : CompletableFuture() { + val handoffStarted = CountDownLatch(1) + val releaseHandoff = CountDownLatch(1) + + init { + complete(value) + } + + override fun whenComplete(action: BiConsumer): CompletableFuture { + handoffStarted.countDown() + releaseHandoff.await() + return super.whenComplete(action) + } +} + +private class SequenceResponseClient(responses: List) : HttpClient { + private val responses = ConcurrentLinkedQueue(responses) + val executeCount = AtomicInteger() + var closed = false + private set + + override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse = + error("Unexpected synchronous call") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture { + executeCount.incrementAndGet() + return CompletableFuture.completedFuture( + responses.poll() ?: error("No X.509 test response remains") + ) + } + + override fun close() { + closed = true + } +} + +private class BlockingResponse : HttpResponse { + val bodyStarted = CountDownLatch(1) + val closed = CountDownLatch(1) + val closeCount = AtomicInteger() + + @Volatile var bodyThreadName: String? = null + + override fun statusCode(): Int = 200 + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream = + object : InputStream() { + override fun read(): Int { + bodyThreadName = Thread.currentThread().name + bodyStarted.countDown() + closed.await() + throw IOException("response closed") + } + } + + override fun close() { + closeCount.incrementAndGet() + closed.countDown() + } +} + +private class GatedResponse(body: String, private val release: CountDownLatch) : HttpResponse { + private val bytes = body.toByteArray() + val bodyStarted = CountDownLatch(1) + var closed = false + private set + + override fun statusCode(): Int = 200 + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream = + object : ByteArrayInputStream(bytes) { + override fun read(): Int { + awaitRelease() + return super.read() + } + + override fun read(buffer: ByteArray, offset: Int, length: Int): Int { + awaitRelease() + return super.read(buffer, offset, length) + } + + private fun awaitRelease() { + bodyStarted.countDown() + try { + release.await() + } catch (error: InterruptedException) { + Thread.currentThread().interrupt() + throw IOException("Interrupted while waiting to read test response", error) + } + } + } + + override fun close() { + closed = true + } +} + +private class TestResponse( + private val statusCode: Int, + body: String, + private val responseHeaders: Headers = Headers.builder().build(), +) : HttpResponse { + private val bytes = body.toByteArray() + val bytesRead = AtomicInteger() + val closeCount = AtomicInteger() + var bodyRead = false + private set + + var closed = false + private set + + override fun statusCode(): Int = statusCode + + override fun headers(): Headers = responseHeaders + + override fun body(): InputStream { + bodyRead = true + return object : ByteArrayInputStream(bytes) { + override fun read(): Int = + super.read().also { if (it >= 0) bytesRead.incrementAndGet() } + + override fun read(buffer: ByteArray, offset: Int, length: Int): Int = + super.read(buffer, offset, length).also { if (it > 0) bytesRead.addAndGet(it) } + } + } + + override fun close() { + closeCount.incrementAndGet() + closed = true + } +} + +private class FailingBodyResponse(private val failure: IOException, private val statusCode: Int) : + HttpResponse { + var closed = false + private set + + override fun statusCode(): Int = statusCode + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream = + object : InputStream() { + override fun read(): Int = throw failure + } + + override fun close() { + closed = true + } +} diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509WireTestInfrastructure.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509WireTestInfrastructure.kt index 14749d382..fa6932d00 100644 --- a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509WireTestInfrastructure.kt +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509WireTestInfrastructure.kt @@ -2,14 +2,17 @@ package com.openai.client.okhttp import java.net.Proxy import java.net.Socket +import java.security.KeyStore import java.security.SecureRandom import java.security.cert.X509Certificate import java.util.concurrent.CopyOnWriteArrayList import javax.net.ssl.ExtendedSSLSession +import javax.net.ssl.KeyManagerFactory import javax.net.ssl.SNIHostName import javax.net.ssl.SSLContext import javax.net.ssl.SSLEngine import javax.net.ssl.SSLSocket +import javax.net.ssl.X509ExtendedKeyManager import javax.net.ssl.X509ExtendedTrustManager import javax.net.ssl.X509TrustManager import okhttp3.OkHttpClient @@ -110,6 +113,27 @@ internal inline fun OkHttpClient.useTestClient(block: (OkHttpClient) -> T): cache?.close() } +internal fun x509TestKeyManager(identities: Map): X509ExtendedKeyManager { + val password = "test password".toCharArray() + val keyStore = + KeyStore.getInstance("PKCS12").apply { + load(null, null) + identities.forEach { (alias, identity) -> + setKeyEntry( + alias, + identity.leaf.keyPair.private, + password, + arrayOf(identity.leaf.certificate, identity.root.certificate), + ) + } + } + val keyManagerFactory = + KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm()).apply { + init(keyStore, password) + } + return keyManagerFactory.keyManagers.filterIsInstance().single() +} + private class RecordingClientTrustManager(private val delegate: X509TrustManager) : X509ExtendedTrustManager() { val requestedServerNames = CopyOnWriteArrayList()