cmd/catalyst/main.go - infra/infra - Git at Google

 // Copyright 2019 The Fuchsia Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can
 // found in the LICENSE file.

 package main

 import (
 	"context"
 	"flag"
 	"fmt"
 	"io"
 	"log"
 	"net/http"
 	"net/url"
 	"os"
 	"os/exec"
 	"os/signal"
 	"path/filepath"
 	"strconv"
 	"strings"
 	"syscall"
 	"time"

 	"cloud.google.com/go/storage"
 	devicePkg "go.fuchsia.dev/infra/devices"
 )

 const (
 	failureExitCode = 1
 	portEnvVar      = "NUC_SERVER_PORT"
 	// chainloaderExitStr instructs a NUC to exit IPXE and boot from disk
 	chainloaderExitStr = "#!ipxe\necho Catalyst chain loader\nexit 0\n"
 	rebootDuration     = 1 * time.Minute
 	zedbootPath        = "zedboot.zbi"
 )

 var (
 	imagesManifest      string
 	deviceConfigPath    string
 	bootserverPath      string
 	disableRebootServer bool
 )

 func init() {
 	flag.StringVar(&imagesManifest, "images", "", "Path to the images manifest json")
 	flag.StringVar(&deviceConfigPath, "config", "", "Path to the device config file")
 	flag.StringVar(&bootserverPath, "bootserver", "", "Path to the bootserver binary")
 	flag.BoolVar(&disableRebootServer, "disable-reboot-server", false, "Disables the NUC soft reboot server")
 }

 // Runs a subprocess and sets up a handler that propagates SIGTERM on context cancel.
 func runSubprocess(ctx context.Context, command []string) int {
 	if len(command) == 0 {
 		return 0
 	}
 	cmd := exec.Command(command[0], command[1:]...)

 	// Spin off handler to exit subprocesses cleanly via SIGTERM.
 	processDone := make(chan bool, 1)
 	go func() {
 		select {
 		case <-processDone:
 		case <-ctx.Done():
 			if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
 				log.Printf("exited cmd with error %v", err)
 			}
 		}
 	}()

 	// Ensure that the context still exists before running the subprocess.
 	if ctx.Err() != nil {
 		return failureExitCode
 	}

 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 	cmd.Run()
 	processDone <- true
 	return cmd.ProcessState.ExitCode()
 }

 // runNUCServer runs an http server that delivers chainloaders/images to NUCs based on their mac.
 // These chainloaders deliver a build's version of zedboot and an exit chainloader (to allow reboot from disk).
 func runNUCServer(ctx context.Context, devices []devicePkg.Device, disabled bool) (*http.Server, error) {
 	// Parse the port from the environment variable.
 	portStr := os.Getenv(portEnvVar)
 	if portStr == "" || disabled {
 		return nil, nil
 	}
 	port, err := strconv.Atoi(portStr)
 	if err != nil {
 		return nil, err
 	}

 	mux := http.NewServeMux()
 	for _, device := range devices {
 		// Add endpoint for exit chainloader. This allows NUCs to reboot from disk.
 		mux.HandleFunc(fmt.Sprintf("/%s.ipxe", device.Mac()), func(w http.ResponseWriter, r *http.Request) {
 			w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s.ipxe", device.Mac()))
 			chainloader := strings.NewReader(chainloaderExitStr)
 			io.Copy(w, chainloader)
 		})

 		// Add single-use endpoint to deliver build's version of zedboot.
 		zedbootFile, err := os.Open(zedbootPath)
 		if err != nil {
 			zedbootFile = nil
 		}
 		mux.HandleFunc(fmt.Sprintf("/zedboot/%s", device.Mac()), func(w http.ResponseWriter, r *http.Request) {
 			if zedbootFile == nil {
 				w.WriteHeader(http.StatusNotFound)
 				return
 			}
 			w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", device.Mac()))
 			io.Copy(w, zedbootFile)
 			zedbootFile.Close()
 			zedbootFile = nil
 		})
 	}

 	srv := &http.Server{
 		Addr:    fmt.Sprintf(":%d", port),
 		Handler: mux,
 	}
 	go srv.ListenAndServe()
 	return srv, nil
 }

 // prepDevices puts all devices into TaskState.
 func prepDevices(ctx context.Context, devices []devicePkg.Device) error {
 	// Put each device into TaskState.
 	errorChannel := make(chan error)
 	for _, device := range devices {
 		go func(device devicePkg.Device) {
 			errorChannel <- device.ToTaskState(ctx)
 		}(device)
 	}
 	for i := 0; i < len(devices); i++ {
 		if err := <-errorChannel; err != nil {
 			return err
 		}
 	}
 	return nil
 }

 // hasNUC returns true iff this testbed contains a NUC.
 func hasNUC(devices []devicePkg.Device) bool {
 	for _, device := range devices {
 		if _, ok := device.(*devicePkg.Nuc); ok {
 			return true
 		}
 	}
 	return false
 }

 // downloadZedboot retrieves zedboot.zbi from GCS so that we can use the proper image
 // in NUCs.
 func downloadZedboot(ctx context.Context, imageURL *url.URL, devices []devicePkg.Device) error {
 	// Ensure that we have a NUC that needs a local zedboot.zbi.
 	if !hasNUC(devices) {
 		return nil
 	}

 	// If we are not getting images from GCS, then this is a no-op.
 	if imageURL.Scheme != "gs" {
 		return nil
 	}
 	// Connect to GCS.
 	bucket := imageURL.Host
 	client, err := storage.NewClient(ctx)
 	if err != nil {
 		return err
 	}
 	bkt := client.Bucket(bucket)

 	// Construct path to zedboot given images manifest path.
 	gcsZedbootPath := strings.TrimLeft(
 		fmt.Sprintf("%s/%s", filepath.Dir(imageURL.Path), zedbootPath), "/",
 	)

 	// Get reader to GCS object.
 	r, err := bkt.Object(gcsZedbootPath).NewReader(ctx)
 	if err != nil {
 		return err
 	}
 	defer r.Close()

 	// Open a local zedboot file and download remote data into it.
 	file, err := os.Create(zedbootPath)
 	if err != nil {
 		return err
 	}
 	defer file.Close()

 	if _, err := io.Copy(file, r); err != nil {
 		return err
 	}
 	return nil
 }

 func execute(ctx context.Context, subcommandArgs []string) (int, error) {
 	// If this is a QEMU test bench, skip the device setup and just run the subprocess.
 	if deviceConfigPath == "" {
 		return runSubprocess(ctx, subcommandArgs), nil
 	}
 	// Contains all necessary bootserver flags except device nodename.
 	bootserverCmdStub := []string{
 		bootserverPath,
 		"--images", imagesManifest,
 		"--mode", "pave-zedboot",
 		"-n",
 	}

 	// Create devicePkg.Devices for each of the devices in the config file
 	configs, err := devicePkg.LoadDeviceConfigs(deviceConfigPath)
 	if err != nil {
 		return failureExitCode, err
 	}
 	devices, err := devicePkg.CreateDevices(ctx, configs, bootserverCmdStub)
 	if err != nil {
 		return failureExitCode, err
 	}

 	// If the image manifest provided is a gcs url and we have a NUC, we need to download
 	// zedboot. If the image manifest is a file path, then zedboot.zbi should exist in this
 	// directory.
 	imageURL, err := url.Parse(imagesManifest)
 	if err != nil {
 		return failureExitCode, err
 	}
 	if imageURL.Scheme == "gs" {
 		// This is a no-op if no NUCs exist on this testbed.
 		if err := downloadZedboot(ctx, imageURL, devices); err != nil {
 			return failureExitCode, err
 		}
 	}

 	// Set up a NUC server. If the port environment variable is not set, or if
 	// the disableRebootServer flag is set, this is a no-op.
 	srv, err := runNUCServer(ctx, devices, disableRebootServer)
 	if err != nil {
 		return failureExitCode, err
 	}

 	// Clean up after tests
 	defer func() {
 		if srv != nil {
 			if err := srv.Shutdown(ctx); err != nil {
 				log.Printf("NUC server shutdown failed: %v", err)
 			}
 		}
 	}()

 	if err := prepDevices(ctx, devices); err != nil {
 		return failureExitCode, err
 	}

 	time.Sleep(rebootDuration)
 	// Execute the passed in subcommand
 	return runSubprocess(ctx, subcommandArgs), nil
 }

 func main() {
 	// Initialize
 	flag.Parse()

 	// Handle SIGTERM
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	signals := make(chan os.Signal)
 	defer func() {
 		signal.Stop(signals)
 		close(signals)
 	}()
 	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)

 	go func() {
 		select {
 		case <-signals:
 			cancel()
 		case <-ctx.Done():
 		}
 	}()

 	// Resolve environment variables within the subcommand's command line, as
 	// Swarming does not do it automatically, and best to have this done once and
 	// uniformly across all possible test task commands.
 	var subcmdArgs []string
 	for _, arg := range flag.Args() {
 		subcmdArgs = append(subcmdArgs, os.ExpandEnv(arg))
 	}
 	exitCode, err := execute(ctx, subcmdArgs)

 	if err != nil {
 		log.Printf("Exit code: %d, Err: %s\n", exitCode, err)
 	}
 	os.Exit(exitCode)
 }
	// Copyright 2019 The Fuchsia Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style license that can
	// found in the LICENSE file.

	package main

	import (
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	"cloud.google.com/go/storage"
	devicePkg "go.fuchsia.dev/infra/devices"
	)

	const (
	failureExitCode = 1
	portEnvVar = "NUC_SERVER_PORT"
	// chainloaderExitStr instructs a NUC to exit IPXE and boot from disk
	chainloaderExitStr = "#!ipxe\necho Catalyst chain loader\nexit 0\n"
	rebootDuration = 1 * time.Minute
	zedbootPath = "zedboot.zbi"
	)

	var (
	imagesManifest string
	deviceConfigPath string
	bootserverPath string
	disableRebootServer bool
	)

	func init() {
	flag.StringVar(&imagesManifest, "images", "", "Path to the images manifest json")
	flag.StringVar(&deviceConfigPath, "config", "", "Path to the device config file")
	flag.StringVar(&bootserverPath, "bootserver", "", "Path to the bootserver binary")
	flag.BoolVar(&disableRebootServer, "disable-reboot-server", false, "Disables the NUC soft reboot server")
	}

	// Runs a subprocess and sets up a handler that propagates SIGTERM on context cancel.
	func runSubprocess(ctx context.Context, command []string) int {
	if len(command) == 0 {
	return 0
	}
	cmd := exec.Command(command[0], command[1:]...)

	// Spin off handler to exit subprocesses cleanly via SIGTERM.
	processDone := make(chan bool, 1)
	go func() {
	select {
	case <-processDone:
	case <-ctx.Done():
	if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
	log.Printf("exited cmd with error %v", err)
	}
	}
	}()

	// Ensure that the context still exists before running the subprocess.
	if ctx.Err() != nil {
	return failureExitCode
	}

	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()
	processDone <- true
	return cmd.ProcessState.ExitCode()
	}

	// runNUCServer runs an http server that delivers chainloaders/images to NUCs based on their mac.
	// These chainloaders deliver a build's version of zedboot and an exit chainloader (to allow reboot from disk).
	func runNUCServer(ctx context.Context, devices []devicePkg.Device, disabled bool) (*http.Server, error) {
	// Parse the port from the environment variable.
	portStr := os.Getenv(portEnvVar)
	if portStr == "" \|\| disabled {
	return nil, nil
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
	return nil, err
	}

	mux := http.NewServeMux()
	for _, device := range devices {
	// Add endpoint for exit chainloader. This allows NUCs to reboot from disk.
	mux.HandleFunc(fmt.Sprintf("/%s.ipxe", device.Mac()), func(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s.ipxe", device.Mac()))
	chainloader := strings.NewReader(chainloaderExitStr)
	io.Copy(w, chainloader)
	})

	// Add single-use endpoint to deliver build's version of zedboot.
	zedbootFile, err := os.Open(zedbootPath)
	if err != nil {
	zedbootFile = nil
	}
	mux.HandleFunc(fmt.Sprintf("/zedboot/%s", device.Mac()), func(w http.ResponseWriter, r *http.Request) {
	if zedbootFile == nil {
	w.WriteHeader(http.StatusNotFound)
	return
	}
	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", device.Mac()))
	io.Copy(w, zedbootFile)
	zedbootFile.Close()
	zedbootFile = nil
	})
	}

	srv := &http.Server{
	Addr: fmt.Sprintf(":%d", port),
	Handler: mux,
	}
	go srv.ListenAndServe()
	return srv, nil
	}

	// prepDevices puts all devices into TaskState.
	func prepDevices(ctx context.Context, devices []devicePkg.Device) error {
	// Put each device into TaskState.
	errorChannel := make(chan error)
	for _, device := range devices {
	go func(device devicePkg.Device) {
	errorChannel <- device.ToTaskState(ctx)
	}(device)
	}
	for i := 0; i < len(devices); i++ {
	if err := <-errorChannel; err != nil {
	return err
	}
	}
	return nil
	}

	// hasNUC returns true iff this testbed contains a NUC.
	func hasNUC(devices []devicePkg.Device) bool {
	for _, device := range devices {
	if _, ok := device.(*devicePkg.Nuc); ok {
	return true
	}
	}
	return false
	}

	// downloadZedboot retrieves zedboot.zbi from GCS so that we can use the proper image
	// in NUCs.
	func downloadZedboot(ctx context.Context, imageURL *url.URL, devices []devicePkg.Device) error {
	// Ensure that we have a NUC that needs a local zedboot.zbi.
	if !hasNUC(devices) {
	return nil
	}

	// If we are not getting images from GCS, then this is a no-op.
	if imageURL.Scheme != "gs" {
	return nil
	}
	// Connect to GCS.
	bucket := imageURL.Host
	client, err := storage.NewClient(ctx)
	if err != nil {
	return err
	}
	bkt := client.Bucket(bucket)

	// Construct path to zedboot given images manifest path.
	gcsZedbootPath := strings.TrimLeft(
	fmt.Sprintf("%s/%s", filepath.Dir(imageURL.Path), zedbootPath), "/",
	)

	// Get reader to GCS object.
	r, err := bkt.Object(gcsZedbootPath).NewReader(ctx)
	if err != nil {
	return err
	}
	defer r.Close()

	// Open a local zedboot file and download remote data into it.
	file, err := os.Create(zedbootPath)
	if err != nil {
	return err
	}
	defer file.Close()

	if _, err := io.Copy(file, r); err != nil {
	return err
	}
	return nil
	}

	func execute(ctx context.Context, subcommandArgs []string) (int, error) {
	// If this is a QEMU test bench, skip the device setup and just run the subprocess.
	if deviceConfigPath == "" {
	return runSubprocess(ctx, subcommandArgs), nil
	}
	// Contains all necessary bootserver flags except device nodename.
	bootserverCmdStub := []string{
	bootserverPath,
	"--images", imagesManifest,
	"--mode", "pave-zedboot",
	"-n",
	}

	// Create devicePkg.Devices for each of the devices in the config file
	configs, err := devicePkg.LoadDeviceConfigs(deviceConfigPath)
	if err != nil {
	return failureExitCode, err
	}
	devices, err := devicePkg.CreateDevices(ctx, configs, bootserverCmdStub)
	if err != nil {
	return failureExitCode, err
	}

	// If the image manifest provided is a gcs url and we have a NUC, we need to download
	// zedboot. If the image manifest is a file path, then zedboot.zbi should exist in this
	// directory.
	imageURL, err := url.Parse(imagesManifest)
	if err != nil {
	return failureExitCode, err
	}
	if imageURL.Scheme == "gs" {
	// This is a no-op if no NUCs exist on this testbed.
	if err := downloadZedboot(ctx, imageURL, devices); err != nil {
	return failureExitCode, err
	}
	}

	// Set up a NUC server. If the port environment variable is not set, or if
	// the disableRebootServer flag is set, this is a no-op.
	srv, err := runNUCServer(ctx, devices, disableRebootServer)
	if err != nil {
	return failureExitCode, err
	}

	// Clean up after tests
	defer func() {
	if srv != nil {
	if err := srv.Shutdown(ctx); err != nil {
	log.Printf("NUC server shutdown failed: %v", err)
	}
	}
	}()

	if err := prepDevices(ctx, devices); err != nil {
	return failureExitCode, err
	}

	time.Sleep(rebootDuration)
	// Execute the passed in subcommand
	return runSubprocess(ctx, subcommandArgs), nil
	}

	func main() {
	// Initialize
	flag.Parse()

	// Handle SIGTERM
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	signals := make(chan os.Signal)
	defer func() {
	signal.Stop(signals)
	close(signals)
	}()
	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)

	go func() {
	select {
	case <-signals:
	cancel()
	case <-ctx.Done():
	}
	}()

	// Resolve environment variables within the subcommand's command line, as
	// Swarming does not do it automatically, and best to have this done once and
	// uniformly across all possible test task commands.
	var subcmdArgs []string
	for _, arg := range flag.Args() {
	subcmdArgs = append(subcmdArgs, os.ExpandEnv(arg))
	}
	exitCode, err := execute(ctx, subcmdArgs)

	if err != nil {
	log.Printf("Exit code: %d, Err: %s\n", exitCode, err)
	}
	os.Exit(exitCode)
	}