From 34f4109fbd6b5c5da322eb31ac1dbe1b5d9d74e7 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Sun, 28 Nov 2021 09:17:18 +0000
Subject: [PATCH] Add back privatekey, but automatically generate it if it does
 not exist

---
 api.go                     |  2 +-
 app.go                     | 67 +++++++++++++++++++++++++++++++-------
 cmd/headscale/cli/utils.go |  9 ++---
 utils.go                   | 11 +++++++
 4 files changed, 73 insertions(+), 16 deletions(-)

diff --git a/api.go b/api.go
index 076594dc..96e22bd5 100644
--- a/api.go
+++ b/api.go
@@ -34,7 +34,7 @@ func (h *Headscale) KeyHandler(ctx *gin.Context) {
 	ctx.Data(
 		http.StatusOK,
 		"text/plain; charset=utf-8",
-		[]byte(MachinePublicKeyStripPrefix(*h.publicKey)),
+		[]byte(MachinePublicKeyStripPrefix(h.privateKey.Public())),
 	)
 }
 
diff --git a/app.go b/app.go
index 640803b6..db788908 100644
--- a/app.go
+++ b/app.go
@@ -47,11 +47,12 @@ import (
 )
 
 const (
-	AuthPrefix      = "Bearer "
-	Postgres        = "postgres"
-	Sqlite          = "sqlite3"
-	updateInterval  = 5000
-	HTTPReadTimeout = 30 * time.Second
+	AuthPrefix         = "Bearer "
+	Postgres           = "postgres"
+	Sqlite             = "sqlite3"
+	updateInterval     = 5000
+	HTTPReadTimeout    = 30 * time.Second
+	privateKeyFileMode = 0o600
 
 	requestedExpiryCacheExpiration      = time.Minute * 5
 	requestedExpiryCacheCleanupInterval = time.Minute * 10
@@ -68,6 +69,7 @@ type Config struct {
 	Addr                           string
 	EphemeralNodeInactivityTimeout time.Duration
 	IPPrefix                       netaddr.IPPrefix
+	PrivateKeyPath                 string
 	BaseDomain                     string
 
 	DERP DERPConfig
@@ -128,7 +130,6 @@ type Headscale struct {
 	dbString   string
 	dbType     string
 	dbDebug    bool
-	publicKey  *key.MachinePublic
 	privateKey *key.MachinePrivate
 
 	DERPMap *tailcfg.DERPMap
@@ -147,8 +148,10 @@ type Headscale struct {
 
 // NewHeadscale returns the Headscale app.
 func NewHeadscale(cfg Config) (*Headscale, error) {
-	privKey := key.NewMachine()
-	pubKey := privKey.Public()
+	privKey, err := readOrCreatePrivateKey(cfg.PrivateKeyPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read or create private key: %w", err)
+	}
 
 	var dbString string
 	switch cfg.DBtype {
@@ -176,13 +179,12 @@ func NewHeadscale(cfg Config) (*Headscale, error) {
 		cfg:                  cfg,
 		dbType:               cfg.DBtype,
 		dbString:             dbString,
-		privateKey:           &privKey,
-		publicKey:            &pubKey,
+		privateKey:           privKey,
 		aclRules:             tailcfg.FilterAllowAll, // default allowall
 		requestedExpiryCache: requestedExpiryCache,
 	}
 
-	err := app.initDB()
+	err = app.initDB()
 	if err != nil {
 		return nil, err
 	}
@@ -694,3 +696,46 @@ func stdoutHandler(ctx *gin.Context) {
 		Bytes("body", body).
 		Msg("Request did not match")
 }
+
+func readOrCreatePrivateKey(path string) (*key.MachinePrivate, error) {
+	privateKey, err := os.ReadFile(path)
+	if errors.Is(err, os.ErrNotExist) {
+		log.Info().Str("path", path).Msg("No private key file at path, creating...")
+
+		machineKey := key.NewMachine()
+
+		machineKeyStr, err := machineKey.MarshalText()
+		if err != nil {
+			return nil, fmt.Errorf(
+				"failed to convert private key to string for saving: %w",
+				err,
+			)
+		}
+		err = os.WriteFile(path, machineKeyStr, privateKeyFileMode)
+		if err != nil {
+			return nil, fmt.Errorf(
+				"failed to save private key to disk: %w",
+				err,
+			)
+		}
+
+		return &machineKey, nil
+	} else if err != nil {
+		return nil, fmt.Errorf("failed to read private key file: %w", err)
+	}
+
+	privateKeyEnsurePrefix := PrivateKeyEnsurePrefix(string(privateKey))
+
+	var machineKey key.MachinePrivate
+	if err = machineKey.UnmarshalText([]byte(privateKeyEnsurePrefix)); err != nil {
+		log.Info().
+			Str("path", path).
+			Msg("This might be due to a legacy (headscale pre-0.12) private key. " +
+				"If the key is in WireGuard format, delete the key and restart headscale. " +
+				"A new key will automatically be generated. All Tailscale clients will have to be restarted")
+
+		return nil, fmt.Errorf("failed to parse private key: %w", err)
+	}
+
+	return &machineKey, nil
+}
diff --git a/cmd/headscale/cli/utils.go b/cmd/headscale/cli/utils.go
index d47a5ca9..a4856642 100644
--- a/cmd/headscale/cli/utils.go
+++ b/cmd/headscale/cli/utils.go
@@ -222,10 +222,11 @@ func getHeadscaleConfig() headscale.Config {
 	derpConfig := GetDERPConfig()
 
 	return headscale.Config{
-		ServerURL:  viper.GetString("server_url"),
-		Addr:       viper.GetString("listen_addr"),
-		IPPrefix:   netaddr.MustParseIPPrefix(viper.GetString("ip_prefix")),
-		BaseDomain: baseDomain,
+		ServerURL:      viper.GetString("server_url"),
+		Addr:           viper.GetString("listen_addr"),
+		IPPrefix:       netaddr.MustParseIPPrefix(viper.GetString("ip_prefix")),
+		PrivateKeyPath: absPath(viper.GetString("private_key_path")),
+		BaseDomain:     baseDomain,
 
 		DERP: derpConfig,
 
diff --git a/utils.go b/utils.go
index c9971390..7011e63d 100644
--- a/utils.go
+++ b/utils.go
@@ -46,6 +46,9 @@ const (
 	// This prefix is used in the control protocol, so cannot be
 	// changed.
 	discoPublicHexPrefix = "discokey:"
+
+	// privateKey prefix.
+	privateHexPrefix = "privkey:"
 )
 
 func MachinePublicKeyStripPrefix(machineKey key.MachinePublic) string {
@@ -84,6 +87,14 @@ func DiscoPublicKeyEnsurePrefix(discoKey string) string {
 	return discoKey
 }
 
+func PrivateKeyEnsurePrefix(privateKey string) string {
+	if !strings.HasPrefix(privateKey, privateHexPrefix) {
+		return privateHexPrefix + privateKey
+	}
+
+	return privateKey
+}
+
 // Error is used to compare errors as per https://dave.cheney.net/2016/04/07/constant-errors
 type Error string