control/controlclient: send load balancing hint HTTP request header

Updates tailscale/corp#1297

Change-Id: I0b102081e81dfc1261f4b05521ab248a2e4a1298
Signed-off-by: Brad Fitzpatrick <[email protected]>

control/controlclient/direct.go

@@ -641,6 +641,9 @@ func (c *Direct) doLogin(ctx context.Context, opt loginOpt) (mustRegen bool, new
 	if err != nil {
 		return regen, opt.URL, nil, err
 	}
+	addLBHeader(req, request.OldNodeKey)
+	addLBHeader(req, request.NodeKey)
+
 	res, err := httpc.Do(req)
 	if err != nil {
 		return regen, opt.URL, nil, fmt.Errorf("register request: %w", err)
@@ -884,10 +887,11 @@ func (c *Direct) sendMapRequest(ctx context.Context, isStreaming bool, nu Netmap
 		vlogf = c.logf
 	}
 
+	nodeKey := persist.PublicNodeKey()
 	request := &tailcfg.MapRequest{
 		Version:       tailcfg.CurrentCapabilityVersion,
 		KeepAlive:     true,
-		NodeKey:       persist.PublicNodeKey(),
+		NodeKey:       nodeKey,
 		DiscoKey:      c.discoPubKey,
 		Endpoints:     eps,
 		EndpointTypes: epTypes,
@@ -946,6 +950,7 @@ func (c *Direct) sendMapRequest(ctx context.Context, isStreaming bool, nu Netmap
 	if err != nil {
 		return err
 	}
+	addLBHeader(req, nodeKey)
 
 	res, err := httpc.Do(req)
 	if err != nil {
@@ -1537,7 +1542,7 @@ func (c *Direct) setDNSNoise(ctx context.Context, req *tailcfg.SetDNSRequest) er
 	if err != nil {
 		return err
 	}
-	res, err := nc.post(ctx, "/machine/set-dns", &newReq)
+	res, err := nc.post(ctx, "/machine/set-dns", newReq.NodeKey, &newReq)
 	if err != nil {
 		return err
 	}
@@ -1714,8 +1719,10 @@ func (c *Direct) ReportHealthChange(sys health.Subsystem, sysErr error) {
 		// Don't report errors to control if the server doesn't support noise.
 		return
 	}
+	nodeKey := c.GetPersist().PublicNodeKey()
 	req := &tailcfg.HealthChangeRequest{
-		Subsys: string(sys),
+		Subsys:  string(sys),
+		NodeKey: nodeKey,
 	}
 	if sysErr != nil {
 		req.Error = sysErr.Error()
@@ -1724,7 +1731,7 @@ func (c *Direct) ReportHealthChange(sys health.Subsystem, sysErr error) {
 	// Best effort, no logging:
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
-	res, err := np.post(ctx, "/machine/update-health", req)
+	res, err := np.post(ctx, "/machine/update-health", nodeKey, req)
 	if err != nil {
 		return
 	}
@@ -1768,6 +1775,12 @@ func decodeWrappedAuthkey(key string, logf logger.Logf) (authKey string, isWrapp
 	return authKey, true, sig, priv
 }
 
+func addLBHeader(req *http.Request, nodeKey key.NodePublic) {
+	if !nodeKey.IsZero() {
+		req.Header.Add(tailcfg.LBHeader, nodeKey.String())
+	}
+}
+
 var (
 	metricMapRequestsActive = clientmetric.NewGauge("controlclient_map_requests_active")
 

control/controlclient/noise.go

@@ -484,7 +484,9 @@ func (nc *NoiseClient) dial(ctx context.Context) (*noiseConn, error) {
 	return ncc, nil
 }
 
-func (nc *NoiseClient) post(ctx context.Context, path string, body any) (*http.Response, error) {
+// post does a POST to the control server at the given path, JSON-encoding body.
+// The provided nodeKey is an optional load balancing hint.
+func (nc *NoiseClient) post(ctx context.Context, path string, nodeKey key.NodePublic, body any) (*http.Response, error) {
 	jbody, err := json.Marshal(body)
 	if err != nil {
 		return nil, err
@@ -493,6 +495,7 @@ func (nc *NoiseClient) post(ctx context.Context, path string, body any) (*http.R
 	if err != nil {
 		return nil, err
 	}
+	addLBHeader(req, nodeKey)
 	req.Header.Set("Content-Type", "application/json")
 
 	conn, err := nc.getConn(ctx)

control/controlclient/noise_test.go

@@ -128,7 +128,7 @@ func (tt noiseClientTest) run(t *testing.T) {
 	checkRes(t, res)
 
 	// And try using the high-level nc.post API as well.
-	res, err = nc.post(context.Background(), "/", nil)
+	res, err = nc.post(context.Background(), "/", key.NodePublic{}, nil)
 	if err != nil {
 		t.Fatal(err)
 	}

tailcfg/tailcfg.go

@@ -2266,6 +2266,10 @@ type SetDNSResponse struct{}
 type HealthChangeRequest struct {
 	Subsys string // a health.Subsystem value in string form
 	Error  string // or empty if cleared
+
+	// NodeKey is the client's current node key.
+	// In clients <= 1.62.0 it was always the zero value.
+	NodeKey key.NodePublic
 }
 
 // SSHPolicy is the policy for how to handle incoming SSH connections
@@ -2683,3 +2687,21 @@ type EarlyNoise struct {
 	// the client to prove possession of a wireguard private key.
 	NodeKeyChallenge key.ChallengePublic `json:"nodeKeyChallenge"`
 }
+
+// LBHeader is the HTTP request header used to provide a load balancer or
+// internal reverse proxy with information about the request body without the
+// reverse proxy needing to read the body to parse it out. Think of it akin to
+// an HTTP Host header or SNI. The value may be absent (notably for old clients)
+// but if present, it should match the request. A non-empty value that doesn't
+// match the request body's.
+//
+// The possible values depend on the request path, but for /machine (Noise)
+// requests, they'll usually be a node public key (in key.NodePublic.String
+// format), matching the Request JSON body's NodeKey.
+//
+// Note that this is not a security or authentication header; it's strictly
+// denormalized redundant data as an optimization.
+//
+// For some request types, the header may have multiple values. (e.g. OldNodeKey
+// vs NodeKey)
+const LBHeader = "Ts-Lb"