Merge branch 'cassandra-6.0' into trunk

* cassandra-6.0:
  Ensure a Message's Response streamId is always set

src/java/org/apache/cassandra/transport/CQLMessageHandler.java

@@ -106,7 +106,7 @@ public class CQLMessageHandler<M extends Message> extends AbstractMessageHandler
 
     interface MessageConsumer<M extends Message>
     {
-        void dispatch(Channel channel, M message, Dispatcher.FlushItemConverter toFlushItem, Overload backpressure);
+        <P> void dispatch(Channel channel, M message, Dispatcher.FlushItemConverter<P> toFlushItem, P param, Overload backpressure);
         boolean hasQueueCapacity();
     }
 
@@ -395,7 +395,7 @@ public class CQLMessageHandler<M extends Message> extends AbstractMessageHandler
         try
         {
             message = messageDecoder.decode(channel, request);
-            dispatcher.dispatch(channel, message, this::toFlushItem, backpressure);
+            dispatcher.dispatch(channel, message, CQLMessageHandler::toFlushItem, this, backpressure);
             
             // sucessfully delivered a CQL message to the execution
             // stage, so reset the counter of consecutive errors
@@ -491,7 +491,8 @@ public class CQLMessageHandler<M extends Message> extends AbstractMessageHandler
         // The Dispatcher will call this to obtain the FlushItem to enqueue with its Flusher once
         // a dispatched request has been processed.
 
-        Envelope responseFrame = response.encode(request.getSource().header.version);
+        Envelope.Header header = request.getSource().header;
+        Envelope responseFrame = response.encode(header.version, header.streamId);
         int responseSize = envelopeSize(responseFrame.header);
         ClientMessageSizeMetrics.bytesSent.inc(responseSize);
         ClientMessageSizeMetrics.bytesSentPerResponse.update(responseSize);
@@ -506,9 +507,9 @@ public class CQLMessageHandler<M extends Message> extends AbstractMessageHandler
     @Override
     public void cleanup(Flusher.FlushItem<Envelope> flushItem)
     {
-        release(flushItem.request.header);
+        release(flushItem.requestEnvelope.header);
-        flushItem.request.release();
+        flushItem.requestEnvelope.release();
-        flushItem.response.release();
+        flushItem.responseEnvelope.release();
     }
 
     private void release(Envelope.Header header)
@@ -530,8 +531,9 @@ public class CQLMessageHandler<M extends Message> extends AbstractMessageHandler
             if (!extracted.isSuccess())
             {
                 // Hard fail on any decoding error as we can't trust the subsequent frames of
-                // the large message
+                // the large message. The stream id is a best-effort value read before extraction
-                handleError(ProtocolException.toFatalException(extracted.error()));
+                // failed, so route it back where possible rather than defaulting.
+                handleError(ProtocolException.toFatalException(extracted.error()), extracted.streamId());
                 return false;
             }
 
@@ -548,7 +550,7 @@ public class CQLMessageHandler<M extends Message> extends AbstractMessageHandler
                 // not make sense to continue processing subsequent frames
                 handleError(ProtocolException.toFatalException(new OversizedAuthMessageException(
                             MULTI_FRAME_AUTH_ERROR_MESSAGE_PREFIX +
-                            "type = " + header.type + ", size = " + header.bodySizeInBytes)));
+                            "type = " + header.type + ", size = " + header.bodySizeInBytes)), header.streamId);
                 ClientMetrics.instance.markRequestDiscarded();
                 return false;
             }

src/java/org/apache/cassandra/transport/Dispatcher.java

@@ -93,10 +93,9 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
      * The instances of these FlushItem subclasses are specialized to release resources in the
      * right way for the specific pipeline that produced them.
      */
-    // TODO parameterize with FlushItem subclass
+    public interface FlushItemConverter<P>
-    interface FlushItemConverter
     {
-        FlushItem<?> toFlushItem(Channel channel, Message.Request request, Message.Response response);
+        FlushItem<?> toFlushItem(P param, Channel channel, Message.Request request, Message.Response response);
     }
 
     public Dispatcher(boolean useLegacyFlusher)
@@ -105,18 +104,17 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
     }
 
     @Override
-    public void dispatch(Channel channel, Message.Request request, FlushItemConverter forFlusher, Overload backpressure)
+    public <P> void dispatch(Channel channel, Message.Request request, FlushItemConverter<P> forFlusher, P param, Overload backpressure)
     {
         if (!request.connection().getTracker().isRunning())
         {
             // We can not respond with a custom, transport, or server exceptions since, given current implementation of clients,
             // they will defunct the connection. Without a protocol version bump that introduces an "I am going away message",
             // we have to stick to an existing error code.
-            Message.Response response = ErrorMessage.fromException(new OverloadedException("Server is shutting down"));
+            Message.Response response = ErrorMessage.fromTransportException(new OverloadedException("Server is shutting down"));
-            response.setStreamId(request.getStreamId());
             response.setWarnings(ClientWarn.instance.getWarnings());
             response.attach(request.connection);
-            FlushItem<?> toFlush = forFlusher.toFlushItem(channel, request, response);
+            FlushItem<?> toFlush = forFlusher.toFlushItem(param, channel, request, response);
             flush(toFlush);
             return;
         }
@@ -128,7 +126,7 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
         // Importantly, the authExecutor will handle the AUTHENTICATE message which may be CPU intensive.
         LocalAwareExecutorPlus executor = isAuthQuery ? authExecutor : requestExecutor;
 
-        executor.submit(new RequestProcessor(channel, request, forFlusher, backpressure));
+        executor.submit(new RequestProcessor<>(channel, request, forFlusher, param, backpressure));
         ClientMetrics.instance.markRequestDispatched();
     }
 
@@ -293,20 +291,22 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
      * is the only way we can keep it not wrapped into a callable on SEPExecutor submission path. And we need this
      * functionality for tracking time purposes.
      */
-    public class RequestProcessor implements DebuggableTask.RunnableDebuggableTask
+    public class RequestProcessor<P> implements DebuggableTask.RunnableDebuggableTask
     {
         private final Channel channel;
         private final Message.Request request;
-        private final FlushItemConverter forFlusher;
+        private final FlushItemConverter<P> forFlusher;
+        private final P flusherParam;
         private final Overload backpressure;
 
         private volatile long startTimeNanos;
 
-        public RequestProcessor(Channel channel, Message.Request request, FlushItemConverter forFlusher, Overload backpressure)
+        public RequestProcessor(Channel channel, Message.Request request, FlushItemConverter<P> forFlusher, P flusherParam, Overload backpressure)
         {
             this.channel = channel;
             this.request = request;
             this.forFlusher = forFlusher;
+            this.flusherParam = flusherParam;
             this.backpressure = backpressure;
         }
 
@@ -314,7 +314,7 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
         public void run()
         {
             startTimeNanos = MonotonicClock.Global.preciseTime.now();
-            processRequest(channel, request, forFlusher, backpressure, new RequestTime(request.createdAtNanos, startTimeNanos));
+            processRequest(channel, request, forFlusher, flusherParam, backpressure, new RequestTime(request.createdAtNanos, startTimeNanos));
         }
 
         @Override
@@ -374,7 +374,7 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
         if (queueTime > DatabaseDescriptor.getNativeTransportTimeout(TimeUnit.NANOSECONDS))
         {
             ClientMetrics.instance.markTimedOutBeforeProcessing();
-            return ErrorMessage.fromException(new OverloadedException("Query timed out before it could start"));
+            return ErrorMessage.fromTransportException(new OverloadedException("Query timed out before it could start"));
         }
 
         if (connection.getVersion().isGreaterOrEqualTo(ProtocolVersion.V4))
@@ -434,8 +434,6 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
             CoordinatorWriteWarnings.done();
         }
 
-        response.setStreamId(request.getStreamId());
-        response.setWarnings(ClientWarn.instance.getWarnings());
         response.attach(connection);
         connection.applyStateTransition(request.type, response.type);
         return response;
@@ -446,9 +444,10 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
      */
     static Message.Response processRequest(Channel channel, Message.Request request, Overload backpressure, RequestTime requestTime)
     {
+        Message.Response response = null;
         try
         {
-            return processRequest((ServerConnection) request.connection(), request, backpressure, requestTime);
+            response = processRequest((ServerConnection) request.connection(), request, backpressure, requestTime);
         }
         catch (Throwable t)
         {
@@ -461,26 +460,26 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
             }
 
             Predicate<Throwable> handler = ExceptionHandlers.getUnexpectedExceptionHandler(channel, true);
-            ErrorMessage error = ErrorMessage.fromException(t, handler);
+            response = ErrorMessage.fromExceptionNoStreamId(t, handler);
-            error.setStreamId(request.getStreamId());
-            error.setWarnings(ClientWarn.instance.getWarnings());
-            return error;
         }
         finally
         {
+            if (response != null)
+                response.setWarnings(ClientWarn.instance.getWarnings());
             CoordinatorWarnings.reset();
             CoordinatorWriteWarnings.reset();
             ClientWarn.instance.resetWarnings();
         }
+        return response;
     }
 
     /**
      * Note: this method is not expected to execute on the netty event loop.
      */
-    void processRequest(Channel channel, Message.Request request, FlushItemConverter forFlusher, Overload backpressure, RequestTime requestTime)
+    <P> void processRequest(Channel channel, Message.Request request, FlushItemConverter<P> forFlusher, P param, Overload backpressure, RequestTime requestTime)
     {
         Message.Response response = processRequest(channel, request, backpressure, requestTime);
-        FlushItem<?> toFlush = forFlusher.toFlushItem(channel, request, response);
+        FlushItem<?> toFlush = forFlusher.toFlushItem(param, channel, request, response);
         Message.logger.trace("Responding: {}, v={}", response, request.connection().getVersion());
         flush(toFlush);
     }
@@ -517,7 +516,7 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
      * for delivering events to registered clients is dependent on protocol version and the configuration
      * of the pipeline. For v5 and newer connections, the event message is encoded into an Envelope,
      * wrapped in a FlushItem and then delivered via the pipeline's flusher, in a similar way to
-     * a Response returned from {@link #processRequest(Channel, Message.Request, FlushItemConverter, Overload, RequestTime)}.
+     * a Response returned from {@link #processRequest(Channel, Message.Request, FlushItemConverter, Object, Overload, RequestTime)}.
      * It's worth noting that events are not generally fired as a direct response to a client request,
      * so this flush item has a null request attribute. The dispatcher itself is created when the
      * pipeline is first configured during protocol negotiation and is attached to the channel for
@@ -531,9 +530,9 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
                                            final FrameEncoder.PayloadAllocator allocator)
     {
         return eventMessage -> flush(new FlushItem.Framed(channel,
-                                                          eventMessage.encode(version),
+                                                          eventMessage.encode(version, EventMessage.EVENT_MESSAGE_STREAM_ID), // -1 was set in EventMessage previously
                                                           null,
                                                           allocator,
-                                                          f -> f.response.release()));
+                                                          f -> f.responseEnvelope.release()));
     }
 }

src/java/org/apache/cassandra/transport/Envelope.java

@@ -146,6 +146,11 @@ public class Envelope
         public final Message.Type type;
         public final long bodySizeInBytes;
 
+        public static Header dummy(int streamId, Message.Type type)
+        {
+            return new Header(ProtocolVersion.CURRENT, 0, streamId, type, 0);
+        }
+
         private Header(ProtocolVersion version, int flags, int streamId, Message.Type type, long bodySizeInBytes)
         {
             this.version = version;
@@ -258,7 +263,7 @@ public class Envelope
                 // This throws a protocol exception if the version number is unsupported,
                 // the opcode is unknown or invalid flags are set for the version
                 version = ProtocolVersion.decode(versionNum, DatabaseDescriptor.getNativeTransportAllowOlderProtocols());
-                validateFlags(version, flags);
+                validateFlags(version, flags, streamId);
                 type = Message.Type.fromOpcode(opcode, direction);
                 return new HeaderExtractionResult.Success(new Header(version, flags, streamId, type, bodyLength));
             }
@@ -272,6 +277,10 @@ public class Envelope
                 // cause the channel to be closed.
                 return new HeaderExtractionResult.Error(e, streamId, bodyLength);
             }
+            catch (ErrorMessage.WrappedException e)
+            {
+                return new HeaderExtractionResult.Error((ProtocolException) e.getCause(), e.getStreamId(), bodyLength);
+            }
         }
 
         public static abstract class HeaderExtractionResult
@@ -370,7 +379,8 @@ public class Envelope
             Message.Direction direction = Message.Direction.extractFromVersion(firstByte);
             int versionNum = firstByte & PROTOCOL_VERSION_MASK;
 
-            ProtocolVersion version;
+            ProtocolVersion version = null;
+            ProtocolException protocolException = null;
             
             try
             {
@@ -378,19 +388,43 @@ public class Envelope
             }
             catch (ProtocolException e)
             {
-                // Skip the remaining useless bytes. Otherwise the channel closing logic may try to decode again. 
+                // defer throw to attempt to extract the stream id
-                buffer.skipBytes(readableBytes);
+                protocolException = e;
-                throw e;
             }
 
             // Wait until we have the complete header
             if (readableBytes < Header.LENGTH)
+            {
+                if (protocolException != null)
+                {
+                    // Skip the remaining useless bytes. Otherwise the channel closing logic may try to decode again.
+                    buffer.skipBytes(readableBytes);
+                    // The header is incomplete, so there is no stream id to recover. Wrap with the unset
+                    // sentinel; the channel-level exception handler sees it has no routable stream id and
+                    // closes the connection rather than emit an unroutable error frame.
+                    throw protocolException;
+                }
                 return null;
+            }
 
             int flags = buffer.getByte(idx++);
-            validateFlags(version, flags);
-
             int streamId = buffer.getShort(idx);
+
+            if (protocolException != null)
+            {
+                // Protocol versions 1 and 2 use a shorter header with a single-byte stream id. Reading a
+                // 16-bit stream id from such a header splices the stream id byte together with the opcode
+                // byte and recovers a bogus id, routing the error to a stream the client never used
+                // (CASSANDRA-21508). A v1/v2 client that downgrades would then never see the error and time
+                // out, so recover the stream id using the attempted version's header layout.
+                int recoveredStreamId = versionNum < ProtocolVersion.V3.asInt() ? buffer.getByte(idx) : streamId;
+                // Skip the remaining useless bytes. Otherwise the channel closing logic may try to decode again.
+                buffer.skipBytes(readableBytes);
+                throw ErrorMessage.wrap(protocolException, recoveredStreamId);
+            }
+
+            validateFlags(version, flags, streamId);
+
             idx += 2;
 
             // This throws a protocol exceptions if the opcode is unknown
@@ -436,11 +470,12 @@ public class Envelope
             return new Envelope(new Header(version, flags, streamId, type, bodyLength), body);
         }
 
-        private void validateFlags(ProtocolVersion version, int flags)
+        private void validateFlags(ProtocolVersion version, int flags, int streamId)
         {
             if (version.isBeta() && !Header.Flag.contains(flags, Header.Flag.USE_BETA))
-                throw new ProtocolException(String.format("Beta version of the protocol used (%s), but USE_BETA flag is unset", version),
+                throw ErrorMessage.wrap(new ProtocolException(String.format("Beta version of the protocol used (%s), but USE_BETA flag is unset", version),
-                                            version);
+                                                              version),
+                                        streamId);
         }
 
         @Override

src/java/org/apache/cassandra/transport/ExceptionHandlers.java

@@ -38,6 +38,7 @@ import org.apache.cassandra.exceptions.OversizedCQLMessageException;
 import org.apache.cassandra.metrics.ClientMetrics;
 import org.apache.cassandra.net.FrameEncoder;
 import org.apache.cassandra.transport.messages.ErrorMessage;
+import org.apache.cassandra.transport.messages.ErrorMessage.WithStreamId;
 import org.apache.cassandra.utils.JVMStabilityInspector;
 import org.apache.cassandra.utils.NoSpamLogger;
 import org.apache.cassandra.utils.Throwables;
@@ -76,23 +77,53 @@ public class ExceptionHandlers
             if (ctx.channel().isOpen())
             {
                 Predicate<Throwable> handler = getUnexpectedExceptionHandler(ctx.channel(), false);
-                ErrorMessage errorMessage = ErrorMessage.fromException(cause, handler);
+                // No request in scope at the channel level; a WrappedException cause carries the frame's
-                Envelope response = errorMessage.encode(version);
+                // stream id and overrides this fallback.
-                FrameEncoder.Payload payload = allocator.allocate(true, CQLMessageHandler.envelopeSize(response.header));
+                WithStreamId withStreamId = ErrorMessage.fromException(cause, handler);
+                ErrorMessage errorMessage = withStreamId.message;
                 try
                 {
-                    response.encodeInto(payload.buffer);
+                    int streamId = withStreamId.streamId;
-                    response.release();
+                    boolean isFatal = isFatal(cause);
-                    payload.finish();
+                    if (streamId == Message.UNSET_STREAM_ID)
-                    ChannelPromise promise = ctx.newPromise();
+                    {
-                    // On protocol exception, close the channel as soon as the message has been sent
+                        // No stream id could be recovered, so we have no request to route a response to.
-                    if (isFatal(cause))
+                        // Close the connection rather than emit an unroutable frame (CASSANDRA-21508).
-                        promise.addListener(future -> ctx.close());
+                        isFatal = true;
-                    ctx.writeAndFlush(payload, promise);
+                        streamId = 0;
+                    }
+
+                    Envelope response = errorMessage.encode(version, streamId);
+                    FrameEncoder.Payload payload = allocator.allocate(true, CQLMessageHandler.envelopeSize(response.header));
+                    try
+                    {
+                        response.encodeInto(payload.buffer);
+                        response.release();
+                        payload.finish();
+                        ChannelPromise promise = ctx.newPromise();
+                        // On a fatal error, close the channel only once the error frame has been written,
+                        // so the client receives the diagnostic before the connection is torn down. Closing
+                        // synchronously here can abort the in-flight flush and drop the frame when the socket
+                        // can't drain it immediately (TCP backpressure, TLS buffering, or a large frame).
+                        // Matches PreV5Handlers.ExceptionHandler and InitialConnectionHandler.
+                        //
+                        // Trade-off of deferring the close (CASSANDRA-21508):
+                        //  - There is a slim chance we send two frames with the same streamId. Responses
+                        //    already queued on this connection will have a chance to flush before the close
+                        //    fires. For the majority of cases, each frame will carry its own unique stream
+                        //    id, so nothing is misrouted. These are valid responses to requests that were
+                        //    correctly-decoded earlier.
+                        if (isFatal)
+                            promise.addListener(future -> ctx.close());
+                        ctx.writeAndFlush(payload, promise);
+                    }
+                    finally
+                    {
+                        payload.release();
+                    }
                 }
                 finally
                 {
-                    payload.release();
                     JVMStabilityInspector.inspectThrowable(cause);
                 }
             }

src/java/org/apache/cassandra/transport/Flusher.java

@@ -36,7 +36,6 @@ import org.slf4j.LoggerFactory;
 import org.apache.cassandra.net.FrameEncoder;
 import org.apache.cassandra.net.FrameEncoderCrc;
 import org.apache.cassandra.net.FrameEncoderLZ4;
-import org.apache.cassandra.transport.Message.Response;
 import org.apache.cassandra.utils.FBUtilities;
 import org.apache.cassandra.utils.memory.BufferPool;
 
@@ -57,22 +56,23 @@ abstract class Flusher implements Runnable
     interface OnFlushCleanup<T> {
         void cleanup(FlushItem<T> item);
     }
-    static class FlushItem<T>
+
+    public static class FlushItem<T>
     {
         enum Kind {FRAMED, UNFRAMED}
 
         final Kind kind;
         final Channel channel;
-        final T response;
+        final T responseEnvelope;
-        final Envelope request;
+        final Envelope requestEnvelope;
         final OnFlushCleanup<T> tidy;
 
-        FlushItem(Kind kind, Channel channel, T response, Envelope request, OnFlushCleanup<T> tidy)
+        FlushItem(Kind kind, Channel channel, T responseEnvelope, Envelope requestEnvelope, OnFlushCleanup<T> tidy)
         {
             this.kind = kind;
             this.channel = channel;
-            this.request = request;
+            this.requestEnvelope = requestEnvelope;
-            this.response = response;
+            this.responseEnvelope = responseEnvelope;
             this.tidy = tidy;
         }
 
@@ -85,21 +85,21 @@ abstract class Flusher implements Runnable
         {
             final FrameEncoder.PayloadAllocator allocator;
             Framed(Channel channel,
-                   Envelope response,
+                   Envelope responseEnvelope,
-                   Envelope request,
+                   Envelope requestEnvelope,
                    FrameEncoder.PayloadAllocator allocator,
                    OnFlushCleanup<Envelope> tidy)
             {
-                super(Kind.FRAMED, channel, response, request, tidy);
+                super(Kind.FRAMED, channel, responseEnvelope, requestEnvelope, tidy);
                 this.allocator = allocator;
             }
         }
 
-        static class Unframed extends FlushItem<Response>
+        static class Unframed extends FlushItem<Envelope>
         {
-            Unframed(Channel channel, Response response, Envelope request, OnFlushCleanup<Response> tidy)
+            Unframed(Channel channel, Envelope responseEnvelope, Envelope requestEnvelope, OnFlushCleanup<Envelope> tidy)
             {
-                super(Kind.UNFRAMED, channel, response, request, tidy);
+                super(Kind.UNFRAMED, channel, responseEnvelope, requestEnvelope, tidy);
             }
         }
     }
@@ -151,13 +151,13 @@ abstract class Flusher implements Runnable
 
     private void processUnframedResponse(FlushItem.Unframed flush)
     {
-        flush.channel.write(flush.response, flush.channel.voidPromise());
+        flush.channel.write(flush.responseEnvelope, flush.channel.voidPromise());
         channels.add(flush.channel);
     }
 
     private void processFramedResponse(FlushItem.Framed flush)
     {
-        Envelope outbound = flush.response;
+        Envelope outbound = flush.responseEnvelope;
         if (envelopeSize(outbound.header) >= MAX_FRAMED_PAYLOAD_SIZE)
         {
             flushLargeMessage(flush.channel, outbound, flush.allocator);
@@ -171,7 +171,7 @@ abstract class Flusher implements Runnable
                 payloads.put(flushBuffer.channel, flushBuffer);
             }
 
-            flushBuffer.add(flush.response);
+            flushBuffer.add(flush.responseEnvelope);
         }
     }
 

src/java/org/apache/cassandra/transport/InitialConnectionHandler.java

@@ -36,6 +36,7 @@ import org.apache.cassandra.transport.messages.StartupMessage;
 import org.apache.cassandra.transport.messages.SupportedMessage;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelFutureListener;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPromise;
 import io.netty.channel.VoidChannelPromise;
@@ -91,8 +92,7 @@ public class InitialConnectionHandler extends ByteToMessageDecoder
                     supportedOptions.put(StartupMessage.COMPRESSION, compressions);
                     supportedOptions.put(StartupMessage.PROTOCOL_VERSIONS, ProtocolVersion.supportedVersions());
                     SupportedMessage supported = new SupportedMessage(supportedOptions);
-                    supported.setStreamId(inbound.header.streamId);
+                    outbound = supported.encode(inbound.header.version, inbound.header.streamId);
-                    outbound = supported.encode(inbound.header.version);
                     ctx.writeAndFlush(outbound);
                     break;
 
@@ -131,8 +131,8 @@ public class InitialConnectionHandler extends ByteToMessageDecoder
                                 if (null == cause)
                                     cause = new ServerError("Unexpected error establishing connection");
                                 logger.warn("Writing response to STARTUP failed, unable to configure pipeline", cause);
-                                ErrorMessage error = ErrorMessage.fromException(cause);
+                                ErrorMessage error = ErrorMessage.fromExceptionNoStreamId(cause);
-                                Envelope response = error.encode(inbound.header.version);
+                                Envelope response = error.encode(inbound.header.version, inbound.header.streamId);
                                 ChannelPromise closeChannel = AsyncChannelPromise.withListener(ctx, f -> ctx.close());
                                 ctx.writeAndFlush(response, closeChannel);
                                 if (ctx.channel().isOpen())
@@ -152,18 +152,20 @@ public class InitialConnectionHandler extends ByteToMessageDecoder
 
                     final Message.Response response = Dispatcher.processRequest(ctx.channel(), startup, Overload.NONE, Dispatcher.RequestTime.forImmediateExecution());
 
-                    outbound = response.encode(inbound.header.version);
+                    outbound = response.encode(inbound.header.version, inbound.header.streamId);
                     ctx.writeAndFlush(outbound, promise);
                     logger.trace("Configured pipeline: {}", ctx.pipeline());
                     break;
 
                 default:
                     ErrorMessage error =
-                        ErrorMessage.fromException(
+                        ErrorMessage.fromTransportException(
                             new ProtocolException(String.format("Unexpected message %s, expecting STARTUP or OPTIONS",
                                                                 inbound.header.type)));
-                    outbound = error.encode(inbound.header.version);
+                    outbound = error.encode(inbound.header.version, inbound.header.streamId);
-                    ctx.writeAndFlush(outbound);
+                    // An unexpected message during initial connection setup leaves the connection in a
+                    // corrupted state; send the error, then close the connection.
+                    ctx.writeAndFlush(outbound).addListener(ChannelFutureListener.CLOSE);
             }
         }
         finally

src/java/org/apache/cassandra/transport/Message.java

@@ -65,6 +65,12 @@ public abstract class Message
 {
     protected static final Logger logger = LoggerFactory.getLogger(Message.class);
 
+    /**
+     * Sentinel default for a {@link Response}'s stream id;
+     * must be overwritten before {@link #encode}.
+     **/
+    public static final int UNSET_STREAM_ID = Integer.MIN_VALUE;
+
     public interface Codec<M extends Message> extends CBCodec<M> {}
 
     public enum Direction
@@ -160,7 +166,6 @@ public abstract class Message
 
     public final Type type;
     protected Connection connection;
-    private int streamId;
     private Envelope source;
     private Map<String, ByteBuffer> customPayload;
     protected ProtocolVersion forcedProtocolVersion = null;
@@ -180,17 +185,6 @@ public abstract class Message
         return connection;
     }
 
-    public Message setStreamId(int streamId)
-    {
-        this.streamId = streamId;
-        return this;
-    }
-
-    public int getStreamId()
-    {
-        return streamId;
-    }
-
     public void setSource(Envelope source)
     {
         this.source = source;
@@ -214,7 +208,7 @@ public abstract class Message
     @Override
     public String toString()
     {
-        return String.format("(%s:%s:%s)", type, streamId, connection == null ? "null" :  connection.getVersion().asInt());
+        return String.format("(%s:%s)", type, connection == null ? "null" :  connection.getVersion().asInt());
     }
 
     public static abstract class Request extends Message
@@ -342,8 +336,16 @@ public abstract class Message
         }
     }
 
-    public Envelope encode(ProtocolVersion version)
+    public Envelope encode(ProtocolVersion version, int streamId)
     {
+        // A Response's stream id must be stamped before it is serialized to the wire. UNSET_STREAM_ID here
+        // means a server code path produced a response without routing information; sending it would risk
+        // delivering it to an unrelated in-flight request (CASSANDRA-21508). Fail fatally so the connection
+        // is torn down rather than mis-route a response. Checked before the try below so it is not caught and
+        // re-wrapped (which would carry the unset id forward).
+        if (streamId == UNSET_STREAM_ID)
+            throw ProtocolException.toFatalException(new ProtocolException("Attempted to encode a response with an unset stream id: " + this));
+
         int flags = Flag.none();
         @SuppressWarnings("unchecked")
         Codec<Message> codec = (Codec<Message>)this.type.codec;
@@ -430,11 +432,11 @@ public abstract class Message
             if (responseVersion.isBeta())
                 flags = Flag.add(flags, Flag.USE_BETA);
 
-            return Envelope.create(type, getStreamId(), responseVersion, flags, body);
+            return Envelope.create(type, streamId, responseVersion, flags, body);
         }
         catch (Throwable e)
         {
-            throw ErrorMessage.wrap(e, getStreamId());
+            throw ErrorMessage.wrap(e, streamId);
         }
     }
 
@@ -455,7 +457,6 @@ public abstract class Message
                 throw new ProtocolException("Received frame with CUSTOM_PAYLOAD flag for native protocol version < 4");
 
             Message message = inbound.header.type.codec.decode(inbound.body, inbound.header.version);
-            message.setStreamId(inbound.header.streamId);
             message.setSource(inbound);
             message.setCustomPayload(customPayload);
 

src/java/org/apache/cassandra/transport/PipelineConfigurator.java

@@ -409,7 +409,7 @@ public class PipelineConfigurator
         pipeline.addBefore(INITIAL_HANDLER, MESSAGE_DECOMPRESSOR, Envelope.Decompressor.instance);
         pipeline.addBefore(INITIAL_HANDLER, MESSAGE_COMPRESSOR, Envelope.Compressor.instance);
         pipeline.addBefore(INITIAL_HANDLER, MESSAGE_DECODER, PreV5Handlers.ProtocolDecoder.instance);
-        pipeline.addBefore(INITIAL_HANDLER, MESSAGE_ENCODER, PreV5Handlers.ProtocolEncoder.instance);
+        pipeline.addBefore(INITIAL_HANDLER, MESSAGE_ENCODER, PreV5Handlers.EventMessageEncoder.instance);
         pipeline.addBefore(INITIAL_HANDLER, LEGACY_MESSAGE_PROCESSOR, new PreV5Handlers.LegacyDispatchHandler(dispatcher, queueBackpressure, limits));
         pipeline.remove(INITIAL_HANDLER);
         onNegotiationComplete(pipeline);

src/java/org/apache/cassandra/transport/PreV5Handlers.java

@@ -24,6 +24,7 @@ import java.util.List;
 
 import javax.net.ssl.SSLException;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Predicate;
 
 import org.apache.commons.lang3.exception.ExceptionUtils;
@@ -40,6 +41,7 @@ import org.apache.cassandra.service.ClientState;
 import org.apache.cassandra.service.QueryState;
 import org.apache.cassandra.transport.ClientResourceLimits.Overload;
 import org.apache.cassandra.transport.messages.ErrorMessage;
+import org.apache.cassandra.transport.messages.EventMessage;
 import org.apache.cassandra.utils.JVMStabilityInspector;
 
 import io.netty.channel.Channel;
@@ -93,23 +95,26 @@ public class PreV5Handlers
             // The only reason we won't process this message is if checkLimits() throws an OverloadedException.
             // (i.e. Even if backpressure is applied, the current request is allowed to finish.)
             checkLimits(ctx, request);
-            dispatcher.dispatch(ctx.channel(), request, this::toFlushItem, backpressure);
+            dispatcher.dispatch(ctx.channel(), request, this::toFlushItem, ctx, backpressure);
         }
 
         // Acts as a Dispatcher.FlushItemConverter
-        private Flusher.FlushItem.Unframed toFlushItem(Channel channel, Message.Request request, Message.Response response)
+        private Flusher.FlushItem.Unframed toFlushItem(ChannelHandlerContext ctx, Channel channel, Message.Request request, Message.Response response)
         {
-            return new Flusher.FlushItem.Unframed(channel, response, request.getSource(), this::releaseItem);
+            ProtocolVersion version = getConnectionVersion(ctx);
+            Envelope requestEnvelope = request.getSource();
+            Envelope responseEnvelope = response.encode(version, requestEnvelope.header.streamId);
+            return new Flusher.FlushItem.Unframed(channel, responseEnvelope, requestEnvelope, this::releaseItem);
         }
 
-        private void releaseItem(Flusher.FlushItem<Message.Response> item)
+        private void releaseItem(Flusher.FlushItem<Envelope> item)
         {
             // Note: in contrast to the equivalent for V5 protocol, CQLMessageHandler::release(FlushItem item),
             // this does not release the FlushItem's Message.Response. In V4, the buffers for the response's body
             // and serialised header are emitted directly down the Netty pipeline from Envelope.Encoder, so
             // releasing them is handled by the pipeline itself.
-            long itemSize = item.request.header.bodySizeInBytes;
+            long itemSize = item.requestEnvelope.header.bodySizeInBytes;
-            item.request.release();
+            item.requestEnvelope.release();
 
             // since the request has been processed, decrement inflight payload at channel, endpoint and global levels
             channelPayloadBytesInFlight -= itemSize;
@@ -307,14 +312,14 @@ public class PreV5Handlers
      * Simple adaptor to plug CQL message encoding into pre-V5 pipelines
      */
     @ChannelHandler.Sharable
-    public static class ProtocolEncoder extends MessageToMessageEncoder<Message>
+    public static class EventMessageEncoder extends MessageToMessageEncoder<EventMessage>
     {
-        public static final ProtocolEncoder instance = new ProtocolEncoder();
+        public static final EventMessageEncoder instance = new EventMessageEncoder();
-        private ProtocolEncoder(){}
+        private EventMessageEncoder(){}
-        public void encode(ChannelHandlerContext ctx, Message source, List<Object> results)
+        public void encode(ChannelHandlerContext ctx, EventMessage source, List<Object> results)
         {
             ProtocolVersion version = getConnectionVersion(ctx);
-            results.add(source.encode(version));
+            results.add(source.encode(version, EventMessage.EVENT_MESSAGE_STREAM_ID));
         }
     }
 
@@ -337,13 +342,28 @@ public class PreV5Handlers
             if (ctx.channel().isOpen())
             {
                 Predicate<Throwable> handler = ExceptionHandlers.getUnexpectedExceptionHandler(ctx.channel(), false);
-                ErrorMessage errorMessage = ErrorMessage.fromException(cause, handler);
+                // No request in scope at the channel level; a WrappedException cause carries the frame's
-                ChannelFuture future = ctx.writeAndFlush(errorMessage.encode(getConnectionVersion(ctx)));
+                // stream id and overrides this fallback.
-                // On protocol exception, close the channel as soon as the message have been sent.
+                ErrorMessage.WithStreamId withStreamId = ErrorMessage.fromException(cause, handler);
-                // Most cases of PE are wrapped so the type check below is expected to fail more often than not.
+
-                // At this moment Fatal exceptions are not thrown in v4, but just as a precaustion we check for them here
+                if (withStreamId.streamId == Message.UNSET_STREAM_ID)
-                if (isFatal(cause))
+                {
-                    future.addListener((ChannelFutureListener) f -> ctx.close());
+                    // No stream id could be recovered, so we have no request to route a response to.
+                    // Close the connection rather than emit an unroutable frame (CASSANDRA-21508).
+                    ctx.close();
+                }
+                else
+                {
+                    ErrorMessage errorMessage = withStreamId.message;
+                    int streamId = withStreamId.streamId;
+
+                    ChannelFuture future = ctx.writeAndFlush(errorMessage.encode(getConnectionVersion(ctx), streamId));
+                    // On protocol exception, close the channel as soon as the message have been sent.
+                    // Most cases of PE are wrapped so the type check below is expected to fail more often than not.
+                    // At this moment Fatal exceptions are not thrown in v4, but just as a precaustion we check for them here
+                    if (isFatal(cause))
+                        future.addListener((ChannelFutureListener) f -> ctx.close());
+                }
             }
 
             SocketAddress remoteAddress = ctx.channel().remoteAddress();
@@ -393,7 +413,8 @@ public class PreV5Handlers
         }
     }
 
-    private static ProtocolVersion getConnectionVersion(ChannelHandlerContext ctx)
+    @VisibleForTesting
+    static ProtocolVersion getConnectionVersion(ChannelHandlerContext ctx)
     {
         Connection connection = ctx.channel().attr(Connection.attributeKey).get();
         // The only case the connection can be null is when we send the initial STARTUP message

src/java/org/apache/cassandra/transport/SimpleClient.java

@@ -352,7 +352,7 @@ public class SimpleClient implements Closeable
                 for (int i = 0; i < requests.size(); i++)
                 {
                     Message.Request message = requests.get(i);
-                    message.setStreamId(i);
+                    message.setSource(new Envelope(Envelope.Header.dummy(i, message.type), null));
                     message.attach(connection);
                 }
                 lastWriteFuture = channel.writeAndFlush(requests);
@@ -365,7 +365,7 @@ public class SimpleClient implements Closeable
                         throw new RuntimeException("timeout");
                     if (msg instanceof ErrorMessage)
                         throw new RuntimeException((Throwable) ((ErrorMessage) msg).error);
-                    rrMap.put(requests.get(msg.getStreamId()), msg);
+                    rrMap.put(requests.get(msg.getSource().header.streamId), msg);
                 }
             }
             else
@@ -383,6 +383,18 @@ public class SimpleClient implements Closeable
         }
     }
 
+    /**
+     * The stream id to frame an outbound client request with. SimpleClient carries the intended id on the
+     * request's (dummy) source envelope (see {@link #execute(List)} and callers that pipeline requests).
+     * When no source has been assigned, we fall back to 0, which is sufficient for the non-pipelined path
+     * where only a single request is ever in flight.
+     */
+    private static int outboundStreamId(Message message)
+    {
+        Envelope source = message.getSource();
+        return source == null ? 0 : source.header.streamId;
+    }
+
     public interface EventHandler
     {
         void onEvent(Event event);
@@ -438,37 +450,38 @@ public class SimpleClient implements Closeable
             this.largeMessageThreshold = largeMessageThreshold;
         }
 
-        protected void decode(ChannelHandlerContext ctx, Envelope response, List<Object> results)
+        @Override
+        protected void decode(ChannelHandlerContext ctx, Envelope request, List<Object> results)
         {
-            switch(response.header.type)
+            switch(request.header.type)
             {
                 case READY:
                 case AUTHENTICATE:
-                    if (response.header.version.isGreaterOrEqualTo(ProtocolVersion.V5))
+                    if (request.header.version.isGreaterOrEqualTo(ProtocolVersion.V5))
                     {
-                        configureModernPipeline(ctx, response, largeMessageThreshold);
+                        configureModernPipeline(ctx, request, largeMessageThreshold);
                         // consuming the message is done when setting up the pipeline
                     }
                     else
                     {
                         configureLegacyPipeline(ctx);
                         // really just removes self from the pipeline, so pass this message on
-                        ctx.pipeline().context(Envelope.Decoder.class).fireChannelRead(response);
+                        ctx.pipeline().context(Envelope.Decoder.class).fireChannelRead(request);
                     }
                     break;
                 case SUPPORTED:
                 case ERROR:
                     // just pass through
-                    results.add(response);
+                    results.add(request);
                     break;
                 default:
-                    throw new ProtocolException(String.format("Unexpected %s response expecting " +
+                    throw new ProtocolException(String.format("Unexpected %s request expecting " +
                                                               "READY, AUTHENTICATE, ERROR or SUPPORTED",
-                                                              response.header.type));
+                                                              request.header.type));
             }
         }
 
-        private void configureModernPipeline(ChannelHandlerContext ctx, Envelope response, int largeMessageThreshold)
+        private void configureModernPipeline(ChannelHandlerContext ctx, Envelope request, int largeMessageThreshold)
         {
             logger.info("Configuring modern pipeline");
             ChannelPipeline pipeline = ctx.pipeline();
@@ -490,7 +503,7 @@ public class SimpleClient implements Closeable
 
             CQLMessageHandler.MessageConsumer<Message.Response> responseConsumer = new CQLMessageHandler.MessageConsumer<Message.Response>()
             {
-                public void dispatch(Channel channel, Message.Response message, Dispatcher.FlushItemConverter toFlushItem, Overload backpressure)
+                public <P> void dispatch(Channel channel, Message.Response message, Dispatcher.FlushItemConverter<P> toFlushItem, P param, Overload backpressure)
                 {
                     responseHandler.handleResponse(channel, message);
                 }
@@ -585,15 +598,15 @@ public class SimpleClient implements Closeable
                     ProtocolVersion version = connection == null ? ProtocolVersion.CURRENT : connection.getVersion();
                     SimpleFlusher flusher = new SimpleFlusher(frameEncoder, largeMessageThreshold);
                     for (Message message : (List<Message>) msg)
-                        flusher.enqueue(message.encode(version));
+                        flusher.enqueue(message.encode(version, outboundStreamId(message)));
 
                     flusher.maybeWrite(ctx, promise);
                 }
             });
             pipeline.remove(this);
 
-            Message.Response message = messageDecoder.decode(ctx.channel(), response);
+            Message.Response message = messageDecoder.decode(ctx.channel(), request);
-            responseConsumer.dispatch(channel, message, (ch, req, resp) -> null, Overload.NONE);
+            responseConsumer.dispatch(channel, message, (p, ch, req, resp) -> null, null, Overload.NONE);
         }
 
         private FrameDecoder frameDecoder(ChannelHandlerContext ctx, BufferPoolAllocator allocator)
@@ -638,7 +651,8 @@ public class SimpleClient implements Closeable
             // The only case the connection can be null is when we send the initial STARTUP message (client side thus)
             ProtocolVersion version = connection == null ? ProtocolVersion.CURRENT : connection.getVersion();
             assert messages.size() == 1;
-            results.add(messages.get(0).encode(version));
+            Message message = messages.get(0);
+            results.add(message.encode(version, outboundStreamId(message)));
         }
     }
 
@@ -662,7 +676,7 @@ public class SimpleClient implements Closeable
             pipeline.addLast(HandlerNames.INITIAL_HANDLER, new InitialHandler(version, responseHandler, largeMessageThreshold));
             pipeline.addLast(HandlerNames.MESSAGE_DECODER, PreV5Handlers.ProtocolDecoder.instance);
             pipeline.addLast(HandlerNames.MESSAGE_ENCODER, MessageBatchEncoder.instance);
-            pipeline.addLast(HandlerNames.RESPONSE_HANDLER,  responseHandler);
+            pipeline.addLast(HandlerNames.RESPONSE_HANDLER, responseHandler);
         }
     }
 

src/java/org/apache/cassandra/transport/messages/AuthUtil.java

@@ -87,7 +87,7 @@ public final class AuthUtil
         {
             ClientMetrics.instance.markAuthFailure(negotiator.getAuthenticationMode());
             AuthEvents.instance.notifyAuthFailure(queryState, e);
-            return ErrorMessage.fromException(e);
+            return ErrorMessage.fromTransportException(e);
         }
     }
 

src/java/org/apache/cassandra/transport/messages/BatchMessage.java

@@ -236,7 +236,7 @@ public class BatchMessage extends Message.Request
         {
             QueryEvents.instance.notifyBatchFailure(prepared, batchType, queryOrIdList, values, options, state, e);
             JVMStabilityInspector.inspectThrowable(e);
-            return ErrorMessage.fromException(e);
+            return ErrorMessage.fromExceptionNoStreamId(e);
         }
     }
 

src/java/org/apache/cassandra/transport/messages/ErrorMessage.java

@@ -423,25 +423,49 @@ public class ErrorMessage extends Message.Response
         this.error = error;
     }
 
-    private ErrorMessage(TransportException error, int streamId)
+    public static ErrorMessage fromTransportException(TransportException e)
     {
-        this(error);
+        ErrorMessage message = new ErrorMessage(e);
-        setStreamId(streamId);
+        if (e instanceof ProtocolException)
+        {
+            // if the driver attempted to connect with a protocol version not supported then
+            // respond with the appropiate version, see ProtocolVersion.decode()
+            ProtocolVersion forcedProtocolVersion = ((ProtocolException) e).getForcedProtocolVersion();
+            if (forcedProtocolVersion != null)
+                message.forcedProtocolVersion = forcedProtocolVersion;
+        }
+        return message;
     }
 
-    public static ErrorMessage fromException(Throwable e)
+    public static ErrorMessage fromExceptionNoStreamId(Throwable e)
     {
-        return fromException(e, null);
+        return fromExceptionNoStreamId(e, null);
     }
 
+    public static ErrorMessage fromExceptionNoStreamId(Throwable e, Predicate<Throwable> unexpectedExceptionHandler)
+    {
+        return fromException(e, unexpectedExceptionHandler).message;
+    }
+
+    public static class WithStreamId
+    {
+        public final ErrorMessage message;
+        public final int streamId;
+
+        WithStreamId(ErrorMessage message, int streamId)
+        {
+            this.message = message;
+            this.streamId = streamId;
+        }
+    }
     /**
      * @param e the exception
      * @param unexpectedExceptionHandler a callback for handling unexpected exceptions. If null, or if this
      *                                   returns false, the error is logged at ERROR level via sl4fj
      */
-    public static ErrorMessage fromException(Throwable e, Predicate<Throwable> unexpectedExceptionHandler)
+    public static WithStreamId fromException(Throwable e, Predicate<Throwable> unexpectedExceptionHandler)
     {
-        int streamId = 0;
+        int streamId = UNSET_STREAM_ID;
 
         // Netty will wrap exceptions during decoding in a CodecException. If the cause was one of our ProtocolExceptions
         // or some other internal exception, extract that and use it.
@@ -469,23 +493,15 @@ public class ErrorMessage extends Message.Response
 
         if (e instanceof TransportException)
         {
-            ErrorMessage message = new ErrorMessage((TransportException) e, streamId);
+            ErrorMessage message = fromTransportException((TransportException) e);
-            if (e instanceof ProtocolException)
+            return new WithStreamId(message, streamId);
-            {
-                // if the driver attempted to connect with a protocol version not supported then
-                // respond with the appropiate version, see ProtocolVersion.decode()
-                ProtocolVersion forcedProtocolVersion = ((ProtocolException) e).getForcedProtocolVersion();
-                if (forcedProtocolVersion != null)
-                    message.forcedProtocolVersion = forcedProtocolVersion;
-            }
-            return message;
         }
 
         // Unexpected exception
         if (unexpectedExceptionHandler == null || !unexpectedExceptionHandler.apply(e))
             logger.error("Unexpected exception during request", e);
 
-        return new ErrorMessage(new ServerError(e), streamId);
+        return new WithStreamId(new ErrorMessage(new ServerError(e)), streamId);
     }
 
     @Override

src/java/org/apache/cassandra/transport/messages/EventMessage.java

@@ -25,6 +25,8 @@ import io.netty.buffer.ByteBuf;
 
 public class EventMessage extends Message.Response
 {
+    public static final int EVENT_MESSAGE_STREAM_ID = -1;
+
     public static final Message.Codec<EventMessage> codec = new Message.Codec<EventMessage>()
     {
         public EventMessage decode(ByteBuf body, ProtocolVersion version)
@@ -49,7 +51,6 @@ public class EventMessage extends Message.Response
     {
         super(Message.Type.EVENT);
         this.event = event;
-        this.setStreamId(-1);
     }
 
     @Override

src/java/org/apache/cassandra/transport/messages/ExecuteMessage.java

@@ -207,7 +207,7 @@ public class ExecuteMessage extends Message.Request
         {
             QueryEvents.instance.notifyExecuteFailure(prepared, options, state, e);
             JVMStabilityInspector.inspectThrowable(e);
-            return ErrorMessage.fromException(e);
+            return ErrorMessage.fromExceptionNoStreamId(e);
         }
     }
 

src/java/org/apache/cassandra/transport/messages/PrepareMessage.java

@@ -51,12 +51,13 @@ public class PrepareMessage extends Message.Request
         {
             String query = CBUtil.readLongString(body);
             String keyspace = null;
-            if (version.isGreaterOrEqualTo(ProtocolVersion.V5)) {
+            if (version.isGreaterOrEqualTo(ProtocolVersion.V5))
+            {
                 // If flags grows, we may want to consider creating a PrepareOptions class with an internal codec
                 // class that handles flags and options of the prepare message. Since there's only one right now,
                 // we just take care of business here.
 
-                int flags = (int)body.readUnsignedInt();
+                int flags = (int) body.readUnsignedInt();
                 if ((flags & 0x1) == 0x1)
                 {
                     keyspace = CBUtil.readString(body);
@@ -74,8 +75,11 @@ public class PrepareMessage extends Message.Request
             {
                 // If we have no keyspace, write out a 0-valued flag field.
                 if (msg.keyspace == null)
+                {
                     dest.writeInt(0x0);
-                else {
+                }
+                else
+                {
                     dest.writeInt(0x1);
                     CBUtil.writeAsciiString(msg.keyspace, dest);
                 }
@@ -134,7 +138,7 @@ public class PrepareMessage extends Message.Request
         {
             QueryEvents.instance.notifyPrepareFailure(null, query, state, e);
             JVMStabilityInspector.inspectThrowable(e);
-            return ErrorMessage.fromException(e);
+            return ErrorMessage.fromExceptionNoStreamId(e);
         }
     }
 

src/java/org/apache/cassandra/transport/messages/QueryMessage.java

@@ -129,7 +129,7 @@ public class QueryMessage extends Message.Request
             JVMStabilityInspector.inspectThrowable(e);
             if (!((e instanceof RequestValidationException) || (e instanceof RequestExecutionException)))
                 logger.error("Unexpected error during query", e);
-            return ErrorMessage.fromException(e);
+            return ErrorMessage.fromExceptionNoStreamId(e);
         }
     }
 

test/burn/org/apache/cassandra/transport/DriverBurnTest.java

@@ -340,10 +340,10 @@ public class DriverBurnTest extends CQLTester
         SimpleStatement request = generateQueryStatement(0, requestCaps);
         ResultMessage.Rows response = generateRows(0, responseCaps);
         QueryMessage requestMessage = generateQueryMessage(0, requestCaps, version);
-        Envelope message = requestMessage.encode(version);
+        Envelope message = requestMessage.encode(version, 0);
         int requestSize = message.body.readableBytes();
         message.release();
-        message = response.encode(version);
+        message = response.encode(version, 0);
         int responseSize = message.body.readableBytes();
         message.release();
         Message.Type.QUERY.unsafeSetCodec(new Message.Codec<QueryMessage>() {

test/burn/org/apache/cassandra/transport/SimpleClientPerfTest.java

@@ -155,10 +155,10 @@ public class SimpleClientPerfTest
     {
         ResultMessage.Rows response = generateRows(0, responseCaps);
         QueryMessage requestMessage = generateQueryMessage(0, requestCaps, version);
-        Envelope message = requestMessage.encode(version);
+        Envelope message = requestMessage.encode(version, 0);
         int requestSize = message.body.readableBytes();
         message.release();
-        message = response.encode(version);
+        message = response.encode(version, 0);
         int responseSize = message.body.readableBytes();
         message.release();
 

test/distributed/org/apache/cassandra/distributed/test/StreamIdMisrouteTest.java

@@ -0,0 +1,254 @@
+/*
+ * 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.cassandra.distributed.test;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
+import net.bytebuddy.implementation.MethodDelegation;
+import net.bytebuddy.implementation.bind.annotation.SuperCall;
+
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import org.apache.cassandra.cql3.QueryOptions;
+import org.apache.cassandra.cql3.statements.SelectStatement;
+import org.apache.cassandra.distributed.Cluster;
+import org.apache.cassandra.exceptions.OverloadedException;
+import org.apache.cassandra.service.QueryState;
+import org.apache.cassandra.transport.Dispatcher;
+import org.apache.cassandra.transport.Envelope;
+import org.apache.cassandra.transport.Message;
+import org.apache.cassandra.transport.ProtocolVersion;
+import org.apache.cassandra.transport.SimpleClient;
+import org.apache.cassandra.transport.messages.ErrorMessage;
+import org.apache.cassandra.transport.messages.QueryMessage;
+import org.apache.cassandra.transport.messages.ResultMessage;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+import static org.apache.cassandra.config.DatabaseDescriptor.clientInitialization;
+import static org.apache.cassandra.config.DatabaseDescriptor.getNativeTransportPort;
+import static org.apache.cassandra.distributed.api.Feature.GOSSIP;
+import static org.apache.cassandra.distributed.api.Feature.NATIVE_PROTOCOL;
+import static org.apache.cassandra.distributed.api.Feature.NETWORK;
+
+/**
+ * Regression test for CASSANDRA-21508. A coordinator that load-sheds a request which waited past its
+ * native-transport queue deadline must stamp the OVERLOADED error with the timed-out request's own
+ * stream id. Before the fix the error went out on stream id 0, which misroutes it to an unrelated
+ * in-flight request and can escalate into the client-side "column-shift" read corruption (a value
+ * from one table decoded against another query's column definitions under the v5 skip-metadata
+ * optimization).
+ *
+ * <p>The original defect: {@code Dispatcher.processRequest} load-sheds a request that has waited in the
+ * Native-Transport-Requests queue longer than {@code native_transport_timeout} by returning
+ *
+ * <pre>
+ *   ErrorMessage.fromException(new OverloadedException("Query timed out before it could start"))
+ * </pre>
+ * <p>
+ * without calling {@code setStreamId(request.getStreamId())}. {@code ErrorMessage.streamId} therefore
+ * kept its default of 0, and the error frame was written to the client on stream id 0 instead of the
+ * timed-out request's real stream id. The fix routes every response through a central stamping step
+ * and makes {@code ErrorMessage.fromException} require the stream id at the call site.
+ *
+ * <p>{@link #loadShedErrorIsStampedWithRequestStreamId()} proves the fix on the wire: a request sent on
+ * a non-zero stream id is load-shed and the OVERLOADED error comes back on that same stream id (not 0).
+ * Against the unpatched server this assertion fails with the error arriving on stream id 0.
+ *
+ * <p>Why the stream id matters: on a busy connection stream id 0 is almost always in use, so an error
+ * mis-stamped with 0 is applied to an unrelated in-flight request. The client then frees and reuses
+ * stream id 0 for a new query, and when the original query's real rows arrive on it they are decoded
+ * positionally against the reusing query's cached column definitions (skip-metadata carries no column
+ * info) - shifting each value into the wrong column and either throwing in a codec or silently
+ * returning a plausible-but-wrong value.
+ */
+public class StreamIdMisrouteTest extends TestBaseImpl
+{
+    private static final int TIMED_OUT_STREAM_ID = 42; // any non-zero id; before the fix the reply was forced to 0
+
+    @BeforeClass
+    public static void initClientSide()
+    {
+        // SimpleClient runs in the test classloader; it needs DatabaseDescriptor initialized here to
+        // build the native-protocol pipeline. Without this, connect() fails with a ClosedChannelException.
+        clientInitialization();
+    }
+
+    /**
+     * Proves the fix directly: under native-transport queue backlog, the load-shed OVERLOADED error is
+     * returned on the stream id of the request that actually timed out, not on stream id 0. Against the
+     * unpatched server this fails with the error arriving on stream id 0.
+     */
+    @Test
+    public void loadShedErrorIsStampedWithRequestStreamId() throws Throwable
+    {
+        try (Cluster cluster = init(Cluster.build().withNodes(1)
+                                           .withInstanceInitializer(SlowSelect::install)
+                                           .withConfig(config -> config.with(GOSSIP, NETWORK, NATIVE_PROTOCOL)
+                                                                       // one NTR thread so a slow query blocks the queue head
+                                                                       .set("native_transport_max_threads", 1)
+                                                                       // short deadline so a queued request is shed quickly
+                                                                       .set("native_transport_timeout", "150ms")
+                                                                       .set("read_request_timeout", "500ms")
+                                                                       .set("range_request_timeout", "500ms"))
+                                           .start()))
+        {
+            cluster.schemaChange(withKeyspace("CREATE TABLE %s.tbl (pk int PRIMARY KEY, v int)"));
+
+            InetSocketHost host = hostAndPort(cluster);
+            ExecutorService executor = Executors.newCachedThreadPool();
+
+            // Two independent connections. A SimpleClient's response queue is a SynchronousQueue, so
+            // concurrent execute() calls on ONE client would race; we use one client to build the
+            // backlog and a second, dedicated client for the request whose stream id we assert on.
+            try (SimpleClient filler = SimpleClient.builder(host.address, host.port)
+                                                   .protocolVersion(ProtocolVersion.V5).build().connect(false);
+                 SimpleClient victim = SimpleClient.builder(host.address, host.port)
+                                                   .protocolVersion(ProtocolVersion.V5).build().connect(false))
+            {
+                // Warm up both connections while the server is still fast.
+                filler.execute(query(withKeyspace("SELECT * FROM %s.tbl"), 1));
+                victim.execute(query(withKeyspace("SELECT * FROM %s.tbl"), 1));
+
+                // From here every non-internal SELECT sleeps 1s, far past the 150ms queue deadline.
+                cluster.get(1).runOnInstance(() -> Assert.assertTrue(SlowSelect.enabled.compareAndSet(false, true)));
+
+                // Saturate the single NTR worker and pile several requests behind it, so anything
+                // enqueued now waits well beyond native_transport_timeout before a worker frees up.
+                List<Future<?>> backlog = new ArrayList<>();
+                for (int i = 0; i < 8; i++)
+                {
+                    int streamId = 10 + i;
+                    backlog.add(executor.submit(() ->
+                                                filler.execute(query(withKeyspace("SELECT * FROM %s.tbl"), streamId), false)));
+                }
+
+                // Let the backlog form and the queue-time clock run past the 150ms deadline.
+                TimeUnit.MILLISECONDS.sleep(800);
+
+                // This request enqueues behind a queue that is already older than the deadline, so the
+                // worker load-sheds it immediately when it is dequeued. It goes out on stream id 42.
+                Message.Response shed =
+                victim.execute(query(withKeyspace("SELECT * FROM %s.tbl"), TIMED_OUT_STREAM_ID), false);
+
+                Assert.assertTrue("Expected an OVERLOADED error, got: " + shed,
+                                  shed instanceof ErrorMessage
+                                  && ((ErrorMessage) shed).error instanceof OverloadedException);
+
+                // The fix: the request went out on stream id 42, and the load-shed error now comes back
+                // on stream id 42 too because Dispatcher stamps every response with the request's id.
+                Assert.assertEquals("Load-shed OVERLOADED error should carry the request's own stream id (" +
+                                    TIMED_OUT_STREAM_ID + ')',
+                                    TIMED_OUT_STREAM_ID, shed.getSource().header.streamId);
+
+                // Drain the backlog so the connection can close cleanly.
+                cluster.get(1).runOnInstance(() -> SlowSelect.enabled.set(false));
+                for (Future<?> f : backlog)
+                {
+                    try
+                    {
+                        f.get(30, TimeUnit.SECONDS);
+                    }
+                    catch (Exception ignored)
+                    { /* shed/timed out */ }
+                }
+            }
+            finally
+            {
+                cluster.get(1).runOnInstance(() -> SlowSelect.enabled.set(false));
+                executor.shutdownNow();
+            }
+        }
+    }
+
+    // ------------------------------------------------------------------------------------------------
+    // Helpers
+    // ------------------------------------------------------------------------------------------------
+
+    private static QueryMessage query(String cql, int streamId)
+    {
+        QueryMessage msg = new QueryMessage(cql, QueryOptions.forInternalCalls(
+        org.apache.cassandra.db.ConsistencyLevel.ONE, Collections.emptyList()));
+        msg.setSource(new Envelope(Envelope.Header.dummy(streamId, Message.Type.QUERY), null));
+        return msg;
+    }
+
+    private static InetSocketHost hostAndPort(Cluster cluster)
+    {
+        // Node 1's native transport binds on its broadcast address; the port is whatever the in-JVM
+        // provisioning assigned (9042 for the default multi-interface strategy). Derive both from the
+        // instance config so this works regardless of provisioning strategy.
+        String address = cluster.get(1).config().broadcastAddress().getAddress().getHostAddress();
+        int port = cluster.get(1).callOnInstance(
+        () -> getNativeTransportPort());
+        return new InetSocketHost(address, port);
+    }
+
+    private static final class InetSocketHost
+    {
+        final String address;
+        final int port;
+
+        InetSocketHost(String address, int port)
+        {
+            this.address = address;
+            this.port = port;
+        }
+    }
+
+    /**
+     * ByteBuddy interceptor that makes client-issued SELECTs sleep past the queue deadline, so a
+     * request queued behind one is load-shed by Dispatcher. Mirrors OverloadTest.SlowSelect.
+     */
+    public static class SlowSelect
+    {
+        static final AtomicBoolean enabled = new AtomicBoolean(false);
+
+        static void install(ClassLoader cl, int nodeNumber)
+        {
+            new ByteBuddy().rebase(SelectStatement.class)
+                           .method(named("execute").and(takesArguments(QueryState.class, QueryOptions.class, Dispatcher.RequestTime.class)))
+                           .intercept(MethodDelegation.to(SlowSelect.class))
+                           .make()
+                           .load(cl, ClassLoadingStrategy.Default.INJECTION);
+        }
+
+        @SuppressWarnings("unused")
+        public static ResultMessage.Rows execute(QueryState state, QueryOptions options,
+                                                 Dispatcher.RequestTime requestTime,
+                                                 @SuperCall Callable<ResultMessage.Rows> zuper) throws Exception
+        {
+            if (enabled.get() && !state.getClientState().isInternal)
+                TimeUnit.SECONDS.sleep(1);
+            return zuper.call();
+        }
+    }
+}

test/distributed/org/apache/cassandra/distributed/test/UnableToParseClientMessageTest.java

@@ -185,9 +185,9 @@ public class UnableToParseClientMessageTest extends TestBaseImpl
         }
 
         @Override
-        public Envelope encode(ProtocolVersion version)
+        public Envelope encode(ProtocolVersion version, int streamId)
         {
-            Envelope base = super.encode(version);
+            Envelope base = super.encode(version, streamId);
             return new CustomHeaderEnvelope(base.header, base.body, headerEncoded);
         }
     }
@@ -228,7 +228,7 @@ public class UnableToParseClientMessageTest extends TestBaseImpl
         }
 
         @Override
-        public Envelope encode(ProtocolVersion version)
+        public Envelope encode(ProtocolVersion version, int streamId)
         {
             Codec<?> originalCodec = type.codec;
             try
@@ -255,7 +255,7 @@ public class UnableToParseClientMessageTest extends TestBaseImpl
                     }
                 });
 
-                return super.encode(version);
+                return super.encode(version, streamId);
             }
             finally
             {

test/unit/org/apache/cassandra/transport/CQLConnectionTest.java

@@ -20,7 +20,9 @@ package org.apache.cassandra.transport;
 
 import java.io.IOException;
 import java.net.InetAddress;
+import java.net.InetSocketAddress;
 import java.net.ServerSocket;
+import java.net.SocketAddress;
 import java.nio.ByteBuffer;
 import java.security.SecureRandom;
 import java.util.HashMap;
@@ -81,18 +83,24 @@ import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelFuture;
+import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.channel.ChannelInitializer;
 import io.netty.channel.ChannelOption;
+import io.netty.channel.ChannelOutboundHandlerAdapter;
 import io.netty.channel.ChannelPipeline;
+import io.netty.channel.ChannelPromise;
+import io.netty.channel.embedded.EmbeddedChannel;
 import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.handler.codec.MessageToMessageDecoder;
+import io.netty.handler.codec.MessageToMessageEncoder;
 
 import static org.apache.cassandra.config.EncryptionOptions.TlsEncryptionPolicy.UNENCRYPTED;
 import static org.apache.cassandra.io.util.FileUtils.ONE_MIB;
 import static org.apache.cassandra.net.FramingTest.randomishBytes;
 import static org.apache.cassandra.transport.Flusher.MAX_FRAMED_PAYLOAD_SIZE;
+import static org.apache.cassandra.transport.PreV5Handlers.getConnectionVersion;
 import static org.apache.cassandra.utils.concurrent.Condition.newOneTimeCondition;
 import static org.apache.cassandra.utils.concurrent.NonBlockingRateLimiter.NO_OP_LIMITER;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -176,6 +184,45 @@ public class CQLConnectionTest
         observer.verifier().accept(0);
     }
 
+    @Test
+    public void handleIncompleteHeaderErrorDuringNegotiation() throws Throwable
+    {
+        // CASSANDRA-21508: a partial header (fewer than Header.LENGTH=9 bytes) carrying an unsupported
+        // protocol version gives the server no stream id to route an error back to. Rather than emit an
+        // unroutable error frame (which could be misapplied to another request), the server closes the
+        // connection. Here the client observes a closed connection with no error frame.
+        int messageCount = 0;
+        Codec codec = Codec.crc(alloc);
+        AllocationObserver observer = new AllocationObserver();
+        InboundProxyHandler.Controller controller = new InboundProxyHandler.Controller();
+        // Truncate the client's STARTUP to a partial header (< 9 bytes) and set an unsupported version.
+        controller.withPayloadTransform(msg -> {
+            ByteBuf bb = (ByteBuf) msg;
+            ByteBuf truncated = bb.copy(0, 4);
+            truncated.setByte(0, 99 & Envelope.PROTOCOL_VERSION_MASK);
+            bb.release();
+            return truncated;
+        });
+
+        ServerConfigurator configurator = ServerConfigurator.builder()
+                                                            .withAllocationObserver(observer)
+                                                            .withProxyController(controller)
+                                                            .build();
+        Server server = server(configurator);
+        Client client = new Client(codec, messageCount);
+        server.start();
+        client.connect(address, port);
+
+        // The server cannot recover a stream id from the incomplete header, so it closes the connection
+        // without sending an error frame.
+        assertFalse(client.isConnected());
+        assertThat(client.getConnectionError()).isNull();
+        server.stop();
+
+        // the failure happens before any capacity is allocated
+        observer.verifier().accept(0);
+    }
+
     @Test
     public void handleFrameCorruptionAfterNegotiation() throws Throwable
     {
@@ -247,6 +294,43 @@ public class CQLConnectionTest
         testFrameCorruption(1, Codec.crc(alloc), envelopeProvider, corruptor, totalBytesPerEnvelope, errorCheck);
     }
 
+    @Test
+    public void fatalErrorFrameIsFlushedBeforeChannelClose()
+    {
+        // CASSANDRA-21508: the post-V5 exception handler must close the connection only AFTER the diagnostic
+        // error frame has been flushed. If it closes synchronously right after writeAndFlush, a write that
+        // can't drain immediately (TCP backpressure, TLS, large frame) is aborted and the client never sees
+        // the error. Simulate "write cannot complete synchronously" with an outbound handler that captures
+        // the write promise without completing it - the handler must not have closed the channel yet.
+        StallingOutboundHandler stall = new StallingOutboundHandler();
+        EmbeddedChannel channel = new EmbeddedChannel()
+        {
+            // client_error_reporting_exclusions (SubnetGroups.contains) requires a real InetSocketAddress;
+            // EmbeddedChannel's default remote address is not one, so supply a loopback address.
+            @Override
+            protected SocketAddress remoteAddress0()
+            {
+                return new InetSocketAddress("127.0.0.1", 9042);
+            }
+        };
+        channel.pipeline().addLast(stall);
+        channel.pipeline().addLast(ExceptionHandlers.postV5Handler(FrameEncoderCrc.instance.allocator(),
+                                                                   ProtocolVersion.V5));
+
+        // Fatal protocol error, no recoverable stream id -> handler encodes an error frame and fatally closes.
+        channel.pipeline().fireExceptionCaught(ProtocolException.toFatalException(new ProtocolException("boom")));
+
+        // The error frame was handed to the outbound pipeline...
+        assertNotNull("expected the handler to write an error frame", stall.writePromise);
+        // ...but since that write has not completed, the channel must still be open. The buggy synchronous
+        // ctx.close() would have already closed it here.
+        assertTrue("channel must not close before the error frame is flushed", channel.isOpen());
+
+        // Once the write completes, the deferred close fires.
+        stall.writePromise.setSuccess();
+        assertFalse("channel should close after the error frame is flushed", channel.isOpen());
+    }
+
     @Test
     public void testAquireAndRelease()
     {
@@ -322,6 +406,48 @@ public class CQLConnectionTest
         runTest(configurator, codec, messageCount, envelopeProvider, responseMatcher, observer.verifier());
     }
 
+    @Test
+    public void testRecoverableBetaFlagEnvelopeErrors()
+    {
+        // CASSANDRA-21508: a V6 (beta) envelope header with the USE_BETA flag unset is an
+        // invalid but *recoverable* protocol error, exactly like an unknown opcode. The server should
+        // return an ERROR on the offending stream id and continue processing subsequent envelopes on the
+        // same connection.
+
+        // every other message advertises V6 with USE_BETA unset and should error while extracting the header
+        IntPredicate shouldError = i -> i % 2 == 0;
+        testBetaFlagEnvelopeErrors(10, shouldError, Codec.crc(alloc));
+        testBetaFlagEnvelopeErrors(10, shouldError, Codec.lz4(alloc));
+
+        testBetaFlagEnvelopeErrors(100, shouldError, Codec.crc(alloc));
+    }
+
+    private void testBetaFlagEnvelopeErrors(int messageCount, IntPredicate shouldError, Codec codec)
+    {
+        TestConsumer consumer = new TestConsumer(new ResultMessage.Void(), codec.encoder);
+        AllocationObserver observer = new AllocationObserver(false);
+        Message.Decoder<Message.Request> decoder = new FixedDecoder();
+
+        // Mutate the erroring streams' headers to advertise the V6 (beta) version with the USE_BETA flag
+        // cleared. In the envelope header, byte 0 is the (direction & version) byte and byte 1 is the flags.
+        int betaBit = 1 << Envelope.Header.Flag.USE_BETA.ordinal();
+        IntFunction<Envelope> envelopeProvider = mutatedEnvelopeProvider(shouldError, b -> {
+            b.put(0, (byte) ProtocolVersion.V6.asInt());     // REQUEST direction, V6 (beta)
+            b.put(1, (byte) (b.get(1) & ~betaBit));          // clear USE_BETA
+        });
+
+        Predicate<Envelope.Header> responseMatcher =
+        h -> (shouldError.test(h.streamId) && h.type == Message.Type.ERROR) || h.type == Message.Type.RESULT;
+
+        ServerConfigurator configurator = ServerConfigurator.builder()
+                                                            .withConsumer(consumer)
+                                                            .withAllocationObserver(observer)
+                                                            .withDecoder(decoder)
+                                                            .build();
+
+        runTest(configurator, codec, messageCount, envelopeProvider, responseMatcher, observer.verifier());
+    }
+
     @Test
     public void testUnrecoverableEnvelopeDecodingErrors()
     {
@@ -580,6 +706,19 @@ public class CQLConnectionTest
         buf.setByte(index, buf.getByte(index) ^ (1 << 4));
     }
 
+    private static class StallingOutboundHandler extends ChannelOutboundHandlerAdapter
+    {
+        volatile ChannelPromise writePromise;
+
+        @Override
+        public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise)
+        {
+            // Capture the write but neither complete nor forward it: simulates a socket that cannot drain
+            // synchronously. The handler's own finally releases the payload, so we don't touch msg.
+            writePromise = promise;
+        }
+    }
+
     private static class MutableEnvelope extends Envelope
     {
         public MutableEnvelope(Envelope source)
@@ -639,7 +778,6 @@ public class CQLConnectionTest
 
             Message.Request request = new OptionsMessage();
             request.setSource(source);
-            request.setStreamId(source.header.streamId);
             Connection connection = channel.attr(Connection.attributeKey).get();
             request.attach(connection);
 
@@ -658,17 +796,17 @@ public class CQLConnectionTest
         TestConsumer(Message.Response fixedResponse, FrameEncoder frameEncoder)
         {
             this.fixedResponse = fixedResponse;
-            this.responseTemplate = fixedResponse.encode(ProtocolVersion.V5);
+            this.responseTemplate = fixedResponse.encode(ProtocolVersion.V5, 0);
             this.frameEncoder = frameEncoder;
         }
 
-        public void dispatch(Channel channel, Message.Request message, Dispatcher.FlushItemConverter toFlushItem, Overload backpressure)
+        public <P> void dispatch(Channel channel, Message.Request message, Dispatcher.FlushItemConverter<P> toFlushItem, P param, Overload backpressure)
         {
             if (flusher == null)
                 flusher = new SimpleClient.SimpleFlusher(frameEncoder);
 
             Envelope response = Envelope.create(responseTemplate.header.type,
-                                                message.getStreamId(),
+                                                message.getSource().header.streamId,
                                                 ProtocolVersion.V5,
                                                 responseTemplate.header.flags,
                                                 responseTemplate.body.copy());
@@ -677,7 +815,7 @@ public class CQLConnectionTest
             // and flush them to the outbound pipeline
             flusher.schedule(channel.pipeline().lastContext());
             // this simulates the release of the allocated resources that a real flusher would do
-            Flusher.FlushItem.Framed item = (Flusher.FlushItem.Framed)toFlushItem.toFlushItem(channel, message, fixedResponse);
+            Flusher.FlushItem.Framed item = (Flusher.FlushItem.Framed)toFlushItem.toFlushItem(param, channel, message, fixedResponse);
             item.release();
         }
 
@@ -984,6 +1122,21 @@ public class CQLConnectionTest
         }
     }
 
+    /**
+     * Simple adaptor to plug CQL message encoding into pre-V5 pipelines
+     */
+    @ChannelHandler.Sharable
+    public static class ProtocolEncoder extends MessageToMessageEncoder<Message>
+    {
+        public static final ProtocolEncoder instance = new ProtocolEncoder();
+        private ProtocolEncoder(){}
+        public void encode(ChannelHandlerContext ctx, Message source, List<Object> results)
+        {
+            ProtocolVersion version = getConnectionVersion(ctx);
+            results.add(source.encode(version, 0));
+        }
+    }
+
     static class Client
     {
         private final Codec codec;
@@ -1022,7 +1175,7 @@ public class CQLConnectionTest
                     ChannelPipeline pipeline = channel.pipeline();
                     // Outbound handlers to enable us to send the initial STARTUP
                     pipeline.addLast("envelopeEncoder", Envelope.Encoder.instance);
-                    pipeline.addLast("messageEncoder", PreV5Handlers.ProtocolEncoder.instance);
+                    pipeline.addLast("messageEncoder", ProtocolEncoder.instance);
                     pipeline.addLast("envelopeDecoder", new Envelope.Decoder());
                     // Inbound handler to perform the handshake & modify the pipeline on receipt of a READY
                     pipeline.addLast("handshake", new MessageToMessageDecoder<Envelope>()
@@ -1100,6 +1253,17 @@ public class CQLConnectionTest
                             flusher.schedule(channel.pipeline().lastContext());
                             ready.countDown();
                         }
+
+                        @Override
+                        public void channelInactive(ChannelHandlerContext ctx)
+                        {
+                            // If the server closes the connection during negotiation (e.g. an unroutable
+                            // protocol error that cannot be answered with a stream id), unblock connect()
+                            // so the test observes the disconnection rather than hanging for a READY/ERROR.
+                            connected = false;
+                            ready.countDown();
+                            ctx.fireChannelInactive();
+                        }
                     });
                 }
             });

test/unit/org/apache/cassandra/transport/ErrorMessageTest.java

@@ -67,7 +67,7 @@ public class ErrorMessageTest extends EncodeAndDecodeTestBase<ErrorMessage>
         boolean dataPresent = false;
         ReadFailureException rfe = new ReadFailureException(consistencyLevel, receivedBlockFor, receivedBlockFor, dataPresent, failureReasonMap1);
 
-        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromException(rfe), ProtocolVersion.V5);
+        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromTransportException(rfe), ProtocolVersion.V5);
         ReadFailureException deserializedRfe = (ReadFailureException) deserialized.error;
 
         assertEquals(failureReasonMap1, deserializedRfe.failureReasonByEndpoint);
@@ -85,7 +85,7 @@ public class ErrorMessageTest extends EncodeAndDecodeTestBase<ErrorMessage>
         WriteType writeType = WriteType.SIMPLE;
         WriteFailureException wfe = new WriteFailureException(consistencyLevel, receivedBlockFor, receivedBlockFor, writeType, failureReasonMap2);
 
-        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromException(wfe), ProtocolVersion.V5);
+        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromTransportException(wfe), ProtocolVersion.V5);
         WriteFailureException deserializedWfe = (WriteFailureException) deserialized.error;
 
         assertEquals(failureReasonMap2, deserializedWfe.failureReasonByEndpoint);
@@ -103,7 +103,7 @@ public class ErrorMessageTest extends EncodeAndDecodeTestBase<ErrorMessage>
         ConsistencyLevel consistencyLevel = ConsistencyLevel.SERIAL;
         CasWriteTimeoutException ex = new CasWriteTimeoutException(WriteType.CAS, consistencyLevel, 0, receivedBlockFor, contentions);
 
-        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromException(ex), ProtocolVersion.V5);
+        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromTransportException(ex), ProtocolVersion.V5);
         assertTrue(deserialized.error instanceof CasWriteTimeoutException);
         CasWriteTimeoutException deserializedEx = (CasWriteTimeoutException) deserialized.error;
 
@@ -124,7 +124,7 @@ public class ErrorMessageTest extends EncodeAndDecodeTestBase<ErrorMessage>
         ConsistencyLevel consistencyLevel = ConsistencyLevel.SERIAL;
         CasWriteTimeoutException ex = new CasWriteTimeoutException(WriteType.CAS, consistencyLevel, receivedBlockFor, receivedBlockFor, contentions);
 
-        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromException(ex), ProtocolVersion.V4);
+        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromTransportException(ex), ProtocolVersion.V4);
         assertTrue(deserialized.error instanceof WriteTimeoutException);
         assertFalse(deserialized.error instanceof CasWriteTimeoutException);
         WriteTimeoutException deserializedEx = (WriteTimeoutException) deserialized.error;
@@ -142,7 +142,7 @@ public class ErrorMessageTest extends EncodeAndDecodeTestBase<ErrorMessage>
         ConsistencyLevel consistencyLevel = ConsistencyLevel.SERIAL;
         CasWriteUnknownResultException ex = new CasWriteUnknownResultException(consistencyLevel, receivedBlockFor, receivedBlockFor);
 
-        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromException(ex), ProtocolVersion.V5);
+        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromTransportException(ex), ProtocolVersion.V5);
         assertTrue(deserialized.error instanceof CasWriteUnknownResultException);
         CasWriteUnknownResultException deserializedEx = (CasWriteUnknownResultException) deserialized.error;
 
@@ -160,7 +160,7 @@ public class ErrorMessageTest extends EncodeAndDecodeTestBase<ErrorMessage>
         ConsistencyLevel consistencyLevel = ConsistencyLevel.SERIAL;
         CasWriteUnknownResultException ex = new CasWriteUnknownResultException(consistencyLevel, receivedBlockFor, receivedBlockFor);
 
-        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromException(ex), ProtocolVersion.V4);
+        ErrorMessage deserialized = encodeThenDecode(ErrorMessage.fromTransportException(ex), ProtocolVersion.V4);
         assertTrue(deserialized.error instanceof WriteTimeoutException);
         assertFalse(deserialized.error instanceof CasWriteUnknownResultException);
         WriteTimeoutException deserializedEx = (WriteTimeoutException) deserialized.error;

test/unit/org/apache/cassandra/transport/MessageDispatcherTest.java

@@ -157,7 +157,7 @@ public class MessageDispatcherTest
     public long tryAuth(Callable<Long> check, Message.Request request) throws Exception
     {
         long start = check.call();
-        dispatch.dispatch(null, request, (channel,req,response) -> null, ClientResourceLimits.Overload.NONE);
+        dispatch.dispatch(null, request, (p, channel, req, response) -> null, null, ClientResourceLimits.Overload.NONE);
 
         // While this is timeout based, we should be *well below* a full second on any of this processing in any sane environment.
         long timeout = System.currentTimeMillis();
@@ -176,9 +176,10 @@ public class MessageDispatcherTest
         }
 
         @Override
-        void processRequest(Channel channel,
+        <P> void processRequest(Channel channel,
                             Message.Request request,
-                            FlushItemConverter forFlusher,
+                            FlushItemConverter<P> forFlusher,
+                            P param,
                             ClientResourceLimits.Overload backpressure,
                             RequestTime requestTime)
         {

test/unit/org/apache/cassandra/transport/ProtocolErrorTest.java

@@ -17,6 +17,7 @@
  */
 package org.apache.cassandra.transport;
 
+import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -73,7 +74,12 @@ public class ProtocolErrorTest {
         try {
             dec.decode(null, buf, results);
             Assert.fail("Expected protocol error");
-        } catch (ProtocolException e) {
+        } catch (ErrorMessage.WrappedException e) {
+            // The stream id is recovered using the attempted version's header layout (CASSANDRA-21508):
+            // v1/v2 use a single-byte stream id (header byte 2, which is 0x00 here), whereas v3+ use a
+            // two-byte id (0x0001). See decodeRecoversSingleByteStreamIdForOldProtocolVersions.
+            int expectedStreamId = version < ProtocolVersion.V3.asInt() ? 0x00 : 0x01;
+            Assert.assertEquals(expectedStreamId, e.getStreamId());
             Assert.assertTrue(e.getMessage().contains("Invalid or unsupported protocol version"));
         }
     }
@@ -101,6 +107,108 @@ public class ProtocolErrorTest {
         }
     }
 
+    @Test
+    public void testIncompleteHeaderWithInvalidProtocolVersion() throws Exception
+    {
+        // CASSANDRA-21508: when fewer than a full header's worth of bytes have arrived AND the protocol
+        // version is unsupported, decode cannot trust/recover a stream id. It defers the protocol error and
+        // wraps it with Message.UNSET_STREAM_ID, so the channel-level exception handler closes the connection
+        // rather than emit an unroutable error frame. Exercises Envelope.Decoder.decode() lines 387-398.
+        Envelope.Decoder dec = new Envelope.Decoder();
+
+        List<Object> results = new ArrayList<>();
+        // Unsupported version (two above CURRENT) with only part of a header - fewer than Header.LENGTH (9) bytes.
+        byte[] bytes = new byte[] {
+                (byte) REQUEST.addToVersion(ProtocolVersion.CURRENT.asInt() + 2),  // direction & unsupported version
+                0x00,        // flags
+                0x00, 0x2a,  // partial header, truncated before the length field
+        };
+        ByteBuf buf = Unpooled.wrappedBuffer(bytes);
+        try {
+            dec.decode(null, buf, results);
+            Assert.fail("Expected protocol error");
+        } catch (ProtocolException e) {
+            Assert.assertTrue(e.getMessage().contains("Invalid or unsupported protocol version"));
+        }
+    }
+
+    @Test
+    public void extractHeaderReturnsRecoverableErrorOnBetaFlagViolation() throws Exception
+    {
+        // CASSANDRA-21508: on the framed (V5+) path, a V6 (beta) envelope whose USE_BETA flag is
+        // unset is an invalid but *recoverable* protocol error. Envelope.Decoder.extractHeader documents that
+        // it never throws, and instead returns a HeaderExtractionResult carrying the frame's stream id, so the
+        // caller (CQLMessageHandler.processOneContainedMessage) can route an ERROR back on that stream and keep
+        // processing subsequent envelopes.
+        Envelope.Decoder dec = new Envelope.Decoder();
+
+        int streamId = 42;
+        // A complete 9-byte header advertising V6 (beta) as a REQUEST, with USE_BETA deliberately unset.
+        byte[] header = new byte[] {
+                (byte) REQUEST.addToVersion(ProtocolVersion.V6.asInt()),  // direction & beta version
+                0x00,        // flags - USE_BETA deliberately unset
+                0x00, 0x2a,  // stream id = 42
+                0x05,        // opcode = OPTIONS
+                0x00, 0x00, 0x00, 0x00,  // body length = 0
+        };
+        ByteBuffer buf = ByteBuffer.wrap(header);
+
+        Envelope.Decoder.HeaderExtractionResult result;
+        try
+        {
+            result = dec.extractHeader(buf);
+        }
+        catch (ErrorMessage.WrappedException e)
+        {
+            throw new AssertionError("extractHeader must not throw on a USE_BETA violation; a recoverable " +
+                                     "protocol error escaped as a WrappedException and would tear down the " +
+                                     "connection instead of returning a routable error", e);
+        }
+
+        Assert.assertFalse("A USE_BETA violation should be reported as a (recoverable) extraction error",
+                           result.isSuccess());
+        Assert.assertEquals("the frame's stream id must be preserved so the error can be routed back",
+                            streamId, result.streamId());
+        Assert.assertTrue("expected a USE_BETA protocol error, got: " + result.error().getMessage(),
+                          result.error().getMessage().contains("USE_BETA flag is unset"));
+    }
+
+    @Test
+    public void decodeRecoversSingleByteStreamIdForOldProtocolVersions() throws Exception
+    {
+        // CASSANDRA-21508: protocol versions 1 and 2 use a shorter header whose stream id is a single
+        // byte. When such a version is rejected, decode must recover the stream id using that layout.
+        // Reading a 16-bit value would splice the stream id byte together with the following opcode byte and
+        // recover a bogus id (e.g. stream 0 + STARTUP opcode 1 -> stream 1), routing the error to a stream the
+        // client never used. A v1/v2 client (which downgrades on such an error) would then never see it and
+        // its connection would time out - as observed in ProtocolNegotiationTest#olderVersionsAreUnsupported.
+        Envelope.Decoder dec = new Envelope.Decoder();
+
+        int streamId = 0x07;
+        List<Object> results = new ArrayList<>();
+        // A full (>= Header.LENGTH bytes) v1 frame so decode reaches the stream-id recovery, not the
+        // incomplete-header path: single-byte stream id, STARTUP opcode, and one body byte.
+        byte[] bytes = new byte[] {
+                (byte) REQUEST.addToVersion(1),  // unsupported v1
+                0x00,                            // flags
+                (byte) streamId,                 // v1/v2 single-byte stream id
+                0x01,                            // opcode = STARTUP (would be spliced in as the low byte)
+                0x00, 0x00, 0x00, 0x01,          // body length = 1
+                0x00                             // 1 body byte, so readableBytes == Header.LENGTH
+        };
+        ByteBuf buf = Unpooled.wrappedBuffer(bytes);
+        try {
+            dec.decode(null, buf, results);
+            Assert.fail("Expected protocol error");
+        } catch (ErrorMessage.WrappedException e) {
+            Assert.assertEquals("a v1/v2 stream id must be read as a single byte, not spliced with the opcode",
+                                streamId, e.getStreamId());
+            Assert.assertTrue("expected a ProtocolException cause, got: " + e.getCause(),
+                              e.getCause() instanceof ProtocolException);
+            Assert.assertTrue(e.getMessage().contains("Invalid or unsupported protocol version"));
+        }
+    }
+
     @Test
     public void testInvalidDirection() throws Exception
     {
@@ -161,7 +269,7 @@ public class ProtocolErrorTest {
     public void testErrorMessageWithNullString()
     {
         // test for CASSANDRA-11167
-        ErrorMessage msg = ErrorMessage.fromException(new ServerError((String) null));
+        ErrorMessage msg = ErrorMessage.fromTransportException(new ServerError((String) null));
         assert msg.toString().endsWith("null") : msg.toString();
         int size = ErrorMessage.codec.encodedSize(msg, ProtocolVersion.CURRENT);
         ByteBuf buf = Unpooled.buffer(size);

test/unit/org/apache/cassandra/transport/ProtocolNegotiationTest.java

@@ -126,29 +126,31 @@ public class ProtocolNegotiationTest extends CQLTester
             for (int i = 0; i < 100; i++)
             {
                 int streamId = random.nextInt(254) + 1;
-                options.setStreamId(streamId);
+                options.setSource(new Envelope(Envelope.Header.dummy(streamId, options.type), null));
                 Message.Response response = client.execute(options);
+                // The stream id is stamped onto the frame at encoding time; verify it round-trips by
+                // reading it back off the response's source envelope.
                 assertEquals(String.format("StreamId mismatch; version: %s, seed: %s, iter: %s, expected: %s, actual: %s",
-                                           version, seed, i, streamId, response.getStreamId()),
+                                           version, seed, i, streamId, response.getSource().header.streamId),
-                             streamId, response.getStreamId());
+                             streamId, response.getSource().header.streamId);
             }
 
             int streamId = random.nextInt(254) + 1;
             // STARTUP messages are handled by the initial connection handler
             StartupMessage startup = new StartupMessage(ImmutableMap.of(CQL_VERSION, QueryProcessor.CQL_VERSION.toString()));
-            startup.setStreamId(streamId);
+            startup.setSource(new Envelope(Envelope.Header.dummy(streamId, startup.type), null));
             Message.Response response = client.execute(startup);
             assertEquals(String.format("StreamId mismatch after negotiation; version: %s, expected: %s, actual %s",
-                                       version, streamId, response.getStreamId()),
+                                       version, streamId, response.getSource().header.streamId),
-                         streamId, response.getStreamId());
+                         streamId, response.getSource().header.streamId);
 
             // Following STARTUP, the version specific handlers are fully responsible for processing messages
             QueryMessage query = new QueryMessage("SELECT * FROM system.local", QueryOptions.DEFAULT);
-            query.setStreamId(streamId);
+            query.setSource(new Envelope(Envelope.Header.dummy(streamId, query.type), null));
             response = client.execute(query);
             assertEquals(String.format("StreamId mismatch after negotiation; version: %s, expected: %s, actual %s",
-                                       version, streamId, response.getStreamId()),
+                                       version, streamId, response.getSource().header.streamId),
-                         streamId, response.getStreamId());
+                         streamId, response.getSource().header.streamId);
         }
         catch (IOException e)
         {
@@ -215,9 +217,9 @@ public class ProtocolNegotiationTest extends CQLTester
             QueryMessage query = new QueryMessage("SELECT * FROM system.local", QueryOptions.DEFAULT)
             {
                 @Override
-                public Envelope encode(ProtocolVersion originalVersion)
+                public Envelope encode(ProtocolVersion originalVersion, int streamId)
                 {
-                    return super.encode(v);
+                    return super.encode(v, streamId);
                 }
             };
             try