mirror of
				https://github.com/juanfont/headscale.git
				synced 2025-10-28 10:51:44 +01:00 
			
		
		
		
	
		
			
				
	
	
		
			308 lines
		
	
	
		
			8.3 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			308 lines
		
	
	
		
			8.3 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
| package routes
 | |
| 
 | |
| import (
 | |
| 	"fmt"
 | |
| 	"net/netip"
 | |
| 	"slices"
 | |
| 	"sort"
 | |
| 	"strings"
 | |
| 	"sync"
 | |
| 
 | |
| 	"github.com/juanfont/headscale/hscontrol/types"
 | |
| 	"github.com/juanfont/headscale/hscontrol/util"
 | |
| 	"github.com/rs/zerolog/log"
 | |
| 	xmaps "golang.org/x/exp/maps"
 | |
| 	"tailscale.com/net/tsaddr"
 | |
| 	"tailscale.com/util/set"
 | |
| )
 | |
| 
 | |
| type PrimaryRoutes struct {
 | |
| 	mu sync.Mutex
 | |
| 
 | |
| 	// routes is a map of prefixes that are adverties and approved and available
 | |
| 	// in the global headscale state.
 | |
| 	routes map[types.NodeID]set.Set[netip.Prefix]
 | |
| 
 | |
| 	// primaries is a map of prefixes to the node that is the primary for that prefix.
 | |
| 	primaries map[netip.Prefix]types.NodeID
 | |
| 	isPrimary map[types.NodeID]bool
 | |
| }
 | |
| 
 | |
| func New() *PrimaryRoutes {
 | |
| 	return &PrimaryRoutes{
 | |
| 		routes:    make(map[types.NodeID]set.Set[netip.Prefix]),
 | |
| 		primaries: make(map[netip.Prefix]types.NodeID),
 | |
| 		isPrimary: make(map[types.NodeID]bool),
 | |
| 	}
 | |
| }
 | |
| 
 | |
| // updatePrimaryLocked recalculates the primary routes and updates the internal state.
 | |
| // It returns true if the primary routes have changed.
 | |
| // It is assumed that the caller holds the lock.
 | |
| // The algorithm is as follows:
 | |
| // 1. Reset the primaries map.
 | |
| // 2. Iterate over the routes and count the number of times a prefix is advertised.
 | |
| // 3. If a prefix is advertised by at least two nodes, it is a primary route.
 | |
| // 4. If the primary routes have changed, update the internal state and return true.
 | |
| // 5. Otherwise, return false.
 | |
| func (pr *PrimaryRoutes) updatePrimaryLocked() bool {
 | |
| 	log.Debug().Caller().Msg("updatePrimaryLocked starting")
 | |
| 
 | |
| 	// reset the primaries map, as we are going to recalculate it.
 | |
| 	allPrimaries := make(map[netip.Prefix][]types.NodeID)
 | |
| 	pr.isPrimary = make(map[types.NodeID]bool)
 | |
| 	changed := false
 | |
| 
 | |
| 	// sort the node ids so we can iterate over them in a deterministic order.
 | |
| 	// this is important so the same node is chosen two times in a row
 | |
| 	// as the primary route.
 | |
| 	ids := types.NodeIDs(xmaps.Keys(pr.routes))
 | |
| 	sort.Sort(ids)
 | |
| 
 | |
| 	// Create a map of prefixes to nodes that serve them so we
 | |
| 	// can determine the primary route for each prefix.
 | |
| 	for _, id := range ids {
 | |
| 		routes := pr.routes[id]
 | |
| 		for route := range routes {
 | |
| 			if _, ok := allPrimaries[route]; !ok {
 | |
| 				allPrimaries[route] = []types.NodeID{id}
 | |
| 			} else {
 | |
| 				allPrimaries[route] = append(allPrimaries[route], id)
 | |
| 			}
 | |
| 		}
 | |
| 	}
 | |
| 
 | |
| 	// Go through all prefixes and determine the primary route for each.
 | |
| 	// If the number of routes is below the minimum, remove the primary.
 | |
| 	// If the current primary is still available, continue.
 | |
| 	// If the current primary is not available, select a new one.
 | |
| 	for prefix, nodes := range allPrimaries {
 | |
| 		log.Debug().
 | |
| 			Caller().
 | |
| 			Str("prefix", prefix.String()).
 | |
| 			Uints64("availableNodes", func() []uint64 {
 | |
| 				ids := make([]uint64, len(nodes))
 | |
| 				for i, id := range nodes {
 | |
| 					ids[i] = id.Uint64()
 | |
| 				}
 | |
| 
 | |
| 				return ids
 | |
| 			}()).
 | |
| 			Msg("Processing prefix for primary route selection")
 | |
| 
 | |
| 		if node, ok := pr.primaries[prefix]; ok {
 | |
| 			// If the current primary is still available, continue.
 | |
| 			if slices.Contains(nodes, node) {
 | |
| 				log.Debug().
 | |
| 					Caller().
 | |
| 					Str("prefix", prefix.String()).
 | |
| 					Uint64("currentPrimary", node.Uint64()).
 | |
| 					Msg("Current primary still available, keeping it")
 | |
| 
 | |
| 				continue
 | |
| 			} else {
 | |
| 				log.Debug().
 | |
| 					Caller().
 | |
| 					Str("prefix", prefix.String()).
 | |
| 					Uint64("oldPrimary", node.Uint64()).
 | |
| 					Msg("Current primary no longer available")
 | |
| 			}
 | |
| 		}
 | |
| 		if len(nodes) >= 1 {
 | |
| 			pr.primaries[prefix] = nodes[0]
 | |
| 			changed = true
 | |
| 			log.Debug().
 | |
| 				Caller().
 | |
| 				Str("prefix", prefix.String()).
 | |
| 				Uint64("newPrimary", nodes[0].Uint64()).
 | |
| 				Msg("Selected new primary for prefix")
 | |
| 		}
 | |
| 	}
 | |
| 
 | |
| 	// Clean up any remaining primaries that are no longer valid.
 | |
| 	for prefix := range pr.primaries {
 | |
| 		if _, ok := allPrimaries[prefix]; !ok {
 | |
| 			log.Debug().
 | |
| 				Caller().
 | |
| 				Str("prefix", prefix.String()).
 | |
| 				Msg("Cleaning up primary route that no longer has available nodes")
 | |
| 			delete(pr.primaries, prefix)
 | |
| 			changed = true
 | |
| 		}
 | |
| 	}
 | |
| 
 | |
| 	// Populate the quick lookup index for primary routes
 | |
| 	for _, nodeID := range pr.primaries {
 | |
| 		pr.isPrimary[nodeID] = true
 | |
| 	}
 | |
| 
 | |
| 	log.Debug().
 | |
| 		Caller().
 | |
| 		Bool("changed", changed).
 | |
| 		Str("finalState", pr.stringLocked()).
 | |
| 		Msg("updatePrimaryLocked completed")
 | |
| 
 | |
| 	return changed
 | |
| }
 | |
| 
 | |
| // SetRoutes sets the routes for a given Node ID and recalculates the primary routes
 | |
| // of the headscale.
 | |
| // It returns true if there was a change in primary routes.
 | |
| // All exit routes are ignored as they are not used in primary route context.
 | |
| func (pr *PrimaryRoutes) SetRoutes(node types.NodeID, prefixes ...netip.Prefix) bool {
 | |
| 	pr.mu.Lock()
 | |
| 	defer pr.mu.Unlock()
 | |
| 
 | |
| 	log.Debug().
 | |
| 		Caller().
 | |
| 		Uint64("node.id", node.Uint64()).
 | |
| 		Strs("prefixes", util.PrefixesToString(prefixes)).
 | |
| 		Msg("PrimaryRoutes.SetRoutes called")
 | |
| 
 | |
| 	// If no routes are being set, remove the node from the routes map.
 | |
| 	if len(prefixes) == 0 {
 | |
| 		wasPresent := false
 | |
| 		if _, ok := pr.routes[node]; ok {
 | |
| 			delete(pr.routes, node)
 | |
| 			wasPresent = true
 | |
| 			log.Debug().
 | |
| 				Caller().
 | |
| 				Uint64("node.id", node.Uint64()).
 | |
| 				Msg("Removed node from primary routes (no prefixes)")
 | |
| 		}
 | |
| 		changed := pr.updatePrimaryLocked()
 | |
| 		log.Debug().
 | |
| 			Caller().
 | |
| 			Uint64("node.id", node.Uint64()).
 | |
| 			Bool("wasPresent", wasPresent).
 | |
| 			Bool("changed", changed).
 | |
| 			Str("newState", pr.stringLocked()).
 | |
| 			Msg("SetRoutes completed (remove)")
 | |
| 
 | |
| 		return changed
 | |
| 	}
 | |
| 
 | |
| 	rs := make(set.Set[netip.Prefix], len(prefixes))
 | |
| 	for _, prefix := range prefixes {
 | |
| 		if !tsaddr.IsExitRoute(prefix) {
 | |
| 			rs.Add(prefix)
 | |
| 		}
 | |
| 	}
 | |
| 
 | |
| 	if rs.Len() != 0 {
 | |
| 		pr.routes[node] = rs
 | |
| 		log.Debug().
 | |
| 			Caller().
 | |
| 			Uint64("node.id", node.Uint64()).
 | |
| 			Strs("routes", util.PrefixesToString(rs.Slice())).
 | |
| 			Msg("Updated node routes in primary route manager")
 | |
| 	} else {
 | |
| 		delete(pr.routes, node)
 | |
| 		log.Debug().
 | |
| 			Caller().
 | |
| 			Uint64("node.id", node.Uint64()).
 | |
| 			Msg("Removed node from primary routes (only exit routes)")
 | |
| 	}
 | |
| 
 | |
| 	changed := pr.updatePrimaryLocked()
 | |
| 	log.Debug().
 | |
| 		Caller().
 | |
| 		Uint64("node.id", node.Uint64()).
 | |
| 		Bool("changed", changed).
 | |
| 		Str("newState", pr.stringLocked()).
 | |
| 		Msg("SetRoutes completed (update)")
 | |
| 
 | |
| 	return changed
 | |
| }
 | |
| 
 | |
| func (pr *PrimaryRoutes) PrimaryRoutes(id types.NodeID) []netip.Prefix {
 | |
| 	if pr == nil {
 | |
| 		return nil
 | |
| 	}
 | |
| 
 | |
| 	pr.mu.Lock()
 | |
| 	defer pr.mu.Unlock()
 | |
| 
 | |
| 	// Short circuit if the node is not a primary for any route.
 | |
| 	if _, ok := pr.isPrimary[id]; !ok {
 | |
| 		return nil
 | |
| 	}
 | |
| 
 | |
| 	var routes []netip.Prefix
 | |
| 
 | |
| 	for prefix, node := range pr.primaries {
 | |
| 		if node == id {
 | |
| 			routes = append(routes, prefix)
 | |
| 		}
 | |
| 	}
 | |
| 
 | |
| 	tsaddr.SortPrefixes(routes)
 | |
| 
 | |
| 	return routes
 | |
| }
 | |
| 
 | |
| func (pr *PrimaryRoutes) String() string {
 | |
| 	pr.mu.Lock()
 | |
| 	defer pr.mu.Unlock()
 | |
| 
 | |
| 	return pr.stringLocked()
 | |
| }
 | |
| 
 | |
| func (pr *PrimaryRoutes) stringLocked() string {
 | |
| 	var sb strings.Builder
 | |
| 
 | |
| 	fmt.Fprintln(&sb, "Available routes:")
 | |
| 
 | |
| 	ids := types.NodeIDs(xmaps.Keys(pr.routes))
 | |
| 	sort.Sort(ids)
 | |
| 	for _, id := range ids {
 | |
| 		prefixes := pr.routes[id]
 | |
| 		fmt.Fprintf(&sb, "\nNode %d: %s", id, strings.Join(util.PrefixesToString(prefixes.Slice()), ", "))
 | |
| 	}
 | |
| 
 | |
| 	fmt.Fprintln(&sb, "\n\nCurrent primary routes:")
 | |
| 	for route, nodeID := range pr.primaries {
 | |
| 		fmt.Fprintf(&sb, "\nRoute %s: %d", route, nodeID)
 | |
| 	}
 | |
| 
 | |
| 	return sb.String()
 | |
| }
 | |
| 
 | |
| // DebugRoutes represents the primary routes state in a structured format for JSON serialization.
 | |
| type DebugRoutes struct {
 | |
| 	// AvailableRoutes maps node IDs to their advertised routes
 | |
| 	// In the context of primary routes, this represents the routes that are available
 | |
| 	// for each node. A route will only be available if it is advertised by the node
 | |
| 	// AND approved.
 | |
| 	// Only routes by nodes currently connected to the headscale server are included.
 | |
| 	AvailableRoutes map[types.NodeID][]netip.Prefix `json:"available_routes"`
 | |
| 
 | |
| 	// PrimaryRoutes maps route prefixes to the primary node serving them
 | |
| 	PrimaryRoutes map[string]types.NodeID `json:"primary_routes"`
 | |
| }
 | |
| 
 | |
| // DebugJSON returns a structured representation of the primary routes state suitable for JSON serialization.
 | |
| func (pr *PrimaryRoutes) DebugJSON() DebugRoutes {
 | |
| 	pr.mu.Lock()
 | |
| 	defer pr.mu.Unlock()
 | |
| 
 | |
| 	debug := DebugRoutes{
 | |
| 		AvailableRoutes: make(map[types.NodeID][]netip.Prefix),
 | |
| 		PrimaryRoutes:   make(map[string]types.NodeID),
 | |
| 	}
 | |
| 
 | |
| 	// Populate available routes
 | |
| 	for nodeID, routes := range pr.routes {
 | |
| 		prefixes := routes.Slice()
 | |
| 		tsaddr.SortPrefixes(prefixes)
 | |
| 		debug.AvailableRoutes[nodeID] = prefixes
 | |
| 	}
 | |
| 
 | |
| 	// Populate primary routes
 | |
| 	for prefix, nodeID := range pr.primaries {
 | |
| 		debug.PrimaryRoutes[prefix.String()] = nodeID
 | |
| 	}
 | |
| 
 | |
| 	return debug
 | |
| }
 |