improve searching in CIDR ranges using cidranger

go.mod

@@ -25,6 +25,7 @@ require (
 	github.com/spf13/afero v1.11.0
 	github.com/spf13/cobra v1.8.1
 	github.com/stretchr/testify v1.9.0
+	github.com/yl2chen/cidranger v1.0.2
 	go.opentelemetry.io/otel v1.30.0
 	go.uber.org/multierr v1.11.0
 	go.uber.org/zap v1.27.0

go.sum

@@ -1339,6 +1339,8 @@ github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2 h1:eY9dn8+vbi4tKz5
 github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU=
 github.com/xo/terminfo v0.0.0-20210125001918-ca9a967f8778 h1:QldyIu/L63oPpyvQmHgvgickp1Yw510KJOqX7H24mg8=
 github.com/xo/terminfo v0.0.0-20210125001918-ca9a967f8778/go.mod h1:2MuV+tbUrU1zIOPMxZ5EncGwgmMJsa+9ucAQZXxsObs=
+github.com/yl2chen/cidranger v1.0.2 h1:lbOWZVCG1tCRX4u24kuM1Tb4nHqWkDxwLdoS+SevawU=
+github.com/yl2chen/cidranger v1.0.2/go.mod h1:9U1yz7WPYDwf0vpNWFaeRh0bjwz5RVgRy/9UEQfHl0g=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

pkg/apis/softwarecomposition/knownserversfinder.go

@@ -0,0 +1,77 @@
+package softwarecomposition
+
+import (
+	"net"
+
+	"github.com/yl2chen/cidranger"
+)
+
+var _ IKnownServersFinder = (*KnownServersFinderImpl)(nil)
+var _ IKnownServerEntry = (*KnownServersFinderEntry)(nil)
+
+type KnownServersFinderImpl struct {
+	ranger cidranger.Ranger
+}
+
+type KnownServersFinderEntry struct {
+	knownServer KnownServerEntry
+	network     net.IPNet
+}
+
+func NewKnownServersFinderImpl(knownServers []KnownServer) IKnownServersFinder {
+	// build the ranger for searching
+	ranger := cidranger.NewPCTrieRanger()
+	for _, knownServer := range knownServers {
+		for _, knownServerEntry := range knownServer.Spec {
+			if v := NewKnownServersFinderEntry(knownServerEntry); v != nil {
+				_ = ranger.Insert(v)
+			}
+		}
+	}
+	return &KnownServersFinderImpl{
+		ranger: ranger,
+	}
+}
+
+func (k *KnownServersFinderImpl) Contains(ip net.IP) ([]IKnownServerEntry, bool) {
+	if k.ranger == nil {
+		return nil, false
+	}
+	contains, _ := k.ranger.Contains(ip)
+	if contains {
+		if entries, err := k.ranger.ContainingNetworks(ip); err == nil && len(entries) > 0 {
+			knownServersEntries := make([]IKnownServerEntry, 0, len(entries))
+			for _, entry := range entries {
+				if v, ok := entry.(*KnownServersFinderEntry); ok {
+					knownServersEntries = append(knownServersEntries, v)
+				}
+			}
+			return knownServersEntries, true
+		}
+	}
+	return nil, false
+}
+
+func NewKnownServersFinderEntry(kse KnownServerEntry) *KnownServersFinderEntry {
+	_, res, err := net.ParseCIDR(kse.IPBlock)
+	if err != nil || res == nil {
+		return nil
+	}
+	return &KnownServersFinderEntry{knownServer: kse, network: *res}
+}
+
+func (k *KnownServersFinderEntry) GetServer() string {
+	return k.knownServer.Server
+}
+
+func (k *KnownServersFinderEntry) GetName() string {
+	return k.knownServer.Name
+}
+
+func (k *KnownServersFinderEntry) GetIPBlock() string {
+	return k.knownServer.IPBlock
+}
+
+func (k *KnownServersFinderEntry) Network() net.IPNet {
+	return k.network
+}

pkg/apis/softwarecomposition/knownserversfinder_test.go

@@ -0,0 +1,57 @@
+package softwarecomposition
+
+import (
+	"net"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestContains(t *testing.T) {
+	knownServers := []KnownServer{
+		{
+			Spec: []KnownServerEntry{
+				{Server: "server1", Name: "name1", IPBlock: "192.168.1.0/24"},
+				{Server: "server2", Name: "name2", IPBlock: "10.0.0.0/8"},
+				{Server: "server3", Name: "name3", IPBlock: ""},
+				{Server: "server4", Name: "name4", IPBlock: "invalid"},
+				{Server: "server5", Name: "name5", IPBlock: "192.168.1.128/25"},
+			},
+		},
+	}
+
+	finder := NewKnownServersFinderImpl(knownServers)
+
+	tests := []struct {
+		ip               string
+		expected         bool
+		expectedIPBlocks []string
+		expectedServers  []string
+	}{
+		{"192.168.1.1", true, []string{"192.168.1.0/24"}, []string{"server1"}},
+		{"10.0.0.1", true, []string{"10.0.0.0/8"}, []string{"server2"}},
+		{"172.16.0.1", false, nil, nil},
+		{"192.168.1.200", true, []string{"192.168.1.0/24", "192.168.1.128/25"}, []string{"server1", "server5"}},
+	}
+
+	for _, test := range tests {
+		ip := net.ParseIP(test.ip)
+		entries, contains := finder.Contains(ip)
+		assert.Equal(t, test.expected, contains)
+		if contains {
+			assert.NotEmpty(t, entries)
+			servers := []string{}
+			for _, entry := range entries {
+				servers = append(servers, entry.GetServer())
+			}
+			ipBlocks := []string{}
+			for _, entry := range entries {
+				ipBlocks = append(ipBlocks, entry.GetIPBlock())
+			}
+			assert.ElementsMatch(t, test.expectedIPBlocks, ipBlocks)
+			assert.ElementsMatch(t, test.expectedServers, servers)
+		} else {
+			assert.Empty(t, entries)
+		}
+	}
+}

pkg/apis/softwarecomposition/network_types.go

@@ -1,6 +1,8 @@
 package softwarecomposition
 
 import (
+	"net"
+
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -163,3 +165,13 @@ type KnownServerEntry struct {
 	Server  string
 	Name    string
 }
+
+type IKnownServerEntry interface {
+	GetIPBlock() string
+	GetName() string
+	GetServer() string
+}
+
+type IKnownServersFinder interface {
+	Contains(ip net.IP) ([]IKnownServerEntry, bool)
+}

pkg/apis/softwarecomposition/networkpolicy/v1/networkpolicy.go

@@ -12,16 +12,14 @@ import (
 
 	helpersv1 "github.com/kubescape/k8s-interface/instanceidhandler/v1/helpers"
 
-	"github.com/kubescape/go-logger"
-	"github.com/kubescape/go-logger/helpers"
 	"github.com/kubescape/storage/pkg/apis/softwarecomposition"
 	"github.com/kubescape/storage/pkg/apis/softwarecomposition/networkpolicy"
 	"golang.org/x/exp/maps"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
-func GenerateNetworkPolicy(networkNeighbors softwarecomposition.NetworkNeighbors, knownServers []softwarecomposition.KnownServer, timeProvider metav1.Time) (softwarecomposition.GeneratedNetworkPolicy, error) {
+func GenerateNetworkPolicy(networkNeighbors softwarecomposition.NetworkNeighbors, knownServers softwarecomposition.IKnownServersFinder, timeProvider metav1.Time) (softwarecomposition.GeneratedNetworkPolicy, error) {
 	if !IsAvailable(networkNeighbors) {
 		return softwarecomposition.GeneratedNetworkPolicy{}, fmt.Errorf("networkNeighbors %s/%s status annotation is not ready nor completed", networkNeighbors.Namespace, networkNeighbors.Name)
 	}
@@ -264,7 +262,7 @@ func mergeEgressRulesByPorts(rules []softwarecomposition.NetworkPolicyEgressRule
 	return mergedRules
 }
 
-func generateEgressRule(neighbor softwarecomposition.NetworkNeighbor, KnownServer []softwarecomposition.KnownServer) (softwarecomposition.NetworkPolicyEgressRule, []softwarecomposition.PolicyRef) {
+func generateEgressRule(neighbor softwarecomposition.NetworkNeighbor, knownServers softwarecomposition.IKnownServersFinder) (softwarecomposition.NetworkPolicyEgressRule, []softwarecomposition.PolicyRef) {
 	egressRule := softwarecomposition.NetworkPolicyEgressRule{}
 	policyRefs := []softwarecomposition.PolicyRef{}
 
@@ -288,41 +286,32 @@ func generateEgressRule(neighbor softwarecomposition.NetworkNeighbor, KnownServe
 	}
 
 	if neighbor.IPAddress != "" {
-		isKnownServer := false
-		// look if this IP is part of any known server
-		for _, knownServer := range KnownServer {
-			for _, entry := range knownServer.Spec {
-				_, subNet, err := net.ParseCIDR(entry.IPBlock)
-				if err != nil {
-					logger.L().Error("error parsing cidr", helpers.Error(err))
-					continue
+		if entries, contains := knownServers.Contains(net.ParseIP(neighbor.IPAddress)); contains {
+			// look if this IP is part of any known server
+			for _, entry := range entries {
+
+				egressRule.To = append(egressRule.To, softwarecomposition.NetworkPolicyPeer{
+					IPBlock: &softwarecomposition.IPBlock{
+						CIDR: entry.GetIPBlock(),
+					},
+				})
+
+				policyRef := softwarecomposition.PolicyRef{
+					Name:       entry.GetName(),
+					OriginalIP: neighbor.IPAddress,
+					IPBlock:    entry.GetIPBlock(),
+					Server:     entry.GetServer(),
 				}
-				if subNet.Contains(net.ParseIP(neighbor.IPAddress)) {
-					egressRule.To = append(egressRule.To, softwarecomposition.NetworkPolicyPeer{
-						IPBlock: &softwarecomposition.IPBlock{
-							CIDR: entry.IPBlock,
-						},
-					})
-					isKnownServer = true
-
-					policyRef := softwarecomposition.PolicyRef{
-						Name:       entry.Name,
-						OriginalIP: neighbor.IPAddress,
-						IPBlock:    entry.IPBlock,
-						Server:     entry.Server,
-					}
-
-					if neighbor.DNS != "" {
-						policyRef.DNS = neighbor.DNS
-					}
-
-					policyRefs = append(policyRefs, policyRef)
-					break
+
+				if neighbor.DNS != "" {
+					policyRef.DNS = neighbor.DNS
 				}
+
+				policyRefs = append(policyRefs, policyRef)
+
 			}
-		}
 
-		if !isKnownServer {
+		} else {
 			ipBlock := getSingleIP(neighbor.IPAddress)
 			egressRule.To = append(egressRule.To, softwarecomposition.NetworkPolicyPeer{
 				IPBlock: ipBlock,
@@ -361,7 +350,7 @@ func hash(s any) (string, error) {
 	return hex.EncodeToString(vv[:]), nil
 }
 
-func generateIngressRule(neighbor softwarecomposition.NetworkNeighbor, KnownServer []softwarecomposition.KnownServer) (softwarecomposition.NetworkPolicyIngressRule, []softwarecomposition.PolicyRef) {
+func generateIngressRule(neighbor softwarecomposition.NetworkNeighbor, knownServers softwarecomposition.IKnownServersFinder) (softwarecomposition.NetworkPolicyIngressRule, []softwarecomposition.PolicyRef) {
 	ingressRule := softwarecomposition.NetworkPolicyIngressRule{}
 	policyRefs := []softwarecomposition.PolicyRef{}
 
@@ -384,41 +373,30 @@ func generateIngressRule(neighbor softwarecomposition.NetworkNeighbor, KnownServ
 	}
 
 	if neighbor.IPAddress != "" {
-		isKnownServer := false
 		// look if this IP is part of any known server
-		for _, knownServer := range KnownServer {
-			for _, entry := range knownServer.Spec {
-				_, subNet, err := net.ParseCIDR(entry.IPBlock)
-				if err != nil {
-					logger.L().Error("error parsing cidr", helpers.Error(err))
-					continue
+		if entries, contains := knownServers.Contains(net.ParseIP(neighbor.IPAddress)); contains {
+			for _, entry := range entries {
+				ingressRule.From = append(ingressRule.From, softwarecomposition.NetworkPolicyPeer{
+					IPBlock: &softwarecomposition.IPBlock{
+						CIDR: entry.GetIPBlock(),
+					},
+				})
+
+				policyRef := softwarecomposition.PolicyRef{
+					Name:       entry.GetName(),
+					OriginalIP: neighbor.IPAddress,
+					IPBlock:    entry.GetIPBlock(),
+					Server:     entry.GetServer(),
 				}
-				if subNet.Contains(net.ParseIP(neighbor.IPAddress)) {
-					ingressRule.From = append(ingressRule.From, softwarecomposition.NetworkPolicyPeer{
-						IPBlock: &softwarecomposition.IPBlock{
-							CIDR: entry.IPBlock,
-						},
-					})
-					isKnownServer = true
-
-					policyRef := softwarecomposition.PolicyRef{
-						Name:       entry.Name,
-						OriginalIP: neighbor.IPAddress,
-						IPBlock:    entry.IPBlock,
-						Server:     entry.Server,
-					}
-
-					if neighbor.DNS != "" {
-						policyRef.DNS = neighbor.DNS
-					}
-
-					policyRefs = append(policyRefs, policyRef)
-					break
+
+				if neighbor.DNS != "" {
+					policyRef.DNS = neighbor.DNS
 				}
-			}
-		}
 
-		if !isKnownServer {
+				policyRefs = append(policyRefs, policyRef)
+
+			}
+		} else {
 			ipBlock := getSingleIP(neighbor.IPAddress)
 			ingressRule.From = append(ingressRule.From, softwarecomposition.NetworkPolicyPeer{
 				IPBlock: ipBlock,
@@ -432,6 +410,7 @@ func generateIngressRule(neighbor softwarecomposition.NetworkNeighbor, KnownServ
 				})
 			}
 		}
+
 	}
 
 	for _, networkPort := range neighbor.Ports {

pkg/apis/softwarecomposition/networkpolicy/v1/networkpolicy_test.go

@@ -2010,7 +2010,7 @@ func TestGenerateNetworkPolicy(t *testing.T) {
 
 	for _, test := range tests {
 
-		got, err := GenerateNetworkPolicy(test.networkNeighbors, test.KnownServer, timeProvider)
+		got, err := GenerateNetworkPolicy(test.networkNeighbors, softwarecomposition.NewKnownServersFinderImpl(test.KnownServer), timeProvider)
 
 		assert.NoError(t, err)