Improvements on ts2021 implementation, so the connection is reused.

This commit renames ts2021App to noiseServer to better match Tailscale code (they talk about NoiseClient).

It also adds more information to the struct, now that is better kept from registration -> map.

And it also adds initial support for the EarlyNoise functionality (capVer 49)

noise.go

@@ -1,6 +1,9 @@
 package headscale
 
 import (
+	"encoding/binary"
+	"encoding/json"
+	"io"
 	"net/http"
 
 	"github.com/gorilla/mux"
@@ -9,18 +12,37 @@ import (
 	"golang.org/x/net/http2/h2c"
 	"tailscale.com/control/controlbase"
 	"tailscale.com/control/controlhttp"
-	"tailscale.com/net/netutil"
+	"tailscale.com/tailcfg"
+	"tailscale.com/types/key"
 )
 
 const (
 	// ts2021UpgradePath is the path that the server listens on for the WebSockets upgrade.
 	ts2021UpgradePath = "/ts2021"
+
+	// The first 9 bytes from the server to client over Noise are either an HTTP/2
+	// settings frame (a normal HTTP/2 setup) or, as Tailscale added later, an "early payload"
+	// header that's also 9 bytes long: 5 bytes (earlyPayloadMagic) followed by 4 bytes
+	// of length. Then that many bytes of JSON-encoded tailcfg.EarlyNoise.
+	// The early payload is optional. Some servers may not send it... But we do!
+	earlyPayloadMagic = "\xff\xff\xffTS"
+
+	// EarlyNoise was added in protocol version 49.
+	earlyNoiseCapabilityVersion = 49
 )
 
 type noiseServer struct {
 	headscale *Headscale
 
-	conn *controlbase.Conn
+	httpBaseConfig *http.Server
+	http2Server    *http2.Server
+	conn           *controlbase.Conn
+	machineKey     key.MachinePublic
+	nodeKey        key.NodePublic
+
+	// EarlyNoise-related stuff
+	challenge       key.ChallengePrivate
+	protocolVersion int
 }
 
 // NoiseUpgradeHandler is to upgrade the connection and hijack the net.Conn
@@ -44,18 +66,28 @@ func (h *Headscale) NoiseUpgradeHandler(
 		return
 	}
 
-	noiseConn, err := controlhttp.AcceptHTTP(req.Context(), writer, req, *h.noisePrivateKey, nil)
+	noiseServer := noiseServer{
+		headscale: h,
+		challenge: key.NewChallenge(),
+	}
+
+	noiseConn, err := controlhttp.AcceptHTTP(
+		req.Context(),
+		writer,
+		req,
+		*h.noisePrivateKey,
+		noiseServer.earlyNoise,
+	)
 	if err != nil {
 		log.Error().Err(err).Msg("noise upgrade failed")
 		http.Error(writer, err.Error(), http.StatusInternalServerError)
 
 		return
 	}
 
-	noiseServer := noiseServer{
-		headscale: h,
-		conn:      noiseConn,
-	}
+	noiseServer.conn = noiseConn
+	noiseServer.machineKey = noiseServer.conn.Peer()
+	noiseServer.protocolVersion = noiseServer.conn.ProtocolVersion()
 
 	// This router is served only over the Noise connection, and exposes only the new API.
 	//
@@ -70,9 +102,63 @@ func (h *Headscale) NoiseUpgradeHandler(
 	server := http.Server{
 		ReadTimeout: HTTPReadTimeout,
 	}
-	server.Handler = h2c.NewHandler(router, &http2.Server{})
-	err = server.Serve(netutil.NewOneConnListener(noiseConn, nil))
+
+	noiseServer.httpBaseConfig = &http.Server{
+		Handler:           router,
+		ReadHeaderTimeout: HTTPReadTimeout,
+	}
+	noiseServer.http2Server = &http2.Server{}
+
+	server.Handler = h2c.NewHandler(router, noiseServer.http2Server)
+
+	noiseServer.http2Server.ServeConn(
+		noiseConn,
+		&http2.ServeConnOpts{
+			BaseConfig: noiseServer.httpBaseConfig,
+		},
+	)
+}
+
+func (ns *noiseServer) earlyNoise(protocolVersion int, writer io.Writer) error {
+	log.Trace().
+		Caller().
+		Int("protocol_version", protocolVersion).
+		Str("challenge", ns.challenge.Public().String()).
+		Msg("earlyNoise called")
+
+	if protocolVersion < earlyNoiseCapabilityVersion {
+		log.Trace().
+			Caller().
+			Msgf("protocol version %d does not support early noise", protocolVersion)
+
+		return nil
+	}
+
+	earlyJSON, err := json.Marshal(&tailcfg.EarlyNoise{
+		NodeKeyChallenge: ns.challenge.Public(),
+	})
 	if err != nil {
-		log.Info().Err(err).Msg("The HTTP2 server was closed")
+		return err
+	}
+
+	// 5 bytes that won't be mistaken for an HTTP/2 frame:
+	// https://httpwg.org/specs/rfc7540.html#rfc.section.4.1 (Especially not
+	// an HTTP/2 settings frame, which isn't of type 'T')
+	var notH2Frame [5]byte
+	copy(notH2Frame[:], earlyPayloadMagic)
+	var lenBuf [4]byte
+	binary.BigEndian.PutUint32(lenBuf[:], uint32(len(earlyJSON)))
+	// These writes are all buffered by caller, so fine to do them
+	// separately:
+	if _, err := writer.Write(notH2Frame[:]); err != nil {
+		return err
 	}
+	if _, err := writer.Write(lenBuf[:]); err != nil {
+		return err
+	}
+	if _, err := writer.Write(earlyJSON); err != nil {
+		return err
+	}
+
+	return nil
 }

protocol_noise.go

@@ -38,5 +38,7 @@ func (ns *noiseServer) NoiseRegistrationHandler(
 		return
 	}
 
+	ns.nodeKey = registerRequest.NodeKey
+
 	ns.headscale.handleRegisterCommon(writer, req, registerRequest, ns.conn.Peer(), true)
 }

protocol_noise_poll.go

@@ -46,6 +46,8 @@ func (ns *noiseServer) NoisePollNetMapHandler(
 		return
 	}
 
+	ns.nodeKey = mapRequest.NodeKey
+
 	machine, err := ns.headscale.GetMachineByAnyKey(ns.conn.Peer(), mapRequest.NodeKey, key.NodePublic{})
 	if err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {