Minor rework of disabled hooks

Signed-off-by: Evan Lezar <[email protected]>

Author: elezar

cmd/nvidia-ctk/cdi/generate/generate.go

@@ -270,7 +270,7 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 		deviceNamers = append(deviceNamers, deviceNamer)
 	}
 
-	initOpts := []nvcdi.Option{
+	cdiOptions := []nvcdi.Option{
 		nvcdi.WithLogger(m.logger),
 		nvcdi.WithDriverRoot(opts.driverRoot),
 		nvcdi.WithDevRoot(opts.devRoot),
@@ -287,10 +287,10 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 	}
 
 	for _, hook := range opts.disabledHooks.Value() {
-		initOpts = append(initOpts, nvcdi.WithDisabledHook(hook))
+		cdiOptions = append(cdiOptions, nvcdi.WithDisabledHook(hook))
 	}
 
-	cdilib, err := nvcdi.New(initOpts...)
+	cdilib, err := nvcdi.New(cdiOptions...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create CDI library: %v", err)
 	}

cmd/nvidia-ctk/cdi/generate/generate_test.go

@@ -37,9 +37,6 @@ func TestGenerateSpec(t *testing.T) {
 	require.NoError(t, err)
 
 	driverRoot := filepath.Join(moduleRoot, "testdata", "lookup", "rootfs-1")
-	disableHook1 := cli.NewStringSlice("enable-cuda-compat")
-	disableHook2 := cli.NewStringSlice("enable-cuda-compat", "update-ldcache")
-	disableHook3 := cli.NewStringSlice("all")
 
 	logger, _ := testlog.NewNullLogger()
 	testCases := []struct {
@@ -127,7 +124,7 @@ containerEdits:
 				vendor:        "example.com",
 				class:         "device",
 				driverRoot:    driverRoot,
-				disabledHooks: *disableHook1,
+				disabledHooks: valueOf(cli.NewStringSlice("enable-cuda-compat")),
 			},
 			expectedOptions: options{
 				format:            "yaml",
@@ -136,7 +133,7 @@ containerEdits:
 				class:             "device",
 				nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
 				driverRoot:        driverRoot,
-				disabledHooks:     *disableHook1,
+				disabledHooks:     valueOf(cli.NewStringSlice("enable-cuda-compat")),
 			},
 			expectedSpec: `---
 cdiVersion: 0.5.0
@@ -192,7 +189,7 @@ containerEdits:
 				vendor:        "example.com",
 				class:         "device",
 				driverRoot:    driverRoot,
-				disabledHooks: *disableHook2,
+				disabledHooks: valueOf(cli.NewStringSlice("enable-cuda-compat", "update-ldcache")),
 			},
 			expectedOptions: options{
 				format:            "yaml",
@@ -201,7 +198,7 @@ containerEdits:
 				class:             "device",
 				nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
 				driverRoot:        driverRoot,
-				disabledHooks:     *disableHook2,
+				disabledHooks:     valueOf(cli.NewStringSlice("enable-cuda-compat", "update-ldcache")),
 			},
 			expectedSpec: `---
 cdiVersion: 0.5.0
@@ -250,7 +247,7 @@ containerEdits:
 				vendor:        "example.com",
 				class:         "device",
 				driverRoot:    driverRoot,
-				disabledHooks: *disableHook3,
+				disabledHooks: valueOf(cli.NewStringSlice("all")),
 			},
 			expectedOptions: options{
 				format:            "yaml",
@@ -259,7 +256,7 @@ containerEdits:
 				class:             "device",
 				nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
 				driverRoot:        driverRoot,
-				disabledHooks:     *disableHook3,
+				disabledHooks:     valueOf(cli.NewStringSlice("all")),
 			},
 			expectedSpec: `---
 cdiVersion: 0.5.0
@@ -333,3 +330,9 @@ containerEdits:
 		})
 	}
 }
+
+// valueOf returns the value of a pointer.
+// Note that this does not check for a nil pointer and is only used for testing.
+func valueOf[T any](v *T) T {
+	return *v
+}

internal/discover/graphics_test.go

@@ -25,7 +25,7 @@ import (
 
 func TestGraphicsLibrariesDiscoverer(t *testing.T) {
 	logger, _ := testlog.NewNullLogger()
-	hookCreator := NewHookCreator("/usr/bin/nvidia-cdi-hook")
+	hookCreator := NewHookCreator()
 
 	testCases := []struct {
 		description    string

internal/discover/hooks.go

@@ -22,6 +22,28 @@ import (
 	"tags.cncf.io/container-device-interface/pkg/cdi"
 )
 
+// A HookName represents a supported CDI hooks.
+type HookName string
+
+const (
+	// AllHooks is a special hook name that allows all hooks to be matched.
+	AllHooks = HookName("all")
+
+	// A ChmodHook is used to set the file mode of the specified paths.
+	// Deprecated: The chmod hook is deprecated and will be removed in a future release.
+	ChmodHook = HookName("chmod")
+	// A CreateSymlinksHook is used to create symlinks in the container.
+	CreateSymlinksHook = HookName("create-symlinks")
+	// An EnableCudaCompatHook is used to enabled CUDA Forward Compatibility.
+	// Added in v1.17.5
+	EnableCudaCompatHook = HookName("enable-cuda-compat")
+	// An UpdateLDCacheHook is the hook used to update the ldcache in the
+	// container. This allows injected libraries to be discoverable.
+	UpdateLDCacheHook = HookName("update-ldcache")
+
+	defaultNvidiaCDIHookPath = "/usr/bin/nvidia-cdi-hook"
+)
+
 var _ Discover = (*Hook)(nil)
 
 // Devices returns an empty list of devices for a Hook discoverer.
@@ -44,89 +66,128 @@ func (h *Hook) Hooks() ([]Hook, error) {
 	return []Hook{*h}, nil
 }
 
-type HookName string
-
-// disabledHooks allows individual hooks to be disabled.
-type disabledHooks map[HookName]bool
+type Option func(*cdiHookCreator)
 
-const (
-	// HookEnableCudaCompat refers to the hook used to enable CUDA Forward Compatibility.
-	// This was added with v1.17.5 of the NVIDIA Container Toolkit.
-	HookEnableCudaCompat = HookName("enable-cuda-compat")
-	// directory path to be mounted into a container.
-	HookCreateSymlinks = HookName("create-symlinks")
-	// HookUpdateLDCache refers to the hook used to  Update the dynamic linker
-	// cache inside the directory path to be mounted into a container.
-	HookUpdateLDCache = HookName("update-ldcache")
-)
+type cdiHookCreator struct {
+	nvidiaCDIHookPath string
+	disabledHooks     map[HookName]bool
 
-// AllHooks maintains a future-proof list of all defined hooks.
-var AllHooks = []HookName{
-	HookEnableCudaCompat,
-	HookCreateSymlinks,
-	HookUpdateLDCache,
+	fixedArgs []string
 }
 
-type Option func(*CDIHook)
+// An allDisabledHookCreator is a HookCreator that does not create any hooks.
+type allDisabledHookCreator struct{}
 
-type CDIHook struct {
-	nvidiaCDIHookPath string
-	disabledHooks     disabledHooks
+// Create returns nil for all hooks for an allDisabledHookCreator.
+func (a *allDisabledHookCreator) Create(name HookName, args ...string) *Hook {
+	return nil
 }
 
+// A HookCreator defines an interface for creating discover hooks.
 type HookCreator interface {
 	Create(HookName, ...string) *Hook
 }
 
+// WithDisabledHooks sets the set of hooks that are disabled for the CDI hook creator.
+// This can be specified multiple times.
 func WithDisabledHooks(hooks ...HookName) Option {
-	return func(c *CDIHook) {
+	return func(c *cdiHookCreator) {
 		for _, hook := range hooks {
 			c.disabledHooks[hook] = true
 		}
 	}
 }
 
-func NewHookCreator(nvidiaCDIHookPath string, opts ...Option) HookCreator {
-	CDIHook := &CDIHook{
-		nvidiaCDIHookPath: nvidiaCDIHookPath,
-		disabledHooks:     disabledHooks{},
+// WithNVIDIACDIHookPath sets the path to the nvidia-cdi-hook binary.
+func WithNVIDIACDIHookPath(nvidiaCDIHookPath string) Option {
+	return func(c *cdiHookCreator) {
+		c.nvidiaCDIHookPath = nvidiaCDIHookPath
 	}
+}
 
+func NewHookCreator(opts ...Option) HookCreator {
+	cdiHookCreator := &cdiHookCreator{
+		nvidiaCDIHookPath: defaultNvidiaCDIHookPath,
+		disabledHooks:     make(map[HookName]bool),
+	}
 	for _, opt := range opts {
-		opt(CDIHook)
+		opt(cdiHookCreator)
+	}
+
+	if cdiHookCreator.disabledHooks[AllHooks] {
+		return &allDisabledHookCreator{}
 	}
 
-	return CDIHook
+	cdiHookCreator.fixedArgs = getFixedArgsForCDIHookCLI(cdiHookCreator.nvidiaCDIHookPath)
+
+	return cdiHookCreator
 }
 
-func (c CDIHook) Create(name HookName, args ...string) *Hook {
-	if c.disabledHooks[name] {
+// Create creates a new hook with the given name and arguments.
+// If a hook is disabled, a nil hook is returned.
+func (c cdiHookCreator) Create(name HookName, args ...string) *Hook {
+	if c.isDisabled(name, args...) {
 		return nil
 	}
 
-	if name == "create-symlinks" {
+	return &Hook{
+		Lifecycle: cdi.CreateContainerHook,
+		Path:      c.nvidiaCDIHookPath,
+		Args:      append(c.requiredArgs(name), c.transformArgs(name, args...)...),
+	}
+}
+
+// isDisabled checks if the specified hook name is disabled.
+func (c cdiHookCreator) isDisabled(name HookName, args ...string) bool {
+	if c.disabledHooks[name] {
+		return true
+	}
+
+	switch name {
+	case CreateSymlinksHook:
 		if len(args) == 0 {
-			return nil
+			return true
 		}
-
-		links := []string{}
-		for _, arg := range args {
-			links = append(links, "--link", arg)
+	case ChmodHook:
+		if len(args) == 0 {
+			return true
 		}
-		args = links
 	}
+	return false
+}
 
-	return &Hook{
-		Lifecycle: cdi.CreateContainerHook,
-		Path:      c.nvidiaCDIHookPath,
-		Args:      append(c.requiredArgs(string(name)), args...),
+func (c cdiHookCreator) requiredArgs(name HookName) []string {
+	return append(c.fixedArgs, string(name))
+}
+
+func (c cdiHookCreator) transformArgs(name HookName, args ...string) []string {
+	switch name {
+	case CreateSymlinksHook:
+		var transformedArgs []string
+		for _, arg := range args {
+			transformedArgs = append(transformedArgs, "--link", arg)
+		}
+		return transformedArgs
+	case ChmodHook:
+		var transformedArgs = []string{"--mode", "755"}
+		for _, arg := range args {
+			transformedArgs = append(transformedArgs, "--path", arg)
+		}
+		return transformedArgs
+	default:
+		return args
 	}
 }
 
-func (c CDIHook) requiredArgs(name string) []string {
-	base := filepath.Base(c.nvidiaCDIHookPath)
+// getFixedArgsForCDIHookCLI returns the fixed arguments for the hook CLI.
+// If the nvidia-ctk binary is used, hooks are implemented under the hook
+// subcommand.
+// For the nvidia-cdi-hook binary, the hooks are implemented as subcommands of
+// the top-level CLI.
+func getFixedArgsForCDIHookCLI(nvidiaCDIHookPath string) []string {
+	base := filepath.Base(nvidiaCDIHookPath)
 	if base == "nvidia-ctk" {
-		return []string{base, "hook", name}
+		return []string{base, "hook"}
 	}
-	return []string{base, name}
+	return []string{base}
 }

internal/discover/ldconfig.go

@@ -72,7 +72,7 @@ func createLDCacheUpdateHook(hookCreator HookCreator, ldconfig string, libraries
 		args = append(args, "--folder", f)
 	}
 
-	return hookCreator.Create(HookUpdateLDCache, args...)
+	return hookCreator.Create(UpdateLDCacheHook, args...)
 }
 
 // getLibraryPaths extracts the library dirs from the specified mounts

internal/discover/ldconfig_test.go

@@ -31,7 +31,7 @@ const (
 
 func TestLDCacheUpdateHook(t *testing.T) {
 	logger, _ := testlog.NewNullLogger()
-	hookCreator := NewHookCreator(testNvidiaCDIHookPath)
+	hookCreator := NewHookCreator(WithNVIDIACDIHookPath(testNvidiaCDIHookPath))
 
 	testCases := []struct {
 		description   string

internal/discover/symlinks_test.go

@@ -113,7 +113,7 @@ func TestWithWithDriverDotSoSymlinks(t *testing.T) {
 			expectedHooks: []Hook{
 				{
 					Lifecycle: "createContainer",
-					Path:      "/path/to/nvidia-cdi-hook",
+					Path:      "/usr/bin/nvidia-cdi-hook",
 					Args:      []string{"nvidia-cdi-hook", "create-symlinks", "--link", "libcuda.so.1::/usr/lib/libcuda.so"},
 				},
 			},
@@ -145,7 +145,7 @@ func TestWithWithDriverDotSoSymlinks(t *testing.T) {
 			expectedHooks: []Hook{
 				{
 					Lifecycle: "createContainer",
-					Path:      "/path/to/nvidia-cdi-hook",
+					Path:      "/usr/bin/nvidia-cdi-hook",
 					Args:      []string{"nvidia-cdi-hook", "create-symlinks", "--link", "libcuda.so.1::/usr/lib/libcuda.so"},
 				},
 			},
@@ -176,7 +176,7 @@ func TestWithWithDriverDotSoSymlinks(t *testing.T) {
 			expectedHooks: []Hook{
 				{
 					Lifecycle: "createContainer",
-					Path:      "/path/to/nvidia-cdi-hook",
+					Path:      "/usr/bin/nvidia-cdi-hook",
 					Args:      []string{"nvidia-cdi-hook", "create-symlinks", "--link", "libcuda.so.1::/usr/lib/libcuda.so"},
 				},
 			},
@@ -245,7 +245,7 @@ func TestWithWithDriverDotSoSymlinks(t *testing.T) {
 				},
 				{
 					Lifecycle: "createContainer",
-					Path:      "/path/to/nvidia-cdi-hook",
+					Path:      "/usr/bin/nvidia-cdi-hook",
 					Args:      []string{"nvidia-cdi-hook", "create-symlinks", "--link", "libcuda.so.1::/usr/lib/libcuda.so"},
 				},
 			},
@@ -294,7 +294,7 @@ func TestWithWithDriverDotSoSymlinks(t *testing.T) {
 			expectedHooks: []Hook{
 				{
 					Lifecycle: "createContainer",
-					Path:      "/path/to/nvidia-cdi-hook",
+					Path:      "/usr/bin/nvidia-cdi-hook",
 					Args: []string{
 						"nvidia-cdi-hook", "create-symlinks",
 						"--link", "libcuda.so.1::/usr/lib/libcuda.so",
@@ -306,7 +306,7 @@ func TestWithWithDriverDotSoSymlinks(t *testing.T) {
 		},
 	}
 
-	hookCreator := NewHookCreator("/path/to/nvidia-cdi-hook")
+	hookCreator := NewHookCreator()
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
 			d := WithDriverDotSoSymlinks(

internal/platform-support/tegra/csv_test.go

@@ -181,7 +181,7 @@ func TestDiscovererFromCSVFiles(t *testing.T) {
 		},
 	}
 
-	hookCreator := discover.NewHookCreator("/usr/bin/nvidia-cdi-hook")
+	hookCreator := discover.NewHookCreator()
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
 			defer setGetTargetsFromCSVFiles(tc.moutSpecs)()

internal/runtime/runtime_factory.go

@@ -75,7 +75,7 @@ func newSpecModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Sp
 		return nil, err
 	}
 
-	hookCreator := discover.NewHookCreator(cfg.NVIDIACTKConfig.Path)
+	hookCreator := discover.NewHookCreator(discover.WithNVIDIACDIHookPath(cfg.NVIDIACTKConfig.Path))
 
 	mode := info.ResolveAutoMode(logger, cfg.NVIDIAContainerRuntimeConfig.Mode, image)
 	// We update the mode here so that we can continue passing just the config to other functions.