fix: #0 rilot to take dynamic value

maninderpreetsingh · maninderpreetsingh · commit c65797b1cc14 · 2025-04-28T21:11:18.000-04:00
diff --git a/examples/src/lib.rs b/examples/src/lib.rs
@@ -2,7 +2,7 @@ use serde::Serialize;
 
 #[derive(Serialize)]
 struct WasmResponse {
-    app_url: String,
+    app_url: Option<String>,
     headers_to_update: std::collections::HashMap<String, String>,
     headers_to_remove: Vec<String>,
 }
@@ -21,9 +21,10 @@ pub extern "C" fn modify_request(ptr: *mut u8, len: usize) -> i32 {
         let mut headers = std::collections::HashMap::new();
         headers.insert("X-Debug-Path".to_string(), path.to_string());
         headers.insert("X-Debug-Method".to_string(), method.to_string());
+        headers.insert("X-Via-Rilot".to_string(), "5".to_string());
 
         let response = WasmResponse {
-            app_url: "http://127.0.0.1:5501/".to_string(), // change to any backend dynamically
+            app_url: None, // change to any backend dynamically if needed else default based on config.json
             headers_to_update: headers,
             headers_to_remove: vec![],
         };
diff --git a/src/proxy.rs b/src/proxy.rs
@@ -118,11 +118,9 @@ async fn handle_request(mut req: Request<Body>, config: Arc<config::Config>) ->
                     }
                 }
             }
-
-            // TODO:: update body?
         }
 
-        // Direct forwarding
+        // Prepare new request
         let method = req.method().clone();
         let headers = req.headers().clone();
         let body = req.into_body();
@@ -143,12 +141,20 @@ async fn handle_request(mut req: Request<Body>, config: Arc<config::Config>) ->
         let resp_result = client.request(new_req).await;
 
         match resp_result {
-            Ok(mut backend_resp) => {
-                let backend_body = hyper::body::to_bytes(backend_resp.body_mut()).await.unwrap_or_default();
-                Ok(Response::builder()
-                    .status(backend_resp.status())
-                    .body(Body::from(backend_body))
-                    .expect("Failed to build final response"))
+            Ok(backend_resp) => {
+                // Forward body as-is, with original headers
+                let status = backend_resp.status();
+                let headers = backend_resp.headers().clone();
+                let body = backend_resp.into_body();
+
+                let mut response = Response::builder()
+                    .status(status)
+                    .body(body)
+                    .expect("Failed to build final response");
+
+                *response.headers_mut() = headers;
+
+                Ok(response)
             }
             Err(_) => Ok(Response::builder()
                 .status(502)
diff --git a/src/wasm_engine.rs b/src/wasm_engine.rs
@@ -9,28 +9,28 @@ pub struct WasmOutput {
 }
 
 pub fn run_modify_request(wasm_path: &str, input: &str) -> Result<WasmOutput, String> {
-    let engine = Engine::default();
-    let module = Module::from_file(&engine, wasm_path).map_err(|e| e.to_string())?;
-    let mut store = Store::new(&engine, ());
-    let instance = Instance::new(&mut store, &module, &[]).map_err(|e| e.to_string())?;
+    let engine: Engine = Engine::default();
+    let module: Module = Module::from_file(&engine, wasm_path).map_err(|e| e.to_string())?;
+    let mut store: Store<()> = Store::new(&engine, ());
+    let instance: Instance = Instance::new(&mut store, &module, &[]).map_err(|e| e.to_string())?;
 
-    let func = instance.get_typed_func::<(i32, i32), i32>(&mut store, "modify_request")
+    let func: TypedFunc<(i32, i32), i32> = instance.get_typed_func::<(i32, i32), i32>(&mut store, "modify_request")
         .map_err(|e| e.to_string())?;
 
-    let memory = instance.get_memory(&mut store, "memory")
+    let memory: Memory = instance.get_memory(&mut store, "memory")
         .ok_or_else(|| "failed to find memory export".to_string())?;
 
-    let input_bytes = input.as_bytes();
-    let ptr = 0;
+    let input_bytes: &[u8] = input.as_bytes();
+    let ptr: usize = 0;
     memory.write(&mut store, ptr, input_bytes).map_err(|e| e.to_string())?;
 
-    let result_ptr = func.call(&mut store, (ptr as i32, input_bytes.len() as i32))
+    let result_ptr: i32 = func.call(&mut store, (ptr as i32, input_bytes.len() as i32))
         .map_err(|e| e.to_string())?;
 
-    let mut buffer = [0u8; 4096];
+    let mut buffer: [u8; 4096] = [0u8; 4096];
     memory.read(&mut store, result_ptr as usize, &mut buffer).map_err(|e| e.to_string())?;
     let nul_pos = buffer.iter().position(|&b| b == 0).unwrap_or(buffer.len());
 
-    let output = String::from_utf8_lossy(&buffer[..nul_pos]).to_string();
+    let output: String = String::from_utf8_lossy(&buffer[..nul_pos]).to_string();
     serde_json::from_str(&output).map_err(|e| e.to_string())
 }
