diff --git a/core/server.go b/core/server.go index b64a724..8c355e0 100644 --- a/core/server.go +++ b/core/server.go @@ -13,6 +13,7 @@ import ( type SearchEngine interface { Search(Query) ([]SearchResult, error) + SearchImage(Query) ([]SearchResult, error) IsInitialized() bool Name() string GetRateLimiter() *rate.Limiter @@ -53,7 +54,42 @@ func NewServer(host string, port int, searchEngines ...SearchEngine) *Server { case ErrCaptcha: err = errors.New(fmt.Sprintf("Captcha found, please stop sending requests for a while\n%s", err)) case ErrSearchTimeout: - err = errors.New(fmt.Sprintf("Error: %s\nProbably need to update CSS selector", err)) + err = errors.New(fmt.Sprintf("%s", err)) + } + + logrus.Errorf("Error during %s search: %s", locEngine.Name(), err) + return fiber.NewError(fiber.StatusServiceUnavailable, err.Error()) + } + + return c.JSON(res) + }) + + serv.app.Get(fmt.Sprintf("/%s/image", 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.SearchImage(q) + + if err != nil && len(res) > 0 { + c.Status(503) + return c.JSON(res) + } + + if err != nil { + switch err { + case ErrCaptcha: + err = errors.New(fmt.Sprintf("Captcha found, please stop sending requests for a while\n%s", err)) + case ErrSearchTimeout: + err = errors.New(fmt.Sprintf("%s", err)) } logrus.Errorf("Error during %s search: %s", locEngine.Name(), err) diff --git a/core/server_test.go b/core/server_test.go index b438ab7..a431265 100644 --- a/core/server_test.go +++ b/core/server_test.go @@ -27,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) SearchImage(q Query) (res []SearchResult, err error) { + return []SearchResult{{Title: s.EngineName}}, nil +} func (s SeMock) GetRateLimiter() *rate.Limiter { return nil }