Cleanup ipset create function + restore (#177)

* Cleanup ipset create function + restore. Fix #176

Switch from array to map to simplify set use

* Add sets initialization on create

Author: ryarnyah

app/controllers/network_routes_controller.go

@@ -676,9 +676,9 @@ func deletePodSubnetIpSet() error {
 	if err != nil {
 		return err
 	}
-	ipset.Sets = append(ipset.Sets, &utils.Set{
+	ipset.Sets[podSubnetIpSetName] = &utils.Set{
 		Name: podSubnetIpSetName,
-	})
+	}
 	err = ipset.Destroy()
 	if err != nil {
 		return errors.New("Failure deleting Pod egress ipset: " + err.Error())

utils/ipset.go

@@ -84,7 +84,7 @@ type IPSet struct {
 	// ipset bianry path.
 	ipSetPath *string
 	// Sets maintainted by ipset.
-	Sets []*Set
+	Sets map[string]*Set
 }
 
 // Set reprensent a ipset set entry.
@@ -161,6 +161,7 @@ func NewIPSet() (*IPSet, error) {
 	}
 	ipSet := &IPSet{
 		ipSetPath: ipSetPath,
+		Sets:      make(map[string]*Set),
 	}
 	return ipSet, nil
 }
@@ -170,8 +171,9 @@ func (ipSet *IPSet) Create(setName string, createOptions ...string) (*Set, error
 	set := &Set{
 		Name:    setName,
 		Options: createOptions,
+		IPSet:   ipSet,
 	}
-	ipSet.Sets = append(ipSet.Sets, set)
+	ipSet.Sets[setName] = set
 	_, err := ipSet.run(append([]string{"create", "-exist", set.Name}, createOptions...)...)
 	if err != nil {
 		return nil, err
@@ -234,20 +236,20 @@ func (set *IPSet) Destroy() error {
 // ex:
 // 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 parseIPSetSave(ipSet *IPSet, result string) []*Set {
-	sets := make([]*Set, 0)
+func parseIPSetSave(ipSet *IPSet, result string) map[string]*Set {
+	sets := make(map[string]*Set)
 	// Save is always in order
 	lines := strings.Split(result, "\n")
 	for _, line := range lines {
 		content := strings.Split(line, " ")
 		if content[0] == "create" {
-			sets = append(sets, &Set{
+			sets[content[1]] = &Set{
 				IPSet:   ipSet,
 				Name:    content[1],
 				Options: content[2:],
-			})
+			}
 		} else if content[0] == "add" {
-			set := sets[len(sets)-1]
+			set := sets[content[1]]
 			set.Entries = append(set.Entries, &Entry{
 				Set:     set,
 				Options: content[2:],
@@ -265,15 +267,16 @@ func parseIPSetSave(ipSet *IPSet, result string) []*Set {
 func buildIPSetRestore(ipSet *IPSet) string {
 	ipSetRestore := ""
 	for _, set := range ipSet.Sets {
-		ipSetRestore += fmt.Sprintf("create %s %v\n", set.Name, set.Options)
+		ipSetRestore += fmt.Sprintf("create %s %s\n", set.Name, strings.Join(set.Options[:], " "))
 		for _, entry := range set.Entries {
-			ipSetRestore += fmt.Sprintf("add %s %v\n", set.Name, entry.Options)
+			ipSetRestore += fmt.Sprintf("add %s %s\n", set.Name, strings.Join(entry.Options[:], " "))
 		}
 	}
 	return ipSetRestore
 }
 
 // Save the given set, or all sets if none is given to stdout in a format that restore can read. The option -file can be used to specify a filename instead of stdout.
+// save "ipset save" command output to ipset.sets.
 func (set *IPSet) Save() error {
 	stdout, err := set.run("save")
 	if err != nil {
@@ -284,6 +287,7 @@ func (set *IPSet) Save() error {
 }
 
 // Restore a saved session generated by save. The saved session can be fed from stdin or the option -file can be used to specify a filename instead of stdin. Please note, existing sets and elements are not erased by restore unless specified so in the restore file. All commands are allowed in restore mode except list, help, version, interactive mode and restore itself.
+// Send formated ipset.sets into stdin of "ipset restore" command.
 func (set *IPSet) Restore() error {
 	stdin := bytes.NewBufferString(buildIPSetRestore(set))
 	_, err := set.runWithStdin(stdin, "restore", "-exist")
@@ -313,12 +317,7 @@ func (set *IPSet) Flush() error {
 
 // Get Set by Name.
 func (ipset *IPSet) Get(setName string) *Set {
-	for _, set := range ipset.Sets {
-		if set.Name == setName {
-			return set
-		}
-	}
-	return nil
+	return ipset.Sets[setName]
 }
 
 // Rename a set. Set identified by SETNAME-TO must not exist.
@@ -354,7 +353,7 @@ func (set *Set) Refresh(entries []string, extraOptions ...string) error {
 			Options: append([]string{entry}, extraOptions...),
 		})
 	}
-	set.IPSet.Sets = append(set.IPSet.Sets, s)
+	set.IPSet.Sets[tempName] = s
 	err := set.IPSet.Restore()
 	if err != nil {
 		return err
@@ -370,5 +369,9 @@ func (set *Set) Refresh(entries []string, extraOptions ...string) error {
 		return err
 	}
 
+	s.Name = set.Name
+	set.IPSet.Sets[set.Name] = s
+	delete(set.IPSet.Sets, tempName)
+
 	return nil
 }