Added routes for testcontrol ping communication, added integration tests for injecting ControlPingRequestsimenghe/add-ping-route-testcontrol-mux

Signed-off-by: Simeng He <simeng@tailscale.com>

3 files changed, 142 insertions, 5 deletions

diff --git a/tailcfg/tailcfg.go b/tailcfg/tailcfg.go
index 87836d10c..b2273d528 100644
--- a/tailcfg/tailcfg.go
+++ b/tailcfg/tailcfg.go
@@ -886,6 +886,30 @@ type PingRequest struct {
 	Log bool `json:",omitempty"`
 }
 
+// ControlPingRequest is a request to have the client ping another client.
+// It will use the lower level ping TSMP or disco.
+type ControlPingRequest struct {
+	IP               netaddr.IP // IP address that we are going to ping
+	Peer             *Node      // Peer is a pointer to the Node that we want to Ping
+	StopAfterNDirect int        // StopAfterNDirect 1 means stop on 1st direct ping; 4 means 4 direct pings; 0 means do MaxPings and stop
+	MaxPings         int        // MaxPings total, direct or DERPed
+	Types            string     // Types empty means all: TSMP+ICMP+disco
+	URL              string     // URL of where we should stream responses back via HTTP
+}
+
+// StreamedPingResult sends the results of the LowLevelPing requested by a
+// ControlPingRequest in a MapResponse.
+type StreamedPingResult struct {
+	IP      netaddr.IP // IP Address that was Pinged
+	SeqNum  int        // SeqNum is somewhat redundant with TxID but for clarity
+	SentTo  NodeID     // SentTo for exit/subnet relays
+	TxID    string     // TxID N hex bytes random
+	Dir     string     // Dir "in"/"out"
+	Type    string     // Type of ping : ICMP, disco, TSMP, ...
+	Via     string     // Via "direct", "derp-nyc", ...
+	Seconds float64    // Seconds for Dir "in" only
+}
+
 type MapResponse struct {
 	// KeepAlive, if set, represents an empty message just to keep
 	// the connection alive. When true, all other fields except
@@ -899,6 +923,12 @@ type MapResponse struct {
 	// KeepAlive true or false).
 	PingRequest *PingRequest `json:",omitempty"`
 
+	// ControlPingRequest, if non-empty, is a request to the client to ping another node
+	// The IP given in the ControlPingRequest using either TSMP or disco pings.
+	// ControlPingRequest may be sent on any MapResponse (ones with
+	// KeepAlive true or false).
+	ControlPingRequest *ControlPingRequest `json:",omitempty"`
+
 	// Networking
 
 	// Node describes the node making the map request.
diff --git a/tstest/integration/integration_test.go b/tstest/integration/integration_test.go
index 442633856..f5ca90dcf 100644
--- a/tstest/integration/integration_test.go
+++ b/tstest/integration/integration_test.go
@@ -223,6 +223,55 @@ func TestNodeAddressIPFields(t *testing.T) {
 
 	d1.MustCleanShutdown(t)
 }
+func TestControlSelectivePing(t *testing.T) {
+	t.Parallel()
+	bins := BuildTestBinaries(t)
+
+	env := newTestEnv(t, bins)
+	defer env.Close()
+
+	// Create two nodes:
+	n1 := newTestNode(t, env)
+	d1 := n1.StartDaemon(t)
+	defer d1.Kill()
+
+	n2 := newTestNode(t, env)
+	d2 := n2.StartDaemon(t)
+	defer d2.Kill()
+
+	n1.AwaitListening(t)
+	n2.AwaitListening(t)
+	n1.MustUp()
+	n2.MustUp()
+	n1.AwaitRunning(t)
+	n2.AwaitRunning(t)
+
+	// Wait for server to start serveMap
+	if err := tstest.WaitFor(2*time.Second, func() error {
+		env.Control.QueueControlPingRequest()
+		if len(env.Control.PingRequestC) == 0 {
+			return errors.New("failed to add to PingRequestC")
+		}
+		return nil
+	}); err != nil {
+		t.Error(err)
+	}
+
+	// Wait for a MapResponse by Simulating
+	// the time needed for MapResponse method call.
+	if err := tstest.WaitFor(20*time.Second, func() error {
+		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)
+	}
+	d1.MustCleanShutdown(t)
+	d2.MustCleanShutdown(t)
+}
 
 // testEnv contains the test environment (set of servers) used by one
 // or more nodes.
diff --git a/tstest/integration/testcontrol/testcontrol.go b/tstest/integration/testcontrol/testcontrol.go
index 4683cbc0b..6fa905aa3 100644
--- a/tstest/integration/testcontrol/testcontrol.go
+++ b/tstest/integration/testcontrol/testcontrol.go
@@ -36,15 +36,18 @@ import (
 // Server is a control plane server. Its zero value is ready for use.
 // Everything is stored in-memory in one tailnet.
 type Server struct {
-	Logf        logger.Logf      // nil means to use the log package
-	DERPMap     *tailcfg.DERPMap // nil means to use prod DERP map
-	RequireAuth bool
-	BaseURL     string // must be set to e.g. "http://127.0.0.1:1234" with no trailing URL
-	Verbose     bool
+	Logf         logger.Logf      // nil means to use the log package
+	DERPMap      *tailcfg.DERPMap // nil means to use prod DERP map
+	RequireAuth  bool
+	BaseURL      string // must be set to e.g. "http://127.0.0.1:1234" with no trailing URL
+	Verbose      bool
+	PingRequestC chan bool
 
 	initMuxOnce sync.Once
 	mux         *http.ServeMux
 
+	initPRchannelOnce sync.Once
+
 	mu            sync.Mutex
 	pubKey        wgkey.Key
 	privKey       wgkey.Private
@@ -99,10 +102,16 @@ func (s *Server) initMux() {
 	s.mux.HandleFunc("/", s.serveUnhandled)
 	s.mux.HandleFunc("/key", s.serveKey)
 	s.mux.HandleFunc("/machine/", s.serveMachine)
+	s.mux.HandleFunc("/ping", s.servePingInfo)
+}
+
+func (s *Server) initPingRequestC() {
+	s.PingRequestC = make(chan bool, 1)
 }
 
 func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	s.initMuxOnce.Do(s.initMux)
+	s.initPRchannelOnce.Do(s.initPingRequestC)
 	s.mux.ServeHTTP(w, r)
 }
 
@@ -112,6 +121,26 @@ func (s *Server) serveUnhandled(w http.ResponseWriter, r *http.Request) {
 	go panic(fmt.Sprintf("testcontrol.Server received unhandled request: %s", got.Bytes()))
 }
 
+// addControlPingRequest adds a ControlPingRequest pointer
+func (s *Server) addControlPingRequest(res *tailcfg.MapResponse) error {
+	if len(res.Peers) == 0 {
+		return errors.New("MapResponse has no peers to ping")
+	}
+
+	if len(res.Peers[0].Addresses) == 0 {
+		return errors.New("peer has no Addresses")
+	}
+
+	if len(res.Peers[0].AllowedIPs) == 0 {
+		return errors.New("peer has no AllowedIPs")
+	}
+
+	targetNode := res.Peers[0]
+	targetIP := res.Peers[0].AllowedIPs[0].IP()
+	res.ControlPingRequest = &tailcfg.ControlPingRequest{URL: s.BaseURL + "/ping", Peer: targetNode, IP: targetIP, Types: "tsmp"}
+	return nil
+}
+
 func (s *Server) publicKey() wgkey.Key {
 	pub, _ := s.keyPair()
 	return pub
@@ -172,6 +201,29 @@ func (s *Server) serveMachine(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+// servePingInfo TODO, determine the correct response to client
+func (s *Server) servePingInfo(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "text/plain")
+	if r.Method != "PUT" {
+		http.Error(w, "Only PUT requests are supported currently", http.StatusMethodNotAllowed)
+	}
+	_, err := ioutil.ReadAll(r.Body)
+	defer r.Body.Close()
+	if err != nil {
+		http.Error(w, "Failed to read request body", http.StatusInternalServerError)
+	}
+}
+
+// QueueControlPingRequest enqueues a bool to PingRequestC.
+// in serveMap this will result to a ControlPingRequest
+// added to the next MapResponse sent to the client
+func (s *Server) QueueControlPingRequest() {
+	// Redundant check to avoid errors when called multiple times
+	if len(s.PingRequestC) == 0 {
+		s.PingRequestC <- true
+	}
+}
+
 // Node returns the node for nodeKey. It's always nil or cloned memory.
 func (s *Server) Node(nodeKey tailcfg.NodeKey) *tailcfg.Node {
 	s.mu.Lock()
@@ -467,6 +519,12 @@ func (s *Server) serveMap(w http.ResponseWriter, r *http.Request, mkey tailcfg.M
 	w.WriteHeader(200)
 	for {
 		res, err := s.MapResponse(req)
+
+		select {
+		case <-s.PingRequestC:
+			s.addControlPingRequest(res)
+		default:
+		}
 		if err != nil {
 			// TODO: log
 			return