From 3ed2124356c56badbdbcfa3826e7abf9a678d048 Mon Sep 17 00:00:00 2001
From: Avery Pennarun <apenwarr@tailscale.com>
Date: Thu, 21 May 2020 16:30:20 -0400
Subject: [PATCH] ipn: Resolve some resource leaks in test.

Updates tailscale/corp#255.

Signed-off-by: Avery Pennarun <apenwarr@tailscale.com>
---
 ipn/backend.go  |  4 ++++
 ipn/e2e_test.go | 29 ++++++++++++++++++++---------
 ipn/local.go    |  4 +++-
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/ipn/backend.go b/ipn/backend.go
index 5dae95e2c..f717eab81 100644
--- a/ipn/backend.go
+++ b/ipn/backend.go
@@ -5,6 +5,7 @@
 package ipn
 
 import (
+	"net/http"
 	"time"
 
 	"tailscale.com/control/controlclient"
@@ -104,6 +105,9 @@ type Options struct {
 	LegacyConfigPath string
 	// Notify is called when backend events happen.
 	Notify func(Notify) `json:"-"`
+	// HTTPTestClient is an optional HTTP client to pass to controlclient
+	// (for tests only).
+	HTTPTestClient *http.Client
 }
 
 // Backend is the interface between Tailscale frontends
diff --git a/ipn/e2e_test.go b/ipn/e2e_test.go
index 8308593e3..a3b7acd9a 100644
--- a/ipn/e2e_test.go
+++ b/ipn/e2e_test.go
@@ -8,6 +8,7 @@ package ipn
 
 import (
 	"bytes"
+	"context"
 	"io/ioutil"
 	"net/http"
 	"net/http/cookiejar"
@@ -40,6 +41,10 @@ func init() {
 
 func TestIPN(t *testing.T) {
 	tstest.PanicOnLog()
+	rc := tstest.NewResourceCheck()
+	defer rc.Assert(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
 
 	// This gets reassigned inside every test, so that the connections
 	// all log using the "current" t.Logf function. Sigh.
@@ -66,27 +71,32 @@ func TestIPN(t *testing.T) {
 		ctl.ServeHTTP(w, r)
 	}
 	https := httptest.NewServer(http.HandlerFunc(ctlHandler))
+	https.Config.ErrorLog = logger.StdLogger(logf)
 	serverURL := https.URL
-	defer https.Close()
-	defer https.CloseClientConnections()
 
 	tmpdir, err := ioutil.TempDir("", "ipntest")
 	if err != nil {
 		t.Fatalf("create tempdir: %v\n", err)
 	}
+
 	ctl, err = control.New(tmpdir, tmpdir, tmpdir, serverURL, true, logf)
 	if err != nil {
 		t.Fatalf("create control server: %v\n", ctl)
 	}
+	defer ctl.Shutdown()
+	defer https.Close()
+	defer https.CloseClientConnections()
+	defer cancel()
+
 	if _, err := ctl.DB().FindOrCreateUser("google", "test1@example.com", "", ""); err != nil {
 		t.Fatal(err)
 	}
 
-	n1 := newNode(t, logf, "n1", https, false)
+	n1 := newNode(t, ctx, logf, "n1", https, false)
 	defer n1.Backend.Shutdown()
 	n1.Backend.StartLoginInteractive()
 
-	n2 := newNode(t, logf, "n2", https, true)
+	n2 := newNode(t, ctx, logf, "n2", https, true)
 	defer n2.Backend.Shutdown()
 	n2.Backend.StartLoginInteractive()
 
@@ -126,7 +136,7 @@ func TestIPN(t *testing.T) {
 						authNode(t, ctl, n2.Backend)
 					}
 				}
-			case <-time.After(3 * time.Second):
+			case <-time.After(5 * time.Second):
 				t.Fatalf("\n\n\nFATAL: timed out waiting for notifications.\n\n\n")
 			}
 		}
@@ -210,7 +220,7 @@ type testNode struct {
 }
 
 // Create a new IPN node.
-func newNode(t *testing.T, logfx logger.Logf, prefix string, https *httptest.Server, weirdPrefs bool) testNode {
+func newNode(t *testing.T, ctx context.Context, logfx logger.Logf, prefix string, https *httptest.Server, weirdPrefs bool) testNode {
 	t.Helper()
 
 	logfe := logger.WithPrefix(logfx, prefix+"e: ")
@@ -254,8 +264,9 @@ func newNode(t *testing.T, logfx logger.Logf, prefix string, https *httptest.Ser
 	}
 
 	n.Start(Options{
-		FrontendLogID: prefix + "-f",
-		Prefs:         prefs,
+		HTTPTestClient: httpc,
+		FrontendLogID:  prefix + "-f",
+		Prefs:          prefs,
 		Notify: func(n Notify) {
 			// Automatically visit auth URLs
 			if n.BrowseToURL != nil {
@@ -275,7 +286,7 @@ func newNode(t *testing.T, logfx logger.Logf, prefix string, https *httptest.Ser
 				}
 				req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
 
-				if _, err := httpc.Do(req); err != nil {
+				if _, err := httpc.Do(req.WithContext(ctx)); err != nil {
 					logf("BrowseToURL: %v\n", err)
 				}
 			}
diff --git a/ipn/local.go b/ipn/local.go
index 41b4a467f..73c67e1a2 100644
--- a/ipn/local.go
+++ b/ipn/local.go
@@ -116,13 +116,14 @@ func NewLocalBackend(logf logger.Logf, logid string, store StateStore, e wgengin
 // Shutdown halts the backend and all its sub-components. The backend
 // can no longer be used after Shutdown returns.
 func (b *LocalBackend) Shutdown() {
-	b.ctxCancel()
 	b.mu.Lock()
 	cli := b.c
 	b.mu.Unlock()
+
 	if cli != nil {
 		cli.Shutdown()
 	}
+	b.ctxCancel()
 	b.e.Close()
 	b.e.Wait()
 }
@@ -258,6 +259,7 @@ func (b *LocalBackend) Start(opts Options) error {
 		Hostinfo:        hi,
 		KeepAlive:       true,
 		NewDecompressor: b.newDecompressor,
+		HTTPTestClient:  opts.HTTPTestClient,
 	})
 	if err != nil {
 		return err