From a37954be581e1bba0314964bbe14939010c113d9 Mon Sep 17 00:00:00 2001
From: Miguel Bautista <bautista.miguel@gmail.com>
Date: Thu, 15 Sep 2022 11:04:57 -0400
Subject: [PATCH] feat: add panic propagation option to gin middleware

---
 module/apmgin/middleware.go      | 32 +++++++++++++++++++++++---------
 module/apmgin/middleware_test.go | 19 +++++++++++++++++++
 2 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/module/apmgin/middleware.go b/module/apmgin/middleware.go
index 908cd0cdb..1721c572e 100644
--- a/module/apmgin/middleware.go
+++ b/module/apmgin/middleware.go
@@ -44,8 +44,9 @@ func init() {
 // Use WithTracer to specify an alternative tracer.
 func Middleware(engine *gin.Engine, o ...Option) gin.HandlerFunc {
 	m := &middleware{
-		engine: engine,
-		tracer: apm.DefaultTracer(),
+		engine:           engine,
+		tracer:           apm.DefaultTracer(),
+		panicPropagation: false,
 	}
 	for _, o := range o {
 		o(m)
@@ -57,9 +58,10 @@ func Middleware(engine *gin.Engine, o ...Option) gin.HandlerFunc {
 }
 
 type middleware struct {
-	engine         *gin.Engine
-	tracer         *apm.Tracer
-	requestIgnorer apmhttp.RequestIgnorerFunc
+	engine           *gin.Engine
+	tracer           *apm.Tracer
+	requestIgnorer   apmhttp.RequestIgnorerFunc
+	panicPropagation bool
 }
 
 func (m *middleware) handle(c *gin.Context) {
@@ -75,17 +77,20 @@ func (m *middleware) handle(c *gin.Context) {
 
 	defer func() {
 		if v := recover(); v != nil {
-			if !c.Writer.Written() {
-				c.AbortWithStatus(http.StatusInternalServerError)
+			if m.panicPropagation {
+				defer panic(v)
 			} else {
-				c.Abort()
+				if !c.Writer.Written() {
+					c.AbortWithStatus(http.StatusInternalServerError)
+				} else {
+					c.Abort()
+				}
 			}
 			e := m.tracer.Recovered(v)
 			e.SetTransaction(tx)
 			setContext(&e.Context, c, body)
 			e.Send()
 		}
-		c.Writer.WriteHeaderNow()
 		tx.Result = apmhttp.StatusCodeResult(c.Writer.Status())
 
 		if tx.Sampled() {
@@ -144,3 +149,12 @@ func WithRequestIgnorer(r apmhttp.RequestIgnorerFunc) Option {
 		m.requestIgnorer = r
 	}
 }
+
+// WithPanicPropagation returns an Option which enable panic propagation.
+// Any panic will be recovered and recorded as an error in a transaction, then
+// panic will be caused again.
+func WithPanicPropagation() Option {
+	return func(m *middleware) {
+		m.panicPropagation = true
+	}
+}
diff --git a/module/apmgin/middleware_test.go b/module/apmgin/middleware_test.go
index c53113a19..b8d6dd0f8 100644
--- a/module/apmgin/middleware_test.go
+++ b/module/apmgin/middleware_test.go
@@ -219,6 +219,25 @@ func TestMiddlewarePanicHeadersSent(t *testing.T) {
 	assert.NotContains(t, debugOutput.String(), "[WARNING] Headers were already written")
 }
 
+func TestMiddlewarePanicPropagation(t *testing.T) {
+	debugOutput.Reset()
+	tracer, transport := transporttest.NewRecorderTracer()
+	defer tracer.Close()
+
+	e := gin.New()
+	e.Use(
+		gin.CustomRecovery(func(c *gin.Context, err interface{}) {
+			c.AbortWithStatus(http.StatusNotImplemented)
+		}),
+		apmgin.Middleware(e, apmgin.WithTracer(tracer), apmgin.WithPanicPropagation()))
+	e.GET("/panic", handlePanic)
+
+	w := doRequest(e, "GET", "http://server.testing/panic")
+	assert.Equal(t, http.StatusNotImplemented, w.Code)
+	tracer.Flush(nil)
+	assertError(t, transport.Payloads(), "handlePanic", "boom", false)
+}
+
 func TestMiddlewareError(t *testing.T) {
 	debugOutput.Reset()
 	tracer, transport := transporttest.NewRecorderTracer()