Fix /machine/map endpoint vulnerability

hscontrol/noise.go

@@ -100,6 +100,10 @@ func (h *Headscale) NoiseUpgradeHandler(
 
 	router.HandleFunc("/machine/register", noiseServer.NoiseRegistrationHandler).
 		Methods(http.MethodPost)
+
+	// Endpoints outside of the register endpoint must use getAndValidateNode to
+	// get the node to ensure that the MachineKey matches the Node setting up the
+	// connection.
 	router.HandleFunc("/machine/map", noiseServer.NoisePollNetMapHandler)
 
 	noiseServer.httpBaseConfig = &http.Server{
@@ -209,18 +213,14 @@ func (ns *noiseServer) NoisePollNetMapHandler(
 		return
 	}
 
-	ns.nodeKey = mapRequest.NodeKey
-
-	node, err := ns.headscale.db.GetNodeByNodeKey(mapRequest.NodeKey)
+	node, err := ns.getAndValidateNode(mapRequest)
 	if err != nil {
-		if errors.Is(err, gorm.ErrRecordNotFound) {
-			httpError(writer, NewHTTPError(http.StatusNotFound, "node not found", nil))
-			return
-		}
 		httpError(writer, err)
 		return
 	}
 
+	ns.nodeKey = node.NodeKey
+
 	sess := ns.headscale.newMapSession(req.Context(), mapRequest, writer, node)
 	sess.tracef("a node sending a MapRequest with Noise protocol")
 	if !sess.isStreaming() {
@@ -266,8 +266,8 @@ func (ns *noiseServer) NoiseRegistrationHandler(
 					Error: httpErr.Msg,
 				}
 				return &regReq, resp
-			} else {
 			}
+
 			return &regReq, regErr(err)
 		}
 
@@ -289,3 +289,22 @@ func (ns *noiseServer) NoiseRegistrationHandler(
 	writer.WriteHeader(http.StatusOK)
 	writer.Write(respBody)
 }
+
+// getAndValidateNode retrieves the node from the database using the NodeKey
+// and validates that it matches the MachineKey from the Noise session.
+func (ns *noiseServer) getAndValidateNode(mapRequest tailcfg.MapRequest) (*types.Node, error) {
+	node, err := ns.headscale.db.GetNodeByNodeKey(mapRequest.NodeKey)
+	if err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, NewHTTPError(http.StatusNotFound, "node not found", nil)
+		}
+		return nil, err
+	}
+
+	// Validate that the MachineKey in the Noise session matches the one associated with the NodeKey.
+	if ns.machineKey != node.MachineKey {
+		return nil, NewHTTPError(http.StatusNotFound, "node key in request does not match the one associated with this machine key", nil)
+	}
+
+	return node, nil
+}