From 7cf469c5cc7313b214411e99d1822dd2dc0f7d2b Mon Sep 17 00:00:00 2001
From: Oleg Kalnichevski <olegk@apache.org>
Date: Mon, 15 Jan 2024 16:41:59 +0100
Subject: [PATCH] Connect executors to update the execution context with the
 protocol and TLS details upon successful connect

---
 .../http/impl/async/AsyncConnectExec.java     | 33 +++++++++++++------
 .../http/impl/classic/ConnectExec.java        |  7 ++++
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/httpclient5/src/main/java/org/apache/hc/client5/http/impl/async/AsyncConnectExec.java b/httpclient5/src/main/java/org/apache/hc/client5/http/impl/async/AsyncConnectExec.java
index 493ab6c9e..5fa79ceb3 100644
--- a/httpclient5/src/main/java/org/apache/hc/client5/http/impl/async/AsyncConnectExec.java
+++ b/httpclient5/src/main/java/org/apache/hc/client5/http/impl/async/AsyncConnectExec.java
@@ -34,6 +34,7 @@ import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.hc.client5.http.AuthenticationStrategy;
+import org.apache.hc.client5.http.EndpointInfo;
 import org.apache.hc.client5.http.HttpRoute;
 import org.apache.hc.client5.http.RouteTracker;
 import org.apache.hc.client5.http.SchemePortResolver;
@@ -176,11 +177,7 @@ public final class AsyncConnectExec implements AsyncExecChainHandler {
                     }));
         } else {
             if (execRuntime.isEndpointConnected()) {
-                try {
-                    chain.proceed(request, entityProducer, scope, asyncExecCallback);
-                } catch (final HttpException | IOException ex) {
-                    asyncExecCallback.failed(ex);
-                }
+                proceedConnected(request, entityProducer, scope, chain, asyncExecCallback);
             } else {
                 proceedToNextHop(state, request, entityProducer, scope, chain, asyncExecCallback);
             }
@@ -363,11 +360,7 @@ public final class AsyncConnectExec implements AsyncExecChainHandler {
                 if (LOG.isDebugEnabled()) {
                     LOG.debug("{} route fully established", exchangeId);
                 }
-                try {
-                    chain.proceed(request, entityProducer, scope, asyncExecCallback);
-                } catch (final HttpException | IOException ex) {
-                    asyncExecCallback.failed(ex);
-                }
+                proceedConnected(request, entityProducer, scope, chain, asyncExecCallback);
                 break;
 
             default:
@@ -543,4 +536,24 @@ public final class AsyncConnectExec implements AsyncExecChainHandler {
         return false;
     }
 
+    private void proceedConnected(
+            final HttpRequest request,
+            final AsyncEntityProducer entityProducer,
+            final AsyncExecChain.Scope scope,
+            final AsyncExecChain chain,
+            final AsyncExecCallback asyncExecCallback) {
+        final AsyncExecRuntime execRuntime = scope.execRuntime;
+        final HttpClientContext clientContext = scope.clientContext;
+        final EndpointInfo endpointInfo = execRuntime.getEndpointInfo();
+        if (endpointInfo != null) {
+            clientContext.setProtocolVersion(endpointInfo.getProtocol());
+            clientContext.setAttribute(HttpCoreContext.SSL_SESSION, endpointInfo.getSslSession());
+        }
+        try {
+            chain.proceed(request, entityProducer, scope, asyncExecCallback);
+        } catch (final HttpException | IOException ex) {
+            asyncExecCallback.failed(ex);
+        }
+    }
+
 }
diff --git a/httpclient5/src/main/java/org/apache/hc/client5/http/impl/classic/ConnectExec.java b/httpclient5/src/main/java/org/apache/hc/client5/http/impl/classic/ConnectExec.java
index 56fee6aad..54551666b 100644
--- a/httpclient5/src/main/java/org/apache/hc/client5/http/impl/classic/ConnectExec.java
+++ b/httpclient5/src/main/java/org/apache/hc/client5/http/impl/classic/ConnectExec.java
@@ -30,6 +30,7 @@ package org.apache.hc.client5.http.impl.classic;
 import java.io.IOException;
 
 import org.apache.hc.client5.http.AuthenticationStrategy;
+import org.apache.hc.client5.http.EndpointInfo;
 import org.apache.hc.client5.http.HttpRoute;
 import org.apache.hc.client5.http.RouteTracker;
 import org.apache.hc.client5.http.SchemePortResolver;
@@ -62,6 +63,7 @@ import org.apache.hc.core5.http.Method;
 import org.apache.hc.core5.http.io.entity.EntityUtils;
 import org.apache.hc.core5.http.message.BasicClassicHttpRequest;
 import org.apache.hc.core5.http.message.StatusLine;
+import org.apache.hc.core5.http.protocol.HttpCoreContext;
 import org.apache.hc.core5.http.protocol.HttpProcessor;
 import org.apache.hc.core5.util.Args;
 import org.slf4j.Logger;
@@ -185,6 +187,11 @@ public final class ConnectExec implements ExecChainHandler {
 
                 } while (step > HttpRouteDirector.COMPLETE);
             }
+            final EndpointInfo endpointInfo = execRuntime.getEndpointInfo();
+            if (endpointInfo != null) {
+                context.setProtocolVersion(endpointInfo.getProtocol());
+                context.setAttribute(HttpCoreContext.SSL_SESSION, endpointInfo.getSslSession());
+            }
             return chain.proceed(request, scope);
 
         } catch (final IOException | HttpException | RuntimeException ex) {