all: use syncs.AtomicValue

Signed-off-by: Maisem Ali <maisem@tailscale.com>

wgengine/magicsock/magicsock.go

@@ -264,7 +264,7 @@ type Conn struct {
 
 	// stunReceiveFunc holds the current STUN packet processing func.
 	// Its Loaded value is always non-nil.
-	stunReceiveFunc atomic.Value // of func(p []byte, fromAddr *net.UDPAddr)
+	stunReceiveFunc syncs.AtomicValue[func(p []byte, fromAddr netip.AddrPort)]
 
 	// derpRecvCh is used by receiveDERP to read DERP messages.
 	// It must have buffer size > 0; see issue 3736.
@@ -300,7 +300,7 @@ type Conn struct {
 
 	// havePrivateKey is whether privateKey is non-zero.
 	havePrivateKey  atomic.Bool
-	publicKeyAtomic atomic.Value // of key.NodePublic (or NodeKey zero value if !havePrivateKey)
+	publicKeyAtomic syncs.AtomicValue[key.NodePublic] // or NodeKey zero value if !havePrivateKey
 
 	// derpMapAtomic is the same as derpMap, but without requiring
 	// sync.Mutex. For use with NewRegionClient's callback, to avoid
@@ -1668,7 +1668,7 @@ func (c *Conn) receiveIPv4(b []byte) (n int, ep conn.Endpoint, err error) {
 // caller).
 func (c *Conn) receiveIP(b []byte, ipp netip.AddrPort, cache *ippEndpointCache) (ep *endpoint, ok bool) {
 	if stun.Is(b) {
-		c.stunReceiveFunc.Load().(func([]byte, netip.AddrPort))(b, ipp)
+		c.stunReceiveFunc.Load()(b, ipp)
 		return nil, false
 	}
 	if c.handleDiscoMessage(b, ipp, key.NodePublic{}) {
@@ -2979,10 +2979,8 @@ type RebindingUDPConn struct {
 	// check pconn (after acquiring mu) to see if there's been a rebind
 	// meanwhile.
 	// pconn isn't really needed, but makes some of the code simpler
-	// to keep it in a type safe form. TODO(bradfitz): really we should make a generic
-	// atomic.Value. Unfortunately Go 1.19's atomic.Pointer[T] is only for pointers,
-	// not interfaces.
-	pconnAtomic atomic.Value // of nettype.PacketConn
+	// to keep it in a type safe form.
+	pconnAtomic syncs.AtomicValue[nettype.PacketConn]
 
 	mu    sync.Mutex // held while changing pconn (and pconnAtomic)
 	pconn nettype.PacketConn
@@ -3004,7 +3002,7 @@ func (c *RebindingUDPConn) currentConn() nettype.PacketConn {
 // It returns the number of bytes copied and the source address.
 func (c *RebindingUDPConn) ReadFrom(b []byte) (int, net.Addr, error) {
 	for {
-		pconn := c.pconnAtomic.Load().(nettype.PacketConn)
+		pconn := c.pconnAtomic.Load()
 		n, addr, err := pconn.ReadFrom(b)
 		if err != nil && pconn != c.currentConn() {
 			continue
@@ -3022,7 +3020,7 @@ func (c *RebindingUDPConn) ReadFrom(b []byte) (int, net.Addr, error) {
 // when c's underlying connection is a net.UDPConn.
 func (c *RebindingUDPConn) ReadFromNetaddr(b []byte) (n int, ipp netip.AddrPort, err error) {
 	for {
-		pconn := c.pconnAtomic.Load().(nettype.PacketConn)
+		pconn := c.pconnAtomic.Load()
 
 		// Optimization: Treat *net.UDPConn specially.
 		// This lets us avoid allocations by calling ReadFromUDPAddrPort.
@@ -3081,7 +3079,7 @@ func (c *RebindingUDPConn) closeLocked() error {
 
 func (c *RebindingUDPConn) WriteTo(b []byte, addr net.Addr) (int, error) {
 	for {
-		pconn := c.pconnAtomic.Load().(nettype.PacketConn)
+		pconn := c.pconnAtomic.Load()
 
 		n, err := pconn.WriteTo(b, addr)
 		if err != nil {
@@ -3095,7 +3093,7 @@ func (c *RebindingUDPConn) WriteTo(b []byte, addr net.Addr) (int, error) {
 
 func (c *RebindingUDPConn) WriteToUDPAddrPort(b []byte, addr netip.AddrPort) (int, error) {
 	for {
-		pconn := c.pconnAtomic.Load().(nettype.PacketConn)
+		pconn := c.pconnAtomic.Load()
 
 		n, err := pconn.WriteToUDPAddrPort(b, addr)
 		if err != nil {
@@ -3643,10 +3641,9 @@ func (de *endpoint) removeSentPingLocked(txid stun.TxID, sp sentPing) {
 // The caller should use de.discoKey as the discoKey argument.
 // It is passed in so that sendDiscoPing doesn't need to lock de.mu.
 func (de *endpoint) sendDiscoPing(ep netip.AddrPort, discoKey key.DiscoPublic, txid stun.TxID, logLevel discoLogLevel) {
-	selfPubKey, _ := de.c.publicKeyAtomic.Load().(key.NodePublic)
 	sent, _ := de.c.sendDiscoMessage(ep, de.publicKey, discoKey, &disco.Ping{
 		TxID:    [12]byte(txid),
-		NodeKey: selfPubKey,
+		NodeKey: de.c.publicKeyAtomic.Load(),
 	}, logLevel)
 	if !sent {
 		de.forgetPing(txid)

wgengine/netstack/netstack.go

@@ -117,7 +117,7 @@ type Impl struct {
 	// is a local (non-subnet) Tailscale IP address of this
 	// machine. It's always a non-nil func. It's changed on netmap
 	// updates.
-	atomicIsLocalIPFunc atomic.Value // of func(netip.Addr) bool
+	atomicIsLocalIPFunc syncs.AtomicValue[func(netip.Addr) bool]
 
 	mu sync.Mutex
 	// connsOpenBySubnetIP keeps track of number of connections open
@@ -513,7 +513,7 @@ func (ns *Impl) inject() {
 // isLocalIP reports whether ip is a Tailscale IP assigned to this
 // node directly (but not a subnet-routed IP).
 func (ns *Impl) isLocalIP(ip netip.Addr) bool {
-	return ns.atomicIsLocalIPFunc.Load().(func(netip.Addr) bool)(ip)
+	return ns.atomicIsLocalIPFunc.Load()(ip)
 }
 
 func (ns *Impl) processSSH() bool {

wgengine/userspace.go

@@ -17,7 +17,6 @@ import (
 	"runtime"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"go4.org/mem"
@@ -36,6 +35,7 @@ import (
 	"tailscale.com/net/tsdial"
 	"tailscale.com/net/tshttpproxy"
 	"tailscale.com/net/tstun"
+	"tailscale.com/syncs"
 	"tailscale.com/tailcfg"
 	"tailscale.com/tstime/mono"
 	"tailscale.com/types/dnstype"
@@ -107,11 +107,11 @@ type userspaceEngine struct {
 	// isLocalAddr reports the whether an IP is assigned to the local
 	// tunnel interface. It's used to reflect local packets
 	// incorrectly sent to us.
-	isLocalAddr atomic.Value // of func(netip.Addr)bool
+	isLocalAddr syncs.AtomicValue[func(netip.Addr) bool]
 
 	// isDNSIPOverTailscale reports the whether a DNS resolver's IP
 	// is being routed over Tailscale.
-	isDNSIPOverTailscale atomic.Value // of func(netip.Addr)bool
+	isDNSIPOverTailscale syncs.AtomicValue[func(netip.Addr) bool]
 
 	wgLock              sync.Mutex // serializes all wgdev operations; see lock order comment below
 	lastCfgFull         wgcfg.Config
@@ -497,7 +497,7 @@ func (e *userspaceEngine) handleLocalPackets(p *packet.Parsed, t *tstun.Wrapper)
 	}
 
 	if runtime.GOOS == "darwin" || runtime.GOOS == "ios" {
-		isLocalAddr, ok := e.isLocalAddr.Load().(func(netip.Addr) bool)
+		isLocalAddr, ok := e.isLocalAddr.LoadOk()
 		if !ok {
 			e.logf("[unexpected] e.isLocalAddr was nil, can't check for loopback packet")
 		} else if isLocalAddr(p.Dst.Addr()) {
@@ -1621,7 +1621,7 @@ type fwdDNSLinkSelector struct {
 }
 
 func (ls fwdDNSLinkSelector) PickLink(ip netip.Addr) (linkName string) {
-	if ls.ue.isDNSIPOverTailscale.Load().(func(netip.Addr) bool)(ip) {
+	if ls.ue.isDNSIPOverTailscale.Load()(ip) {
 		return ls.tunName
 	}
 	return ""

wgengine/wglog/wglog.go

@@ -9,9 +9,9 @@ import (
 	"fmt"
 	"strings"
 	"sync"
-	"sync/atomic"
 
 	"golang.zx2c4.com/wireguard/device"
+	"tailscale.com/syncs"
 	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
 	"tailscale.com/wgengine/wgcfg"
@@ -21,7 +21,7 @@ import (
 // It can be modified at run time to adjust to new wireguard-go configurations.
 type Logger struct {
 	DeviceLogger *device.Logger
-	replace      atomic.Value                 // of map[string]string
+	replace      syncs.AtomicValue[map[string]string]
 	mu           sync.Mutex                   // protects strs
 	strs         map[key.NodePublic]*strCache // cached strs used to populate replace
 }
@@ -52,7 +52,7 @@ func NewLogger(logf logger.Logf) *Logger {
 			// See https://github.com/tailscale/tailscale/issues/1388.
 			return
 		}
-		replace, _ := ret.replace.Load().(map[string]string)
+		replace := ret.replace.Load()
 		if replace == nil {
 			// No replacements specified; log as originally planned.
 			logf(format, args...)