add option to randomly assign IPs from prefix

Fixes #968
Updates #1828

Signed-off-by: Kristoffer Dalby <[email protected]>

CHANGELOG.md

@@ -39,6 +39,7 @@ after improving the test harness as part of adopting [#1460](https://github.com/
   - `/var/lib/headscale` and `/var/run/headscale` is no longer created automatically, see [container docs](./docs/running-headscale-container.md)
 - Prefixes are now defined per v4 and v6 range. [#1756](https://github.com/juanfont/headscale/pull/1756)
   - `ip_prefixes` option is now `prefixes.v4` and `prefixes.v6`
+  - `prefixes.random_allocation` can be set to assign IPs randomly from the prefix, rather than sequentially. [#1869](https://github.com/juanfont/headscale/pull/1869)
 
 ### Changes
 

config-example.yaml

@@ -61,6 +61,9 @@ prefixes:
   v6: fd7a:115c:a1e0::/48
   v4: 100.64.0.0/10
 
+  # Allocate IP addresses randomly from the prefixes above instead of (best effort) sequentially.
+  random_allocation: false
+
 # DERP is a relay system that Tailscale uses when a direct
 # connection cannot be established.
 # https://tailscale.com/blog/how-tailscale-works/#encrypted-tcp-relays-derp

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, cfg.RandomAllocation)
 	if err != nil {
 		return nil, err
 	}

hscontrol/db/ip.go

@@ -1,8 +1,10 @@
 package db
 
 import (
+	"crypto/rand"
 	"errors"
 	"fmt"
+	"math/big"
 	"net/netip"
 	"sync"
 
@@ -27,6 +29,10 @@ type IPAllocator struct {
 	prev4 *netip.Addr
 	prev6 *netip.Addr
 
+	// randomAllocation changes the allocation
+	// strategy from sequential to random.
+	randomAllocation bool
+
 	// Set of all IPs handed out.
 	// This might not be in sync with the database,
 	// but it is more conservative. If saves to the
@@ -40,10 +46,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, randomAlloc bool) (*IPAllocator, error) {
 	ret := IPAllocator{
 		prefix4: prefix4,
 		prefix6: prefix6,
+
+		randomAllocation: randomAlloc,
 	}
 
 	var v4s []string
@@ -143,8 +151,18 @@ func (i *IPAllocator) Next() (*netip.Addr, *netip.Addr, error) {
 var ErrCouldNotAllocateIP = errors.New("failed to allocate IP")
 
 func (i *IPAllocator) next(prev *netip.Addr, prefix *netip.Prefix) (*netip.Addr, error) {
-	// Get the first IP in our prefix
-	ip := prev.Next()
+	var err error
+	var ip netip.Addr
+
+	if i.randomAllocation {
+		ip, err = randomNext(*prefix)
+		if err != nil {
+			return nil, fmt.Errorf("getting random IP: %w", err)
+		}
+	} else {
+		// Get the first IP in our prefix
+		ip = prev.Next()
+	}
 
 	// TODO(kradalby): maybe this can be done less often.
 	set, err := i.usedIPs.IPSet()
@@ -159,7 +177,15 @@ func (i *IPAllocator) next(prev *netip.Addr, prefix *netip.Prefix) (*netip.Addr,
 
 		// Check if the IP has already been allocated.
 		if set.Contains(ip) {
-			ip = ip.Next()
+			if i.randomAllocation {
+				ip, err = randomNext(*prefix)
+				if err != nil {
+					return nil, fmt.Errorf("getting random IP: %w", err)
+				}
+			} else {
+				// Get the first IP in our prefix
+				ip = prev.Next()
+			}
 
 			continue
 		}
@@ -169,3 +195,36 @@ func (i *IPAllocator) next(prev *netip.Addr, prefix *netip.Prefix) (*netip.Addr,
 		return &ip, nil
 	}
 }
+
+func randomNext(pfx netip.Prefix) (netip.Addr, error) {
+	rang := netipx.RangeOfPrefix(pfx)
+	fromIP, toIP := rang.From(), rang.To()
+
+	var from, to big.Int
+
+	from.SetBytes(fromIP.AsSlice())
+	to.SetBytes(toIP.AsSlice())
+
+	// Find the max, this is how we can do "random range",
+	// get the "max" as 0 -> to - from and then add back from
+	// after.
+	tempMax := big.NewInt(0).Sub(&to, &from)
+
+	out, err := rand.Int(rand.Reader, tempMax)
+	if err != nil {
+		return netip.Addr{}, fmt.Errorf("generating random IP: %w", err)
+	}
+
+	valInRange := big.NewInt(0).Add(&from, out)
+
+	ip, ok := netip.AddrFromSlice(valInRange.Bytes())
+	if !ok {
+		return netip.Addr{}, fmt.Errorf("generated ip bytes are invalid ip")
+	}
+
+	if !pfx.Contains(ip) {
+		return netip.Addr{}, fmt.Errorf("generated ip(%s) not in prefix(%s)", ip.String(), pfx.String())
+	}
+
+	return ip, nil
+}

hscontrol/db/ip_test.go

@@ -3,7 +3,6 @@ package db
 import (
 	"database/sql"
 	"net/netip"
-	"os"
 	"testing"
 
 	"github.com/davecgh/go-spew/spew"
@@ -12,32 +11,15 @@ import (
 	"github.com/juanfont/headscale/hscontrol/util"
 )
 
-func TestIPAllocator(t *testing.T) {
-	mpp := func(pref string) *netip.Prefix {
-		p := netip.MustParsePrefix(pref)
-		return &p
-	}
-	na := func(pref string) netip.Addr {
-		return netip.MustParseAddr(pref)
-	}
-	newDb := func() *HSDatabase {
-		tmpDir, err := os.MkdirTemp("", "headscale-db-test-*")
-		if err != nil {
-			t.Fatalf("creating temp dir: %s", err)
-		}
-		db, _ = NewHeadscaleDatabase(
-			types.DatabaseConfig{
-				Type: "sqlite3",
-				Sqlite: types.SqliteConfig{
-					Path: tmpDir + "/headscale_test.db",
-				},
-			},
-			"",
-		)
-
-		return db
-	}
+var mpp = func(pref string) *netip.Prefix {
+	p := netip.MustParsePrefix(pref)
+	return &p
+}
+var na = func(pref string) netip.Addr {
+	return netip.MustParseAddr(pref)
+}
 
+func TestIPAllocatorSequential(t *testing.T) {
 	tests := []struct {
 		name   string
 		dbFunc func() *HSDatabase
@@ -97,7 +79,7 @@ func TestIPAllocator(t *testing.T) {
 		{
 			name: "simple-with-db",
 			dbFunc: func() *HSDatabase {
-				db := newDb()
+				db := dbForTest(t, "simple-with-db")
 
 				db.DB.Save(&types.Node{
 					IPv4DatabaseField: sql.NullString{
@@ -128,7 +110,7 @@ func TestIPAllocator(t *testing.T) {
 		{
 			name: "before-after-free-middle-in-db",
 			dbFunc: func() *HSDatabase {
-				db := newDb()
+				db := dbForTest(t, "before-after-free-middle-in-db")
 
 				db.DB.Save(&types.Node{
 					IPv4DatabaseField: sql.NullString{
@@ -164,7 +146,7 @@ func TestIPAllocator(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			db := tt.dbFunc()
 
-			alloc, _ := NewIPAllocator(db, tt.prefix4, tt.prefix6)
+			alloc, _ := NewIPAllocator(db, tt.prefix4, tt.prefix6, false)
 
 			spew.Dump(alloc)
 
@@ -195,3 +177,103 @@ func TestIPAllocator(t *testing.T) {
 		})
 	}
 }
+
+func TestIPAllocatorRandom(t *testing.T) {
+	tests := []struct {
+		name   string
+		dbFunc func() *HSDatabase
+
+		getCount int
+
+		prefix4 *netip.Prefix
+		prefix6 *netip.Prefix
+		want4   bool
+		want6   bool
+	}{
+		{
+			name: "simple",
+			dbFunc: func() *HSDatabase {
+				return nil
+			},
+
+			prefix4: mpp("100.64.0.0/10"),
+			prefix6: mpp("fd7a:115c:a1e0::/48"),
+
+			getCount: 1,
+
+			want4: true,
+			want6: true,
+		},
+		{
+			name: "simple-v4",
+			dbFunc: func() *HSDatabase {
+				return nil
+			},
+
+			prefix4: mpp("100.64.0.0/10"),
+
+			getCount: 1,
+
+			want4: true,
+			want6: false,
+		},
+		{
+			name: "simple-v6",
+			dbFunc: func() *HSDatabase {
+				return nil
+			},
+
+			prefix6: mpp("fd7a:115c:a1e0::/48"),
+
+			getCount: 1,
+
+			want4: false,
+			want6: true,
+		},
+		{
+			name: "generate-lots-of-random",
+			dbFunc: func() *HSDatabase {
+				return nil
+			},
+
+			prefix4: mpp("100.64.0.0/10"),
+			prefix6: mpp("fd7a:115c:a1e0::/48"),
+
+			getCount: 1000,
+
+			want4: true,
+			want6: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			db := tt.dbFunc()
+
+			alloc, _ := NewIPAllocator(db, tt.prefix4, tt.prefix6, true)
+
+			spew.Dump(alloc)
+
+			for range tt.getCount {
+				got4, got6, err := alloc.Next()
+				if err != nil {
+					t.Fatalf("allocating next IP: %s", err)
+				}
+
+				t.Logf("addrs ipv4: %v, ipv6: %v", got4, got6)
+
+				if tt.want4 {
+					if got4 == nil {
+						t.Fatalf("expected ipv4 addr, got nil")
+					}
+				}
+
+				if tt.want6 {
+					if got6 == nil {
+						t.Fatalf("expected ipv4 addr, got nil")
+					}
+				}
+			}
+		})
+	}
+}

hscontrol/types/config.go

@@ -42,6 +42,7 @@ type Config struct {
 	NodeUpdateCheckInterval        time.Duration
 	PrefixV4                       *netip.Prefix
 	PrefixV6                       *netip.Prefix
+	RandomAllocation               bool
 	NoisePrivateKeyPath            string
 	BaseDomain                     string
 	Log                            LogConfig
@@ -678,8 +679,9 @@ func GetHeadscaleConfig() (*Config, error) {
 		GRPCAllowInsecure:  viper.GetBool("grpc_allow_insecure"),
 		DisableUpdateCheck: viper.GetBool("disable_check_updates"),
 
-		PrefixV4: prefix4,
-		PrefixV6: prefix6,
+		PrefixV4:         prefix4,
+		PrefixV6:         prefix6,
+		RandomAllocation: viper.GetBool("prefixes.random_allocation"),
 
 		NoisePrivateKeyPath: util.AbsolutePathFromConfigPath(
 			viper.GetString("noise.private_key_path"),

integration/hsic/config.go

@@ -117,6 +117,7 @@ func DefaultConfigEnv() map[string]string {
 		"HEADSCALE_NODE_UPDATE_CHECK_INTERVAL":        "10s",
 		"HEADSCALE_PREFIXES_V4":                       "100.64.0.0/10",
 		"HEADSCALE_PREFIXES_V6":                       "fd7a:115c:a1e0::/48",
+		"HEADSCALE_PREFIXES_RANDOM_ALLOCATION":        "1",
 		"HEADSCALE_DNS_CONFIG_BASE_DOMAIN":            "headscale.net",
 		"HEADSCALE_DNS_CONFIG_MAGIC_DNS":              "true",
 		"HEADSCALE_DNS_CONFIG_DOMAINS":                "",