Made ByteExec use static filename

byteexec.go

@@ -2,77 +2,164 @@
 // supplied as byte arrays, which is handy when used with
 // github.com/jteeuwen/go-bindata.
 //
-// ByteExec works by storing the provided command in a temp file.  A ByteExec
-// should always be closed using its Close() method to clean up the temp file.
+// byteexec works by storing the provided command in a file.
 //
 // Example Usage:
 //
 //    programBytes := // read bytes from somewhere
-//    be, err := NewByteExec(programBytes)
+//    be, err := byteexec.New(programBytes)
 //    if err != nil {
 //      log.Fatalf("Uh oh: %s", err)
 //    }
-//    defer be.Close()
 //    cmd := be.Command("arg1", "arg2")
 //    // cmd is an os/exec.Cmd
 //    err = cmd.Run()
+//
+// Note - byteexec.New is somewhat expensive, and Exec is safe for concurrent
+// use, so it's advisable to create only one Exec for each executable.
 package byteexec
 
 import (
-	"io/ioutil"
+	"bytes"
+	"crypto/sha256"
+	"fmt"
+	"io"
 	"os"
 	"os/exec"
+	"os/user"
+	"path"
+	"path/filepath"
+	"sync"
+
+	"github.com/getlantern/golog"
+)
+
+const (
+	fileMode = 0744
+)
+
+var (
+	log = golog.LoggerFor("Exec")
+
+	initMutex sync.Mutex
 )
 
-type ByteExec struct {
-	fileName string
+// Exec is a handle to an executable that can be used to create an exec.Cmd
+// using the Command method. Exec is safe for concurrent use.
+type Exec struct {
+	filename string
 }
 
-// NewByteExec creates a new ByteExec using the program stored in the provided
-// bytes.
-func NewByteExec(bytes []byte) (be *ByteExec, err error) {
-	return NewNamedByteExec(bytes, "byteexec")
+// New creates a new Exec using the program stored in the provided data, at the
+// provided filename (relative or absolute path allowed). If the path given is
+// a relative path, the executable will be placed in the user's home directory
+// in a subfolder named ".byteexec".
+//
+// Creating a new Exec can be somewhat expensive, so it's best to create only
+// one Exec per executable and reuse that.
+//
+// WARNING - if a file already exists at this location and its contents differ
+// from data, Exec will attempt to overwrite it.
+func New(data []byte, filename string) (*Exec, error) {
+	// Use initMutex to synchronize file operations by this process
+	initMutex.Lock()
+	defer initMutex.Unlock()
+
+	var err error
+	if !path.IsAbs(filename) {
+		filename, err = inUserDir(filename)
+		if err != nil {
+			return nil, err
+		}
+	}
+	filename = renameExecutable(filename)
+	log.Tracef("Placing executable in %s", filename)
+
+	log.Trace("Attempting to open file for creating, but only if it doesn't already exist")
+	file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_EXCL, fileMode)
+	if err != nil {
+		if !os.IsExist(err) {
+			return nil, fmt.Errorf("Unexpected error opening %s: %s", filename, err)
+		}
+
+		log.Tracef("%s already exists, check to make sure contents is the same", filename)
+		if checksumsMatch(filename, data) {
+			log.Tracef("Data in %s matches expected, using existing", filename)
+			return newExecFromExisting(filename)
+		}
+
+		log.Tracef("Data in %s doesn't match expected, truncating file", filename)
+		file, err = os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_TRUNC, fileMode)
+		if err != nil {
+			return nil, fmt.Errorf("Unable to truncate %s: %s", err)
+		}
+	}
+
+	log.Tracef("Created new file at %s, saving executable", filename)
+	_, err = file.Write(data)
+	if err != nil {
+		os.Remove(filename)
+		return nil, fmt.Errorf("Unable to write to file at %s: %s", filename, err)
+	}
+	file.Sync()
+	file.Close()
+
+	log.Trace("File saved, returning new Exec")
+	return newExec(filename)
+}
+
+// Command creates an exec.Cmd using the supplied args.
+func (be *Exec) Command(args ...string) *exec.Cmd {
+	return exec.Command(be.filename, args...)
 }
 
-// NewNamedByteExec creates a new ByteExec using the program stored in the
-// provided bytes and uses the given prefix to name the temporary file that gets
-// executed.
-func NewNamedByteExec(bytes []byte, prefix string) (be *ByteExec, err error) {
-	var tmpFile *os.File
-	tmpFile, err = ioutil.TempFile("", prefix+"_")
+func checksumsMatch(filename string, data []byte) bool {
+	shasum := sha256.New()
+	file, err := os.OpenFile(filename, os.O_RDONLY, 0)
 	if err != nil {
-		return
+		log.Tracef("Unable to open existing file at %s for reading: %s", filename, err)
+		return false
 	}
-	_, err = tmpFile.Write(bytes)
+	_, err = io.Copy(shasum, file)
 	if err != nil {
-		return
+		log.Tracef("Unable to read bytes to calculate sha sum: %s", err)
+		return false
 	}
-	tmpFile.Sync()
-	tmpFile.Chmod(0755)
-	tmpFile.Close()
-
-	orig := tmpFile.Name()
-	renamed := renameExecutable(orig)
-	if renamed != orig {
-		err = os.Rename(orig, renamed)
+	checksumOnDisk := shasum.Sum(nil)
+	expectedChecksum := sha256.Sum256(data)
+	return bytes.Equal(checksumOnDisk, expectedChecksum[:])
+}
+
+func newExecFromExisting(filename string) (*Exec, error) {
+	fi, err := os.Stat(filename)
+	if err != nil || fi.Mode() != fileMode {
+		log.Tracef("Chmodding %s", filename)
+		err = os.Chmod(filename, fileMode)
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("Unable to chmod file %s: %s", filename, err)
 		}
 	}
-	be = &ByteExec{fileName: renamed}
-	return
+	return newExec(filename)
 }
 
-// Command creates an exec.Cmd using the supplied args.
-func (be *ByteExec) Command(args ...string) *exec.Cmd {
-	return exec.Command(be.fileName, args...)
+func newExec(filename string) (*Exec, error) {
+	absolutePath, err := filepath.Abs(filename)
+	if err != nil {
+		return nil, err
+	}
+	return &Exec{filename: absolutePath}, nil
 }
 
-// Close() closes the ByteExec, cleaning up the associated temp file.
-func (be *ByteExec) Close() error {
-	if be.fileName == "" {
-		return nil
-	} else {
-		return os.Remove(be.fileName)
+func inUserDir(filename string) (string, error) {
+	log.Tracef("Determining user's home directory")
+	usr, err := user.Current()
+	if err != nil {
+		return filename, fmt.Errorf("Unable to determine user's home directory: %s", err)
+	}
+	folder := path.Join(usr.HomeDir, ".byteexec")
+	err = os.MkdirAll(folder, fileMode)
+	if err != nil {
+		return filename, fmt.Errorf("Unable to make folder %s: %s", folder, err)
 	}
+	return path.Join(folder, filename), nil
 }

byteexec_test.go

@@ -1,26 +1,84 @@
 package byteexec
 
 import (
+	"io/ioutil"
+	"os"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/getlantern/testify/assert"
 )
 
+const (
+	program = "helloworld"
+
+	concurrency = 10
+)
+
 func TestExec(t *testing.T) {
-	bytes, err := Asset("helloworld")
+	data, err := Asset(program)
 	if err != nil {
 		t.Fatalf("Unable to read helloworld program: %s", err)
 	}
-	be, err := NewByteExec(bytes)
+	be := createByteExec(t, data)
+
+	// Concurrently create some other BEs and make sure they don't get errors
+	var wg sync.WaitGroup
+	wg.Add(concurrency)
+	for i := 0; i < concurrency; i++ {
+		_, err := New(data, program)
+		assert.NoError(t, err, "Concurrent New should have succeeded")
+		wg.Done()
+	}
+	wg.Wait()
+
+	originalInfo := testByteExec(t, be)
+
+	// Recreate be and make sure file is reused
+	be = createByteExec(t, data)
+	updatedInfo := testByteExec(t, be)
+	assert.Equal(t, originalInfo.ModTime(), updatedInfo.ModTime(), "File modification time should be unchanged after creating new ByteExec")
+
+	// Now mess with the file permissions and make sure that we can still run
+	err = os.Chmod(be.filename, 0655)
+	if err != nil {
+		t.Fatalf("Unable to chmod test executable %s: %s", be.filename, err)
+	}
+	be = createByteExec(t, data)
+	updatedInfo = testByteExec(t, be)
+	assert.Equal(t, fileMode, updatedInfo.Mode(), "File mode is changed back to %v", fileMode)
+
+	// Now mess with the file contents and make sure it gets overwritten on next
+	// ByteExec
+	ioutil.WriteFile(be.filename, []byte("Junk"), 0755)
+	be = createByteExec(t, data)
+	updatedInfo = testByteExec(t, be)
+	assert.NotEqual(t, originalInfo.ModTime(), updatedInfo.ModTime(), "File modification time should be changed after creating new ByteExec on bad data")
+}
+
+func createByteExec(t *testing.T, data []byte) *Exec {
+	// Sleep 1 second to give file timestamp a chance to increase
+	time.Sleep(1 * time.Second)
+
+	be, err := New(data, program)
 	if err != nil {
 		t.Fatalf("Unable to create new ByteExec: %s", err)
 	}
-	defer be.Close()
+	return be
+}
+
+func testByteExec(t *testing.T, be *Exec) os.FileInfo {
 	cmd := be.Command()
 	out, err := cmd.CombinedOutput()
 	if err != nil {
 		t.Errorf("Unable to run helloworld program: %s", err)
 	}
+	assert.Equal(t, "Hello world\n", string(out), "Should receive expected output from helloworld program")
 
-	assert.Equal(t, "Hello world\n", string(out), "Did not receive expected output from helloworld program")
+	fileInfo, err := os.Stat(be.filename)
+	if err != nil {
+		t.Fatalf("Unable to re-stat file %s: %s", be.filename, err)
+	}
+	return fileInfo
 }

rename_darwin.go → rename_stub.go

@@ -1,3 +1,5 @@
+// +build !windows
+
 package byteexec
 
 func renameExecutable(orig string) string {