watch for conn changes

ethanndickson · ethanndickson · commit 0bebad47aec9 · 2024-07-09T07:23:27.000Z
diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -264,6 +264,7 @@ func NewConn(options *Options) (conn *Conn, err error) {
 	nodeUp.setAddresses(options.Addresses)
 	nodeUp.setBlockEndpoints(options.BlockEndpoints)
 
+	ctx, ctxCancel := context.WithCancel(context.Background())
 	server := &Conn{
 		id:               uuid.New(),
 		closed:           make(chan struct{}),
@@ -283,6 +284,8 @@ func NewConn(options *Options) (conn *Conn, err error) {
 		telemetrySink:   options.TelemetrySink,
 		telemetryStore:  telemetryStore,
 		createdAt:       time.Now(),
+		watchCtx:        ctx,
+		watchCancel:     ctxCancel,
 	}
 	defer func() {
 		if err != nil {
@@ -293,8 +296,17 @@ func NewConn(options *Options) (conn *Conn, err error) {
 		server.wireguardEngine.SetNetInfoCallback(func(ni *tailcfg.NetInfo) {
 			server.telemetryStore.setNetInfo(ni)
 			nodeUp.setNetInfo(ni)
+			if server.telemetryStore.connectedIP != nil {
+				_, _, _, _ = server.Ping(ctx, *server.telemetryStore.connectedIP)
+			}
 		})
-		server.wireguardEngine.AddNetworkMapCallback(server.networkMapCallback)
+		server.wireguardEngine.AddNetworkMapCallback(func(nm *netmap.NetworkMap) {
+			server.telemetryStore.updateNetworkMap(nm)
+			if server.telemetryStore.connectedIP != nil {
+				_, _, _, _ = server.Ping(ctx, *server.telemetryStore.connectedIP)
+			}
+		})
+		go server.watchConnChange()
 	} else {
 		server.wireguardEngine.SetNetInfoCallback(nodeUp.setNetInfo)
 	}
@@ -361,6 +373,9 @@ type Conn struct {
 	telemetryStore *TelemetryStore
 	telemetryWg    sync.WaitGroup
 
+	watchCtx    context.Context
+	watchCancel func()
+
 	trafficStats *connstats.Statistics
 }
 
@@ -542,6 +557,7 @@ func (c *Conn) Closed() <-chan struct{} {
 // Close shuts down the Wireguard connection.
 func (c *Conn) Close() error {
 	c.logger.Info(context.Background(), "closing tailnet Conn")
+	c.watchCancel()
 	c.telemetryWg.Wait()
 	c.configMaps.close()
 	c.nodeUpdater.close()
@@ -771,13 +787,6 @@ func (c *Conn) newTelemetryEvent() *proto.TelemetryEvent {
 	return event
 }
 
-func (c *Conn) networkMapCallback(nm *netmap.NetworkMap) {
-	c.telemetryStore.updateNetworkMap(nm)
-	if c.telemetryStore.connectedIP != nil {
-		_, _, _, _ = c.Ping(context.Background(), *c.telemetryStore.connectedIP)
-	}
-}
-
 func (c *Conn) sendTelemetryBackground(e *proto.TelemetryEvent) {
 	c.telemetryWg.Add(1)
 	go func() {
@@ -786,6 +795,32 @@ func (c *Conn) sendTelemetryBackground(e *proto.TelemetryEvent) {
 	}()
 }
 
+// Watch for changes in the connection type (P2P<->DERP) and send telemetry events.
+func (c *Conn) watchConnChange() {
+	ticker := time.NewTicker(time.Millisecond * 50)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-c.watchCtx.Done():
+			return
+		case <-ticker.C:
+		}
+		status := c.Status()
+		peers := status.Peers()
+		if len(peers) > 1 {
+			// Not a CLI<->agent connection, stop watching
+			return
+		} else if len(peers) == 0 {
+			continue
+		}
+		peer := status.Peer[peers[0]]
+		// If the connection type has changed, send a telemetry event with the latest ping stats
+		if c.telemetryStore.checkConnType(peer.Relay) && c.telemetryStore.connectedIP != nil {
+			_, _, _, _ = c.Ping(c.watchCtx, *c.telemetryStore.connectedIP)
+		}
+	}
+}
+
 // PeerDiagnostics is a checklist of human-readable conditions necessary to establish an encrypted
 // tunnel to a peer via a Conn
 type PeerDiagnostics struct {
diff --git a/tailnet/telemetry.go b/tailnet/telemetry.go
@@ -42,6 +42,7 @@ type TelemetryStore struct {
 	connectedIP   *netip.Addr
 	// 0 if not connected
 	nodeIDRemote uint64
+	p2p          bool
 }
 
 func newTelemetryStore() (*TelemetryStore, error) {
@@ -84,6 +85,23 @@ func (b *TelemetryStore) markConnected(ip *netip.Addr, connCreatedAt time.Time,
 	b.application = application
 }
 
+// Return whether we've changed to/from a P2P connection
+func (b *TelemetryStore) checkConnType(relay string) bool {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	if b.p2p && relay == "" {
+		return false
+	} else if !b.p2p && relay == "" {
+		b.p2p = true
+		return true
+	} else if b.p2p && relay != "" {
+		b.p2p = false
+		return true
+	}
+	return false
+}
+
 func (b *TelemetryStore) updateRemoteNodeIDLocked(nm *netmap.NetworkMap) {
 	if b.connectedIP == nil {
 		return
