Implement ExpireMachine rpc

This commit is contained in:
Kristoffer Dalby 2021-11-21 13:40:19 +00:00
parent f1c05f8010
commit bd1d1b1a3b
3 changed files with 64 additions and 0 deletions

View File

@ -201,6 +201,27 @@ func (api headscaleV1APIServer) DeleteMachine(
return &v1.DeleteMachineResponse{}, nil 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( func (api headscaleV1APIServer) ListMachines(
ctx context.Context, ctx context.Context,
request *v1.ListMachinesRequest, request *v1.ListMachinesRequest,

View File

@ -262,6 +262,14 @@ func (h *Headscale) UpdateMachine(machine *Machine) error {
return nil 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. // DeleteMachine softs deletes a Machine from the database.
func (h *Headscale) DeleteMachine(machine *Machine) error { func (h *Headscale) DeleteMachine(machine *Machine) error {
err := h.RemoveSharedMachineFromAllNamespaces(machine) err := h.RemoveSharedMachineFromAllNamespaces(machine)

View File

@ -3,6 +3,7 @@ package headscale
import ( import (
"encoding/json" "encoding/json"
"strconv" "strconv"
"time"
"gopkg.in/check.v1" "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[5].Name, check.Equals, "testmachine7")
c.Assert(peersOfMachine0[8].Name, check.Equals, "testmachine10") 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)
}