feat: support errors format for the recovery middleware (#304)

pkg/app/middlewares/server/recovery/option.go

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2022 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package recovery
+
+import (
+	"context"
+	"time"
+
+	"github.com/cloudwego/hertz/pkg/app"
+	"github.com/cloudwego/hertz/pkg/common/hlog"
+	"github.com/cloudwego/hertz/pkg/protocol/consts"
+)
+
+type (
+	options struct {
+		recoveryHandler func(c context.Context, ctx *app.RequestContext, err interface{}, stack []byte)
+	}
+
+	Option func(o *options)
+)
+
+func defaultRecoveryHandler(c context.Context, ctx *app.RequestContext, err interface{}, stack []byte) {
+	hlog.SystemLogger().CtxErrorf(c, "[Recovery] %s panic recovered:\n%s\n%s\n",
+		timeFormat(time.Now()), err, stack)
+	ctx.AbortWithStatus(consts.StatusInternalServerError)
+}
+
+func newOptions(opts ...Option) *options {
+	cfg := &options{
+		recoveryHandler: defaultRecoveryHandler,
+	}
+
+	for _, opt := range opts {
+		opt(cfg)
+	}
+
+	return cfg
+}
+
+func WithRecoveryHandler(f func(c context.Context, ctx *app.RequestContext, err interface{}, stack []byte)) Option {
+	return func(o *options) {
+		o.recoveryHandler = f
+	}
+}

pkg/app/middlewares/server/recovery/option_test.go

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2022 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package recovery
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/cloudwego/hertz/pkg/app"
+	"github.com/cloudwego/hertz/pkg/common/hlog"
+	"github.com/cloudwego/hertz/pkg/common/test/assert"
+	"github.com/cloudwego/hertz/pkg/common/utils"
+)
+
+func TestDefaultOption(t *testing.T) {
+	opts := newOptions()
+	assert.DeepEqual(t, fmt.Sprintf("%p", defaultRecoveryHandler), fmt.Sprintf("%p", opts.recoveryHandler))
+}
+
+func newRecoveryHandler(c context.Context, ctx *app.RequestContext, err interface{}, stack []byte) {
+	hlog.SystemLogger().CtxErrorf(c, "[New Recovery] %s panic recovered:\n%s\n%s\n",
+		timeFormat(time.Now()), err, stack)
+	ctx.JSON(501, utils.H{"msg": err.(string)})
+}
+
+func TestOption(t *testing.T) {
+	opts := newOptions(WithRecoveryHandler(newRecoveryHandler))
+	assert.DeepEqual(t, fmt.Sprintf("%p", newRecoveryHandler), fmt.Sprintf("%p", opts.recoveryHandler))
+}

pkg/app/middlewares/server/recovery/recovery.go

@@ -25,8 +25,6 @@ import (
 	"time"
 
 	"github.com/cloudwego/hertz/pkg/app"
-	"github.com/cloudwego/hertz/pkg/common/hlog"
-	"github.com/cloudwego/hertz/pkg/protocol/consts"
 )
 
 var (
@@ -36,16 +34,18 @@ var (
 	slash     = []byte("/")
 )
 
-// Recovery returns a middleware that recovers from any panic and writes a 500 if there was one.
-func Recovery() app.HandlerFunc {
+// Recovery returns a middleware that recovers from any panic.
+// By default, it will print the time, content, and stack information of the error and write a 500.
+// Overriding the Config configuration, you can customize the error printing logic.
+func Recovery(opts ...Option) app.HandlerFunc {
+	cfg := newOptions(opts...)
+
 	return func(c context.Context, ctx *app.RequestContext) {
 		defer func() {
 			if err := recover(); err != nil {
 				stack := stack(3)
 
-				hlog.SystemLogger().CtxErrorf(c, "[Recovery] %s panic recovered:\n%s\n%s\n",
-					timeFormat(time.Now()), err, stack)
-				ctx.AbortWithStatus(consts.StatusInternalServerError)
+				cfg.recoveryHandler(c, ctx, err, stack)
 			}
 		}()
 		ctx.Next(c)

pkg/app/middlewares/server/recovery/recovery_test.go

@@ -22,6 +22,7 @@ import (
 	"testing"
 
 	"github.com/cloudwego/hertz/pkg/app"
+	"github.com/cloudwego/hertz/pkg/common/test/assert"
 )
 
 func TestRecovery(t *testing.T) {
@@ -39,3 +40,20 @@ func TestRecovery(t *testing.T) {
 		t.Fatalf("unexpected %v. Expecting %v", ctx.Response.StatusCode(), 500)
 	}
 }
+
+func TestWithRecoveryHandler(t *testing.T) {
+	ctx := app.NewContext(0)
+	var hc app.HandlersChain
+	hc = append(hc, func(c context.Context, ctx *app.RequestContext) {
+		fmt.Println("this is test")
+		panic("test")
+	})
+	ctx.SetHandlers(hc)
+
+	Recovery(WithRecoveryHandler(newRecoveryHandler))(context.Background(), ctx)
+
+	if ctx.Response.StatusCode() != 501 {
+		t.Fatalf("unexpected %v. Expecting %v", ctx.Response.StatusCode(), 501)
+	}
+	assert.DeepEqual(t, "{\"msg\":\"test\"}", string(ctx.Response.Body()))
+}