close streams when Session.Close() is called

session.go

@@ -88,7 +88,7 @@ func (s *Session) Run() {
 			switch err {
 			// Can happen e.g. when packets thought missing arrive late
 			case ackhandler.ErrDuplicateOrOutOfOrderAck:
-				// Can happen when RST_STREAMs arrive early or late (?)
+			// Can happen when RST_STREAMs arrive early or late (?)
 			case errRstStreamOnInvalidStream:
 				fmt.Printf("Ignoring error in session: %s\n", err.Error())
 			default:
@@ -205,6 +205,7 @@ func (s *Session) Close(e error) error {
 	if ok {
 		errorCode = quicError.ErrorCode
 	}
+	s.closeStreamsWithError(e)
 	// TODO: Don't queue, but send immediately
 	return s.QueueFrame(&frames.ConnectionCloseFrame{
 		ErrorCode:    errorCode,
@@ -268,6 +269,9 @@ func (s *Session) garbageCollectStreams() {
 	s.streamsMutex.Lock()
 	defer s.streamsMutex.Unlock()
 	for k, v := range s.streams {
+		if v == nil {
+			continue
+		}
 		// Strictly speaking, this is not thread-safe. However it doesn't matter
 		// if the stream is deleted just shortly later, so we don't care.
 		if v.finishedReading() {

session_test.go

@@ -3,11 +3,16 @@ package quic
 import (
 	"errors"
 	"io"
+	"os"
+	"runtime"
+	"time"
 
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 
+	"github.com/lucas-clemente/quic-go/crypto"
 	"github.com/lucas-clemente/quic-go/frames"
+	"github.com/lucas-clemente/quic-go/handshake"
 	"github.com/lucas-clemente/quic-go/protocol"
 	"github.com/lucas-clemente/quic-go/utils"
 )
@@ -155,4 +160,27 @@ var _ = Describe("Session", func() {
 			Expect(err).To(MatchError("RST_STREAM received for unknown stream"))
 		})
 	})
+
+	Context("closing", func() {
+		var (
+			nGoRoutinesBefore int
+		)
+
+		BeforeEach(func() {
+			nGoRoutinesBefore = runtime.NumGoroutine()
+			path := os.Getenv("GOPATH") + "/src/github.com/lucas-clemente/quic-go/example/"
+			signer, err := crypto.NewRSASigner(path+"cert.der", path+"key.der")
+			Expect(err).ToNot(HaveOccurred())
+			scfg := handshake.NewServerConfig(crypto.NewCurve25519KEX(), signer)
+			session = NewSession(nil, 0, 0, scfg, nil).(*Session)
+		})
+
+		It("shuts down without error", func() {
+			// crypto stream is running in separate go routine
+			Expect(runtime.NumGoroutine()).To(Equal(nGoRoutinesBefore + 1))
+			session.Close(nil)
+			time.Sleep(1 * time.Millisecond)
+			Expect(runtime.NumGoroutine()).To(Equal(nGoRoutinesBefore))
+		})
+	})
 })