From dd5e0a2069c4062289fcb8ef01eb691a3943bc2e Mon Sep 17 00:00:00 2001 From: Rustem Kamalov Date: Wed, 28 Jun 2023 22:20:19 +0300 Subject: [PATCH] Default ratelimeter --- README.md | 8 +++----- baidu/search.go | 18 +++++++++++++----- cmd/root.go | 17 ++++++++++------- cmd/search.go | 6 +++--- cmd/serve.go | 6 +++--- core/common.go | 23 +++++++++++++++++++++++ core/server.go | 11 ++++++++++- core/server_test.go | 5 +++++ go.mod | 1 + go.sum | 2 ++ google/search.go | 21 +++++++++++++++------ google/search_test.go | 2 +- yandex/search.go | 22 +++++++++++++++------- yandex/search_test.go | 2 +- 14 files changed, 105 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 3dc098a..44f9d1a 100644 --- a/README.md +++ b/README.md @@ -35,14 +35,14 @@ You can replace `google` to `yandex` or `baidu` in query to change search engine | lang | Search pages in selected language (`EN`, `DE`, `RU`...) | | date | Date in `YYYYMMDD..YYYYMMDD` format (e.g. 20181010..20231010) | | file | File extension to search (e.g. `PDF`, `DOC`) | -| site | Search only in selected site | +| site | Search within a specific website | | limit | Limit the number of results | ### *Example response* ```JSON [ { - "rank": 0, + "rank": 1, "url": "https://en.wikipedia.org/wiki/%22Hello,_World!%22_program", "title": "\"Hello, World!\" program", "description": "A \"Hello, World!\" program is generally a computer program that ignores any input, and outputs or displays a message similar to \"Hello, World!\"." @@ -50,8 +50,6 @@ You can replace `google` to `yandex` or `baidu` in query to change search engine ] ``` - - ## CLI ⌨️ * Use `-h` flag to see commands. * You can use `serve` command to serve API: @@ -66,7 +64,7 @@ As a result you should get JSON output containting search results: ```json [ { - "rank": 0, + "rank": 1, "url": "https://www.cyberoptik.net/blog/6-sure-fire-ways-to-get-banned-from-google/", "title": "11 Sure-Fire Ways to Get Banned From Google | CyberOptik", "description": "How To Get Banned From Google · 1. Cloaking: The Art of Deception · 2. Plagiarism: Because Originality is Overrated · 3. Keyword Stuffing: More is Always Better · 4 ..." diff --git a/baidu/search.go b/baidu/search.go index 2cc0244..909b3d1 100644 --- a/baidu/search.go +++ b/baidu/search.go @@ -7,24 +7,32 @@ import ( "github.com/go-rod/rod" "github.com/karust/openserp/core" "github.com/sirupsen/logrus" + "golang.org/x/time/rate" ) type Baidu struct { core.Browser - checkTimeout time.Duration + core.SearchEngineOptions } -func New(browser core.Browser) *Baidu { +func New(browser core.Browser, opts core.SearchEngineOptions) *Baidu { baid := Baidu{Browser: browser} - baid.checkTimeout = time.Second * 2 + opts.Init() + baid.SearchEngineOptions = opts return &baid } + func (baid *Baidu) Name() string { return "baidu" } +func (baid *Baidu) GetRateLimiter() *rate.Limiter { + ratelimit := rate.Every(baid.RateTime / time.Duration(baid.RateRequests)) + return rate.NewLimiter(ratelimit, baid.RateBurst) +} + func (baid *Baidu) isCaptcha(page *rod.Page) bool { - _, err := page.Timeout(baid.checkTimeout).Search("div.passMod_dialog-body") + _, err := page.Timeout(baid.SelectorTimeout).Search("div.passMod_dialog-body") if err != nil { return false } @@ -32,7 +40,7 @@ func (baid *Baidu) isCaptcha(page *rod.Page) bool { } func (baid *Baidu) isTimeout(page *rod.Page) bool { - _, err := page.Timeout(baid.checkTimeout).Search("button.timeout-button") + _, err := page.Timeout(baid.SelectorTimeout).Search("button.timeout-button") if err != nil { return false } diff --git a/cmd/root.go b/cmd/root.go index 9b1272b..05042f2 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -11,7 +11,7 @@ import ( ) const ( - version = "0.1.2" + version = "0.2.1" defaultConfigFilename = "config" envPrefix = "OPENSERP" replaceHyphenWithCamelCase = false @@ -22,12 +22,15 @@ type AppConfig struct { Port int Timeout int ConfigPath string - IsBrowserHead bool `mapstructure:"head"` - IsLeaveHead bool `mapstructure:"leave_head"` - IsLeakless bool `mapstructure:"leakless"` - IsDebug bool `mapstructure:"debug"` - IsVerbose bool `mapstructure:"verbose"` - IsRawRequests bool `mapstructure:"raw_requests"` + IsBrowserHead bool `mapstructure:"head"` + IsLeaveHead bool `mapstructure:"leave_head"` + IsLeakless bool `mapstructure:"leakless"` + IsDebug bool `mapstructure:"debug"` + IsVerbose bool `mapstructure:"verbose"` + IsRawRequests bool `mapstructure:"raw_requests"` + GoogleConfig core.SearchEngineOptions `mapstructure:"google"` + YandexConfig core.SearchEngineOptions `mapstructure:"yandex"` + BaiduConfig core.SearchEngineOptions `mapstructure:"baidu"` } var appConf = AppConfig{} diff --git a/cmd/search.go b/cmd/search.go index a4c6c24..8fe5090 100644 --- a/cmd/search.go +++ b/cmd/search.go @@ -71,11 +71,11 @@ func searchBrowser(engineType string, query core.Query) ([]core.SearchResult, er switch strings.ToLower(engineType) { case "yandex": - engine = yandex.New(*browser) + engine = yandex.New(*browser, appConf.YandexConfig) case "google": - engine = google.New(*browser) + engine = google.New(*browser, appConf.GoogleConfig) case "baidu": - engine = baidu.New(*browser) + engine = baidu.New(*browser, appConf.BaiduConfig) default: logrus.Infof("No `%s` search engine found", engineType) } diff --git a/cmd/serve.go b/cmd/serve.go index 006bdd3..2c5503a 100644 --- a/cmd/serve.go +++ b/cmd/serve.go @@ -36,9 +36,9 @@ func serve(cmd *cobra.Command, args []string) { logrus.Error(err) } - yand := yandex.New(*browser) - gogl := google.New(*browser) - baidu := baidu.New(*browser) + yand := yandex.New(*browser, appConf.YandexConfig) + gogl := google.New(*browser, appConf.GoogleConfig) + baidu := baidu.New(*browser, appConf.BaiduConfig) serv := core.NewServer(appConf.Host, appConf.Port, gogl, yand, baidu) serv.Listen() diff --git a/core/common.go b/core/common.go index 89688bb..0e1e7ec 100644 --- a/core/common.go +++ b/core/common.go @@ -3,6 +3,7 @@ package core import ( "errors" "strconv" + "time" "github.com/gofiber/fiber/v2" ) @@ -52,3 +53,25 @@ func (q *Query) InitFromContext(c *fiber.Ctx) error { return nil } + +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"` +} + +func (o *SearchEngineOptions) Init() { + if o.RateRequests == 0 { + o.RateRequests = 1 + } + if o.RateTime == 0 { + o.RateTime = time.Second * 10 + } + if o.RateBurst == 0 { + o.RateBurst = 1 + } + if o.SelectorTimeout == 0 { + o.SelectorTimeout = time.Second * 5 + } +} diff --git a/core/server.go b/core/server.go index 00636f9..b64a724 100644 --- a/core/server.go +++ b/core/server.go @@ -1,18 +1,21 @@ package core import ( + "context" "errors" "fmt" "strings" "github.com/gofiber/fiber/v2" "github.com/sirupsen/logrus" + "golang.org/x/time/rate" ) type SearchEngine interface { Search(Query) ([]SearchResult, error) IsInitialized() bool Name() string + GetRateLimiter() *rate.Limiter } type Server struct { @@ -29,15 +32,21 @@ func NewServer(host string, port int, searchEngines ...SearchEngine) *Server { for _, engine := range searchEngines { locEngine := engine + limiter := engine.GetRateLimiter() + serv.app.Get(fmt.Sprintf("/%s/search", strings.ToLower(locEngine.Name())), func(c *fiber.Ctx) error { q := Query{} err := q.InitFromContext(c) - if err != nil { logrus.Errorf("Error while setting %s query: %s", locEngine.Name(), err) return err } + err = limiter.Wait(context.Background()) + if err != nil { + logrus.Errorf("Ratelimiter error during %s query: %s", locEngine.Name(), err) + } + res, err := locEngine.Search(q) if err != nil { switch err { diff --git a/core/server_test.go b/core/server_test.go index 50b383e..b438ab7 100644 --- a/core/server_test.go +++ b/core/server_test.go @@ -4,6 +4,8 @@ import ( "fmt" "testing" "time" + + "golang.org/x/time/rate" ) var ( @@ -25,6 +27,9 @@ func (SeMock) IsInitialized() bool { func (s SeMock) Search(q Query) (res []SearchResult, err error) { return []SearchResult{{Title: s.EngineName}}, nil } +func (s SeMock) GetRateLimiter() *rate.Limiter { + return nil +} func TestCreateServer(t *testing.T) { se1 := SeMock{"mock_engine_1"} diff --git a/go.mod b/go.mod index c66eda8..e4ef690 100644 --- a/go.mod +++ b/go.mod @@ -12,6 +12,7 @@ require ( github.com/spf13/cobra v1.7.0 github.com/spf13/pflag v1.0.5 github.com/spf13/viper v1.16.0 + golang.org/x/time v0.3.0 ) require ( diff --git a/go.sum b/go.sum index 6cb1a03..e79b704 100644 --- a/go.sum +++ b/go.sum @@ -421,6 +421,8 @@ golang.org/x/text v0.10.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= +golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4= +golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY= diff --git a/google/search.go b/google/search.go index 619c022..363f077 100644 --- a/google/search.go +++ b/google/search.go @@ -10,26 +10,35 @@ import ( "github.com/go-rod/rod" "github.com/karust/openserp/core" "github.com/sirupsen/logrus" + "golang.org/x/time/rate" ) type Google struct { core.Browser - findNumRgxp *regexp.Regexp - checkTimeout time.Duration + core.SearchEngineOptions + findNumRgxp *regexp.Regexp } -func New(browser core.Browser) *Google { +func New(browser core.Browser, opts core.SearchEngineOptions) *Google { gogl := Google{Browser: browser} - gogl.checkTimeout = time.Second * 5 + opts.Init() + gogl.SearchEngineOptions = opts + gogl.findNumRgxp = regexp.MustCompile("\\d") return &gogl } + func (gogl *Google) Name() string { return "google" } +func (gogl *Google) GetRateLimiter() *rate.Limiter { + ratelimit := rate.Every(gogl.RateTime / time.Duration(gogl.RateRequests)) + return rate.NewLimiter(ratelimit, gogl.RateBurst) +} + func (gogl *Google) FindTotalResults(page *rod.Page) (int, error) { - resultsStats, err := page.Timeout(gogl.checkTimeout).Search("div#result-stats") + resultsStats, err := page.Timeout(gogl.SelectorTimeout).Search("div#result-stats") if err != nil { return 0, errors.New("Result stats not found: " + err.Error()) } @@ -50,7 +59,7 @@ func (gogl *Google) FindTotalResults(page *rod.Page) (int, error) { } func (gogl *Google) isCaptcha(page *rod.Page) bool { - _, err := page.Timeout(gogl.checkTimeout).Search("form#captcha-form") + _, err := page.Timeout(gogl.SelectorTimeout).Search("form#captcha-form") if err != nil { return false } diff --git a/google/search_test.go b/google/search_test.go index 55202ec..eb4344e 100644 --- a/google/search_test.go +++ b/google/search_test.go @@ -15,7 +15,7 @@ func init() { } func TestSearchGoogle(t *testing.T) { - gogl := New(*browser) + gogl := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "HEY", Limit: 10} results, err := gogl.Search(query) diff --git a/yandex/search.go b/yandex/search.go index 6c95fba..1301d10 100644 --- a/yandex/search.go +++ b/yandex/search.go @@ -6,17 +6,20 @@ import ( "github.com/go-rod/rod" "github.com/karust/openserp/core" "github.com/sirupsen/logrus" + "golang.org/x/time/rate" ) type Yandex struct { core.Browser - checkTimeout time.Duration // Timeout for secondary elements check - pageSleep time.Duration // Sleep between pages + core.SearchEngineOptions + pageSleep time.Duration // Sleep between pages } -func New(browser core.Browser) *Yandex { +func New(browser core.Browser, opts core.SearchEngineOptions) *Yandex { yand := Yandex{Browser: browser} - yand.checkTimeout = time.Second * 2 + opts.Init() + yand.SearchEngineOptions = opts + yand.pageSleep = time.Second * 1 return &yand } @@ -25,8 +28,13 @@ func (yand *Yandex) Name() string { return "yandex" } +func (yand *Yandex) GetRateLimiter() *rate.Limiter { + ratelimit := rate.Every(yand.RateTime / time.Duration(yand.RateRequests)) + return rate.NewLimiter(ratelimit, yand.RateBurst) +} + func (yand *Yandex) isCaptcha(page *rod.Page) bool { - _, err := page.Timeout(yand.checkTimeout).Search("form#checkbox-captcha-form") + _, err := page.Timeout(yand.SelectorTimeout).Search("form#checkbox-captcha-form") if err != nil { return false } @@ -37,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.checkTimeout).Search("div.EmptySearchResults-Title") + _, err := page.Timeout(yand.SelectorTimeout).Search("div.EmptySearchResults-Title") if err == nil { noResFound = true } - _, err = page.Timeout(yand.checkTimeout).Search("div>div.RequestMeta-Message") + _, err = page.Timeout(yand.SelectorTimeout).Search("div>div.RequestMeta-Message") if err == nil { noResFound = true } diff --git a/yandex/search_test.go b/yandex/search_test.go index 068aded..8cacf2a 100644 --- a/yandex/search_test.go +++ b/yandex/search_test.go @@ -16,7 +16,7 @@ func init() { func TestSearchYandex(t *testing.T) { - yand := New(*browser) + yand := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "HEY", Limit: 10} results, err := yand.Search(query)