Merge pull request #2343 from smallseacreature/fix-endless-stream-oom

Fix endless stream oom

README.md

@@ -307,4 +307,4 @@ Probing feature is inspired by [@tomnomnom/httprobe](https://github.com/tomnomno
 
 <a href="https://discord.gg/projectdiscovery"><img src="https://raw.githubusercontent.com/projectdiscovery/nuclei-burp-plugin/main/static/join-discord.png" width="300" alt="Join Discord"></a>
 
-</div>
+</div>

common/httpx/httpx.go

@@ -235,6 +235,14 @@ get_response:
 
 	resp.Headers = httpresp.Header.Clone()
 
+	if h.Options.MaxResponseBodySizeToRead > 0 {
+		httpresp.Body = io.NopCloser(io.LimitReader(httpresp.Body, h.Options.MaxResponseBodySizeToRead))
+		defer func() {
+			_, _ = io.Copy(io.Discard, httpresp.Body)
+			_ = httpresp.Body.Close()
+		}()
+	}
+
 	// httputil.DumpResponse does not handle websockets
 	headers, rawResp, err := pdhttputil.DumpResponseHeadersAndRaw(httpresp)
 	if err != nil {

common/httpx/option.go

@@ -5,10 +5,19 @@ import (
 	"strings"
 	"time"
 
+	"github.com/dustin/go-humanize"
 	"github.com/projectdiscovery/cdncheck"
 	"github.com/projectdiscovery/networkpolicy"
 )
 
+// DefaultMaxResponseBodySize is the default maximum response body size
+var DefaultMaxResponseBodySize int64
+
+func init() {
+	maxResponseBodySize, _ := humanize.ParseBytes("512Mb")
+	DefaultMaxResponseBodySize = int64(maxResponseBodySize)
+}
+
 // Options contains configuration options for the client
 type Options struct {
 	RandomAgent      bool
@@ -66,7 +75,7 @@ var DefaultOptions = Options{
 	Unsafe:                    false,
 	CdnCheck:                  "true",
 	ExcludeCdn:                false,
-	MaxResponseBodySizeToRead: 1024 * 1024 * 10,
+	MaxResponseBodySizeToRead: DefaultMaxResponseBodySize,
 	// VHOSTs options
 	VHostIgnoreStatusCode:    false,
 	VHostIgnoreContentLength: true,

common/stringz/stringz.go

@@ -85,9 +85,10 @@ func AddURLDefaultPort(rawURL string) string {
 	}
 	// Force default port to be added if not present
 	if u.Port() == "" {
-		if u.Scheme == urlutil.HTTP {
+		switch u.Scheme {
+		case urlutil.HTTP:
 			u.UpdatePort("80")
-		} else if u.Scheme == urlutil.HTTPS {
+		case urlutil.HTTPS:
 			u.UpdatePort("443")
 		}
 	}

go.mod

@@ -51,6 +51,7 @@ require (
 
 require (
 	github.com/JohannesKaufmann/html-to-markdown/v2 v2.5.0
+	github.com/dustin/go-humanize v1.0.1
 	github.com/go-viper/mapstructure/v2 v2.4.0
 	github.com/gocarina/gocsv v0.0.0-20240520201108-78e41c74b4b1
 	github.com/weppos/publicsuffix-go v0.50.1

go.sum

@@ -116,6 +116,8 @@ github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDD
 github.com/dsnet/compress v0.0.2-0.20230904184137-39efe44ab707 h1:2tV76y6Q9BB+NEBasnqvs7e49aEBFI8ejC89PSnWH+4=
 github.com/dsnet/compress v0.0.2-0.20230904184137-39efe44ab707/go.mod h1:qssHWj60/X5sZFNxpG4HBPDHVqxNm4DfnCKgrbZOT+s=
 github.com/dsnet/golib v0.0.0-20171103203638-1ea166775780/go.mod h1:Lj+Z9rebOhdfkVLjJ8T6VcRQv3SXugXy999NBtR9aFY=
+github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
+github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
 github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=

runner/options.go

@@ -2,7 +2,6 @@ package runner
 
 import (
 	"fmt"
-	"math"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -23,7 +22,7 @@ import (
 	"github.com/projectdiscovery/httpx/common/customlist"
 	customport "github.com/projectdiscovery/httpx/common/customports"
 	fileutilz "github.com/projectdiscovery/httpx/common/fileutil"
-	"github.com/projectdiscovery/httpx/common/httpx"
+	httpxcommon "github.com/projectdiscovery/httpx/common/httpx"
 	"github.com/projectdiscovery/httpx/common/stringz"
 	"github.com/projectdiscovery/networkpolicy"
 	pdcpauth "github.com/projectdiscovery/utils/auth/pdcp"
@@ -541,8 +540,8 @@ func ParseOptions() *Options {
 		flagSet.IntVar(&options.Retries, "retries", 0, "number of retries"),
 		flagSet.IntVar(&options.Timeout, "timeout", 10, "timeout in seconds"),
 		flagSet.DurationVar(&options.Delay, "delay", -1, "duration between each http request (eg: 200ms, 1s)"),
-		flagSet.IntVarP(&options.MaxResponseBodySizeToSave, "response-size-to-save", "rsts", math.MaxInt32, "max response size to save in bytes"),
-		flagSet.IntVarP(&options.MaxResponseBodySizeToRead, "response-size-to-read", "rstr", math.MaxInt32, "max response size to read in bytes"),
+		flagSet.IntVarP(&options.MaxResponseBodySizeToSave, "response-size-to-save", "rsts", int(httpxcommon.DefaultMaxResponseBodySize), "max response size to save in bytes"),
+		flagSet.IntVarP(&options.MaxResponseBodySizeToRead, "response-size-to-read", "rstr", int(httpxcommon.DefaultMaxResponseBodySize), "max response size to read in bytes"),
 	)
 
 	flagSet.CreateGroup("cloud", "Cloud",
@@ -772,7 +771,7 @@ func (options *Options) ValidateOptions() error {
 		options.OutputCDN = "true"
 	}
 
-	if !stringsutil.EqualFoldAny(options.Protocol, string(httpx.UNKNOWN), string(httpx.HTTP11)) {
+	if !stringsutil.EqualFoldAny(options.Protocol, string(httpxcommon.UNKNOWN), string(httpxcommon.HTTP11)) {
 		return fmt.Errorf("invalid protocol: %s", options.Protocol)
 	}