Limiting body reading support implemented.

This commit is contained in:
Musab Gültekin 2019-06-11 16:19:30 +03:00
parent 3790295658
commit b8305d5e1a
2 changed files with 12 additions and 2 deletions

View File

@ -6,6 +6,7 @@ import (
"github.com/fpfeng/httpcache" "github.com/fpfeng/httpcache"
"github.com/geziyor/geziyor/exporter" "github.com/geziyor/geziyor/exporter"
"golang.org/x/net/html/charset" "golang.org/x/net/html/charset"
"io"
"io/ioutil" "io/ioutil"
"log" "log"
"math/rand" "math/rand"
@ -70,6 +71,9 @@ func NewGeziyor(opt Options) *Geziyor {
if len(opt.Exporters) == 0 { if len(opt.Exporters) == 0 {
geziyor.opt.Exporters = []Exporter{exporter.JSONExporter{}} geziyor.opt.Exporters = []Exporter{exporter.JSONExporter{}}
} }
if opt.MaxBodySize == 0 {
geziyor.opt.MaxBodySize = 1024 * 1024 * 1024 // 1GB
}
return geziyor return geziyor
} }
@ -148,8 +152,11 @@ func (g *Geziyor) Do(req *http.Request, callback func(resp *Response)) {
return return
} }
// Limit response body reading
bodyReader := io.LimitReader(resp.Body, g.opt.MaxBodySize)
// Start reading body and determine encoding // Start reading body and determine encoding
reader, err := charset.NewReader(resp.Body, resp.Header.Get("Content-Type")) bodyReader, err = charset.NewReader(bodyReader, resp.Header.Get("Content-Type"))
if err != nil { if err != nil {
log.Printf("Determine encoding error: %v\n", err) log.Printf("Determine encoding error: %v\n", err)
g.releaseSem(req) g.releaseSem(req)
@ -157,7 +164,7 @@ func (g *Geziyor) Do(req *http.Request, callback func(resp *Response)) {
} }
// Continue reading body // Continue reading body
body, err := ioutil.ReadAll(reader) body, err := ioutil.ReadAll(bodyReader)
if err != nil { if err != nil {
log.Printf("Reading Body error: %v\n", err) log.Printf("Reading Body error: %v\n", err)
g.releaseSem(req) g.releaseSem(req)

View File

@ -43,4 +43,7 @@ type Options struct {
// For extracting data // For extracting data
Exporters []Exporter Exporters []Exporter
// Max body reading size in bytes
MaxBodySize int64
} }