Skip to content

Commit 40ea4a8

Browse files
sfc-gh-rkowalskisfc-gh-rkowalski
committed
SNOW-2041976 Add HTTP header customizer
1 parent 9a99cdf commit 40ea4a8

16 files changed

+959
-17
lines changed

‎src/main/java/net/snowflake/client/core/AttributeEnhancingHttpRequestRetryHandler.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package net.snowflake.client.core;
2+
3+
import java.io.IOException;
4+
import org.apache.http.impl.client.DefaultHttpRequestRetryHandler;
5+
import org.apache.http.protocol.HttpContext;
6+
7+
/**
8+
* Extends {@link DefaultHttpRequestRetryHandler} to store the current execution count (attempt
9+
* number) in the {@link HttpContext}. This allows interceptors to identify retry attempts.
10+
*
11+
* <p>The execution count is stored using the key defined by {@link #EXECUTION_COUNT_ATTRIBUTE}.
12+
*/
13+
public class AttributeEnhancingHttpRequestRetryHandler extends DefaultHttpRequestRetryHandler {
14+
/**
15+
* The key used to store the current execution count (attempt number) in the {@link HttpContext}.
16+
* Interceptors can use this key to retrieve the count. The value stored will be an {@link
17+
* Integer}.
18+
*/
19+
public static final String EXECUTION_COUNT_ATTRIBUTE =
20+
"net.snowflake.client.core.execution-count";
21+
22+
@Override
23+
public boolean retryRequest(IOException exception, int executionCount, HttpContext context) {
24+
context.setAttribute(EXECUTION_COUNT_ATTRIBUTE, executionCount);
25+
return super.retryRequest(exception, executionCount, context);
26+
}
27+
}

‎src/main/java/net/snowflake/client/core/HeaderCustomizerHttpRequestInterceptor.java‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
package net.snowflake.client.core;
2+
3+
import com.amazonaws.Request;
4+
import com.amazonaws.handlers.RequestHandler2;
5+
import java.io.IOException;
6+
import java.util.ArrayList;
7+
import java.util.HashMap;
8+
import java.util.List;
9+
import java.util.Map;
10+
import java.util.Set;
11+
import java.util.stream.Collectors;
12+
import net.snowflake.client.jdbc.HttpHeadersCustomizer;
13+
import net.snowflake.client.log.SFLogger;
14+
import net.snowflake.client.log.SFLoggerFactory;
15+
import org.apache.http.Header;
16+
import org.apache.http.HttpException;
17+
import org.apache.http.HttpRequest;
18+
import org.apache.http.HttpRequestInterceptor;
19+
import org.apache.http.protocol.HttpContext;
20+
21+
/**
22+
* Implements Apache HttpClient's {@link HttpRequestInterceptor} to provide a mechanism for adding
23+
* custom HTTP headers to outgoing requests made by the Snowflake JDBC driver.
24+
*
25+
* <p>This class iterates through a list of user-provided {@link HttpHeadersCustomizer}
26+
* implementations. For each customizer, it checks if it applies to the current request. If it does,
27+
* it retrieves new headers from the customizer and adds them to the request, ensuring that existing
28+
* driver-set headers are not overridden.
29+
*
30+
* <p>For Apache HttpClient, retry detection is handled by checking the {@link
31+
* AttributeEnhancingHttpRequestRetryHandler#EXECUTION_COUNT_ATTRIBUTE} attribute in the {@link
32+
* HttpContext} set by {@link AttributeEnhancingHttpRequestRetryHandler} to honor the {@code
33+
* invokeOnce()} contract of the customizer.
34+
*
35+
* @see HttpHeadersCustomizer
36+
*/
37+
public class HeaderCustomizerHttpRequestInterceptor extends RequestHandler2
38+
implements HttpRequestInterceptor {
39+
private static final SFLogger logger =
40+
SFLoggerFactory.getLogger(HeaderCustomizerHttpRequestInterceptor.class);
41+
private final List<HttpHeadersCustomizer> headersCustomizers;
42+
43+
public HeaderCustomizerHttpRequestInterceptor(List<HttpHeadersCustomizer> headersCustomizers) {
44+
if (headersCustomizers != null) {
45+
this.headersCustomizers = new ArrayList<>(headersCustomizers); // Defensive copy
46+
} else {
47+
this.headersCustomizers = new ArrayList<>();
48+
}
49+
}
50+
51+
/**
52+
* Processes an Apache HttpClient {@link HttpRequest} before it is sent. It iterates through
53+
* registered {@link HttpHeadersCustomizer}s, checks applicability, retrieves new headers,
54+
* verifies against overriding driver headers, and adds them to the request. Handles the {@code
55+
* invokeOnce()} flag based on the "execution-count" attribute in the {@link HttpContext}.
56+
*
57+
* @param httpRequest The HTTP request to process.
58+
* @param httpContext The context for the HTTP request execution, used to retrieve retry count.
59+
* @throws DriverHeaderOverridingNotAllowedException If a customizer attempts to override a
60+
* driver-set header.
61+
*/
62+
@Override
63+
public void process(HttpRequest httpRequest, HttpContext httpContext)
64+
throws HttpException, IOException {
65+
if (this.headersCustomizers.isEmpty()) {
66+
return;
67+
}
68+
String httpMethod = httpRequest.getRequestLine().getMethod();
69+
String uri = httpRequest.getRequestLine().getUri();
70+
Map<String, List<String>> currentHeaders = extractHeaders(httpRequest);
71+
// convert header names to lower case for case in-sensitive lookup
72+
Set<String> protectedHeaders =
73+
currentHeaders.keySet().stream().map(String::toLowerCase).collect(Collectors.toSet());
74+
Object executionCount =
75+
httpContext.getAttribute(
76+
AttributeEnhancingHttpRequestRetryHandler.EXECUTION_COUNT_ATTRIBUTE);
77+
// If count is null or 0, it's the first attempt. Otherwise, it's a retry.
78+
boolean isRetry = (executionCount != null && (Integer) executionCount > 0);
79+
80+
for (HttpHeadersCustomizer customizer : this.headersCustomizers) {
81+
if (customizer.applies(httpMethod, uri, currentHeaders)) {
82+
if (customizer.invokeOnce() && isRetry) {
83+
logger.debug(
84+
"{} customizer should only run on the first attempt and this is a {} retry. Skipping.",
85+
customizer.getClass(),
86+
executionCount);
87+
continue;
88+
}
89+
Map<String, List<String>> newHeaders = customizer.newHeaders();
90+
91+
throwIfExistingHeadersAreModified(protectedHeaders, newHeaders.keySet(), customizer);
92+
93+
for (Map.Entry<String, List<String>> entry : newHeaders.entrySet()) {
94+
for (String value : entry.getValue()) {
95+
httpRequest.addHeader(entry.getKey(), value);
96+
}
97+
}
98+
}
99+
}
100+
}
101+
102+
@Override
103+
public void beforeRequest(Request<?> request) {
104+
super.beforeRequest(request);
105+
if (this.headersCustomizers.isEmpty()) {
106+
return;
107+
}
108+
String httpMethod = request.getHttpMethod().name();
109+
String uri = request.getEndpoint().toString();
110+
Map<String, List<String>> currentHeaders = extractHeaders(request);
111+
Set<String> protectedHeaders =
112+
currentHeaders.keySet().stream()
113+
.map(String::toLowerCase)
114+
.collect(Collectors.toSet()); // convert to lower case for case in-sensitive lookup
115+
116+
for (HttpHeadersCustomizer customizer : this.headersCustomizers) {
117+
if (customizer.applies(httpMethod, uri, currentHeaders)) {
118+
Map<String, List<String>> newHeaders = customizer.newHeaders();
119+
120+
throwIfExistingHeadersAreModified(protectedHeaders, newHeaders.keySet(), customizer);
121+
122+
for (Map.Entry<String, List<String>> entry : newHeaders.entrySet()) {
123+
for (String value : entry.getValue()) {
124+
request.addHeader(entry.getKey(), value);
125+
}
126+
}
127+
}
128+
}
129+
}
130+
131+
private static Map<String, List<String>> extractHeaders(HttpRequest request) {
132+
Map<String, List<String>> headerMap = new HashMap<>();
133+
for (Header header : request.getAllHeaders()) {
134+
headerMap.computeIfAbsent(header.getName(), k -> new ArrayList<>()).add(header.getValue());
135+
}
136+
return headerMap;
137+
}
138+
139+
private static Map<String, List<String>> extractHeaders(Request<?> request) {
140+
Map<String, List<String>> headerMap = new HashMap<>();
141+
for (Map.Entry<String, String> entry : request.getHeaders().entrySet()) {
142+
headerMap.computeIfAbsent(entry.getKey(), k -> new ArrayList<>()).add(entry.getValue());
143+
}
144+
return headerMap;
145+
}
146+
147+
/**
148+
* Checks if any header names from the customizer's new headers attempt to override existing
149+
* driver-set headers. Compares header names case-insensitively.
150+
*
151+
* @param protectedHeaders A set of lowercase header names initially present on the request.
152+
* @param newHeaders A set of header names provided by the customizer.
153+
* @param customizer The customizer attempting to add headers, for logging/exception messages.
154+
* @throws DriverHeaderOverridingNotAllowedException If an override is detected.
155+
*/
156+
private static void throwIfExistingHeadersAreModified(
157+
Set<String> protectedHeaders, Set<String> newHeaders, HttpHeadersCustomizer customizer) {
158+
for (String headerName : newHeaders) {
159+
if (headerName != null && protectedHeaders.contains(headerName.toLowerCase())) {
160+
logger.debug(
161+
"Customizer {} attempted to override existing driver header: {}",
162+
customizer.getClass().getName(),
163+
headerName);
164+
throw new DriverHeaderOverridingNotAllowedException(headerName);
165+
}
166+
}
167+
}
168+
169+
public static class DriverHeaderOverridingNotAllowedException extends RuntimeException {
170+
public DriverHeaderOverridingNotAllowedException(String header) {
171+
super(String.format("Driver headers overriding not allowed. Tried for header: %s", header));
172+
}
173+
}
174+
}

‎src/main/java/net/snowflake/client/core/HttpUtil.java‎

Lines changed: 92 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.security.KeyManagementException;
2222
import java.security.NoSuchAlgorithmException;
2323
import java.time.Duration;
24+
import java.util.List;
2425
import java.util.Map;
2526
import java.util.Properties;
2627
import java.util.concurrent.ConcurrentHashMap;
@@ -29,6 +30,7 @@
2930
import javax.annotation.Nullable;
3031
import javax.net.ssl.TrustManager;
3132
import net.snowflake.client.jdbc.ErrorCode;
33+
import net.snowflake.client.jdbc.HttpHeadersCustomizer;
3234
import net.snowflake.client.jdbc.RestRequest;
3335
import net.snowflake.client.jdbc.RetryContextManager;
3436
import net.snowflake.client.jdbc.SnowflakeDriver;
@@ -290,6 +292,24 @@ static String buildUserAgent(String customSuffix) {
290292
*/
291293
public static CloseableHttpClient buildHttpClient(
292294
@Nullable HttpClientSettingsKey key, File ocspCacheFile, boolean downloadUnCompressed) {
295+
return buildHttpClient(key, ocspCacheFile, downloadUnCompressed, null);
296+
}
297+
298+
/**
299+
* Build an Http client using our set of default.
300+
*
301+
* @param key Key to HttpClient hashmap containing OCSP mode and proxy information, could be null
302+
* @param ocspCacheFile OCSP response cache file. If null, the default OCSP response file will be
303+
* used.
304+
* @param downloadUnCompressed Whether the HTTP client should be built requesting no decompression
305+
* @param httpHeadersCustomizers List of HTTP headers customizers
306+
* @return HttpClient object
307+
*/
308+
public static CloseableHttpClient buildHttpClient(
309+
@Nullable HttpClientSettingsKey key,
310+
File ocspCacheFile,
311+
boolean downloadUnCompressed,
312+
List<HttpHeadersCustomizer> httpHeadersCustomizers) {
293313
logger.debug(
294314
"Building http client with client settings key: {}, ocsp cache file: {}, download uncompressed: {}",
295315
key != null ? key.toString() : null,
@@ -422,6 +442,12 @@ public static CloseableHttpClient buildHttpClient(
422442
logger.debug("Disabling content compression for http client");
423443
httpClientBuilder.disableContentCompression();
424444
}
445+
if (httpHeadersCustomizers != null && !httpHeadersCustomizers.isEmpty()) {
446+
logger.debug("Setting up http headers customizers");
447+
httpClientBuilder.setRetryHandler(new AttributeEnhancingHttpRequestRetryHandler());
448+
httpClientBuilder.addInterceptorLast(
449+
new HeaderCustomizerHttpRequestInterceptor(httpHeadersCustomizers));
450+
}
425451
return httpClientBuilder.build();
426452
} catch (NoSuchAlgorithmException | KeyManagementException ex) {
427453
throw new SSLInitializationException(ex.getMessage(), ex);
@@ -464,7 +490,7 @@ public static void updateRoutePlanner(HttpClientSettingsKey key) {
464490
* @return HttpClient object shared across all connections
465491
*/
466492
public static CloseableHttpClient getHttpClient(HttpClientSettingsKey ocspAndProxyKey) {
467-
return initHttpClient(ocspAndProxyKey, null);
493+
return initHttpClient(ocspAndProxyKey, null, null);
468494
}
469495

470496
/**
@@ -475,7 +501,31 @@ public static CloseableHttpClient getHttpClient(HttpClientSettingsKey ocspAndPro
475501
*/
476502
public static CloseableHttpClient getHttpClientWithoutDecompression(
477503
HttpClientSettingsKey ocspAndProxyKey) {
478-
return initHttpClientWithoutDecompression(ocspAndProxyKey, null);
504+
return initHttpClientWithoutDecompression(ocspAndProxyKey, null, null);
505+
}
506+
507+
/**
508+
* Gets HttpClient with insecureMode false
509+
*
510+
* @param ocspAndProxyKey OCSP mode and proxy settings for httpclient
511+
* @param httpHeadersCustomizers List of HTTP headers customizers
512+
* @return HttpClient object shared across all connections
513+
*/
514+
public static CloseableHttpClient getHttpClient(
515+
HttpClientSettingsKey ocspAndProxyKey, List<HttpHeadersCustomizer> httpHeadersCustomizers) {
516+
return initHttpClient(ocspAndProxyKey, null, httpHeadersCustomizers);
517+
}
518+
519+
/**
520+
* Gets HttpClient with insecureMode false and disabling decompression
521+
*
522+
* @param ocspAndProxyKey OCSP mode and proxy settings for httpclient
523+
* @param httpHeadersCustomizers List of HTTP headers customizers
524+
* @return HttpClient object shared across all connections
525+
*/
526+
public static CloseableHttpClient getHttpClientWithoutDecompression(
527+
HttpClientSettingsKey ocspAndProxyKey, List<HttpHeadersCustomizer> httpHeadersCustomizers) {
528+
return initHttpClientWithoutDecompression(ocspAndProxyKey, null, httpHeadersCustomizers);
479529
}
480530

481531
/**
@@ -489,7 +539,7 @@ public static CloseableHttpClient initHttpClientWithoutDecompression(
489539
HttpClientSettingsKey key, File ocspCacheFile) {
490540
updateRoutePlanner(key);
491541
return httpClientWithoutDecompression.computeIfAbsent(
492-
key, k -> buildHttpClient(key, ocspCacheFile, true));
542+
key, k -> buildHttpClient(key, ocspCacheFile, true, null));
493543
}
494544

495545
/**
@@ -502,7 +552,42 @@ public static CloseableHttpClient initHttpClientWithoutDecompression(
502552
public static CloseableHttpClient initHttpClient(HttpClientSettingsKey key, File ocspCacheFile) {
503553
updateRoutePlanner(key);
504554
return httpClient.computeIfAbsent(
505-
key, k -> buildHttpClient(key, ocspCacheFile, key.getGzipDisabled()));
555+
key, k -> buildHttpClient(key, ocspCacheFile, key.getGzipDisabled(), null));
556+
}
557+
558+
/**
559+
* Accessor for the HTTP client singleton.
560+
*
561+
* @param key contains information needed to build specific HttpClient
562+
* @param ocspCacheFile OCSP response cache file name. if null, the default file will be used.
563+
* @param httpHeadersCustomizers List of HTTP headers customizers
564+
* @return HttpClient object shared across all connections
565+
*/
566+
public static CloseableHttpClient initHttpClientWithoutDecompression(
567+
HttpClientSettingsKey key,
568+
File ocspCacheFile,
569+
List<HttpHeadersCustomizer> httpHeadersCustomizers) {
570+
updateRoutePlanner(key);
571+
return httpClientWithoutDecompression.computeIfAbsent(
572+
key, k -> buildHttpClient(key, ocspCacheFile, true, httpHeadersCustomizers));
573+
}
574+
575+
/**
576+
* Accessor for the HTTP client singleton.
577+
*
578+
* @param key contains information needed to build specific HttpClient
579+
* @param ocspCacheFile OCSP response cache file name. if null, the default file will be used.
580+
* @param httpHeadersCustomizers List of HTTP headers customizers
581+
* @return HttpClient object shared across all connections
582+
*/
583+
public static CloseableHttpClient initHttpClient(
584+
HttpClientSettingsKey key,
585+
File ocspCacheFile,
586+
List<HttpHeadersCustomizer> httpHeadersCustomizers) {
587+
updateRoutePlanner(key);
588+
return httpClient.computeIfAbsent(
589+
key,
590+
k -> buildHttpClient(key, ocspCacheFile, key.getGzipDisabled(), httpHeadersCustomizers));
506591
}
507592

508593
/**
@@ -622,7 +707,7 @@ static String executeRequestWithoutCookies(
622707
false, // no retry parameter
623708
true, // guid? (do we need this?)
624709
false, // no retry on HTTP 403
625-
getHttpClient(ocspAndProxyKey),
710+
getHttpClient(ocspAndProxyKey, null),
626711
new ExecTimeTelemetryData(),
627712
null);
628713
}
@@ -680,7 +765,7 @@ public static String executeGeneralRequestOmitRequestGuid(
680765
false,
681766
false,
682767
false,
683-
getHttpClient(ocspAndProxyAndGzipKey),
768+
getHttpClient(ocspAndProxyAndGzipKey, null),
684769
new ExecTimeTelemetryData(),
685770
null);
686771
}
@@ -860,7 +945,7 @@ public static String executeRequest(
860945
includeRetryParameters,
861946
true, // include request GUID
862947
retryOnHTTP403,
863-
getHttpClient(ocspAndProxyKey),
948+
getHttpClient(ocspAndProxyKey, null),
864949
execTimeData,
865950
retryContextManager);
866951
}

0 commit comments

Comments
 (0)