1
0
mirror of https://github.com/MontFerret/ferret.git synced 2025-01-04 03:02:02 +02:00

changed AllowedHTTPCodes from []int to map[int]struct{}

This commit is contained in:
Gabriel Marinkovic 2019-10-16 00:07:01 +02:00
parent f4c63c0c6c
commit d0a3218b9a
2 changed files with 7 additions and 14 deletions

View File

@ -152,7 +152,7 @@ func (drv *Driver) Open(ctx context.Context, params drivers.Params) (drivers.HTM
defer resp.Body.Close() defer resp.Body.Close()
if !drv.responseCodeAllowed(resp) { if _, exists := drv.options.AllowedHTTPCodes[resp.StatusCode]; !exists {
return nil, errors.New(resp.Status) return nil, errors.New(resp.Status)
} }
@ -177,15 +177,6 @@ func (drv *Driver) Open(ctx context.Context, params drivers.Params) (drivers.HTM
return NewHTMLPage(doc, params.URL, &r, cookies) return NewHTMLPage(doc, params.URL, &r, cookies)
} }
func (drv *Driver) responseCodeAllowed(resp *http.Response) bool {
for _, code := range drv.options.AllowedHTTPCodes {
if resp.StatusCode == code {
return true
}
}
return false
}
func (drv *Driver) Parse(_ context.Context, str values.String) (drivers.HTMLPage, error) { func (drv *Driver) Parse(_ context.Context, str values.String) (drivers.HTMLPage, error) {
buf := bytes.NewBuffer([]byte(str)) buf := bytes.NewBuffer([]byte(str))

View File

@ -19,7 +19,7 @@ type (
UserAgent string UserAgent string
Headers drivers.HTTPHeaders Headers drivers.HTTPHeaders
Cookies drivers.HTTPCookies Cookies drivers.HTTPCookies
AllowedHTTPCodes []int AllowedHTTPCodes map[int]struct{}
} }
) )
@ -29,7 +29,7 @@ func newOptions(setters []Option) *Options {
opts.Backoff = pester.ExponentialBackoff opts.Backoff = pester.ExponentialBackoff
opts.Concurrency = 3 opts.Concurrency = 3
opts.MaxRetries = 5 opts.MaxRetries = 5
opts.AllowedHTTPCodes = []int{stdhttp.StatusOK} opts.AllowedHTTPCodes = map[int]struct{}{stdhttp.StatusOK: struct{}{}}
for _, setter := range setters { for _, setter := range setters {
setter(opts) setter(opts)
@ -132,12 +132,14 @@ func WithCookies(cookies []drivers.HTTPCookie) Option {
func WithAllowedHTTPCode(httpCode int) Option { func WithAllowedHTTPCode(httpCode int) Option {
return func(opts *Options) { return func(opts *Options) {
opts.AllowedHTTPCodes = append(opts.AllowedHTTPCodes, httpCode) opts.AllowedHTTPCodes[httpCode] = struct{}{}
} }
} }
func WithAllowedHTTPCodes(httpCodes []int) Option { func WithAllowedHTTPCodes(httpCodes []int) Option {
return func(opts *Options) { return func(opts *Options) {
opts.AllowedHTTPCodes = append(opts.AllowedHTTPCodes, httpCodes...) for _, code := range httpCodes {
opts.AllowedHTTPCodes[code] = struct{}{}
}
} }
} }