Added ErrorHandler and ErrorHandlerWithContext in CSRF middleware

middleware/csrf.go

@@ -61,7 +61,17 @@ type (
 		// Indicates SameSite mode of the CSRF cookie.
 		// Optional. Default value SameSiteDefaultMode.
 		CookieSameSite http.SameSite `yaml:"cookie_same_site"`
+
+		// ErrorHandler defines a function which is executed for returning custom errors.
+		ErrorHandler CSRFErrorHandler
+
+		// ErrorHandlerWithContext is almost identical to ErrorHandler, but it's passed the current context.
+		ErrorHandlerWithContext CSRFErrorHandlerWithContext
 	}
+
+	CSRFErrorHandler func(err error) error
+
+	CSRFErrorHandlerWithContext func(err error, c echo.Context) error
 )
 
 // ErrCSRFInvalid is returned when CSRF check fails
@@ -154,8 +164,9 @@ func CSRFWithConfig(config CSRFConfig) echo.MiddlewareFunc {
 						lastTokenErr = ErrCSRFInvalid
 					}
 				}
+				var finalErr error
 				if lastTokenErr != nil {
-					return lastTokenErr
+					finalErr = lastTokenErr
 				} else if lastExtractorErr != nil {
 					// ugly part to preserve backwards compatible errors. someone could rely on them
 					if lastExtractorErr == errQueryExtractorValueMissing {
@@ -167,7 +178,17 @@ func CSRFWithConfig(config CSRFConfig) echo.MiddlewareFunc {
 					} else {
 						lastExtractorErr = echo.NewHTTPError(http.StatusBadRequest, lastExtractorErr.Error())
 					}
-					return lastExtractorErr
+					finalErr = lastExtractorErr
+				}
+
+				if finalErr != nil {
+					if config.ErrorHandler != nil {
+						return config.ErrorHandler(finalErr)
+					}
+					if config.ErrorHandlerWithContext != nil {
+						return config.ErrorHandlerWithContext(finalErr, c)
+					}
+					return finalErr
 				}
 			}
 

middleware/csrf_test.go

@@ -1,6 +1,7 @@
 package middleware
 
 import (
+	"encoding/json"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -358,3 +359,53 @@ func TestCSRFConfig_skipper(t *testing.T) {
 		})
 	}
 }
+
+func TestCSRFErrorHandling(t *testing.T) {
+	testCases := []struct {
+		name        string
+		cfg         CSRFConfig
+		expectedErr *echo.HTTPError
+	}{
+		{
+			name: "ok, ErrorHandler is executed",
+			cfg: CSRFConfig{
+				ErrorHandler: func(err error) error {
+					return echo.NewHTTPError(http.StatusTeapot, "error_handler_executed")
+				},
+			},
+			expectedErr: echo.NewHTTPError(http.StatusTeapot, "error_handler_executed"),
+		},
+		{
+			name: "ok, ErrorHandlerWithContext is executed",
+			cfg: CSRFConfig{
+				ErrorHandlerWithContext: func(err error, c echo.Context) error {
+					return echo.NewHTTPError(http.StatusTeapot, "error_handler_with_context_executed")
+				},
+			},
+			expectedErr: echo.NewHTTPError(http.StatusTeapot, "error_handler_with_context_executed"),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			e := echo.New()
+			e.POST("/", func(c echo.Context) error {
+				return c.String(http.StatusNotImplemented, "should not end up here")
+			})
+
+			e.Use(CSRFWithConfig(tc.cfg))
+
+			req := httptest.NewRequest(http.MethodPost, "/", nil)
+			res := httptest.NewRecorder()
+			e.ServeHTTP(res, req)
+
+			var response struct {
+				Message string `json:"message"`
+			}
+
+			assert.NoError(t, json.Unmarshal(res.Body.Bytes(), &response))
+			assert.Equal(t, tc.expectedErr.Code, res.Code)
+			assert.Equal(t, tc.expectedErr.Message, response.Message)
+		})
+	}
+}