implement h2quic.Server.Close()

fixes #155

h2quic/server.go

@@ -30,6 +30,7 @@ type Server struct {
 	CloseAfterFirstRequest bool
 
 	port   int
+	server *quic.Server
 }
 
 // ListenAndServe listens on the UDP address s.Addr and calls s.Handler to handle HTTP/2 requests on incoming connections.
@@ -37,11 +38,15 @@ func (s *Server) ListenAndServe() error {
 	if s.Server == nil {
 		return errors.New("use of h2quic.Server without http.Server")
 	}
-	server, err := quic.NewServer(s.Addr, s.TLSConfig, s.handleStreamCb)
+	if s.server != nil {
+		return errors.New("ListenAndServe may only be called once")
+	}
+	var err error
+	s.server, err = quic.NewServer(s.Addr, s.TLSConfig, s.handleStreamCb)
 	if err != nil {
 		return err
 	}
-	return server.ListenAndServe()
+	return s.server.ListenAndServe()
 }
 
 // ListenAndServeTLS listens on the UDP address s.Addr and calls s.Handler to handle HTTP/2 requests on incoming connections.
@@ -52,18 +57,19 @@ func (s *Server) ListenAndServeTLS(certFile, keyFile string) error {
 	if err != nil {
 		return err
 	}
-
 	// We currently only use the cert-related stuff from tls.Config,
 	// so we don't need to make a full copy.
 	config := &tls.Config{
 		Certificates: certs,
 	}
-	server, err := quic.NewServer(s.Addr, config, s.handleStreamCb)
+	if s.server != nil {
+		return errors.New("ListenAndServe may only be called once")
+	}
+	s.server, err = quic.NewServer(s.Addr, config, s.handleStreamCb)
 	if err != nil {
 		return err
 	}
-
-	return server.ListenAndServe()
+	return s.server.ListenAndServe()
 }
 
 // Serve should not be called, since it only works properly for TCP listeners.
@@ -149,7 +155,9 @@ func (s *Server) handleRequest(session streamCreator, headerStream utils.Stream,
 
 // Close the server immediately, aborting requests and sending CONNECTION_CLOSE frames to connected clients
 func (s *Server) Close() error {
-	// TODO: implement
+	if s.server != nil {
+		return s.server.Close()
+	}
 	return nil
 }
 

server.go

@@ -17,6 +17,7 @@ import (
 type packetHandler interface {
 	handlePacket(addr interface{}, hdr *publicHeader, data []byte)
 	run()
+	Close(error) error
 }
 
 // A Server of QUIC
@@ -89,6 +90,15 @@ func (s *Server) ListenAndServe() error {
 
 // Close the server
 func (s *Server) Close() error {
+	s.sessionsMutex.Lock()
+	for _, session := range s.sessions {
+		if session != nil {
+			s.sessionsMutex.Unlock()
+			_ = session.Close(nil)
+			s.sessionsMutex.Lock()
+		}
+	}
+	s.sessionsMutex.Unlock()
 	if s.conn == nil {
 		return nil
 	}

server_test.go

@@ -16,14 +16,15 @@ import (
 type mockSession struct {
 	connectionID protocol.ConnectionID
 	packetCount  int
+	closed       bool
 }
 
 func (s *mockSession) handlePacket(addr interface{}, hdr *publicHeader, data []byte) {
 	s.packetCount++
 }
 
-func (s *mockSession) run() {
-}
+func (s *mockSession) run()              {}
+func (s *mockSession) Close(error) error { s.closed = true; return nil }
 
 func newMockSession(conn connection, v protocol.VersionNumber, connectionID protocol.ConnectionID, sCfg *handshake.ServerConfig, streamCallback StreamCallback, closeCallback closeCallback) (packetHandler, error) {
 	return &mockSession{
@@ -81,6 +82,13 @@ var _ = Describe("Server", func() {
 			Expect(server.sessions[0x4cfa9f9b668619f6]).To(BeNil())
 		})
 
+		It("closes sessions when Close is called", func() {
+			session := &mockSession{}
+			server.sessions[1] = session
+			err := server.Close()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(session.closed).To(BeTrue())
+		})
 	})
 
 	It("setups and responds with version negotiation", func(done Done) {