From ba4ee2631328e6b3c1ed8175b6dbb5109ff50ed8 Mon Sep 17 00:00:00 2001
From: Simeng He <simeng@tailscale.com>
Date: Wed, 26 May 2021 14:41:35 -0400
Subject: [PATCH] Enough changes for a PR

---
 tstest/integration/integration_test.go        | 33 ++++++++++++++-----
 tstest/integration/testcontrol/testcontrol.go | 17 +++++++---
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/tstest/integration/integration_test.go b/tstest/integration/integration_test.go
index 63ef2bcf0..a5eadccd1 100644
--- a/tstest/integration/integration_test.go
+++ b/tstest/integration/integration_test.go
@@ -276,10 +276,12 @@ func newTestEnv(t testing.TB, bins *testBinaries) *testEnv {
 	derpMap, derpShutdown := runDERPAndStun(t, logger.Discard)
 	logc := new(logCatcher)
 	control := &testcontrol.Server{
-		DERPMap:      derpMap,
-		PingRequestC: make(chan bool, 1),
+		DERPMap: derpMap,
 	}
 	trafficTrap := new(trafficTrap)
+	log.Println("SERVER ATTACHED")
+	log.Println(len(control.PingRequestC))
+	// go func() { control.PingRequestC <- true }()
 	e := &testEnv{
 		t:                 t,
 		Binaries:          bins,
@@ -767,6 +769,7 @@ func TestControlSelectivePing(t *testing.T) {
 	bins := buildTestBinaries(t)
 
 	env := newTestEnv(t, bins)
+	log.Println("POSTSTARTUP")
 	defer env.Close()
 
 	// Create two nodes:
@@ -785,14 +788,26 @@ func TestControlSelectivePing(t *testing.T) {
 	n1.AwaitRunning(t)
 	n2.AwaitRunning(t)
 
-	req := new(tailcfg.MapRequest)
-	env.Control.MapResponse(req)
+	// Wait for server to start serveMap
 	if err := tstest.WaitFor(2*time.Second, func() error {
-		st := n1.MustStatus(t)
-		req.NodeKey = tailcfg.NodeKey(st.Self.PublicKey)
-		// env.Control.AddControlPingRequest()
-		env.Control.PingRequestC <- true
-		t.Log("CHANNEL LENGTH", len(env.Control.PingRequestC))
+		t.Log("ENOUGHTIME")
+		env.Control.AddControlPingRequest()
+		if len(env.Control.PingRequestC) == 0 {
+			return errors.New("failed to add to PingRequestC")
+		}
+		log.Println("CHANNEL LENGTH", len(env.Control.PingRequestC))
+		return nil
+	}); err != nil {
+		t.Error(err)
+	}
+
+	// Wait for a MapResponse
+	if err := tstest.WaitFor(20*time.Second, func() error {
+		// Simulate the time needed for MapResponse method call.
+		time.Sleep(500 * time.Millisecond)
+		if len(env.Control.PingRequestC) == 1 {
+			t.Error("Expected PingRequestC to be empty")
+		}
 		return nil
 	}); err != nil {
 		t.Error(err)
diff --git a/tstest/integration/testcontrol/testcontrol.go b/tstest/integration/testcontrol/testcontrol.go
index dfebb342d..314270e28 100644
--- a/tstest/integration/testcontrol/testcontrol.go
+++ b/tstest/integration/testcontrol/testcontrol.go
@@ -43,8 +43,9 @@ type Server struct {
 	Verbose      bool
 	PingRequestC chan bool
 
-	initMuxOnce sync.Once
-	mux         *http.ServeMux
+	initMuxOnce       sync.Once
+	mux               *http.ServeMux
+	initPRchannelOnce sync.Once
 
 	mu            sync.Mutex
 	pubKey        wgkey.Key
@@ -94,10 +95,17 @@ func (s *Server) initMux() {
 	s.mux.HandleFunc("/mockpingrequest", s.serveMockPing)
 }
 
+func (s *Server) initPingRequestC() {
+	log.Println("Channel created")
+	s.PingRequestC = make(chan bool, 1)
+	// s.AddControlPingRequest()
+	// log.Println("Channel length : ", len(s.PingRequestC))
+}
+
 func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	log.Println("HTTPSERVE")
-	s.PingRequestC = make(chan bool, 1)
 	s.initMuxOnce.Do(s.initMux)
+	s.initPRchannelOnce.Do(s.initPingRequestC)
 	s.mux.ServeHTTP(w, r)
 }
 
@@ -196,6 +204,7 @@ func (s *Server) AllNodes() (nodes []*tailcfg.Node) {
 // in serveMap this will result to a ControlPingRequest
 // added to the next MapResponse sent to the client
 func (s *Server) AddControlPingRequest() {
+	// Redundant check to avoid errors when called multiple times
 	if len(s.PingRequestC) == 0 {
 		s.PingRequestC <- true
 	}
@@ -503,7 +512,7 @@ func (s *Server) serveMap(w http.ResponseWriter, r *http.Request, mkey tailcfg.M
 		log.Println("LENGTHER", len(s.PingRequestC))
 		select {
 		case <-s.PingRequestC:
-			log.Println("PINGADD")
+			log.Println("PINGADD", len(s.PingRequestC))
 			s.addPingRequest(res)
 		default:
 			log.Println("NOTEXIST")