agent: restore tokens for services and checks in config

command/agent/agent.go

@@ -523,7 +523,7 @@ func (a *Agent) ResumeSync() {
 func (a *Agent) persistService(service *structs.NodeService) error {
 	svcPath := filepath.Join(a.config.DataDir, servicesDir, stringHash(service.ID))
 	if _, err := os.Stat(svcPath); os.IsNotExist(err) {
-		wrapped := &persistedService{
+		wrapped := persistedService{
 			Token:   a.state.ServiceToken(service.ID),
 			Service: service,
 		}
@@ -563,9 +563,13 @@ func (a *Agent) persistCheck(check *structs.HealthCheck, chkType *CheckType) err
 	}
 
 	// Create the persisted check
-	p := persistedCheck{check, chkType, a.state.CheckToken(check.CheckID)}
+	wrapped := persistedCheck{
+		Check:   check,
+		ChkType: chkType,
+		Token:   a.state.CheckToken(check.CheckID),
+	}
 
-	encoded, err := json.Marshal(p)
+	encoded, err := json.Marshal(wrapped)
 	if err != nil {
 		return nil
 	}
@@ -924,6 +928,7 @@ func (a *Agent) loadServices(conf *Config) error {
 	for _, service := range conf.Services {
 		ns := service.NodeService()
 		chkTypes := service.CheckTypes()
+		a.state.SetServiceToken(service.ID, service.Token)
 		if err := a.AddService(ns, chkTypes, false); err != nil {
 			return fmt.Errorf("Failed to register service '%s': %v", service.ID, err)
 		}
@@ -1002,6 +1007,7 @@ func (a *Agent) loadChecks(conf *Config) error {
 	for _, check := range conf.Checks {
 		health := check.HealthCheck(conf.NodeName)
 		chkType := &check.CheckType
+		a.state.SetCheckToken(check.ID, check.Token)
 		if err := a.AddCheck(health, chkType, false); err != nil {
 			return fmt.Errorf("Failed to register check '%s': %v %v", check.Name, err, check)
 		}

command/agent/agent_test.go

@@ -541,7 +541,7 @@ func TestAgent_PersistService(t *testing.T) {
 		t.Fatalf("err: %s", err)
 	}
 
-	expected, err := json.Marshal(&persistedService{
+	expected, err := json.Marshal(persistedService{
 		Token:   "hello",
 		Service: svc,
 	})
@@ -572,7 +572,7 @@ func TestAgent_PersistService(t *testing.T) {
 	}
 }
 
-func TestAgent_persistedCheck_compat(t *testing.T) {
+func TestAgent_persistedService_compat(t *testing.T) {
 	// Tests backwards compatibility of persisted services from pre-0.5.1
 	config := nextConfig()
 	dir, agent := makeAgent(t, config)
@@ -742,8 +742,11 @@ func TestAgent_PersistCheck(t *testing.T) {
 		t.Fatalf("err: %s", err)
 	}
 
-	p := persistedCheck{check, chkType, "hello"}
-	expected, err := json.Marshal(p)
+	expected, err := json.Marshal(persistedCheck{
+		Check:   check,
+		ChkType: chkType,
+		Token:   "hello",
+	})
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
@@ -763,7 +766,7 @@ func TestAgent_PersistCheck(t *testing.T) {
 	}
 	defer agent2.Shutdown()
 
-	result, ok := agent2.state.checks[p.Check.CheckID]
+	result, ok := agent2.state.checks[check.CheckID]
 	if !ok {
 		t.Fatalf("bad: %#v", agent2.state.checks)
 	}
@@ -772,11 +775,11 @@ func TestAgent_PersistCheck(t *testing.T) {
 	}
 
 	// Should have restored the monitor
-	if _, ok := agent2.checkMonitors[p.Check.CheckID]; !ok {
+	if _, ok := agent2.checkMonitors[check.CheckID]; !ok {
 		t.Fatalf("bad: %#v", agent2.checkMonitors)
 	}
-	if agent2.state.checkTokens[p.Check.CheckID] != "hello" {
-		t.Fatalf("bad: %s", agent2.state.checkTokens[p.Check.CheckID])
+	if agent2.state.checkTokens[check.CheckID] != "hello" {
+		t.Fatalf("bad: %s", agent2.state.checkTokens[check.CheckID])
 	}
 }
 
@@ -867,6 +870,29 @@ func TestAgent_PurgeCheckOnDuplicate(t *testing.T) {
 	}
 }
 
+func TestAgent_loadChecks_token(t *testing.T) {
+	config := nextConfig()
+	config.Checks = append(config.Checks, &CheckDefinition{
+		ID:    "rabbitmq",
+		Name:  "rabbitmq",
+		Token: "abc123",
+		CheckType: CheckType{
+			TTL: 10 * time.Second,
+		},
+	})
+	dir, agent := makeAgent(t, config)
+	defer os.RemoveAll(dir)
+	defer agent.Shutdown()
+
+	checks := agent.state.Checks()
+	if _, ok := checks["rabbitmq"]; !ok {
+		t.Fatalf("missing check")
+	}
+	if token := agent.state.CheckToken("rabbitmq"); token != "abc123" {
+		t.Fatalf("bad: %s", token)
+	}
+}
+
 func TestAgent_unloadChecks(t *testing.T) {
 	config := nextConfig()
 	dir, agent := makeAgent(t, config)
@@ -920,6 +946,27 @@ func TestAgent_unloadChecks(t *testing.T) {
 	}
 }
 
+func TestAgent_loadServices_token(t *testing.T) {
+	config := nextConfig()
+	config.Services = append(config.Services, &ServiceDefinition{
+		ID:    "rabbitmq",
+		Name:  "rabbitmq",
+		Port:  5672,
+		Token: "abc123",
+	})
+	dir, agent := makeAgent(t, config)
+	defer os.RemoveAll(dir)
+	defer agent.Shutdown()
+
+	services := agent.state.Services()
+	if _, ok := services["rabbitmq"]; !ok {
+		t.Fatalf("missing service")
+	}
+	if token := agent.state.ServiceToken("rabbitmq"); token != "abc123" {
+		t.Fatalf("bad: %s", token)
+	}
+}
+
 func TestAgent_unloadServices(t *testing.T) {
 	config := nextConfig()
 	dir, agent := makeAgent(t, config)

command/agent/structs.go

@@ -46,6 +46,7 @@ type CheckDefinition struct {
 	Name      string
 	Notes     string
 	ServiceID string
+	Token     string
 	CheckType `mapstructure:",squash"`
 }