From c741b6fc092ed77a5cd3669bb248b208c32d62af Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 1 Mar 2021 11:23:19 +0800
Subject: [PATCH] drop STREAM and *_BLOCKED frames from queue when 0-RTT is
 rejected

---
 framer.go      | 29 +++++++++++++++++++++++++++++
 framer_test.go | 34 +++++++++++++++++++++++++++++++---
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/framer.go b/framer.go
index 4fbb9ef9..29d36b85 100644
--- a/framer.go
+++ b/framer.go
@@ -1,6 +1,7 @@
 package quic
 
 import (
+	"errors"
 	"sync"
 
 	"github.com/lucas-clemente/quic-go/internal/ackhandler"
@@ -17,6 +18,8 @@ type framer interface {
 
 	AddActiveStream(protocol.StreamID)
 	AppendStreamFrames([]ackhandler.Frame, protocol.ByteCount) ([]ackhandler.Frame, protocol.ByteCount)
+
+	Handle0RTTRejection() error
 }
 
 type framerI struct {
@@ -140,3 +143,29 @@ func (f *framerI) AppendStreamFrames(frames []ackhandler.Frame, maxLen protocol.
 	}
 	return frames, length
 }
+
+func (f *framerI) Handle0RTTRejection() error {
+	f.mutex.Lock()
+	defer f.mutex.Unlock()
+
+	f.controlFrameMutex.Lock()
+	f.streamQueue = f.streamQueue[:0]
+	for id := range f.activeStreams {
+		delete(f.activeStreams, id)
+	}
+	var j int
+	for i, frame := range f.controlFrames {
+		switch frame.(type) {
+		case *wire.MaxDataFrame, *wire.MaxStreamDataFrame, *wire.MaxStreamsFrame:
+			return errors.New("didn't expect MAX_DATA / MAX_STREAM_DATA / MAX_STREAMS frame to be sent in 0-RTT")
+		case *wire.DataBlockedFrame, *wire.StreamDataBlockedFrame, *wire.StreamsBlockedFrame:
+			continue
+		default:
+			f.controlFrames[j] = f.controlFrames[i]
+			j++
+		}
+	}
+	f.controlFrames = f.controlFrames[:j]
+	f.controlFrameMutex.Unlock()
+	return nil
+}
diff --git a/framer_test.go b/framer_test.go
index 28590275..71f7b2aa 100644
--- a/framer_test.go
+++ b/framer_test.go
@@ -2,14 +2,14 @@ package quic
 
 import (
 	"bytes"
+	"math/rand"
 
 	"github.com/lucas-clemente/quic-go/internal/ackhandler"
-
-	"github.com/golang/mock/gomock"
-
 	"github.com/lucas-clemente/quic-go/internal/protocol"
 	"github.com/lucas-clemente/quic-go/internal/wire"
 
+	"github.com/golang/mock/gomock"
+
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 )
@@ -86,6 +86,26 @@ var _ = Describe("Framer", func() {
 			Expect(frames).To(HaveLen(1))
 			Expect(length).To(Equal(bfLen))
 		})
+
+		It("drops *_BLOCKED frames when 0-RTT is rejected", func() {
+			ping := &wire.PingFrame{}
+			ncid := &wire.NewConnectionIDFrame{SequenceNumber: 10, ConnectionID: protocol.ConnectionID{0xde, 0xad, 0xbe, 0xef}}
+			frames := []wire.Frame{
+				&wire.DataBlockedFrame{MaximumData: 1337},
+				&wire.StreamDataBlockedFrame{StreamID: 42, MaximumStreamData: 1337},
+				&wire.StreamsBlockedFrame{StreamLimit: 13},
+				ping,
+				ncid,
+			}
+			rand.Shuffle(len(frames), func(i, j int) { frames[i], frames[j] = frames[j], frames[i] })
+			for _, f := range frames {
+				framer.QueueControlFrame(f)
+			}
+			Expect(framer.Handle0RTTRejection()).To(Succeed())
+			fs, length := framer.AppendControlFrames(nil, protocol.MaxByteCount)
+			Expect(fs).To(HaveLen(2))
+			Expect(length).To(Equal(ping.Length(version) + ncid.Length(version)))
+		})
 	})
 
 	Context("popping STREAM frames", func() {
@@ -353,5 +373,13 @@ var _ = Describe("Framer", func() {
 			Expect(fs[0].Frame).To(Equal(f))
 			Expect(length).To(Equal(f.Length(version)))
 		})
+
+		It("drops all STREAM frames when 0-RTT is rejected", func() {
+			framer.AddActiveStream(id1)
+			Expect(framer.Handle0RTTRejection()).To(Succeed())
+			fs, length := framer.AppendStreamFrames(nil, protocol.MaxByteCount)
+			Expect(fs).To(BeEmpty())
+			Expect(length).To(BeZero())
+		})
 	})
 })