From 23fa38daec6bf30f3180bda6b16aaef78b599fb2 Mon Sep 17 00:00:00 2001
From: Brad Davidson <brad.davidson@rancher.com>
Date: Wed, 4 Sep 2024 22:13:00 +0000
Subject: [PATCH] Only clean up containerd hosts dirs managed by k3s

Signed-off-by: Brad Davidson <brad.davidson@rancher.com>
(cherry picked from commit ce0de686f029d331a554549fa824aab9f5a7e467)
Signed-off-by: Brad Davidson <brad.davidson@rancher.com>
---
 pkg/agent/containerd/config.go   | 42 +++++++++++++++++++++++++++++---
 pkg/agent/templates/templates.go |  3 +++
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/pkg/agent/containerd/config.go b/pkg/agent/containerd/config.go
index f2d4545606c4..c97e0495a3c2 100644
--- a/pkg/agent/containerd/config.go
+++ b/pkg/agent/containerd/config.go
@@ -1,6 +1,7 @@
 package containerd
 
 import (
+	"bufio"
 	"fmt"
 	"net"
 	"net/url"
@@ -42,12 +43,14 @@ func writeContainerdConfig(cfg *config.Node, containerdConfig templates.Containe
 
 // writeContainerdHosts merges registry mirrors/configs, and renders and saves hosts.toml from the filled template
 func writeContainerdHosts(cfg *config.Node, containerdConfig templates.ContainerdConfig) error {
+	// Clean up previous configuration templates
+	if err := cleanContainerdHosts(cfg.Containerd.Registry); err != nil {
+		return err
+	}
+
 	mirrorAddr := net.JoinHostPort(spegel.DefaultRegistry.InternalAddress, spegel.DefaultRegistry.RegistryPort)
 	hosts := getHostConfigs(containerdConfig.PrivateRegistryConfig, containerdConfig.NoDefaultEndpoint, mirrorAddr)
 
-	// Clean up previous configuration templates
-	os.RemoveAll(cfg.Containerd.Registry)
-
 	// Write out new templates
 	for host, config := range hosts {
 		hostDir := filepath.Join(cfg.Containerd.Registry, host)
@@ -67,6 +70,39 @@ func writeContainerdHosts(cfg *config.Node, containerdConfig templates.Container
 	return nil
 }
 
+// cleanContainerdHosts removes any registry host config dirs containing a hosts.toml
+// file with a header that indicates it was created by k3s.  Directories not
+// containing this file, or containing a file without the header, are left alone.
+func cleanContainerdHosts(dir string) error {
+	ents, err := os.ReadDir(dir)
+	if err != nil && !os.IsNotExist(err) {
+		return err
+	}
+
+	for _, ent := range ents {
+		if !ent.IsDir() {
+			continue
+		}
+		hostsFile := filepath.Join(dir, ent.Name(), "hosts.toml")
+		file, err := os.Open(hostsFile)
+		if err != nil {
+			if os.IsNotExist(err) {
+				continue
+			}
+			return err
+		}
+		line, err := bufio.NewReader(file).ReadString('\n')
+		if err != nil {
+			continue
+		}
+		if line == templates.HostsTomlHeader {
+			hostsDir := filepath.Join(dir, ent.Name())
+			os.RemoveAll(hostsDir)
+		}
+	}
+	return nil
+}
+
 // getHostConfigs merges the registry mirrors/configs into HostConfig template structs
 func getHostConfigs(registry *registries.Registry, noDefaultEndpoint bool, mirrorAddr string) HostConfigs {
 	hosts := map[string]templates.HostConfig{}
diff --git a/pkg/agent/templates/templates.go b/pkg/agent/templates/templates.go
index c910fb61b914..b233066e376e 100644
--- a/pkg/agent/templates/templates.go
+++ b/pkg/agent/templates/templates.go
@@ -8,6 +8,7 @@ import (
 	"github.com/rancher/wharfie/pkg/registries"
 
 	"github.com/k3s-io/k3s/pkg/daemons/config"
+	"github.com/k3s-io/k3s/pkg/version"
 )
 
 type ContainerdRuntimeConfig struct {
@@ -40,6 +41,8 @@ type HostConfig struct {
 	Endpoints []RegistryEndpoint
 }
 
+var HostsTomlHeader = "# File generated by " + version.Program + ". DO NOT EDIT."
+
 const HostsTomlTemplate = `
 {{- /* */ -}}
 # File generated by {{ .Program }}. DO NOT EDIT.