Put envd, ptys, socats, and commands into their own cgroups (#1580)

Author: djeebus

packages/envd/internal/port/forward.go

@@ -14,6 +14,8 @@ import (
 	"syscall"
 
 	"github.com/rs/zerolog"
+
+	"github.com/e2b-dev/infra/packages/envd/internal/services/cgroups"
 )
 
 type PortState string
@@ -36,7 +38,8 @@ type PortToForward struct {
 }
 
 type Forwarder struct {
-	logger *zerolog.Logger
+	logger        *zerolog.Logger
+	cgroupManager cgroups.Manager
 	// Map of ports that are being currently forwarded.
 	ports             map[string]*PortToForward
 	scannerSubscriber *ScannerSubscriber
@@ -46,6 +49,7 @@ type Forwarder struct {
 func NewForwarder(
 	logger *zerolog.Logger,
 	scanner *Scanner,
+	cgroupManager cgroups.Manager,
 ) *Forwarder {
 	scannerSub := scanner.AddSubscriber(
 		logger,
@@ -62,6 +66,7 @@ func NewForwarder(
 		sourceIP:          defaultGatewayIP,
 		ports:             make(map[string]*PortToForward),
 		scannerSubscriber: scannerSub,
+		cgroupManager:     cgroupManager,
 	}
 }
 
@@ -135,8 +140,13 @@ func (f *Forwarder) startPortForwarding(ctx context.Context, p *PortToForward) {
 		fmt.Sprintf("TCP4-LISTEN:%v,bind=%s,reuseaddr,fork", p.port, f.sourceIP.To4()),
 		fmt.Sprintf("TCP%d:localhost:%v", p.family, p.port),
 	)
+
+	cgroupFD, ok := f.cgroupManager.GetFileDescriptor(cgroups.ProcessTypeSocat)
+
 	cmd.SysProcAttr = &syscall.SysProcAttr{
-		Setpgid: true,
+		Setpgid:     true,
+		CgroupFD:    cgroupFD,
+		UseCgroupFD: ok,
 	}
 
 	f.logger.Debug().

packages/envd/internal/services/cgroups/cgroup2.go

@@ -0,0 +1,127 @@
+package cgroups
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"golang.org/x/sys/unix"
+)
+
+type Cgroup2Manager struct {
+	cgroupFDs map[ProcessType]int
+}
+
+var _ Manager = (*Cgroup2Manager)(nil)
+
+type cgroup2Config struct {
+	rootPath     string
+	processTypes map[ProcessType]Cgroup2Config
+}
+
+type Cgroup2ManagerOption func(*cgroup2Config)
+
+func WithCgroup2RootSysFSPath(path string) Cgroup2ManagerOption {
+	return func(config *cgroup2Config) {
+		config.rootPath = path
+	}
+}
+
+func WithCgroup2ProcessType(processType ProcessType, path string, properties map[string]string) Cgroup2ManagerOption {
+	return func(config *cgroup2Config) {
+		if config.processTypes == nil {
+			config.processTypes = make(map[ProcessType]Cgroup2Config)
+		}
+		config.processTypes[processType] = Cgroup2Config{Path: path, Properties: properties}
+	}
+}
+
+type Cgroup2Config struct {
+	Path       string
+	Properties map[string]string
+}
+
+func NewCgroup2Manager(opts ...Cgroup2ManagerOption) (*Cgroup2Manager, error) {
+	config := cgroup2Config{
+		rootPath: "/sys/fs/cgroup",
+	}
+
+	for _, opt := range opts {
+		opt(&config)
+	}
+
+	cgroupFDs, err := createCgroups(config)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create cgroups: %w", err)
+	}
+
+	return &Cgroup2Manager{cgroupFDs: cgroupFDs}, nil
+}
+
+func createCgroups(configs cgroup2Config) (map[ProcessType]int, error) {
+	var (
+		results = make(map[ProcessType]int)
+		errs    []error
+	)
+
+	for procType, config := range configs.processTypes {
+		fullPath := filepath.Join(configs.rootPath, config.Path)
+		fd, err := createCgroup(fullPath, config.Properties)
+		if err != nil {
+			errs = append(errs, fmt.Errorf("failed to create %s cgroup: %w", procType, err))
+
+			continue
+		}
+		results[procType] = fd
+	}
+
+	if len(errs) > 0 {
+		for procType, fd := range results {
+			err := unix.Close(fd)
+			if err != nil {
+				errs = append(errs, fmt.Errorf("failed to close cgroup fd for %s: %w", procType, err))
+			}
+		}
+
+		return nil, errors.Join(errs...)
+	}
+
+	return results, nil
+}
+
+func createCgroup(fullPath string, properties map[string]string) (int, error) {
+	if err := os.MkdirAll(fullPath, 0o755); err != nil {
+		return -1, fmt.Errorf("failed to create cgroup root: %w", err)
+	}
+
+	var errs []error
+	for name, value := range properties {
+		if err := os.WriteFile(filepath.Join(fullPath, name), []byte(value), 0o644); err != nil {
+			errs = append(errs, fmt.Errorf("failed to write cgroup property: %w", err))
+		}
+	}
+	if len(errs) > 0 {
+		return -1, errors.Join(errs...)
+	}
+
+	return unix.Open(fullPath, unix.O_RDONLY, 0)
+}
+
+func (c Cgroup2Manager) GetFileDescriptor(procType ProcessType) (int, bool) {
+	fd, ok := c.cgroupFDs[procType]
+
+	return fd, ok
+}
+
+func (c Cgroup2Manager) Close() error {
+	var errs []error
+	for procType, fd := range c.cgroupFDs {
+		if err := unix.Close(fd); err != nil {
+			errs = append(errs, fmt.Errorf("failed to close cgroup fd for %s: %w", procType, err))
+		}
+		delete(c.cgroupFDs, procType)
+	}
+
+	return errors.Join(errs...)
+}

packages/envd/internal/services/cgroups/cgroup2_test.go

@@ -0,0 +1,185 @@
+package cgroups
+
+import (
+	"context"
+	"fmt"
+	"math/rand"
+	"os"
+	"os/exec"
+	"strconv"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+const (
+	oneByte  = 1
+	kilobyte = 1024 * oneByte
+	megabyte = 1024 * kilobyte
+)
+
+func TestCgroupRoundTrip(t *testing.T) {
+	t.Parallel()
+
+	if os.Geteuid() != 0 {
+		t.Skip("must run as root")
+
+		return
+	}
+
+	maxTimeout := time.Second * 5
+
+	t.Run("process does not die without cgroups", func(t *testing.T) {
+		t.Parallel()
+
+		// create manager
+		m, err := NewCgroup2Manager()
+		require.NoError(t, err)
+
+		// create new child process
+		cmd := startProcess(t, m, "not-a-real-one")
+
+		// wait for child process to die
+		err = waitForProcess(t, cmd, maxTimeout)
+
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
+	t.Run("process dies with cgroups", func(t *testing.T) {
+		t.Parallel()
+
+		cgroupPath := createCgroupPath(t, "real-one")
+
+		// create manager
+		m, err := NewCgroup2Manager(
+			WithCgroup2ProcessType(ProcessTypePTY, cgroupPath, map[string]string{
+				"memory.max": strconv.Itoa(1 * megabyte),
+			}),
+		)
+		require.NoError(t, err)
+
+		t.Cleanup(func() {
+			err := m.Close()
+			assert.NoError(t, err)
+		})
+
+		// create new child process
+		cmd := startProcess(t, m, ProcessTypePTY)
+
+		// wait for child process to die
+		err = waitForProcess(t, cmd, maxTimeout)
+
+		// verify process exited correctly
+		var exitErr *exec.ExitError
+		require.ErrorAs(t, err, &exitErr)
+		assert.Equal(t, "signal: killed", exitErr.Error())
+		assert.False(t, exitErr.Exited())
+		assert.False(t, exitErr.Success())
+		assert.Equal(t, -1, exitErr.ExitCode())
+
+		// dig a little deeper
+		ws, ok := exitErr.Sys().(syscall.WaitStatus)
+		require.True(t, ok)
+		assert.Equal(t, syscall.SIGKILL, ws.Signal())
+		assert.True(t, ws.Signaled())
+		assert.False(t, ws.Stopped())
+		assert.False(t, ws.Continued())
+		assert.False(t, ws.CoreDump())
+		assert.False(t, ws.Exited())
+		assert.Equal(t, -1, ws.ExitStatus())
+	})
+
+	t.Run("process cannot be spawned because memory limit is too low", func(t *testing.T) {
+		t.Parallel()
+
+		cgroupPath := createCgroupPath(t, "real-one")
+
+		// create manager
+		m, err := NewCgroup2Manager(
+			WithCgroup2ProcessType(ProcessTypeSocat, cgroupPath, map[string]string{
+				"memory.max": strconv.Itoa(1 * kilobyte),
+			}),
+		)
+		require.NoError(t, err)
+
+		t.Cleanup(func() {
+			err := m.Close()
+			assert.NoError(t, err)
+		})
+
+		// create new child process
+		cmd := startProcess(t, m, ProcessTypeSocat)
+
+		// wait for child process to die
+		err = waitForProcess(t, cmd, maxTimeout)
+
+		// verify process exited correctly
+		var exitErr *exec.ExitError
+		require.ErrorAs(t, err, &exitErr)
+		assert.Equal(t, "exit status 253", exitErr.Error())
+		assert.True(t, exitErr.Exited())
+		assert.False(t, exitErr.Success())
+		assert.Equal(t, 253, exitErr.ExitCode())
+
+		// dig a little deeper
+		ws, ok := exitErr.Sys().(syscall.WaitStatus)
+		require.True(t, ok)
+		assert.Equal(t, syscall.Signal(-1), ws.Signal())
+		assert.False(t, ws.Signaled())
+		assert.False(t, ws.Stopped())
+		assert.False(t, ws.Continued())
+		assert.False(t, ws.CoreDump())
+		assert.True(t, ws.Exited())
+		assert.Equal(t, 253, ws.ExitStatus())
+	})
+}
+
+func createCgroupPath(t *testing.T, s string) string {
+	t.Helper()
+
+	randPart := rand.Int()
+
+	return fmt.Sprintf("envd-test-%s-%d", s, randPart)
+}
+
+func startProcess(t *testing.T, m *Cgroup2Manager, pt ProcessType) *exec.Cmd {
+	t.Helper()
+
+	cmdName, args := "bash", []string{"-c", `sleep 1 && tail /dev/zero`}
+	cmd := exec.CommandContext(t.Context(), cmdName, args...)
+
+	fd, ok := m.GetFileDescriptor(pt)
+	cmd.SysProcAttr = &syscall.SysProcAttr{
+		UseCgroupFD: ok,
+		CgroupFD:    fd,
+	}
+
+	err := cmd.Start()
+	require.NoError(t, err)
+
+	return cmd
+}
+
+func waitForProcess(t *testing.T, cmd *exec.Cmd, timeout time.Duration) error {
+	t.Helper()
+
+	done := make(chan error, 1)
+
+	go func() {
+		defer close(done)
+		done <- cmd.Wait()
+	}()
+
+	ctx, cancel := context.WithTimeout(t.Context(), timeout)
+	t.Cleanup(cancel)
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case err := <-done:
+		return err
+	}
+}

packages/envd/internal/services/cgroups/iface.go

@@ -0,0 +1,14 @@
+package cgroups
+
+type ProcessType string
+
+const (
+	ProcessTypePTY   ProcessType = "pty"
+	ProcessTypeUser  ProcessType = "user"
+	ProcessTypeSocat ProcessType = "socat"
+)
+
+type Manager interface {
+	GetFileDescriptor(procType ProcessType) (int, bool)
+	Close() error
+}

packages/envd/internal/services/cgroups/noop.go

@@ -0,0 +1,17 @@
+package cgroups
+
+type NoopManager struct{}
+
+var _ Manager = (*NoopManager)(nil)
+
+func NewNoopManager() *NoopManager {
+	return &NoopManager{}
+}
+
+func (n NoopManager) GetFileDescriptor(ProcessType) (int, bool) {
+	return 0, false
+}
+
+func (n NoopManager) Close() error {
+	return nil
+}