Swap temporary IPSets during ipset restore (#1068)

* ipset restore: use temporary sets and swap them into the real ones

* move const

* switch to shared tmp ipsets

* preemptively flush tmp set in case it already existed

pkg/utils/ipset.go

@@ -2,9 +2,12 @@ package utils
 
 import (
 	"bytes"
+	"crypto/sha1"
+	"encoding/base32"
 	"errors"
 	"fmt"
 	"os/exec"
+	"sort"
 	"strings"
 )
 
@@ -77,6 +80,9 @@ const (
 	OptionNoMatch = "nomatch"
 	// OptionForceAdd All hash set types support the optional forceadd parameter when creating a set. When sets created with this option become full the next addition to the set may succeed and evict a random entry from the set.
 	OptionForceAdd = "forceadd"
+
+	// tmpIPSetPrefix Is the prefix added to temporary ipset names used in the atomic swap operations during ipset restore. You should never see these on your system because they only exist during the restore.
+	tmpIPSetPrefix = "TMP-"
 )
 
 // IPSet represent ipset sets managed by.
@@ -400,15 +406,53 @@ func parseIPSetSave(ipset *IPSet, result string) map[string]*Set {
 // create KUBE-DST-3YNVZWWGX3UQQ4VQ hash:ip family inet hashsize 1024 maxelem 65536 timeout 0
 // add KUBE-DST-3YNVZWWGX3UQQ4VQ 100.96.1.6 timeout 0
 func buildIPSetRestore(ipset *IPSet) string {
-	ipSetRestore := ""
-	for _, set := range ipset.Sets {
-		ipSetRestore += fmt.Sprintf("create %s %s\n", set.Name, strings.Join(set.Options[:], " "))
-		ipSetRestore += fmt.Sprintf("flush %s\n", set.Name)
+	setNames := make([]string, 0, len(ipset.Sets))
+	for setName := range ipset.Sets {
+		// we need setNames in some consistent order so that we can unit-test this method has a predictable output:
+		setNames = append(setNames, setName)
+	}
+
+	sort.Strings(setNames)
+
+	tmpSets := map[string]string{}
+	ipSetRestore := &strings.Builder{}
+	for _, setName := range setNames {
+		set := ipset.Sets[setName]
+		setOptions := strings.Join(set.Options, " ")
+
+		tmpSetName := tmpSets[setOptions]
+		if tmpSetName == "" {
+			// create a temporary set per unique set-options:
+			hash := sha1.Sum([]byte("tmp:" + setOptions))
+			tmpSetName = tmpIPSetPrefix + base32.StdEncoding.EncodeToString(hash[:10])
+			ipSetRestore.WriteString(fmt.Sprintf("create %s %s\n", tmpSetName, setOptions))
+			// just in case we are starting up after a crash, we should flush the TMP ipset to be safe if it
+			// already existed, so we do not pollute other ipsets:
+			ipSetRestore.WriteString(fmt.Sprintf("flush %s\n", tmpSetName))
+			tmpSets[setOptions] = tmpSetName
+		}
+
 		for _, entry := range set.Entries {
-			ipSetRestore += fmt.Sprintf("add %s %s\n", set.Name, strings.Join(entry.Options[:], " "))
+			// add entries to the tmp set:
+			ipSetRestore.WriteString(fmt.Sprintf("add %s %s\n", tmpSetName, strings.Join(entry.Options, " ")))
 		}
+
+		// now create the actual IPSet (this is a noop if it already exists, because we run with -exists):
+		ipSetRestore.WriteString(fmt.Sprintf("create %s %s\n", set.Name, setOptions))
+
+		// now that both exist, we can swap them:
+		ipSetRestore.WriteString(fmt.Sprintf("swap %s %s\n", tmpSetName, set.Name))
+
+		// empty the tmp set (which is actually the old one now):
+		ipSetRestore.WriteString(fmt.Sprintf("flush %s\n", tmpSetName))
+	}
+
+	for _, tmpSetName := range tmpSets {
+		// finally, destroy the tmp sets.
+		ipSetRestore.WriteString(fmt.Sprintf("destroy %s\n", tmpSetName))
 	}
-	return ipSetRestore
+
+	return ipSetRestore.String()
 }
 
 // Save the given set, or all sets if none is given to stdout in a format that

pkg/utils/ipset_test.go

@@ -0,0 +1,74 @@
+package utils
+
+import "testing"
+
+func Test_buildIPSetRestore(t *testing.T) {
+	type args struct {
+		ipset *IPSet
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			name: "simple-restore",
+			args: args{
+				ipset: &IPSet{Sets: map[string]*Set{
+					"foo": {
+						Name:    "foo",
+						Options: []string{"hash:ip", "yolo", "things", "12345"},
+						Entries: []*Entry{
+							{Options: []string{"1.2.3.4"}},
+						},
+					},
+					"google-dns-servers": {
+						Name:    "google-dns-servers",
+						Options: []string{"hash:ip", "lol"},
+						Entries: []*Entry{
+							{Options: []string{"4.4.4.4"}},
+							{Options: []string{"8.8.8.8"}},
+						},
+					},
+					// this one and the one above share the same exact options -- and therefore will reuse the same
+					// tmp ipset:
+					"more-ip-addresses": {
+						Name:    "google-dns-servers",
+						Options: []string{"hash:ip", "lol"},
+						Entries: []*Entry{
+							{Options: []string{"5.5.5.5"}},
+							{Options: []string{"6.6.6.6"}},
+						},
+					},
+				}},
+			},
+			want: "create TMP-7NOTZDOMLXBX6DAJ hash:ip yolo things 12345\n" +
+				"flush TMP-7NOTZDOMLXBX6DAJ\n" +
+				"add TMP-7NOTZDOMLXBX6DAJ 1.2.3.4\n" +
+				"create foo hash:ip yolo things 12345\n" +
+				"swap TMP-7NOTZDOMLXBX6DAJ foo\n" +
+				"flush TMP-7NOTZDOMLXBX6DAJ\n" +
+				"create TMP-XD7BSSQZELS7TP35 hash:ip lol\n" +
+				"flush TMP-XD7BSSQZELS7TP35\n" +
+				"add TMP-XD7BSSQZELS7TP35 4.4.4.4\n" +
+				"add TMP-XD7BSSQZELS7TP35 8.8.8.8\n" +
+				"create google-dns-servers hash:ip lol\n" +
+				"swap TMP-XD7BSSQZELS7TP35 google-dns-servers\n" +
+				"flush TMP-XD7BSSQZELS7TP35\n" +
+				"add TMP-XD7BSSQZELS7TP35 5.5.5.5\n" +
+				"add TMP-XD7BSSQZELS7TP35 6.6.6.6\n" +
+				"create google-dns-servers hash:ip lol\n" +
+				"swap TMP-XD7BSSQZELS7TP35 google-dns-servers\n" +
+				"flush TMP-XD7BSSQZELS7TP35\n" +
+				"destroy TMP-7NOTZDOMLXBX6DAJ\n" +
+				"destroy TMP-XD7BSSQZELS7TP35\n",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := buildIPSetRestore(tt.args.ipset); got != tt.want {
+				t.Errorf("buildIPSetRestore() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}