From 9198bb5a68ebcdbf890cceb22c336579c360e6b0 Mon Sep 17 00:00:00 2001
From: Martynas Pumputis <martynasp@gmail.com>
Date: Tue, 22 May 2018 10:59:05 +0200
Subject: [PATCH] npc: Refactor npc/ipblock.go

- No need to validate networkingv1.IPBlock fields as k8s does it.
- Fix provisioning.
---
 npc/ipblock.go | 133 +++++++++++++++++++++++--------------------------
 1 file changed, 63 insertions(+), 70 deletions(-)

diff --git a/npc/ipblock.go b/npc/ipblock.go
index fd5df974..594af2b4 100644
--- a/npc/ipblock.go
+++ b/npc/ipblock.go
@@ -1,40 +1,32 @@
 package npc
 
 import (
-	"github.com/weaveworks/weave/common"
-	"github.com/weaveworks/weave/npc/ipset"
-	networkingv1 "k8s.io/api/networking/v1"
-	"k8s.io/apimachinery/pkg/types"
-	"net"
 	"sort"
 	"strings"
+
+	networkingv1 "k8s.io/api/networking/v1"
+	"k8s.io/apimachinery/pkg/types"
+
+	"github.com/weaveworks/weave/common"
+	"github.com/weaveworks/weave/npc/ipset"
 )
 
-type ipBlock struct {
-	ipsetName ipset.Name // generated ipset name
-	spec      *networkingv1.IPBlock
+type ipBlockSpec struct {
+	key       string
+	ipsetName ipset.Name // ipset for storing excepted CIDRs
+	ipBlock   *networkingv1.IPBlock
 }
 
-func newIPBlock(ipb *networkingv1.IPBlock, ns string) (b *ipBlock, err error) {
-	if _, _, err = net.ParseCIDR(ipb.CIDR); err != nil {
-		return
-	}
-
-	for _, ex := range ipb.Except {
-		if _, _, err = net.ParseCIDR(ex); err != nil {
-			return
-		}
-	}
-
-	b = &ipBlock{spec: ipb}
+func newIPBlockSpec(ipb *networkingv1.IPBlock) *ipBlockSpec {
+	spec := &ipBlockSpec{ipBlock: ipb}
 
 	if len(ipb.Except) > 0 {
-		sortedExcept := ipb.Except
-		sort.Strings(sortedExcept)
-		b.ipsetName = ipset.Name(IpsetNamePrefix + shortName("except:"+strings.Join(sortedExcept, "")))
+		sort.Strings(ipb.Except)
+		spec.key = strings.Join(ipb.Except, " ")
+		spec.ipsetName = ipset.Name(IpsetNamePrefix + shortName("ipblock-except:"+spec.key))
 	}
 
-	return
+	return spec
 }
 
 //func (ipb ipBlock) GetRuleSpec() (spec *ruleSpec) {
@@ -56,65 +48,66 @@ func newIPBlock(ipb *networkingv1.IPBlock, ns string) (b *ipBlock, err error) {
 
 type ipBlockSet struct {
 	ips   ipset.Interface
-	users map[string]int
+	users map[string]map[types.UID]struct{}
 }
 
 func newIPBlockSet(ips ipset.Interface) *ipBlockSet {
 	return &ipBlockSet{
 		ips:   ips,
-		users: make(map[string]int),
+		users: make(map[string]map[types.UID]struct{}),
 	}
 }
 
-func (s *ipBlockSet) deprovision(uid types.UID, current *ipBlock) (err error) {
-	if current == nil || len(current.ipsetName) == 0 {
-		return
-	}
-
-	key := string(current.ipsetName)
-	if _, found := s.users[key]; !found {
-		return
-	}
-
-	s.users[key]--
-	if s.users[key] == 0 {
-		if err = s.ips.Destroy(current.ipsetName); err != nil {
-			return
+func (s *ipBlockSet) deprovision(user types.UID, current, desired map[string]*ipBlockSpec) error {
+	for key, spec := range current {
+		if key == "" {
+			continue
 		}
 
-		delete(s.users, key)
-	}
+		if _, found := desired[key]; !found {
+			delete(s.users[key], user)
+			if len(s.users[key]) == 0 {
+				common.Log.Infof("destroying ipset: %#v", spec)
+				if err := s.ips.Destroy(spec.ipsetName); err != nil {
+					return err
+				}
 
-	return
-}
-
-func (s *ipBlockSet) provision(uid types.UID, desired *ipBlock) (err error) {
-	if desired == nil || len(desired.ipsetName) == 0 {
-		return
-	}
-
-	key := string(desired.ipsetName)
-	if _, found := s.users[key]; found {
-		common.Log.Debug()
-		s.users[key]++
-		return
-	}
-
-	err = s.ips.Create(desired.ipsetName, ipset.HashNet)
-	if err != nil {
-		return
-	}
-
-	comment := `excepted from ` + desired.spec.CIDR
-	for _, excepted := range desired.spec.Except {
-		if err = s.ips.AddEntry(uid, desired.ipsetName, excepted, comment); err != nil {
-			if details := s.ips.Destroy(desired.ipsetName); details != nil {
-				common.Log.Warn("can't destroy ipset ", desired.ipsetName, ":", details)
+				delete(s.users, key)
 			}
-			return
 		}
 	}
 
-	s.users[key] = 1
-	return
+	return nil
+}
+
+func (s *ipBlockSet) provision(user types.UID, current, desired map[string]*ipBlockSpec) (err error) {
+	for key, spec := range desired {
+		if key == "" {
+			// No need to provision an ipBlock with empty list of excepted CIDRs
+			// (i.e. no ipset is needed for the related iptables rule).
+			continue
+		}
+
+		if _, found := current[key]; !found {
+			if _, found := s.users[key]; !found {
+				common.Log.Infof("creating ipset: %#v", spec)
+				if err := s.ips.Create(spec.ipsetName, ipset.HashNet); err != nil {
+					return err
+				}
+
+				// TODO(brb) Pass comment to ips.Create instead.
+				comment := "excepted from " + spec.ipBlock.CIDR
+				for _, cidr := range spec.ipBlock.Except {
+					if err = s.ips.AddEntry(user, spec.ipsetName, cidr, comment); err != nil {
+						return err
+					}
+				}
+
+				s.users[key] = make(map[types.UID]struct{})
+			}
+			s.users[key][user] = struct{}{}
+		}
+	}
+
+	return nil
 }