diff --git a/internal/nodes/http_writer.go b/internal/nodes/http_writer.go index be2bc79..7cc111e 100644 --- a/internal/nodes/http_writer.go +++ b/internal/nodes/http_writer.go @@ -132,7 +132,7 @@ func (this *HTTPWriter) Prepare(resp *http.Response, size int64, status int, ena this.req.web.RequestLimit != nil && this.req.web.RequestLimit.IsOn && this.req.web.RequestLimit.OutBandwidthPerConnBytes() > 0 { - this.writer = writers.NewRateLimitWriter(this.writer, this.req.web.RequestLimit.OutBandwidthPerConnBytes()) + this.writer = writers.NewRateLimitWriter(this.req.RawReq.Context(), this.writer, this.req.web.RequestLimit.OutBandwidthPerConnBytes()) } return diff --git a/internal/utils/writers/writer_rate_limit.go b/internal/utils/writers/writer_rate_limit.go index 154cc9b..8cb4fe0 100644 --- a/internal/utils/writers/writer_rate_limit.go +++ b/internal/utils/writers/writer_rate_limit.go @@ -3,6 +3,7 @@ package writers import ( + "context" "github.com/iwind/TeaGo/types" "io" "time" @@ -11,6 +12,7 @@ import ( // RateLimitWriter 限速写入 type RateLimitWriter struct { rawWriter io.WriteCloser + ctx context.Context rateBytes int @@ -18,9 +20,10 @@ type RateLimitWriter struct { before time.Time } -func NewRateLimitWriter(rawWriter io.WriteCloser, rateBytes int64) io.WriteCloser { +func NewRateLimitWriter(ctx context.Context, rawWriter io.WriteCloser, rateBytes int64) io.WriteCloser { return &RateLimitWriter{ rawWriter: rawWriter, + ctx: ctx, rateBytes: types.Int(rateBytes), before: time.Now(), } @@ -71,6 +74,14 @@ func (this *RateLimitWriter) write(p []byte) (n int, err error) { n, err = this.rawWriter.Write(p) if err == nil { + select { + case <-this.ctx.Done(): + err = io.EOF + return + default: + + } + this.written += n if this.written >= this.rateBytes {