Merge pull request #47 from LyricTian/develop

optimization of error handling

Author: LyricTian

README.md

@@ -43,6 +43,7 @@ import (
 	"log"
 	"net/http"
 
+	"gopkg.in/oauth2.v3/errors"
 	"gopkg.in/oauth2.v3/manage"
 	"gopkg.in/oauth2.v3/models"
 	"gopkg.in/oauth2.v3/server"
@@ -67,8 +68,13 @@ func main() {
 	srv.SetAllowGetAccessRequest(true)
 	srv.SetClientInfoHandler(server.ClientFormHandler)
 
-	srv.SetInternalErrorHandler(func(err error) {
-		log.Println("OAuth2 Error:", err.Error())
+	srv.SetInternalErrorHandler(func(err error) (re *errors.Response) {
+		log.Println("Internal Error:", err.Error())
+		return
+	})
+
+	srv.SetResponseErrorHandler(func(re *errors.Response) {
+		log.Println("Response Error:", re.Error.Error())
 	})
 
 	http.HandleFunc("/authorize", func(w http.ResponseWriter, r *http.Request) {
@@ -82,8 +88,9 @@ func main() {
 		srv.HandleTokenRequest(w, r)
 	})
 
-	http.ListenAndServe(":9096", nil)
+	log.Fatal(http.ListenAndServe(":9096", nil))
 }
+
 ```
 
 ### Build and run
@@ -139,8 +146,8 @@ Copyright (c) 2016 Lyric
 [License-Image]: https://img.shields.io/npm/l/express.svg
 [Build-Status-Url]: https://travis-ci.org/go-oauth2/oauth2
 [Build-Status-Image]: https://travis-ci.org/go-oauth2/oauth2.svg?branch=master
-[Release-Url]: https://github.com/go-oauth2/oauth2/releases/tag/v3.6.3
-[Release-image]: http://img.shields.io/badge/release-v3.6.3-1eb0fc.svg
+[Release-Url]: https://github.com/go-oauth2/oauth2/releases/tag/v3.7.0
+[Release-image]: http://img.shields.io/badge/release-v3.7.0-1eb0fc.svg
 [ReportCard-Url]: https://goreportcard.com/report/gopkg.in/oauth2.v3
 [ReportCard-Image]: https://goreportcard.com/badge/gopkg.in/oauth2.v3
 [GoDoc-Url]: https://godoc.org/gopkg.in/oauth2.v3

errors/response.go

@@ -1,6 +1,9 @@
 package errors
 
-import "errors"
+import (
+	"errors"
+	"net/http"
+)
 
 // Response error response
 type Response struct {
@@ -9,6 +12,24 @@ type Response struct {
 	Description string
 	URI         string
 	StatusCode  int
+	Header      http.Header
+}
+
+// NewResponse create the response pointer
+func NewResponse(err error, statusCode int) *Response {
+	return &Response{
+		Error:      err,
+		StatusCode: statusCode,
+	}
+}
+
+// SetHeader sets the header entries associated with key to
+// the single element value.
+func (r *Response) SetHeader(key, value string) {
+	if r.Header == nil {
+		r.Header = make(http.Header)
+	}
+	r.Header.Set(key, value)
 }
 
 // https://tools.ietf.org/html/rfc6749#section-5.2

example/server/server.go

@@ -6,6 +6,7 @@ import (
 	"net/url"
 	"os"
 
+	"gopkg.in/oauth2.v3/errors"
 	"gopkg.in/oauth2.v3/manage"
 	"gopkg.in/oauth2.v3/models"
 	"gopkg.in/oauth2.v3/server"
@@ -37,8 +38,14 @@ func main() {
 
 	srv := server.NewServer(server.NewConfig(), manager)
 	srv.SetUserAuthorizationHandler(userAuthorizeHandler)
-	srv.SetInternalErrorHandler(func(err error) {
-		log.Println("[oauth2] error:", err.Error())
+
+	srv.SetInternalErrorHandler(func(err error) (re *errors.Response) {
+		log.Println("Internal Error:", err.Error())
+		return
+	})
+
+	srv.SetResponseErrorHandler(func(re *errors.Response) {
+		log.Println("Response Error:", re.Error.Error())
 	})
 
 	http.HandleFunc("/login", loginHandler)

server/handler.go

@@ -28,10 +28,10 @@ type (
 	RefreshingScopeHandler func(newScope, oldScope string) (allowed bool, err error)
 
 	// ResponseErrorHandler response error handing
-	ResponseErrorHandler func(err error) (re *errors.Response)
+	ResponseErrorHandler func(re *errors.Response)
 
 	// InternalErrorHandler internal error handing
-	InternalErrorHandler func(err error)
+	InternalErrorHandler func(err error) (re *errors.Response)
 
 	// AuthorizeScopeHandler set the authorized scope
 	AuthorizeScopeHandler func(w http.ResponseWriter, r *http.Request) (scope string, err error)

server/server.go

@@ -65,7 +65,7 @@ func (s *Server) redirectError(w http.ResponseWriter, req *AuthorizeRequest, err
 		uerr = err
 		return
 	}
-	data, _ := s.GetErrorData(err)
+	data, _, _ := s.GetErrorData(err)
 	err = s.redirect(w, req, data)
 	return
 }
@@ -81,16 +81,21 @@ func (s *Server) redirect(w http.ResponseWriter, req *AuthorizeRequest, data map
 }
 
 func (s *Server) tokenError(w http.ResponseWriter, err error) (uerr error) {
-	data, statusCode := s.GetErrorData(err)
-	uerr = s.token(w, data, statusCode)
+	data, statusCode, header := s.GetErrorData(err)
+
+	uerr = s.token(w, data, header, statusCode)
 	return
 }
 
-func (s *Server) token(w http.ResponseWriter, data map[string]interface{}, statusCode ...int) (err error) {
+func (s *Server) token(w http.ResponseWriter, data map[string]interface{}, header http.Header, statusCode ...int) (err error) {
 	w.Header().Set("Content-Type", "application/json;charset=UTF-8")
 	w.Header().Set("Cache-Control", "no-store")
 	w.Header().Set("Pragma", "no-cache")
 
+	for key := range header {
+		w.Header().Set(key, header.Get(key))
+	}
+
 	status := http.StatusOK
 	if len(statusCode) > 0 && statusCode[0] > 0 {
 		status = statusCode[0]
@@ -490,38 +495,45 @@ func (s *Server) HandleTokenRequest(w http.ResponseWriter, r *http.Request) (err
 		err = s.tokenError(w, verr)
 		return
 	}
-	err = s.token(w, s.GetTokenData(ti))
+
+	err = s.token(w, s.GetTokenData(ti), nil)
 	return
 }
 
 // GetErrorData get error response data
-func (s *Server) GetErrorData(err error) (data map[string]interface{}, statusCode int) {
-	re := &errors.Response{}
+func (s *Server) GetErrorData(err error) (data map[string]interface{}, statusCode int, header http.Header) {
+	re := new(errors.Response)
 
 	if v, ok := errors.Descriptions[err]; ok {
 		re.Error = err
 		re.Description = v
 		re.StatusCode = errors.StatusCodes[err]
 	} else {
 		if fn := s.InternalErrorHandler; fn != nil {
-			fn(err)
+			if vre := fn(err); vre != nil {
+				re = vre
+			}
+		}
+
+		if re.Error == nil {
+			re.Error = errors.ErrServerError
+			re.Description = errors.Descriptions[errors.ErrServerError]
+			re.StatusCode = errors.StatusCodes[errors.ErrServerError]
 		}
 	}
 
 	if fn := s.ResponseErrorHandler; fn != nil {
-		if vre := fn(err); vre != nil {
-			re = vre
+		fn(re)
+
+		if re == nil {
+			re = new(errors.Response)
 		}
 	}
 
-	if re.Error == nil {
-		re.Error = errors.ErrServerError
-		re.Description = errors.Descriptions[errors.ErrServerError]
-		re.StatusCode = errors.StatusCodes[errors.ErrServerError]
-	}
+	data = make(map[string]interface{})
 
-	data = map[string]interface{}{
-		"error": re.Error.Error(),
+	if err := re.Error; err != nil {
+		data["error"] = err.Error()
 	}
 
 	if v := re.ErrorCode; v != 0 {
@@ -536,11 +548,13 @@ func (s *Server) GetErrorData(err error) (data map[string]interface{}, statusCod
 		data["error_uri"] = v
 	}
 
-	statusCode = 400
+	header = re.Header
 
+	statusCode = http.StatusInternalServerError
 	if v := re.StatusCode; v > 0 {
 		statusCode = v
 	}
+
 	return
 }
 

server/server_test.go

@@ -1,14 +1,15 @@
 package server_test
 
 import (
-	"errors"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"testing"
 
 	"github.com/gavv/httpexpect"
 	"gopkg.in/oauth2.v3"
+	"gopkg.in/oauth2.v3/errors"
 	"gopkg.in/oauth2.v3/manage"
 	"gopkg.in/oauth2.v3/models"
 	"gopkg.in/oauth2.v3/server"
@@ -144,7 +145,7 @@ func TestPasswordCredentials(t *testing.T) {
 			userID = "000000"
 			return
 		}
-		err = errors.New("user not found")
+		err = fmt.Errorf("user not found")
 		return
 	})
 
@@ -174,9 +175,16 @@ func TestClientCredentials(t *testing.T) {
 
 	srv = server.NewDefaultServer(manager)
 	srv.SetClientInfoHandler(server.ClientFormHandler)
-	srv.SetInternalErrorHandler(func(err error) {
+
+	srv.SetInternalErrorHandler(func(err error) (re *errors.Response) {
 		t.Log("OAuth 2.0 Error:", err.Error())
+		return
 	})
+
+	srv.SetResponseErrorHandler(func(re *errors.Response) {
+		t.Log("Response Error:", re.Error)
+	})
+
 	srv.SetAllowedGrantType(oauth2.ClientCredentials)
 	srv.SetAllowGetAccessRequest(false)
 	srv.SetExtensionFieldsHandler(func(ti oauth2.TokenInfo) (fieldsValue map[string]interface{}) {