Merge pull request #1 from delvelabs/exception-handler-callback

Add a method to register a custom exception handler.

Author: GuilloOme

flask_oauthlib/provider/oauth2.py

@@ -72,6 +72,7 @@ def user():
     def __init__(self, app=None):
         self._before_request_funcs = []
         self._after_request_funcs = []
+        self._exception_handler = None
         self._invalid_response = None
         if app:
             self.init_app(app)
@@ -85,6 +86,13 @@ def init_app(self, app):
         app.extensions = getattr(app, 'extensions', {})
         app.extensions['oauthlib.provider.oauth2'] = self
 
+    def _on_exception(self, error, redirect_content=None):
+
+        if self._exception_handler:
+            return self._exception_handler(error, redirect_content)
+        else:
+            return redirect(redirect_content)
+
     @cached_property
     def error_uri(self):
         """The error page URI.
@@ -208,6 +216,34 @@ def valid_after_request(valid, oauth):
         self._after_request_funcs.append(f)
         return f
 
+    def exception_handler(self, f):
+        """Register a function as custom exception handler.
+
+        **As the default error handling is leaking error to the client, it is
+        STRONGLY RECOMMENDED to implement your own handler to mask
+        the server side errors in production environment.**
+
+        When an error occur during execution, we can
+        handle the error with with the registered function. The function
+        accepts two parameters:
+            - error: the error raised
+            - redirect_content: the content used in the redirect by default
+
+        usage with the flask error handler ::
+            @oauth.exception_handler
+            def custom_exception_handler(error, *args):
+                raise error
+
+            @app.errorhandler(Exception)
+            def all_exception_handler(*args):
+                # any treatment you need for the error
+                return "Server error", 500
+
+        If no function is registered, it will do a redirect with ``redirect_content`` as content.
+        """
+        self._exception_handler = f
+        return f
+
     def invalid_response(self, f):
         """Register a function for responsing with invalid request.
 
@@ -391,13 +427,13 @@ def decorated(*args, **kwargs):
                     kwargs.update(credentials)
                 except oauth2.FatalClientError as e:
                     log.debug('Fatal client error %r', e, exc_info=True)
-                    return redirect(e.in_uri(self.error_uri))
+                    return self._on_exception(e, e.in_uri(self.error_uri))
                 except oauth2.OAuth2Error as e:
                     log.debug('OAuth2Error: %r', e, exc_info=True)
-                    return redirect(e.in_uri(redirect_uri))
+                    return self._on_exception(e, e.in_uri(redirect_uri))
                 except Exception as e:
                     log.exception(e)
-                    return redirect(add_params_to_uri(
+                    return self._on_exception(e, add_params_to_uri(
                         self.error_uri, {'error': str(e)}
                     ))
 
@@ -410,10 +446,10 @@ def decorated(*args, **kwargs):
                 rv = f(*args, **kwargs)
             except oauth2.FatalClientError as e:
                 log.debug('Fatal client error %r', e, exc_info=True)
-                return redirect(e.in_uri(self.error_uri))
+                return self._on_exception(e, e.in_uri(self.error_uri))
             except oauth2.OAuth2Error as e:
                 log.debug('OAuth2Error: %r', e, exc_info=True)
-                return redirect(e.in_uri(redirect_uri))
+                return self._on_exception(e, e.in_uri(redirect_uri))
 
             if not isinstance(rv, bool):
                 # if is a response or redirect
@@ -422,7 +458,7 @@ def decorated(*args, **kwargs):
             if not rv:
                 # denied by user
                 e = oauth2.AccessDeniedError()
-                return redirect(e.in_uri(redirect_uri))
+                return self._on_exception(e, e.in_uri(redirect_uri))
             return self.confirm_authorization_request()
         return decorated
 
@@ -449,13 +485,13 @@ def confirm_authorization_request(self):
             return create_response(*ret)
         except oauth2.FatalClientError as e:
             log.debug('Fatal client error %r', e, exc_info=True)
-            return redirect(e.in_uri(self.error_uri))
+            return self._on_exception(e, e.in_uri(self.error_uri))
         except oauth2.OAuth2Error as e:
             log.debug('OAuth2Error: %r', e, exc_info=True)
-            return redirect(e.in_uri(redirect_uri or self.error_uri))
+            return self._on_exception(e, e.in_uri(redirect_uri or self.error_uri))
         except Exception as e:
             log.exception(e)
-            return redirect(add_params_to_uri(
+            return self._on_exception(e, add_params_to_uri(
                 self.error_uri, {'error': str(e)}
             ))
 

tests/test_oauth2/test_code.py

@@ -2,7 +2,7 @@
 
 from datetime import datetime, timedelta
 from .._base import to_base64
-from .base import TestCase
+from .base import TestCase, default_provider
 from .base import create_server, sqlalchemy_provider, cache_provider
 from .base import db, Client, User, Grant
 
@@ -159,3 +159,25 @@ def test_get_token(self):
         url += '&client_secret=' + self.oauth_client.client_secret
         rv = self.client.get(url)
         assert b'access_token' in rv.data
+
+
+class TestProviderWithExceptionHandler(TestCase):
+
+    def prepare_data(self):
+        oauth = default_provider(self.app)
+
+        @oauth.exception_handler
+        def custom_exception_handler(error, *args):
+            raise error
+
+        @self.app.errorhandler(Exception)
+        def all_exception_handler(*args):
+            return "Testing server error", 500
+
+        create_server(self.app, oauth=oauth)
+
+    def test_exception_handler(self):
+        rv = self.client.get('/oauth/authorize')
+
+        assert rv.status_code == 500
+        assert rv.data.decode("utf-8") == "Testing server error"