Skip to content

Commit d2d9c39

Browse files
ethanndicksonethanndickson
committed
use netmap callback
1 parent 612ae5f commit d2d9c39

File tree

4 files changed

+50
-44
lines changed

4 files changed

+50
-44
lines changed

‎tailnet/configmaps.go

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -284,16 +284,15 @@ func (c *configMaps) getBlockEndpoints() bool {
284284
// setDERPMap sets the DERP map, triggering a configuration of the engine if it has changed.
285285
// c.L MUST NOT be held.
286286
// Returns if the derpMap is dirty.
287-
func (c *configMaps) setDERPMap(derpMap *tailcfg.DERPMap) bool {
287+
func (c *configMaps) setDERPMap(derpMap *tailcfg.DERPMap) {
288288
c.L.Lock()
289289
defer c.L.Unlock()
290290
if CompareDERPMaps(c.derpMap, derpMap) {
291-
return false
291+
return
292292
}
293293
c.derpMap = derpMap
294294
c.derpMapDirty = true
295295
c.Broadcast()
296-
return true
297296
}
298297

299298
// derMapLocked returns the current DERPMap. c.L must be held

‎tailnet/conn.go

Lines changed: 18 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import (
3131
"tailscale.com/types/key"
3232
tslogger "tailscale.com/types/logger"
3333
"tailscale.com/types/netlogtype"
34+
"tailscale.com/types/netmap"
3435
"tailscale.com/wgengine"
3536
"tailscale.com/wgengine/capture"
3637
"tailscale.com/wgengine/magicsock"
@@ -262,16 +263,6 @@ func NewConn(options *Options) (conn *Conn, err error) {
262263
)
263264
nodeUp.setAddresses(options.Addresses)
264265
nodeUp.setBlockEndpoints(options.BlockEndpoints)
265-
wireguardEngine.SetStatusCallback(nodeUp.setStatus)
266-
magicConn.SetDERPForcedWebsocketCallback(nodeUp.setDERPForcedWebsocket)
267-
if telemetryStore != nil {
268-
wireguardEngine.SetNetInfoCallback(func(ni *tailcfg.NetInfo) {
269-
telemetryStore.setNetInfo(ni)
270-
nodeUp.setNetInfo(ni)
271-
})
272-
} else {
273-
wireguardEngine.SetNetInfoCallback(nodeUp.setNetInfo)
274-
}
275266

276267
server := &Conn{
277268
id: uuid.New(),
@@ -298,7 +289,21 @@ func NewConn(options *Options) (conn *Conn, err error) {
298289
_ = server.Close()
299290
}
300291
}()
301-
server.SetNodeCallback(nil)
292+
if server.telemetryStore != nil {
293+
server.wireguardEngine.SetNetInfoCallback(func(ni *tailcfg.NetInfo) {
294+
server.telemetryStore.setNetInfo(ni)
295+
nodeUp.setNetInfo(ni)
296+
})
297+
server.wireguardEngine.AddNetworkMapCallback(func(nm *netmap.NetworkMap) {
298+
if server.telemetryStore.updateNetworkMap(nm) {
299+
server.sendUpdatedTelemetry()
300+
}
301+
})
302+
} else {
303+
server.wireguardEngine.SetNetInfoCallback(nodeUp.setNetInfo)
304+
}
305+
server.wireguardEngine.SetStatusCallback(nodeUp.setStatus)
306+
server.magicConn.SetDERPForcedWebsocketCallback(nodeUp.setDERPForcedWebsocket)
302307

303308
netStack.GetTCPHandlerForFlow = server.forwardTCP
304309

@@ -391,23 +396,12 @@ func (c *Conn) SetAddresses(ips []netip.Prefix) error {
391396
// If telemetry is enabled, the callback will first update the telemetry store,
392397
// send the updated telemetry, and then call the provided callback.
393398
func (c *Conn) SetNodeCallback(callback func(node *Node)) {
394-
if c.telemetryStore != nil {
395-
c.nodeUpdater.setCallback(func(node *Node) {
396-
if c.telemetryStore.updateByNode(node) {
397-
c.sendUpdatedTelemetry()
398-
}
399-
callback(node)
400-
})
401-
} else {
402-
c.nodeUpdater.setCallback(callback)
403-
}
399+
c.nodeUpdater.setCallback(callback)
404400
}
405401

406402
// SetDERPMap updates the DERPMap of a connection.
407403
func (c *Conn) SetDERPMap(derpMap *tailcfg.DERPMap) {
408-
if c.configMaps.setDERPMap(derpMap) && c.telemetryStore != nil {
409-
c.telemetryStore.updateDerpMap(derpMap)
410-
}
404+
c.configMaps.setDERPMap(derpMap)
411405
}
412406

413407
func (c *Conn) SetDERPForceWebSockets(v bool) {
@@ -729,7 +723,6 @@ func (c *Conn) SendConnectedTelemetry(ip netip.Addr, application string) {
729723
return
730724
}
731725
c.telemetryStore.markConnected(&ip, c.createdAt, application)
732-
c.telemetryStore.updateRemoteNodeID(c.wireguardEngine)
733726
e := c.newTelemetryEvent()
734727
e.Status = proto.TelemetryEvent_CONNECTED
735728
c.telemetryWg.Add(1)
@@ -745,7 +738,6 @@ func (c *Conn) sendUpdatedTelemetry() {
745738
if c.telemetrySink == nil {
746739
return
747740
}
748-
c.telemetryStore.updateRemoteNodeID(c.wireguardEngine)
749741
e := c.newTelemetryEvent()
750742
e.Status = proto.TelemetryEvent_CONNECTED
751743
c.telemetryWg.Add(1)
@@ -761,7 +753,6 @@ func (c *Conn) SendDisconnectedTelemetry() {
761753
}
762754
e := c.newTelemetryEvent()
763755
e.Status = proto.TelemetryEvent_DISCONNECTED
764-
c.telemetryStore.updateRemoteNodeID(c.wireguardEngine)
765756
c.telemetryWg.Add(1)
766757
go func() {
767758
defer c.telemetryWg.Done()

‎tailnet/telemetry.go

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import (
1212
"google.golang.org/protobuf/types/known/timestamppb"
1313
"google.golang.org/protobuf/types/known/wrapperspb"
1414
"tailscale.com/tailcfg"
15-
"tailscale.com/wgengine"
15+
"tailscale.com/types/netmap"
1616

1717
"github.com/coder/coder/v2/cryptorand"
1818
"github.com/coder/coder/v2/tailnet/proto"
@@ -83,26 +83,39 @@ func (b *TelemetryStore) markConnected(ip *netip.Addr, connCreatedAt time.Time,
8383
b.application = application
8484
}
8585

86-
func (b *TelemetryStore) updateRemoteNodeID(engine wgengine.Engine) {
86+
func (b *TelemetryStore) updateRemoteNodeIDLocked(nm *netmap.NetworkMap) {
8787
b.mu.Lock()
8888
defer b.mu.Unlock()
8989

9090
if b.connectedIP == nil {
9191
return
9292
}
9393

94-
pip, ok := engine.PeerForIP(*b.connectedIP)
95-
if ok {
96-
b.connectedNodeID = uint64(pip.Node.ID)
94+
ip := *b.connectedIP
95+
96+
for _, p := range nm.Peers {
97+
for _, a := range p.Addresses {
98+
if a.Addr() == ip && a.IsSingleIP() {
99+
b.connectedNodeID = uint64(p.ID)
100+
}
101+
}
97102
}
98103
}
99104

105+
// Returning whether a new telemetry event should be sent
106+
func (b *TelemetryStore) updateNetworkMap(nm *netmap.NetworkMap) bool {
107+
b.mu.Lock()
108+
defer b.mu.Unlock()
109+
110+
b.updateDerpMapLocked(nm.DERPMap)
111+
b.updateRemoteNodeIDLocked(nm)
112+
return b.updateByNodeLocked(nm.SelfNode)
113+
}
114+
100115
// Given a DERPMap, anonymise all IPs and hostnames.
101116
// Keep track of seen hostnames/cert names to anonymize them from future logs.
102117
// b.mu must NOT be held.
103-
func (b *TelemetryStore) updateDerpMap(cur *tailcfg.DERPMap) {
104-
b.mu.Lock()
105-
defer b.mu.Unlock()
118+
func (b *TelemetryStore) updateDerpMapLocked(cur *tailcfg.DERPMap) {
106119
cleanMap := cur.Clone()
107120
for _, r := range cleanMap.Regions {
108121
for _, n := range r.Nodes {
@@ -123,12 +136,13 @@ func (b *TelemetryStore) updateDerpMap(cur *tailcfg.DERPMap) {
123136

124137
// Update the telemetry store with the current self node state.
125138
// Returns true if the home DERP has changed.
126-
func (b *TelemetryStore) updateByNode(n *Node) bool {
127-
b.mu.Lock()
128-
defer b.mu.Unlock()
129-
139+
func (b *TelemetryStore) updateByNodeLocked(n *tailcfg.Node) bool {
130140
b.nodeID = uint64(n.ID)
131-
newHome := int32(n.PreferredDERP)
141+
derpIP, err := netip.ParseAddrPort(n.DERP)
142+
if err != nil {
143+
return false
144+
}
145+
newHome := int32(derpIP.Port())
132146
if b.homeDerp != newHome {
133147
b.homeDerp = newHome
134148
return true

‎tailnet/telemetry_internal_test.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,9 @@ func TestTelemetryStore(t *testing.T) {
133133
},
134134
},
135135
}
136-
telemetry.updateDerpMap(derpMap)
136+
telemetry.mu.Lock()
137+
telemetry.updateDerpMapLocked(derpMap)
138+
telemetry.mu.Unlock()
137139

138140
event := telemetry.newEvent()
139141
require.Len(t, event.DerpMap.Regions[999].Nodes, 1)

0 commit comments

Comments
 (0)