From 87b0d792b443bec400b985b38786616967b03f35 Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Sun, 3 Dec 2023 11:29:43 +0800
Subject: [PATCH] 1

---
 main.go | 34 +++++++++++++++++++++++-----------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/main.go b/main.go
index 17783db..6598f98 100644
--- a/main.go
+++ b/main.go
@@ -330,25 +330,37 @@ func wsDealer(ctx context.Context, w http.ResponseWriter, r *http.Request, route
 			fin := make(chan error)
 			reqc := req.NetConn()
 			resc := res.NetConn()
+			ctx, cancle := pctx.WithWait(ctx, 2, time.Minute)
 			defer func() {
-				reqc.Close()
-				resc.Close()
+				if errors.Is(cancle(), pctx.ErrWaitTo) {
+					logger.L(`E:`, "退出超时")
+				}
 			}()
 			go func() {
+				ctx1, done1 := pctx.WaitCtx(ctx)
+				defer done1()
 				_, e := io.Copy(reqc, resc)
-				fin <- e
+				select {
+				case fin <- e:
+				case <-ctx1.Done():
+					fin <- nil
+				}
+				reqc.Close()
 			}()
 			go func() {
+				ctx1, done1 := pctx.WaitCtx(ctx)
+				defer done1()
 				_, e := io.Copy(resc, reqc)
-				fin <- e
-			}()
-			select {
-			case e := <-fin:
-				if e != nil {
-					logger.L(`E:`, fmt.Sprintf("%s=>%s %v", routePath, back.Name, e))
-					return ErrCopy
+				select {
+				case fin <- e:
+				case <-ctx1.Done():
+					fin <- nil
 				}
-			case <-ctx.Done():
+				resc.Close()
+			}()
+			if e := <-fin; e != nil {
+				logger.L(`E:`, fmt.Sprintf("%s=>%s %v", routePath, back.Name, e))
+				return ErrCopy
 			}
 			return nil
 		}
-- 
2.39.2