From abcb28e5068eeac1b60f7fe16f38d971f3737276 Mon Sep 17 00:00:00 2001
From: kevin <wanjunfeng@gmail.com>
Date: Tue, 17 Nov 2020 17:11:06 +0800
Subject: [PATCH] support error customization

---
 go.mod                  |  2 +-
 rest/httpx/responses.go | 22 +++++++++++++++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/go.mod b/go.mod
index ebd064ac..e02204bf 100644
--- a/go.mod
+++ b/go.mod
@@ -55,7 +55,7 @@ require (
 	golang.org/x/tools v0.0.0-20200410132612-ae9902aceb98 // indirect
 	google.golang.org/genproto v0.0.0-20200729003335-053ba62fc06f // indirect
 	google.golang.org/grpc v1.29.1
-	google.golang.org/protobuf v1.25.0 // indirect
+	google.golang.org/protobuf v1.25.0
 	gopkg.in/cheggaaa/pb.v1 v1.0.28
 	gopkg.in/h2non/gock.v1 v1.0.15
 	gopkg.in/yaml.v2 v2.3.0
diff --git a/rest/httpx/responses.go b/rest/httpx/responses.go
index 134c9466..f2b932fa 100644
--- a/rest/httpx/responses.go
+++ b/rest/httpx/responses.go
@@ -3,12 +3,22 @@ package httpx
 import (
 	"encoding/json"
 	"net/http"
+	"sync"
 
 	"github.com/tal-tech/go-zero/core/logx"
 )
 
+var (
+	errorHandler = defaultErrorHandler
+	lock         sync.RWMutex
+)
+
 func Error(w http.ResponseWriter, err error) {
-	http.Error(w, err.Error(), http.StatusBadRequest)
+	lock.RLock()
+	code, body := errorHandler(err)
+	lock.RUnlock()
+
+	WriteJson(w, code, body)
 }
 
 func Ok(w http.ResponseWriter) {
@@ -19,6 +29,12 @@ func OkJson(w http.ResponseWriter, v interface{}) {
 	WriteJson(w, http.StatusOK, v)
 }
 
+func SetErrorHandler(handler func(error) (int, interface{})) {
+	lock.Lock()
+	defer lock.Unlock()
+	errorHandler = handler
+}
+
 func WriteJson(w http.ResponseWriter, code int, v interface{}) {
 	w.Header().Set(ContentType, ApplicationJson)
 	w.WriteHeader(code)
@@ -35,3 +51,7 @@ func WriteJson(w http.ResponseWriter, code int, v interface{}) {
 		logx.Errorf("actual bytes: %d, written bytes: %d", len(bs), n)
 	}
 }
+
+func defaultErrorHandler(err error) (int, interface{}) {
+	return http.StatusBadRequest, err
+}