pkg/fileutil: move from fileutil

2 files changed, 175 insertions, 0 deletions

diff --git a/pkg/fileutil/fileutil.go b/pkg/fileutil/fileutil.go
new file mode 100644
index 000000000..c244204a9
--- /dev/null
+++ b/pkg/fileutil/fileutil.go
@@ -0,0 +1,110 @@
+// Copyright 2015 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package fileutil
+
+import (
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strconv"
+	"syscall"
+	"unsafe"
+)
+
+// CopyFile copies oldFile to newFile.
+func CopyFile(oldFile, newFile string) error {
+	oldf, err := os.Open(oldFile)
+	if err != nil {
+		return err
+	}
+	defer oldf.Close()
+	newf, err := os.Create(newFile)
+	if err != nil {
+		return err
+	}
+	defer newf.Close()
+	_, err = io.Copy(newf, oldf)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// WriteTempFile writes data to a temp file and returns its name.
+func WriteTempFile(data []byte) (string, error) {
+	f, err := ioutil.TempFile("", "syzkaller")
+	if err != nil {
+		return "", fmt.Errorf("failed to create a temp file: %v", err)
+	}
+	if _, err := f.Write(data); err != nil {
+		f.Close()
+		os.Remove(f.Name())
+		return "", fmt.Errorf("failed to write a temp file: %v", err)
+	}
+	f.Close()
+	return f.Name(), nil
+}
+
+// ProcessTempDir creates a new temp dir in where and returns its path and an unique index.
+// It also cleans up old, unused temp dirs after dead processes.
+func ProcessTempDir(where string) (string, error) {
+	lk := filepath.Join(where, "instance-lock")
+	lkf, err := syscall.Open(lk, syscall.O_RDWR|syscall.O_CREAT, 0600)
+	if err != nil {
+		return "", err
+	}
+	defer syscall.Close(lkf)
+	if err := syscall.Flock(lkf, syscall.LOCK_EX); err != nil {
+		return "", err
+	}
+	defer syscall.Flock(lkf, syscall.LOCK_UN)
+
+	for i := 0; i < 1e3; i++ {
+		path := filepath.Join(where, fmt.Sprintf("instance-%v", i))
+		pidfile := filepath.Join(path, ".pid")
+		err := os.Mkdir(path, 0700)
+		if os.IsExist(err) {
+			// Try to clean up.
+			data, err := ioutil.ReadFile(pidfile)
+			if err == nil && len(data) > 0 {
+				pid, err := strconv.Atoi(string(data))
+				if err == nil && pid > 1 {
+					if err := syscall.Kill(pid, 0); err == syscall.ESRCH {
+						if os.Remove(pidfile) == nil {
+							if os.RemoveAll(path) == nil {
+								i--
+								continue
+							}
+						}
+					}
+				}
+			}
+			// If err != nil, assume that the pid file is not created yet.
+			continue
+		}
+		if err != nil {
+			return "", err
+		}
+		if err := ioutil.WriteFile(pidfile, []byte(strconv.Itoa(syscall.Getpid())), 0600); err != nil {
+			return "", err
+		}
+		return path, nil
+	}
+	return "", fmt.Errorf("too many live instances")
+}
+
+// UmountAll recurusively unmounts all mounts in dir.
+func UmountAll(dir string) {
+	files, _ := ioutil.ReadDir(dir)
+	for _, f := range files {
+		name := filepath.Join(dir, f.Name())
+		if f.IsDir() {
+			UmountAll(name)
+		}
+		fn := []byte(name + "\x00")
+		syscall.Syscall(syscall.SYS_UMOUNT2, uintptr(unsafe.Pointer(&fn[0])), syscall.MNT_FORCE, 0)
+	}
+}
diff --git a/pkg/fileutil/fileutil_test.go b/pkg/fileutil/fileutil_test.go
new file mode 100644
index 000000000..d432ef5e7
--- /dev/null
+++ b/pkg/fileutil/fileutil_test.go
@@ -0,0 +1,65 @@
+// Copyright 2015 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package fileutil
+
+import (
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strconv"
+	"sync"
+	"testing"
+)
+
+func TestProcessTempDir(t *testing.T) {
+	for try := 0; try < 10; try++ {
+		func() {
+			tmp, err := ioutil.TempDir("", "syz")
+			if err != nil {
+				t.Fatalf("failed to create a temp dir: %v", err)
+			}
+			defer os.RemoveAll(tmp)
+			const P = 16
+			// Pre-create half of the instances with stale pid.
+			var dirs []string
+			for i := 0; i < P/2; i++ {
+				dir, err := ProcessTempDir(tmp)
+				if err != nil {
+					t.Fatalf("failed to create process temp dir")
+				}
+				dirs = append(dirs, dir)
+			}
+			for _, dir := range dirs {
+				if err := ioutil.WriteFile(filepath.Join(dir, ".pid"), []byte(strconv.Itoa(999999999)), 0600); err != nil {
+					t.Fatalf("failed to write pid file: %v", err)
+				}
+			}
+			// Now request a bunch of instances concurrently.
+			done := make(chan bool)
+			allDirs := make(map[string]bool)
+			var mu sync.Mutex
+			for p := 0; p < P; p++ {
+				go func() {
+					defer func() {
+						done <- true
+					}()
+					dir, err := ProcessTempDir(tmp)
+					if err != nil {
+						t.Fatalf("failed to create process temp dir")
+					}
+					mu.Lock()
+					present := allDirs[dir]
+					allDirs[dir] = true
+					mu.Unlock()
+					if present {
+						t.Fatalf("duplicate dir %v", dir)
+					}
+				}()
+			}
+			for p := 0; p < P; p++ {
+				<-done
+			}
+		}()
+	}
+}