Fix app server goroutine leak (#9332)

Fixes gravitational/teleport-private#79
LAT-APP21-4: DOS - Goroutine leak in app server

Prevent the app server's HandleConnection from blocking for 
every connection until the server closes. This change blocks 
only until the connection is closed.

lib/srv/app/listener.go

@@ -18,11 +18,16 @@ package app
 
 import (
 	"context"
+	"errors"
 	"net"
 
 	"github.com/gravitational/trace"
 )
 
+// errListenerConnServed is used as a signal between listener and Server.
+// See listener.Accept for details.
+var errListenerConnServed = errors.New("ok: listener conn served")
+
 // listener wraps a net.Conn in a net.Listener interface. This allows passing
 // a channel connection from the reverse tunnel subsystem to an HTTP server.
 type listener struct {
@@ -48,10 +53,20 @@ func newListener(ctx context.Context, conn net.Conn) *listener {
 	}
 }
 
-// Accept returns the connection.
+// Accept returns the connection. An error is returned when this listener
+// is closed, its parent context is closed, or the second time it is called.
+//
+// On the second call, this method returns errListenerConnServed. This will
+// trigger the calling http.Serve function to exit gracefully, close this
+// listener, and return control to the http.Serve caller. The caller should
+// ignore errListenerConnServed and handle all other errors.
 func (l *listener) Accept() (net.Conn, error) {
 	select {
-	case conn := <-l.connCh:
+	case conn, more := <-l.connCh:
+		if !more {
+			return nil, errListenerConnServed // normal operation signal
+		}
+		close(l.connCh)
 		return conn, nil
 	case <-l.closeContext.Done():
 		return nil, trace.BadParameter("closing context")

lib/srv/app/server.go

@@ -22,6 +22,7 @@ package app
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"net"
 	"net/http"
 	"strconv"
@@ -544,13 +545,27 @@ func (s *Server) ForceHeartbeat() error {
 // HandleConnection takes a connection and wraps it in a listener so it can
 // be passed to http.Serve to process as a HTTP request.
 func (s *Server) HandleConnection(conn net.Conn) {
-	// Wrap the listener in a TLS authorizing listener.
-	listener := newListener(s.closeContext, conn)
-	tlsListener := tls.NewListener(listener, s.tlsConfig)
-
-	if err := s.httpServer.Serve(tlsListener); err != nil {
+	// Wrap conn in a CloserConn to detect when it is closed.
+	// Returning early will close conn before it has been serviced.
+	// httpServer will initiate the close call.
+	closerConn := utils.NewCloserConn(conn)
+
+	// Wrap a TLS authorizing conn in a single-use listener.
+	tlsConn := tls.Server(closerConn, s.tlsConfig)
+	listener := newListener(s.closeContext, tlsConn)
+
+	// Serve will return as soon as tlsConn is running in its own goroutine
+	err := s.httpServer.Serve(listener)
+	if err != nil && !errors.Is(err, errListenerConnServed) {
+		// okay to ignore errListenerConnServed; it is a signal that our
+		// single-use listener has passed the connection to http.Serve
+		// and conn is being served. See listener.Accept for details.
 		s.log.Warnf("Failed to handle connection: %v.", err)
+		return
 	}
+
+	// Wait for connection to close.
+	closerConn.Wait()
 }
 
 // ServeHTTP will forward the *http.Request to the target application.

lib/srv/app/server_test.go

@@ -569,10 +569,8 @@ func (s *Suite) checkHTTPResponse(t *testing.T, clientCert tls.Certificate, chec
 	checkResp(resp)
 	require.NoError(t, resp.Body.Close())
 
-	// Context will close because of the net.Pipe, expect a context canceled
-	// error here.
-	err = s.appServer.Close()
-	require.NotNil(t, err)
+	// Close should not trigger an error.
+	require.NoError(t, s.appServer.Close())
 
 	// Wait for the application server to actually stop serving before
 	// closing the test. This will make sure the server removes the listeners

lib/utils/conn.go

@@ -18,6 +18,7 @@ package utils
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -31,16 +32,20 @@ import (
 // NewCloserConn returns new connection wrapper that
 // when closed will also close passed closers
 func NewCloserConn(conn net.Conn, closers ...io.Closer) *CloserConn {
-	return &CloserConn{
+	c := &CloserConn{
 		Conn:    conn,
 		closers: closers,
 	}
+	c.ctx, c.ctxCancel = context.WithCancel(context.Background())
+	return c
 }
 
 // CloserConn wraps connection and attaches additional closers to it
 type CloserConn struct {
 	net.Conn
-	closers []io.Closer
+	closers   []io.Closer
+	ctx       context.Context
+	ctxCancel context.CancelFunc
 }
 
 // AddCloser adds any closer in ctx that will be called
@@ -49,16 +54,27 @@ func (c *CloserConn) AddCloser(closer io.Closer) {
 	c.closers = append(c.closers, closer)
 }
 
-// Close closes connection and all closers
+// Close connection, all closers, and cancel context.
 func (c *CloserConn) Close() error {
 	var errors []error
 	for _, closer := range c.closers {
 		errors = append(errors, closer.Close())
 	}
 	errors = append(errors, c.Conn.Close())
+	c.ctxCancel()
 	return trace.NewAggregate(errors...)
 }
 
+// Context returns a context that is cancelled once the connection is closed.
+func (c *CloserConn) Context() context.Context {
+	return c.ctx
+}
+
+// Wait for connection to close.
+func (c *CloserConn) Wait() {
+	<-c.ctx.Done()
+}
+
 // Roundtrip is a single connection simplistic HTTP client
 // that allows us to bypass a connection pool to test load balancing
 // used in tests, as it only supports GET request on /