use pre-release cdn and image repo for aws integration flows

lib/integrations/awsoidc/deployservice.go

@@ -30,13 +30,14 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/ecs"
 	ecsTypes "github.com/aws/aws-sdk-go-v2/service/ecs/types"
 	"github.com/aws/aws-sdk-go-v2/service/sts"
+	"github.com/coreos/go-semver/semver"
 	"github.com/gravitational/trace"
 
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/api/utils/retryutils"
 	"github.com/gravitational/teleport/lib/integrations/awsoidc/tags"
-	"github.com/gravitational/teleport/lib/modules"
+	"github.com/gravitational/teleport/lib/utils/teleportassets"
 )
 
 var (
@@ -56,11 +57,6 @@ var (
 )
 
 const (
-	// distrolessTeleportOSS is the distroless image of the OSS version of Teleport
-	distrolessTeleportOSS = "public.ecr.aws/gravitational/teleport-distroless"
-	// distrolessTeleportEnt is the distroless image of the Enterprise version of Teleport
-	distrolessTeleportEnt = "public.ecr.aws/gravitational/teleport-ent-distroless"
-
 	// clusterStatusActive is the string representing an ACTIVE ECS Cluster.
 	clusterStatusActive = "ACTIVE"
 	// clusterStatusInactive is the string representing an INACTIVE ECS Cluster.
@@ -479,7 +475,10 @@ type upsertTaskRequest struct {
 
 // upsertTask ensures a TaskDefinition with TaskName exists
 func upsertTask(ctx context.Context, clt DeployServiceClient, req upsertTaskRequest) (*ecsTypes.TaskDefinition, error) {
-	taskAgentContainerImage := getDistrolessTeleportImage(req.TeleportVersionTag)
+	taskAgentContainerImage, err := getDistrolessTeleportImage(req.TeleportVersionTag)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
 
 	taskDefIn := &ecs.RegisterTaskDefinitionInput{
 		Family: aws.String(req.TaskName),
@@ -767,10 +766,11 @@ func upsertService(ctx context.Context, clt DeployServiceClient, req upsertServi
 }
 
 // getDistrolessTeleportImage returns the distroless teleport image string
-func getDistrolessTeleportImage(version string) string {
-	teleportImage := distrolessTeleportOSS
-	if modules.GetModules().BuildType() == modules.BuildEnterprise {
-		teleportImage = distrolessTeleportEnt
+func getDistrolessTeleportImage(version string) (string, error) {
+	semVer, err := semver.NewVersion(version)
+	if err != nil {
+		return "", trace.BadParameter("invalid version tag %s", version)
 	}
-	return fmt.Sprintf("%s:%s", teleportImage, version)
+
+	return teleportassets.DistrolessImage(*semVer), nil
 }

lib/integrations/awsoidc/deployservice_test.go

@@ -238,7 +238,9 @@ func TestUpsertTask(t *testing.T) {
 		},
 	}
 
-	taskDefinition, err := upsertTask(ctx, mockClient, upsertTaskRequest{})
+	semVer := *teleport.SemVersion
+	semVer.PreRelease = ""
+	taskDefinition, err := upsertTask(ctx, mockClient, upsertTaskRequest{TeleportVersionTag: semVer.String()})
 	require.NoError(t, err)
 	require.Equal(t, expected, taskDefinition.ContainerDefinitions[0].Environment)
 }

lib/integrations/awsoidc/deployservice_update.go

@@ -71,7 +71,10 @@ func UpdateDeployService(ctx context.Context, clt DeployServiceClient, log *slog
 		return trace.Wrap(err)
 	}
 
-	teleportImage := getDistrolessTeleportImage(req.TeleportVersionTag)
+	teleportImage, err := getDistrolessTeleportImage(req.TeleportVersionTag)
+	if err != nil {
+		return trace.Wrap(err)
+	}
 	services, err := getManagedServices(ctx, clt, log, req.TeleportClusterName, req.OwnershipTags)
 	if err != nil {
 		return trace.Wrap(err)

lib/integrations/awsoidc/deployservice_update_test.go

@@ -162,7 +162,9 @@ func TestUpdateDeployServices(t *testing.T) {
 	clusterName := "my-cluster"
 	integrationName := "my-integration"
 	ownershipTags := tags.DefaultResourceCreationTags(clusterName, integrationName)
-	teleportVersion := teleport.Version
+	semVer := *teleport.SemVersion
+	semVer.PreRelease = ""
+	teleportVersion := semVer.String()
 	log := utils.NewSlogLoggerForTests().With("test", t.Name())
 
 	t.Run("only legacy service present", func(t *testing.T) {

lib/utils/teleportassets/teleportassets.go

@@ -0,0 +1,84 @@
+/*
+ * Teleport
+ * Copyright (C) 2024  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package teleportassets
+
+import (
+	"fmt"
+
+	"github.com/coreos/go-semver/semver"
+
+	"github.com/gravitational/teleport"
+	"github.com/gravitational/teleport/lib/modules"
+)
+
+const (
+	// teleportReleaseCDN is the Teleport CDN URL for release builds.
+	// This can be used to download the Teleport binary for release builds.
+	teleportReleaseCDN = "https://cdn.teleport.dev"
+	// teleportPreReleaseCDN is the Teleport CDN URL for pre-release builds.
+	// This can be used to download the Teleport binary for pre-release builds.
+	teleportPreReleaseCDN = "https://cdn.cloud.gravitational.io"
+)
+
+// CDNBaseURL returns the URL of the CDN that can be used to download Teleport
+// binary assets.
+func CDNBaseURL() string {
+	return cdnBaseURL(*teleport.SemVersion)
+}
+
+// cdnBaseURL returns the base URL of the CDN that can be used to download
+// Teleport binary assets.
+func cdnBaseURL(version semver.Version) string {
+	if version.PreRelease != "" {
+		return teleportPreReleaseCDN
+	}
+	return teleportReleaseCDN
+}
+
+const (
+	// teleportReleaseECR is the official release repo for Teleport images.
+	teleportReleaseECR = "public.ecr.aws/gravitational"
+	// teleportReleaseECR is the pre-release repo for Teleport images.
+	teleportPreReleaseECR = "public.ecr.aws/gravitational-staging"
+	// distrolessTeleportOSSImage is the distroless image of the OSS version of Teleport
+	distrolessTeleportOSSImage = "teleport-distroless"
+	// distrolessTeleportEntImage is the distroless image of the Enterprise version of Teleport
+	distrolessTeleportEntImage = "teleport-ent-distroless"
+)
+
+// DistrolessImage returns the distroless teleport image repo.
+func DistrolessImage(version semver.Version) string {
+	repo := distrolessImageRepo(version)
+	name := distrolessImageName(modules.GetModules().BuildType())
+	return fmt.Sprintf("%s/%s:%s", repo, name, version)
+}
+
+func distrolessImageRepo(version semver.Version) string {
+	if version.PreRelease != "" {
+		return teleportPreReleaseECR
+	}
+	return teleportReleaseECR
+}
+
+func distrolessImageName(buildType string) string {
+	if buildType == modules.BuildEnterprise {
+		return distrolessTeleportEntImage
+	}
+	return distrolessTeleportOSSImage
+}

lib/utils/teleportassets/teleportassets_test.go

@@ -0,0 +1,82 @@
+/*
+ * Teleport
+ * Copyright (C) 2024  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package teleportassets
+
+import (
+	"testing"
+
+	"github.com/coreos/go-semver/semver"
+	"github.com/stretchr/testify/require"
+
+	"github.com/gravitational/teleport/lib/modules"
+)
+
+func TestDistrolessTeleportImageRepo(t *testing.T) {
+	tests := []struct {
+		desc      string
+		buildType string
+		version   string
+		want      string
+	}{
+		{
+			desc:      "ent release",
+			buildType: modules.BuildEnterprise,
+			version:   "16.0.0",
+			want:      "public.ecr.aws/gravitational/teleport-ent-distroless:16.0.0",
+		},
+		{
+			desc:      "oss release",
+			buildType: modules.BuildOSS,
+			version:   "16.0.0",
+			want:      "public.ecr.aws/gravitational/teleport-distroless:16.0.0",
+		},
+		{
+			desc:      "community release",
+			buildType: modules.BuildCommunity,
+			version:   "16.0.0",
+			want:      "public.ecr.aws/gravitational/teleport-distroless:16.0.0",
+		},
+		{
+			desc:      "ent pre-release",
+			buildType: modules.BuildEnterprise,
+			version:   "16.0.0-alpha.1",
+			want:      "public.ecr.aws/gravitational-staging/teleport-ent-distroless:16.0.0-alpha.1",
+		},
+		{
+			desc:      "oss pre-release",
+			buildType: modules.BuildOSS,
+			version:   "16.0.0-alpha.1",
+			want:      "public.ecr.aws/gravitational-staging/teleport-distroless:16.0.0-alpha.1",
+		},
+		{
+			desc:      "community pre-release",
+			buildType: modules.BuildCommunity,
+			version:   "16.0.0-alpha.1",
+			want:      "public.ecr.aws/gravitational-staging/teleport-distroless:16.0.0-alpha.1",
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			semVer, err := semver.NewVersion(test.version)
+			require.NoError(t, err)
+			modules.SetTestModules(t, &modules.TestModules{TestBuildType: test.buildType})
+			require.Equal(t, test.want, DistrolessImage(*semVer))
+		})
+	}
+}

lib/web/scripts/oneoff/oneoff.go

@@ -29,13 +29,10 @@ import (
 	"github.com/gravitational/teleport/api"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/lib/modules"
+	"github.com/gravitational/teleport/lib/utils/teleportassets"
 )
 
 const (
-	// teleportCDNLocation is the Teleport's CDN URL
-	// This is used to download the Teleport Binary
-	teleportCDNLocation = "https://cdn.teleport.dev"
-
 	// binUname is the default binary name for inspecting the host's OS.
 	binUname = "uname"
 
@@ -113,14 +110,14 @@ func (p *OneOffScriptParams) CheckAndSetDefaults() error {
 		p.binSudo = "sudo"
 	}
 
-	if p.CDNBaseURL == "" {
-		p.CDNBaseURL = teleportCDNLocation
-	}
-
 	if p.TeleportVersion == "" {
 		p.TeleportVersion = "v" + api.Version
 	}
 
+	if p.CDNBaseURL == "" {
+		p.CDNBaseURL = teleportassets.CDNBaseURL()
+	}
+
 	if p.TeleportFlavor == "" {
 		p.TeleportFlavor = types.PackageNameOSS
 		if modules.GetModules().BuildType() == modules.BuildEnterprise {