feat: use request host for hot-reload instead of window.location.host

Cargo.toml

@@ -49,6 +49,8 @@ reqwest = { version = "0.11", default-features = false, features = [
   "stream",
   "trust-dns",
 ] }
+http-body = "0.4"
+hyper = "0.14"
 seahash = "4"
 serde = { version = "1", features = ["derive"] }
 tar = "0.4"

src/autoreload.js

@@ -1,6 +1,6 @@
 (function () {
     var protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
-    var url = protocol + '//' + window.location.host + '/_trunk/ws';
+    var url = protocol + '//' + '{{__TRUNK_ADDRESS__}}' + '/_trunk/ws';
     var poll_interval = 5000;
     var reload_upon_connect = () => {
         window.setTimeout(

src/serve.rs

@@ -3,12 +3,18 @@ use std::path::PathBuf;
 use std::sync::Arc;
 
 use anyhow::{Context, Result};
-use axum::body::{self, Body};
+use axum::body::{self, Body, Bytes};
 use axum::extract::ws::{WebSocket, WebSocketUpgrade};
-use axum::http::StatusCode;
+use axum::http::response::Parts;
+use axum::http::{
+    header::{CONTENT_TYPE, HOST},
+    Request, StatusCode,
+};
+use axum::middleware::Next;
 use axum::response::Response;
 use axum::routing::{get, get_service, Router};
 use axum::Server;
+use hyper::header::CONTENT_LENGTH;
 use tokio::sync::broadcast;
 use tokio::task::JoinHandle;
 use tower_http::services::{ServeDir, ServeFile};
@@ -227,7 +233,8 @@ fn router(state: Arc<State>, cfg: Arc<RtcServe>) -> Router {
                     tracing::error!(?error, "failed serving static file");
                     StatusCode::INTERNAL_SERVER_ERROR
                 })
-                .layer(TraceLayer::new_for_http()),
+                .layer(TraceLayer::new_for_http())
+                .layer(axum::middleware::from_fn(html_address_middleware)),
             ),
         )
         .route(
@@ -303,6 +310,45 @@ fn router(state: Arc<State>, cfg: Arc<RtcServe>) -> Router {
     router
 }
 
+async fn html_address_middleware<B: std::fmt::Debug>(
+    request: Request<B>,
+    next: Next<B>,
+) -> (Parts, Bytes) {
+    let uri = request.headers().get(HOST).cloned();
+    let response = next.run(request).await;
+    let (parts, body) = response.into_parts();
+
+    match hyper::body::to_bytes(body).await {
+        Err(_) => (parts, Bytes::default()),
+        Ok(bytes) => {
+            let (mut parts, mut bytes) = (parts, bytes);
+
+            if let Some(uri) = uri {
+                if parts
+                    .headers
+                    .get(CONTENT_TYPE)
+                    .map(|t| t == "text/html")
+                    .unwrap_or(false)
+                {
+                    if let Ok(data_str) = std::str::from_utf8(&bytes) {
+                        let data_str = data_str.replace(
+                            "'{{__TRUNK_ADDRESS__}}'",
+                            &uri.to_str()
+                                .map(|s| format!("'{}'", s))
+                                .unwrap_or_else(|_| "window.location.href".into()),
+                        );
+                        let bytes_vec = data_str.as_bytes().to_vec();
+                        parts.headers.insert(CONTENT_LENGTH, bytes_vec.len().into());
+                        bytes = Bytes::from(bytes_vec);
+                    }
+                }
+            }
+
+            (parts, bytes)
+        }
+    }
+}
+
 async fn handle_ws(mut ws: WebSocket, state: Arc<State>) {
     let mut rx = state.build_done_chan.subscribe();
     tracing::debug!("autoreload websocket opened");