map = Maps.newMapWithExpectedSize(headers.size());
for (String header : headers) {
final String value = context.requestHeaders.get(header);
if (value != null) {
map.put(header, value);
}
}
return map;
}
/**
* Captures the current thread context as writeable, allowing it to be serialized out later
*/
public Writeable captureAsWriteable() {
final ThreadContextStruct context = threadLocal.get();
return out -> context.writeTo(out, defaultHeader);
}
/**
* Removes the current context and resets a default context marked with as
* originating from the supplied string. The removed context can be
* restored by closing the returned {@link StoredContext}. Callers should
* be careful to save the current context before calling this method and
* restore it any listeners, likely with
* {@link ContextPreservingActionListener}. Use {@link OriginSettingClient}
* which can be used to do this automatically.
*
* Without security the origin is ignored, but security uses it to authorize
* actions that are made up of many sub-actions. These actions call
* {@link #stashWithOrigin} before performing on behalf of a user that
* should be allowed even if the user doesn't have permission to perform
* those actions on their own.
*
* For example, a user might not have permission to GET from the tasks index
* but the tasks API will perform a get on their behalf using this method
* if it can't find the task in memory.
*/
public StoredContext stashWithOrigin(String origin) {
final ThreadContext.StoredContext storedContext = stashContext();
putTransient(ACTION_ORIGIN_TRANSIENT_NAME, origin);
return storedContext;
}
/**
* Removes the current context and resets a new context that contains a merge of the current headers and the given headers.
* The removed context can be restored when closing the returned {@link StoredContext}. The merge strategy is that headers
* that are already existing are preserved unless they are defaults.
*/
public StoredContext stashAndMergeHeaders(Map headers) {
final ThreadContextStruct context = threadLocal.get();
Map newHeader = new HashMap<>(headers);
newHeader.putAll(context.requestHeaders);
threadLocal.set(DEFAULT_CONTEXT.putHeaders(newHeader));
return storedOriginalContext(context);
}
/**
* Just like {@link #stashContext()} but no default context is set and the response headers of the restore thread will be preserved.
*/
public StoredContext newStoredContextPreservingResponseHeaders() {
final ThreadContextStruct originalContext = threadLocal.get();
return () -> {
var found = threadLocal.get();
if (found != originalContext) {
threadLocal.set(originalContext.putResponseHeaders(found.responseHeaders));
}
};
}
/**
* Capture the current context and then restore the given context, returning a {@link StoredContext} that reverts back to the current
* context again. Equivalent to using {@link #newStoredContext()} and then calling {@code existingContext.restore()}.
*/
public StoredContext restoreExistingContext(StoredContext existingContext) {
final var originalContext = threadLocal.get();
existingContext.restore();
return storedOriginalContext(originalContext);
}
/**
* Just like {@link #stashContext()} but no default context is set.
*/
public StoredContext newStoredContext() {
final ThreadContextStruct originalContext = threadLocal.get();
return storedOriginalContext(originalContext);
}
/**
* Just like {@link #stashContext()} but no default context is set. Instead, the {@code transientHeadersToClear} argument can be used
* to clear specific transient headers in the new context and {@code requestHeadersToClear} can be used to clear specific request
* headers. All original headers (without the {@code responseHeaders}) are restored by closing the returned {@link StoredContext}.
*/
public StoredContext newStoredContext(Collection transientHeadersToClear, Collection requestHeadersToClear) {
return newStoredContext(false, transientHeadersToClear, requestHeadersToClear);
}
/**
* Just like {@link #newStoredContext(Collection, Collection)} but all headers are restored to original,
* except of {@code responseHeaders} which will be preserved from the restore thread.
*/
public StoredContext newStoredContextPreservingResponseHeaders(
Collection transientHeadersToClear,
Collection requestHeadersToClear
) {
return newStoredContext(true, transientHeadersToClear, requestHeadersToClear);
}
private StoredContext newStoredContext(
boolean preserveResponseHeaders,
Collection transientHeadersToClear,
Collection requestHeadersToClear
) {
final ThreadContextStruct originalContext = threadLocal.get();
// clear specific transient headers from the current context
Map newTransientHeaders = null;
for (String transientHeaderToClear : transientHeadersToClear) {
if (originalContext.transientHeaders.containsKey(transientHeaderToClear)) {
if (newTransientHeaders == null) {
newTransientHeaders = new HashMap<>(originalContext.transientHeaders);
}
newTransientHeaders.remove(transientHeaderToClear);
}
}
Map newRequestHeaders = null;
for (String requestHeaderToClear : requestHeadersToClear) {
if (originalContext.requestHeaders.containsKey(requestHeaderToClear)) {
if (newRequestHeaders == null) {
newRequestHeaders = new HashMap<>(originalContext.requestHeaders);
}
newRequestHeaders.remove(requestHeaderToClear);
}
}
if (newTransientHeaders != null || newRequestHeaders != null) {
ThreadContextStruct threadContextStruct = new ThreadContextStruct(
newRequestHeaders != null ? newRequestHeaders : originalContext.requestHeaders,
originalContext.responseHeaders,
newTransientHeaders != null ? newTransientHeaders : originalContext.transientHeaders,
originalContext.isSystemContext,
originalContext.warningHeadersSize
);
threadLocal.set(threadContextStruct);
}
// this is the context when this method returns
final ThreadContextStruct newContext = threadLocal.get();
return () -> {
if (preserveResponseHeaders && threadLocal.get() != newContext) {
threadLocal.set(originalContext.putResponseHeaders(threadLocal.get().responseHeaders));
} else {
threadLocal.set(originalContext);
}
};
}
/**
* Returns a supplier that gathers a {@link #newStoredContextPreservingResponseHeaders()} and restores it once the
* returned supplier is invoked. The context returned from the supplier is a stored version of the
* suppliers callers context that should be restored once the originally gathered context is not needed anymore.
* For instance this method should be used like this:
*
*
* Supplier<ThreadContext.StoredContext> restorable = context.newRestorableContext(true);
* new Thread() {
* public void run() {
* try (ThreadContext.StoredContext ctx = restorable.get()) {
* // execute with the parents context and restore the threads context afterwards
* }
* }
*
* }.start();
*
*
* @param preserveResponseHeaders if set to true the response headers of the restore thread will be preserved.
* @return a restorable context supplier
*/
public Supplier newRestorableContext(boolean preserveResponseHeaders) {
return wrapRestorable(preserveResponseHeaders ? newStoredContextPreservingResponseHeaders() : newStoredContext());
}
/**
* Same as {@link #newRestorableContext(boolean)} but wraps an existing context to restore.
* @param storedContext the context to restore
*/
public Supplier wrapRestorable(StoredContext storedContext) {
return () -> {
StoredContext context = newStoredContext();
storedContext.restore();
return context;
};
}
@Override
public void writeTo(StreamOutput out) throws IOException {
threadLocal.get().writeTo(out, defaultHeader);
}
/**
* Reads the headers from the stream into the current context
*/
public void readHeaders(StreamInput in) throws IOException {
setHeaders(readHeadersFromStream(in));
}
public void setHeaders(Tuple, Map>> headerTuple) {
final Map requestHeaders = headerTuple.v1();
final Map> responseHeaders = headerTuple.v2();
final ThreadContextStruct struct;
if (requestHeaders.isEmpty() && responseHeaders.isEmpty()) {
struct = ThreadContextStruct.EMPTY;
} else {
struct = new ThreadContextStruct(requestHeaders, responseHeaders, Collections.emptyMap(), false);
}
threadLocal.set(struct);
}
public static Tuple, Map>> readHeadersFromStream(StreamInput in) throws IOException {
final Map requestHeaders = in.readMap(StreamInput::readString);
final Map> responseHeaders = in.readMap(input -> {
final int size = input.readVInt();
if (size == 0) {
return Collections.emptySet();
} else if (size == 1) {
return Collections.singleton(input.readString());
} else {
// use a linked hash set to preserve order
final LinkedHashSet values = Sets.newLinkedHashSetWithExpectedSize(size);
for (int i = 0; i < size; i++) {
final String value = input.readString();
final boolean added = values.add(value);
assert added : value;
}
return values;
}
});
return new Tuple<>(requestHeaders, responseHeaders);
}
/**
* Returns the header for the given key or null if not present
*/
public String getHeader(String key) {
String value = threadLocal.get().requestHeaders.get(key);
if (value == null) {
return defaultHeader.get(key);
}
return value;
}
/**
* Returns the header for the given key or defaultValue if not present
*/
public String getHeaderOrDefault(String key, String defaultValue) {
String value = getHeader(key);
if (value == null) {
return defaultValue;
}
return value;
}
/**
* Returns all of the request headers from the thread's context.
* Be advised, headers might contain credentials.
* In order to avoid storing, and erroneously exposing, such headers,
* it is recommended to instead store security headers that prove
* the credentials have been verified successfully, and which are
* internal to the system, in the sense that they cannot be sent
* by the clients.
*/
public Map getHeaders() {
HashMap map = new HashMap<>(defaultHeader);
map.putAll(threadLocal.get().requestHeaders);
return Collections.unmodifiableMap(map);
}
/**
* Returns the request headers, without the default headers
*/
public Map getRequestHeadersOnly() {
return Collections.unmodifiableMap(new HashMap<>(threadLocal.get().requestHeaders));
}
/**
* Get a copy of all response headers.
*
* @return Never {@code null}.
*/
public Map> getResponseHeaders() {
Map> responseHeaders = threadLocal.get().responseHeaders;
Map> map = Maps.newMapWithExpectedSize(responseHeaders.size());
for (Map.Entry> entry : responseHeaders.entrySet()) {
map.put(entry.getKey(), List.copyOf(entry.getValue()));
}
return Collections.unmodifiableMap(map);
}
/**
* Copies all header key, value pairs into the current context
*/
public void copyHeaders(Iterable> headers) {
threadLocal.set(threadLocal.get().copyHeaders(headers));
}
/**
* Puts a header into the context
*/
public void putHeader(String key, String value) {
threadLocal.set(threadLocal.get().putRequest(key, value));
}
/**
* Puts all of the given headers into this context
*/
public void putHeader(Map header) {
threadLocal.set(threadLocal.get().putHeaders(header));
}
public void setErrorTraceTransportHeader(RestRequest r) {
// set whether data nodes should send back stack trace based on the `error_trace` query parameter
if (r.paramAsBoolean("error_trace", RestController.ERROR_TRACE_DEFAULT)) {
// We only set it if error_trace is true (defaults to false) to avoid sending useless bytes
putHeader("error_trace", "true");
}
}
/**
* Puts a transient header object into this context
*/
public void putTransient(String key, Object value) {
threadLocal.set(threadLocal.get().putTransient(key, value));
}
/**
* Returns a transient header object or null if there is no header for the given key
*/
@SuppressWarnings("unchecked") // (T)object
public T getTransient(String key) {
return (T) threadLocal.get().transientHeaders.get(key);
}
/**
* Returns unmodifiable copy of all transient headers.
*/
public Map getTransientHeaders() {
return Collections.unmodifiableMap(threadLocal.get().transientHeaders);
}
/**
* Add the {@code value} for the specified {@code key}. Any duplicate {@code value} is ignored.
*
* @param key the header name
* @param value the header value
*/
public void addResponseHeader(final String key, final String value) {
addResponseHeader(key, value, v -> v);
}
/**
* Add the {@code value} for the specified {@code key} with the specified {@code uniqueValue} used for de-duplication. Any duplicate
* {@code value} after applying {@code uniqueValue} is ignored.
*
* @param key the header name
* @param value the header value
* @param uniqueValue the function that produces de-duplication values
*/
public void addResponseHeader(final String key, final String value, final Function uniqueValue) {
threadLocal.set(threadLocal.get().putResponse(key, value, uniqueValue, maxWarningHeaderCount, maxWarningHeaderSize));
}
/**
* Saves the current thread context and wraps command in a Runnable that restores that context before running command. If
* command has already been passed through this method then it is returned unaltered rather than wrapped twice.
*/
public Runnable preserveContext(Runnable command) {
return doPreserveContext(command, false);
}
/**
* Saves the current thread context and wraps command in a Runnable that restores that context before running command. Also
* starts a new tracing context durin executing. If command has already been wrapped then it is returned unaltered.
*/
public Runnable preserveContextWithTracing(Runnable command) {
return doPreserveContext(command, true);
}
private Runnable doPreserveContext(Runnable command, boolean preserveContext) {
if (command instanceof ContextPreservingAbstractRunnable) {
return command;
}
if (command instanceof ContextPreservingRunnable) {
return command;
}
if (command instanceof AbstractRunnable abstractRunnable) {
return new ContextPreservingAbstractRunnable(abstractRunnable, preserveContext);
}
return new ContextPreservingRunnable(command);
}
/**
* Unwraps a command that was previously wrapped by {@link #preserveContext(Runnable)}.
*/
public static Runnable unwrap(Runnable command) {
if (command instanceof WrappedRunnable) {
return ((WrappedRunnable) command).unwrap();
}
return command;
}
/**
* Returns true if the current context is the default context.
*/
public boolean isDefaultContext() {
return threadLocal.get() == DEFAULT_CONTEXT;
}
/**
* Marks this thread context as an internal system context. This signals that actions in this context are issued
* by the system itself rather than by a user action.
*/
public void markAsSystemContext() {
threadLocal.set(threadLocal.get().setSystemContext());
}
/**
* Returns true iff this context is a system context
*/
public boolean isSystemContext() {
return threadLocal.get().isSystemContext;
}
/**
* Remove unwanted and unneeded headers from the thread context. Does not store prior context.
*/
public void sanitizeHeaders() {
final ThreadContextStruct originalContext = threadLocal.get();
final Map newRequestHeaders = new HashMap<>(originalContext.requestHeaders);
newRequestHeaders.entrySet()
.removeIf(
entry -> entry.getKey().equalsIgnoreCase("authorization")
|| entry.getKey().equalsIgnoreCase("es-secondary-authorization")
|| entry.getKey().equalsIgnoreCase("ES-Client-Authentication")
);
final ThreadContextStruct newContext = new ThreadContextStruct(
newRequestHeaders,
originalContext.responseHeaders,
originalContext.transientHeaders,
originalContext.isSystemContext,
originalContext.warningHeadersSize
);
threadLocal.set(newContext);
// intentionally not storing prior context to avoid restoring unwanted headers
}
@FunctionalInterface
public interface StoredContext extends AutoCloseable, Releasable {
default void restore() {
close();
}
}
public static Map buildDefaultHeaders(Settings settings) {
Settings headers = DEFAULT_HEADERS_SETTING.get(settings);
if (headers == null) {
return Map.of();
} else {
Map defaultHeader = new HashMap<>();
for (String key : headers.names()) {
defaultHeader.put(key, headers.get(key));
}
return Map.copyOf(defaultHeader);
}
}
private static final class ThreadContextStruct {
private static final ThreadContextStruct EMPTY = new ThreadContextStruct(
Collections.emptyMap(),
Collections.emptyMap(),
Collections.emptyMap(),
false
);
private final Map requestHeaders;
private final Map transientHeaders;
private final Map> responseHeaders;
private final boolean isSystemContext;
// saving current warning headers' size not to recalculate the size with every new warning header
private final long warningHeadersSize;
private ThreadContextStruct setSystemContext() {
if (isSystemContext) {
return this;
}
return new ThreadContextStruct(requestHeaders, responseHeaders, transientHeaders, true);
}
private ThreadContextStruct(
Map requestHeaders,
Map> responseHeaders,
Map transientHeaders,
boolean isSystemContext
) {
this(requestHeaders, responseHeaders, transientHeaders, isSystemContext, 0L);
}
private ThreadContextStruct(
Map requestHeaders,
Map> responseHeaders,
Map transientHeaders,
boolean isSystemContext,
long warningHeadersSize
) {
this.requestHeaders = requestHeaders;
this.responseHeaders = responseHeaders;
this.transientHeaders = transientHeaders;
this.isSystemContext = isSystemContext;
this.warningHeadersSize = warningHeadersSize;
}
/**
* This represents the default context and it should only ever be called by {@link #DEFAULT_CONTEXT}.
*/
private ThreadContextStruct() {
this(Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap(), false);
}
private ThreadContextStruct putRequest(String key, String value) {
Map newRequestHeaders = new HashMap<>(this.requestHeaders);
putSingleHeader(key, value, newRequestHeaders);
return new ThreadContextStruct(newRequestHeaders, responseHeaders, transientHeaders, isSystemContext);
}
private static void putSingleHeader(String key, T value, Map newHeaders) {
if (newHeaders.putIfAbsent(key, value) != null) {
throw new IllegalArgumentException("value for key [" + key + "] already present");
}
}
private ThreadContextStruct putHeaders(Map headers) {
if (headers.isEmpty()) {
return this;
} else {
final Map newHeaders = new HashMap<>(this.requestHeaders);
for (Map.Entry entry : headers.entrySet()) {
putSingleHeader(entry.getKey(), entry.getValue(), newHeaders);
}
return new ThreadContextStruct(newHeaders, responseHeaders, transientHeaders, isSystemContext);
}
}
private ThreadContextStruct putResponseHeaders(Map> headers) {
assert headers != null;
if (headers.isEmpty()) {
return this;
}
final Map> newResponseHeaders = new HashMap<>(this.responseHeaders);
for (Map.Entry> entry : headers.entrySet()) {
newResponseHeaders.merge(entry.getKey(), entry.getValue(), (existing, added) -> {
final Set updated = new LinkedHashSet<>(added);
updated.addAll(existing);
return Collections.unmodifiableSet(updated);
});
}
return new ThreadContextStruct(requestHeaders, newResponseHeaders, transientHeaders, isSystemContext);
}
private void logWarningHeaderThresholdExceeded(long threshold, Setting thresholdSetting) {
// If available, log some selected headers to help identifying the source of the request.
// Note: Only Task.HEADERS_TO_COPY are guaranteed to be preserved at this point.
Map selectedHeaders = new HashMap<>(requestHeaders);
selectedHeaders.keySet().retainAll(List.of(Task.X_OPAQUE_ID_HTTP_HEADER, Task.X_ELASTIC_PRODUCT_ORIGIN_HTTP_HEADER));
boolean sizeExceeded = HttpTransportSettings.SETTING_HTTP_MAX_WARNING_HEADER_SIZE.equals(thresholdSetting);
StringBuilder message = new StringBuilder().append("Dropping a warning header");
if (selectedHeaders.isEmpty() == false) {
message.append(" for request [").append(selectedHeaders).append("]");
}
message.append(", as their total").append(sizeExceeded ? " size" : " count");
message.append(" reached the maximum allowed of [").append(threshold).append("]").append(sizeExceeded ? " bytes" : "");
message.append(" set in [").append(thresholdSetting.getKey()).append("]!");
if (selectedHeaders.containsKey(Task.X_OPAQUE_ID_HTTP_HEADER) == false) {
message.append(" Add X-Opaque-Id headers to identify the source of this warning");
message.append(", see ").append(ReferenceDocs.X_OPAQUE_ID).append(" for more information.");
}
logger.warn(message);
}
private ThreadContextStruct putResponse(
final String key,
final String value,
final Function uniqueValue,
final int maxWarningHeaderCount,
final long maxWarningHeaderSize
) {
assert value != null;
long newWarningHeaderSize = warningHeadersSize;
// check if we can add another warning header - if max size within limits
if (key.equals("Warning") && (maxWarningHeaderSize != -1)) { // if size is NOT unbounded, check its limits
if (warningHeadersSize > maxWarningHeaderSize) { // if max size has already been reached before
logWarningHeaderThresholdExceeded(maxWarningHeaderSize, HttpTransportSettings.SETTING_HTTP_MAX_WARNING_HEADER_SIZE);
return this;
}
newWarningHeaderSize += "Warning".getBytes(StandardCharsets.UTF_8).length + value.getBytes(StandardCharsets.UTF_8).length;
if (newWarningHeaderSize > maxWarningHeaderSize) {
logWarningHeaderThresholdExceeded(maxWarningHeaderSize, HttpTransportSettings.SETTING_HTTP_MAX_WARNING_HEADER_SIZE);
return new ThreadContextStruct(
requestHeaders,
responseHeaders,
transientHeaders,
isSystemContext,
newWarningHeaderSize
);
}
}
final Map> newResponseHeaders;
final Set existingValues = responseHeaders.get(key);
if (existingValues != null && existingValues.contains(uniqueValue.apply(value))) {
return this;
}
newResponseHeaders = new HashMap<>(responseHeaders);
if (existingValues != null) {
// preserve insertion order
final Set newValues = new LinkedHashSet<>(existingValues);
newValues.add(value);
newResponseHeaders.put(key, Collections.unmodifiableSet(newValues));
} else {
newResponseHeaders.put(key, Collections.singleton(value));
}
// check if we can add another warning header - if max count within limits
if ((key.equals("Warning")) && (maxWarningHeaderCount != -1)) { // if count is NOT unbounded, check its limits
final int warningHeaderCount = newResponseHeaders.containsKey("Warning") ? newResponseHeaders.get("Warning").size() : 0;
if (warningHeaderCount > maxWarningHeaderCount) {
logWarningHeaderThresholdExceeded(maxWarningHeaderCount, HttpTransportSettings.SETTING_HTTP_MAX_WARNING_HEADER_COUNT);
return this;
}
}
return new ThreadContextStruct(requestHeaders, newResponseHeaders, transientHeaders, isSystemContext, newWarningHeaderSize);
}
private ThreadContextStruct putTransient(String key, Object value) {
Map newTransient = new HashMap<>(this.transientHeaders);
putSingleHeader(key, value, newTransient);
return new ThreadContextStruct(requestHeaders, responseHeaders, newTransient, isSystemContext);
}
private ThreadContextStruct copyHeaders(Iterable> headers) {
Map newHeaders = new HashMap<>();
for (Map.Entry header : headers) {
newHeaders.put(header.getKey(), header.getValue());
}
return putHeaders(newHeaders);
}
private void writeTo(StreamOutput out, Map defaultHeaders) throws IOException {
final Map requestHeaders;
if (defaultHeaders.isEmpty()) {
requestHeaders = this.requestHeaders;
} else {
requestHeaders = new HashMap<>(defaultHeaders);
requestHeaders.putAll(this.requestHeaders);
}
out.writeMap(requestHeaders, StreamOutput::writeString);
out.writeMap(responseHeaders, StreamOutput::writeStringCollection);
}
}
/**
* Wraps a Runnable to preserve the thread context.
*/
private class ContextPreservingRunnable implements WrappedRunnable {
private final Runnable in;
private final ThreadContext.StoredContext ctx;
private ContextPreservingRunnable(Runnable in) {
this.ctx = newStoredContext();
this.in = in;
}
@Override
public void run() {
try (var ignore = restoreExistingContext(ctx)) {
in.run();
}
}
@Override
public String toString() {
return in.toString();
}
@Override
public Runnable unwrap() {
return in;
}
}
/**
* Wraps an AbstractRunnable to preserve the thread context, optionally creating a new trace context before
* executing.
*/
private class ContextPreservingAbstractRunnable extends AbstractRunnable implements WrappedRunnable {
private final AbstractRunnable in;
private final ThreadContext.StoredContext creatorsContext;
private final boolean useNewTraceContext;
private ThreadContext.StoredContext threadsOriginalContext = null;
private ContextPreservingAbstractRunnable(AbstractRunnable in, boolean useNewTraceContext) {
creatorsContext = newStoredContext();
this.in = in;
this.useNewTraceContext = useNewTraceContext;
}
@Override
public boolean isForceExecution() {
return in.isForceExecution();
}
@Override
public void onAfter() {
try {
in.onAfter();
} finally {
if (threadsOriginalContext != null) {
threadsOriginalContext.restore();
}
}
}
@Override
public void onFailure(Exception e) {
in.onFailure(e);
}
@Override
public void onRejection(Exception e) {
in.onRejection(e);
}
@Override
protected void doRun() throws Exception {
threadsOriginalContext = stashContext();
creatorsContext.restore();
if (useNewTraceContext) {
// Discard the return value - we'll restore threadsOriginalContext in `onAfter()`.
// noinspection resource
newTraceContext();
}
in.doRun();
}
@Override
public String toString() {
return in.toString();
}
@Override
public AbstractRunnable unwrap() {
return in;
}
}
}