From 27b5817d5e1292cf1ae8792267ee72516e8ac779 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E7=A5=A5=E8=B6=85?= <iwind.liu@gmail.com>
Date: Tue, 29 Nov 2022 19:14:46 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E8=AF=B7=E6=B1=82=E9=99=90?=
 =?UTF-8?q?=E5=88=B6=E9=80=BB=E8=BE=91=EF=BC=8C=E8=BF=9E=E6=8E=A5=E5=85=B3?=
 =?UTF-8?q?=E9=97=AD=E6=97=B6=E8=87=AA=E5=8A=A8=E7=BB=88=E6=AD=A2=E5=86=85?=
 =?UTF-8?q?=E5=AE=B9=E5=8F=91=E9=80=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/nodes/http_writer.go               |  2 +-
 internal/utils/writers/writer_rate_limit.go | 13 ++++++++++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

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 {