From 73a52bad6c866acba1d1dacb47498c83f089795f Mon Sep 17 00:00:00 2001
From: zhulongcheng <zhulongcheng.me@gmail.com>
Date: Fri, 21 Dec 2018 14:48:58 +0800
Subject: [PATCH] add panic handler

---
 http/router.go      | 19 ++++++++-
 http/router_test.go | 98 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 115 insertions(+), 2 deletions(-)

diff --git a/http/router.go b/http/router.go
index 559f50ca8c..44ff45510d 100644
--- a/http/router.go
+++ b/http/router.go
@@ -1,20 +1,22 @@
 package http
 
 import (
+	"fmt"
 	"net/http"
 
 	"github.com/influxdata/platform"
 	"github.com/julienschmidt/httprouter"
 )
 
-// NewRouter returns a new router with a 404 handler that returns a JSON response.
+// NewRouter returns a new router with a 404 handler and a panic handler.
 func NewRouter() *httprouter.Router {
 	router := httprouter.New()
 	router.NotFound = http.HandlerFunc(notFoundHandler)
+	router.PanicHandler = panicHandler
 	return router
 }
 
-// notFoundHandler represents a 404 handler that return a JSON response
+// notFoundHandler represents a 404 handler that return a JSON response.
 func notFoundHandler(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 	pe := &platform.Error{
@@ -24,3 +26,16 @@ func notFoundHandler(w http.ResponseWriter, r *http.Request) {
 
 	EncodeError(ctx, pe, w)
 }
+
+// panicHandler handles panics recovered from http handlers.
+// It returns a json response with http status code 500 and the recovered error message.
+func panicHandler(w http.ResponseWriter, r *http.Request, rcv interface{}) {
+	ctx := r.Context()
+	pe := &platform.Error{
+		Code: platform.EInternal,
+		Msg:  "a panic has occurred",
+		Err:  fmt.Errorf("%v", rcv),
+	}
+
+	EncodeError(ctx, pe, w)
+}
diff --git a/http/router_test.go b/http/router_test.go
index f69fae494a..a707f816c8 100644
--- a/http/router_test.go
+++ b/http/router_test.go
@@ -97,3 +97,101 @@ func TestRouter_NotFound(t *testing.T) {
 		})
 	}
 }
+
+func TestRouter_Panic(t *testing.T) {
+	type fields struct {
+		method    string
+		path      string
+		handlerFn http.HandlerFunc
+	}
+	type args struct {
+		method string
+		path   string
+	}
+	type wants struct {
+		statusCode  int
+		contentType string
+		body        string
+	}
+
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		wants  wants
+	}{
+		{
+			name: "no panic",
+			fields: fields{
+				method: "GET",
+				path:   "/ping",
+				handlerFn: func(w http.ResponseWriter, r *http.Request) {
+					encodeResponse(r.Context(), w, http.StatusOK, map[string]string{"message": "pong"})
+				},
+			},
+			args: args{
+				method: "GET",
+				path:   "/ping",
+			},
+			wants: wants{
+				statusCode:  http.StatusOK,
+				contentType: "application/json; charset=utf-8",
+				body: `
+{
+  "message": "pong"
+}
+`,
+			},
+		},
+		{
+			name: "panic",
+			fields: fields{
+				method: "GET",
+				path:   "/ping",
+				handlerFn: func(w http.ResponseWriter, r *http.Request) {
+					panic("not implemented")
+				},
+			},
+			args: args{
+				method: "GET",
+				path:   "/ping",
+			},
+			wants: wants{
+				statusCode:  http.StatusInternalServerError,
+				contentType: "application/json; charset=utf-8",
+				body: `
+{
+  "code": "internal error",
+  "msg": "a panic has occurred",
+  "err": "not implemented"
+}`,
+			},
+		},
+	}
+
+	for _, tt := range tests[1:] {
+		t.Run(tt.name, func(t *testing.T) {
+			router := NewRouter()
+			router.HandlerFunc(tt.fields.method, tt.fields.path, tt.fields.handlerFn)
+
+			r := httptest.NewRequest(tt.args.method, tt.args.path, nil)
+			w := httptest.NewRecorder()
+			router.ServeHTTP(w, r)
+
+			res := w.Result()
+			content := res.Header.Get("Content-Type")
+			body, _ := ioutil.ReadAll(res.Body)
+
+			if res.StatusCode != tt.wants.statusCode {
+				t.Errorf("%q. get %v, want %v", tt.name, res.StatusCode, tt.wants.statusCode)
+			}
+			if tt.wants.contentType != "" && content != tt.wants.contentType {
+				t.Errorf("%q. get %v, want %v", tt.name, content, tt.wants.contentType)
+			}
+			if eq, _ := jsonEqual(string(body), tt.wants.body); tt.wants.body != "" && !eq {
+				t.Errorf("%q. get\n***%v***\n,\nwant\n***%v***", tt.name, string(body), tt.wants.body)
+			}
+
+		})
+	}
+}