Add bootstrap support for AutoUpdateConfig/AutoUpdateVersion resources (#51904) (#53049)

* Add bootstrap support for AutoUpdateConfig/AutoUpdateVersion

* Add support for `applyResources`
Reduce nested switch for bootstrap new type of resources

* Transform test to table test to check same resources for both `bootstrapResources` and `applyResources`

Author: vapopov

lib/auth/autoupdate/autoupdatev1/service.go

@@ -55,6 +55,11 @@ type ServiceConfig struct {
 	Emitter apievents.Emitter
 }
 
+// Backend interface for manipulating AutoUpdate resources.
+type Backend interface {
+	services.AutoUpdateService
+}
+
 // Service implements the gRPC API layer for the AutoUpdate.
 type Service struct {
 	autoupdate.UnimplementedAutoUpdateServiceServer
@@ -212,6 +217,18 @@ func (s *Service) UpsertAutoUpdateConfig(ctx context.Context, req *autoupdate.Up
 	return config, trace.Wrap(err)
 }
 
+// UpsertAutoUpdateConfig creates a new AutoUpdateConfig or forcefully updates an existing AutoUpdateConfig.
+// This is a function rather than a method so that it can be used by the gRPC service
+// and the auth server init code when dealing with resources to be applied at startup.
+func UpsertAutoUpdateConfig(
+	ctx context.Context,
+	backend Backend,
+	config *autoupdate.AutoUpdateConfig,
+) (*autoupdate.AutoUpdateConfig, error) {
+	out, err := backend.UpsertAutoUpdateConfig(ctx, config)
+	return out, trace.Wrap(err)
+}
+
 // DeleteAutoUpdateConfig deletes AutoUpdateConfig singleton.
 func (s *Service) DeleteAutoUpdateConfig(ctx context.Context, req *autoupdate.DeleteAutoUpdateConfigRequest) (*emptypb.Empty, error) {
 	authCtx, err := s.authorizer.Authorize(ctx)
@@ -390,6 +407,18 @@ func (s *Service) UpsertAutoUpdateVersion(ctx context.Context, req *autoupdate.U
 	return autoUpdateVersion, trace.Wrap(err)
 }
 
+// UpsertAutoUpdateVersion creates a new AutoUpdateVersion or forcefully updates an existing AutoUpdateVersion.
+// This is a function rather than a method so that it can be used by the gRPC service
+// and the auth server init code when dealing with resources to be applied at startup.
+func UpsertAutoUpdateVersion(
+	ctx context.Context,
+	backend Backend,
+	version *autoupdate.AutoUpdateVersion,
+) (*autoupdate.AutoUpdateVersion, error) {
+	out, err := backend.UpsertAutoUpdateVersion(ctx, version)
+	return out, trace.Wrap(err)
+}
+
 // DeleteAutoUpdateVersion deletes AutoUpdateVersion singleton.
 func (s *Service) DeleteAutoUpdateVersion(ctx context.Context, req *autoupdate.DeleteAutoUpdateVersionRequest) (*emptypb.Empty, error) {
 	authCtx, err := s.authorizer.Authorize(ctx)

lib/auth/init.go

@@ -38,11 +38,13 @@ import (
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/api/client/proto"
 	apidefaults "github.com/gravitational/teleport/api/defaults"
+	autoupdatev1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/autoupdate/v1"
 	"github.com/gravitational/teleport/api/types"
 	apievents "github.com/gravitational/teleport/api/types/events"
 	"github.com/gravitational/teleport/lib"
 	"github.com/gravitational/teleport/lib/ai"
 	"github.com/gravitational/teleport/lib/ai/embedding"
+	"github.com/gravitational/teleport/lib/auth/autoupdate/autoupdatev1"
 	"github.com/gravitational/teleport/lib/auth/keystore"
 	"github.com/gravitational/teleport/lib/auth/migration"
 	"github.com/gravitational/teleport/lib/auth/native"
@@ -1133,13 +1135,23 @@ func migrateRemoteClusters(ctx context.Context, asrv *Server) error {
 func applyResources(ctx context.Context, service *Services, resources []types.Resource) error {
 	var err error
 	for _, resource := range resources {
-		switch r := resource.(type) {
+		// Unwrap "new style" resources.
+		// We always want to switch over the underlying type.
+		var res any = resource
+		if w, ok := res.(interface{ Unwrap() types.Resource153 }); ok {
+			res = w.Unwrap()
+		}
+		switch r := res.(type) {
 		case types.ProvisionToken:
 			err = service.Provisioner.UpsertToken(ctx, r)
 		case types.ClusterNetworkingConfig:
 			err = service.ClusterConfiguration.SetClusterNetworkingConfig(ctx, r)
 		case types.AuthPreference:
 			err = service.ClusterConfiguration.SetAuthPreference(ctx, r)
+		case *autoupdatev1pb.AutoUpdateConfig:
+			_, err = autoupdatev1.UpsertAutoUpdateConfig(ctx, service, r)
+		case *autoupdatev1pb.AutoUpdateVersion:
+			_, err = autoupdatev1.UpsertAutoUpdateVersion(ctx, service, r)
 		default:
 			return trace.NotImplemented("cannot apply resource of type %T", resource)
 		}

lib/auth/init_test.go

@@ -33,6 +33,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/gravitational/trace"
 	"github.com/jonboulle/clockwork"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/crypto/ssh"
@@ -1766,3 +1767,54 @@ func TestTeleportProcessAuthVersionUpgradeCheck(t *testing.T) {
 		})
 	}
 }
+
+// TestInitWithAutoUpdateResources verifies that auth init support bootstrapping and apply
+// `AutoUpdateConfig` and `AutoUpdateVersion` resources as well as unmarshalling them from
+// yaml configuration.
+func TestInitWithAutoUpdateResources(t *testing.T) {
+	t.Parallel()
+
+	const autoUpdateConfigYAML = `kind: autoupdate_config
+metadata:
+  name: autoupdate-config
+spec:
+  tools:
+    mode: enabled
+version: v1`
+	const autoUpdateVersionYAML = `kind: autoupdate_version
+metadata:
+  name: autoupdate-version
+spec:
+  tools:
+    target_version: 1.2.3
+version: v1`
+
+	ctx := context.Background()
+	resources := []types.Resource{
+		resourceFromYAML(t, autoUpdateConfigYAML),
+		resourceFromYAML(t, autoUpdateVersionYAML),
+	}
+
+	for _, test := range []struct {
+		name string
+		fn   func(cfg *InitConfig)
+	}{
+		{name: "bootstrap", fn: func(cfg *InitConfig) { cfg.BootstrapResources = resources }},
+		{name: "apply", fn: func(cfg *InitConfig) { cfg.ApplyOnStartupResources = resources }},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			cfg := setupConfig(t)
+			test.fn(&cfg)
+			auth, err := Init(ctx, cfg)
+			require.NoError(t, err)
+
+			config, err := auth.GetAutoUpdateConfig(ctx)
+			assert.NoError(t, err)
+			assert.Equal(t, "enabled", config.GetSpec().GetTools().GetMode())
+
+			version, err := auth.GetAutoUpdateVersion(ctx)
+			assert.NoError(t, err)
+			assert.Equal(t, "1.2.3", version.GetSpec().GetTools().GetTargetVersion())
+		})
+	}
+}

lib/services/local/autoupdate.go

@@ -156,3 +156,55 @@ func (s *AutoUpdateService) GetAutoUpdateVersion(ctx context.Context) (*autoupda
 func (s *AutoUpdateService) DeleteAutoUpdateVersion(ctx context.Context) error {
 	return trace.Wrap(s.version.DeleteResource(ctx, types.MetaNameAutoUpdateVersion))
 }
+
+// itemFromAutoUpdateConfig generates `backend.Item` from `AutoUpdateConfig` resource type.
+func itemFromAutoUpdateConfig(config *autoupdate.AutoUpdateConfig) (*backend.Item, error) {
+	if err := update.ValidateAutoUpdateConfig(config); err != nil {
+		return nil, trace.Wrap(err)
+	}
+	rev, err := types.GetRevision(config)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	value, err := services.MarshalProtoResource[*autoupdate.AutoUpdateConfig](config)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	expires, err := types.GetExpiry(config)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	item := &backend.Item{
+		Key:      backend.NewKey(autoUpdateConfigPrefix, types.MetaNameAutoUpdateConfig),
+		Value:    value,
+		Expires:  expires,
+		Revision: rev,
+	}
+	return item, nil
+}
+
+// itemFromAutoUpdateVersion generates `backend.Item` from `AutoUpdateVersion` resource type.
+func itemFromAutoUpdateVersion(version *autoupdate.AutoUpdateVersion) (*backend.Item, error) {
+	if err := update.ValidateAutoUpdateVersion(version); err != nil {
+		return nil, trace.Wrap(err)
+	}
+	rev, err := types.GetRevision(version)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	value, err := services.MarshalProtoResource[*autoupdate.AutoUpdateVersion](version)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	expires, err := types.GetExpiry(version)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	item := &backend.Item{
+		Key:      backend.NewKey(autoUpdateVersionPrefix, types.MetaNameAutoUpdateVersion),
+		Value:    value,
+		Expires:  expires,
+		Revision: rev,
+	}
+	return item, nil
+}

lib/services/local/resource.go

@@ -24,6 +24,7 @@ import (
 
 	"github.com/gravitational/trace"
 
+	autoupdatev1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/autoupdate/v1"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/services"
@@ -82,7 +83,15 @@ func itemsFromResource(resource types.Resource) ([]backend.Item, error) {
 	var item *backend.Item
 	var extItems []backend.Item
 	var err error
-	switch r := resource.(type) {
+
+	// Unwrap "new style" resources.
+	// We always want to switch over the underlying type.
+	var res any = resource
+	if w, ok := res.(types.Resource153Unwrapper); ok {
+		res = w.Unwrap()
+	}
+
+	switch r := res.(type) {
 	case types.User:
 		item, err = itemFromUser(r)
 		if auth := r.GetLocalAuth(); err == nil && auth != nil {
@@ -108,6 +117,10 @@ func itemsFromResource(resource types.Resource) ([]backend.Item, error) {
 		item, err = itemFromClusterNetworkingConfig(r)
 	case types.AuthPreference:
 		item, err = itemFromAuthPreference(r)
+	case *autoupdatev1pb.AutoUpdateConfig:
+		item, err = itemFromAutoUpdateConfig(r)
+	case *autoupdatev1pb.AutoUpdateVersion:
+		item, err = itemFromAutoUpdateVersion(r)
 	default:
 		return nil, trace.NotImplemented("cannot itemFrom resource of type %T", resource)
 	}

lib/services/resource.go

@@ -28,6 +28,7 @@ import (
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
+	autoupdatev1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/autoupdate/v1"
 	headerv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/header/v1"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/lib/utils"
@@ -644,6 +645,20 @@ func init() {
 		}
 		return ap, nil
 	})
+	RegisterResourceUnmarshaler(types.KindAutoUpdateConfig, func(bytes []byte, option ...MarshalOption) (types.Resource, error) {
+		c := &autoupdatev1pb.AutoUpdateConfig{}
+		if err := protojson.Unmarshal(bytes, c); err != nil {
+			return nil, trace.Wrap(err)
+		}
+		return types.Resource153ToLegacy(c), nil
+	})
+	RegisterResourceUnmarshaler(types.KindAutoUpdateVersion, func(bytes []byte, option ...MarshalOption) (types.Resource, error) {
+		v := &autoupdatev1pb.AutoUpdateVersion{}
+		if err := protojson.Unmarshal(bytes, v); err != nil {
+			return nil, trace.Wrap(err)
+		}
+		return types.Resource153ToLegacy(v), nil
+	})
 }
 
 // CheckAndSetDefaults calls [r.CheckAndSetDefaults] if r implements the method.