diff --git a/core/cs/frag.go b/core/cs/frag.go index c9e491d..82275e3 100644 --- a/core/cs/frag.go +++ b/core/cs/frag.go @@ -41,7 +41,7 @@ func (d *defragger) Feed(m udpMessage) *udpMessage { // wtf is this? return nil } - if m.MsgID != d.msgID { + if m.MsgID != d.msgID || m.FragCount != uint8(len(d.frags)) { // new message, clear previous state d.msgID = m.MsgID d.frags = make([]*udpMessage, m.FragCount) diff --git a/core/cs/frag_test.go b/core/cs/frag_test.go index b0c27cb..884a322 100644 --- a/core/cs/frag_test.go +++ b/core/cs/frag_test.go @@ -196,6 +196,50 @@ func Test_defragger_Feed(t *testing.T) { Data: []byte("hello"), }, }, + { + "frag 0 - 1/2", + args{ + udpMessage{ + SessionID: 123, + HostLen: 4, + Host: "test", + Port: 123, + MsgID: 0, + FragID: 0, + FragCount: 2, + DataLen: 5, + Data: []byte("ilove"), + }, + }, + nil, + }, + { + "frag 0 - 2/2", + args{ + udpMessage{ + SessionID: 123, + HostLen: 4, + Host: "test", + Port: 123, + MsgID: 0, + FragID: 1, + FragCount: 2, + DataLen: 6, + Data: []byte("nobody"), + }, + }, + &udpMessage{ + SessionID: 123, + HostLen: 4, + Host: "test", + Port: 123, + MsgID: 0, + FragID: 0, + FragCount: 1, + DataLen: 11, + Data: []byte("ilovenobody"), + }, + }, { "frag 1 - 1/3", args{