Merge branch 'add-config-set-command' into 'main'

Allow config options to be set usign the nvidia-ctk config command

See merge request nvidia/container-toolkit/container-toolkit!464

CHANGELOG.md

@@ -4,9 +4,9 @@
 * Added support for generating OCI hook JSON file to `nvidia-ctk runtime configure` command.
 * Remove installation of OCI hook JSON from RPM package.
 * Refactored config for `nvidia-container-runtime-hook`.
+* Added a `nvidia-ctk config` command which supports setting config options using a `--set` flag.
 
 ## v1.14.0-rc.2
-
 * Fix bug causing incorrect nvidia-smi symlink to be created on WSL2 systems with multiple driver roots.
 * Remove dependency on coreutils when installing package on RPM-based systems.
 * Create ouput folders if required when running `nvidia-ctk runtime configure`

cmd/nvidia-container-runtime-hook/hook_config.go

@@ -43,13 +43,15 @@ func loadConfig() (*config.Config, error) {
 	}
 
 	for _, p := range configPaths {
-		cfg, err := config.Load(p)
+		cfg, err := config.New(
+			config.WithConfigFile(p),
+		)
 		if err == nil {
-			return cfg, nil
+			return cfg.Config()
 		} else if os.IsNotExist(err) && !required {
 			continue
 		}
-		return nil, fmt.Errorf("couldn't open configuration file: %v", err)
+		return nil, fmt.Errorf("couldn't open required configuration file: %v", err)
 	}
 
 	return config.GetDefault()

cmd/nvidia-ctk/README.md

@@ -16,6 +16,31 @@ nvidia-ctk runtime configure --set-as-default
 will ensure that the NVIDIA Container Runtime is added as the default runtime to the default container
 engine.
 
+## Configure the NVIDIA Container Toolkit
+
+The `config` command of the `nvidia-ctk` CLI allows a user to display and manipulate the NVIDIA Container Toolkit
+configuration.
+
+For example, running the following command:
+```bash
+nvidia-ctk config default
+```
+will display the default config for the detected platform.
+
+Whereas
+```bash
+nvidia-ctk config
+```
+will display the effective NVIDIA Container Toolkit config using the configured config file, and running:
+
+Individual config options can be set by specifying these are key-value pairs to the `--set` argument:
+
+```bash
+nvidia-ctk config --set nvidia-container-cli.no-cgroups=true
+```
+
+By default, all commands output to `STDOUT`, but specifying the `--output` flag writes the config to the specified file.
+
 ### Generate CDI specifications
 
 The [Container Device Interface (CDI)](https://github.com/container-orchestrated-devices/container-device-interface) provides

cmd/nvidia-ctk/config/config.go

@@ -17,7 +17,15 @@
 package config
 
 import (
+	"errors"
+	"fmt"
+	"os"
+	"strconv"
+	"strings"
+
 	createdefault "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk/config/create-default"
+	"github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk/config/flags"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 	"github.com/urfave/cli/v2"
 )
@@ -26,6 +34,12 @@ type command struct {
 	logger logger.Interface
 }
 
+// options stores the subcommand options
+type options struct {
+	flags.Options
+	sets cli.StringSlice
+}
+
 // NewCommand constructs an config command with the specified logger
 func NewCommand(logger logger.Interface) *cli.Command {
 	c := command{
@@ -36,10 +50,42 @@ func NewCommand(logger logger.Interface) *cli.Command {
 
 // build
 func (m command) build() *cli.Command {
+	opts := options{}
+
 	// Create the 'config' command
 	c := cli.Command{
 		Name:  "config",
 		Usage: "Interact with the NVIDIA Container Toolkit configuration",
+		Action: func(ctx *cli.Context) error {
+			return run(ctx, &opts)
+		},
+	}
+
+	c.Flags = []cli.Flag{
+		&cli.StringFlag{
+			Name:        "config-file",
+			Aliases:     []string{"config", "c"},
+			Usage:       "Specify the config file to modify.",
+			Value:       config.GetConfigFilePath(),
+			Destination: &opts.Config,
+		},
+		&cli.StringSliceFlag{
+			Name:        "set",
+			Usage:       "Set a config value using the pattern key=value. If value is empty, this is equivalent to specifying the same key in unset. This flag can be specified multiple times",
+			Destination: &opts.sets,
+		},
+		&cli.BoolFlag{
+			Name:        "in-place",
+			Aliases:     []string{"i"},
+			Usage:       "Modify the config file in-place",
+			Destination: &opts.InPlace,
+		},
+		&cli.StringFlag{
+			Name:        "output",
+			Aliases:     []string{"o"},
+			Usage:       "Specify the output file to write to; If not specified, the output is written to stdout",
+			Destination: &opts.Output,
+		},
 	}
 
 	c.Subcommands = []*cli.Command{
@@ -48,3 +94,71 @@ func (m command) build() *cli.Command {
 
 	return &c
 }
+
+func run(c *cli.Context, opts *options) error {
+	cfgToml, err := config.New(
+		config.WithConfigFile(opts.Config),
+	)
+	if err != nil {
+		return fmt.Errorf("unable to create config: %v", err)
+	}
+
+	for _, set := range opts.sets.Value() {
+		key, value, err := (*configToml)(cfgToml).setFlagToKeyValue(set)
+		if err != nil {
+			return fmt.Errorf("invalid --set option %v: %w", set, err)
+		}
+		cfgToml.Set(key, value)
+	}
+
+	cfgToml.Save(os.Stdout)
+	return nil
+}
+
+type configToml config.Toml
+
+var errInvalidConfigOption = errors.New("invalid config option")
+var errInvalidFormat = errors.New("invalid format")
+
+// setFlagToKeyValue converts a --set flag to a key-value pair.
+// The set flag is of the form key[=value], with the value being optional if key refers to a
+// boolean config option.
+func (c *configToml) setFlagToKeyValue(setFlag string) (string, interface{}, error) {
+	if c == nil {
+		return "", nil, errInvalidConfigOption
+	}
+
+	setParts := strings.SplitN(setFlag, "=", 2)
+	key := setParts[0]
+
+	v := (*config.Toml)(c).Get(key)
+	if v == nil {
+		return key, nil, errInvalidConfigOption
+	}
+	switch v.(type) {
+	case bool:
+		if len(setParts) == 1 {
+			return key, true, nil
+		}
+	}
+
+	if len(setParts) != 2 {
+		return key, nil, fmt.Errorf("%w: expected key=value; got %v", errInvalidFormat, setFlag)
+	}
+
+	value := setParts[1]
+	switch vt := v.(type) {
+	case bool:
+		b, err := strconv.ParseBool(value)
+		if err != nil {
+			return key, value, fmt.Errorf("%w: %w", errInvalidFormat, err)
+		}
+		return key, b, err
+	case string:
+		return key, value, nil
+	case []string:
+		return key, strings.Split(value, ","), nil
+	default:
+		return key, nil, fmt.Errorf("unsupported type for %v (%v)", setParts, vt)
+	}
+}

cmd/nvidia-ctk/config/config_test.go

@@ -0,0 +1,173 @@
+/**
+# 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 config
+
+import (
+	"testing"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
+	"github.com/pelletier/go-toml"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSetFlagToKeyValue(t *testing.T) {
+	testCases := []struct {
+		description   string
+		config        map[string]interface{}
+		setFlag       string
+		expectedKey   string
+		expectedValue interface{}
+		expectedError error
+	}{
+		{
+			description:   "empty config returns an error",
+			setFlag:       "anykey=value",
+			expectedKey:   "anykey",
+			expectedError: errInvalidConfigOption,
+		},
+		{
+			description: "option not present returns an error",
+			config: map[string]interface{}{
+				"defined": "defined-value",
+			},
+			setFlag:       "undefined=new-value",
+			expectedKey:   "undefined",
+			expectedError: errInvalidConfigOption,
+		},
+		{
+			description: "boolean option assumes true",
+			config: map[string]interface{}{
+				"boolean": false,
+			},
+			setFlag:       "boolean",
+			expectedKey:   "boolean",
+			expectedValue: true,
+		},
+		{
+			description: "boolean option returns true",
+			config: map[string]interface{}{
+				"boolean": false,
+			},
+			setFlag:       "boolean=true",
+			expectedKey:   "boolean",
+			expectedValue: true,
+		},
+		{
+			description: "boolean option returns false",
+			config: map[string]interface{}{
+				"boolean": false,
+			},
+			setFlag:       "boolean=false",
+			expectedKey:   "boolean",
+			expectedValue: false,
+		},
+		{
+			description: "invalid boolean option returns error",
+			config: map[string]interface{}{
+				"boolean": false,
+			},
+			setFlag:       "boolean=something",
+			expectedKey:   "boolean",
+			expectedValue: "something",
+			expectedError: errInvalidFormat,
+		},
+		{
+			description: "string option requires value",
+			config: map[string]interface{}{
+				"string": "value",
+			},
+			setFlag:       "string",
+			expectedKey:   "string",
+			expectedValue: nil,
+			expectedError: errInvalidFormat,
+		},
+		{
+			description: "string option returns value",
+			config: map[string]interface{}{
+				"string": "value",
+			},
+			setFlag:       "string=string-value",
+			expectedKey:   "string",
+			expectedValue: "string-value",
+		},
+		{
+			description: "string option returns value with equals",
+			config: map[string]interface{}{
+				"string": "value",
+			},
+			setFlag:       "string=string-value=more",
+			expectedKey:   "string",
+			expectedValue: "string-value=more",
+		},
+		{
+			description: "string option treats bool value as string",
+			config: map[string]interface{}{
+				"string": "value",
+			},
+			setFlag:       "string=true",
+			expectedKey:   "string",
+			expectedValue: "true",
+		},
+		{
+			description: "string option treats int value as string",
+			config: map[string]interface{}{
+				"string": "value",
+			},
+			setFlag:       "string=5",
+			expectedKey:   "string",
+			expectedValue: "5",
+		},
+		{
+			description: "[]string option returns single value",
+			config: map[string]interface{}{
+				"string": []string{"value"},
+			},
+			setFlag:       "string=string-value",
+			expectedKey:   "string",
+			expectedValue: []string{"string-value"},
+		},
+		{
+			description: "[]string option returns multiple values",
+			config: map[string]interface{}{
+				"string": []string{"value"},
+			},
+			setFlag:       "string=first,second",
+			expectedKey:   "string",
+			expectedValue: []string{"first", "second"},
+		},
+		{
+			description: "[]string option returns values with equals",
+			config: map[string]interface{}{
+				"string": []string{"value"},
+			},
+			setFlag:       "string=first=1,second=2",
+			expectedKey:   "string",
+			expectedValue: []string{"first=1", "second=2"},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			tree, _ := toml.TreeFromMap(tc.config)
+			cfgToml := (*config.Toml)(tree)
+			k, v, err := (*configToml)(cfgToml).setFlagToKeyValue(tc.setFlag)
+			require.ErrorIs(t, err, tc.expectedError)
+			require.EqualValues(t, tc.expectedKey, k)
+			require.EqualValues(t, tc.expectedValue, v)
+		})
+	}
+}

cmd/nvidia-ctk/config/create-default/create-default.go

@@ -17,13 +17,9 @@
 package defaultsubcommand
 
 import (
-	"bytes"
 	"fmt"
-	"io"
-	"os"
-	"path/filepath"
-	"regexp"
 
+	"github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk/config/flags"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 	"github.com/urfave/cli/v2"
@@ -33,13 +29,6 @@ type command struct {
 	logger logger.Interface
 }
 
-// options stores the subcommand options
-type options struct {
-	config  string
-	output  string
-	inPlace bool
-}
-
 // NewCommand constructs a default command with the specified logger
 func NewCommand(logger logger.Interface) *cli.Command {
 	c := command{
@@ -50,12 +39,12 @@ func NewCommand(logger logger.Interface) *cli.Command {
 
 // build creates the CLI command
 func (m command) build() *cli.Command {
-	opts := options{}
+	opts := flags.Options{}
 
 	// Create the 'default' command
 	c := cli.Command{
-		Name:    "generate-default",
+		Name:    "default",
-		Aliases: []string{"default"},
+		Aliases: []string{"create-default", "generate-default"},
 		Usage:   "Generate the default NVIDIA Container Toolkit configuration file",
 		Before: func(c *cli.Context) error {
 			return m.validateFlags(c, &opts)
@@ -66,118 +55,40 @@ func (m command) build() *cli.Command {
 	}
 
 	c.Flags = []cli.Flag{
-		&cli.StringFlag{
-			Name:        "config",
-			Usage:       "Specify the config file to process; The contents of this file overrides the default config",
-			Destination: &opts.config,
-		},
-		&cli.BoolFlag{
-			Name:        "in-place",
-			Aliases:     []string{"i"},
-			Usage:       "Modify the config file in-place",
-			Destination: &opts.inPlace,
-		},
 		&cli.StringFlag{
 			Name:        "output",
+			Aliases:     []string{"o"},
 			Usage:       "Specify the output file to write to; If not specified, the output is written to stdout",
-			Destination: &opts.output,
+			Destination: &opts.Output,
 		},
 	}
 
 	return &c
 }
 
-func (m command) validateFlags(c *cli.Context, opts *options) error {
+func (m command) validateFlags(c *cli.Context, opts *flags.Options) error {
-	if opts.inPlace {
+	return opts.Validate()
-		if opts.output != "" {
-			return fmt.Errorf("cannot specify both --in-place and --output")
-		}
-		opts.output = opts.config
-	}
-	return nil
 }
 
-func (m command) run(c *cli.Context, opts *options) error {
+func (m command) run(c *cli.Context, opts *flags.Options) error {
-	if err := opts.ensureOutputFolder(); err != nil {
+	cfgToml, err := config.New()
-		return fmt.Errorf("unable to create output directory: %v", err)
-	}
-
-	contents, err := opts.getFormattedConfig()
 	if err != nil {
-		return fmt.Errorf("unable to fix comments: %v", err)
+		return fmt.Errorf("unable to load or create config: %v", err)
 	}
 
-	if _, err := opts.Write(contents); err != nil {
+	if err := opts.EnsureOutputFolder(); err != nil {
-		return fmt.Errorf("unable to write to output: %v", err)
+		return fmt.Errorf("failed to create output directory: %v", err)
+	}
+	output, err := opts.CreateOutput()
+	if err != nil {
+		return fmt.Errorf("failed to open output file: %v", err)
+	}
+	defer output.Close()
+
+	_, err = cfgToml.Save(output)
+	if err != nil {
+		return fmt.Errorf("failed to write output: %v", err)
 	}
 
 	return nil
 }
-
-// getFormattedConfig returns the default config formatted as required from the specified config file.
-// The config is then formatted as required.
-// No indentation is used and comments are modified so that there is no space
-// after the '#' character.
-func (opts options) getFormattedConfig() ([]byte, error) {
-	cfg, err := config.Load(opts.config)
-	if err != nil {
-		return nil, fmt.Errorf("unable to load or create config: %v", err)
-	}
-
-	buffer := bytes.NewBuffer(nil)
-
-	if _, err := cfg.Save(buffer); err != nil {
-		return nil, fmt.Errorf("unable to save config: %v", err)
-	}
-	return fixComments(buffer.Bytes())
-}
-
-func fixComments(contents []byte) ([]byte, error) {
-	r, err := regexp.Compile(`(\n*)\s*?#\s*(\S.*)`)
-	if err != nil {
-		return nil, fmt.Errorf("unable to compile regexp: %v", err)
-	}
-	replaced := r.ReplaceAll(contents, []byte("$1#$2"))
-
-	return replaced, nil
-}
-
-func (opts options) outputExists() (bool, error) {
-	if opts.output == "" {
-		return false, nil
-	}
-	_, err := os.Stat(opts.output)
-	if err == nil {
-		return true, nil
-	} else if !os.IsNotExist(err) {
-		return false, fmt.Errorf("unable to stat output file: %v", err)
-	}
-	return false, nil
-}
-
-func (opts options) ensureOutputFolder() error {
-	if opts.output == "" {
-		return nil
-	}
-	if dir := filepath.Dir(opts.output); dir != "" {
-		return os.MkdirAll(dir, 0755)
-	}
-	return nil
-}
-
-// Write writes the contents to the output file specified in the options.
-func (opts options) Write(contents []byte) (int, error) {
-	var output io.Writer
-	if opts.output == "" {
-		output = os.Stdout
-	} else {
-		outputFile, err := os.Create(opts.output)
-		if err != nil {
-			return 0, fmt.Errorf("unable to create output file: %v", err)
-		}
-		defer outputFile.Close()
-		output = outputFile
-	}
-
-	return output.Write(contents)
-}

cmd/nvidia-ctk/config/create-default/create-default_test.go

@@ -1,82 +0,0 @@
-/**
-# 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 defaultsubcommand
-
-import (
-	"strings"
-	"testing"
-
-	"github.com/stretchr/testify/require"
-)
-
-func TestFixComment(t *testing.T) {
-	testCases := []struct {
-		input    string
-		expected string
-	}{
-		{
-			input:    "# comment",
-			expected: "#comment",
-		},
-		{
-			input:    " #comment",
-			expected: "#comment",
-		},
-		{
-			input:    " # comment",
-			expected: "#comment",
-		},
-		{
-			input: strings.Join([]string{
-				"some",
-				"# comment",
-				" # comment",
-				" #comment",
-				"other"}, "\n"),
-			expected: strings.Join([]string{
-				"some",
-				"#comment",
-				"#comment",
-				"#comment",
-				"other"}, "\n"),
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.input, func(t *testing.T) {
-			actual, _ := fixComments([]byte(tc.input))
-			require.Equal(t, tc.expected, string(actual))
-		})
-	}
-}
-
-func TestGetFormattedConfig(t *testing.T) {
-	expectedLines := []string{
-		"#no-cgroups = false",
-		"#debug = \"/var/log/nvidia-container-toolkit.log\"",
-		"#debug = \"/var/log/nvidia-container-runtime.log\"",
-	}
-
-	opts := &options{}
-	contents, err := opts.getFormattedConfig()
-	require.NoError(t, err)
-	lines := strings.Split(string(contents), "\n")
-
-	for _, line := range expectedLines {
-		require.Contains(t, lines, line)
-	}
-}

cmd/nvidia-ctk/config/flags/options.go

@@ -0,0 +1,71 @@
+/**
+# 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 flags
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+)
+
+// Options stores options for the config commands
+type Options struct {
+	Config  string
+	Output  string
+	InPlace bool
+}
+
+// Validate checks whether the options are valid.
+func (o Options) Validate() error {
+	if o.InPlace && o.Output != "" {
+		return fmt.Errorf("cannot specify both --in-place and --output")
+	}
+
+	return nil
+}
+
+// EnsureOutputFolder creates the output folder if it does not exist.
+// If the output folder is not specified (i.e. output to STDOUT), it is ignored.
+func (o Options) EnsureOutputFolder() error {
+	if o.Output == "" {
+		return nil
+	}
+	if dir := filepath.Dir(o.Output); dir != "" {
+		return os.MkdirAll(dir, 0755)
+	}
+	return nil
+}
+
+// CreateOutput creates the writer for the output.
+func (o Options) CreateOutput() (io.WriteCloser, error) {
+	if o.Output != "" {
+		return os.Create(o.Output)
+	}
+
+	return nullCloser{os.Stdout}, nil
+}
+
+// nullCloser is a writer that does nothing on Close.
+type nullCloser struct {
+	io.Writer
+}
+
+// Close is a no-op for a nullCloser.
+func (d nullCloser) Close() error {
+	return nil
+}

internal/config/config.go

@@ -18,10 +18,7 @@ package config
 
 import (
 	"bufio"
-	"fmt"
-	"io"
 	"os"
-	"path"
 	"path/filepath"
 	"strings"
 
@@ -29,7 +26,6 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup"
 	"github.com/container-orchestrated-devices/container-device-interface/pkg/cdi"
-	"github.com/pelletier/go-toml"
 )
 
 const (
@@ -51,8 +47,6 @@ var (
 	NVIDIAContainerRuntimeHookExecutable = "nvidia-container-runtime-hook"
 	// NVIDIAContainerToolkitExecutable is the executable name for the NVIDIA Container Toolkit (an alias for the NVIDIA Container Runtime Hook)
 	NVIDIAContainerToolkitExecutable = "nvidia-container-toolkit"
-
-	configDir = "/etc/"
 )
 
 // Config represents the contents of the config.toml file for the NVIDIA Container Toolkit
@@ -70,67 +64,26 @@ type Config struct {
 	NVIDIAContainerRuntimeHookConfig RuntimeHookConfig  `toml:"nvidia-container-runtime-hook"`
 }
 
+// GetConfigFilePath returns the path to the config file for the configured system
+func GetConfigFilePath() string {
+	if XDGConfigDir := os.Getenv(configOverride); len(XDGConfigDir) != 0 {
+		return filepath.Join(XDGConfigDir, configFilePath)
+	}
+
+	return filepath.Join("/etc", configFilePath)
+}
+
 // GetConfig sets up the config struct. Values are read from a toml file
 // or set via the environment.
 func GetConfig() (*Config, error) {
-	if XDGConfigDir := os.Getenv(configOverride); len(XDGConfigDir) != 0 {
+	cfg, err := New(
-		configDir = XDGConfigDir
+		WithConfigFile(GetConfigFilePath()),
-	}
+	)
-
-	configFilePath := path.Join(configDir, configFilePath)
-
-	return Load(configFilePath)
-}
-
-// Load loads the config from the specified file path.
-func Load(configFilePath string) (*Config, error) {
-	if configFilePath == "" {
-		return GetDefault()
-	}
-
-	tomlFile, err := os.Open(configFilePath)
-	if err != nil {
-		return GetDefault()
-	}
-	defer tomlFile.Close()
-
-	cfg, err := LoadFrom(tomlFile)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read config values: %v", err)
-	}
-
-	return cfg, nil
-}
-
-// LoadFrom reads the config from the specified Reader
-func LoadFrom(reader io.Reader) (*Config, error) {
-	var tree *toml.Tree
-	if reader != nil {
-		toml, err := toml.LoadReader(reader)
-		if err != nil {
-			return nil, err
-		}
-		tree = toml
-	}
-
-	return getFromTree(tree)
-}
-
-// getFromTree reads the nvidia container runtime config from the specified toml Tree.
-func getFromTree(toml *toml.Tree) (*Config, error) {
-	cfg, err := GetDefault()
 	if err != nil {
 		return nil, err
 	}
-	if toml == nil {
-		return cfg, nil
-	}
 
-	if err := toml.Unmarshal(cfg); err != nil {
+	return cfg.Config()
-		return nil, fmt.Errorf("failed to unmarshal config: %v", err)
-	}
-
-	return cfg, nil
 }
 
 // GetDefault defines the default values for the config
@@ -260,64 +213,3 @@ func resolveWithDefault(logger logger.Interface, label string, path string, defa
 
 	return resolvedPath
 }
-
-func (c Config) asCommentedToml() (*toml.Tree, error) {
-	contents, err := toml.Marshal(c)
-	if err != nil {
-		return nil, err
-	}
-	asToml, err := toml.LoadBytes(contents)
-	if err != nil {
-		return nil, err
-	}
-
-	commentedDefaults := map[string]interface{}{
-		"swarm-resource": "DOCKER_RESOURCE_GPU",
-		"accept-nvidia-visible-devices-envvar-when-unprivileged": true,
-		"accept-nvidia-visible-devices-as-volume-mounts":         false,
-		"nvidia-container-cli.root":                              "/run/nvidia/driver",
-		"nvidia-container-cli.path":                              "/usr/bin/nvidia-container-cli",
-		"nvidia-container-cli.debug":                             "/var/log/nvidia-container-toolkit.log",
-		"nvidia-container-cli.ldcache":                           "/etc/ld.so.cache",
-		"nvidia-container-cli.no-cgroups":                        false,
-		"nvidia-container-cli.user":                              "root:video",
-		"nvidia-container-runtime.debug":                         "/var/log/nvidia-container-runtime.log",
-	}
-	for k, v := range commentedDefaults {
-		set := asToml.Get(k)
-		if !shouldComment(k, v, set) {
-			continue
-		}
-		asToml.SetWithComment(k, "", true, v)
-	}
-
-	return asToml, nil
-}
-
-func shouldComment(key string, defaultValue interface{}, setTo interface{}) bool {
-	if key == "nvidia-container-cli.user" && !getCommentedUserGroup() {
-		return false
-	}
-	if key == "nvidia-container-runtime.debug" && setTo == "/dev/null" {
-		return true
-	}
-	if setTo == nil || defaultValue == setTo || setTo == "" {
-		return true
-	}
-
-	return false
-}
-
-// Save writes the config to the specified writer.
-func (c Config) Save(w io.Writer) (int64, error) {
-	asToml, err := c.asCommentedToml()
-	if err != nil {
-		return 0, err
-	}
-
-	enc := toml.NewEncoder(w).Indentation("")
-	if err := enc.Encode(asToml); err != nil {
-		return 0, fmt.Errorf("invalid config: %v", err)
-	}
-	return 0, nil
-}

internal/config/config_test.go

@@ -17,8 +17,6 @@
 package config
 
 import (
-	"bytes"
-	"io/ioutil"
 	"os"
 	"path/filepath"
 	"strings"
@@ -28,24 +26,20 @@ import (
 )
 
 func TestGetConfigWithCustomConfig(t *testing.T) {
-	wd, err := os.Getwd()
+	testDir := t.TempDir()
-	require.NoError(t, err)
+	t.Setenv(configOverride, testDir)
+
+	filename := filepath.Join(testDir, configFilePath)
 
 	// By default debug is disabled
 	contents := []byte("[nvidia-container-runtime]\ndebug = \"/nvidia-container-toolkit.log\"")
-	testDir := filepath.Join(wd, "test")
-	filename := filepath.Join(testDir, configFilePath)
-
-	os.Setenv(configOverride, testDir)
 
 	require.NoError(t, os.MkdirAll(filepath.Dir(filename), 0766))
-	require.NoError(t, ioutil.WriteFile(filename, contents, 0766))
+	require.NoError(t, os.WriteFile(filename, contents, 0766))
-
-	defer func() { require.NoError(t, os.RemoveAll(testDir)) }()
 
 	cfg, err := GetConfig()
 	require.NoError(t, err)
-	require.Equal(t, cfg.NVIDIAContainerRuntimeConfig.DebugFilePath, "/nvidia-container-toolkit.log")
+	require.Equal(t, "/nvidia-container-toolkit.log", cfg.NVIDIAContainerRuntimeConfig.DebugFilePath)
 }
 
 func TestGetConfig(t *testing.T) {
@@ -219,12 +213,14 @@ func TestGetConfig(t *testing.T) {
 		t.Run(tc.description, func(t *testing.T) {
 			reader := strings.NewReader(strings.Join(tc.contents, "\n"))
 
-			cfg, err := LoadFrom(reader)
+			tomlCfg, err := loadConfigTomlFrom(reader)
 			if tc.expectedError != nil {
 				require.Error(t, err)
 			} else {
 				require.NoError(t, err)
 			}
+			cfg, err := tomlCfg.Config()
+			require.NoError(t, err)
 
 			// We first handle the ldconfig path since this is currently system-dependent.
 			if tc.inspectLdconfig {
@@ -240,47 +236,3 @@ func TestGetConfig(t *testing.T) {
 		})
 	}
 }
-
-func TestConfigDefault(t *testing.T) {
-	config, err := GetDefault()
-	require.NoError(t, err)
-
-	buffer := new(bytes.Buffer)
-	_, err = config.Save(buffer)
-	require.NoError(t, err)
-
-	var lines []string
-	for _, l := range strings.Split(buffer.String(), "\n") {
-		l = strings.TrimSpace(l)
-		if strings.HasPrefix(l, "# ") {
-			l = "#" + strings.TrimPrefix(l, "# ")
-		}
-		lines = append(lines, l)
-	}
-
-	// We take the lines from the config that was included in previous packages.
-	expectedLines := []string{
-		"disable-require = false",
-		"#swarm-resource = \"DOCKER_RESOURCE_GPU\"",
-		"#accept-nvidia-visible-devices-envvar-when-unprivileged = true",
-		"#accept-nvidia-visible-devices-as-volume-mounts = false",
-
-		"#root = \"/run/nvidia/driver\"",
-		"#path = \"/usr/bin/nvidia-container-cli\"",
-		"environment = []",
-		"#debug = \"/var/log/nvidia-container-toolkit.log\"",
-		"#ldcache = \"/etc/ld.so.cache\"",
-		"load-kmods = true",
-		"#no-cgroups = false",
-		"#user = \"root:video\"",
-
-		"[nvidia-container-runtime]",
-		"#debug = \"/var/log/nvidia-container-runtime.log\"",
-		"log-level = \"info\"",
-		"mode = \"auto\"",
-
-		"mount-spec-path = \"/etc/nvidia-container-runtime/host-files-for-container.d\"",
-	}
-
-	require.Subset(t, lines, expectedLines)
-}

internal/config/toml.go

@@ -0,0 +1,203 @@
+/**
+# 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 config
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"os"
+	"regexp"
+
+	"github.com/pelletier/go-toml"
+)
+
+// Toml is a type for the TOML representation of a config.
+type Toml toml.Tree
+
+type options struct {
+	configFile string
+}
+
+// Option is a functional option for loading TOML config files.
+type Option func(*options)
+
+// WithConfigFile sets the config file option.
+func WithConfigFile(configFile string) Option {
+	return func(o *options) {
+		o.configFile = configFile
+	}
+}
+
+// New creates a new toml tree based on the provided options
+func New(opts ...Option) (*Toml, error) {
+	o := &options{}
+	for _, opt := range opts {
+		opt(o)
+	}
+
+	return loadConfigToml(o.configFile)
+}
+
+func loadConfigToml(filename string) (*Toml, error) {
+	if filename == "" {
+		return defaultToml()
+	}
+
+	tomlFile, err := os.Open(filename)
+	if os.IsNotExist(err) {
+		return defaultToml()
+	} else if err != nil {
+		return nil, fmt.Errorf("failed to load specified config file: %v", err)
+	}
+	defer tomlFile.Close()
+
+	return loadConfigTomlFrom(tomlFile)
+
+}
+
+func defaultToml() (*Toml, error) {
+	cfg, err := GetDefault()
+	if err != nil {
+		return nil, err
+	}
+	contents, err := toml.Marshal(cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	return loadConfigTomlFrom(bytes.NewReader(contents))
+}
+
+func loadConfigTomlFrom(reader io.Reader) (*Toml, error) {
+	tree, err := toml.LoadReader(reader)
+	if err != nil {
+		return nil, err
+	}
+	return (*Toml)(tree), nil
+}
+
+// Config returns the typed config associated with the toml tree.
+func (t *Toml) Config() (*Config, error) {
+	cfg, err := GetDefault()
+	if err != nil {
+		return nil, err
+	}
+	if t == nil {
+		return cfg, nil
+	}
+	if err := t.Unmarshal(cfg); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal config: %v", err)
+	}
+	return cfg, nil
+}
+
+// Unmarshal wraps the toml.Tree Unmarshal function.
+func (t *Toml) Unmarshal(v interface{}) error {
+	return (*toml.Tree)(t).Unmarshal(v)
+}
+
+// Save saves the config to the specified Writer.
+func (t *Toml) Save(w io.Writer) (int64, error) {
+	contents, err := t.contents()
+	if err != nil {
+		return 0, err
+	}
+
+	n, err := w.Write(contents)
+	return int64(n), err
+}
+
+// contents returns the config TOML as a byte slice.
+// Any required formatting is applied.
+func (t Toml) contents() ([]byte, error) {
+	commented := t.commentDefaults()
+
+	buffer := bytes.NewBuffer(nil)
+
+	enc := toml.NewEncoder(buffer).Indentation("")
+	if err := enc.Encode((*toml.Tree)(commented)); err != nil {
+		return nil, fmt.Errorf("invalid config: %v", err)
+	}
+	return t.format(buffer.Bytes())
+}
+
+// format fixes the comments for the config to ensure that they start in column
+// 1 and are not followed by a space.
+func (t Toml) format(contents []byte) ([]byte, error) {
+	r, err := regexp.Compile(`(\n*)\s*?#\s*(\S.*)`)
+	if err != nil {
+		return nil, fmt.Errorf("unable to compile regexp: %v", err)
+	}
+	replaced := r.ReplaceAll(contents, []byte("$1#$2"))
+
+	return replaced, nil
+}
+
+// Delete deletes the specified key from the TOML config.
+func (t *Toml) Delete(key string) error {
+	return (*toml.Tree)(t).Delete(key)
+}
+
+// Get returns the value for the specified key.
+func (t *Toml) Get(key string) interface{} {
+	return (*toml.Tree)(t).Get(key)
+}
+
+// Set sets the specified key to the specified value in the TOML config.
+func (t *Toml) Set(key string, value interface{}) {
+	(*toml.Tree)(t).Set(key, value)
+}
+
+// commentDefaults applies the required comments for default values to the Toml.
+func (t *Toml) commentDefaults() *Toml {
+	asToml := (*toml.Tree)(t)
+	commentedDefaults := map[string]interface{}{
+		"swarm-resource": "DOCKER_RESOURCE_GPU",
+		"accept-nvidia-visible-devices-envvar-when-unprivileged": true,
+		"accept-nvidia-visible-devices-as-volume-mounts":         false,
+		"nvidia-container-cli.root":                              "/run/nvidia/driver",
+		"nvidia-container-cli.path":                              "/usr/bin/nvidia-container-cli",
+		"nvidia-container-cli.debug":                             "/var/log/nvidia-container-toolkit.log",
+		"nvidia-container-cli.ldcache":                           "/etc/ld.so.cache",
+		"nvidia-container-cli.no-cgroups":                        false,
+		"nvidia-container-cli.user":                              "root:video",
+		"nvidia-container-runtime.debug":                         "/var/log/nvidia-container-runtime.log",
+	}
+	for k, v := range commentedDefaults {
+		set := asToml.Get(k)
+		if !shouldComment(k, v, set) {
+			continue
+		}
+		asToml.SetWithComment(k, "", true, v)
+	}
+	return (*Toml)(asToml)
+}
+
+func shouldComment(key string, defaultValue interface{}, setTo interface{}) bool {
+	if key == "nvidia-container-cli.user" && !getCommentedUserGroup() {
+		return false
+	}
+	if key == "nvidia-container-runtime.debug" && setTo == "/dev/null" {
+		return true
+	}
+	if setTo == nil || defaultValue == setTo || setTo == "" {
+		return true
+	}
+
+	return false
+}

internal/config/toml_test.go

@@ -0,0 +1,248 @@
+/**
+# 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 config
+
+import (
+	"bytes"
+	"strings"
+	"testing"
+
+	"github.com/pelletier/go-toml"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTomlSave(t *testing.T) {
+	testCases := []struct {
+		description string
+		config      *Toml
+		expected    string
+	}{
+		{
+			description: "defaultConfig",
+			config: func() *Toml {
+				t, _ := defaultToml()
+				// TODO: We handle the ldconfig path specifically, since this is platform
+				// dependent.
+				(*toml.Tree)(t).Set("nvidia-container-cli.ldconfig", "OVERRIDDEN")
+				return t
+			}(),
+			expected: `
+#accept-nvidia-visible-devices-as-volume-mounts = false
+#accept-nvidia-visible-devices-envvar-when-unprivileged = true
+disable-require = false
+supported-driver-capabilities = "compat32,compute,display,graphics,ngx,utility,video"
+#swarm-resource = "DOCKER_RESOURCE_GPU"
+
+[nvidia-container-cli]
+#debug = "/var/log/nvidia-container-toolkit.log"
+environment = []
+#ldcache = "/etc/ld.so.cache"
+ldconfig = "OVERRIDDEN"
+load-kmods = true
+#no-cgroups = false
+#path = "/usr/bin/nvidia-container-cli"
+#root = "/run/nvidia/driver"
+#user = "root:video"
+
+[nvidia-container-runtime]
+#debug = "/var/log/nvidia-container-runtime.log"
+log-level = "info"
+mode = "auto"
+runtimes = ["docker-runc", "runc"]
+
+[nvidia-container-runtime.modes]
+
+[nvidia-container-runtime.modes.cdi]
+annotation-prefixes = ["cdi.k8s.io/"]
+default-kind = "nvidia.com/gpu"
+spec-dirs = ["/etc/cdi", "/var/run/cdi"]
+
+[nvidia-container-runtime.modes.csv]
+mount-spec-path = "/etc/nvidia-container-runtime/host-files-for-container.d"
+
+[nvidia-container-runtime-hook]
+path = "nvidia-container-runtime-hook"
+skip-mode-detection = false
+
+[nvidia-ctk]
+path = "nvidia-ctk"
+`,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			buffer := new(bytes.Buffer)
+			_, err := tc.config.Save(buffer)
+			require.NoError(t, err)
+
+			require.EqualValues(t,
+				strings.TrimSpace(tc.expected),
+				strings.TrimSpace(buffer.String()),
+			)
+		})
+	}
+}
+
+func TestFormat(t *testing.T) {
+	testCases := []struct {
+		input    string
+		expected string
+	}{
+		{
+			input:    "# comment",
+			expected: "#comment",
+		},
+		{
+			input:    " #comment",
+			expected: "#comment",
+		},
+		{
+			input:    " # comment",
+			expected: "#comment",
+		},
+		{
+			input: strings.Join([]string{
+				"some",
+				"# comment",
+				" # comment",
+				" #comment",
+				"other"}, "\n"),
+			expected: strings.Join([]string{
+				"some",
+				"#comment",
+				"#comment",
+				"#comment",
+				"other"}, "\n"),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.input, func(t *testing.T) {
+			actual, _ := (Toml{}).format([]byte(tc.input))
+			require.Equal(t, tc.expected, string(actual))
+		})
+	}
+}
+
+func TestGetFormattedConfig(t *testing.T) {
+	expectedLines := []string{
+		"#no-cgroups = false",
+		"#debug = \"/var/log/nvidia-container-toolkit.log\"",
+		"#debug = \"/var/log/nvidia-container-runtime.log\"",
+	}
+
+	contents, err := createEmpty().contents()
+	require.NoError(t, err)
+	lines := strings.Split(string(contents), "\n")
+
+	for _, line := range expectedLines {
+		require.Contains(t, lines, line)
+	}
+}
+
+func TestTomlContents(t *testing.T) {
+	testCases := []struct {
+		description string
+		contents    map[string]interface{}
+		expected    string
+	}{
+		{
+			description: "empty config returns commented defaults",
+			expected: `
+#accept-nvidia-visible-devices-as-volume-mounts = false
+#accept-nvidia-visible-devices-envvar-when-unprivileged = true
+#swarm-resource = "DOCKER_RESOURCE_GPU"
+
+[nvidia-container-cli]
+#debug = "/var/log/nvidia-container-toolkit.log"
+#ldcache = "/etc/ld.so.cache"
+#no-cgroups = false
+#path = "/usr/bin/nvidia-container-cli"
+#root = "/run/nvidia/driver"
+#user = "root:video"
+
+[nvidia-container-runtime]
+#debug = "/var/log/nvidia-container-runtime.log"`,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			tree, err := toml.TreeFromMap(tc.contents)
+			require.NoError(t, err)
+			cfg := (*Toml)(tree)
+			contents, err := cfg.contents()
+			require.NoError(t, err)
+
+			require.EqualValues(t,
+				strings.TrimSpace(tc.expected),
+				strings.TrimSpace(string(contents)),
+			)
+		})
+	}
+}
+
+func TestConfigFromToml(t *testing.T) {
+	testCases := []struct {
+		description    string
+		contents       map[string]interface{}
+		expectedConfig *Config
+	}{
+		{
+			description: "empty config returns default config",
+			contents:    nil,
+			expectedConfig: func() *Config {
+				c, _ := GetDefault()
+				return c
+			}(),
+		},
+		{
+			description: "contents overrides default",
+			contents: map[string]interface{}{
+				"nvidia-container-runtime": map[string]interface{}{
+					"debug": "/some/log/file.log",
+					"mode":  "csv",
+				},
+			},
+			expectedConfig: func() *Config {
+				c, _ := GetDefault()
+				c.NVIDIAContainerRuntimeConfig.DebugFilePath = "/some/log/file.log"
+				c.NVIDIAContainerRuntimeConfig.Mode = "csv"
+				return c
+			}(),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			tomlCfg := fromMap(tc.contents)
+			config, err := tomlCfg.Config()
+			require.NoError(t, err)
+			require.EqualValues(t, tc.expectedConfig, config)
+		})
+	}
+}
+
+func fromMap(c map[string]interface{}) *Toml {
+	tree, _ := toml.TreeFromMap(c)
+	return (*Toml)(tree)
+}
+
+func createEmpty() *Toml {
+	return fromMap(nil)
+}

packaging/debian/nvidia-container-toolkit-base.postinst

@@ -4,7 +4,7 @@ set -e
 
 case "$1" in
     configure)
-        /usr/bin/nvidia-ctk --quiet config default --in-place --config=/etc/nvidia-container-runtime/config.toml
+        /usr/bin/nvidia-ctk --quiet config --config-file=/etc/nvidia-container-runtime/config.toml --in-place
     ;;
 
     abort-upgrade|abort-remove|abort-deconfigure)

packaging/rpm/SPECS/nvidia-container-toolkit.spec

@@ -59,7 +59,7 @@ rm -rf %{_localstatedir}/lib/rpm-state/nvidia-container-toolkit
 ln -sf %{_bindir}/nvidia-container-runtime-hook %{_bindir}/nvidia-container-toolkit
 
 # Generate the default config; If this file already exists no changes are made.
-%{_bindir}/nvidia-ctk --quiet config default --in-place --config=%{_sysconfdir}/nvidia-container-runtime/config.toml
+%{_bindir}/nvidia-ctk --quiet config --config-file=%{_sysconfdir}/nvidia-container-runtime/config.toml --in-place
 
 %postun
 if [ "$1" = 0 ]; then  # package is uninstalled, not upgraded