Implement security module with support for cookies (GH-10)

Author: ArtyomVancyan

README.md

@@ -13,7 +13,6 @@ the [social-core](https://github.com/python-social-auth/social-core) authenticat
 
 - Use multiple OAuth2 providers at the same time
     * There need to be provided a way to configure the OAuth2 for multiple providers
-- Provide `fastapi.security.*` implementations that use cookies
 - Token -> user data, user data -> token easy conversion
 - Customizable OAuth2 routes
 - Registration support

examples/demonstration/router.py

@@ -1,11 +1,13 @@
 import json
 
 from fastapi import Depends
-from fastapi import Request, APIRouter
+from fastapi import Request
+from fastapi import APIRouter
 from fastapi.responses import HTMLResponse
-from fastapi.security import OAuth2
 from fastapi.templating import Jinja2Templates
 
+from fastapi_oauth2.security import OAuth2
+
 oauth2 = OAuth2()
 router = APIRouter()
 templates = Jinja2Templates(directory="templates")

src/fastapi_oauth2/middleware.py

@@ -95,7 +95,7 @@ async def authenticate(self, request: Request) -> Optional[Tuple["Auth", "User"]
             return Auth(), User()
 
         user = Auth.jwt_decode(param)
-        return Auth(user.pop("scope")), User(user)
+        return Auth(user.pop("scope", [])), User(user)
 
 
 class OAuth2Middleware:

src/fastapi_oauth2/security.py

@@ -0,0 +1,34 @@
+from fastapi.security import OAuth2 as FastAPIOAuth2
+from fastapi.security import OAuth2AuthorizationCodeBearer as FastAPICodeBearer
+from fastapi.security import OAuth2PasswordBearer as FastAPIPasswordBearer
+from starlette.datastructures import Headers
+from starlette.requests import Request
+
+
+def use_cookie(cls: FastAPIOAuth2):
+    def _use_cookie(*args, **kwargs):
+        async def __call__(self, request: Request):
+            authorization = request.headers.get("Authorization", request.cookies.get("Authorization"))
+            if authorization:
+                request._headers = Headers({**request.headers, "Authorization": authorization})
+            return await super(cls, self).__call__(request)
+
+        cls.__call__ = __call__
+        return cls(*args, **kwargs)
+
+    return _use_cookie
+
+
+@use_cookie
+class OAuth2(FastAPIOAuth2):
+    ...
+
+
+@use_cookie
+class OAuth2PasswordBearer(FastAPIPasswordBearer):
+    ...
+
+
+@use_cookie
+class OAuth2AuthorizationCodeBearer(FastAPICodeBearer):
+    ...

tests/test_oauth2_middleware.py

@@ -1,15 +1,47 @@
 import pytest
+from fastapi import APIRouter
+from fastapi import Depends
 from fastapi import FastAPI
+from fastapi import Request
 from httpx import AsyncClient
 from social_core.backends.github import GithubOAuth2
+from starlette.responses import Response
 
 from fastapi_oauth2.client import OAuth2Client
 from fastapi_oauth2.core import OAuth2Core
 from fastapi_oauth2.middleware import OAuth2Middleware
 from fastapi_oauth2.router import router as oauth2_router
+from fastapi_oauth2.security import OAuth2
 
 app = FastAPI()
+oauth2 = OAuth2()
+app_router = APIRouter()
 
+
+@app_router.get("/user")
+def user(request: Request, _: str = Depends(oauth2)):
+    return request.user
+
+
+@app_router.get("/auth")
+def auth(request: Request):
+    access_token = request.auth.jwt_create({
+        "name": "test",
+        "sub": "test",
+        "id": "test",
+    })
+    response = Response()
+    response.set_cookie(
+        "Authorization",
+        value=f"Bearer {access_token}",
+        max_age=request.auth.expires,
+        expires=request.auth.expires,
+        httponly=request.auth.http,
+    )
+    return response
+
+
+app.include_router(app_router)
 app.include_router(oauth2_router)
 app.add_middleware(OAuth2Middleware, config={
     "allow_http": True,
@@ -30,6 +62,18 @@ async def test_auth_redirect():
         assert response.status_code == 303  # Redirect
 
 
+@pytest.mark.anyio
+async def test_authenticated_request():
+    async with AsyncClient(app=app, base_url="http://test") as client:
+        response = await client.get("/user")
+        assert response.status_code == 403  # Forbidden
+
+        await client.get("/auth")  # Simulate login
+
+        response = await client.get("/user")
+        assert response.status_code == 200  # OK
+
+
 @pytest.mark.anyio
 async def test_core_init(backends):
     for backend in backends: