agent: fix data races with registerEndpoint

Only register a different endpoint after it has been
fully created.

agent/acl_test.go

@@ -46,16 +46,17 @@ func TestACL_Version8(t *testing.T) {
 	a := NewTestAgent(t.Name(), cfg)
 	defer a.Shutdown()
 
-	m := MockServer{}
+	m := MockServer{
+		// With version 8 enforcement off, this should not get called.
+		getPolicyFn: func(*structs.ACLPolicyRequest, *structs.ACLPolicy) error {
+			t.Fatalf("should not have called to server")
+			return nil
+		},
+	}
 	if err := a.registerEndpoint("ACL", &m); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 
-	// With version 8 enforcement off, this should not get called.
-	m.getPolicyFn = func(*structs.ACLPolicyRequest, *structs.ACLPolicy) error {
-		t.Fatalf("should not have called to server")
-		return nil
-	}
 	if token, err := a.resolveToken("nope"); token != nil || err != nil {
 		t.Fatalf("bad: %v err: %v", token, err)
 	}
@@ -69,15 +70,16 @@ func TestACL_Disabled(t *testing.T) {
 	a := NewTestAgent(t.Name(), cfg)
 	defer a.Shutdown()
 
-	m := MockServer{}
+	m := MockServer{
+		// Fetch a token without ACLs enabled and make sure the manager sees it.
+		getPolicyFn: func(*structs.ACLPolicyRequest, *structs.ACLPolicy) error {
+			return errors.New(aclDisabled)
+		},
+	}
 	if err := a.registerEndpoint("ACL", &m); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 
-	// Fetch a token without ACLs enabled and make sure the manager sees it.
-	m.getPolicyFn = func(*structs.ACLPolicyRequest, *structs.ACLPolicy) error {
-		return errors.New(aclDisabled)
-	}
 	if a.acls.isDisabled() {
 		t.Fatalf("should not be disabled yet")
 	}
@@ -122,17 +124,17 @@ func TestACL_Special_IDs(t *testing.T) {
 	a := NewTestAgent(t.Name(), cfg)
 	defer a.Shutdown()
 
-	m := MockServer{}
+	m := MockServer{
-	if err := a.registerEndpoint("ACL", &m); err != nil {
-		t.Fatalf("err: %v", err)
-	}
-
 		// An empty ID should get mapped to the anonymous token.
-	m.getPolicyFn = func(req *structs.ACLPolicyRequest, reply *structs.ACLPolicy) error {
+		getPolicyFn: func(req *structs.ACLPolicyRequest, reply *structs.ACLPolicy) error {
 			if req.ACL != "anonymous" {
 				t.Fatalf("bad: %#v", *req)
 			}
 			return errors.New(aclNotFound)
+		},
+	}
+	if err := a.registerEndpoint("ACL", &m); err != nil {
+		t.Fatalf("err: %v", err)
 	}
 	_, err := a.resolveToken("")
 	if err == nil || !strings.Contains(err.Error(), aclNotFound) {
@@ -175,15 +177,16 @@ func TestACL_Down_Deny(t *testing.T) {
 	a := NewTestAgent(t.Name(), cfg)
 	defer a.Shutdown()
 
-	m := MockServer{}
+	m := MockServer{
+		// Resolve with ACLs down.
+		getPolicyFn: func(*structs.ACLPolicyRequest, *structs.ACLPolicy) error {
+			return fmt.Errorf("ACLs are broken")
+		},
+	}
 	if err := a.registerEndpoint("ACL", &m); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 
-	// Resolve with ACLs down.
-	m.getPolicyFn = func(*structs.ACLPolicyRequest, *structs.ACLPolicy) error {
-		return fmt.Errorf("ACLs are broken")
-	}
 	acl, err := a.resolveToken("nope")
 	if err != nil {
 		t.Fatalf("err: %v", err)
@@ -205,15 +208,16 @@ func TestACL_Down_Allow(t *testing.T) {
 	a := NewTestAgent(t.Name(), cfg)
 	defer a.Shutdown()
 
-	m := MockServer{}
+	m := MockServer{
+		// Resolve with ACLs down.
+		getPolicyFn: func(*structs.ACLPolicyRequest, *structs.ACLPolicy) error {
+			return fmt.Errorf("ACLs are broken")
+		},
+	}
 	if err := a.registerEndpoint("ACL", &m); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 
-	// Resolve with ACLs down.
-	m.getPolicyFn = func(*structs.ACLPolicyRequest, *structs.ACLPolicy) error {
-		return fmt.Errorf("ACLs are broken")
-	}
 	acl, err := a.resolveToken("nope")
 	if err != nil {
 		t.Fatalf("err: %v", err)
@@ -235,13 +239,9 @@ func TestACL_Down_Extend(t *testing.T) {
 	a := NewTestAgent(t.Name(), cfg)
 	defer a.Shutdown()
 
-	m := MockServer{}
+	m := MockServer{
-	if err := a.registerEndpoint("ACL", &m); err != nil {
-		t.Fatalf("err: %v", err)
-	}
-
 		// Populate the cache for one of the tokens.
-	m.getPolicyFn = func(req *structs.ACLPolicyRequest, reply *structs.ACLPolicy) error {
+		getPolicyFn: func(req *structs.ACLPolicyRequest, reply *structs.ACLPolicy) error {
 			*reply = structs.ACLPolicy{
 				Parent: "allow",
 				Policy: &rawacl.Policy{
@@ -254,7 +254,12 @@ func TestACL_Down_Extend(t *testing.T) {
 				},
 			}
 			return nil
+		},
 	}
+	if err := a.registerEndpoint("ACL", &m); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
 	acl, err := a.resolveToken("yep")
 	if err != nil {
 		t.Fatalf("err: %v", err)
@@ -312,13 +317,9 @@ func TestACL_Cache(t *testing.T) {
 	a := NewTestAgent(t.Name(), cfg)
 	defer a.Shutdown()
 
-	m := MockServer{}
+	m := MockServer{
-	if err := a.registerEndpoint("ACL", &m); err != nil {
-		t.Fatalf("err: %v", err)
-	}
-
 		// Populate the cache for one of the tokens.
-	m.getPolicyFn = func(req *structs.ACLPolicyRequest, reply *structs.ACLPolicy) error {
+		getPolicyFn: func(req *structs.ACLPolicyRequest, reply *structs.ACLPolicy) error {
 			*reply = structs.ACLPolicy{
 				ETag:   "hash1",
 				Parent: "deny",
@@ -333,7 +334,12 @@ func TestACL_Cache(t *testing.T) {
 				TTL: 10 * time.Millisecond,
 			}
 			return nil
+		},
 	}
+	if err := a.registerEndpoint("ACL", &m); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
 	acl, err := a.resolveToken("yep")
 	if err != nil {
 		t.Fatalf("err: %v", err)

agent/dns_test.go

@@ -3919,15 +3919,16 @@ func TestDNS_PreparedQuery_AllowStale(t *testing.T) {
 	a := NewTestAgent(t.Name(), cfg)
 	defer a.Shutdown()
 
-	m := MockPreparedQuery{}
+	m := MockPreparedQuery{
-	if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+		executeFn: func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
-		t.Fatalf("err: %v", err)
-	}
-
-	m.executeFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
 			// Return a response that's perpetually too stale.
 			reply.LastContact = 2 * time.Second
 			return nil
+		},
+	}
+
+	if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+		t.Fatalf("err: %v", err)
 	}
 
 	// Make sure that the lookup terminates and results in an SOA since
@@ -4000,12 +4001,8 @@ func TestDNS_PreparedQuery_AgentSource(t *testing.T) {
 	a := NewTestAgent(t.Name(), nil)
 	defer a.Shutdown()
 
-	m := MockPreparedQuery{}
+	m := MockPreparedQuery{
-	if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+		executeFn: func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
-		t.Fatalf("err: %v", err)
-	}
-
-	m.executeFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
 			// Check that the agent inserted its self-name and datacenter to
 			// the RPC request body.
 			if args.Agent.Datacenter != a.Config.Datacenter ||
@@ -4013,6 +4010,11 @@ func TestDNS_PreparedQuery_AgentSource(t *testing.T) {
 				t.Fatalf("bad: %#v", args.Agent)
 			}
 			return nil
+		},
+	}
+
+	if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+		t.Fatalf("err: %v", err)
 	}
 
 	{

agent/prepared_query_endpoint_test.go

@@ -73,12 +73,8 @@ func TestPreparedQuery_Create(t *testing.T) {
 	a := NewTestAgent(t.Name(), nil)
 	defer a.Shutdown()
 
-	m := MockPreparedQuery{}
+	m := MockPreparedQuery{
-	if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+		applyFn: func(args *structs.PreparedQueryRequest, reply *string) error {
-		t.Fatalf("err: %v", err)
-	}
-
-	m.applyFn = func(args *structs.PreparedQueryRequest, reply *string) error {
 			expected := &structs.PreparedQueryRequest{
 				Datacenter: "dc1",
 				Op:         structs.PreparedQueryCreate,
@@ -109,6 +105,10 @@ func TestPreparedQuery_Create(t *testing.T) {
 
 			*reply = "my-id"
 			return nil
+		},
+	}
+	if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+		t.Fatalf("err: %v", err)
 	}
 
 	body := bytes.NewBuffer(nil)
@@ -158,14 +158,14 @@ func TestPreparedQuery_List(t *testing.T) {
 		a := NewTestAgent(t.Name(), nil)
 		defer a.Shutdown()
 
-		m := MockPreparedQuery{}
+		m := MockPreparedQuery{
-		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			listFn: func(args *structs.DCSpecificRequest, reply *structs.IndexedPreparedQueries) error {
-			t.Fatalf("err: %v", err)
-		}
-
-		m.listFn = func(args *structs.DCSpecificRequest, reply *structs.IndexedPreparedQueries) error {
 				// Return an empty response.
 				return nil
+			},
+		}
+		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			t.Fatalf("err: %v", err)
 		}
 
 		body := bytes.NewBuffer(nil)
@@ -191,12 +191,8 @@ func TestPreparedQuery_List(t *testing.T) {
 		a := NewTestAgent(t.Name(), nil)
 		defer a.Shutdown()
 
-		m := MockPreparedQuery{}
+		m := MockPreparedQuery{
-		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			listFn: func(args *structs.DCSpecificRequest, reply *structs.IndexedPreparedQueries) error {
-			t.Fatalf("err: %v", err)
-		}
-
-		m.listFn = func(args *structs.DCSpecificRequest, reply *structs.IndexedPreparedQueries) error {
 				expected := &structs.DCSpecificRequest{
 					Datacenter: "dc1",
 					QueryOptions: structs.QueryOptions{
@@ -213,6 +209,10 @@ func TestPreparedQuery_List(t *testing.T) {
 				}
 				reply.Queries = append(reply.Queries, query)
 				return nil
+			},
+		}
+		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			t.Fatalf("err: %v", err)
 		}
 
 		body := bytes.NewBuffer(nil)
@@ -241,14 +241,14 @@ func TestPreparedQuery_Execute(t *testing.T) {
 		a := NewTestAgent(t.Name(), nil)
 		defer a.Shutdown()
 
-		m := MockPreparedQuery{}
+		m := MockPreparedQuery{
-		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			executeFn: func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
-			t.Fatalf("err: %v", err)
-		}
-
-		m.executeFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
 				// Just return an empty response.
 				return nil
+			},
+		}
+		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			t.Fatalf("err: %v", err)
 		}
 
 		body := bytes.NewBuffer(nil)
@@ -274,12 +274,8 @@ func TestPreparedQuery_Execute(t *testing.T) {
 		a := NewTestAgent(t.Name(), nil)
 		defer a.Shutdown()
 
-		m := MockPreparedQuery{}
+		m := MockPreparedQuery{
-		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			executeFn: func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
-			t.Fatalf("err: %v", err)
-		}
-
-		m.executeFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
 				expected := &structs.PreparedQueryExecuteRequest{
 					Datacenter:    "dc1",
 					QueryIDOrName: "my-id",
@@ -304,6 +300,10 @@ func TestPreparedQuery_Execute(t *testing.T) {
 				// Just set something so we can tell this is returned.
 				reply.Failovers = 99
 				return nil
+			},
+		}
+		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			t.Fatalf("err: %v", err)
 		}
 
 		body := bytes.NewBuffer(nil)
@@ -330,12 +330,8 @@ func TestPreparedQuery_Execute(t *testing.T) {
 		a := NewTestAgent(t.Name(), nil)
 		defer a.Shutdown()
 
-		m := MockPreparedQuery{}
+		m := MockPreparedQuery{
-		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			executeFn: func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
-			t.Fatalf("err: %v", err)
-		}
-
-		m.executeFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
 				if args.Source.Node != "" {
 					t.Fatalf("expect node to be empty, got %q", args.Source.Node)
 				}
@@ -347,6 +343,10 @@ func TestPreparedQuery_Execute(t *testing.T) {
 					t.Fatalf("expect: %#v\nactual: %#v", expect, args.Agent)
 				}
 				return nil
+			},
+		}
+		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			t.Fatalf("err: %v", err)
 		}
 
 		req, _ := http.NewRequest("GET", "/v1/query/my-id/execute", nil)
@@ -364,12 +364,8 @@ func TestPreparedQuery_Execute(t *testing.T) {
 		a := NewTestAgent(t.Name(), cfg)
 		defer a.Shutdown()
 
-		m := MockPreparedQuery{}
+		m := MockPreparedQuery{
-		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			executeFn: func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
-			t.Fatalf("err: %v", err)
-		}
-
-		m.executeFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
 				nodesResponse := make(structs.CheckServiceNodes, 1)
 				nodesResponse[0].Node = &structs.Node{
 					Node: "foo", Address: "127.0.0.1",
@@ -380,6 +376,10 @@ func TestPreparedQuery_Execute(t *testing.T) {
 				reply.Nodes = nodesResponse
 				reply.Datacenter = "dc2"
 				return nil
+			},
+		}
+		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			t.Fatalf("err: %v", err)
 		}
 
 		body := bytes.NewBuffer(nil)
@@ -414,12 +414,8 @@ func TestPreparedQuery_Execute(t *testing.T) {
 		a := NewTestAgent(t.Name(), cfg)
 		defer a.Shutdown()
 
-		m := MockPreparedQuery{}
+		m := MockPreparedQuery{
-		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			executeFn: func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
-			t.Fatalf("err: %v", err)
-		}
-
-		m.executeFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
 				nodesResponse := make(structs.CheckServiceNodes, 1)
 				nodesResponse[0].Node = &structs.Node{
 					Node: "foo", Address: "127.0.0.1",
@@ -430,6 +426,10 @@ func TestPreparedQuery_Execute(t *testing.T) {
 				reply.Nodes = nodesResponse
 				reply.Datacenter = "dc1"
 				return nil
+			},
+		}
+		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			t.Fatalf("err: %v", err)
 		}
 
 		body := bytes.NewBuffer(nil)
@@ -478,12 +478,8 @@ func TestPreparedQuery_Explain(t *testing.T) {
 		a := NewTestAgent(t.Name(), nil)
 		defer a.Shutdown()
 
-		m := MockPreparedQuery{}
+		m := MockPreparedQuery{
-		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			explainFn: func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExplainResponse) error {
-			t.Fatalf("err: %v", err)
-		}
-
-		m.explainFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExplainResponse) error {
 				expected := &structs.PreparedQueryExecuteRequest{
 					Datacenter:    "dc1",
 					QueryIDOrName: "my-id",
@@ -508,6 +504,10 @@ func TestPreparedQuery_Explain(t *testing.T) {
 				// Just set something so we can tell this is returned.
 				reply.Query.Name = "hello"
 				return nil
+			},
+		}
+		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			t.Fatalf("err: %v", err)
 		}
 
 		body := bytes.NewBuffer(nil)
@@ -551,12 +551,8 @@ func TestPreparedQuery_Get(t *testing.T) {
 		a := NewTestAgent(t.Name(), nil)
 		defer a.Shutdown()
 
-		m := MockPreparedQuery{}
+		m := MockPreparedQuery{
-		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			getFn: func(args *structs.PreparedQuerySpecificRequest, reply *structs.IndexedPreparedQueries) error {
-			t.Fatalf("err: %v", err)
-		}
-
-		m.getFn = func(args *structs.PreparedQuerySpecificRequest, reply *structs.IndexedPreparedQueries) error {
 				expected := &structs.PreparedQuerySpecificRequest{
 					Datacenter: "dc1",
 					QueryID:    "my-id",
@@ -574,6 +570,10 @@ func TestPreparedQuery_Get(t *testing.T) {
 				}
 				reply.Queries = append(reply.Queries, query)
 				return nil
+			},
+		}
+		if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+			t.Fatalf("err: %v", err)
 		}
 
 		body := bytes.NewBuffer(nil)
@@ -616,12 +616,8 @@ func TestPreparedQuery_Update(t *testing.T) {
 	a := NewTestAgent(t.Name(), nil)
 	defer a.Shutdown()
 
-	m := MockPreparedQuery{}
+	m := MockPreparedQuery{
-	if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+		applyFn: func(args *structs.PreparedQueryRequest, reply *string) error {
-		t.Fatalf("err: %v", err)
-	}
-
-	m.applyFn = func(args *structs.PreparedQueryRequest, reply *string) error {
 			expected := &structs.PreparedQueryRequest{
 				Datacenter: "dc1",
 				Op:         structs.PreparedQueryUpdate,
@@ -653,6 +649,10 @@ func TestPreparedQuery_Update(t *testing.T) {
 
 			*reply = "don't care"
 			return nil
+		},
+	}
+	if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+		t.Fatalf("err: %v", err)
 	}
 
 	body := bytes.NewBuffer(nil)
@@ -694,12 +694,8 @@ func TestPreparedQuery_Delete(t *testing.T) {
 	a := NewTestAgent(t.Name(), nil)
 	defer a.Shutdown()
 
-	m := MockPreparedQuery{}
+	m := MockPreparedQuery{
-	if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+		applyFn: func(args *structs.PreparedQueryRequest, reply *string) error {
-		t.Fatalf("err: %v", err)
-	}
-
-	m.applyFn = func(args *structs.PreparedQueryRequest, reply *string) error {
 			expected := &structs.PreparedQueryRequest{
 				Datacenter: "dc1",
 				Op:         structs.PreparedQueryDelete,
@@ -716,6 +712,10 @@ func TestPreparedQuery_Delete(t *testing.T) {
 
 			*reply = "don't care"
 			return nil
+		},
+	}
+	if err := a.registerEndpoint("PreparedQuery", &m); err != nil {
+		t.Fatalf("err: %v", err)
 	}
 
 	body := bytes.NewBuffer(nil)