cmd/catalyst/main.go

// 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"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	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"
)

var (
	imagesManifest   string
	deviceConfigPath string
	bootserverPath   string
)

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")
}

func rebootDevices(ctx context.Context, devices []*devicePkg.DeviceTarget) {
	errs := make(chan error)
	defer close(errs)
	for _, device := range devices {
		go func() {
			errs <- device.Restart(ctx)
		}()
	}
	for i := 0; i < len(devices); i++ {
		if err := <-errs; err != nil {
			log.Printf("%s\n", err)
		}
	}
}

// Runs a subprocess and sets up a handler that propagates SIGTERM on context cancel.
func runSubprocess(ctx context.Context, command []string) int {
	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()
}

func runNUCServer(ctx context.Context, devices []*devicePkg.DeviceTarget) (*http.Server, error) {
	// Parse the port from the environment variable.
	portStr := os.Getenv(portEnvVar)
	if portStr == "" {
		return nil, nil
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, err
	}

	mux := http.NewServeMux()
	for _, device := range devices {
		endpoint := fmt.Sprintf("/%s.ipxe", device.Mac())
		mux.HandleFunc(endpoint, 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)
		})
	}

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

func runBootservers(ctx context.Context, devices []*devicePkg.DeviceTarget) int {
	// Execute bootserver for each node
	exitCodes := make(chan int)
	for _, device := range devices {
		go func(device *devicePkg.DeviceTarget) {
			exitCodes <- runSubprocess(ctx, device.BootserverCmd)
		}(device)
	}

	// Wait for all of the bootservers to finish running and ensure success
	numErrs := 0
	for i := 0; i < len(devices); i++ {
		if exitCode := <-exitCodes; exitCode != 0 {
			log.Printf("bootserver exited with exit code: %d\n", exitCode)
			numErrs += 1
		}
	}

	if numErrs > 0 {
		return failureExitCode
	}
	return 0
}

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.DeviceTargets for each of the devices in the config file
	devices, err := devicePkg.CreateDeviceTargets(ctx, deviceConfigPath, bootserverCmdStub)
	if err != nil {
		return failureExitCode, err
	}

	// Set up a NUC server. If the port environment variable is not set, this is a no-op.
	srv, err := runNUCServer(ctx, devices)
	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)
			}
		}
		rebootDevices(ctx, devices)
	}()

	if exitCode := runBootservers(ctx, devices); exitCode != 0 {
		return exitCode, nil
	}

	// 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():
		}
	}()

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