Except AuthCodeNotFound error in AuthHandler

Author: wndhydrnt

oauth2/grant.py

@@ -28,7 +28,8 @@
 
 """
 from oauth2.error import OAuthInvalidError, UserNotAuthenticated, \
-    AccessTokenNotFound, UserIdentifierMissingError, RedirectUriUnknown
+    AccessTokenNotFound, UserIdentifierMissingError, RedirectUriUnknown, \
+    AuthCodeNotFound
 from oauth2.compatibility import urlencode, quote
 import json
 import time
@@ -527,25 +528,27 @@ def _read_params(self, request):
                 explanation="Invalid redirect_uri parameter")
 
     def _validate_code(self):
-        stored_code = self.auth_code_store.fetch_by_code(self.code)
-
-        if stored_code is None:
-            raise OAuthInvalidError(error="invalid_request",
-                                    explanation="Invalid authorization code " \
-                                                "parameter")
+        try:
+            stored_code = self.auth_code_store.fetch_by_code(self.code)
+        except AuthCodeNotFound:
+            raise OAuthInvalidError(
+                error="invalid_request",
+                explanation="Invalid authorization code parameter")
 
         if stored_code.code != self.code:
-            raise OAuthInvalidError(error="invalid_grant",
-                                    explanation="Invalid code parameter in " \
-                                                "request")
+            raise OAuthInvalidError(
+                error="invalid_grant",
+                explanation="Invalid code parameter in request")
 
         if stored_code.redirect_uri != self.redirect_uri:
-            raise OAuthInvalidError(error="invalid_request",
-                                    explanation="Invalid redirect_uri parameter")
+            raise OAuthInvalidError(
+                error="invalid_request",
+                explanation="Invalid redirect_uri parameter")
 
         if stored_code.is_expired():
-            raise OAuthInvalidError(error="invalid_grant",
-                                    explanation="Authorization code has expired")
+            raise OAuthInvalidError(
+                error="invalid_grant",
+                explanation="Authorization code has expired")
 
         self.data = stored_code.data
         self.scopes = stored_code.scopes

oauth2/test/test_grant.py

@@ -10,7 +10,7 @@
     ClientCredentialsHandler, AuthorizeMixin
 from oauth2.store import ClientStore, AuthCodeStore, AccessTokenStore
 from oauth2.error import OAuthInvalidError, UserNotAuthenticated, \
-    AccessTokenNotFound, UserIdentifierMissingError
+    AccessTokenNotFound, UserIdentifierMissingError, AuthCodeNotFound
 from oauth2 import Provider
 from oauth2.datatype import Client, AuthorizationCode, AccessToken
 from oauth2.tokengenerator import TokenGenerator
@@ -410,7 +410,7 @@ def test_read_validate_params_no_auth_code_found(self):
         redirect_uri = "http://callback"
 
         auth_code_store_mock = Mock(spec=AuthCodeStore)
-        auth_code_store_mock.fetch_by_code.return_value = None
+        auth_code_store_mock.fetch_by_code.side_effect = AuthCodeNotFound
 
         client = Client(identifier=client_id, secret=client_secret,
                         redirect_uris=[redirect_uri])