From 2ae7cb07cf71fa17d8d5aa3f8c3d3f70574eb46d Mon Sep 17 00:00:00 2001
From: Kevin Klues <kklues@nvidia.com>
Date: Thu, 23 Jul 2020 13:50:42 +0000
Subject: [PATCH] Add ability to consider container mounts to generate
 nvidiaConfig

Signed-off-by: Kevin Klues <kklues@nvidia.com>
---
 pkg/container_config.go | 14 ++++++++++++--
 pkg/container_test.go   |  4 ++--
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/pkg/container_config.go b/pkg/container_config.go
index 4d129acd..3d51698a 100644
--- a/pkg/container_config.go
+++ b/pkg/container_config.go
@@ -73,6 +73,15 @@ type LinuxCapabilities struct {
 	Ambient     []string `json:"ambient,omitempty" platform:"linux"`
 }
 
+// Mount from OCI runtime spec
+// https://github.com/opencontainers/runtime-spec/blob/v1.0.0/specs-go/config.go#L103
+type Mount struct {
+	Destination string   `json:"destination"`
+	Type        string   `json:"type,omitempty" platform:"linux,solaris"`
+	Source      string   `json:"source,omitempty"`
+	Options     []string `json:"options,omitempty"`
+}
+
 // Spec from OCI runtime spec
 // We use pointers to structs, similarly to the latest version of runtime-spec:
 // https://github.com/opencontainers/runtime-spec/blob/v1.0.0/specs-go/config.go#L5-L28
@@ -80,6 +89,7 @@ type Spec struct {
 	Version *string  `json:"ociVersion"`
 	Process *Process `json:"process,omitempty"`
 	Root    *Root    `json:"root,omitempty"`
+	Mounts  []Mount  `json:"mounts,omitempty"`
 }
 
 // HookState holds state information about the hook
@@ -285,7 +295,7 @@ func getRequirements(env map[string]string, legacyImage bool) []string {
 	return requirements
 }
 
-func getNvidiaConfig(env map[string]string, privileged bool) *nvidiaConfig {
+func getNvidiaConfig(env map[string]string, mounts []Mount, privileged bool) *nvidiaConfig {
 	legacyImage := isLegacyCUDAImage(env)
 
 	var devices string
@@ -353,6 +363,6 @@ func getContainerConfig(hook HookConfig) (config containerConfig) {
 		Pid:    h.Pid,
 		Rootfs: s.Root.Path,
 		Env:    env,
-		Nvidia: getNvidiaConfig(env, privileged),
+		Nvidia: getNvidiaConfig(env, s.Mounts, privileged),
 	}
 }
diff --git a/pkg/container_test.go b/pkg/container_test.go
index 365e9c03..9c197a30 100644
--- a/pkg/container_test.go
+++ b/pkg/container_test.go
@@ -5,7 +5,7 @@ import (
 	"testing"
 )
 
-func TestGetNvidiaConfig(t *testing.T) {
+func TestGetNvidiaConfigEnvvar(t *testing.T) {
 	var tests = []struct {
 		description    string
 		env            map[string]string
@@ -407,7 +407,7 @@ func TestGetNvidiaConfig(t *testing.T) {
 			// Wrap the call to getNvidiaConfig() in a closure.
 			var config *nvidiaConfig
 			getConfig := func() {
-				config = getNvidiaConfig(tc.env, tc.privileged)
+				config = getNvidiaConfig(tc.env, nil, tc.privileged)
 			}
 
 			// For any tests that are expected to panic, make sure they do.