Prevent ThreadContext header leak when sending response

docs/changelog/68649.yaml

@@ -0,0 +1,6 @@
+pr: 68649
+summary: Prevent `ThreadContext` header leak when sending response
+area: Infra/Core
+type: bug
+issues:
+ - 68278

modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestHandler.java

@@ -14,6 +14,7 @@
 
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.http.HttpPipelinedRequest;
+import org.elasticsearch.transport.Transports;
 
 @ChannelHandler.Sharable
 class Netty4HttpRequestHandler extends SimpleChannelInboundHandler<HttpPipelinedRequest> {
@@ -26,6 +27,8 @@ class Netty4HttpRequestHandler extends SimpleChannelInboundHandler<HttpPipelined
 
     @Override
     protected void channelRead0(ChannelHandlerContext ctx, HttpPipelinedRequest httpRequest) {
+        assert Transports.assertDefaultThreadContext(serverTransport.getThreadPool().getThreadContext());
+        assert Transports.assertTransportThread();
         final Netty4HttpChannel channel = ctx.channel().attr(Netty4HttpServerTransport.HTTP_CHANNEL_KEY).get();
         boolean success = false;
         try {
@@ -41,6 +44,8 @@ protected void channelRead0(ChannelHandlerContext ctx, HttpPipelinedRequest http
     @Override
     public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
         ExceptionsHelper.maybeDieOnAnotherThread(cause);
+        assert Transports.assertDefaultThreadContext(serverTransport.getThreadPool().getThreadContext());
+
         Netty4HttpChannel channel = ctx.channel().attr(Netty4HttpServerTransport.HTTP_CHANNEL_KEY).get();
         if (cause instanceof Error) {
             serverTransport.onException(channel, new Exception(cause));

modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java

@@ -304,7 +304,7 @@ protected HttpChannelHandler(final Netty4HttpServerTransport transport, final Ht
         protected void initChannel(Channel ch) throws Exception {
             Netty4HttpChannel nettyHttpChannel = new Netty4HttpChannel(ch);
             ch.attr(HTTP_CHANNEL_KEY).set(nettyHttpChannel);
-            ch.pipeline().addLast("chunked_writer", new Netty4WriteThrottlingHandler());
+            ch.pipeline().addLast("chunked_writer", new Netty4WriteThrottlingHandler(transport.getThreadPool().getThreadContext()));
             ch.pipeline().addLast("byte_buf_sizer", NettyByteBufSizer.INSTANCE);
             ch.pipeline().addLast("read_timeout", new ReadTimeoutHandler(transport.readTimeoutMillis, TimeUnit.MILLISECONDS));
             final HttpRequestDecoder decoder = new HttpRequestDecoder(

modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4Transport.java

@@ -374,7 +374,7 @@ private void setupPipeline(Channel ch) {
         ch.pipeline()
             .addLast("byte_buf_sizer", NettyByteBufSizer.INSTANCE)
             .addLast("logging", ESLoggingHandler.INSTANCE)
-            .addLast("chunked_writer", new Netty4WriteThrottlingHandler())
+            .addLast("chunked_writer", new Netty4WriteThrottlingHandler(getThreadPool().getThreadContext()))
             .addLast("dispatcher", new Netty4MessageInboundHandler(this, recycler));
     }
 

modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4WriteThrottlingHandler.java

@@ -15,6 +15,9 @@
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPromise;
 
+import org.elasticsearch.common.util.concurrent.ThreadContext;
+import org.elasticsearch.transport.Transports;
+
 import java.nio.channels.ClosedChannelException;
 import java.util.ArrayDeque;
 import java.util.Queue;
@@ -28,13 +31,18 @@ public final class Netty4WriteThrottlingHandler extends ChannelDuplexHandler {
 
     private final Queue<WriteOperation> queuedWrites = new ArrayDeque<>();
 
+    private final ThreadContext threadContext;
     private WriteOperation currentWrite;
 
-    public Netty4WriteThrottlingHandler() {}
+    public Netty4WriteThrottlingHandler(ThreadContext threadContext) {
+        this.threadContext = threadContext;
+    }
 
     @Override
     public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) {
         assert msg instanceof ByteBuf;
+        assert Transports.assertDefaultThreadContext(threadContext);
+        assert Transports.assertTransportThread();
         final boolean queued = queuedWrites.offer(new WriteOperation((ByteBuf) msg, promise));
         assert queued;
     }

server/src/main/java/org/elasticsearch/http/AbstractHttpServerTransport.java

@@ -492,4 +492,8 @@ private static ActionListener<Void> earlyResponseListener(HttpRequest request, H
             return ActionListener.noop();
         }
     }
+
+    public ThreadPool getThreadPool() {
+        return threadPool;
+    }
 }

server/src/main/java/org/elasticsearch/http/DefaultRestChannel.java

@@ -131,7 +131,9 @@ public void sendResponse(RestResponse restResponse) {
             addCookies(httpResponse);
 
             ActionListener<Void> listener = ActionListener.wrap(() -> Releasables.close(toClose));
-            httpChannel.sendResponse(httpResponse, listener);
+            try (ThreadContext.StoredContext existing = threadContext.stashContext()) {
+                httpChannel.sendResponse(httpResponse, listener);
+            }
             success = true;
         } finally {
             if (success == false) {