Merge pull request #1054 from thaJeztah/fix-host-rm-being-too-greedy

Fix service update --host-rm not being granular enough

cli/command/service/update.go

@@ -963,45 +963,99 @@ func updateReplicas(flags *pflag.FlagSet, serviceMode *swarm.ServiceMode) error
 	return nil
 }
 
+type hostMapping struct {
+	IPAddr string
+	Host   string
+}
+
 // updateHosts performs a diff between existing host entries, entries to be
 // removed, and entries to be added. Host entries preserve the order in which they
 // were added, as the specification mentions that in case multiple entries for a
 // host exist, the first entry should be used (by default).
+//
+// Note that, even though unsupported by the the CLI, the service specs format
+// allow entries with both a _canonical_ hostname, and one or more aliases
+// in an entry (IP-address canonical_hostname [alias ...])
+//
+// Entries can be removed by either a specific `<host-name>:<ip-address>` mapping,
+// or by `<host>` alone:
+//
+// - If both IP-address and host-name is provided, the hostname is removed only
+//   from entries that match the given IP-address.
+// - If only a host-name is provided, the hostname is removed from any entry it
+//   is part of (either as canonical host-name, or as alias).
+// - If, after removing the host-name from an entry, no host-names remain in
+//   the entry, the entry itself is removed.
+//
+// For example, the list of host-entries before processing could look like this:
+//
+//    hosts = &[]string{
+//        "127.0.0.2 host3 host1 host2 host4",
+//        "127.0.0.1 host1 host4",
+//        "127.0.0.3 host1",
+//        "127.0.0.1 host1",
+//    }
+//
+// Removing `host1` removes every occurrence:
+//
+//    hosts = &[]string{
+//        "127.0.0.2 host3 host2 host4",
+//        "127.0.0.1 host4",
+//    }
+//
+// Removing `host1:127.0.0.1` on the other hand, only remove the host if the
+// IP-address matches:
+//
+//    hosts = &[]string{
+//        "127.0.0.2 host3 host1 host2 host4",
+//        "127.0.0.1 host4",
+//        "127.0.0.3 host1",
+//    }
 func updateHosts(flags *pflag.FlagSet, hosts *[]string) error {
-	// Combine existing Hosts (in swarmkit format) with the host to add (convert to swarmkit format)
-	if flags.Changed(flagHostAdd) {
-		values := convertExtraHostsToSwarmHosts(flags.Lookup(flagHostAdd).Value.(*opts.ListOpts).GetAll())
-		*hosts = append(*hosts, values...)
-	}
-	// Remove duplicate
-	*hosts = removeDuplicates(*hosts)
-
-	keysToRemove := make(map[string]struct{})
+	var toRemove []hostMapping
 	if flags.Changed(flagHostRemove) {
-		var empty struct{}
 		extraHostsToRemove := flags.Lookup(flagHostRemove).Value.(*opts.ListOpts).GetAll()
 		for _, entry := range extraHostsToRemove {
-			key := strings.SplitN(entry, ":", 2)[0]
-			keysToRemove[key] = empty
+			v := strings.SplitN(entry, ":", 2)
+			if len(v) > 1 {
+				toRemove = append(toRemove, hostMapping{IPAddr: v[1], Host: v[0]})
+			} else {
+				toRemove = append(toRemove, hostMapping{Host: v[0]})
+			}
 		}
 	}
 
-	newHosts := []string{}
+	var newHosts []string
 	for _, entry := range *hosts {
-		// Since this is in swarmkit format, we need to find the key, which is canonical_hostname of:
+		// Since this is in SwarmKit format, we need to find the key, which is canonical_hostname of:
 		// IP_address canonical_hostname [aliases...]
 		parts := strings.Fields(entry)
-		if len(parts) > 1 {
-			key := parts[1]
-			if _, exists := keysToRemove[key]; !exists {
-				newHosts = append(newHosts, entry)
+		if len(parts) == 0 {
+			continue
+		}
+		ip := parts[0]
+		hostNames := parts[1:]
+		for _, rm := range toRemove {
+			if rm.IPAddr != "" && rm.IPAddr != ip {
+				continue
 			}
-		} else {
-			newHosts = append(newHosts, entry)
+			for i, h := range hostNames {
+				if h == rm.Host {
+					hostNames = append(hostNames[:i], hostNames[i+1:]...)
+				}
+			}
+		}
+		if len(hostNames) > 0 {
+			newHosts = append(newHosts, fmt.Sprintf("%s %s", ip, strings.Join(hostNames, " ")))
 		}
 	}
 
-	*hosts = newHosts
+	// Append new hosts (in SwarmKit format)
+	if flags.Changed(flagHostAdd) {
+		values := convertExtraHostsToSwarmHosts(flags.Lookup(flagHostAdd).Value.(*opts.ListOpts).GetAll())
+		newHosts = append(newHosts, values...)
+	}
+	*hosts = removeDuplicates(newHosts)
 	return nil
 }
 

cli/command/service/update_test.go

@@ -366,7 +366,7 @@ func TestUpdateHosts(t *testing.T) {
 	assert.ErrorContains(t, flags.Set("host-add", "$example.com$"), `bad format for add-host: "$example.com$"`)
 
 	hosts := []string{"1.2.3.4 example.com", "4.3.2.1 example.org", "2001:db8:abc8::1 example.net"}
-	expected := []string{"1.2.3.4 example.com", "4.3.2.1 example.org", "2001:db8:abc8::1 ipv6.net"}
+	expected := []string{"1.2.3.4 example.com", "4.3.2.1 example.org", "2.2.2.2 example.net", "2001:db8:abc8::1 ipv6.net"}
 
 	err := updateHosts(flags, &hosts)
 	assert.NilError(t, err)
@@ -385,6 +385,60 @@ func TestUpdateHostsPreservesOrder(t *testing.T) {
 	assert.Check(t, is.DeepEqual([]string{"127.0.0.2 foobar", "127.0.0.1 foobar", "127.0.0.3 foobar"}, hosts))
 }
 
+func TestUpdateHostsReplaceEntry(t *testing.T) {
+	flags := newUpdateCommand(nil).Flags()
+	flags.Set("host-add", "foobar:127.0.0.4")
+	flags.Set("host-rm", "foobar:127.0.0.2")
+
+	hosts := []string{"127.0.0.2 foobar", "127.0.0.1 foobar", "127.0.0.3 foobar"}
+
+	err := updateHosts(flags, &hosts)
+	assert.NilError(t, err)
+	assert.Check(t, is.DeepEqual([]string{"127.0.0.1 foobar", "127.0.0.3 foobar", "127.0.0.4 foobar"}, hosts))
+}
+
+func TestUpdateHostsRemoveHost(t *testing.T) {
+	flags := newUpdateCommand(nil).Flags()
+	flags.Set("host-rm", "host1")
+
+	hosts := []string{"127.0.0.2 host3 host1 host2 host4", "127.0.0.1 host1 host4", "127.0.0.3 host1"}
+
+	err := updateHosts(flags, &hosts)
+	assert.NilError(t, err)
+
+	// Removing host `host1` should remove the entry from each line it appears in.
+	// If there are no other hosts in the entry, the entry itself should be removed.
+	assert.Check(t, is.DeepEqual([]string{"127.0.0.2 host3 host2 host4", "127.0.0.1 host4"}, hosts))
+}
+
+func TestUpdateHostsRemoveHostIP(t *testing.T) {
+	flags := newUpdateCommand(nil).Flags()
+	flags.Set("host-rm", "host1:127.0.0.1")
+
+	hosts := []string{"127.0.0.2 host3 host1 host2 host4", "127.0.0.1 host1 host4", "127.0.0.3 host1", "127.0.0.1 host1"}
+
+	err := updateHosts(flags, &hosts)
+	assert.NilError(t, err)
+
+	// Removing host `host1` should remove the entry from each line it appears in,
+	// but only if the IP-address matches. If there are no other hosts in the entry,
+	// the entry itself should be removed.
+	assert.Check(t, is.DeepEqual([]string{"127.0.0.2 host3 host1 host2 host4", "127.0.0.1 host4", "127.0.0.3 host1"}, hosts))
+}
+
+func TestUpdateHostsRemoveAll(t *testing.T) {
+	flags := newUpdateCommand(nil).Flags()
+	flags.Set("host-add", "host-three:127.0.0.4")
+	flags.Set("host-add", "host-one:127.0.0.5")
+	flags.Set("host-rm", "host-one")
+
+	hosts := []string{"127.0.0.1 host-one", "127.0.0.2 host-two", "127.0.0.3 host-one"}
+
+	err := updateHosts(flags, &hosts)
+	assert.NilError(t, err)
+	assert.Check(t, is.DeepEqual([]string{"127.0.0.2 host-two", "127.0.0.4 host-three", "127.0.0.5 host-one"}, hosts))
+}
+
 func TestUpdatePortsRmWithProtocol(t *testing.T) {
 	flags := newUpdateCommand(nil).Flags()
 	flags.Set("publish-add", "8081:81")