From 7a010f541616a4cc0822250754ef4786b0674da2 Mon Sep 17 00:00:00 2001
From: Derek Parker <parkerderek86@gmail.com>
Date: Tue, 2 May 2023 14:43:57 -0700
Subject: [PATCH] pkg/proc: remove memlock limit for ebpf trace backend

Fixes #3283
---
 pkg/proc/internal/ebpf/helpers.go             |   4 +
 .../github.com/cilium/ebpf/rlimit/rlimit.go   | 113 ++++++++++++++++++
 vendor/modules.txt                            |   3 +-
 3 files changed, 118 insertions(+), 2 deletions(-)
 create mode 100644 vendor/github.com/cilium/ebpf/rlimit/rlimit.go

diff --git a/pkg/proc/internal/ebpf/helpers.go b/pkg/proc/internal/ebpf/helpers.go
index d9b5e7dfc0..b856cef1cd 100644
--- a/pkg/proc/internal/ebpf/helpers.go
+++ b/pkg/proc/internal/ebpf/helpers.go
@@ -18,6 +18,7 @@ import (
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/link"
 	"github.com/cilium/ebpf/ringbuf"
+	"github.com/cilium/ebpf/rlimit"
 )
 
 //lint:file-ignore U1000 some fields are used by the C program
@@ -109,6 +110,9 @@ func LoadEBPFTracingProgram(path string) (*EBPFContext, error) {
 		objs traceObjects
 	)
 
+	if err = rlimit.RemoveMemlock(); err != nil {
+		return nil, err
+	}
 	ctx.executable, err = link.OpenExecutable(path)
 	if err != nil {
 		return nil, err
diff --git a/vendor/github.com/cilium/ebpf/rlimit/rlimit.go b/vendor/github.com/cilium/ebpf/rlimit/rlimit.go
new file mode 100644
index 0000000000..22fc8177f2
--- /dev/null
+++ b/vendor/github.com/cilium/ebpf/rlimit/rlimit.go
@@ -0,0 +1,113 @@
+// Package rlimit allows raising RLIMIT_MEMLOCK if necessary for the use of BPF.
+package rlimit
+
+import (
+	"errors"
+	"fmt"
+	"sync"
+
+	"github.com/cilium/ebpf/internal"
+	"github.com/cilium/ebpf/internal/unix"
+)
+
+var (
+	unsupportedMemcgAccounting = &internal.UnsupportedFeatureError{
+		MinimumVersion: internal.Version{5, 11, 0},
+		Name:           "memcg-based accounting for BPF memory",
+	}
+	haveMemcgAccounting error
+
+	rlimitMu sync.Mutex
+)
+
+func init() {
+	// We have to run this feature test at init, since it relies on changing
+	// RLIMIT_MEMLOCK. Doing so is not safe in a concurrent program. Instead,
+	// we rely on the initialization order guaranteed by the Go runtime to
+	// execute the test in a safe environment:
+	//
+	//    the invocation of init functions happens in a single goroutine,
+	//    sequentially, one package at a time.
+	//
+	// This is also the reason why RemoveMemlock is in its own package:
+	// we only want to run the initializer if RemoveMemlock is called
+	// from somewhere.
+	haveMemcgAccounting = detectMemcgAccounting()
+}
+
+func detectMemcgAccounting() error {
+	// Reduce the limit to zero and store the previous limit. This should always succeed.
+	var oldLimit unix.Rlimit
+	zeroLimit := unix.Rlimit{Cur: 0, Max: oldLimit.Max}
+	if err := unix.Prlimit(0, unix.RLIMIT_MEMLOCK, &zeroLimit, &oldLimit); err != nil {
+		return fmt.Errorf("lowering memlock rlimit: %s", err)
+	}
+
+	attr := internal.BPFMapCreateAttr{
+		MapType:    2, /* Array */
+		KeySize:    4,
+		ValueSize:  4,
+		MaxEntries: 1,
+	}
+
+	// Creating a map allocates shared (and locked) memory that counts against
+	// the rlimit on pre-5.11 kernels, but against the memory cgroup budget on
+	// kernels 5.11 and over. If this call succeeds with the process' memlock
+	// rlimit set to 0, we can reasonably assume memcg accounting is supported.
+	fd, mapErr := internal.BPFMapCreate(&attr)
+
+	// Restore old limits regardless of what happened.
+	if err := unix.Prlimit(0, unix.RLIMIT_MEMLOCK, &oldLimit, nil); err != nil {
+		return fmt.Errorf("restoring old memlock rlimit: %s", err)
+	}
+
+	// Map creation successful, memcg accounting supported.
+	if mapErr == nil {
+		fd.Close()
+		return nil
+	}
+
+	// EPERM shows up when map creation would exceed the memory budget.
+	if errors.Is(mapErr, unix.EPERM) {
+		return unsupportedMemcgAccounting
+	}
+
+	// This shouldn't happen really.
+	return fmt.Errorf("unexpected error detecting memory cgroup accounting: %s", mapErr)
+}
+
+// RemoveMemlock removes the limit on the amount of memory the current
+// process can lock into RAM, if necessary.
+//
+// This is not required to load eBPF resources on kernel versions 5.11+
+// due to the introduction of cgroup-based memory accounting. On such kernels
+// the function is a no-op.
+//
+// Since the function may change global per-process limits it should be invoked
+// at program start up, in main() or init().
+//
+// This function exists as a convenience and should only be used when
+// permanently raising RLIMIT_MEMLOCK to infinite is appropriate. Consider
+// invoking prlimit(2) directly with a more reasonable limit if desired.
+//
+// Requires CAP_SYS_RESOURCE on kernels < 5.11.
+func RemoveMemlock() error {
+	if haveMemcgAccounting == nil {
+		return nil
+	}
+
+	if !errors.Is(haveMemcgAccounting, unsupportedMemcgAccounting) {
+		return haveMemcgAccounting
+	}
+
+	rlimitMu.Lock()
+	defer rlimitMu.Unlock()
+
+	// pid 0 affects the current process. Requires CAP_SYS_RESOURCE.
+	newLimit := unix.Rlimit{Cur: unix.RLIM_INFINITY, Max: unix.RLIM_INFINITY}
+	if err := unix.Prlimit(0, unix.RLIMIT_MEMLOCK, &newLimit, nil); err != nil {
+		return fmt.Errorf("failed to set memlock rlimit: %w", err)
+	}
+
+	return nil
+}
diff --git a/vendor/modules.txt b/vendor/modules.txt
index ac1d2a3ab4..161a09fe33 100644
--- a/vendor/modules.txt
+++ b/vendor/modules.txt
@@ -7,6 +7,7 @@ github.com/cilium/ebpf/internal/btf
 github.com/cilium/ebpf/internal/unix
 github.com/cilium/ebpf/link
 github.com/cilium/ebpf/ringbuf
+github.com/cilium/ebpf/rlimit
 # github.com/cosiner/argv v0.1.0
 ## explicit
 github.com/cosiner/argv
@@ -91,8 +92,6 @@ golang.org/x/tools/internal/gocommand
 golang.org/x/tools/internal/packagesinternal
 golang.org/x/tools/internal/typeparams
 golang.org/x/tools/internal/typesinternal
-# golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1
-## explicit
 # gopkg.in/yaml.v2 v2.4.0
 ## explicit
 gopkg.in/yaml.v2