From 7157e14aff341b3b643bdbeae6a3a17373f5a27f Mon Sep 17 00:00:00 2001
From: Even Holthe <even.holthe@bekk.no>
Date: Thu, 15 Dec 2022 01:10:26 +0100
Subject: [PATCH] add expiration from OIDC token to machine

---
 grpcv1.go  | 1 +
 machine.go | 5 +++++
 oidc.go    | 4 +++-
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/grpcv1.go b/grpcv1.go
index 7b12fa05..c3936df6 100644
--- a/grpcv1.go
+++ b/grpcv1.go
@@ -176,6 +176,7 @@ func (api headscaleV1APIServer) RegisterMachine(
 	machine, err := api.h.RegisterMachineFromAuthCallback(
 		request.GetKey(),
 		request.GetNamespace(),
+		nil,
 		RegisterMethodCLI,
 	)
 	if err != nil {
diff --git a/machine.go b/machine.go
index 3ddf471c..79485f7d 100644
--- a/machine.go
+++ b/machine.go
@@ -852,6 +852,7 @@ func getTags(
 func (h *Headscale) RegisterMachineFromAuthCallback(
 	nodeKeyStr string,
 	namespaceName string,
+	machineExpiry *time.Time,
 	registrationMethod string,
 ) (*Machine, error) {
 	nodeKey := key.NodePublic{}
@@ -885,6 +886,10 @@ func (h *Headscale) RegisterMachineFromAuthCallback(
 			registrationMachine.NamespaceID = namespace.ID
 			registrationMachine.RegisterMethod = registrationMethod
 
+			if machineExpiry != nil {
+				registrationMachine.Expiry = machineExpiry
+			}
+
 			machine, err := h.RegisterMachine(
 				registrationMachine,
 			)
diff --git a/oidc.go b/oidc.go
index 3eed9187..8c7e8304 100644
--- a/oidc.go
+++ b/oidc.go
@@ -236,7 +236,7 @@ func (h *Headscale) OIDCCallback(
 		return
 	}
 
-	if err := h.registerMachineForOIDCCallback(writer, namespace, nodeKey); err != nil {
+	if err := h.registerMachineForOIDCCallback(writer, namespace, nodeKey, idToken.Expiry); err != nil {
 		return
 	}
 
@@ -679,10 +679,12 @@ func (h *Headscale) registerMachineForOIDCCallback(
 	writer http.ResponseWriter,
 	namespace *Namespace,
 	nodeKey *key.NodePublic,
+	expiry time.Time,
 ) error {
 	if _, err := h.RegisterMachineFromAuthCallback(
 		nodeKey.String(),
 		namespace.Name,
+		&expiry,
 		RegisterMethodOIDC,
 	); err != nil {
 		log.Error().