put back packet buffers after processing a packet

This introduces a reference counter in the packet buffer, which will be
used to process coalesced packets.

buffer_pool.go

@@ -6,22 +6,42 @@ import (
 	"github.com/lucas-clemente/quic-go/internal/protocol"
 )
 
+type packetBuffer struct {
+	Slice []byte
+
+	// refCount counts how many packets the Slice is used in.
+	// It doesn't support concurrent use.
+	// It is > 1 when used for coalesced packet.
+	refCount int
+}
+
 var bufferPool sync.Pool
 
-func getPacketBuffer() *[]byte {
-	return bufferPool.Get().(*[]byte)
+func getPacketBuffer() *packetBuffer {
+	buf := bufferPool.Get().(*packetBuffer)
+	buf.refCount = 1
+	buf.Slice = buf.Slice[:protocol.MaxReceivePacketSize]
+	return buf
 }
 
-func putPacketBuffer(buf *[]byte) {
-	if cap(*buf) != int(protocol.MaxReceivePacketSize) {
+func putPacketBuffer(buf *packetBuffer) {
+	if cap(buf.Slice) != int(protocol.MaxReceivePacketSize) {
 		panic("putPacketBuffer called with packet of wrong size!")
 	}
+	buf.refCount--
+	if buf.refCount < 0 {
+		panic("negative packetBuffer refCount")
+	}
+	// only put the packetBuffer back if it's not used any more
+	if buf.refCount == 0 {
 		bufferPool.Put(buf)
 	}
+}
 
 func init() {
 	bufferPool.New = func() interface{} {
-		b := make([]byte, 0, protocol.MaxReceivePacketSize)
-		return &b
+		return &packetBuffer{
+			Slice: make([]byte, 0, protocol.MaxReceivePacketSize),
+		}
 	}
 }

buffer_pool_test.go

@@ -9,13 +9,24 @@ import (
 
 var _ = Describe("Buffer Pool", func() {
 	It("returns buffers of cap", func() {
-		buf := *getPacketBuffer()
-		Expect(buf).To(HaveCap(int(protocol.MaxReceivePacketSize)))
+		buf := getPacketBuffer()
+		Expect(buf.Slice).To(HaveCap(int(protocol.MaxReceivePacketSize)))
+	})
+
+	It("puts buffers back", func() {
+		buf := getPacketBuffer()
+		putPacketBuffer(buf)
 	})
 
 	It("panics if wrong-sized buffers are passed", func() {
-		Expect(func() {
-			putPacketBuffer(&[]byte{0})
-		}).To(Panic())
+		buf := getPacketBuffer()
+		buf.Slice = make([]byte, 10)
+		Expect(func() { putPacketBuffer(buf) }).To(Panic())
+	})
+
+	It("panics if it is put pack twice", func() {
+		buf := getPacketBuffer()
+		putPacketBuffer(buf)
+		Expect(func() { putPacketBuffer(buf) }).To(Panic())
 	})
 })

packet_handler_map.go

@@ -144,8 +144,8 @@ func (h *packetHandlerMap) close(e error) error {
 
 func (h *packetHandlerMap) listen() {
 	for {
-		data := *getPacketBuffer()
-		data = data[:protocol.MaxReceivePacketSize]
+		buffer := getPacketBuffer()
+		data := buffer.Slice
 		// The packet size should not exceed protocol.MaxReceivePacketSize bytes
 		// If it does, we only read a truncated packet, which will then end up undecryptable
 		n, addr, err := h.conn.ReadFrom(data)
@@ -155,13 +155,17 @@ func (h *packetHandlerMap) listen() {
 		}
 		data = data[:n]
 
-		if err := h.handlePacket(addr, data); err != nil {
+		if err := h.handlePacket(addr, buffer, data); err != nil {
 			h.logger.Debugf("error handling packet from %s: %s", addr, err)
 		}
 	}
 }
 
-func (h *packetHandlerMap) handlePacket(addr net.Addr, data []byte) error {
+func (h *packetHandlerMap) handlePacket(
+	addr net.Addr,
+	buffer *packetBuffer,
+	data []byte,
+) error {
 	r := bytes.NewReader(data)
 	hdr, err := wire.ParseHeader(r, h.connIDLen)
 	// drop the packet if we can't parse the header
@@ -172,8 +176,9 @@ func (h *packetHandlerMap) handlePacket(addr net.Addr, data []byte) error {
 	p := &receivedPacket{
 		remoteAddr: addr,
 		hdr:        hdr,
-		data:       data,
 		rcvTime:    time.Now(),
+		data:       data,
+		buffer:     buffer,
 	}
 
 	h.mutex.RLock()

packet_handler_map_test.go

@@ -81,7 +81,7 @@ var _ = Describe("Packet Handler Map", func() {
 		})
 
 		It("drops unparseable packets", func() {
-			err := handler.handlePacket(nil, []byte{0, 1, 2, 3})
+			err := handler.handlePacket(nil, nil, []byte{0, 1, 2, 3})
 			Expect(err).To(HaveOccurred())
 			Expect(err.Error()).To(ContainSubstring("error parsing header:"))
 		})
@@ -91,7 +91,7 @@ var _ = Describe("Packet Handler Map", func() {
 			connID := protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8}
 			handler.Add(connID, NewMockPacketHandler(mockCtrl))
 			handler.Remove(connID)
-			Expect(handler.handlePacket(nil, getPacket(connID))).To(MatchError("received a packet with an unexpected connection ID 0x0102030405060708"))
+			Expect(handler.handlePacket(nil, nil, getPacket(connID))).To(MatchError("received a packet with an unexpected connection ID 0x0102030405060708"))
 		})
 
 		It("deletes retired session entries after a wait time", func() {
@@ -100,7 +100,7 @@ var _ = Describe("Packet Handler Map", func() {
 			handler.Add(connID, NewMockPacketHandler(mockCtrl))
 			handler.Retire(connID)
 			time.Sleep(scaleDuration(30 * time.Millisecond))
-			Expect(handler.handlePacket(nil, getPacket(connID))).To(MatchError("received a packet with an unexpected connection ID 0x0102030405060708"))
+			Expect(handler.handlePacket(nil, nil, getPacket(connID))).To(MatchError("received a packet with an unexpected connection ID 0x0102030405060708"))
 		})
 
 		It("passes packets arriving late for closed sessions to that session", func() {
@@ -110,13 +110,13 @@ var _ = Describe("Packet Handler Map", func() {
 			packetHandler.EXPECT().handlePacket(gomock.Any())
 			handler.Add(connID, packetHandler)
 			handler.Retire(connID)
-			err := handler.handlePacket(nil, getPacket(connID))
+			err := handler.handlePacket(nil, nil, getPacket(connID))
 			Expect(err).ToNot(HaveOccurred())
 		})
 
 		It("drops packets for unknown receivers", func() {
 			connID := protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8}
-			err := handler.handlePacket(nil, getPacket(connID))
+			err := handler.handlePacket(nil, nil, getPacket(connID))
 			Expect(err).To(MatchError("received a packet with an unexpected connection ID 0x0102030405060708"))
 		})
 
@@ -171,10 +171,10 @@ var _ = Describe("Packet Handler Map", func() {
 			handler.AddWithResetToken(connID, NewMockPacketHandler(mockCtrl), token)
 			handler.Retire(connID)
 			time.Sleep(scaleDuration(30 * time.Millisecond))
-			Expect(handler.handlePacket(nil, getPacket(connID))).To(MatchError("received a packet with an unexpected connection ID 0xdeadbeef42"))
+			Expect(handler.handlePacket(nil, nil, getPacket(connID))).To(MatchError("received a packet with an unexpected connection ID 0xdeadbeef42"))
 			packet := append([]byte{0x40, 0xde, 0xca, 0xfb, 0xad, 0x99} /* short header packet */, make([]byte, 50)...)
 			packet = append(packet, token[:]...)
-			Expect(handler.handlePacket(nil, packet)).To(MatchError("received a short header packet with an unexpected connection ID 0xdecafbad99"))
+			Expect(handler.handlePacket(nil, nil, packet)).To(MatchError("received a short header packet with an unexpected connection ID 0xdecafbad99"))
 			Expect(handler.resetTokens).To(BeEmpty())
 		})
 	})
@@ -188,7 +188,7 @@ var _ = Describe("Packet Handler Map", func() {
 				Expect(p.hdr.DestConnectionID).To(Equal(connID))
 			})
 			handler.SetServer(server)
-			Expect(handler.handlePacket(nil, p)).To(Succeed())
+			Expect(handler.handlePacket(nil, nil, p)).To(Succeed())
 		})
 
 		It("closes all server sessions", func() {
@@ -209,7 +209,7 @@ var _ = Describe("Packet Handler Map", func() {
 			server := NewMockUnknownPacketHandler(mockCtrl)
 			handler.SetServer(server)
 			handler.CloseServer()
-			Expect(handler.handlePacket(nil, p)).To(MatchError("received a packet with an unexpected connection ID 0x1122334455667788"))
+			Expect(handler.handlePacket(nil, nil, p)).To(MatchError("received a packet with an unexpected connection ID 0x1122334455667788"))
 		})
 	})
 })

packet_packer.go

@@ -28,6 +28,8 @@ type packedPacket struct {
 	header *wire.ExtendedHeader
 	raw    []byte
 	frames []wire.Frame
+
+	buffer *packetBuffer
 }
 
 func (p *packedPacket) EncryptionLevel() protocol.EncryptionLevel {
@@ -374,8 +376,8 @@ func (p *packetPacker) writeAndSealPacket(
 	frames []wire.Frame,
 	sealer handshake.Sealer,
 ) (*packedPacket, error) {
-	raw := *getPacketBuffer()
-	buffer := bytes.NewBuffer(raw[:0])
+	packetBuffer := getPacketBuffer()
+	buffer := bytes.NewBuffer(packetBuffer.Slice[:0])
 
 	addPaddingForInitial := p.perspective == protocol.PerspectiveClient && header.Type == protocol.PacketTypeInitial
 
@@ -436,7 +438,7 @@ func (p *packetPacker) writeAndSealPacket(
 		return nil, fmt.Errorf("PacketPacker BUG: packet too large (%d bytes, allowed %d bytes)", size, p.maxPacketSize)
 	}
 
-	raw = raw[0:buffer.Len()]
+	raw := buffer.Bytes()
 	_ = sealer.Seal(raw[payloadOffset:payloadOffset], raw[payloadOffset:], header.PacketNumber, raw[:payloadOffset])
 	raw = raw[0 : buffer.Len()+sealer.Overhead()]
 
@@ -455,6 +457,7 @@ func (p *packetPacker) writeAndSealPacket(
 		header: header,
 		raw:    raw,
 		frames: frames,
+		buffer: packetBuffer,
 	}, nil
 }
 

server.go

@@ -318,21 +318,27 @@ func (s *server) handlePacket(p *receivedPacket) {
 	}
 	if hdr.Type == protocol.PacketTypeInitial {
 		go s.handleInitial(p)
+		return
 	}
+
+	putPacketBuffer(p.buffer)
 	// TODO(#943): send Stateless Reset
 }
 
 func (s *server) handleInitial(p *receivedPacket) {
-	// TODO: add a check that DestConnID == SrcConnID
 	s.logger.Debugf("<- Received Initial packet.")
 	sess, connID, err := s.handleInitialImpl(p)
 	if err != nil {
+		putPacketBuffer(p.buffer)
 		s.logger.Errorf("Error occurred handling initial packet: %s", err)
 		return
 	}
 	if sess == nil { // a retry was done
+		putPacketBuffer(p.buffer)
 		return
 	}
+	// Don't put the packet buffer back if a new session was created.
+	// The session will handle the packet and take of that.
 	serverSession := newServerSession(sess, s.config, s.logger)
 	s.sessionHandler.Add(connID, serverSession)
 }
@@ -455,6 +461,7 @@ func (s *server) sendRetry(remoteAddr net.Addr, hdr *wire.Header) error {
 }
 
 func (s *server) sendVersionNegotiationPacket(p *receivedPacket) {
+	defer putPacketBuffer(p.buffer)
 	hdr := p.hdr
 	s.logger.Debugf("Client offered version %s, sending Version Negotiation", hdr.Version)
 	data, err := wire.ComposeVersionNegotiation(hdr.SrcConnectionID, hdr.DestConnectionID, s.config.Versions)

server_test.go

@@ -122,19 +122,19 @@ var _ = Describe("Server", func() {
 		}
 
 		It("drops Initial packets with a too short connection ID", func() {
-			serv.handlePacket(&receivedPacket{
+			serv.handlePacket(insertPacketBuffer(&receivedPacket{
 				hdr: &wire.Header{
 					IsLongHeader:     true,
 					Type:             protocol.PacketTypeInitial,
 					DestConnectionID: protocol.ConnectionID{1, 2, 3, 4},
 					Version:          serv.config.Versions[0],
 				},
-			})
+			}))
 			Consistently(conn.dataWritten).ShouldNot(Receive())
 		})
 
 		It("drops too small Initial", func() {
-			serv.handlePacket(&receivedPacket{
+			serv.handlePacket(insertPacketBuffer(&receivedPacket{
 				hdr: &wire.Header{
 					IsLongHeader:     true,
 					Type:             protocol.PacketTypeInitial,
@@ -142,12 +142,12 @@ var _ = Describe("Server", func() {
 					Version:          serv.config.Versions[0],
 				},
 				data: bytes.Repeat([]byte{0}, protocol.MinInitialPacketSize-100),
-			})
+			}))
 			Consistently(conn.dataWritten).ShouldNot(Receive())
 		})
 
 		It("drops packets with a too short connection ID", func() {
-			serv.handlePacket(&receivedPacket{
+			serv.handlePacket(insertPacketBuffer(&receivedPacket{
 				hdr: &wire.Header{
 					IsLongHeader:     true,
 					Type:             protocol.PacketTypeInitial,
@@ -156,19 +156,19 @@ var _ = Describe("Server", func() {
 					Version:          serv.config.Versions[0],
 				},
 				data: bytes.Repeat([]byte{0}, protocol.MinInitialPacketSize),
-			})
+			}))
 			Consistently(conn.dataWritten).ShouldNot(Receive())
 		})
 
 		It("drops non-Initial packets", func() {
 			serv.logger.SetLogLevel(utils.LogLevelDebug)
-			serv.handlePacket(&receivedPacket{
+			serv.handlePacket(insertPacketBuffer(&receivedPacket{
 				hdr: &wire.Header{
 					Type:    protocol.PacketTypeHandshake,
 					Version: serv.config.Versions[0],
 				},
 				data: []byte("invalid"),
-			})
+			}))
 		})
 
 		It("decodes the cookie from the Token field", func() {
@@ -185,7 +185,7 @@ var _ = Describe("Server", func() {
 			}
 			token, err := serv.cookieGenerator.NewToken(raddr, nil)
 			Expect(err).ToNot(HaveOccurred())
-			serv.handlePacket(&receivedPacket{
+			serv.handlePacket(insertPacketBuffer(&receivedPacket{
 				remoteAddr: raddr,
 				hdr: &wire.Header{
 					Type:    protocol.PacketTypeInitial,
@@ -193,7 +193,7 @@ var _ = Describe("Server", func() {
 					Version: serv.config.Versions[0],
 				},
 				data: bytes.Repeat([]byte{0}, protocol.MinInitialPacketSize),
-			})
+			}))
 			Eventually(done).Should(BeClosed())
 		})
 
@@ -209,7 +209,7 @@ var _ = Describe("Server", func() {
 				close(done)
 				return false
 			}
-			serv.handlePacket(&receivedPacket{
+			serv.handlePacket(insertPacketBuffer(&receivedPacket{
 				remoteAddr: raddr,
 				hdr: &wire.Header{
 					Type:    protocol.PacketTypeInitial,
@@ -217,14 +217,14 @@ var _ = Describe("Server", func() {
 					Version: serv.config.Versions[0],
 				},
 				data: bytes.Repeat([]byte{0}, protocol.MinInitialPacketSize),
-			})
+			}))
 			Eventually(done).Should(BeClosed())
 		})
 
 		It("sends a Version Negotiation Packet for unsupported versions", func() {
 			srcConnID := protocol.ConnectionID{1, 2, 3, 4, 5}
 			destConnID := protocol.ConnectionID{1, 2, 3, 4, 5, 6}
-			serv.handlePacket(&receivedPacket{
+			serv.handlePacket(insertPacketBuffer(&receivedPacket{
 				remoteAddr: &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 1337},
 				hdr: &wire.Header{
 					IsLongHeader:     true,
@@ -233,7 +233,7 @@ var _ = Describe("Server", func() {
 					DestConnectionID: destConnID,
 					Version:          0x42,
 				},
-			})
+			}))
 			var write mockPacketConnWrite
 			Eventually(conn.dataWritten).Should(Receive(&write))
 			Expect(write.to.String()).To(Equal("127.0.0.1:1337"))
@@ -253,11 +253,11 @@ var _ = Describe("Server", func() {
 				DestConnectionID: protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
 				Version:          protocol.VersionTLS,
 			}
-			serv.handleInitial(&receivedPacket{
+			serv.handleInitial(insertPacketBuffer(&receivedPacket{
 				remoteAddr: &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 1337},
 				hdr:        hdr,
 				data:       bytes.Repeat([]byte{0}, protocol.MinInitialPacketSize),
-			})
+			}))
 			var write mockPacketConnWrite
 			Eventually(conn.dataWritten).Should(Receive(&write))
 			Expect(write.to.String()).To(Equal("127.0.0.1:1337"))
@@ -308,7 +308,7 @@ var _ = Describe("Server", func() {
 			done := make(chan struct{})
 			go func() {
 				defer GinkgoRecover()
-				serv.handlePacket(p)
+				serv.handlePacket(insertPacketBuffer(p))
 				// the Handshake packet is written by the session
 				Consistently(conn.dataWritten).ShouldNot(Receive())
 				close(done)

session.go

@@ -53,8 +53,10 @@ type cryptoStreamHandler interface {
 type receivedPacket struct {
 	remoteAddr net.Addr
 	hdr        *wire.Header
-	data       []byte
 	rcvTime    time.Time
+	data       []byte
+
+	buffer *packetBuffer
 }
 
 type closeError struct {
@@ -368,9 +370,6 @@ runLoop:
 			if wasProcessed := s.handlePacketImpl(p); !wasProcessed {
 				continue
 			}
-			// This is a bit unclean, but works properly, since the packet always
-			// begins with the public header and we never copy it.
-			// TODO: putPacketBuffer(&p.extHdr.Raw)
 		case <-s.handshakeCompleteChan:
 			s.handleHandshakeComplete()
 		}
@@ -475,6 +474,15 @@ func (s *session) handleHandshakeComplete() {
 }
 
 func (s *session) handlePacketImpl(p *receivedPacket) bool /* was the packet successfully processed */ {
+	var wasQueued bool
+
+	defer func() {
+		// Put back the packet buffer if the packet wasn't queued for later decryption.
+		if !wasQueued {
+			putPacketBuffer(p.buffer)
+		}
+	}()
+
 	// The server can change the source connection ID with the first Handshake packet.
 	// After this, all packets with a different source connection have to be ignored.
 	if s.receivedFirstPacket && p.hdr.IsLongHeader && !p.hdr.SrcConnectionID.Equal(s.destConnID) {
@@ -490,6 +498,7 @@ func (s *session) handlePacketImpl(p *receivedPacket) bool /* was the packet suc
 	// if the decryption failed, this might be a packet sent by an attacker
 	if err != nil {
 		if err == handshake.ErrOpenerNotYetAvailable {
+			wasQueued = true
 			s.tryQueueingUndecryptablePacket(p)
 			return false
 		}
@@ -953,7 +962,7 @@ func (s *session) sendPacket() (bool, error) {
 }
 
 func (s *session) sendPackedPacket(packet *packedPacket) error {
-	defer putPacketBuffer(&packet.raw)
+	defer putPacketBuffer(packet.buffer)
 	s.logPacket(packet)
 	return s.conn.Write(packet.raw)
 }

session_test.go

@@ -61,6 +61,11 @@ func areSessionsRunning() bool {
 	return strings.Contains(b.String(), "quic-go.(*session).run")
 }
 
+func insertPacketBuffer(p *receivedPacket) *receivedPacket {
+	p.buffer = getPacketBuffer()
+	return p
+}
+
 var _ = Describe("Session", func() {
 	var (
 		sess          *session
@@ -496,11 +501,11 @@ var _ = Describe("Session", func() {
 			rph := mockackhandler.NewMockReceivedPacketHandler(mockCtrl)
 			rph.EXPECT().ReceivedPacket(protocol.PacketNumber(0x1337), rcvTime, false)
 			sess.receivedPacketHandler = rph
-			Expect(sess.handlePacketImpl(&receivedPacket{
+			Expect(sess.handlePacketImpl(insertPacketBuffer(&receivedPacket{
 				rcvTime: rcvTime,
 				hdr:     &hdr.Header,
 				data:    getData(hdr),
-			})).To(BeTrue())
+			}))).To(BeTrue())
 		})
 
 		It("closes when handling a packet fails", func() {
@@ -518,7 +523,10 @@ var _ = Describe("Session", func() {
 				close(done)
 			}()
 			sessionRunner.EXPECT().retireConnectionID(gomock.Any())
-			sess.handlePacket(&receivedPacket{hdr: &wire.Header{}, data: getData(&wire.ExtendedHeader{PacketNumberLen: protocol.PacketNumberLen1})})
+			sess.handlePacket(insertPacketBuffer(&receivedPacket{
+				hdr:  &wire.Header{},
+				data: getData(&wire.ExtendedHeader{PacketNumberLen: protocol.PacketNumberLen1}),
+			}))
 			Eventually(done).Should(BeClosed())
 		})
 
@@ -528,18 +536,18 @@ var _ = Describe("Session", func() {
 				PacketNumberLen: protocol.PacketNumberLen1,
 			}
 			unpacker.EXPECT().Unpack(gomock.Any(), gomock.Any()).Return(&unpackedPacket{hdr: hdr}, nil).Times(2)
-			Expect(sess.handlePacketImpl(&receivedPacket{hdr: &hdr.Header, data: getData(hdr)})).To(BeTrue())
-			Expect(sess.handlePacketImpl(&receivedPacket{hdr: &hdr.Header, data: getData(hdr)})).To(BeTrue())
+			Expect(sess.handlePacketImpl(insertPacketBuffer(&receivedPacket{hdr: &hdr.Header, data: getData(hdr)}))).To(BeTrue())
+			Expect(sess.handlePacketImpl(insertPacketBuffer(&receivedPacket{hdr: &hdr.Header, data: getData(hdr)}))).To(BeTrue())
 		})
 
 		It("ignores 0-RTT packets", func() {
-			Expect(sess.handlePacketImpl(&receivedPacket{
+			Expect(sess.handlePacketImpl(insertPacketBuffer(&receivedPacket{
 				hdr: &wire.Header{
 					IsLongHeader:     true,
 					Type:             protocol.PacketType0RTT,
 					DestConnectionID: sess.srcConnID,
 				},
-			})).To(BeFalse())
+			}))).To(BeFalse())
 		})
 
 		It("ignores packets with a different source connection ID", func() {
@@ -552,12 +560,12 @@ var _ = Describe("Session", func() {
 			// Send one packet, which might change the connection ID.
 			// only EXPECT one call to the unpacker
 			unpacker.EXPECT().Unpack(gomock.Any(), gomock.Any()).Return(&unpackedPacket{hdr: &wire.ExtendedHeader{Header: *hdr}}, nil)
-			Expect(sess.handlePacketImpl(&receivedPacket{
+			Expect(sess.handlePacketImpl(insertPacketBuffer(&receivedPacket{
 				hdr:  hdr,
 				data: getData(&wire.ExtendedHeader{PacketNumberLen: protocol.PacketNumberLen1}),
-			})).To(BeTrue())
+			}))).To(BeTrue())
 			// The next packet has to be ignored, since the source connection ID doesn't match.
-			Expect(sess.handlePacketImpl(&receivedPacket{
+			Expect(sess.handlePacketImpl(insertPacketBuffer(&receivedPacket{
 				hdr: &wire.Header{
 					IsLongHeader:     true,
 					DestConnectionID: sess.destConnID,
@@ -565,7 +573,7 @@ var _ = Describe("Session", func() {
 					Length:           1,
 				},
 				data: getData(&wire.ExtendedHeader{PacketNumberLen: protocol.PacketNumberLen1}),
-			})).To(BeFalse())
+			}))).To(BeFalse())
 		})
 
 		Context("updating the remote address", func() {
@@ -574,12 +582,11 @@ var _ = Describe("Session", func() {
 				origAddr := sess.conn.(*mockConnection).remoteAddr
 				remoteIP := &net.IPAddr{IP: net.IPv4(192, 168, 0, 100)}
 				Expect(origAddr).ToNot(Equal(remoteIP))
-				p := receivedPacket{
+				Expect(sess.handlePacketImpl(insertPacketBuffer(&receivedPacket{
 					remoteAddr: remoteIP,
 					hdr:        &wire.Header{},
 					data:       getData(&wire.ExtendedHeader{PacketNumberLen: protocol.PacketNumberLen1}),
-				}
-				Expect(sess.handlePacketImpl(&p)).To(BeTrue())
+				}))).To(BeTrue())
 				Expect(sess.conn.(*mockConnection).remoteAddr).To(Equal(origAddr))
 			})
 		})
@@ -587,10 +594,12 @@ var _ = Describe("Session", func() {
 
 	Context("sending packets", func() {
 		getPacket := func(pn protocol.PacketNumber) *packedPacket {
-			data := *getPacketBuffer()
+			buffer := getPacketBuffer()
+			data := buffer.Slice[:0]
 			data = append(data, []byte("foobar")...)
 			return &packedPacket{
 				raw:    data,
+				buffer: buffer,
 				header: &wire.ExtendedHeader{PacketNumber: pn},
 			}
 		}
@@ -963,7 +972,7 @@ var _ = Describe("Session", func() {
 				defer close(done)
 				return &packedPacket{
 					header: &wire.ExtendedHeader{},
-					raw:    *getPacketBuffer(),
+					buffer: getPacketBuffer(),
 				}, nil
 			}),
 			packer.EXPECT().PackPacket().AnyTimes(),
@@ -1352,7 +1361,7 @@ var _ = Describe("Client Session", func() {
 		}()
 		newConnID := protocol.ConnectionID{1, 3, 3, 7, 1, 3, 3, 7}
 		packer.EXPECT().ChangeDestConnectionID(newConnID)
-		Expect(sess.handlePacketImpl(&receivedPacket{
+		Expect(sess.handlePacketImpl(insertPacketBuffer(&receivedPacket{
 			hdr: &wire.Header{
 				IsLongHeader:     true,
 				Type:             protocol.PacketTypeHandshake,
@@ -1361,7 +1370,7 @@ var _ = Describe("Client Session", func() {
 				Length:           1,
 			},
 			data: []byte{0},
-		})).To(BeTrue())
+		}))).To(BeTrue())
 		// make sure the go routine returns
 		packer.EXPECT().PackConnectionClose(gomock.Any()).Return(&packedPacket{}, nil)
 		sessionRunner.EXPECT().retireConnectionID(gomock.Any())