Add distribution validation with structured error handling

toolkit/tools/imagecreator/main.go

@@ -27,8 +27,8 @@ type ImageCreatorCmd struct {
 	ToolsTar          string   `name:"tools-file" help:"Path to tdnf worker tarball" required:""`
 	OutputImageFile   string   `name:"output-image-file" help:"Path to write the customized image to."`
 	OutputImageFormat string   `name:"output-image-format" placeholder:"(vhd|vhd-fixed|vhdx|qcow2|raw)" help:"Format of output image." enum:"${imageformat}" default:""`
-	Distro            string   `name:"distro" help:"Target distribution for the image." enum:"azurelinux,fedora" default:"azurelinux"`
-	DistroVersion     string   `name:"distro-version" help:"Target distribution version (e.g., 3.0 for Azure Linux, 42 for Fedora)." default:""`
+	Distro            string   `name:"distro" help:"Target distribution for the image." enum:"azurelinux,fedora" required:""`
+	DistroVersion     string   `name:"distro-version" help:"Target distribution version (e.g., 3.0 for Azure Linux, 42 for Fedora)." required:""`
 	exekong.LogFlags
 	PackageSnapshotTime string `name:"package-snapshot-time" help:"Only packages published before this snapshot time will be available during customization. Supports 'YYYY-MM-DD' or full RFC3339 timestamp (e.g., 2024-05-20T23:59:59Z)."`
 }
@@ -54,9 +54,13 @@ func main() {
 
 	logger.InitBestEffort(ptrutils.PtrTo(cli.LogFlags.AsLoggerFlags()))
 
+	distro := imagecreatorlib.Distribution{
+		Name:    cli.Distro,
+		Version: cli.DistroVersion,
+	}
+
 	err := imagecreatorlib.CreateImageWithConfigFile(ctx, cli.BuildDir, cli.ConfigFile, cli.RpmSources,
-		cli.ToolsTar, cli.OutputImageFile, cli.OutputImageFormat, cli.Distro, cli.DistroVersion,
-		cli.PackageSnapshotTime)
+		cli.ToolsTar, cli.OutputImageFile, cli.OutputImageFormat, distro, cli.PackageSnapshotTime)
 	if err != nil {
 		log.Fatalf("image creation failed:\n%v", err)
 	}

toolkit/tools/imagecustomizerapi/distribution.go

@@ -0,0 +1,17 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerapi
+
+const (
+	DistroNameAzureLinux string = "azurelinux"
+	DistroNameFedora     string = "fedora"
+)
+
+func GetSupportedDistros() map[string][]string {
+	// supportedDistros defines valid distribution and version combinations
+	return map[string][]string{
+		DistroNameAzureLinux: {"2.0", "3.0"},
+		DistroNameFedora:     {"42"},
+	}
+}

toolkit/tools/pkg/imagecreatorlib/distribution_validation.go

@@ -0,0 +1,50 @@
+package imagecreatorlib
+
+import (
+	"fmt"
+	"maps"
+	"slices"
+	"strings"
+
+	"github.com/microsoft/azure-linux-image-tools/toolkit/tools/imagecustomizerapi"
+	"github.com/microsoft/azure-linux-image-tools/toolkit/tools/pkg/imagecustomizerlib"
+)
+
+var (
+	// ErrUnsupportedDistribution indicates an unsupported Linux distribution
+	ErrUnsupportedDistribution = imagecustomizerlib.NewImageCustomizerError("Distribution:UnsupportedDistribution",
+		"unsupported distro")
+
+	// ErrUnsupportedVersion indicates an unsupported version for a given distribution
+	ErrUnsupportedVersion = imagecustomizerlib.NewImageCustomizerError("Distribution:UnsupportedVersion",
+		"unsupported distro-version")
+)
+
+// distribution represents a supported Linux distribution and version combination
+type Distribution struct {
+	Name    string
+	Version string
+}
+
+// Validate ensures the distribution and version combination is supported
+func (d *Distribution) Validate() error {
+	// Get all supported distributions and their versions
+	supportedDistros := imagecustomizerapi.GetSupportedDistros()
+
+	// Check if the distribution is supported
+	validVersions, exists := supportedDistros[d.Name]
+	if !exists {
+		// Get list of supported distributions
+		distros := slices.Collect(maps.Keys(supportedDistros))
+		return fmt.Errorf("%w (%q)\nsupported distributions are: (%s)",
+			ErrUnsupportedDistribution, d.Name, strings.Join(distros, ", "))
+	}
+
+	// Check if the version is supported for this distribution
+	if slices.Contains(validVersions, d.Version) {
+		return nil
+	}
+
+	return fmt.Errorf("%w (%q)\nsupported versions for %q distro are: (%s)",
+		ErrUnsupportedVersion, d.Version, d.Name, strings.Join(validVersions, ", "))
+}

toolkit/tools/pkg/imagecreatorlib/imagecreator.go

@@ -48,7 +48,7 @@ type ImageCreatorParameters struct {
 }
 
 func CreateImageWithConfigFile(ctx context.Context, buildDir string, configFile string, rpmsSources []string,
-	toolsTar string, outputImageFile string, outputImageFormat string, distro string, distroVersion string,
+	toolsTar string, outputImageFile string, outputImageFormat string, distro Distribution,
 	packageSnapshotTime string,
 ) error {
 	var config imagecustomizerapi.Config
@@ -66,7 +66,7 @@ func CreateImageWithConfigFile(ctx context.Context, buildDir string, configFile
 
 	err = createNewImage(
 		ctx, buildDir, absBaseConfigPath, config, rpmsSources, outputImageFile,
-		outputImageFormat, toolsTar, distro, distroVersion, packageSnapshotTime)
+		outputImageFormat, toolsTar, distro, packageSnapshotTime)
 	if err != nil {
 		return err
 	}
@@ -75,10 +75,15 @@ func CreateImageWithConfigFile(ctx context.Context, buildDir string, configFile
 }
 
 func createNewImage(ctx context.Context, buildDir string, baseConfigPath string, config imagecustomizerapi.Config,
-	rpmsSources []string, outputImageFile string, outputImageFormat string, toolsTar string, distro string,
-	distroVersion string, packageSnapshotTime string,
+	rpmsSources []string, outputImageFile string, outputImageFormat string, toolsTar string, distro Distribution,
+	packageSnapshotTime string,
 ) error {
-	err := validateConfig(
+	err := distro.Validate()
+	if err != nil {
+		return fmt.Errorf("invalid distribution arguments:\n%w", err)
+	}
+
+	err = validateConfig(
 		ctx, baseConfigPath, &config, rpmsSources, toolsTar, outputImageFile,
 		outputImageFormat, packageSnapshotTime)
 	if err != nil {
@@ -119,7 +124,7 @@ func createNewImage(ctx context.Context, buildDir string, baseConfigPath string,
 	logger.Log.Infof("Creating new image with parameters: %+v\n", imageCreatorParameters)
 
 	// Create distro config from distro name and version
-	distroHandler := imagecustomizerlib.NewDistroHandler(distro, distroVersion)
+	distroHandler := imagecustomizerlib.NewDistroHandler(distro.Name, distro.Version)
 
 	partIdToPartUuid, err := imagecustomizerlib.CreateNewImage(
 		distroHandler.GetTargetOs(), imageCreatorParameters.rawImageFile,

toolkit/tools/pkg/imagecreatorlib/imagecreator_test.go

@@ -28,9 +28,13 @@ func TestCreateImageRaw(t *testing.T) {
 	rpmSources := []string{downloadedRpmsRepoFile}
 	toolsFile := testutils.GetDownloadedToolsFile(t, testutilsDir, "3.0", true)
 
+	distro := Distribution{
+		Name:    "azurelinux",
+		Version: "3.0",
+	}
 	err := CreateImageWithConfigFile(
 		t.Context(), buildDir, partitionsConfigFile, rpmSources, toolsFile,
-		outputImageFilePath, outputImageFormat, "azurelinux", "3.0", "")
+		outputImageFilePath, outputImageFormat, distro, "")
 	if !assert.NoError(t, err) {
 		return
 	}
@@ -73,8 +77,12 @@ func TestCreateImageRawNoTar(t *testing.T) {
 	downloadedRpmsRepoFile := testutils.GetDownloadedRpmsRepoFile(t, testutilsDir, "3.0", false, true)
 	rpmSources := []string{downloadedRpmsRepoFile}
 
+	distro := Distribution{
+		Name:    "azurelinux",
+		Version: "3.0",
+	}
 	err := CreateImageWithConfigFile(t.Context(), buildDir, partitionsConfigFile, rpmSources, "",
-		outputImageFilePath, "raw", "azurelinux", "3.0", "")
+		outputImageFilePath, "raw", distro, "")
 
 	assert.ErrorContains(t, err, "tools tar file is required for image creation")
 }
@@ -87,12 +95,16 @@ func TestCreateImageEmptyConfig(t *testing.T) {
 	// create an empty config file
 	emptyConfigFile := filepath.Join(testDir, "empty-config.yaml")
 
+	distro := Distribution{
+		Name:    "azurelinux",
+		Version: "3.0",
+	}
 	err := CreateImageWithConfigFile(t.Context(), buildDir, "", []string{}, "", outputImageFilePath, "raw",
-		"azurelinux", "3.0", "")
+		distro, "")
 	assert.ErrorContains(t, err, "failed to unmarshal config file")
 
 	err = CreateImageWithConfigFile(t.Context(), buildDir, emptyConfigFile, []string{}, "", outputImageFilePath, "raw",
-		"azurelinux", "3.0", "")
+		distro, "")
 	assert.ErrorContains(t, err, "failed to unmarshal config file")
 }
 
@@ -123,8 +135,12 @@ func TestCreateImage_OutputImageFileAsRelativePath(t *testing.T) {
 
 	// Pass the output image file relative to the current working directory through the argument.
 	// This will create the file at the absolute path.
+	distro := Distribution{
+		Name:    "azurelinux",
+		Version: "3.0",
+	}
 	err = createNewImage(t.Context(), buildDir, baseConfigPath, config, rpmSources, outputImageFile,
-		outputImageFormat, toolsFile, "azurelinux", "3.0", "")
+		outputImageFormat, toolsFile, distro, "")
 	assert.NoError(t, err)
 	assert.FileExists(t, outputImageFileAbsolute)
 	err = os.Remove(outputImageFileAbsolute)
@@ -136,7 +152,7 @@ func TestCreateImage_OutputImageFileAsRelativePath(t *testing.T) {
 	// Pass the output image file relative to the config file through the config. This will create
 	// the file at the absolute path.
 	err = createNewImage(t.Context(), buildDir, baseConfigPath, config, rpmSources, outputImageFile,
-		outputImageFormat, toolsFile, "azurelinux", "3.0", "")
+		outputImageFormat, toolsFile, distro, "")
 	assert.NoError(t, err)
 	assert.FileExists(t, outputImageFileAbsolute)
 	err = os.Remove(outputImageFileAbsolute)

toolkit/tools/pkg/imagecustomizerlib/distrohandler.go

@@ -12,25 +12,6 @@ import (
 	"github.com/microsoft/azure-linux-image-tools/toolkit/tools/internal/targetos"
 )
 
-// PackageManagerType represents the type of package manager
-type PackageManagerType string
-
-const (
-	packageManagerTDNF PackageManagerType = "tdnf"
-	packageManagerDNF  PackageManagerType = "dnf"
-)
-
-// PackageType represents the type of package format
-type PackageType string
-
-// DistroName represents the distribution name
-type DistroName string
-
-const (
-	distroNameAzureLinux DistroName = "azurelinux"
-	distroNameFedora     DistroName = "fedora"
-)
-
 // distroHandler represents the interface for distribution-specific configuration
 type distroHandler interface {
 	GetTargetOs() targetos.TargetOs
@@ -68,9 +49,9 @@ func NewDistroHandlerFromImageConnection(imageConnection *imageconnection.ImageC
 // NewDistroHandler creates the appropriate distro handler with version support (legacy)
 func NewDistroHandler(distroName string, version string) distroHandler {
 	switch distroName {
-	case string(distroNameFedora):
+	case string(imagecustomizerapi.DistroNameFedora):
 		return newFedoraDistroHandler(version)
-	case string(distroNameAzureLinux):
+	case string(imagecustomizerapi.DistroNameAzureLinux):
 		return newAzureLinuxDistroHandler(version)
 	default:
 		panic("unsupported distro name: " + distroName)