From db5ee0e4d21c039a990889fcc7d59ec6bf7d4673 Mon Sep 17 00:00:00 2001
From: "Chris S. Kim" <ckim@hashicorp.com>
Date: Fri, 19 Nov 2021 11:50:44 -0500
Subject: [PATCH] Merge from ent (#11506)

---
 agent/agent_endpoint.go      | 37 ++++++++++++++++--------------------
 agent/agent_endpoint_test.go |  9 +++------
 2 files changed, 19 insertions(+), 27 deletions(-)

diff --git a/agent/agent_endpoint.go b/agent/agent_endpoint.go
index bb45fc541c..e1ad1f2a11 100644
--- a/agent/agent_endpoint.go
+++ b/agent/agent_endpoint.go
@@ -664,22 +664,16 @@ func (s *HTTPHandlers) AgentRegisterCheck(resp http.ResponseWriter, req *http.Re
 	}
 
 	if err := decodeBody(req.Body, &args); err != nil {
-		resp.WriteHeader(http.StatusBadRequest)
-		fmt.Fprintf(resp, "Request decode failed: %v", err)
-		return nil, nil
+		return nil, BadRequestError{fmt.Sprintf("Request decode failed: %v", err)}
 	}
 
 	// Verify the check has a name.
 	if args.Name == "" {
-		resp.WriteHeader(http.StatusBadRequest)
-		fmt.Fprint(resp, "Missing check name")
-		return nil, nil
+		return nil, BadRequestError{"Missing check name"}
 	}
 
 	if args.Status != "" && !structs.ValidStatus(args.Status) {
-		resp.WriteHeader(http.StatusBadRequest)
-		fmt.Fprint(resp, "Bad check status")
-		return nil, nil
+		return nil, BadRequestError{"Bad check status"}
 	}
 
 	authz, err := s.agent.delegate.ResolveTokenAndDefaultMeta(token, &args.EnterpriseMeta, nil)
@@ -698,19 +692,20 @@ func (s *HTTPHandlers) AgentRegisterCheck(resp http.ResponseWriter, req *http.Re
 	chkType := args.CheckType()
 	err = chkType.Validate()
 	if err != nil {
-		resp.WriteHeader(http.StatusBadRequest)
-		fmt.Fprint(resp, fmt.Errorf("Invalid check: %v", err))
-		return nil, nil
+		return nil, BadRequestError{fmt.Sprintf("Invalid check: %v", err)}
 	}
 
 	// Store the type of check based on the definition
 	health.Type = chkType.Type()
 
 	if health.ServiceID != "" {
+		cid := health.CompoundServiceID()
 		// fixup the service name so that vetCheckRegister requires the right ACLs
-		service := s.agent.State.Service(health.CompoundServiceID())
+		service := s.agent.State.Service(cid)
 		if service != nil {
 			health.ServiceName = service.Service
+		} else {
+			return nil, NotFoundError{fmt.Sprintf("ServiceID %q does not exist", cid.String())}
 		}
 	}
 
@@ -746,14 +741,14 @@ func (s *HTTPHandlers) AgentDeregisterCheck(resp http.ResponseWriter, req *http.
 
 	checkID.Normalize()
 
-	if err := s.agent.vetCheckUpdateWithAuthorizer(authz, checkID); err != nil {
-		return nil, err
-	}
-
 	if !s.validateRequestPartition(resp, &checkID.EnterpriseMeta) {
 		return nil, nil
 	}
 
+	if err := s.agent.vetCheckUpdateWithAuthorizer(authz, checkID); err != nil {
+		return nil, err
+	}
+
 	if err := s.agent.RemoveCheck(checkID, true); err != nil {
 		return nil, err
 	}
@@ -1205,14 +1200,14 @@ func (s *HTTPHandlers) AgentDeregisterService(resp http.ResponseWriter, req *htt
 
 	sid.Normalize()
 
-	if err := s.agent.vetServiceUpdateWithAuthorizer(authz, sid); err != nil {
-		return nil, err
-	}
-
 	if !s.validateRequestPartition(resp, &sid.EnterpriseMeta) {
 		return nil, nil
 	}
 
+	if err := s.agent.vetServiceUpdateWithAuthorizer(authz, sid); err != nil {
+		return nil, err
+	}
+
 	if err := s.agent.RemoveService(sid); err != nil {
 		return nil, err
 	}
diff --git a/agent/agent_endpoint_test.go b/agent/agent_endpoint_test.go
index 7e8f3d7ffa..9c2da7a9e2 100644
--- a/agent/agent_endpoint_test.go
+++ b/agent/agent_endpoint_test.go
@@ -2516,12 +2516,9 @@ func TestAgent_RegisterCheck_BadStatus(t *testing.T) {
 	}
 	req, _ := http.NewRequest("PUT", "/v1/agent/check/register", jsonReader(args))
 	resp := httptest.NewRecorder()
-	if _, err := a.srv.AgentRegisterCheck(resp, req); err != nil {
-		t.Fatalf("err: %v", err)
-	}
-	if resp.Code != 400 {
-		t.Fatalf("accepted bad status")
-	}
+	a.srv.h.ServeHTTP(resp, req)
+	require.Equalf(t, http.StatusBadRequest, resp.Code, "resp: %v", resp.Body.String())
+	require.Contains(t, resp.Body.String(), "Bad check status")
 }
 
 func TestAgent_RegisterCheck_ACLDeny(t *testing.T) {