Fix default ratelimiter

This commit is contained in:
Rustem Kamalov
2023-06-30 22:04:26 +03:00
parent dd5e0a2069
commit df972585c1
5 changed files with 31 additions and 24 deletions

View File

@@ -40,7 +40,7 @@ func TestUrlBuild(t *testing.T) {
// }
func TestSearchBaidu(t *testing.T) {
baid := New(*browser)
baid := New(*browser, core.SearchEngineOptions{})
results, err := baid.Search(testQuery)
if err != nil {
t.Fatal(err)

View File

@@ -2,7 +2,6 @@ package baidu
import (
"strings"
"time"
"github.com/go-rod/rod"
"github.com/karust/openserp/core"
@@ -27,12 +26,12 @@ func (baid *Baidu) Name() string {
}
func (baid *Baidu) GetRateLimiter() *rate.Limiter {
ratelimit := rate.Every(baid.RateTime / time.Duration(baid.RateRequests))
ratelimit := rate.Every(baid.GetRatelimit())
return rate.NewLimiter(ratelimit, baid.RateBurst)
}
func (baid *Baidu) isCaptcha(page *rod.Page) bool {
_, err := page.Timeout(baid.SelectorTimeout).Search("div.passMod_dialog-body")
_, err := page.Timeout(baid.GetSelectorTimeout()).Search("div.passMod_dialog-body")
if err != nil {
return false
}
@@ -40,7 +39,7 @@ func (baid *Baidu) isCaptcha(page *rod.Page) bool {
}
func (baid *Baidu) isTimeout(page *rod.Page) bool {
_, err := page.Timeout(baid.SelectorTimeout).Search("button.timeout-button")
_, err := page.Timeout(baid.GetSelectorTimeout()).Search("button.timeout-button")
if err != nil {
return false
}

View File

@@ -55,23 +55,31 @@ func (q *Query) InitFromContext(c *fiber.Ctx) error {
}
type SearchEngineOptions struct {
RateRequests int `mapstructure:"rate_requests"`
RateTime time.Duration `mapstructure:"rate_seconds"`
RateBurst int `mapstructure:"rate_burst"`
SelectorTimeout time.Duration `mapstructure:"selector_timeout"`
RateRequests int `mapstructure:"rate_requests"`
RateTime int64 `mapstructure:"rate_seconds"`
RateBurst int `mapstructure:"rate_burst"`
SelectorTimeout int64 `mapstructure:"selector_timeout"` // CSS selector timeout in seconds
}
func (o *SearchEngineOptions) Init() {
if o.RateRequests == 0 {
o.RateRequests = 1
o.RateRequests = 6
}
if o.RateTime == 0 {
o.RateTime = time.Second * 10
o.RateTime = 60
}
if o.RateBurst == 0 {
o.RateBurst = 1
}
if o.SelectorTimeout == 0 {
o.SelectorTimeout = time.Second * 5
o.SelectorTimeout = 5
}
}
func (o *SearchEngineOptions) GetRatelimit() time.Duration {
return (time.Duration(o.RateTime) * time.Second) / time.Duration(o.RateRequests)
}
func (o *SearchEngineOptions) GetSelectorTimeout() time.Duration {
return time.Duration(o.SelectorTimeout) * time.Second
}

View File

@@ -5,7 +5,6 @@ import (
"regexp"
"strconv"
"strings"
"time"
"github.com/go-rod/rod"
"github.com/karust/openserp/core"
@@ -33,15 +32,16 @@ func (gogl *Google) Name() string {
}
func (gogl *Google) GetRateLimiter() *rate.Limiter {
ratelimit := rate.Every(gogl.RateTime / time.Duration(gogl.RateRequests))
ratelimit := rate.Every(gogl.GetRatelimit())
return rate.NewLimiter(ratelimit, gogl.RateBurst)
}
func (gogl *Google) FindTotalResults(page *rod.Page) (int, error) {
resultsStats, err := page.Timeout(gogl.SelectorTimeout).Search("div#result-stats")
func (gogl *Google) findTotalResults(page *rod.Page) (int, error) {
resultsStats, err := page.Timeout(gogl.GetSelectorTimeout()).Search("div#result-stats")
if err != nil {
return 0, errors.New("Result stats not found: " + err.Error())
}
stats, err := resultsStats.First.Text()
if err != nil {
return 0, errors.New("Cannot extract result stats text: " + err.Error())
@@ -59,7 +59,7 @@ func (gogl *Google) FindTotalResults(page *rod.Page) (int, error) {
}
func (gogl *Google) isCaptcha(page *rod.Page) bool {
_, err := page.Timeout(gogl.SelectorTimeout).Search("form#captcha-form")
_, err := page.Timeout(gogl.GetSelectorTimeout()).Search("form#captcha-form")
if err != nil {
return false
}
@@ -103,11 +103,11 @@ func (gogl *Google) Search(query core.Query) ([]core.SearchResult, error) {
return nil, err
}
totalResults, err := gogl.FindTotalResults(page)
totalResults, err := gogl.findTotalResults(page)
if err != nil {
return nil, err
logrus.Errorf("Error capturing total results: %v", err)
}
logrus.Tracef("%d total results found", totalResults)
logrus.Infof("%d total results found", totalResults)
resultElements, err := results.All()
if err != nil {

View File

@@ -29,12 +29,12 @@ func (yand *Yandex) Name() string {
}
func (yand *Yandex) GetRateLimiter() *rate.Limiter {
ratelimit := rate.Every(yand.RateTime / time.Duration(yand.RateRequests))
ratelimit := rate.Every(yand.GetRatelimit())
return rate.NewLimiter(ratelimit, yand.RateBurst)
}
func (yand *Yandex) isCaptcha(page *rod.Page) bool {
_, err := page.Timeout(yand.SelectorTimeout).Search("form#checkbox-captcha-form")
_, err := page.Timeout(yand.GetSelectorTimeout()).Search("form#checkbox-captcha-form")
if err != nil {
return false
}
@@ -45,12 +45,12 @@ func (yand *Yandex) isCaptcha(page *rod.Page) bool {
func (yand *Yandex) isNoResults(page *rod.Page) bool {
noResFound := false
_, err := page.Timeout(yand.SelectorTimeout).Search("div.EmptySearchResults-Title")
_, err := page.Timeout(yand.GetSelectorTimeout()).Search("div.EmptySearchResults-Title")
if err == nil {
noResFound = true
}
_, err = page.Timeout(yand.SelectorTimeout).Search("div>div.RequestMeta-Message")
_, err = page.Timeout(yand.GetSelectorTimeout()).Search("div>div.RequestMeta-Message")
if err == nil {
noResFound = true
}