From 3eeb2757e0d2660b0e1f502d9ee39f2ca5b8bd04 Mon Sep 17 00:00:00 2001
From: Josef Citrine <josef@citrine.dev>
Date: Fri, 26 Nov 2021 15:31:39 +0000
Subject: [PATCH] Parse IP with netaddr

---
 cmd/headscale/cli/preauthkeys.go | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/cmd/headscale/cli/preauthkeys.go b/cmd/headscale/cli/preauthkeys.go
index 7426d4bc..b7c63182 100644
--- a/cmd/headscale/cli/preauthkeys.go
+++ b/cmd/headscale/cli/preauthkeys.go
@@ -3,7 +3,6 @@ package cli
 import (
 	"fmt"
 	"strconv"
-	"strings"
 	"time"
 
 	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
@@ -11,6 +10,7 @@ import (
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 	"google.golang.org/protobuf/types/known/timestamppb"
+	"inet.af/netaddr"
 )
 
 const (
@@ -142,15 +142,30 @@ var createPreAuthKeyCmd = &cobra.Command{
 
 		subnet, _ := cmd.Flags().GetString("subnet")
 
+		if subnet != "" {
+			ipPrefix, err := netaddr.ParseIPPrefix(subnet)
+
+			if err != nil {
+				ErrorOutput(err, fmt.Sprintf("Error parsing subnet: %s", err), output)
+				return
+			}
+
+			subnet = ipPrefix.String()
+		}
+
 		if !reusable && subnet == "" {
 			ip, _ := cmd.Flags().GetString("ip")
 			if ip != "" {
-				// If IP is in CIDR notation, strip the last octet
-				if strings.Contains(ip, "/") {
-					ip = strings.Split(ip, "/")[0]
+				// Parse IP and convert to prefix
+				ipAddr, err := netaddr.ParseIP(ip)
+
+				if err != nil {
+					ErrorOutput(err, fmt.Sprintf("Error parsing IP address: %s", err), output)
+					return
 				}
 
-				subnet = ip + "/32"
+				ipPrefix := netaddr.IPPrefixFrom(ipAddr, 32)
+				subnet = ipPrefix.String()
 			}
 		}