From df972585c12ae684bae645d3391ae89d47ccf24f Mon Sep 17 00:00:00 2001 From: Rustem Kamalov Date: Fri, 30 Jun 2023 22:04:26 +0300 Subject: [PATCH] Fix default ratelimiter --- baidu/all_test.go | 2 +- baidu/search.go | 7 +++---- core/common.go | 22 +++++++++++++++------- google/search.go | 16 ++++++++-------- yandex/search.go | 8 ++++---- 5 files changed, 31 insertions(+), 24 deletions(-) diff --git a/baidu/all_test.go b/baidu/all_test.go index 8974458..4ce6902 100644 --- a/baidu/all_test.go +++ b/baidu/all_test.go @@ -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) diff --git a/baidu/search.go b/baidu/search.go index 909b3d1..c09a3b7 100644 --- a/baidu/search.go +++ b/baidu/search.go @@ -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 } diff --git a/core/common.go b/core/common.go index 0e1e7ec..1ae7115 100644 --- a/core/common.go +++ b/core/common.go @@ -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 +} diff --git a/google/search.go b/google/search.go index 363f077..94bdb4a 100644 --- a/google/search.go +++ b/google/search.go @@ -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 { diff --git a/yandex/search.go b/yandex/search.go index 1301d10..732e7e5 100644 --- a/yandex/search.go +++ b/yandex/search.go @@ -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 }