Bound large HTTP Flow bodies by cgroup memory

pkg/protocols/http/body_admission.go

@@ -0,0 +1,98 @@
+package http
+
+import (
+	"context"
+	"sync"
+
+	"github.com/projectdiscovery/gologger"
+	"github.com/projectdiscovery/utils/memguardian"
+	"golang.org/x/sync/semaphore"
+)
+
+const (
+	largeBodyThreshold           = 1 << 20
+	largeBodyMemoryAmplification = 64
+)
+
+type bodyAdmission struct {
+	limit int64
+	slots *semaphore.Weighted
+}
+
+var (
+	sharedBodyAdmission     *bodyAdmission
+	sharedBodyAdmissionOnce sync.Once
+	sharedFlowAdmission     *bodyAdmission
+	sharedFlowAdmissionOnce sync.Once
+)
+
+func newBodyAdmission(totalMemory uint64) *bodyAdmission {
+	limit := int64(totalMemory / largeBodyMemoryAmplification)
+	if limit < largeBodyThreshold {
+		limit = largeBodyThreshold
+	}
+	return &bodyAdmission{limit: limit, slots: semaphore.NewWeighted(limit)}
+}
+
+func systemBodyAdmission() *bodyAdmission {
+	availableMemory := uint64(4 << 30)
+	info, err := memguardian.GetSysInfo()
+	if err != nil {
+		gologger.Warning().Msgf("could not read memory limit for large HTTP body admission: %s", err)
+	} else {
+		availableMemory = info.TotalRam()
+	}
+	if limit, limitErr := cgroupMemoryLimit("/proc/self/cgroup", "/sys/fs/cgroup"); limitErr != nil {
+		gologger.Warning().Msgf("could not read cgroup memory limit for large HTTP body admission: %s", limitErr)
+	} else if limit > 0 && limit < availableMemory {
+		availableMemory = limit
+	}
+	return newBodyAdmission(availableMemory)
+}
+
+func bodySize(req *generatedRequest) int64 {
+	if req.rawRequest != nil {
+		return int64(len(req.rawRequest.Data))
+	}
+	if req.request != nil {
+		return req.request.ContentLength
+	}
+	return 0
+}
+
+func (a *bodyAdmission) acquire(ctx context.Context, size int64) (func(), error) {
+	if size < largeBodyThreshold {
+		return func() {}, nil
+	}
+	weight := min(size, a.limit)
+	if err := a.slots.Acquire(ctx, weight); err != nil {
+		return nil, err
+	}
+	return func() { a.slots.Release(weight) }, nil
+}
+
+func admitLargeBody(ctx context.Context, req *generatedRequest) (func(), error) {
+	size := bodySize(req)
+	return admitBodySize(ctx, size)
+}
+
+func admitBodySize(ctx context.Context, size int64) (func(), error) {
+	if size < largeBodyThreshold {
+		return func() {}, nil
+	}
+	sharedBodyAdmissionOnce.Do(func() {
+		sharedBodyAdmission = systemBodyAdmission()
+	})
+	return sharedBodyAdmission.acquire(ctx, size)
+}
+
+// AdmitDeclaredBody bounds Flow executions before a declared large body is built.
+func AdmitDeclaredBody(ctx context.Context, size int64) (func(), error) {
+	if size < largeBodyThreshold {
+		return func() {}, nil
+	}
+	sharedFlowAdmissionOnce.Do(func() {
+		sharedFlowAdmission = systemBodyAdmission()
+	})
+	return sharedFlowAdmission.acquire(ctx, size)
+}

pkg/protocols/http/body_admission_memory.go

@@ -0,0 +1,73 @@
+package http
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+)
+
+func cgroupMemoryLimit(procPath, root string) (uint64, error) {
+	contents, err := os.ReadFile(procPath)
+	if errors.Is(err, os.ErrNotExist) {
+		return 0, nil
+	}
+	if err != nil {
+		return 0, err
+	}
+
+	var lowest uint64
+	for line := range strings.SplitSeq(string(contents), "\n") {
+		fields := strings.SplitN(line, ":", 3)
+		if len(fields) != 3 {
+			continue
+		}
+		controllerRoot := root
+		limitFile := "memory.max"
+		if fields[1] != "" {
+			if !strings.Contains(","+fields[1]+",", ",memory,") {
+				continue
+			}
+			controllerRoot = filepath.Join(root, "memory")
+			limitFile = "memory.limit_in_bytes"
+		}
+		limit, err := memoryLimitInHierarchy(controllerRoot, fields[2], limitFile)
+		if err != nil {
+			return 0, err
+		}
+		if limit > 0 && (lowest == 0 || limit < lowest) {
+			lowest = limit
+		}
+	}
+	return lowest, nil
+}
+
+func memoryLimitInHierarchy(root, cgroupPath, limitFile string) (uint64, error) {
+	root = filepath.Clean(root)
+	current := filepath.Join(root, strings.TrimPrefix(filepath.Clean("/"+cgroupPath), "/"))
+	var lowest uint64
+	for {
+		contents, err := os.ReadFile(filepath.Join(current, limitFile))
+		if err == nil {
+			value := strings.TrimSpace(string(contents))
+			if value != "max" {
+				limit, parseErr := strconv.ParseUint(value, 10, 64)
+				if parseErr != nil || limit == 0 {
+					return 0, fmt.Errorf("invalid cgroup memory limit %q in %s", value, current)
+				}
+				if lowest == 0 || limit < lowest {
+					lowest = limit
+				}
+			}
+		} else if !errors.Is(err, os.ErrNotExist) {
+			return 0, err
+		}
+		if current == root {
+			break
+		}
+		current = filepath.Dir(current)
+	}
+	return lowest, nil
+}

pkg/protocols/http/body_admission_test.go

@@ -0,0 +1,100 @@
+package http
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/projectdiscovery/utils/memguardian"
+	"github.com/stretchr/testify/require"
+)
+
+func TestLargeBodyAdmissionUsesMemoryBudget(t *testing.T) {
+	admission := newBodyAdmission(3 << 30)
+	require.EqualValues(t, 48<<20, admission.limit)
+
+	releases := make([]func(), 0, 2)
+	for range 2 {
+		release, err := admission.acquire(t.Context(), 20<<20)
+		require.NoError(t, err)
+		releases = append(releases, release)
+	}
+
+	ctx, cancel := context.WithTimeout(t.Context(), 20*time.Millisecond)
+	defer cancel()
+	_, err := admission.acquire(ctx, 20<<20)
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+
+	releases[0]()
+	release, err := admission.acquire(t.Context(), 20<<20)
+	require.NoError(t, err)
+	release()
+	releases[1]()
+}
+
+func TestSystemBodyAdmissionRespectsCgroupLimit(t *testing.T) {
+	raw, err := os.ReadFile("/sys/fs/cgroup/memory.max")
+	if err != nil {
+		t.Skipf("cgroup v2 memory limit unavailable: %v", err)
+	}
+	limit, err := strconv.ParseUint(strings.TrimSpace(string(raw)), 10, 64)
+	if err != nil {
+		t.Skipf("cgroup v2 has no finite memory limit: %v", err)
+	}
+	info, err := memguardian.GetSysInfo()
+	require.NoError(t, err)
+	expected := newBodyAdmission(min(limit, info.TotalRam())).limit
+	actual := systemBodyAdmission().limit
+	require.Equalf(t, expected, actual, "body admission must use the lower of host RAM and %d-byte cgroup limit", limit)
+}
+
+func TestCgroupMemoryLimit(t *testing.T) {
+	for _, test := range []struct {
+		name, membership string
+		files            map[string]string
+		want             uint64
+		wantError        bool
+	}{
+		{name: "v2 namespaced root", membership: "0::/\n", files: map[string]string{"memory.max": "536870912\n"}, want: 512 << 20},
+		{name: "v2 parent bound", membership: "0::/tenant/pod\n", files: map[string]string{"memory.max": "max\n", "tenant/memory.max": "1073741824\n", "tenant/pod/memory.max": "2147483648\n"}, want: 1 << 30},
+		{name: "v2 unlimited", membership: "0::/\n", files: map[string]string{"memory.max": "max\n"}},
+		{name: "v1 memory controller", membership: "5:cpu,memory:/pod\n4:cpu:/pod\n", files: map[string]string{"memory/memory.limit_in_bytes": "max\n", "memory/pod/memory.limit_in_bytes": "268435456\n"}, want: 256 << 20},
+		{name: "invalid bound", membership: "0::/\n", files: map[string]string{"memory.max": "invalid\n"}, wantError: true},
+		{name: "no membership", membership: ""},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			root := t.TempDir()
+			membershipPath := filepath.Join(root, "self-cgroup")
+			require.NoError(t, os.WriteFile(membershipPath, []byte(test.membership), 0o600))
+			for name, contents := range test.files {
+				path := filepath.Join(root, name)
+				require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
+				require.NoError(t, os.WriteFile(path, []byte(contents), 0o600))
+			}
+			limit, err := cgroupMemoryLimit(membershipPath, root)
+			if test.wantError {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			require.Equal(t, test.want, limit)
+		})
+	}
+}
+
+func TestSmallBodyBypassesLargeBodyAdmission(t *testing.T) {
+	admission := newBodyAdmission(3 << 30)
+	fullRelease, err := admission.acquire(t.Context(), admission.limit)
+	require.NoError(t, err)
+	defer fullRelease()
+
+	ctx, cancel := context.WithTimeout(t.Context(), 20*time.Millisecond)
+	defer cancel()
+	release, err := admission.acquire(ctx, largeBodyThreshold-1)
+	require.NoError(t, err)
+	release()
+}

pkg/protocols/http/request.go

@@ -1058,6 +1058,11 @@ func (request *Request) executeRequest(input *contextargs.Context, generatedRequ
 	if request.isUnresponsiveAddress(contextargs.GetCopyIfHostOutdated(input, generatedRequest.URL())) {
 		return fmt.Errorf("hostErrorsCache : host %s is unresponsive", input.MetaInput.Input)
 	}
+	releaseBody, err := admitLargeBody(input.Context(), generatedRequest)
+	if err != nil {
+		return err
+	}
+	defer releaseBody()
 
 	// wrap one more callback for validation and fixing event
 	callback := func(event *output.InternalWrappedEvent) {

pkg/tmplexec/flow/body_admission.go

@@ -0,0 +1,19 @@
+package flow
+
+import (
+	"fmt"
+
+	"github.com/spf13/cast"
+)
+
+func declaredRequestBodySize(metadata map[string]interface{}) (int64, bool, error) {
+	value, exists := metadata["max-request-body-bytes"]
+	if !exists {
+		return 0, false, nil
+	}
+	size, err := cast.ToInt64E(value)
+	if err != nil || size < 0 {
+		return 0, true, fmt.Errorf("invalid max-request-body-bytes metadata: %v", value)
+	}
+	return size, true, nil
+}

pkg/tmplexec/flow/body_admission_test.go

@@ -0,0 +1,23 @@
+package flow
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestDeclaredRequestBodySize(t *testing.T) {
+	size, exists, err := declaredRequestBodySize(nil)
+	require.NoError(t, err)
+	require.False(t, exists)
+	require.Zero(t, size)
+
+	size, exists, err = declaredRequestBodySize(map[string]interface{}{"max-request-body-bytes": 20 << 20})
+	require.NoError(t, err)
+	require.True(t, exists)
+	require.EqualValues(t, 20<<20, size)
+
+	_, exists, err = declaredRequestBodySize(map[string]interface{}{"max-request-body-bytes": -1})
+	require.True(t, exists)
+	require.Error(t, err)
+}

pkg/tmplexec/flow/flow_executor.go

@@ -12,6 +12,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/js/compiler"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/generators"
+	httpbody "github.com/projectdiscovery/nuclei/v3/pkg/protocols/http"
 	"github.com/projectdiscovery/nuclei/v3/pkg/scan"
 	templateTypes "github.com/projectdiscovery/nuclei/v3/pkg/templates/types"
 
@@ -220,6 +221,17 @@ func (f *FlowExecutor) ExecuteWithResults(ctx *scan.ScanContext) error {
 		return ctx.Context().Err()
 	default:
 	}
+	size, declared, err := declaredRequestBodySize(f.options.TemplateInfo.Metadata)
+	if err != nil {
+		return fmt.Errorf("template %s: %w", f.options.TemplateID, err)
+	}
+	if declared {
+		release, err := httpbody.AdmitDeclaredBody(ctx.Context(), size)
+		if err != nil {
+			return err
+		}
+		defer release()
+	}
 
 	f.ctx.Input = ctx.Input
 	// -----Load all types of variables-----