use new gomock feature to generate type-safe methods in mocks (#4057)

2025-04-04 04:37:36 +03:00 · 2023-10-21 12:55:33 +07:00 · 2023-10-21 12:55:33 +07:00 · a263164d9f
commit a263164d9f
parent b344940f06
53 changed files with 7933 additions and 754 deletions
--- a/internal/ackhandler/mock_sent_packet_tracker_test.go
+++ b/internal/ackhandler/mock_sent_packet_tracker_test.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package ackhandler -destination mock_sent_packet_tracker_test.go github.com/quic-go/quic-go/internal/ackhandler SentPacketTracker
+//	mockgen -typed -build_flags=-tags=gomock -package ackhandler -destination mock_sent_packet_tracker_test.go github.com/quic-go/quic-go/internal/ackhandler SentPacketTracker
 //
 // Package ackhandler is a generated GoMock package.
 package ackhandler
@ -47,9 +47,33 @@ func (m *MockSentPacketTracker) GetLowestPacketNotConfirmedAcked() protocol.Pack
 }

 // GetLowestPacketNotConfirmedAcked indicates an expected call of GetLowestPacketNotConfirmedAcked.
-func (mr *MockSentPacketTrackerMockRecorder) GetLowestPacketNotConfirmedAcked() *gomock.Call {
+func (mr *MockSentPacketTrackerMockRecorder) GetLowestPacketNotConfirmedAcked() *SentPacketTrackerGetLowestPacketNotConfirmedAckedCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLowestPacketNotConfirmedAcked", reflect.TypeOf((*MockSentPacketTracker)(nil).GetLowestPacketNotConfirmedAcked))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLowestPacketNotConfirmedAcked", reflect.TypeOf((*MockSentPacketTracker)(nil).GetLowestPacketNotConfirmedAcked))
+	return &SentPacketTrackerGetLowestPacketNotConfirmedAckedCall{Call: call}
+}
+
+// SentPacketTrackerGetLowestPacketNotConfirmedAckedCall wrap *gomock.Call
+type SentPacketTrackerGetLowestPacketNotConfirmedAckedCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketTrackerGetLowestPacketNotConfirmedAckedCall) Return(arg0 protocol.PacketNumber) *SentPacketTrackerGetLowestPacketNotConfirmedAckedCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketTrackerGetLowestPacketNotConfirmedAckedCall) Do(f func() protocol.PacketNumber) *SentPacketTrackerGetLowestPacketNotConfirmedAckedCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketTrackerGetLowestPacketNotConfirmedAckedCall) DoAndReturn(f func() protocol.PacketNumber) *SentPacketTrackerGetLowestPacketNotConfirmedAckedCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // ReceivedPacket mocks base method.
@ -59,7 +83,31 @@ func (m *MockSentPacketTracker) ReceivedPacket(arg0 protocol.EncryptionLevel) {
 }

 // ReceivedPacket indicates an expected call of ReceivedPacket.
-func (mr *MockSentPacketTrackerMockRecorder) ReceivedPacket(arg0 any) *gomock.Call {
+func (mr *MockSentPacketTrackerMockRecorder) ReceivedPacket(arg0 any) *SentPacketTrackerReceivedPacketCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReceivedPacket", reflect.TypeOf((*MockSentPacketTracker)(nil).ReceivedPacket), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReceivedPacket", reflect.TypeOf((*MockSentPacketTracker)(nil).ReceivedPacket), arg0)
+	return &SentPacketTrackerReceivedPacketCall{Call: call}
+}
+
+// SentPacketTrackerReceivedPacketCall wrap *gomock.Call
+type SentPacketTrackerReceivedPacketCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketTrackerReceivedPacketCall) Return() *SentPacketTrackerReceivedPacketCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketTrackerReceivedPacketCall) Do(f func(protocol.EncryptionLevel)) *SentPacketTrackerReceivedPacketCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketTrackerReceivedPacketCall) DoAndReturn(f func(protocol.EncryptionLevel)) *SentPacketTrackerReceivedPacketCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/ackhandler/mockgen.go
+++ b/internal/ackhandler/mockgen.go
@ -2,7 +2,7 @@

 package ackhandler

-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\"  -package ackhandler -destination mock_sent_packet_tracker_test.go github.com/quic-go/quic-go/internal/ackhandler SentPacketTracker"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\"  -package ackhandler -destination mock_sent_packet_tracker_test.go github.com/quic-go/quic-go/internal/ackhandler SentPacketTracker"
 type SentPacketTracker = sentPacketTracker

 //go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\"  -package ackhandler -destination mock_ecn_handler_test.go github.com/quic-go/quic-go/internal/ackhandler ECNHandler"
--- a/internal/mocks/ackhandler/received_packet_handler.go
+++ b/internal/mocks/ackhandler/received_packet_handler.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mockackhandler -destination ackhandler/received_packet_handler.go github.com/quic-go/quic-go/internal/ackhandler ReceivedPacketHandler
+//	mockgen -typed -build_flags=-tags=gomock -package mockackhandler -destination ackhandler/received_packet_handler.go github.com/quic-go/quic-go/internal/ackhandler ReceivedPacketHandler
 //
 // Package mockackhandler is a generated GoMock package.
 package mockackhandler
@ -47,9 +47,33 @@ func (m *MockReceivedPacketHandler) DropPackets(arg0 protocol.EncryptionLevel) {
 }

 // DropPackets indicates an expected call of DropPackets.
-func (mr *MockReceivedPacketHandlerMockRecorder) DropPackets(arg0 any) *gomock.Call {
+func (mr *MockReceivedPacketHandlerMockRecorder) DropPackets(arg0 any) *ReceivedPacketHandlerDropPacketsCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DropPackets", reflect.TypeOf((*MockReceivedPacketHandler)(nil).DropPackets), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DropPackets", reflect.TypeOf((*MockReceivedPacketHandler)(nil).DropPackets), arg0)
+	return &ReceivedPacketHandlerDropPacketsCall{Call: call}
+}
+
+// ReceivedPacketHandlerDropPacketsCall wrap *gomock.Call
+type ReceivedPacketHandlerDropPacketsCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ReceivedPacketHandlerDropPacketsCall) Return() *ReceivedPacketHandlerDropPacketsCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ReceivedPacketHandlerDropPacketsCall) Do(f func(protocol.EncryptionLevel)) *ReceivedPacketHandlerDropPacketsCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ReceivedPacketHandlerDropPacketsCall) DoAndReturn(f func(protocol.EncryptionLevel)) *ReceivedPacketHandlerDropPacketsCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // GetAckFrame mocks base method.
@ -61,9 +85,33 @@ func (m *MockReceivedPacketHandler) GetAckFrame(arg0 protocol.EncryptionLevel, a
 }

 // GetAckFrame indicates an expected call of GetAckFrame.
-func (mr *MockReceivedPacketHandlerMockRecorder) GetAckFrame(arg0, arg1 any) *gomock.Call {
+func (mr *MockReceivedPacketHandlerMockRecorder) GetAckFrame(arg0, arg1 any) *ReceivedPacketHandlerGetAckFrameCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAckFrame", reflect.TypeOf((*MockReceivedPacketHandler)(nil).GetAckFrame), arg0, arg1)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAckFrame", reflect.TypeOf((*MockReceivedPacketHandler)(nil).GetAckFrame), arg0, arg1)
+	return &ReceivedPacketHandlerGetAckFrameCall{Call: call}
+}
+
+// ReceivedPacketHandlerGetAckFrameCall wrap *gomock.Call
+type ReceivedPacketHandlerGetAckFrameCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ReceivedPacketHandlerGetAckFrameCall) Return(arg0 *wire.AckFrame) *ReceivedPacketHandlerGetAckFrameCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ReceivedPacketHandlerGetAckFrameCall) Do(f func(protocol.EncryptionLevel, bool) *wire.AckFrame) *ReceivedPacketHandlerGetAckFrameCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ReceivedPacketHandlerGetAckFrameCall) DoAndReturn(f func(protocol.EncryptionLevel, bool) *wire.AckFrame) *ReceivedPacketHandlerGetAckFrameCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // GetAlarmTimeout mocks base method.
@ -75,9 +123,33 @@ func (m *MockReceivedPacketHandler) GetAlarmTimeout() time.Time {
 }

 // GetAlarmTimeout indicates an expected call of GetAlarmTimeout.
-func (mr *MockReceivedPacketHandlerMockRecorder) GetAlarmTimeout() *gomock.Call {
+func (mr *MockReceivedPacketHandlerMockRecorder) GetAlarmTimeout() *ReceivedPacketHandlerGetAlarmTimeoutCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAlarmTimeout", reflect.TypeOf((*MockReceivedPacketHandler)(nil).GetAlarmTimeout))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAlarmTimeout", reflect.TypeOf((*MockReceivedPacketHandler)(nil).GetAlarmTimeout))
+	return &ReceivedPacketHandlerGetAlarmTimeoutCall{Call: call}
+}
+
+// ReceivedPacketHandlerGetAlarmTimeoutCall wrap *gomock.Call
+type ReceivedPacketHandlerGetAlarmTimeoutCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ReceivedPacketHandlerGetAlarmTimeoutCall) Return(arg0 time.Time) *ReceivedPacketHandlerGetAlarmTimeoutCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ReceivedPacketHandlerGetAlarmTimeoutCall) Do(f func() time.Time) *ReceivedPacketHandlerGetAlarmTimeoutCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ReceivedPacketHandlerGetAlarmTimeoutCall) DoAndReturn(f func() time.Time) *ReceivedPacketHandlerGetAlarmTimeoutCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // IsPotentiallyDuplicate mocks base method.
@ -89,9 +161,33 @@ func (m *MockReceivedPacketHandler) IsPotentiallyDuplicate(arg0 protocol.PacketN
 }

 // IsPotentiallyDuplicate indicates an expected call of IsPotentiallyDuplicate.
-func (mr *MockReceivedPacketHandlerMockRecorder) IsPotentiallyDuplicate(arg0, arg1 any) *gomock.Call {
+func (mr *MockReceivedPacketHandlerMockRecorder) IsPotentiallyDuplicate(arg0, arg1 any) *ReceivedPacketHandlerIsPotentiallyDuplicateCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsPotentiallyDuplicate", reflect.TypeOf((*MockReceivedPacketHandler)(nil).IsPotentiallyDuplicate), arg0, arg1)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsPotentiallyDuplicate", reflect.TypeOf((*MockReceivedPacketHandler)(nil).IsPotentiallyDuplicate), arg0, arg1)
+	return &ReceivedPacketHandlerIsPotentiallyDuplicateCall{Call: call}
+}
+
+// ReceivedPacketHandlerIsPotentiallyDuplicateCall wrap *gomock.Call
+type ReceivedPacketHandlerIsPotentiallyDuplicateCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ReceivedPacketHandlerIsPotentiallyDuplicateCall) Return(arg0 bool) *ReceivedPacketHandlerIsPotentiallyDuplicateCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ReceivedPacketHandlerIsPotentiallyDuplicateCall) Do(f func(protocol.PacketNumber, protocol.EncryptionLevel) bool) *ReceivedPacketHandlerIsPotentiallyDuplicateCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ReceivedPacketHandlerIsPotentiallyDuplicateCall) DoAndReturn(f func(protocol.PacketNumber, protocol.EncryptionLevel) bool) *ReceivedPacketHandlerIsPotentiallyDuplicateCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // ReceivedPacket mocks base method.
@ -103,7 +199,31 @@ func (m *MockReceivedPacketHandler) ReceivedPacket(arg0 protocol.PacketNumber, a
 }

 // ReceivedPacket indicates an expected call of ReceivedPacket.
-func (mr *MockReceivedPacketHandlerMockRecorder) ReceivedPacket(arg0, arg1, arg2, arg3, arg4 any) *gomock.Call {
+func (mr *MockReceivedPacketHandlerMockRecorder) ReceivedPacket(arg0, arg1, arg2, arg3, arg4 any) *ReceivedPacketHandlerReceivedPacketCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReceivedPacket", reflect.TypeOf((*MockReceivedPacketHandler)(nil).ReceivedPacket), arg0, arg1, arg2, arg3, arg4)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReceivedPacket", reflect.TypeOf((*MockReceivedPacketHandler)(nil).ReceivedPacket), arg0, arg1, arg2, arg3, arg4)
+	return &ReceivedPacketHandlerReceivedPacketCall{Call: call}
+}
+
+// ReceivedPacketHandlerReceivedPacketCall wrap *gomock.Call
+type ReceivedPacketHandlerReceivedPacketCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ReceivedPacketHandlerReceivedPacketCall) Return(arg0 error) *ReceivedPacketHandlerReceivedPacketCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ReceivedPacketHandlerReceivedPacketCall) Do(f func(protocol.PacketNumber, protocol.ECN, protocol.EncryptionLevel, time.Time, bool) error) *ReceivedPacketHandlerReceivedPacketCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ReceivedPacketHandlerReceivedPacketCall) DoAndReturn(f func(protocol.PacketNumber, protocol.ECN, protocol.EncryptionLevel, time.Time, bool) error) *ReceivedPacketHandlerReceivedPacketCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/ackhandler/sent_packet_handler.go
+++ b/internal/mocks/ackhandler/sent_packet_handler.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mockackhandler -destination ackhandler/sent_packet_handler.go github.com/quic-go/quic-go/internal/ackhandler SentPacketHandler
+//	mockgen -typed -build_flags=-tags=gomock -package mockackhandler -destination ackhandler/sent_packet_handler.go github.com/quic-go/quic-go/internal/ackhandler SentPacketHandler
 //
 // Package mockackhandler is a generated GoMock package.
 package mockackhandler
@ -48,9 +48,33 @@ func (m *MockSentPacketHandler) DropPackets(arg0 protocol.EncryptionLevel) {
 }

 // DropPackets indicates an expected call of DropPackets.
-func (mr *MockSentPacketHandlerMockRecorder) DropPackets(arg0 any) *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) DropPackets(arg0 any) *SentPacketHandlerDropPacketsCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DropPackets", reflect.TypeOf((*MockSentPacketHandler)(nil).DropPackets), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DropPackets", reflect.TypeOf((*MockSentPacketHandler)(nil).DropPackets), arg0)
+	return &SentPacketHandlerDropPacketsCall{Call: call}
+}
+
+// SentPacketHandlerDropPacketsCall wrap *gomock.Call
+type SentPacketHandlerDropPacketsCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerDropPacketsCall) Return() *SentPacketHandlerDropPacketsCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerDropPacketsCall) Do(f func(protocol.EncryptionLevel)) *SentPacketHandlerDropPacketsCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerDropPacketsCall) DoAndReturn(f func(protocol.EncryptionLevel)) *SentPacketHandlerDropPacketsCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // ECNMode mocks base method.
@ -62,9 +86,33 @@ func (m *MockSentPacketHandler) ECNMode(arg0 bool) protocol.ECN {
 }

 // ECNMode indicates an expected call of ECNMode.
-func (mr *MockSentPacketHandlerMockRecorder) ECNMode(arg0 any) *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) ECNMode(arg0 any) *SentPacketHandlerECNModeCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ECNMode", reflect.TypeOf((*MockSentPacketHandler)(nil).ECNMode), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ECNMode", reflect.TypeOf((*MockSentPacketHandler)(nil).ECNMode), arg0)
+	return &SentPacketHandlerECNModeCall{Call: call}
+}
+
+// SentPacketHandlerECNModeCall wrap *gomock.Call
+type SentPacketHandlerECNModeCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerECNModeCall) Return(arg0 protocol.ECN) *SentPacketHandlerECNModeCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerECNModeCall) Do(f func(bool) protocol.ECN) *SentPacketHandlerECNModeCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerECNModeCall) DoAndReturn(f func(bool) protocol.ECN) *SentPacketHandlerECNModeCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // GetLossDetectionTimeout mocks base method.
@ -76,9 +124,33 @@ func (m *MockSentPacketHandler) GetLossDetectionTimeout() time.Time {
 }

 // GetLossDetectionTimeout indicates an expected call of GetLossDetectionTimeout.
-func (mr *MockSentPacketHandlerMockRecorder) GetLossDetectionTimeout() *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) GetLossDetectionTimeout() *SentPacketHandlerGetLossDetectionTimeoutCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLossDetectionTimeout", reflect.TypeOf((*MockSentPacketHandler)(nil).GetLossDetectionTimeout))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLossDetectionTimeout", reflect.TypeOf((*MockSentPacketHandler)(nil).GetLossDetectionTimeout))
+	return &SentPacketHandlerGetLossDetectionTimeoutCall{Call: call}
+}
+
+// SentPacketHandlerGetLossDetectionTimeoutCall wrap *gomock.Call
+type SentPacketHandlerGetLossDetectionTimeoutCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerGetLossDetectionTimeoutCall) Return(arg0 time.Time) *SentPacketHandlerGetLossDetectionTimeoutCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerGetLossDetectionTimeoutCall) Do(f func() time.Time) *SentPacketHandlerGetLossDetectionTimeoutCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerGetLossDetectionTimeoutCall) DoAndReturn(f func() time.Time) *SentPacketHandlerGetLossDetectionTimeoutCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // OnLossDetectionTimeout mocks base method.
@ -90,9 +162,33 @@ func (m *MockSentPacketHandler) OnLossDetectionTimeout() error {
 }

 // OnLossDetectionTimeout indicates an expected call of OnLossDetectionTimeout.
-func (mr *MockSentPacketHandlerMockRecorder) OnLossDetectionTimeout() *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) OnLossDetectionTimeout() *SentPacketHandlerOnLossDetectionTimeoutCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnLossDetectionTimeout", reflect.TypeOf((*MockSentPacketHandler)(nil).OnLossDetectionTimeout))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnLossDetectionTimeout", reflect.TypeOf((*MockSentPacketHandler)(nil).OnLossDetectionTimeout))
+	return &SentPacketHandlerOnLossDetectionTimeoutCall{Call: call}
+}
+
+// SentPacketHandlerOnLossDetectionTimeoutCall wrap *gomock.Call
+type SentPacketHandlerOnLossDetectionTimeoutCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerOnLossDetectionTimeoutCall) Return(arg0 error) *SentPacketHandlerOnLossDetectionTimeoutCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerOnLossDetectionTimeoutCall) Do(f func() error) *SentPacketHandlerOnLossDetectionTimeoutCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerOnLossDetectionTimeoutCall) DoAndReturn(f func() error) *SentPacketHandlerOnLossDetectionTimeoutCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // PeekPacketNumber mocks base method.
@ -105,9 +201,33 @@ func (m *MockSentPacketHandler) PeekPacketNumber(arg0 protocol.EncryptionLevel)
 }

 // PeekPacketNumber indicates an expected call of PeekPacketNumber.
-func (mr *MockSentPacketHandlerMockRecorder) PeekPacketNumber(arg0 any) *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) PeekPacketNumber(arg0 any) *SentPacketHandlerPeekPacketNumberCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PeekPacketNumber", reflect.TypeOf((*MockSentPacketHandler)(nil).PeekPacketNumber), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PeekPacketNumber", reflect.TypeOf((*MockSentPacketHandler)(nil).PeekPacketNumber), arg0)
+	return &SentPacketHandlerPeekPacketNumberCall{Call: call}
+}
+
+// SentPacketHandlerPeekPacketNumberCall wrap *gomock.Call
+type SentPacketHandlerPeekPacketNumberCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerPeekPacketNumberCall) Return(arg0 protocol.PacketNumber, arg1 protocol.PacketNumberLen) *SentPacketHandlerPeekPacketNumberCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerPeekPacketNumberCall) Do(f func(protocol.EncryptionLevel) (protocol.PacketNumber, protocol.PacketNumberLen)) *SentPacketHandlerPeekPacketNumberCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerPeekPacketNumberCall) DoAndReturn(f func(protocol.EncryptionLevel) (protocol.PacketNumber, protocol.PacketNumberLen)) *SentPacketHandlerPeekPacketNumberCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // PopPacketNumber mocks base method.
@ -119,9 +239,33 @@ func (m *MockSentPacketHandler) PopPacketNumber(arg0 protocol.EncryptionLevel) p
 }

 // PopPacketNumber indicates an expected call of PopPacketNumber.
-func (mr *MockSentPacketHandlerMockRecorder) PopPacketNumber(arg0 any) *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) PopPacketNumber(arg0 any) *SentPacketHandlerPopPacketNumberCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PopPacketNumber", reflect.TypeOf((*MockSentPacketHandler)(nil).PopPacketNumber), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PopPacketNumber", reflect.TypeOf((*MockSentPacketHandler)(nil).PopPacketNumber), arg0)
+	return &SentPacketHandlerPopPacketNumberCall{Call: call}
+}
+
+// SentPacketHandlerPopPacketNumberCall wrap *gomock.Call
+type SentPacketHandlerPopPacketNumberCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerPopPacketNumberCall) Return(arg0 protocol.PacketNumber) *SentPacketHandlerPopPacketNumberCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerPopPacketNumberCall) Do(f func(protocol.EncryptionLevel) protocol.PacketNumber) *SentPacketHandlerPopPacketNumberCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerPopPacketNumberCall) DoAndReturn(f func(protocol.EncryptionLevel) protocol.PacketNumber) *SentPacketHandlerPopPacketNumberCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // QueueProbePacket mocks base method.
@ -133,9 +277,33 @@ func (m *MockSentPacketHandler) QueueProbePacket(arg0 protocol.EncryptionLevel)
 }

 // QueueProbePacket indicates an expected call of QueueProbePacket.
-func (mr *MockSentPacketHandlerMockRecorder) QueueProbePacket(arg0 any) *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) QueueProbePacket(arg0 any) *SentPacketHandlerQueueProbePacketCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "QueueProbePacket", reflect.TypeOf((*MockSentPacketHandler)(nil).QueueProbePacket), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "QueueProbePacket", reflect.TypeOf((*MockSentPacketHandler)(nil).QueueProbePacket), arg0)
+	return &SentPacketHandlerQueueProbePacketCall{Call: call}
+}
+
+// SentPacketHandlerQueueProbePacketCall wrap *gomock.Call
+type SentPacketHandlerQueueProbePacketCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerQueueProbePacketCall) Return(arg0 bool) *SentPacketHandlerQueueProbePacketCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerQueueProbePacketCall) Do(f func(protocol.EncryptionLevel) bool) *SentPacketHandlerQueueProbePacketCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerQueueProbePacketCall) DoAndReturn(f func(protocol.EncryptionLevel) bool) *SentPacketHandlerQueueProbePacketCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // ReceivedAck mocks base method.
@ -148,9 +316,33 @@ func (m *MockSentPacketHandler) ReceivedAck(arg0 *wire.AckFrame, arg1 protocol.E
 }

 // ReceivedAck indicates an expected call of ReceivedAck.
-func (mr *MockSentPacketHandlerMockRecorder) ReceivedAck(arg0, arg1, arg2 any) *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) ReceivedAck(arg0, arg1, arg2 any) *SentPacketHandlerReceivedAckCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReceivedAck", reflect.TypeOf((*MockSentPacketHandler)(nil).ReceivedAck), arg0, arg1, arg2)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReceivedAck", reflect.TypeOf((*MockSentPacketHandler)(nil).ReceivedAck), arg0, arg1, arg2)
+	return &SentPacketHandlerReceivedAckCall{Call: call}
+}
+
+// SentPacketHandlerReceivedAckCall wrap *gomock.Call
+type SentPacketHandlerReceivedAckCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerReceivedAckCall) Return(arg0 bool, arg1 error) *SentPacketHandlerReceivedAckCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerReceivedAckCall) Do(f func(*wire.AckFrame, protocol.EncryptionLevel, time.Time) (bool, error)) *SentPacketHandlerReceivedAckCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerReceivedAckCall) DoAndReturn(f func(*wire.AckFrame, protocol.EncryptionLevel, time.Time) (bool, error)) *SentPacketHandlerReceivedAckCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // ReceivedBytes mocks base method.
@ -160,9 +352,33 @@ func (m *MockSentPacketHandler) ReceivedBytes(arg0 protocol.ByteCount) {
 }

 // ReceivedBytes indicates an expected call of ReceivedBytes.
-func (mr *MockSentPacketHandlerMockRecorder) ReceivedBytes(arg0 any) *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) ReceivedBytes(arg0 any) *SentPacketHandlerReceivedBytesCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReceivedBytes", reflect.TypeOf((*MockSentPacketHandler)(nil).ReceivedBytes), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReceivedBytes", reflect.TypeOf((*MockSentPacketHandler)(nil).ReceivedBytes), arg0)
+	return &SentPacketHandlerReceivedBytesCall{Call: call}
+}
+
+// SentPacketHandlerReceivedBytesCall wrap *gomock.Call
+type SentPacketHandlerReceivedBytesCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerReceivedBytesCall) Return() *SentPacketHandlerReceivedBytesCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerReceivedBytesCall) Do(f func(protocol.ByteCount)) *SentPacketHandlerReceivedBytesCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerReceivedBytesCall) DoAndReturn(f func(protocol.ByteCount)) *SentPacketHandlerReceivedBytesCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // ResetForRetry mocks base method.
@ -174,9 +390,33 @@ func (m *MockSentPacketHandler) ResetForRetry(arg0 time.Time) error {
 }

 // ResetForRetry indicates an expected call of ResetForRetry.
-func (mr *MockSentPacketHandlerMockRecorder) ResetForRetry(arg0 any) *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) ResetForRetry(arg0 any) *SentPacketHandlerResetForRetryCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResetForRetry", reflect.TypeOf((*MockSentPacketHandler)(nil).ResetForRetry), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResetForRetry", reflect.TypeOf((*MockSentPacketHandler)(nil).ResetForRetry), arg0)
+	return &SentPacketHandlerResetForRetryCall{Call: call}
+}
+
+// SentPacketHandlerResetForRetryCall wrap *gomock.Call
+type SentPacketHandlerResetForRetryCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerResetForRetryCall) Return(arg0 error) *SentPacketHandlerResetForRetryCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerResetForRetryCall) Do(f func(time.Time) error) *SentPacketHandlerResetForRetryCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerResetForRetryCall) DoAndReturn(f func(time.Time) error) *SentPacketHandlerResetForRetryCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SendMode mocks base method.
@ -188,9 +428,33 @@ func (m *MockSentPacketHandler) SendMode(arg0 time.Time) ackhandler.SendMode {
 }

 // SendMode indicates an expected call of SendMode.
-func (mr *MockSentPacketHandlerMockRecorder) SendMode(arg0 any) *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) SendMode(arg0 any) *SentPacketHandlerSendModeCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendMode", reflect.TypeOf((*MockSentPacketHandler)(nil).SendMode), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendMode", reflect.TypeOf((*MockSentPacketHandler)(nil).SendMode), arg0)
+	return &SentPacketHandlerSendModeCall{Call: call}
+}
+
+// SentPacketHandlerSendModeCall wrap *gomock.Call
+type SentPacketHandlerSendModeCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerSendModeCall) Return(arg0 ackhandler.SendMode) *SentPacketHandlerSendModeCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerSendModeCall) Do(f func(time.Time) ackhandler.SendMode) *SentPacketHandlerSendModeCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerSendModeCall) DoAndReturn(f func(time.Time) ackhandler.SendMode) *SentPacketHandlerSendModeCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SentPacket mocks base method.
@ -200,9 +464,33 @@ func (m *MockSentPacketHandler) SentPacket(arg0 time.Time, arg1, arg2 protocol.P
 }

 // SentPacket indicates an expected call of SentPacket.
-func (mr *MockSentPacketHandlerMockRecorder) SentPacket(arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8 any) *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) SentPacket(arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8 any) *SentPacketHandlerSentPacketCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SentPacket", reflect.TypeOf((*MockSentPacketHandler)(nil).SentPacket), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SentPacket", reflect.TypeOf((*MockSentPacketHandler)(nil).SentPacket), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8)
+	return &SentPacketHandlerSentPacketCall{Call: call}
+}
+
+// SentPacketHandlerSentPacketCall wrap *gomock.Call
+type SentPacketHandlerSentPacketCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerSentPacketCall) Return() *SentPacketHandlerSentPacketCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerSentPacketCall) Do(f func(time.Time, protocol.PacketNumber, protocol.PacketNumber, []ackhandler.StreamFrame, []ackhandler.Frame, protocol.EncryptionLevel, protocol.ECN, protocol.ByteCount, bool)) *SentPacketHandlerSentPacketCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerSentPacketCall) DoAndReturn(f func(time.Time, protocol.PacketNumber, protocol.PacketNumber, []ackhandler.StreamFrame, []ackhandler.Frame, protocol.EncryptionLevel, protocol.ECN, protocol.ByteCount, bool)) *SentPacketHandlerSentPacketCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SetHandshakeConfirmed mocks base method.
@ -212,9 +500,33 @@ func (m *MockSentPacketHandler) SetHandshakeConfirmed() {
 }

 // SetHandshakeConfirmed indicates an expected call of SetHandshakeConfirmed.
-func (mr *MockSentPacketHandlerMockRecorder) SetHandshakeConfirmed() *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) SetHandshakeConfirmed() *SentPacketHandlerSetHandshakeConfirmedCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetHandshakeConfirmed", reflect.TypeOf((*MockSentPacketHandler)(nil).SetHandshakeConfirmed))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetHandshakeConfirmed", reflect.TypeOf((*MockSentPacketHandler)(nil).SetHandshakeConfirmed))
+	return &SentPacketHandlerSetHandshakeConfirmedCall{Call: call}
+}
+
+// SentPacketHandlerSetHandshakeConfirmedCall wrap *gomock.Call
+type SentPacketHandlerSetHandshakeConfirmedCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerSetHandshakeConfirmedCall) Return() *SentPacketHandlerSetHandshakeConfirmedCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerSetHandshakeConfirmedCall) Do(f func()) *SentPacketHandlerSetHandshakeConfirmedCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerSetHandshakeConfirmedCall) DoAndReturn(f func()) *SentPacketHandlerSetHandshakeConfirmedCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SetMaxDatagramSize mocks base method.
@ -224,9 +536,33 @@ func (m *MockSentPacketHandler) SetMaxDatagramSize(arg0 protocol.ByteCount) {
 }

 // SetMaxDatagramSize indicates an expected call of SetMaxDatagramSize.
-func (mr *MockSentPacketHandlerMockRecorder) SetMaxDatagramSize(arg0 any) *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) SetMaxDatagramSize(arg0 any) *SentPacketHandlerSetMaxDatagramSizeCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetMaxDatagramSize", reflect.TypeOf((*MockSentPacketHandler)(nil).SetMaxDatagramSize), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetMaxDatagramSize", reflect.TypeOf((*MockSentPacketHandler)(nil).SetMaxDatagramSize), arg0)
+	return &SentPacketHandlerSetMaxDatagramSizeCall{Call: call}
+}
+
+// SentPacketHandlerSetMaxDatagramSizeCall wrap *gomock.Call
+type SentPacketHandlerSetMaxDatagramSizeCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerSetMaxDatagramSizeCall) Return() *SentPacketHandlerSetMaxDatagramSizeCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerSetMaxDatagramSizeCall) Do(f func(protocol.ByteCount)) *SentPacketHandlerSetMaxDatagramSizeCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerSetMaxDatagramSizeCall) DoAndReturn(f func(protocol.ByteCount)) *SentPacketHandlerSetMaxDatagramSizeCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // TimeUntilSend mocks base method.
@ -238,7 +574,31 @@ func (m *MockSentPacketHandler) TimeUntilSend() time.Time {
 }

 // TimeUntilSend indicates an expected call of TimeUntilSend.
-func (mr *MockSentPacketHandlerMockRecorder) TimeUntilSend() *gomock.Call {
+func (mr *MockSentPacketHandlerMockRecorder) TimeUntilSend() *SentPacketHandlerTimeUntilSendCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TimeUntilSend", reflect.TypeOf((*MockSentPacketHandler)(nil).TimeUntilSend))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TimeUntilSend", reflect.TypeOf((*MockSentPacketHandler)(nil).TimeUntilSend))
+	return &SentPacketHandlerTimeUntilSendCall{Call: call}
+}
+
+// SentPacketHandlerTimeUntilSendCall wrap *gomock.Call
+type SentPacketHandlerTimeUntilSendCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SentPacketHandlerTimeUntilSendCall) Return(arg0 time.Time) *SentPacketHandlerTimeUntilSendCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SentPacketHandlerTimeUntilSendCall) Do(f func() time.Time) *SentPacketHandlerTimeUntilSendCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SentPacketHandlerTimeUntilSendCall) DoAndReturn(f func() time.Time) *SentPacketHandlerTimeUntilSendCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/congestion.go
+++ b/internal/mocks/congestion.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mocks -destination congestion.go github.com/quic-go/quic-go/internal/congestion SendAlgorithmWithDebugInfos
+//	mockgen -typed -build_flags=-tags=gomock -package mocks -destination congestion.go github.com/quic-go/quic-go/internal/congestion SendAlgorithmWithDebugInfos
 //
 // Package mocks is a generated GoMock package.
 package mocks
@ -48,9 +48,33 @@ func (m *MockSendAlgorithmWithDebugInfos) CanSend(arg0 protocol.ByteCount) bool
 }

 // CanSend indicates an expected call of CanSend.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) CanSend(arg0 any) *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) CanSend(arg0 any) *SendAlgorithmWithDebugInfosCanSendCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CanSend", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).CanSend), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CanSend", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).CanSend), arg0)
+	return &SendAlgorithmWithDebugInfosCanSendCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosCanSendCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosCanSendCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosCanSendCall) Return(arg0 bool) *SendAlgorithmWithDebugInfosCanSendCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosCanSendCall) Do(f func(protocol.ByteCount) bool) *SendAlgorithmWithDebugInfosCanSendCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosCanSendCall) DoAndReturn(f func(protocol.ByteCount) bool) *SendAlgorithmWithDebugInfosCanSendCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // GetCongestionWindow mocks base method.
@ -62,9 +86,33 @@ func (m *MockSendAlgorithmWithDebugInfos) GetCongestionWindow() protocol.ByteCou
 }

 // GetCongestionWindow indicates an expected call of GetCongestionWindow.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) GetCongestionWindow() *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) GetCongestionWindow() *SendAlgorithmWithDebugInfosGetCongestionWindowCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetCongestionWindow", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).GetCongestionWindow))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetCongestionWindow", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).GetCongestionWindow))
+	return &SendAlgorithmWithDebugInfosGetCongestionWindowCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosGetCongestionWindowCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosGetCongestionWindowCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosGetCongestionWindowCall) Return(arg0 protocol.ByteCount) *SendAlgorithmWithDebugInfosGetCongestionWindowCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosGetCongestionWindowCall) Do(f func() protocol.ByteCount) *SendAlgorithmWithDebugInfosGetCongestionWindowCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosGetCongestionWindowCall) DoAndReturn(f func() protocol.ByteCount) *SendAlgorithmWithDebugInfosGetCongestionWindowCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // HasPacingBudget mocks base method.
@ -76,9 +124,33 @@ func (m *MockSendAlgorithmWithDebugInfos) HasPacingBudget(arg0 time.Time) bool {
 }

 // HasPacingBudget indicates an expected call of HasPacingBudget.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) HasPacingBudget(arg0 any) *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) HasPacingBudget(arg0 any) *SendAlgorithmWithDebugInfosHasPacingBudgetCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasPacingBudget", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).HasPacingBudget), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasPacingBudget", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).HasPacingBudget), arg0)
+	return &SendAlgorithmWithDebugInfosHasPacingBudgetCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosHasPacingBudgetCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosHasPacingBudgetCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosHasPacingBudgetCall) Return(arg0 bool) *SendAlgorithmWithDebugInfosHasPacingBudgetCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosHasPacingBudgetCall) Do(f func(time.Time) bool) *SendAlgorithmWithDebugInfosHasPacingBudgetCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosHasPacingBudgetCall) DoAndReturn(f func(time.Time) bool) *SendAlgorithmWithDebugInfosHasPacingBudgetCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // InRecovery mocks base method.
@ -90,9 +162,33 @@ func (m *MockSendAlgorithmWithDebugInfos) InRecovery() bool {
 }

 // InRecovery indicates an expected call of InRecovery.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) InRecovery() *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) InRecovery() *SendAlgorithmWithDebugInfosInRecoveryCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InRecovery", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).InRecovery))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InRecovery", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).InRecovery))
+	return &SendAlgorithmWithDebugInfosInRecoveryCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosInRecoveryCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosInRecoveryCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosInRecoveryCall) Return(arg0 bool) *SendAlgorithmWithDebugInfosInRecoveryCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosInRecoveryCall) Do(f func() bool) *SendAlgorithmWithDebugInfosInRecoveryCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosInRecoveryCall) DoAndReturn(f func() bool) *SendAlgorithmWithDebugInfosInRecoveryCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // InSlowStart mocks base method.
@ -104,9 +200,33 @@ func (m *MockSendAlgorithmWithDebugInfos) InSlowStart() bool {
 }

 // InSlowStart indicates an expected call of InSlowStart.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) InSlowStart() *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) InSlowStart() *SendAlgorithmWithDebugInfosInSlowStartCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InSlowStart", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).InSlowStart))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InSlowStart", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).InSlowStart))
+	return &SendAlgorithmWithDebugInfosInSlowStartCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosInSlowStartCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosInSlowStartCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosInSlowStartCall) Return(arg0 bool) *SendAlgorithmWithDebugInfosInSlowStartCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosInSlowStartCall) Do(f func() bool) *SendAlgorithmWithDebugInfosInSlowStartCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosInSlowStartCall) DoAndReturn(f func() bool) *SendAlgorithmWithDebugInfosInSlowStartCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // MaybeExitSlowStart mocks base method.
@ -116,9 +236,33 @@ func (m *MockSendAlgorithmWithDebugInfos) MaybeExitSlowStart() {
 }

 // MaybeExitSlowStart indicates an expected call of MaybeExitSlowStart.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) MaybeExitSlowStart() *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) MaybeExitSlowStart() *SendAlgorithmWithDebugInfosMaybeExitSlowStartCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MaybeExitSlowStart", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).MaybeExitSlowStart))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MaybeExitSlowStart", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).MaybeExitSlowStart))
+	return &SendAlgorithmWithDebugInfosMaybeExitSlowStartCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosMaybeExitSlowStartCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosMaybeExitSlowStartCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosMaybeExitSlowStartCall) Return() *SendAlgorithmWithDebugInfosMaybeExitSlowStartCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosMaybeExitSlowStartCall) Do(f func()) *SendAlgorithmWithDebugInfosMaybeExitSlowStartCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosMaybeExitSlowStartCall) DoAndReturn(f func()) *SendAlgorithmWithDebugInfosMaybeExitSlowStartCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // OnCongestionEvent mocks base method.
@ -128,9 +272,33 @@ func (m *MockSendAlgorithmWithDebugInfos) OnCongestionEvent(arg0 protocol.Packet
 }

 // OnCongestionEvent indicates an expected call of OnCongestionEvent.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) OnCongestionEvent(arg0, arg1, arg2 any) *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) OnCongestionEvent(arg0, arg1, arg2 any) *SendAlgorithmWithDebugInfosOnCongestionEventCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnCongestionEvent", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).OnCongestionEvent), arg0, arg1, arg2)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnCongestionEvent", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).OnCongestionEvent), arg0, arg1, arg2)
+	return &SendAlgorithmWithDebugInfosOnCongestionEventCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosOnCongestionEventCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosOnCongestionEventCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosOnCongestionEventCall) Return() *SendAlgorithmWithDebugInfosOnCongestionEventCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosOnCongestionEventCall) Do(f func(protocol.PacketNumber, protocol.ByteCount, protocol.ByteCount)) *SendAlgorithmWithDebugInfosOnCongestionEventCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosOnCongestionEventCall) DoAndReturn(f func(protocol.PacketNumber, protocol.ByteCount, protocol.ByteCount)) *SendAlgorithmWithDebugInfosOnCongestionEventCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // OnPacketAcked mocks base method.
@ -140,9 +308,33 @@ func (m *MockSendAlgorithmWithDebugInfos) OnPacketAcked(arg0 protocol.PacketNumb
 }

 // OnPacketAcked indicates an expected call of OnPacketAcked.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) OnPacketAcked(arg0, arg1, arg2, arg3 any) *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) OnPacketAcked(arg0, arg1, arg2, arg3 any) *SendAlgorithmWithDebugInfosOnPacketAckedCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPacketAcked", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).OnPacketAcked), arg0, arg1, arg2, arg3)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPacketAcked", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).OnPacketAcked), arg0, arg1, arg2, arg3)
+	return &SendAlgorithmWithDebugInfosOnPacketAckedCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosOnPacketAckedCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosOnPacketAckedCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosOnPacketAckedCall) Return() *SendAlgorithmWithDebugInfosOnPacketAckedCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosOnPacketAckedCall) Do(f func(protocol.PacketNumber, protocol.ByteCount, protocol.ByteCount, time.Time)) *SendAlgorithmWithDebugInfosOnPacketAckedCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosOnPacketAckedCall) DoAndReturn(f func(protocol.PacketNumber, protocol.ByteCount, protocol.ByteCount, time.Time)) *SendAlgorithmWithDebugInfosOnPacketAckedCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // OnPacketSent mocks base method.
@ -152,9 +344,33 @@ func (m *MockSendAlgorithmWithDebugInfos) OnPacketSent(arg0 time.Time, arg1 prot
 }

 // OnPacketSent indicates an expected call of OnPacketSent.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) OnPacketSent(arg0, arg1, arg2, arg3, arg4 any) *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) OnPacketSent(arg0, arg1, arg2, arg3, arg4 any) *SendAlgorithmWithDebugInfosOnPacketSentCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPacketSent", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).OnPacketSent), arg0, arg1, arg2, arg3, arg4)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPacketSent", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).OnPacketSent), arg0, arg1, arg2, arg3, arg4)
+	return &SendAlgorithmWithDebugInfosOnPacketSentCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosOnPacketSentCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosOnPacketSentCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosOnPacketSentCall) Return() *SendAlgorithmWithDebugInfosOnPacketSentCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosOnPacketSentCall) Do(f func(time.Time, protocol.ByteCount, protocol.PacketNumber, protocol.ByteCount, bool)) *SendAlgorithmWithDebugInfosOnPacketSentCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosOnPacketSentCall) DoAndReturn(f func(time.Time, protocol.ByteCount, protocol.PacketNumber, protocol.ByteCount, bool)) *SendAlgorithmWithDebugInfosOnPacketSentCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // OnRetransmissionTimeout mocks base method.
@ -164,9 +380,33 @@ func (m *MockSendAlgorithmWithDebugInfos) OnRetransmissionTimeout(arg0 bool) {
 }

 // OnRetransmissionTimeout indicates an expected call of OnRetransmissionTimeout.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) OnRetransmissionTimeout(arg0 any) *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) OnRetransmissionTimeout(arg0 any) *SendAlgorithmWithDebugInfosOnRetransmissionTimeoutCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnRetransmissionTimeout", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).OnRetransmissionTimeout), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnRetransmissionTimeout", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).OnRetransmissionTimeout), arg0)
+	return &SendAlgorithmWithDebugInfosOnRetransmissionTimeoutCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosOnRetransmissionTimeoutCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosOnRetransmissionTimeoutCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosOnRetransmissionTimeoutCall) Return() *SendAlgorithmWithDebugInfosOnRetransmissionTimeoutCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosOnRetransmissionTimeoutCall) Do(f func(bool)) *SendAlgorithmWithDebugInfosOnRetransmissionTimeoutCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosOnRetransmissionTimeoutCall) DoAndReturn(f func(bool)) *SendAlgorithmWithDebugInfosOnRetransmissionTimeoutCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SetMaxDatagramSize mocks base method.
@ -176,9 +416,33 @@ func (m *MockSendAlgorithmWithDebugInfos) SetMaxDatagramSize(arg0 protocol.ByteC
 }

 // SetMaxDatagramSize indicates an expected call of SetMaxDatagramSize.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) SetMaxDatagramSize(arg0 any) *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) SetMaxDatagramSize(arg0 any) *SendAlgorithmWithDebugInfosSetMaxDatagramSizeCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetMaxDatagramSize", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).SetMaxDatagramSize), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetMaxDatagramSize", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).SetMaxDatagramSize), arg0)
+	return &SendAlgorithmWithDebugInfosSetMaxDatagramSizeCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosSetMaxDatagramSizeCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosSetMaxDatagramSizeCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosSetMaxDatagramSizeCall) Return() *SendAlgorithmWithDebugInfosSetMaxDatagramSizeCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosSetMaxDatagramSizeCall) Do(f func(protocol.ByteCount)) *SendAlgorithmWithDebugInfosSetMaxDatagramSizeCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosSetMaxDatagramSizeCall) DoAndReturn(f func(protocol.ByteCount)) *SendAlgorithmWithDebugInfosSetMaxDatagramSizeCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // TimeUntilSend mocks base method.
@ -190,7 +454,31 @@ func (m *MockSendAlgorithmWithDebugInfos) TimeUntilSend(arg0 protocol.ByteCount)
 }

 // TimeUntilSend indicates an expected call of TimeUntilSend.
-func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) TimeUntilSend(arg0 any) *gomock.Call {
+func (mr *MockSendAlgorithmWithDebugInfosMockRecorder) TimeUntilSend(arg0 any) *SendAlgorithmWithDebugInfosTimeUntilSendCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TimeUntilSend", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).TimeUntilSend), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TimeUntilSend", reflect.TypeOf((*MockSendAlgorithmWithDebugInfos)(nil).TimeUntilSend), arg0)
+	return &SendAlgorithmWithDebugInfosTimeUntilSendCall{Call: call}
+}
+
+// SendAlgorithmWithDebugInfosTimeUntilSendCall wrap *gomock.Call
+type SendAlgorithmWithDebugInfosTimeUntilSendCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *SendAlgorithmWithDebugInfosTimeUntilSendCall) Return(arg0 time.Time) *SendAlgorithmWithDebugInfosTimeUntilSendCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *SendAlgorithmWithDebugInfosTimeUntilSendCall) Do(f func(protocol.ByteCount) time.Time) *SendAlgorithmWithDebugInfosTimeUntilSendCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *SendAlgorithmWithDebugInfosTimeUntilSendCall) DoAndReturn(f func(protocol.ByteCount) time.Time) *SendAlgorithmWithDebugInfosTimeUntilSendCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/connection_flow_controller.go
+++ b/internal/mocks/connection_flow_controller.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mocks -destination connection_flow_controller.go github.com/quic-go/quic-go/internal/flowcontrol ConnectionFlowController
+//	mockgen -typed -build_flags=-tags=gomock -package mocks -destination connection_flow_controller.go github.com/quic-go/quic-go/internal/flowcontrol ConnectionFlowController
 //
 // Package mocks is a generated GoMock package.
 package mocks
@ -45,9 +45,33 @@ func (m *MockConnectionFlowController) AddBytesRead(arg0 protocol.ByteCount) {
 }

 // AddBytesRead indicates an expected call of AddBytesRead.
-func (mr *MockConnectionFlowControllerMockRecorder) AddBytesRead(arg0 any) *gomock.Call {
+func (mr *MockConnectionFlowControllerMockRecorder) AddBytesRead(arg0 any) *ConnectionFlowControllerAddBytesReadCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddBytesRead", reflect.TypeOf((*MockConnectionFlowController)(nil).AddBytesRead), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddBytesRead", reflect.TypeOf((*MockConnectionFlowController)(nil).AddBytesRead), arg0)
+	return &ConnectionFlowControllerAddBytesReadCall{Call: call}
+}
+
+// ConnectionFlowControllerAddBytesReadCall wrap *gomock.Call
+type ConnectionFlowControllerAddBytesReadCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ConnectionFlowControllerAddBytesReadCall) Return() *ConnectionFlowControllerAddBytesReadCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ConnectionFlowControllerAddBytesReadCall) Do(f func(protocol.ByteCount)) *ConnectionFlowControllerAddBytesReadCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ConnectionFlowControllerAddBytesReadCall) DoAndReturn(f func(protocol.ByteCount)) *ConnectionFlowControllerAddBytesReadCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // AddBytesSent mocks base method.
@ -57,9 +81,33 @@ func (m *MockConnectionFlowController) AddBytesSent(arg0 protocol.ByteCount) {
 }

 // AddBytesSent indicates an expected call of AddBytesSent.
-func (mr *MockConnectionFlowControllerMockRecorder) AddBytesSent(arg0 any) *gomock.Call {
+func (mr *MockConnectionFlowControllerMockRecorder) AddBytesSent(arg0 any) *ConnectionFlowControllerAddBytesSentCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddBytesSent", reflect.TypeOf((*MockConnectionFlowController)(nil).AddBytesSent), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddBytesSent", reflect.TypeOf((*MockConnectionFlowController)(nil).AddBytesSent), arg0)
+	return &ConnectionFlowControllerAddBytesSentCall{Call: call}
+}
+
+// ConnectionFlowControllerAddBytesSentCall wrap *gomock.Call
+type ConnectionFlowControllerAddBytesSentCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ConnectionFlowControllerAddBytesSentCall) Return() *ConnectionFlowControllerAddBytesSentCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ConnectionFlowControllerAddBytesSentCall) Do(f func(protocol.ByteCount)) *ConnectionFlowControllerAddBytesSentCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ConnectionFlowControllerAddBytesSentCall) DoAndReturn(f func(protocol.ByteCount)) *ConnectionFlowControllerAddBytesSentCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // GetWindowUpdate mocks base method.
@ -71,9 +119,33 @@ func (m *MockConnectionFlowController) GetWindowUpdate() protocol.ByteCount {
 }

 // GetWindowUpdate indicates an expected call of GetWindowUpdate.
-func (mr *MockConnectionFlowControllerMockRecorder) GetWindowUpdate() *gomock.Call {
+func (mr *MockConnectionFlowControllerMockRecorder) GetWindowUpdate() *ConnectionFlowControllerGetWindowUpdateCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWindowUpdate", reflect.TypeOf((*MockConnectionFlowController)(nil).GetWindowUpdate))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWindowUpdate", reflect.TypeOf((*MockConnectionFlowController)(nil).GetWindowUpdate))
+	return &ConnectionFlowControllerGetWindowUpdateCall{Call: call}
+}
+
+// ConnectionFlowControllerGetWindowUpdateCall wrap *gomock.Call
+type ConnectionFlowControllerGetWindowUpdateCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ConnectionFlowControllerGetWindowUpdateCall) Return(arg0 protocol.ByteCount) *ConnectionFlowControllerGetWindowUpdateCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ConnectionFlowControllerGetWindowUpdateCall) Do(f func() protocol.ByteCount) *ConnectionFlowControllerGetWindowUpdateCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ConnectionFlowControllerGetWindowUpdateCall) DoAndReturn(f func() protocol.ByteCount) *ConnectionFlowControllerGetWindowUpdateCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // IsNewlyBlocked mocks base method.
@ -86,9 +158,33 @@ func (m *MockConnectionFlowController) IsNewlyBlocked() (bool, protocol.ByteCoun
 }

 // IsNewlyBlocked indicates an expected call of IsNewlyBlocked.
-func (mr *MockConnectionFlowControllerMockRecorder) IsNewlyBlocked() *gomock.Call {
+func (mr *MockConnectionFlowControllerMockRecorder) IsNewlyBlocked() *ConnectionFlowControllerIsNewlyBlockedCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsNewlyBlocked", reflect.TypeOf((*MockConnectionFlowController)(nil).IsNewlyBlocked))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsNewlyBlocked", reflect.TypeOf((*MockConnectionFlowController)(nil).IsNewlyBlocked))
+	return &ConnectionFlowControllerIsNewlyBlockedCall{Call: call}
+}
+
+// ConnectionFlowControllerIsNewlyBlockedCall wrap *gomock.Call
+type ConnectionFlowControllerIsNewlyBlockedCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ConnectionFlowControllerIsNewlyBlockedCall) Return(arg0 bool, arg1 protocol.ByteCount) *ConnectionFlowControllerIsNewlyBlockedCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ConnectionFlowControllerIsNewlyBlockedCall) Do(f func() (bool, protocol.ByteCount)) *ConnectionFlowControllerIsNewlyBlockedCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ConnectionFlowControllerIsNewlyBlockedCall) DoAndReturn(f func() (bool, protocol.ByteCount)) *ConnectionFlowControllerIsNewlyBlockedCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Reset mocks base method.
@ -100,9 +196,33 @@ func (m *MockConnectionFlowController) Reset() error {
 }

 // Reset indicates an expected call of Reset.
-func (mr *MockConnectionFlowControllerMockRecorder) Reset() *gomock.Call {
+func (mr *MockConnectionFlowControllerMockRecorder) Reset() *ConnectionFlowControllerResetCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Reset", reflect.TypeOf((*MockConnectionFlowController)(nil).Reset))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Reset", reflect.TypeOf((*MockConnectionFlowController)(nil).Reset))
+	return &ConnectionFlowControllerResetCall{Call: call}
+}
+
+// ConnectionFlowControllerResetCall wrap *gomock.Call
+type ConnectionFlowControllerResetCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ConnectionFlowControllerResetCall) Return(arg0 error) *ConnectionFlowControllerResetCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ConnectionFlowControllerResetCall) Do(f func() error) *ConnectionFlowControllerResetCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ConnectionFlowControllerResetCall) DoAndReturn(f func() error) *ConnectionFlowControllerResetCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SendWindowSize mocks base method.
@ -114,9 +234,33 @@ func (m *MockConnectionFlowController) SendWindowSize() protocol.ByteCount {
 }

 // SendWindowSize indicates an expected call of SendWindowSize.
-func (mr *MockConnectionFlowControllerMockRecorder) SendWindowSize() *gomock.Call {
+func (mr *MockConnectionFlowControllerMockRecorder) SendWindowSize() *ConnectionFlowControllerSendWindowSizeCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendWindowSize", reflect.TypeOf((*MockConnectionFlowController)(nil).SendWindowSize))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendWindowSize", reflect.TypeOf((*MockConnectionFlowController)(nil).SendWindowSize))
+	return &ConnectionFlowControllerSendWindowSizeCall{Call: call}
+}
+
+// ConnectionFlowControllerSendWindowSizeCall wrap *gomock.Call
+type ConnectionFlowControllerSendWindowSizeCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ConnectionFlowControllerSendWindowSizeCall) Return(arg0 protocol.ByteCount) *ConnectionFlowControllerSendWindowSizeCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ConnectionFlowControllerSendWindowSizeCall) Do(f func() protocol.ByteCount) *ConnectionFlowControllerSendWindowSizeCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ConnectionFlowControllerSendWindowSizeCall) DoAndReturn(f func() protocol.ByteCount) *ConnectionFlowControllerSendWindowSizeCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // UpdateSendWindow mocks base method.
@ -126,7 +270,31 @@ func (m *MockConnectionFlowController) UpdateSendWindow(arg0 protocol.ByteCount)
 }

 // UpdateSendWindow indicates an expected call of UpdateSendWindow.
-func (mr *MockConnectionFlowControllerMockRecorder) UpdateSendWindow(arg0 any) *gomock.Call {
+func (mr *MockConnectionFlowControllerMockRecorder) UpdateSendWindow(arg0 any) *ConnectionFlowControllerUpdateSendWindowCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateSendWindow", reflect.TypeOf((*MockConnectionFlowController)(nil).UpdateSendWindow), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateSendWindow", reflect.TypeOf((*MockConnectionFlowController)(nil).UpdateSendWindow), arg0)
+	return &ConnectionFlowControllerUpdateSendWindowCall{Call: call}
+}
+
+// ConnectionFlowControllerUpdateSendWindowCall wrap *gomock.Call
+type ConnectionFlowControllerUpdateSendWindowCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ConnectionFlowControllerUpdateSendWindowCall) Return() *ConnectionFlowControllerUpdateSendWindowCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ConnectionFlowControllerUpdateSendWindowCall) Do(f func(protocol.ByteCount)) *ConnectionFlowControllerUpdateSendWindowCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ConnectionFlowControllerUpdateSendWindowCall) DoAndReturn(f func(protocol.ByteCount)) *ConnectionFlowControllerUpdateSendWindowCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/crypto_setup.go
+++ b/internal/mocks/crypto_setup.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mocks -destination crypto_setup_tmp.go github.com/quic-go/quic-go/internal/handshake CryptoSetup
+//	mockgen -typed -build_flags=-tags=gomock -package mocks -destination crypto_setup_tmp.go github.com/quic-go/quic-go/internal/handshake CryptoSetup
 //
 // Package mocks is a generated GoMock package.
 package mocks
@ -46,9 +46,33 @@ func (m *MockCryptoSetup) ChangeConnectionID(arg0 protocol.ConnectionID) {
 }

 // ChangeConnectionID indicates an expected call of ChangeConnectionID.
-func (mr *MockCryptoSetupMockRecorder) ChangeConnectionID(arg0 any) *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) ChangeConnectionID(arg0 any) *CryptoSetupChangeConnectionIDCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ChangeConnectionID", reflect.TypeOf((*MockCryptoSetup)(nil).ChangeConnectionID), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ChangeConnectionID", reflect.TypeOf((*MockCryptoSetup)(nil).ChangeConnectionID), arg0)
+	return &CryptoSetupChangeConnectionIDCall{Call: call}
+}
+
+// CryptoSetupChangeConnectionIDCall wrap *gomock.Call
+type CryptoSetupChangeConnectionIDCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupChangeConnectionIDCall) Return() *CryptoSetupChangeConnectionIDCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupChangeConnectionIDCall) Do(f func(protocol.ConnectionID)) *CryptoSetupChangeConnectionIDCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupChangeConnectionIDCall) DoAndReturn(f func(protocol.ConnectionID)) *CryptoSetupChangeConnectionIDCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Close mocks base method.
@ -60,9 +84,33 @@ func (m *MockCryptoSetup) Close() error {
 }

 // Close indicates an expected call of Close.
-func (mr *MockCryptoSetupMockRecorder) Close() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) Close() *CryptoSetupCloseCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockCryptoSetup)(nil).Close))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockCryptoSetup)(nil).Close))
+	return &CryptoSetupCloseCall{Call: call}
+}
+
+// CryptoSetupCloseCall wrap *gomock.Call
+type CryptoSetupCloseCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupCloseCall) Return(arg0 error) *CryptoSetupCloseCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupCloseCall) Do(f func() error) *CryptoSetupCloseCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupCloseCall) DoAndReturn(f func() error) *CryptoSetupCloseCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // ConnectionState mocks base method.
@ -74,9 +122,33 @@ func (m *MockCryptoSetup) ConnectionState() handshake.ConnectionState {
 }

 // ConnectionState indicates an expected call of ConnectionState.
-func (mr *MockCryptoSetupMockRecorder) ConnectionState() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) ConnectionState() *CryptoSetupConnectionStateCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ConnectionState", reflect.TypeOf((*MockCryptoSetup)(nil).ConnectionState))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ConnectionState", reflect.TypeOf((*MockCryptoSetup)(nil).ConnectionState))
+	return &CryptoSetupConnectionStateCall{Call: call}
+}
+
+// CryptoSetupConnectionStateCall wrap *gomock.Call
+type CryptoSetupConnectionStateCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupConnectionStateCall) Return(arg0 handshake.ConnectionState) *CryptoSetupConnectionStateCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupConnectionStateCall) Do(f func() handshake.ConnectionState) *CryptoSetupConnectionStateCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupConnectionStateCall) DoAndReturn(f func() handshake.ConnectionState) *CryptoSetupConnectionStateCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // DiscardInitialKeys mocks base method.
@ -86,9 +158,33 @@ func (m *MockCryptoSetup) DiscardInitialKeys() {
 }

 // DiscardInitialKeys indicates an expected call of DiscardInitialKeys.
-func (mr *MockCryptoSetupMockRecorder) DiscardInitialKeys() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) DiscardInitialKeys() *CryptoSetupDiscardInitialKeysCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscardInitialKeys", reflect.TypeOf((*MockCryptoSetup)(nil).DiscardInitialKeys))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscardInitialKeys", reflect.TypeOf((*MockCryptoSetup)(nil).DiscardInitialKeys))
+	return &CryptoSetupDiscardInitialKeysCall{Call: call}
+}
+
+// CryptoSetupDiscardInitialKeysCall wrap *gomock.Call
+type CryptoSetupDiscardInitialKeysCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupDiscardInitialKeysCall) Return() *CryptoSetupDiscardInitialKeysCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupDiscardInitialKeysCall) Do(f func()) *CryptoSetupDiscardInitialKeysCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupDiscardInitialKeysCall) DoAndReturn(f func()) *CryptoSetupDiscardInitialKeysCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Get0RTTOpener mocks base method.
@ -101,9 +197,33 @@ func (m *MockCryptoSetup) Get0RTTOpener() (handshake.LongHeaderOpener, error) {
 }

 // Get0RTTOpener indicates an expected call of Get0RTTOpener.
-func (mr *MockCryptoSetupMockRecorder) Get0RTTOpener() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) Get0RTTOpener() *CryptoSetupGet0RTTOpenerCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get0RTTOpener", reflect.TypeOf((*MockCryptoSetup)(nil).Get0RTTOpener))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get0RTTOpener", reflect.TypeOf((*MockCryptoSetup)(nil).Get0RTTOpener))
+	return &CryptoSetupGet0RTTOpenerCall{Call: call}
+}
+
+// CryptoSetupGet0RTTOpenerCall wrap *gomock.Call
+type CryptoSetupGet0RTTOpenerCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupGet0RTTOpenerCall) Return(arg0 handshake.LongHeaderOpener, arg1 error) *CryptoSetupGet0RTTOpenerCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupGet0RTTOpenerCall) Do(f func() (handshake.LongHeaderOpener, error)) *CryptoSetupGet0RTTOpenerCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupGet0RTTOpenerCall) DoAndReturn(f func() (handshake.LongHeaderOpener, error)) *CryptoSetupGet0RTTOpenerCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Get0RTTSealer mocks base method.
@ -116,9 +236,33 @@ func (m *MockCryptoSetup) Get0RTTSealer() (handshake.LongHeaderSealer, error) {
 }

 // Get0RTTSealer indicates an expected call of Get0RTTSealer.
-func (mr *MockCryptoSetupMockRecorder) Get0RTTSealer() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) Get0RTTSealer() *CryptoSetupGet0RTTSealerCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get0RTTSealer", reflect.TypeOf((*MockCryptoSetup)(nil).Get0RTTSealer))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get0RTTSealer", reflect.TypeOf((*MockCryptoSetup)(nil).Get0RTTSealer))
+	return &CryptoSetupGet0RTTSealerCall{Call: call}
+}
+
+// CryptoSetupGet0RTTSealerCall wrap *gomock.Call
+type CryptoSetupGet0RTTSealerCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupGet0RTTSealerCall) Return(arg0 handshake.LongHeaderSealer, arg1 error) *CryptoSetupGet0RTTSealerCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupGet0RTTSealerCall) Do(f func() (handshake.LongHeaderSealer, error)) *CryptoSetupGet0RTTSealerCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupGet0RTTSealerCall) DoAndReturn(f func() (handshake.LongHeaderSealer, error)) *CryptoSetupGet0RTTSealerCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Get1RTTOpener mocks base method.
@ -131,9 +275,33 @@ func (m *MockCryptoSetup) Get1RTTOpener() (handshake.ShortHeaderOpener, error) {
 }

 // Get1RTTOpener indicates an expected call of Get1RTTOpener.
-func (mr *MockCryptoSetupMockRecorder) Get1RTTOpener() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) Get1RTTOpener() *CryptoSetupGet1RTTOpenerCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get1RTTOpener", reflect.TypeOf((*MockCryptoSetup)(nil).Get1RTTOpener))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get1RTTOpener", reflect.TypeOf((*MockCryptoSetup)(nil).Get1RTTOpener))
+	return &CryptoSetupGet1RTTOpenerCall{Call: call}
+}
+
+// CryptoSetupGet1RTTOpenerCall wrap *gomock.Call
+type CryptoSetupGet1RTTOpenerCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupGet1RTTOpenerCall) Return(arg0 handshake.ShortHeaderOpener, arg1 error) *CryptoSetupGet1RTTOpenerCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupGet1RTTOpenerCall) Do(f func() (handshake.ShortHeaderOpener, error)) *CryptoSetupGet1RTTOpenerCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupGet1RTTOpenerCall) DoAndReturn(f func() (handshake.ShortHeaderOpener, error)) *CryptoSetupGet1RTTOpenerCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Get1RTTSealer mocks base method.
@ -146,9 +314,33 @@ func (m *MockCryptoSetup) Get1RTTSealer() (handshake.ShortHeaderSealer, error) {
 }

 // Get1RTTSealer indicates an expected call of Get1RTTSealer.
-func (mr *MockCryptoSetupMockRecorder) Get1RTTSealer() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) Get1RTTSealer() *CryptoSetupGet1RTTSealerCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get1RTTSealer", reflect.TypeOf((*MockCryptoSetup)(nil).Get1RTTSealer))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get1RTTSealer", reflect.TypeOf((*MockCryptoSetup)(nil).Get1RTTSealer))
+	return &CryptoSetupGet1RTTSealerCall{Call: call}
+}
+
+// CryptoSetupGet1RTTSealerCall wrap *gomock.Call
+type CryptoSetupGet1RTTSealerCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupGet1RTTSealerCall) Return(arg0 handshake.ShortHeaderSealer, arg1 error) *CryptoSetupGet1RTTSealerCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupGet1RTTSealerCall) Do(f func() (handshake.ShortHeaderSealer, error)) *CryptoSetupGet1RTTSealerCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupGet1RTTSealerCall) DoAndReturn(f func() (handshake.ShortHeaderSealer, error)) *CryptoSetupGet1RTTSealerCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // GetHandshakeOpener mocks base method.
@ -161,9 +353,33 @@ func (m *MockCryptoSetup) GetHandshakeOpener() (handshake.LongHeaderOpener, erro
 }

 // GetHandshakeOpener indicates an expected call of GetHandshakeOpener.
-func (mr *MockCryptoSetupMockRecorder) GetHandshakeOpener() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) GetHandshakeOpener() *CryptoSetupGetHandshakeOpenerCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetHandshakeOpener", reflect.TypeOf((*MockCryptoSetup)(nil).GetHandshakeOpener))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetHandshakeOpener", reflect.TypeOf((*MockCryptoSetup)(nil).GetHandshakeOpener))
+	return &CryptoSetupGetHandshakeOpenerCall{Call: call}
+}
+
+// CryptoSetupGetHandshakeOpenerCall wrap *gomock.Call
+type CryptoSetupGetHandshakeOpenerCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupGetHandshakeOpenerCall) Return(arg0 handshake.LongHeaderOpener, arg1 error) *CryptoSetupGetHandshakeOpenerCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupGetHandshakeOpenerCall) Do(f func() (handshake.LongHeaderOpener, error)) *CryptoSetupGetHandshakeOpenerCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupGetHandshakeOpenerCall) DoAndReturn(f func() (handshake.LongHeaderOpener, error)) *CryptoSetupGetHandshakeOpenerCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // GetHandshakeSealer mocks base method.
@ -176,9 +392,33 @@ func (m *MockCryptoSetup) GetHandshakeSealer() (handshake.LongHeaderSealer, erro
 }

 // GetHandshakeSealer indicates an expected call of GetHandshakeSealer.
-func (mr *MockCryptoSetupMockRecorder) GetHandshakeSealer() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) GetHandshakeSealer() *CryptoSetupGetHandshakeSealerCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetHandshakeSealer", reflect.TypeOf((*MockCryptoSetup)(nil).GetHandshakeSealer))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetHandshakeSealer", reflect.TypeOf((*MockCryptoSetup)(nil).GetHandshakeSealer))
+	return &CryptoSetupGetHandshakeSealerCall{Call: call}
+}
+
+// CryptoSetupGetHandshakeSealerCall wrap *gomock.Call
+type CryptoSetupGetHandshakeSealerCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupGetHandshakeSealerCall) Return(arg0 handshake.LongHeaderSealer, arg1 error) *CryptoSetupGetHandshakeSealerCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupGetHandshakeSealerCall) Do(f func() (handshake.LongHeaderSealer, error)) *CryptoSetupGetHandshakeSealerCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupGetHandshakeSealerCall) DoAndReturn(f func() (handshake.LongHeaderSealer, error)) *CryptoSetupGetHandshakeSealerCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // GetInitialOpener mocks base method.
@ -191,9 +431,33 @@ func (m *MockCryptoSetup) GetInitialOpener() (handshake.LongHeaderOpener, error)
 }

 // GetInitialOpener indicates an expected call of GetInitialOpener.
-func (mr *MockCryptoSetupMockRecorder) GetInitialOpener() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) GetInitialOpener() *CryptoSetupGetInitialOpenerCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInitialOpener", reflect.TypeOf((*MockCryptoSetup)(nil).GetInitialOpener))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInitialOpener", reflect.TypeOf((*MockCryptoSetup)(nil).GetInitialOpener))
+	return &CryptoSetupGetInitialOpenerCall{Call: call}
+}
+
+// CryptoSetupGetInitialOpenerCall wrap *gomock.Call
+type CryptoSetupGetInitialOpenerCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupGetInitialOpenerCall) Return(arg0 handshake.LongHeaderOpener, arg1 error) *CryptoSetupGetInitialOpenerCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupGetInitialOpenerCall) Do(f func() (handshake.LongHeaderOpener, error)) *CryptoSetupGetInitialOpenerCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupGetInitialOpenerCall) DoAndReturn(f func() (handshake.LongHeaderOpener, error)) *CryptoSetupGetInitialOpenerCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // GetInitialSealer mocks base method.
@ -206,9 +470,33 @@ func (m *MockCryptoSetup) GetInitialSealer() (handshake.LongHeaderSealer, error)
 }

 // GetInitialSealer indicates an expected call of GetInitialSealer.
-func (mr *MockCryptoSetupMockRecorder) GetInitialSealer() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) GetInitialSealer() *CryptoSetupGetInitialSealerCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInitialSealer", reflect.TypeOf((*MockCryptoSetup)(nil).GetInitialSealer))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInitialSealer", reflect.TypeOf((*MockCryptoSetup)(nil).GetInitialSealer))
+	return &CryptoSetupGetInitialSealerCall{Call: call}
+}
+
+// CryptoSetupGetInitialSealerCall wrap *gomock.Call
+type CryptoSetupGetInitialSealerCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupGetInitialSealerCall) Return(arg0 handshake.LongHeaderSealer, arg1 error) *CryptoSetupGetInitialSealerCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupGetInitialSealerCall) Do(f func() (handshake.LongHeaderSealer, error)) *CryptoSetupGetInitialSealerCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupGetInitialSealerCall) DoAndReturn(f func() (handshake.LongHeaderSealer, error)) *CryptoSetupGetInitialSealerCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // GetSessionTicket mocks base method.
@ -221,9 +509,33 @@ func (m *MockCryptoSetup) GetSessionTicket() ([]byte, error) {
 }

 // GetSessionTicket indicates an expected call of GetSessionTicket.
-func (mr *MockCryptoSetupMockRecorder) GetSessionTicket() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) GetSessionTicket() *CryptoSetupGetSessionTicketCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSessionTicket", reflect.TypeOf((*MockCryptoSetup)(nil).GetSessionTicket))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSessionTicket", reflect.TypeOf((*MockCryptoSetup)(nil).GetSessionTicket))
+	return &CryptoSetupGetSessionTicketCall{Call: call}
+}
+
+// CryptoSetupGetSessionTicketCall wrap *gomock.Call
+type CryptoSetupGetSessionTicketCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupGetSessionTicketCall) Return(arg0 []byte, arg1 error) *CryptoSetupGetSessionTicketCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupGetSessionTicketCall) Do(f func() ([]byte, error)) *CryptoSetupGetSessionTicketCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupGetSessionTicketCall) DoAndReturn(f func() ([]byte, error)) *CryptoSetupGetSessionTicketCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // HandleMessage mocks base method.
@ -235,9 +547,33 @@ func (m *MockCryptoSetup) HandleMessage(arg0 []byte, arg1 protocol.EncryptionLev
 }

 // HandleMessage indicates an expected call of HandleMessage.
-func (mr *MockCryptoSetupMockRecorder) HandleMessage(arg0, arg1 any) *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) HandleMessage(arg0, arg1 any) *CryptoSetupHandleMessageCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HandleMessage", reflect.TypeOf((*MockCryptoSetup)(nil).HandleMessage), arg0, arg1)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HandleMessage", reflect.TypeOf((*MockCryptoSetup)(nil).HandleMessage), arg0, arg1)
+	return &CryptoSetupHandleMessageCall{Call: call}
+}
+
+// CryptoSetupHandleMessageCall wrap *gomock.Call
+type CryptoSetupHandleMessageCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupHandleMessageCall) Return(arg0 error) *CryptoSetupHandleMessageCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupHandleMessageCall) Do(f func([]byte, protocol.EncryptionLevel) error) *CryptoSetupHandleMessageCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupHandleMessageCall) DoAndReturn(f func([]byte, protocol.EncryptionLevel) error) *CryptoSetupHandleMessageCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // NextEvent mocks base method.
@ -249,9 +585,33 @@ func (m *MockCryptoSetup) NextEvent() handshake.Event {
 }

 // NextEvent indicates an expected call of NextEvent.
-func (mr *MockCryptoSetupMockRecorder) NextEvent() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) NextEvent() *CryptoSetupNextEventCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "NextEvent", reflect.TypeOf((*MockCryptoSetup)(nil).NextEvent))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "NextEvent", reflect.TypeOf((*MockCryptoSetup)(nil).NextEvent))
+	return &CryptoSetupNextEventCall{Call: call}
+}
+
+// CryptoSetupNextEventCall wrap *gomock.Call
+type CryptoSetupNextEventCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupNextEventCall) Return(arg0 handshake.Event) *CryptoSetupNextEventCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupNextEventCall) Do(f func() handshake.Event) *CryptoSetupNextEventCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupNextEventCall) DoAndReturn(f func() handshake.Event) *CryptoSetupNextEventCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SetHandshakeConfirmed mocks base method.
@ -261,9 +621,33 @@ func (m *MockCryptoSetup) SetHandshakeConfirmed() {
 }

 // SetHandshakeConfirmed indicates an expected call of SetHandshakeConfirmed.
-func (mr *MockCryptoSetupMockRecorder) SetHandshakeConfirmed() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) SetHandshakeConfirmed() *CryptoSetupSetHandshakeConfirmedCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetHandshakeConfirmed", reflect.TypeOf((*MockCryptoSetup)(nil).SetHandshakeConfirmed))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetHandshakeConfirmed", reflect.TypeOf((*MockCryptoSetup)(nil).SetHandshakeConfirmed))
+	return &CryptoSetupSetHandshakeConfirmedCall{Call: call}
+}
+
+// CryptoSetupSetHandshakeConfirmedCall wrap *gomock.Call
+type CryptoSetupSetHandshakeConfirmedCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupSetHandshakeConfirmedCall) Return() *CryptoSetupSetHandshakeConfirmedCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupSetHandshakeConfirmedCall) Do(f func()) *CryptoSetupSetHandshakeConfirmedCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupSetHandshakeConfirmedCall) DoAndReturn(f func()) *CryptoSetupSetHandshakeConfirmedCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SetLargest1RTTAcked mocks base method.
@ -275,9 +659,33 @@ func (m *MockCryptoSetup) SetLargest1RTTAcked(arg0 protocol.PacketNumber) error
 }

 // SetLargest1RTTAcked indicates an expected call of SetLargest1RTTAcked.
-func (mr *MockCryptoSetupMockRecorder) SetLargest1RTTAcked(arg0 any) *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) SetLargest1RTTAcked(arg0 any) *CryptoSetupSetLargest1RTTAckedCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetLargest1RTTAcked", reflect.TypeOf((*MockCryptoSetup)(nil).SetLargest1RTTAcked), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetLargest1RTTAcked", reflect.TypeOf((*MockCryptoSetup)(nil).SetLargest1RTTAcked), arg0)
+	return &CryptoSetupSetLargest1RTTAckedCall{Call: call}
+}
+
+// CryptoSetupSetLargest1RTTAckedCall wrap *gomock.Call
+type CryptoSetupSetLargest1RTTAckedCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupSetLargest1RTTAckedCall) Return(arg0 error) *CryptoSetupSetLargest1RTTAckedCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupSetLargest1RTTAckedCall) Do(f func(protocol.PacketNumber) error) *CryptoSetupSetLargest1RTTAckedCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupSetLargest1RTTAckedCall) DoAndReturn(f func(protocol.PacketNumber) error) *CryptoSetupSetLargest1RTTAckedCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // StartHandshake mocks base method.
@ -289,7 +697,31 @@ func (m *MockCryptoSetup) StartHandshake() error {
 }

 // StartHandshake indicates an expected call of StartHandshake.
-func (mr *MockCryptoSetupMockRecorder) StartHandshake() *gomock.Call {
+func (mr *MockCryptoSetupMockRecorder) StartHandshake() *CryptoSetupStartHandshakeCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StartHandshake", reflect.TypeOf((*MockCryptoSetup)(nil).StartHandshake))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StartHandshake", reflect.TypeOf((*MockCryptoSetup)(nil).StartHandshake))
+	return &CryptoSetupStartHandshakeCall{Call: call}
+}
+
+// CryptoSetupStartHandshakeCall wrap *gomock.Call
+type CryptoSetupStartHandshakeCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *CryptoSetupStartHandshakeCall) Return(arg0 error) *CryptoSetupStartHandshakeCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *CryptoSetupStartHandshakeCall) Do(f func() error) *CryptoSetupStartHandshakeCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *CryptoSetupStartHandshakeCall) DoAndReturn(f func() error) *CryptoSetupStartHandshakeCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/logging/internal/connection_tracer.go
+++ b/internal/mocks/logging/internal/connection_tracer.go
--- a/internal/mocks/logging/internal/tracer.go
+++ b/internal/mocks/logging/internal/tracer.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package internal -destination internal/tracer.go github.com/quic-go/quic-go/internal/mocks/logging Tracer
+//	mockgen -typed -build_flags=-tags=gomock -package internal -destination internal/tracer.go github.com/quic-go/quic-go/internal/mocks/logging Tracer
 //
 // Package internal is a generated GoMock package.
 package internal
@ -48,9 +48,33 @@ func (m *MockTracer) DroppedPacket(arg0 net.Addr, arg1 logging.PacketType, arg2
 }

 // DroppedPacket indicates an expected call of DroppedPacket.
-func (mr *MockTracerMockRecorder) DroppedPacket(arg0, arg1, arg2, arg3 any) *gomock.Call {
+func (mr *MockTracerMockRecorder) DroppedPacket(arg0, arg1, arg2, arg3 any) *TracerDroppedPacketCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DroppedPacket", reflect.TypeOf((*MockTracer)(nil).DroppedPacket), arg0, arg1, arg2, arg3)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DroppedPacket", reflect.TypeOf((*MockTracer)(nil).DroppedPacket), arg0, arg1, arg2, arg3)
+	return &TracerDroppedPacketCall{Call: call}
+}
+
+// TracerDroppedPacketCall wrap *gomock.Call
+type TracerDroppedPacketCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *TracerDroppedPacketCall) Return() *TracerDroppedPacketCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *TracerDroppedPacketCall) Do(f func(net.Addr, logging.PacketType, protocol.ByteCount, logging.PacketDropReason)) *TracerDroppedPacketCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *TracerDroppedPacketCall) DoAndReturn(f func(net.Addr, logging.PacketType, protocol.ByteCount, logging.PacketDropReason)) *TracerDroppedPacketCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SentPacket mocks base method.
@ -60,9 +84,33 @@ func (m *MockTracer) SentPacket(arg0 net.Addr, arg1 *wire.Header, arg2 protocol.
 }

 // SentPacket indicates an expected call of SentPacket.
-func (mr *MockTracerMockRecorder) SentPacket(arg0, arg1, arg2, arg3 any) *gomock.Call {
+func (mr *MockTracerMockRecorder) SentPacket(arg0, arg1, arg2, arg3 any) *TracerSentPacketCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SentPacket", reflect.TypeOf((*MockTracer)(nil).SentPacket), arg0, arg1, arg2, arg3)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SentPacket", reflect.TypeOf((*MockTracer)(nil).SentPacket), arg0, arg1, arg2, arg3)
+	return &TracerSentPacketCall{Call: call}
+}
+
+// TracerSentPacketCall wrap *gomock.Call
+type TracerSentPacketCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *TracerSentPacketCall) Return() *TracerSentPacketCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *TracerSentPacketCall) Do(f func(net.Addr, *wire.Header, protocol.ByteCount, []logging.Frame)) *TracerSentPacketCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *TracerSentPacketCall) DoAndReturn(f func(net.Addr, *wire.Header, protocol.ByteCount, []logging.Frame)) *TracerSentPacketCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SentVersionNegotiationPacket mocks base method.
@ -72,7 +120,31 @@ func (m *MockTracer) SentVersionNegotiationPacket(arg0 net.Addr, arg1, arg2 prot
 }

 // SentVersionNegotiationPacket indicates an expected call of SentVersionNegotiationPacket.
-func (mr *MockTracerMockRecorder) SentVersionNegotiationPacket(arg0, arg1, arg2, arg3 any) *gomock.Call {
+func (mr *MockTracerMockRecorder) SentVersionNegotiationPacket(arg0, arg1, arg2, arg3 any) *TracerSentVersionNegotiationPacketCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SentVersionNegotiationPacket", reflect.TypeOf((*MockTracer)(nil).SentVersionNegotiationPacket), arg0, arg1, arg2, arg3)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SentVersionNegotiationPacket", reflect.TypeOf((*MockTracer)(nil).SentVersionNegotiationPacket), arg0, arg1, arg2, arg3)
+	return &TracerSentVersionNegotiationPacketCall{Call: call}
+}
+
+// TracerSentVersionNegotiationPacketCall wrap *gomock.Call
+type TracerSentVersionNegotiationPacketCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *TracerSentVersionNegotiationPacketCall) Return() *TracerSentVersionNegotiationPacketCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *TracerSentVersionNegotiationPacketCall) Do(f func(net.Addr, protocol.ArbitraryLenConnectionID, protocol.ArbitraryLenConnectionID, []protocol.VersionNumber)) *TracerSentVersionNegotiationPacketCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *TracerSentVersionNegotiationPacketCall) DoAndReturn(f func(net.Addr, protocol.ArbitraryLenConnectionID, protocol.ArbitraryLenConnectionID, []protocol.VersionNumber)) *TracerSentVersionNegotiationPacketCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/logging/mockgen.go
+++ b/internal/mocks/logging/mockgen.go
@ -9,14 +9,14 @@ import (
 	"github.com/quic-go/quic-go/logging"
 )

-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package internal -destination internal/tracer.go github.com/quic-go/quic-go/internal/mocks/logging Tracer"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package internal -destination internal/tracer.go github.com/quic-go/quic-go/internal/mocks/logging Tracer"
 type Tracer interface {
 	SentPacket(net.Addr, *logging.Header, logging.ByteCount, []logging.Frame)
 	SentVersionNegotiationPacket(_ net.Addr, dest, src logging.ArbitraryLenConnectionID, _ []logging.VersionNumber)
 	DroppedPacket(net.Addr, logging.PacketType, logging.ByteCount, logging.PacketDropReason)
 }

-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package internal -destination internal/connection_tracer.go github.com/quic-go/quic-go/internal/mocks/logging ConnectionTracer"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package internal -destination internal/connection_tracer.go github.com/quic-go/quic-go/internal/mocks/logging ConnectionTracer"
 type ConnectionTracer interface {
 	StartedConnection(local, remote net.Addr, srcConnID, destConnID logging.ConnectionID)
 	NegotiatedVersion(chosen logging.VersionNumber, clientVersions, serverVersions []logging.VersionNumber)
--- a/internal/mocks/long_header_opener.go
+++ b/internal/mocks/long_header_opener.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mocks -destination long_header_opener.go github.com/quic-go/quic-go/internal/handshake LongHeaderOpener
+//	mockgen -typed -build_flags=-tags=gomock -package mocks -destination long_header_opener.go github.com/quic-go/quic-go/internal/handshake LongHeaderOpener
 //
 // Package mocks is a generated GoMock package.
 package mocks
@ -47,9 +47,33 @@ func (m *MockLongHeaderOpener) DecodePacketNumber(arg0 protocol.PacketNumber, ar
 }

 // DecodePacketNumber indicates an expected call of DecodePacketNumber.
-func (mr *MockLongHeaderOpenerMockRecorder) DecodePacketNumber(arg0, arg1 any) *gomock.Call {
+func (mr *MockLongHeaderOpenerMockRecorder) DecodePacketNumber(arg0, arg1 any) *LongHeaderOpenerDecodePacketNumberCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DecodePacketNumber", reflect.TypeOf((*MockLongHeaderOpener)(nil).DecodePacketNumber), arg0, arg1)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DecodePacketNumber", reflect.TypeOf((*MockLongHeaderOpener)(nil).DecodePacketNumber), arg0, arg1)
+	return &LongHeaderOpenerDecodePacketNumberCall{Call: call}
+}
+
+// LongHeaderOpenerDecodePacketNumberCall wrap *gomock.Call
+type LongHeaderOpenerDecodePacketNumberCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *LongHeaderOpenerDecodePacketNumberCall) Return(arg0 protocol.PacketNumber) *LongHeaderOpenerDecodePacketNumberCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *LongHeaderOpenerDecodePacketNumberCall) Do(f func(protocol.PacketNumber, protocol.PacketNumberLen) protocol.PacketNumber) *LongHeaderOpenerDecodePacketNumberCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *LongHeaderOpenerDecodePacketNumberCall) DoAndReturn(f func(protocol.PacketNumber, protocol.PacketNumberLen) protocol.PacketNumber) *LongHeaderOpenerDecodePacketNumberCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // DecryptHeader mocks base method.
@ -59,9 +83,33 @@ func (m *MockLongHeaderOpener) DecryptHeader(arg0 []byte, arg1 *byte, arg2 []byt
 }

 // DecryptHeader indicates an expected call of DecryptHeader.
-func (mr *MockLongHeaderOpenerMockRecorder) DecryptHeader(arg0, arg1, arg2 any) *gomock.Call {
+func (mr *MockLongHeaderOpenerMockRecorder) DecryptHeader(arg0, arg1, arg2 any) *LongHeaderOpenerDecryptHeaderCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DecryptHeader", reflect.TypeOf((*MockLongHeaderOpener)(nil).DecryptHeader), arg0, arg1, arg2)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DecryptHeader", reflect.TypeOf((*MockLongHeaderOpener)(nil).DecryptHeader), arg0, arg1, arg2)
+	return &LongHeaderOpenerDecryptHeaderCall{Call: call}
+}
+
+// LongHeaderOpenerDecryptHeaderCall wrap *gomock.Call
+type LongHeaderOpenerDecryptHeaderCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *LongHeaderOpenerDecryptHeaderCall) Return() *LongHeaderOpenerDecryptHeaderCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *LongHeaderOpenerDecryptHeaderCall) Do(f func([]byte, *byte, []byte)) *LongHeaderOpenerDecryptHeaderCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *LongHeaderOpenerDecryptHeaderCall) DoAndReturn(f func([]byte, *byte, []byte)) *LongHeaderOpenerDecryptHeaderCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Open mocks base method.
@ -74,7 +122,31 @@ func (m *MockLongHeaderOpener) Open(arg0, arg1 []byte, arg2 protocol.PacketNumbe
 }

 // Open indicates an expected call of Open.
-func (mr *MockLongHeaderOpenerMockRecorder) Open(arg0, arg1, arg2, arg3 any) *gomock.Call {
+func (mr *MockLongHeaderOpenerMockRecorder) Open(arg0, arg1, arg2, arg3 any) *LongHeaderOpenerOpenCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Open", reflect.TypeOf((*MockLongHeaderOpener)(nil).Open), arg0, arg1, arg2, arg3)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Open", reflect.TypeOf((*MockLongHeaderOpener)(nil).Open), arg0, arg1, arg2, arg3)
+	return &LongHeaderOpenerOpenCall{Call: call}
+}
+
+// LongHeaderOpenerOpenCall wrap *gomock.Call
+type LongHeaderOpenerOpenCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *LongHeaderOpenerOpenCall) Return(arg0 []byte, arg1 error) *LongHeaderOpenerOpenCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *LongHeaderOpenerOpenCall) Do(f func([]byte, []byte, protocol.PacketNumber, []byte) ([]byte, error)) *LongHeaderOpenerOpenCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *LongHeaderOpenerOpenCall) DoAndReturn(f func([]byte, []byte, protocol.PacketNumber, []byte) ([]byte, error)) *LongHeaderOpenerOpenCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/mockgen.go
+++ b/internal/mocks/mockgen.go
@ -2,18 +2,18 @@

 package mocks

-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mockquic -destination quic/stream.go github.com/quic-go/quic-go Stream"
-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mockquic -destination quic/early_conn_tmp.go github.com/quic-go/quic-go EarlyConnection && sed 's/qtls.ConnectionState/quic.ConnectionState/g' quic/early_conn_tmp.go > quic/early_conn.go && rm quic/early_conn_tmp.go && go run golang.org/x/tools/cmd/goimports -w quic/early_conn.go"
-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mocks -destination short_header_sealer.go github.com/quic-go/quic-go/internal/handshake ShortHeaderSealer"
-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mocks -destination short_header_opener.go github.com/quic-go/quic-go/internal/handshake ShortHeaderOpener"
-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mocks -destination long_header_opener.go github.com/quic-go/quic-go/internal/handshake LongHeaderOpener"
-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mocks -destination crypto_setup_tmp.go github.com/quic-go/quic-go/internal/handshake CryptoSetup && sed -E 's~github.com/quic-go/qtls[[:alnum:]_-]*~github.com/quic-go/quic-go/internal/qtls~g; s~qtls.ConnectionStateWith0RTT~qtls.ConnectionState~g' crypto_setup_tmp.go > crypto_setup.go && rm crypto_setup_tmp.go && go run golang.org/x/tools/cmd/goimports -w crypto_setup.go"
-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mocks -destination stream_flow_controller.go github.com/quic-go/quic-go/internal/flowcontrol StreamFlowController"
-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mocks -destination congestion.go github.com/quic-go/quic-go/internal/congestion SendAlgorithmWithDebugInfos"
-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mocks -destination connection_flow_controller.go github.com/quic-go/quic-go/internal/flowcontrol ConnectionFlowController"
-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mockackhandler -destination ackhandler/sent_packet_handler.go github.com/quic-go/quic-go/internal/ackhandler SentPacketHandler"
-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mockackhandler -destination ackhandler/received_packet_handler.go github.com/quic-go/quic-go/internal/ackhandler ReceivedPacketHandler"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mockquic -destination quic/stream.go github.com/quic-go/quic-go Stream"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mockquic -destination quic/early_conn_tmp.go github.com/quic-go/quic-go EarlyConnection && sed 's/qtls.ConnectionState/quic.ConnectionState/g' quic/early_conn_tmp.go > quic/early_conn.go && rm quic/early_conn_tmp.go && go run golang.org/x/tools/cmd/goimports -w quic/early_conn.go"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mocks -destination short_header_sealer.go github.com/quic-go/quic-go/internal/handshake ShortHeaderSealer"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mocks -destination short_header_opener.go github.com/quic-go/quic-go/internal/handshake ShortHeaderOpener"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mocks -destination long_header_opener.go github.com/quic-go/quic-go/internal/handshake LongHeaderOpener"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mocks -destination crypto_setup_tmp.go github.com/quic-go/quic-go/internal/handshake CryptoSetup && sed -E 's~github.com/quic-go/qtls[[:alnum:]_-]*~github.com/quic-go/quic-go/internal/qtls~g; s~qtls.ConnectionStateWith0RTT~qtls.ConnectionState~g' crypto_setup_tmp.go > crypto_setup.go && rm crypto_setup_tmp.go && go run golang.org/x/tools/cmd/goimports -w crypto_setup.go"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mocks -destination stream_flow_controller.go github.com/quic-go/quic-go/internal/flowcontrol StreamFlowController"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mocks -destination congestion.go github.com/quic-go/quic-go/internal/congestion SendAlgorithmWithDebugInfos"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mocks -destination connection_flow_controller.go github.com/quic-go/quic-go/internal/flowcontrol ConnectionFlowController"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mockackhandler -destination ackhandler/sent_packet_handler.go github.com/quic-go/quic-go/internal/ackhandler SentPacketHandler"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mockackhandler -destination ackhandler/received_packet_handler.go github.com/quic-go/quic-go/internal/ackhandler ReceivedPacketHandler"

 // The following command produces a warning message on OSX, however, it still generates the correct mock file.
 // See https://github.com/golang/mock/issues/339 for details.
-//go:generate sh -c "go run go.uber.org/mock/mockgen -build_flags=\"-tags=gomock\" -package mocktls -destination tls/client_session_cache.go crypto/tls ClientSessionCache"
+//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mocktls -destination tls/client_session_cache.go crypto/tls ClientSessionCache"
--- a/internal/mocks/quic/early_conn.go
+++ b/internal/mocks/quic/early_conn.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mockquic -destination quic/early_conn_tmp.go github.com/quic-go/quic-go EarlyConnection
+//	mockgen -typed -build_flags=-tags=gomock -package mockquic -destination quic/early_conn_tmp.go github.com/quic-go/quic-go EarlyConnection
 //
 // Package mockquic is a generated GoMock package.
 package mockquic
@ -51,9 +51,33 @@ func (m *MockEarlyConnection) AcceptStream(arg0 context.Context) (quic.Stream, e
 }

 // AcceptStream indicates an expected call of AcceptStream.
-func (mr *MockEarlyConnectionMockRecorder) AcceptStream(arg0 any) *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) AcceptStream(arg0 any) *EarlyConnectionAcceptStreamCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AcceptStream", reflect.TypeOf((*MockEarlyConnection)(nil).AcceptStream), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AcceptStream", reflect.TypeOf((*MockEarlyConnection)(nil).AcceptStream), arg0)
+	return &EarlyConnectionAcceptStreamCall{Call: call}
+}
+
+// EarlyConnectionAcceptStreamCall wrap *gomock.Call
+type EarlyConnectionAcceptStreamCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionAcceptStreamCall) Return(arg0 quic.Stream, arg1 error) *EarlyConnectionAcceptStreamCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionAcceptStreamCall) Do(f func(context.Context) (quic.Stream, error)) *EarlyConnectionAcceptStreamCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionAcceptStreamCall) DoAndReturn(f func(context.Context) (quic.Stream, error)) *EarlyConnectionAcceptStreamCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // AcceptUniStream mocks base method.
@ -66,9 +90,33 @@ func (m *MockEarlyConnection) AcceptUniStream(arg0 context.Context) (quic.Receiv
 }

 // AcceptUniStream indicates an expected call of AcceptUniStream.
-func (mr *MockEarlyConnectionMockRecorder) AcceptUniStream(arg0 any) *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) AcceptUniStream(arg0 any) *EarlyConnectionAcceptUniStreamCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AcceptUniStream", reflect.TypeOf((*MockEarlyConnection)(nil).AcceptUniStream), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AcceptUniStream", reflect.TypeOf((*MockEarlyConnection)(nil).AcceptUniStream), arg0)
+	return &EarlyConnectionAcceptUniStreamCall{Call: call}
+}
+
+// EarlyConnectionAcceptUniStreamCall wrap *gomock.Call
+type EarlyConnectionAcceptUniStreamCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionAcceptUniStreamCall) Return(arg0 quic.ReceiveStream, arg1 error) *EarlyConnectionAcceptUniStreamCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionAcceptUniStreamCall) Do(f func(context.Context) (quic.ReceiveStream, error)) *EarlyConnectionAcceptUniStreamCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionAcceptUniStreamCall) DoAndReturn(f func(context.Context) (quic.ReceiveStream, error)) *EarlyConnectionAcceptUniStreamCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // CloseWithError mocks base method.
@ -80,9 +128,33 @@ func (m *MockEarlyConnection) CloseWithError(arg0 qerr.ApplicationErrorCode, arg
 }

 // CloseWithError indicates an expected call of CloseWithError.
-func (mr *MockEarlyConnectionMockRecorder) CloseWithError(arg0, arg1 any) *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) CloseWithError(arg0, arg1 any) *EarlyConnectionCloseWithErrorCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CloseWithError", reflect.TypeOf((*MockEarlyConnection)(nil).CloseWithError), arg0, arg1)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CloseWithError", reflect.TypeOf((*MockEarlyConnection)(nil).CloseWithError), arg0, arg1)
+	return &EarlyConnectionCloseWithErrorCall{Call: call}
+}
+
+// EarlyConnectionCloseWithErrorCall wrap *gomock.Call
+type EarlyConnectionCloseWithErrorCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionCloseWithErrorCall) Return(arg0 error) *EarlyConnectionCloseWithErrorCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionCloseWithErrorCall) Do(f func(qerr.ApplicationErrorCode, string) error) *EarlyConnectionCloseWithErrorCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionCloseWithErrorCall) DoAndReturn(f func(qerr.ApplicationErrorCode, string) error) *EarlyConnectionCloseWithErrorCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // ConnectionState mocks base method.
@ -94,9 +166,33 @@ func (m *MockEarlyConnection) ConnectionState() quic.ConnectionState {
 }

 // ConnectionState indicates an expected call of ConnectionState.
-func (mr *MockEarlyConnectionMockRecorder) ConnectionState() *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) ConnectionState() *EarlyConnectionConnectionStateCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ConnectionState", reflect.TypeOf((*MockEarlyConnection)(nil).ConnectionState))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ConnectionState", reflect.TypeOf((*MockEarlyConnection)(nil).ConnectionState))
+	return &EarlyConnectionConnectionStateCall{Call: call}
+}
+
+// EarlyConnectionConnectionStateCall wrap *gomock.Call
+type EarlyConnectionConnectionStateCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionConnectionStateCall) Return(arg0 quic.ConnectionState) *EarlyConnectionConnectionStateCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionConnectionStateCall) Do(f func() quic.ConnectionState) *EarlyConnectionConnectionStateCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionConnectionStateCall) DoAndReturn(f func() quic.ConnectionState) *EarlyConnectionConnectionStateCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Context mocks base method.
@ -108,9 +204,33 @@ func (m *MockEarlyConnection) Context() context.Context {
 }

 // Context indicates an expected call of Context.
-func (mr *MockEarlyConnectionMockRecorder) Context() *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) Context() *EarlyConnectionContextCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Context", reflect.TypeOf((*MockEarlyConnection)(nil).Context))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Context", reflect.TypeOf((*MockEarlyConnection)(nil).Context))
+	return &EarlyConnectionContextCall{Call: call}
+}
+
+// EarlyConnectionContextCall wrap *gomock.Call
+type EarlyConnectionContextCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionContextCall) Return(arg0 context.Context) *EarlyConnectionContextCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionContextCall) Do(f func() context.Context) *EarlyConnectionContextCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionContextCall) DoAndReturn(f func() context.Context) *EarlyConnectionContextCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // HandshakeComplete mocks base method.
@ -122,9 +242,33 @@ func (m *MockEarlyConnection) HandshakeComplete() <-chan struct{} {
 }

 // HandshakeComplete indicates an expected call of HandshakeComplete.
-func (mr *MockEarlyConnectionMockRecorder) HandshakeComplete() *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) HandshakeComplete() *EarlyConnectionHandshakeCompleteCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HandshakeComplete", reflect.TypeOf((*MockEarlyConnection)(nil).HandshakeComplete))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HandshakeComplete", reflect.TypeOf((*MockEarlyConnection)(nil).HandshakeComplete))
+	return &EarlyConnectionHandshakeCompleteCall{Call: call}
+}
+
+// EarlyConnectionHandshakeCompleteCall wrap *gomock.Call
+type EarlyConnectionHandshakeCompleteCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionHandshakeCompleteCall) Return(arg0 <-chan struct{}) *EarlyConnectionHandshakeCompleteCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionHandshakeCompleteCall) Do(f func() <-chan struct{}) *EarlyConnectionHandshakeCompleteCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionHandshakeCompleteCall) DoAndReturn(f func() <-chan struct{}) *EarlyConnectionHandshakeCompleteCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // LocalAddr mocks base method.
@ -136,9 +280,33 @@ func (m *MockEarlyConnection) LocalAddr() net.Addr {
 }

 // LocalAddr indicates an expected call of LocalAddr.
-func (mr *MockEarlyConnectionMockRecorder) LocalAddr() *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) LocalAddr() *EarlyConnectionLocalAddrCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LocalAddr", reflect.TypeOf((*MockEarlyConnection)(nil).LocalAddr))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LocalAddr", reflect.TypeOf((*MockEarlyConnection)(nil).LocalAddr))
+	return &EarlyConnectionLocalAddrCall{Call: call}
+}
+
+// EarlyConnectionLocalAddrCall wrap *gomock.Call
+type EarlyConnectionLocalAddrCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionLocalAddrCall) Return(arg0 net.Addr) *EarlyConnectionLocalAddrCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionLocalAddrCall) Do(f func() net.Addr) *EarlyConnectionLocalAddrCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionLocalAddrCall) DoAndReturn(f func() net.Addr) *EarlyConnectionLocalAddrCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // NextConnection mocks base method.
@ -150,9 +318,33 @@ func (m *MockEarlyConnection) NextConnection() quic.Connection {
 }

 // NextConnection indicates an expected call of NextConnection.
-func (mr *MockEarlyConnectionMockRecorder) NextConnection() *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) NextConnection() *EarlyConnectionNextConnectionCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "NextConnection", reflect.TypeOf((*MockEarlyConnection)(nil).NextConnection))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "NextConnection", reflect.TypeOf((*MockEarlyConnection)(nil).NextConnection))
+	return &EarlyConnectionNextConnectionCall{Call: call}
+}
+
+// EarlyConnectionNextConnectionCall wrap *gomock.Call
+type EarlyConnectionNextConnectionCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionNextConnectionCall) Return(arg0 quic.Connection) *EarlyConnectionNextConnectionCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionNextConnectionCall) Do(f func() quic.Connection) *EarlyConnectionNextConnectionCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionNextConnectionCall) DoAndReturn(f func() quic.Connection) *EarlyConnectionNextConnectionCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // OpenStream mocks base method.
@ -165,9 +357,33 @@ func (m *MockEarlyConnection) OpenStream() (quic.Stream, error) {
 }

 // OpenStream indicates an expected call of OpenStream.
-func (mr *MockEarlyConnectionMockRecorder) OpenStream() *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) OpenStream() *EarlyConnectionOpenStreamCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OpenStream", reflect.TypeOf((*MockEarlyConnection)(nil).OpenStream))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OpenStream", reflect.TypeOf((*MockEarlyConnection)(nil).OpenStream))
+	return &EarlyConnectionOpenStreamCall{Call: call}
+}
+
+// EarlyConnectionOpenStreamCall wrap *gomock.Call
+type EarlyConnectionOpenStreamCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionOpenStreamCall) Return(arg0 quic.Stream, arg1 error) *EarlyConnectionOpenStreamCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionOpenStreamCall) Do(f func() (quic.Stream, error)) *EarlyConnectionOpenStreamCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionOpenStreamCall) DoAndReturn(f func() (quic.Stream, error)) *EarlyConnectionOpenStreamCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // OpenStreamSync mocks base method.
@ -180,9 +396,33 @@ func (m *MockEarlyConnection) OpenStreamSync(arg0 context.Context) (quic.Stream,
 }

 // OpenStreamSync indicates an expected call of OpenStreamSync.
-func (mr *MockEarlyConnectionMockRecorder) OpenStreamSync(arg0 any) *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) OpenStreamSync(arg0 any) *EarlyConnectionOpenStreamSyncCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OpenStreamSync", reflect.TypeOf((*MockEarlyConnection)(nil).OpenStreamSync), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OpenStreamSync", reflect.TypeOf((*MockEarlyConnection)(nil).OpenStreamSync), arg0)
+	return &EarlyConnectionOpenStreamSyncCall{Call: call}
+}
+
+// EarlyConnectionOpenStreamSyncCall wrap *gomock.Call
+type EarlyConnectionOpenStreamSyncCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionOpenStreamSyncCall) Return(arg0 quic.Stream, arg1 error) *EarlyConnectionOpenStreamSyncCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionOpenStreamSyncCall) Do(f func(context.Context) (quic.Stream, error)) *EarlyConnectionOpenStreamSyncCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionOpenStreamSyncCall) DoAndReturn(f func(context.Context) (quic.Stream, error)) *EarlyConnectionOpenStreamSyncCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // OpenUniStream mocks base method.
@ -195,9 +435,33 @@ func (m *MockEarlyConnection) OpenUniStream() (quic.SendStream, error) {
 }

 // OpenUniStream indicates an expected call of OpenUniStream.
-func (mr *MockEarlyConnectionMockRecorder) OpenUniStream() *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) OpenUniStream() *EarlyConnectionOpenUniStreamCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OpenUniStream", reflect.TypeOf((*MockEarlyConnection)(nil).OpenUniStream))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OpenUniStream", reflect.TypeOf((*MockEarlyConnection)(nil).OpenUniStream))
+	return &EarlyConnectionOpenUniStreamCall{Call: call}
+}
+
+// EarlyConnectionOpenUniStreamCall wrap *gomock.Call
+type EarlyConnectionOpenUniStreamCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionOpenUniStreamCall) Return(arg0 quic.SendStream, arg1 error) *EarlyConnectionOpenUniStreamCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionOpenUniStreamCall) Do(f func() (quic.SendStream, error)) *EarlyConnectionOpenUniStreamCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionOpenUniStreamCall) DoAndReturn(f func() (quic.SendStream, error)) *EarlyConnectionOpenUniStreamCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // OpenUniStreamSync mocks base method.
@ -210,9 +474,33 @@ func (m *MockEarlyConnection) OpenUniStreamSync(arg0 context.Context) (quic.Send
 }

 // OpenUniStreamSync indicates an expected call of OpenUniStreamSync.
-func (mr *MockEarlyConnectionMockRecorder) OpenUniStreamSync(arg0 any) *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) OpenUniStreamSync(arg0 any) *EarlyConnectionOpenUniStreamSyncCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OpenUniStreamSync", reflect.TypeOf((*MockEarlyConnection)(nil).OpenUniStreamSync), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OpenUniStreamSync", reflect.TypeOf((*MockEarlyConnection)(nil).OpenUniStreamSync), arg0)
+	return &EarlyConnectionOpenUniStreamSyncCall{Call: call}
+}
+
+// EarlyConnectionOpenUniStreamSyncCall wrap *gomock.Call
+type EarlyConnectionOpenUniStreamSyncCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionOpenUniStreamSyncCall) Return(arg0 quic.SendStream, arg1 error) *EarlyConnectionOpenUniStreamSyncCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionOpenUniStreamSyncCall) Do(f func(context.Context) (quic.SendStream, error)) *EarlyConnectionOpenUniStreamSyncCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionOpenUniStreamSyncCall) DoAndReturn(f func(context.Context) (quic.SendStream, error)) *EarlyConnectionOpenUniStreamSyncCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // ReceiveMessage mocks base method.
@ -225,9 +513,33 @@ func (m *MockEarlyConnection) ReceiveMessage(arg0 context.Context) ([]byte, erro
 }

 // ReceiveMessage indicates an expected call of ReceiveMessage.
-func (mr *MockEarlyConnectionMockRecorder) ReceiveMessage(arg0 any) *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) ReceiveMessage(arg0 any) *EarlyConnectionReceiveMessageCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReceiveMessage", reflect.TypeOf((*MockEarlyConnection)(nil).ReceiveMessage), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReceiveMessage", reflect.TypeOf((*MockEarlyConnection)(nil).ReceiveMessage), arg0)
+	return &EarlyConnectionReceiveMessageCall{Call: call}
+}
+
+// EarlyConnectionReceiveMessageCall wrap *gomock.Call
+type EarlyConnectionReceiveMessageCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionReceiveMessageCall) Return(arg0 []byte, arg1 error) *EarlyConnectionReceiveMessageCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionReceiveMessageCall) Do(f func(context.Context) ([]byte, error)) *EarlyConnectionReceiveMessageCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionReceiveMessageCall) DoAndReturn(f func(context.Context) ([]byte, error)) *EarlyConnectionReceiveMessageCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // RemoteAddr mocks base method.
@ -239,9 +551,33 @@ func (m *MockEarlyConnection) RemoteAddr() net.Addr {
 }

 // RemoteAddr indicates an expected call of RemoteAddr.
-func (mr *MockEarlyConnectionMockRecorder) RemoteAddr() *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) RemoteAddr() *EarlyConnectionRemoteAddrCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoteAddr", reflect.TypeOf((*MockEarlyConnection)(nil).RemoteAddr))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoteAddr", reflect.TypeOf((*MockEarlyConnection)(nil).RemoteAddr))
+	return &EarlyConnectionRemoteAddrCall{Call: call}
+}
+
+// EarlyConnectionRemoteAddrCall wrap *gomock.Call
+type EarlyConnectionRemoteAddrCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionRemoteAddrCall) Return(arg0 net.Addr) *EarlyConnectionRemoteAddrCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionRemoteAddrCall) Do(f func() net.Addr) *EarlyConnectionRemoteAddrCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionRemoteAddrCall) DoAndReturn(f func() net.Addr) *EarlyConnectionRemoteAddrCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SendMessage mocks base method.
@ -253,7 +589,31 @@ func (m *MockEarlyConnection) SendMessage(arg0 []byte) error {
 }

 // SendMessage indicates an expected call of SendMessage.
-func (mr *MockEarlyConnectionMockRecorder) SendMessage(arg0 any) *gomock.Call {
+func (mr *MockEarlyConnectionMockRecorder) SendMessage(arg0 any) *EarlyConnectionSendMessageCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendMessage", reflect.TypeOf((*MockEarlyConnection)(nil).SendMessage), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendMessage", reflect.TypeOf((*MockEarlyConnection)(nil).SendMessage), arg0)
+	return &EarlyConnectionSendMessageCall{Call: call}
+}
+
+// EarlyConnectionSendMessageCall wrap *gomock.Call
+type EarlyConnectionSendMessageCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *EarlyConnectionSendMessageCall) Return(arg0 error) *EarlyConnectionSendMessageCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *EarlyConnectionSendMessageCall) Do(f func([]byte) error) *EarlyConnectionSendMessageCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *EarlyConnectionSendMessageCall) DoAndReturn(f func([]byte) error) *EarlyConnectionSendMessageCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/quic/stream.go
+++ b/internal/mocks/quic/stream.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mockquic -destination quic/stream.go github.com/quic-go/quic-go Stream
+//	mockgen -typed -build_flags=-tags=gomock -package mockquic -destination quic/stream.go github.com/quic-go/quic-go Stream
 //
 // Package mockquic is a generated GoMock package.
 package mockquic
@ -48,9 +48,33 @@ func (m *MockStream) CancelRead(arg0 qerr.StreamErrorCode) {
 }

 // CancelRead indicates an expected call of CancelRead.
-func (mr *MockStreamMockRecorder) CancelRead(arg0 any) *gomock.Call {
+func (mr *MockStreamMockRecorder) CancelRead(arg0 any) *StreamCancelReadCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CancelRead", reflect.TypeOf((*MockStream)(nil).CancelRead), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CancelRead", reflect.TypeOf((*MockStream)(nil).CancelRead), arg0)
+	return &StreamCancelReadCall{Call: call}
+}
+
+// StreamCancelReadCall wrap *gomock.Call
+type StreamCancelReadCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamCancelReadCall) Return() *StreamCancelReadCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamCancelReadCall) Do(f func(qerr.StreamErrorCode)) *StreamCancelReadCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamCancelReadCall) DoAndReturn(f func(qerr.StreamErrorCode)) *StreamCancelReadCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // CancelWrite mocks base method.
@ -60,9 +84,33 @@ func (m *MockStream) CancelWrite(arg0 qerr.StreamErrorCode) {
 }

 // CancelWrite indicates an expected call of CancelWrite.
-func (mr *MockStreamMockRecorder) CancelWrite(arg0 any) *gomock.Call {
+func (mr *MockStreamMockRecorder) CancelWrite(arg0 any) *StreamCancelWriteCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CancelWrite", reflect.TypeOf((*MockStream)(nil).CancelWrite), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CancelWrite", reflect.TypeOf((*MockStream)(nil).CancelWrite), arg0)
+	return &StreamCancelWriteCall{Call: call}
+}
+
+// StreamCancelWriteCall wrap *gomock.Call
+type StreamCancelWriteCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamCancelWriteCall) Return() *StreamCancelWriteCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamCancelWriteCall) Do(f func(qerr.StreamErrorCode)) *StreamCancelWriteCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamCancelWriteCall) DoAndReturn(f func(qerr.StreamErrorCode)) *StreamCancelWriteCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Close mocks base method.
@ -74,9 +122,33 @@ func (m *MockStream) Close() error {
 }

 // Close indicates an expected call of Close.
-func (mr *MockStreamMockRecorder) Close() *gomock.Call {
+func (mr *MockStreamMockRecorder) Close() *StreamCloseCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockStream)(nil).Close))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockStream)(nil).Close))
+	return &StreamCloseCall{Call: call}
+}
+
+// StreamCloseCall wrap *gomock.Call
+type StreamCloseCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamCloseCall) Return(arg0 error) *StreamCloseCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamCloseCall) Do(f func() error) *StreamCloseCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamCloseCall) DoAndReturn(f func() error) *StreamCloseCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Context mocks base method.
@ -88,9 +160,33 @@ func (m *MockStream) Context() context.Context {
 }

 // Context indicates an expected call of Context.
-func (mr *MockStreamMockRecorder) Context() *gomock.Call {
+func (mr *MockStreamMockRecorder) Context() *StreamContextCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Context", reflect.TypeOf((*MockStream)(nil).Context))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Context", reflect.TypeOf((*MockStream)(nil).Context))
+	return &StreamContextCall{Call: call}
+}
+
+// StreamContextCall wrap *gomock.Call
+type StreamContextCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamContextCall) Return(arg0 context.Context) *StreamContextCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamContextCall) Do(f func() context.Context) *StreamContextCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamContextCall) DoAndReturn(f func() context.Context) *StreamContextCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Read mocks base method.
@ -103,9 +199,33 @@ func (m *MockStream) Read(arg0 []byte) (int, error) {
 }

 // Read indicates an expected call of Read.
-func (mr *MockStreamMockRecorder) Read(arg0 any) *gomock.Call {
+func (mr *MockStreamMockRecorder) Read(arg0 any) *StreamReadCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Read", reflect.TypeOf((*MockStream)(nil).Read), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Read", reflect.TypeOf((*MockStream)(nil).Read), arg0)
+	return &StreamReadCall{Call: call}
+}
+
+// StreamReadCall wrap *gomock.Call
+type StreamReadCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamReadCall) Return(arg0 int, arg1 error) *StreamReadCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamReadCall) Do(f func([]byte) (int, error)) *StreamReadCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamReadCall) DoAndReturn(f func([]byte) (int, error)) *StreamReadCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SetDeadline mocks base method.
@ -117,9 +237,33 @@ func (m *MockStream) SetDeadline(arg0 time.Time) error {
 }

 // SetDeadline indicates an expected call of SetDeadline.
-func (mr *MockStreamMockRecorder) SetDeadline(arg0 any) *gomock.Call {
+func (mr *MockStreamMockRecorder) SetDeadline(arg0 any) *StreamSetDeadlineCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetDeadline", reflect.TypeOf((*MockStream)(nil).SetDeadline), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetDeadline", reflect.TypeOf((*MockStream)(nil).SetDeadline), arg0)
+	return &StreamSetDeadlineCall{Call: call}
+}
+
+// StreamSetDeadlineCall wrap *gomock.Call
+type StreamSetDeadlineCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamSetDeadlineCall) Return(arg0 error) *StreamSetDeadlineCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamSetDeadlineCall) Do(f func(time.Time) error) *StreamSetDeadlineCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamSetDeadlineCall) DoAndReturn(f func(time.Time) error) *StreamSetDeadlineCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SetReadDeadline mocks base method.
@ -131,9 +275,33 @@ func (m *MockStream) SetReadDeadline(arg0 time.Time) error {
 }

 // SetReadDeadline indicates an expected call of SetReadDeadline.
-func (mr *MockStreamMockRecorder) SetReadDeadline(arg0 any) *gomock.Call {
+func (mr *MockStreamMockRecorder) SetReadDeadline(arg0 any) *StreamSetReadDeadlineCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetReadDeadline", reflect.TypeOf((*MockStream)(nil).SetReadDeadline), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetReadDeadline", reflect.TypeOf((*MockStream)(nil).SetReadDeadline), arg0)
+	return &StreamSetReadDeadlineCall{Call: call}
+}
+
+// StreamSetReadDeadlineCall wrap *gomock.Call
+type StreamSetReadDeadlineCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamSetReadDeadlineCall) Return(arg0 error) *StreamSetReadDeadlineCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamSetReadDeadlineCall) Do(f func(time.Time) error) *StreamSetReadDeadlineCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamSetReadDeadlineCall) DoAndReturn(f func(time.Time) error) *StreamSetReadDeadlineCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SetWriteDeadline mocks base method.
@ -145,9 +313,33 @@ func (m *MockStream) SetWriteDeadline(arg0 time.Time) error {
 }

 // SetWriteDeadline indicates an expected call of SetWriteDeadline.
-func (mr *MockStreamMockRecorder) SetWriteDeadline(arg0 any) *gomock.Call {
+func (mr *MockStreamMockRecorder) SetWriteDeadline(arg0 any) *StreamSetWriteDeadlineCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetWriteDeadline", reflect.TypeOf((*MockStream)(nil).SetWriteDeadline), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetWriteDeadline", reflect.TypeOf((*MockStream)(nil).SetWriteDeadline), arg0)
+	return &StreamSetWriteDeadlineCall{Call: call}
+}
+
+// StreamSetWriteDeadlineCall wrap *gomock.Call
+type StreamSetWriteDeadlineCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamSetWriteDeadlineCall) Return(arg0 error) *StreamSetWriteDeadlineCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamSetWriteDeadlineCall) Do(f func(time.Time) error) *StreamSetWriteDeadlineCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamSetWriteDeadlineCall) DoAndReturn(f func(time.Time) error) *StreamSetWriteDeadlineCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // StreamID mocks base method.
@ -159,9 +351,33 @@ func (m *MockStream) StreamID() protocol.StreamID {
 }

 // StreamID indicates an expected call of StreamID.
-func (mr *MockStreamMockRecorder) StreamID() *gomock.Call {
+func (mr *MockStreamMockRecorder) StreamID() *StreamStreamIDCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StreamID", reflect.TypeOf((*MockStream)(nil).StreamID))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StreamID", reflect.TypeOf((*MockStream)(nil).StreamID))
+	return &StreamStreamIDCall{Call: call}
+}
+
+// StreamStreamIDCall wrap *gomock.Call
+type StreamStreamIDCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamStreamIDCall) Return(arg0 protocol.StreamID) *StreamStreamIDCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamStreamIDCall) Do(f func() protocol.StreamID) *StreamStreamIDCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamStreamIDCall) DoAndReturn(f func() protocol.StreamID) *StreamStreamIDCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Write mocks base method.
@ -174,7 +390,31 @@ func (m *MockStream) Write(arg0 []byte) (int, error) {
 }

 // Write indicates an expected call of Write.
-func (mr *MockStreamMockRecorder) Write(arg0 any) *gomock.Call {
+func (mr *MockStreamMockRecorder) Write(arg0 any) *StreamWriteCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Write", reflect.TypeOf((*MockStream)(nil).Write), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Write", reflect.TypeOf((*MockStream)(nil).Write), arg0)
+	return &StreamWriteCall{Call: call}
+}
+
+// StreamWriteCall wrap *gomock.Call
+type StreamWriteCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamWriteCall) Return(arg0 int, arg1 error) *StreamWriteCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamWriteCall) Do(f func([]byte) (int, error)) *StreamWriteCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamWriteCall) DoAndReturn(f func([]byte) (int, error)) *StreamWriteCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/short_header_opener.go
+++ b/internal/mocks/short_header_opener.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mocks -destination short_header_opener.go github.com/quic-go/quic-go/internal/handshake ShortHeaderOpener
+//	mockgen -typed -build_flags=-tags=gomock -package mocks -destination short_header_opener.go github.com/quic-go/quic-go/internal/handshake ShortHeaderOpener
 //
 // Package mocks is a generated GoMock package.
 package mocks
@ -48,9 +48,33 @@ func (m *MockShortHeaderOpener) DecodePacketNumber(arg0 protocol.PacketNumber, a
 }

 // DecodePacketNumber indicates an expected call of DecodePacketNumber.
-func (mr *MockShortHeaderOpenerMockRecorder) DecodePacketNumber(arg0, arg1 any) *gomock.Call {
+func (mr *MockShortHeaderOpenerMockRecorder) DecodePacketNumber(arg0, arg1 any) *ShortHeaderOpenerDecodePacketNumberCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DecodePacketNumber", reflect.TypeOf((*MockShortHeaderOpener)(nil).DecodePacketNumber), arg0, arg1)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DecodePacketNumber", reflect.TypeOf((*MockShortHeaderOpener)(nil).DecodePacketNumber), arg0, arg1)
+	return &ShortHeaderOpenerDecodePacketNumberCall{Call: call}
+}
+
+// ShortHeaderOpenerDecodePacketNumberCall wrap *gomock.Call
+type ShortHeaderOpenerDecodePacketNumberCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ShortHeaderOpenerDecodePacketNumberCall) Return(arg0 protocol.PacketNumber) *ShortHeaderOpenerDecodePacketNumberCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ShortHeaderOpenerDecodePacketNumberCall) Do(f func(protocol.PacketNumber, protocol.PacketNumberLen) protocol.PacketNumber) *ShortHeaderOpenerDecodePacketNumberCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ShortHeaderOpenerDecodePacketNumberCall) DoAndReturn(f func(protocol.PacketNumber, protocol.PacketNumberLen) protocol.PacketNumber) *ShortHeaderOpenerDecodePacketNumberCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // DecryptHeader mocks base method.
@ -60,9 +84,33 @@ func (m *MockShortHeaderOpener) DecryptHeader(arg0 []byte, arg1 *byte, arg2 []by
 }

 // DecryptHeader indicates an expected call of DecryptHeader.
-func (mr *MockShortHeaderOpenerMockRecorder) DecryptHeader(arg0, arg1, arg2 any) *gomock.Call {
+func (mr *MockShortHeaderOpenerMockRecorder) DecryptHeader(arg0, arg1, arg2 any) *ShortHeaderOpenerDecryptHeaderCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DecryptHeader", reflect.TypeOf((*MockShortHeaderOpener)(nil).DecryptHeader), arg0, arg1, arg2)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DecryptHeader", reflect.TypeOf((*MockShortHeaderOpener)(nil).DecryptHeader), arg0, arg1, arg2)
+	return &ShortHeaderOpenerDecryptHeaderCall{Call: call}
+}
+
+// ShortHeaderOpenerDecryptHeaderCall wrap *gomock.Call
+type ShortHeaderOpenerDecryptHeaderCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ShortHeaderOpenerDecryptHeaderCall) Return() *ShortHeaderOpenerDecryptHeaderCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ShortHeaderOpenerDecryptHeaderCall) Do(f func([]byte, *byte, []byte)) *ShortHeaderOpenerDecryptHeaderCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ShortHeaderOpenerDecryptHeaderCall) DoAndReturn(f func([]byte, *byte, []byte)) *ShortHeaderOpenerDecryptHeaderCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Open mocks base method.
@ -75,7 +123,31 @@ func (m *MockShortHeaderOpener) Open(arg0, arg1 []byte, arg2 time.Time, arg3 pro
 }

 // Open indicates an expected call of Open.
-func (mr *MockShortHeaderOpenerMockRecorder) Open(arg0, arg1, arg2, arg3, arg4, arg5 any) *gomock.Call {
+func (mr *MockShortHeaderOpenerMockRecorder) Open(arg0, arg1, arg2, arg3, arg4, arg5 any) *ShortHeaderOpenerOpenCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Open", reflect.TypeOf((*MockShortHeaderOpener)(nil).Open), arg0, arg1, arg2, arg3, arg4, arg5)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Open", reflect.TypeOf((*MockShortHeaderOpener)(nil).Open), arg0, arg1, arg2, arg3, arg4, arg5)
+	return &ShortHeaderOpenerOpenCall{Call: call}
+}
+
+// ShortHeaderOpenerOpenCall wrap *gomock.Call
+type ShortHeaderOpenerOpenCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ShortHeaderOpenerOpenCall) Return(arg0 []byte, arg1 error) *ShortHeaderOpenerOpenCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ShortHeaderOpenerOpenCall) Do(f func([]byte, []byte, time.Time, protocol.PacketNumber, protocol.KeyPhaseBit, []byte) ([]byte, error)) *ShortHeaderOpenerOpenCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ShortHeaderOpenerOpenCall) DoAndReturn(f func([]byte, []byte, time.Time, protocol.PacketNumber, protocol.KeyPhaseBit, []byte) ([]byte, error)) *ShortHeaderOpenerOpenCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/short_header_sealer.go
+++ b/internal/mocks/short_header_sealer.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mocks -destination short_header_sealer.go github.com/quic-go/quic-go/internal/handshake ShortHeaderSealer
+//	mockgen -typed -build_flags=-tags=gomock -package mocks -destination short_header_sealer.go github.com/quic-go/quic-go/internal/handshake ShortHeaderSealer
 //
 // Package mocks is a generated GoMock package.
 package mocks
@ -45,9 +45,33 @@ func (m *MockShortHeaderSealer) EncryptHeader(arg0 []byte, arg1 *byte, arg2 []by
 }

 // EncryptHeader indicates an expected call of EncryptHeader.
-func (mr *MockShortHeaderSealerMockRecorder) EncryptHeader(arg0, arg1, arg2 any) *gomock.Call {
+func (mr *MockShortHeaderSealerMockRecorder) EncryptHeader(arg0, arg1, arg2 any) *ShortHeaderSealerEncryptHeaderCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EncryptHeader", reflect.TypeOf((*MockShortHeaderSealer)(nil).EncryptHeader), arg0, arg1, arg2)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EncryptHeader", reflect.TypeOf((*MockShortHeaderSealer)(nil).EncryptHeader), arg0, arg1, arg2)
+	return &ShortHeaderSealerEncryptHeaderCall{Call: call}
+}
+
+// ShortHeaderSealerEncryptHeaderCall wrap *gomock.Call
+type ShortHeaderSealerEncryptHeaderCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ShortHeaderSealerEncryptHeaderCall) Return() *ShortHeaderSealerEncryptHeaderCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ShortHeaderSealerEncryptHeaderCall) Do(f func([]byte, *byte, []byte)) *ShortHeaderSealerEncryptHeaderCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ShortHeaderSealerEncryptHeaderCall) DoAndReturn(f func([]byte, *byte, []byte)) *ShortHeaderSealerEncryptHeaderCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // KeyPhase mocks base method.
@ -59,9 +83,33 @@ func (m *MockShortHeaderSealer) KeyPhase() protocol.KeyPhaseBit {
 }

 // KeyPhase indicates an expected call of KeyPhase.
-func (mr *MockShortHeaderSealerMockRecorder) KeyPhase() *gomock.Call {
+func (mr *MockShortHeaderSealerMockRecorder) KeyPhase() *ShortHeaderSealerKeyPhaseCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "KeyPhase", reflect.TypeOf((*MockShortHeaderSealer)(nil).KeyPhase))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "KeyPhase", reflect.TypeOf((*MockShortHeaderSealer)(nil).KeyPhase))
+	return &ShortHeaderSealerKeyPhaseCall{Call: call}
+}
+
+// ShortHeaderSealerKeyPhaseCall wrap *gomock.Call
+type ShortHeaderSealerKeyPhaseCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ShortHeaderSealerKeyPhaseCall) Return(arg0 protocol.KeyPhaseBit) *ShortHeaderSealerKeyPhaseCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ShortHeaderSealerKeyPhaseCall) Do(f func() protocol.KeyPhaseBit) *ShortHeaderSealerKeyPhaseCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ShortHeaderSealerKeyPhaseCall) DoAndReturn(f func() protocol.KeyPhaseBit) *ShortHeaderSealerKeyPhaseCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Overhead mocks base method.
@ -73,9 +121,33 @@ func (m *MockShortHeaderSealer) Overhead() int {
 }

 // Overhead indicates an expected call of Overhead.
-func (mr *MockShortHeaderSealerMockRecorder) Overhead() *gomock.Call {
+func (mr *MockShortHeaderSealerMockRecorder) Overhead() *ShortHeaderSealerOverheadCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Overhead", reflect.TypeOf((*MockShortHeaderSealer)(nil).Overhead))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Overhead", reflect.TypeOf((*MockShortHeaderSealer)(nil).Overhead))
+	return &ShortHeaderSealerOverheadCall{Call: call}
+}
+
+// ShortHeaderSealerOverheadCall wrap *gomock.Call
+type ShortHeaderSealerOverheadCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ShortHeaderSealerOverheadCall) Return(arg0 int) *ShortHeaderSealerOverheadCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ShortHeaderSealerOverheadCall) Do(f func() int) *ShortHeaderSealerOverheadCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ShortHeaderSealerOverheadCall) DoAndReturn(f func() int) *ShortHeaderSealerOverheadCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Seal mocks base method.
@ -87,7 +159,31 @@ func (m *MockShortHeaderSealer) Seal(arg0, arg1 []byte, arg2 protocol.PacketNumb
 }

 // Seal indicates an expected call of Seal.
-func (mr *MockShortHeaderSealerMockRecorder) Seal(arg0, arg1, arg2, arg3 any) *gomock.Call {
+func (mr *MockShortHeaderSealerMockRecorder) Seal(arg0, arg1, arg2, arg3 any) *ShortHeaderSealerSealCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Seal", reflect.TypeOf((*MockShortHeaderSealer)(nil).Seal), arg0, arg1, arg2, arg3)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Seal", reflect.TypeOf((*MockShortHeaderSealer)(nil).Seal), arg0, arg1, arg2, arg3)
+	return &ShortHeaderSealerSealCall{Call: call}
+}
+
+// ShortHeaderSealerSealCall wrap *gomock.Call
+type ShortHeaderSealerSealCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ShortHeaderSealerSealCall) Return(arg0 []byte) *ShortHeaderSealerSealCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ShortHeaderSealerSealCall) Do(f func([]byte, []byte, protocol.PacketNumber, []byte) []byte) *ShortHeaderSealerSealCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ShortHeaderSealerSealCall) DoAndReturn(f func([]byte, []byte, protocol.PacketNumber, []byte) []byte) *ShortHeaderSealerSealCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/stream_flow_controller.go
+++ b/internal/mocks/stream_flow_controller.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mocks -destination stream_flow_controller.go github.com/quic-go/quic-go/internal/flowcontrol StreamFlowController
+//	mockgen -typed -build_flags=-tags=gomock -package mocks -destination stream_flow_controller.go github.com/quic-go/quic-go/internal/flowcontrol StreamFlowController
 //
 // Package mocks is a generated GoMock package.
 package mocks
@ -45,9 +45,33 @@ func (m *MockStreamFlowController) Abandon() {
 }

 // Abandon indicates an expected call of Abandon.
-func (mr *MockStreamFlowControllerMockRecorder) Abandon() *gomock.Call {
+func (mr *MockStreamFlowControllerMockRecorder) Abandon() *StreamFlowControllerAbandonCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Abandon", reflect.TypeOf((*MockStreamFlowController)(nil).Abandon))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Abandon", reflect.TypeOf((*MockStreamFlowController)(nil).Abandon))
+	return &StreamFlowControllerAbandonCall{Call: call}
+}
+
+// StreamFlowControllerAbandonCall wrap *gomock.Call
+type StreamFlowControllerAbandonCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamFlowControllerAbandonCall) Return() *StreamFlowControllerAbandonCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamFlowControllerAbandonCall) Do(f func()) *StreamFlowControllerAbandonCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamFlowControllerAbandonCall) DoAndReturn(f func()) *StreamFlowControllerAbandonCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // AddBytesRead mocks base method.
@ -57,9 +81,33 @@ func (m *MockStreamFlowController) AddBytesRead(arg0 protocol.ByteCount) {
 }

 // AddBytesRead indicates an expected call of AddBytesRead.
-func (mr *MockStreamFlowControllerMockRecorder) AddBytesRead(arg0 any) *gomock.Call {
+func (mr *MockStreamFlowControllerMockRecorder) AddBytesRead(arg0 any) *StreamFlowControllerAddBytesReadCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddBytesRead", reflect.TypeOf((*MockStreamFlowController)(nil).AddBytesRead), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddBytesRead", reflect.TypeOf((*MockStreamFlowController)(nil).AddBytesRead), arg0)
+	return &StreamFlowControllerAddBytesReadCall{Call: call}
+}
+
+// StreamFlowControllerAddBytesReadCall wrap *gomock.Call
+type StreamFlowControllerAddBytesReadCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamFlowControllerAddBytesReadCall) Return() *StreamFlowControllerAddBytesReadCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamFlowControllerAddBytesReadCall) Do(f func(protocol.ByteCount)) *StreamFlowControllerAddBytesReadCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamFlowControllerAddBytesReadCall) DoAndReturn(f func(protocol.ByteCount)) *StreamFlowControllerAddBytesReadCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // AddBytesSent mocks base method.
@ -69,9 +117,33 @@ func (m *MockStreamFlowController) AddBytesSent(arg0 protocol.ByteCount) {
 }

 // AddBytesSent indicates an expected call of AddBytesSent.
-func (mr *MockStreamFlowControllerMockRecorder) AddBytesSent(arg0 any) *gomock.Call {
+func (mr *MockStreamFlowControllerMockRecorder) AddBytesSent(arg0 any) *StreamFlowControllerAddBytesSentCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddBytesSent", reflect.TypeOf((*MockStreamFlowController)(nil).AddBytesSent), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddBytesSent", reflect.TypeOf((*MockStreamFlowController)(nil).AddBytesSent), arg0)
+	return &StreamFlowControllerAddBytesSentCall{Call: call}
+}
+
+// StreamFlowControllerAddBytesSentCall wrap *gomock.Call
+type StreamFlowControllerAddBytesSentCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamFlowControllerAddBytesSentCall) Return() *StreamFlowControllerAddBytesSentCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamFlowControllerAddBytesSentCall) Do(f func(protocol.ByteCount)) *StreamFlowControllerAddBytesSentCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamFlowControllerAddBytesSentCall) DoAndReturn(f func(protocol.ByteCount)) *StreamFlowControllerAddBytesSentCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // GetWindowUpdate mocks base method.
@ -83,9 +155,33 @@ func (m *MockStreamFlowController) GetWindowUpdate() protocol.ByteCount {
 }

 // GetWindowUpdate indicates an expected call of GetWindowUpdate.
-func (mr *MockStreamFlowControllerMockRecorder) GetWindowUpdate() *gomock.Call {
+func (mr *MockStreamFlowControllerMockRecorder) GetWindowUpdate() *StreamFlowControllerGetWindowUpdateCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWindowUpdate", reflect.TypeOf((*MockStreamFlowController)(nil).GetWindowUpdate))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWindowUpdate", reflect.TypeOf((*MockStreamFlowController)(nil).GetWindowUpdate))
+	return &StreamFlowControllerGetWindowUpdateCall{Call: call}
+}
+
+// StreamFlowControllerGetWindowUpdateCall wrap *gomock.Call
+type StreamFlowControllerGetWindowUpdateCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamFlowControllerGetWindowUpdateCall) Return(arg0 protocol.ByteCount) *StreamFlowControllerGetWindowUpdateCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamFlowControllerGetWindowUpdateCall) Do(f func() protocol.ByteCount) *StreamFlowControllerGetWindowUpdateCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamFlowControllerGetWindowUpdateCall) DoAndReturn(f func() protocol.ByteCount) *StreamFlowControllerGetWindowUpdateCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // IsNewlyBlocked mocks base method.
@ -98,9 +194,33 @@ func (m *MockStreamFlowController) IsNewlyBlocked() (bool, protocol.ByteCount) {
 }

 // IsNewlyBlocked indicates an expected call of IsNewlyBlocked.
-func (mr *MockStreamFlowControllerMockRecorder) IsNewlyBlocked() *gomock.Call {
+func (mr *MockStreamFlowControllerMockRecorder) IsNewlyBlocked() *StreamFlowControllerIsNewlyBlockedCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsNewlyBlocked", reflect.TypeOf((*MockStreamFlowController)(nil).IsNewlyBlocked))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsNewlyBlocked", reflect.TypeOf((*MockStreamFlowController)(nil).IsNewlyBlocked))
+	return &StreamFlowControllerIsNewlyBlockedCall{Call: call}
+}
+
+// StreamFlowControllerIsNewlyBlockedCall wrap *gomock.Call
+type StreamFlowControllerIsNewlyBlockedCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamFlowControllerIsNewlyBlockedCall) Return(arg0 bool, arg1 protocol.ByteCount) *StreamFlowControllerIsNewlyBlockedCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamFlowControllerIsNewlyBlockedCall) Do(f func() (bool, protocol.ByteCount)) *StreamFlowControllerIsNewlyBlockedCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamFlowControllerIsNewlyBlockedCall) DoAndReturn(f func() (bool, protocol.ByteCount)) *StreamFlowControllerIsNewlyBlockedCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // SendWindowSize mocks base method.
@ -112,9 +232,33 @@ func (m *MockStreamFlowController) SendWindowSize() protocol.ByteCount {
 }

 // SendWindowSize indicates an expected call of SendWindowSize.
-func (mr *MockStreamFlowControllerMockRecorder) SendWindowSize() *gomock.Call {
+func (mr *MockStreamFlowControllerMockRecorder) SendWindowSize() *StreamFlowControllerSendWindowSizeCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendWindowSize", reflect.TypeOf((*MockStreamFlowController)(nil).SendWindowSize))
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendWindowSize", reflect.TypeOf((*MockStreamFlowController)(nil).SendWindowSize))
+	return &StreamFlowControllerSendWindowSizeCall{Call: call}
+}
+
+// StreamFlowControllerSendWindowSizeCall wrap *gomock.Call
+type StreamFlowControllerSendWindowSizeCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamFlowControllerSendWindowSizeCall) Return(arg0 protocol.ByteCount) *StreamFlowControllerSendWindowSizeCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamFlowControllerSendWindowSizeCall) Do(f func() protocol.ByteCount) *StreamFlowControllerSendWindowSizeCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamFlowControllerSendWindowSizeCall) DoAndReturn(f func() protocol.ByteCount) *StreamFlowControllerSendWindowSizeCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // UpdateHighestReceived mocks base method.
@ -126,9 +270,33 @@ func (m *MockStreamFlowController) UpdateHighestReceived(arg0 protocol.ByteCount
 }

 // UpdateHighestReceived indicates an expected call of UpdateHighestReceived.
-func (mr *MockStreamFlowControllerMockRecorder) UpdateHighestReceived(arg0, arg1 any) *gomock.Call {
+func (mr *MockStreamFlowControllerMockRecorder) UpdateHighestReceived(arg0, arg1 any) *StreamFlowControllerUpdateHighestReceivedCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateHighestReceived", reflect.TypeOf((*MockStreamFlowController)(nil).UpdateHighestReceived), arg0, arg1)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateHighestReceived", reflect.TypeOf((*MockStreamFlowController)(nil).UpdateHighestReceived), arg0, arg1)
+	return &StreamFlowControllerUpdateHighestReceivedCall{Call: call}
+}
+
+// StreamFlowControllerUpdateHighestReceivedCall wrap *gomock.Call
+type StreamFlowControllerUpdateHighestReceivedCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamFlowControllerUpdateHighestReceivedCall) Return(arg0 error) *StreamFlowControllerUpdateHighestReceivedCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamFlowControllerUpdateHighestReceivedCall) Do(f func(protocol.ByteCount, bool) error) *StreamFlowControllerUpdateHighestReceivedCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamFlowControllerUpdateHighestReceivedCall) DoAndReturn(f func(protocol.ByteCount, bool) error) *StreamFlowControllerUpdateHighestReceivedCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // UpdateSendWindow mocks base method.
@ -138,7 +306,31 @@ func (m *MockStreamFlowController) UpdateSendWindow(arg0 protocol.ByteCount) {
 }

 // UpdateSendWindow indicates an expected call of UpdateSendWindow.
-func (mr *MockStreamFlowControllerMockRecorder) UpdateSendWindow(arg0 any) *gomock.Call {
+func (mr *MockStreamFlowControllerMockRecorder) UpdateSendWindow(arg0 any) *StreamFlowControllerUpdateSendWindowCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateSendWindow", reflect.TypeOf((*MockStreamFlowController)(nil).UpdateSendWindow), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateSendWindow", reflect.TypeOf((*MockStreamFlowController)(nil).UpdateSendWindow), arg0)
+	return &StreamFlowControllerUpdateSendWindowCall{Call: call}
+}
+
+// StreamFlowControllerUpdateSendWindowCall wrap *gomock.Call
+type StreamFlowControllerUpdateSendWindowCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *StreamFlowControllerUpdateSendWindowCall) Return() *StreamFlowControllerUpdateSendWindowCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *StreamFlowControllerUpdateSendWindowCall) Do(f func(protocol.ByteCount)) *StreamFlowControllerUpdateSendWindowCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *StreamFlowControllerUpdateSendWindowCall) DoAndReturn(f func(protocol.ByteCount)) *StreamFlowControllerUpdateSendWindowCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }
--- a/internal/mocks/tls/client_session_cache.go
+++ b/internal/mocks/tls/client_session_cache.go
@ -3,7 +3,7 @@
 //
 // Generated by this command:
 //
-//	mockgen -build_flags=-tags=gomock -package mocktls -destination tls/client_session_cache.go crypto/tls ClientSessionCache
+//	mockgen -typed -build_flags=-tags=gomock -package mocktls -destination tls/client_session_cache.go crypto/tls ClientSessionCache
 //
 // Package mocktls is a generated GoMock package.
 package mocktls
@ -48,9 +48,33 @@ func (m *MockClientSessionCache) Get(arg0 string) (*tls.ClientSessionState, bool
 }

 // Get indicates an expected call of Get.
-func (mr *MockClientSessionCacheMockRecorder) Get(arg0 any) *gomock.Call {
+func (mr *MockClientSessionCacheMockRecorder) Get(arg0 any) *ClientSessionCacheGetCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockClientSessionCache)(nil).Get), arg0)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockClientSessionCache)(nil).Get), arg0)
+	return &ClientSessionCacheGetCall{Call: call}
+}
+
+// ClientSessionCacheGetCall wrap *gomock.Call
+type ClientSessionCacheGetCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ClientSessionCacheGetCall) Return(arg0 *tls.ClientSessionState, arg1 bool) *ClientSessionCacheGetCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ClientSessionCacheGetCall) Do(f func(string) (*tls.ClientSessionState, bool)) *ClientSessionCacheGetCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ClientSessionCacheGetCall) DoAndReturn(f func(string) (*tls.ClientSessionState, bool)) *ClientSessionCacheGetCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }

 // Put mocks base method.
@ -60,7 +84,31 @@ func (m *MockClientSessionCache) Put(arg0 string, arg1 *tls.ClientSessionState)
 }

 // Put indicates an expected call of Put.
-func (mr *MockClientSessionCacheMockRecorder) Put(arg0, arg1 any) *gomock.Call {
+func (mr *MockClientSessionCacheMockRecorder) Put(arg0, arg1 any) *ClientSessionCachePutCall {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Put", reflect.TypeOf((*MockClientSessionCache)(nil).Put), arg0, arg1)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Put", reflect.TypeOf((*MockClientSessionCache)(nil).Put), arg0, arg1)
+	return &ClientSessionCachePutCall{Call: call}
+}
+
+// ClientSessionCachePutCall wrap *gomock.Call
+type ClientSessionCachePutCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *ClientSessionCachePutCall) Return() *ClientSessionCachePutCall {
+	c.Call = c.Call.Return()
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *ClientSessionCachePutCall) Do(f func(string, *tls.ClientSessionState)) *ClientSessionCachePutCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *ClientSessionCachePutCall) DoAndReturn(f func(string, *tls.ClientSessionState)) *ClientSessionCachePutCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
 }