From d8550deb7f1f72d632f4eaeb10149abaca3db79d Mon Sep 17 00:00:00 2001
From: Piotr Tabor <ptab@google.com>
Date: Fri, 14 May 2021 10:20:38 +0200
Subject: [PATCH] Fix pkg/proxy tests such that they don't leek goroutines and
 do close transports.

---
 pkg/proxy/server.go      |  3 +++
 pkg/proxy/server_test.go | 18 +++++++++++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/pkg/proxy/server.go b/pkg/proxy/server.go
index 480a9492b4a..9a7b105f9a7 100644
--- a/pkg/proxy/server.go
+++ b/pkg/proxy/server.go
@@ -401,13 +401,16 @@ func (s *server) listenAndServe() {
 			continue
 		}
 
+		s.closeWg.Add(2)
 		go func() {
+			defer s.closeWg.Done()
 			// read incoming bytes from listener, dispatch to outgoing connection
 			s.transmit(out, in)
 			out.Close()
 			in.Close()
 		}()
 		go func() {
+			defer s.closeWg.Done()
 			// read response from outgoing connection, write back to listener
 			s.receive(in, out)
 			in.Close()
diff --git a/pkg/proxy/server_test.go b/pkg/proxy/server_test.go
index 26cd157b153..686a8c362b3 100644
--- a/pkg/proxy/server_test.go
+++ b/pkg/proxy/server_test.go
@@ -477,6 +477,7 @@ func testServerHTTP(t *testing.T, secure, delayTx bool) {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/hello", func(w http.ResponseWriter, req *http.Request) {
 		d, err := ioutil.ReadAll(req.Body)
+		req.Body.Close()
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -505,13 +506,12 @@ func testServerHTTP(t *testing.T, secure, delayTx bool) {
 		<-donec
 	}()
 	go func() {
-		defer close(donec)
 		if !secure {
 			srv.ListenAndServe()
 		} else {
 			srv.ListenAndServeTLS(tlsInfo.CertFile, tlsInfo.KeyFile)
 		}
-		defer srv.Close()
+		defer close(donec)
 	}()
 	time.Sleep(200 * time.Millisecond)
 
@@ -525,7 +525,11 @@ func testServerHTTP(t *testing.T, secure, delayTx bool) {
 	}
 	p := NewServer(cfg)
 	<-p.Ready()
-	defer p.Close()
+	defer func() {
+		lg.Info("closing Proxy server...")
+		p.Close()
+		lg.Info("closed Proxy server.")
+	}()
 
 	data := "Hello World!"
 
@@ -537,14 +541,18 @@ func testServerHTTP(t *testing.T, secure, delayTx bool) {
 		assert.NoError(t, terr)
 		cli := &http.Client{Transport: tp}
 		resp, err = cli.Post("https://"+srcAddr+"/hello", "", strings.NewReader(data))
+		defer cli.CloseIdleConnections()
+		defer tp.CloseIdleConnections()
 	} else {
 		resp, err = http.Post("http://"+srcAddr+"/hello", "", strings.NewReader(data))
+		defer http.DefaultClient.CloseIdleConnections()
 	}
 	assert.NoError(t, err)
 	d, err := ioutil.ReadAll(resp.Body)
 	if err != nil {
 		t.Fatal(err)
 	}
+	resp.Body.Close()
 	took1 := time.Since(now)
 	t.Logf("took %v with no latency", took1)
 
@@ -571,8 +579,11 @@ func testServerHTTP(t *testing.T, secure, delayTx bool) {
 		}
 		cli := &http.Client{Transport: tp}
 		resp, err = cli.Post("https://"+srcAddr+"/hello", "", strings.NewReader(data))
+		defer cli.CloseIdleConnections()
+		defer tp.CloseIdleConnections()
 	} else {
 		resp, err = http.Post("http://"+srcAddr+"/hello", "", strings.NewReader(data))
+		defer http.DefaultClient.CloseIdleConnections()
 	}
 	if err != nil {
 		t.Fatal(err)
@@ -581,6 +592,7 @@ func testServerHTTP(t *testing.T, secure, delayTx bool) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	resp.Body.Close()
 	took2 := time.Since(now)
 	t.Logf("took %v with latency %v±%v", took2, lat, rv)