Fixes for known bugs in v1.0.0-rc1 (#241)

* Fix #238 for same-site cookies

* Redirect in Python + bump to rc2

custom_components/auth_oidc/endpoints/injected_auth_page.py

@@ -1,12 +1,18 @@
 """Injected authorization page, replacing the original"""
 
+import base64
 import logging
 from functools import partial
-from homeassistant.components.http import HomeAssistantView, StaticPathConfig
-from homeassistant.core import HomeAssistant
+from urllib.parse import quote, unquote
 from aiohttp import web
 from aiofiles import open as async_open
 
+from homeassistant.components.http import HomeAssistantView, StaticPathConfig
+from homeassistant.core import HomeAssistant
+
+from .welcome import PATH as WELCOME_PATH
+from ..tools.helpers import get_url
+
 PATH = "/auth/authorize"
 
 _LOGGER = logging.getLogger(__name__)
@@ -18,7 +24,7 @@ async def read_file(path: str) -> str:
         return await f.read()
 
 
-async def frontend_injection(hass: HomeAssistant) -> None:
+async def frontend_injection(hass: HomeAssistant, force_https: bool) -> None:
     """Inject new frontend code into /auth/authorize."""
     router = hass.http.app.router
     frontend_path = None
@@ -61,7 +67,7 @@ async def frontend_injection(hass: HomeAssistant) -> None:
     frontend_code = await read_file(frontend_path)
 
     # Inject JS and register that route
-    injection_js = "<script src='/auth/oidc/static/injection.js?v=4'></script>"
+    injection_js = "<script src='/auth/oidc/static/injection.js?v=6'></script>"
     frontend_code = frontend_code.replace("</body>", f"{injection_js}</body>")
 
     await hass.http.async_register_static_paths(
@@ -80,7 +86,7 @@ async def frontend_injection(hass: HomeAssistant) -> None:
     )
 
     # If everything is succesful, register a fake view that just returns the modified HTML
-    hass.http.register_view(OIDCInjectedAuthPage(frontend_code))
+    hass.http.register_view(OIDCInjectedAuthPage(frontend_code, force_https))
     _LOGGER.info("Performed OIDC frontend injection")
 
 
@@ -91,18 +97,46 @@ class OIDCInjectedAuthPage(HomeAssistantView):
     url = PATH
     name = "auth:oidc:authorize_page"
 
-    def __init__(self, html: str) -> None:
+    def __init__(self, html: str, force_https: bool) -> None:
         """Initialize the injected auth page."""
         self.html = html
+        self.force_https = force_https
 
     @staticmethod
-    async def inject(hass: HomeAssistant) -> None:
+    async def inject(hass: HomeAssistant, force_https: bool) -> None:
         """Inject the OIDC auth page into the frontend."""
         try:
-            await frontend_injection(hass)
+            await frontend_injection(hass, force_https)
         except Exception as e:  # pylint: disable=broad-except
             _LOGGER.error("Failed to inject OIDC auth page: %s", e)
 
-    async def get(self, _) -> web.Response:
-        """Return the screen"""
+    @staticmethod
+    def _should_do_oidc_redirect(req: web.Request) -> bool:
+        """Check if we should redirect to the OIDC flow."""
+        if req.query.get("skip_oidc_redirect") == "true":
+            return False
+
+        redirect_uri = req.query.get("redirect_uri")
+        if not redirect_uri:
+            return False
+
+        # Handle both encoded and plain redirect_uri values.
+        decoded_redirect_uri = unquote(redirect_uri)
+        return "skip_oidc_redirect=true" not in decoded_redirect_uri
+
+    def _get_welcome_redirect_location(self, req: web.Request) -> str:
+        """Build the welcome URL for the injected auth page redirect."""
+        encoded_current_url = quote(
+            base64.b64encode(str(req.url).encode("utf-8")).decode("ascii")
+        )
+        return get_url(
+            f"{WELCOME_PATH}?redirect_uri={encoded_current_url}",
+            self.force_https,
+        )
+
+    async def get(self, req: web.Request) -> web.Response:
+        """Return the original page or redirect into the OIDC flow."""
+        if self._should_do_oidc_redirect(req):
+            raise web.HTTPFound(location=self._get_welcome_redirect_location(req))
+
         return web.Response(text=self.html, content_type="text/html")