Merge pull request #1557 from lucas-clemente/crypto-setup-close

close the TLS crypto setup when the session closes

internal/handshake/crypto_setup_tls.go

@@ -65,8 +65,12 @@ type cryptoSetupTLS struct {
 	messageErrChan chan error
 	// handshakeComplete is closed when the handshake completes
 	handshakeComplete chan<- struct{}
+	// handshakeDone is closed as soon as the go routine running qtls.Handshake() returns
+	handshakeDone chan struct{}
 	// transport parameters are sent on the receivedTransportParams, as soon as they are received
 	receivedTransportParams <-chan TransportParameters
+	// is closed when Close() is called
+	closeChan chan struct{}
 
 	clientHelloWritten     bool
 	clientHelloWrittenChan chan struct{}
@@ -190,12 +194,14 @@ func newCryptoSetupTLS(
 		handshakeComplete:      handshakeComplete,
 		logger:                 logger,
 		perspective:            perspective,
+		handshakeDone:          make(chan struct{}),
 		handshakeErrChan:       make(chan struct{}),
 		messageErrChan:         make(chan error, 1),
 		clientHelloWrittenChan: make(chan struct{}),
 		messageChan:            make(chan []byte, 100),
 		receivedReadKey:        make(chan struct{}),
 		receivedWriteKey:       make(chan struct{}),
+		closeChan:              make(chan struct{}),
 	}
 	var extHandler tlsExtensionHandler
 	switch perspective {
@@ -235,6 +241,7 @@ func (h *cryptoSetupTLS) RunHandshake() error {
 	handshakeErrChan := make(chan error, 1)
 	handshakeComplete := make(chan struct{})
 	go func() {
+		defer close(h.handshakeDone)
 		if err := conn.Handshake(); err != nil {
 			handshakeErrChan <- err
 			return
@@ -243,6 +250,11 @@ func (h *cryptoSetupTLS) RunHandshake() error {
 	}()
 
 	select {
+	case <-h.closeChan:
+		close(h.messageChan)
+		// wait until the Handshake() go routine has returned
+		<-handshakeErrChan
+		return errors.New("Handshake aborted")
 	case <-handshakeComplete: // return when the handshake is done
 		close(h.handshakeComplete)
 		return nil
@@ -261,6 +273,13 @@ func (h *cryptoSetupTLS) RunHandshake() error {
 	}
 }
 
+func (h *cryptoSetupTLS) Close() error {
+	close(h.closeChan)
+	// wait until qtls.Handshake() actually returned
+	<-h.handshakeDone
+	return nil
+}
+
 // handleMessage handles a TLS handshake message.
 // It is called by the crypto streams when a new message is available.
 // It returns if it is done with messages on the same encryption level.

internal/handshake/crypto_setup_tls_test.go

@@ -116,6 +116,34 @@ var _ = Describe("Crypto Setup TLS", func() {
 		Eventually(done).Should(BeClosed())
 	})
 
+	It("returns Handshake() when it is closed", func() {
+		_, sInitialStream, sHandshakeStream := initStreams()
+		server, err := NewCryptoSetupTLSServer(
+			sInitialStream,
+			sHandshakeStream,
+			protocol.ConnectionID{},
+			&TransportParameters{},
+			func(p *TransportParameters) {},
+			make(chan struct{}),
+			testdata.GetTLSConfig(),
+			[]protocol.VersionNumber{protocol.VersionTLS},
+			protocol.VersionTLS,
+			utils.DefaultLogger.WithPrefix("server"),
+			protocol.PerspectiveServer,
+		)
+		Expect(err).ToNot(HaveOccurred())
+
+		done := make(chan struct{})
+		go func() {
+			defer GinkgoRecover()
+			err := server.RunHandshake()
+			Expect(err).To(MatchError("Handshake aborted"))
+			close(done)
+		}()
+		Expect(server.Close()).To(Succeed())
+		Eventually(done).Should(BeClosed())
+	})
+
 	Context("doing the handshake", func() {
 		generateCert := func() tls.Certificate {
 			priv, err := rsa.GenerateKey(rand.Reader, 2048)

internal/handshake/interface.go

@@ -2,6 +2,7 @@ package handshake
 
 import (
 	"crypto/x509"
+	"io"
 
 	"github.com/lucas-clemente/quic-go/internal/protocol"
 	"github.com/marten-seemann/qtls"
@@ -44,6 +45,7 @@ type CryptoSetup interface {
 type CryptoSetupTLS interface {
 	baseCryptoSetup
 
+	io.Closer
 	HandleMessage([]byte, protocol.EncryptionLevel) bool
 	OpenInitial(dst, src []byte, pn protocol.PacketNumber, ad []byte) ([]byte, error)
 	OpenHandshake(dst, src []byte, pn protocol.PacketNumber, ad []byte) ([]byte, error)

session.go

@@ -6,6 +6,7 @@ import (
 	"crypto/tls"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"sync"
 	"time"
@@ -603,6 +604,9 @@ runLoop:
 	}
 	s.logger.Infof("Connection %s closed.", s.srcConnID)
 	s.sessionRunner.removeConnectionID(s.srcConnID)
+	if s.version.UsesTLS() {
+		s.cryptoStreamHandler.(io.Closer).Close()
+	}
 	return closeErr.err
 }
 

session_test.go

@@ -1638,6 +1638,7 @@ var _ = Describe("Client Session", func() {
 		})
 		Expect(err).ToNot(HaveOccurred())
 		// make sure the go routine returns
+		sess.version = protocol.Version39
 		packer.EXPECT().PackConnectionClose(gomock.Any()).Return(&packedPacket{}, nil)
 		sessionRunner.EXPECT().removeConnectionID(gomock.Any())
 		Expect(sess.Close()).To(Succeed())