Move SentPacketHandler.CheckForError into SentPacket

ackhandler/interfaces.go

@@ -19,9 +19,6 @@ type SentPacketHandler interface {
 
 	GetAlarmTimeout() time.Time
 	OnAlarm()
-
-	// TODO(lclemente): Remove this now that the logic is simpler
-	CheckForError() error
 }
 
 // ReceivedPacketHandler handles ACKs needed to send for incoming packets

ackhandler/sent_packet_handler.go

@@ -94,6 +94,10 @@ func (h *sentPacketHandler) SentPacket(packet *Packet) error {
 		return errPacketNumberNotIncreasing
 	}
 
+	if protocol.PacketNumber(len(h.retransmissionQueue)+h.packetHistory.Len()+1) > protocol.MaxTrackedSentPackets {
+		return ErrTooManyTrackedSentPackets
+	}
+
 	for p := h.lastSentPacketNumber + 1; p < packet.PacketNumber; p++ {
 		h.skippedPackets = append(h.skippedPackets, p)
 
@@ -334,14 +338,6 @@ func (h *sentPacketHandler) SendingAllowed() bool {
 	return !(congestionLimited || maxTrackedLimited)
 }
 
-func (h *sentPacketHandler) CheckForError() error {
-	length := len(h.retransmissionQueue) + h.packetHistory.Len()
-	if protocol.PacketNumber(length) > protocol.MaxTrackedSentPackets {
-		return ErrTooManyTrackedSentPackets
-	}
-	return nil
-}
-
 func (h *sentPacketHandler) retransmitOldestTwoPackets() {
 	if p := h.packetHistory.Front(); p != nil {
 		h.queueRTO(p)

ackhandler/sent_packet_handler_test.go

@@ -207,12 +207,14 @@ var _ = Describe("SentPacketHandler", func() {
 
 	Context("DoS mitigation", func() {
 		It("checks the size of the packet history, for unacked packets", func() {
-			for i := protocol.PacketNumber(1); i < protocol.MaxTrackedSentPackets+10; i++ {
-				packet := Packet{PacketNumber: protocol.PacketNumber(i), Frames: []frames.Frame{&streamFrame}, Length: 1}
+			i := protocol.PacketNumber(1)
+			for ; i <= protocol.MaxTrackedSentPackets; i++ {
+				packet := Packet{PacketNumber: protocol.PacketNumber(i), Length: 1}
 				err := handler.SentPacket(&packet)
 				Expect(err).ToNot(HaveOccurred())
 			}
-			err := handler.CheckForError()
+			packet := Packet{PacketNumber: protocol.PacketNumber(i), Length: 1}
+			err := handler.SentPacket(&packet)
 			Expect(err).To(MatchError(ErrTooManyTrackedSentPackets))
 		})