Replace shell wrapper with a Go wrapper

tools/container/toolkit/executable.go

@@ -18,123 +18,133 @@ package toolkit
 
 import (
 	"fmt"
-	"io"
 	"os"
 	"path/filepath"
 	"sort"
-	"strings"
 
 	log "github.com/sirupsen/logrus"
 )
 
 type executableTarget struct {
-	dotfileName string
 	wrapperName string
 }
 
 type executable struct {
-	source   string
-	target   executableTarget
-	env      map[string]string
-	preLines []string
-	argLines []string
+	source string
+	target executableTarget
+	argv   []string
+	envm   map[string]string
 }
 
 // install installs an executable component of the NVIDIA container toolkit. The source executable
 // is copied to a `.real` file and a wapper is created to set up the environment as required.
 func (e executable) install(destFolder string) (string, error) {
+	if destFolder == "" {
+		return "", fmt.Errorf("destination folder must be specified")
+	}
+	if e.source == "" {
+		return "", fmt.Errorf("source executable must be specified")
+	}
 	log.Infof("Installing executable '%v' to %v", e.source, destFolder)
-
-	dotfileName := e.dotfileName()
-
-	installedDotfileName, err := installFileToFolderWithName(destFolder, dotfileName, e.source)
+	dotRealFilename := e.dotRealFilename()
+	dotRealPath, err := installFileToFolderWithName(destFolder, dotRealFilename, e.source)
 	if err != nil {
-		return "", fmt.Errorf("error installing file '%v' as '%v': %v", e.source, dotfileName, err)
+		return "", fmt.Errorf("error installing file '%v' as '%v': %v", e.source, dotRealFilename, err)
 	}
-	log.Infof("Installed '%v'", installedDotfileName)
+	log.Infof("Installed '%v'", dotRealPath)
 
-	wrapperFilename, err := e.installWrapper(destFolder, installedDotfileName)
+	wrapperPath, err := e.installWrapper(destFolder)
 	if err != nil {
-		return "", fmt.Errorf("error wrapping '%v': %v", installedDotfileName, err)
+		return "", fmt.Errorf("error installing wrapper: %v", err)
 	}
-	log.Infof("Installed wrapper '%v'", wrapperFilename)
-
-	return wrapperFilename, nil
+	log.Infof("Installed wrapper '%v'", wrapperPath)
+	return wrapperPath, nil
 }
 
-func (e executable) dotfileName() string {
-	return e.target.dotfileName
+func (e executable) dotRealFilename() string {
+	return e.wrapperName() + ".real"
 }
 
 func (e executable) wrapperName() string {
+	if e.target.wrapperName == "" {
+		return filepath.Base(e.source)
+	}
 	return e.target.wrapperName
 }
 
-func (e executable) installWrapper(destFolder string, dotfileName string) (string, error) {
-	wrapperPath := filepath.Join(destFolder, e.wrapperName())
-	wrapper, err := os.Create(wrapperPath)
+func (e executable) installWrapper(destFolder string) (string, error) {
+	currentExe, err := os.Executable()
 	if err != nil {
-		return "", fmt.Errorf("error creating executable wrapper: %v", err)
+		return "", fmt.Errorf("error getting current executable: %v", err)
 	}
-	defer wrapper.Close()
-
-	err = e.writeWrapperTo(wrapper, destFolder, dotfileName)
+	src := filepath.Join(filepath.Dir(currentExe), "wrapper")
+	wrapperPath, err := installFileToFolderWithName(destFolder, e.wrapperName(), src)
 	if err != nil {
-		return "", fmt.Errorf("error writing wrapper contents: %v", err)
+		return "", fmt.Errorf("error installing wrapper program: %v", err)
+	}
+	err = e.writeWrapperArgv(wrapperPath, destFolder)
+	if err != nil {
+		return "", fmt.Errorf("error writing wrapper argv: %v", err)
+	}
+	err = e.writeWrapperEnvv(wrapperPath, destFolder)
+	if err != nil {
+		return "", fmt.Errorf("error writing wrapper envv: %v", err)
 	}
-
 	err = ensureExecutable(wrapperPath)
 	if err != nil {
 		return "", fmt.Errorf("error making wrapper executable: %v", err)
 	}
 	return wrapperPath, nil
 }
 
-func (e executable) writeWrapperTo(wrapper io.Writer, destFolder string, dotfileName string) error {
+func (e executable) writeWrapperArgv(wrapperPath string, destFolder string) error {
+	if e.argv == nil {
+		return nil
+	}
 	r := newReplacements(destDirPattern, destFolder)
-
-	// Add the shebang
-	fmt.Fprintln(wrapper, "#! /bin/sh")
-
-	// Add the preceding lines if any
-	for _, line := range e.preLines {
-		fmt.Fprintf(wrapper, "%s\n", r.apply(line))
+	f, err := os.OpenFile(wrapperPath+".argv", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0440)
+	if err != nil {
+		return err
 	}
-
-	// Update the path to include the destination folder
-	var env map[string]string
-	if e.env == nil {
-		env = make(map[string]string)
-	} else {
-		env = e.env
+	defer f.Close()
+	for _, arg := range e.argv {
+		fmt.Fprintf(f, "%s\n", r.apply(arg))
 	}
+	return nil
+}
 
-	path, specified := env["PATH"]
-	if !specified {
-		path = "$PATH"
+func (e executable) writeWrapperEnvv(wrapperPath string, destFolder string) error {
+	r := newReplacements(destDirPattern, destFolder)
+	f, err := os.OpenFile(wrapperPath+".envv", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0440)
+	if err != nil {
+		return err
 	}
-	env["PATH"] = strings.Join([]string{destFolder, path}, ":")
+	defer f.Close()
 
-	var sortedEnvvars []string
-	for e := range env {
-		sortedEnvvars = append(sortedEnvvars, e)
+	// Update PATH to insert the destination folder at the head.
+	var envm map[string]string
+	if e.envm == nil {
+		envm = make(map[string]string)
+	} else {
+		envm = e.envm
 	}
-	sort.Strings(sortedEnvvars)
-
-	for _, e := range sortedEnvvars {
-		v := env[e]
-		fmt.Fprintf(wrapper, "%s=%s \\\n", e, r.apply(v))
+	if path, ok := envm["PATH"]; ok {
+		envm["PATH"] = destFolder + ":" + path
+	} else {
+		// Replace PATH with <PATH, which instructs wrapper to insert the value at the head of a
+		// colon-separated environment variable list.
+		delete(envm, "PATH")
+		envm["<PATH"] = destFolder
 	}
-	// Add the call to the target executable
-	fmt.Fprintf(wrapper, "%s \\\n", dotfileName)
 
-	// Insert additional lines in the `arg` list
-	for _, line := range e.argLines {
-		fmt.Fprintf(wrapper, "\t%s \\\n", r.apply(line))
+	var envv []string
+	for k, v := range envm {
+		envv = append(envv, k+"="+r.apply(v))
+	}
+	sort.Strings(envv)
+	for _, e := range envv {
+		fmt.Fprintf(f, "%s\n", e)
 	}
-	// Add the script arguments "$@"
-	fmt.Fprintln(wrapper, "\t\"$@\"")
-
 	return nil
 }
 

tools/container/toolkit/executable_test.go

@@ -17,102 +17,102 @@
 package toolkit
 
 import (
-	"bytes"
+	"bufio"
+	"fmt"
+	"io/fs"
 	"os"
 	"path/filepath"
-	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
 func TestWrapper(t *testing.T) {
-	const shebang = "#! /bin/sh"
-	const destFolder = "/dest/folder"
-	const dotfileName = "source.real"
+	createTestWrapperProgram(t)
 
 	testCases := []struct {
-		e             executable
-		expectedLines []string
+		e            executable
+		expectedArgv []string
+		expectedEnvv []string
 	}{
 		{
-			e: executable{},
-			expectedLines: []string{
-				shebang,
-				"PATH=/dest/folder:$PATH \\",
-				"source.real \\",
-				"\t\"$@\"",
-				"",
+			e: executable{source: "source"},
+			expectedEnvv: []string{
+				fmt.Sprintf("<PATH=%s", destDirPattern),
 			},
 		},
 		{
 			e: executable{
-				env: map[string]string{
-					"PATH": "some-path",
+				source: "source",
+				envm: map[string]string{
+					"FOO": "BAR",
 				},
 			},
-			expectedLines: []string{
-				shebang,
-				"PATH=/dest/folder:some-path \\",
-				"source.real \\",
-				"\t\"$@\"",
-				"",
+			expectedEnvv: []string{
+				fmt.Sprintf("<PATH=%s", destDirPattern),
+				"FOO=BAR",
 			},
 		},
 		{
 			e: executable{
-				preLines: []string{
-					"preline1",
-					"preline2",
+				source: "source",
+				envm: map[string]string{
+					"PATH": "some-path",
+					"FOO":  "BAR",
 				},
 			},
-			expectedLines: []string{
-				shebang,
-				"preline1",
-				"preline2",
-				"PATH=/dest/folder:$PATH \\",
-				"source.real \\",
-				"\t\"$@\"",
-				"",
+			expectedEnvv: []string{
+				"FOO=BAR",
+				fmt.Sprintf("PATH=%s:some-path", destDirPattern),
 			},
 		},
 		{
 			e: executable{
-				argLines: []string{
-					"argline1",
-					"argline2",
+				source: "source",
+				argv: []string{
+					"argb",
+					"arga",
+					"argc",
 				},
 			},
-			expectedLines: []string{
-				shebang,
-				"PATH=/dest/folder:$PATH \\",
-				"source.real \\",
-				"\targline1 \\",
-				"\targline2 \\",
-				"\t\"$@\"",
-				"",
+			expectedArgv: []string{
+				"argb",
+				"arga",
+				"argc",
+			},
+			expectedEnvv: []string{
+				fmt.Sprintf("<PATH=%s", destDirPattern),
 			},
 		},
 	}
 
-	for i, tc := range testCases {
-		buf := &bytes.Buffer{}
-
-		err := tc.e.writeWrapperTo(buf, destFolder, dotfileName)
+	for _, tc := range testCases {
+		destFolder := t.TempDir()
+		r := newReplacements(destDirPattern, destFolder)
+		for k, v := range tc.expectedEnvv {
+			tc.expectedEnvv[k] = r.apply(v)
+		}
+		path, err := tc.e.installWrapper(destFolder)
 		require.NoError(t, err)
-
-		exepectedContents := strings.Join(tc.expectedLines, "\n")
-		require.Equal(t, exepectedContents, buf.String(), "%v: %v", i, tc)
+		require.FileExists(t, path)
+		envv, err := readAllLines(path + ".envv")
+		require.NoError(t, err)
+		require.Equal(t, tc.expectedEnvv, envv)
+		argv, err := readAllLines(path + ".argv")
+		if tc.expectedArgv == nil {
+			require.ErrorAs(t, err, &fs.ErrNotExist)
+		} else {
+			require.Equal(t, tc.expectedArgv, argv)
+
+		}
 	}
 }
 
 func TestInstallExecutable(t *testing.T) {
-	inputFolder, err := os.MkdirTemp("", "")
-	require.NoError(t, err)
-	defer os.RemoveAll(inputFolder)
+	createTestWrapperProgram(t)
 
 	// Create the source file
-	source := filepath.Join(inputFolder, "input")
+	source := filepath.Join(t.TempDir(), "input")
 	sourceFile, err := os.Create(source)
 
 	base := filepath.Base(source)
@@ -123,7 +123,6 @@ func TestInstallExecutable(t *testing.T) {
 	e := executable{
 		source: source,
 		target: executableTarget{
-			dotfileName: "input.real",
 			wrapperName: "input",
 		},
 	}
@@ -150,3 +149,31 @@ func TestInstallExecutable(t *testing.T) {
 	require.NoError(t, err)
 	require.NotEqual(t, 0, wrapperInfo.Mode()&0111)
 }
+
+func createTestWrapperProgram(t *testing.T) {
+	t.Helper()
+	currentExe, err := os.Executable()
+	if err != nil {
+		t.Fatalf("error getting current executable: %v", err)
+	}
+	wrapperPath := filepath.Join(filepath.Dir(currentExe), "wrapper")
+	f, err := os.Create(wrapperPath)
+	if err != nil {
+		t.Fatalf("error creating test wrapper: %v", err)
+	}
+	f.Close()
+}
+
+func readAllLines(path string) (s []string, err error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return
+	}
+	defer f.Close()
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		s = append(s, scanner.Text())
+	}
+	err = scanner.Err()
+	return
+}