Replace shell wrapper with a Go wrapper

.golangci.yml

@@ -52,6 +52,11 @@ linters:
       # TODO: We should address each of the following integer overflows.
       - path: (.+)\.go$
         text: 'G115: integer overflow conversion(.+)'
+      # Specify syntax required for CGO symbol lookup in Go wrapper program.
+      - linters:
+          - gocritic
+        path: cmd/nvidia-ctk-installer-wrapper/main.go
+        text: could simplify
     paths:
       - third_party$
       - builtin$

cmd/nvidia-ctk-installer-wrapper/main.go

@@ -0,0 +1,103 @@
+/**
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+**/
+
+package main
+
+/*
+__attribute__((section(".nvctkiwc")))
+char wrapper_config[4096] = {0};
+extern char wrapper_config[4096];
+*/
+import "C"
+
+import (
+	"log"
+	"os"
+	"os/exec"
+	"strings"
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore"
+)
+
+func main() {
+	config := loadConfig()
+	if config.RequiresKernelModule && !isNvidiaModuleLoaded() {
+		log.Println("nvidia driver modules are not yet loaded, invoking default runtime")
+		runtimePath := config.DefaultRuntimeExecutablePath
+		if runtimePath == "" {
+			runtimePath = "runc"
+		}
+		program, err := exec.LookPath(runtimePath)
+		if err != nil {
+			log.Fatalf("failed to find %s: %v", runtimePath, err)
+		}
+		argv := []string{runtimePath}
+		argv = append(argv, os.Args[1:]...)
+		if err := unix.Exec(program, argv, os.Environ()); err != nil {
+			log.Fatalf("failed to exec %s: %v", program, err)
+		}
+	}
+	program, err := os.Executable()
+	if err != nil {
+		log.Fatalf("failed to get executable: %v", err)
+	}
+	argv := makeArgv(config)
+	envv := makeEnvv(config)
+	if err := unix.Exec(program+".real", argv, envv); err != nil {
+		log.Fatalf("failed to exec %s: %v", program+".real", err)
+	}
+}
+
+func loadConfig() *wrappercore.WrapperConfig {
+	ptr := unsafe.Pointer(&C.wrapper_config[0])
+	sectionData := unsafe.Slice((*byte)(ptr), 4096)
+	config, err := wrappercore.ReadWrapperConfigSection(sectionData)
+	if err != nil {
+		log.Fatalf("failed to load wrapper config: %v", err)
+	}
+	return config
+}
+
+func isNvidiaModuleLoaded() bool {
+	_, err := os.Stat("/proc/driver/nvidia/version")
+	return err == nil
+}
+
+func makeArgv(config *wrappercore.WrapperConfig) []string {
+	argv := []string{os.Args[0] + ".real"}
+	argv = append(argv, config.Argv...)
+	return append(argv, os.Args[1:]...)
+}
+
+func makeEnvv(config *wrappercore.WrapperConfig) []string {
+	var env []string
+	for k, v := range config.Envm {
+		value := v
+		if strings.HasPrefix(k, "<") {
+			k = k[1:]
+			value = value + ":" + os.Getenv(k)
+		} else if strings.HasPrefix(k, ">") {
+			k = k[1:]
+			value = os.Getenv(k) + ":" + value
+		}
+		env = append(env, k+"="+value)
+	}
+	return append(env, os.Environ()...)
+}

cmd/nvidia-ctk-installer/toolkit/installer/executables.go

@@ -18,24 +18,21 @@
 package installer
 
 import (
-	"bytes"
-	"fmt"
-	"html/template"
-	"io"
 	"path/filepath"
-	"strings"
 
 	log "github.com/sirupsen/logrus"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/container/operator"
+	"github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
 )
 
 type executable struct {
 	requiresKernelModule bool
 	path                 string
 	symlink              string
-	env                  map[string]string
+	argv                 []string
+	envm                 map[string]string
 }
 
 func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, error) {
@@ -53,7 +50,7 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
 		e := executable{
 			path:                 runtime.Path,
 			requiresKernelModule: true,
-			env: map[string]string{
+			envm: map[string]string{
 				config.FilePathOverrideEnvVar: configFilePath,
 			},
 		}
@@ -62,15 +59,15 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
 	executables = append(executables,
 		executable{
 			path: "nvidia-container-cli",
-			env:  map[string]string{"LD_LIBRARY_PATH": destDir + ":$LD_LIBRARY_PATH"},
+			envm: map[string]string{"<LD_LIBRARY_PATH": destDir},
 		},
 	)
 
 	executables = append(executables,
 		executable{
 			path:    "nvidia-container-runtime-hook",
 			symlink: "nvidia-container-toolkit",
-			env: map[string]string{
+			envm: map[string]string{
 				config.FilePathOverrideEnvVar: configFilePath,
 			},
 		},
@@ -87,25 +84,25 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
 			return nil, err
 		}
 
-		wrappedExecutableFilename := filepath.Base(executablePath)
-		dotRealFilename := wrappedExecutableFilename + ".real"
-
 		w := &wrapper{
-			Source:            executablePath,
-			WrappedExecutable: dotRealFilename,
-			CheckModules:      executable.requiresKernelModule,
-			Envvars: map[string]string{
-				"PATH": strings.Join([]string{destDir, "$PATH"}, ":"),
+			Source:             executablePath,
+			WrapperProgramPath: t.wrapperProgramPath,
+			Config: wrappercore.WrapperConfig{
+				Argv: executable.argv,
+				Envm: map[string]string{
+					"<PATH": destDir,
+				},
+				RequiresKernelModule: executable.requiresKernelModule,
 			},
 		}
-		for k, v := range executable.env {
-			w.Envvars[k] = v
+		for k, v := range executable.envm {
+			w.Config.Envm[k] = v
 		}
 
 		if len(t.defaultRuntimeExecutablePath) > 0 {
-			w.DefaultRuntimeExecutablePath = t.defaultRuntimeExecutablePath
+			w.Config.DefaultRuntimeExecutablePath = t.defaultRuntimeExecutablePath
 		} else {
-			w.DefaultRuntimeExecutablePath = "runc"
+			w.Config.DefaultRuntimeExecutablePath = "runc"
 		}
 
 		installers = append(installers, w)
@@ -121,66 +118,4 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
 	}
 
 	return installers, nil
-
-}
-
-type wrapper struct {
-	Source                       string
-	Envvars                      map[string]string
-	WrappedExecutable            string
-	CheckModules                 bool
-	DefaultRuntimeExecutablePath string
-}
-
-type render struct {
-	*wrapper
-	DestDir string
-}
-
-func (w *wrapper) Install(destDir string) error {
-	// Copy the executable with a .real extension.
-	mode, err := installFile(w.Source, filepath.Join(destDir, w.WrappedExecutable))
-	if err != nil {
-		return err
-	}
-
-	// Create a wrapper file.
-	r := render{
-		wrapper: w,
-		DestDir: destDir,
-	}
-	content, err := r.render()
-	if err != nil {
-		return fmt.Errorf("failed to render wrapper: %w", err)
-	}
-	wrapperFile := filepath.Join(destDir, filepath.Base(w.Source))
-	return installContent(content, wrapperFile, mode|0111)
-}
-
-func (w *render) render() (io.Reader, error) {
-	wrapperTemplate := `#! /bin/sh
-{{- if (.CheckModules) }}
-cat /proc/modules | grep -e "^nvidia " >/dev/null 2>&1
-if [ "${?}" != "0" ]; then
-	echo "nvidia driver modules are not yet loaded, invoking {{ .DefaultRuntimeExecutablePath }} directly" >&2
-	exec {{ .DefaultRuntimeExecutablePath }} "$@"
-fi
-{{- end }}
-{{- range $key, $value := .Envvars }}
-{{$key}}={{$value}} \
-{{- end }}
-	{{ .DestDir }}/{{ .WrappedExecutable }} \
-		"$@"
-`
-
-	var content bytes.Buffer
-	tmpl, err := template.New("wrapper").Parse(wrapperTemplate)
-	if err != nil {
-		return nil, err
-	}
-	if err := tmpl.Execute(&content, w); err != nil {
-		return nil, err
-	}
-
-	return &content, nil
 }