improve prefix errors, allow only v4 or v6

Fixes juanfont#1785 Fixes juanfont#1827 Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
kradalby · Apr 12, 2024 · b9fe96a · b9fe96a
1 parent a329442
commit b9fe96a
Show file tree

Hide file tree

Showing 4 changed files with 118 additions and 46 deletions.
diff --git a/hscontrol/app.go b/hscontrol/app.go
@@ -148,7 +148,7 @@ func NewHeadscale(cfg *types.Config) (*Headscale, error) {
 		return nil, err
 	}
 
-	app.ipAlloc, err = db.NewIPAllocator(app.db, *cfg.PrefixV4, *cfg.PrefixV6)
+	app.ipAlloc, err = db.NewIPAllocator(app.db, cfg.PrefixV4, cfg.PrefixV6)
 	if err != nil {
 		return nil, err
 	}

diff --git a/hscontrol/db/ip.go b/hscontrol/db/ip.go
@@ -20,12 +20,12 @@ import (
 type IPAllocator struct {
 	mu sync.Mutex
 
-	prefix4 netip.Prefix
-	prefix6 netip.Prefix
+	prefix4 *netip.Prefix
+	prefix6 *netip.Prefix
 
 	// Previous IPs handed out
-	prev4 netip.Addr
-	prev6 netip.Addr
+	prev4 *netip.Addr
+	prev6 *netip.Addr
 
 	// Set of all IPs handed out.
 	// This might not be in sync with the database,
@@ -40,7 +40,12 @@ type IPAllocator struct {
 // provided IPv4 and IPv6 prefix. It needs to be created
 // when headscale starts and needs to finish its read
 // transaction before any writes to the database occur.
-func NewIPAllocator(db *HSDatabase, prefix4, prefix6 netip.Prefix) (*IPAllocator, error) {
+func NewIPAllocator(db *HSDatabase, prefix4, prefix6 *netip.Prefix) (*IPAllocator, error) {
+	ret := IPAllocator{
+		prefix4: prefix4,
+		prefix6: prefix6,
+	}
+
 	var addressesSlices []string
 
 	if db != nil {
@@ -53,12 +58,24 @@ func NewIPAllocator(db *HSDatabase, prefix4, prefix6 netip.Prefix) (*IPAllocator
 
 	// Add network and broadcast addrs to used pool so they
 	// are not handed out to nodes.
-	network4, broadcast4 := util.GetIPPrefixEndpoints(prefix4)
-	network6, broadcast6 := util.GetIPPrefixEndpoints(prefix6)
-	ips.Add(network4)
-	ips.Add(broadcast4)
-	ips.Add(network6)
-	ips.Add(broadcast6)
+	if prefix4 != nil {
+		network4, broadcast4 := util.GetIPPrefixEndpoints(*prefix4)
+		ips.Add(network4)
+		ips.Add(broadcast4)
+
+		// Use network as starting point, it will be used to call .Next()
+		// TODO(kradalby): Could potentially take all the IPs loaded from
+		// the database into account to start at a more "educated" location.
+		ret.prev4 = &network4
+	}
+
+	if prefix6 != nil {
+		network6, broadcast6 := util.GetIPPrefixEndpoints(*prefix6)
+		ips.Add(network6)
+		ips.Add(broadcast6)
+
+		ret.prev6 = &network6
+	}
 
 	// Fetch all the IP Addresses currently handed out from the Database
 	// and add them to the used IP set.
@@ -86,40 +103,41 @@ func NewIPAllocator(db *HSDatabase, prefix4, prefix6 netip.Prefix) (*IPAllocator
 		)
 	}
 
-	return &IPAllocator{
-		usedIPs: ips,
+	ret.usedIPs = ips
 
-		prefix4: prefix4,
-		prefix6: prefix6,
-
-		// Use network as starting point, it will be used to call .Next()
-		// TODO(kradalby): Could potentially take all the IPs loaded from
-		// the database into account to start at a more "educated" location.
-		prev4: network4,
-		prev6: network6,
-	}, nil
+	return &ret, nil
 }
 
 func (i *IPAllocator) Next() (types.NodeAddresses, error) {
 	i.mu.Lock()
 	defer i.mu.Unlock()
 
-	v4, err := i.next(i.prev4, i.prefix4)
-	if err != nil {
-		return nil, fmt.Errorf("allocating IPv4 address: %w", err)
+	var ret types.NodeAddresses
+
+	if i.prefix4 != nil {
+		v4, err := i.next(i.prev4, i.prefix4)
+		if err != nil {
+			return nil, fmt.Errorf("allocating IPv4 address: %w", err)
+		}
+
+		ret = append(ret, *v4)
 	}
 
-	v6, err := i.next(i.prev6, i.prefix6)
-	if err != nil {
-		return nil, fmt.Errorf("allocating IPv6 address: %w", err)
+	if i.prefix6 != nil {
+		v6, err := i.next(i.prev6, i.prefix6)
+		if err != nil {
+			return nil, fmt.Errorf("allocating IPv6 address: %w", err)
+		}
+
+		ret = append(ret, *v6)
 	}
 
-	return types.NodeAddresses{*v4, *v6}, nil
+	return ret, nil
 }
 
 var ErrCouldNotAllocateIP = errors.New("failed to allocate IP")
 
-func (i *IPAllocator) next(prev netip.Addr, prefix netip.Prefix) (*netip.Addr, error) {
+func (i *IPAllocator) next(prev *netip.Addr, prefix *netip.Prefix) (*netip.Addr, error) {
 	// Get the first IP in our prefix
 	ip := prev.Next()
 

diff --git a/hscontrol/db/ip_test.go b/hscontrol/db/ip_test.go
@@ -12,8 +12,9 @@ import (
 )
 
 func TestIPAllocator(t *testing.T) {
-	mpp := func(pref string) netip.Prefix {
-		return netip.MustParsePrefix(pref)
+	mpp := func(pref string) *netip.Prefix {
+		p := netip.MustParsePrefix(pref)
+		return &p
 	}
 	na := func(pref string) netip.Addr {
 		return netip.MustParseAddr(pref)
@@ -40,8 +41,8 @@ func TestIPAllocator(t *testing.T) {
 		name   string
 		dbFunc func() *HSDatabase
 
-		prefix4  netip.Prefix
-		prefix6  netip.Prefix
+		prefix4  *netip.Prefix
+		prefix6  *netip.Prefix
 		getCount int
 		want     []types.NodeAddresses
 	}{
@@ -63,6 +64,38 @@ func TestIPAllocator(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "simple-v4",
+			dbFunc: func() *HSDatabase {
+				return nil
+			},
+
+			prefix4: mpp("100.64.0.0/10"),
+
+			getCount: 1,
+
+			want: []types.NodeAddresses{
+				{
+					na("100.64.0.1"),
+				},
+			},
+		},
+		{
+			name: "simple-v6",
+			dbFunc: func() *HSDatabase {
+				return nil
+			},
+
+			prefix6: mpp("fd7a:115c:a1e0::/48"),
+
+			getCount: 1,
+
+			want: []types.NodeAddresses{
+				{
+					na("fd7a:115c:a1e0::1"),
+				},
+			},
+		},
 		{
 			name: "simple-with-db",
 			dbFunc: func() *HSDatabase {
@@ -131,8 +164,6 @@ func TestIPAllocator(t *testing.T) {
 
 			spew.Dump(alloc)
 
-			t.Logf("prefixes: %q, %q", tt.prefix4.String(), tt.prefix6.String())
-
 			var got []types.NodeAddresses
 
 			for range tt.getCount {

diff --git a/hscontrol/types/config.go b/hscontrol/types/config.go
@@ -579,18 +579,16 @@ func GetDNSConfig() (*tailcfg.DNSConfig, string) {
 	return nil, ""
 }
 
-func Prefixes() (*netip.Prefix, *netip.Prefix, error) {
+func PrefixV4() (*netip.Prefix, error) {
 	prefixV4Str := viper.GetString("prefixes.v4")
-	prefixV6Str := viper.GetString("prefixes.v6")
 
-	prefixV4, err := netip.ParsePrefix(prefixV4Str)
-	if err != nil {
-		return nil, nil, err
+	if prefixV4Str == "" {
+		return nil, nil
 	}
 
-	prefixV6, err := netip.ParsePrefix(prefixV6Str)
+	prefixV4, err := netip.ParsePrefix(prefixV4Str)
 	if err != nil {
-		return nil, nil, err
+		return nil, fmt.Errorf("parsing IPv4 prefix from config: %w", err)
 	}
 
 	builder := netipx.IPSetBuilder{}
@@ -603,13 +601,33 @@ func Prefixes() (*netip.Prefix, *netip.Prefix, error) {
 				prefixV4Str, tsaddr.CGNATRange())
 	}
 
+	return &prefixV4, nil
+}
+
+func PrefixV6() (*netip.Prefix, error) {
+	prefixV6Str := viper.GetString("prefixes.v6")
+
+	if prefixV6Str == "" {
+		return nil, nil
+	}
+
+	prefixV6, err := netip.ParsePrefix(prefixV6Str)
+	if err != nil {
+		return nil, fmt.Errorf("parsing IPv6 prefix from config: %w", err)
+	}
+
+	builder := netipx.IPSetBuilder{}
+	builder.AddPrefix(tsaddr.CGNATRange())
+	builder.AddPrefix(tsaddr.TailscaleULARange())
+	ipSet, _ := builder.IPSet()
+
 	if !ipSet.ContainsPrefix(prefixV6) {
 		log.Warn().
 			Msgf("Prefix %s is not in the %s range. This is an unsupported configuration.",
 				prefixV6Str, tsaddr.TailscaleULARange())
 	}
 
-	return &prefixV4, &prefixV6, nil
+	return &prefixV6, nil
 }
 
 func GetHeadscaleConfig() (*Config, error) {
@@ -624,7 +642,12 @@ func GetHeadscaleConfig() (*Config, error) {
 		}, nil
 	}
 
-	prefix4, prefix6, err := Prefixes()
+	prefix4, err := PrefixV4()
+	if err != nil {
+		return nil, err
+	}
+
+	prefix6, err := PrefixV6()
 	if err != nil {
 		return nil, err
 	}