From bdf8a355e1d34544566d32eb80702412f7cba30b Mon Sep 17 00:00:00 2001
From: galal-hussein <hussein.galal.ahmed.11@gmail.com>
Date: Fri, 19 Apr 2019 23:08:05 +0200
Subject: [PATCH] Add containerd config go template

---
 pkg/agent/config/config.go         |  1 +
 pkg/agent/containerd/containerd.go | 58 ++++++++++++++----------------
 pkg/agent/templates/templates.go   | 42 ++++++++++++++++++++++
 pkg/daemons/config/types.go        | 13 +++----
 4 files changed, 76 insertions(+), 38 deletions(-)
 create mode 100644 pkg/agent/templates/templates.go

diff --git a/pkg/agent/config/config.go b/pkg/agent/config/config.go
index a7d6d7a19f..545f61b6d5 100644
--- a/pkg/agent/config/config.go
+++ b/pkg/agent/config/config.go
@@ -240,6 +240,7 @@ func get(envInfo *cmds.Agent) (*config.Node, error) {
 	}
 	nodeConfig.Containerd.State = "/run/k3s/containerd"
 	nodeConfig.Containerd.Address = filepath.Join(nodeConfig.Containerd.State, "containerd.sock")
+	nodeConfig.Containerd.Template = filepath.Join(envInfo.DataDir, "etc/containerd/config.toml.tmpl")
 	nodeConfig.ServerAddress = serverURLParsed.Host
 	nodeConfig.Certificate = nodeCert
 	if !nodeConfig.NoFlannel {
diff --git a/pkg/agent/containerd/containerd.go b/pkg/agent/containerd/containerd.go
index c38631dc29..d36b13019e 100644
--- a/pkg/agent/containerd/containerd.go
+++ b/pkg/agent/containerd/containerd.go
@@ -8,7 +8,6 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
-	"strings"
 	"syscall"
 	"time"
 
@@ -16,6 +15,7 @@ import (
 	"github.com/containerd/containerd/namespaces"
 	"github.com/natefinch/lumberjack"
 	"github.com/opencontainers/runc/libcontainer/system"
+	"github.com/rancher/k3s/pkg/agent/templates"
 	util2 "github.com/rancher/k3s/pkg/agent/util"
 	"github.com/rancher/k3s/pkg/daemons/config"
 	"github.com/sirupsen/logrus"
@@ -26,23 +26,6 @@ import (
 
 const (
 	maxMsgSize = 1024 * 1024 * 16
-	configToml = `
-[plugins.opt]
-  path = "%OPT%"
-[plugins.cri]
-  stream_server_address = "%NODE%"
-  stream_server_port = "10010"
-`
-	configUserNSToml = `
-  disable_cgroup = true
-  disable_apparmor = true
-  restrict_oom_score_adj = true
-`
-	configCNIToml = `
-  [plugins.cri.cni]
-    bin_dir = "%CNIBIN%"
-    conf_dir = "%CNICFG%"
-`
 )
 
 func Run(ctx context.Context, cfg *config.Node) error {
@@ -54,20 +37,7 @@ func Run(ctx context.Context, cfg *config.Node) error {
 		"--root", cfg.Containerd.Root,
 	}
 
-	template := configToml
-	if system.RunningInUserNS() {
-		template += configUserNSToml
-	}
-	if !cfg.NoFlannel {
-		template += configCNIToml
-	}
-
-	template = strings.Replace(template, "%OPT%", cfg.Containerd.Opt, -1)
-	template = strings.Replace(template, "%CNIBIN%", cfg.AgentConfig.CNIBinDir, -1)
-	template = strings.Replace(template, "%CNICFG%", cfg.AgentConfig.CNIConfDir, -1)
-	template = strings.Replace(template, "%NODE%", cfg.AgentConfig.NodeName, -1)
-
-	if err := util2.WriteFile(cfg.Containerd.Config, template); err != nil {
+	if err := setupContainerdConfig(ctx, cfg); err != nil {
 		return err
 	}
 
@@ -186,3 +156,27 @@ func preloadImages(cfg *config.Node) error {
 	}
 	return nil
 }
+
+func setupContainerdConfig(ctx context.Context, cfg *config.Node) error {
+	var containerdTemplate string
+	containerdConfig := templates.ContainerdConfig{
+		NodeConfig:        cfg,
+		IsRunningInUserNS: system.RunningInUserNS(),
+	}
+
+	containerdTemplateBytes, err := ioutil.ReadFile(cfg.Containerd.Template)
+	if err == nil {
+		logrus.Infof("Using containerd template at %s", cfg.Containerd.Template)
+		containerdTemplate = string(containerdTemplateBytes)
+	} else if os.IsNotExist(err) {
+		containerdTemplate = templates.ContainerdConfigTemplate
+	} else {
+		return err
+	}
+	parsedTemplate, err := templates.ParseTemplateFromConfig(containerdTemplate, containerdConfig)
+	if err != nil {
+		return err
+	}
+
+	return util2.WriteFile(cfg.Containerd.Config, parsedTemplate)
+}
diff --git a/pkg/agent/templates/templates.go b/pkg/agent/templates/templates.go
new file mode 100644
index 0000000000..5094c1d5ee
--- /dev/null
+++ b/pkg/agent/templates/templates.go
@@ -0,0 +1,42 @@
+package templates
+
+import (
+	"bytes"
+	"text/template"
+
+	"github.com/rancher/k3s/pkg/daemons/config"
+)
+
+type ContainerdConfig struct {
+	NodeConfig        *config.Node
+	IsRunningInUserNS bool
+}
+
+const ContainerdConfigTemplate = `
+[plugins.opt]
+path = "{{ .NodeConfig.Containerd.Opt }}"
+
+[plugins.cri]
+stream_server_address = "{{ .NodeConfig.AgentConfig.NodeName }}"
+stream_server_port = "10010"
+{{ if .IsRunningInUserNS }}
+disable_cgroup = true
+disable_apparmor = true
+restrict_oom_score_adj = true
+{{ end }}
+
+{{ if not .NodeConfig.NoFlannel }}
+  [plugins.cri.cni]
+    bin_dir = "{{ .NodeConfig.AgentConfig.CNIBinDir }}"
+    conf_dir = "{{ .NodeConfig.AgentConfig.CNIConfDir }}"
+{{ end }}
+`
+
+func ParseTemplateFromConfig(templateBuffer string, config interface{}) (string, error) {
+	out := new(bytes.Buffer)
+	t := template.Must(template.New("compiled_template").Parse(templateBuffer))
+	if err := t.Execute(out, config); err != nil {
+		return "", err
+	}
+	return out.String(), nil
+}
diff --git a/pkg/daemons/config/types.go b/pkg/daemons/config/types.go
index 577f26eb3b..6fb555f936 100644
--- a/pkg/daemons/config/types.go
+++ b/pkg/daemons/config/types.go
@@ -26,12 +26,13 @@ type Node struct {
 }
 
 type Containerd struct {
-	Address string
-	Log     string
-	Root    string
-	State   string
-	Config  string
-	Opt     string
+	Address  string
+	Log      string
+	Root     string
+	State    string
+	Config   string
+	Opt      string
+	Template string
 }
 
 type Agent struct {