From c8d74f0bcfd273fd3a245bbc48b7a76b5a905ddb Mon Sep 17 00:00:00 2001 From: WithoutPants <53250216+WithoutPants@users.noreply.github.com> Date: Thu, 27 Mar 2025 11:54:00 +1100 Subject: [PATCH] Add rate limit to stashbox connection (#5764) * Add max requests per minute stashbox option * Implement rate limiting * Add requests per minute to stashbox config * Add UI setting --- go.mod | 1 + go.sum | 2 + graphql/schema/types/stash-box.graphql | 3 + internal/api/stash_box.go | 2 +- internal/manager/config/config.go | 7 +- internal/manager/task_identify.go | 2 +- internal/manager/task_stash_box_tag.go | 4 +- pkg/models/stash_box.go | 7 +- pkg/stashbox/client.go | 73 ++++++++++++++++--- ui/v2.5/graphql/data/config.graphql | 1 + .../Settings/StashBoxConfiguration.tsx | 34 +++++++++ ui/v2.5/src/locales/en-GB.json | 2 + 12 files changed, 118 insertions(+), 20 deletions(-) diff --git a/go.mod b/go.mod index 0053c1e1a..368f98340 100644 --- a/go.mod +++ b/go.mod @@ -57,6 +57,7 @@ require ( golang.org/x/sys v0.28.0 golang.org/x/term v0.27.0 golang.org/x/text v0.21.0 + golang.org/x/time v0.10.0 gopkg.in/guregu/null.v4 v4.0.0 gopkg.in/yaml.v2 v2.4.0 ) diff --git a/go.sum b/go.sum index 91f460c91..70dfe5556 100644 --- a/go.sum +++ b/go.sum @@ -954,6 +954,8 @@ golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= 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.10.0 h1:3usCWA8tQn0L8+hFJQNgzpWbd89begxN66o1Ojdn5L4= +golang.org/x/time v0.10.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= 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/graphql/schema/types/stash-box.graphql b/graphql/schema/types/stash-box.graphql index d1da8c74a..c3c2867e9 100644 --- a/graphql/schema/types/stash-box.graphql +++ b/graphql/schema/types/stash-box.graphql @@ -2,12 +2,15 @@ type StashBox { endpoint: String! api_key: String! name: String! + max_requests_per_minute: Int! } input StashBoxInput { endpoint: String! api_key: String! name: String! + # defaults to 240 + max_requests_per_minute: Int } type StashID { diff --git a/internal/api/stash_box.go b/internal/api/stash_box.go index 7c656368a..852c219a5 100644 --- a/internal/api/stash_box.go +++ b/internal/api/stash_box.go @@ -11,7 +11,7 @@ import ( ) func (r *Resolver) newStashBoxClient(box models.StashBox) *stashbox.Client { - return stashbox.NewClient(box, manager.GetInstance().Config.GetScraperExcludeTagPatterns()) + return stashbox.NewClient(box, stashbox.ExcludeTagPatterns(manager.GetInstance().Config.GetScraperExcludeTagPatterns())) } func resolveStashBoxFn(indexField, endpointField string) func(index *int, endpoint *string) (*models.StashBox, error) { diff --git a/internal/manager/config/config.go b/internal/manager/config/config.go index 632c63065..65ba111ac 100644 --- a/internal/manager/config/config.go +++ b/internal/manager/config/config.go @@ -1105,9 +1105,10 @@ func stashBoxValidate(str string) bool { } type StashBoxInput struct { - Endpoint string `json:"endpoint"` - APIKey string `json:"api_key"` - Name string `json:"name"` + Endpoint string `json:"endpoint"` + APIKey string `json:"api_key"` + Name string `json:"name"` + MaxRequestsPerMinute int `json:"max_requests_per_minute"` } func (i *Config) ValidateStashBoxes(boxes []*StashBoxInput) error { diff --git a/internal/manager/task_identify.go b/internal/manager/task_identify.go index 8d5993f54..137842928 100644 --- a/internal/manager/task_identify.go +++ b/internal/manager/task_identify.go @@ -180,7 +180,7 @@ func (j *IdentifyJob) getSources() ([]identify.ScraperSource, error) { src = identify.ScraperSource{ Name: "stash-box: " + stashBox.Endpoint, Scraper: stashboxSource{ - Client: stashbox.NewClient(*stashBox, instance.Config.GetScraperExcludeTagPatterns()), + Client: stashbox.NewClient(*stashBox, stashbox.ExcludeTagPatterns(instance.Config.GetScraperExcludeTagPatterns())), endpoint: stashBox.Endpoint, txnManager: instance.Repository.TxnManager, sceneFingerprintGetter: instance.SceneService, diff --git a/internal/manager/task_stash_box_tag.go b/internal/manager/task_stash_box_tag.go index 3325653e8..d20b71f06 100644 --- a/internal/manager/task_stash_box_tag.go +++ b/internal/manager/task_stash_box_tag.go @@ -96,7 +96,7 @@ func (t *StashBoxBatchTagTask) findStashBoxPerformer(ctx context.Context) (*mode r := instance.Repository - client := stashbox.NewClient(*t.box, instance.Config.GetScraperExcludeTagPatterns()) + client := stashbox.NewClient(*t.box, stashbox.ExcludeTagPatterns(instance.Config.GetScraperExcludeTagPatterns())) if t.refresh { var remoteID string @@ -295,7 +295,7 @@ func (t *StashBoxBatchTagTask) findStashBoxStudio(ctx context.Context) (*models. r := instance.Repository - client := stashbox.NewClient(*t.box, instance.Config.GetScraperExcludeTagPatterns()) + client := stashbox.NewClient(*t.box, stashbox.ExcludeTagPatterns(instance.Config.GetScraperExcludeTagPatterns())) if t.refresh { var remoteID string diff --git a/pkg/models/stash_box.go b/pkg/models/stash_box.go index 9c9d7ed2e..6a254a3f9 100644 --- a/pkg/models/stash_box.go +++ b/pkg/models/stash_box.go @@ -7,7 +7,8 @@ type StashBoxFingerprint struct { } type StashBox struct { - Endpoint string `json:"endpoint"` - APIKey string `json:"api_key"` - Name string `json:"name"` + Endpoint string `json:"endpoint"` + APIKey string `json:"api_key"` + Name string `json:"name"` + MaxRequestsPerMinute int `json:"max_requests_per_minute" koanf:"max_requests_per_minute"` } diff --git a/pkg/stashbox/client.go b/pkg/stashbox/client.go index 4ec0cbead..ca789aa59 100644 --- a/pkg/stashbox/client.go +++ b/pkg/stashbox/client.go @@ -10,33 +10,86 @@ import ( "github.com/stashapp/stash/pkg/models" "github.com/stashapp/stash/pkg/scraper" "github.com/stashapp/stash/pkg/stashbox/graphql" + + "golang.org/x/time/rate" ) +// DefaultMaxRequestsPerMinute is the default maximum number of requests per minute. +const DefaultMaxRequestsPerMinute = 240 + // Client represents the client interface to a stash-box server instance. type Client struct { client *graphql.Client box models.StashBox + maxRequestsPerMinute int + // tag patterns to be excluded excludeTagRE []*regexp.Regexp } -// NewClient returns a new instance of a stash-box client. -func NewClient(box models.StashBox, excludeTagPatterns []string) *Client { - authHeader := func(ctx context.Context, req *http.Request, gqlInfo *clientv2.GQLRequestInfo, res interface{}, next clientv2.RequestInterceptorFunc) error { - req.Header.Set("ApiKey", box.APIKey) +type ClientOption func(*Client) + +func ExcludeTagPatterns(patterns []string) ClientOption { + return func(c *Client) { + c.excludeTagRE = scraper.CompileExclusionRegexps(patterns) + } +} + +func MaxRequestsPerMinute(n int) ClientOption { + return func(c *Client) { + if n > 0 { + c.maxRequestsPerMinute = n + } + } +} + +func setApiKeyHeader(apiKey string) clientv2.RequestInterceptor { + return func(ctx context.Context, req *http.Request, gqlInfo *clientv2.GQLRequestInfo, res interface{}, next clientv2.RequestInterceptorFunc) error { + req.Header.Set("ApiKey", apiKey) return next(ctx, req, gqlInfo, res) } +} + +func rateLimit(n int) clientv2.RequestInterceptor { + perSec := float64(n) / 60 + limiter := rate.NewLimiter(rate.Limit(perSec), 1) + + return func(ctx context.Context, req *http.Request, gqlInfo *clientv2.GQLRequestInfo, res interface{}, next clientv2.RequestInterceptorFunc) error { + if err := limiter.Wait(ctx); err != nil { + // should only happen if the context is canceled + return err + } + + return next(ctx, req, gqlInfo, res) + } +} + +// NewClient returns a new instance of a stash-box client. +func NewClient(box models.StashBox, options ...ClientOption) *Client { + ret := &Client{ + box: box, + maxRequestsPerMinute: DefaultMaxRequestsPerMinute, + } + + if box.MaxRequestsPerMinute > 0 { + ret.maxRequestsPerMinute = box.MaxRequestsPerMinute + } + + for _, option := range options { + option(ret) + } + + authHeader := setApiKeyHeader(box.APIKey) + limitRequests := rateLimit(ret.maxRequestsPerMinute) client := &graphql.Client{ - Client: clientv2.NewClient(http.DefaultClient, box.Endpoint, nil, authHeader), + Client: clientv2.NewClient(http.DefaultClient, box.Endpoint, nil, authHeader, limitRequests), } - return &Client{ - client: client, - box: box, - excludeTagRE: scraper.CompileExclusionRegexps(excludeTagPatterns), - } + ret.client = client + + return ret } func (c Client) getHTTPClient() *http.Client { diff --git a/ui/v2.5/graphql/data/config.graphql b/ui/v2.5/graphql/data/config.graphql index ae15aa939..c0bcda821 100644 --- a/ui/v2.5/graphql/data/config.graphql +++ b/ui/v2.5/graphql/data/config.graphql @@ -49,6 +49,7 @@ fragment ConfigGeneralData on ConfigGeneralResult { name endpoint api_key + max_requests_per_minute } pythonPath transcodeInputArgs diff --git a/ui/v2.5/src/components/Settings/StashBoxConfiguration.tsx b/ui/v2.5/src/components/Settings/StashBoxConfiguration.tsx index 9d0c4cfd1..cab8a4e76 100644 --- a/ui/v2.5/src/components/Settings/StashBoxConfiguration.tsx +++ b/ui/v2.5/src/components/Settings/StashBoxConfiguration.tsx @@ -10,6 +10,8 @@ export interface IStashBoxModal { close: (v?: GQL.StashBoxInput) => void; } +const defaultMaxRequestsPerMinute = 240; + export const StashBoxModal: React.FC = ({ value, close }) => { const intl = useIntl(); const endpoint = useRef(null); @@ -114,6 +116,38 @@ export const StashBoxModal: React.FC = ({ value, close }) => { )} + + +
+ {intl.formatMessage({ + id: "config.stashbox.max_requests_per_minute", + })} +
+ = 0 + } + type="number" + onChange={(e: React.ChangeEvent) => + setValue({ + ...v!, + max_requests_per_minute: parseInt(e.currentTarget.value), + }) + } + ref={apiKey} + /> +
+ +
+ )} close={close} diff --git a/ui/v2.5/src/locales/en-GB.json b/ui/v2.5/src/locales/en-GB.json index f8dfe0658..7e5379e18 100644 --- a/ui/v2.5/src/locales/en-GB.json +++ b/ui/v2.5/src/locales/en-GB.json @@ -448,6 +448,8 @@ "description": "Stash-box facilitates automated tagging of scenes and performers based on fingerprints and filenames.\nEndpoint and API key can be found on your account page on the stash-box instance. Names are required when more than one instance is added.", "endpoint": "Endpoint", "graphql_endpoint": "GraphQL endpoint", + "max_requests_per_minute": "Max requests per minute", + "max_requests_per_minute_description": "Uses default value of {defaultValue} if set to 0", "name": "Name", "title": "Stash-box Endpoints" },