From bd1d1b1a3b859f155012710cc2681f23d37707d5 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Sun, 21 Nov 2021 13:40:19 +0000
Subject: [PATCH] Implement ExpireMachine rpc

---
 grpcv1.go       | 21 +++++++++++++++++++++
 machine.go      |  8 ++++++++
 machine_test.go | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+)

diff --git a/grpcv1.go b/grpcv1.go
index 40419c3c..1850ce7a 100644
--- a/grpcv1.go
+++ b/grpcv1.go
@@ -201,6 +201,27 @@ func (api headscaleV1APIServer) DeleteMachine(
 	return &v1.DeleteMachineResponse{}, nil
 }
 
+func (api headscaleV1APIServer) ExpireMachine(
+	ctx context.Context,
+	request *v1.ExpireMachineRequest,
+) (*v1.ExpireMachineResponse, error) {
+	machine, err := api.h.GetMachineByID(request.GetMachineId())
+	if err != nil {
+		return nil, err
+	}
+
+	api.h.ExpireMachine(
+		machine,
+	)
+
+	log.Trace().
+		Str("machine", machine.Name).
+		Time("expiry", *machine.Expiry).
+		Msg("machine expired")
+
+	return &v1.ExpireMachineResponse{Machine: machine.toProto()}, nil
+}
+
 func (api headscaleV1APIServer) ListMachines(
 	ctx context.Context,
 	request *v1.ListMachinesRequest,
diff --git a/machine.go b/machine.go
index 813da35c..e34bff87 100644
--- a/machine.go
+++ b/machine.go
@@ -262,6 +262,14 @@ func (h *Headscale) UpdateMachine(machine *Machine) error {
 	return nil
 }
 
+// ExpireMachine takes a Machine struct and sets the expire field to now.
+func (h *Headscale) ExpireMachine(machine *Machine) {
+	now := time.Now()
+	machine.Expiry = &now
+
+	h.db.Save(machine)
+}
+
 // DeleteMachine softs deletes a Machine from the database.
 func (h *Headscale) DeleteMachine(machine *Machine) error {
 	err := h.RemoveSharedMachineFromAllNamespaces(machine)
diff --git a/machine_test.go b/machine_test.go
index fdbc3cb5..eb090072 100644
--- a/machine_test.go
+++ b/machine_test.go
@@ -3,6 +3,7 @@ package headscale
 import (
 	"encoding/json"
 	"strconv"
+	"time"
 
 	"gopkg.in/check.v1"
 )
@@ -164,3 +165,37 @@ func (s *Suite) TestGetDirectPeers(c *check.C) {
 	c.Assert(peersOfMachine0[5].Name, check.Equals, "testmachine7")
 	c.Assert(peersOfMachine0[8].Name, check.Equals, "testmachine10")
 }
+
+func (s *Suite) TestExpireMachine(c *check.C) {
+	namespace, err := app.CreateNamespace("test")
+	c.Assert(err, check.IsNil)
+
+	pak, err := app.CreatePreAuthKey(namespace.Name, false, false, nil)
+	c.Assert(err, check.IsNil)
+
+	_, err = app.GetMachine("test", "testmachine")
+	c.Assert(err, check.NotNil)
+
+	machine := &Machine{
+		ID:             0,
+		MachineKey:     "foo",
+		NodeKey:        "bar",
+		DiscoKey:       "faa",
+		Name:           "testmachine",
+		NamespaceID:    namespace.ID,
+		Registered:     true,
+		RegisterMethod: RegisterMethodAuthKey,
+		AuthKeyID:      uint(pak.ID),
+		Expiry:         &time.Time{},
+	}
+	app.db.Save(machine)
+
+	machineFromDB, err := app.GetMachine("test", "testmachine")
+	c.Assert(err, check.IsNil)
+
+	c.Assert(machineFromDB.isExpired(), check.Equals, false)
+
+	app.ExpireMachine(machineFromDB)
+
+	c.Assert(machineFromDB.isExpired(), check.Equals, true)
+}