Refactor spawn-daemon so it can be used by all OSes and make it write exit code to a file

command/spawn_daemon.go

@@ -2,19 +2,19 @@ package command
 
 import (
 	"encoding/json"
+	"fmt"
+	"io"
 	"os"
+	"os/exec"
+	"strconv"
 	"strings"
+	"syscall"
 )
 
 type SpawnDaemonCommand struct {
 	Meta
-}
-
-// Status of executing the user's command.
-type SpawnStartStatus struct {
-	// ErrorMsg will be empty if the user command was started successfully.
-	// Otherwise it will have an error message.
-	ErrorMsg string
+	config   *DaemonConfig
+	exitFile io.WriteCloser
 }
 
 func (c *SpawnDaemonCommand) Help() string {
@@ -23,15 +23,15 @@ Usage: nomad spawn-daemon [options] <daemon_config>
 
   INTERNAL ONLY
 
-  Spawns a daemon process optionally inside a cgroup. The required daemon_config is a json
-  encoding of the DaemonConfig struct containing the isolation configuration and command to run.
-  SpawnStartStatus is json serialized to Stdout upon running the user command or if any error
-  prevents its execution. If there is no error, the process waits on the users
-  command and then json serializes  SpawnExitStatus to Stdout after its termination.
-
-General Options:
-
-  ` + generalOptionsUsage()
+  Spawns a daemon process by double forking. The required daemon_config is a
+  json encoding of the DaemonConfig struct containing the isolation
+  configuration and command to run. SpawnStartStatus is json serialized to
+  stdout upon running the user command or if any error prevents its execution.
+  If there is no error, the process waits on the users command. Once the user
+  command exits, the exit code is written to a file specified in the
+  daemon_config and this process exits with the same exit status as the user
+  command.
+  `
 
 	return strings.TrimSpace(helpText)
 }
@@ -40,6 +40,147 @@ func (c *SpawnDaemonCommand) Synopsis() string {
 	return "Spawn a daemon command with configurable isolation."
 }
 
+// Status of executing the user's command.
+type SpawnStartStatus struct {
+	// The PID of the user's command.
+	UserPID int
+
+	// ErrorMsg will be empty if the user command was started successfully.
+	// Otherwise it will have an error message.
+	ErrorMsg string
+}
+
+// Exit status of the user's command.
+type SpawnExitStatus struct {
+	// The exit code of the user's command.
+	ExitCode int
+}
+
+// Configuration for the command to start as a daemon.
+type DaemonConfig struct {
+	exec.Cmd
+
+	// The filepath to write the exit status to.
+	ExitStatusFile string
+
+	// The paths, if not /dev/null, must be either in the tasks root directory
+	// or in the shared alloc directory.
+	StdoutFile string
+	StdinFile  string
+	StderrFile string
+
+	// An optional path specifying the directory to chroot the process in.
+	Chroot string
+}
+
+// Whether to start the user command or abort.
+type TaskStart bool
+
+// parseConfig reads the DaemonConfig from the passed arguments. If not
+// successful, an error is returned.
+func (c *SpawnDaemonCommand) parseConfig(args []string) (*DaemonConfig, error) {
+	flags := c.Meta.FlagSet("spawn-daemon", FlagSetClient)
+	flags.Usage = func() { c.Ui.Output(c.Help()) }
+	if err := flags.Parse(args); err != nil {
+		return nil, fmt.Errorf("failed to parse args: %v", err)
+	}
+
+	// Check that we got json input.
+	args = flags.Args()
+	if len(args) != 1 {
+		return nil, fmt.Errorf("incorrect number of args; got %v; want 1", len(args))
+	}
+	jsonInput, err := strconv.Unquote(args[0])
+	if err != nil {
+		return nil, fmt.Errorf("Failed to unquote json input: %v", err)
+	}
+
+	// De-serialize the passed command.
+	var config DaemonConfig
+	dec := json.NewDecoder(strings.NewReader(jsonInput))
+	if err := dec.Decode(&config); err != nil {
+		return nil, err
+	}
+
+	return &config, nil
+}
+
+// configureLogs creates the log files and redirects the process
+// stdin/stderr/stdout to them. If unsuccessful, an error is returned.
+func (c *SpawnDaemonCommand) configureLogs() error {
+	stdo, err := os.OpenFile(c.config.StdoutFile, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
+	if err != nil {
+		return fmt.Errorf("Error opening file to redirect stdout: %v", err)
+	}
+
+	stde, err := os.OpenFile(c.config.StderrFile, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
+	if err != nil {
+		return fmt.Errorf("Error opening file to redirect stderr: %v", err)
+	}
+
+	stdi, err := os.OpenFile(c.config.StdinFile, os.O_CREATE|os.O_RDONLY, 0666)
+	if err != nil {
+		return fmt.Errorf("Error opening file to redirect stdin: %v", err)
+	}
+
+	c.config.Cmd.Stdout = stdo
+	c.config.Cmd.Stderr = stde
+	c.config.Cmd.Stdin = stdi
+	return nil
+}
+
+func (c *SpawnDaemonCommand) Run(args []string) int {
+	var err error
+	c.config, err = c.parseConfig(args)
+	if err != nil {
+		return c.outputStartStatus(err, 1)
+	}
+
+	// Open the file we will be using to write exit codes to. We do this early
+	// to ensure that we don't start the user process when we can't capture its
+	// exit status.
+	c.exitFile, err = os.OpenFile(c.config.ExitStatusFile, os.O_CREATE|os.O_RDWR, 0666)
+	if err != nil {
+		return c.outputStartStatus(fmt.Errorf("Error opening file to store exit status: %v", err), 1)
+	}
+
+	// Isolate the user process.
+	if err := c.isolateCmd(); err != nil {
+		return c.outputStartStatus(err, 1)
+	}
+
+	// Redirect logs.
+	if err := c.configureLogs(); err != nil {
+		return c.outputStartStatus(err, 1)
+	}
+
+	// Chroot jail the process and set its working directory.
+	c.configureChroot()
+
+	// Wait to get the start command.
+	var start TaskStart
+	dec := json.NewDecoder(os.Stdin)
+	if err := dec.Decode(&start); err != nil {
+		return c.outputStartStatus(err, 1)
+	}
+
+	// Aborted by Nomad process.
+	if !start {
+		return 0
+	}
+
+	// Spawn the user process.
+	if err := c.config.Cmd.Start(); err != nil {
+		return c.outputStartStatus(fmt.Errorf("Error starting user command: %v", err), 1)
+	}
+
+	// Indicate that the command was started successfully.
+	c.outputStartStatus(nil, 0)
+
+	// Wait and then output the exit status.
+	return c.writeExitStatus(c.config.Cmd.Wait())
+}
+
 // outputStartStatus is a helper function that outputs a SpawnStartStatus to
 // Stdout with the passed error, which may be nil to indicate no error. It
 // returns the passed status.
@@ -51,6 +192,36 @@ func (c *SpawnDaemonCommand) outputStartStatus(err error, status int) int {
 		startStatus.ErrorMsg = err.Error()
 	}
 
+	if c.config != nil && c.config.Process == nil {
+		startStatus.UserPID = c.config.Process.Pid
+	}
+
 	enc.Encode(startStatus)
 	return status
 }
+
+// writeExitStatus takes in the error result from calling wait and writes out
+// the exit status to a file. It returns the same exit status as the user
+// command.
+func (c *SpawnDaemonCommand) writeExitStatus(exit error) int {
+	// Parse the exit code.
+	exitStatus := &SpawnExitStatus{}
+	if exit != nil {
+		// Default to exit code 1 if we can not get the actual exit code.
+		exitStatus.ExitCode = 1
+
+		if exiterr, ok := exit.(*exec.ExitError); ok {
+			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
+				exitStatus.ExitCode = status.ExitStatus()
+			}
+		}
+	}
+
+	if c.exitFile != nil {
+		enc := json.NewEncoder(c.exitFile)
+		enc.Encode(exitStatus)
+		c.exitFile.Close()
+	}
+
+	return exitStatus.ExitCode
+}

command/spawn_daemon_darwin.go

@@ -0,0 +1,4 @@
+package command
+
+// No chroot on darwin.
+func (c *SpawnDaemonCommand) configureChroot() {}

command/spawn_daemon_linux.go

@@ -1,115 +1,16 @@
 package command
 
-import (
-	"encoding/json"
-	"fmt"
-	"os"
-	"os/exec"
-	"strconv"
-	"strings"
-	"syscall"
-)
+import "syscall"
 
-// Configuration for the command to start as a daemon.
-type DaemonConfig struct {
-	exec.Cmd
+// configureChroot enters the user command into a chroot if specified in the
+// config and on an OS that supports Chroots.
+func (c *SpawnDaemonCommand) configureChroot() {
+	if len(c.config.Chroot) != 0 {
+		if c.config.Cmd.SysProcAttr == nil {
+			c.config.Cmd.SysProcAttr = &syscall.SysProcAttr{}
+		}
 
-	// The paths, if not /dev/null, must be either in the tasks root directory
-	// or in the shared alloc directory.
-	StdoutFile string
-	StdinFile  string
-	StderrFile string
-
-	Chroot string
-}
-
-// Whether to start the user command or abort.
-type TaskStart bool
-
-func (c *SpawnDaemonCommand) Run(args []string) int {
-	flags := c.Meta.FlagSet("spawn-daemon", FlagSetClient)
-	flags.Usage = func() { c.Ui.Output(c.Help()) }
-
-	if err := flags.Parse(args); err != nil {
-		return 1
-	}
-
-	// Check that we got json input.
-	args = flags.Args()
-	if len(args) != 1 {
-		c.Ui.Error(c.Help())
-		return 1
-	}
-	jsonInput, err := strconv.Unquote(args[0])
-	if err != nil {
-		return c.outputStartStatus(fmt.Errorf("Failed to unquote json input: %v", err), 1)
-	}
-
-	// De-serialize the passed command.
-	var cmd DaemonConfig
-	dec := json.NewDecoder(strings.NewReader(jsonInput))
-	if err := dec.Decode(&cmd); err != nil {
-		return c.outputStartStatus(err, 1)
-	}
-
-	// Isolate the user process.
-	if _, err := syscall.Setsid(); err != nil {
-		return c.outputStartStatus(fmt.Errorf("Failed setting sid: %v", err), 1)
-	}
-
-	syscall.Umask(0)
-
-	// Redirect logs.
-	stdo, err := os.OpenFile(cmd.StdoutFile, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
-	if err != nil {
-		return c.outputStartStatus(fmt.Errorf("Error opening file to redirect Stdout: %v", err), 1)
-	}
-
-	stde, err := os.OpenFile(cmd.StderrFile, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
-	if err != nil {
-		return c.outputStartStatus(fmt.Errorf("Error opening file to redirect Stderr: %v", err), 1)
-	}
-
-	stdi, err := os.OpenFile(cmd.StdinFile, os.O_CREATE|os.O_RDONLY, 0666)
-	if err != nil {
-		return c.outputStartStatus(fmt.Errorf("Error opening file to redirect Stdin: %v", err), 1)
-	}
-
-	cmd.Cmd.Stdout = stdo
-	cmd.Cmd.Stderr = stde
-	cmd.Cmd.Stdin = stdi
-
-	// Chroot jail the process and set its working directory.
-	if cmd.Cmd.SysProcAttr == nil {
-		cmd.Cmd.SysProcAttr = &syscall.SysProcAttr{}
-	}
-
-	cmd.Cmd.SysProcAttr.Chroot = cmd.Chroot
-	cmd.Cmd.Dir = "/"
-
-	// Wait to get the start command.
-	var start TaskStart
-	dec = json.NewDecoder(os.Stdin)
-	if err := dec.Decode(&start); err != nil {
-		return c.outputStartStatus(err, 1)
-	}
-
-	if !start {
-		return 0
-	}
-
-	// Spawn the user process.
-	if err := cmd.Cmd.Start(); err != nil {
-		return c.outputStartStatus(fmt.Errorf("Error starting user command: %v", err), 1)
-	}
-
-	// Indicate that the command was started successfully.
-	c.outputStartStatus(nil, 0)
-
-	// Wait and then output the exit status.
-	if err := cmd.Wait(); err != nil {
-		return 1
-	}
-
-	return 0
+		c.config.Cmd.SysProcAttr.Chroot = c.config.Chroot
+		c.config.Cmd.Dir = "/"
+	}
 }

command/spawn_daemon_test.go

@@ -0,0 +1,48 @@
+package command
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os/exec"
+	"testing"
+)
+
+type nopCloser struct {
+	io.ReadWriter
+}
+
+func (n *nopCloser) Close() error {
+	return nil
+}
+
+func TestSpawnDaemon_WriteExitStatus(t *testing.T) {
+	// Check if there is python.
+	path, err := exec.LookPath("python")
+	if err != nil {
+		t.Skip("python not detected")
+	}
+
+	var b bytes.Buffer
+	daemon := &SpawnDaemonCommand{exitFile: &nopCloser{&b}}
+
+	code := 3
+	cmd := exec.Command(path, "./test-resources/exiter.py", fmt.Sprintf("%d", code))
+	err = cmd.Run()
+	actual := daemon.writeExitStatus(err)
+	if actual != code {
+		t.Fatalf("writeExitStatus(%v) returned %v; want %v", err, actual, code)
+	}
+
+	// De-serialize the passed command.
+	var exitStatus SpawnExitStatus
+	dec := json.NewDecoder(&b)
+	if err := dec.Decode(&exitStatus); err != nil {
+		t.Fatalf("failed to decode exit status: %v", err)
+	}
+
+	if exitStatus.ExitCode != code {
+		t.Fatalf("writeExitStatus(%v) wrote exit status %v; want %v", err, exitStatus.ExitCode, code)
+	}
+}

command/spawn_daemon_universal.go

@@ -1,9 +0,0 @@
-// +build !linux
-
-package command
-
-import "errors"
-
-func (c *SpawnDaemonCommand) Run(args []string) int {
-	return c.outputStartStatus(errors.New("spawn-daemon not supported"), 1)
-}

command/spawn_daemon_unix.go

@@ -0,0 +1,16 @@
+// +build !windows
+
+package command
+
+import "syscall"
+
+// isolateCmd sets the session id for the process and the umask.
+func (c *SpawnDaemonCommand) isolateCmd() error {
+	if c.config.Cmd.SysProcAttr == nil {
+		c.config.Cmd.SysProcAttr = &syscall.SysProcAttr{}
+	}
+
+	c.config.Cmd.SysProcAttr.Setsid = true
+	syscall.Umask(0)
+	return nil
+}

command/spawn_daemon_windows.go

@@ -0,0 +1,7 @@
+// build !linux !darwin
+
+package command
+
+// No isolation on Windows.
+func (c *SpawnDaemonCommand) isolateCmd() error { return nil }
+func (c *SpawnDaemonCommand) configureChroot()  {}

command/test-resources/exiter.py

@@ -0,0 +1,3 @@
+import sys
+
+sys.exit(int(sys.argv[1]))