From 3b5b4c1557760243277ec6667cf5cef8ac7a0f91 Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Fri, 3 Jun 2016 13:30:19 +0200
Subject: [PATCH] add h2quic.Server.Serve(*net.UDPConn) and improve server
 tests

---
 h2quic/server.go      | 67 ++++++++++++++++++++-----------------------
 h2quic/server_test.go | 16 +++--------
 2 files changed, 35 insertions(+), 48 deletions(-)

diff --git a/h2quic/server.go b/h2quic/server.go
index 9c4178e5..72764f3a 100644
--- a/h2quic/server.go
+++ b/h2quic/server.go
@@ -38,6 +38,34 @@ type Server struct {
 
 // ListenAndServe listens on the UDP address s.Addr and calls s.Handler to handle HTTP/2 requests on incoming connections.
 func (s *Server) ListenAndServe() error {
+	if s.Server == nil {
+		return errors.New("use of h2quic.Server without http.Server")
+	}
+	return s.serveImpl(s.TLSConfig, nil)
+}
+
+// ListenAndServeTLS listens on the UDP address s.Addr and calls s.Handler to handle HTTP/2 requests on incoming connections.
+func (s *Server) ListenAndServeTLS(certFile, keyFile string) error {
+	var err error
+	certs := make([]tls.Certificate, 1)
+	certs[0], err = tls.LoadX509KeyPair(certFile, keyFile)
+	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,
+	}
+	return s.serveImpl(config, nil)
+}
+
+// Serve an existing udp conn
+func (s *Server) Serve(conn *net.UDPConn) error {
+	return s.serveImpl(s.TLSConfig, conn)
+}
+
+func (s *Server) serveImpl(tlsConfig *tls.Config, conn *net.UDPConn) error {
 	if s.Server == nil {
 		return errors.New("use of h2quic.Server without http.Server")
 	}
@@ -54,43 +82,10 @@ func (s *Server) ListenAndServe() error {
 	}
 	s.server = server
 	s.serverMutex.Unlock()
-	return server.ListenAndServe()
-}
-
-// ListenAndServeTLS listens on the UDP address s.Addr and calls s.Handler to handle HTTP/2 requests on incoming connections.
-func (s *Server) ListenAndServeTLS(certFile, keyFile string) error {
-	if s.Server == nil {
-		return errors.New("use of h2quic.Server without http.Server")
+	if conn == nil {
+		return server.ListenAndServe()
 	}
-	var err error
-	certs := make([]tls.Certificate, 1)
-	certs[0], err = tls.LoadX509KeyPair(certFile, keyFile)
-	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,
-	}
-	s.serverMutex.Lock()
-	if s.server != nil {
-		s.serverMutex.Unlock()
-		return errors.New("ListenAndServe may only be called once")
-	}
-	server, err := quic.NewServer(s.Addr, config, s.handleStreamCb)
-	if err != nil {
-		s.serverMutex.Unlock()
-		return err
-	}
-	s.server = server
-	s.serverMutex.Unlock()
-	return server.ListenAndServe()
-}
-
-// Serve should not be called, since it only works properly for TCP listeners.
-func (Server) Serve(net.Listener) error {
-	panic("h2quic.Server.Serve should not be called, see https://godoc.org/github.com/lucas-clemente/quic-go/h2quic")
+	return server.Serve(conn)
 }
 
 func (s *Server) handleStreamCb(session *quic.Session, stream utils.Stream) {
diff --git a/h2quic/server_test.go b/h2quic/server_test.go
index ff826016..8f72b610 100644
--- a/h2quic/server_test.go
+++ b/h2quic/server_test.go
@@ -31,9 +31,6 @@ func (s *mockSession) GetOrOpenStream(id protocol.StreamID) (utils.Stream, error
 func (s *mockSession) Close(error) error { s.closed = true; return nil }
 
 var _ = Describe("H2 server", func() {
-	const port = "4826"
-	const addr = "127.0.0.1:" + port
-
 	certPath := os.Getenv("GOPATH")
 	certPath += "/src/github.com/lucas-clemente/quic-go/example/"
 
@@ -163,12 +160,6 @@ var _ = Describe("H2 server", func() {
 		Eventually(func() bool { return handlerCalled }).Should(BeTrue())
 	})
 
-	It("should panic when Serve() is called", func() {
-		Expect(func() {
-			Server{}.Serve(nil)
-		}).To(Panic())
-	})
-
 	Context("setting http headers", func() {
 		expected := http.Header{
 			"Alt-Svc":            {`quic=":443"; ma=2592000; v="33,32,31,30"`},
@@ -218,7 +209,7 @@ var _ = Describe("H2 server", func() {
 	})
 
 	It("should error when ListenAndServeTLS is called with s.Server nil", func() {
-		err := (&Server{}).ListenAndServeTLS("", "")
+		err := (&Server{}).ListenAndServeTLS(certPath+"fullchain.pem", certPath+"privkey.pem")
 		Expect(err).To(MatchError("use of h2quic.Server without http.Server"))
 	})
 
@@ -229,7 +220,7 @@ var _ = Describe("H2 server", func() {
 
 	Context("ListenAndServe", func() {
 		BeforeEach(func() {
-			s.Server.Addr = addr
+			s.Server.Addr = "localhost:0"
 		})
 
 		AfterEach(func() {
@@ -266,7 +257,7 @@ var _ = Describe("H2 server", func() {
 
 	Context("ListenAndServeTLS", func() {
 		BeforeEach(func() {
-			s.Server.Addr = addr
+			s.Server.Addr = "localhost:0"
 		})
 
 		AfterEach(func() {
@@ -311,6 +302,7 @@ var _ = Describe("H2 server", func() {
 		// once it's started. So, we open a socket on the same port before the test,
 		// so that ListenAndServeQUIC definitely fails. This way we know it at least
 		// created a socket on the proper address :)
+		const addr = "127.0.0.1:4826"
 		udpAddr, err := net.ResolveUDPAddr("udp", addr)
 		Expect(err).NotTo(HaveOccurred())
 		c, err := net.ListenUDP("udp", udpAddr)