implement golang wrapper to replace shell scripts

Some platforms and Kubernetes distributions do not include a shell. This
patch replaces the shell wrapper scripts with a small Go program.

Signed-off-by: Jean-Francois Roy <jeroy@nvidia.com>
This commit is contained in:
Jean-Francois Roy 2024-09-18 08:47:21 -07:00
parent ab0a4eaa19
commit dc79fa8513

View File

@ -0,0 +1,115 @@
/**
# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
#
# 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 (
"bufio"
"errors"
"io/fs"
"log"
"os"
"os/exec"
"path/filepath"
"strings"
"golang.org/x/sys/unix"
)
func main() {
program, err := os.Executable()
if err != nil {
log.Fatalf("failed to get executable: %v", err)
}
if isRuntimeWrapper(program) && !isNvidiaModuleLoaded() {
log.Println("nvidia driver modules are not yet loaded, invoking runc directly")
program, err := exec.LookPath("runc")
if err != nil {
log.Fatalf("failed to find runc: %v", err)
}
argv := []string{"runc"}
argv = append(argv, os.Args[1:]...)
execve(program, argv, os.Environ())
}
argv := makeArgv(program)
envv := makeEnvv(program)
execve(program+".real", argv, envv)
}
func isRuntimeWrapper(program string) bool {
return filepath.Base(program) == "nvidia-container-runtime" ||
filepath.Base(program) == "nvidia-container-runtime.cdi" ||
filepath.Base(program) == "nvidia-container-runtime.legacy"
}
func isNvidiaModuleLoaded() bool {
_, err := os.Stat("/proc/driver/nvidia/version")
return err == nil
}
func makeArgv(program string) []string {
argv := []string{os.Args[0] + ".real"}
f, err := os.Open(program + ".argv")
if err != nil {
if !errors.Is(err, fs.ErrNotExist) {
log.Printf("failed to open argv file: %v", err)
}
return append(argv, os.Args[1:]...)
}
defer f.Close()
scanner := bufio.NewScanner(f)
for scanner.Scan() {
argv = append(argv, scanner.Text())
}
if err := scanner.Err(); err != nil {
log.Fatalf("failed to read argv file: %v", err)
}
return append(argv, os.Args[1:]...)
}
func makeEnvv(program string) []string {
f, err := os.Open(program + ".envv")
if err != nil {
if !errors.Is(err, fs.ErrNotExist) {
log.Printf("failed to open env file: %v", err)
}
return os.Environ()
}
defer f.Close()
var env []string
scanner := bufio.NewScanner(f)
for scanner.Scan() {
kv := strings.SplitN(scanner.Text(), "=", 2)
if strings.HasPrefix(kv[0], "<") {
kv[0] = kv[0][1:]
kv[1] = kv[1] + ":" + os.Getenv(kv[0])
} else if strings.HasPrefix(kv[0], ">") {
kv[0] = kv[0][1:]
kv[1] = os.Getenv(kv[0]) + ":" + kv[1]
}
env = append(env, kv[0]+"="+kv[1])
}
if err := scanner.Err(); err != nil {
log.Fatalf("failed to read argv file: %v", err)
}
return append(env, os.Environ()...)
}
func execve(program string, argv []string, envv []string) {
if err := unix.Exec(program, argv, envv); err != nil {
log.Fatalf("failed to exec %s: %v", program, err)
}
}