From 9a43a92b1c7d61e984b0bac56d64ab972a70734b Mon Sep 17 00:00:00 2001
From: Nick Vollmar <nvollmar@cloudflare.com>
Date: Fri, 5 Apr 2019 01:57:00 -0500
Subject: [PATCH] TUN-1577: decompose carrier.StartServer to make
 TestStartServer less flappy

---
 carrier/carrier.go      | 14 +++++++++++---
 carrier/carrier_test.go | 14 +++++++-------
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/carrier/carrier.go b/carrier/carrier.go
index dd710198..f54cd653 100644
--- a/carrier/carrier.go
+++ b/carrier/carrier.go
@@ -38,8 +38,8 @@ func StartClient(logger *logrus.Logger, originURL string, stream io.ReadWriter,
 	return serveStream(logger, originURL, stream, headers)
 }
 
-// StartServer will setup a server on a specified port and copy data over a WebSocket connection
-// to the edge (originURL)
+// StartServer will setup a listener on a specified address/port and then
+// forward connections to the origin by calling `Serve()`.
 func StartServer(logger *logrus.Logger, address, originURL string, shutdownC <-chan struct{}, headers http.Header) error {
 	listener, err := net.Listen("tcp", address)
 	if err != nil {
@@ -47,6 +47,14 @@ func StartServer(logger *logrus.Logger, address, originURL string, shutdownC <-c
 		return err
 	}
 	logger.Info("Started listening on ", address)
+	return Serve(logger, listener, originURL, shutdownC, headers)
+}
+
+// Serve accepts incoming connections on the specified net.Listener.
+// Each connection is handled in a new goroutine: its data is copied over a
+// WebSocket connection to the edge (originURL).
+// `Serve` always closes `listener`.
+func Serve(logger *logrus.Logger, listener net.Listener, originURL string, shutdownC <-chan struct{}, headers http.Header) error {
 	defer listener.Close()
 	for {
 		select {
@@ -62,7 +70,7 @@ func StartServer(logger *logrus.Logger, address, originURL string, shutdownC <-c
 	}
 }
 
-// serveConnection handles connections for the StartServer call
+// serveConnection handles connections for the Serve() call
 func serveConnection(logger *logrus.Logger, c net.Conn, originURL string, headers http.Header) {
 	defer c.Close()
 	serveStream(logger, originURL, c, headers)
diff --git a/carrier/carrier_test.go b/carrier/carrier_test.go
index 0828e45d..e1171175 100644
--- a/carrier/carrier_test.go
+++ b/carrier/carrier_test.go
@@ -58,7 +58,10 @@ func TestStartClient(t *testing.T) {
 }
 
 func TestStartServer(t *testing.T) {
-	listenerAddress := "localhost:1117"
+	listener, err := net.Listen("tcp", "localhost:")
+	if err != nil {
+		t.Fatalf("Error starting listener: %v", err)
+	}
 	message := "Good morning Austin! Time for another sunny day in the great state of Texas."
 	logger := logrus.New()
 	shutdownC := make(chan struct{})
@@ -66,16 +69,13 @@ func TestStartServer(t *testing.T) {
 	defer ts.Close()
 
 	go func() {
-		err := StartServer(logger, listenerAddress, "http://"+ts.Listener.Addr().String(), shutdownC, nil)
+		err := Serve(logger, listener, "http://"+ts.Listener.Addr().String(), shutdownC, nil)
 		if err != nil {
-			t.Fatalf("Error starting server: %v", err)
+			t.Fatalf("Error running server: %v", err)
 		}
 	}()
 
-	conn, err := net.Dial("tcp", listenerAddress)
-	if err != nil {
-		t.Fatalf("Error connecting to server: %v", err)
-	}
+	conn, err := net.Dial("tcp", listener.Addr().String())
 	conn.Write([]byte(message))
 
 	readBuffer := make([]byte, len(message))