Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Respect tokens in maintenance mode #1230

Merged
merged 2 commits into from
Sep 10, 2015
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions command/agent/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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")
}

Expand Down
12 changes: 10 additions & 2 deletions command/agent/agent_endpoint.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
}
Expand Down
18 changes: 14 additions & 4 deletions command/agent/agent_endpoint_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
}

Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand Down
18 changes: 14 additions & 4 deletions command/agent/agent_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand All @@ -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)
Expand All @@ -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)
}

Expand Down Expand Up @@ -1299,14 +1304,19 @@ 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]
if !ok {
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)
Expand All @@ -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]
Expand Down