Added revocation check support (#18)

* Added revokation callback

* minor style cleaning

* reorder

aiohttp_jwt/middleware.py

@@ -1,32 +1,26 @@
-import asyncio
 import logging
 import re
+from functools import partial
 
 import aiohttp
 import jwt
 
+from .utils import check_request, invoke
+
 logger = logging.getLogger(__name__)
 
 __config = dict()
 
 __REQUEST_IDENT = 'request_property'
 
 
-def check_request(request, entries):
-    """Check if request.path match group of certain patterns."""
-    for pattern in entries:
-        if re.match(pattern, request.path):
-            return True
-
-    return False
-
-
 def JWTMiddleware(
     secret_or_pub_key,
     request_property='payload',
     credentials_required=True,
     whitelist=tuple(),
     token_getter=None,
+    is_revoked=None,
     store_token=False,
     algorithms=None,
 ):
@@ -46,9 +40,7 @@ async def middleware(request):
                 token = None
 
                 if callable(token_getter):
-                    token = token_getter(request)
-                    if asyncio.iscoroutine(token):
-                        token = await token
+                    token = await invoke(partial(token_getter, request))
                 elif 'Authorization' in request.headers:
                     try:
                         scheme, token = request.headers.get(
@@ -70,23 +62,36 @@ async def middleware(request):
                     )
 
                 if token is not None:
+                    if not isinstance(token, bytes):
+                        token = token.encode()
+
                     try:
-                        if not isinstance(token, bytes):
-                            token = token.encode()
                         decoded = jwt.decode(
                             token,
                             secret_or_pub_key,
                             algorithms=algorithms,
                         )
-                        request[request_property] = decoded
-                        if store_token and isinstance(store_token, str):
-                            request[store_token] = token
                     except jwt.InvalidTokenError as exc:
                         logger.exception(exc, exc_info=exc)
                         raise aiohttp.web.HTTPForbidden(
                             reason='Invalid authorization token',
                         )
 
+                    if callable(is_revoked):
+                        if await invoke(partial(
+                            is_revoked,
+                            request,
+                            decoded,
+                        )):
+                            raise aiohttp.web.HTTPForbidden(
+                                reason='Token is revoked',
+                            )
+
+                    request[request_property] = decoded
+
+                    if store_token and isinstance(store_token, str):
+                        request[store_token] = token
+
             return await handler(request)
         return middleware
 

aiohttp_jwt/utils.py

@@ -0,0 +1,17 @@
+import asyncio
+import re
+
+
+def check_request(request, entries):
+    for pattern in entries:
+        if re.match(pattern, request.path):
+            return True
+
+    return False
+
+
+async def invoke(func):
+    result = func()
+    if asyncio.iscoroutine(result):
+        result = await result
+    return result

tests/test_middleware.py

@@ -168,3 +168,28 @@ async def handler(request):
         'auth_token': token.decode('utf-8'),
     })
     assert response.status == 200
+
+
+def is_revoked(request, payload):
+    return True
+
+
+async def is_revoked_coro(request, payload):
+    return True
+
+
+@pytest.mark.parametrize('check', [
+    is_revoked,
+    is_revoked_coro,
+])
+async def test_token_revoked(
+        check, create_app, aiohttp_client, fake_payload, token):
+    async def handler(request):
+        return web.json_response({})
+    routes = (('/foo', handler),)
+    client = await aiohttp_client(create_app(routes, is_revoked=check))
+    response = await client.get('/foo', headers={
+        'Authorization': 'Bearer {}'.format(token.decode('utf-8')),
+    })
+    assert response.status == 403
+    assert 'Token is revoked' in response.reason