mKCP transport: Make sure ACKs are limited within MTU

transport/internet/kcp/receiving.go

@@ -47,15 +47,18 @@ type AckList struct {
 
 	flushCandidates []uint32
 	dirty           bool
+
+	mss uint32
 }
 
-func NewAckList(writer SegmentWriter) *AckList {
+func NewAckList(writer SegmentWriter, mss uint32) *AckList {
 	return &AckList{
 		writer:          writer,
 		timestamps:      make([]uint32, 0, 128),
 		numbers:         make([]uint32, 0, 128),
 		nextFlush:       make([]uint32, 0, 128),
 		flushCandidates: make([]uint32, 0, 128),
+		mss:             mss,
 	}
 }
 
@@ -90,7 +93,7 @@ func (l *AckList) Clear(una uint32) {
 func (l *AckList) Flush(current uint32, rto uint32) {
 	l.flushCandidates = l.flushCandidates[:0]
 
-	seg := NewAckSegment()
+	seg := NewAckSegment((int(l.mss) - 17) / 4)
 	for i := 0; i < len(l.numbers); i++ {
 		if l.nextFlush[i] > current {
 			if len(l.flushCandidates) < cap(l.flushCandidates) {
@@ -109,7 +112,7 @@ func (l *AckList) Flush(current uint32, rto uint32) {
 		if seg.IsFull() {
 			l.writer.Write(seg)
 			seg.Release()
-			seg = NewAckSegment()
+			seg = NewAckSegment((int(l.mss) - 17) / 4)
 			l.dirty = false
 		}
 	}
@@ -144,7 +147,7 @@ func NewReceivingWorker(kcp *Connection) *ReceivingWorker {
 		window:     NewReceivingWindow(),
 		windowSize: kcp.Config.GetReceivingInFlightSize(),
 	}
-	worker.acklist = NewAckList(worker)
+	worker.acklist = NewAckList(worker, kcp.mss+DataSegmentOverhead)
 	return worker
 }
 

transport/internet/kcp/segment.go

@@ -131,12 +131,22 @@ type AckSegment struct {
 	ReceivingNext   uint32
 	Timestamp       uint32
 	NumberList      []uint32
+
+	Limit int
 }
 
 const ackNumberLimit = 128
 
-func NewAckSegment() *AckSegment {
-	return new(AckSegment)
+func NewAckSegment(limit int) *AckSegment {
+	if limit <= 0 {
+		limit = 1
+	}
+	if limit > ackNumberLimit {
+		limit = ackNumberLimit
+	}
+	return &AckSegment{
+		Limit: limit,
+	}
 }
 
 func (s *AckSegment) parse(conv uint16, cmd Command, opt SegmentOption, buf []byte) (bool, []byte) {
@@ -188,7 +198,7 @@ func (s *AckSegment) PutNumber(number uint32) {
 }
 
 func (s *AckSegment) IsFull() bool {
-	return len(s.NumberList) == ackNumberLimit
+	return len(s.NumberList) == s.Limit
 }
 
 func (s *AckSegment) IsEmpty() bool {
@@ -290,7 +300,7 @@ func ReadSegment(buf []byte) (Segment, []byte) {
 	case CommandData:
 		seg = NewDataSegment()
 	case CommandACK:
-		seg = NewAckSegment()
+		seg = NewAckSegment(128)
 	default:
 		seg = NewCmdOnlySegment()
 	}

transport/internet/kcp/segment_test.go

@@ -71,6 +71,7 @@ func TestACKSegment(t *testing.T) {
 		ReceivingNext:   3,
 		Timestamp:       10,
 		NumberList:      []uint32{1, 3, 5, 7, 9},
+		Limit:           128,
 	}
 
 	nBytes := seg.ByteSize()