From 309f868a2113c9cf1ba193f744ff135808829eb0 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Mon, 2 Aug 2021 20:06:26 +0100
Subject: [PATCH] Make IP prefix configurable

This commit makes the IP prefix used to generate addresses configurable
to users. This can be useful if you would like to use a smaller range or
if your current setup is overlapping with the current range.

The current range is left as a default
---
 app.go                     | 2 ++
 app_test.go                | 5 ++++-
 cmd/headscale/cli/utils.go | 4 ++++
 utils.go                   | 9 +++++----
 4 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/app.go b/app.go
index fa9b011b..e3978212 100644
--- a/app.go
+++ b/app.go
@@ -13,6 +13,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"golang.org/x/crypto/acme/autocert"
 	"gorm.io/gorm"
+	"inet.af/netaddr"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/wgkey"
 )
@@ -24,6 +25,7 @@ type Config struct {
 	PrivateKeyPath                 string
 	DerpMap                        *tailcfg.DERPMap
 	EphemeralNodeInactivityTimeout time.Duration
+	IPPrefix                       netaddr.IPPrefix
 
 	DBtype string
 	DBpath string
diff --git a/app_test.go b/app_test.go
index ad633334..ff3755e9 100644
--- a/app_test.go
+++ b/app_test.go
@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"gopkg.in/check.v1"
+	"inet.af/netaddr"
 )
 
 func Test(t *testing.T) {
@@ -36,7 +37,9 @@ func (s *Suite) ResetDB(c *check.C) {
 	if err != nil {
 		c.Fatal(err)
 	}
-	cfg := Config{}
+	cfg := Config{
+		IPPrefix: netaddr.MustParseIPPrefix("127.0.0.1/32"),
+	}
 
 	h = Headscale{
 		cfg:      cfg,
diff --git a/cmd/headscale/cli/utils.go b/cmd/headscale/cli/utils.go
index 5e47d157..1c259c74 100644
--- a/cmd/headscale/cli/utils.go
+++ b/cmd/headscale/cli/utils.go
@@ -14,6 +14,7 @@ import (
 	"github.com/juanfont/headscale"
 	"github.com/spf13/viper"
 	"gopkg.in/yaml.v2"
+	"inet.af/netaddr"
 	"tailscale.com/tailcfg"
 )
 
@@ -36,6 +37,8 @@ func LoadConfig(path string) error {
 	viper.SetDefault("tls_letsencrypt_cache_dir", "/var/www/.cache")
 	viper.SetDefault("tls_letsencrypt_challenge_type", "HTTP-01")
 
+	viper.SetDefault("ip_prefix", "100.64.0.0/10")
+
 	err := viper.ReadInConfig()
 	if err != nil {
 		return fmt.Errorf("Fatal error reading config file: %s \n", err)
@@ -97,6 +100,7 @@ func getHeadscaleApp() (*headscale.Headscale, error) {
 		Addr:           viper.GetString("listen_addr"),
 		PrivateKeyPath: absPath(viper.GetString("private_key_path")),
 		DerpMap:        derpMap,
+		IPPrefix:       netaddr.MustParseIPPrefix(viper.GetString("ip_prefix")),
 
 		EphemeralNodeInactivityTimeout: viper.GetDuration("ephemeral_node_inactivity_timeout"),
 
diff --git a/utils.go b/utils.go
index f21063b0..1da25084 100644
--- a/utils.go
+++ b/utils.go
@@ -19,6 +19,7 @@ import (
 
 	"golang.org/x/crypto/nacl/box"
 	"gorm.io/gorm"
+	"inet.af/netaddr"
 	"tailscale.com/types/wgkey"
 )
 
@@ -80,7 +81,7 @@ func encodeMsg(b []byte, pubKey *wgkey.Key, privKey *wgkey.Private) ([]byte, err
 func (h *Headscale) getAvailableIP() (*net.IP, error) {
 	i := 0
 	for {
-		ip, err := getRandomIP()
+		ip, err := getRandomIP(h.cfg.IPPrefix)
 		if err != nil {
 			return nil, err
 		}
@@ -93,12 +94,12 @@ func (h *Headscale) getAvailableIP() (*net.IP, error) {
 			break
 		}
 	}
-	return nil, errors.New("Could not find an available IP address in 100.64.0.0/10")
+	return nil, errors.New(fmt.Sprintf("Could not find an available IP address in %s", h.cfg.IPPrefix.String()))
 }
 
-func getRandomIP() (*net.IP, error) {
+func getRandomIP(ipPrefix netaddr.IPPrefix) (*net.IP, error) {
 	mathrand.Seed(time.Now().Unix())
-	ipo, ipnet, err := net.ParseCIDR("100.64.0.0/10")
+	ipo, ipnet, err := net.ParseCIDR(ipPrefix.String())
 	if err == nil {
 		ip := ipo.To4()
 		// fmt.Println("In Randomize IPAddr: IP ", ip, " IPNET: ", ipnet)