Refactor Toml config handling

Signed-off-by: Evan Lezar <[email protected]>

pkg/config/engine/containerd/config_v1_test.go

@@ -19,9 +19,10 @@ package containerd
 import (
 	"testing"
 
-	"github.com/pelletier/go-toml"
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config"
 )
 
 func TestAddRuntimeV1(t *testing.T) {
@@ -226,20 +227,20 @@ func TestAddRuntimeV1(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
-			config, err := toml.Load(tc.config)
+			cfg, err := config.Load(tc.config)
 			require.NoError(t, err)
-			expectedConfig, err := toml.Load(tc.expectedConfig)
+			expectedConfig, err := config.Load(tc.expectedConfig)
 			require.NoError(t, err)
 
 			c := &ConfigV1{
 				Logger: logger,
-				Tree:   config,
+				Toml:   cfg,
 			}
 
 			err = c.AddRuntime("test", "/usr/bin/test", tc.setAsDefault, tc.configOverrides...)
 			require.NoError(t, err)
 
-			require.EqualValues(t, expectedConfig.String(), config.String())
+			require.EqualValues(t, expectedConfig.String(), cfg.String())
 		})
 	}
 }

pkg/config/engine/containerd/config_v2.go

@@ -20,8 +20,6 @@ import (
 	"fmt"
 
 	"github.com/pelletier/go-toml"
-
-	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
 )
 
 // AddRuntime adds a runtime to the containerd config
@@ -151,15 +149,3 @@ func (c *Config) RemoveRuntime(name string) error {
 	*c.Tree = config
 	return nil
 }
-
-// Save writes the config to the specified path
-func (c Config) Save(path string) (int64, error) {
-	config := c.Tree
-	output, err := config.Marshal()
-	if err != nil {
-		return 0, fmt.Errorf("unable to convert to TOML: %v", err)
-	}
-
-	n, err := engine.Config(path).Write(output)
-	return int64(n), err
-}

pkg/config/engine/containerd/config_v2_test.go

@@ -19,9 +19,10 @@ package containerd
 import (
 	"testing"
 
-	"github.com/pelletier/go-toml"
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config"
 )
 
 func TestAddRuntime(t *testing.T) {
@@ -224,20 +225,20 @@ func TestAddRuntime(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
-			config, err := toml.Load(tc.config)
+			cfg, err := config.Load(tc.config)
 			require.NoError(t, err)
-			expectedConfig, err := toml.Load(tc.expectedConfig)
+			expectedConfig, err := config.Load(tc.expectedConfig)
 			require.NoError(t, err)
 
 			c := &Config{
 				Logger: logger,
-				Tree:   config,
+				Toml:   cfg,
 			}
 
 			err = c.AddRuntime("test", "/usr/bin/test", tc.setAsDefault, tc.configOverrides...)
 			require.NoError(t, err)
 
-			require.EqualValues(t, expectedConfig.String(), config.String())
+			require.EqualValues(t, expectedConfig.String(), cfg.String())
 		})
 	}
 }

pkg/config/engine/containerd/containerd.go

@@ -17,15 +17,16 @@
 package containerd
 
 import (
-	"github.com/pelletier/go-toml"
+	"fmt"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
 )
 
 // Config represents the containerd config
 type Config struct {
-	*toml.Tree
+	*config.Toml
 	Logger                logger.Interface
 	RuntimeType           string
 	UseDefaultRuntimeName bool
@@ -36,14 +37,64 @@ var _ engine.Interface = (*Config)(nil)
 
 // New creates a containerd config with the specified options
 func New(opts ...Option) (engine.Interface, error) {
-	b := &builder{}
+	b := &builder{
+		runtimeType: defaultRuntimeType,
+	}
 	for _, opt := range opts {
 		opt(b)
 	}
-
 	if b.logger == nil {
 		b.logger = logger.New()
 	}
+	if b.reference == nil {
+		b.reference = config.FromFile(b.path)
+	}
+
+	tomlConfig, err := b.reference.Load()
+	if err != nil {
+		return nil, fmt.Errorf("failed to load config: %v", err)
+	}
+
+	cfg := &Config{
+		Toml:                  tomlConfig,
+		Logger:                b.logger,
+		RuntimeType:           b.runtimeType,
+		UseDefaultRuntimeName: b.useLegacyConfig,
+		ContainerAnnotations:  b.containerAnnotations,
+	}
 
-	return b.build()
+	version, err := cfg.parseVersion(b.useLegacyConfig)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse config version: %v", err)
+	}
+	switch version {
+	case 1:
+		return (*ConfigV1)(cfg), nil
+	case 2:
+		return cfg, nil
+	}
+
+	return nil, fmt.Errorf("unsupported config version: %v", version)
+}
+
+// parseVersion returns the version of the config
+func (c *Config) parseVersion(useLegacyConfig bool) (int, error) {
+	defaultVersion := 2
+	if useLegacyConfig {
+		defaultVersion = 1
+	}
+
+	switch v := c.Get("version").(type) {
+	case nil:
+		switch len(c.Keys()) {
+		case 0: // No config exists, or the config file is empty, use version inferred from containerd
+			return defaultVersion, nil
+		default: // A config file exists, has content, and no version is set
+			return 1, nil
+		}
+	case int64:
+		return int(v), nil
+	default:
+		return -1, fmt.Errorf("unsupported type for version field: %v", v)
+	}
 }

pkg/config/engine/containerd/option.go

@@ -17,13 +17,8 @@
 package containerd
 
 import (
-	"fmt"
-	"os"
-
-	"github.com/pelletier/go-toml"
-
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
-	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config"
 )
 
 const (
@@ -32,6 +27,7 @@ const (
 
 type builder struct {
 	logger               logger.Interface
+	reference            config.Loader
 	path                 string
 	runtimeType          string
 	useLegacyConfig      bool
@@ -55,6 +51,13 @@ func WithPath(path string) Option {
 	}
 }
 
+// WithReference sets the TOML source for the config.
+func WithReference(reference config.Loader) Option {
+	return func(b *builder) {
+		b.reference = reference
+	}
+}
+
 // WithRuntimeType sets the runtime type for the config builder
 func WithRuntimeType(runtimeType string) Option {
 	return func(b *builder) {
@@ -75,82 +78,3 @@ func WithContainerAnnotations(containerAnnotations ...string) Option {
 		b.containerAnnotations = containerAnnotations
 	}
 }
-
-func (b *builder) build() (engine.Interface, error) {
-	if b.path == "" {
-		return nil, fmt.Errorf("config path is empty")
-	}
-
-	if b.runtimeType == "" {
-		b.runtimeType = defaultRuntimeType
-	}
-
-	config, err := b.loadConfig(b.path)
-	if err != nil {
-		return nil, fmt.Errorf("failed to load config: %v", err)
-	}
-	config.Logger = b.logger
-	config.RuntimeType = b.runtimeType
-	config.UseDefaultRuntimeName = !b.useLegacyConfig
-	config.ContainerAnnotations = b.containerAnnotations
-
-	version, err := config.parseVersion(b.useLegacyConfig)
-	if err != nil {
-		return nil, fmt.Errorf("failed to parse config version: %v", err)
-	}
-	switch version {
-	case 1:
-		return (*ConfigV1)(config), nil
-	case 2:
-		return config, nil
-	}
-
-	return nil, fmt.Errorf("unsupported config version: %v", version)
-}
-
-// loadConfig loads the containerd config from disk
-func (b *builder) loadConfig(config string) (*Config, error) {
-	info, err := os.Stat(config)
-	if os.IsExist(err) && info.IsDir() {
-		return nil, fmt.Errorf("config file is a directory")
-	}
-
-	if os.IsNotExist(err) {
-		b.logger.Infof("Config file does not exist; using empty config")
-		config = "/dev/null"
-	} else {
-		b.logger.Infof("Loading config from %v", config)
-	}
-
-	tomlConfig, err := toml.LoadFile(config)
-	if err != nil {
-		return nil, err
-	}
-
-	cfg := Config{
-		Tree: tomlConfig,
-	}
-	return &cfg, nil
-}
-
-// parseVersion returns the version of the config
-func (c *Config) parseVersion(useLegacyConfig bool) (int, error) {
-	defaultVersion := 2
-	if useLegacyConfig {
-		defaultVersion = 1
-	}
-
-	switch v := c.Get("version").(type) {
-	case nil:
-		switch len(c.Keys()) {
-		case 0: // No config exists, or the config file is empty, use version inferred from containerd
-			return defaultVersion, nil
-		default: // A config file exists, has content, and no version is set
-			return 1, nil
-		}
-	case int64:
-		return int(v), nil
-	default:
-		return -1, fmt.Errorf("unsupported type for version field: %v", v)
-	}
-}

pkg/config/engine/crio/crio.go

@@ -22,12 +22,13 @@ import (
 	"github.com/pelletier/go-toml"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
 )
 
 // Config represents the cri-o config
 type Config struct {
-	*toml.Tree
+	*config.Toml
 	Logger logger.Interface
 }
 
@@ -39,8 +40,23 @@ func New(opts ...Option) (engine.Interface, error) {
 	for _, opt := range opts {
 		opt(b)
 	}
+	if b.logger == nil {
+		b.logger = logger.New()
+	}
+	if b.reference == nil {
+		b.reference = config.FromFile(b.path)
+	}
+
+	tomlConfig, err := b.reference.Load()
+	if err != nil {
+		return nil, err
+	}
 
-	return b.build()
+	cfg := Config{
+		Toml:   tomlConfig,
+		Logger: b.logger,
+	}
+	return &cfg, nil
 }
 
 // AddRuntime adds a new runtime to the crio config
@@ -115,22 +131,3 @@ func (c *Config) RemoveRuntime(name string) error {
 	*c.Tree = config
 	return nil
 }
-
-// Set sets the specified cri-o option.
-func (c *Config) Set(key string, value interface{}) {
-	config := *c.Tree
-	config.Set(key, value)
-	*c.Tree = config
-}
-
-// Save writes the config to the specified path
-func (c Config) Save(path string) (int64, error) {
-	config := c.Tree
-	output, err := config.Marshal()
-	if err != nil {
-		return 0, fmt.Errorf("unable to convert to TOML: %v", err)
-	}
-
-	n, err := engine.Config(path).Write(output)
-	return int64(n), err
-}