From c5fcc38bf16b2152c0b4655541b75d6894e4259d Mon Sep 17 00:00:00 2001
From: Avery Pennarun <apenwarr@tailscale.com>
Date: Wed, 20 May 2020 11:08:24 -0400
Subject: [PATCH] controlclient tests: fix more memory leaks and add resource
 checking.

I can now run these tests with -count=1000 without running out of RAM.

Signed-off-by: Avery Pennarun <apenwarr@tailscale.com>
---
 control/controlclient/auto_test.go   |  6 +++-
 control/controlclient/direct.go      |  1 -
 control/controlclient/direct_test.go | 48 ++++++++++++++++++++--------
 3 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/control/controlclient/auto_test.go b/control/controlclient/auto_test.go
index 847336852..0c6a0145e 100644
--- a/control/controlclient/auto_test.go
+++ b/control/controlclient/auto_test.go
@@ -1148,7 +1148,11 @@ func (s *server) newClientWithKey(t *testing.T, name, authKey string) *client {
 		},
 		Hostinfo: hi,
 		NewDecompressor: func() (Decompressor, error) {
-			return zstd.NewReader(nil)
+			return zstd.NewReader(nil,
+				zstd.WithDecoderLowmem(true),
+				zstd.WithDecoderConcurrency(1),
+				zstd.WithDecoderMaxMemory(65536),
+			)
 		},
 		KeepAlive: true,
 		AuthKey:   authKey,
diff --git a/control/controlclient/direct.go b/control/controlclient/direct.go
index 7fbf522f3..0e7afb6f8 100644
--- a/control/controlclient/direct.go
+++ b/control/controlclient/direct.go
@@ -654,7 +654,6 @@ func (c *Direct) decodeMsg(msg []byte, v interface{}) error {
 	if c.newDecompressor == nil {
 		b = decrypted
 	} else {
-		//decoder, err := zstd.NewReader(nil)
 		decoder, err := c.newDecompressor()
 		if err != nil {
 			return err
diff --git a/control/controlclient/direct_test.go b/control/controlclient/direct_test.go
index ec01fdebe..0fc2f313d 100644
--- a/control/controlclient/direct_test.go
+++ b/control/controlclient/direct_test.go
@@ -19,12 +19,18 @@
 	"github.com/klauspost/compress/zstd"
 	"github.com/tailscale/wireguard-go/wgcfg"
 	"tailscale.com/tailcfg"
+	"tailscale.com/tstest"
+	"tailscale.com/types/logger"
 	"tailscale.io/control" // not yet released
 )
 
 // Test that when there are two controlclient connections using the
 // same credentials, the later one disconnects the earlier one.
-func TestClientsReusingKeys(t *testing.T) {
+func TestDirectReusingKeys(t *testing.T) {
+	tstest.PanicOnLog()
+	rc := tstest.NewResourceCheck()
+	defer rc.Assert(t)
+
 	tmpdir, err := ioutil.TempDir("", "control-test-")
 	if err != nil {
 		t.Fatal(err)
@@ -33,6 +39,7 @@ func TestClientsReusingKeys(t *testing.T) {
 	httpsrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		server.ServeHTTP(w, r)
 	}))
+	httpsrv.Config.ErrorLog = logger.StdLogger(t.Logf)
 	defer func() {
 		httpsrv.CloseClientConnections()
 		httpsrv.Close()
@@ -50,6 +57,7 @@ func TestClientsReusingKeys(t *testing.T) {
 		t.Fatal(err)
 	}
 	server.QuietLogging = true
+	defer server.Shutdown()
 
 	hi := NewHostinfo()
 	hi.FrontendLogID = "go-test-only"
@@ -112,7 +120,7 @@ func TestClientsReusingKeys(t *testing.T) {
 	// update function periodically, then exit once c2 starts its own
 	// poll below.
 	gotNetmap := make(chan struct{}, 1)
-	pollErrCh := make(chan error)
+	pollErrCh := make(chan error, 1)
 	go func() {
 		pollErrCh <- c1.PollNetMap(ctx, -1, func(netMap *NetworkMap) {
 			select {
@@ -143,7 +151,11 @@ func TestClientsReusingKeys(t *testing.T) {
 		Persist:  c1.GetPersist(),
 		Hostinfo: hi,
 		NewDecompressor: func() (Decompressor, error) {
-			return zstd.NewReader(nil)
+			return zstd.NewReader(nil,
+				zstd.WithDecoderLowmem(true),
+				zstd.WithDecoderConcurrency(1),
+				zstd.WithDecoderMaxMemory(65536),
+			)
 		},
 		KeepAlive: true,
 	})
@@ -177,7 +189,11 @@ func TestClientsReusingKeys(t *testing.T) {
 	}
 }
 
-func TestClientsReusingOldKey(t *testing.T) {
+func TestDirectReusingOldKey(t *testing.T) {
+	tstest.PanicOnLog()
+	rc := tstest.NewResourceCheck()
+	defer rc.Assert(t)
+
 	tmpdir, err := ioutil.TempDir("", "control-test-")
 	if err != nil {
 		t.Fatal(err)
@@ -186,22 +202,26 @@ func TestClientsReusingOldKey(t *testing.T) {
 	httpsrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		server.ServeHTTP(w, r)
 	}))
-	httpc := httpsrv.Client()
-	httpc.Jar, err = cookiejar.New(nil)
-	if err != nil {
-		t.Fatal(err)
-	}
-	server, err = control.New(tmpdir, tmpdir, tmpdir, httpsrv.URL, true, t.Logf)
-	if err != nil {
-		t.Fatal(err)
-	}
-	server.QuietLogging = true
+	httpsrv.Config.ErrorLog = logger.StdLogger(t.Logf)
 	defer func() {
 		httpsrv.CloseClientConnections()
 		httpsrv.Close()
 		os.RemoveAll(tmpdir)
 	}()
 
+	httpc := httpsrv.Client()
+	httpc.Jar, err = cookiejar.New(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	server, err = control.New(tmpdir, tmpdir, tmpdir, httpsrv.URL, true, t.Logf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	server.QuietLogging = true
+	defer server.Shutdown()
+
 	hi := NewHostinfo()
 	hi.FrontendLogID = "go-test-only"
 	hi.BackendLogID = "go-test-only"