support x-forwarded-host as source for redirect

Author: erikgaas

fasthtml/_modidx.py

@@ -197,6 +197,7 @@
                                                                               'fasthtml/oauth.py'),
                                 'fasthtml.oauth._AppClient.retr_id': ('api/oauth.html#_appclient.retr_id', 'fasthtml/oauth.py'),
                                 'fasthtml.oauth._AppClient.retr_info': ('api/oauth.html#_appclient.retr_info', 'fasthtml/oauth.py'),
+                                'fasthtml.oauth.get_host': ('api/oauth.html#get_host', 'fasthtml/oauth.py'),
                                 'fasthtml.oauth.load_creds': ('api/oauth.html#load_creds', 'fasthtml/oauth.py'),
                                 'fasthtml.oauth.redir_url': ('api/oauth.html#redir_url', 'fasthtml/oauth.py'),
                                 'fasthtml.oauth.url_match': ('api/oauth.html#url_match', 'fasthtml/oauth.py')},

fasthtml/oauth.py

@@ -4,7 +4,7 @@
 
 # %% auto 0
 __all__ = ['http_patterns', 'GoogleAppClient', 'GitHubAppClient', 'HuggingFaceClient', 'DiscordAppClient', 'Auth0AppClient',
-           'redir_url', 'url_match', 'OAuth', 'load_creds']
+           'get_host', 'redir_url', 'url_match', 'OAuth', 'load_creds']
 
 # %% ../nbs/api/08_oauth.ipynb
 from .common import *
@@ -120,11 +120,18 @@ def login_link(self:WebApplicationClient, redirect_uri, scope=None, state=None,
     if not state: state=getattr(self, 'state', None)
     return self.prepare_request_uri(self.base_url, redirect_uri, scope, state=state, **kwargs)
 
+# %% ../nbs/api/08_oauth.ipynb
+def get_host(request):
+    """Get the host, preferring X-Forwarded-Host if available"""
+    forwarded_host = request.headers.get('x-forwarded-host')
+    return forwarded_host if forwarded_host else request.url.netloc
+
 # %% ../nbs/api/08_oauth.ipynb
 def redir_url(request, redir_path, scheme=None):
     "Get the redir url for the host in `request`"
-    scheme = 'http' if request.url.hostname in ("localhost", "127.0.0.1") else 'https'
-    return f"{scheme}://{request.url.netloc}{redir_path}"
+    host = get_host(request)
+    scheme = 'http' if host.split(':')[0] in ("localhost", "127.0.0.1") else 'https'
+    return f"{scheme}://{host}{redir_path}"
 
 # %% ../nbs/api/08_oauth.ipynb
 @patch
@@ -159,8 +166,8 @@ def retr_id(self:_AppClient, code, redirect_uri):
 
 # %% ../nbs/api/08_oauth.ipynb
 http_patterns = (r'^(localhost|127\.0\.0\.1)(:\d+)?$',)
-def url_match(url, patterns=http_patterns):
-    return any(re.match(pattern, url.netloc.split(':')[0]) for pattern in patterns)
+def url_match(request, patterns=http_patterns):
+    return any(re.match(pattern, get_host(request).split(':')[0]) for pattern in patterns)
 
 # %% ../nbs/api/08_oauth.ipynb
 class OAuth:
@@ -179,7 +186,7 @@ def before(req, session):
         def redirect(req, session, code:str=None, error:str=None, state:str=None):
             if not code: session['oauth_error']=error; return RedirectResponse(self.error_path, status_code=303)
             scheme = 'http' if url_match(req.url,self.http_patterns) or not self.https else 'https'
-            base_url = f"{scheme}://{req.url.netloc}"
+            base_url = f"{scheme}://{get_host(req)}"
             info = AttrDictDefault(cli.retr_info(code, base_url+redir_path))
             ident = info.get(self.cli.id_key)
             if not ident: return self.redir_login(session)

nbs/api/08_oauth.ipynb

@@ -301,6 +301,49 @@
     "server = JupyUvi(app, port=port)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "551454f9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#|export\n",
+    "def get_host(request):\n",
+    "    \"\"\"Get the host, preferring X-Forwarded-Host if available\"\"\"\n",
+    "    forwarded_host = request.headers.get('x-forwarded-host')\n",
+    "    return forwarded_host if forwarded_host else request.url.netloc"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3cc9b978",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Without X-Forwarded-Host: localhost:8000\n",
+      "With X-Forwarded-Host: example.com\n"
+     ]
+    }
+   ],
+   "source": [
+    "from types import SimpleNamespace\n",
+    "from urllib.parse import urlparse\n",
+    "\n",
+    "mock_request_localhost = SimpleNamespace(headers={}, url=SimpleNamespace(netloc='localhost:8000'))\n",
+    "mock_request_with_forward = SimpleNamespace(\n",
+    "    headers={'x-forwarded-host': 'example.com'}, \n",
+    "    url=SimpleNamespace(netloc='localhost:8000', hostname='localhost')\n",
+    ")\n",
+    "\n",
+    "print(\"Without X-Forwarded-Host:\", get_host(mock_request_localhost))\n",
+    "print(\"With X-Forwarded-Host:\", get_host(mock_request_with_forward))"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -311,8 +354,36 @@
     "#| export\n",
     "def redir_url(request, redir_path, scheme=None):\n",
     "    \"Get the redir url for the host in `request`\"\n",
-    "    scheme = 'http' if request.url.hostname in (\"localhost\", \"127.0.0.1\") else 'https'\n",
-    "    return f\"{scheme}://{request.url.netloc}{redir_path}\""
+    "    host = get_host(request)\n",
+    "    scheme = 'http' if host.split(':')[0] in (\"localhost\", \"127.0.0.1\") else 'https'\n",
+    "    return f\"{scheme}://{host}{redir_path}\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fee2db6c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Localhost: http://localhost:8000/redirect\n",
+      "With X-Forwarded-Host: https://example.com/redirect\n",
+      "Production: https://myapp.com/redirect\n"
+     ]
+    }
+   ],
+   "source": [
+    "from types import SimpleNamespace\n",
+    "from urllib.parse import urlparse\n",
+    "\n",
+    "mock_request_prod = SimpleNamespace(headers={}, url=SimpleNamespace(netloc='myapp.com', hostname='myapp.com'))\n",
+    "\n",
+    "print(\"Localhost:\", redir_url(mock_request_localhost, '/redirect'))\n",
+    "print(\"With X-Forwarded-Host:\", redir_url(mock_request_with_forward, '/redirect'))\n",
+    "print(\"Production:\", redir_url(mock_request_prod, '/redirect'))"
    ]
   },
   {
@@ -447,8 +518,8 @@
    "source": [
     "#| export\n",
     "http_patterns = (r'^(localhost|127\\.0\\.0\\.1)(:\\d+)?$',)\n",
-    "def url_match(url, patterns=http_patterns):\n",
-    "    return any(re.match(pattern, url.netloc.split(':')[0]) for pattern in patterns)"
+    "def url_match(request, patterns=http_patterns):\n",
+    "    return any(re.match(pattern, get_host(request).split(':')[0]) for pattern in patterns)"
    ]
   },
   {
@@ -475,7 +546,7 @@
     "        def redirect(req, session, code:str=None, error:str=None, state:str=None):\n",
     "            if not code: session['oauth_error']=error; return RedirectResponse(self.error_path, status_code=303)\n",
     "            scheme = 'http' if url_match(req.url,self.http_patterns) or not self.https else 'https'\n",
-    "            base_url = f\"{scheme}://{req.url.netloc}\"\n",
+    "            base_url = f\"{scheme}://{get_host(req)}\"\n",
     "            info = AttrDictDefault(cli.retr_info(code, base_url+redir_path))\n",
     "            ident = info.get(self.cli.id_key)\n",
     "            if not ident: return self.redir_login(session)\n",