From 2caa785e17a03d978f998ac8cc5e9f632c911b31 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 270f85e468853633674e2f1feed7428efde6aa55)
Signed-off-by: Brad Davidson <brad.davidson@rancher.com>
---
 pkg/agent/containerd/config.go   | 47 +++++++++++++++++++++++++++++++-
 pkg/agent/templates/templates.go |  3 ++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/pkg/agent/containerd/config.go b/pkg/agent/containerd/config.go
index f2d4545606c4..9b01ac06f88e 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"
@@ -46,7 +47,9 @@ func writeContainerdHosts(cfg *config.Node, containerdConfig templates.Container
 	hosts := getHostConfigs(containerdConfig.PrivateRegistryConfig, containerdConfig.NoDefaultEndpoint, mirrorAddr)
 
 	// Clean up previous configuration templates
-	os.RemoveAll(cfg.Containerd.Registry)
+	if err := cleanContainerdHosts(cfg.Containerd.Registry, hosts); err != nil {
+		return err
+	}
 
 	// Write out new templates
 	for host, config := range hosts {
@@ -67,6 +70,48 @@ 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, or directories where a hosts.toml
+// is about to be written.  Unmanaged directories not containing this file, or containing
+// a file without the header, are left alone.
+func cleanContainerdHosts(dir string, hosts HostConfigs) error {
+	// clean directories for any registries that we are about to generate a hosts.toml for
+	for host := range hosts {
+		hostsDir := filepath.Join(dir, host)
+		os.RemoveAll(hostsDir)
+	}
+
+	// clean directories that contain a hosts.toml with a header indicating it was  created by k3s
+	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.