updates from code review

api.go

@@ -65,7 +65,7 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) {
 			Str("handler", "Registration").
 			Err(err).
 			Msg("Cannot parse machine key")
-		machineRegistrations.WithLabelValues("unkown", "web", "error", "unknown").Inc()
+		machineRegistrations.WithLabelValues("unknown", "web", "error", "unknown").Inc()
 		c.String(http.StatusInternalServerError, "Sad!")
 		return
 	}
@@ -76,34 +76,33 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) {
 			Str("handler", "Registration").
 			Err(err).
 			Msg("Cannot decode message")
-		machineRegistrations.WithLabelValues("unkown", "web", "error", "unknown").Inc()
+		machineRegistrations.WithLabelValues("unknown", "web", "error", "unknown").Inc()
 		c.String(http.StatusInternalServerError, "Very sad!")
 		return
 	}
 
 	now := time.Now().UTC()
-	var m Machine
+	m, err := h.GetMachineByMachineKey(mKey.HexString())
-	if result := h.db.Preload("Namespace").First(&m, "machine_key = ?", mKey.HexString()); errors.Is(result.Error, gorm.ErrRecordNotFound) {
+	if errors.Is(err, gorm.ErrRecordNotFound) {
 		log.Info().Str("machine", req.Hostinfo.Hostname).Msg("New machine")
-		m = Machine{
+		newMachine := Machine{
-			Expiry:               &time.Time{},
+			Expiry:     &time.Time{},
-			MachineKey:           mKey.HexString(),
+			MachineKey: mKey.HexString(),
-			Name:                 req.Hostinfo.Hostname,
+			Name:       req.Hostinfo.Hostname,
-			NodeKey:              wgkey.Key(req.NodeKey).HexString(),
-			LastSuccessfulUpdate: &now,
 		}
-		if err := h.db.Create(&m).Error; err != nil {
+		if err := h.db.Create(&newMachine).Error; err != nil {
 			log.Error().
 				Str("handler", "Registration").
 				Err(err).
 				Msg("Could not create row")
-			machineRegistrations.WithLabelValues("unkown", "web", "error", m.Namespace.Name).Inc()
+			machineRegistrations.WithLabelValues("unknown", "web", "error", m.Namespace.Name).Inc()
 			return
 		}
+		m = &newMachine
 	}
 
 	if !m.Registered && req.Auth.AuthKey != "" {
-		h.handleAuthKey(c, h.db, mKey, req, m)
+		h.handleAuthKey(c, h.db, mKey, req, *m)
 		return
 	}
 
@@ -112,13 +111,14 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) {
 	// We have the updated key!
 	if m.NodeKey == wgkey.Key(req.NodeKey).HexString() {
 
+		// The client sends an Expiry in the past if the client is requesting a logout
 		if !req.Expiry.IsZero() && req.Expiry.UTC().Before(now) {
-			log.Debug().
+			log.Info().
 				Str("handler", "Registration").
 				Str("machine", m.Name).
 				Msg("Client requested logout")
 
-			m.Expiry = &req.Expiry
+			m.Expiry = &req.Expiry // save the expiry so that the machine is marked as expired
 			h.db.Save(&m)
 
 			resp.AuthURL = ""
@@ -138,6 +138,7 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) {
 		}
 
 		if m.Registered && m.Expiry.UTC().After(now) {
+			// The machine registration is valid, respond with redirect to /map
 			log.Debug().
 				Str("handler", "Registration").
 				Str("machine", m.Name).
@@ -161,10 +162,11 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) {
 			return
 		}
 
+		// The client has registered before, but has expired
 		log.Debug().
 			Str("handler", "Registration").
 			Str("machine", m.Name).
-			Msg("Not registered (or expired) and not NodeKey rotation. Sending a authurl to register")
+			Msg("Machine registration has expired. Sending a authurl to register")
 
 		if h.cfg.OIDCIssuer != "" {
 			resp.AuthURL = fmt.Sprintf("%s/oidc/register/%s",
@@ -174,7 +176,7 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) {
 				strings.TrimSuffix(h.cfg.ServerURL, "/"), mKey.HexString())
 		}
 
-		m.Expiry = &req.Expiry // save the requested expiry time for retrieval later
+		m.RequestedExpiry = &req.Expiry // save the requested expiry time for retrieval later in the authentication flow
 		h.db.Save(&m)
 
 		respBody, err := encode(resp, &mKey, h.privateKey)
@@ -216,34 +218,7 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) {
 		return
 	}
 
-	// We arrive here after a client is restarted without finalizing the authentication flow or
+	// The machine registration is new, redirect the client to the registration URL
-	// when headscale is stopped in the middle of the auth process.
-	if m.Registered && m.Expiry.UTC().After(now) {
-		log.Debug().
-			Str("handler", "Registration").
-			Str("machine", m.Name).
-			Msg("The node is sending us a new NodeKey, but machine is registered. All clear for /map")
-
-		m.NodeKey = wgkey.Key(req.NodeKey).HexString()
-		h.db.Save(&m)
-
-		resp.AuthURL = ""
-		resp.MachineAuthorized = true
-		resp.User = *m.Namespace.toUser()
-
-		respBody, err := encode(resp, &mKey, h.privateKey)
-		if err != nil {
-			log.Error().
-				Str("handler", "Registration").
-				Err(err).
-				Msg("Cannot encode message")
-			c.String(http.StatusInternalServerError, "")
-			return
-		}
-		c.Data(200, "application/json; charset=utf-8", respBody)
-		return
-	}
-
 	log.Debug().
 		Str("handler", "Registration").
 		Str("machine", m.Name).
@@ -255,8 +230,8 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) {
 			strings.TrimSuffix(h.cfg.ServerURL, "/"), mKey.HexString())
 	}
 
-	m.Expiry = &req.Expiry                         // save the requested expiry time for retrieval later
+	m.RequestedExpiry = &req.Expiry                // save the requested expiry time for retrieval later in the authentication flow
-	m.NodeKey = wgkey.Key(req.NodeKey).HexString() // save the new nodekey
+	m.NodeKey = wgkey.Key(req.NodeKey).HexString() // save the NodeKey
 	h.db.Save(&m)
 
 	respBody, err := encode(resp, &mKey, h.privateKey)
@@ -436,6 +411,8 @@ func (h *Headscale) handleAuthKey(c *gin.Context, db *gorm.DB, idKey wgkey.Key,
 	m.RegisterMethod = "authKey"
 	db.Save(&m)
 
+	h.updateMachineExpiry(&m) // TODO: do we want to do different expiry times for AuthKeys?
+
 	resp.MachineAuthorized = true
 	resp.User = *pak.Namespace.toUser()
 	respBody, err := encode(resp, &idKey, h.privateKey)

app.go

@@ -59,8 +59,8 @@ type Config struct {
 	OIDCClientID     string
 	OIDCClientSecret string
 
-	MaxMachineExpiry     time.Duration
+	MaxMachineRegistrationDuration     time.Duration
-	DefaultMachineExpiry time.Duration
+	DefaultMachineRegistrationDuration time.Duration
 }
 
 // Headscale represents the base app of the service

cli.go

@@ -23,6 +23,8 @@ func (h *Headscale) RegisterMachine(key string, namespace string) (*Machine, err
 		return nil, errors.New("Machine not found")
 	}
 
+	h.updateMachineExpiry(&m) // update the machine's expiry before bailing if its already registered
+
 	if m.isAlreadyRegistered() {
 		return nil, errors.New("Machine already registered")
 	}
@@ -36,5 +38,6 @@ func (h *Headscale) RegisterMachine(key string, namespace string) (*Machine, err
 	m.Registered = true
 	m.RegisterMethod = "cli"
 	h.db.Save(&m)
+
 	return &m, nil
 }

cmd/headscale/cli/utils.go

@@ -144,14 +144,16 @@ func getHeadscaleApp() (*headscale.Headscale, error) {
 		return nil, err
 	}
 
-	maxMachineExpiry, _ := time.ParseDuration("8h")
+	// maxMachineRegistrationDuration is the maximum time a client can request for a client registration
-	if viper.GetDuration("max_machine_expiry") >= time.Second {
+	maxMachineRegistrationDuration, _ := time.ParseDuration("10h")
-		maxMachineExpiry = viper.GetDuration("max_machine_expiry")
+	if viper.GetDuration("max_machine_registration_duration") >= time.Second {
+		maxMachineRegistrationDuration = viper.GetDuration("max_machine_registration_duration")
 	}
 
-	defaultMachineExpiry, _ := time.ParseDuration("8h")
+	// defaultMachineRegistrationDuration is the default time assigned to a client registration if one is not specified by the client
-	if viper.GetDuration("default_machine_expiry") >= time.Second {
+	defaultMachineRegistrationDuration, _ := time.ParseDuration("8h")
-		defaultMachineExpiry = viper.GetDuration("default_machine_expiry")
+	if viper.GetDuration("default_machine_registration_duration") >= time.Second {
+		defaultMachineRegistrationDuration = viper.GetDuration("default_machine_registration_duration")
 	}
 
 	cfg := headscale.Config{
@@ -188,8 +190,8 @@ func getHeadscaleApp() (*headscale.Headscale, error) {
 		OIDCClientID:     viper.GetString("oidc_client_id"),
 		OIDCClientSecret: viper.GetString("oidc_client_secret"),
 
-		MaxMachineExpiry:     maxMachineExpiry,
+		MaxMachineRegistrationDuration:     maxMachineRegistrationDuration,     // the maximum duration a client may request for expiry time
-		DefaultMachineExpiry: defaultMachineExpiry,
+		DefaultMachineRegistrationDuration: defaultMachineRegistrationDuration, // if a client does not request a specific expiry time, use this duration
 	}
 
 	h, err := headscale.NewHeadscale(cfg)

go.mod

@@ -11,7 +11,7 @@ require (
 	github.com/docker/cli v20.10.8+incompatible // indirect
 	github.com/docker/docker v20.10.8+incompatible // indirect
 	github.com/efekarakus/termcolor v1.0.1
-	github.com/fatih/set v0.2.1 // indirect
+	github.com/fatih/set v0.2.1
 	github.com/gin-gonic/gin v1.7.4
 	github.com/gofrs/uuid v4.0.0+incompatible
 	github.com/google/go-github v17.0.0+incompatible // indirect
@@ -23,7 +23,7 @@ require (
 	github.com/opencontainers/runc v1.0.2 // indirect
 	github.com/ory/dockertest/v3 v3.7.0
 	github.com/patrickmn/go-cache v2.1.0+incompatible
-	github.com/prometheus/client_golang v1.11.0 // indirect
+	github.com/prometheus/client_golang v1.11.0
 	github.com/pterm/pterm v0.12.30
 	github.com/rs/zerolog v1.25.0
 	github.com/s12v/go-jwks v0.2.1
@@ -33,7 +33,7 @@ require (
 	github.com/tailscale/hujson v0.0.0-20210818175511-7360507a6e88
 	github.com/tcnksm/go-latest v0.0.0-20170313132115-e3007ae9052e
 	github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb // indirect
-	github.com/zsais/go-gin-prometheus v0.1.0 // indirect
+	github.com/zsais/go-gin-prometheus v0.1.0
 	golang.org/x/crypto v0.0.0-20210817164053-32db794688a5
 	golang.org/x/net v0.0.0-20210913180222-943fd674d43e // indirect
 	golang.org/x/oauth2 v0.0.0-20210402161424-2e8d93401602

machine.go

@@ -36,6 +36,7 @@ type Machine struct {
 	LastSeen             *time.Time
 	LastSuccessfulUpdate *time.Time
 	Expiry               *time.Time
+	RequestedExpiry      *time.Time // when a client connects, it may request a specific expiry time, use this field to store it
 
 	HostInfo      datatypes.JSON
 	Endpoints     datatypes.JSON
@@ -61,6 +62,31 @@ func (m Machine) isExpired() bool {
 	return time.Now().UTC().After(*m.Expiry)
 }
 
+// If the Machine is expired, updateMachineExpiry updates the Machine Expiry time to the maximum allowed duration,
+// or the default duration if no Expiry time was requested by the client
+func (h *Headscale) updateMachineExpiry(m *Machine) {
+
+	if m.isExpired() {
+		now := time.Now().UTC()
+		maxExpiry := now.Add(h.cfg.MaxMachineRegistrationDuration)         // calculate the maximum expiry
+		defaultExpiry := now.Add(h.cfg.DefaultMachineRegistrationDuration) // calculate the default expiry
+
+		// clamp the expiry time of the machine registration to the maximum allowed, or use the default if none supplied
+		if maxExpiry.Before(*m.RequestedExpiry) {
+			log.Debug().Msgf("Clamping registration expiry time to maximum: %v (%v)", maxExpiry, h.cfg.MaxMachineRegistrationDuration)
+			m.Expiry = &maxExpiry
+		} else if m.RequestedExpiry.IsZero() {
+			log.Debug().Msgf("Using default machine registration expiry time: %v (%v)", defaultExpiry, h.cfg.DefaultMachineRegistrationDuration)
+			m.Expiry = &defaultExpiry
+		} else {
+			log.Debug().Msgf("Using requested machine registration expiry time: %v", m.RequestedExpiry)
+			m.Expiry = m.RequestedExpiry
+		}
+
+		h.db.Save(&m)
+	}
+}
+
 func (h *Headscale) getDirectPeers(m *Machine) (Machines, error) {
 	log.Trace().
 		Str("func", "getDirectPeers").

oidc.go

@@ -4,14 +4,12 @@ import (
 	"context"
 	"crypto/rand"
 	"encoding/hex"
-	"errors"
 	"fmt"
 	"github.com/coreos/go-oidc/v3/oidc"
 	"github.com/gin-gonic/gin"
 	"github.com/patrickmn/go-cache"
 	"github.com/rs/zerolog/log"
 	"golang.org/x/oauth2"
-	"gorm.io/gorm"
 	"net/http"
 	"strings"
 	"time"
@@ -103,6 +101,8 @@ func (h *Headscale) OIDCCallback(c *gin.Context) {
 		return
 	}
 
+	log.Debug().Msgf("AccessToken: %v", oauth2Token.AccessToken)
+
 	rawIDToken, rawIDTokenOK := oauth2Token.Extra("id_token").(string)
 	if !rawIDTokenOK {
 		c.String(http.StatusBadRequest, "Could not extract ID Token")
@@ -117,16 +117,17 @@ func (h *Headscale) OIDCCallback(c *gin.Context) {
 		return
 	}
 
+	// TODO: we can use userinfo at some point to grab additional information about the user (groups membership, etc)
 	//userInfo, err := oidcProvider.UserInfo(context.Background(), oauth2.StaticTokenSource(oauth2Token))
 	//if err != nil {
-	//	c.String(http.StatusBadRequest, "Failed to retrieve userinfo: "+err.Error())
+	//	c.String(http.StatusBadRequest, fmt.Sprintf("Failed to retrieve userinfo: %s", err))
 	//	return
 	//}
 
 	// Extract custom claims
 	var claims IDTokenClaims
 	if err = idToken.Claims(&claims); err != nil {
-		c.String(http.StatusBadRequest, "Failed to decode id token claims: "+err.Error())
+		c.String(http.StatusBadRequest, fmt.Sprintf("Failed to decode id token claims: %s", err))
 		return
 	}
 
@@ -134,39 +135,44 @@ func (h *Headscale) OIDCCallback(c *gin.Context) {
 	mKeyIf, mKeyFound := h.oidcStateCache.Get(state)
 
 	if !mKeyFound {
+		log.Error().Msg("requested machine state key expired before authorisation completed")
 		c.String(http.StatusBadRequest, "state has expired")
 		return
 	}
 	mKeyStr, mKeyOK := mKeyIf.(string)
 
 	if !mKeyOK {
+		log.Error().Msg("could not get machine key from cache")
 		c.String(http.StatusInternalServerError, "could not get machine key from cache")
 		return
 	}
 
 	// retrieve machine information
-	var m Machine
+	m, err := h.GetMachineByMachineKey(mKeyStr)
-	if result := h.db.Preload("Namespace").First(&m, "machine_key = ?", mKeyStr); errors.Is(result.Error, gorm.ErrRecordNotFound) {
+
+	if err != nil {
 		log.Error().Msg("machine key not found in database")
 		c.String(http.StatusInternalServerError, "could not get machine info from database")
 		return
 	}
 
-	//look for a namespace of the users email for now
+	now := time.Now().UTC()
+
+	// register the machine if it's new
 	if !m.Registered {
+		nsName := strings.ReplaceAll(claims.Email, "@", "-") // TODO: Implement a better email sanitisation
 
 		log.Debug().Msg("Registering new machine after successful callback")
 
-		ns, err := h.GetNamespace(claims.Email)
+		ns, err := h.GetNamespace(nsName)
 		if err != nil {
-			ns, err = h.CreateNamespace(claims.Email)
+			ns, err = h.CreateNamespace(nsName)
 
 			if err != nil {
 				log.Error().Msgf("could not create new namespace '%s'", claims.Email)
 				c.String(http.StatusInternalServerError, "could not create new namespace")
 				return
 			}
-
 		}
 
 		ip, err := h.getAvailableIP()
@@ -179,24 +185,11 @@ func (h *Headscale) OIDCCallback(c *gin.Context) {
 		m.NamespaceID = ns.ID
 		m.Registered = true
 		m.RegisterMethod = "oidc"
+		m.LastSuccessfulUpdate = &now
 		h.db.Save(&m)
 	}
 
-	if m.isExpired() {
+	h.updateMachineExpiry(m)
-		maxExpiry := time.Now().UTC().Add(h.cfg.MaxMachineExpiry)
-
-		// use the maximum expiry if it's sooner than the requested expiry
-		if maxExpiry.Before(*m.Expiry) {
-			log.Debug().Msgf("Clamping expiry time to maximum: %v (%v)", maxExpiry, h.cfg.MaxMachineExpiry)
-			m.Expiry = &maxExpiry
-			h.db.Save(&m)
-		} else if m.Expiry.IsZero() {
-			log.Debug().Msgf("Using default machine expiry time: %v (%v)", maxExpiry, h.cfg.MaxMachineExpiry)
-			defaultExpiry := time.Now().UTC().Add(h.cfg.DefaultMachineExpiry)
-			m.Expiry = &defaultExpiry
-			h.db.Save(&m)
-		}
-	}
 
 	c.Data(http.StatusOK, "text/html; charset=utf-8", []byte(fmt.Sprintf(`
 <html>