From fc54aec7e49a8e9492f0920138e3fdaea0a33a5a Mon Sep 17 00:00:00 2001
From: cnderrauber <zengjie9004@gmail.com>
Date: Fri, 28 Jun 2024 10:08:58 +0800
Subject: [PATCH] Fix gap block in edge condition

Fix gap block in edge condition
---
 receive_payload_queue.go      | 24 ++++++++++++++----------
 receive_payload_queue_test.go | 23 +++++++++++++----------
 2 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/receive_payload_queue.go b/receive_payload_queue.go
index 6f5416eb..2a4f2bbf 100644
--- a/receive_payload_queue.go
+++ b/receive_payload_queue.go
@@ -127,21 +127,25 @@ func (q *receivePayloadQueue) getGapAckBlocks() (gapAckBlocks []gapAckBlock) {
 			if zeroBit, ok := getFirstZeroBit(q.tsnBitmask[index], offset, 64); ok {
 				b.end = uint16(tsn + uint32(zeroBit-offset) - 1 - q.cumulativeTSN)
 				tsn += uint32(zeroBit - offset)
-				gapAckBlocks = append(gapAckBlocks, gapAckBlock{
-					start: b.start,
-					end:   b.end,
-				})
-				findEnd = false
-			} else {
-				tsn += uint32(64 - offset)
-				if tsn > endTSN {
-					b.end = uint16(endTSN - q.cumulativeTSN)
+				if sna32LTE(tsn, endTSN) {
 					gapAckBlocks = append(gapAckBlocks, gapAckBlock{
 						start: b.start,
 						end:   b.end,
 					})
-					break
 				}
+				findEnd = false
+			} else {
+				tsn += uint32(64 - offset)
+			}
+
+			// no zero bit at the end, close and append the last gap
+			if sna32GT(tsn, endTSN) {
+				b.end = uint16(endTSN - q.cumulativeTSN)
+				gapAckBlocks = append(gapAckBlocks, gapAckBlock{
+					start: b.start,
+					end:   b.end,
+				})
+				break
 			}
 		}
 	}
diff --git a/receive_payload_queue_test.go b/receive_payload_queue_test.go
index 0000a122..43056715 100644
--- a/receive_payload_queue_test.go
+++ b/receive_payload_queue_test.go
@@ -26,22 +26,25 @@ func TestReceivePayloadQueue(t *testing.T) {
 	assert.Zero(t, q.size())
 	assert.Empty(t, q.getGapAckBlocks())
 
+	nextTSN := initTSN + maxOffset - 1
+	assert.True(t, q.push(nextTSN))
+	assert.Equal(t, 1, q.size())
+	lastTSN, ok := q.getLastTSNReceived()
+	assert.True(t, lastTSN == nextTSN && ok, "lastTSN:%d, ok:%t", lastTSN, ok)
+	assert.True(t, q.hasChunk(nextTSN))
+
 	assert.True(t, q.push(initTSN))
 	assert.False(t, q.canPush(initTSN-1))
-	assert.True(t, q.canPush(initTSN+maxOffset-1))
 	assert.False(t, q.canPush(initTSN+maxOffset))
 	assert.False(t, q.push(initTSN+maxOffset))
-	assert.Equal(t, 1, q.size())
+	assert.True(t, q.canPush(nextTSN-1))
+	assert.Equal(t, 2, q.size())
 
 	gaps := q.getGapAckBlocks()
-	assert.EqualValues(t, []gapAckBlock{{start: uint16(1), end: uint16(1)}}, gaps)
-
-	nextTSN := initTSN + maxOffset - 1
-	assert.True(t, q.push(nextTSN))
-	assert.Equal(t, 2, q.size())
-	lastTSN, ok := q.getLastTSNReceived()
-	assert.True(t, lastTSN == nextTSN && ok, "lastTSN:%d, ok:%t", lastTSN, ok)
-	assert.True(t, q.hasChunk(nextTSN))
+	assert.EqualValues(t, []gapAckBlock{
+		{start: uint16(1), end: uint16(1)},
+		{start: uint16(maxOffset), end: uint16(maxOffset)},
+	}, gaps)
 
 	assert.True(t, q.pop(false))
 	assert.Equal(t, 1, q.size())