#457: allow to rewrite system error response

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
 }

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.
@@ -1835,7 +1846,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
 		}
@@ -1865,7 +1876,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
 			}
 		}
@@ -2314,12 +2325,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)
 	}

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")