add option to randomly assign IPs from prefix

Fixes juanfont#968 Updates juanfont#1828 Signed-off-by: Kristoffer Dalby <[email protected]>
kradalby · Apr 15, 2024 · b6861b6 · b6861b6
1 parent c1a05f4
commit b6861b6
Show file tree

Hide file tree

Showing 7 changed files with 203 additions and 57 deletions.
diff --git a/CHANGELOG.md b/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
 

diff --git a/config-example.yaml b/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

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

diff --git a/hscontrol/db/ip.go b/hscontrol/db/ip.go
@@ -1,9 +1,11 @@
 package db
 
 import (
+	"crypto/rand"
 	"database/sql"
 	"errors"
 	"fmt"
+	"math/big"
 	"net/netip"
 	"sync"
 
@@ -25,8 +27,12 @@ type IPAllocator struct {
 	prefix6 *netip.Prefix
 
 	// Previous IPs handed out
-	prev4 *netip.Addr
-	prev6 *netip.Addr
+	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,
@@ -41,10 +47,16 @@ 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 []sql.NullString
@@ -79,15 +91,15 @@ func NewIPAllocator(db *HSDatabase, prefix4, prefix6 *netip.Prefix) (*IPAllocato
 		// 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
+		ret.prev4 = network4
 	}
 
 	if prefix6 != nil {
 		network6, broadcast6 := util.GetIPPrefixEndpoints(*prefix6)
 		ips.Add(network6)
 		ips.Add(broadcast6)
 
-		ret.prev6 = &network6
+		ret.prev6 = network6
 	}
 
 	// Fetch all the IP Addresses currently handed out from the Database
@@ -130,14 +142,15 @@ func (i *IPAllocator) Next() (*netip.Addr, *netip.Addr, error) {
 		if err != nil {
 			return nil, nil, fmt.Errorf("allocating IPv4 address: %w", err)
 		}
-
+		i.prev4 = *ret4
 	}
 
 	if i.prefix6 != nil {
 		ret6, err = i.next(i.prev6, i.prefix6)
 		if err != nil {
 			return nil, nil, fmt.Errorf("allocating IPv6 address: %w", err)
 		}
+		i.prev6 = *ret6
 	}
 
 	return ret4, ret6, nil
@@ -146,8 +159,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()
@@ -162,7 +185,14 @@ 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 {
+				ip = ip.Next()
+			}
 
 			continue
 		}
@@ -172,3 +202,40 @@ 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
+}
diff --git a/hscontrol/db/ip_test.go b/hscontrol/db/ip_test.go
@@ -2,8 +2,9 @@ package db
 
 import (
 	"database/sql"
+	"fmt"
 	"net/netip"
-	"os"
+	"strings"
 	"testing"
 
 	"github.com/davecgh/go-spew/spew"
@@ -12,32 +13,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,17 +81,11 @@ 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{
-						Valid:  true,
-						String: "100.64.0.1",
-					},
-					IPv6DatabaseField: sql.NullString{
-						Valid:  true,
-						String: "fd7a:115c:a1e0::1",
-					},
+					IPv4: nap("100.64.0.1"),
+					IPv6: nap("fd7a:115c:a1e0::1"),
 				})
 
 				return db
@@ -128,17 +106,11 @@ 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{
-						Valid:  true,
-						String: "100.64.0.2",
-					},
-					IPv6DatabaseField: sql.NullString{
-						Valid:  true,
-						String: "fd7a:115c:a1e0::2",
-					},
+					IPv4: nap("100.64.0.2"),
+					IPv6: nap("fd7a:115c:a1e0::2"),
 				})
 
 				return db
@@ -164,7 +136,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 +167,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")
+					}
+				}
+			}
+		})
+	}
+}