fix packet buffer usage when handling coalesced packets

buffer_pool.go

@@ -22,23 +22,41 @@ func (b *packetBuffer) Split() {
 	b.refCount++
 }
 
-// Release decreases the refCount.
-// It should be called when processing the packet is finished.
-// When the refCount reaches 0, the packet buffer is put back into the pool.
-func (b *packetBuffer) Release() {
-	if cap(b.Slice) != int(protocol.MaxReceivePacketSize) {
-		panic("putPacketBuffer called with packet of wrong size!")
-	}
+// Decrement decrements the reference counter.
+// It doesn't put the buffer back into the pool.
+func (b *packetBuffer) Decrement() {
 	b.refCount--
 	if b.refCount < 0 {
 		panic("negative packetBuffer refCount")
 	}
+}
+
+// MaybeRelease puts the packet buffer back into the pool,
+// if the reference counter already reached 0.
+func (b *packetBuffer) MaybeRelease() {
 	// only put the packetBuffer back if it's not used any more
 	if b.refCount == 0 {
-		bufferPool.Put(b)
+		b.putBack()
 	}
 }
 
+// Release puts back the packet buffer into the pool.
+// It should be called when processing is definitely finished.
+func (b *packetBuffer) Release() {
+	b.Decrement()
+	if b.refCount != 0 {
+		panic("packetBuffer refCount not zero")
+	}
+	b.putBack()
+}
+
+func (b *packetBuffer) putBack() {
+	if cap(b.Slice) != int(protocol.MaxReceivePacketSize) {
+		panic("putPacketBuffer called with packet of wrong size!")
+	}
+	bufferPool.Put(b)
+}
+
 var bufferPool sync.Pool
 
 func getPacketBuffer() *packetBuffer {

buffer_pool_test.go

@@ -30,14 +30,20 @@ var _ = Describe("Buffer Pool", func() {
 		Expect(func() { buf.Release() }).To(Panic())
 	})
 
+	It("panics if it is decremented too many times", func() {
+		buf := getPacketBuffer()
+		buf.Decrement()
+		Expect(func() { buf.Decrement() }).To(Panic())
+	})
+
 	It("waits until all parts have been released", func() {
 		buf := getPacketBuffer()
 		buf.Split()
 		buf.Split()
 		// now we have 3 parts
-		buf.Release()
-		buf.Release()
-		buf.Release()
-		Expect(func() { buf.Release() }).To(Panic())
+		buf.Decrement()
+		buf.Decrement()
+		buf.Decrement()
+		Expect(func() { buf.Decrement() }).To(Panic())
 	})
 })

session.go

@@ -515,6 +515,7 @@ func (s *session) handlePacketImpl(p *receivedPacket) bool {
 		}
 		p.data = rest
 	}
+	p.buffer.MaybeRelease()
 	return processed
 }
 
@@ -524,7 +525,7 @@ func (s *session) handleSinglePacket(p *receivedPacket, hdr *wire.Header) bool /
 	defer func() {
 		// Put back the packet buffer if the packet wasn't queued for later decryption.
 		if !wasQueued {
-			p.buffer.Release()
+			p.buffer.Decrement()
 		}
 	}()