Repository: valyala/fasthttp
Pull request #458 (2018-11-09 10:48:11)

#457: allow to rewrite system error response

erikdubbelboer: I think this is a feature other people have requested before. I would like to see some changes to make this more usable before I can accept it.

I think it would be useful to add some documentation on what kind of errors the user can expect as argument to the `ErrorHandler` function. I think it's error that can be returned by Go while reading from a socket such as `io.EOF` or`io.ErrUnexpectedEOF`. It can be `ErrBodyTooLarge`.
But it can also be `errGetOnly`, so I think its best if you change [`errGetOnly`](https://github.com/valyala/fasthttp/blob/8f4c5d5d1dbb772a29f059645f0f03a4556c1dad/http.go#L850) to be exported.
Can you also turn [this](https://github.com/valyala/fasthttp/blob/8f4c5d5d1dbb772a29f059645f0f03a4556c1dad/http.go#L1674) into a properly exported error and add it to the list.
Then I think only [these](https://github.com/valyala/fasthttp/blob/8f4c5d5d1dbb772a29f059645f0f03a4556c1dad/http.go#L1691-L1707) errors remain. I think its best if you turned these into a custom error struct that is exported so users can do a type test for it. Since they contain formatted messages we can't just have an exported error like the others. The custom error struct only has to contain an unexported string and a `Error() string` function that returns that string.
derElektrobesen: Thanks for review.

All fixes are done

Patch:
diff --git a/http.go b/http.go
index 6bae63c..59f9dd9 100644
--- a/http.go
+++ b/http.go
@@ -847,7 +847,9 @@ func (req *Request) Read(r *bufio.Reader) error {
 
 const defaultMaxInMemoryFileSize = 16 * 1024 * 1024
 
-var errGetOnly = errors.New("non-GET request received")
+// ErrGetOnly is returned when server expects only GET requests,
+// but some other type of request came (Server.GetOnly option is true).
+var ErrGetOnly = errors.New("non-GET request received")
 
 // ReadLimitBody reads request from the given r, limiting the body size.
 //
@@ -881,7 +883,7 @@ func (req *Request) readLimitBody(r *bufio.Reader, maxBodySize int, getOnly bool
 		return err
 	}
 	if getOnly && !req.Header.IsGet() {
-		return errGetOnly
+		return ErrGetOnly
 	}
 
 	if req.MayContinue() {
@@ -1652,6 +1654,11 @@ func appendBodyFixedSize(r *bufio.Reader, dst []byte, n int) ([]byte, error) {
 	}
 }
 
+// ErrBrokenChunk is returned when server receives a broken chunked body (Transfer-Encoding: chunked).
+type ErrBrokenChunk struct {
+	error
+}
+
 func readBodyChunked(r *bufio.Reader, maxBodySize int, dst []byte) ([]byte, error) {
 	if len(dst) > 0 {
 		panic("BUG: expected zero-length buffer")
@@ -1671,7 +1678,9 @@ func readBodyChunked(r *bufio.Reader, maxBodySize int, dst []byte) ([]byte, erro
 			return dst, err
 		}
 		if !bytes.Equal(dst[len(dst)-strCRLFLen:], strCRLF) {
-			return dst, fmt.Errorf("cannot find crlf at the end of chunk")
+			return dst, ErrBrokenChunk{
+				error: fmt.Errorf("cannot find crlf at the end of chunk"),
+			}
 		}
 		dst = dst[:len(dst)-strCRLFLen]
 		if chunkSize == 0 {
@@ -1688,23 +1697,31 @@ func parseChunkSize(r *bufio.Reader) (int, error) {
 	for {
 		c, err := r.ReadByte()
 		if err != nil {
-			return -1, fmt.Errorf("cannot read '\r' char at the end of chunk size: %s", err)
+			return -1, ErrBrokenChunk{
+				error: fmt.Errorf("cannot read '\r' char at the end of chunk size: %s", err),
+			}
 		}
 		// Skip any trailing whitespace after chunk size.
 		if c == ' ' {
 			continue
 		}
 		if c != '\r' {
-			return -1, fmt.Errorf("unexpected char %q at the end of chunk size. Expected %q", c, '\r')
+			return -1, ErrBrokenChunk{
+				error: fmt.Errorf("unexpected char %q at the end of chunk size. Expected %q", c, '\r'),
+			}
 		}
 		break
 	}
 	c, err := r.ReadByte()
 	if err != nil {
-		return -1, fmt.Errorf("cannot read '\n' char at the end of chunk size: %s", err)
+		return -1, ErrBrokenChunk{
+			error: fmt.Errorf("cannot read '\n' char at the end of chunk size: %s", err),
+		}
 	}
 	if c != '\n' {
-		return -1, fmt.Errorf("unexpected char %q at the end of chunk size. Expected %q", c, '\n')
+		return -1, ErrBrokenChunk{
+			error: fmt.Errorf("unexpected char %q at the end of chunk size. Expected %q", c, '\n'),
+		}
 	}
 	return n, nil
 }
diff --git a/server.go b/server.go
index 33347fc..ce70513 100644
--- a/server.go
+++ b/server.go
@@ -152,6 +152,17 @@ type Server struct {
 	// Handler for processing incoming requests.
 	Handler RequestHandler
 
+	// ErrorHandler for returning a response in case of an error while receiving or parsing the request.
+	//
+	// The following is a non-exhaustive list of errors that can be expected as argument:
+	//   * io.EOF
+	//   * io.ErrUnexpectedEOF
+	//   * ErrGetOnly
+	//   * ErrSmallBuffer
+	//   * ErrBodyTooLarge
+	//   * ErrBrokenChunks
+	ErrorHandler func(ctx *RequestCtx, err error)
+
 	// Server name for sending in response headers.
 	//
 	// Default server name is used if left blank.
@@ -1840,7 +1851,7 @@ func (s *Server) serveConn(c net.Conn) error {
 			if err == io.EOF {
 				err = nil
 			} else {
-				bw = writeErrorResponse(bw, ctx, serverName, err)
+				bw = s.writeErrorResponse(bw, ctx, serverName, err)
 			}
 			break
 		}
@@ -1870,7 +1881,7 @@ func (s *Server) serveConn(c net.Conn) error {
 				br = nil
 			}
 			if err != nil {
-				bw = writeErrorResponse(bw, ctx, serverName, err)
+				bw = s.writeErrorResponse(bw, ctx, serverName, err)
 				break
 			}
 		}
@@ -2319,12 +2330,22 @@ func (s *Server) writeFastError(w io.Writer, statusCode int, msg string) {
 		serverDate.Load(), len(msg), msg)
 }
 
-func writeErrorResponse(bw *bufio.Writer, ctx *RequestCtx, serverName []byte, err error) *bufio.Writer {
+func defaultErrorHandler(ctx *RequestCtx, err error) {
 	if _, ok := err.(*ErrSmallBuffer); ok {
 		ctx.Error("Too big request header", StatusRequestHeaderFieldsTooLarge)
 	} else {
 		ctx.Error("Error when parsing request", StatusBadRequest)
 	}
+}
+
+func (s *Server) writeErrorResponse(bw *bufio.Writer, ctx *RequestCtx, serverName []byte, err error) *bufio.Writer {
+	errorHandler := defaultErrorHandler
+	if s.ErrorHandler != nil {
+		errorHandler = s.ErrorHandler
+	}
+
+	errorHandler(ctx, err)
+
 	if serverName != nil {
 		ctx.Response.Header.SetServerBytes(serverName)
 	}


Test patch:
diff --git a/server_test.go b/server_test.go
index fea26d7..3a107aa 100644
--- a/server_test.go
+++ b/server_test.go
@@ -2053,8 +2053,8 @@ func TestServerGetOnly(t *testing.T) {
 		if err == nil {
 			t.Fatalf("expecting error")
 		}
-		if err != errGetOnly {
-			t.Fatalf("Unexpected error from serveConn: %s. Expecting %s", err, errGetOnly)
+		if err != ErrGetOnly {
+			t.Fatalf("Unexpected error from serveConn: %s. Expecting %s", err, ErrGetOnly)
 		}
 	case <-time.After(100 * time.Millisecond):
 		t.Fatalf("timeout")