From d547e0d459984450ef7be0d31d7ec753c834b74d Mon Sep 17 00:00:00 2001
From: Jeff Mitchell <jeffrey.mitchell@gmail.com>
Date: Mon, 17 Jun 2024 12:26:55 -0400
Subject: [PATCH] Add test ability to reinitialize worker auth roots on startup
 (#4895)

---
 internal/cmd/commands/server/server.go   | 7 +++++++
 internal/daemon/controller/config.go     | 2 ++
 internal/daemon/controller/controller.go | 2 +-
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/internal/cmd/commands/server/server.go b/internal/cmd/commands/server/server.go
index b54fe6e9b..b41684e9c 100644
--- a/internal/cmd/commands/server/server.go
+++ b/internal/cmd/commands/server/server.go
@@ -69,6 +69,7 @@ type Command struct {
 	presetConfig                         *atomic.String // for tests
 	flagWorkerAuthWorkerRotationInterval time.Duration  // for tests
 	flagWorkerAuthCaCertificateLifetime  time.Duration  // for tests
+	flagWorkerAuthCaReinitialize         bool           // for tests
 }
 
 func (c *Command) Synopsis() string {
@@ -140,6 +141,11 @@ func (c *Command) Flags() *base.FlagSets {
 		Target: &c.flagWorkerAuthCaCertificateLifetime,
 		Hidden: true,
 	})
+	f.BoolVar(&base.BoolVar{
+		Name:   "worker-auth-ca-reinitialize",
+		Target: &c.flagWorkerAuthCaReinitialize,
+		Hidden: true,
+	})
 
 	f.BoolVar(&base.BoolVar{
 		Name:   "skip-plugins",
@@ -629,6 +635,7 @@ func (c *Command) StartController(ctx context.Context) error {
 		RawConfig: c.Config,
 		Server:    c.Server,
 		TestOverrideWorkerAuthCaCertificateLifetime: c.flagWorkerAuthCaCertificateLifetime,
+		TestWorkerAuthCaReinitialize:                c.flagWorkerAuthCaReinitialize,
 	}
 
 	var err error
diff --git a/internal/daemon/controller/config.go b/internal/daemon/controller/config.go
index 337af919d..20a77fb3d 100644
--- a/internal/daemon/controller/config.go
+++ b/internal/daemon/controller/config.go
@@ -21,6 +21,8 @@ type Config struct {
 	DisableAuthorizationFailures bool
 	// Override worker auth CA certificate lifetime for testing
 	TestOverrideWorkerAuthCaCertificateLifetime time.Duration
+	// Reinitialize the roots at startup
+	TestWorkerAuthCaReinitialize bool
 
 	// This is derived from the config.Config. It tracks the state of the
 	// rate limiter's configuration, and is updated if the config changes via a
diff --git a/internal/daemon/controller/controller.go b/internal/daemon/controller/controller.go
index 4d3d30aea..c15e9f260 100644
--- a/internal/daemon/controller/controller.go
+++ b/internal/daemon/controller/controller.go
@@ -475,7 +475,7 @@ func New(ctx context.Context, conf *Config) (*Controller, error) {
 	if err != nil {
 		return nil, fmt.Errorf("unable to instantiate worker auth repository: %w", err)
 	}
-	_, err = server.RotateRoots(ctx, serversRepo, nodeenrollment.WithCertificateLifetime(conf.TestOverrideWorkerAuthCaCertificateLifetime))
+	_, err = server.RotateRoots(ctx, serversRepo, nodeenrollment.WithCertificateLifetime(conf.TestOverrideWorkerAuthCaCertificateLifetime), nodeenrollment.WithReinitializeRoots(conf.TestWorkerAuthCaReinitialize))
 	if err != nil {
 		event.WriteSysEvent(ctx, op, "unable to ensure worker auth roots exist, may be due to multiple controllers starting at once, continuing")
 	}