[agent-smith] tests for Smith.checkEgressTrafficCallback

components/ee/agent-smith/BUILD.yaml

@@ -4,6 +4,7 @@ packages:
   - name: app
     type: go
     srcs:
+      - "pkg/agent/testdata/**"
       - "**/*.go"
       - "go.mod"
       - "go.sum"

components/ee/agent-smith/pkg/agent/agent.go

@@ -22,7 +22,6 @@ import (
 
 	"github.com/cilium/ebpf/perf"
 	"github.com/gitpod-io/gitpod/agent-smith/pkg/bpf"
-	"github.com/gitpod-io/gitpod/agent-smith/pkg/network"
 	"github.com/gitpod-io/gitpod/agent-smith/pkg/signature"
 	"github.com/gitpod-io/gitpod/common-go/log"
 	"github.com/gitpod-io/gitpod/common-go/util"
@@ -52,6 +51,9 @@ type Smith struct {
 	notifiedInfringements *lru.Cache
 	perfHandler           chan perfHandlerFunc
 	pidsMap               sync.Map
+
+	egressTrafficCheckHandler func(pid int) (int64, error)
+	timeElapsedHandler        func(t time.Time) time.Duration
 }
 
 // EgressTraffic configures an upper limit of allowed egress traffic over time
@@ -138,12 +140,13 @@ func NewAgentSmith(cfg Config) (*Smith, error) {
 				GradeKind(InfringementExcessiveEgress, InfringementSeverityVery):      PenaltyStopWorkspace,
 			},
 		},
-		Config:                cfg,
-		GitpodAPI:             api,
-		notifiedInfringements: notificationCache,
-		perfHandler:           make(chan perfHandlerFunc, 10),
-		metrics:               newAgentMetrics(),
-		pidsMap:               sync.Map{},
+		Config:                    cfg,
+		GitpodAPI:                 api,
+		notifiedInfringements:     notificationCache,
+		perfHandler:               make(chan perfHandlerFunc, 10),
+		metrics:                   newAgentMetrics(),
+		egressTrafficCheckHandler: getEgressTraffic,
+		timeElapsedHandler:        time.Since,
 	}
 	if cfg.Enforcement.Default != nil {
 		if err := cfg.Enforcement.Default.Validate(); err != nil {
@@ -246,7 +249,10 @@ type GradedInfringementKind string
 
 // GradeKind produces a graded infringement kind from severity and kind
 func GradeKind(kind InfringementKind, severity InfringementSeverity) GradedInfringementKind {
-	return GradedInfringementKind(fmt.Sprintf("%s %s", severity, kind))
+	if len(severity) > 0 {
+		return GradedInfringementKind(fmt.Sprintf("%s %s", severity, kind))
+	}
+	return GradedInfringementKind(kind)
 }
 
 // Severity returns the severity of the graded infringement kind
@@ -319,7 +325,7 @@ func (agent *Smith) Start(ctx context.Context, callback func(InfringingWorkspace
 
 	defer abpf.Close()
 
-	agent.cleanupDeadPIDS(ctx)
+	go agent.cleanupDeadPIDS(ctx)
 
 	egressTicker := time.NewTicker(30 * time.Second)
 
@@ -402,17 +408,15 @@ func (agent *Smith) Start(ctx context.Context, callback func(InfringingWorkspace
 
 func (agent *Smith) cleanupDeadPIDS(ctx context.Context) {
 	ticker := time.NewTicker(30 * time.Second)
-	go func() {
-		for {
-			select {
-			case <-ticker.C:
-				agent.cleanupDeadPidsCallback()
-			case <-ctx.Done():
-				ticker.Stop()
-				return
-			}
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			agent.cleanupDeadPidsCallback()
+		case <-ctx.Done():
+			return
 		}
-	}()
+	}
 }
 
 func (agent *Smith) cleanupDeadPidsCallback() {
@@ -559,6 +563,7 @@ func parseExecveExit(evtHdr EventHeader, buffer []byte) Execve {
 	dataOffsetPtr := unsafe.Sizeof(evtHdr) + unsafe.Sizeof(i)*uintptr(evtHdr.NParams) - 6 // todo(fntlnz): check why this -6 is necessary
 	scratchHeaderOffset := uint32(dataOffsetPtr)
 
+	//lint:ignore SA4006 this is used with unsafe.Sizeof
 	retval := int64(buffer[scratchHeaderOffset])
 
 	// einfo := bpf.EventTable[bpf.PPME_SYSCALL_EXECVE_19_X]
@@ -765,46 +770,6 @@ func getWorkspaceFromProcess(tid int) (res *InfringingWorkspace, err error) {
 	}, nil
 }
 
-//nolint:deadcode,unused
-func mergeInfringingWorkspaces(vws []InfringingWorkspace) (vw InfringingWorkspace) {
-	for _, r := range vws {
-		if vw.Pod == "" {
-			vw.Pod = r.Pod
-		}
-		if vw.Owner == "" {
-			vw.Owner = r.Owner
-		}
-		if vw.InstanceID == "" {
-			vw.InstanceID = r.InstanceID
-		}
-		if vw.WorkspaceID == "" {
-			vw.WorkspaceID = r.WorkspaceID
-		}
-
-		// Note: the remote URL list is likekly to be very small hence the O^2 complexity is ok
-		//       And just in case the remote URL list isn't small we have a circuit breaker
-		if len(r.GitRemoteURL) < 100 && len(vw.GitRemoteURL) < 100 {
-			for _, rr := range r.GitRemoteURL {
-				var found bool
-				for _, vr := range vw.GitRemoteURL {
-					if rr == vr {
-						found = true
-						break
-					}
-				}
-				if !found {
-					vw.GitRemoteURL = append(vw.GitRemoteURL, rr)
-				}
-			}
-		} else if len(vw.GitRemoteURL) == 0 {
-			vw.GitRemoteURL = r.GitRemoteURL
-		}
-
-		vw.Infringements = append(vw.Infringements, r.Infringements...)
-	}
-	return
-}
-
 // RegisterMetrics registers prometheus metrics for this driver
 func (agent *Smith) RegisterMetrics(reg prometheus.Registerer) error {
 	return agent.metrics.Register(reg)
@@ -815,8 +780,8 @@ func (agent *Smith) checkEgressTrafficCallback(pid int, pidCreationTime time.Tim
 		return nil, nil
 	}
 
-	podLifetime := time.Since(pidCreationTime)
-	resp, err := network.GetEgressTraffic(pid)
+	podLifetime := agent.timeElapsedHandler(pidCreationTime)
+	resp, err := agent.egressTrafficCheckHandler(pid)
 	if err != nil {
 		return nil, err
 	}

components/ee/agent-smith/pkg/agent/agent_test.go

@@ -5,8 +5,14 @@
 package agent
 
 import (
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"path"
+	"reflect"
 	"sort"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 )
@@ -93,3 +99,135 @@ func BenchmarkFindEnforcementRules(b *testing.B) {
 		findEnforcementRules(rules, "foobar")
 	}
 }
+
+func TestCeckEgressTrafficCallback(t *testing.T) {
+	simpleTime, _ := time.Parse(time.RFC3339, "2021-07-05T15:16:17+02:00")
+
+	type args struct {
+		pid             int
+		pidCreationTime time.Time
+	}
+	tests := map[string]struct {
+		args                      args
+		want                      *Infringement
+		egressTrafficCheckHandler func(pid int) (int64, error)
+		timeElapsedHandler        func(t time.Time) time.Duration
+		wantErr                   bool
+	}{
+		"no_infringement": {
+			args: args{
+				pid:             1234,
+				pidCreationTime: simpleTime,
+			},
+			want: nil,
+			egressTrafficCheckHandler: func(pid int) (int64, error) {
+				return 2000000, nil
+			},
+			timeElapsedHandler: func(t time.Time) time.Duration {
+				d, _ := time.ParseDuration("1m")
+				return d
+			},
+			wantErr: false,
+		},
+		"zero_egress": {
+			args: args{
+				pid:             1234,
+				pidCreationTime: simpleTime,
+			},
+			want: nil,
+			egressTrafficCheckHandler: func(pid int) (int64, error) {
+				return 0, nil
+			},
+			timeElapsedHandler: func(t time.Time) time.Duration {
+				d, _ := time.ParseDuration("1m")
+				return d
+			},
+			wantErr: false,
+		},
+		"excessive_egress": {
+			args: args{
+				pid:             1234,
+				pidCreationTime: simpleTime,
+			},
+			want: &Infringement{
+				Kind:        GradedInfringementKind(InfringementExcessiveEgress),
+				Description: "egress traffic is 12.805 megabytes over limit",
+			},
+			egressTrafficCheckHandler: func(pid int) (int64, error) {
+				return 328000000, nil
+			},
+			timeElapsedHandler: func(t time.Time) time.Duration {
+				d, _ := time.ParseDuration("1m")
+				return d
+			},
+			wantErr: false,
+		},
+		"very_excessive_egress_simple": {
+			args: args{
+				pid:             1234,
+				pidCreationTime: simpleTime,
+			},
+			want: &Infringement{
+				Kind:        GradedInfringementKind(InfringementVeryExcessiveEgress),
+				Description: "egress traffic is 188686.863 megabytes over limit",
+			},
+			egressTrafficCheckHandler: func(pid int) (int64, error) {
+				return 200000000000, nil
+			},
+			timeElapsedHandler: func(t time.Time) time.Duration {
+				d, _ := time.ParseDuration("1s")
+				return d
+			},
+			wantErr: false,
+		},
+		"very_excessive_egress": {
+			args: args{
+				pid:             1234,
+				pidCreationTime: simpleTime,
+			},
+			want: &Infringement{
+				Kind:        GradedInfringementKind(InfringementVeryExcessiveEgress),
+				Description: "egress traffic is 188686.863 megabytes over limit",
+			},
+			egressTrafficCheckHandler: func(pid int) (int64, error) {
+				return 200000000000, nil
+			},
+			timeElapsedHandler: func(t time.Time) time.Duration {
+				d, _ := time.ParseDuration("1m")
+				return d
+			},
+			wantErr: false,
+		},
+	}
+
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			fc, err := ioutil.ReadFile(path.Join("testdata", fmt.Sprintf("agent_check_egress_%s.golden", name)))
+			if err != nil {
+				t.Errorf("cannot read config: %v", err)
+				return
+			}
+			var cfg Config
+			err = json.Unmarshal(fc, &cfg)
+			if err != nil {
+				t.Errorf("cannot unmarshal config: %v", err)
+				return
+			}
+			agent, err := NewAgentSmith(cfg)
+			if err != nil {
+				t.Errorf("cannot create test agent smith from config: %v", err)
+				return
+			}
+			agent.egressTrafficCheckHandler = tt.egressTrafficCheckHandler
+			agent.timeElapsedHandler = tt.timeElapsedHandler
+			got, err := agent.checkEgressTrafficCallback(tt.args.pid, tt.args.pidCreationTime)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("Smith.checkEgressTrafficCallback() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("Smith.checkEgressTrafficCallback() = %s", cmp.Diff(got, tt.want))
+			}
+		})
+	}
+}

components/ee/agent-smith/pkg/network/egress.go → components/ee/agent-smith/pkg/agent/egress.go

@@ -2,13 +2,13 @@
 // Licensed under the Gitpod Enterprise Source Code License,
 // See License.enterprise.txt in the project root folder.
 
-package network
+package agent
 
 import (
 	"github.com/prometheus/procfs"
 )
 
-func GetEgressTraffic(pid int) (int64, error) {
+func getEgressTraffic(pid int) (int64, error) {
 	pproc, err := procfs.NewProc(pid)
 	if err != nil {
 		return -1, err

components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_excessive_egress.golden

@@ -0,0 +1,13 @@
+{
+    "egressTraffic": {
+        "dt": "2m",
+        "excessive": {
+            "baseBudget": "300Mi",
+            "perDtThreshold": "100Mi"
+        },
+        "veryExcessive": {
+            "baseBudget": "2Gi",
+            "perDtThreshold": "250Mi"
+        }
+    }
+}

components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_no_infringement.golden

@@ -0,0 +1,13 @@
+{
+    "egressTraffic": {
+        "dt": "2m",
+        "excessive": {
+            "baseBudget": "300Mi",
+            "perDtThreshold": "100Mi"
+        },
+        "veryExcessive": {
+            "baseBudget": "2Gi",
+            "perDtThreshold": "250Mi"
+        }
+    }
+}

components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_very_excessive_egress.golden

@@ -0,0 +1,13 @@
+{
+    "egressTraffic": {
+        "dt": "1h",
+        "excessive": {
+            "baseBudget": "300Mi",
+            "perDtThreshold": "100Mi"
+        },
+        "veryExcessive": {
+            "baseBudget": "2Gi",
+            "perDtThreshold": "250Mi"
+        }
+    }
+}

components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_very_excessive_egress_simple.golden

@@ -0,0 +1,13 @@
+{
+    "egressTraffic": {
+        "dt": "2m",
+        "excessive": {
+            "baseBudget": "300Mi",
+            "perDtThreshold": "100Mi"
+        },
+        "veryExcessive": {
+            "baseBudget": "2Gi",
+            "perDtThreshold": "250Mi"
+        }
+    }
+}

components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_zero_egress.golden

@@ -0,0 +1,4 @@
+{
+    "egressTraffic": {
+    }
+}