Proxy vite dev server (#10)

* refactor: pass reqwest from tuono_lib crate and prepared vite_reverse_proxy

* feat: add proxy entry point on vite side

* refactor: internal_handlers_folder

* feat: handle vite websocket tunnel

* feat: handle unexpect WebSocket messages

* feat: split dev/prod ssr handling

* feat: split js/rust bundle reloading

* feat: prevent close error on websocket connection

* feat: remove vite proxy on prod server

* feat: update version to v0.4.0

crates/tuono/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tuono"
-version = "0.3.1"
+version = "0.4.0"
 edition = "2021"
 authors = ["V. Ageno <valerioageno@yahoo.it>"]
 description = "The react/rust fullstack framework"

crates/tuono/src/source_builder.rs

@@ -48,7 +48,6 @@ pub const AXUM_ENTRY_POINT: &str = r##"
 // Do not manually change it
 
 use tuono_lib::{tokio, Mode, Server, axum::Router, axum::routing::get};
-use reqwest::Client;
 
 const MODE: Mode = /*MODE*/;
 
@@ -56,11 +55,9 @@ const MODE: Mode = /*MODE*/;
 
 #[tokio::main]
 async fn main() {
-    let fetch = Client::new();
-
     let router = Router::new()
         // ROUTE_BUILDER
-        .with_state(fetch);
+        ;
 
     Server::init(router, MODE).start().await
 }

crates/tuono/src/watch.rs

@@ -57,27 +57,32 @@ pub async fn watch() -> Result<()> {
     build_ssr_bundle.to_wait().await;
 
     let wx = Watchexec::new(move |mut action| {
-        let mut should_reload = false;
+        let mut should_reload_ssr_bundle = false;
+        let mut should_reload_rust_server = false;
 
         for event in action.events.iter() {
             for path in event.paths() {
-                if path.0.to_string_lossy().ends_with(".rs")
-                    // Either tsx and jsx
-                    || path.0.to_string_lossy().ends_with("sx")
-                {
-                    should_reload = true
+                if path.0.to_string_lossy().ends_with(".rs") {
+                    should_reload_rust_server = true
+                }
+
+                // Either tsx and jsx
+                if path.0.to_string_lossy().ends_with("sx") {
+                    should_reload_ssr_bundle = true
                 }
             }
         }
 
-        if should_reload {
+        if should_reload_rust_server {
             println!("Reloading...");
             run_server.stop();
-            build_ssr_bundle.stop();
-            build_ssr_bundle.start();
             bundle_axum_source(Mode::Dev).expect("Failed to bunlde rust source");
             run_server.start();
-            println!("Ready!");
+        }
+
+        if should_reload_ssr_bundle {
+            build_ssr_bundle.stop();
+            build_ssr_bundle.start();
         }
 
         // if Ctrl-C is received, quit

crates/tuono_lib/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tuono_lib"
-version = "0.3.1"
+version = "0.4.0"
 edition = "2021"
 authors = ["V. Ageno <valerioageno@yahoo.it>"]
 description = "The react/rust fullstack framework"
@@ -19,16 +19,20 @@ path = "src/lib.rs"
 
 [dependencies]
 ssr_rs = "0.5.5"
-axum = {version = "0.7.5", features = ["json"]}
+axum = {version = "0.7.5", features = ["json", "ws"]}
 tokio = { version = "1.37.0", features = ["full"] }
 serde = { version = "1.0.202", features = ["derive"] }
 erased-serde = "0.4.5"
 serde_json = "1.0"
-
-tuono_lib_macros = {path = "../tuono_lib_macros", version = "0.3.1"}
+reqwest = {version = "0.12.4", features = ["json", "stream"]}
 once_cell = "1.19.0"
-lazy_static = "1.5.0"
 regex = "1.10.5"
 either = "1.13.0"
 tower-http = {version = "0.5.2", features = ["fs"]}
 
+tuono_lib_macros = {path = "../tuono_lib_macros", version = "0.4.0"}
+# Match the same version used by axum
+tokio-tungstenite = "0.21.0"
+futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
+tungstenite = "0.23.0"
+

crates/tuono_lib/src/internal_handlers.rs

@@ -0,0 +1,7 @@
+mod catch_all;
+mod vite_reverse_proxy;
+mod vite_websocket_proxy;
+
+pub use catch_all::catch_all;
+pub use vite_reverse_proxy::vite_reverse_proxy;
+pub use vite_websocket_proxy::vite_websocket_proxy;

crates/tuono_lib/src/internal_handlers/catch_all.rs

@@ -15,7 +15,7 @@ pub async fn catch_all(
     // TODO: remove unwrap
     let payload = Payload::new(&req, &"").client_payload().unwrap();
 
-    let result = Js::SSR.with(|ssr| ssr.borrow_mut().render_to_string(Some(&payload)));
+    let result = Js::render_to_string(Some(&payload));
 
     match result {
         Ok(html) => Html(html),

crates/tuono_lib/src/internal_handlers/vite_reverse_proxy.rs

@@ -0,0 +1,33 @@
+use axum::body::Body;
+use axum::extract::{Path, State};
+
+use axum::http::{HeaderName, HeaderValue};
+use axum::response::{IntoResponse, Response};
+use reqwest::Client;
+
+const VITE_URL: &str = "http://localhost:3001/vite-server";
+
+pub async fn vite_reverse_proxy(
+    State(client): State<Client>,
+    Path(path): Path<String>,
+) -> impl IntoResponse {
+    match client.get(format!("{VITE_URL}/{path}")).send().await {
+        Ok(res) => {
+            let mut response_builder = Response::builder().status(res.status().as_u16());
+
+            {
+                let headers = response_builder.headers_mut().unwrap();
+                res.headers().into_iter().for_each(|(name, value)| {
+                    let name = HeaderName::from_bytes(name.as_ref()).unwrap();
+                    let value = HeaderValue::from_bytes(value.as_ref()).unwrap();
+                    headers.insert(name, value);
+                });
+            }
+
+            response_builder
+                .body(Body::from_stream(res.bytes_stream()))
+                .unwrap()
+        }
+        Err(_) => todo!(),
+    }
+}

crates/tuono_lib/src/internal_handlers/vite_websocket_proxy.rs

@@ -0,0 +1,104 @@
+use axum::extract::ws::{self, WebSocket, WebSocketUpgrade};
+use axum::response::IntoResponse;
+use futures_util::{SinkExt, StreamExt};
+use tokio_tungstenite::connect_async;
+use tokio_tungstenite::tungstenite::{Error, Message};
+use tungstenite::client::IntoClientRequest;
+use tungstenite::ClientRequestBuilder;
+
+const VITE_WS: &str = "ws://localhost:3001/vite-server/";
+const VITE_WS_PROTOCOL: &str = "vite-hmr";
+
+/// This is the entry point to proxy the vite's WebSocket.
+/// The proxy is needed for allowing all the development requests to point
+/// to localhost:3000/*. This enabled the framework to be developed in a Docker
+/// environment with just the 3000 port exposed.
+pub async fn vite_websocket_proxy(ws: WebSocketUpgrade) -> impl IntoResponse {
+    ws.protocols([VITE_WS_PROTOCOL]).on_upgrade(handle_socket)
+}
+
+async fn handle_socket(mut tuono_socket: WebSocket) {
+    // Send back a message to confirm connection
+    if tuono_socket
+        .send(ws::Message::Ping("tuono connected".into()))
+        .await
+        .is_err()
+    {
+        // If is error close the connection
+        return;
+    }
+
+    let vite_ws_request = ClientRequestBuilder::new(VITE_WS.parse().unwrap())
+        .with_sub_protocol(VITE_WS_PROTOCOL)
+        .into_client_request()
+        .expect("Failed to create vite WS request");
+
+    let vite_socket = match connect_async(vite_ws_request).await {
+        Ok((stream, _)) => {
+            // Connected to vite's WebSocket
+            stream
+        }
+        Err(e) => {
+            eprintln!("Failed to connect to vite's WebSocket. Error: {e}");
+            // As fallback vite automatically connect to port 3001.
+            return;
+        }
+    };
+
+    let (mut vite_sender, mut vite_receiver) = vite_socket.split();
+    let (mut tuono_sender, mut tuono_receiver) = tuono_socket.split();
+
+    // Handle browser messages.
+    // Every message gets forwarded to the vite WebSocket
+    tokio::spawn(async move {
+        while let Some(msg) = tuono_receiver.next().await {
+            if let Ok(msg) = msg {
+                let msg_to_vite = match msg.clone() {
+                    ws::Message::Text(str) => Message::Text(str),
+                    ws::Message::Pong(payload) => Message::Pong(payload),
+                    ws::Message::Ping(payload) => Message::Ping(payload),
+                    ws::Message::Binary(payload) => Message::Binary(payload),
+                    // Hard to match axum and tungstenite close payload.
+                    // Not a priority
+                    ws::Message::Close(_) => Message::Close(None),
+                };
+
+                vite_sender
+                    .send(msg_to_vite)
+                    .await
+                    .expect("Failed to tunnel msg to vite's WebSocket");
+
+                msg
+            } else {
+                // Close browser's WebSocket connection.
+                return;
+            };
+        }
+    });
+
+    // Handle vite messages.
+    // Every message gets forwarded to the browser.
+    tokio::spawn(async move {
+        while let Some(Ok(msg)) = vite_receiver.next().await {
+            let msg_to_browser = match msg {
+                Message::Text(str) => ws::Message::Text(str),
+                Message::Ping(payload) => ws::Message::Ping(payload),
+                Message::Pong(payload) => ws::Message::Pong(payload),
+                Message::Binary(payload) => ws::Message::Binary(payload),
+                // Hard to match axum and tungstenite close payload.
+                // Not a priority
+                Message::Close(_) => ws::Message::Close(None),
+                _ => {
+                    eprintln!("Unexpected message from the vite WebSocket to the browser: {msg:?}");
+                    ws::Message::Text("Unhandled".to_string())
+                }
+            };
+
+            if let Err(err) = tuono_sender.send(msg_to_browser).await {
+                if err.to_string() != Error::AlreadyClosed.to_string() {
+                    eprintln!("Failed to send back message from vite to browser: {err}")
+                }
+            }
+        }
+    });
+}

crates/tuono_lib/src/lib.rs

@@ -1,4 +1,4 @@
-mod catch_all;
+mod internal_handlers;
 mod manifest;
 mod mode;
 mod payload;
@@ -16,4 +16,5 @@ pub use tuono_lib_macros::handler;
 
 // Re-exports
 pub use axum;
+pub use reqwest;
 pub use tokio;

crates/tuono_lib/src/response.rs

@@ -72,7 +72,7 @@ impl Response {
             Self::Props(Props { data, http_code }) => {
                 let payload = Payload::new(&req, data).client_payload().unwrap();
 
-                match Js::SSR.with(|ssr| ssr.borrow_mut().render_to_string(Some(&payload))) {
+                match Js::render_to_string(Some(&payload)) {
                     Ok(html) => (*http_code, Html(html)).into_response(),
                     Err(_) => {
                         (*http_code, Html("500 Internal server error".to_string())).into_response()

crates/tuono_lib/src/server.rs

@@ -5,18 +5,18 @@ use axum::routing::{get, Router};
 use ssr_rs::Ssr;
 use tower_http::services::ServeDir;
 
-use crate::catch_all::catch_all;
+use crate::internal_handlers::{catch_all, vite_reverse_proxy, vite_websocket_proxy};
 
 const DEV_PUBLIC_DIR: &str = "public";
 const PROD_PUBLIC_DIR: &str = "out/client";
 
 pub struct Server {
-    router: Router,
+    router: Router<reqwest::Client>,
     mode: Mode,
 }
 
 impl Server {
-    pub fn init(router: Router, mode: Mode) -> Server {
+    pub fn init(router: Router<reqwest::Client>, mode: Mode) -> Server {
         Ssr::create_platform();
 
         GLOBAL_MODE.set(mode).unwrap();
@@ -31,23 +31,32 @@ impl Server {
     pub async fn start(&self) {
         let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
 
+        let fetch = reqwest::Client::new();
+
         if self.mode == Mode::Dev {
             println!("\nDevelopment app ready at http://localhost:3000/");
+            let router = self
+                .router
+                .to_owned()
+                .route("/vite-server/", get(vite_websocket_proxy))
+                .route("/vite-server/*path", get(vite_reverse_proxy))
+                .fallback_service(ServeDir::new(DEV_PUBLIC_DIR).fallback(get(catch_all)))
+                .with_state(fetch);
+
+            axum::serve(listener, router)
+                .await
+                .expect("Failed to serve development server");
         } else {
             println!("\nProduction app ready at http://localhost:3000/");
+            let router = self
+                .router
+                .to_owned()
+                .fallback_service(ServeDir::new(PROD_PUBLIC_DIR).fallback(get(catch_all)))
+                .with_state(fetch);
+
+            axum::serve(listener, router)
+                .await
+                .expect("Failed to serve production server");
         }
-
-        let public_dir = if self.mode == Mode::Dev {
-            DEV_PUBLIC_DIR
-        } else {
-            PROD_PUBLIC_DIR
-        };
-
-        let router = self
-            .router
-            .to_owned()
-            .fallback_service(ServeDir::new(public_dir).fallback(get(catch_all)));
-
-        axum::serve(listener, router).await.unwrap();
     }
 }

crates/tuono_lib/src/ssr.rs

@@ -1,27 +1,48 @@
 use crate::mode::{Mode, GLOBAL_MODE};
-use lazy_static::lazy_static;
 use ssr_rs::Ssr;
 use std::cell::RefCell;
 use std::fs::read_to_string;
 use std::path::PathBuf;
 
-lazy_static! {
-    static ref BUNDLE_PATH: &'static str = {
-        if GLOBAL_MODE.get().unwrap() == &Mode::Dev {
-            return "./.tuono/server/dev-server.js";
-        }
-        "./out/server/prod-server.js"
-    };
-}
-
+/// For the server side rendering we need to split the implementation between dev and prod.
+/// This completely remove the multi-thread optimization on dev but allow the dev server to
+/// update the SSR result without reloading the whole server.
 pub struct Js;
 
 impl Js {
+    pub fn render_to_string(payload: Option<&str>) -> Result<String, &'static str> {
+        let mode = GLOBAL_MODE.get().expect("Failed to get GLOBAL_MODE");
+
+        if *mode == Mode::Dev {
+            DevJs::render_to_string(payload)
+        } else {
+            ProdJs::SSR.with(|ssr| ssr.borrow_mut().render_to_string(payload))
+        }
+    }
+}
+
+struct ProdJs;
+
+impl ProdJs {
     thread_local! {
         pub static SSR: RefCell<Ssr<'static, 'static>> = RefCell::new(
                 Ssr::from(
-                    read_to_string(PathBuf::from(*BUNDLE_PATH)).expect("Server bundle not found"), ""
+                    read_to_string(PathBuf::from("./out/server/prod-server.js")).expect("Server bundle not found"), ""
                     ).unwrap()
                 )
     }
 }
+
+struct DevJs;
+
+impl DevJs {
+    pub fn render_to_string(params: Option<&str>) -> Result<String, &'static str> {
+        Ssr::from(
+            read_to_string(PathBuf::from("./.tuono/server/dev-server.js"))
+                .expect("Server bundle not found"),
+            "",
+        )
+        .unwrap()
+        .render_to_string(params)
+    }
+}

crates/tuono_lib_macros/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tuono_lib_macros"
-version = "0.3.1"
+version = "0.4.0"
 edition = "2021"
 description = "The react/rust fullstack framework"
 repository = "https://github.com/Valerioageno/tuono"

crates/tuono_lib_macros/src/handler.rs

@@ -11,13 +11,12 @@ pub fn handler_core(_args: TokenStream, item: TokenStream) -> TokenStream {
         use tuono_lib::axum::response::IntoResponse;
         use std::collections::HashMap;
         use tuono_lib::axum::extract::{State, Path};
-        use reqwest::Client;
 
         #item
 
         pub async fn route(
             Path(params): Path<HashMap<String, String>>,
-            State(client): State<Client>,
+            State(client): State<tuono_lib::reqwest::Client>,
             request: tuono_lib::axum::extract::Request
         ) -> impl IntoResponse {
            let pathname = &request.uri();
@@ -30,7 +29,7 @@ pub fn handler_core(_args: TokenStream, item: TokenStream) -> TokenStream {
 
         pub async fn api(
             Path(params): Path<HashMap<String, String>>,
-            State(client): State<Client>,
+            State(client): State<tuono_lib::reqwest::Client>,
             request: tuono_lib::axum::extract::Request
         ) -> impl IntoResponse{
            let pathname = &request.uri();

examples/tutorial/Cargo.toml

@@ -10,5 +10,4 @@ path = ".tuono/main.rs"
 [dependencies]
 tuono_lib = { path = "../../crates/tuono_lib/"}
 serde = { version = "1.0.202", features = ["derive"] }
-reqwest = {version = "0.12.4", features = ["json"]}
 

examples/tutorial/src/routes/index.rs

@@ -1,6 +1,6 @@
 // src/routes/index.rs
-use reqwest::StatusCode;
 use serde::{Deserialize, Serialize};
+use tuono_lib::reqwest::{Client, StatusCode};
 use tuono_lib::{Props, Request, Response};
 
 const ALL_POKEMON: &str = "https://pokeapi.co/api/v2/pokemon?limit=151";
@@ -17,8 +17,8 @@ struct Pokemon {
 }
 
 #[tuono_lib::handler]
-async fn get_all_pokemons(_req: Request<'_>, fetch: reqwest::Client) -> Response {
-    return match fetch.get(ALL_POKEMON).send().await {
+async fn get_all_pokemons(_req: Request<'_>, fetch: Client) -> Response {
+    match fetch.get(ALL_POKEMON).send().await {
         Ok(res) => {
             let data = res.json::<Pokemons>().await.unwrap();
             Response::Props(Props::new(data))
@@ -27,5 +27,5 @@ async fn get_all_pokemons(_req: Request<'_>, fetch: reqwest::Client) -> Response
             "{}", // Return empty JSON
             StatusCode::INTERNAL_SERVER_ERROR,
         )),
-    };
+    }
 }

examples/tutorial/src/routes/pokemons/GOAT.rs

@@ -1,7 +1,7 @@
 // src/routes/pokemons/GOAT.rs
-use tuono_lib::{Request, Response};
+use tuono_lib::{reqwest::Client, Request, Response};
 
 #[tuono_lib::handler]
-async fn redirect_to_goat(_: Request<'_>, _: reqwest::Client) -> Response {
+async fn redirect_to_goat(_: Request<'_>, _: Client) -> Response {
     Response::Redirect("/pokemons/mewtwo".to_string())
 }

examples/tutorial/src/routes/pokemons/[pokemon].rs

@@ -1,6 +1,6 @@
 // src/routes/pokemons/[pokemon].rs
-use reqwest::StatusCode;
 use serde::{Deserialize, Serialize};
+use tuono_lib::reqwest::{Client, StatusCode};
 use tuono_lib::{Props, Request, Response};
 
 const POKEMON_API: &str = "https://pokeapi.co/api/v2/pokemon";
@@ -14,11 +14,11 @@ struct Pokemon {
 }
 
 #[tuono_lib::handler]
-async fn get_pokemon(req: Request<'_>, fetch: reqwest::Client) -> Response {
+async fn get_pokemon(req: Request<'_>, fetch: Client) -> Response {
     // The param `pokemon` is defined in the route filename [pokemon].rs
     let pokemon = req.params.get("pokemon").unwrap();
 
-    return match fetch.get(format!("{POKEMON_API}/{pokemon}")).send().await {
+    match fetch.get(format!("{POKEMON_API}/{pokemon}")).send().await {
         Ok(res) => {
             if res.status() == StatusCode::NOT_FOUND {
                 return Response::Props(Props::new_with_status("{}", StatusCode::NOT_FOUND));
@@ -30,5 +30,5 @@ async fn get_pokemon(req: Request<'_>, fetch: reqwest::Client) -> Response {
             "{}",
             StatusCode::INTERNAL_SERVER_ERROR,
         )),
-    };
+    }
 }

packages/lazy-fn-vite-plugin/package.json

@@ -1,6 +1,6 @@
 {
   "name": "tuono-lazy-fn-vite-plugin",
-  "version": "0.3.1",
+  "version": "0.4.0",
   "description": "Plugin for the tuono's lazy fn. Tuono is the react/rust fullstack framework",
   "scripts": {
     "dev": "vite build --watch",

packages/tuono/package.json

@@ -1,6 +1,6 @@
 {
   "name": "tuono",
-  "version": "0.3.1",
+  "version": "0.4.0",
   "description": "The react/rust fullstack framework",
   "scripts": {
     "dev": "vite build --watch",

packages/tuono/src/build/index.ts

@@ -12,6 +12,8 @@ const BASE_CONFIG: InlineConfig = {
   plugins: [react(), ViteFsRouter(), LazyLoadingPlugin()],
 }
 
+const VITE_PORT = 3001
+
 export function developmentSSRBundle() {
   ;(async () => {
     await build({
@@ -43,8 +45,10 @@ export function developmentCSRWatch() {
   ;(async () => {
     const server = await createServer({
       ...BASE_CONFIG,
+      // Entry point for the development vite proxy
+      base: '/vite-server/',
       server: {
-        port: 3001,
+        port: VITE_PORT,
         strictPort: true,
       },
       build: {

packages/tuono/src/ssr/index.tsx

@@ -7,15 +7,18 @@ import { createRouter } from '../router'
 type RouteTree = any
 type Mode = 'Dev' | 'Prod'
 
+const TUONO_DEV_SERVER_PORT = 3000
+const VITE_PROXY_PATH = '/vite-server'
+
 const VITE_DEV_AND_HMR = `<script type="module">
-import RefreshRuntime from 'http://localhost:3001/@react-refresh'
+import RefreshRuntime from 'http://localhost:${TUONO_DEV_SERVER_PORT}${VITE_PROXY_PATH}/@react-refresh'
 RefreshRuntime.injectIntoGlobalHook(window)
 window.$RefreshReg$ = () => {}
 window.$RefreshSig$ = () => (type) => type
 window.__vite_plugin_react_preamble_installed__ = true
 </script>
-<script type="module" src="http://localhost:3001/@vite/client"></script>
-<script type="module" src="http://localhost:3001/client-main.tsx"></script>`
+<script type="module" src="http://localhost:${TUONO_DEV_SERVER_PORT}${VITE_PROXY_PATH}/@vite/client"></script>
+<script type="module" src="http://localhost:${TUONO_DEV_SERVER_PORT}${VITE_PROXY_PATH}/client-main.tsx"></script>`
 
 function generateCssLinks(cssBundles: string[], mode: Mode): string {
   if (mode === 'Dev') return ''