diff --git a/.golangci.yml b/.golangci.yml index e0204a449..e6d7593f7 100644 --- a/.golangci.yml +++ b/.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$ diff --git a/cmd/nvidia-ctk-installer-wrapper/main.go b/cmd/nvidia-ctk-installer-wrapper/main.go new file mode 100644 index 000000000..319cde5a6 --- /dev/null +++ b/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()...) +} diff --git a/cmd/nvidia-ctk-installer/toolkit/installer/executables.go b/cmd/nvidia-ctk-installer/toolkit/installer/executables.go index 0bedd8511..aac0b7139 100644 --- a/cmd/nvidia-ctk-installer/toolkit/installer/executables.go +++ b/cmd/nvidia-ctk-installer/toolkit/installer/executables.go @@ -18,16 +18,12 @@ 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" ) @@ -35,7 +31,8 @@ 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,7 +59,7 @@ 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{" 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 } diff --git a/cmd/nvidia-ctk-installer/toolkit/installer/executables_test.go b/cmd/nvidia-ctk-installer/toolkit/installer/executables_test.go deleted file mode 100644 index 7c8b0d83c..000000000 --- a/cmd/nvidia-ctk-installer/toolkit/installer/executables_test.go +++ /dev/null @@ -1,94 +0,0 @@ -/** -# 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 installer - -import ( - "bytes" - "testing" - - "github.com/stretchr/testify/require" -) - -func TestWrapperRender(t *testing.T) { - testCases := []struct { - description string - w *wrapper - expected string - }{ - { - description: "executable is added", - w: &wrapper{ - WrappedExecutable: "some-runtime", - DefaultRuntimeExecutablePath: "runc", - }, - expected: `#! /bin/sh - /dest-dir/some-runtime \ - "$@" -`, - }, - { - description: "module check is added", - w: &wrapper{ - WrappedExecutable: "some-runtime", - CheckModules: true, - DefaultRuntimeExecutablePath: "runc", - }, - expected: `#! /bin/sh -cat /proc/modules | grep -e "^nvidia " >/dev/null 2>&1 -if [ "${?}" != "0" ]; then - echo "nvidia driver modules are not yet loaded, invoking runc directly" >&2 - exec runc "$@" -fi - /dest-dir/some-runtime \ - "$@" -`, - }, - { - description: "environment is added", - w: &wrapper{ - WrappedExecutable: "some-runtime", - Envvars: map[string]string{ - "PATH": "/foo/bar/baz", - }, - DefaultRuntimeExecutablePath: "runc", - }, - expected: `#! /bin/sh -PATH=/foo/bar/baz \ - /dest-dir/some-runtime \ - "$@" -`, - }, - } - - for _, tc := range testCases { - t.Run(tc.description, func(t *testing.T) { - r := render{ - wrapper: tc.w, - DestDir: "/dest-dir", - } - reader, err := r.render() - require.NoError(t, err) - - var content bytes.Buffer - _, err = content.ReadFrom(reader) - require.NoError(t, err) - - require.Equal(t, tc.expected, content.String()) - }) - } -} diff --git a/cmd/nvidia-ctk-installer/toolkit/installer/installer.go b/cmd/nvidia-ctk-installer/toolkit/installer/installer.go index 71c3cc60e..913110abc 100644 --- a/cmd/nvidia-ctk-installer/toolkit/installer/installer.go +++ b/cmd/nvidia-ctk-installer/toolkit/installer/installer.go @@ -25,6 +25,7 @@ import ( "os" "path/filepath" + "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore" "github.com/NVIDIA/nvidia-container-toolkit/internal/logger" ) @@ -43,6 +44,8 @@ type ToolkitInstaller struct { ensureTargetDirectory Installer defaultRuntimeExecutablePath string + + wrapperProgramPath string } var _ Installer = (*ToolkitInstaller)(nil) @@ -66,6 +69,22 @@ func New(opts ...Option) (*ToolkitInstaller, error) { } t.artifactRoot = artifactRoot } + if t.wrapperProgramPath == "" { + // The wrapper program will be next to the installer in the release container images. For + // testing, a pre-made test ELF (see mktestelf) is in the testdata artifact root. + executable, err := os.Executable() + if err != nil { + return nil, err + } + wrapperProgramPath := filepath.Join(filepath.Dir(executable), wrappercore.InstallerWrapperFilename) + if _, err := os.Stat(wrapperProgramPath); err != nil { + wrapperProgramPath = filepath.Join(t.artifactRoot.path, wrappercore.InstallerWrapperFilename) + if _, err := os.Stat(wrapperProgramPath); err != nil { + return nil, fmt.Errorf("failed to find wrapper program: %w", err) + } + } + t.wrapperProgramPath = wrapperProgramPath + } if t.ensureTargetDirectory == nil { t.ensureTargetDirectory = t.createDirectory() diff --git a/cmd/nvidia-ctk-installer/toolkit/installer/installer_test.go b/cmd/nvidia-ctk-installer/toolkit/installer/installer_test.go index d308fb3c3..bc40ca91b 100644 --- a/cmd/nvidia-ctk-installer/toolkit/installer/installer_test.go +++ b/cmd/nvidia-ctk-installer/toolkit/installer/installer_test.go @@ -19,6 +19,7 @@ package installer import ( "bytes" + "debug/elf" "fmt" "io" "io/fs" @@ -29,6 +30,8 @@ import ( testlog "github.com/sirupsen/logrus/hooks/test" "github.com/stretchr/testify/require" + "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/testutil" + "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore" "github.com/NVIDIA/nvidia-container-toolkit/internal/lookup" ) @@ -36,9 +39,9 @@ func TestToolkitInstaller(t *testing.T) { logger, _ := testlog.NewNullLogger() type contentCall struct { - wrapper string - path string - mode fs.FileMode + config *wrappercore.WrapperConfig + path string + mode fs.FileMode } var contentCalls []contentCall @@ -51,10 +54,18 @@ func TestToolkitInstaller(t *testing.T) { if _, err := b.ReadFrom(reader); err != nil { return err } + elfFile, err := elf.NewFile(bytes.NewReader(b.Bytes())) + require.NoError(t, err) + section := elfFile.Section(wrappercore.InstallerWrapperConfigELFSectionName) + require.NotNil(t, section) + sectionData, err := section.Data() + require.NoError(t, err) + config, err := wrappercore.ReadWrapperConfigSection(sectionData) + require.NoError(t, err) contents := contentCall{ - wrapper: b.String(), - path: s, - mode: fileMode, + config: config, + path: s, + mode: fileMode, } contentCalls = append(contentCalls, contents) @@ -117,6 +128,7 @@ func TestToolkitInstaller(t *testing.T) { artifactRoot: r, ensureTargetDirectory: createDirectory, defaultRuntimeExecutablePath: "runc", + wrapperProgramPath: testutil.WriteTestELF(t, wrappercore.InstallerWrapperConfigELFSectionName), } err := i.Install("/foo/bar/baz") @@ -167,85 +179,80 @@ func TestToolkitInstaller(t *testing.T) { { path: "/foo/bar/baz/nvidia-container-runtime", mode: 0777, - wrapper: `#! /bin/sh -cat /proc/modules | grep -e "^nvidia " >/dev/null 2>&1 -if [ "${?}" != "0" ]; then - echo "nvidia driver modules are not yet loaded, invoking runc directly" >&2 - exec runc "$@" -fi -NVIDIA_CTK_CONFIG_FILE_PATH=/foo/bar/baz/.config/nvidia-container-runtime/config.toml \ -PATH=/foo/bar/baz:$PATH \ - /foo/bar/baz/nvidia-container-runtime.real \ - "$@" -`, + config: &wrappercore.WrapperConfig{ + RequiresKernelModule: true, + DefaultRuntimeExecutablePath: "runc", + Envm: map[string]string{ + "NVIDIA_CTK_CONFIG_FILE_PATH": "/foo/bar/baz/.config/nvidia-container-runtime/config.toml", + "/dev/null 2>&1 -if [ "${?}" != "0" ]; then - echo "nvidia driver modules are not yet loaded, invoking runc directly" >&2 - exec runc "$@" -fi -NVIDIA_CTK_CONFIG_FILE_PATH=/foo/bar/baz/.config/nvidia-container-runtime/config.toml \ -PATH=/foo/bar/baz:$PATH \ - /foo/bar/baz/nvidia-container-runtime.cdi.real \ - "$@" -`, + config: &wrappercore.WrapperConfig{ + RequiresKernelModule: true, + DefaultRuntimeExecutablePath: "runc", + Envm: map[string]string{ + "NVIDIA_CTK_CONFIG_FILE_PATH": "/foo/bar/baz/.config/nvidia-container-runtime/config.toml", + "/dev/null 2>&1 -if [ "${?}" != "0" ]; then - echo "nvidia driver modules are not yet loaded, invoking runc directly" >&2 - exec runc "$@" -fi -NVIDIA_CTK_CONFIG_FILE_PATH=/foo/bar/baz/.config/nvidia-container-runtime/config.toml \ -PATH=/foo/bar/baz:$PATH \ - /foo/bar/baz/nvidia-container-runtime.legacy.real \ - "$@" -`, + config: &wrappercore.WrapperConfig{ + RequiresKernelModule: true, + DefaultRuntimeExecutablePath: "runc", + Envm: map[string]string{ + "NVIDIA_CTK_CONFIG_FILE_PATH": "/foo/bar/baz/.config/nvidia-container-runtime/config.toml", + "\n", os.Args[0]) + os.Exit(1) + } + t := &testing.T{} + elfBuf := testutil.CreateTestELF(t, wrappercore.InstallerWrapperConfigELFSectionName) + if err := os.WriteFile(os.Args[1], elfBuf.Bytes(), 0755); err != nil { //nolint:gosec + fmt.Fprintf(os.Stderr, "Error writing file: %v\n", err) + os.Exit(1) + } +} diff --git a/cmd/nvidia-ctk-installer/toolkit/installer/testutil/wrappertestutil.go b/cmd/nvidia-ctk-installer/toolkit/installer/testutil/wrappertestutil.go new file mode 100644 index 000000000..3676211b0 --- /dev/null +++ b/cmd/nvidia-ctk-installer/toolkit/installer/testutil/wrappertestutil.go @@ -0,0 +1,138 @@ +/** +# 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 testutil + +import ( + "bytes" + "debug/elf" + "encoding/binary" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +// WriteTestELF returns the path to a minimal ELF file with a "config section" for testing. +func WriteTestELF(t *testing.T, configSectionName string) string { + t.Helper() + elfBuf := CreateTestELF(t, configSectionName) + elfPath := filepath.Join(t.TempDir(), "test-wrapper") + f, err := os.Create(elfPath) + require.NoError(t, err) + defer f.Close() + _, err = f.Write(elfBuf.Bytes()) + require.NoError(t, err) + return elfPath +} + +// CreateTestELF returns an in-memory minimal ELF file with a "config section" for testing. +func CreateTestELF(t *testing.T, configSectionName string) *bytes.Buffer { + t.Helper() + + var buf bytes.Buffer + + // ELF header with 4 sections (null, .shstrtab, wrapper config section, alignment) + ehdr := elf.Header64{ + Ident: [16]byte{0x7f, 'E', 'L', 'F', 2, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0}, + Type: uint16(elf.ET_EXEC), + Machine: uint16(elf.EM_X86_64), + Version: 1, + Entry: 0, + Phoff: 0, + Shoff: 64, // Section headers start after ELF header + Flags: 0, + Ehsize: 64, + Phentsize: 0, + Phnum: 0, + Shentsize: 64, + Shnum: 4, + Shstrndx: 1, // .shstrtab is section 1 + } + require.NoError(t, binary.Write(&buf, binary.LittleEndian, &ehdr)) + + shdrTableOffset := uint64(64) + shdrTableSize := uint64(4 * 64) + shstrtabOffset := shdrTableOffset + shdrTableSize + configSectionOffset := shstrtabOffset + 64 // some space for string table + + // Build section string table + shstrtab := []byte{0} + shstrtab = append(shstrtab, ".shstrtab"...) + shstrtab = append(shstrtab, 0) + configSectionNameOffset := len(shstrtab) + shstrtab = append(shstrtab, configSectionName...) + shstrtab = append(shstrtab, 0) + + // Pad to section headers offset + for buf.Len() < int(shdrTableOffset) { + buf.WriteByte(0) + } + + // Section 0: null + nullShdr := elf.Section64{} + require.NoError(t, binary.Write(&buf, binary.LittleEndian, &nullShdr)) + + // Section 1: .shstrtab + shstrtabShdr := elf.Section64{ + Name: 1, // offset in shstrtab + Type: uint32(elf.SHT_STRTAB), + Flags: 0, + Addr: 0, + Off: shstrtabOffset, + Size: uint64(len(shstrtab)), + Link: 0, + Info: 0, + Addralign: 1, + Entsize: 0, + } + require.NoError(t, binary.Write(&buf, binary.LittleEndian, &shstrtabShdr)) + + // Section 2: wrapper config + configShdr := elf.Section64{ + Name: uint32(configSectionNameOffset), + Type: uint32(elf.SHT_PROGBITS), + Flags: uint64(elf.SHF_ALLOC), + Addr: 0, + Off: configSectionOffset, + Size: 4096, + Link: 0, + Info: 0, + Addralign: 8, + Entsize: 0, + } + require.NoError(t, binary.Write(&buf, binary.LittleEndian, &configShdr)) + + // Section 3: alignment + dummyShdr := elf.Section64{} + require.NoError(t, binary.Write(&buf, binary.LittleEndian, &dummyShdr)) + + // Pad to shstrtab offset and write string table + for buf.Len() < int(shstrtabOffset) { + buf.WriteByte(0) + } + buf.Write(shstrtab) + + // Pad to wrapper config section and zero fill + for buf.Len() < int(configSectionOffset) { + buf.WriteByte(0) + } + buf.Write(make([]byte, 4096)) + + return &buf +} diff --git a/cmd/nvidia-ctk-installer/toolkit/installer/testutil/writeconfig/main.go b/cmd/nvidia-ctk-installer/toolkit/installer/testutil/writeconfig/main.go new file mode 100644 index 000000000..1b4a107fb --- /dev/null +++ b/cmd/nvidia-ctk-installer/toolkit/installer/testutil/writeconfig/main.go @@ -0,0 +1,65 @@ +/** +# 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 + +import ( + "debug/elf" + "encoding/json" + "fmt" + "os" + + "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore" +) + +func main() { + if len(os.Args) < 3 { + fmt.Fprintf(os.Stderr, "Usage: %s \n", os.Args[0]) + os.Exit(1) + } + + var config wrappercore.WrapperConfig + if err := json.Unmarshal([]byte(os.Args[2]), &config); err != nil { + fmt.Fprintf(os.Stderr, "Failed to parse wrapper config: %v\n", err) + os.Exit(1) + } + + file, err := os.OpenFile(os.Args[1], os.O_RDWR, 0) + if err != nil { + fmt.Fprintf(os.Stderr, "Failed to open wrapper program: %v\n", err) + os.Exit(1) + } + elfFile, err := elf.NewFile(file) + if err != nil { + fmt.Fprintf(os.Stderr, "Failed to parse ELF file: %v\n", err) + os.Exit(1) + } + configSection := elfFile.Section(wrappercore.InstallerWrapperConfigELFSectionName) + if configSection == nil { + fmt.Fprintf(os.Stderr, "Wrapper config section %q not found\n", wrappercore.InstallerWrapperConfigELFSectionName) + os.Exit(1) + } + sectionData := make([]byte, configSection.Size) + if err := wrappercore.WriteWrapperConfigSection(sectionData, &config); err != nil { + fmt.Fprintf(os.Stderr, "Failed to write wrapper config: %v\n", err) + os.Exit(1) + } + if _, err := file.WriteAt(sectionData, int64(configSection.Offset)); err != nil { + fmt.Fprintf(os.Stderr, "Failed to write section data to file: %v\n", err) + os.Exit(1) + } +} diff --git a/cmd/nvidia-ctk-installer/toolkit/installer/wrapper.go b/cmd/nvidia-ctk-installer/toolkit/installer/wrapper.go new file mode 100644 index 000000000..b89c086be --- /dev/null +++ b/cmd/nvidia-ctk-installer/toolkit/installer/wrapper.go @@ -0,0 +1,79 @@ +/** +# 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 installer + +import ( + "bytes" + "debug/elf" + "fmt" + "io" + "os" + "path/filepath" + + "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore" +) + +type wrapper struct { + Source string + WrapperProgramPath string + Config wrappercore.WrapperConfig +} + +func (w *wrapper) Install(destDir string) error { + sourceFilename := filepath.Base(w.Source) + sourceFilenameDotReal := sourceFilename + ".real" + + // Copy the source/original executable to the destination with a .real extension. + mode, err := installFile(w.Source, filepath.Join(destDir, sourceFilenameDotReal)) + if err != nil { + return err + } + + // Create a wrapper program with the original's filename. + content, err := w.render() + if err != nil { + return fmt.Errorf("failed to render wrapper: %w", err) + } + wrapperFile := filepath.Join(destDir, sourceFilename) + return installContent(content, wrapperFile, mode|0111) +} + +func (r *wrapper) render() (*bytes.Buffer, error) { + wrapperProgram, err := os.Open(r.WrapperProgramPath) + if err != nil { + return nil, fmt.Errorf("failed to open file: %v", err) + } + defer wrapperProgram.Close() + elfBytes, err := io.ReadAll(wrapperProgram) + if err != nil { + return nil, fmt.Errorf("failed to read wrapper program: %v", err) + } + elfBuf := bytes.NewBuffer(elfBytes) + elfFile, err := elf.NewFile(bytes.NewReader(elfBuf.Bytes())) + if err != nil { + return nil, fmt.Errorf("failed to parse ELF file: %v", err) + } + configSection := elfFile.Section(wrappercore.InstallerWrapperConfigELFSectionName) + if configSection == nil { + return nil, fmt.Errorf("wrapper config section not found") + } + if err := wrappercore.WriteWrapperConfigSection(elfBuf.Bytes()[configSection.Offset:configSection.Offset+configSection.Size], &r.Config); err != nil { + return nil, fmt.Errorf("failed to write wrapper config section: %v", err) + } + return elfBuf, nil +} diff --git a/cmd/nvidia-ctk-installer/toolkit/installer/wrapper_test.go b/cmd/nvidia-ctk-installer/toolkit/installer/wrapper_test.go new file mode 100644 index 000000000..5e151fec2 --- /dev/null +++ b/cmd/nvidia-ctk-installer/toolkit/installer/wrapper_test.go @@ -0,0 +1,196 @@ +/** +# 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 installer + +import ( + "bytes" + "debug/elf" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/testutil" + "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore" +) + +func TestWrapperRender(t *testing.T) { + elfPath := testutil.WriteTestELF(t, wrappercore.InstallerWrapperConfigELFSectionName) + + testCases := []struct { + description string + w *wrapper + }{ + { + description: "executable is added", + w: &wrapper{ + Config: wrappercore.WrapperConfig{ + DefaultRuntimeExecutablePath: "runc", + }, + }, + }, + { + description: "module check is added", + w: &wrapper{ + Config: wrappercore.WrapperConfig{ + RequiresKernelModule: true, + DefaultRuntimeExecutablePath: "runc", + }, + }, + }, + { + description: "environment is added", + w: &wrapper{ + Config: wrappercore.WrapperConfig{ + Envm: map[string]string{ + "PATH": "/foo/bar/baz", + }, + DefaultRuntimeExecutablePath: "runc", + }, + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + tc.w.WrapperProgramPath = elfPath + buffer, err := tc.w.render() + require.NoError(t, err) + elfFile, err := elf.NewFile(bytes.NewReader(buffer.Bytes())) + require.NoError(t, err) + section := elfFile.Section(wrappercore.InstallerWrapperConfigELFSectionName) + require.NotNil(t, section) + sectionData, err := section.Data() + require.NoError(t, err) + config, err := wrappercore.ReadWrapperConfigSection(sectionData) + require.NoError(t, err) + require.Equal(t, tc.w.Config, *config) + }) + } +} + +func TestWrapperRender_Errors(t *testing.T) { + tmpDir := t.TempDir() + + t.Run("section not found", func(t *testing.T) { + elfBuf := testutil.CreateTestELF(t, ".foobar") + elfPath := filepath.Join(tmpDir, "test-wrapper") + f, err := os.Create(elfPath) + require.NoError(t, err) + defer f.Close() + _, err = f.Write(elfBuf.Bytes()) + require.NoError(t, err) + w := &wrapper{ + WrapperProgramPath: elfPath, + } + _, err = w.render() + require.Contains(t, err.Error(), "section not found") + }) +} + +func TestWriteWrapperConfigSection(t *testing.T) { + testCases := []struct { + description string + config wrappercore.WrapperConfig + }{ + { + description: "empty config", + config: wrappercore.WrapperConfig{}, + }, + { + description: "basic config", + config: wrappercore.WrapperConfig{ + DefaultRuntimeExecutablePath: "runc", + }, + }, + { + description: "full config", + config: wrappercore.WrapperConfig{ + Argv: []string{"--flag1", "--flag2"}, + Envm: map[string]string{ + "PATH": "/usr/local/bin", + "CUSTOM_VAR": "/append/path", + }, + RequiresKernelModule: true, + DefaultRuntimeExecutablePath: "/usr/bin/crun", + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + elfBuf := bytes.NewBuffer(make([]byte, 4096)) + err := wrappercore.WriteWrapperConfigSection(elfBuf.Bytes(), &tc.config) + require.NoError(t, err) + readConfig, err := wrappercore.ReadWrapperConfigSection(elfBuf.Bytes()) + require.NoError(t, err) + require.Equal(t, tc.config, *readConfig) + }) + } +} + +func TestWriteWrapperConfigSection_Errors(t *testing.T) { + t.Run("data too large", func(t *testing.T) { + largeConfig := wrappercore.WrapperConfig{ + Envm: make(map[string]string), + } + for i := 0; i < 500; i++ { + largeConfig.Envm[fmt.Sprintf("VAR_%d", i)] = strings.Repeat("x", 100) + } + elfBuf := bytes.NewBuffer(make([]byte, 4096)) + err := wrappercore.WriteWrapperConfigSection(elfBuf.Bytes(), &largeConfig) + require.Error(t, err) + require.Contains(t, err.Error(), "exceeds section size") + }) +} + +func TestReadWrapperConfigSection_Errors(t *testing.T) { + testCases := []struct { + description string + sectionData []byte + expectError string + }{ + { + description: "too small for size header", + sectionData: []byte{0x01, 0x02}, + expectError: "too small", + }, + { + description: "size mismatch", + sectionData: []byte{0xFF, 0xFF, 0xFF, 0xFF, 0x01, 0x02}, + expectError: "size mismatch", + }, + { + description: "invalid JSON", + sectionData: append([]byte{0x05, 0x00, 0x00, 0x00}, []byte("{not json")...), + expectError: "unmarshal", + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + _, err := wrappercore.ReadWrapperConfigSection(tc.sectionData) + require.Error(t, err) + require.Contains(t, err.Error(), tc.expectError) + }) + } +} diff --git a/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore/wrappercore.go b/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore/wrappercore.go new file mode 100644 index 000000000..70c5a2aa6 --- /dev/null +++ b/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore/wrappercore.go @@ -0,0 +1,81 @@ +/** +# 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 wrappercore provides program wrapper functionality shared between the installer package +// and the wrapper program. The installer package imports many additional Go and CGO modules that +// should be excluded from the wrapper program binary to minimize its size. +package wrappercore + +import ( + "encoding/binary" + "encoding/json" + "fmt" +) + +const InstallerWrapperConfigELFSectionName = ".nvctkiwc" +const InstallerWrapperFilename = "nvidia-ctk-installer-wrapper" + +// WrapperConfig is the configuration for the installer wrapper program. +// +// The configuration for a wrapper is embedded in a dedicated ELF section with a short binary prefix +// using the `WriteWrapperConfigSection` function. +type WrapperConfig struct { + Argv []string `json:"argv,omitempty"` + // Envm is the environment variable map for the wrapped executable. + // + // The wrapper supports prepending and appending to path list variables (like PATH). + // + // - '<': prepend (e.g. '': append (e.g. '>PATH') + Envm map[string]string `json:"envm,omitempty"` + RequiresKernelModule bool `json:"requiresKernelModule,omitempty"` + DefaultRuntimeExecutablePath string `json:"defaultRuntimeExecutablePath,omitempty"` +} + +// WriteWrapperConfigSection writes a wrapper config into the provided section buffer. +func WriteWrapperConfigSection(sectionData []byte, config *WrapperConfig) error { + configData, err := json.Marshal(config) + if err != nil { + return fmt.Errorf("failed to marshal wrapper config: %v", err) + } + + if int64(len(configData))+4 > int64(len(sectionData)) { + return fmt.Errorf("data size %d (+ size header) exceeds section size %d", len(configData), len(sectionData)) + } + var sizeBuf [4]byte + binary.NativeEndian.PutUint32(sizeBuf[:], uint32(len(configData))) + copy(sectionData, sizeBuf[:]) + copy(sectionData[4:], configData) + return nil +} + +// ReadWrapperConfigSection unmarshals the wrapper config from the provided section buffer. +func ReadWrapperConfigSection(sectionData []byte) (*WrapperConfig, error) { + if len(sectionData) < 4 { + return nil, fmt.Errorf("wrapper config section too small: %d bytes", len(sectionData)) + } + dataSize := int64(binary.NativeEndian.Uint32(sectionData[:4])) + if int64(len(sectionData)) < 4+dataSize { + return nil, fmt.Errorf("wrapper config section data size mismatch: header says %d but only %d bytes available", dataSize, len(sectionData)-4) + } + configData := sectionData[4 : 4+dataSize] + var config WrapperConfig + if err := json.Unmarshal(configData, &config); err != nil { + return nil, fmt.Errorf("failed to unmarshal wrapper config: %v", err) + } + return &config, nil +} diff --git a/cmd/nvidia-ctk-installer/toolkit/options.go b/cmd/nvidia-ctk-installer/toolkit/options.go index 10e49b95d..6d1271ceb 100644 --- a/cmd/nvidia-ctk-installer/toolkit/options.go +++ b/cmd/nvidia-ctk-installer/toolkit/options.go @@ -38,3 +38,9 @@ func WithSourceRoot(sourceRoot string) Option { i.sourceRoot = sourceRoot } } + +func WithWrapperProgramPath(path string) Option { + return func(i *Installer) { + i.wrapperProgramPath = path + } +} diff --git a/cmd/nvidia-ctk-installer/toolkit/toolkit.go b/cmd/nvidia-ctk-installer/toolkit/toolkit.go index 73d887572..c6f166297 100644 --- a/cmd/nvidia-ctk-installer/toolkit/toolkit.go +++ b/cmd/nvidia-ctk-installer/toolkit/toolkit.go @@ -218,7 +218,8 @@ type Installer struct { sourceRoot string // toolkitRoot specifies the destination path at which the toolkit is installed. - toolkitRoot string + toolkitRoot string + wrapperProgramPath string } // NewInstaller creates an installer for the NVIDIA Container Toolkit. @@ -317,6 +318,7 @@ func (t *Installer) Install(cli *cli.Command, opts *Options, runtime string) err installer.WithSourceRoot(t.sourceRoot), installer.WithIgnoreErrors(opts.ignoreErrors), installer.WithDefaultRuntimeExecutablePath(defaultRuntimeExecutable), + installer.WithWrapperProgramPath(t.wrapperProgramPath), ) if err != nil { if !opts.ignoreErrors { diff --git a/cmd/nvidia-ctk-installer/toolkit/toolkit_test.go b/cmd/nvidia-ctk-installer/toolkit/toolkit_test.go index 5ac6be20d..fed9ef04f 100644 --- a/cmd/nvidia-ctk-installer/toolkit/toolkit_test.go +++ b/cmd/nvidia-ctk-installer/toolkit/toolkit_test.go @@ -27,6 +27,7 @@ import ( "github.com/stretchr/testify/require" "github.com/urfave/cli/v3" + "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore" "github.com/NVIDIA/nvidia-container-toolkit/internal/config" "github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/symlinks" "github.com/NVIDIA/nvidia-container-toolkit/internal/test" @@ -40,6 +41,7 @@ func TestInstall(t *testing.T) { require.NoError(t, err) artifactRoot := filepath.Join(moduleRoot, "testdata", "installer", "artifacts") + wrapperProgramPath := filepath.Join(artifactRoot, "deb", wrappercore.InstallerWrapperFilename) testCases := []struct { description string @@ -153,6 +155,7 @@ containerEdits: WithLogger(logger), WithToolkitRoot(toolkitRoot), WithSourceRoot(sourceRoot), + WithWrapperProgramPath(wrapperProgramPath), ) require.NoError(t, ti.ValidateOptions(&options)) diff --git a/deployments/container/Dockerfile b/deployments/container/Dockerfile index 2617fa8ee..5a62381a7 100644 --- a/deployments/container/Dockerfile +++ b/deployments/container/Dockerfile @@ -52,6 +52,7 @@ RUN mkdir -p /artifacts/bin ARG VERSION="N/A" ARG GIT_COMMIT="unknown" RUN make PREFIX=/artifacts/bin cmd-nvidia-ctk-installer +RUN make PREFIX=/artifacts/bin cmd-nvidia-ctk-installer-wrapper # The rpmdigests stage updates the existing rpm packages to have 256bit digests. # This is done using fpm. diff --git a/testdata/installer/artifacts/deb/nvidia-ctk-installer-wrapper b/testdata/installer/artifacts/deb/nvidia-ctk-installer-wrapper new file mode 100755 index 000000000..dcfd7443a Binary files /dev/null and b/testdata/installer/artifacts/deb/nvidia-ctk-installer-wrapper differ