diff --git a/docs/configuration/index.md b/docs/configuration/index.md index aad78964f199..7446f7e8b4ec 100644 --- a/docs/configuration/index.md +++ b/docs/configuration/index.md @@ -348,6 +348,24 @@ Coordinator and Overlord log changes to lookups, segment load/drop rules, and dy |`druid.audit.manager.maxPayloadSizeBytes`|The maximum size of audit payload to store in Druid's metadata store audit table. If the size of audit payload exceeds this value, the audit log would be stored with a message indicating that the payload was omitted instead. Setting `maxPayloadSizeBytes` to -1 (default value) disables this check, meaning Druid will always store audit payload regardless of it's size. Setting to any negative number other than `-1` is invalid. Human-readable format is supported, see [here](human-readable-byte.md). |-1| |`druid.audit.manager.skipNullField`|If true, the audit payload stored in metadata store will exclude any field with null value. |false| +### Request header propagation + +Druid can capture configured inbound HTTP headers and propagate their values through the query context and into audit events. This is useful for correlating Druid queries and configuration changes with an external distributed-trace or request-tracking system. + +|Property|Description|Default| +|--------|-----------|-------| +|`druid.audit.requestHeaders.headerToContextKey`|JSON map of inbound HTTP header name to the [query context](../querying/query-context.md) key the header value is bound to. For each configured header present on a request, the value is captured by a servlet filter and injected into the query context, from where Druid's native sub-query context propagation carries it to data servers (for example, broker to historical). The captured header values are also recorded, keyed by their context key (for example `traceId`), in the `requestMetadata` map of the audit entry for any configuration change made via an audited API. Mapping a header to a reserved Druid context key (`queryId`, `subQueryId`, or `sqlQueryId`) is rejected at startup, because it would let a client overwrite the server-assigned value. Set to an empty map (`{}`) to disable header propagation.|`{"X-Druid-Trace-Id":"traceId"}`| + +For example, to also capture a tenant identifier: + +```properties +druid.audit.requestHeaders.headerToContextKey={"X-Druid-Trace-Id":"traceId","X-Tenant-Id":"tenantId"} +``` + +A client-supplied value for a configured context key in the query body is ignored (stripped); only a value captured from the corresponding inbound header takes effect. This prevents a client from spoofing the trace identifier by setting it in the query JSON. + +When query-context authorization is enabled (`druid.auth.authorizeQueryContextParams=true`), a value captured from a header is subject to the same `QUERY_CONTEXT` `WRITE` authorization as a value supplied in the query body. To allow the default `traceId` (or any other propagated key) without an explicit grant, add it to `druid.auth.unsecuredContextKeys`. + ### Metadata storage These properties specify the JDBC connection and other configuration around the metadata storage. The only services that connect to the metadata storage with these properties are the [Coordinator](../design/coordinator.md) and [Overlord](../design/overlord.md). diff --git a/extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/RabbitStreamIndexTaskClientFactory.java b/extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/RabbitStreamIndexTaskClientFactory.java index 088dc6b77d53..4f2155db04c2 100644 --- a/extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/RabbitStreamIndexTaskClientFactory.java +++ b/extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/RabbitStreamIndexTaskClientFactory.java @@ -21,19 +21,30 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Inject; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Json; import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory; import org.apache.druid.java.util.http.client.HttpClient; +import javax.annotation.Nullable; + @LazySingleton public class RabbitStreamIndexTaskClientFactory extends SeekableStreamIndexTaskClientFactory { @Inject public RabbitStreamIndexTaskClientFactory( @EscalatedGlobal HttpClient httpClient, - @Json ObjectMapper mapper) + @Json ObjectMapper mapper, + @Nullable RequestHeaderContextConfig requestHeaderContextConfig) + { + super(httpClient, mapper, requestHeaderContextConfig != null ? requestHeaderContextConfig : new RequestHeaderContextConfig()); + } + + public RabbitStreamIndexTaskClientFactory( + HttpClient httpClient, + ObjectMapper mapper) { super(httpClient, mapper); } diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientFactory.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientFactory.java index 10b57dd4757f..7d9619ce30c6 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientFactory.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientFactory.java @@ -21,6 +21,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Inject; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.data.input.kafka.KafkaTopicPartition; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.annotations.EscalatedGlobal; @@ -28,13 +29,24 @@ import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory; import org.apache.druid.java.util.http.client.HttpClient; +import javax.annotation.Nullable; + @LazySingleton public class KafkaIndexTaskClientFactory extends SeekableStreamIndexTaskClientFactory { @Inject public KafkaIndexTaskClientFactory( @EscalatedGlobal HttpClient httpClient, - @Json ObjectMapper mapper + @Json ObjectMapper mapper, + @Nullable RequestHeaderContextConfig requestHeaderContextConfig + ) + { + super(httpClient, mapper, requestHeaderContextConfig != null ? requestHeaderContextConfig : new RequestHeaderContextConfig()); + } + + public KafkaIndexTaskClientFactory( + HttpClient httpClient, + ObjectMapper mapper ) { super(httpClient, mapper); diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClientFactory.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClientFactory.java index 599c32b21584..8008d36f0ae8 100644 --- a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClientFactory.java +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClientFactory.java @@ -21,19 +21,31 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Inject; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Json; import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory; import org.apache.druid.java.util.http.client.HttpClient; +import javax.annotation.Nullable; + @LazySingleton public class KinesisIndexTaskClientFactory extends SeekableStreamIndexTaskClientFactory { @Inject public KinesisIndexTaskClientFactory( @EscalatedGlobal HttpClient httpClient, - @Json ObjectMapper mapper + @Json ObjectMapper mapper, + @Nullable RequestHeaderContextConfig requestHeaderContextConfig + ) + { + super(httpClient, mapper, requestHeaderContextConfig != null ? requestHeaderContextConfig : new RequestHeaderContextConfig()); + } + + public KinesisIndexTaskClientFactory( + HttpClient httpClient, + ObjectMapper mapper ) { super(httpClient, mapper); diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskClientFactory.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskClientFactory.java index 92acb3a33296..456ab18d3f0c 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskClientFactory.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskClientFactory.java @@ -20,6 +20,7 @@ package org.apache.druid.indexing.seekablestream; import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.indexing.common.TaskInfoProvider; import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor; import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorTuningConfig; @@ -36,14 +37,25 @@ public abstract class SeekableStreamIndexTaskClientFactory build( return new SeekableStreamIndexTaskClientAsyncImpl<>( dataSource, - new ServiceClientFactoryImpl(httpClient, connectExec), + new ServiceClientFactoryImpl(httpClient, connectExec, requestHeaderContextConfig), taskInfoProvider, jsonMapper, tuningConfig.getHttpTimeout(), diff --git a/processing/src/main/java/org/apache/druid/audit/AuditManager.java b/processing/src/main/java/org/apache/druid/audit/AuditManager.java index 3ab126e93373..4e39615bbc73 100644 --- a/processing/src/main/java/org/apache/druid/audit/AuditManager.java +++ b/processing/src/main/java/org/apache/druid/audit/AuditManager.java @@ -31,6 +31,7 @@ public interface AuditManager String X_DRUID_AUTHOR = "X-Druid-Author"; String X_DRUID_COMMENT = "X-Druid-Comment"; + String X_DRUID_TRACE_ID = "X-Druid-Trace-Id"; void doAudit(AuditEntry event); diff --git a/processing/src/main/java/org/apache/druid/audit/RequestHeaderContextConfig.java b/processing/src/main/java/org/apache/druid/audit/RequestHeaderContextConfig.java new file mode 100644 index 000000000000..594f6e5b9a6a --- /dev/null +++ b/processing/src/main/java/org/apache/druid/audit/RequestHeaderContextConfig.java @@ -0,0 +1,159 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.audit; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import org.apache.druid.error.DruidException; +import org.apache.druid.query.BaseQuery; + +import java.util.Map; +import java.util.Set; +import java.util.function.BiConsumer; + +/** + * Configures which inbound HTTP headers are captured and where they're stored in the + * {@link org.apache.druid.query.Query} context. The map is keyed by header name (case as + * sent on the wire) and the value is the reserved context key. + * + *

Defaults: {@code X-Druid-Trace-Id → traceId} so trace propagation works out of the + * box. Operators can extend by setting: + * 

+ *   druid.audit.requestHeaders.headerToContextKey={"X-Druid-Trace-Id": "traceId",
+ *                                                  "X-Custom-Foo": "foo"}
+ *
+ * + *

The values in this map are reserved context keys: user-supplied values for these keys + * are stripped from the query context before the filter-captured value is merged + * (anti-spoof). Setting overlapping keys to existing query-context keys is therefore + * discouraged — pick context-key names that don't collide with standard Druid context. + */ +public class RequestHeaderContextConfig +{ + private static final Map DEFAULT_HEADER_TO_CONTEXT_KEY = ImmutableMap.of( + AuditManager.X_DRUID_TRACE_ID, "traceId" + ); + + /** + * Context keys that are server-managed by Druid and must not be configurable as targets + * for header propagation. Mapping a header to one of these would let a client overwrite + * the server-assigned value (e.g. the random per-request {@code queryId}), breaking + * query identification, cancellation, and metric correlation. + */ + private static final Set RESERVED_DRUID_CONTEXT_KEYS = ImmutableSet.of( + BaseQuery.QUERY_ID, + BaseQuery.SUB_QUERY_ID, + BaseQuery.SQL_QUERY_ID + ); + + // @JsonProperty on the field is REQUIRED: JsonConfigurator.verifyClazzIsConfigurable() + // rejects config classes whose bean-property fields lack a Jackson field annotation, which + // would fail server startup (JettyServerModule binds this via JsonConfigProvider). + // Declared as ImmutableMap (not Map) so getHeaderToContextKey() can return the field + // reference directly without exposing a mutable internal representation (CodeQL). + @JsonProperty + private final ImmutableMap headerToContextKey; + + public RequestHeaderContextConfig() + { + this(DEFAULT_HEADER_TO_CONTEXT_KEY); + } + + /** + * @param headerToContextKey explicit mapping; null means "use defaults" (the typical case + * when {@code JsonConfigProvider} deserializes from an empty config block); + * {@link java.util.Collections#emptyMap()} explicitly disables propagation. + */ + @JsonCreator + public RequestHeaderContextConfig( + @JsonProperty("headerToContextKey") Map headerToContextKey + ) + { + final Map raw = headerToContextKey == null + ? DEFAULT_HEADER_TO_CONTEXT_KEY + : headerToContextKey; + for (String contextKey : raw.values()) { + if (RESERVED_DRUID_CONTEXT_KEYS.contains(contextKey)) { + throw DruidException.forPersona(DruidException.Persona.OPERATOR) + .ofCategory(DruidException.Category.INVALID_INPUT) + .build( + "druid.audit.requestHeaders.headerToContextKey maps a header to " + + "reserved Druid context key[%s]; this would let a client overwrite " + + "the server-assigned value. Reserved keys are %s.", + contextKey, RESERVED_DRUID_CONTEXT_KEYS + ); + } + } + this.headerToContextKey = ImmutableMap.copyOf(raw); + } + + /** + * Returns the configured map of {@code header → context-key}. Immutable. Empty means + * no header propagation is performed (the feature is effectively disabled). + */ + @JsonProperty + public Map getHeaderToContextKey() + { + return headerToContextKey; + } + + /** + * For each configured {@code header → contextKey}, if the supplied query context + * carries a value for {@code contextKey}, invokes {@code headerSetter} with + * {@code (headerName, stringifiedValue)}. Used by internal RPC clients + * (broker → historical, etc.) to attach propagated headers to outbound requests so + * the receiving node's filter captures them just as it would for a client-originated + * request. + */ + public void applyToOutboundRequest(Map queryContext, BiConsumer headerSetter) + { + if (queryContext == null || queryContext.isEmpty() || headerToContextKey.isEmpty()) { + return; + } + for (Map.Entry entry : headerToContextKey.entrySet()) { + final Object value = queryContext.get(entry.getValue()); + if (value != null) { + headerSetter.accept(entry.getKey(), value.toString()); + } + } + } + + /** + * For each configured {@code header → contextKey}, if the supplied captured-header map (keyed + * by context key, as held by {@code RequestHeaderContext}) carries a value for {@code contextKey}, + * invokes {@code headerSetter} with {@code (headerName, value)}. Used by the shared outbound + * HTTP client to forward configured headers on ALL inter-service calls made on the request + * thread (not just queries), so the chain of internal calls carries the same header values. + */ + public void applyCapturedHeaders(Map capturedByContextKey, BiConsumer headerSetter) + { + if (capturedByContextKey == null || capturedByContextKey.isEmpty() || headerToContextKey.isEmpty()) { + return; + } + for (Map.Entry entry : headerToContextKey.entrySet()) { + final String value = capturedByContextKey.get(entry.getValue()); + if (value != null) { + headerSetter.accept(entry.getKey(), value); + } + } + } +} diff --git a/processing/src/main/java/org/apache/druid/audit/RequestInfo.java b/processing/src/main/java/org/apache/druid/audit/RequestInfo.java index 706fdca2f01a..d322f6296215 100644 --- a/processing/src/main/java/org/apache/druid/audit/RequestInfo.java +++ b/processing/src/main/java/org/apache/druid/audit/RequestInfo.java @@ -20,8 +20,12 @@ package org.apache.druid.audit; import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.ImmutableMap; +import javax.annotation.Nullable; +import java.util.Map; import java.util.Objects; /** @@ -33,19 +37,30 @@ public class RequestInfo private final String method; private final String uri; private final String queryParams; + @Nullable + private final Map requestMetadata; + + public RequestInfo(String service, String method, String uri, String queryParams) + { + this(service, method, uri, queryParams, null); + } @JsonCreator public RequestInfo( @JsonProperty("service") String service, @JsonProperty("method") String method, @JsonProperty("uri") String uri, - @JsonProperty("queryParams") String queryParams + @JsonProperty("queryParams") String queryParams, + @JsonProperty("requestMetadata") @Nullable Map requestMetadata ) { this.service = service; this.method = method; this.uri = uri; this.queryParams = queryParams; + this.requestMetadata = (requestMetadata == null || requestMetadata.isEmpty()) + ? null + : ImmutableMap.copyOf(requestMetadata); } @JsonProperty @@ -72,6 +87,21 @@ public String getQueryParams() return queryParams; } + /** + * Metadata captured from configured inbound request headers (see + * {@link RequestHeaderContextConfig}), keyed by the operator-configured context key (for + * example {@code traceId}). Lets audit consumers extract whichever fields they need, and new + * mapped headers require no code change. Omitted from the audit JSON when empty, so records + * remain byte-identical when no configured header is sent. + */ + @JsonProperty + @JsonInclude(JsonInclude.Include.NON_EMPTY) + @Nullable + public Map getRequestMetadata() + { + return requestMetadata; + } + @Override public boolean equals(Object o) { @@ -85,13 +115,14 @@ public boolean equals(Object o) return Objects.equals(this.service, that.service) && Objects.equals(this.method, that.method) && Objects.equals(this.uri, that.uri) - && Objects.equals(this.queryParams, that.queryParams); + && Objects.equals(this.queryParams, that.queryParams) + && Objects.equals(this.requestMetadata, that.requestMetadata); } @Override public int hashCode() { - return Objects.hash(service, method, uri, queryParams); + return Objects.hash(service, method, uri, queryParams, requestMetadata); } @Override @@ -102,6 +133,7 @@ public String toString() ", method='" + method + '\'' + ", path='" + uri + '\'' + ", queryParams='" + queryParams + '\'' + + ", requestMetadata=" + requestMetadata + '}'; } } diff --git a/processing/src/main/java/org/apache/druid/common/config/JacksonConfigManager.java b/processing/src/main/java/org/apache/druid/common/config/JacksonConfigManager.java index aa432058f0af..0146e3148553 100644 --- a/processing/src/main/java/org/apache/druid/common/config/JacksonConfigManager.java +++ b/processing/src/main/java/org/apache/druid/common/config/JacksonConfigManager.java @@ -27,6 +27,7 @@ import org.apache.druid.audit.AuditEntry; import org.apache.druid.audit.AuditInfo; import org.apache.druid.audit.AuditManager; +import org.apache.druid.audit.RequestInfo; import org.apache.druid.common.config.ConfigManager.SetResult; import org.apache.druid.guice.annotations.Json; import org.apache.druid.java.util.common.jackson.JacksonUtils; @@ -112,6 +113,26 @@ public SetResult set( T newValue, AuditInfo auditInfo ) + { + return set(key, oldValue, newValue, auditInfo, null); + } + + /** + * Set the config and add an audit entry that also carries request metadata. + * + * @param requestInfo details of the REST request that triggered the change (service, + * method, URI, and trace ID). May be null for non-HTTP-triggered + * changes; when provided it lands on the audit entry's {@code request} + * field so config-change audits can be correlated with distributed + * traces. See {@code AuthorizationUtils.buildRequestInfo}. + */ + public SetResult set( + String key, + @Nullable byte[] oldValue, + T newValue, + AuditInfo auditInfo, + @Nullable RequestInfo requestInfo + ) { ConfigSerde configSerde = create(newValue.getClass(), null); // Audit and actual config change are done in separate transactions @@ -121,6 +142,7 @@ public SetResult set( .key(key) .type(key) .auditInfo(auditInfo) + .request(requestInfo) .payload(newValue) .build() ); diff --git a/processing/src/test/java/org/apache/druid/audit/AuditInfoTest.java b/processing/src/test/java/org/apache/druid/audit/AuditInfoTest.java index b75b531ce9df..1b1254763aaa 100644 --- a/processing/src/test/java/org/apache/druid/audit/AuditInfoTest.java +++ b/processing/src/test/java/org/apache/druid/audit/AuditInfoTest.java @@ -20,6 +20,7 @@ package org.apache.druid.audit; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableMap; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.DateTimes; import org.junit.jupiter.api.Assertions; @@ -89,4 +90,46 @@ public void testRequestInfoEquality() throws IOException Assertions.assertEquals(requestInfo, deserialized); } + @Test + public void testRequestInfoMetadataSerde() throws IOException + { + RequestInfo withMeta = new RequestInfo( + "overlord", "GET", "/uri", "a=b", ImmutableMap.of("traceId", "trace-abc-123") + ); + RequestInfo deserialized = mapper.readValue(mapper.writeValueAsString(withMeta), RequestInfo.class); + Assertions.assertEquals(withMeta, deserialized); + Assertions.assertEquals("trace-abc-123", deserialized.getRequestMetadata().get("traceId")); + } + + @Test + public void testRequestInfoEmptyMetadataOmittedFromJson() throws IOException + { + RequestInfo withoutMeta = new RequestInfo("overlord", "GET", "/uri", "a=b", null); + String json = mapper.writeValueAsString(withoutMeta); + Assertions.assertFalse( + json.contains("requestMetadata"), + "empty requestMetadata should be omitted from JSON for wire-size hygiene; got: " + json + ); + } + + @Test + public void testRequestInfoMetadataBackwardsCompatible() throws IOException + { + // Old audit rows persisted before requestMetadata existed must still deserialize. + String legacyJson = "{\"service\":\"overlord\",\"method\":\"GET\",\"uri\":\"/uri\",\"queryParams\":\"a=b\"}"; + RequestInfo parsed = mapper.readValue(legacyJson, RequestInfo.class); + Assertions.assertEquals("overlord", parsed.getService()); + Assertions.assertNull(parsed.getRequestMetadata()); + } + + @Test + public void testRequestInfoEqualityConsidersMetadata() + { + RequestInfo a = new RequestInfo("s", "GET", "/u", "p", ImmutableMap.of("traceId", "trace-1")); + RequestInfo b = new RequestInfo("s", "GET", "/u", "p", ImmutableMap.of("traceId", "trace-2")); + RequestInfo c = new RequestInfo("s", "GET", "/u", "p", ImmutableMap.of("traceId", "trace-1")); + Assertions.assertNotEquals(a, b); + Assertions.assertEquals(a, c); + } + } diff --git a/processing/src/test/java/org/apache/druid/audit/RequestHeaderContextConfigTest.java b/processing/src/test/java/org/apache/druid/audit/RequestHeaderContextConfigTest.java new file mode 100644 index 000000000000..2f8c1400b0b3 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/audit/RequestHeaderContextConfigTest.java @@ -0,0 +1,164 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.audit; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.io.IOException; + +public class RequestHeaderContextConfigTest +{ + private final ObjectMapper mapper = new DefaultObjectMapper(); + + @Test + public void testDefaultIncludesTraceId() + { + RequestHeaderContextConfig config = new RequestHeaderContextConfig(); + Assertions.assertEquals( + "traceId", + config.getHeaderToContextKey().get(AuditManager.X_DRUID_TRACE_ID) + ); + } + + @Test + public void testDeserFromJson() throws IOException + { + String json = "{\"headerToContextKey\":{\"X-Custom-Foo\":\"foo\",\"X-Druid-Trace-Id\":\"traceId\"}}"; + RequestHeaderContextConfig parsed = mapper.readValue(json, RequestHeaderContextConfig.class); + Assertions.assertEquals("foo", parsed.getHeaderToContextKey().get("X-Custom-Foo")); + Assertions.assertEquals("traceId", parsed.getHeaderToContextKey().get("X-Druid-Trace-Id")); + } + + @Test + public void testEmptyMapDisablesFeature() throws IOException + { + String json = "{\"headerToContextKey\":{}}"; + RequestHeaderContextConfig parsed = mapper.readValue(json, RequestHeaderContextConfig.class); + Assertions.assertTrue(parsed.getHeaderToContextKey().isEmpty()); + } + + @Test + public void testEmptyJsonObjectFallsBackToDefault() throws IOException + { + // Critical: JsonConfigProvider does convertValue({}, RequestHeaderContextConfig.class) + // when nothing is configured under druid.audit.requestHeaders.*; Jackson invokes the + // @JsonCreator with headerToContextKey=null. That must NOT silently disable the feature. + RequestHeaderContextConfig parsed = mapper.readValue("{}", RequestHeaderContextConfig.class); + Assertions.assertEquals( + "traceId", + parsed.getHeaderToContextKey().get(AuditManager.X_DRUID_TRACE_ID) + ); + } + + @Test + public void testMappingToReservedQueryIdRejected() + { + final String json = "{\"headerToContextKey\":{\"X-My-Header\":\"queryId\"}}"; + Assertions.assertThrows( + com.fasterxml.jackson.databind.JsonMappingException.class, + () -> mapper.readValue(json, RequestHeaderContextConfig.class) + ); + } + + @Test + public void testMappingToReservedSqlQueryIdRejected() + { + final String json = "{\"headerToContextKey\":{\"X-Foo\":\"sqlQueryId\"}}"; + Assertions.assertThrows( + com.fasterxml.jackson.databind.JsonMappingException.class, + () -> mapper.readValue(json, RequestHeaderContextConfig.class) + ); + } + + @Test + public void testGetHeaderToContextKeyIsImmutable() + { + RequestHeaderContextConfig config = new RequestHeaderContextConfig(); + Assertions.assertThrows( + UnsupportedOperationException.class, + () -> config.getHeaderToContextKey().put("X-Evil", "evil") + ); + } + + @Test + public void testApplyToOutboundRequest_attachesPresentHeadersFromContext() + { + RequestHeaderContextConfig config = new RequestHeaderContextConfig(); + java.util.Map ctx = com.google.common.collect.ImmutableMap.of( + "traceId", "abc-123", + "otherKey", "ignored" + ); + java.util.Map captured = new java.util.HashMap<>(); + config.applyToOutboundRequest(ctx, captured::put); + Assertions.assertEquals(1, captured.size()); + Assertions.assertEquals("abc-123", captured.get(AuditManager.X_DRUID_TRACE_ID)); + } + + @Test + public void testApplyToOutboundRequest_skipsAbsentKeys() + { + RequestHeaderContextConfig config = new RequestHeaderContextConfig(); + java.util.Map captured = new java.util.HashMap<>(); + config.applyToOutboundRequest(java.util.Collections.emptyMap(), captured::put); + Assertions.assertTrue(captured.isEmpty()); + } + + @Test + public void testApplyToOutboundRequest_nullContextIsNoOp() + { + RequestHeaderContextConfig config = new RequestHeaderContextConfig(); + java.util.Map captured = new java.util.HashMap<>(); + config.applyToOutboundRequest(null, captured::put); + Assertions.assertTrue(captured.isEmpty()); + } + + @Test + public void testApplyCapturedHeaders_attachesByHeaderName() + { + RequestHeaderContextConfig config = new RequestHeaderContextConfig( + com.google.common.collect.ImmutableMap.of( + AuditManager.X_DRUID_TRACE_ID, "traceId", + "X-Tenant-Id", "tenantId" + ) + ); + // keyed by context key, as held by RequestHeaderContext + java.util.Map capturedByContextKey = com.google.common.collect.ImmutableMap.of( + "traceId", "abc-123", + "tenantId", "acme" + ); + java.util.Map outbound = new java.util.HashMap<>(); + config.applyCapturedHeaders(capturedByContextKey, outbound::put); + Assertions.assertEquals("abc-123", outbound.get(AuditManager.X_DRUID_TRACE_ID)); + Assertions.assertEquals("acme", outbound.get("X-Tenant-Id")); + } + + @Test + public void testApplyCapturedHeaders_emptyOrNullIsNoOp() + { + RequestHeaderContextConfig config = new RequestHeaderContextConfig(); + java.util.Map outbound = new java.util.HashMap<>(); + config.applyCapturedHeaders(java.util.Collections.emptyMap(), outbound::put); + config.applyCapturedHeaders(null, outbound::put); + Assertions.assertTrue(outbound.isEmpty()); + } +} diff --git a/processing/src/test/java/org/apache/druid/common/config/JacksonConfigManagerTest.java b/processing/src/test/java/org/apache/druid/common/config/JacksonConfigManagerTest.java index c4cf97a236f5..61cd20de767a 100644 --- a/processing/src/test/java/org/apache/druid/common/config/JacksonConfigManagerTest.java +++ b/processing/src/test/java/org/apache/druid/common/config/JacksonConfigManagerTest.java @@ -23,9 +23,11 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableMap; import org.apache.druid.audit.AuditEntry; import org.apache.druid.audit.AuditInfo; import org.apache.druid.audit.AuditManager; +import org.apache.druid.audit.RequestInfo; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -77,6 +79,39 @@ public void testSet() Assertions.assertNotNull(auditCapture.getValue()); } + @Test + public void testSetWithRequestInfoAttachesRequestToAudit() + { + String key = "key"; + TestConfig val = new TestConfig("version", "string", 3); + AuditInfo auditInfo = new AuditInfo("testAuthor", "testIdentity", "testComment", "127.0.0.1"); + RequestInfo requestInfo = new RequestInfo( + "coordinator", "POST", "/druid/coordinator/v1/config", null, ImmutableMap.of("traceId", "trace-xyz") + ); + + jacksonConfigManager.set(key, null, val, auditInfo, requestInfo); + + ArgumentCaptor auditCapture = ArgumentCaptor.forClass(AuditEntry.class); + Mockito.verify(mockAuditManager).doAudit(auditCapture.capture()); + AuditEntry entry = auditCapture.getValue(); + Assertions.assertNotNull(entry.getRequest()); + Assertions.assertEquals("trace-xyz", entry.getRequest().getRequestMetadata().get("traceId")); + } + + @Test + public void testSetWithoutRequestInfoLeavesRequestNull() + { + String key = "key"; + TestConfig val = new TestConfig("version", "string", 3); + AuditInfo auditInfo = new AuditInfo("testAuthor", "testIdentity", "testComment", "127.0.0.1"); + + jacksonConfigManager.set(key, val, auditInfo); + + ArgumentCaptor auditCapture = ArgumentCaptor.forClass(AuditEntry.class); + Mockito.verify(mockAuditManager).doAudit(auditCapture.capture()); + Assertions.assertNull(auditCapture.getValue().getRequest()); + } + @Test public void testConvertByteToConfigWithNullConfigInByte() { diff --git a/server/src/main/java/org/apache/druid/client/DirectDruidClient.java b/server/src/main/java/org/apache/druid/client/DirectDruidClient.java index 3c746e410701..6effc98f3454 100644 --- a/server/src/main/java/org/apache/druid/client/DirectDruidClient.java +++ b/server/src/main/java/org/apache/druid/client/DirectDruidClient.java @@ -27,6 +27,7 @@ import com.google.common.util.concurrent.FutureCallback; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.java.util.common.RE; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.concurrent.Execs; @@ -70,6 +71,7 @@ import java.io.IOException; import java.io.InputStream; import java.io.SequenceInputStream; +import java.net.URI; import java.net.URL; import java.util.Enumeration; import java.util.concurrent.BlockingQueue; @@ -100,6 +102,7 @@ public class DirectDruidClient implements QueryRunner private final String scheme; private final String host; private final ServiceEmitter emitter; + private final RequestHeaderContextConfig requestHeaderContextConfig; private final AtomicInteger openConnections; private final boolean isSmile; @@ -131,6 +134,22 @@ public DirectDruidClient( ServiceEmitter emitter, ScheduledExecutorService queryCancellationExecutor ) + { + this(conglomerate, queryWatcher, objectMapper, httpClient, scheme, host, emitter, + queryCancellationExecutor, new RequestHeaderContextConfig()); + } + + public DirectDruidClient( + QueryRunnerFactoryConglomerate conglomerate, + QueryWatcher queryWatcher, + ObjectMapper objectMapper, + HttpClient httpClient, + String scheme, + String host, + ServiceEmitter emitter, + ScheduledExecutorService queryCancellationExecutor, + RequestHeaderContextConfig requestHeaderContextConfig + ) { this.conglomerate = conglomerate; this.queryWatcher = queryWatcher; @@ -143,6 +162,7 @@ public DirectDruidClient( this.isSmile = this.objectMapper.getFactory() instanceof SmileFactory; this.openConnections = new AtomicInteger(); this.queryCancellationExecutor = queryCancellationExecutor; + this.requestHeaderContextConfig = requestHeaderContextConfig; } public int getNumOpenConnections() @@ -471,18 +491,18 @@ private void checkTotalBytesLimit(long bytes) // we can increment the count earlier so that we can route the request to a different server openConnections.getAndIncrement(); try { - future = httpClient.go( - new Request( - HttpMethod.POST, - new URL(url) - ).setContent(objectMapper.writeValueAsBytes(Queries.withTimeout(query, timeLeft))) - .setHeader( - HttpHeaders.Names.CONTENT_TYPE, - isSmile ? SmileMediaTypes.APPLICATION_JACKSON_SMILE : MediaType.APPLICATION_JSON - ), - responseHandler, - Duration.millis(timeLeft) - ); + final Request outbound = new Request(HttpMethod.POST, URI.create(url).toURL()) + .setContent(objectMapper.writeValueAsBytes(Queries.withTimeout(query, timeLeft))) + .setHeader( + HttpHeaders.Names.CONTENT_TYPE, + isSmile ? SmileMediaTypes.APPLICATION_JACKSON_SMILE : MediaType.APPLICATION_JSON + ); + // Propagate configured request-context headers onto the inter-Druid RPC so the + // receiving node's RequestHeaderContextFilter captures them just as if a client + // had set them. Enables cross-node trace ID propagation without trusting the + // values already serialized in the query JSON body (anti-spoof). + requestHeaderContextConfig.applyToOutboundRequest(query.getContext(), outbound::setHeader); + future = httpClient.go(outbound, responseHandler, Duration.millis(timeLeft)); } catch (Exception e) { openConnections.getAndDecrement(); diff --git a/server/src/main/java/org/apache/druid/client/DirectDruidClientFactory.java b/server/src/main/java/org/apache/druid/client/DirectDruidClientFactory.java index 4be9f98edb65..416d7b9cdd3d 100644 --- a/server/src/main/java/org/apache/druid/client/DirectDruidClientFactory.java +++ b/server/src/main/java/org/apache/druid/client/DirectDruidClientFactory.java @@ -21,6 +21,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Inject; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.annotations.EscalatedClient; import org.apache.druid.guice.annotations.Smile; @@ -45,6 +46,7 @@ public class DirectDruidClientFactory implements QueryableDruidServer.Maker private final ObjectMapper smileMapper; private final HttpClient httpClient; private final ScheduledExecutorService queryCancellationExecutor; + private final RequestHeaderContextConfig requestHeaderContextConfig; @Inject public DirectDruidClientFactory( @@ -52,7 +54,8 @@ public DirectDruidClientFactory( final QueryRunnerFactoryConglomerate conglomerate, final QueryWatcher queryWatcher, final @Smile ObjectMapper smileMapper, - final @EscalatedClient HttpClient httpClient + final @EscalatedClient HttpClient httpClient, + final RequestHeaderContextConfig requestHeaderContextConfig ) { this.emitter = emitter; @@ -60,6 +63,7 @@ public DirectDruidClientFactory( this.queryWatcher = queryWatcher; this.smileMapper = smileMapper; this.httpClient = httpClient; + this.requestHeaderContextConfig = requestHeaderContextConfig; int threadCount = Math.max(1, JvmUtils.getRuntimeInfo().getAvailableProcessors() / 2); this.queryCancellationExecutor = ScheduledExecutors.fixed(threadCount, "query-cancellation-executor"); @@ -75,7 +79,8 @@ public DirectDruidClient makeDirectClient(DruidServer server) server.getScheme(), server.getHost(), emitter, - queryCancellationExecutor + queryCancellationExecutor, + requestHeaderContextConfig ); } diff --git a/server/src/main/java/org/apache/druid/rpc/ServiceClientFactoryImpl.java b/server/src/main/java/org/apache/druid/rpc/ServiceClientFactoryImpl.java index 5c997f21194d..ba10213ed290 100644 --- a/server/src/main/java/org/apache/druid/rpc/ServiceClientFactoryImpl.java +++ b/server/src/main/java/org/apache/druid/rpc/ServiceClientFactoryImpl.java @@ -19,6 +19,7 @@ package org.apache.druid.rpc; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.java.util.http.client.HttpClient; import java.util.concurrent.ScheduledExecutorService; @@ -30,14 +31,25 @@ public class ServiceClientFactoryImpl implements ServiceClientFactory { private final HttpClient httpClient; private final ScheduledExecutorService connectExec; + private final RequestHeaderContextConfig requestHeaderContextConfig; public ServiceClientFactoryImpl( final HttpClient httpClient, final ScheduledExecutorService connectExec ) + { + this(httpClient, connectExec, new RequestHeaderContextConfig()); + } + + public ServiceClientFactoryImpl( + final HttpClient httpClient, + final ScheduledExecutorService connectExec, + final RequestHeaderContextConfig requestHeaderContextConfig + ) { this.httpClient = httpClient; this.connectExec = connectExec; + this.requestHeaderContextConfig = requestHeaderContextConfig; } @Override @@ -47,6 +59,13 @@ public ServiceClient makeClient( final ServiceRetryPolicy retryPolicy ) { - return new ServiceClientImpl(serviceName, httpClient, serviceLocator, retryPolicy, connectExec); + return new ServiceClientImpl( + serviceName, + httpClient, + serviceLocator, + retryPolicy, + connectExec, + requestHeaderContextConfig + ); } } diff --git a/server/src/main/java/org/apache/druid/rpc/ServiceClientImpl.java b/server/src/main/java/org/apache/druid/rpc/ServiceClientImpl.java index ca7ae0371c6a..8c0916838759 100644 --- a/server/src/main/java/org/apache/druid/rpc/ServiceClientImpl.java +++ b/server/src/main/java/org/apache/druid/rpc/ServiceClientImpl.java @@ -27,6 +27,7 @@ import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.java.util.common.Either; import org.apache.druid.java.util.common.IAE; import org.apache.druid.java.util.common.StringUtils; @@ -37,6 +38,7 @@ import org.apache.druid.java.util.http.client.response.HttpResponseHandler; import org.apache.druid.java.util.http.client.response.ObjectOrErrorResponseHandler; import org.apache.druid.java.util.http.client.response.StringFullResponseHolder; +import org.apache.druid.server.audit.RequestHeaderContext; import org.jboss.netty.handler.codec.http.HttpResponseStatus; import javax.annotation.Nullable; @@ -58,6 +60,7 @@ public class ServiceClientImpl implements ServiceClient private final ServiceLocator serviceLocator; private final ServiceRetryPolicy retryPolicy; private final ScheduledExecutorService connectExec; + private final RequestHeaderContextConfig requestHeaderContextConfig; // Populated when we receive a redirect. The location here has no base path; it only identifies a preferred server. private final AtomicReference preferredLocationNoPath = new AtomicReference<>(); @@ -69,12 +72,26 @@ public ServiceClientImpl( final ServiceRetryPolicy retryPolicy, final ScheduledExecutorService connectExec ) + { + this(serviceName, httpClient, serviceLocator, retryPolicy, connectExec, new RequestHeaderContextConfig()); + } + + public ServiceClientImpl( + final String serviceName, + final HttpClient httpClient, + final ServiceLocator serviceLocator, + final ServiceRetryPolicy retryPolicy, + final ScheduledExecutorService connectExec, + final RequestHeaderContextConfig requestHeaderContextConfig + ) { this.serviceName = Preconditions.checkNotNull(serviceName, "serviceName"); this.httpClient = Preconditions.checkNotNull(httpClient, "httpClient"); this.serviceLocator = Preconditions.checkNotNull(serviceLocator, "serviceLocator"); this.retryPolicy = Preconditions.checkNotNull(retryPolicy, "retryPolicy"); this.connectExec = Preconditions.checkNotNull(connectExec, "connectExec"); + this.requestHeaderContextConfig = + Preconditions.checkNotNull(requestHeaderContextConfig, "requestHeaderContextConfig"); if (retryPolicy.maxAttempts() == 0) { throw new IAE("Invalid maxAttempts[%d] in retry policy", retryPolicy.maxAttempts()); @@ -87,6 +104,13 @@ public ListenableFuture asyncRequest( final HttpResponseHandler handler ) { + // Forward configured request headers (e.g. X-Druid-Trace-Id) captured from the inbound + // request onto this outbound inter-service call, so they propagate to downstream services on + // the normal (non-query) path. This is done HERE, on the caller's thread, where the + // RequestHeaderContextFilter's thread-local is bound — the request is actually built and sent + // later on the connectExec pool (see tryRequest), where the thread-local is not visible, so + // we stamp the RequestBuilder now and it carries the headers into that async send. + requestHeaderContextConfig.applyCapturedHeaders(RequestHeaderContext.current(), requestBuilder::header); final SettableFuture retVal = SettableFuture.create(); tryRequest(requestBuilder, handler, retVal, 0, ImmutableSet.of()); return retVal; @@ -95,7 +119,14 @@ public ListenableFuture asyncRequest( @Override public ServiceClientImpl withRetryPolicy(ServiceRetryPolicy newRetryPolicy) { - return new ServiceClientImpl(serviceName, httpClient, serviceLocator, newRetryPolicy, connectExec); + return new ServiceClientImpl( + serviceName, + httpClient, + serviceLocator, + newRetryPolicy, + connectExec, + requestHeaderContextConfig + ); } /** diff --git a/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java b/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java index 203886f99fea..9647ef97df66 100644 --- a/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java +++ b/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java @@ -22,6 +22,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Binder; import com.google.inject.Provides; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.client.broker.Broker; import org.apache.druid.client.broker.BrokerClient; import org.apache.druid.client.broker.BrokerClientImpl; @@ -47,6 +48,8 @@ import org.apache.druid.rpc.indexing.OverlordClient; import org.apache.druid.rpc.indexing.OverlordClientImpl; +import javax.annotation.Nullable; + import java.util.concurrent.ScheduledExecutorService; public class ServiceClientModule implements DruidModule @@ -63,9 +66,15 @@ public void configure(Binder binder) @Provides @LazySingleton @EscalatedGlobal - public ServiceClientFactory getServiceClientFactory(@EscalatedGlobal final HttpClient httpClient) + public ServiceClientFactory getServiceClientFactory( + @EscalatedGlobal final HttpClient httpClient, + @Nullable final RequestHeaderContextConfig requestHeaderContextConfig + ) { - return makeServiceClientFactory(httpClient); + return makeServiceClientFactory( + httpClient, + requestHeaderContextConfig != null ? requestHeaderContextConfig : new RequestHeaderContextConfig() + ); } @Provides @@ -165,9 +174,17 @@ public BrokerClient makeBrokerClient( } public static ServiceClientFactory makeServiceClientFactory(@EscalatedGlobal final HttpClient httpClient) + { + return makeServiceClientFactory(httpClient, new RequestHeaderContextConfig()); + } + + public static ServiceClientFactory makeServiceClientFactory( + @EscalatedGlobal final HttpClient httpClient, + final RequestHeaderContextConfig requestHeaderContextConfig + ) { final ScheduledExecutorService connectExec = ScheduledExecutors.fixed(CONNECT_EXEC_THREADS, "ServiceClientFactory-%d"); - return new ServiceClientFactoryImpl(httpClient, connectExec); + return new ServiceClientFactoryImpl(httpClient, connectExec, requestHeaderContextConfig); } } diff --git a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java index 9a514b85e2d5..337595f27cc2 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java @@ -23,6 +23,7 @@ import com.google.common.base.Preconditions; import com.google.common.base.Strings; import com.google.common.collect.Iterables; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.client.DirectDruidClient; import org.apache.druid.error.DruidException; import org.apache.druid.java.util.common.DateTimes; @@ -49,6 +50,7 @@ import org.apache.druid.query.QueryToolChest; import org.apache.druid.query.context.ResponseContext; import org.apache.druid.query.policy.PolicyEnforcer; +import org.apache.druid.server.audit.RequestHeaderContext; import org.apache.druid.server.broker.PerSegmentTimeoutConfig; import org.apache.druid.server.log.RequestLogger; import org.apache.druid.server.security.Action; @@ -104,6 +106,7 @@ public class QueryLifecycle private final PolicyEnforcer policyEnforcer; private final List queryBlocklist; private final Map perSegmentTimeoutConfig; + private final RequestHeaderContextConfig requestHeaderContextConfig; private final long startMs; private final long startNs; @@ -128,6 +131,7 @@ public QueryLifecycle( final PolicyEnforcer policyEnforcer, final List queryBlocklist, final Map perSegmentTimeoutConfig, + final RequestHeaderContextConfig requestHeaderContextConfig, final long startMs, final long startNs ) @@ -143,6 +147,7 @@ public QueryLifecycle( this.policyEnforcer = policyEnforcer; this.queryBlocklist = queryBlocklist; this.perSegmentTimeoutConfig = perSegmentTimeoutConfig; + this.requestHeaderContextConfig = requestHeaderContextConfig; this.startMs = startMs; this.startNs = startNs; } @@ -215,7 +220,23 @@ public void initialize(final Query baseQuery) { transition(State.NEW, State.INITIALIZED); + // Values captured from inbound request headers (see RequestHeaderContext). Computed once + // and used both for context-key authorization (below) and for anti-spoof injection (further + // down). + final Map captured = RequestHeaderContext.current(); + userContextKeys = new HashSet<>(baseQuery.getContext().keySet()); + // A body-supplied value for a header-target key is stripped below (anti-spoof) and never + // takes effect, so it must not count toward context-key authorization. + userContextKeys.removeAll(requestHeaderContextConfig.getHeaderToContextKey().values()); + // A value actually captured from an inbound header DOES take effect on the query, and the + // caller controls what headers it sends, so it is subject to the same QUERY_CONTEXT WRITE + // authorization as a value supplied in the query body (see #authorize). This prevents a + // client from setting an operational context key (priority, lane, cache flags, ...) via a + // header to bypass the authorization required for the same key in the body. Inter-node RPCs + // run as the escalated internal identity, which is authorized for all context keys, so + // cross-node propagation is unaffected. + userContextKeys.addAll(captured.keySet()); String queryId = baseQuery.getId(); if (Strings.isNullOrEmpty(queryId)) { queryId = UUID.randomUUID().toString(); @@ -227,6 +248,39 @@ public void initialize(final Query baseQuery) Map finalContext = QueryContexts.override(contextWithDefaults, baseQuery.getContext()); finalContext.put(BaseQuery.QUERY_ID, queryId); + // Anti-spoof + propagation. Reserved keys may ONLY originate from a filter-captured + // header — either a real client request (entry point) or an inter-Druid RPC where the + // upstream Druid node attached the header onto the wire (see + // DirectDruidClient.applyToOutboundRequest). Without sanitizing, a malicious client + // could supply {"context":{"traceId":"forged"}} in the query JSON body and have it + // survive into audit events and any downstream consumers of the query context. + // + // NOTE: withOverriddenContext() MERGES finalContext on top of baseQuery.getContext() + // (see QueryContexts.override) rather than replacing it. So merely removing a reserved + // key from finalContext does NOT drop a client-supplied value — the original (possibly + // forged) value still lives in baseQuery.getContext() and would survive the merge. + // Each reserved key must therefore be overridden explicitly: + // - header captured -> inject the captured value (authoritative) + // - present in body only -> override to null so the spoofed value cannot survive + // - absent everywhere -> remove, keeping the executed context clean + // + // Trade-off: if an intermediate L7 proxy strips the custom X- header between two + // Druid nodes, the propagated value is lost on the receiving node (no fallback to + // the body-context value, which was just stripped). Druid's internal RPCs are + // expected to be direct (broker→historical etc.) rather than mediated by a + // header-rewriting proxy. Operators running a mesh that strips custom X-* headers + // should add the configured headers to their mesh's allow-list. + for (String reservedKey : requestHeaderContextConfig.getHeaderToContextKey().values()) { + final String capturedValue = captured.get(reservedKey); + if (capturedValue != null) { + finalContext.put(reservedKey, capturedValue); + } else if (baseQuery.getContext().containsKey(reservedKey)) { + finalContext.put(reservedKey, null); + } else { + finalContext.remove(reservedKey); + } + } + this.baseQuery = baseQuery.withOverriddenContext(finalContext); this.toolChest = conglomerate.getToolChest(this.baseQuery); } diff --git a/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java b/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java index 57fcd1b47e91..fc0058be5a21 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java @@ -20,6 +20,7 @@ package org.apache.druid.server; import com.google.inject.Inject; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.client.BrokerViewOfBrokerConfig; import org.apache.druid.guice.LazySingleton; import org.apache.druid.java.util.emitter.service.ServiceEmitter; @@ -51,8 +52,16 @@ public class QueryLifecycleFactory private final AuthConfig authConfig; private final PolicyEnforcer policyEnforcer; private final BrokerViewOfBrokerConfig brokerViewOfBrokerConfig; + private final RequestHeaderContextConfig requestHeaderContextConfig; - @Inject + /** + * Convenience constructor for callers (chiefly tests) that don't have access to + * {@link RequestHeaderContextConfig}. Delegates to the full constructor with the default + * config — which has {@code X-Druid-Trace-Id → traceId} enabled. Tests that need the + * propagation feature fully disabled should construct an explicit + * {@code new RequestHeaderContextConfig(java.util.Collections.emptyMap())} and use the + * full constructor instead. + */ public QueryLifecycleFactory( final QueryRunnerFactoryConglomerate conglomerate, final QuerySegmentWalker texasRanger, @@ -65,6 +74,36 @@ public QueryLifecycleFactory( final QueryConfigProvider queryConfigProvider, @Nullable final BrokerViewOfBrokerConfig brokerViewOfBrokerConfig ) + { + this( + conglomerate, + texasRanger, + queryMetricsFactory, + emitter, + requestLogger, + authConfig, + policyEnforcer, + authorizerMapper, + queryConfigProvider, + brokerViewOfBrokerConfig, + null + ); + } + + @Inject + public QueryLifecycleFactory( + final QueryRunnerFactoryConglomerate conglomerate, + final QuerySegmentWalker texasRanger, + final GenericQueryMetricsFactory queryMetricsFactory, + final ServiceEmitter emitter, + final RequestLogger requestLogger, + final AuthConfig authConfig, + final PolicyEnforcer policyEnforcer, + final AuthorizerMapper authorizerMapper, + final QueryConfigProvider queryConfigProvider, + @Nullable final BrokerViewOfBrokerConfig brokerViewOfBrokerConfig, + @Nullable final RequestHeaderContextConfig requestHeaderContextConfig + ) { this.conglomerate = conglomerate; this.texasRanger = texasRanger; @@ -76,6 +115,12 @@ public QueryLifecycleFactory( this.authConfig = authConfig; this.policyEnforcer = policyEnforcer; this.brokerViewOfBrokerConfig = brokerViewOfBrokerConfig; + // Fall back to the default config when not injected (e.g. tests that pass null). + // The default has X-Druid-Trace-Id → traceId enabled; pass an explicit empty-map + // config to disable the feature. Production injection is wired via JettyServerModule + // binding `druid.audit.requestHeaders.*`. + this.requestHeaderContextConfig = + requestHeaderContextConfig != null ? requestHeaderContextConfig : new RequestHeaderContextConfig(); } public QueryLifecycle factorize() @@ -102,6 +147,7 @@ public QueryLifecycle factorize() policyEnforcer, queryBlocklist, perSegmentTimeoutConfig, + requestHeaderContextConfig, System.currentTimeMillis(), System.nanoTime() ); diff --git a/server/src/main/java/org/apache/druid/server/audit/RequestHeaderContext.java b/server/src/main/java/org/apache/druid/server/audit/RequestHeaderContext.java new file mode 100644 index 000000000000..37c7e4735de7 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/audit/RequestHeaderContext.java @@ -0,0 +1,79 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.server.audit; + +import com.google.common.collect.ImmutableMap; + +import java.util.Collections; +import java.util.Map; + +/** + * Holds the per-request map of headers captured from the inbound HTTP request, keyed by + * the operator-configured context key (see {@link RequestHeaderContextConfig}). Bound to + * the request-handling thread by a servlet filter, then read during query construction so + * the captured values can be merged into {@link org.apache.druid.query.Query}'s context. + * + *

From there, Druid's existing native sub-query context propagation carries the headers + * to historicals and peons. Consumers that already read the query context (request logger, + * audit pipeline, lineage emitters, metric dimensions) pick them up automatically. + * + *

Only valid for the lifetime of the request thread. Background tasks and async work + * spawned from the request thread will not see these values unless they explicitly copy them. + */ +public final class RequestHeaderContext +{ + private static final ThreadLocal> CURRENT = new ThreadLocal<>(); + + private RequestHeaderContext() + { + } + + /** + * Returns the captured headers for the current thread (keyed by context-key, not header + * name). Empty if none bound. Never null. + */ + public static Map current() + { + final Map map = CURRENT.get(); + return map == null ? Collections.emptyMap() : map; + } + + /** + * Binds the given map of context-key → value to the current thread. Should be paired with + * {@link #clear()} via try/finally so values don't leak across threads in a pool. + * Null or empty maps are treated as a clear. + */ + public static void bind(Map values) + { + if (values == null || values.isEmpty()) { + CURRENT.remove(); + } else { + CURRENT.set(ImmutableMap.copyOf(values)); + } + } + + /** + * Clears any captured headers bound to the current thread. + */ + public static void clear() + { + CURRENT.remove(); + } +} diff --git a/server/src/main/java/org/apache/druid/server/audit/RequestHeaderContextFilter.java b/server/src/main/java/org/apache/druid/server/audit/RequestHeaderContextFilter.java new file mode 100644 index 000000000000..8e487a5eb7e7 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/audit/RequestHeaderContextFilter.java @@ -0,0 +1,98 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.server.audit; + +import org.apache.druid.audit.RequestHeaderContextConfig; +import org.apache.druid.java.util.common.logger.Logger; + +import javax.servlet.Filter; +import javax.servlet.FilterChain; +import javax.servlet.FilterConfig; +import javax.servlet.ServletException; +import javax.servlet.ServletRequest; +import javax.servlet.ServletResponse; +import javax.servlet.http.HttpServletRequest; +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +/** + * Servlet filter that reads the headers configured via {@link RequestHeaderContextConfig} + * from each inbound HTTP request, captures any present values into a thread-local + * {@link RequestHeaderContext} keyed by the configured context-key, and clears the + * thread-local in a finally block so the values don't leak across Jetty's pooled threads. + */ +public class RequestHeaderContextFilter implements Filter +{ + private static final Logger log = new Logger(RequestHeaderContextFilter.class); + + private final RequestHeaderContextConfig config; + + public RequestHeaderContextFilter(RequestHeaderContextConfig config) + { + this.config = config; + } + + @Override + public void init(FilterConfig filterConfig) + { + // No-op. + } + + @Override + public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) + throws IOException, ServletException + { + // Always pair with clear() on the exit, even when the filter is effectively a no-op for + // this request, so that any stale value bound by other code on this pooled thread is + // wiped before the chain runs (defense in depth). + try { + if (request instanceof HttpServletRequest && !config.getHeaderToContextKey().isEmpty()) { + final HttpServletRequest httpRequest = (HttpServletRequest) request; + Map captured = null; + for (Map.Entry entry : config.getHeaderToContextKey().entrySet()) { + final String value = httpRequest.getHeader(entry.getKey()); + if (value != null && !value.isEmpty()) { + if (captured == null) { + captured = new HashMap<>(); + } + captured.put(entry.getValue(), value); + } + } + if (captured != null) { + RequestHeaderContext.bind(captured); + // Debug-level so operators can confirm header capture/propagation without log spam. + // The keys are the configured context-keys (e.g. traceId); values are caller-supplied. + log.debug("Captured request-header context %s from inbound request", captured.keySet()); + } + } + chain.doFilter(request, response); + } + finally { + RequestHeaderContext.clear(); + } + } + + @Override + public void destroy() + { + // No-op. + } +} diff --git a/server/src/main/java/org/apache/druid/server/audit/RequestHeaderContextFilterHolder.java b/server/src/main/java/org/apache/druid/server/audit/RequestHeaderContextFilterHolder.java new file mode 100644 index 000000000000..503657fd62d6 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/audit/RequestHeaderContextFilterHolder.java @@ -0,0 +1,75 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.server.audit; + +import com.google.inject.Inject; +import org.apache.druid.audit.RequestHeaderContextConfig; +import org.apache.druid.server.initialization.jetty.ServletFilterHolder; + +import javax.servlet.DispatcherType; +import javax.servlet.Filter; +import java.util.Collections; +import java.util.EnumSet; +import java.util.Map; + +/** + * Wires {@link RequestHeaderContextFilter} into the Jetty pipeline on every Druid HTTP + * endpoint. + */ +public class RequestHeaderContextFilterHolder implements ServletFilterHolder +{ + private final RequestHeaderContextConfig config; + + @Inject + public RequestHeaderContextFilterHolder(RequestHeaderContextConfig config) + { + this.config = config; + } + + @Override + public Filter getFilter() + { + return new RequestHeaderContextFilter(config); + } + + @Override + public Class getFilterClass() + { + return null; + } + + @Override + public Map getInitParameters() + { + return Collections.emptyMap(); + } + + @Override + public String getPath() + { + return "/*"; + } + + @Override + public EnumSet getDispatcherType() + { + return null; + } +} diff --git a/server/src/main/java/org/apache/druid/server/coordinator/CoordinatorConfigManager.java b/server/src/main/java/org/apache/druid/server/coordinator/CoordinatorConfigManager.java index a047bf99070d..a2dd0a2c010d 100644 --- a/server/src/main/java/org/apache/druid/server/coordinator/CoordinatorConfigManager.java +++ b/server/src/main/java/org/apache/druid/server/coordinator/CoordinatorConfigManager.java @@ -26,6 +26,7 @@ import org.apache.druid.audit.AuditEntry; import org.apache.druid.audit.AuditInfo; import org.apache.druid.audit.AuditManager; +import org.apache.druid.audit.RequestInfo; import org.apache.druid.common.config.ConfigManager; import org.apache.druid.common.config.Configs; import org.apache.druid.common.config.JacksonConfigManager; @@ -87,12 +88,18 @@ public CoordinatorDynamicConfig getCurrentDynamicConfig() return Preconditions.checkNotNull(dynamicConfig, "Got null config from watcher?!"); } - public ConfigManager.SetResult setDynamicConfig(CoordinatorDynamicConfig config, AuditInfo auditInfo) + public ConfigManager.SetResult setDynamicConfig( + CoordinatorDynamicConfig config, + AuditInfo auditInfo, + @Nullable RequestInfo requestInfo + ) { return jacksonConfigManager.set( CoordinatorDynamicConfig.CONFIG_KEY, + null, config, - auditInfo + auditInfo, + requestInfo ); } diff --git a/server/src/main/java/org/apache/druid/server/http/CoordinatorBrokerConfigsResource.java b/server/src/main/java/org/apache/druid/server/http/CoordinatorBrokerConfigsResource.java index f023b26e31c4..82edf638b736 100644 --- a/server/src/main/java/org/apache/druid/server/http/CoordinatorBrokerConfigsResource.java +++ b/server/src/main/java/org/apache/druid/server/http/CoordinatorBrokerConfigsResource.java @@ -91,8 +91,10 @@ public Response setBrokerDynamicConfig( final SetResult setResult = configManager.set( BrokerDynamicConfig.CONFIG_KEY, + null, newConfig, - AuthorizationUtils.buildAuditInfo(req) + AuthorizationUtils.buildAuditInfo(req), + AuthorizationUtils.buildRequestInfo("coordinator", req) ); if (setResult.isOk()) { diff --git a/server/src/main/java/org/apache/druid/server/http/CoordinatorDynamicConfigsResource.java b/server/src/main/java/org/apache/druid/server/http/CoordinatorDynamicConfigsResource.java index 93feb328a8c2..5364e4f5c7a7 100644 --- a/server/src/main/java/org/apache/druid/server/http/CoordinatorDynamicConfigsResource.java +++ b/server/src/main/java/org/apache/druid/server/http/CoordinatorDynamicConfigsResource.java @@ -90,7 +90,8 @@ public Response setDynamicConfigs( final SetResult setResult = manager.setDynamicConfig( dynamicConfigBuilder.build(current), - AuthorizationUtils.buildAuditInfo(req) + AuthorizationUtils.buildAuditInfo(req), + AuthorizationUtils.buildRequestInfo("coordinator", req) ); if (setResult.isOk()) { diff --git a/server/src/main/java/org/apache/druid/server/initialization/jetty/JettyServerModule.java b/server/src/main/java/org/apache/druid/server/initialization/jetty/JettyServerModule.java index 9baa832ca4ec..083a0fb6d5ef 100644 --- a/server/src/main/java/org/apache/druid/server/initialization/jetty/JettyServerModule.java +++ b/server/src/main/java/org/apache/druid/server/initialization/jetty/JettyServerModule.java @@ -131,9 +131,15 @@ protected void configureServlets() // Add empty binding for Handlers so that the injector returns an empty set if none are provided by extensions. Multibinder.newSetBinder(binder, Handler.class); Multibinder.newSetBinder(binder, JettyBindings.QosFilterHolder.class); - Multibinder.newSetBinder(binder, ServletFilterHolder.class) - .addBinding() - .to(StandardResponseHeaderFilterHolder.class); + JsonConfigProvider.bind( + binder, + "druid.audit.requestHeaders", + org.apache.druid.audit.RequestHeaderContextConfig.class + ); + final Multibinder filterMultibinder = + Multibinder.newSetBinder(binder, ServletFilterHolder.class); + filterMultibinder.addBinding().to(StandardResponseHeaderFilterHolder.class); + filterMultibinder.addBinding().to(org.apache.druid.server.audit.RequestHeaderContextFilterHolder.class); MetricsModule.register(binder, JettyMonitor.class); } diff --git a/server/src/main/java/org/apache/druid/server/security/AuthorizationUtils.java b/server/src/main/java/org/apache/druid/server/security/AuthorizationUtils.java index 9bcff9bdc1de..937a426ba9a9 100644 --- a/server/src/main/java/org/apache/druid/server/security/AuthorizationUtils.java +++ b/server/src/main/java/org/apache/druid/server/security/AuthorizationUtils.java @@ -156,6 +156,14 @@ public static AuditInfo buildAuditInfo(HttpServletRequest request) /** * Builds a RequestInfo object that can be used for auditing purposes. + * + *

The {@link RequestInfo#getRequestMetadata()} map is populated from the thread-local + * {@link org.apache.druid.server.audit.RequestHeaderContext}, which holds the values captured + * from the configured inbound request headers (see {@code RequestHeaderContextFilter} and + * {@code druid.audit.requestHeaders.headerToContextKey}), keyed by their context key (for + * example {@code traceId}). Passing the whole captured map — rather than a single hard-coded + * field — lets operators map additional headers and have them appear in audit records with no + * code change, and lets consumers extract whichever fields they need. */ public static RequestInfo buildRequestInfo(String service, HttpServletRequest request) { @@ -163,7 +171,8 @@ public static RequestInfo buildRequestInfo(String service, HttpServletRequest re service, request.getMethod(), request.getRequestURI(), - request.getQueryString() + request.getQueryString(), + org.apache.druid.server.audit.RequestHeaderContext.current() ); } diff --git a/server/src/test/java/org/apache/druid/client/BrokerServerViewTest.java b/server/src/test/java/org/apache/druid/client/BrokerServerViewTest.java index e39aabf984c7..50f40707d178 100644 --- a/server/src/test/java/org/apache/druid/client/BrokerServerViewTest.java +++ b/server/src/test/java/org/apache/druid/client/BrokerServerViewTest.java @@ -805,7 +805,8 @@ public CallbackAction segmentSchemasAnnounced(SegmentSchemas segmentSchemas) EasyMock.createMock(QueryRunnerFactoryConglomerate.class), EasyMock.createMock(QueryWatcher.class), getSmileMapper(), - EasyMock.createMock(HttpClient.class) + EasyMock.createMock(HttpClient.class), + new org.apache.druid.audit.RequestHeaderContextConfig() ); brokerServerView = new BrokerServerView( diff --git a/server/src/test/java/org/apache/druid/client/SimpleServerView.java b/server/src/test/java/org/apache/druid/client/SimpleServerView.java index a90c2b0ff105..32b2b919cc22 100644 --- a/server/src/test/java/org/apache/druid/client/SimpleServerView.java +++ b/server/src/test/java/org/apache/druid/client/SimpleServerView.java @@ -75,7 +75,8 @@ public SimpleServerView( conglomerate, NOOP_QUERY_WATCHER, objectMapper, - httpClient + httpClient, + new org.apache.druid.audit.RequestHeaderContextConfig() ); } diff --git a/server/src/test/java/org/apache/druid/rpc/ServiceClientImplTest.java b/server/src/test/java/org/apache/druid/rpc/ServiceClientImplTest.java index 0b2bb7fe2640..9240cbf4fb96 100644 --- a/server/src/test/java/org/apache/druid/rpc/ServiceClientImplTest.java +++ b/server/src/test/java/org/apache/druid/rpc/ServiceClientImplTest.java @@ -24,6 +24,7 @@ import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; +import org.apache.druid.audit.AuditManager; import org.apache.druid.java.util.common.Either; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.StringUtils; @@ -32,6 +33,7 @@ import org.apache.druid.java.util.http.client.Request; import org.apache.druid.java.util.http.client.response.ObjectOrErrorResponseHandler; import org.apache.druid.java.util.http.client.response.StringFullResponseHolder; +import org.apache.druid.server.audit.RequestHeaderContext; import org.hamcrest.CoreMatchers; import org.hamcrest.MatcherAssert; import org.jboss.netty.buffer.ChannelBuffers; @@ -45,6 +47,7 @@ import org.junit.Rule; import org.junit.Test; import org.junit.internal.matchers.ThrowableMessageMatcher; +import org.mockito.ArgumentCaptor; import org.mockito.ArgumentMatchers; import org.mockito.Mock; import org.mockito.Mockito; @@ -118,6 +121,42 @@ public void test_request_ok() throws Exception Assert.assertEquals(expectedResponseObject, response); } + @Test + public void test_request_forwardsCapturedRequestHeader() throws Exception + { + // A header captured from the inbound request (here simulated via the thread-local that + // RequestHeaderContextFilter would bind) must be forwarded onto this outbound inter-service + // call, so it propagates to downstream services on the normal (non-query) path. The forwarding + // happens on this caller thread; the request is later sent on connectExec, so the value must be + // stamped onto the RequestBuilder now. + final RequestBuilder requestBuilder = new RequestBuilder(HttpMethod.GET, "/foo"); + final ImmutableMap expectedResponseObject = ImmutableMap.of("foo", "bar"); + + stubLocatorCall(locations(SERVER1)); + expectHttpCall(requestBuilder, SERVER1).thenReturn(valueResponse(expectedResponseObject)); + + serviceClient = makeServiceClient(StandardRetryPolicy.noRetries()); + + RequestHeaderContext.bind(ImmutableMap.of("traceId", "trace-abc")); + try { + doRequest(serviceClient, requestBuilder); + } + finally { + RequestHeaderContext.clear(); + } + + final ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(Request.class); + Mockito.verify(httpClient).go( + requestCaptor.capture(), + ArgumentMatchers.any(ObjectOrErrorResponseHandler.class), + ArgumentMatchers.eq(RequestBuilder.DEFAULT_TIMEOUT) + ); + MatcherAssert.assertThat( + requestCaptor.getValue().getHeaders().get(AuditManager.X_DRUID_TRACE_ID), + CoreMatchers.hasItem("trace-abc") + ); + } + @Test public void test_request_serverError() { diff --git a/server/src/test/java/org/apache/druid/server/PerSegmentTimeoutInjectionTest.java b/server/src/test/java/org/apache/druid/server/PerSegmentTimeoutInjectionTest.java index 440f065ba485..e4800327e17a 100644 --- a/server/src/test/java/org/apache/druid/server/PerSegmentTimeoutInjectionTest.java +++ b/server/src/test/java/org/apache/druid/server/PerSegmentTimeoutInjectionTest.java @@ -216,6 +216,7 @@ private QueryLifecycle createLifecycle(Map perS NoopPolicyEnforcer.instance(), Collections.emptyList(), perSegmentTimeoutConfig, + new org.apache.druid.audit.RequestHeaderContextConfig(), System.currentTimeMillis(), System.nanoTime() ); diff --git a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java index 5f8da0d21a39..520f88963045 100644 --- a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java @@ -59,6 +59,7 @@ import org.apache.druid.query.policy.RestrictAllTablesPolicyEnforcer; import org.apache.druid.query.policy.RowFilterPolicy; import org.apache.druid.query.timeseries.TimeseriesQuery; +import org.apache.druid.server.audit.RequestHeaderContext; import org.apache.druid.server.log.RequestLogger; import org.apache.druid.server.security.Access; import org.apache.druid.server.security.Action; @@ -861,6 +862,7 @@ public void testRunSimple_queryBlocklisted() policyEnforcer, queryBlocklist, Collections.emptyMap(), + new org.apache.druid.audit.RequestHeaderContextConfig(), System.currentTimeMillis(), System.nanoTime() ); @@ -913,6 +915,7 @@ public void testRunSimple_queryNotBlocklisted() policyEnforcer, queryBlocklist, Collections.emptyMap(), + new org.apache.druid.audit.RequestHeaderContextConfig(), System.currentTimeMillis(), System.nanoTime() ); @@ -921,6 +924,132 @@ public void testRunSimple_queryNotBlocklisted() lifecycle.runSimple(query, authenticationResult, AuthorizationResult.ALLOW_NO_RESTRICTION); } + /** + * Anti-spoof regression test: a client supplies a value for a reserved context key + * ("traceId") in the query body but sends NO corresponding header. The forged value + * must NOT survive into the executed query context. This guards against the + * withOverriddenContext() merge re-introducing the body value after the strip. + */ + @Test + public void testInitialize_forgedReservedContextKeyStrippedWhenNoHeader() + { + EasyMock.expect(queryConfig.getContext()).andReturn(ImmutableMap.of()).anyTimes(); + EasyMock.expect(conglomerate.getToolChest(EasyMock.anyObject())).andReturn(toolChest).anyTimes(); + replayAll(); + + final TimeseriesQuery forgedQuery = Druids.newTimeseriesQueryBuilder() + .dataSource(DATASOURCE) + .intervals(ImmutableList.of(Intervals.ETERNITY)) + .aggregators(new CountAggregatorFactory("chocula")) + .context(ImmutableMap.of("traceId", "FORGED", "foo", "bar")) + .build(); + + RequestHeaderContext.clear(); + try { + QueryLifecycle lifecycle = createLifecycle(); + lifecycle.initialize(forgedQuery); + + final Map ctx = lifecycle.getQuery().getContext(); + Assert.assertNull("forged traceId must not survive into the query context", ctx.get("traceId")); + // Non-reserved user context is untouched. + Assert.assertEquals("bar", ctx.get("foo")); + } + finally { + RequestHeaderContext.clear(); + } + } + + /** + * A captured header value is authoritative: even when the body carries a forged value + * for the same reserved key, the filter-captured value wins. + */ + @Test + public void testInitialize_capturedHeaderOverridesForgedBody() + { + EasyMock.expect(queryConfig.getContext()).andReturn(ImmutableMap.of()).anyTimes(); + EasyMock.expect(conglomerate.getToolChest(EasyMock.anyObject())).andReturn(toolChest).anyTimes(); + replayAll(); + + final TimeseriesQuery forgedQuery = Druids.newTimeseriesQueryBuilder() + .dataSource(DATASOURCE) + .intervals(ImmutableList.of(Intervals.ETERNITY)) + .aggregators(new CountAggregatorFactory("chocula")) + .context(ImmutableMap.of("traceId", "FORGED")) + .build(); + + RequestHeaderContext.bind(ImmutableMap.of("traceId", "real-trace-123")); + try { + QueryLifecycle lifecycle = createLifecycle(); + lifecycle.initialize(forgedQuery); + Assert.assertEquals( + "captured header value must override forged body value", + "real-trace-123", + lifecycle.getQuery().getContext().get("traceId") + ); + } + finally { + RequestHeaderContext.clear(); + } + } + + /** + * A context key whose value is captured from an inbound request header is client-influenced + * (a caller controls the headers it sends), so it must be subject to QUERY_CONTEXT WRITE + * authorization just like a body-supplied key. This guards against a client bypassing + * context-key authorization by sending an operational key's value as a header instead of in + * the query body. + */ + @Test + public void testAuthorizeQueryContext_capturedHeaderKeyRequiresAuthorization() + { + EasyMock.expect(queryConfig.getContext()).andReturn(ImmutableMap.of()).anyTimes(); + EasyMock.expect(authenticationResult.getIdentity()).andReturn(IDENTITY).anyTimes(); + EasyMock.expect(authenticationResult.getAuthorizerName()).andReturn(AUTHORIZER).anyTimes(); + EasyMock.expect(authorizer.authorize( + authenticationResult, + new Resource(DATASOURCE, ResourceType.DATASOURCE), + Action.READ + )) + .andReturn(Access.OK) + .times(2); + // "traceId" is supplied only via the captured header (not the body) yet must still be + // authorized as a QUERY_CONTEXT write. + EasyMock.expect(authorizer.authorize( + authenticationResult, + new Resource("traceId", ResourceType.QUERY_CONTEXT), + Action.WRITE + )) + .andReturn(Access.DENIED) + .times(2); + EasyMock.expect(conglomerate.getToolChest(EasyMock.anyObject())) + .andReturn(toolChest) + .times(2); + + replayAll(); + + // Note: no context keys in the query body; "traceId" arrives only via the captured header. + final TimeseriesQuery query = Druids.newTimeseriesQueryBuilder() + .dataSource(DATASOURCE) + .intervals(ImmutableList.of(Intervals.ETERNITY)) + .aggregators(new CountAggregatorFactory("chocula")) + .build(); + + authConfig = AuthConfig.newBuilder().setAuthorizeQueryContextParams(true).build(); + RequestHeaderContext.bind(ImmutableMap.of("traceId", "client-set-trace")); + try { + QueryLifecycle lifecycle = createLifecycle(); + lifecycle.initialize(query); + Assert.assertFalse(lifecycle.authorize(mockRequest()).allowBasicAccess()); + + lifecycle = createLifecycle(); + lifecycle.initialize(query); + Assert.assertFalse(lifecycle.authorize(authenticationResult).allowBasicAccess()); + } + finally { + RequestHeaderContext.clear(); + } + } + private HttpServletRequest mockRequest() { HttpServletRequest request = EasyMock.createNiceMock(HttpServletRequest.class); diff --git a/server/src/test/java/org/apache/druid/server/QueryResourceTest.java b/server/src/test/java/org/apache/druid/server/QueryResourceTest.java index d7a611bfc320..615a79b4f88e 100644 --- a/server/src/test/java/org/apache/druid/server/QueryResourceTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryResourceTest.java @@ -777,6 +777,7 @@ public QueryLifecycle factorize() NoopPolicyEnforcer.instance(), null, Collections.emptyMap(), + new org.apache.druid.audit.RequestHeaderContextConfig(), System.currentTimeMillis(), System.nanoTime() ) diff --git a/server/src/test/java/org/apache/druid/server/audit/RequestHeaderContextFilterTest.java b/server/src/test/java/org/apache/druid/server/audit/RequestHeaderContextFilterTest.java new file mode 100644 index 000000000000..e440e5e3e2b9 --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/audit/RequestHeaderContextFilterTest.java @@ -0,0 +1,163 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.server.audit; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.audit.AuditManager; +import org.apache.druid.audit.RequestHeaderContextConfig; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.junit.After; +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +import javax.servlet.FilterChain; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; +import java.io.IOException; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +public class RequestHeaderContextFilterTest +{ + private final ObjectMapper mapper = new DefaultObjectMapper(); + + @After + public void clearTraceId() + { + RequestHeaderContext.clear(); + } + + @Test + public void testHeaderCapturedBoundForChainAndClearedAfter() throws Exception + { + RequestHeaderContextFilter filter = new RequestHeaderContextFilter(new RequestHeaderContextConfig()); + + HttpServletRequest request = Mockito.mock(HttpServletRequest.class); + HttpServletResponse response = Mockito.mock(HttpServletResponse.class); + Mockito.when(request.getHeader(AuditManager.X_DRUID_TRACE_ID)).thenReturn("trace-xyz"); + + AtomicReference> observed = new AtomicReference<>(); + FilterChain chain = (req, resp) -> observed.set(RequestHeaderContext.current()); + + filter.doFilter(request, response, chain); + + Assert.assertEquals("trace-xyz", observed.get().get("traceId")); + Assert.assertTrue( + "thread-local must be cleared after the filter chain", + RequestHeaderContext.current().isEmpty() + ); + } + + @Test + public void testMultipleConfiguredHeaders() throws Exception + { + String json = "{\"headerToContextKey\":{\"X-Druid-Trace-Id\":\"traceId\",\"X-Custom-Foo\":\"foo\"}}"; + RequestHeaderContextConfig config = mapper.readValue(json, RequestHeaderContextConfig.class); + RequestHeaderContextFilter filter = new RequestHeaderContextFilter(config); + + HttpServletRequest request = Mockito.mock(HttpServletRequest.class); + HttpServletResponse response = Mockito.mock(HttpServletResponse.class); + Mockito.when(request.getHeader(AuditManager.X_DRUID_TRACE_ID)).thenReturn("trace-1"); + Mockito.when(request.getHeader("X-Custom-Foo")).thenReturn("foo-1"); + + AtomicReference> observed = new AtomicReference<>(); + FilterChain chain = (req, resp) -> observed.set(RequestHeaderContext.current()); + + filter.doFilter(request, response, chain); + + Assert.assertEquals("trace-1", observed.get().get("traceId")); + Assert.assertEquals("foo-1", observed.get().get("foo")); + } + + @Test + public void testNoHeadersPresentIsNoOp() throws Exception + { + RequestHeaderContextFilter filter = new RequestHeaderContextFilter(new RequestHeaderContextConfig()); + + HttpServletRequest request = Mockito.mock(HttpServletRequest.class); + HttpServletResponse response = Mockito.mock(HttpServletResponse.class); + Mockito.when(request.getHeader(AuditManager.X_DRUID_TRACE_ID)).thenReturn(null); + + AtomicReference> observed = new AtomicReference<>(); + FilterChain chain = (req, resp) -> observed.set(RequestHeaderContext.current()); + + filter.doFilter(request, response, chain); + + Assert.assertTrue(observed.get().isEmpty()); + Assert.assertTrue(RequestHeaderContext.current().isEmpty()); + } + + @Test + public void testEmptyHeaderValueTreatedAsAbsent() throws Exception + { + RequestHeaderContextFilter filter = new RequestHeaderContextFilter(new RequestHeaderContextConfig()); + + HttpServletRequest request = Mockito.mock(HttpServletRequest.class); + HttpServletResponse response = Mockito.mock(HttpServletResponse.class); + Mockito.when(request.getHeader(AuditManager.X_DRUID_TRACE_ID)).thenReturn(""); + + AtomicReference> observed = new AtomicReference<>(); + FilterChain chain = (req, resp) -> observed.set(RequestHeaderContext.current()); + + filter.doFilter(request, response, chain); + + Assert.assertTrue(observed.get().isEmpty()); + } + + @Test + public void testClearedEvenIfChainThrows() throws IOException + { + RequestHeaderContextFilter filter = new RequestHeaderContextFilter(new RequestHeaderContextConfig()); + + HttpServletRequest request = Mockito.mock(HttpServletRequest.class); + HttpServletResponse response = Mockito.mock(HttpServletResponse.class); + Mockito.when(request.getHeader(AuditManager.X_DRUID_TRACE_ID)).thenReturn("trace-throws"); + + FilterChain chain = (req, resp) -> { + throw new RuntimeException("boom"); + }; + + Assert.assertThrows(RuntimeException.class, () -> filter.doFilter(request, response, chain)); + Assert.assertTrue( + "thread-local must be cleared even on exception", + RequestHeaderContext.current().isEmpty() + ); + } + + @Test + public void testEmptyConfigIsNoOp() throws Exception + { + RequestHeaderContextConfig empty = mapper.readValue("{\"headerToContextKey\":{}}", RequestHeaderContextConfig.class); + RequestHeaderContextFilter filter = new RequestHeaderContextFilter(empty); + + HttpServletRequest request = Mockito.mock(HttpServletRequest.class); + HttpServletResponse response = Mockito.mock(HttpServletResponse.class); + // Even if a header is present, an empty config shouldn't capture anything. + Mockito.when(request.getHeader(AuditManager.X_DRUID_TRACE_ID)).thenReturn("trace-1"); + + AtomicReference> observed = new AtomicReference<>(); + FilterChain chain = (req, resp) -> observed.set(RequestHeaderContext.current()); + + filter.doFilter(request, response, chain); + + Assert.assertTrue(observed.get().isEmpty()); + } +} diff --git a/server/src/test/java/org/apache/druid/server/audit/RequestHeaderContextTest.java b/server/src/test/java/org/apache/druid/server/audit/RequestHeaderContextTest.java new file mode 100644 index 000000000000..4ce72300f046 --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/audit/RequestHeaderContextTest.java @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.server.audit; + +import com.google.common.collect.ImmutableMap; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.Collections; + +public class RequestHeaderContextTest +{ + @AfterEach + public void tearDown() + { + RequestHeaderContext.clear(); + } + + @Test + public void testCurrentEmptyByDefault() + { + Assertions.assertTrue(RequestHeaderContext.current().isEmpty()); + } + + @Test + public void testBindAndCurrent() + { + RequestHeaderContext.bind(ImmutableMap.of("traceId", "abc-123", "foo", "bar")); + Assertions.assertEquals(2, RequestHeaderContext.current().size()); + Assertions.assertEquals("abc-123", RequestHeaderContext.current().get("traceId")); + Assertions.assertEquals("bar", RequestHeaderContext.current().get("foo")); + } + + @Test + public void testBindEmptyTreatedAsClear() + { + RequestHeaderContext.bind(ImmutableMap.of("traceId", "abc")); + RequestHeaderContext.bind(Collections.emptyMap()); + Assertions.assertTrue(RequestHeaderContext.current().isEmpty()); + } + + @Test + public void testBindNullTreatedAsClear() + { + RequestHeaderContext.bind(ImmutableMap.of("traceId", "abc")); + RequestHeaderContext.bind(null); + Assertions.assertTrue(RequestHeaderContext.current().isEmpty()); + } + + @Test + public void testClearRemoves() + { + RequestHeaderContext.bind(ImmutableMap.of("traceId", "abc")); + RequestHeaderContext.clear(); + Assertions.assertTrue(RequestHeaderContext.current().isEmpty()); + } + + @Test + public void testCurrentReturnsImmutableSnapshot() + { + RequestHeaderContext.bind(ImmutableMap.of("traceId", "abc")); + Assertions.assertThrows( + UnsupportedOperationException.class, + () -> RequestHeaderContext.current().put("evil", "value") + ); + } +} diff --git a/server/src/test/java/org/apache/druid/server/http/CoordinatorBrokerConfigsResourceTest.java b/server/src/test/java/org/apache/druid/server/http/CoordinatorBrokerConfigsResourceTest.java index 543f2b7bb9bb..280d60dfcbb8 100644 --- a/server/src/test/java/org/apache/druid/server/http/CoordinatorBrokerConfigsResourceTest.java +++ b/server/src/test/java/org/apache/druid/server/http/CoordinatorBrokerConfigsResourceTest.java @@ -22,6 +22,7 @@ import com.google.common.collect.ImmutableList; import org.apache.druid.audit.AuditInfo; import org.apache.druid.audit.AuditManager; +import org.apache.druid.audit.RequestInfo; import org.apache.druid.common.config.ConfigManager.SetResult; import org.apache.druid.common.config.JacksonConfigManager; import org.apache.druid.server.broker.BrokerDynamicConfig; @@ -166,8 +167,10 @@ public void testSetBrokerDynamicConfig() EasyMock.expect( configManager.set( EasyMock.anyObject(String.class), + EasyMock.anyObject(byte[].class), EasyMock.anyObject(BrokerDynamicConfig.class), - EasyMock.anyObject(AuditInfo.class) + EasyMock.anyObject(AuditInfo.class), + EasyMock.anyObject(RequestInfo.class) ) ).andReturn(SetResult.ok()).once(); diff --git a/server/src/test/java/org/apache/druid/server/security/AuthorizationUtilsTest.java b/server/src/test/java/org/apache/druid/server/security/AuthorizationUtilsTest.java index 50e5fce688ed..ae2e4cbf60c4 100644 --- a/server/src/test/java/org/apache/druid/server/security/AuthorizationUtilsTest.java +++ b/server/src/test/java/org/apache/druid/server/security/AuthorizationUtilsTest.java @@ -20,6 +20,7 @@ package org.apache.druid.server.security; import com.google.common.base.Function; +import org.apache.druid.audit.RequestInfo; import org.apache.druid.error.DruidException; import org.apache.druid.query.filter.EqualityFilter; import org.apache.druid.query.policy.NoRestrictionPolicy; @@ -243,4 +244,57 @@ public void testAuthorizeAllResourceActions_policyForNonReadDatasourceThrows() ); Assert.assertTrue(exception.getMessage().contains("Policy should only present when reading a table")); } + + @Test + public void test_buildRequestInfo_capturesMetadataFromFilterThreadLocal() + { + // RequestInfo.requestMetadata is sourced from the filter's thread-local + // (RequestHeaderContext) so operator header remapping is honored. + MockHttpServletRequest request = newRequest("GET", "/druid/coordinator/v1/datasources"); + org.apache.druid.server.audit.RequestHeaderContext.bind( + java.util.Collections.singletonMap("traceId", "trace-abc-123") + ); + try { + RequestInfo info = AuthorizationUtils.buildRequestInfo("coordinator", request); + + Assert.assertEquals("coordinator", info.getService()); + Assert.assertEquals("GET", info.getMethod()); + Assert.assertEquals("/druid/coordinator/v1/datasources", info.getUri()); + Assert.assertEquals("trace-abc-123", info.getRequestMetadata().get("traceId")); + } + finally { + org.apache.druid.server.audit.RequestHeaderContext.clear(); + } + } + + @Test + public void test_buildRequestInfo_noMetadataWhenContextEmpty() + { + // No header captured, no thread-local bound -> requestMetadata is null. + MockHttpServletRequest request = newRequest("POST", "/druid/indexer/v1/task"); + + RequestInfo info = AuthorizationUtils.buildRequestInfo("overlord", request); + + Assert.assertEquals("overlord", info.getService()); + Assert.assertNull(info.getRequestMetadata()); + } + + /** + * {@link MockHttpServletRequest#getQueryString()} throws by default; override to return null + * since these tests don't exercise query parameters. + */ + private static MockHttpServletRequest newRequest(String method, String uri) + { + MockHttpServletRequest request = new MockHttpServletRequest() + { + @Override + public String getQueryString() + { + return null; + } + }; + request.method = method; + request.requestUri = uri; + return request; + } } diff --git a/sql/src/main/java/org/apache/druid/sql/AbstractStatement.java b/sql/src/main/java/org/apache/druid/sql/AbstractStatement.java index 35b9883b2594..063cce4f7c65 100644 --- a/sql/src/main/java/org/apache/druid/sql/AbstractStatement.java +++ b/sql/src/main/java/org/apache/druid/sql/AbstractStatement.java @@ -21,6 +21,7 @@ import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.query.QueryContexts; +import org.apache.druid.server.audit.RequestHeaderContext; import org.apache.druid.server.security.Action; import org.apache.druid.server.security.AuthorizationResult; import org.apache.druid.server.security.AuthorizationUtils; @@ -87,9 +88,41 @@ public AbstractStatement( this.sqlToolbox = sqlToolbox; this.reporter = new SqlExecutionReporter(this, remoteAddress); this.queryPlus = queryPlus; - this.authContextKeys = queryPlus.authContextKeys(); this.queryContext = new HashMap<>(queryPlus.context()); sqlToolbox.engine.initContextMap(this.queryContext); + // Anti-spoof + propagation for reserved request-header context keys, mirroring + // QueryLifecycle.initialize() on the native path. A client must not be able to set a + // value for one of these keys via the SQL body context; the only legitimate source is + // a filter-captured inbound header (RequestHeaderContext). Unlike the native path, + // queryContext here is a fresh mutable map that IS the authoritative context (it is not + // re-merged over the request body), so we can simply remove uncaptured keys rather than + // overriding them to null. This sanitizes both the SQL request log and the context handed + // to the planner; generated native sub-queries are additionally sanitized by + // QueryLifecycle.initialize(). + final Map capturedHeaders = RequestHeaderContext.current(); + for (String reservedKey : sqlToolbox.requestHeaderContextConfig.getHeaderToContextKey().values()) { + final String capturedValue = capturedHeaders.get(reservedKey); + if (capturedValue != null) { + this.queryContext.put(reservedKey, capturedValue); + } else { + this.queryContext.remove(reservedKey); + } + } + // A value captured from an inbound header is client-influenced (the caller controls the + // headers it sends), so it must be authorized like a body-supplied context key. The body's + // authContextKeys were computed before the captured values were injected above, so union + // the captured keys in here, mirroring QueryLifecycle.initialize() on the native path. This + // stops a client from mapping a header to an operational key (priority, lane, cache flags, + // ...) to bypass the QUERY_CONTEXT WRITE authorization required for that key in the body. + final Set keysToAuthorize = new HashSet<>(queryPlus.authContextKeys()); + // A body-supplied value for a header-target key was just stripped above and never takes + // effect, so it must not count toward context-key authorization (mirrors the removeAll on + // QueryLifecycle.initialize()). Without this, a body context key that collides with a + // configured header-target key (e.g. traceId) would still require QUERY_CONTEXT WRITE and + // could falsely reject the query when authorizeQueryContextParams is enabled. + keysToAuthorize.removeAll(sqlToolbox.requestHeaderContextConfig.getHeaderToContextKey().values()); + keysToAuthorize.addAll(capturedHeaders.keySet()); + this.authContextKeys = Set.copyOf(keysToAuthorize); // "bySegment" results are never valid to use with SQL because the result format is incompatible // so, overwrite any user specified context to avoid exceptions down the line if (this.queryContext.remove(QueryContexts.BY_SEGMENT_KEY) != null) { diff --git a/sql/src/main/java/org/apache/druid/sql/SqlToolbox.java b/sql/src/main/java/org/apache/druid/sql/SqlToolbox.java index 66e55db3caad..c8ed2ce462ef 100644 --- a/sql/src/main/java/org/apache/druid/sql/SqlToolbox.java +++ b/sql/src/main/java/org/apache/druid/sql/SqlToolbox.java @@ -20,6 +20,7 @@ package org.apache.druid.sql; import com.google.common.base.Preconditions; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.server.QueryScheduler; import org.apache.druid.server.log.RequestLogger; @@ -37,7 +38,17 @@ public class SqlToolbox final RequestLogger requestLogger; final QueryScheduler queryScheduler; final SqlLifecycleManager sqlLifecycleManager; + /** + * Drives anti-spoof sanitization and propagation of reserved request-header context keys + * on the SQL path (see {@link AbstractStatement}), mirroring the native-query handling in + * {@code QueryLifecycle.initialize()}. + */ + final RequestHeaderContextConfig requestHeaderContextConfig; + /** + * Convenience constructor that uses the default {@link RequestHeaderContextConfig}. Used + * by tests; production wiring goes through the full constructor so operator config is honored. + */ public SqlToolbox( final SqlEngine engine, final PlannerFactory plannerFactory, @@ -46,6 +57,20 @@ public SqlToolbox( final QueryScheduler queryScheduler, final SqlLifecycleManager sqlLifecycleManager ) + { + this(engine, plannerFactory, emitter, requestLogger, queryScheduler, sqlLifecycleManager, + new RequestHeaderContextConfig()); + } + + public SqlToolbox( + final SqlEngine engine, + final PlannerFactory plannerFactory, + final ServiceEmitter emitter, + final RequestLogger requestLogger, + final QueryScheduler queryScheduler, + final SqlLifecycleManager sqlLifecycleManager, + final RequestHeaderContextConfig requestHeaderContextConfig + ) { this.engine = engine; this.plannerFactory = plannerFactory; @@ -53,6 +78,10 @@ public SqlToolbox( this.requestLogger = requestLogger; this.queryScheduler = queryScheduler; this.sqlLifecycleManager = Preconditions.checkNotNull(sqlLifecycleManager, "sqlLifecycleManager"); + this.requestHeaderContextConfig = Preconditions.checkNotNull( + requestHeaderContextConfig, + "requestHeaderContextConfig" + ); } public SqlToolbox withEngine(final SqlEngine engine) @@ -63,7 +92,8 @@ public SqlToolbox withEngine(final SqlEngine engine) emitter, requestLogger, queryScheduler, - sqlLifecycleManager + sqlLifecycleManager, + requestHeaderContextConfig ); } } diff --git a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java index 1565527f601b..7287f0bd04e0 100644 --- a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java +++ b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java @@ -26,6 +26,7 @@ import com.google.inject.Module; import com.google.inject.Provides; import com.google.inject.multibindings.Multibinder; +import org.apache.druid.audit.RequestHeaderContextConfig; import org.apache.druid.catalog.model.TableDefnRegistry; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.PolyBind; @@ -52,6 +53,8 @@ import org.apache.druid.sql.calcite.view.ViewManager; import org.apache.druid.sql.http.SqlHttpModule; +import javax.annotation.Nullable; + import java.util.Properties; public class SqlModule implements Module @@ -170,7 +173,8 @@ public SqlToolbox makeSqlToolbox( final ServiceEmitter emitter, final RequestLogger requestLogger, final QueryScheduler queryScheduler, - final SqlLifecycleManager sqlLifecycleManager + final SqlLifecycleManager sqlLifecycleManager, + @Nullable final RequestHeaderContextConfig requestHeaderContextConfig ) { return new SqlToolbox( @@ -179,7 +183,8 @@ public SqlToolbox makeSqlToolbox( emitter, requestLogger, queryScheduler, - sqlLifecycleManager + sqlLifecycleManager, + requestHeaderContextConfig != null ? requestHeaderContextConfig : new RequestHeaderContextConfig() ); } diff --git a/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java b/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java index b7152fd9d33b..3a0be00b2c5f 100644 --- a/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java +++ b/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java @@ -45,6 +45,7 @@ import org.apache.druid.server.QueryScheduler; import org.apache.druid.server.QueryStackTests; import org.apache.druid.server.SpecificSegmentsQuerySegmentWalker; +import org.apache.druid.server.audit.RequestHeaderContext; import org.apache.druid.server.initialization.ServerConfig; import org.apache.druid.server.log.TestRequestLogger; import org.apache.druid.server.metrics.NoopServiceEmitter; @@ -84,6 +85,7 @@ import static org.apache.druid.sql.calcite.BaseCalciteQueryTest.assertResultsEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; @@ -206,6 +208,97 @@ public void testDirectHappyPath() stmt.close(); } + @Test + public void testDirectStripsForgedReservedContextKeyWhenNoHeader() + { + // Anti-spoof: a client supplies a value for the reserved "traceId" context key in the + // SQL body but no header was captured. The forged value must be stripped from the + // statement context (which is what gets logged and handed to the planner). + SqlQueryPlus sqlReq = SqlQueryPlus.builder("SELECT COUNT(*) AS cnt FROM druid.foo") + .queryContext(ImmutableMap.of("traceId", "FORGED", "foo", "bar")) + .auth(CalciteTests.REGULAR_USER_AUTH_RESULT) + .build(); + RequestHeaderContext.clear(); + try { + DirectStatement stmt = sqlStatementFactory.directStatement(sqlReq); + assertNull("forged traceId must be stripped from SQL context", stmt.context().get("traceId")); + assertEquals("bar", stmt.context().get("foo")); + stmt.close(); + } + finally { + RequestHeaderContext.clear(); + } + } + + @Test + public void testDirectInjectsCapturedHeaderOverridingForgedBody() + { + SqlQueryPlus sqlReq = SqlQueryPlus.builder("SELECT COUNT(*) AS cnt FROM druid.foo") + .queryContext(ImmutableMap.of("traceId", "FORGED")) + .auth(CalciteTests.REGULAR_USER_AUTH_RESULT) + .build(); + RequestHeaderContext.bind(ImmutableMap.of("traceId", "real-trace")); + try { + DirectStatement stmt = sqlStatementFactory.directStatement(sqlReq); + assertEquals("real-trace", stmt.context().get("traceId")); + stmt.close(); + } + finally { + RequestHeaderContext.clear(); + } + } + + @Test + public void testDirectAuthorizesCapturedHeaderContextKey() + { + // A context key whose value is captured from an inbound header is client-influenced, so it + // must be subject to QUERY_CONTEXT authorization like a body key. authContextKeys is frozen + // from the body before header injection, so the captured key must be unioned in (mirrors + // QueryLifecycle on the native path). Without it, mapping a header to e.g. priority/lane + // would reach planning without the WRITE authorization required for the same body key. + SqlQueryPlus sqlReq = SqlQueryPlus.builder("SELECT COUNT(*) AS cnt FROM druid.foo") + .auth(CalciteTests.REGULAR_USER_AUTH_RESULT) + .build(); + RequestHeaderContext.bind(ImmutableMap.of("traceId", "client-set")); + try { + DirectStatement stmt = sqlStatementFactory.directStatement(sqlReq); + // "traceId" arrived only via the captured header (not the body) yet must be authorized. + assertTrue(stmt.authContextKeys.contains("traceId")); + stmt.close(); + } + finally { + RequestHeaderContext.clear(); + } + } + + @Test + public void testDirectDoesNotAuthorizeStrippedBodyReservedContextKey() + { + // A body-supplied value for a header-target key ("traceId") with no captured header is + // stripped from the executed context (anti-spoof) and never takes effect, so it must NOT + // count toward context-key authorization. Mirrors QueryLifecycle.initialize() on the native + // path; without the removeAll, enabling authorizeQueryContextParams would falsely reject a + // query carrying traceId in its body even though the value is discarded. + SqlQueryPlus sqlReq = SqlQueryPlus.builder("SELECT COUNT(*) AS cnt FROM druid.foo") + .queryContext(ImmutableMap.of("traceId", "FORGED", "foo", "bar")) + .auth(CalciteTests.REGULAR_USER_AUTH_RESULT) + .build(); + RequestHeaderContext.clear(); + try { + DirectStatement stmt = sqlStatementFactory.directStatement(sqlReq); + assertFalse( + "stripped body traceId must not require context-key authorization", + stmt.authContextKeys.contains("traceId") + ); + // A genuine (non header-target) body context key still requires authorization. + assertTrue(stmt.authContextKeys.contains("foo")); + stmt.close(); + } + finally { + RequestHeaderContext.clear(); + } + } + @Test public void testDirectPlanTwice() {