From eb4a252ca3c919838aa5c98417d144e5e34ecde1 Mon Sep 17 00:00:00 2001
From: Tariq Ibrahim <tibrahim@nvidia.com>
Date: Wed, 21 Jan 2026 21:35:49 -0800
Subject: [PATCH] restrict NRI injection of CDI devices to pods in the
 toolkit's namespace

Signed-off-by: Tariq Ibrahim <tibrahim@nvidia.com>
---
 .../container/runtime/nri/plugin.go           | 22 +++++++++++++------
 cmd/nvidia-ctk-installer/main.go              | 13 ++++++++++-
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/cmd/nvidia-ctk-installer/container/runtime/nri/plugin.go b/cmd/nvidia-ctk-installer/container/runtime/nri/plugin.go
index 7289f2d84..43cf01b01 100644
--- a/cmd/nvidia-ctk-installer/container/runtime/nri/plugin.go
+++ b/cmd/nvidia-ctk-installer/container/runtime/nri/plugin.go
@@ -43,14 +43,16 @@ type Plugin struct {
 	ctx    context.Context
 	logger logger.Interface
 
-	stub stub.Stub
+	namespace string
+	stub      stub.Stub
 }
 
 // NewPlugin creates a new NRI plugin for injecting CDI devices
-func NewPlugin(ctx context.Context, logger logger.Interface) *Plugin {
+func NewPlugin(ctx context.Context, logger logger.Interface, namespace string) *Plugin {
 	return &Plugin{
-		ctx:    ctx,
-		logger: logger,
+		ctx:       ctx,
+		logger:    logger,
+		namespace: namespace,
 	}
 }
 
@@ -69,25 +71,31 @@ func (p *Plugin) injectCDIDevices(pod *api.PodSandbox, ctr *api.Container, a *ap
 	ctx := p.ctx
 	pluginLogger := p.stub.Logger()
 
-	devices := parseCDIDevices(pod, nriCDIAnnotationDomain, ctr.Name)
+	devices := p.parseCDIDevices(pod, nriCDIAnnotationDomain, ctr.Name)
 	if len(devices) == 0 {
 		pluginLogger.Debugf(ctx, "%s: no CDI devices annotated...", containerName(pod, ctr))
 		return nil
 	}
 
+	pluginLogger.Infof(ctx, "%s: injecting CDI devices %v...", containerName(pod, ctr), devices)
 	for _, name := range devices {
 		a.AddCDIDevice(
 			&api.CDIDevice{
 				Name: name,
 			},
 		)
-		pluginLogger.Infof(ctx, "%s: injected CDI device %q...", containerName(pod, ctr), name)
 	}
 
 	return nil
 }
 
-func parseCDIDevices(pod *api.PodSandbox, key, container string) []string {
+// parseCDIDevices processes the podSpec and determines which containers which need CDI devices injected to them
+func (p *Plugin) parseCDIDevices(pod *api.PodSandbox, key, container string) []string {
+	if p.namespace != pod.Namespace {
+		p.logger.Debugf("pod %s/%s is not in the toolkit's namespace %s. Skipping CDI device injection...", pod.Namespace, pod.Name, p.namespace)
+		return nil
+	}
+
 	cdiDeviceNames, ok := plugin.GetEffectiveAnnotation(pod, key, container)
 	if !ok {
 		return nil
diff --git a/cmd/nvidia-ctk-installer/main.go b/cmd/nvidia-ctk-installer/main.go
index ece435737..7643cc8b9 100644
--- a/cmd/nvidia-ctk-installer/main.go
+++ b/cmd/nvidia-ctk-installer/main.go
@@ -50,6 +50,7 @@ type options struct {
 	enableNRIPlugin bool
 	nriPluginIndex  uint
 	nriSocket       string
+	nriNamespace    string
 
 	toolkitOptions toolkit.Options
 
@@ -150,6 +151,12 @@ func (a app) build() *cli.Command {
 				Destination: &options.nriSocket,
 				Sources:     cli.EnvVars("NRI_SOCKET"),
 			},
+			&cli.StringFlag{
+				Name:        "nri-namespace",
+				Usage:       "Specify the kubernetes namespace the toolkit's NRI plugin is running in.",
+				Destination: &options.nriNamespace,
+				Sources:     cli.EnvVars("NRI_NAMESPACE"),
+			},
 			&cli.StringFlag{
 				Name:    "runtime",
 				Aliases: []string{"r"},
@@ -226,6 +233,10 @@ func (a *app) validateFlags(c *cli.Command, o *options) error {
 		return fmt.Errorf("invalid toolkit.pid path %v", o.pidFile)
 	}
 
+	if o.enableNRIPlugin && len(o.nriNamespace) == 0 {
+		return fmt.Errorf("the NRI namespace must be specified when the NRI plugin is enabled")
+	}
+
 	if err := a.toolkit.ValidateOptions(&o.toolkitOptions); err != nil {
 		return err
 	}
@@ -353,7 +364,7 @@ func (a *app) startNRIPluginServer(ctx context.Context, opts *options) (*nri.Plu
 		retryBackoff     = 2 * time.Second
 	)
 
-	plugin := nri.NewPlugin(ctx, a.logger)
+	plugin := nri.NewPlugin(ctx, a.logger, opts.nriNamespace)
 	retriable := func() error {
 		return plugin.Start(ctx, opts.nriSocket, fmt.Sprintf("%02d", opts.nriPluginIndex))
 	}