From da75f56f380057ac0f5c12b1b42b5b6fd7beb0e9 Mon Sep 17 00:00:00 2001
From: Thomas Hipp <thomas.hipp@canonical.com>
Date: Tue, 13 Jun 2023 08:46:32 +0200
Subject: [PATCH] windows: Detect Windows architecture from filename

Virtio drivers are built for both amd64 and arm64.

This tries extracting the architecture from the filename, and use the
appropriate virtio drivers.

Fixes #705

Signed-off-by: Thomas Hipp <thomas.hipp@canonical.com>
---
 distrobuilder/main_repack-windows.go      | 41 +++++++++++++++++++++--
 distrobuilder/main_repack-windows_test.go | 30 +++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/distrobuilder/main_repack-windows.go b/distrobuilder/main_repack-windows.go
index 7510cd52..b45e724a 100644
--- a/distrobuilder/main_repack-windows.go
+++ b/distrobuilder/main_repack-windows.go
@@ -27,8 +27,9 @@ import (
 type cmdRepackWindows struct {
 	global *cmdGlobal
 
-	flagDrivers        string
-	flagWindowsVersion string
+	flagDrivers             string
+	flagWindowsVersion      string
+	flagWindowsArchitecture string
 }
 
 func init() {
@@ -89,6 +90,7 @@ func (c *cmdRepackWindows) command() *cobra.Command {
 
 	cmd.Flags().StringVar(&c.flagDrivers, "drivers", "", "Path to drivers ISO"+"``")
 	cmd.Flags().StringVar(&c.flagWindowsVersion, "windows-version", "", "Windows version to repack"+"``")
+	cmd.Flags().StringVar(&c.flagWindowsArchitecture, "windows-arch", "", "Windows architecture to repack"+"``")
 
 	return cmd
 }
@@ -113,6 +115,22 @@ func (c *cmdRepackWindows) preRun(cmd *cobra.Command, args []string) error {
 		}
 	}
 
+	if c.flagWindowsArchitecture == "" {
+		detectedArchitecture := detectWindowsArchitecture(filepath.Base(args[0]))
+
+		if detectedArchitecture == "" {
+			return errors.New("Failed to detect Windows architecture. Please provide the architecture using the --windows-arch flag")
+		}
+
+		c.flagWindowsArchitecture = detectedArchitecture
+	} else {
+		supportedArchitectures := []string{"amd64", "ARM64"}
+
+		if !lxd.StringInSlice(c.flagWindowsArchitecture, supportedArchitectures) {
+			return fmt.Errorf("Architecture must be one of %v", supportedArchitectures)
+		}
+	}
+
 	// Check dependencies
 	err := c.checkDependencies()
 	if err != nil {
@@ -521,7 +539,7 @@ func (c *cmdRepackWindows) injectDrivers(dirs map[string]string) error {
 			"driverName":  driver,
 		}
 
-		sourceDir := filepath.Join(driverPath, driver, c.flagWindowsVersion, "amd64")
+		sourceDir := filepath.Join(driverPath, driver, c.flagWindowsVersion, c.flagWindowsArchitecture)
 		targetBasePath := filepath.Join(dirs["filerepository"], info.PackageName)
 
 		if !lxd.PathExists(targetBasePath) {
@@ -690,6 +708,23 @@ func detectWindowsVersion(fileName string) string {
 	return ""
 }
 
+func detectWindowsArchitecture(fileName string) string {
+	aliases := map[string][]string{
+		"amd64": {"amd64", "x64"},
+		"ARM64": {"arm64"},
+	}
+
+	for k, v := range aliases {
+		for _, alias := range v {
+			if regexp.MustCompile(fmt.Sprintf("(?i)%s", alias)).MatchString(fileName) {
+				return k
+			}
+		}
+	}
+
+	return ""
+}
+
 // toHex is a pongo2 filter which converts the provided value to a hex value understood by the Windows registry.
 func toHex(in *pongo2.Value, param *pongo2.Value) (out *pongo2.Value, err *pongo2.Error) {
 	dst := make([]byte, hex.EncodedLen(len(in.String())))
diff --git a/distrobuilder/main_repack-windows_test.go b/distrobuilder/main_repack-windows_test.go
index 4ecbfdc7..2d9c7864 100644
--- a/distrobuilder/main_repack-windows_test.go
+++ b/distrobuilder/main_repack-windows_test.go
@@ -105,3 +105,33 @@ func Test_detectWindowsVersion(t *testing.T) {
 		})
 	}
 }
+
+func Test_detectWindowsArchitecture(t *testing.T) {
+	type args struct {
+		fileName string
+	}
+
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			"Windows 11 (1)",
+			args{"Win10_22H2_English_x64.iso"},
+			"amd64",
+		},
+		{
+			"Windows 11 (2)",
+			args{"Win10_22H2_English_arm64.iso"},
+			"ARM64",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := detectWindowsArchitecture(tt.args.fileName)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}