From caffbd8956776e97e54bbeb6b18ea30a4c287b40 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Sun, 27 Feb 2022 18:42:43 +0100
Subject: [PATCH] Update cli registration with new method

---
 grpcv1.go | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/grpcv1.go b/grpcv1.go
index 60e181df..b1396bc8 100644
--- a/grpcv1.go
+++ b/grpcv1.go
@@ -159,9 +159,27 @@ func (api headscaleV1APIServer) RegisterMachine(
 		Str("namespace", request.GetNamespace()).
 		Str("machine_key", request.GetKey()).
 		Msg("Registering machine")
+
+	// TODO(kradalby): Currently, if it fails to find a requested expiry, non will be set
+	// This means that if a user is to slow with register a machine, it will possibly not
+	// have the correct expiry.
+	requestedTime := time.Time{}
+	if requestedTimeIf, found := api.h.requestedExpiryCache.Get(request.GetKey()); found {
+		log.Trace().
+			Caller().
+			Str("machine", request.Key).
+			Msg("Expiry time found in cache, assigning to node")
+		if reqTime, ok := requestedTimeIf.(time.Time); ok {
+			requestedTime = reqTime
+		}
+	}
+
 	machine, err := api.h.RegisterMachine(
 		request.GetKey(),
 		request.GetNamespace(),
+		RegisterMethodCLI,
+		&requestedTime,
+		nil, nil, nil,
 	)
 	if err != nil {
 		return nil, err