Migrate cri-o config update to engine.Interface

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

cmd/nvidia-ctk/runtime/configure/configure.go

@@ -173,13 +173,14 @@ func (m command) configureCrio(c *cli.Context, config *config) error {
 		configFilePath = defaultCrioConfigFilePath
 	}
 
-	cfg, err := crio.LoadConfig(configFilePath)
+	cfg, err := crio.New(
+		crio.WithPath(configFilePath),
+	)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
-	err = crio.UpdateConfig(
+	err = cfg.AddRuntime(
-		cfg,
 		config.nvidiaOptions.RuntimeName,
 		config.nvidiaOptions.RuntimePath,
 		config.nvidiaOptions.SetAsDefault,
@@ -196,12 +197,16 @@ func (m command) configureCrio(c *cli.Context, config *config) error {
 		os.Stdout.WriteString(fmt.Sprintf("%s\n", output))
 		return nil
 	}
-	err = crio.FlushConfig(configFilePath, cfg)
+	n, err := cfg.Save(configFilePath)
 	if err != nil {
 		return fmt.Errorf("unable to flush config: %v", err)
 	}
 
-	m.logger.Infof("Wrote updated config to %v", configFilePath)
+	if n == 0 {
+		m.logger.Infof("Removed empty config from %v", configFilePath)
+	} else {
+		m.logger.Infof("Wrote updated config to %v", configFilePath)
+	}
 	m.logger.Infof("It is recommended that the cri-o daemon be restarted.")
 
 	return nil

internal/config/engine/crio/crio.go

@@ -20,62 +20,71 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
 	"github.com/pelletier/go-toml"
-	log "github.com/sirupsen/logrus"
 )
 
-// LoadConfig loads the cri-o config from disk
+// Config represents the cri-o config
-func LoadConfig(config string) (*toml.Tree, error) {
+type Config toml.Tree
-	log.Infof("Loading config: %v", config)
 
-	info, err := os.Stat(config)
+// New creates a cri-o config with the specified options
-	if os.IsExist(err) && info.IsDir() {
+func New(opts ...Option) (engine.Interface, error) {
-		return nil, fmt.Errorf("config file is a directory")
+	b := &builder{}
+	for _, opt := range opts {
+		opt(b)
 	}
 
-	configFile := config
+	return b.build()
-	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
 }
 
-// UpdateConfig updates the cri-o config to include the NVIDIA Container Runtime
+// AddRuntime adds a new runtime to the crio config
-func UpdateConfig(config *toml.Tree, runtimeClass string, runtimePath string, setAsDefault bool) error {
+func (c *Config) AddRuntime(name string, path string, setAsDefault bool) error {
+	if c == nil {
+		return fmt.Errorf("config is nil")
+	}
+
+	config := (toml.Tree)(*c)
+
 	switch runc := config.Get("crio.runtime.runtimes.runc").(type) {
 	case *toml.Tree:
 		runc, _ = toml.Load(runc.String())
-		config.SetPath([]string{"crio", "runtime", "runtimes", runtimeClass}, runc)
+		config.SetPath([]string{"crio", "runtime", "runtimes", name}, runc)
 	}
 
-	config.SetPath([]string{"crio", "runtime", "runtimes", runtimeClass, "runtime_path"}, runtimePath)
+	config.SetPath([]string{"crio", "runtime", "runtimes", name, "runtime_path"}, path)
-	config.SetPath([]string{"crio", "runtime", "runtimes", runtimeClass, "runtime_type"}, "oci")
+	config.SetPath([]string{"crio", "runtime", "runtimes", name, "runtime_type"}, "oci")
 
 	if setAsDefault {
-		config.SetPath([]string{"crio", "runtime", "default_runtime"}, runtimeClass)
+		config.SetPath([]string{"crio", "runtime", "default_runtime"}, name)
 	}
 
+	*c = (Config)(config)
 	return nil
 }
 
-// RevertConfig reverts the cri-o config to remove the NVIDIA Container Runtime
+// DefaultRuntime returns the default runtime for the cri-o config
-func RevertConfig(config *toml.Tree, runtimeClass string) error {
+func (c Config) DefaultRuntime() string {
+	config := (toml.Tree)(c)
 	if runtime, ok := config.GetPath([]string{"crio", "runtime", "default_runtime"}).(string); ok {
-		if runtimeClass == runtime {
+		return runtime
+	}
+	return ""
+}
+
+// RemoveRuntime removes a runtime from the cri-o config
+func (c *Config) RemoveRuntime(name string) error {
+	if c == nil {
+		return nil
+	}
+
+	config := (toml.Tree)(*c)
+	if runtime, ok := config.GetPath([]string{"crio", "runtime", "default_runtime"}).(string); ok {
+		if runtime == name {
 			config.DeletePath([]string{"crio", "runtime", "default_runtime"})
 		}
 	}
 
-	runtimeClassPath := []string{"crio", "runtime", "runtimes", runtimeClass}
+	runtimeClassPath := []string{"crio", "runtime", "runtimes", name}
 	config.DeletePath(runtimeClassPath)
 	for i := 0; i < len(runtimeClassPath); i++ {
 		remainingPath := runtimeClassPath[:len(runtimeClassPath)-i]
@@ -87,39 +96,36 @@ func RevertConfig(config *toml.Tree, runtimeClass string) error {
 		}
 	}
 
+	*c = (Config)(config)
 	return nil
 }
 
-// FlushConfig flushes the updated/reverted config out to disk
+// Save writes the config to the specified path
-func FlushConfig(config string, cfg *toml.Tree) error {
+func (c Config) Save(path string) (int64, error) {
-	log.Infof("Flushing config")
+	config := (toml.Tree)(c)
-
+	output, err := config.ToTomlString()
-	output, err := cfg.ToTomlString()
 	if err != nil {
-		return fmt.Errorf("unable to convert to TOML: %v", err)
+		return 0, fmt.Errorf("unable to convert to TOML: %v", err)
 	}
 
-	switch len(output) {
+	if len(output) == 0 {
-	case 0:
+		err := os.Remove(path)
-		err := os.Remove(config)
 		if err != nil {
-			return fmt.Errorf("unable to remove empty file: %v", err)
+			return 0, 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 0, nil
 	}
 
-	log.Infof("Successfully flushed config")
+	f, err := os.Create(path)
+	if err != nil {
+		return 0, fmt.Errorf("unable to open '%v' for writing: %v", path, err)
+	}
+	defer f.Close()
 
-	return nil
+	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/crio/option.go

@@ -0,0 +1,73 @@
+/**
+# 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 crio
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/pelletier/go-toml"
+	log "github.com/sirupsen/logrus"
+)
+
+type builder struct {
+	path string
+}
+
+// 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
+	}
+}
+
+func (b *builder) build() (*Config, error) {
+	if b.path == "" {
+		empty := toml.Tree{}
+		return (*Config)(&empty), nil
+	}
+
+	return loadConfig(b.path)
+}
+
+// loadConfig loads the cri-o 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")
+	}
+
+	cfg, err := toml.LoadFile(configFile)
+	if err != nil {
+		return nil, err
+	}
+
+	log.Infof("Successfully loaded config")
+
+	return (*Config)(cfg), nil
+}

tools/container/crio/crio.go

@@ -24,8 +24,9 @@ import (
 	"path/filepath"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/crio"
-	"github.com/pelletier/go-toml"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator"
 	log "github.com/sirupsen/logrus"
 	cli "github.com/urfave/cli/v2"
 )
@@ -213,7 +214,9 @@ func setupHook(o *options) error {
 func setupConfig(o *options) error {
 	log.Infof("Updating config file")
 
-	cfg, err := crio.LoadConfig(o.config)
+	cfg, err := crio.New(
+		crio.WithPath(o.config),
+	)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
@@ -223,10 +226,14 @@ func setupConfig(o *options) error {
 		return fmt.Errorf("unable to update config: %v", err)
 	}
 
-	err = crio.FlushConfig(o.config, cfg)
+	log.Infof("Flushing cri-o 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 = RestartCrio(o)
 	if err != nil {
@@ -267,7 +274,9 @@ func cleanupHook(o *options) error {
 func cleanupConfig(o *options) error {
 	log.Infof("Reverting config file modifications")
 
-	cfg, err := crio.LoadConfig(o.config)
+	cfg, err := crio.New(
+		crio.WithPath(o.config),
+	)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
@@ -277,10 +286,14 @@ func cleanupConfig(o *options) error {
 		return fmt.Errorf("unable to update config: %v", err)
 	}
 
-	err = crio.FlushConfig(o.config, cfg)
+	log.Infof("Flushing cri-o 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 = RestartCrio(o)
 	if err != nil {
@@ -345,14 +358,36 @@ func generateOciHook(toolkitDir string) podmanHook {
 }
 
 // UpdateConfig updates the cri-o config to include the NVIDIA Container Runtime
-func UpdateConfig(config *toml.Tree, o *options) error {
+func UpdateConfig(cfg engine.Interface, o *options) error {
-	runtimePath := filepath.Join(o.runtimeDir, "nvidia-container-runtime")
+	runtimes := operator.GetRuntimes(
-	return crio.UpdateConfig(config, o.runtimeClass, runtimePath, o.setAsDefault)
+		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)
+		}
+	}
+
+	return nil
 }
 
 // RevertConfig reverts the cri-o config to remove the NVIDIA Container Runtime
-func RevertConfig(config *toml.Tree, o *options) error {
+func RevertConfig(cfg engine.Interface, o *options) error {
-	return crio.RevertConfig(config, o.runtimeClass)
+	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 revert config for runtime class '%v': %v", class, err)
+		}
+	}
+	return nil
 }
 
 // RestartCrio restarts crio depending on the value of restartModeFlag