Add flag to JWT and KeyAuth middleware to allow continuing execution next(c) when error handler decides to swallow the error (returns nil).

Author: aldas

middleware/jwt.go

@@ -21,7 +21,8 @@ type (
 		// BeforeFunc defines a function which is executed just before the middleware.
 		BeforeFunc BeforeFunc
 
-		// SuccessHandler defines a function which is executed for a valid token.
+		// SuccessHandler defines a function which is executed for a valid token before middleware chain continues with next
+		// middleware or handler.
 		SuccessHandler JWTSuccessHandler
 
 		// ErrorHandler defines a function which is executed for an invalid token.
@@ -31,6 +32,13 @@ type (
 		// ErrorHandlerWithContext is almost identical to ErrorHandler, but it's passed the current context.
 		ErrorHandlerWithContext JWTErrorHandlerWithContext
 
+		// NoErrorContinuesExecution allows next middleware/handler to be called when ErrorHandlerWithContext decides to
+		// swallow the error (returns nil).
+		// This is useful in cases when portion of your site/api is publicly accessible and has extra features for authorized
+		// users. In that case you can use ErrorHandlerWithContext to set default public JWT token value to request and
+		// continue with handler chain. Assuming logic downstream execution chain has to check that (public) token value.
+		NoErrorContinuesExecution bool
+
 		// Signing key to validate token.
 		// This is one of the three options to provide a token validation key.
 		// The order of precedence is a user-defined KeyFunc, SigningKeys and SigningKey.
@@ -228,7 +236,11 @@ func JWTWithConfig(config JWTConfig) echo.MiddlewareFunc {
 				return config.ErrorHandler(err)
 			}
 			if config.ErrorHandlerWithContext != nil {
-				return config.ErrorHandlerWithContext(err, c)
+				tmpErr := config.ErrorHandlerWithContext(err, c)
+				if config.NoErrorContinuesExecution && tmpErr == nil {
+					return next(c)
+				}
+				return tmpErr
 			}
 
 			// backwards compatible errors codes

middleware/jwt_test.go

@@ -703,3 +703,77 @@ func TestJWTConfig_SuccessHandler(t *testing.T) {
 		})
 	}
 }
+
+func TestJWTConfig_NoErrorContinuesExecution(t *testing.T) {
+	var testCases = []struct {
+		name                          string
+		whenNoErrorContinuesExecution bool
+		givenToken                    string
+		expectStatus                  int
+		expectBody                    string
+	}{
+		{
+			name:                          "no error handler is called",
+			whenNoErrorContinuesExecution: true,
+			givenToken:                    "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9.TJVA95OrM7E2cBab30RMHrHDcEfxjoYZgeFONFh7HgQ",
+			expectStatus:                  http.StatusTeapot,
+			expectBody:                    "",
+		},
+		{
+			name:                          "NoErrorContinuesExecution is false and error handler is called for missing token",
+			whenNoErrorContinuesExecution: false,
+			givenToken:                    "",
+			// empty response with 200. This emulates previous behaviour when error handler swallowed the error
+			expectStatus: http.StatusOK,
+			expectBody:   "",
+		},
+		{
+			name:                          "error handler is called for missing token",
+			whenNoErrorContinuesExecution: true,
+			givenToken:                    "",
+			expectStatus:                  http.StatusTeapot,
+			expectBody:                    "public-token",
+		},
+		{
+			name:                          "error handler is called for invalid token",
+			whenNoErrorContinuesExecution: true,
+			givenToken:                    "x.x.x",
+			expectStatus:                  http.StatusUnauthorized,
+			expectBody:                    "{\"message\":\"Unauthorized\"}\n",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			e := echo.New()
+
+			e.GET("/", func(c echo.Context) error {
+				testValue, _ := c.Get("test").(string)
+				return c.String(http.StatusTeapot, testValue)
+			})
+
+			e.Use(JWTWithConfig(JWTConfig{
+				NoErrorContinuesExecution: tc.whenNoErrorContinuesExecution,
+				SigningKey:                []byte("secret"),
+				ErrorHandlerWithContext: func(err error, c echo.Context) error {
+					if err == ErrJWTMissing {
+						c.Set("test", "public-token")
+						return nil
+					}
+					return echo.ErrUnauthorized
+				},
+			}))
+
+			req := httptest.NewRequest(http.MethodGet, "/", nil)
+			if tc.givenToken != "" {
+				req.Header.Set(echo.HeaderAuthorization, "bearer "+tc.givenToken)
+			}
+			res := httptest.NewRecorder()
+
+			e.ServeHTTP(res, req)
+
+			assert.Equal(t, tc.expectStatus, res.Code)
+			assert.Equal(t, tc.expectBody, res.Body.String())
+		})
+	}
+}

middleware/key_auth.go

@@ -35,6 +35,13 @@ type (
 		// ErrorHandler defines a function which is executed for an invalid key.
 		// It may be used to define a custom error.
 		ErrorHandler KeyAuthErrorHandler
+
+		// NoErrorContinuesExecution allows next middleware/handler to be called when ErrorHandler decides to swallow
+		// the error (returns nil).
+		// This is useful in cases when portion of your site/api is publicly accessible and has extra features for valid
+		// requests. In that case you can use ErrorHandler to set default public auth values to request and continue with
+		// handler chain. Assuming logic downstream execution chain has to check that (public) auth value.
+		NoErrorContinuesExecution bool
 	}
 
 	// KeyAuthValidator defines a function to validate KeyAuth credentials.
@@ -53,6 +60,21 @@ var (
 	}
 )
 
+// ErrKeyAuthMissing is error type when KeyAuth middleware is unable to extract value from lookups
+type ErrKeyAuthMissing struct {
+	Err error
+}
+
+// Error returns errors text
+func (e *ErrKeyAuthMissing) Error() string {
+	return e.Err.Error()
+}
+
+// Unwrap unwraps error
+func (e *ErrKeyAuthMissing) Unwrap() error {
+	return e.Err
+}
+
 // KeyAuth returns an KeyAuth middleware.
 //
 // For valid key it calls the next handler.
@@ -131,10 +153,15 @@ func KeyAuthWithConfig(config KeyAuthConfig) echo.MiddlewareFunc {
 				} else {
 					err = lastExtractorErr
 				}
+				err = &ErrKeyAuthMissing{Err: err}
 			}
 
 			if config.ErrorHandler != nil {
-				return config.ErrorHandler(err, c)
+				tmpErr := config.ErrorHandler(err, c)
+				if config.NoErrorContinuesExecution && tmpErr == nil {
+					return next(c)
+				}
+				return tmpErr
 			}
 			if lastValidatorErr != nil { // prioritize validator errors over extracting errors
 				return &echo.HTTPError{

middleware/key_auth_test.go

@@ -288,3 +288,78 @@ func TestKeyAuthWithConfig_panicsOnEmptyValidator(t *testing.T) {
 		},
 	)
 }
+
+func TestKeyAuthWithConfig_NoErrorContinuesExecution(t *testing.T) {
+	var testCases = []struct {
+		name                          string
+		whenNoErrorContinuesExecution bool
+		givenKey                      string
+		expectStatus                  int
+		expectBody                    string
+	}{
+		{
+			name:                          "no error handler is called",
+			whenNoErrorContinuesExecution: true,
+			givenKey:                      "valid-key",
+			expectStatus:                  http.StatusTeapot,
+			expectBody:                    "",
+		},
+		{
+			name:                          "NoErrorContinuesExecution is false and error handler is called for missing token",
+			whenNoErrorContinuesExecution: false,
+			givenKey:                      "",
+			// empty response with 200. This emulates previous behaviour when error handler swallowed the error
+			expectStatus: http.StatusOK,
+			expectBody:   "",
+		},
+		{
+			name:                          "error handler is called for missing token",
+			whenNoErrorContinuesExecution: true,
+			givenKey:                      "",
+			expectStatus:                  http.StatusTeapot,
+			expectBody:                    "public-auth",
+		},
+		{
+			name:                          "error handler is called for invalid token",
+			whenNoErrorContinuesExecution: true,
+			givenKey:                      "x.x.x",
+			expectStatus:                  http.StatusUnauthorized,
+			expectBody:                    "{\"message\":\"Unauthorized\"}\n",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			e := echo.New()
+
+			e.GET("/", func(c echo.Context) error {
+				testValue, _ := c.Get("test").(string)
+				return c.String(http.StatusTeapot, testValue)
+			})
+
+			e.Use(KeyAuthWithConfig(KeyAuthConfig{
+				Validator: testKeyValidator,
+				ErrorHandler: func(err error, c echo.Context) error {
+					if _, ok := err.(*ErrKeyAuthMissing); ok {
+						c.Set("test", "public-auth")
+						return nil
+					}
+					return echo.ErrUnauthorized
+				},
+				KeyLookup:                 "header:X-API-Key",
+				NoErrorContinuesExecution: tc.whenNoErrorContinuesExecution,
+			}))
+
+			req := httptest.NewRequest(http.MethodGet, "/", nil)
+			if tc.givenKey != "" {
+				req.Header.Set("X-API-Key", tc.givenKey)
+			}
+			res := httptest.NewRecorder()
+
+			e.ServeHTTP(res, req)
+
+			assert.Equal(t, tc.expectStatus, res.Code)
+			assert.Equal(t, tc.expectBody, res.Body.String())
+		})
+	}
+}