From dd3302cb9b3b3fcb0def35be4f05e146f17a46c5 Mon Sep 17 00:00:00 2001
From: Rebecca Mahany-Horton <rebeccamahany@gmail.com>
Date: Thu, 21 Dec 2023 11:49:12 -0500
Subject: [PATCH] Add flag for TUF autoupdater rollout

---
 ee/agent/flags/flag_controller.go |  9 +++++++++
 ee/agent/flags/keys/keys.go       |  1 +
 ee/agent/knapsack/knapsack.go     |  7 +++++++
 ee/agent/types/flags.go           |  4 ++++
 ee/agent/types/mocks/flags.go     | 28 ++++++++++++++++++++++++++++
 ee/agent/types/mocks/knapsack.go  | 28 ++++++++++++++++++++++++++++
 6 files changed, 77 insertions(+)

diff --git a/ee/agent/flags/flag_controller.go b/ee/agent/flags/flag_controller.go
index f39c099f2..86d094d00 100644
--- a/ee/agent/flags/flag_controller.go
+++ b/ee/agent/flags/flag_controller.go
@@ -500,6 +500,15 @@ func (fc *FlagController) UpdateDirectory() string {
 	).get(fc.getControlServerValue(keys.UpdateDirectory))
 }
 
+func (fc *FlagController) SetUseTUFAutoupdater(enabled bool) error {
+	return fc.setControlServerValue(keys.UseTUFAutoupdater, boolToBytes(enabled))
+}
+func (fc *FlagController) UseTUFAutoupdater() bool {
+	return NewBoolFlagValue(
+		WithDefaultBool(false),
+	).get(fc.getControlServerValue(keys.UseTUFAutoupdater))
+}
+
 func (fc *FlagController) SetExportTraces(enabled bool) error {
 	return fc.setControlServerValue(keys.ExportTraces, boolToBytes(enabled))
 }
diff --git a/ee/agent/flags/keys/keys.go b/ee/agent/flags/keys/keys.go
index 470c85185..c6aa5a91f 100644
--- a/ee/agent/flags/keys/keys.go
+++ b/ee/agent/flags/keys/keys.go
@@ -44,6 +44,7 @@ const (
 	MirrorServerURL                 FlagKey = "mirror_url"
 	AutoupdateInterval              FlagKey = "autoupdate_interval"
 	UpdateChannel                   FlagKey = "update_channel"
+	UseTUFAutoupdater               FlagKey = "use_tuf_autoupdater"
 	NotaryPrefix                    FlagKey = "notary_prefix"
 	AutoupdateInitialDelay          FlagKey = "autoupdater_initial_delay"
 	UpdateDirectory                 FlagKey = "update_directory"
diff --git a/ee/agent/knapsack/knapsack.go b/ee/agent/knapsack/knapsack.go
index 21a68bf8e..926b03a76 100644
--- a/ee/agent/knapsack/knapsack.go
+++ b/ee/agent/knapsack/knapsack.go
@@ -419,6 +419,13 @@ func (k *knapsack) UpdateDirectory() string {
 	return k.flags.UpdateDirectory()
 }
 
+func (k *knapsack) SetUseTUFAutoupdater(enabled bool) error {
+	return k.flags.SetUseTUFAutoupdater(enabled)
+}
+func (k *knapsack) UseTUFAutoupdater() bool {
+	return k.flags.UseTUFAutoupdater()
+}
+
 func (k *knapsack) SetExportTraces(enabled bool) error {
 	return k.flags.SetExportTraces(enabled)
 }
diff --git a/ee/agent/types/flags.go b/ee/agent/types/flags.go
index 78843f47f..483256b42 100644
--- a/ee/agent/types/flags.go
+++ b/ee/agent/types/flags.go
@@ -185,6 +185,10 @@ type Flags interface {
 	SetUpdateDirectory(directory string) error
 	UpdateDirectory() string
 
+	// UseTUFAutoupdater controls whether launcher uses the new TUF autoupdater instead of the legacy autoupdater
+	SetUseTUFAutoupdater(enabled bool) error
+	UseTUFAutoupdater() bool
+
 	// ExportTraces enables exporting our traces
 	SetExportTraces(enabled bool) error
 	SetExportTracesOverride(value bool, duration time.Duration)
diff --git a/ee/agent/types/mocks/flags.go b/ee/agent/types/mocks/flags.go
index 63a4bfad0..5ca514983 100644
--- a/ee/agent/types/mocks/flags.go
+++ b/ee/agent/types/mocks/flags.go
@@ -1160,6 +1160,20 @@ func (_m *Flags) SetUpdateDirectory(directory string) error {
 	return r0
 }
 
+// SetUseTUFAutoupdater provides a mock function with given fields: enabled
+func (_m *Flags) SetUseTUFAutoupdater(enabled bool) error {
+	ret := _m.Called(enabled)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(bool) error); ok {
+		r0 = rf(enabled)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
 // SetWatchdogDelaySec provides a mock function with given fields: sec
 func (_m *Flags) SetWatchdogDelaySec(sec int) error {
 	ret := _m.Called(sec)
@@ -1314,6 +1328,20 @@ func (_m *Flags) UpdateDirectory() string {
 	return r0
 }
 
+// UseTUFAutoupdater provides a mock function with given fields:
+func (_m *Flags) UseTUFAutoupdater() bool {
+	ret := _m.Called()
+
+	var r0 bool
+	if rf, ok := ret.Get(0).(func() bool); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	return r0
+}
+
 // WatchdogDelaySec provides a mock function with given fields:
 func (_m *Flags) WatchdogDelaySec() int {
 	ret := _m.Called()
diff --git a/ee/agent/types/mocks/knapsack.go b/ee/agent/types/mocks/knapsack.go
index 0772a47b7..46d8438c4 100644
--- a/ee/agent/types/mocks/knapsack.go
+++ b/ee/agent/types/mocks/knapsack.go
@@ -1352,6 +1352,20 @@ func (_m *Knapsack) SetUpdateDirectory(directory string) error {
 	return r0
 }
 
+// SetUseTUFAutoupdater provides a mock function with given fields: enabled
+func (_m *Knapsack) SetUseTUFAutoupdater(enabled bool) error {
+	ret := _m.Called(enabled)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(bool) error); ok {
+		r0 = rf(enabled)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
 // SetWatchdogDelaySec provides a mock function with given fields: sec
 func (_m *Knapsack) SetWatchdogDelaySec(sec int) error {
 	ret := _m.Called(sec)
@@ -1570,6 +1584,20 @@ func (_m *Knapsack) UpdateDirectory() string {
 	return r0
 }
 
+// UseTUFAutoupdater provides a mock function with given fields:
+func (_m *Knapsack) UseTUFAutoupdater() bool {
+	ret := _m.Called()
+
+	var r0 bool
+	if rf, ok := ret.Get(0).(func() bool); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	return r0
+}
+
 // WatchdogDelaySec provides a mock function with given fields:
 func (_m *Knapsack) WatchdogDelaySec() int {
 	ret := _m.Called()