From 2384c112c9a94689ddb5860d867758bebbb793f6 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Tue, 23 Mar 2021 15:16:15 -0700
Subject: [PATCH] net/packet, wgengine/{filter,tstun}: add TSMP ping

Fixes #1467

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 cmd/tailscale/cli/ping.go |  12 ++++-
 ipn/backend.go            |   2 +-
 ipn/fake_test.go          |   2 +-
 ipn/ipnlocal/local.go     |   4 +-
 ipn/ipnstate/ipnstate.go  |  14 ++++--
 ipn/message.go            |  12 +++--
 net/packet/packet.go      |  13 +++++
 net/packet/tsmp.go        |  61 +++++++++++++++++++++++
 wgengine/filter/filter.go |   2 +
 wgengine/tstun/tun.go     |  35 +++++++++++++
 wgengine/userspace.go     | 102 ++++++++++++++++++++++++++++++++++++--
 wgengine/watchdog.go      |   4 +-
 wgengine/wgengine.go      |   2 +-
 13 files changed, 247 insertions(+), 18 deletions(-)

diff --git a/cmd/tailscale/cli/ping.go b/cmd/tailscale/cli/ping.go
index e09a6a75b..c82eb7902 100644
--- a/cmd/tailscale/cli/ping.go
+++ b/cmd/tailscale/cli/ping.go
@@ -48,6 +48,7 @@ relay node.
 		fs := flag.NewFlagSet("ping", flag.ExitOnError)
 		fs.BoolVar(&pingArgs.verbose, "verbose", false, "verbose output")
 		fs.BoolVar(&pingArgs.untilDirect, "until-direct", true, "stop once a direct path is established")
+		fs.BoolVar(&pingArgs.tsmp, "tsmp", false, "do a TSMP-level ping (through IP + wireguard, but not involving host OS stack)")
 		fs.IntVar(&pingArgs.num, "c", 10, "max number of pings to send")
 		fs.DurationVar(&pingArgs.timeout, "timeout", 5*time.Second, "timeout before giving up on a ping")
 		return fs
@@ -58,6 +59,7 @@ var pingArgs struct {
 	num         int
 	untilDirect bool
 	verbose     bool
+	tsmp        bool
 	timeout     time.Duration
 }
 
@@ -120,7 +122,7 @@ func runPing(ctx context.Context, args []string) error {
 	anyPong := false
 	for {
 		n++
-		bc.Ping(ip)
+		bc.Ping(ip, pingArgs.tsmp)
 		timer := time.NewTimer(pingArgs.timeout)
 		select {
 		case <-timer.C:
@@ -135,8 +137,16 @@ func runPing(ctx context.Context, args []string) error {
 			if pr.DERPRegionID != 0 {
 				via = fmt.Sprintf("DERP(%s)", pr.DERPRegionCode)
 			}
+			if pingArgs.tsmp {
+				// TODO(bradfitz): populate the rest of ipnstate.PingResult for TSMP queries?
+				// For now just say it came via TSMP.
+				via = "TSMP"
+			}
 			anyPong = true
 			fmt.Printf("pong from %s (%s) via %v in %v\n", pr.NodeName, pr.NodeIP, via, latency)
+			if pingArgs.tsmp {
+				return nil
+			}
 			if pr.Endpoint != "" && pingArgs.untilDirect {
 				return nil
 			}
diff --git a/ipn/backend.go b/ipn/backend.go
index 5ef54fde1..dee548c54 100644
--- a/ipn/backend.go
+++ b/ipn/backend.go
@@ -165,5 +165,5 @@ type Backend interface {
 	// Ping attempts to start connecting to the given IP and sends a Notify
 	// with its PingResult. If the host is down, there might never
 	// be a PingResult sent. The cmd/tailscale CLI client adds a timeout.
-	Ping(ip string)
+	Ping(ip string, useTSMP bool)
 }
diff --git a/ipn/fake_test.go b/ipn/fake_test.go
index 98685b013..eef580f57 100644
--- a/ipn/fake_test.go
+++ b/ipn/fake_test.go
@@ -91,6 +91,6 @@ func (b *FakeBackend) FakeExpireAfter(x time.Duration) {
 	b.notify(Notify{NetMap: &netmap.NetworkMap{}})
 }
 
-func (b *FakeBackend) Ping(ip string) {
+func (b *FakeBackend) Ping(ip string, useTSMP bool) {
 	b.notify(Notify{PingResult: &ipnstate.PingResult{}})
 }
diff --git a/ipn/ipnlocal/local.go b/ipn/ipnlocal/local.go
index 0d8ef1582..a5bd272a7 100644
--- a/ipn/ipnlocal/local.go
+++ b/ipn/ipnlocal/local.go
@@ -1173,13 +1173,13 @@ func (b *LocalBackend) FakeExpireAfter(x time.Duration) {
 	b.send(ipn.Notify{NetMap: b.netMap})
 }
 
-func (b *LocalBackend) Ping(ipStr string) {
+func (b *LocalBackend) Ping(ipStr string, useTSMP bool) {
 	ip, err := netaddr.ParseIP(ipStr)
 	if err != nil {
 		b.logf("ignoring Ping request to invalid IP %q", ipStr)
 		return
 	}
-	b.e.Ping(ip, func(pr *ipnstate.PingResult) {
+	b.e.Ping(ip, useTSMP, func(pr *ipnstate.PingResult) {
 		b.send(ipn.Notify{PingResult: pr})
 	})
 }
diff --git a/ipn/ipnstate/ipnstate.go b/ipn/ipnstate/ipnstate.go
index a419067ad..a61555760 100644
--- a/ipn/ipnstate/ipnstate.go
+++ b/ipn/ipnstate/ipnstate.go
@@ -407,10 +407,18 @@ type PingResult struct {
 	Err            string
 	LatencySeconds float64
 
-	Endpoint string // ip:port if direct UDP was used
+	// Endpoint is the ip:port if direct UDP was used.
+	// It is not currently set for TSMP pings.
+	Endpoint string
 
-	DERPRegionID   int    // non-zero if DERP was used
-	DERPRegionCode string // three-letter airport/region code if DERP was used
+	// DERPRegionID is non-zero DERP region ID if DERP was used.
+	// It is not currently set for TSMP pings.
+	DERPRegionID int
+
+	// DERPRegionCode is the three-letter region code
+	// corresponding to DERPRegionID.
+	// It is not currently set for TSMP pings.
+	DERPRegionCode string
 
 	// TODO(bradfitz): details like whether port mapping was used on either side? (Once supported)
 }
diff --git a/ipn/message.go b/ipn/message.go
index 96b4708db..905664c93 100644
--- a/ipn/message.go
+++ b/ipn/message.go
@@ -56,7 +56,8 @@ type FakeExpireAfterArgs struct {
 }
 
 type PingArgs struct {
-	IP string
+	IP      string
+	UseTSMP bool
 }
 
 // Command is a command message that is JSON encoded and sent by a
@@ -174,7 +175,7 @@ func (bs *BackendServer) GotCommand(ctx context.Context, cmd *Command) error {
 		bs.b.RequestEngineStatus()
 		return nil
 	} else if c := cmd.Ping; c != nil {
-		bs.b.Ping(c.IP)
+		bs.b.Ping(c.IP, c.UseTSMP)
 		return nil
 	}
 
@@ -320,8 +321,11 @@ func (bc *BackendClient) FakeExpireAfter(x time.Duration) {
 	bc.send(Command{FakeExpireAfter: &FakeExpireAfterArgs{Duration: x}})
 }
 
-func (bc *BackendClient) Ping(ip string) {
-	bc.send(Command{Ping: &PingArgs{IP: ip}})
+func (bc *BackendClient) Ping(ip string, useTSMP bool) {
+	bc.send(Command{Ping: &PingArgs{
+		IP:      ip,
+		UseTSMP: useTSMP,
+	}})
 }
 
 func (bc *BackendClient) SetWantRunning(v bool) {
diff --git a/net/packet/packet.go b/net/packet/packet.go
index 91c92f641..05c4a382f 100644
--- a/net/packet/packet.go
+++ b/net/packet/packet.go
@@ -343,6 +343,19 @@ func (q *Parsed) IP4Header() IP4Header {
 	}
 }
 
+func (q *Parsed) IP6Header() IP6Header {
+	if q.IPVersion != 6 {
+		panic("IP6Header called on non-IPv6 Parsed")
+	}
+	ipid := (binary.BigEndian.Uint32(q.b[:4]) << 12) >> 12
+	return IP6Header{
+		IPID:    ipid,
+		IPProto: q.IPProto,
+		Src:     q.Src.IP,
+		Dst:     q.Dst.IP,
+	}
+}
+
 func (q *Parsed) ICMP4Header() ICMP4Header {
 	if q.IPVersion != 4 {
 		panic("IP4Header called on non-IPv4 Parsed")
diff --git a/net/packet/tsmp.go b/net/packet/tsmp.go
index 8d3c65d56..fb257556c 100644
--- a/net/packet/tsmp.go
+++ b/net/packet/tsmp.go
@@ -70,6 +70,12 @@ type TSMPType uint8
 const (
 	// TSMPTypeRejectedConn is the type byte for a TailscaleRejectedHeader.
 	TSMPTypeRejectedConn TSMPType = '!'
+
+	// TSMPTypePing is the type byte for a TailscalePingRequest.
+	TSMPTypePing TSMPType = 'p'
+
+	// TSMPTypePong is the type byte for a TailscalePongResponse.
+	TSMPTypePong TSMPType = 'o'
 )
 
 type TailscaleRejectReason byte
@@ -195,3 +201,58 @@ func (pp *Parsed) AsTailscaleRejectedHeader() (h TailscaleRejectedHeader, ok boo
 	}
 	return h, true
 }
+
+// TSMPPingRequest is a TSMP message that's like an ICMP ping request.
+//
+// On the wire, after the IP header, it's currently 9 bytes:
+//     * 'p' (TSMPTypePing)
+//     * 8 opaque ping bytes to copy back in the response
+type TSMPPingRequest struct {
+	Data [8]byte
+}
+
+func (pp *Parsed) AsTSMPPing() (h TSMPPingRequest, ok bool) {
+	if pp.IPProto != ipproto.TSMP {
+		return
+	}
+	p := pp.Payload()
+	if len(p) < 9 || p[0] != byte(TSMPTypePing) {
+		return
+	}
+	copy(h.Data[:], p[1:])
+	return h, true
+}
+
+type TSMPPongReply struct {
+	IPHeader Header
+	Data     [8]byte
+}
+
+func (pp *Parsed) AsTSMPPong() (data [8]byte, ok bool) {
+	if pp.IPProto != ipproto.TSMP {
+		return
+	}
+	p := pp.Payload()
+	if len(p) < 9 || p[0] != byte(TSMPTypePong) {
+		return
+	}
+	copy(data[:], p[1:])
+	return data, true
+}
+
+func (h TSMPPongReply) Len() int {
+	return h.IPHeader.Len() + 9
+}
+
+func (h TSMPPongReply) Marshal(buf []byte) error {
+	if len(buf) < h.Len() {
+		return errSmallBuffer
+	}
+	if err := h.IPHeader.Marshal(buf); err != nil {
+		return err
+	}
+	buf = buf[h.IPHeader.Len():]
+	buf[0] = byte(TSMPTypePong)
+	copy(buf[1:], h.Data[:])
+	return nil
+}
diff --git a/wgengine/filter/filter.go b/wgengine/filter/filter.go
index e904b3905..3c4964c34 100644
--- a/wgengine/filter/filter.go
+++ b/wgengine/filter/filter.go
@@ -423,6 +423,8 @@ func (f *Filter) runIn6(q *packet.Parsed) (r Response, why string) {
 		if f.matches6.match(q) {
 			return Accept, "ok"
 		}
+	case ipproto.TSMP:
+		return Accept, "tsmp ok"
 	default:
 		return Drop, "Unknown proto"
 	}
diff --git a/wgengine/tstun/tun.go b/wgengine/tstun/tun.go
index 367c85829..0f0725eb3 100644
--- a/wgengine/tstun/tun.go
+++ b/wgengine/tstun/tun.go
@@ -109,6 +109,9 @@ type TUN struct {
 	// PostFilterOut is the outbound filter function that runs after the main filter.
 	PostFilterOut FilterFunc
 
+	// OnTSMPPongReceived, if non-nil, is called whenever a TSMP pong arrives.
+	OnTSMPPongReceived func(data [8]byte)
+
 	// disableFilter disables all filtering when set. This should only be used in tests.
 	disableFilter bool
 }
@@ -323,6 +326,18 @@ func (t *TUN) filterIn(buf []byte) filter.Response {
 	defer parsedPacketPool.Put(p)
 	p.Decode(buf)
 
+	if p.IPProto == ipproto.TSMP {
+		if pingReq, ok := p.AsTSMPPing(); ok {
+			t.noteActivity()
+			t.injectOutboundPong(p, pingReq)
+			return filter.DropSilently
+		} else if data, ok := p.AsTSMPPong(); ok {
+			if f := t.OnTSMPPongReceived; f != nil {
+				f(data)
+			}
+		}
+	}
+
 	if t.PreFilterIn != nil {
 		if res := t.PreFilterIn(p, t); res.IsDrop() {
 			return res
@@ -440,6 +455,26 @@ func (t *TUN) InjectInboundCopy(packet []byte) error {
 	return t.InjectInboundDirect(buf, PacketStartOffset)
 }
 
+func (t *TUN) injectOutboundPong(pp *packet.Parsed, req packet.TSMPPingRequest) {
+	pong := packet.TSMPPongReply{
+		Data: req.Data,
+	}
+	switch pp.IPVersion {
+	case 4:
+		h4 := pp.IP4Header()
+		h4.ToResponse()
+		pong.IPHeader = h4
+	case 6:
+		h6 := pp.IP6Header()
+		h6.ToResponse()
+		pong.IPHeader = h6
+	default:
+		return
+	}
+
+	t.InjectOutbound(packet.Generate(pong, nil))
+}
+
 // InjectOutbound makes the TUN device behave as if a packet
 // with the given contents was sent to the network.
 // It does not block, but takes ownership of the packet.
diff --git a/wgengine/userspace.go b/wgengine/userspace.go
index 49bc617e0..09347db1a 100644
--- a/wgengine/userspace.go
+++ b/wgengine/userspace.go
@@ -8,6 +8,7 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	crand "crypto/rand"
 	"errors"
 	"fmt"
 	"io"
@@ -128,6 +129,7 @@ type userspaceEngine struct {
 	pendOpen            map[flowtrack.Tuple]*pendingOpenFlow // see pendopen.go
 	networkMapCallbacks map[*someHandle]NetworkMapCallback
 	tsIPByIPPort        map[netaddr.IPPort]netaddr.IP // allows registration of IP:ports as belonging to a certain Tailscale IP for whois lookups
+	pongCallback        map[[8]byte]func()            // for TSMP pong responses
 
 	// Lock ordering: magicsock.Conn.mu, wgLock, then mu.
 }
@@ -351,6 +353,16 @@ func newUserspaceEngine(logf logger.Logf, rawTUNDev tun.Device, conf Config) (_
 		SkipBindUpdate: true,
 	}
 
+	e.tundev.OnTSMPPongReceived = func(data [8]byte) {
+		e.mu.Lock()
+		defer e.mu.Unlock()
+		cb := e.pongCallback[data]
+		e.logf("wgengine: got TSMP pong %02x; cb=%v", data, cb != nil)
+		if cb != nil {
+			go cb()
+		}
+	}
+
 	// wgdev takes ownership of tundev, will close it when closed.
 	e.logf("Creating wireguard device...")
 	e.wgdev = device.NewDevice(e.tundev, opts)
@@ -1342,7 +1354,7 @@ func (e *userspaceEngine) UpdateStatus(sb *ipnstate.StatusBuilder) {
 	e.magicConn.UpdateStatus(sb)
 }
 
-func (e *userspaceEngine) Ping(ip netaddr.IP, cb func(*ipnstate.PingResult)) {
+func (e *userspaceEngine) Ping(ip netaddr.IP, useTSMP bool, cb func(*ipnstate.PingResult)) {
 	res := &ipnstate.PingResult{IP: ip.String()}
 	peer, err := e.peerForIP(ip)
 	if err != nil {
@@ -1357,8 +1369,92 @@ func (e *userspaceEngine) Ping(ip netaddr.IP, cb func(*ipnstate.PingResult)) {
 		cb(res)
 		return
 	}
-	e.logf("ping(%v): sending ping to %v %v ...", ip, peer.Key.ShortString(), peer.ComputedName)
-	e.magicConn.Ping(peer, res, cb)
+	pingType := "disco"
+	if useTSMP {
+		pingType = "TSMP"
+	}
+	e.logf("ping(%v): sending %v ping to %v %v ...", ip, pingType, peer.Key.ShortString(), peer.ComputedName)
+	if useTSMP {
+		e.sendTSMPPing(ip, peer, res, cb)
+	} else {
+		e.magicConn.Ping(peer, res, cb)
+	}
+}
+
+func (e *userspaceEngine) mySelfIPMatchingFamily(dst netaddr.IP) (src netaddr.IP, err error) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	if e.netMap == nil {
+		return netaddr.IP{}, errors.New("no netmap")
+	}
+	for _, a := range e.netMap.Addresses {
+		if a.IsSingleIP() && a.IP.BitLen() == dst.BitLen() {
+			return a.IP, nil
+		}
+	}
+	if len(e.netMap.Addresses) == 0 {
+		return netaddr.IP{}, errors.New("no self address in netmap")
+	}
+	return netaddr.IP{}, errors.New("no self address in netmap matching address family")
+}
+
+func (e *userspaceEngine) sendTSMPPing(ip netaddr.IP, peer *tailcfg.Node, res *ipnstate.PingResult, cb func(*ipnstate.PingResult)) {
+	srcIP, err := e.mySelfIPMatchingFamily(ip)
+	if err != nil {
+		res.Err = err.Error()
+		cb(res)
+		return
+	}
+	var iph packet.Header
+	if srcIP.Is4() {
+		iph = packet.IP4Header{
+			IPProto: ipproto.TSMP,
+			Src:     srcIP,
+			Dst:     ip,
+		}
+	} else {
+		iph = packet.IP6Header{
+			IPProto: ipproto.TSMP,
+			Src:     srcIP,
+			Dst:     ip,
+		}
+	}
+
+	var data [8]byte
+	crand.Read(data[:])
+
+	expireTimer := time.AfterFunc(10*time.Second, func() {
+		e.setTSMPPongCallback(data, nil)
+	})
+	t0 := time.Now()
+	e.setTSMPPongCallback(data, func() {
+		expireTimer.Stop()
+		d := time.Since(t0)
+		res.LatencySeconds = d.Seconds()
+		res.NodeIP = ip.String()
+		res.NodeName = peer.ComputedName
+		cb(res)
+	})
+
+	var tsmpPayload [9]byte
+	tsmpPayload[0] = byte(packet.TSMPTypePing)
+	copy(tsmpPayload[1:], data[:])
+
+	tsmpPing := packet.Generate(iph, tsmpPayload[:])
+	e.tundev.InjectOutbound(tsmpPing)
+}
+
+func (e *userspaceEngine) setTSMPPongCallback(data [8]byte, cb func()) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	if e.pongCallback == nil {
+		e.pongCallback = map[[8]byte]func(){}
+	}
+	if cb == nil {
+		delete(e.pongCallback, data)
+	} else {
+		e.pongCallback[data] = cb
+	}
 }
 
 func (e *userspaceEngine) RegisterIPPortIdentity(ipport netaddr.IPPort, tsIP netaddr.IP) {
diff --git a/wgengine/watchdog.go b/wgengine/watchdog.go
index f4f7d3085..6a607ce03 100644
--- a/wgengine/watchdog.go
+++ b/wgengine/watchdog.go
@@ -117,8 +117,8 @@ func (e *watchdogEngine) DiscoPublicKey() (k tailcfg.DiscoKey) {
 	e.watchdog("DiscoPublicKey", func() { k = e.wrap.DiscoPublicKey() })
 	return k
 }
-func (e *watchdogEngine) Ping(ip netaddr.IP, cb func(*ipnstate.PingResult)) {
-	e.watchdog("Ping", func() { e.wrap.Ping(ip, cb) })
+func (e *watchdogEngine) Ping(ip netaddr.IP, useTSMP bool, cb func(*ipnstate.PingResult)) {
+	e.watchdog("Ping", func() { e.wrap.Ping(ip, useTSMP, cb) })
 }
 func (e *watchdogEngine) RegisterIPPortIdentity(ipp netaddr.IPPort, tsIP netaddr.IP) {
 	e.watchdog("RegisterIPPortIdentity", func() { e.wrap.RegisterIPPortIdentity(ipp, tsIP) })
diff --git a/wgengine/wgengine.go b/wgengine/wgengine.go
index c8e7963db..2e3a8fd36 100644
--- a/wgengine/wgengine.go
+++ b/wgengine/wgengine.go
@@ -136,7 +136,7 @@ type Engine interface {
 
 	// Ping is a request to start a discovery ping with the peer handling
 	// the given IP and then call cb with its ping latency & method.
-	Ping(ip netaddr.IP, cb func(*ipnstate.PingResult))
+	Ping(ip netaddr.IP, useTSMP bool, cb func(*ipnstate.PingResult))
 
 	// RegisterIPPortIdentity registers a given node (identified by its
 	// Tailscale IP) as temporarily having the given IP:port for whois lookups.