99
1010router = APIRouter (tags = ["proxy" ])
1111
12- # Script injected to report content size to parent window
13- # Uses specific origin (pyplots.ai) for postMessage security
14- SIZE_REPORTER_SCRIPT = """
12+ # Allowed origins for postMessage
13+ ALLOWED_ORIGINS = ["https://pyplots.ai" , "http://localhost:3000" ]
14+
15+
16+ def get_size_reporter_script (target_origin : str ) -> str :
17+ """Generate size reporter script with specified target origin."""
18+ return f"""
1519<script>
16- (function() {
17- function reportSize() {
18- try {
20+ (function() {{
21+ function reportSize() {{
22+ try {{
1923 // Find the main content element (try common patterns for different libraries)
2024 var content = document.querySelector(
2125 '.bk-root, .vega-embed, .plotly, .chart-container, #container, .lp-plot, svg, canvas'
3236 height += padding;
3337
3438 // Send to parent with specific origin for security
35- if (width > 0 && height > 0 && window.parent !== window) {
36- window.parent.postMessage({
39+ if (width > 0 && height > 0 && window.parent !== window) {{
40+ window.parent.postMessage({{
3741 type: 'pyplots-size',
3842 width: Math.ceil(width),
3943 height: Math.ceil(height)
40- }, 'https://pyplots.ai ');
41- }
42- } catch (e) {
44+ }} , '{ target_origin } ');
45+ }}
46+ }} catch (e) { {
4347 // Silently fail if postMessage is blocked
44- }
45- }
48+ }}
49+ }}
4650
4751 // Report after load and after delays (for async rendering libraries)
48- if (document.readyState === 'complete') {
52+ if (document.readyState === 'complete') {{
4953 setTimeout(reportSize, 100);
5054 setTimeout(reportSize, 500);
5155 setTimeout(reportSize, 1000);
52- } else {
53- window.addEventListener('load', function() {
56+ }} else { {
57+ window.addEventListener('load', function() {{
5458 setTimeout(reportSize, 100);
5559 setTimeout(reportSize, 500);
5660 setTimeout(reportSize, 1000);
57- });
58- }
59- })();
61+ }} );
62+ }}
63+ }} )();
6064</script>
6165"""
6266
67+
68+ # Legacy constant for backwards compatibility with tests
69+ SIZE_REPORTER_SCRIPT = get_size_reporter_script ("https://pyplots.ai" )
70+
6371# Allowed GCS bucket for security
6472ALLOWED_HOST = "storage.googleapis.com"
6573ALLOWED_BUCKET = "pyplots-images"
@@ -107,7 +115,7 @@ def build_safe_gcs_url(url: str) -> str | None:
107115
108116
109117@router .get ("/proxy/html" , response_class = HTMLResponse )
110- async def proxy_html (url : str ):
118+ async def proxy_html (url : str , origin : str | None = None ):
111119 """
112120 Proxy an HTML file and inject size reporting script.
113121
@@ -118,6 +126,7 @@ async def proxy_html(url: str):
118126
119127 Args:
120128 url: The GCS URL to fetch (must be from allowed bucket)
129+ origin: Target origin for postMessage (must be in ALLOWED_ORIGINS)
121130
122131 Returns:
123132 Modified HTML with size reporting script injected
@@ -127,6 +136,11 @@ async def proxy_html(url: str):
127136 if safe_url is None :
128137 raise HTTPException (status_code = 400 , detail = f"Only URLs from { ALLOWED_HOST } /{ ALLOWED_BUCKET } are allowed" )
129138
139+ # Validate origin parameter - default to production if not specified or invalid
140+ target_origin = "https://pyplots.ai"
141+ if origin and origin in ALLOWED_ORIGINS :
142+ target_origin = origin
143+
130144 # Fetch the HTML with shorter timeout
131145 async with httpx .AsyncClient (timeout = 10.0 ) as client :
132146 try :
@@ -139,13 +153,16 @@ async def proxy_html(url: str):
139153
140154 html_content = response .text
141155
156+ # Generate script with correct target origin
157+ size_script = get_size_reporter_script (target_origin )
158+
142159 # Inject the size reporter script before </body>
143160 if "</body>" in html_content :
144- html_content = html_content .replace ("</body>" , f"{ SIZE_REPORTER_SCRIPT } </body>" )
161+ html_content = html_content .replace ("</body>" , f"{ size_script } </body>" )
145162 elif "</html>" in html_content :
146- html_content = html_content .replace ("</html>" , f"{ SIZE_REPORTER_SCRIPT } </html>" )
163+ html_content = html_content .replace ("</html>" , f"{ size_script } </html>" )
147164 else :
148165 # Fallback: append to end
149- html_content += SIZE_REPORTER_SCRIPT
166+ html_content += size_script
150167
151168 return HTMLResponse (content = html_content )
0 commit comments