1
0
mirror of https://github.com/MontFerret/ferret.git synced 2024-12-14 11:23: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()
if !drv.responseCodeAllowed(resp) {
if _, exists := drv.options.AllowedHTTPCodes[resp.StatusCode]; !exists {
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)
}
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) {
buf := bytes.NewBuffer([]byte(str))

View File

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