Fix QUIC defragger

transport/tuic/packet.go

@@ -27,11 +27,16 @@ var udpMessagePool = sync.Pool{
 	},
 }
 
+func allocMessage() *udpMessage {
+	message := udpMessagePool.Get().(*udpMessage)
+	message.referenced = true
+	return message
+}
+
 func releaseMessages(messages []*udpMessage) {
 	for _, message := range messages {
 		if message != nil {
-			*message = udpMessage{}
-			udpMessagePool.Put(message)
+			message.release()
 		}
 	}
 }
@@ -43,9 +48,13 @@ type udpMessage struct {
 	fragmentID    uint8
 	destination   M.Socksaddr
 	data          *buf.Buffer
+	referenced    bool
 }
 
 func (m *udpMessage) release() {
+	if !m.referenced {
+		return
+	}
 	*m = udpMessage{}
 	udpMessagePool.Put(m)
 }
@@ -83,7 +92,7 @@ func fragUDPMessage(message *udpMessage, maxPacketSize int) []*udpMessage {
 	originPacket := message.data.Bytes()
 	udpMTU := maxPacketSize - message.headerSize()
 	for remaining := len(originPacket); remaining > 0; remaining -= udpMTU {
-		fragment := udpMessagePool.Get().(*udpMessage)
+		fragment := allocMessage()
 		*fragment = *message
 		if remaining > udpMTU {
 			fragment.data = buf.As(originPacket[:udpMTU])
@@ -214,7 +223,7 @@ func (c *udpPacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr)
 		c.packetId.Store(0)
 		packetId = 0
 	}
-	message := udpMessagePool.Get().(*udpMessage)
+	message := allocMessage()
 	*message = udpMessage{
 		sessionID:     c.sessionID,
 		packetID:      uint16(packetId),
@@ -259,7 +268,7 @@ func (c *udpPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 		c.packetId.Store(0)
 		packetId = 0
 	}
-	message := udpMessagePool.Get().(*udpMessage)
+	message := allocMessage()
 	*message = udpMessage{
 		sessionID:     c.sessionID,
 		packetID:      uint16(packetId),
@@ -431,7 +440,7 @@ func (d *udpDefragger) feed(m *udpMessage) *udpMessage {
 	if int(item.count) != len(item.messages) {
 		return nil
 	}
-	newMessage := udpMessagePool.Get().(*udpMessage)
+	newMessage := allocMessage()
 	*newMessage = *item.messages[0]
 	var dataLength uint16
 	for _, message := range item.messages {
@@ -446,6 +455,7 @@ func (d *udpDefragger) feed(m *udpMessage) *udpMessage {
 		item.messages = nil
 		return newMessage
 	}
+	item.messages = nil
 	return nil
 }