From c6ea9b4b8042f53773bd19e1d63f10a938b12dce Mon Sep 17 00:00:00 2001
From: Laurent Marchaud <laurent@marchaud.com>
Date: Fri, 19 Aug 2022 12:38:39 +0200
Subject: [PATCH] Fix cli registration of expired machines

Signed-off-by: Laurent Marchaud <laurent@marchaud.com>
---
 api.go     |  7 +++++++
 machine.go | 14 ++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/api.go b/api.go
index 2d55ccd3..1bf3edb3 100644
--- a/api.go
+++ b/api.go
@@ -346,6 +346,13 @@ func (h *Headscale) RegistrationHandler(
 		// The machine has expired
 		h.handleMachineExpired(writer, req, machineKey, registerRequest, *machine)
 
+		machine.Expiry = &time.Time{}
+		h.registrationCache.Set(
+			machineKeyStr,
+			*machine,
+			registerCacheExpiration,
+		)
+
 		return
 	}
 }
diff --git a/machine.go b/machine.go
index aebfbcef..a877a5f0 100644
--- a/machine.go
+++ b/machine.go
@@ -24,8 +24,9 @@ const (
 	ErrMachineNotFoundRegistrationCache = Error(
 		"machine not found in registration cache",
 	)
-	ErrCouldNotConvertMachineInterface = Error("failed to convert machine interface")
-	ErrHostnameTooLong                 = Error("Hostname too long")
+	errCouldNotConvertMachineInterface = Error("failed to convert machine interface")
+	errHostnameTooLong                 = Error("Hostname too long")
+	errDifferentRegisteredNamespace    = Error("machine was previously registered with a different namespace")
 	MachineGivenNameHashLength         = 8
 	MachineGivenNameTrimSize           = 2
 )
@@ -789,6 +790,11 @@ func (h *Headscale) RegisterMachineFromAuthCallback(
 				)
 			}
 
+			// Registration of expired machine with different namespace
+			if registrationMachine.ID != 0 && registrationMachine.NamespaceID != namespace.ID {
+				return nil, errDifferentRegisteredNamespace
+			}
+
 			registrationMachine.NamespaceID = namespace.ID
 			registrationMachine.RegisterMethod = registrationMethod
 
@@ -796,6 +802,10 @@ func (h *Headscale) RegisterMachineFromAuthCallback(
 				registrationMachine,
 			)
 
+			if err == nil {
+				h.registrationCache.Delete(machineKeyStr)
+			}
+
 			return machine, err
 		} else {
 			return nil, ErrCouldNotConvertMachineInterface