Inject the current binary into the chroot in test mode

client/allocdir/alloc_dir.go

@@ -108,37 +108,53 @@ func (d *AllocDir) Build(tasks []*structs.Task) error {
 	return nil
 }
 
-// Embed takes a mapping of absolute directory paths on the host to their
-// intended, relative location within the task directory. Embed attempts
+// Embed takes a mapping of absolute directory or file paths on the host to
+// their intended, relative location within the task directory. Embed attempts
 // hardlink and then defaults to copying. If the path exists on the host and
 // can't be embeded an error is returned.
-func (d *AllocDir) Embed(task string, dirs map[string]string) error {
+func (d *AllocDir) Embed(task string, entries map[string]string) error {
 	taskdir, ok := d.TaskDirs[task]
 	if !ok {
 		return fmt.Errorf("Task directory doesn't exist for task %v", task)
 	}
 
 	subdirs := make(map[string]string)
-	for source, dest := range dirs {
+	for source, dest := range entries {
 		// Check to see if directory exists on host.
 		s, err := os.Stat(source)
 		if os.IsNotExist(err) {
 			continue
 		}
 
+		// Embedding a single file
+		if !s.IsDir() {
+			destDir := filepath.Join(taskdir, filepath.Dir(dest))
+			if err := os.MkdirAll(destDir, s.Mode().Perm()); err != nil {
+				return fmt.Errorf("Couldn't create destination directory %v: %v", destDir, err)
+			}
+
+			// Copy the file.
+			taskEntry := filepath.Join(destDir, filepath.Base(dest))
+			if err := d.linkOrCopy(source, taskEntry, s.Mode().Perm()); err != nil {
+				return err
+			}
+
+			continue
+		}
+
 		// Create destination directory.
 		destDir := filepath.Join(taskdir, dest)
 		if err := os.MkdirAll(destDir, s.Mode().Perm()); err != nil {
 			return fmt.Errorf("Couldn't create destination directory %v: %v", destDir, err)
 		}
 
 		// Enumerate the files in source.
-		entries, err := ioutil.ReadDir(source)
+		dirEntries, err := ioutil.ReadDir(source)
 		if err != nil {
 			return fmt.Errorf("Couldn't read directory %v: %v", source, err)
 		}
 
-		for _, entry := range entries {
+		for _, entry := range dirEntries {
 			hostEntry := filepath.Join(source, entry.Name())
 			taskEntry := filepath.Join(destDir, filepath.Base(hostEntry))
 			if entry.IsDir() {

client/driver/executor/exec.go

@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"os/exec"
 	"path/filepath"
+	"strings"
 
 	"github.com/hashicorp/nomad/client/allocdir"
 	"github.com/hashicorp/nomad/nomad/structs"
@@ -33,6 +34,11 @@ import (
 
 var errNoResources = fmt.Errorf("No resources are associated with this task")
 
+// If testModeEnvVar is set in a process's environment variables, the
+// LinuxExecutor will detect it and inject the current binary into the chroot.
+// This enables using the test binary in tests to provide portability.
+var testModeEnvVar = "NOMAD_EXECUTOR_TEST_ONLY_13871827980214"
+
 // Executor is an interface that any platform- or capability-specific exec
 // wrapper must implement. You should not need to implement a Java executor.
 // Rather, you would implement a cgroups executor that the Java driver will use.
@@ -106,3 +112,18 @@ func OpenId(id string) (Executor, error) {
 	}
 	return executor, nil
 }
+
+// isTest returns whether the cmd is a test binary.
+func isTest(cmd *exec.Cmd) bool {
+	if cmd == nil {
+		return false
+	}
+
+	for _, env := range cmd.Env {
+		if strings.HasPrefix(env, testModeEnvVar) {
+			return true
+		}
+	}
+
+	return false
+}

client/driver/executor/exec_linux.go

@@ -252,6 +252,13 @@ func (e *LinuxExecutor) ConfigureTaskDir(taskName string, alloc *allocdir.AllocD
 		return err
 	}
 
+	// Embed ourselves if this is a test. This needs to be done so the test
+	// binary is inside the chroot.
+	if isTest(&e.cmd) {
+		bin := e.cmd.Args[0]
+		alloc.Embed(taskName, map[string]string{bin: bin})
+	}
+
 	if err := alloc.Embed(taskName, chrootEnv); err != nil {
 		return err
 	}

client/driver/executor/test_harness_test.go

@@ -17,27 +17,35 @@ import (
 )
 
 // testBinary is the path to the running test binary
-var testBinary = os.Args[0]
+var testBinary = func() string {
+	abs, err := filepath.Abs(os.Args[0])
+	if err != nil {
+		return err.Error()
+	}
+
+	return abs
+}()
 
 func TestMain(m *testing.M) {
 	// The tests in this package recursively execute the test binary produced
 	// by go test. The TEST_MAIN environment variable controls the recursive
 	// execution.
-	switch tm := os.Getenv("TEST_MAIN"); tm {
+	switch tm := os.Getenv(testModeEnvVar); tm {
 	case "":
 		os.Exit(m.Run())
 	case "app":
 		appMain()
 	default:
-		fmt.Fprintf(os.Stderr, "unexpected value for TEST_MAIN, \"%s\"\n", tm)
+		fmt.Fprintf(os.Stderr,
+			"Unexpected value for test mode environment variable, %q\n", tm)
 		os.Exit(1)
 	}
 }
 
 // setTestAppEnv sets the environement of cmd for a recursive call into
 // TestMain.
 func setTestAppEnv(cmd *exec.Cmd) {
-	cmd.Env = append(os.Environ(), "TEST_MAIN=app")
+	cmd.Env = append(os.Environ(), fmt.Sprintf("%v=app", testModeEnvVar))
 }
 
 func appMain() {