agent: implement initial preflight check command

cmd/agent/internal/cmd/cmd.go

@@ -26,6 +26,7 @@ func Run() {
 
 	root.AddCommand(
 		newCmdStart(cmdCtx),
+		newCmdPreflight(cmdCtx),
 		newCmdDaemon(cmdCtx),
 		newCmdReset(cmdCtx),
 		newCmdVersion(),

cmd/agent/internal/cmd/preflight.go

@@ -0,0 +1,189 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package cmd
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/spf13/cobra"
+
+	"github.com/Azure/unbounded/internal/provision"
+	"github.com/Azure/unbounded/pkg/agent/goalstates"
+	"github.com/Azure/unbounded/pkg/agent/phases/host"
+	"github.com/Azure/unbounded/pkg/agent/phases/nodestart"
+	"github.com/Azure/unbounded/pkg/agent/phases/rootfs"
+	"github.com/Azure/unbounded/pkg/agent/preflight"
+)
+
+type preflightHandler struct {
+	cmdCtx                *CommandContext
+	configPath            string
+	ignorePreflightErrors []string
+	failOnWarnings        bool
+	output                string
+	writer                io.Writer
+}
+
+func newCmdPreflight(cmdCtx *CommandContext) *cobra.Command {
+	handler := &preflightHandler{cmdCtx: cmdCtx, writer: os.Stdout}
+
+	cmd := &cobra.Command{
+		Use:   "preflight",
+		Short: "Run non-mutating preflight checks",
+		Long:  "Run non-mutating preflight checks for the host and agent configuration before node bootstrap.",
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			return handler.execute(cmd.Context())
+		},
+	}
+
+	cmd.Flags().StringVar(&handler.configPath, "config", "", "Path to agent config file")
+	cmd.Flags().StringSliceVar(
+		&handler.ignorePreflightErrors,
+		"ignore-preflight-errors",
+		nil,
+		"Comma-separated preflight check names whose errors should be reported as warnings",
+	)
+	cmd.Flags().BoolVar(&handler.failOnWarnings, "fail-on-warnings", false, "Fail when any preflight warning is returned")
+	cmd.Flags().StringVar(&handler.output, "output", "text", "Output format: text or json")
+
+	return cmd
+}
+
+func (h *preflightHandler) execute(ctx context.Context) error {
+	h.cmdCtx.Setup()
+	logger := h.cmdCtx.Logger
+
+	if h.configPath != "" {
+		oldConfigPath := os.Getenv(configFileEnv)
+		defer os.Setenv(configFileEnv, oldConfigPath) //nolint:errcheck // best effort restore
+
+		if err := os.Setenv(configFileEnv, h.configPath); err != nil {
+			return err
+		}
+	}
+
+	cfg, err := loadConfig(logger)
+	if err != nil {
+		return err
+	}
+
+	if err := cfg.Validate(); err != nil {
+		return fmt.Errorf("validate agent config: %w", err)
+	}
+
+	goalState, err := goalstates.ResolveMachine(
+		logger,
+		&cfg.AgentConfig,
+		goalstates.NSpawnMachineKube1,
+		provision.ResolveDownloadOverrides(cfg.Downloads),
+	)
+	if err != nil {
+		return fmt.Errorf("resolve goal state: %w", err)
+	}
+
+	checks := preflight.Flatten(
+		host.Preflight(logger, cfg, goalState),
+		nodestart.Preflight(logger, cfg, goalState),
+		rootfs.Preflight(logger, cfg, goalState),
+	)
+
+	opts := preflight.Options{
+		IgnoreErrors:   h.ignorePreflightErrors,
+		FailOnWarnings: h.failOnWarnings,
+	}
+	report := preflight.Run(ctx, checks, opts)
+
+	switch strings.ToLower(h.output) {
+	case "", "text":
+		if err := writePreflightText(h.writer, report); err != nil {
+			return err
+		}
+	case "json":
+		enc := json.NewEncoder(h.writer)
+		enc.SetIndent("", "  ")
+
+		if err := enc.Encode(report); err != nil {
+			return err
+		}
+	default:
+		return fmt.Errorf("unsupported output format %q", h.output)
+	}
+
+	return report.Err(h.failOnWarnings)
+}
+
+func writePreflightText(w io.Writer, report preflight.Report) error {
+	if _, err := fmt.Fprintln(w, "[preflight] Running unbounded-agent pre-flight checks"); err != nil {
+		return err
+	}
+
+	var errors []preflight.Result
+
+	for _, result := range report.Checks {
+		switch result.Severity {
+		case preflight.SeverityOK:
+			if err := writePreflightResult(w, "OK", result); err != nil {
+				return err
+			}
+		case preflight.SeverityError:
+			errors = append(errors, result)
+		case preflight.SeverityWarning:
+			if err := writePreflightResult(w, "WARNING", result); err != nil {
+				return err
+			}
+		}
+	}
+
+	if len(errors) == 0 {
+		return nil
+	}
+
+	if _, err := fmt.Fprintln(w, "[preflight] Some fatal errors occurred:"); err != nil {
+		return err
+	}
+
+	for _, result := range errors {
+		if _, err := fmt.Fprintf(w, "\t[ERROR %s]: %s", result.Name, result.Message); err != nil {
+			return err
+		}
+
+		if result.Target != "" {
+			if _, err := fmt.Fprintf(w, " (target: %s)", result.Target); err != nil {
+				return err
+			}
+		}
+
+		if _, err := fmt.Fprintln(w); err != nil {
+			return err
+		}
+	}
+
+	_, err := fmt.Fprintln(
+		w,
+		"[preflight] If you know what you are doing, you can make a check non-fatal with `--ignore-preflight-errors=...`",
+	)
+
+	return err
+}
+
+func writePreflightResult(w io.Writer, status string, result preflight.Result) error {
+	if _, err := fmt.Fprintf(w, "\t[%s %s]: %s", status, result.Name, result.Message); err != nil {
+		return err
+	}
+
+	if result.Target != "" {
+		if _, err := fmt.Fprintf(w, " (target: %s)", result.Target); err != nil {
+			return err
+		}
+	}
+
+	_, err := fmt.Fprintln(w)
+
+	return err
+}

cmd/agent/internal/cmd/preflight_test.go

@@ -0,0 +1,108 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package cmd
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/Azure/unbounded/internal/provision"
+)
+
+func preflightConfig(apiServer string) provision.UnboundedAgentConfig {
+	cfg := sampleConfig()
+	cfg.Kubelet.ApiServer = apiServer
+	cfg.OCIImage = "registry.example.com/unbounded/rootfs:v1"
+
+	return cfg
+}
+
+func TestNewCmdPreflight(t *testing.T) {
+	cmd := newCmdPreflight(&CommandContext{LogFormat: "text"})
+
+	assert.Equal(t, "preflight", cmd.Use)
+	assert.NotNil(t, cmd.Flags().Lookup("config"))
+	assert.NotNil(t, cmd.Flags().Lookup("ignore-preflight-errors"))
+	assert.NotNil(t, cmd.Flags().Lookup("fail-on-warnings"))
+	assert.NotNil(t, cmd.Flags().Lookup("output"))
+}
+
+func TestPreflightJSONOutput(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+	t.Cleanup(srv.Close)
+
+	path := writeConfigFile(t, preflightConfig(srv.URL))
+
+	var out bytes.Buffer
+
+	h := &preflightHandler{
+		cmdCtx:                &CommandContext{LogFormat: "text"},
+		configPath:            path,
+		ignorePreflightErrors: []string{"all"},
+		output:                "json",
+		writer:                &out,
+	}
+
+	require.NoError(t, h.execute(context.Background()))
+
+	var report struct {
+		Status string `json:"status"`
+		Checks []struct {
+			Name    string `json:"name"`
+			Ignored bool   `json:"ignored"`
+		} `json:"checks"`
+	}
+	require.NoError(t, json.Unmarshal(out.Bytes(), &report))
+	assert.Equal(t, "ok", report.Status)
+	assert.NotEmpty(t, report.Checks)
+}
+
+func TestPreflightTextOutputError(t *testing.T) {
+	path := writeConfigFile(t, preflightConfig("https://127.0.0.1:1"))
+
+	var out bytes.Buffer
+
+	h := &preflightHandler{
+		cmdCtx:     &CommandContext{LogFormat: "text"},
+		configPath: path,
+		output:     "text",
+		writer:     &out,
+	}
+
+	err := h.execute(context.Background())
+	require.Error(t, err)
+	assert.Contains(t, out.String(), "[ERROR api-server-reachable]")
+	assert.NotContains(t, out.String(), "127.0.0.1")
+}
+
+func TestPreflightTextOutputIncludesOK(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+	t.Cleanup(srv.Close)
+
+	path := writeConfigFile(t, preflightConfig(srv.URL))
+
+	var out bytes.Buffer
+
+	h := &preflightHandler{
+		cmdCtx:                &CommandContext{LogFormat: "text"},
+		configPath:            path,
+		ignorePreflightErrors: []string{"all"},
+		output:                "text",
+		writer:                &out,
+	}
+
+	require.NoError(t, h.execute(context.Background()))
+	assert.Contains(t, out.String(), "[OK goal-state]")
+}

hack/agent/e2e-kind/e2e.py

@@ -213,6 +213,10 @@ def scp_cmd(src: str, dst: str) -> subprocess.CompletedProcess[str]:
     return run(["scp", *SSH_OPTS, src, dst])
 
 
+def scp_from_vm(src: str, dst: Path) -> subprocess.CompletedProcess[str]:
+    return run(["scp", *SSH_OPTS, f"{SSH_TARGET}:{src}", str(dst)])
+
+
 def kubectl(args: list[str], **kw: Any) -> subprocess.CompletedProcess[str]:
     return run([KUBECTL, *args], **kw)
 
@@ -1639,6 +1643,21 @@ def _run_agent_inner(agent_url: str, node_config: NodeConfig) -> None:
     bootstrap_script_path.chmod(0o600)
     log(f"Bootstrap script written to {bootstrap_script_path}")
 
+    bootstrap_preflight_script = bootstrap_script.replace(
+        'echo "Running unbounded-agent start..."\n"${AGENT_BIN}" start ${_START_ARGS}',
+        textwrap.dedent("""\
+        echo "Running unbounded-agent preflight..."
+        "${AGENT_BIN}" preflight ${_START_ARGS} --output text | tee /tmp/unbounded-agent-preflight.txt
+        "${AGENT_BIN}" preflight ${_START_ARGS} --output json > /tmp/unbounded-agent-preflight.json
+        echo "Running unbounded-agent start..."
+        "${AGENT_BIN}" start ${_START_ARGS}"""),
+    )
+    if bootstrap_preflight_script == bootstrap_script:
+        die("failed to inject unbounded-agent preflight into bootstrap script")
+
+    bootstrap_script = bootstrap_preflight_script
+    bootstrap_script_path.write_text(bootstrap_script)
+
     # Wait for cloud-init and verify connectivity
     log("Waiting for cloud-init to complete on VM...")
     subprocess.run(["ssh", *SSH_OPTS, SSH_TARGET, "sudo cloud-init status --wait"],
@@ -1664,6 +1683,10 @@ def _run_agent_inner(agent_url: str, node_config: NodeConfig) -> None:
         f"sudo {env_prefix} /tmp/bootstrap.sh",
     ])
 
+    log("Copying preflight reports from VM...")
+    scp_from_vm("/tmp/unbounded-agent-preflight.txt", VM_DIR / "unbounded-agent-preflight.txt")
+    scp_from_vm("/tmp/unbounded-agent-preflight.json", VM_DIR / "unbounded-agent-preflight.json")
+
 
 # ---------------------------------------------------------------------------
 # wait-for-node
@@ -2881,6 +2904,21 @@ def _collect_one_vm_logs(logs_dir: Path, vm_name: str, vm_ip: str, vm_dir: Path,
     ]
     ssh_target = f"{VM_SSH_USER}@{vm_ip}"
 
+    for name in ("unbounded-agent-preflight.txt", "unbounded-agent-preflight.json"):
+        src = vm_dir / name
+        if src.exists():
+            shutil.copyfile(src, logs_dir / f"{prefix}{name}")
+
+    for name in ("unbounded-agent-preflight.txt", "unbounded-agent-preflight.json"):
+        result = subprocess.run(
+            ["scp", *ssh_opts, f"{ssh_target}:/tmp/{name}", str(logs_dir / f"{prefix}{name}")],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            check=False,
+        )
+        if result.returncode == 0:
+            diag(f"Collected {name} from VM")
+
     def ssh_log(name: str, command: str) -> None:
         _write_command_log(logs_dir / f"{prefix}{name}", ["ssh", *ssh_opts, ssh_target, command])