From 7512afa9494278d46fcba7dabc4e6fe4458f2dee Mon Sep 17 00:00:00 2001
From: Masen Furer <m_github@0x26.net>
Date: Mon, 13 May 2024 21:29:45 -0700
Subject: [PATCH] Proxy backend requests on '/' to the frontend

If the optional extra `proxy` is installed, then the backend can handle all
requests by proxy unrecognized routes to the frontend nextjs server.
---
 pyproject.toml  |  4 +++
 reflex/app.py   | 20 ++++++++++++-
 reflex/proxy.py | 75 +++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 98 insertions(+), 1 deletion(-)
 create mode 100644 reflex/proxy.py

diff --git a/pyproject.toml b/pyproject.toml
index fb3eab724..ba742e107 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -63,6 +63,7 @@ setuptools = ">=69.1.1,<70.0"
 httpx = ">=0.25.1,<1.0"
 twine = ">=4.0.0,<6.0"
 tomlkit = ">=0.12.4,<1.0"
+asgiproxy = { version = "==0.1.1", optional = true }
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=7.1.2,<8.0"
@@ -90,6 +91,9 @@ pytest-benchmark = ">=4.0.0,<5.0"
 [tool.poetry.scripts]
 reflex = "reflex.reflex:cli"
 
+[tool.poetry.extras]
+proxy = ["asgiproxy"]
+
 [build-system]
 requires = ["poetry-core>=1.5.1"]
 build-backend = "poetry.core.masonry.api"
diff --git a/reflex/app.py b/reflex/app.py
index 4d99d6949..bd2dc771f 100644
--- a/reflex/app.py
+++ b/reflex/app.py
@@ -13,6 +13,7 @@ import os
 import platform
 from typing import (
     Any,
+    AsyncGenerator,
     AsyncIterator,
     Callable,
     Coroutine,
@@ -94,6 +95,23 @@ def default_overlay_component() -> Component:
     return Fragment.create(connection_pulser(), connection_modal())
 
 
+@contextlib.asynccontextmanager
+async def lifespan(api: FastAPI) -> AsyncGenerator[None, None]:
+    """Context manager to handle the lifespan of the app.
+
+    Args:
+        api: The FastAPI instance.
+
+    Yields:
+        None
+    """
+    # try to set up proxying if its enabled
+    from .proxy import proxy_middleware
+
+    async with proxy_middleware(api):
+        yield
+
+
 class OverlayFragment(Fragment):
     """Alias for Fragment, used to wrap the overlay_component."""
 
@@ -203,7 +221,7 @@ class App(Base):
         self.middleware.append(HydrateMiddleware())
 
         # Set up the API.
-        self.api = FastAPI()
+        self.api = FastAPI(lifespan=lifespan)
         self._add_cors()
         self._add_default_endpoints()
 
diff --git a/reflex/proxy.py b/reflex/proxy.py
new file mode 100644
index 000000000..c47203bfd
--- /dev/null
+++ b/reflex/proxy.py
@@ -0,0 +1,75 @@
+"""Handle proxying frontend requests from the backend server."""
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator
+from urllib.parse import urlparse
+
+from fastapi import FastAPI
+from starlette.types import ASGIApp
+
+from .config import get_config
+from .utils import console
+
+try:
+    from asgiproxy.config import BaseURLProxyConfigMixin, ProxyConfig
+    from asgiproxy.context import ProxyContext
+    from asgiproxy.simple_proxy import make_simple_proxy_app
+except ImportError:
+
+    @asynccontextmanager
+    async def proxy_middleware(*args, **kwargs) -> AsyncGenerator[None, None]:
+        """A no-op proxy middleware for when asgiproxy is not installed.
+
+        Args:
+            *args: The positional arguments.
+            **kwargs: The keyword arguments.
+
+        Yields:
+            None
+        """
+        yield
+else:
+
+    def _get_proxy_app_with_context(frontend_host: str) -> tuple[ProxyContext, ASGIApp]:
+        """Get the proxy app with the given frontend host.
+
+        Args:
+            frontend_host: The frontend host to proxy requests to.
+
+        Returns:
+            The proxy context and app.
+        """
+
+        class LocalProxyConfig(BaseURLProxyConfigMixin, ProxyConfig):
+            upstream_base_url = frontend_host
+            rewrite_host_header = urlparse(upstream_base_url).netloc
+
+        proxy_context = ProxyContext(LocalProxyConfig())
+        proxy_app = make_simple_proxy_app(proxy_context)
+        return proxy_context, proxy_app
+
+    @asynccontextmanager
+    async def proxy_middleware(  # pyright: ignore[reportGeneralTypeIssues]
+        api: FastAPI,
+    ) -> AsyncGenerator[None, None]:
+        """A middleware to proxy requests to the separate frontend server.
+
+        The proxy is installed on the / endpoint of the FastAPI instance.
+
+        Args:
+            api: The FastAPI instance.
+
+        Yields:
+            None
+        """
+        config = get_config()
+        backend_port = config.backend_port
+        frontend_host = f"http://localhost:{config.frontend_port}"
+        proxy_context, proxy_app = _get_proxy_app_with_context(frontend_host)
+        api.mount("/", proxy_app)
+        console.debug(
+            f"Proxying '/' requests on port {backend_port} to {frontend_host}"
+        )
+        async with proxy_context:
+            yield