[v14] OneOff Script: use ent build if cluster is Enterprise (#33148)

* OneOff Script: use ent build if cluster is Enterprise

We were always using the OSS version of teleport in the one-off scripts.

This PR changes that to pick the correct version depending on the
running version in the Proxy.

* use gzip bestspeed for compressing files

api/types/constants.go

@@ -36,6 +36,11 @@ const (
 	// releases on the release server
 	EnterpriseReleaseEndpoint = "teleport-ent"
 
+	// PackageNameOSS is the teleport package name for the OSS version.
+	PackageNameOSS = "teleport"
+	// PackageNameOSS is the teleport package name for the Enterprise version.
+	PackageNameEnt = "teleport-ent"
+
 	// ActionRead grants read access (get, list)
 	ActionRead = "read"
 

lib/integrations/awsoidc/deployservice.go

@@ -58,12 +58,6 @@ var (
 )
 
 const (
-	// teleportOSS is the prefix for the image name when deploying the OSS version of Teleport
-	teleportOSS = "teleport"
-
-	// teleportEnt is the prefix for the image name when deploying the Enterprise version of Teleport
-	teleportEnt = "teleport-ent"
-
 	// clusterStatusActive is the string representing an ACTIVE ECS Cluster.
 	clusterStatusActive = "ACTIVE"
 	// clusterStatusInactive is the string representing an INACTIVE ECS Cluster.
@@ -433,9 +427,9 @@ func DeployService(ctx context.Context, clt DeployServiceClient, req DeployServi
 
 // upsertTask ensures a TaskDefinition with TaskName exists
 func upsertTask(ctx context.Context, clt DeployServiceClient, req DeployServiceRequest, configB64 string) (*ecsTypes.TaskDefinition, error) {
-	teleportFlavor := teleportOSS
+	teleportFlavor := types.PackageNameOSS
 	if modules.GetModules().BuildType() == modules.BuildEnterprise {
-		teleportFlavor = teleportEnt
+		teleportFlavor = types.PackageNameEnt
 	}
 	taskAgentContainerImage := fmt.Sprintf("public.ecr.aws/gravitational/%s-distroless:%s", teleportFlavor, req.TeleportVersionTag)
 

lib/utils/archive.go

@@ -38,7 +38,10 @@ type ReadStatFS interface {
 func CompressTarGzArchive(files []string, fileReader ReadStatFS) (*bytes.Buffer, error) {
 	archiveBytes := &bytes.Buffer{}
 
-	gzipWriter := gzip.NewWriter(archiveBytes)
+	gzipWriter, err := gzip.NewWriterLevel(archiveBytes, gzip.BestSpeed)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
 	defer gzipWriter.Close()
 
 	tarWriter := tar.NewWriter(gzipWriter)

lib/web/join_tokens.go

@@ -51,9 +51,6 @@ import (
 )
 
 const (
-	teleportOSSPackageName = "teleport"
-	teleportEntPackageName = "teleport-ent"
-
 	stableCloudChannelRepo = "stable/cloud"
 )
 
@@ -384,9 +381,9 @@ func getJoinScript(ctx context.Context, settings scriptSettings, m nodeAPIGetter
 		}
 	}
 
-	packageName := teleportOSSPackageName
+	packageName := types.PackageNameOSS
 	if modules.GetModules().BuildType() == modules.BuildEnterprise {
-		packageName = teleportEntPackageName
+		packageName = types.PackageNameEnt
 	}
 
 	// By default, it will use `stable/v<majorVersion>`, eg stable/v12

lib/web/scripts/oneoff/oneoff.go

@@ -22,8 +22,11 @@ import (
 	"text/template"
 
 	"github.com/gravitational/trace"
+	"golang.org/x/exp/slices"
 
 	"github.com/gravitational/teleport"
+	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/lib/modules"
 )
 
 const (
@@ -66,10 +69,19 @@ type OneOffScriptParams struct {
 	// Eg, v13.1.0
 	TeleportVersion string
 
+	// TeleportFlavor is the teleport flavor to download.
+	// Only OSS or Enterprise versions are allowed.
+	// Possible values:
+	// - teleport
+	// - teleport-ent
+	TeleportFlavor string
+
 	// SuccessMessage is a message shown to the user after the one off is completed.
 	SuccessMessage string
 }
 
+var validPackageNames = []string{types.PackageNameOSS, types.PackageNameEnt}
+
 // CheckAndSetDefaults checks if the required params ara present.
 func (p *OneOffScriptParams) CheckAndSetDefaults() error {
 	if p.TeleportArgs == "" {
@@ -92,6 +104,16 @@ func (p *OneOffScriptParams) CheckAndSetDefaults() error {
 		p.TeleportVersion = "v" + teleport.Version
 	}
 
+	if p.TeleportFlavor == "" {
+		p.TeleportFlavor = types.PackageNameOSS
+		if modules.GetModules().BuildType() == modules.BuildEnterprise {
+			p.TeleportFlavor = types.PackageNameEnt
+		}
+	}
+	if !slices.Contains(validPackageNames, p.TeleportFlavor) {
+		return trace.BadParameter("invalid teleport flavor, only %v are supported", validPackageNames)
+	}
+
 	if p.SuccessMessage == "" {
 		p.SuccessMessage = "Completed successfully."
 	}

lib/web/scripts/oneoff/oneoff.sh

@@ -3,6 +3,7 @@ set -euo pipefail
 
 cdnBaseURL='{{.CDNBaseURL}}'
 teleportVersion='{{.TeleportVersion}}'
+teleportFlavor='{{.TeleportFlavor}}' # teleport or teleport-ent
 successMessage='{{.SuccessMessage}}'
 
 # shellcheck disable=all
@@ -15,7 +16,7 @@ teleportArgs='{{.TeleportArgs}}'
 
 function teleportTarballName(){
     if [[ ${OS} == "Darwin" ]]; then
-        echo teleport-${teleportVersion}-darwin-universal-bin.tar.gz
+        echo ${teleportFlavor}-${teleportVersion}-darwin-universal-bin.tar.gz
         return 0
     fi;
 
@@ -24,10 +25,10 @@ function teleportTarballName(){
         return 1
     fi;
 
-    if [[ ${ARCH} == "armv7l" ]]; then echo "teleport-${teleportVersion}-linux-arm-bin.tar.gz"
-    elif [[ ${ARCH} == "aarch64" ]]; then echo "teleport-${teleportVersion}-linux-arm64-bin.tar.gz"
-    elif [[ ${ARCH} == "x86_64" ]]; then echo "teleport-${teleportVersion}-linux-amd64-bin.tar.gz"
-    elif [[ ${ARCH} == "i686" ]]; then echo "teleport-${teleportVersion}-linux-386-bin.tar.gz"
+    if [[ ${ARCH} == "armv7l" ]]; then echo "${teleportFlavor}-${teleportVersion}-linux-arm-bin.tar.gz"
+    elif [[ ${ARCH} == "aarch64" ]]; then echo "${teleportFlavor}-${teleportVersion}-linux-arm64-bin.tar.gz"
+    elif [[ ${ARCH} == "x86_64" ]]; then echo "${teleportFlavor}-${teleportVersion}-linux-amd64-bin.tar.gz"
+    elif [[ ${ARCH} == "i686" ]]; then echo "${teleportFlavor}-${teleportVersion}-linux-386-bin.tar.gz"
     else
         echo "Invalid Linux architecture ${ARCH}." >&2
         return 1
@@ -43,7 +44,7 @@ function main() {
     tar -xzf ${tarballName}
 
     mkdir -p ./bin
-    mv ./teleport/teleport ./bin/teleport
+    mv ./${teleportFlavor}/teleport ./bin/teleport
     echo "> ./bin/teleport ${teleportArgs}"
     ./bin/teleport ${teleportArgs} && echo $successMessage
 

lib/web/scripts/oneoff/oneoff_test.go

@@ -28,9 +28,11 @@ import (
 	"time"
 
 	"github.com/buildkite/bintest/v3"
+	"github.com/gravitational/trace"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/gravitational/teleport/lib/modules"
 	"github.com/gravitational/teleport/lib/utils"
 )
 
@@ -75,6 +77,7 @@ func TestOneOffScript(t *testing.T) {
 		require.NoError(t, err)
 
 		testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+			assert.Equal(t, "/teleport-v13.1.0-linux-amd64-bin.tar.gz", req.URL.Path)
 			http.ServeContent(w, req, "teleport-v13.1.0-linux-amd64-bin.tar.gz", time.Now(), bytes.NewReader(teleportBinTarball.Bytes()))
 		}))
 		defer func() { testServer.Close() }()
@@ -151,6 +154,76 @@ func TestOneOffScript(t *testing.T) {
 		require.Error(t, err, string(out))
 		require.Contains(t, string(out), "Invalid Linux architecture apple-silicon.")
 	})
+
+	t.Run("invalid flavor should return an error", func(t *testing.T) {
+		_, err := BuildScript(OneOffScriptParams{
+			BinUname:        unameMock.Path,
+			BinMktemp:       mktempMock.Path,
+			CDNBaseURL:      "dummyURL",
+			TeleportVersion: "v13.1.0",
+			TeleportArgs:    "version",
+			SuccessMessage:  "Test was a success.",
+			TeleportFlavor:  "../not-teleport",
+		})
+		require.True(t, trace.IsBadParameter(err), "expected BadParameter, got %+v", err)
+	})
+
+	t.Run("if enterprise build, it uses the enterprise package name", func(t *testing.T) {
+		// set up
+		testWorkingDir := t.TempDir()
+		require.NoError(t, os.Mkdir(testWorkingDir+"/bin/", 0o755))
+		scriptLocation := testWorkingDir + "/" + scriptName
+
+		teleportMock, err := bintest.NewMock(testWorkingDir + "/bin/teleport")
+		require.NoError(t, err)
+		defer func() {
+			assert.NoError(t, teleportMock.Close())
+		}()
+
+		modules.SetTestModules(t, &modules.TestModules{
+			TestBuildType: modules.BuildEnterprise,
+		})
+		teleportBinTarball, err := utils.CompressTarGzArchive([]string{"teleport-ent/teleport"}, singleFileFS{file: teleportMock.Path})
+		require.NoError(t, err)
+
+		testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+			assert.Equal(t, "/teleport-ent-v13.1.0-linux-amd64-bin.tar.gz", req.URL.Path)
+			http.ServeContent(w, req, "teleport-ent-v13.1.0-linux-amd64-bin.tar.gz", time.Now(), bytes.NewReader(teleportBinTarball.Bytes()))
+		}))
+		defer func() { testServer.Close() }()
+
+		script, err := BuildScript(OneOffScriptParams{
+			BinUname:        unameMock.Path,
+			BinMktemp:       mktempMock.Path,
+			CDNBaseURL:      testServer.URL,
+			TeleportVersion: "v13.1.0",
+			TeleportArgs:    "version",
+			SuccessMessage:  "Test was a success.",
+		})
+		require.NoError(t, err)
+
+		unameMock.Expect("-s").AndWriteToStdout("Linux")
+		unameMock.Expect("-m").AndWriteToStdout("x86_64")
+		mktempMock.Expect("-d").AndWriteToStdout(testWorkingDir)
+		teleportMock.Expect("version").AndWriteToStdout(teleportVersionOutput)
+
+		err = os.WriteFile(scriptLocation, []byte(script), 0700)
+		require.NoError(t, err)
+
+		// execute script
+		out, err := exec.Command("bash", scriptLocation).CombinedOutput()
+
+		// validate
+		require.NoError(t, err, string(out))
+
+		require.True(t, unameMock.Check(t))
+		require.True(t, mktempMock.Check(t))
+		require.True(t, teleportMock.Check(t))
+
+		require.Contains(t, string(out), "> ./bin/teleport version")
+		require.Contains(t, string(out), teleportVersionOutput)
+		require.Contains(t, string(out), "Test was a success.")
+	})
 }
 
 type singleFileFS struct {