Add recover() to H2 server's flushWriter.Write()

Fixes #1748
XTLS · Mar 8, 2023 · 836e84b · 836e84b
1 parent 4a0b45d
commit 836e84b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/transport/internet/http/dialer.go b/transport/internet/http/dialer.go
@@ -204,7 +204,7 @@ type WaitReadCloser struct {
 func (w *WaitReadCloser) Set(rc io.ReadCloser) {
 	w.ReadCloser = rc
 	defer func() {
-		if err := recover(); err != nil {
+		if recover() != nil {
 			rc.Close()
 		}
 	}()
@@ -225,10 +225,8 @@ func (w *WaitReadCloser) Close() error {
 		return w.ReadCloser.Close()
 	}
 	defer func() {
-		if err := recover(); err != nil {
-			if w.ReadCloser != nil {
-				w.ReadCloser.Close()
-			}
+		if recover() != nil && w.ReadCloser != nil {
+			w.ReadCloser.Close()
 		}
 	}()
 	close(w.Wait)

diff --git a/transport/internet/http/hub.go b/transport/internet/http/hub.go
@@ -51,6 +51,13 @@ func (fw flushWriter) Write(p []byte) (n int, err error) {
 		return 0, io.ErrClosedPipe
 	}
 
+	defer func() {
+		if recover() != nil {
+			fw.d.Close()
+			err = io.ErrClosedPipe
+		}
+	}()
+
 	n, err = fw.w.Write(p)
 	if f, ok := fw.w.(http.Flusher); ok && err == nil {
 		f.Flush()