Add flavor to Agent PackageVersion (#1286)

components/datadog/agent/host_linuxos.go

@@ -46,6 +46,10 @@ func (am *agentLinuxManager) getInstallCommand(version agentparams.PackageVersio
 		}
 	}
 
+	if version.Flavor != "" {
+		testEnvVars = append(testEnvVars, fmt.Sprintf("DD_AGENT_FLAVOR=%s", version.Flavor))
+	}
+
 	commandLine = strings.Join(testEnvVars, " ")
 
 	return fmt.Sprintf(

components/datadog/agent/host_windowsos.go

@@ -111,6 +111,9 @@ while ($tries -lt 5) {
 }
 
 func getAgentURL(version agentparams.PackageVersion) (string, error) {
+	if version.Flavor == "" {
+		version.Flavor = agentparams.DefaultFlavor
+	}
 	minor := strings.ReplaceAll(version.Minor, "~", "-")
 	fullVersion := fmt.Sprintf("%v.%v", version.Major, minor)
 
@@ -119,22 +122,20 @@ func getAgentURL(version agentparams.PackageVersion) (string, error) {
 	}
 
 	if version.Channel == agentparams.BetaChannel {
-		finder, err := newAgentURLFinder("https://s3.amazonaws.com/dd-agent-mstesting/builds/beta/installers_v2.json")
+		finder, err := newAgentURLFinder("https://s3.amazonaws.com/dd-agent-mstesting/builds/beta/installers_v2.json", version.Flavor)
 		if err != nil {
 			return "", err
 		}
 
 		url, err := finder.findVersion(fullVersion)
 		if err != nil {
-			// Try to handle custom build
-			minor = strings.TrimSuffix(minor, "-1")
-			return fmt.Sprintf("https://s3.amazonaws.com/dd-agent-mstesting/builds/beta/ddagent-cli-%v.%v.msi", version.Major, minor), nil
+			return "", err
 		}
 
 		return url, nil
 	}
 
-	finder, err := newAgentURLFinder("https://ddagent-windows-stable.s3.amazonaws.com/installers_v2.json")
+	finder, err := newAgentURLFinder("https://ddagent-windows-stable.s3.amazonaws.com/installers_v2.json", version.Flavor)
 	if err != nil {
 		return "", err
 	}
@@ -151,6 +152,28 @@ func getAgentURL(version agentparams.PackageVersion) (string, error) {
 }
 
 func getAgentURLFromPipelineID(version agentparams.PackageVersion) (string, error) {
+	url, err := getPipelineArtifact(version.PipelineID, "dd-agent-mstesting", version.Major, func(artifact string) bool {
+		if !strings.Contains(artifact, fmt.Sprintf("%s-%s", version.Flavor, version.Major)) {
+			return false
+		}
+		if !strings.HasSuffix(artifact, ".msi") {
+			return false
+		}
+
+		return true
+	})
+	if err != nil {
+		return "", err
+	}
+
+	return url, nil
+}
+
+// getPipelineArtifact searches a public S3 bucket for a given artifact from a Gitlab pipeline
+// majorVersion = [6,7]
+// predicate = A function taking the artifact name (from github.com/aws/aws-sdk-go-v2/service/s3/types.Object.Key)
+// and that returns true when the artifact matches.
+func getPipelineArtifact(pipelineID, bucket, majorVersion string, predicate func(string) bool) (string, error) {
 	// TODO: Replace context.Background() with a Pulumi context.Context.
 	// dd-agent-mstesting is a public bucket so we can use anonymous credentials
 	config, err := awsConfig.LoadDefaultConfig(context.Background(), awsConfig.WithCredentialsProvider(aws.AnonymousCredentials{}))
@@ -160,27 +183,37 @@ func getAgentURLFromPipelineID(version agentparams.PackageVersion) (string, erro
 
 	s3Client := s3.NewFromConfig(config)
 
+	// Manual URL example: https://s3.amazonaws.com/dd-agent-mstesting?prefix=pipelines/A7/25309493
 	result, err := s3Client.ListObjectsV2(context.Background(), &s3.ListObjectsV2Input{
-		Bucket: aws.String("dd-agent-mstesting"),
-		Prefix: aws.String(fmt.Sprintf("pipelines/A%v/%v", version.Major, version.PipelineID)),
+		Bucket: aws.String(bucket),
+		Prefix: aws.String(fmt.Sprintf("pipelines/A%s/%s", majorVersion, pipelineID)),
 	})
+
 	if err != nil {
 		return "", err
 	}
 
 	if len(result.Contents) <= 0 {
-		return "", fmt.Errorf("no agent MSI found for pipeline %v", version.PipelineID)
+		return "", fmt.Errorf("no artifact found for pipeline %v", pipelineID)
+	}
+
+	for _, obj := range result.Contents {
+		if !predicate(*obj.Key) {
+			continue
+		}
+
+		return fmt.Sprintf("https://s3.amazonaws.com/%s/%s", bucket, *obj.Key), nil
 	}
 
-	return "https://s3.amazonaws.com/dd-agent-mstesting/" + *result.Contents[0].Key, nil
+	return "", fmt.Errorf("no agent artifact found for pipeline %v", pipelineID)
 }
 
 type agentURLFinder struct {
 	versions     map[string]interface{}
 	installerURL string
 }
 
-func newAgentURLFinder(installerURL string) (*agentURLFinder, error) {
+func newAgentURLFinder(installerURL string, flavor string) (*agentURLFinder, error) {
 	resp, err := http.Get(installerURL)
 	if err != nil {
 		return nil, err
@@ -195,7 +228,7 @@ func newAgentURLFinder(installerURL string) (*agentURLFinder, error) {
 		return nil, err
 	}
 
-	versions, err := getKey[map[string]interface{}](values, "datadog-agent")
+	versions, err := getKey[map[string]interface{}](values, flavor)
 	if err != nil {
 		return nil, err
 	}

components/datadog/agentparams/params.go

@@ -108,6 +108,16 @@ func WithVersion(version string) func(*Params) error {
 	}
 }
 
+// WithFlavor use a specific flavor of the Agent. For example: `datadog-fips-agent`
+//
+// See PackageFlavor https://github.com/DataDog/agent-release-management/blob/main/generator/const.py
+func WithFlavor(flavor string) func(*Params) error {
+	return func(p *Params) error {
+		p.Version.Flavor = flavor
+		return nil
+	}
+}
+
 // WithPipeline use a specific version of the Agent by pipeline id
 func WithPipeline(pipelineID string) func(*Params) error {
 	return func(p *Params) error {

components/datadog/agentparams/version.go

@@ -8,9 +8,22 @@ const (
 	NightlyChannel channel = "nightly"
 )
 
+// Agent flavor constants
+//
+// PackageVersion.Flavor is in some cases passed directly to underlying install scripts,
+// so this list is not exhaustive.
+//
+// See PackageFlavor https://github.com/DataDog/agent-release-management/blob/main/generator/const.py
+const (
+	DefaultFlavor = BaseFlavor
+	BaseFlavor    = "datadog-agent"
+	FIPSFlavor    = "datadog-fips-agent"
+)
+
 type PackageVersion struct {
 	Major      string
 	Minor      string // Empty means latest
 	Channel    channel
 	PipelineID string
+	Flavor     string // Empty means default (base)
 }