add bandwidth limiter

This commit is contained in:
wjqserver 2025-05-14 01:33:54 +08:00
parent 3f8d16511e
commit 71bc2aaed7
10 changed files with 157 additions and 22 deletions

View file

@ -7,6 +7,7 @@ import (
"net/http"
"strconv"
"github.com/WJQSERVER-STUDIO/go-utils/limitreader"
"github.com/cloudwego/hertz/pkg/app"
)
@ -107,10 +108,16 @@ func GhcrRequest(ctx context.Context, c *app.RequestContext, u string, cfg *conf
c.Status(resp.StatusCode)
bodyReader := resp.Body
if cfg.RateLimit.BandwidthLimit.Enabled {
bodyReader = limitreader.NewRateLimitedReader(bodyReader, bandwidthLimit, int(bandwidthBurst), ctx)
}
if contentLength != "" {
c.SetBodyStream(resp.Body, bodySize)
c.SetBodyStream(bodyReader, bodySize)
return
}
c.SetBodyStream(resp.Body, -1)
c.SetBodyStream(bodyReader, -1)
}