Backup_Repos/stash
1
0
Fork 0
mirror of https://github.com/stashapp/stash.git synced 2025-12-06 08:26:00 +01:00
Code Issues Projects Releases Wiki Activity

Add rate limit to stashbox connection (#5764)

Browse source 
* Add max requests per minute stashbox option
* Implement rate limiting
* Add requests per minute to stashbox config
* Add UI setting
This commit is contained in:
WithoutPants 2025-03-27 11:54:00 +11:00 committed by GitHub
parent 18381664aa
commit c8d74f0bcf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 118 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
go.mod
View file

@ -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
)

2
go.sum
View file

@ -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=

3
graphql/schema/types/stash-box.graphql
View file

@ -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 {

2
internal/api/stash_box.go
View file

@ -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) {

1
internal/manager/config/config.go
View file

@ -1108,6 +1108,7 @@ type StashBoxInput struct {
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 {

2
internal/manager/task_identify.go
View file

@ -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,

4
internal/manager/task_stash_box_tag.go
View file

@ -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

1
pkg/models/stash_box.go
View file

@ -10,4 +10,5 @@ type StashBox struct {
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"`
}

73
pkg/stashbox/client.go
View file

@ -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 {

1
ui/v2.5/graphql/data/config.graphql
View file

@ -49,6 +49,7 @@ fragment ConfigGeneralData on ConfigGeneralResult {
name
endpoint
api_key
max_requests_per_minute
}
pythonPath
transcodeInputArgs

34
ui/v2.5/src/components/Settings/StashBoxConfiguration.tsx
View file

@ -10,6 +10,8 @@ export interface IStashBoxModal {
close: (v?: GQL.StashBoxInput) => void;
}
const defaultMaxRequestsPerMinute = 240;
export const StashBoxModal: React.FC<IStashBoxModal> = ({ value, close }) => {
const intl = useIntl();
const endpoint = useRef<HTMLInputElement | null>(null);
@ -114,6 +116,38 @@ export const StashBoxModal: React.FC<IStashBoxModal> = ({ value, close }) => {
</b>
)}
</Form.Group>
<Form.Group id="stashbox-max-requests-per-minute">
<h6>
{intl.formatMessage({
id: "config.stashbox.max_requests_per_minute",
})}
</h6>
<Form.Control
placeholder={intl.formatMessage({
id: "config.stashbox.max_requests_per_minute",
})}
className="text-input"
value={v?.max_requests_per_minute ?? defaultMaxRequestsPerMinute}
isValid={
(v?.max_requests_per_minute ?? defaultMaxRequestsPerMinute) >= 0
}
type="number"
onChange={(e: React.ChangeEvent<HTMLInputElement>) =>
setValue({
...v!,
max_requests_per_minute: parseInt(e.currentTarget.value),
})
}
ref={apiKey}
/>
<div className="sub-heading">
<FormattedMessage
id="config.stashbox.max_requests_per_minute_description"
values={{ defaultValue: defaultMaxRequestsPerMinute }}
/>
</div>
</Form.Group>
</>
)}
close={close}

2
ui/v2.5/src/locales/en-GB.json
View file

@ -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"
},