Minor rework of disabled hooks

Signed-off-by: Evan Lezar <elezar@nvidia.com>

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
+type Option func(*cdiHookCreator)
 
-// disabledHooks allows individual hooks to be disabled.
-type disabledHooks map[HookName]bool
-
-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")
-)
-
-// AllHooks maintains a future-proof list of all defined hooks.
-var AllHooks = []HookName{
-	HookEnableCudaCompat,
-	HookCreateSymlinks,
-	HookUpdateLDCache,
-}
-
-type Option func(*CDIHook)
-
-type CDIHook struct {
+type cdiHookCreator struct {
 	nvidiaCDIHookPath string
-	disabledHooks     disabledHooks
+	disabledHooks     map[HookName]bool
+
+	fixedArgs []string
 }
 
+// An allDisabledHookCreator is a HookCreator that does not create any hooks.
+type allDisabledHookCreator struct{}
+
+// 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
 	}
-
-	for _, opt := range opts {
-		opt(CDIHook)
-	}
-
-	return CDIHook
 }
 
-func (c CDIHook) Create(name HookName, args ...string) *Hook {
-	if c.disabledHooks[name] {
-		return nil
+func NewHookCreator(opts ...Option) HookCreator {
+	cdiHookCreator := &cdiHookCreator{
+		nvidiaCDIHookPath: defaultNvidiaCDIHookPath,
+		disabledHooks:     make(map[HookName]bool),
+	}
+	for _, opt := range opts {
+		opt(cdiHookCreator)
 	}
 
-	if name == "create-symlinks" {
-		if len(args) == 0 {
-			return nil
-		}
+	if cdiHookCreator.disabledHooks[AllHooks] {
+		return &allDisabledHookCreator{}
+	}
 
-		links := []string{}
-		for _, arg := range args {
-			links = append(links, "--link", arg)
-		}
-		args = links
+	cdiHookCreator.fixedArgs = getFixedArgsForCDIHookCLI(cdiHookCreator.nvidiaCDIHookPath)
+
+	return cdiHookCreator
+}
+
+// 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
 	}
 
 	return &Hook{
 		Lifecycle: cdi.CreateContainerHook,
 		Path:      c.nvidiaCDIHookPath,
-		Args:      append(c.requiredArgs(string(name)), args...),
+		Args:      append(c.requiredArgs(name), c.transformArgs(name, args...)...),
 	}
 }
 
-func (c CDIHook) requiredArgs(name string) []string {
-	base := filepath.Base(c.nvidiaCDIHookPath)
-	if base == "nvidia-ctk" {
-		return []string{base, "hook", name}
+// isDisabled checks if the specified hook name is disabled.
+func (c cdiHookCreator) isDisabled(name HookName, args ...string) bool {
+	if c.disabledHooks[name] {
+		return true
 	}
-	return []string{base, name}
+
+	switch name {
+	case CreateSymlinksHook:
+		if len(args) == 0 {
+			return true
+		}
+	case ChmodHook:
+		if len(args) == 0 {
+			return true
+		}
+	}
+	return false
+}
+
+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
+	}
+}
+
+// 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"}
+	}
+	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.

pkg/nvcdi/api.go

@@ -37,12 +37,25 @@ type Interface interface {
 	GetDeviceSpecsByID(...string) ([]specs.Device, error)
 }
 
-// HookName is an alias for the discover.HookName type.
+// A HookName represents one of the predefined NVIDIA CDI hooks.
 type HookName = discover.HookName
 
-// Aliases for the discover.HookName constants.
 const (
-	HookEnableCudaCompat = discover.HookEnableCudaCompat
-	HookCreateSymlinks   = discover.HookCreateSymlinks
-	HookUpdateLDCache    = discover.HookUpdateLDCache
+	// AllHooks is a special hook name that allows all hooks to be matched.
+	AllHooks = discover.AllHooks
+
+	// A CreateSymlinksHook is used to create symlinks in the container.
+	CreateSymlinksHook = discover.CreateSymlinksHook
+	// An EnableCudaCompatHook is used to enabled CUDA Forward Compatibility.
+	// Added in v1.17.5
+	EnableCudaCompatHook = discover.EnableCudaCompatHook
+	// An UpdateLDCacheHook is used to update the ldcache in the container.
+	UpdateLDCacheHook = discover.UpdateLDCacheHook
+
+	// Deprecated: Use CreateSymlinksHook instead.
+	HookCreateSymlinks = CreateSymlinksHook
+	// Deprecated: Use EnableCudaCompatHook instead.
+	HookEnableCudaCompat = EnableCudaCompatHook
+	// Deprecated: Use UpdateLDCacheHook instead.
+	HookUpdateLDCache = UpdateLDCacheHook
 )

pkg/nvcdi/driver-wsl.go

@@ -135,7 +135,7 @@ func (m nvidiaSMISimlinkHook) Hooks() ([]discover.Hook, error) {
 	}
 	link := "/usr/bin/nvidia-smi"
 	links := []string{fmt.Sprintf("%s::%s", target, link)}
-	symlinkHook := m.hookCreator.Create(HookCreateSymlinks, links...)
+	symlinkHook := m.hookCreator.Create(CreateSymlinksHook, links...)
 
 	return symlinkHook.Hooks()
 }

pkg/nvcdi/driver-wsl_test.go

@@ -29,7 +29,7 @@ import (
 
 func TestNvidiaSMISymlinkHook(t *testing.T) {
 	logger, _ := testlog.NewNullLogger()
-	hookCreator := discover.NewHookCreator("nvidia-cdi-hook")
+	hookCreator := discover.NewHookCreator()
 
 	errMounts := errors.New("mounts error")
 
@@ -93,7 +93,7 @@ func TestNvidiaSMISymlinkHook(t *testing.T) {
 			expectedHooks: []discover.Hook{
 				{
 					Lifecycle: "createContainer",
-					Path:      "nvidia-cdi-hook",
+					Path:      "/usr/bin/nvidia-cdi-hook",
 					Args: []string{"nvidia-cdi-hook", "create-symlinks",
 						"--link", "nvidia-smi::/usr/bin/nvidia-smi"},
 				},
@@ -113,7 +113,7 @@ func TestNvidiaSMISymlinkHook(t *testing.T) {
 			expectedHooks: []discover.Hook{
 				{
 					Lifecycle: "createContainer",
-					Path:      "nvidia-cdi-hook",
+					Path:      "/usr/bin/nvidia-cdi-hook",
 					Args: []string{"nvidia-cdi-hook", "create-symlinks",
 						"--link", "/some/path/nvidia-smi::/usr/bin/nvidia-smi"},
 				},
@@ -133,7 +133,7 @@ func TestNvidiaSMISymlinkHook(t *testing.T) {
 			expectedHooks: []discover.Hook{
 				{
 					Lifecycle: "createContainer",
-					Path:      "nvidia-cdi-hook",
+					Path:      "/usr/bin/nvidia-cdi-hook",
 					Args: []string{"nvidia-cdi-hook", "create-symlinks",
 						"--link", "/some/path/nvidia-smi::/usr/bin/nvidia-smi"},
 				},

pkg/nvcdi/lib.go

@@ -172,7 +172,10 @@ func New(opts ...Option) (Interface, error) {
 	}
 
 	// create hookCreator
-	l.hookCreator = discover.NewHookCreator(l.nvidiaCDIHookPath, discover.WithDisabledHooks(l.disabledHooks...))
+	l.hookCreator = discover.NewHookCreator(
+		discover.WithNVIDIACDIHookPath(l.nvidiaCDIHookPath),
+		discover.WithDisabledHooks(l.disabledHooks...),
+	)
 
 	w := wrapper{
 		Interface:           lib,

pkg/nvcdi/options.go

@@ -161,10 +161,6 @@ func WithLibrarySearchPaths(paths []string) Option {
 // This option can be specified multiple times for each hook.
 func WithDisabledHook[T string | HookName](hook T) Option {
 	return func(o *nvcdilib) {
-		if hook == "all" {
-			o.disabledHooks = discover.AllHooks
-		} else {
-			o.disabledHooks = append(o.disabledHooks, discover.HookName(hook))
-		}
+		o.disabledHooks = append(o.disabledHooks, discover.HookName(hook))
 	}
 }

pkg/nvcdi/workarounds-device-folder-permissions.go

@@ -61,18 +61,9 @@ func (d *deviceFolderPermissions) Hooks() ([]discover.Hook, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to get device subfolders: %v", err)
 	}
-	if len(folders) == 0 {
-		return nil, nil
-	}
 
-	args := []string{"--mode", "755"}
-	for _, folder := range folders {
-		args = append(args, "--path", folder)
-	}
-
-	hook := d.hookCreator.Create("chmod", args...)
-
-	return []discover.Hook{*hook}, nil
+	//nolint:staticcheck // The ChmodHook is deprecated and will be removed in a future release.
+	return d.hookCreator.Create(discover.ChmodHook, folders...).Hooks()
 }
 
 func (d *deviceFolderPermissions) getDeviceSubfolders() ([]string, error) {