From dee2dbe83f7a495adf957c8d40e1a63d5374158e Mon Sep 17 00:00:00 2001
From: baranowb <baranowb@gmail.com>
Date: Thu, 19 Sep 2024 10:26:26 +0200
Subject: [PATCH] [UNDERTOW-2444] Fix RST scenario violation in H2

---
 .../protocols/http2/Http2Channel.java         | 39 +++++++++++++++++--
 1 file changed, 36 insertions(+), 3 deletions(-)

diff --git a/core/src/main/java/io/undertow/protocols/http2/Http2Channel.java b/core/src/main/java/io/undertow/protocols/http2/Http2Channel.java
index c1c7cca769..da120f0fc9 100644
--- a/core/src/main/java/io/undertow/protocols/http2/Http2Channel.java
+++ b/core/src/main/java/io/undertow/protocols/http2/Http2Channel.java
@@ -1225,6 +1225,7 @@ private StreamHolder handleRstStream(int streamId, boolean receivedRst) {
             resetStreamTracker.store(streamId, holder);
             if(streamId % 2 == (isClient() ? 1 : 0)) {
                 sendConcurrentStreamsAtomicUpdater.getAndDecrement(this);
+                holder.sent = true;
             } else {
                 receiveConcurrentStreamsAtomicUpdater.getAndDecrement(this);
             }
@@ -1239,13 +1240,14 @@ private StreamHolder handleRstStream(int streamId, boolean receivedRst) {
                     //Server side originated, no input from client other than RST
                     //this can happen on page refresh when push happens, but client
                     //still has valid cache entry
+                    //NOTE: this is specific case when its set.
                     holder.resetByPeer = receivedRst;
                 } else {
                     handleRstWindow();
                 }
             }
         } else if(receivedRst){
-            final StreamHolder resetStream = resetStreamTracker.find(streamId);
+            final StreamHolder resetStream = resetStreamTracker.find(streamId, true);
             if(resetStream != null && resetStream.resetByPeer) {
                 //This means other side reset stream at some point.
                 //depending on peer or network latency our frames might be late and
@@ -1378,6 +1380,10 @@ private static final class StreamHolder {
          * This flag is set only in case of short lived server push that was reset by remote end.
          */
         boolean resetByPeer = false;
+        /**
+         * flag indicate whether our side originated. This is done for caching purposes, handlng differs.
+         */
+        boolean sent = false;
         Http2StreamSourceChannel sourceChannel;
         Http2StreamSinkChannel sinkChannel;
 
@@ -1388,6 +1394,13 @@ private static final class StreamHolder {
         StreamHolder(Http2StreamSinkChannel sinkChannel) {
             this.sinkChannel = sinkChannel;
         }
+
+        @Override
+        public String toString() {
+            return "StreamHolder [sourceClosed=" + sourceClosed + ", sinkClosed=" + sinkClosed + ", resetByPeer=" + resetByPeer
+                    + ", sent=" + sent + ", sourceChannel=" + sourceChannel + ", sinkChannel=" + sinkChannel + "]";
+        }
+
     }
 
     // cache that keeps track of streams until they can be evicted @see Http2Channel#RST_STREAM_EVICATION_TIME
@@ -1403,12 +1416,27 @@ private void store(int streamId, StreamHolder streamHolder) {
             streamHolders.put(streamId, streamHolder);
             entries.add(new StreamCacheEntry(streamId));
         }
-        private StreamHolder find(int streamId) {
+
+        /**
+         * Method will return only sent
+         * @param streamId
+         * @return
+         */
+        private StreamHolder find(final int streamId) {
+            return find(streamId, false);
+        }
+
+        private StreamHolder find(final int streamId, final boolean all) {
             for (Iterator<StreamCacheEntry> iterator = entries.iterator(); iterator.hasNext();) {
                 StreamCacheEntry entry = iterator.next();
                 if (entry.shouldEvict()) {
                     iterator.remove();
                     StreamHolder holder = streamHolders.remove(entry.streamId);
+                    if(!holder.sent || holder.resetByPeer) {
+                        //if its not our end of chain, its just cached, so we only cache for purpose of
+                        // handling eager RST
+                        continue;
+                    }
                     AbstractHttp2StreamSourceChannel receiver = holder.sourceChannel;
                     if(receiver != null) {
                         IoUtils.safeClose(receiver);
@@ -1422,7 +1450,12 @@ private StreamHolder find(int streamId) {
                     }
                 } else break;
             }
-            return streamHolders.get(streamId);
+            final StreamHolder holder = streamHolders.get(streamId);
+            if(holder != null && (!all && !holder.sent)) {
+                return null;
+            } else {
+                return holder;
+            }
         }
 
         private Map<Integer, StreamHolder> getStreamHolders() {