Merge pull request #1230 from hashicorp/f-maintfix

Respect tokens in maintenance mode

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

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

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)

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