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

@ -8,15 +8,16 @@ import (
"net/http"
"strconv"
"github.com/WJQSERVER-STUDIO/go-utils/limitreader"
"github.com/cloudwego/hertz/pkg/app"
)
func GitReq(ctx context.Context, c *app.RequestContext, u string, cfg *config.Config, mode string) {
method := string(c.Request.Method())
bodyReader := bytes.NewBuffer(c.Request.Body())
reqBodyReader := bytes.NewBuffer(c.Request.Body())
//bodyReader := c.Request.BodyStream()
//bodyReader := c.Request.BodyStream() // 不可替换为此实现
if cfg.GitClone.Mode == "cache" {
userPath, repoPath, remainingPath, queryParams, err := extractParts(u)
@ -35,7 +36,7 @@ func GitReq(ctx context.Context, c *app.RequestContext, u string, cfg *config.Co
if cfg.GitClone.Mode == "cache" {
rb := gitclient.NewRequestBuilder(method, u)
rb.NoDefaultHeaders()
rb.SetBody(bodyReader)
rb.SetBody(reqBodyReader)
rb.WithContext(ctx)
req, err := rb.Build()
@ -55,7 +56,7 @@ func GitReq(ctx context.Context, c *app.RequestContext, u string, cfg *config.Co
} else {
rb := client.NewRequestBuilder(string(c.Request.Method()), u)
rb.NoDefaultHeaders()
rb.SetBody(bodyReader)
rb.SetBody(reqBodyReader)
rb.WithContext(ctx)
req, err := rb.Build()
@ -91,7 +92,6 @@ func GitReq(ctx context.Context, c *app.RequestContext, u string, cfg *config.Co
for key, values := range resp.Header {
for _, value := range values {
//c.Header(key, value)
c.Response.Header.Add(key, value)
}
}
@ -124,5 +124,11 @@ func GitReq(ctx context.Context, c *app.RequestContext, u string, cfg *config.Co
c.Response.Header.Set("Expires", "0")
}
c.SetBodyStream(resp.Body, -1)
bodyReader := resp.Body
if cfg.RateLimit.BandwidthLimit.Enabled {
bodyReader = limitreader.NewRateLimitedReader(bodyReader, bandwidthLimit, int(bandwidthBurst), ctx)
}
c.SetBodyStream(bodyReader, -1)
}