diff --git a/command/agent/agent.go b/command/agent/agent.go index 397f616a4335..6e11a535cff5 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 a4b5f6ca94b4..a3baa4deb67a 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() } diff --git a/command/agent/agent_endpoint_test.go b/command/agent/agent_endpoint_test.go index b0f7a481c1bb..c89977049771 100644 --- a/command/agent/agent_endpoint_test.go +++ b/command/agent/agent_endpoint_test.go @@ -655,7 +655,7 @@ func TestHTTPAgent_EnableServiceMaintenance(t *testing.T) { } // Force the service into maintenance mode - req, _ := http.NewRequest("PUT", "/v1/agent/service/maintenance/test?enable=true&reason=broken", nil) + req, _ := http.NewRequest("PUT", "/v1/agent/service/maintenance/test?enable=true&reason=broken&token=mytoken", nil) resp := httptest.NewRecorder() if _, err := srv.AgentServiceMaintenance(resp, req); err != nil { t.Fatalf("err: %s", err) @@ -671,6 +671,11 @@ func TestHTTPAgent_EnableServiceMaintenance(t *testing.T) { t.Fatalf("should have registered maintenance check") } + // Ensure the token was added + if token := srv.agent.state.CheckToken(checkID); token != "mytoken" { + t.Fatalf("expected 'mytoken', got '%s'", token) + } + // Ensure the reason was set in notes if check.Notes != "broken" { t.Fatalf("bad: %#v", check) @@ -693,7 +698,7 @@ func TestHTTPAgent_DisableServiceMaintenance(t *testing.T) { } // Force the service into maintenance mode - if err := srv.agent.EnableServiceMaintenance("test", ""); err != nil { + if err := srv.agent.EnableServiceMaintenance("test", "", ""); err != nil { t.Fatalf("err: %s", err) } @@ -749,7 +754,7 @@ func TestHTTPAgent_EnableNodeMaintenance(t *testing.T) { // Force the node into maintenance mode req, _ := http.NewRequest( - "PUT", "/v1/agent/self/maintenance?enable=true&reason=broken", nil) + "PUT", "/v1/agent/self/maintenance?enable=true&reason=broken&token=mytoken", nil) resp := httptest.NewRecorder() if _, err := srv.AgentNodeMaintenance(resp, req); err != nil { t.Fatalf("err: %s", err) @@ -764,6 +769,11 @@ func TestHTTPAgent_EnableNodeMaintenance(t *testing.T) { t.Fatalf("should have registered maintenance check") } + // Check that the token was used + if token := srv.agent.state.CheckToken(nodeMaintCheckID); token != "mytoken" { + t.Fatalf("expected 'mytoken', got '%s'", token) + } + // Ensure the reason was set in notes if check.Notes != "broken" { t.Fatalf("bad: %#v", check) @@ -777,7 +787,7 @@ func TestHTTPAgent_DisableNodeMaintenance(t *testing.T) { defer srv.agent.Shutdown() // Force the node into maintenance mode - srv.agent.EnableNodeMaintenance("") + srv.agent.EnableNodeMaintenance("", "") // Leave maintenance mode req, _ := http.NewRequest("PUT", "/v1/agent/self/maintenance?enable=false", nil) diff --git a/command/agent/agent_test.go b/command/agent/agent_test.go index e5fcfd09db96..d735f60f303d 100644 --- a/command/agent/agent_test.go +++ b/command/agent/agent_test.go @@ -1207,7 +1207,7 @@ func TestAgent_ServiceMaintenanceMode(t *testing.T) { } // Enter maintenance mode for the service - if err := agent.EnableServiceMaintenance("redis", "broken"); err != nil { + if err := agent.EnableServiceMaintenance("redis", "broken", "mytoken"); err != nil { t.Fatalf("err: %s", err) } @@ -1218,6 +1218,11 @@ func TestAgent_ServiceMaintenanceMode(t *testing.T) { t.Fatalf("should have registered critical maintenance check") } + // Check that the token was used to register the check + if token := agent.state.CheckToken(checkID); token != "mytoken" { + t.Fatalf("expected 'mytoken', got: '%s'", token) + } + // Ensure the reason was set in notes if check.Notes != "broken" { t.Fatalf("bad: %#v", check) @@ -1234,7 +1239,7 @@ func TestAgent_ServiceMaintenanceMode(t *testing.T) { } // Enter service maintenance mode without providing a reason - if err := agent.EnableServiceMaintenance("redis", ""); err != nil { + if err := agent.EnableServiceMaintenance("redis", "", ""); err != nil { t.Fatalf("err: %s", err) } @@ -1299,7 +1304,7 @@ func TestAgent_NodeMaintenanceMode(t *testing.T) { defer agent.Shutdown() // Enter maintenance mode for the node - agent.EnableNodeMaintenance("broken") + agent.EnableNodeMaintenance("broken", "mytoken") // Make sure the critical health check was added check, ok := agent.state.Checks()[nodeMaintCheckID] @@ -1307,6 +1312,11 @@ func TestAgent_NodeMaintenanceMode(t *testing.T) { t.Fatalf("should have registered critical node check") } + // Check that the token was used to register the check + if token := agent.state.CheckToken(nodeMaintCheckID); token != "mytoken" { + t.Fatalf("expected 'mytoken', got: '%s'", token) + } + // Ensure the reason was set in notes if check.Notes != "broken" { t.Fatalf("bad: %#v", check) @@ -1321,7 +1331,7 @@ func TestAgent_NodeMaintenanceMode(t *testing.T) { } // Enter maintenance mode without passing a reason - agent.EnableNodeMaintenance("") + agent.EnableNodeMaintenance("", "") // Make sure the check was registered with the default note check, ok = agent.state.Checks()[nodeMaintCheckID]