feature: add header support

2 files changed, 58 insertions, 45 deletions

diff --git a/src/luax/mod.rs b/src/luax/mod.rs
index 70e7b34..994b80e 100644
--- a/src/luax/mod.rs
+++ b/src/luax/mod.rs
@@ -174,6 +174,8 @@ pub fn prepare_lua(dev_mode: bool) -> Result<Lua> {
         reluax.set("html_page", html_page)?;
         let json = ctx.create_function(utils::wrap_json)?;
         reluax.set("json", json)?;
+        let headers = ctx.create_function(utils::wrap_headers)?;
+        reluax.set("headers", headers)?;
         reluax.set("dev_mode", dev_mode)?;
 
         ctx.globals().set("reluax", reluax)?;
@@ -301,27 +303,30 @@ mod utils {
     }
 
     /// Wrap a table in a table to signal that it should be rendered as HTML
-    pub fn wrap_html<'lua>(ctx: Context<'lua>, table: Table<'lua>) -> Result<Table<'lua>> {
-        let html_table = ctx.create_table()?;
-        html_table.set("type", "html")?;
-        html_table.set("value", table)?;
-        Ok(html_table)
+    pub fn wrap_html<'lua>(_: Context<'lua>, table: Table<'lua>) -> Result<Table<'lua>> {
+        table.set("type", "html")?;
+        Ok(table)
     }
 
     /// Wrap a table in a table to signal that it should be rendered as a full HTML page
-    pub fn wrap_html_page<'lua>(ctx: Context<'lua>, table: Table<'lua>) -> Result<Table<'lua>> {
-        let html_table = ctx.create_table()?;
-        html_table.set("type", "html-page")?;
-        html_table.set("value", table)?;
-        Ok(html_table)
+    pub fn wrap_html_page<'lua>(_: Context<'lua>, table: Table<'lua>) -> Result<Table<'lua>> {
+        table.set("type", "html-page")?;
+        Ok(table)
     }
 
     /// Wrap a table in a table to signal that it should be rendered as JSON
-    pub fn wrap_json<'lua>(ctx: Context<'lua>, table: Table<'lua>) -> Result<Table<'lua>> {
-        let json_table = ctx.create_table()?;
-        json_table.set("type", "json")?;
-        json_table.set("value", table)?;
-        Ok(json_table)
+    pub fn wrap_json<'lua>(_: Context<'lua>, table: Table<'lua>) -> Result<Table<'lua>> {
+        table.set("type", "json")?;
+        Ok(table)
+    }
+
+    /// Wrap a table to add headers to the response
+    pub fn wrap_headers<'lua>(
+        _: Context<'lua>,
+        (table, headers): (Table<'lua>, Table<'lua>),
+    ) -> Result<Table<'lua>> {
+        table.set("headers", headers)?;
+        Ok(table)
     }
 
     #[cfg(test)]
diff --git a/src/server.rs b/src/server.rs
index 7aa3cec..f57bcdc 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -89,54 +89,52 @@ fn mk_file_response(path: PathBuf) -> Result<Response<Full<Bytes>>> {
 }
 
 fn decode_luax_response(status: StatusCode, t: rlua::Table) -> Result<Response<Full<Bytes>>> {
-    if t.contains_key("type")? {
+    let lua_headers: Option<rlua::Table> = t.get("headers")?;
+
+    let (response_body, mime_type)  = if t.contains_key("type")? {
         let ty: String = t.get("type")?;
         let mime_type: Option<String> = t.get("mime_type")?;
 
         match ty.as_str() {
             "html" => {
-                let html: rlua::Table = t.get("value")?;
                 let mut buf = Vec::new();
-                table_to_html(html, &mut buf)?;
-                Ok(Response::builder()
-                    .header("Content-Type", mime_type.unwrap_or("text/html".to_string()))
-                    .status(status)
-                    .body(Full::new(Bytes::from(buf)))?)
+                table_to_html(t, &mut buf)?;
+                (buf, mime_type.unwrap_or("text/html".to_string()))
             }
             "json" => {
-                let json: rlua::Table = t.get("value")?;
                 let mut buf = Vec::new();
-                table_to_json(json, &mut buf)?;
-
-                Ok(Response::builder()
-                    .header(
-                        "Content-Type",
-                        mime_type.unwrap_or("application/json".to_string()),
-                    )
-                    .status(status)
-                    .body(Full::new(Bytes::from(buf)))?)
+                table_to_json(t, &mut buf)?;
+                (buf, mime_type.unwrap_or("application/json".to_string()))
             }
             "html-page" => {
-                let html: rlua::Table = t.get("value")?;
                 let mut buf = Vec::new();
                 writeln!(&mut buf, "<!DOCTYPE html>")?;
-                table_to_html(html, &mut buf)?;
-                Ok(Response::builder()
-                    .header("Content-Type", mime_type.unwrap_or("text/html".to_string()))
-                    .status(status)
-                    .body(Full::new(Bytes::from(buf)))?)
+                table_to_html(t, &mut buf)?;
+                (buf, mime_type.unwrap_or("text/html".to_string()))
             }
-            _ => Err(ReluaxError::Server("Unknown response type".to_string()).into()),
+            _ => return Err(ReluaxError::Server("Unknown response type".to_string()).into()),
         }
     } else {
         let mut buf = Vec::new();
         writeln!(&mut buf, "<!DOCTYPE html>")?;
         table_to_html(t, &mut buf)?;
-        Ok(Response::builder()
-            .header("Content-Type", "text/html")
-            .status(status)
-            .body(Full::new(Bytes::from(buf)))?)
+        (buf, "text/html".to_string())
+    };
+
+    let mut response_builder = Response::builder()
+        .status(status)
+        .header("Content-Type", mime_type);
+
+    if let Some(lua_headers) = lua_headers {
+        for r in lua_headers.pairs::<String, String>() {
+            let (k, v) = r?;
+            response_builder = response_builder.header(k, v);
+        }
     }
+
+    let response = response_builder.body(Full::new(Bytes::from(response_body)))?;
+
+    Ok(response)
 }
 
 impl Service<Request<Incoming>> for State {
@@ -150,10 +148,15 @@ impl Service<Request<Incoming>> for State {
         let method = req.method().clone();
         let lua = self.lua.clone();
         let public_dir = self.public_dir.clone();
+        let headers = req
+            .headers()
+            .iter()
+            .map(|(k, v)| (k.to_string(), v.to_str().unwrap().to_string()))
+            .collect();
         Box::pin(async {
             let body = req.into_body().collect().await?;
 
-            Self::serve(lua, public_dir, path, method, body)
+            Self::serve(lua, public_dir, path, method, body, headers)
         })
     }
 }
@@ -165,6 +168,7 @@ impl State {
         path: String,
         method: Method,
         body: Collected<Bytes>,
+        headers: Vec<(String, String)>,
     ) -> Result<Response<Full<Bytes>>> {
         let lua = lua.lock().unwrap();
 
@@ -187,9 +191,13 @@ impl State {
 
             let method = method.as_str();
             let body: rlua::String = ctx.create_string(&body.to_bytes().to_vec())?;
+            let lua_headers: rlua::Table = ctx.create_table()?;
+            for (k, v) in headers {
+                lua_headers.set(k, v)?;
+            }
 
             let res: rlua::Result<(rlua::Integer, rlua::Value)> =
-                route.call((path.clone(), method, body));
+                route.call((path.clone(), method, lua_headers, body));
 
             let res = match res {
                 Ok(r) => r,