From 137b327795ada70918ce934c5475e2d7dc7da46c Mon Sep 17 00:00:00 2001
From: Josef Citrine <josef@citrine.dev>
Date: Fri, 1 Oct 2021 13:10:28 +0100
Subject: [PATCH] Preauthkey can assign machine to subnet

---
 api.go                           | 23 ++++++++++++++++++++++-
 cmd/headscale/cli/preauthkeys.go |  4 +++-
 preauth_keys.go                  | 12 ++++++++++--
 utils.go                         |  4 +++-
 4 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/api.go b/api.go
index e2a56185..16ea8a8d 100644
--- a/api.go
+++ b/api.go
@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/rs/zerolog/log"
+	"inet.af/netaddr"
 
 	"github.com/gin-gonic/gin"
 	"github.com/klauspost/compress/zstd"
@@ -342,7 +343,27 @@ func (h *Headscale) handleAuthKey(c *gin.Context, db *gorm.DB, idKey wgkey.Key,
 		Str("func", "handleAuthKey").
 		Str("machine", m.Name).
 		Msg("Authentication key was valid, proceeding to acquire an IP address")
-	ip, err := h.getAvailableIP()
+
+	var prefix netaddr.IPPrefix
+	var prefixErr error
+	if pak.Subnet != "" {
+		prefix, prefixErr = netaddr.ParseIPPrefix(pak.Subnet)
+
+		if prefixErr != nil {
+			log.Debug().
+				Str("func", "handleAuthKey").
+				Str("machine", m.Name).
+				Msg("Subnet was not valid, using default")
+		}
+	}
+
+	var ip *netaddr.IP
+	if pak.Subnet != "" && prefixErr == nil {
+		ip, err = h.getAvailableIPForPrefix(prefix)
+	} else {
+		ip, err = h.getAvailableIP()
+	}
+
 	if err != nil {
 		log.Error().
 			Str("func", "handleAuthKey").
diff --git a/cmd/headscale/cli/preauthkeys.go b/cmd/headscale/cli/preauthkeys.go
index 1340267e..2b99fa65 100644
--- a/cmd/headscale/cli/preauthkeys.go
+++ b/cmd/headscale/cli/preauthkeys.go
@@ -25,6 +25,7 @@ func init() {
 	createPreAuthKeyCmd.PersistentFlags().Bool("reusable", false, "Make the preauthkey reusable")
 	createPreAuthKeyCmd.PersistentFlags().Bool("ephemeral", false, "Preauthkey for ephemeral nodes")
 	createPreAuthKeyCmd.Flags().StringP("expiration", "e", "", "Human-readable expiration of the key (30m, 24h, 365d...)")
+	createPreAuthKeyCmd.Flags().StringP("subnet", "s", "", "Subnet to assign new nodes to")
 }
 
 var preauthkeysCmd = &cobra.Command{
@@ -116,7 +117,8 @@ var createPreAuthKeyCmd = &cobra.Command{
 			expiration = &exp
 		}
 
-		k, err := h.CreatePreAuthKey(n, reusable, ephemeral, expiration)
+		subnet, _ := cmd.Flags().GetString("subnet")
+		k, err := h.CreatePreAuthKeyWithSubnet(n, reusable, ephemeral, expiration, subnet)
 		if strings.HasPrefix(o, "json") {
 			JsonOutput(k, err, o)
 			return
diff --git a/preauth_keys.go b/preauth_keys.go
index cc849fc0..eda2ee6e 100644
--- a/preauth_keys.go
+++ b/preauth_keys.go
@@ -20,14 +20,20 @@ type PreAuthKey struct {
 	NamespaceID uint
 	Namespace   Namespace
 	Reusable    bool
-	Ephemeral   bool `gorm:"default:false"`
+	Ephemeral   bool   `gorm:"default:false"`
+	Subnet      string `gorm:"default:''"`
 
 	CreatedAt  *time.Time
 	Expiration *time.Time
 }
 
-// CreatePreAuthKey creates a new PreAuthKey in a namespace, and returns it
+// CreatePreAuthKey creates a new PreAuthKey in a namespace for the default subnet, and returns it
 func (h *Headscale) CreatePreAuthKey(namespaceName string, reusable bool, ephemeral bool, expiration *time.Time) (*PreAuthKey, error) {
+	return h.CreatePreAuthKeyWithSubnet(namespaceName, reusable, ephemeral, expiration, "")
+}
+
+// CreatePreAuthKey creates a new PreAuthKey in a namespace with a subnet, and returns it
+func (h *Headscale) CreatePreAuthKeyWithSubnet(namespaceName string, reusable bool, ephemeral bool, expiration *time.Time, subnet string) (*PreAuthKey, error) {
 	n, err := h.GetNamespace(namespaceName)
 	if err != nil {
 		return nil, err
@@ -47,6 +53,7 @@ func (h *Headscale) CreatePreAuthKey(namespaceName string, reusable bool, epheme
 		Ephemeral:   ephemeral,
 		CreatedAt:   &now,
 		Expiration:  expiration,
+		Subnet:      subnet,
 	}
 	h.db.Save(&k)
 
@@ -55,6 +62,7 @@ func (h *Headscale) CreatePreAuthKey(namespaceName string, reusable bool, epheme
 
 // GetPreAuthKeys returns the list of PreAuthKeys for a namespace
 func (h *Headscale) GetPreAuthKeys(namespaceName string) (*[]PreAuthKey, error) {
+	h.getAvailableIP() // temp
 	n, err := h.GetNamespace(namespaceName)
 	if err != nil {
 		return nil, err
diff --git a/utils.go b/utils.go
index cbe1d870..a375d90a 100644
--- a/utils.go
+++ b/utils.go
@@ -75,8 +75,10 @@ func encodeMsg(b []byte, pubKey *wgkey.Key, privKey *wgkey.Private) ([]byte, err
 }
 
 func (h *Headscale) getAvailableIP() (*netaddr.IP, error) {
-	ipPrefix := h.cfg.IPPrefix
+	return h.getAvailableIPForPrefix(h.cfg.IPPrefix)
+}
 
+func (h *Headscale) getAvailableIPForPrefix(ipPrefix netaddr.IPPrefix) (*netaddr.IP, error) {
 	usedIps, err := h.getUsedIPs()
 	if err != nil {
 		return nil, err