TS2021: Expose Register handler in Noise

api.go

@@ -563,8 +563,13 @@ func (h *Headscale) handleAuthKey(
 	machineKey key.MachinePublic,
 	registerRequest tailcfg.RegisterRequest,
 ) {
-	machineKeyStr := MachinePublicKeyStripPrefix(machineKey)
-
+	var machineKeyStr string
+	if machineKey.IsZero() {
+		// We are handling here a Noise auth key
+		machineKeyStr = ""
+	} else {
+		machineKeyStr = MachinePublicKeyStripPrefix(machineKey)
+	}
 	log.Debug().
 		Str("func", "handleAuthKey").
 		Str("machine", registerRequest.Hostinfo.Hostname).

app.go

@@ -513,6 +513,8 @@ func (h *Headscale) createRouter(grpcMux *runtime.ServeMux) *gin.Engine {
 func (h *Headscale) createNoiseRouter() *gin.Engine {
 	router := gin.Default()
 
+	router.POST("/machine/register", h.NoiseRegistrationHandler)
+
 	return router
 }
 

machine.go

@@ -335,7 +335,7 @@ func (h *Headscale) GetMachineByID(id uint64) (*Machine, error) {
 	return &m, nil
 }
 
-// GetMachineByMachineKey finds a Machine by ID and returns the Machine struct.
+// GetMachineByMachineKey finds a Machine by its MachineKey and returns the Machine struct.
 func (h *Headscale) GetMachineByMachineKey(
 	machineKey key.MachinePublic,
 ) (*Machine, error) {
@@ -347,6 +347,19 @@ func (h *Headscale) GetMachineByMachineKey(
 	return &m, nil
 }
 
+// GetMachineByNodeKeys finds a Machine by its current NodeKey or the old one, and returns the Machine struct.
+func (h *Headscale) GetMachineByNodeKeys(
+	nodeKey key.NodePublic, oldNodeKey key.NodePublic,
+) (*Machine, error) {
+	m := Machine{}
+	if result := h.db.Preload("Namespace").First(&m, "node_key = ? OR node_key = ?",
+		NodePublicKeyStripPrefix(nodeKey), NodePublicKeyStripPrefix(oldNodeKey)); result.Error != nil {
+		return nil, result.Error
+	}
+
+	return &m, nil
+}
+
 // UpdateMachine takes a Machine struct pointer (typically already loaded from database
 // and updates it with the latest data from the database.
 func (h *Headscale) UpdateMachine(machine *Machine) error {
@@ -677,7 +690,7 @@ func (h *Headscale) RegisterMachine(machine Machine,
 ) (*Machine, error) {
 	log.Trace().
 		Caller().
-		Str("machine_key", machine.MachineKey).
+		Str("node_key", machine.NodeKey).
 		Msg("Registering machine")
 
 	log.Trace().

machine_test.go

@@ -10,6 +10,7 @@ import (
 	"gopkg.in/check.v1"
 	"inet.af/netaddr"
 	"tailscale.com/tailcfg"
+	"tailscale.com/types/key"
 )
 
 func (s *Suite) TestGetMachine(c *check.C) {
@@ -64,6 +65,35 @@ func (s *Suite) TestGetMachineByID(c *check.C) {
 	c.Assert(err, check.IsNil)
 }
 
+func (s *Suite) TestGetMachineByNodeKeys(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.GetMachineByID(0)
+	c.Assert(err, check.NotNil)
+
+	nodeKey := key.NewNode()
+	oldNodeKey := key.NewNode()
+
+	machine := Machine{
+		ID:             0,
+		MachineKey:     "foo",
+		NodeKey:        NodePublicKeyStripPrefix(nodeKey.Public()),
+		DiscoKey:       "faa",
+		Name:           "testmachine",
+		NamespaceID:    namespace.ID,
+		RegisterMethod: RegisterMethodAuthKey,
+		AuthKeyID:      uint(pak.ID),
+	}
+	app.db.Save(&machine)
+
+	_, err = app.GetMachineByNodeKeys(nodeKey.Public(), oldNodeKey.Public())
+	c.Assert(err, check.IsNil)
+}
+
 func (s *Suite) TestDeleteMachine(c *check.C) {
 	namespace, err := app.CreateNamespace("test")
 	c.Assert(err, check.IsNil)

noise_api.go

@@ -1 +1,217 @@
 package headscale
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/rs/zerolog/log"
+	"gorm.io/gorm"
+	"tailscale.com/tailcfg"
+	"tailscale.com/types/key"
+)
+
+func (h *Headscale) NoiseRegistrationHandler(ctx *gin.Context) {
+	log.Trace().Caller().Msgf("Noise registration handler for client %s", ctx.ClientIP())
+	body, _ := io.ReadAll(ctx.Request.Body)
+	req := tailcfg.RegisterRequest{}
+	if err := json.Unmarshal(body, &req); err != nil {
+		log.Error().
+			Caller().
+			Err(err).
+			Msg("Cannot parse RegisterRequest")
+		machineRegistrations.WithLabelValues("unknown", "web", "error", "unknown").Inc()
+		ctx.String(http.StatusInternalServerError, "Eek!")
+
+		return
+	}
+
+	now := time.Now().UTC()
+	machine, err := h.GetMachineByNodeKeys(req.NodeKey, req.OldNodeKey)
+	if errors.Is(err, gorm.ErrRecordNotFound) {
+		log.Info().Str("machine", req.Hostinfo.Hostname).Msg("New machine via Noise")
+
+		// If the machine has AuthKey set, handle registration via PreAuthKeys
+		if req.Auth.AuthKey != "" {
+			h.handleAuthKey(ctx, key.MachinePublic{}, req)
+
+			return
+		}
+		hname, err := NormalizeToFQDNRules(
+			req.Hostinfo.Hostname,
+			h.cfg.OIDC.StripEmaildomain,
+		)
+		if err != nil {
+			log.Error().
+				Caller().
+				Str("func", "RegistrationHandler").
+				Str("hostinfo.name", req.Hostinfo.Hostname).
+				Err(err)
+
+			return
+		}
+
+		// The machine did not have a key to authenticate, which means
+		// that we rely on a method that calls back some how (OpenID or CLI)
+		// We create the machine and then keep it around until a callback
+		// happens
+		newMachine := Machine{
+			MachineKey: "",
+			Name:       hname,
+			NodeKey:    NodePublicKeyStripPrefix(req.NodeKey),
+			LastSeen:   &now,
+			Expiry:     &time.Time{},
+		}
+
+		if !req.Expiry.IsZero() {
+			log.Trace().
+				Caller().
+				Str("machine", req.Hostinfo.Hostname).
+				Time("expiry", req.Expiry).
+				Msg("Non-zero expiry time requested")
+			newMachine.Expiry = &req.Expiry
+		}
+
+		h.registrationCache.Set(
+			NodePublicKeyStripPrefix(req.NodeKey),
+			newMachine,
+			registerCacheExpiration,
+		)
+
+		h.handleMachineRegistrationNew(ctx, key.MachinePublic{}, req)
+
+		return
+	}
+
+	// The machine is already registered, so we need to pass through reauth or key update.
+	if machine != nil {
+		// If the NodeKey stored in headscale is the same as the key presented in a registration
+		// request, then we have a node that is either:
+		// - Trying to log out (sending a expiry in the past)
+		// - A valid, registered machine, looking for the node map
+		// - Expired machine wanting to reauthenticate
+		if machine.NodeKey == NodePublicKeyStripPrefix(req.NodeKey) {
+			// The client sends an Expiry in the past if the client is requesting to expire the key (aka logout)
+			//   https://github.com/tailscale/tailscale/blob/main/tailcfg/tailcfg.go#L648
+			if !req.Expiry.IsZero() && req.Expiry.UTC().Before(now) {
+				h.handleNoiseNodeLogOut(ctx, *machine)
+
+				return
+			}
+
+			// If machine is not expired, and is register, we have a already accepted this machine,
+			// let it proceed with a valid registration
+			if !machine.isExpired() {
+				h.handleNoiseNodeValidRegistration(ctx, *machine)
+
+				return
+			}
+		}
+
+		// The NodeKey we have matches OldNodeKey, which means this is a refresh after a key expiration
+		if machine.NodeKey == NodePublicKeyStripPrefix(req.OldNodeKey) &&
+			!machine.isExpired() {
+			h.handleNoiseNodeRefreshKey(ctx, req, *machine)
+
+			return
+		}
+
+		// The node has expired
+		h.handleNoiseNodeExpired(ctx, req, *machine)
+
+		return
+	}
+}
+
+func (h *Headscale) handleNoiseNodeValidRegistration(
+	ctx *gin.Context,
+	machine Machine,
+) {
+	resp := tailcfg.RegisterResponse{}
+
+	// The machine registration is valid, respond with redirect to /map
+	log.Debug().
+		Str("machine", machine.Name).
+		Msg("Client is registered and we have the current NodeKey. All clear to /map")
+
+	resp.AuthURL = ""
+	resp.MachineAuthorized = true
+	resp.User = *machine.Namespace.toUser()
+	resp.Login = *machine.Namespace.toLogin()
+
+	machineRegistrations.WithLabelValues("update", "web", "success", machine.Namespace.Name).
+		Inc()
+	ctx.JSON(http.StatusOK, resp)
+}
+
+func (h *Headscale) handleNoiseNodeLogOut(
+	ctx *gin.Context,
+	machine Machine,
+) {
+	resp := tailcfg.RegisterResponse{}
+
+	log.Info().
+		Str("machine", machine.Name).
+		Msg("Client requested logout")
+
+	h.ExpireMachine(&machine)
+
+	resp.AuthURL = ""
+	resp.MachineAuthorized = false
+	resp.User = *machine.Namespace.toUser()
+	ctx.JSON(http.StatusOK, resp)
+}
+
+func (h *Headscale) handleNoiseNodeRefreshKey(
+	ctx *gin.Context,
+	registerRequest tailcfg.RegisterRequest,
+	machine Machine,
+) {
+	resp := tailcfg.RegisterResponse{}
+
+	log.Debug().
+		Str("machine", machine.Name).
+		Msg("We have the OldNodeKey in the database. This is a key refresh")
+	machine.NodeKey = NodePublicKeyStripPrefix(registerRequest.NodeKey)
+	h.db.Save(&machine)
+
+	resp.AuthURL = ""
+	resp.User = *machine.Namespace.toUser()
+	ctx.JSON(http.StatusOK, resp)
+}
+
+func (h *Headscale) handleNoiseNodeExpired(
+	ctx *gin.Context,
+	registerRequest tailcfg.RegisterRequest,
+	machine Machine,
+) {
+	resp := tailcfg.RegisterResponse{}
+
+	// The client has registered before, but has expired
+	log.Debug().
+		Str("machine", machine.Name).
+		Msg("Machine registration has expired. Sending a authurl to register")
+
+	if registerRequest.Auth.AuthKey != "" {
+		h.handleAuthKey(ctx, key.MachinePublic{}, registerRequest)
+
+		return
+	}
+
+	if h.cfg.OIDC.Issuer != "" {
+		resp.AuthURL = fmt.Sprintf("%s/oidc/register/%s",
+			strings.TrimSuffix(h.cfg.ServerURL, "/"), machine.NodeKey)
+	} else {
+		resp.AuthURL = fmt.Sprintf("%s/register?key=%s",
+			strings.TrimSuffix(h.cfg.ServerURL, "/"), machine.NodeKey)
+	}
+
+	machineRegistrations.WithLabelValues("reauth", "web", "success", machine.Namespace.Name).
+		Inc()
+	ctx.JSON(http.StatusOK, resp)
+}