Repository: google/exposure-notifications-server
Pull request #1479 (2021-03-22 19:42:17)

Add recovery middleware for catching panics
Similar to what we do on the verification server already.

**Release Note**

<!-- Enter your extended release note in the below block. If the PR requires
additional action from users switching to the new release, include the string
"action required". If no release note is required, write "NONE". -->

```release-note
Add recovery middleware for catching runtime panics in HTTP handlers.
```


Patch:
diff --git a/internal/debugger/server.go b/internal/debugger/server.go
index 0f0926a..de216b0 100644
--- a/internal/debugger/server.go
+++ b/internal/debugger/server.go
@@ -59,6 +59,7 @@ func (s *Server) Routes(ctx context.Context) *mux.Router {
 	logger := logging.FromContext(ctx).Named("debugger")
 
 	r := mux.NewRouter()
+	r.Use(middleware.Recovery())
 	r.Use(middleware.PopulateRequestID())
 	r.Use(middleware.PopulateObservability())
 	r.Use(middleware.PopulateLogger(logger))
diff --git a/internal/export/server.go b/internal/export/server.go
index e84b6e6..ca998a9 100644
--- a/internal/export/server.go
+++ b/internal/export/server.go
@@ -57,6 +57,7 @@ func (s *Server) Routes(ctx context.Context) *mux.Router {
 	logger := logging.FromContext(ctx).Named("export")
 
 	r := mux.NewRouter()
+	r.Use(middleware.Recovery())
 	r.Use(middleware.PopulateRequestID())
 	r.Use(middleware.PopulateObservability())
 	r.Use(middleware.PopulateLogger(logger))
diff --git a/internal/exportimport/server.go b/internal/exportimport/server.go
index bcaac29..cccdf46 100644
--- a/internal/exportimport/server.go
+++ b/internal/exportimport/server.go
@@ -68,6 +68,7 @@ func (s *Server) Routes(ctx context.Context) *mux.Router {
 	logger := logging.FromContext(ctx).Named("exportimport")
 
 	r := mux.NewRouter()
+	r.Use(middleware.Recovery())
 	r.Use(middleware.PopulateRequestID())
 	r.Use(middleware.PopulateObservability())
 	r.Use(middleware.PopulateLogger(logger))
diff --git a/internal/federationin/server.go b/internal/federationin/server.go
index 30e9cd5..e8d43f2 100644
--- a/internal/federationin/server.go
+++ b/internal/federationin/server.go
@@ -46,6 +46,7 @@ func (s *Server) Routes(ctx context.Context) *mux.Router {
 	logger := logging.FromContext(ctx).Named("federationin")
 
 	r := mux.NewRouter()
+	r.Use(middleware.Recovery())
 	r.Use(middleware.PopulateRequestID())
 	r.Use(middleware.PopulateObservability())
 	r.Use(middleware.PopulateLogger(logger))
diff --git a/internal/generate/server.go b/internal/generate/server.go
index cb2709f..5091d53 100644
--- a/internal/generate/server.go
+++ b/internal/generate/server.go
@@ -57,6 +57,7 @@ func (s *Server) Routes(ctx context.Context) *mux.Router {
 	logger := logging.FromContext(ctx).Named("generate")
 
 	r := mux.NewRouter()
+	r.Use(middleware.Recovery())
 	r.Use(middleware.PopulateRequestID())
 	r.Use(middleware.PopulateObservability())
 	r.Use(middleware.PopulateLogger(logger))
diff --git a/internal/jwks/server.go b/internal/jwks/server.go
index d42044f..3679378 100644
--- a/internal/jwks/server.go
+++ b/internal/jwks/server.go
@@ -56,6 +56,7 @@ func (s *Server) Routes(ctx context.Context) *mux.Router {
 	logger := logging.FromContext(ctx).Named("jwks")
 
 	r := mux.NewRouter()
+	r.Use(middleware.Recovery())
 	r.Use(middleware.PopulateRequestID())
 	r.Use(middleware.PopulateObservability())
 	r.Use(middleware.PopulateLogger(logger))
diff --git a/internal/keyrotation/server.go b/internal/keyrotation/server.go
index b756b35..ea1eac8 100644
--- a/internal/keyrotation/server.go
+++ b/internal/keyrotation/server.go
@@ -69,6 +69,7 @@ func (s *Server) Routes(ctx context.Context) *mux.Router {
 	logger := logging.FromContext(ctx).Named("keyrotation")
 
 	r := mux.NewRouter()
+	r.Use(middleware.Recovery())
 	r.Use(middleware.PopulateRequestID())
 	r.Use(middleware.PopulateObservability())
 	r.Use(middleware.PopulateLogger(logger))
diff --git a/internal/middleware/recovery.go b/internal/middleware/recovery.go
new file mode 100644
index 0000000..663d882
--- /dev/null
+++ b/internal/middleware/recovery.go
@@ -0,0 +1,46 @@
+// Copyright 2021 Google LLC
+//
+// 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 middleware
+
+import (
+	"net/http"
+
+	"github.com/google/exposure-notifications-server/pkg/logging"
+	"github.com/gorilla/mux"
+)
+
+// Recovery recovers from panics and other fatal errors. It keeps the server and
+// service running, returning 500 to the caller while also logging the error in
+// a structured format.
+func Recovery() mux.MiddlewareFunc {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
+
+			logger := logging.FromContext(ctx).Named("middleware.Recover")
+
+			defer func() {
+				if p := recover(); p != nil {
+					logger.Errorw("http handler panic", "panic", p)
+					w.WriteHeader(http.StatusInternalServerError)
+					return
+				}
+			}()
+
+			next.ServeHTTP(w, r)
+			return
+		})
+	}
+}
diff --git a/internal/mirror/server.go b/internal/mirror/server.go
index 01d4ebf..48ce1e9 100644
--- a/internal/mirror/server.go
+++ b/internal/mirror/server.go
@@ -62,6 +62,7 @@ func (s *Server) Routes(ctx context.Context) *mux.Router {
 	logger := logging.FromContext(ctx).Named("mirror")
 
 	r := mux.NewRouter()
+	r.Use(middleware.Recovery())
 	r.Use(middleware.PopulateRequestID())
 	r.Use(middleware.PopulateObservability())
 	r.Use(middleware.PopulateLogger(logger))
diff --git a/internal/publish/publish.go b/internal/publish/publish.go
index 2489d91..34bfbba 100644
--- a/internal/publish/publish.go
+++ b/internal/publish/publish.go
@@ -136,6 +136,7 @@ func (s *Server) Routes(ctx context.Context) *mux.Router {
 	logger := logging.FromContext(ctx).Named("publish")
 
 	r := mux.NewRouter()
+	r.Use(middleware.Recovery())
 	r.Use(middleware.ProcessChaff(s.tracker))
 	r.Use(middleware.PopulateRequestID())
 	r.Use(middleware.PopulateObservability())


Test patch:
diff --git a/internal/middleware/recovery_test.go b/internal/middleware/recovery_test.go
new file mode 100644
index 0000000..28dd6a2
--- /dev/null
+++ b/internal/middleware/recovery_test.go
@@ -0,0 +1,74 @@
+// Copyright 2021 Google LLC
+//
+// 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 middleware_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/google/exposure-notifications-server/internal/middleware"
+	"github.com/google/exposure-notifications-server/internal/project"
+)
+
+func TestRecovery(t *testing.T) {
+	t.Parallel()
+
+	ctx := project.TestContext(t)
+
+	m := middleware.Recovery()
+
+	cases := []struct {
+		name    string
+		handler http.Handler
+		code    int
+	}{
+		{
+			name: "default",
+			handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(200)
+			}),
+			code: http.StatusOK,
+		},
+		{
+			name: "panic",
+			handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				panic("oops")
+			}),
+			code: http.StatusInternalServerError,
+		},
+	}
+
+	for _, tc := range cases {
+		tc := tc
+
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			r, err := http.NewRequestWithContext(ctx, http.MethodGet, "/", nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			w := httptest.NewRecorder()
+
+			m(tc.handler).ServeHTTP(w, r)
+
+			if got, want := w.Code, tc.code; got != want {
+				t.Errorf("expected %d to be %d", got, want)
+			}
+		})
+	}
+}