From e30eaf5b6fc734c0d9ccc77399432f3b24038dde Mon Sep 17 00:00:00 2001
From: Istvan Toth <stoty@cloudera.com>
Date: Mon, 4 Nov 2024 14:42:55 +0100
Subject: [PATCH] [CALCITE-6671] Update httpclient5 to 5.4.1 and httpcore 5.3.1
 in Avatica

---
 .../remote/AvaticaCommonsHttpClientImpl.java  | 38 ++++++++++++++++---
 gradle.properties                             |  4 +-
 2 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaCommonsHttpClientImpl.java b/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaCommonsHttpClientImpl.java
index dcbfe123f..44beaa7b8 100644
--- a/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaCommonsHttpClientImpl.java
+++ b/core/src/main/java/org/apache/calcite/avatica/remote/AvaticaCommonsHttpClientImpl.java
@@ -60,6 +60,8 @@
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.security.Principal;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 
@@ -92,6 +94,8 @@ public class AvaticaCommonsHttpClientImpl implements AvaticaHttpClient, HttpClie
   protected Lookup<AuthSchemeFactory> authRegistry = null;
   protected Object userToken;
   protected HttpClientContext context;
+  protected long connectTimeout;
+  protected long responseTimeout;
 
   public AvaticaCommonsHttpClientImpl(URL url) {
     this.uri = toURI(Objects.requireNonNull(url));
@@ -100,13 +104,11 @@ public AvaticaCommonsHttpClientImpl(URL url) {
   protected void initializeClient(PoolingHttpClientConnectionManager pool,
                                   ConnectionConfig config) {
     this.authCache = new BasicAuthCache();
+    this.connectTimeout = config.getHttpConnectionTimeout();
+    this.responseTimeout = config.getHttpResponseTimeout();
     // A single thread-safe HttpClient, pooling connections via the
     // ConnectionManager
-    RequestConfig.Builder requestConfigBuilder = RequestConfig.custom();
-    RequestConfig requestConfig = requestConfigBuilder
-        .setConnectTimeout(config.getHttpConnectionTimeout(), TimeUnit.MILLISECONDS)
-        .setResponseTimeout(config.getHttpResponseTimeout(), TimeUnit.MILLISECONDS)
-        .build();
+    RequestConfig requestConfig = createRequestConfig();
     HttpClientBuilder httpClientBuilder = HttpClients.custom().setConnectionManager(pool)
         .setDefaultRequestConfig(requestConfig);
     this.client = httpClientBuilder.build();
@@ -124,6 +126,30 @@ protected void initializeClient(PoolingHttpClientConnectionManager pool,
 
   }
 
+  // This is needed because we initialize the client object too early.
+  private RequestConfig createRequestConfig() {
+    RequestConfig.Builder requestConfigBuilder = RequestConfig.custom();
+    requestConfigBuilder
+        .setConnectTimeout(this.connectTimeout, TimeUnit.MILLISECONDS)
+        .setResponseTimeout(this.responseTimeout, TimeUnit.MILLISECONDS);
+    List<String> preferredSchemes = new ArrayList<>();
+    // In HttpClient 5.3+ SPNEGO is not enabled by default
+    if (authRegistry != null) {
+      if (authRegistry.lookup(StandardAuthScheme.DIGEST) != null) {
+        preferredSchemes.add(StandardAuthScheme.DIGEST);
+      }
+      if (authRegistry.lookup(StandardAuthScheme.BASIC) != null) {
+        preferredSchemes.add(StandardAuthScheme.BASIC);
+      }
+      if (authRegistry.lookup(StandardAuthScheme.SPNEGO) != null) {
+        preferredSchemes.add(StandardAuthScheme.SPNEGO);
+      }
+      requestConfigBuilder.setTargetPreferredAuthSchemes(preferredSchemes);
+      requestConfigBuilder.setProxyPreferredAuthSchemes(preferredSchemes);
+    }
+    return requestConfigBuilder.build();
+  }
+
   @Override public byte[] send(byte[] request) {
     while (true) {
       ByteArrayEntity entity = new ByteArrayEntity(request, ContentType.APPLICATION_OCTET_STREAM);
@@ -186,6 +212,7 @@ CloseableHttpResponse execute(HttpPost post, HttpClientContext context)
     this.authRegistry = authRegistryBuilder.build();
     context.setCredentialsProvider(credentialsProvider);
     context.setAuthSchemeRegistry(authRegistry);
+    context.setRequestConfig(createRequestConfig());
   }
 
   @Override public void setGSSCredential(GSSCredential credential) {
@@ -209,6 +236,7 @@ CloseableHttpResponse execute(HttpPost post, HttpClientContext context)
     }
     context.setCredentialsProvider(credentialsProvider);
     context.setAuthSchemeRegistry(authRegistry);
+    context.setRequestConfig(createRequestConfig());
   }
 
   /**
diff --git a/gradle.properties b/gradle.properties
index 621dd523c..af50e0657 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -65,8 +65,8 @@ guava.version=32.1.1-jre
 hamcrest.version=1.3
 hsqldb.version=2.4.1
 h2.version=1.4.197
-httpclient5.version=5.2.1
-httpcore5.version=5.2.3
+httpclient5.version=5.4.1
+httpcore5.version=5.3.1
 jackson.version=2.15.2
 jcip-annotations.version=1.0-1
 jcommander.version=1.72