Repository: valyala/fasthttp
Pull request #687 (2019-11-03 13:04:12)

Recover from panic in body write
Possibly a successor of https://github.com/valyala/fasthttp/pull/646.
Went as per @erikdubbelboer suggestion (if I understood it correctly). If the general idea is acceptable, will add tests.
Thanks.

Bobochka: @erikdubbelboer done.

Unfortunately this doesn't cover `SetBodyStreamWriter` as writer invocation happens in another goroutine.
erikdubbelboer: I think we can catch the `SetBodyStreamWriter` as well if we do a similar trick as [`io.PipeWriter.CloseWithError()`](https://golang.org/pkg/io/#PipeWriter.CloseWithError). But this will be a bit complicated as [`fasthttputil.NewPipeConns()`](https://github.com/valyala/fasthttp/blob/e5f51c11919d4f66400334047b897ef0a94c6f3c/stream.go#L29) deals with `net.Conn` and not some interface we can just add `CloseWithError()` to.

So lets save that for another pull request.

Patch:
diff --git a/http.go b/http.go
index 9a55805..65a008a 100644
--- a/http.go
+++ b/http.go
@@ -1462,8 +1462,19 @@ func (req *Request) writeBodyStream(w *bufio.Writer) error {
 	return err
 }
 
-func (resp *Response) writeBodyStream(w *bufio.Writer, sendBody bool) error {
-	var err error
+// ErrBodyStreamWritePanic is returned when panic happens during writing body stream.
+type ErrBodyStreamWritePanic struct {
+	error
+}
+
+func (resp *Response) writeBodyStream(w *bufio.Writer, sendBody bool) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = &ErrBodyStreamWritePanic{
+				error: fmt.Errorf("panic while writing body stream: %+v", r),
+			}
+		}
+	}()
 
 	contentLength := resp.Header.ContentLength()
 	if contentLength < 0 {


Test patch:
diff --git a/http_test.go b/http_test.go
index 7244211..be0b8cb 100644
--- a/http_test.go
+++ b/http_test.go
@@ -2189,3 +2189,64 @@ func TestResponseImmediateHeaderFlushChunked(t *testing.T) {
 
 	<-waitForIt
 }
+
+type ErroneousBodyStream struct {
+	errOnRead  bool
+	errOnClose bool
+}
+
+func (ebs *ErroneousBodyStream) Read(p []byte) (n int, err error) {
+	if ebs.errOnRead {
+		panic("reading erroneous body stream")
+	}
+	return 0, io.EOF
+}
+
+func (ebs *ErroneousBodyStream) Close() error {
+	if ebs.errOnClose {
+		panic("closing erroneous body stream")
+	}
+	return nil
+}
+
+func TestResponseBodyStreamErrorOnPanicDuringRead(t *testing.T) {
+	t.Parallel()
+	var resp Response
+	var w bytes.Buffer
+	bw := bufio.NewWriter(&w)
+
+	ebs := &ErroneousBodyStream{errOnRead: true, errOnClose: false}
+	resp.SetBodyStream(ebs, 42)
+	err := resp.Write(bw)
+	if err == nil {
+		t.Fatalf("expected error when writing response.")
+	}
+	e, ok := err.(*ErrBodyStreamWritePanic)
+	if !ok {
+		t.Fatalf("expected error struct to be *ErrBodyStreamWritePanic, got: %+v.", e)
+	}
+	if e.Error() != "panic while writing body stream: reading erroneous body stream" {
+		t.Fatalf("unexpected error value, got: %+v.", e.Error())
+	}
+}
+
+func TestResponseBodyStreamErrorOnPanicDuringClose(t *testing.T) {
+	t.Parallel()
+	var resp Response
+	var w bytes.Buffer
+	bw := bufio.NewWriter(&w)
+
+	ebs := &ErroneousBodyStream{errOnRead: false, errOnClose: true}
+	resp.SetBodyStream(ebs, 42)
+	err := resp.Write(bw)
+	if err == nil {
+		t.Fatalf("expected error when writing response.")
+	}
+	e, ok := err.(*ErrBodyStreamWritePanic)
+	if !ok {
+		t.Fatalf("expected error struct to be *ErrBodyStreamWritePanic, got: %+v.", e)
+	}
+	if e.Error() != "panic while writing body stream: closing erroneous body stream" {
+		t.Fatalf("unexpected error value, got: %+v.", e.Error())
+	}
+}