Migrate containerd config to engine.Interface

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

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,36 +91,46 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
-		o := &options{
-			useLegacyConfig: tc.legacyConfig,
-			setAsDefault:    tc.setAsDefault,
-			runtimeClass:    tc.runtimeClass,
-			runtimeType:     runtimeType,
-			runtimeDir:      runtimeDir,
-		}
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			o := &options{
+				useLegacyConfig: tc.legacyConfig,
+				setAsDefault:    tc.setAsDefault,
+				runtimeClass:    tc.runtimeClass,
+				runtimeType:     runtimeType,
+				runtimeDir:      runtimeDir,
+			}
 
-		config, err := toml.TreeFromMap(map[string]interface{}{})
-		require.NoError(t, err, "%d: %v", i, tc)
-
-		err = UpdateV1Config(config, o)
-		require.NoError(t, err, "%d: %v", i, tc)
-
-		defaultRuntimeName := config.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"})
-		if tc.expectedDefaultRuntimeBinary == nil {
-			require.Nil(t, defaultRuntime, "%d: %v", i, tc)
-		} else {
-			expected, err := defaultRuntimeTomlConfigV1(tc.expectedDefaultRuntimeBinary.(string))
+			config, err := toml.TreeFromMap(map[string]interface{}{})
 			require.NoError(t, err, "%d: %v", i, tc)
 
-			configContents, _ := toml.Marshal(defaultRuntime.(*toml.Tree))
-			expectedContents, _ := toml.Marshal(expected)
+			v1 := &containerd.ConfigV1{
+				Tree:                  config,
+				UseDefaultRuntimeName: !tc.legacyConfig,
+				RuntimeType:           runtimeType,
+			}
 
-			require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, tc)
-		}
+			err = UpdateConfig(v1, o)
+			require.NoError(t, err, "%d: %v", i, tc)
 
+			defaultRuntimeName := v1.GetPath([]string{"plugins", "cri", "containerd", "default_runtime_name"})
+			require.EqualValues(t, tc.expectedDefaultRuntimeName, defaultRuntimeName, "%d: %v", i, tc)
+
+			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)
+
+				configContents, _ := toml.Marshal(defaultRuntime.(*toml.Tree))
+				expectedContents, _ := toml.Marshal(expected)
+
+				require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, tc)
+			}
+
+		})
 	}
 }
 
@@ -150,40 +162,48 @@ func TestUpdateV1Config(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
-		o := &options{
-			runtimeClass: tc.runtimeClass,
-			runtimeType:  runtimeType,
-			runtimeDir:   runtimeDir,
-		}
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			o := &options{
+				runtimeClass: tc.runtimeClass,
+				runtimeType:  runtimeType,
+				runtimeDir:   runtimeDir,
+			}
 
-		config, err := toml.TreeFromMap(map[string]interface{}{})
-		require.NoError(t, err, "%d: %v", i, tc)
-
-		err = UpdateV1Config(config, o)
-		require.NoError(t, err, "%d: %v", i, tc)
-
-		version, ok := config.Get("version").(int64)
-		require.True(t, ok)
-		require.EqualValues(t, expectedVersion, version)
-
-		runtimes, ok := config.GetPath([]string{"plugins", "cri", "containerd", "runtimes"}).(*toml.Tree)
-		require.True(t, ok)
-
-		runtimeClasses := runtimes.Keys()
-		require.ElementsMatch(t, tc.expectedRuntimes, runtimeClasses, "%d: %v", i, tc)
-
-		for i, r := range tc.expectedRuntimes {
-			runtimeConfig := runtimes.Get(r)
-
-			expected, err := runtimeTomlConfigV1(expectedBinaries[i])
+			config, err := toml.TreeFromMap(map[string]interface{}{})
 			require.NoError(t, err, "%d: %v", i, tc)
 
-			configContents, _ := toml.Marshal(runtimeConfig)
-			expectedContents, _ := toml.Marshal(expected)
+			v1 := &containerd.ConfigV1{
+				Tree:                  config,
+				UseDefaultRuntimeName: true,
+				RuntimeType:           runtimeType,
+			}
 
-			require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, r, tc)
+			err = UpdateConfig(v1, o)
+			require.NoError(t, err, "%d: %v", i, tc)
 
-		}
+			version, ok := config.Get("version").(int64)
+			require.True(t, ok)
+			require.EqualValues(t, expectedVersion, version)
+
+			runtimes, ok := config.GetPath([]string{"plugins", "cri", "containerd", "runtimes"}).(*toml.Tree)
+			require.True(t, ok)
+
+			runtimeClasses := runtimes.Keys()
+			require.ElementsMatch(t, tc.expectedRuntimes, runtimeClasses, "%d: %v", i, tc)
+
+			for i, r := range tc.expectedRuntimes {
+				runtimeConfig := runtimes.Get(r)
+
+				expected, err := runtimeTomlConfigV1(expectedBinaries[i])
+				require.NoError(t, err, "%d: %v", i, tc)
+
+				configContents, _ := toml.Marshal(runtimeConfig)
+				expectedContents, _ := toml.Marshal(expected)
+
+				require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, r, tc)
+
+			}
+		})
 	}
 }
 
@@ -217,40 +237,48 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
-		o := &options{
-			runtimeClass: tc.runtimeClass,
-			runtimeType:  runtimeType,
-			runtimeDir:   runtimeDir,
-		}
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			o := &options{
+				runtimeClass: tc.runtimeClass,
+				runtimeType:  runtimeType,
+				runtimeDir:   runtimeDir,
+			}
 
-		config, err := toml.TreeFromMap(runcConfigMapV1("/runc-binary"))
-		require.NoError(t, err, "%d: %v", i, tc)
-
-		err = UpdateV1Config(config, o)
-		require.NoError(t, err, "%d: %v", i, tc)
-
-		version, ok := config.Get("version").(int64)
-		require.True(t, ok)
-		require.EqualValues(t, expectedVersion, version)
-
-		runtimes, ok := config.GetPath([]string{"plugins", "cri", "containerd", "runtimes"}).(*toml.Tree)
-		require.True(t, ok)
-
-		runtimeClasses := runtimes.Keys()
-		require.ElementsMatch(t, tc.expectedRuntimes, runtimeClasses, "%d: %v", i, tc)
-
-		for i, r := range tc.expectedRuntimes {
-			runtimeConfig := runtimes.Get(r)
-
-			expected, err := toml.TreeFromMap(runcRuntimeConfigMapV1(expectedBinaries[i]))
+			config, err := toml.TreeFromMap(runcConfigMapV1("/runc-binary"))
 			require.NoError(t, err, "%d: %v", i, tc)
 
-			configContents, _ := toml.Marshal(runtimeConfig)
-			expectedContents, _ := toml.Marshal(expected)
+			v1 := &containerd.ConfigV1{
+				Tree:                  config,
+				UseDefaultRuntimeName: true,
+				RuntimeType:           runtimeType,
+			}
 
-			require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, r, tc)
+			err = UpdateConfig(v1, o)
+			require.NoError(t, err, "%d: %v", i, tc)
 
-		}
+			version, ok := v1.Get("version").(int64)
+			require.True(t, ok)
+			require.EqualValues(t, expectedVersion, version)
+
+			runtimes, ok := v1.GetPath([]string{"plugins", "cri", "containerd", "runtimes"}).(*toml.Tree)
+			require.True(t, ok)
+
+			runtimeClasses := runtimes.Keys()
+			require.ElementsMatch(t, tc.expectedRuntimes, runtimeClasses, "%d: %v", i, tc)
+
+			for i, r := range tc.expectedRuntimes {
+				runtimeConfig := runtimes.Get(r)
+
+				expected, err := toml.TreeFromMap(runcRuntimeConfigMapV1(expectedBinaries[i]))
+				require.NoError(t, err, "%d: %v", i, tc)
+
+				configContents, _ := toml.Marshal(runtimeConfig)
+				expectedContents, _ := toml.Marshal(expected)
+
+				require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, r, tc)
+
+			}
+		})
 	}
 }
 
@@ -301,23 +329,31 @@ func TestRevertV1Config(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
-		o := &options{
-			runtimeClass: "nvidia",
-		}
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			o := &options{
+				runtimeClass: "nvidia",
+			}
 
-		config, err := toml.TreeFromMap(tc.config)
-		require.NoError(t, err, "%d: %v", i, tc)
+			config, err := toml.TreeFromMap(tc.config)
+			require.NoError(t, err, "%d: %v", i, tc)
 
-		expected, err := toml.TreeFromMap(tc.expected)
-		require.NoError(t, err, "%d: %v", i, tc)
+			expected, err := toml.TreeFromMap(tc.expected)
+			require.NoError(t, err, "%d: %v", i, tc)
 
-		err = RevertV1Config(config, o)
-		require.NoError(t, err, "%d: %v", i, tc)
+			v1 := &containerd.ConfigV1{
+				Tree:                  config,
+				UseDefaultRuntimeName: true,
+				RuntimeType:           runtimeType,
+			}
 
-		configContents, _ := toml.Marshal(config)
-		expectedContents, _ := toml.Marshal(expected)
+			err = RevertConfig(v1, o)
+			require.NoError(t, err, "%d: %v", i, tc)
 
-		require.Equal(t, string(expectedContents), string(configContents), "%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,40 +194,47 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
-		o := &options{
-			runtimeClass: tc.runtimeClass,
-			runtimeType:  runtimeType,
-			runtimeDir:   runtimeDir,
-		}
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			o := &options{
+				runtimeClass: tc.runtimeClass,
+				runtimeType:  runtimeType,
+				runtimeDir:   runtimeDir,
+			}
 
-		config, err := toml.TreeFromMap(runcConfigMapV2("/runc-binary"))
-		require.NoError(t, err, "%d: %v", i, tc)
-
-		err = UpdateV2Config(config, o)
-		require.NoError(t, err, "%d: %v", i, tc)
-
-		version, ok := config.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)
-		require.True(t, ok, "%d: %v", i, tc)
-
-		runtimeClasses := runtimes.Keys()
-		require.ElementsMatch(t, tc.expectedRuntimes, runtimeClasses, "%d: %v", i, tc)
-
-		for i, r := range tc.expectedRuntimes {
-			runtimeConfig := runtimes.Get(r)
-
-			expected, err := toml.TreeFromMap(runcRuntimeConfigMapV2(expectedBinaries[i]))
+			config, err := toml.TreeFromMap(runcConfigMapV2("/runc-binary"))
 			require.NoError(t, err, "%d: %v", i, tc)
 
-			configContents, _ := toml.Marshal(runtimeConfig)
-			expectedContents, _ := toml.Marshal(expected)
+			v2 := &containerd.Config{
+				Tree:        config,
+				RuntimeType: runtimeType,
+			}
 
-			require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, r, tc)
+			err = UpdateConfig(v2, o)
+			require.NoError(t, err, "%d: %v", i, tc)
 
-		}
+			version, ok := v2.Get("version").(int64)
+			require.True(t, ok)
+			require.EqualValues(t, expectedVersion, version)
+
+			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()
+			require.ElementsMatch(t, tc.expectedRuntimes, runtimeClasses, "%d: %v", i, tc)
+
+			for i, r := range tc.expectedRuntimes {
+				runtimeConfig := runtimes.Get(r)
+
+				expected, err := toml.TreeFromMap(runcRuntimeConfigMapV2(expectedBinaries[i]))
+				require.NoError(t, err, "%d: %v", i, tc)
+
+				configContents, _ := toml.Marshal(runtimeConfig)
+				expectedContents, _ := toml.Marshal(expected)
+
+				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)
-	if err != nil {
-		return fmt.Errorf("unable to parse version: %v", err)
-	}
-
-	err = UpdateConfig(cfg, o, version)
+	err = UpdateConfig(cfg, o)
 	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)
-	if err != nil {
-		return fmt.Errorf("unable to parse version: %v", err)
-	}
-
-	err = RevertConfig(cfg, o, version)
+	err = RevertConfig(cfg, o)
 	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 {
-	var err error
-
-	log.Infof("Updating config")
-	switch version {
-	case 1:
-		err = UpdateV1Config(config, o)
-	case 2:
-		err = UpdateV2Config(config, o)
-	default:
-		err = fmt.Errorf("unsupported containerd config version: %v", version)
+func UpdateConfig(cfg engine.Interface, o *options) error {
+	runtimes := operator.GetRuntimes(
+		operator.WithNvidiaRuntimeName(o.runtimeClass),
+		operator.WithSetAsDefault(o.setAsDefault),
+		operator.WithRoot(o.runtimeDir),
+	)
+	for class, runtime := range runtimes {
+		err := cfg.AddRuntime(class, runtime.Path, runtime.SetAsDefault)
+		if err != nil {
+			return fmt.Errorf("unable to update config for runtime class '%v': %v", class, err)
+		}
 	}
-	if err != nil {
-		return 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 {
-	var err error
-
-	log.Infof("Reverting config")
-	switch version {
-	case 1:
-		err = RevertV1Config(config, o)
-	case 2:
-		err = RevertV2Config(config, o)
-	default:
-		err = fmt.Errorf("unsupported containerd config version: %v", version)
-	}
-	if err != nil {
-		return 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)
+func RevertConfig(cfg engine.Interface, o *options) error {
+	runtimes := operator.GetRuntimes(
+		operator.WithNvidiaRuntimeName(o.runtimeClass),
+		operator.WithSetAsDefault(o.setAsDefault),
+		operator.WithRoot(o.runtimeDir),
+	)
+	for class := range runtimes {
+		err := cfg.RemoveRuntime(class)
 		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)
+			return fmt.Errorf("unable to revert config for runtime class '%v': %v", class, 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
-}