roachprod: stage arm64 and FIPS-enabled binaries

Add `--arch` to override binary's architecture and refactor.
As of this change, `roachprod stage` is able to stage both
amd64 and arm64 on linux and darwin, as well as FIPS-enabled
binary built for amd64.
In conjunction with the previous change [1], roachprod
now uses arm64-based AMI for graviton2/graviton3 machines.

Below is an example of how to create a VM with graviton3,
```
roachprod create -n1 --clouds aws --aws-machine-type m7g.2xlarge --local-ssd=false $CRL_USERNAME-test
roachprod stage --arch arm64 $CRL_USERNAME-test release v23.1.0-rc.2
roachprod start $CRL_USERNAME-test
```

[1] cockroachdb#103236

Epic: none
Release note: None

pkg/cmd/roachprod/flags.go

@@ -62,6 +62,7 @@ var (
 	createVMOpts          = vm.DefaultCreateOpts()
 	startOpts             = roachprod.DefaultStartOpts()
 	stageOS               string
+	stageArch             string
 	stageDir              string
 	logsDir               string
 	logsFilter            string
@@ -106,6 +107,8 @@ func initFlags() {
 			vm.AllProviderNames()))
 	createCmd.Flags().BoolVar(&createVMOpts.GeoDistributed,
 		"geo", false, "Create geo-distributed cluster")
+	createCmd.Flags().BoolVar(&createVMOpts.EnableFIPS,
+		"fips", false, "Enable FIPS mode (uses custom AMI)")
 	// N.B. We set "usage=roachprod" as the default, custom label for billing tracking.
 	createCmd.Flags().StringToStringVar(&createVMOpts.CustomLabels,
 		"label", map[string]string{"usage": "roachprod"},
@@ -217,9 +220,14 @@ Default is "RECURRING '*/15 * * * *' FULL BACKUP '@hourly' WITH SCHEDULE OPTIONS
 	putCmd.Flags().BoolVar(&useTreeDist, "treedist", useTreeDist, "use treedist copy algorithm")
 
 	stageCmd.Flags().StringVar(&stageOS, "os", "", "operating system override for staged binaries")
-	stageCmd.Flags().StringVar(&stageDir, "dir", "", "destination for staged binaries")
+	stageCmd.Flags().StringVar(&stageArch, "arch", "",
+		"architecture override for staged binaries [amd64, arm64, fips]; N.B. fips implies amd64 with openssl")
 
+	stageCmd.Flags().StringVar(&stageDir, "dir", "", "destination for staged binaries")
+	// N.B. stageURLCmd just prints the URL that stageCmd would use.
 	stageURLCmd.Flags().StringVar(&stageOS, "os", "", "operating system override for staged binaries")
+	stageURLCmd.Flags().StringVar(&stageArch, "arch", "",
+		"architecture override for staged binaries [amd64, arm64, fips]; N.B. fips implies amd64 with openssl")
 
 	logsCmd.Flags().StringVar(&logsFilter,
 		"filter", "", "re to filter log messages")

pkg/cmd/roachprod/main.go

@@ -56,7 +56,8 @@ The above commands will create a "local" 3 node cluster, start a cockroach
 cluster on these nodes, run a sql command on the 2nd node, stop, wipe and
 destroy the cluster.
 `,
-	Version: "details:\n" + build.GetInfo().Long(),
+	Version:          "details:\n" + build.GetInfo().Long(),
+	PersistentPreRun: validateAndConfigure,
 }
 
 // Provide `cobra.Command` functions with a standard return code handler.
@@ -714,7 +715,7 @@ Currently available application options are:
 		if len(args) == 2 {
 			versionArg = args[1]
 		}
-		urls, err := roachprod.StageURL(config.Logger, args[0], versionArg, stageOS)
+		urls, err := roachprod.StageURL(config.Logger, args[0], versionArg, stageOS, stageArch)
 		if err != nil {
 			return err
 		}
@@ -753,7 +754,7 @@ Some examples of usage:
 		if len(args) == 3 {
 			versionArg = args[2]
 		}
-		return roachprod.Stage(context.Background(), config.Logger, args[0], stageOS, stageDir, args[1], versionArg)
+		return roachprod.Stage(context.Background(), config.Logger, args[0], stageOS, stageArch, stageDir, args[1], versionArg)
 	}),
 }
 
@@ -1076,6 +1077,35 @@ var fixLongRunningAWSHostnamesCmd = &cobra.Command{
 	}),
 }
 
+// Before executing any command, validate and canonicalize args.
+func validateAndConfigure(cmd *cobra.Command, args []string) {
+	// Skip validation for commands that are self-sufficient.
+	switch cmd.Name() {
+	case "help", "version", "list":
+		return
+	}
+
+	printErrAndExit := func(err error) {
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "%s\n", err)
+			os.Exit(1)
+		}
+	}
+
+	// Validate architecture flag, if set.
+	if archOpt := cmd.Flags().Lookup("arch"); archOpt != nil && archOpt.Changed {
+		arch := strings.ToLower(archOpt.Value.String())
+
+		if arch != "amd64" && arch != "arm64" && arch != "fips" {
+			printErrAndExit(fmt.Errorf("unsupported architecture %q", arch))
+		}
+		if arch != archOpt.Value.String() {
+			// Set the canonical value.
+			_ = cmd.Flags().Set("arch", arch)
+		}
+	}
+}
+
 func main() {
 	_ = roachprod.InitProviders()
 	providerOptsContainer = vm.CreateProviderOptionsContainer()

pkg/cmd/roachtest/cluster.go

@@ -1723,7 +1723,8 @@ func (c *clusterImpl) Stage(
 	}
 	c.status("staging binary")
 	defer c.status("")
-	return errors.Wrap(roachprod.Stage(ctx, l, c.MakeNodes(opts...), "" /* stageOS */, dir, application, versionOrSHA), "cluster.Stage")
+	return errors.Wrap(roachprod.Stage(ctx, l, c.MakeNodes(opts...),
+		"" /* stageOS */, "" /* stageArch */, dir, application, versionOrSHA), "cluster.Stage")
 }
 
 // Get gets files from remote hosts.

pkg/roachprod/install/staging.go

@@ -43,21 +43,49 @@ type archInfo struct {
 	ReleaseArchiveExtension string
 }
 
+// N.B. DebugArchitecture must correspond to 'SuffixFromPlatform' followed by stripping the literal (os) version,
+// in 'MakeCRDBBinaryNonReleaseFile' (see pkg/release/build.go and pkg/release/release.go).
+//
+//	ReleaseArchitecture must correspond to 'SuffixFromPlatform' followed by stripping "gnu-" in 'makeArchiveKeys'
+//
+// (see pkg/release/upload.go).
+// TODO(srosenberg): refactor to use the above, directly from pkg/release/ which is the source of truth.
 var (
-	linuxArchInfo = archInfo{
+	linux_x86_64_ArchInfo = archInfo{
 		DebugArchitecture:       "linux-gnu-amd64",
 		ReleaseArchitecture:     "linux-amd64",
 		LibraryExtension:        ".so",
 		ExecutableExtension:     "",
 		ReleaseArchiveExtension: "tgz",
 	}
-	darwinArchInfo = archInfo{
+	linux_x86_64_fips_ArchInfo = archInfo{
+		DebugArchitecture:       "linux-gnu-amd64-fips",
+		ReleaseArchitecture:     "linux-amd64-fips",
+		LibraryExtension:        ".so",
+		ExecutableExtension:     "",
+		ReleaseArchiveExtension: "tgz",
+	}
+	linux_arm64_ArchInfo = archInfo{
+		DebugArchitecture:       "linux-gnu-arm64",
+		ReleaseArchitecture:     "linux-arm64",
+		LibraryExtension:        ".so",
+		ExecutableExtension:     "",
+		ReleaseArchiveExtension: "tgz",
+	}
+	darwin_x86_64_ArchInfo = archInfo{
 		DebugArchitecture:       "darwin-amd64",
 		ReleaseArchitecture:     "darwin-10.9-amd64",
 		LibraryExtension:        ".dylib",
 		ExecutableExtension:     "",
 		ReleaseArchiveExtension: "tgz",
 	}
+	darwin_arm64_ArchInfo = archInfo{
+		DebugArchitecture:       "darwin-arm64.unsigned",
+		ReleaseArchitecture:     "darwin-11.0-arm64",
+		LibraryExtension:        ".dylib",
+		ExecutableExtension:     "",
+		ReleaseArchiveExtension: "tgz",
+	}
 	windowsArchInfo = archInfo{
 		DebugArchitecture:       "windows-amd64",
 		ReleaseArchitecture:     "windows-6.2-amd64",
@@ -69,18 +97,36 @@ var (
 	crdbLibraries = []string{"libgeos", "libgeos_c"}
 )
 
-// ArchInfoForOS returns an ArchInfo for the given OS if the OS is
-// currently supported.
-func archInfoForOS(os string) (archInfo, error) {
+// ArchInfoForOS returns an ArchInfo for the given OS and Architecture if currently supported.
+func ArchInfoForOS(os string, arch string) (archInfo, error) {
+	if arch != "" && arch != "amd64" && arch != "arm64" && arch != "fips" {
+		return archInfo{}, errors.Errorf("unsupported architecture %q", arch)
+	}
+
 	switch os {
 	case "linux":
-		return linuxArchInfo, nil
+		if arch == "arm64" {
+			return linux_arm64_ArchInfo, nil
+		}
+		if arch == "fips" {
+			return linux_x86_64_fips_ArchInfo, nil
+		}
+		return linux_x86_64_ArchInfo, nil
 	case "darwin":
-		return darwinArchInfo, nil
+		if arch == "arm64" {
+			return darwin_arm64_ArchInfo, nil
+		}
+		if arch == "fips" {
+			return archInfo{}, errors.Errorf("%q is not supported on %q", arch, os)
+		}
+		return darwin_x86_64_ArchInfo, nil
 	case "windows":
+		if arch == "fips" || arch == "arm64" {
+			return archInfo{}, errors.Errorf("%q is not supported on %q", arch, os)
+		}
 		return windowsArchInfo, nil
 	default:
-		return archInfo{}, errors.Errorf("no release architecture information for %q", os)
+		return archInfo{}, errors.Errorf("unsupported OS %q", os)
 	}
 }
 
@@ -130,16 +176,17 @@ func StageApplication(
 	applicationName string,
 	version string,
 	os string,
+	arch string,
 	destDir string,
 ) error {
-	archInfo, err := archInfoForOS(os)
+	archInfo, err := ArchInfoForOS(os, arch)
 	if err != nil {
 		return err
 	}
 
 	switch applicationName {
 	case "cockroach":
-		err := StageRemoteBinary(
+		err := stageRemoteBinary(
 			ctx, l, c, applicationName, "cockroach/cockroach", version, archInfo.DebugArchitecture, destDir,
 		)
 		if err != nil {
@@ -164,7 +211,8 @@ func StageApplication(
 		}
 		return nil
 	case "workload":
-		err := StageRemoteBinary(
+		// N.B. workload binary is only available for linux amd64: https://github.com/cockroachdb/cockroach/issues/103563
+		err := stageRemoteBinary(
 			ctx, l, c, applicationName, "cockroach/workload", version, "" /* arch */, destDir,
 		)
 		return err
@@ -177,8 +225,10 @@ func StageApplication(
 
 // URLsForApplication returns a slice of URLs that should be
 // downloaded for the given application.
-func URLsForApplication(application string, version string, os string) ([]*url.URL, error) {
-	archInfo, err := archInfoForOS(os)
+func URLsForApplication(
+	application string, version string, os string, arch string,
+) ([]*url.URL, error) {
+	archInfo, err := ArchInfoForOS(os, arch)
 	if err != nil {
 		return nil, err
 	}
@@ -205,6 +255,7 @@ func URLsForApplication(application string, version string, os string) ([]*url.U
 		}
 		return urls, nil
 	case "workload":
+		// N.B. workload binary is only available for linux amd64: https://github.com/cockroachdb/cockroach/issues/103563
 		u, err := getEdgeURL("cockroach/workload", version, "" /* arch */, "" /* extension */)
 		if err != nil {
 			return nil, err
@@ -225,7 +276,7 @@ func URLsForApplication(application string, version string, os string) ([]*url.U
 // application path to each specified by the cluster to the specified directory.
 // If no SHA is specified, the latest build of the binary is used instead.
 // Returns the SHA of the resolved binary.
-func StageRemoteBinary(
+func stageRemoteBinary(
 	ctx context.Context,
 	l *logger.Logger,
 	c *SyncedCluster,