From 3e8f0e99849538568f329303f032b3655f4add99 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Sat, 13 Aug 2022 11:24:05 +0200
Subject: [PATCH] Added support for Noise clients in /key handler

---
 api.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/api.go b/api.go
index 561545b8..52149d58 100644
--- a/api.go
+++ b/api.go
@@ -9,6 +9,7 @@ import (
 	"html/template"
 	"io"
 	"net/http"
+	"strconv"
 	"strings"
 	"time"
 
@@ -30,6 +31,12 @@ const (
 	ErrRegisterMethodCLIDoesNotSupportExpire = Error(
 		"machines registered with CLI does not support expire",
 	)
+
+	// The CapabilityVersion is used by Tailscale clients to indicate
+	// their codebase version. Tailscale clients can communicate over TS2021
+	// from CapabilityVersion 28.
+	// See https://github.com/tailscale/tailscale/blob/main/tailcfg/tailcfg.go
+	NoiseCapabilityVersion = 28
 )
 
 func (h *Headscale) HealthHandler(
@@ -76,6 +83,45 @@ func (h *Headscale) KeyHandler(
 	writer http.ResponseWriter,
 	req *http.Request,
 ) {
+	// New Tailscale clients send a 'v' parameter to indicate the CurrentCapabilityVersion
+	clientCapabilityStr := req.URL.Query().Get("v")
+	if clientCapabilityStr != "" {
+		clientCapabilityVersion, err := strconv.Atoi(clientCapabilityStr)
+		if err != nil {
+			writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
+			writer.WriteHeader(http.StatusBadRequest)
+			_, err := writer.Write([]byte("Wrong params"))
+			if err != nil {
+				log.Error().
+					Caller().
+					Err(err).
+					Msg("Failed to write response")
+			}
+
+			return
+		}
+
+		if clientCapabilityVersion >= NoiseCapabilityVersion {
+			// Tailscale has a different key for the TS2021 protocol
+			resp := tailcfg.OverTLSPublicKeyResponse{
+				LegacyPublicKey: h.privateKey.Public(),
+				PublicKey:       h.noisePrivateKey.Public(),
+			}
+			writer.Header().Set("Content-Type", "application/json")
+			writer.WriteHeader(http.StatusOK)
+			err = json.NewEncoder(writer).Encode(resp)
+			if err != nil {
+				log.Error().
+					Caller().
+					Err(err).
+					Msg("Failed to write response")
+			}
+
+			return
+		}
+	}
+
+	// Old clients don't send a 'v' parameter, so we send the legacy public key
 	writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 	writer.WriteHeader(http.StatusOK)
 	_, err := writer.Write([]byte(MachinePublicKeyStripPrefix(h.privateKey.Public())))