From e129a5931639e94af0151d59e6d53cdecbb18e42 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Thu, 10 Sep 2015 11:43:59 -0700
Subject: [PATCH] agent: thread tokens through for maintenance mode

---
 command/agent/agent.go          |  8 ++++----
 command/agent/agent_endpoint.go | 12 ++++++++++--
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/command/agent/agent.go b/command/agent/agent.go
index 397f616a4..6e11a535c 100644
--- a/command/agent/agent.go
+++ b/command/agent/agent.go
@@ -1276,7 +1276,7 @@ func serviceMaintCheckID(serviceID string) string {
 
 // EnableServiceMaintenance will register a false health check against the given
 // service ID with critical status. This will exclude the service from queries.
-func (a *Agent) EnableServiceMaintenance(serviceID, reason string) error {
+func (a *Agent) EnableServiceMaintenance(serviceID, reason, token string) error {
 	service, ok := a.state.Services()[serviceID]
 	if !ok {
 		return fmt.Errorf("No service registered with ID %q", serviceID)
@@ -1303,7 +1303,7 @@ func (a *Agent) EnableServiceMaintenance(serviceID, reason string) error {
 		ServiceName: service.Service,
 		Status:      structs.HealthCritical,
 	}
-	a.AddCheck(check, nil, true, "")
+	a.AddCheck(check, nil, true, token)
 	a.logger.Printf("[INFO] agent: Service %q entered maintenance mode", serviceID)
 
 	return nil
@@ -1330,7 +1330,7 @@ func (a *Agent) DisableServiceMaintenance(serviceID string) error {
 }
 
 // EnableNodeMaintenance places a node into maintenance mode.
-func (a *Agent) EnableNodeMaintenance(reason string) {
+func (a *Agent) EnableNodeMaintenance(reason, token string) {
 	// Ensure node maintenance is not already enabled
 	if _, ok := a.state.Checks()[nodeMaintCheckID]; ok {
 		return
@@ -1349,7 +1349,7 @@ func (a *Agent) EnableNodeMaintenance(reason string) {
 		Notes:   reason,
 		Status:  structs.HealthCritical,
 	}
-	a.AddCheck(check, nil, true, "")
+	a.AddCheck(check, nil, true, token)
 	a.logger.Printf("[INFO] agent: Node entered maintenance mode")
 }
 
diff --git a/command/agent/agent_endpoint.go b/command/agent/agent_endpoint.go
index a4b5f6ca9..a3baa4deb 100644
--- a/command/agent/agent_endpoint.go
+++ b/command/agent/agent_endpoint.go
@@ -266,9 +266,13 @@ func (s *HTTPServer) AgentServiceMaintenance(resp http.ResponseWriter, req *http
 		return nil, nil
 	}
 
+	// Get the provided token, if any
+	var token string
+	s.parseToken(req, &token)
+
 	if enable {
 		reason := params.Get("reason")
-		if err = s.agent.EnableServiceMaintenance(serviceID, reason); err != nil {
+		if err = s.agent.EnableServiceMaintenance(serviceID, reason, token); err != nil {
 			resp.WriteHeader(404)
 			resp.Write([]byte(err.Error()))
 			return nil, nil
@@ -307,8 +311,12 @@ func (s *HTTPServer) AgentNodeMaintenance(resp http.ResponseWriter, req *http.Re
 		return nil, nil
 	}
 
+	// Get the provided token, if any
+	var token string
+	s.parseToken(req, &token)
+
 	if enable {
-		s.agent.EnableNodeMaintenance(params.Get("reason"))
+		s.agent.EnableNodeMaintenance(params.Get("reason"), token)
 	} else {
 		s.agent.DisableNodeMaintenance()
 	}