mirror of
https://github.com/juanfont/headscale.git
synced 2026-07-08 00:50:20 +09:00
mapper: consolidate builder and connection helpers
This commit is contained in:
@@ -315,11 +315,7 @@ func (b *Batcher) AddNode(
|
||||
initialMap, err := b.MapResponseFromChange(id, change.FullSelf(id))
|
||||
if err != nil {
|
||||
nlog.Error().Err(err).Msg("initial map generation failed")
|
||||
nodeConn.removeConnectionByChannel(c)
|
||||
|
||||
if !nodeConn.hasActiveConnections() {
|
||||
nodeConn.markDisconnected()
|
||||
}
|
||||
nodeConn.detach(c)
|
||||
|
||||
return fmt.Errorf("generating initial map for node %d: %w", id, err)
|
||||
}
|
||||
@@ -343,11 +339,7 @@ func (b *Batcher) AddNode(
|
||||
nlog.Error().Err(ErrInitialMapSendTimeout).Msg("initial map send timeout")
|
||||
nlog.Debug().Caller().Dur("timeout.duration", 5*time.Second). //nolint:mnd
|
||||
Msg("initial map send timed out because channel was blocked or receiver not ready")
|
||||
nodeConn.removeConnectionByChannel(c)
|
||||
|
||||
if !nodeConn.hasActiveConnections() {
|
||||
nodeConn.markDisconnected()
|
||||
}
|
||||
nodeConn.detach(c)
|
||||
|
||||
return fmt.Errorf("%w for node %d", ErrInitialMapSendTimeout, id)
|
||||
}
|
||||
|
||||
+15
-10
@@ -61,6 +61,16 @@ func (b *MapResponseBuilder) hasErrors() bool {
|
||||
return len(b.errs) > 0
|
||||
}
|
||||
|
||||
// node looks up the requesting node, recording ErrNodeNotFoundMapper on miss.
|
||||
func (b *MapResponseBuilder) node() (types.NodeView, bool) {
|
||||
nv, ok := b.mapper.state.GetNodeByID(b.nodeID)
|
||||
if !ok {
|
||||
b.addError(ErrNodeNotFoundMapper)
|
||||
}
|
||||
|
||||
return nv, ok
|
||||
}
|
||||
|
||||
// WithCapabilityVersion sets the capability version for the response.
|
||||
func (b *MapResponseBuilder) WithCapabilityVersion(capVer tailcfg.CapabilityVersion) *MapResponseBuilder {
|
||||
b.capVer = capVer
|
||||
@@ -69,9 +79,8 @@ func (b *MapResponseBuilder) WithCapabilityVersion(capVer tailcfg.CapabilityVers
|
||||
|
||||
// WithSelfNode adds the requesting node to the response.
|
||||
func (b *MapResponseBuilder) WithSelfNode() *MapResponseBuilder {
|
||||
nv, ok := b.mapper.state.GetNodeByID(b.nodeID)
|
||||
nv, ok := b.node()
|
||||
if !ok {
|
||||
b.addError(ErrNodeNotFoundMapper)
|
||||
return b
|
||||
}
|
||||
|
||||
@@ -136,9 +145,8 @@ func (b *MapResponseBuilder) WithDebugConfig() *MapResponseBuilder {
|
||||
|
||||
// WithSSHPolicy adds SSH policy configuration for the requesting node.
|
||||
func (b *MapResponseBuilder) WithSSHPolicy() *MapResponseBuilder {
|
||||
node, ok := b.mapper.state.GetNodeByID(b.nodeID)
|
||||
node, ok := b.node()
|
||||
if !ok {
|
||||
b.addError(ErrNodeNotFoundMapper)
|
||||
return b
|
||||
}
|
||||
|
||||
@@ -155,9 +163,8 @@ func (b *MapResponseBuilder) WithSSHPolicy() *MapResponseBuilder {
|
||||
|
||||
// WithDNSConfig adds DNS configuration for the requesting node.
|
||||
func (b *MapResponseBuilder) WithDNSConfig() *MapResponseBuilder {
|
||||
node, ok := b.mapper.state.GetNodeByID(b.nodeID)
|
||||
node, ok := b.node()
|
||||
if !ok {
|
||||
b.addError(ErrNodeNotFoundMapper)
|
||||
return b
|
||||
}
|
||||
|
||||
@@ -168,9 +175,8 @@ func (b *MapResponseBuilder) WithDNSConfig() *MapResponseBuilder {
|
||||
|
||||
// WithUserProfiles adds user profiles for the requesting node and given peers.
|
||||
func (b *MapResponseBuilder) WithUserProfiles(peers views.Slice[types.NodeView]) *MapResponseBuilder {
|
||||
node, ok := b.mapper.state.GetNodeByID(b.nodeID)
|
||||
node, ok := b.node()
|
||||
if !ok {
|
||||
b.addError(ErrNodeNotFoundMapper)
|
||||
return b
|
||||
}
|
||||
|
||||
@@ -185,9 +191,8 @@ func (b *MapResponseBuilder) WithUserProfiles(peers views.Slice[types.NodeView])
|
||||
// For autogroup:self policies, it returns per-node compiled rules.
|
||||
// For global policies, it returns the global filter reduced for this node.
|
||||
func (b *MapResponseBuilder) WithPacketFilters() *MapResponseBuilder {
|
||||
node, ok := b.mapper.state.GetNodeByID(b.nodeID)
|
||||
node, ok := b.node()
|
||||
if !ok {
|
||||
b.addError(ErrNodeNotFoundMapper)
|
||||
return b
|
||||
}
|
||||
|
||||
|
||||
+20
-14
@@ -499,15 +499,9 @@ func (m *mapper) filterVisiblePeerPatches(
|
||||
return nil
|
||||
}
|
||||
|
||||
var filtered []*tailcfg.PeerChange
|
||||
|
||||
for _, patch := range patches {
|
||||
if _, vis := visible[patch.NodeID]; vis {
|
||||
filtered = append(filtered, patch)
|
||||
}
|
||||
}
|
||||
|
||||
return filtered
|
||||
return filterByVisible(visible, patches, func(p *tailcfg.PeerChange) tailcfg.NodeID {
|
||||
return p.NodeID
|
||||
})
|
||||
}
|
||||
|
||||
// filterVisibleNodes restricts a peer slice to the nodes the recipient can see
|
||||
@@ -524,15 +518,27 @@ func (m *mapper) filterVisibleNodes(
|
||||
return views.SliceOf([]types.NodeView{})
|
||||
}
|
||||
|
||||
var filtered []types.NodeView
|
||||
return views.SliceOf(filterByVisible(visible, peers.AsSlice(), func(p types.NodeView) tailcfg.NodeID {
|
||||
return p.ID().NodeID()
|
||||
}))
|
||||
}
|
||||
|
||||
for _, peer := range peers.All() {
|
||||
if _, vis := visible[peer.ID().NodeID()]; vis {
|
||||
filtered = append(filtered, peer)
|
||||
// filterByVisible keeps only the items whose key resolves to a NodeID present
|
||||
// in the visible set, preserving input order.
|
||||
func filterByVisible[T any](
|
||||
visible map[tailcfg.NodeID]struct{},
|
||||
items []T,
|
||||
key func(T) tailcfg.NodeID,
|
||||
) []T {
|
||||
var filtered []T
|
||||
|
||||
for _, it := range items {
|
||||
if _, ok := visible[key(it)]; ok {
|
||||
filtered = append(filtered, it)
|
||||
}
|
||||
}
|
||||
|
||||
return views.SliceOf(filtered)
|
||||
return filtered
|
||||
}
|
||||
|
||||
func writeDebugMapResponse(
|
||||
|
||||
@@ -183,6 +183,16 @@ func (mc *multiChannelNodeConn) removeConnectionByChannel(c chan<- *tailcfg.MapR
|
||||
return false
|
||||
}
|
||||
|
||||
// detach removes the connection for the given channel and marks the node
|
||||
// disconnected if no active connections remain.
|
||||
func (mc *multiChannelNodeConn) detach(c chan<- *tailcfg.MapResponse) {
|
||||
mc.removeConnectionByChannel(c)
|
||||
|
||||
if !mc.hasActiveConnections() {
|
||||
mc.markDisconnected()
|
||||
}
|
||||
}
|
||||
|
||||
// hasActiveConnections checks if the node has any active connections.
|
||||
func (mc *multiChannelNodeConn) hasActiveConnections() bool {
|
||||
mc.mutex.RLock()
|
||||
|
||||
Reference in New Issue
Block a user