Fix thread context handling of headers overriding (#26068)

Previously collisions in headers between old and new contexts could be silently ignored, allowing the original context's headers to "win". This commit fixes the headers to require they are disjoint.
elastic · Oct 9, 2017 · 3ab27d1 · 3ab27d1
1 parent 8474269
commit 3ab27d1
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 7 deletions.
diff --git a/core/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContext.java b/core/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContext.java
@@ -407,11 +407,10 @@ private ThreadContextStruct putHeaders(Map<String, String> headers) {
             if (headers.isEmpty()) {
                 return this;
             } else {
-                final Map<String, String> newHeaders = new HashMap<>();
+                final Map<String, String> newHeaders = new HashMap<>(this.requestHeaders);
                 for (Map.Entry<String, String> entry : headers.entrySet()) {
                     putSingleHeader(entry.getKey(), entry.getValue(), newHeaders);
                 }
-                newHeaders.putAll(this.requestHeaders);
                 return new ThreadContextStruct(newHeaders, responseHeaders, transientHeaders, isSystemContext);
             }
         }

diff --git a/core/src/test/java/org/elasticsearch/common/util/concurrent/ThreadContextTests.java b/core/src/test/java/org/elasticsearch/common/util/concurrent/ThreadContextTests.java
@@ -29,7 +29,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
-
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasItem;
 import static org.hamcrest.Matchers.hasSize;
@@ -215,8 +214,8 @@ public void testResponseHeaders() {
     public void testCopyHeaders() {
         Settings build = Settings.builder().put("request.headers.default", "1").build();
         ThreadContext threadContext = new ThreadContext(build);
-        threadContext.copyHeaders(Collections.<String,String>emptyMap().entrySet());
-        threadContext.copyHeaders(Collections.<String,String>singletonMap("foo", "bar").entrySet());
+        threadContext.copyHeaders(Collections.<String, String>emptyMap().entrySet());
+        threadContext.copyHeaders(Collections.<String, String>singletonMap("foo", "bar").entrySet());
         assertEquals("bar", threadContext.getHeader("foo"));
     }
 
@@ -443,7 +442,7 @@ public void onAfter() {
                         assertEquals("bar", threadContext.getHeader("foo"));
                         assertEquals("bar_transient", threadContext.getTransient("foo"));
                         assertNotNull(threadContext.getTransient("failure"));
-                        assertEquals("exception from doRun", ((RuntimeException)threadContext.getTransient("failure")).getMessage());
+                        assertEquals("exception from doRun", ((RuntimeException) threadContext.getTransient("failure")).getMessage());
                         assertFalse(threadContext.isDefaultContext());
                         threadContext.putTransient("after", "after");
                     }
@@ -604,7 +603,7 @@ protected void doRun() throws Exception {
     public void testMarkAsSystemContext() throws IOException {
         try (ThreadContext threadContext = new ThreadContext(Settings.EMPTY)) {
             assertFalse(threadContext.isSystemContext());
-            try(ThreadContext.StoredContext context = threadContext.stashContext()){
+            try (ThreadContext.StoredContext context = threadContext.stashContext()) {
                 assertFalse(threadContext.isSystemContext());
                 threadContext.markAsSystemContext();
                 assertTrue(threadContext.isSystemContext());
@@ -613,6 +612,17 @@ public void testMarkAsSystemContext() throws IOException {
         }
     }
 
+    public void testPutHeaders() {
+        Settings build = Settings.builder().put("request.headers.default", "1").build();
+        ThreadContext threadContext = new ThreadContext(build);
+        threadContext.putHeader(Collections.<String, String>emptyMap());
+        threadContext.putHeader(Collections.<String, String>singletonMap("foo", "bar"));
+        assertEquals("bar", threadContext.getHeader("foo"));
+        IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () ->
+            threadContext.putHeader(Collections.<String, String>singletonMap("foo", "boom")));
+        assertEquals("value for key [foo] already present", e.getMessage());
+    }
+
     /**
      * Sometimes wraps a Runnable in an AbstractRunnable.
      */