From 9f5e141437f2df75204b07c75e6c0662c1577657 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Wed, 22 Feb 2023 16:40:27 +0200
Subject: [PATCH] Expose vendor and class as options

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 pkg/nvcdi/lib.go     | 15 ++++++++++++++-
 pkg/nvcdi/options.go | 14 ++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/pkg/nvcdi/lib.go b/pkg/nvcdi/lib.go
index 451c1831..c13a2ab5 100644
--- a/pkg/nvcdi/lib.go
+++ b/pkg/nvcdi/lib.go
@@ -26,6 +26,9 @@ import (
 
 type wrapper struct {
 	Interface
+
+	vendor string
+	class  string
 }
 
 type nvcdilib struct {
@@ -37,6 +40,9 @@ type nvcdilib struct {
 	driverRoot    string
 	nvidiaCTKPath string
 
+	vendor string
+	class  string
+
 	infolib info.Interface
 }
 
@@ -83,7 +89,12 @@ func New(opts ...Option) Interface {
 		panic("Unknown mode")
 	}
 
-	return &wrapper{Interface: lib}
+	w := wrapper{
+		Interface: lib,
+		vendor:    l.vendor,
+		class:     l.class,
+	}
+	return &w
 }
 
 // GetSpec combines the device specs and common edits from the wrapped Interface to a single spec.Interface.
@@ -101,6 +112,8 @@ func (l *wrapper) GetSpec() (spec.Interface, error) {
 	return spec.New(
 		spec.WithDeviceSpecs(deviceSpecs),
 		spec.WithEdits(*edits.ContainerEdits),
+		spec.WithVendor(l.vendor),
+		spec.WithClass(l.class),
 	)
 
 }
diff --git a/pkg/nvcdi/options.go b/pkg/nvcdi/options.go
index 317cace2..19aef93c 100644
--- a/pkg/nvcdi/options.go
+++ b/pkg/nvcdi/options.go
@@ -73,3 +73,17 @@ func WithMode(mode string) Option {
 		l.mode = mode
 	}
 }
+
+// WithVendor sets the vendor for the library
+func WithVendor(vendor string) Option {
+	return func(o *nvcdilib) {
+		o.vendor = vendor
+	}
+}
+
+// WithClass sets the class for the library
+func WithClass(class string) Option {
+	return func(o *nvcdilib) {
+		o.class = class
+	}
+}