diff --git a/components/ee/agent-smith/BUILD.yaml b/components/ee/agent-smith/BUILD.yaml index c15611a5d2c050..8740af0a2e0c06 100644 --- a/components/ee/agent-smith/BUILD.yaml +++ b/components/ee/agent-smith/BUILD.yaml @@ -4,6 +4,7 @@ packages: - name: app type: go srcs: + - "pkg/agent/testdata/**" - "**/*.go" - "go.mod" - "go.sum" diff --git a/components/ee/agent-smith/pkg/agent/agent.go b/components/ee/agent-smith/pkg/agent/agent.go index 20d4e4ee035a59..cfddea81a612de 100644 --- a/components/ee/agent-smith/pkg/agent/agent.go +++ b/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 } diff --git a/components/ee/agent-smith/pkg/agent/agent_test.go b/components/ee/agent-smith/pkg/agent/agent_test.go index abc83ec17772a3..58ea4e73194aff 100644 --- a/components/ee/agent-smith/pkg/agent/agent_test.go +++ b/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)) + } + }) + } +} diff --git a/components/ee/agent-smith/pkg/network/egress.go b/components/ee/agent-smith/pkg/agent/egress.go similarity index 86% rename from components/ee/agent-smith/pkg/network/egress.go rename to components/ee/agent-smith/pkg/agent/egress.go index 98fd2f0e013990..bd6c07f1546e62 100644 --- a/components/ee/agent-smith/pkg/network/egress.go +++ b/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 diff --git a/components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_excessive_egress.golden b/components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_excessive_egress.golden new file mode 100644 index 00000000000000..d1c4e0b77e6004 --- /dev/null +++ b/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" + } + } +} diff --git a/components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_no_infringement.golden b/components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_no_infringement.golden new file mode 100644 index 00000000000000..d1c4e0b77e6004 --- /dev/null +++ b/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" + } + } +} diff --git a/components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_very_excessive_egress.golden b/components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_very_excessive_egress.golden new file mode 100644 index 00000000000000..b4d32612430169 --- /dev/null +++ b/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" + } + } +} diff --git a/components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_very_excessive_egress_simple.golden b/components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_very_excessive_egress_simple.golden new file mode 100644 index 00000000000000..d1c4e0b77e6004 --- /dev/null +++ b/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" + } + } +} diff --git a/components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_zero_egress.golden b/components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_zero_egress.golden new file mode 100644 index 00000000000000..617daefadbcac8 --- /dev/null +++ b/components/ee/agent-smith/pkg/agent/testdata/agent_check_egress_zero_egress.golden @@ -0,0 +1,4 @@ +{ + "egressTraffic": { + } +}