Fixing memory leak on large file dump via max-dump-body-size

internal/runner/options.go

@@ -12,25 +12,26 @@ import (
 
 // Options of the tool
 type Options struct {
-	ListenAddress  string
-	Folder         string
-	BasicAuth      string
-	username       string
-	password       string
-	Realm          string
-	TLSCertificate string
-	TLSKey         string
-	TLSDomain      string
-	HTTPS          bool
-	Verbose        bool
-	EnableUpload   bool
-	EnableTCP      bool
-	RulesFile      string
-	TCPWithTLS     bool
-	Version        bool
-	Silent         bool
-	Sandbox        bool
-	MaxFileSize    int
+	ListenAddress   string
+	Folder          string
+	BasicAuth       string
+	username        string
+	password        string
+	Realm           string
+	TLSCertificate  string
+	TLSKey          string
+	TLSDomain       string
+	HTTPS           bool
+	Verbose         bool
+	EnableUpload    bool
+	EnableTCP       bool
+	RulesFile       string
+	TCPWithTLS      bool
+	Version         bool
+	Silent          bool
+	Sandbox         bool
+	MaxFileSize     int
+	MaxDumpBodySize int
 }
 
 // ParseOptions parses the command line options for application
@@ -57,6 +58,7 @@ func ParseOptions() *Options {
 	flag.BoolVar(&options.Silent, "silent", false, "Show only results in the output")
 	flag.BoolVar(&options.Sandbox, "sandbox", false, "Enable sandbox mode")
 	flag.IntVar(&options.MaxFileSize, "max-file-size", 50, "Max Upload File Size")
+	flag.IntVar(&options.MaxDumpBodySize, "max-dump-body-size", -1, "Max Dump Body Size")
 
 	flag.Parse()
 

internal/runner/runner.go

@@ -5,6 +5,7 @@ import (
 	"github.com/projectdiscovery/simplehttpserver/pkg/binder"
 	"github.com/projectdiscovery/simplehttpserver/pkg/httpserver"
 	"github.com/projectdiscovery/simplehttpserver/pkg/tcpserver"
+	"github.com/projectdiscovery/simplehttpserver/pkg/unit"
 )
 
 // Runner is a client for running the enumeration process.
@@ -59,6 +60,7 @@ func New(options *Options) (*Runner, error) {
 		Verbose:           r.options.Verbose,
 		Sandbox:           r.options.Sandbox,
 		MaxFileSize:       r.options.MaxFileSize,
+		MaxDumpBodySize:   unit.ToMb(r.options.MaxDumpBodySize),
 	})
 	if err != nil {
 		return nil, err

pkg/httpserver/httpserver.go

@@ -24,6 +24,7 @@ type Options struct {
 	Verbose           bool
 	Sandbox           bool
 	MaxFileSize       int // 50Mb
+	MaxDumpBodySize   int64
 }
 
 // HTTPServer instance

pkg/httpserver/loglayer.go

@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 
 	"github.com/projectdiscovery/gologger"
+	"github.com/projectdiscovery/simplehttpserver/pkg/unit"
 )
 
 // Convenience globals
@@ -17,10 +18,19 @@ var (
 	EnableVerbose bool
 )
 
+func (t *HTTPServer) shouldDumpBody(bodysize int64) bool {
+	return t.options.MaxDumpBodySize > 0 && bodysize > t.options.MaxDumpBodySize
+}
+
 func (t *HTTPServer) loglayer(handler http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		fullRequest, _ := httputil.DumpRequest(r, true)
-		lrw := newLoggingResponseWriter(w)
+		var fullRequest []byte
+		if t.shouldDumpBody(r.ContentLength) {
+			fullRequest, _ = httputil.DumpRequest(r, false)
+		} else {
+			fullRequest, _ = httputil.DumpRequest(r, true)
+		}
+		lrw := newLoggingResponseWriter(w, t.options.MaxDumpBodySize)
 		handler.ServeHTTP(lrw, r)
 
 		// Handles file write if enabled
@@ -52,7 +62,7 @@ func (t *HTTPServer) loglayer(handler http.Handler) http.Handler {
 				err  error
 			)
 			if t.options.Sandbox {
-				maxFileSize := toMb(t.options.MaxFileSize)
+				maxFileSize := unit.ToMb(t.options.MaxFileSize)
 				// check header content length
 				if r.ContentLength > maxFileSize {
 					gologger.Print().Msg("request too large")
@@ -81,24 +91,29 @@ func (t *HTTPServer) loglayer(handler http.Handler) http.Handler {
 			lrw.Header().Write(headers) //nolint
 			gologger.Print().Msgf("\nRemote Address: %s\n%s\n%s %d %s\n%s\n%s\n", r.RemoteAddr, string(fullRequest), r.Proto, lrw.statusCode, http.StatusText(lrw.statusCode), headers.String(), string(lrw.Data))
 		} else {
-			gologger.Print().Msgf("%s \"%s %s %s\" %d %d", r.RemoteAddr, r.Method, r.URL, r.Proto, lrw.statusCode, len(lrw.Data))
+			gologger.Print().Msgf("%s \"%s %s %s\" %d %d", r.RemoteAddr, r.Method, r.URL, r.Proto, lrw.statusCode, lrw.Size)
 		}
 	})
 }
 
 type loggingResponseWriter struct {
 	http.ResponseWriter
-	statusCode int
-	Data       []byte
+	statusCode  int
+	Data        []byte
+	Size        int
+	MaxDumpSize int64
 }
 
-func newLoggingResponseWriter(w http.ResponseWriter) *loggingResponseWriter {
-	return &loggingResponseWriter{w, http.StatusOK, []byte{}}
+func newLoggingResponseWriter(w http.ResponseWriter, maxSize int64) *loggingResponseWriter {
+	return &loggingResponseWriter{w, http.StatusOK, []byte{}, 0, maxSize}
 }
 
 // Write the data
 func (lrw *loggingResponseWriter) Write(data []byte) (int, error) {
-	lrw.Data = append(lrw.Data, data...)
+	if len(lrw.Data) < int(lrw.MaxDumpSize) {
+		lrw.Data = append(lrw.Data, data...)
+	}
+	lrw.Size += len(data)
 	return lrw.ResponseWriter.Write(data)
 }
 

pkg/unit/unit.go

@@ -0,0 +1,6 @@
+package unit
+
+// ToMb converts bytes to megabytes
+func ToMb(n int) int64 {
+	return int64(n) * 1024 * 1024
+}