wgengine: configure wireguard peers lazily, as needed

wireguard-go uses 3 goroutines per peer (with reasonably large stacks & buffers). Rather than tell wireguard-go about all our peers, only tell it about peers we're actively communicating with. That means we need hooks into magicsock's packet receiving path and tstun's packet sending path to lazily create a wireguard peer on demand from the network map. This frees up lots of memory for iOS (where we have almost nothing left for larger domains with many users). We should ideally do this in wireguard-go itself one day, but that'd be a pretty big change. Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
2025-12-08 22:12:21 +00:00 · 2020-07-23 15:15:28 -07:00
parent 5066b824a6
commit 16a9cfe2f4
4 changed files with 410 additions and 78 deletions
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -58,14 +58,15 @@ import (
 // A Conn routes UDP packets and actively manages a list of its endpoints.
 // It implements wireguard/conn.Bind.
 type Conn struct {
-	pconnPort    uint16 // the preferred port from opts.Port; 0 means auto
-	pconn4       *RebindingUDPConn
-	pconn6       *RebindingUDPConn // non-nil if IPv6 available
-	epFunc       func(endpoints []string)
-	logf         logger.Logf
-	sendLogLimit *rate.Limiter
-	netChecker   *netcheck.Client
-	idleFunc     func() time.Duration // nil means unknown
+	pconnPort        uint16 // the preferred port from opts.Port; 0 means auto
+	pconn4           *RebindingUDPConn
+	pconn6           *RebindingUDPConn // non-nil if IPv6 available
+	epFunc           func(endpoints []string)
+	logf             logger.Logf
+	sendLogLimit     *rate.Limiter
+	netChecker       *netcheck.Client
+	idleFunc         func() time.Duration   // nil means unknown
+	noteRecvActivity func(tailcfg.DiscoKey) // or nil, see Options.NoteRecvActivity

 	// bufferedIPv4From and bufferedIPv4Packet are owned by
 	// ReceiveIPv4, and used when both a DERP and IPv4 packet arrive
@@ -89,6 +90,13 @@ type Conn struct {
 	// ============================================================
 	mu sync.Mutex // guards all following fields

+	// canCreateEPUnlocked tracks at one place whether mu is
+	// already held. It's then checked in CreateEndpoint to avoid
+	// double-locking mu and thus deadlocking. mu should be held
+	// while setting this; but can be read without mu held.
+	// TODO(bradfitz): delete this shameful hack; refactor the one use
+	canCreateEPUnlocked syncs.AtomicBool
+
 	started bool // Start was called
 	closed  bool // Close was called

@@ -104,8 +112,8 @@ type Conn struct {
 	nodeOfDisco     map[tailcfg.DiscoKey]*tailcfg.Node
 	discoOfNode     map[tailcfg.NodeKey]tailcfg.DiscoKey
 	discoOfAddr     map[netaddr.IPPort]tailcfg.DiscoKey // validated non-DERP paths only
-	endpointOfDisco map[tailcfg.DiscoKey]*discoEndpoint
-	sharedDiscoKey  map[tailcfg.DiscoKey]*[32]byte // nacl/box precomputed key
+	endpointOfDisco map[tailcfg.DiscoKey]*discoEndpoint // those with activity only
+	sharedDiscoKey  map[tailcfg.DiscoKey]*[32]byte      // nacl/box precomputed key

 	// addrsByUDP is a map of every remote ip:port to a priority
 	// list of endpoint addresses for a peer.
@@ -235,6 +243,17 @@ type Options struct {
 	// PacketListener optionally specifies how to create PacketConns.
 	// It's meant for testing.
 	PacketListener nettype.PacketListener
+
+	// NoteRecvActivity, if provided, is a func for magicsock to
+	// call whenever it receives a packet from a a
+	// discovery-capable peer if it's been more than ~10 seconds
+	// since the last one. (10 seconds is somewhat arbitrary; the
+	// sole user just doesn't need or want it called on every
+	// packet, just every minute or two for Wireguard timeouts,
+	// and 10 seconds seems like a good trade-off between often
+	// enough and not too often.) The provided func likely calls
+	// Conn.CreateEndpoint, which acquires Conn.mu.
+	NoteRecvActivity func(tailcfg.DiscoKey)
 }

 func (o *Options) logf() logger.Logf {
@@ -282,6 +301,7 @@ func NewConn(opts Options) (*Conn, error) {
 	c.epFunc = opts.endpointsFunc()
 	c.idleFunc = opts.IdleFunc
 	c.packetListener = opts.PacketListener
+	c.noteRecvActivity = opts.NoteRecvActivity

 	if err := c.initialBind(); err != nil {
 		return nil, err
@@ -1300,6 +1320,16 @@ func wgRecvAddr(e conn.Endpoint, ipp netaddr.IPPort, addr *net.UDPAddr) *net.UDP
 	return ipp.UDPAddr()
 }

+// noteRecvActivity calls the magicsock.Conn.noteRecvActivity hook if
+// e is a discovery-capable peer.
+//
+// This should be called whenever a packet arrives from e.
+func noteRecvActivity(e conn.Endpoint) {
+	if de, ok := e.(*discoEndpoint); ok {
+		de.onRecvActivity()
+	}
+}
+
 func (c *Conn) ReceiveIPv4(b []byte) (n int, ep conn.Endpoint, addr *net.UDPAddr, err error) {
 Top:
 	// First, process any buffered packet from earlier.
@@ -1307,6 +1337,7 @@ Top:
 		c.bufferedIPv4From = netaddr.IPPort{}
 		addr = from.UDPAddr()
 		ep := c.findEndpoint(from, addr)
+		noteRecvActivity(ep)
 		return copy(b, c.bufferedIPv4Packet), ep, wgRecvAddr(ep, from, addr), nil
 	}

@@ -1319,6 +1350,7 @@ Top:
 	var addrSet *AddrSet
 	var discoEp *discoEndpoint
 	var ipp netaddr.IPPort
+	var didNoteRecvActivity bool

 	select {
 	case dm := <-c.derpRecvCh:
@@ -1360,6 +1392,24 @@ Top:
 		c.mu.Lock()
 		if dk, ok := c.discoOfNode[tailcfg.NodeKey(dm.src)]; ok {
 			discoEp = c.endpointOfDisco[dk]
+			// If we know about the node (it's in discoOfNode) but don't know about the
+			// endpoint, that's because it's an idle peer that doesn't yet exist in the
+			// wireguard config. So run the receive hook, if defined, which should
+			// create the wireguard peer.
+			if discoEp == nil && c.noteRecvActivity != nil {
+				didNoteRecvActivity = true
+				c.mu.Unlock()          // release lock before calling noteRecvActivity
+				c.noteRecvActivity(dk) // (calls back into CreateEndpoint)
+				// Now require the lock. No invariants need to be rechecked; just
+				// 1-2 map lookups follow that are harmless if, say, the peer has
+				// been deleted during this time.  In that case we'll treate it as a
+				// legacy pre-disco UDP receive and hand it to wireguard which'll
+				// likely just drop it.
+				c.mu.Lock()
+
+				discoEp = c.endpointOfDisco[dk]
+				c.logf("magicsock: DERP packet received from idle peer %v; created=%v", dm.src.ShortString(), discoEp != nil)
+			}
 		}
 		if discoEp == nil {
 			addrSet = c.addrsByKey[dm.src]
@@ -1398,6 +1448,9 @@ Top:
 	} else {
 		ep = c.findEndpoint(ipp, addr)
 	}
+	if !didNoteRecvActivity {
+		noteRecvActivity(ep)
+	}
 	return n, ep, wgRecvAddr(ep, ipp, addr), nil
 }

@@ -1424,6 +1477,7 @@ func (c *Conn) ReceiveIPv6(b []byte) (int, conn.Endpoint, *net.UDPAddr, error) {
 		}

 		ep := c.findEndpoint(ipp, addr)
+		noteRecvActivity(ep)
 		return n, ep, wgRecvAddr(ep, ipp, addr), nil
 	}
 }
@@ -1440,7 +1494,7 @@ const (
 	discoVerboseLog
 )

-func (c *Conn) sendDiscoMessage(dst netaddr.IPPort, dstKey key.Public, dstDisco tailcfg.DiscoKey, m disco.Message, logLevel discoLogLevel) (sent bool, err error) {
+func (c *Conn) sendDiscoMessage(dst netaddr.IPPort, dstKey tailcfg.NodeKey, dstDisco tailcfg.DiscoKey, m disco.Message, logLevel discoLogLevel) (sent bool, err error) {
 	c.mu.Lock()
 	if c.closed {
 		c.mu.Unlock()
@@ -1458,7 +1512,7 @@ func (c *Conn) sendDiscoMessage(dst netaddr.IPPort, dstKey key.Public, dstDisco
 	c.mu.Unlock()

 	pkt = box.SealAfterPrecomputation(pkt, m.AppendMarshal(nil), &nonce, sharedKey)
-	sent, err = c.sendAddr(dst, dstKey, pkt)
+	sent, err = c.sendAddr(dst, key.Public(dstKey), pkt)
 	if sent {
 		if logLevel == discoLog || (logLevel == discoVerboseLog && debugDisco) {
 			c.logf("magicsock: disco: %v->%v (%v, %v) sent %v", c.discoShort, dstDisco.ShortString(), dstKey.ShortString(), derpStr(dst.String()), disco.MessageSummary(m))
@@ -1507,16 +1561,45 @@ func (c *Conn) handleDiscoMessage(msg []byte, src netaddr.IPPort) bool {
 		return false
 	}

-	de, ok := c.endpointOfDisco[sender]
+	peerNode, ok := c.nodeOfDisco[sender]
 	if !ok {
 		if debugDisco {
-			c.logf("magicsock: disco: ignoring disco-looking frame, don't know about %v", sender.ShortString())
+			c.logf("magicsock: disco: ignoring disco-looking frame, don't know node for %v", sender.ShortString())
 		}
 		// Returning false keeps passing it down, to WireGuard.
 		// WireGuard will almost surely reject it, but give it a chance.
 		return false
 	}

+	de, ok := c.endpointOfDisco[sender]
+	if !ok {
+		// We don't have an active endpoint for this sender but we knew about the node, so
+		// it's an idle endpoint that doesn't yet exist in the wireguard config. We now have
+		// to notify the userspace engine (via noteRecvActivity) so wireguard-go can create
+		// an Endpoint (ultimately calling our CreateEndpoint).
+		if debugDisco {
+			c.logf("magicsock: disco: got message from inactive peer %v", sender.ShortString())
+		}
+		if c.noteRecvActivity == nil {
+			c.logf("magicsock: [unexpected] have node without endpoint, without c.noteRecvActivity hook")
+			return false
+		}
+		// noteRecvActivity calls back into CreateEndpoint, which we can't easily control,
+		// and CreateEndpoint expects to be called with c.mu held, but we hold it here, and
+		// it's too invasive for now to release it here and recheck invariants.  So instead,
+		// use this unfortunate hack: set canCreateEPUnlocked which CreateEndpoint then
+		// checks to conditionally acquire the mutex. I'm so sorry.
+		c.canCreateEPUnlocked.Set(true)
+		c.noteRecvActivity(sender)
+		c.canCreateEPUnlocked.Set(false)
+		de, ok = c.endpointOfDisco[sender]
+		if !ok {
+			c.logf("magicsock: [unexpected] lazy endpoint not created for %v, %v", peerNode.Key.ShortString(), sender.ShortString())
+			return false
+		}
+		c.logf("magicsock: lazy endpoint created via disco message for %v, %v", peerNode.Key.ShortString(), sender.ShortString())
+	}
+
 	// First, do we even know (and thus care) about this sender? If not,
 	// don't bother decrypting it.

@@ -1556,8 +1639,11 @@ func (c *Conn) handleDiscoMessage(msg []byte, src netaddr.IPPort) bool {

 	switch dm := dm.(type) {
 	case *disco.Ping:
-		c.handlePingLocked(dm, de, src)
+		c.handlePingLocked(dm, de, src, sender, peerNode)
 	case *disco.Pong:
+		if de == nil {
+			return true
+		}
 		de.handlePongConnLocked(dm, src)
 	case disco.CallMeMaybe:
 		if src.IP != derpMagicIPAddr {
@@ -1565,26 +1651,40 @@ func (c *Conn) handleDiscoMessage(msg []byte, src netaddr.IPPort) bool {
 			c.logf("[unexpected] CallMeMaybe packets should only come via DERP")
 			return true
 		}
-		c.logf("magicsock: disco: %v<-%v (%v, %v)  got call-me-maybe", c.discoShort, de.discoShort, de.publicKey.ShortString(), derpStr(src.String()))
-		go de.handleCallMeMaybe()
+		if de != nil {
+			c.logf("magicsock: disco: %v<-%v (%v, %v)  got call-me-maybe", c.discoShort, de.discoShort, de.publicKey.ShortString(), derpStr(src.String()))
+			go de.handleCallMeMaybe()
+		}
 	}

 	return true
 }

-func (c *Conn) handlePingLocked(dm *disco.Ping, de *discoEndpoint, src netaddr.IPPort) {
-	likelyHeartBeat := src == de.lastPingFrom && time.Since(de.lastPingTime) < 5*time.Second
-	de.lastPingFrom = src
-	de.lastPingTime = time.Now()
+// de may be nil
+func (c *Conn) handlePingLocked(dm *disco.Ping, de *discoEndpoint, src netaddr.IPPort, sender tailcfg.DiscoKey, peerNode *tailcfg.Node) {
+	if peerNode == nil {
+		c.logf("magicsock: disco: [unexpected] ignoring ping from unknown peer Node")
+		return
+	}
+	likelyHeartBeat := de != nil && src == de.lastPingFrom && time.Since(de.lastPingTime) < 5*time.Second
+	var discoShort string
+	if de != nil {
+		discoShort = de.discoShort
+		de.lastPingFrom = src
+		de.lastPingTime = time.Now()
+	} else {
+		discoShort = sender.ShortString()
+	}
 	if !likelyHeartBeat || debugDisco {
-		c.logf("magicsock: disco: %v<-%v (%v, %v)  got ping tx=%x", c.discoShort, de.discoShort, de.publicKey.ShortString(), src, dm.TxID[:6])
+		c.logf("magicsock: disco: %v<-%v (%v, %v)  got ping tx=%x", c.discoShort, discoShort, peerNode.Key.ShortString(), src, dm.TxID[:6])
 	}

 	// Remember this route if not present.
-	c.setAddrToDiscoLocked(src, de.discoKey, nil)
+	c.setAddrToDiscoLocked(src, sender, nil)

-	pongDst := src
-	go de.sendDiscoMessage(pongDst, &disco.Pong{
+	ipDst := src
+	discoDest := sender
+	go c.sendDiscoMessage(ipDst, peerNode.Key, discoDest, &disco.Pong{
 		TxID: dm.TxID,
 		Src:  src,
 	}, discoVerboseLog)
@@ -2455,17 +2555,30 @@ func (c *Conn) CreateEndpoint(pubKey [32]byte, addrs string) (conn.Endpoint, err
 		if err != nil {
 			return nil, fmt.Errorf("magicsock: invalid discokey endpoint %q for %v: %w", addrs, pk.ShortString(), err)
 		}
-		c.mu.Lock()
-		defer c.mu.Unlock()
+		if !c.canCreateEPUnlocked.Get() { // sorry
+			c.mu.Lock()
+			defer c.mu.Unlock()
+		}
 		de := &discoEndpoint{
 			c:                  c,
-			publicKey:          pk,                         // peer public key (for WireGuard + DERP)
+			publicKey:          tailcfg.NodeKey(pk),        // peer public key (for WireGuard + DERP)
 			discoKey:           tailcfg.DiscoKey(discoKey), // for discovery mesages
 			discoShort:         tailcfg.DiscoKey(discoKey).ShortString(),
 			wgEndpointHostPort: addrs,
 			sentPing:           map[stun.TxID]sentPing{},
 			endpointState:      map[netaddr.IPPort]*endpointState{},
 		}
+		lastRecvTime := new(int64) // atomic
+		de.onRecvActivity = func() {
+			now := time.Now().Unix()
+			old := atomic.LoadInt64(lastRecvTime)
+			if old == 0 || old <= now-10 {
+				atomic.StoreInt64(lastRecvTime, now)
+				if c.noteRecvActivity != nil {
+					c.noteRecvActivity(de.discoKey)
+				}
+			}
+		}
 		de.initFakeUDPAddr()
 		de.updateFromNode(c.nodeOfDisco[de.discoKey])
 		c.endpointOfDisco[de.discoKey] = de
@@ -2694,14 +2807,14 @@ func (c *Conn) UpdateStatus(sb *ipnstate.StatusBuilder) {
 	c.mu.Lock()
 	defer c.mu.Unlock()

-	for dk, de := range c.endpointOfDisco {
+	for dk, n := range c.nodeOfDisco {
 		ps := &ipnstate.PeerStatus{InMagicSock: true}
-		if node, ok := c.nodeOfDisco[dk]; ok {
-			ps.Addrs = append(ps.Addrs, node.Endpoints...)
-			ps.Relay = c.derpRegionCodeOfAddrLocked(node.DERP)
+		ps.Addrs = append(ps.Addrs, n.Endpoints...)
+		ps.Relay = c.derpRegionCodeOfAddrLocked(n.DERP)
+		if de, ok := c.endpointOfDisco[dk]; ok {
+			de.populatePeerStatus(ps)
 		}
-		de.populatePeerStatus(ps)
-		sb.AddPeer(de.publicKey, ps)
+		sb.AddPeer(key.Public(n.Key), ps)
 	}
 	// Old-style (pre-disco) peers:
 	for k, as := range c.addrsByKey {
@@ -2731,12 +2844,13 @@ func udpAddrDebugString(ua net.UDPAddr) string {
 type discoEndpoint struct {
 	// These fields are initialized once and never modified.
 	c                  *Conn
-	publicKey          key.Public       // peer public key (for WireGuard + DERP)
+	publicKey          tailcfg.NodeKey  // peer public key (for WireGuard + DERP)
 	discoKey           tailcfg.DiscoKey // for discovery mesages
 	discoShort         string           // ShortString of discoKey
 	fakeWGAddr         netaddr.IPPort   // the UDP address we tell wireguard-go we're using
 	fakeWGAddrStd      *net.UDPAddr     // the *net.UDPAddr form of fakeWGAddr
 	wgEndpointHostPort string           // string from CreateEndpoint: "<hex-discovery-key>.disco.tailscale:12345"
+	onRecvActivity     func()

 	// Owned by Conn.mu:
 	lastPingFrom netaddr.IPPort
@@ -2958,10 +3072,10 @@ func (de *discoEndpoint) send(b []byte) error {
 	}
 	var err error
 	if !udpAddr.IsZero() {
-		_, err = de.c.sendAddr(udpAddr, de.publicKey, b)
+		_, err = de.c.sendAddr(udpAddr, key.Public(de.publicKey), b)
 	}
 	if !derpAddr.IsZero() {
-		if ok, _ := de.c.sendAddr(derpAddr, de.publicKey, b); ok && err != nil {
+		if ok, _ := de.c.sendAddr(derpAddr, key.Public(de.publicKey), b); ok && err != nil {
 			// UDP failed but DERP worked, so good enough:
 			return nil
 		}