From 42c02541d18e97b0eac54e57e37825807d6ff3d9 Mon Sep 17 00:00:00 2001
From: Daniel Bedrood <danielb@scribd.com>
Date: Tue, 18 Jun 2024 16:32:38 +0200
Subject: [PATCH] feat: Add recovery middleware for GRPC server

---
 pkg/interceptors/recovery.go | 47 ++++++++++++++++++++++++++++++++++++
 pkg/middleware/recovery.go   | 46 +++++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)
 create mode 100644 pkg/interceptors/recovery.go
 create mode 100644 pkg/middleware/recovery.go

diff --git a/pkg/interceptors/recovery.go b/pkg/interceptors/recovery.go
new file mode 100644
index 0000000..759bed5
--- /dev/null
+++ b/pkg/interceptors/recovery.go
@@ -0,0 +1,47 @@
+package interceptors
+
+import (
+	"context"
+	stdlog "log"
+	"runtime/debug"
+	"time"
+
+	"github.com/getsentry/sentry-go"
+	grpcrecovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	sdkloggercontext "github.com/scribd/go-sdk/pkg/context/logger"
+)
+
+// RecoveryUnaryServerInterceptor returns a unary server interceptor that recovers from panics,
+// sends a sentry event, log in fatal level and halts the service.
+// IMPORTANT: This interceptor should be the last one in the interceptor chain.
+func RecoveryUnaryServerInterceptor() grpc.UnaryServerInterceptor {
+	return grpcrecovery.UnaryServerInterceptor(recoveryOption...)
+}
+
+// RecoveryStreamServerInterceptor returns a streaming server interceptor that recovers from panics,
+// sends a sentry event, log in fatal level and halts the service.
+// IMPORTANT: This interceptor should be the last one in the interceptor chain.
+func RecoveryStreamServerInterceptor() grpc.StreamServerInterceptor {
+	return grpcrecovery.StreamServerInterceptor(recoveryOption...)
+}
+
+var recoveryOption = []grpcrecovery.Option{
+	grpcrecovery.WithRecoveryHandlerContext(func(ctx context.Context, rec interface{}) (err error) {
+		sentry.CurrentHub().Recover(rec)
+		sentry.Flush(time.Second * 5)
+
+		l, err := sdkloggercontext.Extract(ctx)
+		if err != nil {
+			debug.PrintStack()
+			stdlog.Printf("logger not found in context: %v\n", err)
+			stdlog.Fatalf("grpc: panic error: %v", rec)
+		}
+
+		l.Fatalf("panic error: %v", rec)
+		return status.Errorf(codes.Internal, "")
+	}),
+}
diff --git a/pkg/middleware/recovery.go b/pkg/middleware/recovery.go
new file mode 100644
index 0000000..7bd05ef
--- /dev/null
+++ b/pkg/middleware/recovery.go
@@ -0,0 +1,46 @@
+package middleware
+
+import (
+	"log"
+	"net/http"
+	"runtime/debug"
+	"time"
+
+	"github.com/getsentry/sentry-go"
+
+	sdkloggercontext "github.com/scribd/go-sdk/pkg/context/logger"
+)
+
+// RecoveryMiddleware is a middleware that recovers from panics and logs them.
+type RecoveryMiddleware struct{}
+
+// NewRecoveryMiddleware is a constructor used to build a RecoveryMiddleware.
+// IMPORTANT: This middleware should be the last one in the middleware chain.
+func NewRecoveryMiddleware() RecoveryMiddleware {
+	return RecoveryMiddleware{}
+}
+
+// Handler implements the middlewares.Handlerer interface: it returns a
+// http.Handler to be mounted as middleware. The Handler recovers from a panic,
+// sends a sentry event, sends fatal error log and halts the service.
+func (rm RecoveryMiddleware) Handler(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer func() {
+			if rec := recover(); rec != nil {
+				sentry.CurrentHub().Recover(rec)
+				sentry.Flush(time.Second * 5)
+
+				l, err := sdkloggercontext.Extract(r.Context())
+				if err != nil {
+					debug.PrintStack()
+					log.Printf("logger not found in context: %v\n", err)
+					log.Fatalf("http: panic serving URI %s: %v", r.URL.RequestURI(), rec)
+				}
+
+				l.Fatalf("http: panic serving URI %s: %v", r.URL.RequestURI(), rec)
+			}
+		}()
+
+		next.ServeHTTP(w, r)
+	})
+}