From 8b5e8b7dfc2b8915cbb2172a0bdf579bbea5122a Mon Sep 17 00:00:00 2001
From: Mevan Samaratunga <mevansam@gmail.com>
Date: Thu, 10 Mar 2022 08:59:28 -0500
Subject: [PATCH] Refresh expired machine on re-auth - closes #489

---
 api.go | 79 +++++++++++++++++++++++++++++++++++++---------------------
 1 file changed, 51 insertions(+), 28 deletions(-)

diff --git a/api.go b/api.go
index 1023e6f..323e3e7 100644
--- a/api.go
+++ b/api.go
@@ -568,8 +568,13 @@ func (h *Headscale) handleAuthKey(
 			Str("func", "handleAuthKey").
 			Str("machine", registerRequest.Hostinfo.Hostname).
 			Msg("Failed authentication via AuthKey")
-		machineRegistrations.WithLabelValues("new", RegisterMethodAuthKey, "error", pak.Namespace.Name).
-			Inc()
+
+		if pak != nil {
+			machineRegistrations.WithLabelValues("new", RegisterMethodAuthKey, "error", pak.Namespace.Name).
+				Inc()
+		} else {
+			machineRegistrations.WithLabelValues("new", RegisterMethodAuthKey, "error").Inc()
+		}
 
 		return
 	}
@@ -580,35 +585,53 @@ func (h *Headscale) handleAuthKey(
 		Msg("Authentication key was valid, proceeding to acquire IP addresses")
 
 	nodeKey := NodePublicKeyStripPrefix(registerRequest.NodeKey)
-	now := time.Now().UTC()
 
-	machineToRegister := Machine{
-		Name:           registerRequest.Hostinfo.Hostname,
-		NamespaceID:    pak.Namespace.ID,
-		MachineKey:     machineKeyStr,
-		RegisterMethod: RegisterMethodAuthKey,
-		Expiry:         &registerRequest.Expiry,
-		NodeKey:        nodeKey,
-		LastSeen:       &now,
-		AuthKeyID:      uint(pak.ID),
-	}
-
-	machine, err := h.RegisterMachine(
-		machineToRegister,
-	)
-	if err != nil {
-		log.Error().
+	// retrieve machine information if it exist
+	// The error is not important, because if it does not
+	// exist, then this is a new machine and we will move
+	// on to registration.
+	machine, _ := h.GetMachineByMachineKey(machineKey)
+	if machine != nil {
+		log.Trace().
 			Caller().
-			Err(err).
-			Msg("could not register machine")
-		machineRegistrations.WithLabelValues("new", RegisterMethodAuthKey, "error", pak.Namespace.Name).
-			Inc()
-		ctx.String(
-			http.StatusInternalServerError,
-			"could not register machine",
-		)
+			Str("machine", machine.Name).
+			Msg("machine already registered, refreshing with new auth key")
 
-		return
+		machine.NodeKey = nodeKey
+		machine.AuthKeyID = uint(pak.ID)
+		h.RefreshMachine(machine, registerRequest.Expiry)
+
+	} else {
+
+		now := time.Now().UTC()
+		machineToRegister := Machine{
+			Name:           registerRequest.Hostinfo.Hostname,
+			NamespaceID:    pak.Namespace.ID,
+			MachineKey:     machineKeyStr,
+			RegisterMethod: RegisterMethodAuthKey,
+			Expiry:         &registerRequest.Expiry,
+			NodeKey:        nodeKey,
+			LastSeen:       &now,
+			AuthKeyID:      uint(pak.ID),
+		}
+	
+		machine, err = h.RegisterMachine(
+			machineToRegister,
+		)
+		if err != nil {
+			log.Error().
+				Caller().
+				Err(err).
+				Msg("could not register machine")
+			machineRegistrations.WithLabelValues("new", RegisterMethodAuthKey, "error", pak.Namespace.Name).
+				Inc()
+			ctx.String(
+				http.StatusInternalServerError,
+				"could not register machine",
+			)
+	
+			return
+		}			
 	}
 
 	h.UsePreAuthKey(pak)