agent: thread tokens through for maintenance mode

This commit is contained in:
Ryan Uber 2015-09-10 11:43:59 -07:00
parent 10c38f8c3a
commit e129a59316
2 changed files with 14 additions and 6 deletions

View File

@ -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")
}

View File

@ -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()
}