Implement ReadLimit to optionally limit reading from response body #71

options.go

@@ -15,6 +15,7 @@ const (
 	DefaultCrawlDelay         time.Duration             = 5 * time.Second
 	DefaultIdleTTL            time.Duration             = 10 * time.Second
 	DefaultNormalizationFlags purell.NormalizationFlags = purell.FlagsAllGreedy
+	DefaultReadLimit          int64                     = -1
 )
 
 // Options contains the configuration for a Crawler to customize the
@@ -50,6 +51,9 @@ type Options struct {
 	// further by implementing the ComputeDelay extender function.
 	CrawlDelay time.Duration
 
+	// Max count of bytes to read from the response body. -1 for no limit.
+	ReadLimit int64
+
 	// WorkerIdleTTL is the idle time-to-live allowed for a worker
 	// before it is cleared (its goroutine terminated). The crawl
 	// delay is not part of idle time, this is specifically the time
@@ -90,6 +94,7 @@ func NewOptions(ext Extender) *Options {
 		DefaultEnqueueChanBuffer,
 		DefaultHostBufferFactor,
 		DefaultCrawlDelay,
+		DefaultReadLimit,
 		DefaultIdleTTL,
 		true,
 		false,

tblrun_test.go

@@ -78,6 +78,7 @@ func runTestCase(t *testing.T, tc *testCase) {
 			tc.opts.LogFlags = LogAll
 		}
 		tc.opts.Extender = spy
+		tc.opts.ReadLimit = -1
 		c := NewCrawlerWithOptions(tc.opts)
 		if tc.funcs != nil {
 			for emk, f := range tc.funcs {

worker.go

@@ -329,7 +329,11 @@ func (w *worker) visitURL(ctx *URLContext, res *http.Response) interface{} {
 	var doLinks bool
 
 	// Load a goquery document and call the visitor function
-	if bd, e := ioutil.ReadAll(res.Body); e != nil {
+	r := io.Reader(res.Body)
+	if w.opts.ReadLimit != -1 {
+		r = io.LimitReader(r, w.opts.ReadLimit)
+	}
+	if bd, e := ioutil.ReadAll(r); e != nil {
 		w.opts.Extender.Error(newCrawlError(ctx, e, CekReadBody))
 		w.logFunc(LogError, "ERROR reading body %s: %s", ctx.url, e)
 	} else {