From 272eee79fd56b45722e7cb9eaa2124c02d6f9a96 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Sun, 21 Feb 2021 21:34:28 +0100
Subject: [PATCH] Added support for DiscoKey

(from https://github.com/awsong/headscale)
---
 handlers.go | 4 ++--
 machine.go  | 6 ++++++
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/handlers.go b/handlers.go
index d805ec2d..80217b55 100644
--- a/handlers.go
+++ b/handlers.go
@@ -137,6 +137,7 @@ func (h *Headscale) PollNetMapHandler(c *gin.Context) {
 	hostinfo, _ := json.Marshal(req.Hostinfo)
 	m.Endpoints = postgres.Jsonb{RawMessage: json.RawMessage(endpoints)}
 	m.HostInfo = postgres.Jsonb{RawMessage: json.RawMessage(hostinfo)}
+	m.DiscoKey = wgcfg.Key(req.DiscoKey).HexString()
 	now := time.Now().UTC()
 	m.LastSeen = &now
 	db.Save(&m)
@@ -254,8 +255,7 @@ func (h *Headscale) getMapKeepAliveResponse(mKey wgcfg.Key, req tailcfg.MapReque
 	return &data, nil
 }
 
-// RegisterWebAPI attaches the machine to a user or team.
-// Currently this is a rather temp implementation, as it just registers it.
+// RegisterWebAPI shows a simple message in the browser to point to the CLI
 func (h *Headscale) RegisterWebAPI(c *gin.Context) {
 	mKeyStr := c.Query("key")
 	if mKeyStr == "" {
diff --git a/machine.go b/machine.go
index cee6289e..ecb8e899 100644
--- a/machine.go
+++ b/machine.go
@@ -16,6 +16,7 @@ type Machine struct {
 	ID         uint64 `gorm:"primary_key"`
 	MachineKey string `gorm:"type:varchar(64);unique_index"`
 	NodeKey    string
+	DiscoKey   string
 	IPAddress  string
 
 	Registered bool // temp
@@ -44,6 +45,10 @@ func (m Machine) toNode() (*tailcfg.Node, error) {
 	if err != nil {
 		return nil, err
 	}
+	dKey, err := wgcfg.ParseHexKey(m.DiscoKey)
+	if err != nil {
+		return nil, err
+	}
 	addrs := []netaddr.IPPrefix{}
 	allowedIPs := []netaddr.IPPrefix{}
 
@@ -81,6 +86,7 @@ func (m Machine) toNode() (*tailcfg.Node, error) {
 		Key:        tailcfg.NodeKey(nKey),
 		KeyExpiry:  *m.Expiry,
 		Machine:    tailcfg.MachineKey(mKey),
+		DiscoKey:   tailcfg.DiscoKey(dKey),
 		Addresses:  addrs,
 		AllowedIPs: allowedIPs,
 		Endpoints:  endpoints,