add a function to reset the connection flow controller

internal/flowcontrol/connection_flow_controller.go

@@ -1,6 +1,7 @@
 package flowcontrol
 
 import (
+	"errors"
 	"fmt"
 	"time"
 
@@ -86,3 +87,18 @@ func (c *connectionFlowController) EnsureMinimumWindowSize(inc protocol.ByteCoun
 	}
 	c.mutex.Unlock()
 }
+
+// The flow controller is reset when 0-RTT is rejected.
+// All stream data is invalidated, it's if we had never opened a stream and never sent any data.
+// At that point, we only have sent stream data, but we didn't have the keys to open 1-RTT keys yet.
+func (c *connectionFlowController) Reset() error {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	if c.bytesRead > 0 || c.highestReceived > 0 || !c.epochStartTime.IsZero() {
+		return errors.New("flow controller reset after reading data")
+	}
+	c.bytesSent = 0
+	c.lastBlockedAt = 0
+	return nil
+}

internal/flowcontrol/connection_flow_controller_test.go

@@ -129,4 +129,28 @@ var _ = Describe("Connection Flow controller", func() {
 			Expect(controller.epochStartTime).To(BeTemporally("~", time.Now(), 100*time.Millisecond))
 		})
 	})
+
+	Context("resetting", func() {
+		It("resets", func() {
+			const initialWindow protocol.ByteCount = 1337
+			controller.UpdateSendWindow(initialWindow)
+			controller.AddBytesSent(1000)
+			Expect(controller.SendWindowSize()).To(Equal(initialWindow - 1000))
+			Expect(controller.Reset()).To(Succeed())
+			Expect(controller.SendWindowSize()).To(Equal(initialWindow))
+		})
+
+		It("says if is blocked after resetting", func() {
+			const initialWindow protocol.ByteCount = 1337
+			controller.UpdateSendWindow(initialWindow)
+			controller.AddBytesSent(initialWindow)
+			blocked, _ := controller.IsNewlyBlocked()
+			Expect(blocked).To(BeTrue())
+			Expect(controller.Reset()).To(Succeed())
+			controller.AddBytesSent(initialWindow)
+			blocked, blockedAt := controller.IsNewlyBlocked()
+			Expect(blocked).To(BeTrue())
+			Expect(blockedAt).To(Equal(initialWindow))
+		})
+	})
 })

internal/flowcontrol/interface.go

@@ -29,6 +29,7 @@ type StreamFlowController interface {
 // The ConnectionFlowController is the flow controller for the connection.
 type ConnectionFlowController interface {
 	flowController
+	Reset() error
 }
 
 type connectionFlowControllerI interface {

internal/mocks/connection_flow_controller.go

@@ -87,6 +87,20 @@ func (mr *MockConnectionFlowControllerMockRecorder) IsNewlyBlocked() *gomock.Cal
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsNewlyBlocked", reflect.TypeOf((*MockConnectionFlowController)(nil).IsNewlyBlocked))
 }
 
+// Reset mocks base method.
+func (m *MockConnectionFlowController) Reset() error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Reset")
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Reset indicates an expected call of Reset.
+func (mr *MockConnectionFlowControllerMockRecorder) Reset() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Reset", reflect.TypeOf((*MockConnectionFlowController)(nil).Reset))
+}
+
 // SendWindowSize mocks base method.
 func (m *MockConnectionFlowController) SendWindowSize() protocol.ByteCount {
 	m.ctrl.T.Helper()