From 22d4ebcc14a82938fb4d7e69b40e577daedb253a Mon Sep 17 00:00:00 2001
From: Sukun <sukunrt@gmail.com>
Date: Fri, 23 Feb 2024 17:33:57 +0530
Subject: [PATCH] Consider padding when bundling chunks in packets

---
 association.go      |  9 +++++++--
 association_test.go | 16 ++++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/association.go b/association.go
index cf493ef7..6cf086eb 100644
--- a/association.go
+++ b/association.go
@@ -2238,14 +2238,19 @@ func (a *Association) bundleDataChunksIntoPackets(chunks []*chunkPayloadData) []
 		//   single packet.  Furthermore, DATA chunks being retransmitted MAY be
 		//   bundled with new DATA chunks, as long as the resulting packet size
 		//   does not exceed the path MTU.
-		if bytesInPacket+len(c.userData) > int(a.MTU()) {
+
+		chunkSizeInPacket := int(dataChunkHeaderSize) + len(c.userData)
+		if chunkSizeInPacket%paddingMultiple != 0 {
+			chunkSizeInPacket += paddingMultiple - chunkSizeInPacket%paddingMultiple
+		}
+		if bytesInPacket+chunkSizeInPacket > int(a.MTU()) {
 			packets = append(packets, a.createPacket(chunksToSend))
 			chunksToSend = []chunk{}
 			bytesInPacket = int(commonHeaderSize)
 		}
 
 		chunksToSend = append(chunksToSend, c)
-		bytesInPacket += int(dataChunkHeaderSize) + len(c.userData)
+		bytesInPacket += chunkSizeInPacket
 	}
 
 	if len(chunksToSend) > 0 {
diff --git a/association_test.go b/association_test.go
index 2b6fb733..a12503d7 100644
--- a/association_test.go
+++ b/association_test.go
@@ -3156,3 +3156,19 @@ func TestAssociation_ZeroChecksum(t *testing.T) {
 		require.NoError(t, a2.Close())
 	}
 }
+
+func TestDataChunkBundlingIntoPacket(t *testing.T) {
+	a := &Association{mtu: initialMTU}
+	chunks := make([]*chunkPayloadData, 300)
+	for i := 0; i < 300; i++ {
+		chunks[i] = &chunkPayloadData{userData: []byte{1}}
+	}
+	packets := a.bundleDataChunksIntoPackets(chunks)
+	for _, p := range packets {
+		raw, err := p.marshal(false)
+		require.NoError(t, err)
+		if len(raw) > int(initialMTU) {
+			t.Error("packet too long: ", len(raw))
+		}
+	}
+}