Merge pull request #282 from kubescape/multiplier

Adding resources multiplier
kubescape · Aug 15, 2024 · 503ddd1 · 503ddd1
2 parents 841ea43 + ee276d5
commit 503ddd1
Show file tree

Hide file tree

Showing 5 changed files with 280 additions and 15 deletions.
diff --git a/pkg/containerwatcher/v1/container_watcher_private.go b/pkg/containerwatcher/v1/container_watcher_private.go
@@ -4,7 +4,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"os"
 	"runtime"
+	"strings"
 	"time"
 
 	containercollection "github.com/inspektor-gadget/inspektor-gadget/pkg/container-collection"
@@ -381,6 +383,12 @@ func (ch *IGContainerWatcher) ignoreContainer(namespace, name string) bool {
 	if name == ch.podName && namespace == ch.namespace {
 		return true
 	}
+	// do not trace the node-agent pods if MULTIPLY is set
+	if m := os.Getenv("MULTIPLY"); m == "true" {
+		if strings.HasPrefix(name, "node-agent") {
+			return true
+		}
+	}
 	// check if config excludes the namespace
 	return ch.cfg.SkipNamespace(namespace)
 }
diff --git a/pkg/storage/v1/applicationprofile.go b/pkg/storage/v1/applicationprofile.go
@@ -14,10 +14,17 @@ import (
 )
 
 func (sc Storage) GetApplicationProfile(namespace, name string) (*v1beta1.ApplicationProfile, error) {
-	return sc.StorageClient.ApplicationProfiles(namespace).Get(context.Background(), name, v1.GetOptions{})
+	ap, err := sc.StorageClient.ApplicationProfiles(namespace).Get(context.Background(), sc.modifyName(name), v1.GetOptions{})
+	if ap != nil {
+		sc.revertNameP(&ap.Name)
+	}
+	return ap, err
 }
 
 func (sc Storage) CreateApplicationProfile(profile *v1beta1.ApplicationProfile, namespace string) error {
+	sc.modifyNameP(&profile.Name)
+	defer sc.revertNameP(&profile.Name)
+
 	// unset resourceVersion
 	profile.ResourceVersion = ""
 	_, err := sc.StorageClient.ApplicationProfiles(namespace).Create(context.Background(), profile, v1.CreateOptions{})
@@ -42,7 +49,7 @@ func (sc Storage) patchApplicationProfile(name, namespace string, operations []u
 	if err != nil {
 		return fmt.Errorf("marshal patch: %w", err)
 	}
-	profile, err := sc.StorageClient.ApplicationProfiles(namespace).Patch(context.Background(), name, types.JSONPatchType, patch, v1.PatchOptions{})
+	profile, err := sc.StorageClient.ApplicationProfiles(namespace).Patch(context.Background(), sc.modifyName(name), types.JSONPatchType, patch, v1.PatchOptions{})
 	if err != nil {
 		return fmt.Errorf("patch application profile: %w", err)
 	}
@@ -84,7 +91,7 @@ func (sc Storage) patchApplicationProfile(name, namespace string, operations []u
 			if err != nil {
 				return fmt.Errorf("create patch for annotations: %w", err)
 			}
-			_, err = sc.StorageClient.ApplicationProfiles(namespace).Patch(context.Background(), name, types.JSONPatchType, annotationsPatch, v1.PatchOptions{})
+			_, err = sc.StorageClient.ApplicationProfiles(namespace).Patch(context.Background(), sc.modifyName(name), types.JSONPatchType, annotationsPatch, v1.PatchOptions{})
 			if err != nil {
 				return fmt.Errorf("patch application profile annotations: %w", err)
 			}

diff --git a/pkg/storage/v1/network.go b/pkg/storage/v1/network.go
@@ -15,10 +15,17 @@ import (
 )
 
 func (sc Storage) GetNetworkNeighborhood(namespace, name string) (*v1beta1.NetworkNeighborhood, error) {
-	return sc.StorageClient.NetworkNeighborhoods(namespace).Get(context.Background(), name, v1.GetOptions{})
+	nn, err := sc.StorageClient.NetworkNeighborhoods(namespace).Get(context.Background(), sc.modifyName(name), v1.GetOptions{})
+	if nn != nil {
+		sc.revertNameP(&nn.Name)
+	}
+	return nn, err
 }
 
 func (sc Storage) CreateNetworkNeighborhood(neighborhood *v1beta1.NetworkNeighborhood, namespace string) error {
+	sc.modifyNameP(&neighborhood.Name)
+	defer sc.revertNameP(&neighborhood.Name)
+
 	// unset resourceVersion
 	neighborhood.ResourceVersion = ""
 	_, err := sc.StorageClient.NetworkNeighborhoods(namespace).Create(context.Background(), neighborhood, v1.CreateOptions{})
@@ -43,7 +50,7 @@ func (sc Storage) patchNetworkNeighborhood(name, namespace string, operations []
 	if err != nil {
 		return fmt.Errorf("marshal patch: %w", err)
 	}
-	neighborhood, err := sc.StorageClient.NetworkNeighborhoods(namespace).Patch(context.Background(), name, types.JSONPatchType, patch, v1.PatchOptions{})
+	neighborhood, err := sc.StorageClient.NetworkNeighborhoods(namespace).Patch(context.Background(), sc.modifyName(name), types.JSONPatchType, patch, v1.PatchOptions{})
 	if err != nil {
 		return fmt.Errorf("patch application neighborhood: %w", err)
 	}
@@ -83,7 +90,7 @@ func (sc Storage) patchNetworkNeighborhood(name, namespace string, operations []
 			if err != nil {
 				return fmt.Errorf("create patch for annotations: %w", err)
 			}
-			_, err = sc.StorageClient.NetworkNeighborhoods(namespace).Patch(context.Background(), name, types.JSONPatchType, annotationsPatch, v1.PatchOptions{})
+			_, err = sc.StorageClient.NetworkNeighborhoods(namespace).Patch(context.Background(), sc.modifyName(name), types.JSONPatchType, annotationsPatch, v1.PatchOptions{})
 			if err != nil {
 				return fmt.Errorf("patch application neighborhood annotations: %w", err)
 			}

diff --git a/pkg/storage/v1/storage.go b/pkg/storage/v1/storage.go
@@ -6,8 +6,10 @@ import (
 	"fmt"
 	"os"
 	"strconv"
+	"strings"
 	"time"
 
+	"github.com/kubescape/node-agent/pkg/config"
 	"github.com/kubescape/node-agent/pkg/storage"
 
 	"github.com/cenkalti/backoff/v4"
@@ -35,23 +37,24 @@ type Storage struct {
 	maxNetworkNeighborhoodSize int
 	maxJsonPatchOperations     int
 	namespace                  string
+	multiplier                 *int // used for testing to multiply the resources by this
 }
 
 var _ storage.StorageClient = (*Storage)(nil)
 
 func CreateStorage(namespace string) (*Storage, error) {
-	var config *rest.Config
+	var cfg *rest.Config
 	kubeconfig := os.Getenv(KubeConfig)
 	// use the current context in kubeconfig
-	config, err := clientcmd.BuildConfigFromFlags("", kubeconfig)
+	cfg, err := clientcmd.BuildConfigFromFlags("", kubeconfig)
 	if err != nil {
-		config, err = rest.InClusterConfig()
+		cfg, err = rest.InClusterConfig()
 		if err != nil {
 			return nil, fmt.Errorf("failed to create K8S Aggregated API Client with err: %v", err)
 		}
 	}
 
-	clientset, err := versioned.NewForConfig(config)
+	clientset, err := versioned.NewForConfig(cfg)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create K8S Aggregated API Client with err: %v", err)
 	}
@@ -84,6 +87,7 @@ func CreateStorage(namespace string) (*Storage, error) {
 		maxNetworkNeighborhoodSize: maxNetworkNeighborhoodSize,
 		maxJsonPatchOperations:     9999,
 		namespace:                  namespace,
+		multiplier:                 getMultiplier(),
 	}, nil
 }
 
@@ -95,6 +99,9 @@ func CreateFakeStorage(namespace string) (*Storage, error) {
 }
 
 func (sc Storage) CreateNetworkNeighbors(networkNeighbors *v1beta1.NetworkNeighbors, namespace string) error {
+	sc.modifyNameP(&networkNeighbors.Name)
+	defer sc.modifyNameP(&networkNeighbors.Name)
+
 	_, err := sc.StorageClient.NetworkNeighborses(namespace).Create(context.Background(), networkNeighbors, metav1.CreateOptions{})
 	if err != nil {
 		return err
@@ -103,37 +110,58 @@ func (sc Storage) CreateNetworkNeighbors(networkNeighbors *v1beta1.NetworkNeighb
 }
 
 func (sc Storage) GetNetworkNeighbors(namespace, name string) (*v1beta1.NetworkNeighbors, error) {
-	return sc.StorageClient.NetworkNeighborses(namespace).Get(context.Background(), name, metav1.GetOptions{})
+	nn, err := sc.StorageClient.NetworkNeighborses(namespace).Get(context.Background(), sc.modifyName(name), metav1.GetOptions{})
+	if nn != nil {
+		sc.revertNameP(&nn.Name)
+	}
+	return nn, err
 }
 
 func (sc Storage) PatchNetworkNeighborsIngressAndEgress(name, namespace string, networkNeighbors *v1beta1.NetworkNeighbors) error {
+	sc.modifyNameP(&networkNeighbors.Name)
+	defer sc.revertNameP(&networkNeighbors.Name)
+
 	bytes, err := json.Marshal(networkNeighbors)
 	if err != nil {
 		return err
 	}
 
-	_, err = sc.StorageClient.NetworkNeighborses(namespace).Patch(context.Background(), name, types.StrategicMergePatchType, bytes, metav1.PatchOptions{})
+	_, err = sc.StorageClient.NetworkNeighborses(namespace).Patch(context.Background(), sc.modifyName(name), types.StrategicMergePatchType, bytes, metav1.PatchOptions{})
 	if err != nil {
 		return err
 	}
 
 	return nil
 }
 
-func (sc Storage) PatchNetworkNeighborsMatchLabels(_, namespace string, networkNeighbors *v1beta1.NetworkNeighbors) error {
+func (sc Storage) PatchNetworkNeighborsMatchLabels(name, namespace string, networkNeighbors *v1beta1.NetworkNeighbors) error {
+	sc.modifyNameP(&networkNeighbors.Name)
+	defer sc.revertNameP(&networkNeighbors.Name)
 	_, err := sc.StorageClient.NetworkNeighborses(namespace).Update(context.Background(), networkNeighbors, metav1.UpdateOptions{})
 
 	return err
 }
 
 func (sc Storage) CreateApplicationActivity(activity *v1beta1.ApplicationActivity, namespace string) error {
+	sc.modifyNameP(&activity.Name)
+	defer sc.revertNameP(&activity.Name)
+
 	_, err := sc.StorageClient.ApplicationActivities(namespace).Create(context.Background(), activity, metav1.CreateOptions{})
 	if err != nil {
 		return err
 	}
 	return nil
 }
 
+func (sc Storage) GetApplicationActivity(namespace, name string) (*v1beta1.ApplicationActivity, error) {
+
+	aa, err := sc.StorageClient.ApplicationActivities(namespace).Get(context.Background(), sc.modifyName(name), metav1.GetOptions{})
+	if aa != nil {
+		sc.revertNameP(&aa.Name)
+	}
+	return aa, err
+}
+
 func (sc Storage) CreateFilteredSBOM(SBOM *v1beta1.SBOMSyftFiltered) error {
 	_, err := sc.StorageClient.SBOMSyftFiltereds(sc.namespace).Create(context.Background(), SBOM, metav1.CreateOptions{})
 	if err != nil {
@@ -150,8 +178,9 @@ func (sc Storage) GetSBOM(name string) (*v1beta1.SBOMSyft, error) {
 	return sc.StorageClient.SBOMSyfts(sc.namespace).Get(context.Background(), name, metav1.GetOptions{})
 }
 
-func (sc Storage) PatchFilteredSBOM(name string, SBOM *v1beta1.SBOMSyftFiltered) error {
-	bytes, err := json.Marshal(SBOM)
+func (sc Storage) PatchFilteredSBOM(name string, sbom *v1beta1.SBOMSyftFiltered) error {
+
+	bytes, err := json.Marshal(sbom)
 	if err != nil {
 		return err
 	}
@@ -169,3 +198,34 @@ func (sc Storage) IncrementImageUse(_ string) {
 func (sc Storage) DecrementImageUse(_ string) {
 	// noop
 }
+
+func (sc Storage) modifyName(n string) string {
+	if sc.multiplier != nil {
+		return fmt.Sprintf("%s-%d", n, *sc.multiplier)
+	}
+	return n
+}
+func (sc Storage) modifyNameP(n *string) {
+	if sc.multiplier != nil {
+		*n = fmt.Sprintf("%s-%d", *n, *sc.multiplier)
+	}
+}
+
+func (sc Storage) revertNameP(n *string) {
+	if sc.multiplier != nil {
+		*n = strings.TrimSuffix(*n, fmt.Sprintf("-%d", *sc.multiplier))
+	}
+}
+func getMultiplier() *int {
+	if m := os.Getenv("MULTIPLY"); m != "true" {
+		return nil
+	}
+	podName := os.Getenv(config.PodNameEnvVar)
+	s := strings.Split(podName, "-")
+	if len(s) > 0 {
+		if m, err := strconv.Atoi(s[len(s)-1]); err == nil {
+			return &m
+		}
+	}
+	return nil
+}