Added Noise upgrade handler and Noise mux

app.go

@@ -79,7 +79,7 @@ type Headscale struct {
 	privateKey      *key.MachinePrivate
 	noisePrivateKey *key.MachinePrivate
 
-	noiseRouter *gin.Engine
+	noiseMux *http.ServeMux
 
 	DERPMap    *tailcfg.DERPMap
 	DERPServer *DERPServer
@@ -406,6 +406,7 @@ func (h *Headscale) createPrometheusRouter() *gin.Engine {
 func (h *Headscale) createRouter(grpcMux *runtime.ServeMux) *gin.Engine {
 	router := gin.Default()
 
+	router.POST(ts2021UpgradePath, h.NoiseUpgradeHandler)
 	router.GET(
 		"/health",
 		func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"healthy": "ok"}) },
@@ -440,6 +441,15 @@ func (h *Headscale) createRouter(grpcMux *runtime.ServeMux) *gin.Engine {
 	return router
 }
 
+func (h *Headscale) createNoiseMux() *http.ServeMux {
+	mux := http.NewServeMux()
+
+	// mux.HandleFunc("/machine/register", h.NoiseRegistrationHandler)
+	// mux.HandleFunc("/machine/map", h.NoisePollNetMapHandler)
+
+	return mux
+}
+
 // Serve launches a GIN server with the Headscale API.
 func (h *Headscale) Serve() error {
 	var err error
@@ -592,8 +602,14 @@ func (h *Headscale) Serve() error {
 	// HTTP setup
 	//
 
+	// This is the regular router that we expose
+	// over our main Addr. It also serves the legacy Tailcale API
 	router := h.createRouter(grpcGatewayMux)
 
+	// This router is served only over the Noise connection,
+	// and exposes only the new API
+	h.noiseMux = h.createNoiseMux()
+
 	httpServer := &http.Server{
 		Addr:        h.cfg.Addr,
 		Handler:     router,

noise.go

@@ -0,0 +1,125 @@
+package headscale
+
+import (
+	"encoding/base64"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+	"github.com/rs/zerolog/log"
+	"golang.org/x/net/http2"
+	"golang.org/x/net/http2/h2c"
+	"tailscale.com/control/controlbase"
+	"tailscale.com/net/netutil"
+)
+
+const (
+	errWrongConnectionUpgrade = Error("wrong connection upgrade")
+	errCannotHijack           = Error("cannot hijack connection")
+	errNoiseHandshakeFailed   = Error("noise handshake failed")
+)
+
+const (
+	// ts2021UpgradePath is the path that the server listens on for the WebSockets upgrade.
+	ts2021UpgradePath = "/ts2021"
+
+	// upgradeHeader is the value of the Upgrade HTTP header used to
+	// indicate the Tailscale control protocol.
+	upgradeHeaderValue = "tailscale-control-protocol"
+
+	// handshakeHeaderName is the HTTP request header that can
+	// optionally contain base64-encoded initial handshake
+	// payload, to save an RTT.
+	handshakeHeaderName = "X-Tailscale-Handshake"
+)
+
+// NoiseUpgradeHandler is to upgrade the connection and hijack the net.Conn
+// in order to use the Noise-based TS2021 protocol. Listens in /ts2021.
+func (h *Headscale) NoiseUpgradeHandler(ctx *gin.Context) {
+	log.Trace().Caller().Msgf("Noise upgrade handler for client %s", ctx.ClientIP())
+
+	// Under normal circumpstances, we should be able to use the controlhttp.AcceptHTTP()
+	// function to do this - kindly left there by the Tailscale authors for us to use.
+	// (https://github.com/tailscale/tailscale/blob/main/control/controlhttp/server.go)
+	//
+	// However, Gin seems to be doing something funny/different with its writer (see AcceptHTTP code).
+	// This causes problems when the upgrade headers are sent in AcceptHTTP.
+	// So have getNoiseConnection() that is essentially an AcceptHTTP but using the native Gin methods.
+	noiseConn, err := h.getNoiseConnection(ctx)
+	if err != nil {
+		log.Error().Err(err).Msg("noise upgrade failed")
+		ctx.AbortWithError(http.StatusInternalServerError, err)
+
+		return
+	}
+
+	server := http.Server{}
+	server.Handler = h2c.NewHandler(h.noiseMux, &http2.Server{})
+	server.Serve(netutil.NewOneConnListener(noiseConn, nil))
+}
+
+// getNoiseConnection is basically AcceptHTTP from tailscale, but more _alla_ Gin
+// TODO(juan): Figure out why we need to do this at all.
+func (h *Headscale) getNoiseConnection(ctx *gin.Context) (*controlbase.Conn, error) {
+	next := ctx.GetHeader("Upgrade")
+	if next == "" {
+		ctx.String(http.StatusBadRequest, "missing next protocol")
+
+		return nil, errWrongConnectionUpgrade
+	}
+	if next != upgradeHeaderValue {
+		ctx.String(http.StatusBadRequest, "unknown next protocol")
+
+		return nil, errWrongConnectionUpgrade
+	}
+
+	initB64 := ctx.GetHeader(handshakeHeaderName)
+	if initB64 == "" {
+		ctx.String(http.StatusBadRequest, "missing Tailscale handshake header")
+
+		return nil, errWrongConnectionUpgrade
+	}
+	init, err := base64.StdEncoding.DecodeString(initB64)
+	if err != nil {
+		ctx.String(http.StatusBadRequest, "invalid tailscale handshake header")
+
+		return nil, errWrongConnectionUpgrade
+	}
+
+	hijacker, ok := ctx.Writer.(http.Hijacker)
+	if !ok {
+		log.Error().Caller().Err(err).Msgf("Hijack failed")
+		ctx.String(http.StatusInternalServerError, "HTTP does not support general TCP support")
+
+		return nil, errCannotHijack
+	}
+
+	// This is what changes from the original AcceptHTTP() function.
+	ctx.Header("Upgrade", upgradeHeaderValue)
+	ctx.Header("Connection", "upgrade")
+	ctx.Status(http.StatusSwitchingProtocols)
+	ctx.Writer.WriteHeaderNow()
+	// end
+
+	netConn, conn, err := hijacker.Hijack()
+	if err != nil {
+		log.Error().Caller().Err(err).Msgf("Hijack failed")
+		ctx.String(http.StatusInternalServerError, "HTTP does not support general TCP support")
+
+		return nil, errCannotHijack
+	}
+	if err := conn.Flush(); err != nil {
+		netConn.Close()
+
+		return nil, errCannotHijack
+	}
+	netConn = netutil.NewDrainBufConn(netConn, conn.Reader)
+
+	nc, err := controlbase.Server(ctx.Request.Context(), netConn, *h.noisePrivateKey, init)
+	if err != nil {
+		netConn.Close()
+
+		return nil, errNoiseHandshakeFailed
+	}
+
+	return nc, nil
+}