Migrate containerd config to engine.Interface

Signed-off-by: Evan Lezar <elezar@nvidia.com>

internal/config/engine/containerd/config_v1.go

@@ -0,0 +1,130 @@
+/**
+# 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 containerd
+
+import (
+	"fmt"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
+	"github.com/pelletier/go-toml"
+)
+
+// ConfigV1 represents a version 1 containerd config
+type ConfigV1 Config
+
+var _ engine.Interface = (*ConfigV1)(nil)
+
+// AddRuntime adds a runtime to the containerd config
+func (c *ConfigV1) AddRuntime(name string, path string, setAsDefault bool) error {
+	if c == nil || c.Tree == nil {
+		return fmt.Errorf("config is nil")
+	}
+
+	config := *c.Tree
+
+	config.Set("version", int64(1))
+
+	switch runc := config.GetPath([]string{"plugins", "cri", "containerd", "runtimes", "runc"}).(type) {
+	case *toml.Tree:
+		runc, _ = toml.Load(runc.String())
+		config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name}, runc)
+	}
+
+	if config.GetPath([]string{"plugins", "cri", "containerd", "runtimes", name}) == nil {
+		config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "runtime_type"}, c.RuntimeType)
+		config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "runtime_root"}, "")
+		config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "runtime_engine"}, "")
+		config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "privileged_without_host_devices"}, false)
+	}
+	config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "options", "BinaryName"}, path)
+	config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "options", "Runtime"}, path)
+
+	if setAsDefault && c.UseDefaultRuntimeName {
+		config.SetPath([]string{"plugins", "cri", "containerd", "default_runtime_name"}, name)
+	} else if setAsDefault {
+		// Note: This is deprecated in containerd 1.4.0 and will be removed in 1.5.0
+		if config.GetPath([]string{"plugins", "cri", "containerd", "default_runtime"}) == nil {
+			config.SetPath([]string{"plugins", "cri", "containerd", "default_runtime", "runtime_type"}, c.RuntimeType)
+			config.SetPath([]string{"plugins", "cri", "containerd", "default_runtime", "runtime_root"}, "")
+			config.SetPath([]string{"plugins", "cri", "containerd", "default_runtime", "runtime_engine"}, "")
+			config.SetPath([]string{"plugins", "cri", "containerd", "default_runtime", "privileged_without_host_devices"}, false)
+		}
+		config.SetPath([]string{"plugins", "cri", "containerd", "default_runtime", "options", "BinaryName"}, path)
+		config.SetPath([]string{"plugins", "cri", "containerd", "default_runtime", "options", "Runtime"}, path)
+	}
+
+	*c.Tree = config
+	return nil
+}
+
+// DefaultRuntime returns the default runtime for the cri-o config
+func (c ConfigV1) DefaultRuntime() string {
+	if runtime, ok := c.GetPath([]string{"plugins", "cri", "containerd", "default_runtime_name"}).(string); ok {
+		return runtime
+	}
+	return ""
+}
+
+// RemoveRuntime removes a runtime from the docker config
+func (c *ConfigV1) RemoveRuntime(name string) error {
+	if c == nil || c.Tree == nil {
+		return nil
+	}
+
+	config := *c.Tree
+
+	// If the specified runtime was set as the default runtime we need to remove the default runtime too.
+	runtimePath, ok := config.GetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "options", "BinaryName"}).(string)
+	if !ok || runtimePath == "" {
+		runtimePath, _ = config.GetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "options", "Runtime"}).(string)
+	}
+	defaultRuntimePath, ok := config.GetPath([]string{"plugins", "cri", "containerd", "default_runtime", "options", "BinaryName"}).(string)
+	if !ok || defaultRuntimePath == "" {
+		defaultRuntimePath, _ = config.GetPath([]string{"plugins", "cri", "containerd", "default_runtime", "options", "Runtime"}).(string)
+	}
+	if runtimePath != "" && defaultRuntimePath != "" && runtimePath == defaultRuntimePath {
+		config.DeletePath([]string{"plugins", "cri", "containerd", "default_runtime"})
+	}
+
+	config.DeletePath([]string{"plugins", "cri", "containerd", "runtimes", name})
+	if runtime, ok := config.GetPath([]string{"plugins", "cri", "containerd", "default_runtime_name"}).(string); ok {
+		if runtime == name {
+			config.DeletePath([]string{"plugins", "cri", "containerd", "default_runtime_name"})
+		}
+	}
+
+	runtimeConfigPath := []string{"plugins", "cri", "containerd", "runtimes", name}
+	for i := 0; i < len(runtimeConfigPath); i++ {
+		if runtimes, ok := config.GetPath(runtimeConfigPath[:len(runtimeConfigPath)-i]).(*toml.Tree); ok {
+			if len(runtimes.Keys()) == 0 {
+				config.DeletePath(runtimeConfigPath[:len(runtimeConfigPath)-i])
+			}
+		}
+	}
+
+	if len(config.Keys()) == 1 && config.Keys()[0] == "version" {
+		config.Delete("version")
+	}
+
+	*c.Tree = config
+	return nil
+}
+
+// Save wrotes the config to a file
+func (c ConfigV1) Save(path string) (int64, error) {
+	return (Config)(c).Save(path)
+}

internal/config/engine/containerd/config_v2.go

@@ -0,0 +1,125 @@
+/**
+# 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 containerd
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/pelletier/go-toml"
+)
+
+// AddRuntime adds a runtime to the containerd config
+func (c *Config) AddRuntime(name string, path string, setAsDefault bool) error {
+	if c == nil || c.Tree == nil {
+		return fmt.Errorf("config is nil")
+	}
+	config := *c.Tree
+
+	config.Set("version", int64(2))
+
+	switch runc := config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", "runc"}).(type) {
+	case *toml.Tree:
+		runc, _ = toml.Load(runc.String())
+		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name}, runc)
+	}
+
+	if config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name}) == nil {
+		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "runtime_type"}, c.RuntimeType)
+		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "runtime_root"}, "")
+		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "runtime_engine"}, "")
+		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "privileged_without_host_devices"}, false)
+	}
+	config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "options", "BinaryName"}, path)
+
+	if setAsDefault {
+		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"}, name)
+	}
+
+	*c.Tree = config
+	return nil
+}
+
+// DefaultRuntime returns the default runtime for the cri-o config
+func (c Config) DefaultRuntime() string {
+	if runtime, ok := c.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"}).(string); ok {
+		return runtime
+	}
+	return ""
+}
+
+// RemoveRuntime removes a runtime from the docker config
+func (c *Config) RemoveRuntime(name string) error {
+	if c == nil || c.Tree == nil {
+		return nil
+	}
+
+	config := *c.Tree
+
+	config.DeletePath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name})
+	if runtime, ok := config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"}).(string); ok {
+		if runtime == name {
+			config.DeletePath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"})
+		}
+	}
+
+	runtimePath := []string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name}
+	for i := 0; i < len(runtimePath); i++ {
+		if runtimes, ok := config.GetPath(runtimePath[:len(runtimePath)-i]).(*toml.Tree); ok {
+			if len(runtimes.Keys()) == 0 {
+				config.DeletePath(runtimePath[:len(runtimePath)-i])
+			}
+		}
+	}
+
+	if len(config.Keys()) == 1 && config.Keys()[0] == "version" {
+		config.Delete("version")
+	}
+
+	*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.ToTomlString()
+	if err != nil {
+		return 0, fmt.Errorf("unable to convert to TOML: %v", err)
+	}
+
+	if len(output) == 0 {
+		err := os.Remove(path)
+		if err != nil {
+			return 0, fmt.Errorf("unable to remove empty file: %v", err)
+		}
+		return 0, nil
+	}
+
+	f, err := os.Create(path)
+	if err != nil {
+		return 0, fmt.Errorf("unable to open '%v' for writing: %v", path, err)
+	}
+	defer f.Close()
+
+	n, err := f.WriteString(output)
+	if err != nil {
+		return 0, fmt.Errorf("unable to write output: %v", err)
+	}
+
+	return int64(n), err
+}

internal/config/engine/containerd/containerd.go

@@ -0,0 +1,39 @@
+/**
+# Copyright (c) 2022, 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 containerd
+
+import (
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
+	"github.com/pelletier/go-toml"
+)
+
+// Config represents the containerd config
+type Config struct {
+	*toml.Tree
+	RuntimeType           string
+	UseDefaultRuntimeName bool
+}
+
+// New creates a containerd config with the specified options
+func New(opts ...Option) (engine.Interface, error) {
+	b := &builder{}
+	for _, opt := range opts {
+		opt(b)
+	}
+
+	return b.build()
+}

internal/config/engine/containerd/option.go

@@ -0,0 +1,140 @@
+/**
+# 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 containerd
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
+	"github.com/pelletier/go-toml"
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	defaultRuntimeType = "io.containerd.runc.v2"
+)
+
+type builder struct {
+	path            string
+	runtimeType     string
+	useLegacyConfig bool
+}
+
+// Option defines a function that can be used to configure the config builder
+type Option func(*builder)
+
+// WithPath sets the path for the config builder
+func WithPath(path string) Option {
+	return func(b *builder) {
+		b.path = path
+	}
+}
+
+// WithRuntimeType sets the runtime type for the config builder
+func WithRuntimeType(runtimeType string) Option {
+	return func(b *builder) {
+		b.runtimeType = runtimeType
+	}
+}
+
+// WithUseLegacyConfig sets the useLegacyConfig flag for the config builder
+func WithUseLegacyConfig(useLegacyConfig bool) Option {
+	return func(b *builder) {
+		b.useLegacyConfig = useLegacyConfig
+	}
+}
+
+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 := loadConfig(b.path)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load config: %v", err)
+	}
+	config.RuntimeType = b.runtimeType
+	config.UseDefaultRuntimeName = !b.useLegacyConfig
+
+	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 loadConfig(config string) (*Config, error) {
+	log.Infof("Loading config: %v", config)
+
+	info, err := os.Stat(config)
+	if os.IsExist(err) && info.IsDir() {
+		return nil, fmt.Errorf("config file is a directory")
+	}
+
+	configFile := config
+	if os.IsNotExist(err) {
+		configFile = "/dev/null"
+		log.Infof("Config file does not exist, creating new one")
+	}
+
+	tomlConfig, err := toml.LoadFile(configFile)
+	if err != nil {
+		return nil, err
+	}
+
+	log.Infof("Successfully loaded config")
+
+	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)
+	}
+}

tools/container/containerd/config.go

@@ -1,114 +0,0 @@
-/**
-# Copyright (c) 2020-2021, 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 main
-
-import (
-	"github.com/pelletier/go-toml"
-)
-
-// UpdateReverter defines the interface for applying and reverting configurations
-type UpdateReverter interface {
-	Update(o *options) error
-	Revert(o *options) error
-}
-
-type config struct {
-	*toml.Tree
-	version int64
-	cri     string
-}
-
-// update adds the specified runtime class to the the containerd config.
-// if set-as default is specified, the runtime class is also set as the
-// default runtime.
-func (config *config) update(runtimeClass string, runtimeType string, runtimeBinary string, setAsDefault bool) {
-	config.Set("version", config.version)
-
-	runcPath := config.runcPath()
-	runtimeClassPath := config.runtimeClassPath(runtimeClass)
-
-	switch runc := config.GetPath(runcPath).(type) {
-	case *toml.Tree:
-		runc, _ = toml.Load(runc.String())
-		config.SetPath(runtimeClassPath, runc)
-	}
-
-	config.initRuntime(runtimeClassPath, runtimeType, "BinaryName", runtimeBinary)
-	if config.version == 1 {
-		config.initRuntime(runtimeClassPath, runtimeType, "Runtime", runtimeBinary)
-	}
-
-	if setAsDefault {
-		defaultRuntimeNamePath := config.defaultRuntimeNamePath()
-		config.SetPath(defaultRuntimeNamePath, runtimeClass)
-	}
-}
-
-// revert removes the configuration applied in an update call.
-func (config *config) revert(runtimeClass string) {
-	runtimeClassPath := config.runtimeClassPath(runtimeClass)
-	defaultRuntimeNamePath := config.defaultRuntimeNamePath()
-
-	config.DeletePath(runtimeClassPath)
-	if runtime, ok := config.GetPath(defaultRuntimeNamePath).(string); ok {
-		if runtimeClass == runtime {
-			config.DeletePath(defaultRuntimeNamePath)
-		}
-	}
-
-	for i := 0; i < len(runtimeClassPath); i++ {
-		if runtimes, ok := config.GetPath(runtimeClassPath[:len(runtimeClassPath)-i]).(*toml.Tree); ok {
-			if len(runtimes.Keys()) == 0 {
-				config.DeletePath(runtimeClassPath[:len(runtimeClassPath)-i])
-			}
-		}
-	}
-
-	if len(config.Keys()) == 1 && config.Keys()[0] == "version" {
-		config.Delete("version")
-	}
-}
-
-// initRuntime creates a runtime config if it does not exist and ensures that the
-// runtimes binary path is specified.
-func (config *config) initRuntime(path []string, runtimeType string, binaryKey string, binary string) {
-	if config.GetPath(path) == nil {
-		config.SetPath(append(path, "runtime_type"), runtimeType)
-		config.SetPath(append(path, "runtime_root"), "")
-		config.SetPath(append(path, "runtime_engine"), "")
-		config.SetPath(append(path, "privileged_without_host_devices"), false)
-	}
-
-	binaryPath := append(path, "options", binaryKey)
-	config.SetPath(binaryPath, binary)
-}
-
-func (config config) runcPath() []string {
-	return config.runtimeClassPath("runc")
-}
-
-func (config config) runtimeClassPath(runtimeClass string) []string {
-	return append(config.containerdPath(), "runtimes", runtimeClass)
-}
-
-func (config config) defaultRuntimeNamePath() []string {
-	return append(config.containerdPath(), "default_runtime_name")
-}
-
-func (config config) containerdPath() []string {
-	return []string{"plugins", config.cri, "containerd"}
-}

tools/container/containerd/config_v1.go

@@ -1,134 +0,0 @@
-/**
-# Copyright (c) 2020-2021, 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 main
-
-import (
-	"path"
-
-	"github.com/pelletier/go-toml"
-	log "github.com/sirupsen/logrus"
-)
-
-// configV1 represents a V1 containerd config
-type configV1 struct {
-	config
-}
-
-func newConfigV1(cfg *toml.Tree) UpdateReverter {
-	c := configV1{
-		config: config{
-			Tree:    cfg,
-			version: 1,
-			cri:     "cri",
-		},
-	}
-
-	return &c
-}
-
-// Update performs an update specific to v1 of the containerd config
-func (config *configV1) Update(o *options) error {
-
-	// For v1 config, the `default_runtime_name` setting is only supported
-	// for containerd version at least v1.3
-	supportsDefaultRuntimeName := !o.useLegacyConfig
-
-	defaultRuntime := o.getDefaultRuntime()
-
-	for runtimeClass, runtimeBinary := range o.getRuntimeBinaries() {
-		isDefaultRuntime := runtimeClass == defaultRuntime
-		config.update(runtimeClass, o.runtimeType, runtimeBinary, isDefaultRuntime && supportsDefaultRuntimeName)
-
-		if !isDefaultRuntime {
-			continue
-		}
-
-		if supportsDefaultRuntimeName {
-			defaultRuntimePath := append(config.containerdPath(), "default_runtime")
-			if config.GetPath(defaultRuntimePath) != nil {
-				log.Warnf("The setting of default_runtime (%v) in containerd is deprecated", defaultRuntimePath)
-			}
-			continue
-		}
-
-		log.Warnf("Setting default_runtime is deprecated")
-		defaultRuntimePath := append(config.containerdPath(), "default_runtime")
-		config.initRuntime(defaultRuntimePath, o.runtimeType, "Runtime", runtimeBinary)
-		config.initRuntime(defaultRuntimePath, o.runtimeType, "BinaryName", runtimeBinary)
-	}
-	return nil
-}
-
-// Revert performs a revert specific to v1 of the containerd config
-func (config *configV1) Revert(o *options) error {
-	defaultRuntimePath := append(config.containerdPath(), "default_runtime")
-	defaultRuntimeOptionsPath := append(defaultRuntimePath, "options")
-	if runtime, ok := config.GetPath(append(defaultRuntimeOptionsPath, "Runtime")).(string); ok {
-		for _, runtimeBinary := range o.getRuntimeBinaries() {
-			if path.Base(runtimeBinary) == path.Base(runtime) {
-				config.DeletePath(append(defaultRuntimeOptionsPath, "Runtime"))
-				break
-			}
-		}
-	}
-	if runtime, ok := config.GetPath(append(defaultRuntimeOptionsPath, "BinaryName")).(string); ok {
-		for _, runtimeBinary := range o.getRuntimeBinaries() {
-			if path.Base(runtimeBinary) == path.Base(runtime) {
-				config.DeletePath(append(defaultRuntimeOptionsPath, "BinaryName"))
-				break
-			}
-		}
-	}
-
-	if options, ok := config.GetPath(defaultRuntimeOptionsPath).(*toml.Tree); ok {
-		if len(options.Keys()) == 0 {
-			config.DeletePath(defaultRuntimeOptionsPath)
-		}
-	}
-
-	if runtime, ok := config.GetPath(defaultRuntimePath).(*toml.Tree); ok {
-		fields := []string{"runtime_type", "runtime_root", "runtime_engine", "privileged_without_host_devices"}
-		if len(runtime.Keys()) <= len(fields) {
-			matches := []string{}
-			for _, f := range fields {
-				e := runtime.Get(f)
-				if e != nil {
-					matches = append(matches, f)
-				}
-			}
-			if len(matches) == len(runtime.Keys()) {
-				for _, m := range matches {
-					runtime.Delete(m)
-				}
-			}
-		}
-	}
-
-	for i := 0; i < len(defaultRuntimePath); i++ {
-		if runtimes, ok := config.GetPath(defaultRuntimePath[:len(defaultRuntimePath)-i]).(*toml.Tree); ok {
-			if len(runtimes.Keys()) == 0 {
-				config.DeletePath(defaultRuntimePath[:len(defaultRuntimePath)-i])
-			}
-		}
-	}
-
-	for runtimeClass := range nvidiaRuntimeBinaries {
-		config.revert(runtimeClass)
-	}
-
-	return nil
-}

tools/container/containerd/config_v1_test.go

@@ -17,8 +17,10 @@
 package main
 
 import (
+	"fmt"
 	"testing"
 
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/containerd"
 	"github.com/pelletier/go-toml"
 	"github.com/stretchr/testify/require"
 )
@@ -89,6 +91,7 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
 				useLegacyConfig: tc.legacyConfig,
 				setAsDefault:    tc.setAsDefault,
@@ -100,16 +103,24 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 			config, err := toml.TreeFromMap(map[string]interface{}{})
 			require.NoError(t, err, "%d: %v", i, tc)
 
-		err = UpdateV1Config(config, o)
+			v1 := &containerd.ConfigV1{
+				Tree:                  config,
+				UseDefaultRuntimeName: !tc.legacyConfig,
+				RuntimeType:           runtimeType,
+			}
+
+			err = UpdateConfig(v1, o)
 			require.NoError(t, err, "%d: %v", i, tc)
 
-		defaultRuntimeName := config.GetPath([]string{"plugins", "cri", "containerd", "default_runtime_name"})
+			defaultRuntimeName := v1.GetPath([]string{"plugins", "cri", "containerd", "default_runtime_name"})
 			require.EqualValues(t, tc.expectedDefaultRuntimeName, defaultRuntimeName, "%d: %v", i, tc)
 
-		defaultRuntime := config.GetPath([]string{"plugins", "cri", "containerd", "default_runtime"})
+			defaultRuntime := v1.GetPath([]string{"plugins", "cri", "containerd", "default_runtime"})
 			if tc.expectedDefaultRuntimeBinary == nil {
 				require.Nil(t, defaultRuntime, "%d: %v", i, tc)
 			} else {
+				require.NotNil(t, defaultRuntime)
+
 				expected, err := defaultRuntimeTomlConfigV1(tc.expectedDefaultRuntimeBinary.(string))
 				require.NoError(t, err, "%d: %v", i, tc)
 
@@ -119,6 +130,7 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 				require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, tc)
 			}
 
+		})
 	}
 }
 
@@ -150,6 +162,7 @@ func TestUpdateV1Config(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
 				runtimeClass: tc.runtimeClass,
 				runtimeType:  runtimeType,
@@ -159,7 +172,13 @@ func TestUpdateV1Config(t *testing.T) {
 			config, err := toml.TreeFromMap(map[string]interface{}{})
 			require.NoError(t, err, "%d: %v", i, tc)
 
-		err = UpdateV1Config(config, o)
+			v1 := &containerd.ConfigV1{
+				Tree:                  config,
+				UseDefaultRuntimeName: true,
+				RuntimeType:           runtimeType,
+			}
+
+			err = UpdateConfig(v1, o)
 			require.NoError(t, err, "%d: %v", i, tc)
 
 			version, ok := config.Get("version").(int64)
@@ -184,6 +203,7 @@ func TestUpdateV1Config(t *testing.T) {
 				require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, r, tc)
 
 			}
+		})
 	}
 }
 
@@ -217,6 +237,7 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
 				runtimeClass: tc.runtimeClass,
 				runtimeType:  runtimeType,
@@ -226,14 +247,20 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 			config, err := toml.TreeFromMap(runcConfigMapV1("/runc-binary"))
 			require.NoError(t, err, "%d: %v", i, tc)
 
-		err = UpdateV1Config(config, o)
+			v1 := &containerd.ConfigV1{
+				Tree:                  config,
+				UseDefaultRuntimeName: true,
+				RuntimeType:           runtimeType,
+			}
+
+			err = UpdateConfig(v1, o)
 			require.NoError(t, err, "%d: %v", i, tc)
 
-		version, ok := config.Get("version").(int64)
+			version, ok := v1.Get("version").(int64)
 			require.True(t, ok)
 			require.EqualValues(t, expectedVersion, version)
 
-		runtimes, ok := config.GetPath([]string{"plugins", "cri", "containerd", "runtimes"}).(*toml.Tree)
+			runtimes, ok := v1.GetPath([]string{"plugins", "cri", "containerd", "runtimes"}).(*toml.Tree)
 			require.True(t, ok)
 
 			runtimeClasses := runtimes.Keys()
@@ -251,6 +278,7 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 				require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, r, tc)
 
 			}
+		})
 	}
 }
 
@@ -301,6 +329,7 @@ func TestRevertV1Config(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
 				runtimeClass: "nvidia",
 			}
@@ -311,13 +340,20 @@ func TestRevertV1Config(t *testing.T) {
 			expected, err := toml.TreeFromMap(tc.expected)
 			require.NoError(t, err, "%d: %v", i, tc)
 
-		err = RevertV1Config(config, o)
+			v1 := &containerd.ConfigV1{
+				Tree:                  config,
+				UseDefaultRuntimeName: true,
+				RuntimeType:           runtimeType,
+			}
+
+			err = RevertConfig(v1, o)
 			require.NoError(t, err, "%d: %v", i, tc)
 
 			configContents, _ := toml.Marshal(config)
 			expectedContents, _ := toml.Marshal(expected)
 
 			require.Equal(t, string(expectedContents), string(configContents), "%d: %v", i, tc)
+		})
 	}
 }
 

tools/container/containerd/config_v2.go

@@ -1,58 +0,0 @@
-/**
-# Copyright (c) 2020-2021, 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 main
-
-import (
-	"github.com/pelletier/go-toml"
-)
-
-// configV2 represents a V2 containerd config
-type configV2 struct {
-	config
-}
-
-func newConfigV2(cfg *toml.Tree) UpdateReverter {
-	c := configV2{
-		config: config{
-			Tree:    cfg,
-			version: 2,
-			cri:     "io.containerd.grpc.v1.cri",
-		},
-	}
-
-	return &c
-}
-
-// Update performs an update specific to v2 of the containerd config
-func (config *configV2) Update(o *options) error {
-	defaultRuntime := o.getDefaultRuntime()
-	for runtimeClass, runtimeBinary := range o.getRuntimeBinaries() {
-		setAsDefault := defaultRuntime == runtimeClass
-		config.update(runtimeClass, o.runtimeType, runtimeBinary, setAsDefault)
-	}
-
-	return nil
-}
-
-// Revert performs a revert specific to v2 of the containerd config
-func (config *configV2) Revert(o *options) error {
-	for runtimeClass := range o.getRuntimeBinaries() {
-		config.revert(runtimeClass)
-	}
-
-	return nil
-}

tools/container/containerd/config_v2_test.go

@@ -17,8 +17,10 @@
 package main
 
 import (
+	"fmt"
 	"testing"
 
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/containerd"
 	"github.com/pelletier/go-toml"
 	"github.com/stretchr/testify/require"
 )
@@ -78,7 +80,12 @@ func TestUpdateV2ConfigDefaultRuntime(t *testing.T) {
 		config, err := toml.TreeFromMap(map[string]interface{}{})
 		require.NoError(t, err, "%d: %v", i, tc)
 
-		err = UpdateV2Config(config, o)
+		v2 := &containerd.Config{
+			Tree:        config,
+			RuntimeType: runtimeType,
+		}
+
+		err = UpdateConfig(v2, o)
 		require.NoError(t, err, "%d: %v", i, tc)
 
 		defaultRuntimeName := config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"})
@@ -123,7 +130,12 @@ func TestUpdateV2Config(t *testing.T) {
 		config, err := toml.TreeFromMap(map[string]interface{}{})
 		require.NoError(t, err, "%d: %v", i, tc)
 
-		err = UpdateV2Config(config, o)
+		v2 := &containerd.Config{
+			Tree:        config,
+			RuntimeType: runtimeType,
+		}
+
+		err = UpdateConfig(v2, o)
 		require.NoError(t, err, "%d: %v", i, tc)
 
 		version, ok := config.Get("version").(int64)
@@ -182,6 +194,7 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
 				runtimeClass: tc.runtimeClass,
 				runtimeType:  runtimeType,
@@ -191,14 +204,19 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 			config, err := toml.TreeFromMap(runcConfigMapV2("/runc-binary"))
 			require.NoError(t, err, "%d: %v", i, tc)
 
-		err = UpdateV2Config(config, o)
+			v2 := &containerd.Config{
+				Tree:        config,
+				RuntimeType: runtimeType,
+			}
+
+			err = UpdateConfig(v2, o)
 			require.NoError(t, err, "%d: %v", i, tc)
 
-		version, ok := config.Get("version").(int64)
+			version, ok := v2.Get("version").(int64)
 			require.True(t, ok)
 			require.EqualValues(t, expectedVersion, version)
 
-		runtimes, ok := config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes"}).(*toml.Tree)
+			runtimes, ok := v2.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes"}).(*toml.Tree)
 			require.True(t, ok, "%d: %v", i, tc)
 
 			runtimeClasses := runtimes.Keys()
@@ -216,6 +234,7 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 				require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, r, tc)
 
 			}
+		})
 	}
 }
 
@@ -275,7 +294,12 @@ func TestRevertV2Config(t *testing.T) {
 		expected, err := toml.TreeFromMap(tc.expected)
 		require.NoError(t, err, "%d: %v", i, tc)
 
-		err = RevertV2Config(config, o)
+		v2 := &containerd.Config{
+			Tree:        config,
+			RuntimeType: runtimeType,
+		}
+
+		err = RevertConfig(v2, o)
 		require.NoError(t, err, "%d: %v", i, tc)
 
 		configContents, _ := toml.Marshal(config)

tools/container/containerd/containerd.go

@@ -21,11 +21,12 @@ import (
 	"net"
 	"os"
 	"os/exec"
-	"path/filepath"
 	"syscall"
 	"time"
 
-	toml "github.com/pelletier/go-toml"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/containerd"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator"
 	log "github.com/sirupsen/logrus"
 	cli "github.com/urfave/cli/v2"
 )
@@ -194,25 +195,28 @@ func Setup(c *cli.Context, o *options) error {
 	}
 	o.runtimeDir = runtimeDir
 
-	cfg, err := LoadConfig(o.config)
+	cfg, err := containerd.New(
+		containerd.WithPath(o.config),
+		containerd.WithRuntimeType(o.runtimeType),
+		containerd.WithUseLegacyConfig(o.useLegacyConfig),
+	)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
-	version, err := ParseVersion(cfg, o.useLegacyConfig)
+	err = UpdateConfig(cfg, o)
-	if err != nil {
-		return fmt.Errorf("unable to parse version: %v", err)
-	}
-
-	err = UpdateConfig(cfg, o, version)
 	if err != nil {
 		return fmt.Errorf("unable to update config: %v", err)
 	}
 
-	err = FlushConfig(o.config, cfg)
+	log.Infof("Flushing containerd config to %v", o.config)
+	n, err := cfg.Save(o.config)
 	if err != nil {
 		return fmt.Errorf("unable to flush config: %v", err)
 	}
+	if n == 0 {
+		log.Infof("Config file is empty, removed")
+	}
 
 	err = RestartContainerd(o)
 	if err != nil {
@@ -233,25 +237,28 @@ func Cleanup(c *cli.Context, o *options) error {
 		return fmt.Errorf("unable to parse args: %v", err)
 	}
 
-	cfg, err := LoadConfig(o.config)
+	cfg, err := containerd.New(
+		containerd.WithPath(o.config),
+		containerd.WithRuntimeType(o.runtimeType),
+		containerd.WithUseLegacyConfig(o.useLegacyConfig),
+	)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
-	version, err := ParseVersion(cfg, o.useLegacyConfig)
+	err = RevertConfig(cfg, o)
-	if err != nil {
-		return fmt.Errorf("unable to parse version: %v", err)
-	}
-
-	err = RevertConfig(cfg, o, version)
 	if err != nil {
 		return fmt.Errorf("unable to update config: %v", err)
 	}
 
-	err = FlushConfig(o.config, cfg)
+	log.Infof("Flushing containerd config to %v", o.config)
+	n, err := cfg.Save(o.config)
 	if err != nil {
 		return fmt.Errorf("unable to flush config: %v", err)
 	}
+	if n == 0 {
+		log.Infof("Config file is empty, removed")
+	}
 
 	err = RestartContainerd(o)
 	if err != nil {
@@ -277,160 +284,36 @@ func ParseArgs(c *cli.Context) (string, error) {
 	return runtimeDir, nil
 }
 
-// LoadConfig loads the containerd config from disk
-func LoadConfig(config string) (*toml.Tree, error) {
-	log.Infof("Loading config: %v", config)
-
-	info, err := os.Stat(config)
-	if os.IsExist(err) && info.IsDir() {
-		return nil, fmt.Errorf("config file is a directory")
-	}
-
-	configFile := config
-	if os.IsNotExist(err) {
-		configFile = "/dev/null"
-		log.Infof("Config file does not exist, creating new one")
-	}
-
-	cfg, err := toml.LoadFile(configFile)
-	if err != nil {
-		return nil, err
-	}
-
-	log.Infof("Successfully loaded config")
-
-	return cfg, nil
-}
-
-// ParseVersion parses the version field out of the containerd config
-func ParseVersion(config *toml.Tree, useLegacyConfig bool) (int, error) {
-	var defaultVersion int
-	if !useLegacyConfig {
-		defaultVersion = 2
-	} else {
-		defaultVersion = 1
-	}
-
-	var version int
-	switch v := config.Get("version").(type) {
-	case nil:
-		switch len(config.Keys()) {
-		case 0: // No config exists, or the config file is empty, use version inferred from containerd
-			version = defaultVersion
-		default: // A config file exists, has content, and no version is set
-			version = 1
-		}
-	case int64:
-		version = int(v)
-	default:
-		return -1, fmt.Errorf("unsupported type for version field: %v", v)
-	}
-	log.Infof("Config version: %v", version)
-
-	if version == 1 {
-		log.Warnf("Support for containerd config version 1 is deprecated")
-	}
-
-	return version, nil
-}
-
 // UpdateConfig updates the containerd config to include the nvidia-container-runtime
-func UpdateConfig(config *toml.Tree, o *options, version int) error {
+func UpdateConfig(cfg engine.Interface, o *options) error {
-	var err error
+	runtimes := operator.GetRuntimes(
-
+		operator.WithNvidiaRuntimeName(o.runtimeClass),
-	log.Infof("Updating config")
+		operator.WithSetAsDefault(o.setAsDefault),
-	switch version {
+		operator.WithRoot(o.runtimeDir),
-	case 1:
+	)
-		err = UpdateV1Config(config, o)
+	for class, runtime := range runtimes {
-	case 2:
+		err := cfg.AddRuntime(class, runtime.Path, runtime.SetAsDefault)
-		err = UpdateV2Config(config, o)
-	default:
-		err = fmt.Errorf("unsupported containerd config version: %v", version)
-	}
 		if err != nil {
-		return err
+			return fmt.Errorf("unable to update config for runtime class '%v': %v", class, err)
+		}
 	}
-	log.Infof("Successfully updated config")
 
 	return nil
 }
 
 // RevertConfig reverts the containerd config to remove the nvidia-container-runtime
-func RevertConfig(config *toml.Tree, o *options, version int) error {
+func RevertConfig(cfg engine.Interface, o *options) error {
-	var err error
+	runtimes := operator.GetRuntimes(
-
+		operator.WithNvidiaRuntimeName(o.runtimeClass),
-	log.Infof("Reverting config")
+		operator.WithSetAsDefault(o.setAsDefault),
-	switch version {
+		operator.WithRoot(o.runtimeDir),
-	case 1:
+	)
-		err = RevertV1Config(config, o)
+	for class := range runtimes {
-	case 2:
+		err := cfg.RemoveRuntime(class)
-		err = RevertV2Config(config, o)
-	default:
-		err = fmt.Errorf("unsupported containerd config version: %v", version)
-	}
 		if err != nil {
-		return err
+			return fmt.Errorf("unable to revert config for runtime class '%v': %v", class, err)
-	}
-	log.Infof("Successfully reverted config")
-
-	return nil
-}
-
-// UpdateV1Config performs an update specific to v1 of the containerd config
-func UpdateV1Config(config *toml.Tree, o *options) error {
-	c := newConfigV1(config)
-	return c.Update(o)
-}
-
-// RevertV1Config performs a revert specific to v1 of the containerd config
-func RevertV1Config(config *toml.Tree, o *options) error {
-	c := newConfigV1(config)
-	return c.Revert(o)
-}
-
-// UpdateV2Config performs an update specific to v2 of the containerd config
-func UpdateV2Config(config *toml.Tree, o *options) error {
-	c := newConfigV2(config)
-	return c.Update(o)
-}
-
-// RevertV2Config performs a revert specific to v2 of the containerd config
-func RevertV2Config(config *toml.Tree, o *options) error {
-	c := newConfigV2(config)
-	return c.Revert(o)
-}
-
-// FlushConfig flushes the updated/reverted config out to disk
-func FlushConfig(config string, cfg *toml.Tree) error {
-	log.Infof("Flushing config")
-
-	output, err := cfg.ToTomlString()
-	if err != nil {
-		return fmt.Errorf("unable to convert to TOML: %v", err)
-	}
-
-	switch len(output) {
-	case 0:
-		err := os.Remove(config)
-		if err != nil {
-			return fmt.Errorf("unable to remove empty file: %v", err)
-		}
-		log.Infof("Config empty, removing file")
-	default:
-		f, err := os.Create(config)
-		if err != nil {
-			return fmt.Errorf("unable to open '%v' for writing: %v", config, err)
-		}
-		defer f.Close()
-
-		_, err = f.WriteString(output)
-		if err != nil {
-			return fmt.Errorf("unable to write output: %v", err)
 		}
 	}
-
-	log.Infof("Successfully flushed config")
-
 	return nil
 }
 
@@ -551,36 +434,3 @@ func RestartContainerdSystemd(hostRootMount string) error {
 
 	return nil
 }
-
-// getDefaultRuntime returns the default runtime for the configured options.
-// If the configuration is invalid or the default runtimes should not be set
-// the empty string is returned.
-func (o options) getDefaultRuntime() string {
-	if o.setAsDefault {
-		if o.runtimeClass == nvidiaExperimentalRuntimeName {
-			return nvidiaExperimentalRuntimeName
-		}
-		if o.runtimeClass == "" {
-			return defaultRuntimeClass
-		}
-		return o.runtimeClass
-	}
-	return ""
-}
-
-// getRuntimeBinaries returns a map of runtime names to binary paths. This includes the
-// renaming of the `nvidia` runtime as per the --runtime-class command line flag.
-func (o options) getRuntimeBinaries() map[string]string {
-	runtimeBinaries := make(map[string]string)
-
-	for rt, bin := range nvidiaRuntimeBinaries {
-		runtime := rt
-		if o.runtimeClass != "" && o.runtimeClass != nvidiaExperimentalRuntimeName && runtime == defaultRuntimeClass {
-			runtime = o.runtimeClass
-		}
-
-		runtimeBinaries[runtime] = filepath.Join(o.runtimeDir, bin)
-	}
-
-	return runtimeBinaries
-}