Create an internal/hooks pkg to centralize hook management

Signed-off-by: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>

internal/hooks/hooks.go

@@ -0,0 +1,116 @@
+/**
+# Copyright (c) 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 hooks
+
+import (
+	"path/filepath"
+
+	"tags.cncf.io/container-device-interface/pkg/cdi"
+)
+
+// A HookName refers to one of the predefined set of CDI hooks that may be
+// included in the generated CDI specification.
+type HookName string
+
+// DisabledHooks allows individual hooks to be disabled.
+type DisabledHooks map[HookName]bool
+
+const (
+	// EnableCudaCompat refers to the hook used to enable CUDA Forward Compatibility.
+	// This was added with v1.17.5 of the NVIDIA Container Toolkit.
+	EnableCudaCompat = HookName("enable-cuda-compat")
+	// CreateSymlinks refers to the hook used to create symlinks for the NVIDIA
+	// Container Toolkit. This was added with v1.17.5 of the NVIDIA Container Toolkit.
+	CreateSymlinks = HookName("create-symlinks")
+	// UpdateLDCache refers to the hook used to update the LD cache for the NVIDIA
+	// Container Toolkit. This was added with v1.17.5 of the NVIDIA Container Toolkit.
+	UpdateLDCache = HookName("update-ldcache")
+)
+
+// Hook represents an OCI container hook.
+type Hook struct {
+	Lifecycle string
+	Path      string
+	Args      []string
+}
+
+// Option is a function that configures the nvcdilib
+type Option func(*CDIHook)
+
+type CDIHook struct {
+	nvidiaCDIHookPath string
+	disabledHooks     DisabledHooks
+}
+
+type HookCreator interface {
+	Create(HookName, ...string) *Hook
+}
+
+func NewHookCreator(nvidiaCDIHookPath string, disabledHooks DisabledHooks) HookCreator {
+	if disabledHooks == nil {
+		disabledHooks = make(DisabledHooks)
+	}
+
+	CDIHook := &CDIHook{
+		nvidiaCDIHookPath: nvidiaCDIHookPath,
+		disabledHooks:     disabledHooks,
+	}
+
+	return CDIHook
+}
+
+func (c CDIHook) Create(name HookName, args ...string) *Hook {
+	if c.disabledHooks[name] {
+		return nil
+	}
+
+	if name == CreateSymlinks {
+		if len(args) == 0 {
+			return nil
+		}
+
+		links := []string{}
+		for _, arg := range args {
+			links = append(links, "--link", arg)
+		}
+		args = links
+	}
+
+	return &Hook{
+		Lifecycle: cdi.CreateContainerHook,
+		Path:      c.nvidiaCDIHookPath,
+		Args:      append(c.requiredArgs(name), args...),
+	}
+}
+
+func (c CDIHook) requiredArgs(name HookName) []string {
+	base := filepath.Base(c.nvidiaCDIHookPath)
+	if base == "nvidia-ctk" {
+		return []string{base, "hook", string(name)}
+	}
+	return []string{base, string(name)}
+}
+
+// IsSupported checks whether a hook of the specified name is supported.
+// Hooks must be explicitly disabled, meaning that if no disabled hooks are
+// all hooks are supported.
+func (c CDIHook) IsSupported(h HookName) bool {
+	if len(c.disabledHooks) == 0 {
+		return true
+	}
+	return !c.disabledHooks[h]
+}

internal/hooks/hooks_test.go

@@ -0,0 +1,177 @@
+package hooks
+
+import (
+	"testing"
+)
+
+func TestNewHookCreator(t *testing.T) {
+	tests := []struct {
+		name            string
+		hookPath        string
+		disabledHooks   DisabledHooks
+		expectedCreator HookCreator
+	}{
+		{
+			name:     "nil disabled hooks",
+			hookPath: "/usr/bin/nvidia-ctk",
+			expectedCreator: &CDIHook{
+				nvidiaCDIHookPath: "/usr/bin/nvidia-ctk",
+				disabledHooks:     DisabledHooks{},
+			},
+		},
+		{
+			name:     "with disabled hooks",
+			hookPath: "/usr/bin/nvidia-ctk",
+			disabledHooks: DisabledHooks{
+				EnableCudaCompat: true,
+			},
+			expectedCreator: &CDIHook{
+				nvidiaCDIHookPath: "/usr/bin/nvidia-ctk",
+				disabledHooks: DisabledHooks{
+					EnableCudaCompat: true,
+				},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			creator := NewHookCreator(tt.hookPath, tt.disabledHooks)
+			if creator == nil {
+				t.Fatal("NewHookCreator returned nil")
+			}
+		})
+	}
+}
+
+func TestCDIHook_Create(t *testing.T) {
+	tests := []struct {
+		name          string
+		hookPath      string
+		disabledHooks DisabledHooks
+		hookName      HookName
+		args          []string
+		expectedHook  *Hook
+	}{
+		{
+			name:     "disabled hook returns nil",
+			hookPath: "/usr/bin/nvidia-ctk",
+			disabledHooks: DisabledHooks{
+				EnableCudaCompat: true,
+			},
+			hookName:     EnableCudaCompat,
+			expectedHook: nil,
+		},
+		{
+			name:         "create symlinks with no args returns nil",
+			hookPath:     "/usr/bin/nvidia-ctk",
+			hookName:     CreateSymlinks,
+			expectedHook: nil,
+		},
+		{
+			name:     "create symlinks with args",
+			hookPath: "/usr/bin/nvidia-ctk",
+			hookName: CreateSymlinks,
+			args:     []string{"/path/to/lib1", "/path/to/lib2"},
+			expectedHook: &Hook{
+				Lifecycle: "createContainer",
+				Path:      "/usr/bin/nvidia-ctk",
+				Args:      []string{"nvidia-ctk", "hook", "create-symlinks", "--link", "/path/to/lib1", "--link", "/path/to/lib2"},
+			},
+		},
+		{
+			name:     "enable cuda compat",
+			hookPath: "/usr/bin/nvidia-ctk",
+			hookName: EnableCudaCompat,
+			expectedHook: &Hook{
+				Lifecycle: "createContainer",
+				Path:      "/usr/bin/nvidia-ctk",
+				Args:      []string{"nvidia-ctk", "hook", "enable-cuda-compat"},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			hook := &CDIHook{
+				nvidiaCDIHookPath: tt.hookPath,
+				disabledHooks:     tt.disabledHooks,
+			}
+
+			result := hook.Create(tt.hookName, tt.args...)
+			if tt.expectedHook == nil {
+				if result != nil {
+					t.Errorf("expected nil hook, got %v", result)
+				}
+				return
+			}
+
+			if result == nil {
+				t.Fatal("expected non-nil hook, got nil")
+			}
+
+			if result.Lifecycle != tt.expectedHook.Lifecycle {
+				t.Errorf("expected lifecycle %q, got %q", tt.expectedHook.Lifecycle, result.Lifecycle)
+			}
+
+			if result.Path != tt.expectedHook.Path {
+				t.Errorf("expected path %q, got %q", tt.expectedHook.Path, result.Path)
+			}
+
+			if len(result.Args) != len(tt.expectedHook.Args) {
+				t.Errorf("expected %d args, got %d", len(tt.expectedHook.Args), len(result.Args))
+				return
+			}
+
+			for i, arg := range tt.expectedHook.Args {
+				if result.Args[i] != arg {
+					t.Errorf("expected arg[%d] %q, got %q", i, arg, result.Args[i])
+				}
+			}
+		})
+	}
+}
+
+func TestCDIHook_IsSupported(t *testing.T) {
+	tests := []struct {
+		name          string
+		disabledHooks DisabledHooks
+		hookName      HookName
+		expected      bool
+	}{
+		{
+			name:     "no disabled hooks",
+			hookName: EnableCudaCompat,
+			expected: true,
+		},
+		{
+			name: "disabled hook",
+			disabledHooks: DisabledHooks{
+				EnableCudaCompat: true,
+			},
+			hookName: EnableCudaCompat,
+			expected: false,
+		},
+		{
+			name: "non-disabled hook",
+			disabledHooks: DisabledHooks{
+				EnableCudaCompat: true,
+			},
+			hookName: CreateSymlinks,
+			expected: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			hook := &CDIHook{
+				disabledHooks: tt.disabledHooks,
+			}
+
+			result := hook.IsSupported(tt.hookName)
+			if result != tt.expected {
+				t.Errorf("expected %v, got %v", tt.expected, result)
+			}
+		})
+	}
+}