Linux Embed and AllocDir unit tests

client/allocdir/alloc_dir_test.go

@@ -0,0 +1,162 @@
+package allocdir
+
+import (
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/hashicorp/nomad/nomad/structs"
+)
+
+var (
+	t1 = &structs.Task{
+		Name:   "web",
+		Driver: "exec",
+		Config: map[string]string{
+			"command": "/bin/date",
+			"args":    "+%s",
+		},
+		Resources: &structs.Resources{
+			DiskMB: 1,
+		},
+	}
+
+	t2 = &structs.Task{
+		Name:   "web2",
+		Driver: "exec",
+		Config: map[string]string{
+			"command": "/bin/date",
+			"args":    "+%s",
+		},
+		Resources: &structs.Resources{
+			DiskMB: 1,
+		},
+	}
+)
+
+// Test that given a set of tasks, each task gets a directory and that directory
+// has the shared alloc dir inside of it.
+func TestAllocDir_BuildAlloc(t *testing.T) {
+	tmp, err := ioutil.TempDir("", "AllocDir")
+	if err != nil {
+		t.Fatalf("Couldn't create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tmp)
+
+	d := NewAllocDir(tmp)
+	tasks := []*structs.Task{t1, t2}
+	if err := d.Build(tasks); err != nil {
+		t.Fatalf("Build(%v) failed: %v", tasks, err)
+	}
+
+	// Check that the AllocDir and each of the task directories exist.
+	if _, err := os.Stat(d.AllocDir); os.IsNotExist(err) {
+		t.Fatalf("Build(%v) didn't create AllocDir %v", tasks, d.AllocDir)
+	}
+
+	// Create a file in the alloc dir and then check it exists in each of the
+	// task dirs.
+	allocFile := "foo"
+	allocFileData := []byte{'b', 'a', 'r'}
+	if err := ioutil.WriteFile(filepath.Join(d.AllocDir, allocFile), allocFileData, 0777); err != nil {
+		t.Fatalf("Couldn't create file in alloc dir: %v", err)
+	}
+
+	for _, task := range tasks {
+		tDir, err := d.TaskDir(task.Name)
+		if err != nil {
+			t.Fatalf("TaskDir(%v) failed: %v", task.Name, err)
+		}
+
+		if _, err := os.Stat(tDir); os.IsNotExist(err) {
+			t.Fatalf("Build(%v) didn't create TaskDir %v", tasks, tDir)
+		}
+
+		// TODO: Enable once mount is done.
+		//allocExpected := filepath.Join(tDir, SharedAllocName, allocFile)
+		//if _, err := os.Stat(allocExpected); os.IsNotExist(err) {
+		//t.Fatalf("File in shared alloc dir not accessible from task dir %v: %v", tDir, err)
+		//}
+	}
+}
+
+func TestAllocDir_EmbedNonExistent(t *testing.T) {
+	tmp, err := ioutil.TempDir("", "AllocDir")
+	if err != nil {
+		t.Fatalf("Couldn't create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tmp)
+
+	d := NewAllocDir(tmp)
+	tasks := []*structs.Task{t1, t2}
+	if err := d.Build(tasks); err != nil {
+		t.Fatalf("Build(%v) failed: %v", tasks, err)
+	}
+
+	fakeDir := "/foobarbaz"
+	task := tasks[0].Name
+	mapping := map[string]string{fakeDir: fakeDir}
+	if err := d.Embed(task, mapping); err == nil {
+		t.Fatalf("Embed(%v, %v) should have failed. %v does not exist", task, mapping, fakeDir)
+	}
+}
+
+func TestAllocDir_EmbedDirs(t *testing.T) {
+	tmp, err := ioutil.TempDir("", "AllocDir")
+	if err != nil {
+		t.Fatalf("Couldn't create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tmp)
+
+	d := NewAllocDir(tmp)
+	tasks := []*structs.Task{t1, t2}
+	if err := d.Build(tasks); err != nil {
+		t.Fatalf("Build(%v) failed: %v", tasks, err)
+	}
+
+	// Create a fake host directory, with a file, and a subfolder that contains
+	// a file.
+	host, err := ioutil.TempDir("", "AllocDirHost")
+	if err != nil {
+		t.Fatalf("Couldn't create temp dir: %v", err)
+	}
+	defer os.RemoveAll(host)
+
+	subDirName := "subdir"
+	subDir := filepath.Join(host, subDirName)
+	if err := os.Mkdir(subDir, 0777); err != nil {
+		t.Fatalf("Failed to make subdir %v: %v", subDir, err)
+	}
+
+	file := "foo"
+	subFile := "bar"
+	if err := ioutil.WriteFile(filepath.Join(host, file), []byte{'a'}, 0777); err != nil {
+		t.Fatalf("Coudn't create file in host dir %v: %v", host, err)
+	}
+
+	if err := ioutil.WriteFile(filepath.Join(subDir, subFile), []byte{'a'}, 0777); err != nil {
+		t.Fatalf("Coudn't create file in host subdir %v: %v", subDir, err)
+	}
+
+	// Create mapping from host dir to task dir.
+	task := tasks[0].Name
+	taskDest := "bin/test/"
+	mapping := map[string]string{host: taskDest}
+	if err := d.Embed(task, mapping); err != nil {
+		t.Fatalf("Embed(%v, %v) failed: %v", task, mapping, err)
+	}
+
+	// Check that the embedding was done properly.
+	taskDir, err := d.TaskDir(task)
+	if err != nil {
+		t.Fatalf("TaskDir(%v) failed: %v", task, err)
+	}
+
+	exp := []string{filepath.Join(taskDir, taskDest, file), filepath.Join(taskDir, taskDest, subDirName, subFile)}
+	for _, e := range exp {
+		if _, err := os.Stat(e); os.IsNotExist(err) {
+			t.Fatalf("File %v not embeded: %v", e, err)
+		}
+	}
+}

client/allocdir/builder.go

@@ -2,42 +2,42 @@ package allocdir
 
 import (
 	"fmt"
-	"github.com/hashicorp/nomad/nomad/structs"
+	"io"
+	"io/ioutil"
 	"os"
 	"os/user"
 	"path/filepath"
 	"strconv"
+	"syscall"
+
+	"github.com/hashicorp/nomad/nomad/structs"
 )
 
 func (d *AllocDir) Build(tasks []*structs.Task) error {
 	// Make the alloc directory, owned by the nomad process.
-	if err := os.Mkdir(d.AllocDir, 0700); err != nil {
+	if err := os.MkdirAll(d.AllocDir, 0700); err != nil {
 		return fmt.Errorf("Failed to make the alloc directory %v: %v", d.AllocDir, err)
 	}
 
-	nobody, err := user.Lookup("nobody")
-	if err != nil {
-		return fmt.Errorf("Could not set owner/group on shared alloc directory: %v", err)
-	}
-
-	uid, err := getUid(nobody)
-	if err != nil {
-		return err
-	}
-
-	gid, err := getGid(nobody)
-	if err != nil {
-		return err
+	// Check if the process is has root capabilities and if so set the
+	// user/group to nobody.
+	var u *user.User
+	if syscall.Geteuid() == 0 {
+		nobody, err := user.Lookup("nobody")
+		if err != nil {
+			return fmt.Errorf("Could not set owner/group on shared alloc directory: %v", err)
+		}
+		u = nobody
 	}
 
 	// Make the shared directory and make it availabe to all user/groups.
-	if err := mkOwnedDir(d.SharedDir, uid, gid, 0777); err != nil {
+	if err := mkOwnedDir(d.SharedDir, u, 0777); err != nil {
 		return err
 	}
 
 	for _, dir := range SharedAllocDirs {
 		p := filepath.Join(d.SharedDir, dir)
-		if err := mkOwnedDir(p, uid, gid, 0777); err != nil {
+		if err := mkOwnedDir(p, u, 0777); err != nil {
 			return err
 		}
 	}
@@ -45,16 +45,16 @@ func (d *AllocDir) Build(tasks []*structs.Task) error {
 	// Make the task directories.
 	for _, t := range tasks {
 		p := filepath.Join(d.AllocDir, t.Name)
-		if err := mkOwnedDir(p, uid, gid, 0777); err != nil {
+		if err := mkOwnedDir(p, u, 0777); err != nil {
 			return err
 		}
 
 		// Create a local directory that each task can use.
 		local := filepath.Join(p, TaskLocal)
-		if err := mkOwnedDir(local, uid, gid, 0777); err != nil {
+		if err := mkOwnedDir(local, u, 0777); err != nil {
 			return err
 		}
-		d.TaskDirs[t.Name] = local
+		d.TaskDirs[t.Name] = p
 
 		// TODO: Mount the shared alloc dir into each task dir.
 	}
@@ -63,17 +63,88 @@ func (d *AllocDir) Build(tasks []*structs.Task) error {
 }
 
 func (d *AllocDir) Embed(task string, dirs 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 {
+		// Enumerate the files in source.
+		entries, err := ioutil.ReadDir(source)
+		if err != nil {
+			return fmt.Errorf("Couldn't read directory %v: %v", source, err)
+		}
+
+		// Create destination directory.
+		destDir := filepath.Join(taskdir, dest)
+		if err := os.MkdirAll(destDir, 0777); err != nil {
+			return fmt.Errorf("Couldn't create destination directory %v: %v", destDir, err)
+		}
+
+		for _, entry := range entries {
+			hostEntry := filepath.Join(source, entry.Name())
+			if entry.IsDir() {
+				subdirs[hostEntry] = filepath.Join(dest, filepath.Base(hostEntry))
+				continue
+			} else if !entry.Mode().IsRegular() {
+				return fmt.Errorf("Can't embed non-regular file: %v", hostEntry)
+			}
+
+			taskEntry := filepath.Join(destDir, filepath.Base(hostEntry))
+
+			// Attempt to hardlink.
+			if err := os.Link(hostEntry, taskEntry); err == nil {
+				continue
+			}
+
+			// Do a simple copy.
+			src, err := os.Open(hostEntry)
+			if err != nil {
+				return fmt.Errorf("Couldn't open host file %v: %v", hostEntry, err)
+			}
+
+			dst, err := os.OpenFile(taskEntry, os.O_WRONLY|os.O_CREATE, 0777)
+			if err != nil {
+				return fmt.Errorf("Couldn't create task file %v: %v", taskEntry, err)
+			}
+
+			if _, err := io.Copy(dst, src); err != nil {
+				return fmt.Errorf("Couldn't copy %v to %v: %v", hostEntry, taskEntry, err)
+			}
+		}
+	}
+
+	// Recurse on self to copy subdirectories.
+	if len(subdirs) != 0 {
+		return d.Embed(task, subdirs)
+	}
+
 	return nil
 }
 
 // mkOwnedDir creates the directory specified by the path with the passed
-// permissions. It also sets the passed uid and gid. It returns an error if any
-// of these operations fail.
-func mkOwnedDir(path string, uid, gid int, perm os.FileMode) error {
+// permissions. It also sets the user/group based on the passed user if it is
+// non-nil. It returns an error if any of these operations fail.
+func mkOwnedDir(path string, user *user.User, perm os.FileMode) error {
 	if err := os.Mkdir(path, perm); err != nil {
 		return fmt.Errorf("Failed to make directory %v: %v", path, err)
 	}
 
+	if user == nil {
+		return nil
+	}
+
+	uid, err := getUid(user)
+	if err != nil {
+		return err
+	}
+
+	gid, err := getGid(user)
+	if err != nil {
+		return err
+	}
+
 	if err := os.Chown(path, uid, gid); err != nil {
 		return fmt.Errorf("Couldn't change owner/group of %v to (uid: %v, gid: %v): %v", path, uid, gid, err)
 	}

client/allocdir/builder_windows.go

@@ -2,13 +2,17 @@
 
 package allocdir
 
-import "github.com/hashicorp/nomad/nomad/structs"
+import (
+	"errors"
+
+	"github.com/hashicorp/nomad/nomad/structs"
+)
 
 func (d *AllocDir) Build(tasks []*structs.Task) error {
 	// TODO: Need to figure out how to do mounts on windows.
-	return nil
+	return errors.New("Not implemented")
 }
 
 func (d *AllocDir) Embed(task string, dirs map[string]string) error {
-	return nil
+	return errors.New("Not implemented")
 }