implement a Serve method for the Listener

h2quic/server.go

@@ -84,17 +84,20 @@ func (s *Server) serveImpl(tlsConfig *tls.Config, conn *net.UDPConn) error {
 			}
 		},
 	}
-	ln, err := quic.NewListener(&config)
+	var ln quic.Listener
+	var err error
+	if conn == nil {
+		ln, err = quic.ListenAddr(s.Addr, &config)
+	} else {
+		ln, err = quic.Listen(conn, &config)
+	}
 	if err != nil {
 		s.listenerMutex.Unlock()
 		return err
 	}
 	s.listener = ln
 	s.listenerMutex.Unlock()
-	if conn == nil {
-		return ln.ListenAddr(s.Addr)
-	}
-	return ln.Listen(conn)
+	return ln.Serve()
 }
 
 func (s *Server) handleHeaderStream(session streamCreator) {

interface.go

@@ -59,6 +59,5 @@ type Config struct {
 type Listener interface {
 	Close() error
 	Addr() net.Addr
-	ListenAddr(addr string) error
-	Listen(conn net.PacketConn) error
+	Serve() error
 }

server.go

@@ -41,10 +41,22 @@ type server struct {
 
 var _ Listener = &server{}
 
-// NewListener makes a new listener
-func NewListener(config *Config) (Listener, error) {
-	certChain := crypto.NewCertChain(config.TLSConfig)
+// ListenAddr listens for QUIC connections on a given address
+func ListenAddr(addr string, config *Config) (Listener, error) {
+	udpAddr, err := net.ResolveUDPAddr("udp", addr)
+	if err != nil {
+		return nil, err
+	}
+	conn, err := net.ListenUDP("udp", udpAddr)
+	if err != nil {
+		return nil, err
+	}
+	return Listen(conn, config)
+}
 
+// Listen listens for QUIC connections on a given net.PacketConn
+func Listen(conn net.PacketConn, config *Config) (Listener, error) {
+	certChain := crypto.NewCertChain(config.TLSConfig)
 	kex, err := crypto.NewCurve25519KEX()
 	if err != nil {
 		return nil, err
@@ -55,6 +67,7 @@ func NewListener(config *Config) (Listener, error) {
 	}
 
 	return &server{
+		conn:                      conn,
 		config:                    config,
 		certChain:                 certChain,
 		scfg:                      scfg,
@@ -65,15 +78,11 @@ func NewListener(config *Config) (Listener, error) {
 }
 
 // Listen listens on an existing PacketConn
-func (s *server) Listen(conn net.PacketConn) error {
-	s.connMutex.Lock()
-	s.conn = conn
-	s.connMutex.Unlock()
-
+func (s *server) Serve() error {
 	for {
 		data := getPacketBuffer()
 		data = data[:protocol.MaxPacketSize]
-		n, remoteAddr, err := conn.ReadFrom(data)
+		n, remoteAddr, err := s.conn.ReadFrom(data)
 		if err != nil {
 			if strings.HasSuffix(err.Error(), "use of closed network connection") {
 				return nil
@@ -81,24 +90,12 @@ func (s *server) Listen(conn net.PacketConn) error {
 			return err
 		}
 		data = data[:n]
-		if err := s.handlePacket(conn, remoteAddr, data); err != nil {
+		if err := s.handlePacket(s.conn, remoteAddr, data); err != nil {
 			utils.Errorf("error handling packet: %s", err.Error())
 		}
 	}
 }
 
-func (s *server) ListenAddr(addr string) error {
-	udpAddr, err := net.ResolveUDPAddr("udp", addr)
-	if err != nil {
-		return err
-	}
-	conn, err := net.ListenUDP("udp", udpAddr)
-	if err != nil {
-		return err
-	}
-	return s.Listen(conn)
-}
-
 // Close the server
 func (s *server) Close() error {
 	s.sessionsMutex.Lock()
@@ -111,15 +108,10 @@ func (s *server) Close() error {
 	}
 	s.sessionsMutex.Unlock()
 
-	s.connMutex.Lock()
-	conn := s.conn
-	s.conn = nil
-	s.connMutex.Unlock()
-
-	if conn == nil {
+	if s.conn == nil {
 		return nil
 	}
-	return conn.Close()
+	return s.conn.Close()
 }
 
 // Addr returns the server's network address
@@ -201,7 +193,9 @@ func (s *server) handlePacket(pconn net.PacketConn, remoteAddr net.Addr, packet
 		s.sessionsMutex.Lock()
 		s.sessions[hdr.ConnectionID] = session
 		s.sessionsMutex.Unlock()
-		go s.config.ConnState(session, ConnStateVersionNegotiated)
+		if s.config.ConnState != nil {
+			go s.config.ConnState(session, ConnStateVersionNegotiated)
+		}
 	}
 	if session == nil {
 		// Late packet for closed session

server_test.go

@@ -55,35 +55,43 @@ func newMockSession(conn connection, v protocol.VersionNumber, connectionID prot
 
 var _ = Describe("Server", func() {
 	var (
-		serv            *server
 		conn            *mockPacketConn
+		config          *Config
 		connStateStatus ConnState
 		connStateCalled bool
-		firstPacket     []byte // a valid first packet for a new connection with connectionID 0x4cfa9f9b668619f6 (= connID)
-		connID          = protocol.ConnectionID(0x4cfa9f9b668619f6)
 		udpAddr         = &net.UDPAddr{IP: net.IPv4(192, 168, 100, 200), Port: 1337}
 	)
 
 	BeforeEach(func() {
-		serv = &server{
-			sessions:   map[protocol.ConnectionID]packetHandler{},
-			newSession: newMockSession,
-			conn:       &mockPacketConn{},
-			config: &Config{
-				ConnState: func(_ Session, cs ConnState) {
-					connStateStatus = cs
-					connStateCalled = true
-				},
+		conn = &mockPacketConn{}
+		config = &Config{
+			ConnState: func(_ Session, cs ConnState) {
+				connStateStatus = cs
+				connStateCalled = true
 			},
 		}
-		conn = serv.conn.(*mockPacketConn)
-		b := &bytes.Buffer{}
-		utils.WriteUint32(b, protocol.VersionNumberToTag(protocol.SupportedVersions[0]))
-		firstPacket = []byte{0x09, 0xf6, 0x19, 0x86, 0x66, 0x9b, 0x9f, 0xfa, 0x4c}
-		firstPacket = append(append(firstPacket, b.Bytes()...), 0x01)
 	})
 
 	Context("with mock session", func() {
+		var (
+			serv        *server
+			firstPacket []byte // a valid first packet for a new connection with connectionID 0x4cfa9f9b668619f6 (= connID)
+			connID      = protocol.ConnectionID(0x4cfa9f9b668619f6)
+		)
+
+		BeforeEach(func() {
+			serv = &server{
+				sessions:   make(map[protocol.ConnectionID]packetHandler),
+				newSession: newMockSession,
+				conn:       conn,
+				config:     config,
+			}
+			b := &bytes.Buffer{}
+			utils.WriteUint32(b, protocol.VersionNumberToTag(protocol.SupportedVersions[0]))
+			firstPacket = []byte{0x09, 0xf6, 0x19, 0x86, 0x66, 0x9b, 0x9f, 0xfa, 0x4c}
+			firstPacket = append(append(firstPacket, b.Bytes()...), 0x01)
+		})
+
 		It("returns the address", func() {
 			conn.addr = &net.UDPAddr{
 				IP:   net.IPv4(192, 168, 13, 37),
@@ -218,7 +226,7 @@ var _ = Describe("Server", func() {
 		config := Config{
 			ConnState: func(_ Session, _ ConnState) {},
 		}
-		ln, err := NewListener(&config)
+		ln, err := Listen(conn, &config)
 		server := ln.(*server)
 		Expect(err).ToNot(HaveOccurred())
 		Expect(server.deleteClosedSessionsAfter).To(Equal(protocol.ClosedSessionDeleteTimeout))
@@ -228,26 +236,22 @@ var _ = Describe("Server", func() {
 	})
 
 	It("listens on a given address", func() {
-		var listenReturned bool
 		addr := "127.0.0.1:13579"
-		serv.conn = nil
-		go func() {
-			defer GinkgoRecover()
-			err := serv.ListenAddr(addr)
-			Expect(err).ToNot(HaveOccurred())
-			listenReturned = true
-		}()
-		Eventually(func() net.PacketConn { return serv.conn }).ShouldNot(BeNil())
+		ln, err := ListenAddr(addr, config)
+		Expect(err).ToNot(HaveOccurred())
+		serv := ln.(*server)
 		Expect(serv.Addr().String()).To(Equal(addr))
-		Consistently(func() bool { return listenReturned }).Should(BeFalse())
 	})
 
 	It("setups and responds with version negotiation", func() {
 		conn.dataToRead = []byte{0x09, 0x01, 0, 0, 0, 0, 0, 0, 0, 0x01, 0x01, 'Q', '0', '0', '0', 0x01}
 		conn.dataReadFrom = udpAddr
+		ln, err := Listen(conn, config)
+		Expect(err).ToNot(HaveOccurred())
+
 		go func() {
 			defer GinkgoRecover()
-			err := serv.Listen(conn)
+			err := ln.Serve()
 			Expect(err).ToNot(HaveOccurred())
 		}()
 
@@ -263,15 +267,17 @@ var _ = Describe("Server", func() {
 	It("sends a PublicReset for new connections that don't have the VersionFlag set", func() {
 		conn.dataReadFrom = udpAddr
 		conn.dataToRead = []byte{0x08, 0xf6, 0x19, 0x86, 0x66, 0x9b, 0x9f, 0xfa, 0x4c, 0x01}
+		ln, err := Listen(conn, config)
+		Expect(err).ToNot(HaveOccurred())
 		go func() {
 			defer GinkgoRecover()
-			err := serv.Listen(conn)
+			err := ln.Serve()
 			Expect(err).ToNot(HaveOccurred())
 		}()
 
 		Eventually(func() int { return conn.dataWritten.Len() }).ShouldNot(BeZero())
 		Expect(conn.dataWrittenTo).To(Equal(udpAddr))
 		Expect(conn.dataWritten.Bytes()[0] & 0x02).ToNot(BeZero()) // check that the ResetFlag is set
-		Expect(serv.sessions).To(BeEmpty())
+		Expect(ln.(*server).sessions).To(BeEmpty())
 	})
 })