✨ Add swagger UI OAuth2 redirect page for implicit/code auth flows in API docs (#198)

6 years ago · 325edd5f00
5 changed files with 185 additions and 5 deletions
--- a/fastapi/applications.py
+++ b/fastapi/applications.py
@ -1,7 +1,11 @@
 from typing import Any, Callable, Dict, List, Optional, Type, Union

 from fastapi import routing
-from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
+from fastapi.openapi.docs import (
+    get_redoc_html,
+    get_swagger_ui_html,
+    get_swagger_ui_oauth2_redirect_html,
+)
 from fastapi.openapi.utils import get_openapi
 from fastapi.params import Depends
 from pydantic import BaseModel
@ -36,6 +40,7 @@ class FastAPI(Starlette):
        openapi_prefix: str = "",
        docs_url: Optional[str] = "/docs",
        redoc_url: Optional[str] = "/redoc",
+        swagger_ui_oauth2_redirect_url: Optional[str] = "/docs/oauth2-redirect",
        **extra: Dict[str, Any],
    ) -> None:
        self._debug = debug
@ -52,6 +57,7 @@ class FastAPI(Starlette):
        self.openapi_prefix = openapi_prefix.rstrip("/")
        self.docs_url = docs_url
        self.redoc_url = redoc_url
+        self.swagger_ui_oauth2_redirect_url = swagger_ui_oauth2_redirect_url
        self.extra = extra

        self.openapi_version = "3.0.2"
@ -89,10 +95,23 @@ class FastAPI(Starlette):

            async def swagger_ui_html(req: Request) -> HTMLResponse:
                return get_swagger_ui_html(
-                    openapi_url=openapi_url, title=self.title + " - Swagger UI"
+                    openapi_url=openapi_url,
+                    title=self.title + " - Swagger UI",
+                    oauth2_redirect_url=self.swagger_ui_oauth2_redirect_url,
                )

            self.add_route(self.docs_url, swagger_ui_html, include_in_schema=False)
+
+            if self.swagger_ui_oauth2_redirect_url:
+
+                async def swagger_ui_redirect(req: Request) -> HTMLResponse:
+                    return get_swagger_ui_oauth2_redirect_html()
+
+                self.add_route(
+                    self.swagger_ui_oauth2_redirect_url,
+                    swagger_ui_redirect,
+                    include_in_schema=False,
+                )
        if self.openapi_url and self.redoc_url:

            async def redoc_html(req: Request) -> HTMLResponse:
--- a/fastapi/openapi/docs.py
+++ b/fastapi/openapi/docs.py
@ -1,3 +1,5 @@
+from typing import Optional
+
 from starlette.responses import HTMLResponse


@ -8,7 +10,9 @@ def get_swagger_ui_html(
    swagger_js_url: str = "https://cdn.jsdelivr.net/npm/swagger-ui-dist@3/swagger-ui-bundle.js",
    swagger_css_url: str = "https://cdn.jsdelivr.net/npm/swagger-ui-dist@3/swagger-ui.css",
    swagger_favicon_url: str = "https://fastapi.tiangolo.com/img/favicon.png",
+    oauth2_redirect_url: Optional[str] = None,
 ) -> HTMLResponse:
+
    html = f"""
    <! doctype html>
    <html>
@ -25,14 +29,19 @@ def get_swagger_ui_html(
    <script>
    const ui = SwaggerUIBundle({{
        url: '{openapi_url}',
+    """
+
+    if oauth2_redirect_url:
+        html += f"oauth2RedirectUrl: window.location.origin + '{oauth2_redirect_url}',"
+
+    html += """
        dom_id: '#swagger-ui',
        presets: [
        SwaggerUIBundle.presets.apis,
        SwaggerUIBundle.SwaggerUIStandalonePreset
        ],
        layout: "BaseLayout"
-
-    }})
+    })
    </script>
    </body>
    </html>
@ -47,7 +56,6 @@ def get_redoc_html(
    redoc_js_url: str = "https://cdn.jsdelivr.net/npm/redoc@next/bundles/redoc.standalone.js",
    redoc_favicon_url: str = "https://fastapi.tiangolo.com/img/favicon.png",
 ) -> HTMLResponse:
-
    html = f"""
    <!DOCTYPE html>
    <html>
@ -75,3 +83,76 @@ def get_redoc_html(
    </html>
    """
    return HTMLResponse(html)
+
+
+def get_swagger_ui_oauth2_redirect_html() -> HTMLResponse:
+    html = """
+    <!doctype html>
+    <html lang="en-US">
+    <body onload="run()">
+    </body>
+    </html>
+    <script>
+        'use strict';
+        function run () {
+            var oauth2 = window.opener.swaggerUIRedirectOauth2;
+            var sentState = oauth2.state;
+            var redirectUrl = oauth2.redirectUrl;
+            var isValid, qp, arr;
+
+            if (/code|token|error/.test(window.location.hash)) {
+                qp = window.location.hash.substring(1);
+            } else {
+                qp = location.search.substring(1);
+            }
+
+            arr = qp.split("&")
+            arr.forEach(function (v,i,_arr) { _arr[i] = '"' + v.replace('=', '":"') + '"';})
+            qp = qp ? JSON.parse('{' + arr.join() + '}',
+                    function (key, value) {
+                        return key === "" ? value : decodeURIComponent(value)
+                    }
+            ) : {}
+
+            isValid = qp.state === sentState
+
+            if ((
+            oauth2.auth.schema.get("flow") === "accessCode"||
+            oauth2.auth.schema.get("flow") === "authorizationCode"
+            ) && !oauth2.auth.code) {
+                if (!isValid) {
+                    oauth2.errCb({
+                        authId: oauth2.auth.name,
+                        source: "auth",
+                        level: "warning",
+                        message: "Authorization may be unsafe, passed state was changed in server Passed state wasn't returned from auth server"
+                    });
+                }
+
+                if (qp.code) {
+                    delete oauth2.state;
+                    oauth2.auth.code = qp.code;
+                    oauth2.callback({auth: oauth2.auth, redirectUrl: redirectUrl});
+                } else {
+                    let oauthErrorMsg
+                    if (qp.error) {
+                        oauthErrorMsg = "["+qp.error+"]: " +
+                            (qp.error_description ? qp.error_description+ ". " : "no accessCode received from the server. ") +
+                            (qp.error_uri ? "More info: "+qp.error_uri : "");
+                    }
+
+                    oauth2.errCb({
+                        authId: oauth2.auth.name,
+                        source: "auth",
+                        level: "error",
+                        message: oauthErrorMsg || "[Authorization failed]: no accessCode received from the server"
+                    });
+                }
+            } else {
+                oauth2.callback({auth: oauth2.auth, token: qp, isValid: isValid, redirectUrl: redirectUrl});
+            }
+            window.close();
+        }
+    </script>
+        """
+    return HTMLResponse(content=html)
--- a/tests/test_application.py
+++ b/tests/test_application.py
@ -1131,6 +1131,17 @@ def test_swagger_ui():
    assert response.status_code == 200
    assert response.headers["content-type"] == "text/html; charset=utf-8"
    assert "swagger-ui-dist" in response.text
+    assert (
+        f"oauth2RedirectUrl: window.location.origin + '/docs/oauth2-redirect'"
+        in response.text
+    )
+
+
+def test_swagger_ui_oauth2_redirect():
+    response = client.get("/docs/oauth2-redirect")
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "text/html; charset=utf-8"
+    assert "window.opener.swaggerUIRedirectOauth2" in response.text


 def test_redoc():
--- a/tests/test_custom_swagger_ui_redirect.py
+++ b/tests/test_custom_swagger_ui_redirect.py
@ -0,0 +1,38 @@
+from fastapi import FastAPI
+from starlette.testclient import TestClient
+
+swagger_ui_oauth2_redirect_url = "/docs/redirect"
+
+app = FastAPI(swagger_ui_oauth2_redirect_url=swagger_ui_oauth2_redirect_url)
+
+
+@app.get("/items/")
+async def read_items():
+    return {"id": "foo"}
+
+
+client = TestClient(app)
+
+
+def test_swagger_ui():
+    response = client.get("/docs")
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "text/html; charset=utf-8"
+    assert "swagger-ui-dist" in response.text
+    print(client.base_url)
+    assert (
+        f"oauth2RedirectUrl: window.location.origin + '{swagger_ui_oauth2_redirect_url}'"
+        in response.text
+    )
+
+
+def test_swagger_ui_oauth2_redirect():
+    response = client.get(swagger_ui_oauth2_redirect_url)
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "text/html; charset=utf-8"
+    assert "window.opener.swaggerUIRedirectOauth2" in response.text
+
+
+def test_response():
+    response = client.get("/items/")
+    assert response.json() == {"id": "foo"}
--- a/tests/test_no_swagger_ui_redirect.py
+++ b/tests/test_no_swagger_ui_redirect.py
@ -0,0 +1,31 @@
+from fastapi import FastAPI
+from starlette.testclient import TestClient
+
+app = FastAPI(swagger_ui_oauth2_redirect_url=None)
+
+
+@app.get("/items/")
+async def read_items():
+    return {"id": "foo"}
+
+
+client = TestClient(app)
+
+
+def test_swagger_ui():
+    response = client.get("/docs")
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "text/html; charset=utf-8"
+    assert "swagger-ui-dist" in response.text
+    print(client.base_url)
+    assert "oauth2RedirectUrl" not in response.text
+
+
+def test_swagger_ui_no_oauth2_redirect():
+    response = client.get("/docs/oauth2-redirect")
+    assert response.status_code == 404
+
+
+def test_response():
+    response = client.get("/items/")
+    assert response.json() == {"id": "foo"}