luxolus/open-consul
2
0
Fork 0
Code Issues 1 Pull requests Packages Releases Wiki Activity

Merge pull request #1230 from hashicorp/f-maintfix

Browse source 
Respect tokens in maintenance mode
This commit is contained in:
Ryan Uber 2015-09-10 12:30:07 -07:00
parent 40e6527720 948bd57d6a
commit 08d12e978f
4 changed files with 42 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
command/agent/agent.go
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")
}

12
command/agent/agent_endpoint.go
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()
}

18
command/agent/agent_endpoint_test.go
View file

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

18
command/agent/agent_test.go
View file

@ -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]