wgengine/netstack: use tailscale IPs instead of a hardcoded one

Signed-off-by: Naman Sood <mail@nsood.in>

wgengine/netstack/netstack.go

@@ -28,6 +28,7 @@
 	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
 	"gvisor.dev/gvisor/pkg/waiter"
 	"inet.af/netaddr"
+	"tailscale.com/control/controlclient"
 	"tailscale.com/net/packet"
 	"tailscale.com/types/logger"
 	"tailscale.com/wgengine"
@@ -62,7 +63,46 @@ func Impl(logf logger.Logf, tundev *tstun.TUN, e wgengine.Engine, mc *magicsock.
 		log.Fatal(err)
 	}
 
-	ipstack.AddAddress(nicID, ipv4.ProtocolNumber, tcpip.Address(net.ParseIP("100.96.188.101").To4()))
+	e.AddNetworkMapCallback(func(nm *controlclient.NetworkMap) {
+		oldIPs := make(map[tcpip.Address]bool)
+		for _, ip := range ipstack.AllAddresses()[nicID] {
+			oldIPs[ip.AddressWithPrefix.Address] = true
+		}
+		newIPs := make(map[tcpip.Address]bool)
+		for _, ip := range nm.Addresses {
+			newIPs[tcpip.Address(ip.IPNet().IP)] = true
+		}
+
+		ipsToBeAdded := make(map[tcpip.Address]bool)
+		for ip := range newIPs {
+			if !oldIPs[ip] {
+				ipsToBeAdded[ip] = true
+			}
+		}
+		ipsToBeRemoved := make(map[tcpip.Address]bool)
+		for ip := range oldIPs {
+			if !newIPs[ip] {
+				ipsToBeRemoved[ip] = true
+			}
+		}
+
+		for ip := range ipsToBeRemoved {
+			err := ipstack.RemoveAddress(nicID, ip)
+			if err != nil {
+				logf("netstack: could not deregister IP %v: %v", ip.String(), err)
+			} else {
+				logf("netstack: deregistered IP %v", ip.String())
+			}
+		}
+		for ip := range ipsToBeAdded {
+			err := ipstack.AddAddress(nicID, ipv4.ProtocolNumber, ip)
+			if err != nil {
+				logf("netstack: could not register IP %v: %v", ip.String(), err)
+			} else {
+				logf("netstack: registered IP %v", ip.String())
+			}
+		}
+	})
 
 	// Add 0.0.0.0/0 default route.
 	subnet, _ := tcpip.NewSubnet(tcpip.Address(strings.Repeat("\x00", 4)), tcpip.AddressMask(strings.Repeat("\x00", 4)))

wgengine/userspace.go

@@ -109,6 +109,7 @@ type userspaceEngine struct {
 	trimmedDisco        map[tailcfg.DiscoKey]bool // set of disco keys of peers currently excluded from wireguard config
 	sentActivityAt      map[netaddr.IP]*int64     // value is atomic int64 of unixtime
 	destIPActivityFuncs map[netaddr.IP]func()
+	networkMapCallbacks map[*networkMapCallbackHandle]NetworkMapCallback
 
 	mu                 sync.Mutex // guards following; see lock order comment below
 	closing            bool       // Close was called (even if we're still closing)
@@ -208,13 +209,14 @@ func newUserspaceEngineAdvanced(conf EngineConfig) (_ Engine, reterr error) {
 		Forward: true,
 	}
 	e := &userspaceEngine{
-		timeNow:  time.Now,
-		logf:     logf,
-		reqCh:    make(chan struct{}, 1),
-		waitCh:   make(chan struct{}),
-		tundev:   tstun.WrapTUN(logf, conf.TUN),
-		resolver: tsdns.NewResolver(rconf),
-		pingers:  make(map[wgkey.Key]*pinger),
+		timeNow:             time.Now,
+		logf:                logf,
+		reqCh:               make(chan struct{}, 1),
+		waitCh:              make(chan struct{}),
+		tundev:              tstun.WrapTUN(logf, conf.TUN),
+		resolver:            tsdns.NewResolver(rconf),
+		pingers:             make(map[wgkey.Key]*pinger),
+		networkMapCallbacks: make(map[*networkMapCallbackHandle]NetworkMapCallback),
 	}
 	e.localAddrs.Store(map[netaddr.IP]bool{})
 	e.linkState, _ = getLinkState()
@@ -1263,6 +1265,14 @@ func (e *userspaceEngine) SetLinkChangeCallback(cb func(major bool, newState *in
 	}
 }
 
+func (e *userspaceEngine) AddNetworkMapCallback(cb NetworkMapCallback) func() {
+	handle := new(networkMapCallbackHandle)
+	e.networkMapCallbacks[handle] = cb
+	return func() {
+		delete(e.networkMapCallbacks, handle)
+	}
+}
+
 func getLinkState() (*interfaces.State, error) {
 	s, err := interfaces.GetState()
 	if s != nil {
@@ -1281,6 +1291,9 @@ func (e *userspaceEngine) SetDERPMap(dm *tailcfg.DERPMap) {
 
 func (e *userspaceEngine) SetNetworkMap(nm *controlclient.NetworkMap) {
 	e.magicConn.SetNetworkMap(nm)
+	for _, fn := range e.networkMapCallbacks {
+		fn(nm)
+	}
 }
 
 func (e *userspaceEngine) DiscoPublicKey() tailcfg.DiscoKey {

wgengine/watchdog.go

@@ -110,6 +110,11 @@ func (e *watchdogEngine) SetDERPMap(m *tailcfg.DERPMap) {
 func (e *watchdogEngine) SetNetworkMap(nm *controlclient.NetworkMap) {
 	e.watchdog("SetNetworkMap", func() { e.wrap.SetNetworkMap(nm) })
 }
+func (e *watchdogEngine) AddNetworkMapCallback(callback NetworkMapCallback) func() {
+	var fn func()
+	e.watchdog("AddNetworkMapCallback", func() { fn = e.wrap.AddNetworkMapCallback(callback) })
+	return func() { e.watchdog("RemoveNetworkMapCallback", fn) }
+}
 func (e *watchdogEngine) DiscoPublicKey() (k tailcfg.DiscoKey) {
 	e.watchdog("DiscoPublicKey", func() { k = e.wrap.DiscoPublicKey() })
 	return k

wgengine/wgengine.go

@@ -49,6 +49,11 @@ type Status struct {
 // NetInfoCallback is the type used by Engine.SetNetInfoCallback.
 type NetInfoCallback func(*tailcfg.NetInfo)
 
+// NetworkMapCallback is the type used by callbacks that hook
+// into network map updates.
+type NetworkMapCallback func(*controlclient.NetworkMap)
+type networkMapCallbackHandle struct{ _ byte }
+
 // ErrNoChanges is returned by Engine.Reconfig if no changes were made.
 var ErrNoChanges = errors.New("no changes made to Engine config")
 
@@ -114,6 +119,12 @@ type Engine interface {
 	// The network map should only be read from.
 	SetNetworkMap(*controlclient.NetworkMap)
 
+	// AddNetworkMapCallback adds a function to a list of callbacks
+	// that are called when the network map updates. It returns a
+	// function that when called would remove the function from the
+	// list of callbacks.
+	AddNetworkMapCallback(NetworkMapCallback) func()
+
 	// SetNetInfoCallback sets the function to call when a
 	// new NetInfo summary is available.
 	SetNetInfoCallback(NetInfoCallback)