feat: allow adjustment of QPS/Burst values (#263)

* feat: allow adjustment of QPS/Burst values * chore: add function to combine multiple functions * chore: rename symbols to generate cleaner diff * chore: address PR review comments * chore: increase QPS/burst to 100/100
aquasecurity · Jan 8, 2024 · a47c86a · a47c86a
1 parent f921fb0
commit a47c86a
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 8 deletions.
diff --git a/examples/trivy.go b/examples/trivy.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log"
+	"time"
 
 	"github.com/aquasecurity/trivy-kubernetes/pkg/artifacts"
 	"github.com/aquasecurity/trivy-kubernetes/pkg/k8s"
@@ -23,7 +24,7 @@ func main() {
 
 	ctx := context.Background()
 
-	cluster, err := k8s.GetCluster()
+	cluster, err := k8s.GetCluster(k8s.WithBurst(100), k8s.WithQPS(100))
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -34,10 +35,13 @@ func main() {
 	fmt.Println("Scanning cluster")
 
 	//trivy k8s #cluster
+	start := time.Now()
 	artifacts, err := trivyk8s.ListArtifacts(ctx)
+	end := time.Now()
 	if err != nil {
 		log.Fatal(err)
 	}
+	fmt.Printf("Scan took %v\n", end.Sub(start))
 	printArtifacts(artifacts)
 
 	fmt.Println("Scanning kind 'pods' with exclude-owned=true")

diff --git a/pkg/k8s/k8s.go b/pkg/k8s/k8s.go
@@ -138,6 +138,36 @@ func WithKubeConfig(kubeConfig string) ClusterOption {
 		c.KubeConfig = &kubeConfig
 	}
 }
+func WithQPS(qps float32) ClusterOption {
+	return func(o *genericclioptions.ConfigFlags) {
+		o.WrapConfigFn = combineConfigFns(o.WrapConfigFn, func(c *rest.Config) *rest.Config {
+			c.QPS = qps
+			return c
+		})
+	}
+}
+
+func WithBurst(burst int) ClusterOption {
+	return func(o *genericclioptions.ConfigFlags) {
+		o.WrapConfigFn = combineConfigFns(o.WrapConfigFn, func(c *rest.Config) *rest.Config {
+			c.Burst = burst
+			return c
+		})
+	}
+}
+
+// Helper function to combine multiple config functions
+func combineConfigFns(existing, newFn func(*rest.Config) *rest.Config) func(*rest.Config) *rest.Config {
+	if existing == nil {
+		return newFn
+	}
+	return func(c *rest.Config) *rest.Config {
+		if modified := existing(c); modified != nil {
+			return newFn(modified)
+		}
+		return newFn(c)
+	}
+}
 
 // GetCluster returns a current configured cluster,
 func GetCluster(opts ...ClusterOption) (Cluster, error) {
@@ -156,15 +186,16 @@ func GetCluster(opts ...ClusterOption) (Cluster, error) {
 		return nil, err
 	}
 
-	return getCluster(clientConfig, restMapper, *cf.Context, false)
-}
-
-func getCluster(clientConfig clientcmd.ClientConfig, restMapper meta.RESTMapper, currentContext string, fakeConfig bool) (*cluster, error) {
-	kubeConfig, err := clientConfig.ClientConfig()
+	kubeConfig, err := cf.ToRESTConfig()
 	if err != nil {
 		return nil, err
 	}
 
+	return getCluster(clientConfig, kubeConfig, restMapper, *cf.Context, false)
+}
+
+func getCluster(clientConfig clientcmd.ClientConfig, kubeConfig *rest.Config, restMapper meta.RESTMapper, currentContext string, fakeConfig bool) (*cluster, error) {
+
 	k8sDynamicClient, err := dynamic.NewForConfig(kubeConfig)
 	if err != nil {
 		return nil, err

diff --git a/pkg/k8s/k8s_test.go b/pkg/k8s/k8s_test.go
@@ -34,7 +34,9 @@ func TestGetCurrentNamespace(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.Name, func(t *testing.T) {
 			fakeConfig := createValidTestConfig(test.Namespace)
-			cluster, err := getCluster(fakeConfig, nil, "", true)
+			fakeKubeconfig, err := fakeConfig.ClientConfig()
+			assert.NoError(t, err)
+			cluster, err := getCluster(fakeConfig, fakeKubeconfig, nil, "", true)
 			assert.NoError(t, err)
 			assert.Equal(t, test.ExpectedNamespace, cluster.GetCurrentNamespace())
 		})
@@ -66,7 +68,10 @@ func TestGetGVR(t *testing.T) {
 
 		fakeConfig := createValidTestConfig("")
 
-		cluster, err := getCluster(fakeConfig, mapper, "", true)
+		fakeKubeconfig, err := fakeConfig.ClientConfig()
+		assert.NoError(t, err)
+
+		cluster, err := getCluster(fakeConfig, fakeKubeconfig, mapper, "", true)
 		assert.NoError(t, err)
 
 		gvr, err := cluster.GetGVR(test.Resource)