From c56829eff7fcc37cbdf061d9626de40cfc00e158 Mon Sep 17 00:00:00 2001
From: Simeng He <simeng@tailscale.com>
Date: Thu, 13 May 2021 13:05:28 -0400
Subject: [PATCH] Direct tests

---
 control/controlclient/direct.go      |  9 ++++----
 control/controlclient/direct_test.go | 33 ++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/control/controlclient/direct.go b/control/controlclient/direct.go
index da01f96d9..dfdce39c5 100644
--- a/control/controlclient/direct.go
+++ b/control/controlclient/direct.go
@@ -765,7 +765,7 @@ func (c *Direct) sendMapRequest(ctx context.Context, maxPolls int, cb func(*netm
 
 		if pr := resp.PingRequest; pr != nil {
 			log.Println("Ping Triggered")
-			CustomPing(pr)
+			go CustomPing(&resp)
 			// go answerPing(c.logf, c.httpc, pr)
 		}
 
@@ -1220,7 +1220,8 @@ func sleepAsRequested(ctx context.Context, logf logger.Logf, timeoutReset chan<-
 
 // Run the ping suite from this client to another one
 // Send the ping results via http to the adminhttp handlers.
-func CustomPing(pr *tailcfg.PingRequest) {
-	log.Println("Custom Ping Triggered")
-	log.Println(pr)
+func CustomPing(mr *tailcfg.MapResponse) bool {
+	log.Printf("Custom Ping Triggered with %d number of peers\n", len(mr.Peers))
+	log.Println(mr.PingRequest)
+	return len(mr.Peers) > 0
 }
diff --git a/control/controlclient/direct_test.go b/control/controlclient/direct_test.go
index 4d2607616..8fbe55aea 100644
--- a/control/controlclient/direct_test.go
+++ b/control/controlclient/direct_test.go
@@ -103,3 +103,36 @@ func TestNewHostinfo(t *testing.T) {
 	}
 	t.Logf("Got: %s", j)
 }
+
+func TestPingFromMapResponse(t *testing.T) {
+	hi := NewHostinfo()
+	ni := tailcfg.NetInfo{LinkType: "wired"}
+	hi.NetInfo = &ni
+
+	key, err := wgkey.NewPrivate()
+	if err != nil {
+		t.Error(err)
+	}
+	opts := Options{
+		ServerURL: "https://example.com",
+		Hostinfo:  hi,
+		GetMachinePrivateKey: func() (wgkey.Private, error) {
+			return key, nil
+		},
+	}
+	c, err := NewDirect(opts)
+	if c == nil || err != nil {
+		t.Errorf("Direct not created %w", err)
+	}
+	peers := []*tailcfg.Node{
+		{ID: 1},
+		{ID: 2},
+		{ID: 3},
+	}
+	pingRequest := tailcfg.PingRequest{URL: "localhost:3040", Log: true, PayloadSize: 10}
+	mr := &tailcfg.MapResponse{Peers: peers, Domain: "DumbTest", PingRequest: &pingRequest}
+	if !CustomPing(mr) {
+		t.Errorf("Custom ping failed!\n")
+	}
+	t.Log("Successfull ping")
+}