Add rewrite support to hosts.toml loader

Signed-off-by: Brad Davidson <brad.davidson@rancher.com>

core/remotes/docker/config/hosts.go

@@ -54,6 +54,7 @@ type hostConfig struct {
 
 	header http.Header
 
+	rewrites map[string]string
 	// TODO: Add credential configuration (domain alias, username)
 }
 
@@ -263,6 +264,7 @@ func ConfigureHosts(ctx context.Context, options HostOptions) docker.RegistryHos
 			rhosts[i].Path = host.path
 			rhosts[i].Capabilities = host.capabilities
 			rhosts[i].Header = host.header
+			rhosts[i].Rewrites = host.rewrites
 		}
 
 		return rhosts, nil
@@ -352,6 +354,10 @@ type hostFileConfig struct {
 	// API root endpoint.
 	OverridePath bool `toml:"override_path"`
 
+	// Rewrites contains a map of regex/replacement values, used to modify
+	// the image name when pulling.
+	Rewrites map[string]string `toml:"rewrite"`
+
 	// TODO: Credentials: helper? name? username? alternate domain? token?
 }
 
@@ -389,13 +395,22 @@ 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)
+	tree, err := getTopLevelKeyValues(b)
 	if err != nil {
 		return nil, err
 	}
-	hosts = append(hosts, parsed)
 
+	// If the server key is set at the root of the tree, or no other hosts are configured,
+	// parse the root host config and append it as the last element.
+	// This allows not using upstream by including hosts and excluding the server key as documented at
+	// https://github.com/containerd/containerd/blob/release/1.7/docs/hosts.md#setup-a-local-mirror-for-docker
+	if _, ok := tree["server"]; ok || len(hosts) == 0 {
+		parsed, err := parseHostConfig(c.Server, baseDir, c.hostFileConfig)
+		if err != nil {
+			return nil, err
+		}
+		hosts = append(hosts, parsed)
+	}
 	return hosts, nil
 }
 
@@ -427,6 +442,7 @@ func parseHostConfig(server string, baseDir string, config hostFileConfig) (host
 	}
 
 	result.skipVerify = config.SkipVerify
+	result.rewrites = config.Rewrites
 
 	if len(config.Capabilities) > 0 {
 		for _, c := range config.Capabilities {
@@ -515,6 +531,32 @@ func parseHostConfig(server string, baseDir string, config hostFileConfig) (host
 	return result, nil
 }
 
+func getTopLevelKeyValues(b []byte) (map[string]string, error) {
+	keyVals := map[string]string{}
+
+	// Use toml unstable package for directly parsing toml
+	// See https://github.com/pelletier/go-toml/discussions/801#discussioncomment-7083586
+	p := tomlu.Parser{}
+	p.Reset(b)
+
+	// iterate over top level expressions until we find something that's not a simple KeyValue
+	for p.NextExpression() {
+		e := p.Expression()
+
+		if e.Kind != tomlu.KeyValue {
+			break
+		}
+
+		k := e.Key()
+		v := e.Value()
+		if keyNode := k.Node(); keyNode != nil && v != nil {
+			keyVals[string(keyNode.Data)] = string(v.Data)
+		}
+	}
+
+	return keyVals, p.Error()
+}
+
 // getSortedHosts returns the list of hosts in the order are they defined in the file.
 func getSortedHosts(b []byte) ([]string, error) {
 	var hostsInOrder []string

internal/cri/config/config.go

@@ -197,7 +197,7 @@ type Mirror struct {
 	// The scheme, host and path from the endpoint URL will be used.
 	Endpoints []string `toml:"endpoint" json:"endpoint"`
 
-	// Rewrites are repository rewrite rules for a namespace. When fetching image resources
+	// Rewrites is a map of repository rewrite rules for a namespace. When fetching image resources
 	// from an endpoint and a key matches the repository via regular expression matching
 	// it will be replaced with the corresponding value from the map in the resource request.
 	//

internal/cri/server/images/image_pull.go

@@ -437,14 +437,14 @@ func (c *CRIImageService) registryHosts(ctx context.Context, credentials func(ho
 	return func(host string) ([]docker.RegistryHost, error) {
 		var registries []docker.RegistryHost
 
-		endpoints, err := c.registryEndpoints(host)
-		if err != nil {
-			return nil, fmt.Errorf("get registry endpoints: %w", err)
-		}
 		rewrites, err := c.registryRewrites(host)
 		if err != nil {
 			return nil, fmt.Errorf("get registry rewrites: %w", err)
 		}
+		endpoints, err := c.registryEndpoints(host)
+		if err != nil {
+			return nil, fmt.Errorf("get registry endpoints: %w", err)
+		}
 		for _, e := range endpoints {
 			u, err := url.Parse(e)
 			if err != nil {
@@ -570,17 +570,13 @@ func (c *CRIImageService) registryEndpoints(host string) ([]string, error) {
 }
 
 func (c *CRIImageService) registryRewrites(host string) (map[string]string, error) {
-	var rewrites map[string]string
-	_, ok := c.config.Registry.Mirrors[host]
-	if ok {
-		rewrites = c.config.Registry.Mirrors[host].Rewrites
-	} else {
-		rewrites = c.config.Registry.Mirrors["*"].Rewrites
+	hosts := []string{host, "*"}
+	for _, host := range hosts {
+		if host, ok := c.config.Registry.Mirrors[host]; ok {
+			return host.Rewrites, nil
+		}
 	}
-	if rewrites == nil {
-		rewrites = map[string]string{}
-	}
-	return rewrites, nil
+	return nil, nil
 }
 
 // newTransport returns a new HTTP transport used to pull image.