Skip to content

Commit 3cf7d8e

Browse files
committed
Replace shell script wrappers with a Go wrapper
This patch replaces the shell script wrappers emitted by the CTK installer with a Go program. This allows the CTK to work on systems without a shell, such as [Talos Linux](https://www.talos.dev/). The Go wrapper reads its configuration from a dedicated ELF section, which although brings some complexity, allows for single-file standalone wrappers with room to expand the configuration options. Signed-off-by: Jean-Francois Roy <[email protected]>
1 parent b28c301 commit 3cf7d8e

File tree

17 files changed

+830
-239
lines changed

17 files changed

+830
-239
lines changed

.golangci.yml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,11 @@ linters:
5252
# TODO: We should address each of the following integer overflows.
5353
- path: (.+)\.go$
5454
text: 'G115: integer overflow conversion(.+)'
55+
# Specify syntax required for CGO symbol lookup in Go wrapper program.
56+
- linters:
57+
- gocritic
58+
path: cmd/nvidia-ctk-installer-wrapper/main.go
59+
text: could simplify
5560
paths:
5661
- third_party$
5762
- builtin$
Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
/**
2+
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
# SPDX-License-Identifier: Apache-2.0
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
**/
17+
18+
package main
19+
20+
/*
21+
__attribute__((section(".nvctkiwc")))
22+
char wrapper_config[4096] = {0};
23+
extern char wrapper_config[4096];
24+
*/
25+
import "C"
26+
27+
import (
28+
"log"
29+
"os"
30+
"os/exec"
31+
"strings"
32+
"unsafe"
33+
34+
"golang.org/x/sys/unix"
35+
36+
"github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore"
37+
)
38+
39+
func main() {
40+
config := loadConfig()
41+
if config.RequiresKernelModule && !isNvidiaModuleLoaded() {
42+
log.Println("nvidia driver modules are not yet loaded, invoking default runtime")
43+
runtimePath := config.DefaultRuntimeExecutablePath
44+
if runtimePath == "" {
45+
runtimePath = "runc"
46+
}
47+
program, err := exec.LookPath(runtimePath)
48+
if err != nil {
49+
log.Fatalf("failed to find %s: %v", runtimePath, err)
50+
}
51+
argv := []string{runtimePath}
52+
argv = append(argv, os.Args[1:]...)
53+
if err := unix.Exec(program, argv, os.Environ()); err != nil {
54+
log.Fatalf("failed to exec %s: %v", program, err)
55+
}
56+
}
57+
program, err := os.Executable()
58+
if err != nil {
59+
log.Fatalf("failed to get executable: %v", err)
60+
}
61+
argv := makeArgv(config)
62+
envv := makeEnvv(config)
63+
if err := unix.Exec(program+".real", argv, envv); err != nil {
64+
log.Fatalf("failed to exec %s: %v", program+".real", err)
65+
}
66+
}
67+
68+
func loadConfig() *wrappercore.WrapperConfig {
69+
ptr := unsafe.Pointer(&C.wrapper_config[0])
70+
sectionData := unsafe.Slice((*byte)(ptr), 4096)
71+
config, err := wrappercore.ReadWrapperConfigSection(sectionData)
72+
if err != nil {
73+
log.Fatalf("failed to load wrapper config: %v", err)
74+
}
75+
return config
76+
}
77+
78+
func isNvidiaModuleLoaded() bool {
79+
_, err := os.Stat("/proc/driver/nvidia/version")
80+
return err == nil
81+
}
82+
83+
func makeArgv(config *wrappercore.WrapperConfig) []string {
84+
argv := []string{os.Args[0] + ".real"}
85+
argv = append(argv, config.Argv...)
86+
return append(argv, os.Args[1:]...)
87+
}
88+
89+
func makeEnvv(config *wrappercore.WrapperConfig) []string {
90+
var env []string
91+
for k, v := range config.Envm {
92+
value := v
93+
if strings.HasPrefix(k, "<") {
94+
k = k[1:]
95+
value = value + ":" + os.Getenv(k)
96+
} else if strings.HasPrefix(k, ">") {
97+
k = k[1:]
98+
value = os.Getenv(k) + ":" + value
99+
}
100+
env = append(env, k+"="+value)
101+
}
102+
return append(env, os.Environ()...)
103+
}

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

Lines changed: 18 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -18,24 +18,21 @@
1818
package installer
1919

2020
import (
21-
"bytes"
22-
"fmt"
23-
"html/template"
24-
"io"
2521
"path/filepath"
26-
"strings"
2722

2823
log "github.com/sirupsen/logrus"
2924

3025
"github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/container/operator"
26+
"github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore"
3127
"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
3228
)
3329

3430
type executable struct {
3531
requiresKernelModule bool
3632
path string
3733
symlink string
38-
env map[string]string
34+
argv []string
35+
envm map[string]string
3936
}
4037

4138
func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, error) {
@@ -53,7 +50,7 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
5350
e := executable{
5451
path: runtime.Path,
5552
requiresKernelModule: true,
56-
env: map[string]string{
53+
envm: map[string]string{
5754
config.FilePathOverrideEnvVar: configFilePath,
5855
},
5956
}
@@ -62,15 +59,15 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
6259
executables = append(executables,
6360
executable{
6461
path: "nvidia-container-cli",
65-
env: map[string]string{"LD_LIBRARY_PATH": destDir + ":$LD_LIBRARY_PATH"},
62+
envm: map[string]string{"<LD_LIBRARY_PATH": destDir},
6663
},
6764
)
6865

6966
executables = append(executables,
7067
executable{
7168
path: "nvidia-container-runtime-hook",
7269
symlink: "nvidia-container-toolkit",
73-
env: map[string]string{
70+
envm: map[string]string{
7471
config.FilePathOverrideEnvVar: configFilePath,
7572
},
7673
},
@@ -87,25 +84,25 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
8784
return nil, err
8885
}
8986

90-
wrappedExecutableFilename := filepath.Base(executablePath)
91-
dotRealFilename := wrappedExecutableFilename + ".real"
92-
9387
w := &wrapper{
94-
Source: executablePath,
95-
WrappedExecutable: dotRealFilename,
96-
CheckModules: executable.requiresKernelModule,
97-
Envvars: map[string]string{
98-
"PATH": strings.Join([]string{destDir, "$PATH"}, ":"),
88+
Source: executablePath,
89+
WrapperProgramPath: t.wrapperProgramPath,
90+
Config: wrappercore.WrapperConfig{
91+
Argv: executable.argv,
92+
Envm: map[string]string{
93+
"<PATH": destDir,
94+
},
95+
RequiresKernelModule: executable.requiresKernelModule,
9996
},
10097
}
101-
for k, v := range executable.env {
102-
w.Envvars[k] = v
98+
for k, v := range executable.envm {
99+
w.Config.Envm[k] = v
103100
}
104101

105102
if len(t.defaultRuntimeExecutablePath) > 0 {
106-
w.DefaultRuntimeExecutablePath = t.defaultRuntimeExecutablePath
103+
w.Config.DefaultRuntimeExecutablePath = t.defaultRuntimeExecutablePath
107104
} else {
108-
w.DefaultRuntimeExecutablePath = "runc"
105+
w.Config.DefaultRuntimeExecutablePath = "runc"
109106
}
110107

111108
installers = append(installers, w)
@@ -121,66 +118,4 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
121118
}
122119

123120
return installers, nil
124-
125-
}
126-
127-
type wrapper struct {
128-
Source string
129-
Envvars map[string]string
130-
WrappedExecutable string
131-
CheckModules bool
132-
DefaultRuntimeExecutablePath string
133-
}
134-
135-
type render struct {
136-
*wrapper
137-
DestDir string
138-
}
139-
140-
func (w *wrapper) Install(destDir string) error {
141-
// Copy the executable with a .real extension.
142-
mode, err := installFile(w.Source, filepath.Join(destDir, w.WrappedExecutable))
143-
if err != nil {
144-
return err
145-
}
146-
147-
// Create a wrapper file.
148-
r := render{
149-
wrapper: w,
150-
DestDir: destDir,
151-
}
152-
content, err := r.render()
153-
if err != nil {
154-
return fmt.Errorf("failed to render wrapper: %w", err)
155-
}
156-
wrapperFile := filepath.Join(destDir, filepath.Base(w.Source))
157-
return installContent(content, wrapperFile, mode|0111)
158-
}
159-
160-
func (w *render) render() (io.Reader, error) {
161-
wrapperTemplate := `#! /bin/sh
162-
{{- if (.CheckModules) }}
163-
cat /proc/modules | grep -e "^nvidia " >/dev/null 2>&1
164-
if [ "${?}" != "0" ]; then
165-
echo "nvidia driver modules are not yet loaded, invoking {{ .DefaultRuntimeExecutablePath }} directly" >&2
166-
exec {{ .DefaultRuntimeExecutablePath }} "$@"
167-
fi
168-
{{- end }}
169-
{{- range $key, $value := .Envvars }}
170-
{{$key}}={{$value}} \
171-
{{- end }}
172-
{{ .DestDir }}/{{ .WrappedExecutable }} \
173-
"$@"
174-
`
175-
176-
var content bytes.Buffer
177-
tmpl, err := template.New("wrapper").Parse(wrapperTemplate)
178-
if err != nil {
179-
return nil, err
180-
}
181-
if err := tmpl.Execute(&content, w); err != nil {
182-
return nil, err
183-
}
184-
185-
return &content, nil
186121
}

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

Lines changed: 0 additions & 94 deletions
This file was deleted.

0 commit comments

Comments
 (0)