[v14] Fix aws oidc deployment metrics

lib/automaticupgrades/config.go

@@ -27,6 +27,13 @@ import (
 )
 
 const (
+	// EnvUpgrader environment variable specifies the external upgrader type
+	// Ex: unit, kube
+	EnvUpgrader = "TELEPORT_EXT_UPGRADER"
+	// EnvUpgraderVersion environment variable specifies the external upgrader version
+	// Ex: v14.3.6
+	EnvUpgraderVersion = "TELEPORT_EXT_UPGRADER_VERSION"
+
 	// automaticUpgradesEnvar defines the env var to lookup when deciding whether to enable AutomaticUpgrades feature.
 	automaticUpgradesEnvar = "TELEPORT_AUTOMATIC_UPGRADES"
 
@@ -65,7 +72,7 @@ func GetChannel() string {
 
 // GetUpgraderVersion returns the teleport upgrader version
 func GetUpgraderVersion(ctx context.Context) string {
-	if os.Getenv("TELEPORT_EXT_UPGRADER") == "unit" {
+	if os.Getenv(EnvUpgrader) == "unit" {
 		out, err := exec.CommandContext(ctx, teleportUpgradeScript, "version").Output()
 		if err != nil {
 			log.WithError(err).Debug("Failed to exec /usr/sbin/teleport-upgrade version command.")
@@ -75,5 +82,5 @@ func GetUpgraderVersion(ctx context.Context) string {
 			}
 		}
 	}
-	return os.Getenv("TELEPORT_EXT_UPGRADER_VERSION")
+	return os.Getenv(EnvUpgraderVersion)
 }

lib/integrations/awsoidc/deployservice.go

@@ -477,7 +477,7 @@ type upsertTaskRequest struct {
 func upsertTask(ctx context.Context, clt DeployServiceClient, req upsertTaskRequest) (*ecsTypes.TaskDefinition, error) {
 	taskAgentContainerImage := getDistrolessTeleportImage(req.TeleportVersionTag)
 
-	taskDefOut, err := clt.RegisterTaskDefinition(ctx, &ecs.RegisterTaskDefinitionInput{
+	taskDefIn := &ecs.RegisterTaskDefinitionInput{
 		Family: aws.String(req.TaskName),
 		RequiresCompatibilities: []ecsTypes.Compatibility{
 			ecsTypes.CompatibilityFargate,
@@ -489,10 +489,12 @@ func upsertTask(ctx context.Context, clt DeployServiceClient, req upsertTaskRequ
 		TaskRoleArn:      &req.TaskRoleARN,
 		ExecutionRoleArn: &req.TaskRoleARN,
 		ContainerDefinitions: []ecsTypes.ContainerDefinition{{
-			Environment: []ecsTypes.KeyValuePair{{
-				Name:  aws.String(types.InstallMethodAWSOIDCDeployServiceEnvVar),
-				Value: aws.String("true"),
-			}},
+			Environment: []ecsTypes.KeyValuePair{
+				{
+					Name:  aws.String(types.InstallMethodAWSOIDCDeployServiceEnvVar),
+					Value: aws.String("true"),
+				},
+			},
 			Command: []string{
 				// --rewrite 15:3 rewrites SIGTERM -> SIGQUIT. This enables graceful shutdown of teleport
 				"--rewrite",
@@ -517,7 +519,15 @@ func upsertTask(ctx context.Context, clt DeployServiceClient, req upsertTaskRequ
 			},
 		}},
 		Tags: req.ResourceCreationTags.ToECSTags(),
-	})
+	}
+
+	// Ensure that the upgrader environment variables are set.
+	// These will ensure that the instance reports Teleport upgrader metrics.
+	if err := ensureUpgraderEnvironmentVariables(taskDefIn); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	taskDefOut, err := clt.RegisterTaskDefinition(ctx, taskDefIn)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}

lib/integrations/awsoidc/deployservice_test.go

@@ -17,14 +17,19 @@ limitations under the License.
 package awsoidc
 
 import (
+	"context"
 	"regexp"
 	"testing"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	ecstypes "github.com/aws/aws-sdk-go-v2/service/ecs/types"
 	"github.com/google/go-cmp/cmp"
 	"github.com/gravitational/trace"
 	"github.com/stretchr/testify/require"
 
 	"github.com/gravitational/teleport"
+	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/lib/automaticupgrades"
 )
 
 func TestDeployServiceRequest(t *testing.T) {
@@ -207,3 +212,30 @@ func TestNormalizeECSResourceName(t *testing.T) {
 		})
 	}
 }
+
+func TestUpsertTask(t *testing.T) {
+	ctx := context.Background()
+
+	mockClient := &mockDeployServiceClient{
+		clusters:        map[string]*ecstypes.Cluster{},
+		taskDefinitions: map[string]*ecstypes.TaskDefinition{},
+		services:        map[string]*ecstypes.Service{},
+		accountId:       aws.String("123456789012"),
+		iamTokenMissing: true,
+	}
+
+	expected := []ecstypes.KeyValuePair{
+		{
+			Name:  aws.String(types.InstallMethodAWSOIDCDeployServiceEnvVar),
+			Value: aws.String("true"),
+		},
+		{
+			Name:  aws.String(automaticupgrades.EnvUpgraderVersion),
+			Value: aws.String(teleport.Version),
+		},
+	}
+
+	taskDefinition, err := upsertTask(ctx, mockClient, upsertTaskRequest{})
+	require.NoError(t, err)
+	require.Equal(t, expected, taskDefinition.ContainerDefinitions[0].Environment)
+}

lib/integrations/awsoidc/deployservice_update.go

@@ -27,6 +27,8 @@ import (
 	"github.com/gravitational/trace"
 	"github.com/sirupsen/logrus"
 
+	"github.com/gravitational/teleport"
+	"github.com/gravitational/teleport/lib/automaticupgrades"
 	awslib "github.com/gravitational/teleport/lib/cloud/aws"
 )
 
@@ -108,6 +110,12 @@ func updateServiceContainerImage(ctx context.Context, clt DeployServiceClient, l
 		return trace.Wrap(err)
 	}
 
+	// Ensure that the upgrader variables are set.
+	// These will ensure that the instance reports Teleport upgrader metrics.
+	if err := ensureUpgraderEnvironmentVariables(registerTaskDefinitionIn); err != nil {
+		return trace.Wrap(err)
+	}
+
 	registerTaskDefinitionOut, err := clt.RegisterTaskDefinition(ctx, registerTaskDefinitionIn)
 	if err != nil {
 		return trace.Wrap(err)
@@ -279,6 +287,36 @@ func generateTaskDefinitionWithImage(taskDefinition *ecsTypes.TaskDefinition, te
 	return registerTaskDefinitionIn, nil
 }
 
+// ensureUpgraderEnvironmentVariables modifies the taskDefinition and ensures that
+// the upgrader specific environment variables are set.
+func ensureUpgraderEnvironmentVariables(taskDefinition *ecs.RegisterTaskDefinitionInput) error {
+	containerDefinitions := []ecsTypes.ContainerDefinition{}
+	for _, containerDefinition := range taskDefinition.ContainerDefinitions {
+		environment := []ecsTypes.KeyValuePair{}
+
+		// Copy non-upgrader specific environemt variables as is
+		for _, env := range containerDefinition.Environment {
+			if aws.ToString(env.Name) == automaticupgrades.EnvUpgrader ||
+				aws.ToString(env.Name) == automaticupgrades.EnvUpgraderVersion {
+				continue
+			}
+			environment = append(environment, env)
+		}
+
+		// Ensure ugprader specific environment variables are set
+		environment = append(environment,
+			ecsTypes.KeyValuePair{
+				Name:  aws.String(automaticupgrades.EnvUpgraderVersion),
+				Value: aws.String(teleport.Version),
+			},
+		)
+		containerDefinition.Environment = environment
+		containerDefinitions = append(containerDefinitions, containerDefinition)
+	}
+	taskDefinition.ContainerDefinitions = containerDefinitions
+	return nil
+}
+
 // generateServiceWithTaskDefinition returns new update service input with the desired task definition
 func generateServiceWithTaskDefinition(service *ecsTypes.Service, taskDefinitionName string) *ecs.UpdateServiceInput {
 	updateServiceIn := new(ecs.UpdateServiceInput)

lib/integrations/awsoidc/deployservice_update_test.go

@@ -28,6 +28,7 @@ import (
 
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/lib/automaticupgrades"
 )
 
 func TestGenerateServiceWithTaskDefinition(t *testing.T) {
@@ -374,3 +375,89 @@ func TestUpdateDeployServices(t *testing.T) {
 		require.Empty(t, m.taskDefinitions)
 	})
 }
+
+func TestEnsureUpgraderEnvironmentVariables(t *testing.T) {
+	tts := []struct {
+		desc     string
+		env      []ecsTypes.KeyValuePair
+		expected []ecsTypes.KeyValuePair
+	}{
+		{
+			desc: "expected values are unmodified",
+			env: []ecsTypes.KeyValuePair{
+				{
+					Name:  aws.String("EXAMPLE_ENV"),
+					Value: aws.String("EXAMPLE"),
+				},
+				{
+					Name:  aws.String(automaticupgrades.EnvUpgraderVersion),
+					Value: aws.String(teleport.Version),
+				},
+			},
+			expected: []ecsTypes.KeyValuePair{
+				{
+					Name:  aws.String("EXAMPLE_ENV"),
+					Value: aws.String("EXAMPLE"),
+				},
+				{
+					Name:  aws.String(automaticupgrades.EnvUpgraderVersion),
+					Value: aws.String(teleport.Version),
+				},
+			},
+		},
+		{
+			desc: "Teleport upgrader env variables are added",
+			env: []ecsTypes.KeyValuePair{
+				{
+					Name:  aws.String("EXAMPLE_ENV"),
+					Value: aws.String("EXAMPLE"),
+				},
+			},
+			expected: []ecsTypes.KeyValuePair{
+				{
+					Name:  aws.String("EXAMPLE_ENV"),
+					Value: aws.String("EXAMPLE"),
+				},
+				{
+					Name:  aws.String(automaticupgrades.EnvUpgraderVersion),
+					Value: aws.String(teleport.Version),
+				},
+			},
+		},
+		{
+			desc: "Teleport upgrader env variables are updated",
+			env: []ecsTypes.KeyValuePair{
+				{
+					Name:  aws.String("EXAMPLE_ENV"),
+					Value: aws.String("EXAMPLE"),
+				},
+				{
+					Name:  aws.String(automaticupgrades.EnvUpgraderVersion),
+					Value: aws.String("v1.0.0"),
+				},
+			},
+			expected: []ecsTypes.KeyValuePair{
+				{
+					Name:  aws.String("EXAMPLE_ENV"),
+					Value: aws.String("EXAMPLE"),
+				},
+				{
+					Name:  aws.String(automaticupgrades.EnvUpgraderVersion),
+					Value: aws.String(teleport.Version),
+				},
+			},
+		},
+	}
+
+	for _, tt := range tts {
+		t.Run(tt.desc, func(t *testing.T) {
+			taskDefinition := &ecs.RegisterTaskDefinitionInput{
+				ContainerDefinitions: []ecsTypes.ContainerDefinition{
+					{Environment: tt.env},
+				},
+			}
+			require.NoError(t, ensureUpgraderEnvironmentVariables(taskDefinition))
+			require.Equal(t, tt.expected, taskDefinition.ContainerDefinitions[0].Environment)
+		})
+	}
+}

lib/service/service.go

@@ -994,17 +994,24 @@ func NewTeleport(cfg *servicecfg.Config) (*TeleportProcess, error) {
 		return nil, trace.Wrap(err)
 	}
 
-	upgraderKind := os.Getenv("TELEPORT_EXT_UPGRADER")
+	upgraderKind := os.Getenv(automaticupgrades.EnvUpgrader)
 	upgraderVersion := automaticupgrades.GetUpgraderVersion(process.GracefulExitContext())
 
+	// Instances deployed using the AWS OIDC integration are automatically updated
+	// by the proxy. The instance heartbeat should properly reflect that.
+	externalUpgrader := upgraderKind
+	if externalUpgrader == "" && os.Getenv(types.InstallMethodAWSOIDCDeployServiceEnvVar) == "true" {
+		externalUpgrader = types.OriginIntegrationAWSOIDC
+	}
+
 	// note: we must create the inventory handle *after* registerExpectedServices because that function determines
 	// the list of services (instance roles) to be included in the heartbeat.
 	process.inventoryHandle = inventory.NewDownstreamHandle(process.makeInventoryControlStreamWhenReady, proto.UpstreamInventoryHello{
 		ServerID:                cfg.HostUUID,
 		Version:                 teleport.Version,
 		Services:                process.getInstanceRoles(),
 		Hostname:                cfg.Hostname,
-		ExternalUpgrader:        upgraderKind,
+		ExternalUpgrader:        externalUpgrader,
 		ExternalUpgraderVersion: vc.Normalize(upgraderVersion),
 	})