From aebf2386d02abe307b42245b7c11a9627a448cd2 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 4 May 2020 19:13:24 +0700
Subject: [PATCH] fix buffer use after it was released when sending an
 INVALID_TOKEN error

---
 server.go      | 32 ++++++++++++++------------------
 server_test.go | 11 +++++------
 2 files changed, 19 insertions(+), 24 deletions(-)

diff --git a/server.go b/server.go
index 96b59cbb..69962e72 100644
--- a/server.go
+++ b/server.go
@@ -292,7 +292,7 @@ func (s *baseServer) handlePacket(p *receivedPacket) {
 	}
 }
 
-func (s *baseServer) handlePacketImpl(p *receivedPacket) bool /* was the packet handled */ {
+func (s *baseServer) handlePacketImpl(p *receivedPacket) bool /* should the buffer be released */ {
 	// If we're creating a new session, the packet will be passed to the session.
 	// The header will then be parsed again.
 	hdr, _, _, err := wire.ParsePacket(p.data, s.config.ConnectionIDLength)
@@ -328,24 +328,18 @@ func (s *baseServer) handlePacketImpl(p *receivedPacket) bool /* was the packet
 
 	s.logger.Debugf("<- Received Initial packet.")
 
-	sess, err := s.handleInitialImpl(p, hdr)
-	if err != nil {
+	if err := s.handleInitialImpl(p, hdr); err != nil {
 		s.logger.Errorf("Error occurred handling initial packet: %s", err)
-		return false
 	}
-	// A retry was done, or the connection attempt was rejected,
-	// or if the Initial was a duplicate.
-	if sess == nil {
-		return false
-	}
-	// Don't put the packet buffer back if a new session was created.
-	// The session will handle the packet and take of that.
+	// Don't put the packet buffer back.
+	// handleInitialImpl deals with the buffer.
 	return true
 }
 
-func (s *baseServer) handleInitialImpl(p *receivedPacket, hdr *wire.Header) (quicSession, error) {
+func (s *baseServer) handleInitialImpl(p *receivedPacket, hdr *wire.Header) error {
 	if len(hdr.Token) == 0 && hdr.DestConnectionID.Len() < protocol.MinConnectionIDLenInitial {
-		return nil, errors.New("too short connection ID")
+		p.buffer.Release()
+		return errors.New("too short connection ID")
 	}
 
 	var token *Token
@@ -363,6 +357,7 @@ func (s *baseServer) handleInitialImpl(p *receivedPacket, hdr *wire.Header) (qui
 	}
 	if !s.config.AcceptToken(p.remoteAddr, token) {
 		go func() {
+			defer p.buffer.Release()
 			if token != nil && token.IsRetryToken {
 				if err := s.maybeSendInvalidToken(p, hdr); err != nil {
 					s.logger.Debugf("Error sending INVALID_TOKEN error: %s", err)
@@ -373,7 +368,7 @@ func (s *baseServer) handleInitialImpl(p *receivedPacket, hdr *wire.Header) (qui
 				s.logger.Debugf("Error sending Retry: %s", err)
 			}
 		}()
-		return nil, nil
+		return nil
 	}
 
 	if queueLen := atomic.LoadInt32(&s.sessionQueueLen); queueLen >= protocol.MaxAcceptQueueSize {
@@ -383,12 +378,12 @@ func (s *baseServer) handleInitialImpl(p *receivedPacket, hdr *wire.Header) (qui
 				s.logger.Debugf("Error rejecting connection: %s", err)
 			}
 		}()
-		return nil, nil
+		return nil
 	}
 
 	connID, err := protocol.GenerateConnectionID(s.config.ConnectionIDLength)
 	if err != nil {
-		return nil, err
+		return err
 	}
 	s.logger.Debugf("Changing connection ID to %s.", connID)
 	sess := s.createNewSession(
@@ -400,7 +395,8 @@ func (s *baseServer) handleInitialImpl(p *receivedPacket, hdr *wire.Header) (qui
 		hdr.Version,
 	)
 	if sess == nil {
-		return nil, nil
+		p.buffer.Release()
+		return nil
 	}
 	sess.handlePacket(p)
 	for {
@@ -410,7 +406,7 @@ func (s *baseServer) handleInitialImpl(p *receivedPacket, hdr *wire.Header) (qui
 		}
 		sess.handlePacket(p)
 	}
-	return sess, nil
+	return nil
 }
 
 func (s *baseServer) createNewSession(
diff --git a/server_test.go b/server_test.go
index 4120fa02..9803a638 100644
--- a/server_test.go
+++ b/server_test.go
@@ -14,18 +14,17 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/lucas-clemente/quic-go/internal/qerr"
-
-	"github.com/lucas-clemente/quic-go/qlog"
-
-	"github.com/golang/mock/gomock"
 	"github.com/lucas-clemente/quic-go/internal/handshake"
 	"github.com/lucas-clemente/quic-go/internal/protocol"
+	"github.com/lucas-clemente/quic-go/internal/qerr"
 	"github.com/lucas-clemente/quic-go/internal/testdata"
 	"github.com/lucas-clemente/quic-go/internal/utils"
 	"github.com/lucas-clemente/quic-go/internal/wire"
+	"github.com/lucas-clemente/quic-go/qlog"
 	"github.com/lucas-clemente/quic-go/quictrace"
 
+	"github.com/golang/mock/gomock"
+
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 )
@@ -580,7 +579,7 @@ var _ = Describe("Server", func() {
 				p := getInitial(protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8, 9})
 				phm.EXPECT().GetStatelessResetToken(gomock.Any())
 				phm.EXPECT().Add(protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8, 9}, sess).Return(false)
-				Expect(serv.handlePacketImpl(p)).To(BeFalse())
+				Expect(serv.handlePacketImpl(p)).To(BeTrue())
 				Expect(createdSession).To(BeTrue())
 			})