From 18b2da625033f52fd869f12582ce5c31ee248e3c Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@golang.org>
Date: Wed, 2 Mar 2016 18:17:55 +0000
Subject: [PATCH] Simplify server setup when testing the http.Handler-based
 server transport

This simplifies the code that's used for wiring up the http.Handler-based
transport for testing. It's not used outside of tests.

http2.Server.ServeConn was added somewhat recently for other reasons and
I just realized it would simplify this code as well.
---
 server.go | 55 ++++++++-----------------------------------------------
 1 file changed, 8 insertions(+), 47 deletions(-)

diff --git a/server.go b/server.go
index 79124e11..0050d57c 100644
--- a/server.go
+++ b/server.go
@@ -247,7 +247,6 @@ func (s *Server) Serve(lis net.Listener) error {
 		delete(s.lis, lis)
 		s.mu.Unlock()
 	}()
-	listenerAddr := lis.Addr()
 	for {
 		rawConn, err := lis.Accept()
 		if err != nil {
@@ -258,13 +257,13 @@ func (s *Server) Serve(lis net.Listener) error {
 		}
 		// Start a new goroutine to deal with rawConn
 		// so we don't stall this Accept loop goroutine.
-		go s.handleRawConn(listenerAddr, rawConn)
+		go s.handleRawConn(rawConn)
 	}
 }
 
 // handleRawConn is run in its own goroutine and handles a just-accepted
 // connection that has not had any I/O performed on it yet.
-func (s *Server) handleRawConn(listenerAddr net.Addr, rawConn net.Conn) {
+func (s *Server) handleRawConn(rawConn net.Conn) {
 	conn, authInfo, err := s.useTransportAuthenticator(rawConn)
 	if err != nil {
 		s.mu.Lock()
@@ -284,7 +283,7 @@ func (s *Server) handleRawConn(listenerAddr net.Addr, rawConn net.Conn) {
 	s.mu.Unlock()
 
 	if s.opts.useHandlerImpl {
-		s.serveUsingHandler(listenerAddr, conn)
+		s.serveUsingHandler(conn)
 	} else {
 		s.serveNewHTTP2Transport(conn, authInfo)
 	}
@@ -340,29 +339,18 @@ var _ http.Handler = (*Server)(nil)
 // method as one of the environment types.
 //
 // conn is the *tls.Conn that's already been authenticated.
-func (s *Server) serveUsingHandler(listenerAddr net.Addr, conn net.Conn) {
+func (s *Server) serveUsingHandler(conn net.Conn) {
 	if !s.addConn(conn) {
 		conn.Close()
 		return
 	}
 	defer s.removeConn(conn)
-	connDone := make(chan struct{})
-	hs := &http.Server{
-		Handler: s,
-		ConnState: func(c net.Conn, cs http.ConnState) {
-			if cs == http.StateClosed {
-				close(connDone)
-			}
-		},
-	}
-	if err := http2.ConfigureServer(hs, &http2.Server{
+	h2s := &http2.Server{
 		MaxConcurrentStreams: s.opts.maxConcurrentStreams,
-	}); err != nil {
-		grpclog.Fatalf("grpc: http2.ConfigureServer: %v", err)
-		return
 	}
-	hs.Serve(&singleConnListener{addr: listenerAddr, conn: conn})
-	<-connDone
+	h2s.ServeConn(conn, &http2.ServeConnOpts{
+		Handler: s,
+	})
 }
 
 func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -754,30 +742,3 @@ func SetTrailer(ctx context.Context, md metadata.MD) error {
 	}
 	return stream.SetTrailer(md)
 }
-
-// singleConnListener is a net.Listener that yields a single conn.
-type singleConnListener struct {
-	mu   sync.Mutex
-	addr net.Addr
-	conn net.Conn // nil if done
-}
-
-func (ln *singleConnListener) Addr() net.Addr { return ln.addr }
-
-func (ln *singleConnListener) Close() error {
-	ln.mu.Lock()
-	defer ln.mu.Unlock()
-	ln.conn = nil
-	return nil
-}
-
-func (ln *singleConnListener) Accept() (net.Conn, error) {
-	ln.mu.Lock()
-	defer ln.mu.Unlock()
-	c := ln.conn
-	if c == nil {
-		return nil, io.EOF
-	}
-	ln.conn = nil
-	return c, nil
-}