From 5ada2f74a7a6f9d4150789239e5192e23963963d Mon Sep 17 00:00:00 2001
From: Maksym Pavlenko <pavlenko.maksym@gmail.com>
Date: Thu, 1 Apr 2021 09:29:16 -0700
Subject: [PATCH] Keep host order as defined in TOML file

Signed-off-by: Maksym Pavlenko <pavlenko.maksym@gmail.com>
---
 pkg/cri/config/config.go            |  2 +-
 remotes/docker/config/hosts.go      | 49 ++++++++++++++++++++++-------
 remotes/docker/config/hosts_test.go | 12 +++++--
 3 files changed, 48 insertions(+), 15 deletions(-)

diff --git a/pkg/cri/config/config.go b/pkg/cri/config/config.go
index 4280d0a34..7cd65c784 100644
--- a/pkg/cri/config/config.go
+++ b/pkg/cri/config/config.go
@@ -48,7 +48,7 @@ type Runtime struct {
 	// This only works for runtime type "io.containerd.runtime.v1.linux".
 	Root string `toml:"runtime_root" json:"runtimeRoot"`
 	// Options are config options for the runtime. If options is loaded
-	// from toml config, it will be toml.Primitive.
+	// from toml config, it will be toml.Tree.
 	Options *toml.Tree `toml:"options" json:"options"`
 	// PrivilegedWithoutHostDevices overloads the default behaviour for adding host devices to the
 	// runtime spec when the container is privileged. Defaults to false.
diff --git a/remotes/docker/config/hosts.go b/remotes/docker/config/hosts.go
index b01518471..b24ba33f5 100644
--- a/remotes/docker/config/hosts.go
+++ b/remotes/docker/config/hosts.go
@@ -27,15 +27,15 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"sort"
 	"strings"
 	"time"
 
-	"github.com/pelletier/go-toml"
-	"github.com/pkg/errors"
-
 	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/log"
 	"github.com/containerd/containerd/remotes/docker"
+	"github.com/pelletier/go-toml"
+	"github.com/pkg/errors"
 )
 
 // UpdateClientFunc is a function that lets you to amend http Client behavior used by registry clients.
@@ -317,6 +317,11 @@ func parseHostsFile(baseDir string, b []byte) ([]hostConfig, error) {
 		HostConfigs map[string]hostFileConfig `toml:"host"`
 	}{}
 
+	orderedHosts, err := getSortedHosts(tree)
+	if err != nil {
+		return nil, err
+	}
+
 	var (
 		hosts []hostConfig
 	)
@@ -325,15 +330,10 @@ func parseHostsFile(baseDir string, b []byte) ([]hostConfig, error) {
 		return nil, err
 	}
 
-	// Parse root host config
-	parsed, err := parseHostConfig(c.Server, baseDir, c.HostFileConfig)
-	if err != nil {
-		return nil, err
-	}
-	hosts = append(hosts, parsed)
-
 	// Parse hosts array
-	for host, config := range c.HostConfigs {
+	for _, host := range orderedHosts {
+		config := c.HostConfigs[host]
+
 		parsed, err := parseHostConfig(host, baseDir, config)
 		if err != nil {
 			return nil, err
@@ -341,6 +341,13 @@ func parseHostsFile(baseDir string, b []byte) ([]hostConfig, error) {
 		hosts = append(hosts, parsed)
 	}
 
+	// Parse root host config and append it as the last element
+	parsed, err := parseHostConfig(c.Server, baseDir, c.HostFileConfig)
+	if err != nil {
+		return nil, err
+	}
+	hosts = append(hosts, parsed)
+
 	return hosts, nil
 }
 
@@ -464,6 +471,26 @@ func parseHostConfig(server string, baseDir string, config hostFileConfig) (host
 	return result, nil
 }
 
+// getSortedHosts returns the list of hosts as they defined in the file.
+func getSortedHosts(root *toml.Tree) ([]string, error) {
+	iter, ok := root.Get("host").(*toml.Tree)
+	if !ok {
+		return nil, errors.Errorf("invalid `host` tree")
+	}
+
+	list := append([]string{}, iter.Keys()...)
+
+	// go-toml stores TOML sections in the map object, so no order guaranteed.
+	// We retrieve line number for each key and sort the keys by position.
+	sort.Slice(list, func(i, j int) bool {
+		h1 := iter.GetPath([]string{list[i]}).(*toml.Tree)
+		h2 := iter.GetPath([]string{list[j]}).(*toml.Tree)
+		return h1.Position().Line < h2.Position().Line
+	})
+
+	return list, nil
+}
+
 // makeStringSlice is a helper func to convert from []interface{} to []string.
 // Additionally an optional cb func may be passed to perform string mapping.
 func makeStringSlice(slice []interface{}, cb func(string) string) ([]string, error) {
diff --git a/remotes/docker/config/hosts_test.go b/remotes/docker/config/hosts_test.go
index df3f3328c..18dc1c659 100644
--- a/remotes/docker/config/hosts_test.go
+++ b/remotes/docker/config/hosts_test.go
@@ -26,8 +26,6 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-
 	"github.com/containerd/containerd/log/logtest"
 	"github.com/containerd/containerd/remotes/docker"
 )
@@ -181,7 +179,15 @@ ca = "/etc/path/default"
 		}
 	}()
 
-	assert.ElementsMatch(t, expected, hosts)
+	if len(hosts) != len(expected) {
+		t.Fatalf("Unexpected number of hosts %d, expected %d", len(hosts), len(expected))
+	}
+
+	for i := range hosts {
+		if !compareHostConfig(hosts[i], expected[i]) {
+			t.Fatalf("Mismatch at host %d", i)
+		}
+	}
 }
 
 func TestLoadCertFiles(t *testing.T) {