From b76e215256fa8d6d7cc914910eeed95d995458c5 Mon Sep 17 00:00:00 2001 From: Stanisław Zagórowski Date: Mon, 29 Jan 2024 17:07:54 +0100 Subject: feature: add header support --- src/luax/mod.rs | 35 ++++++++++++++++------------- src/server.rs | 68 ++++++++++++++++++++++++++++++++------------------------- 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 { 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> { - 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.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> { - 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.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> { - 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.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.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>> { } fn decode_luax_response(status: StatusCode, t: rlua::Table) -> Result>> { - if t.contains_key("type")? { + let lua_headers: Option = t.get("headers")?; + + let (response_body, mime_type) = if t.contains_key("type")? { let ty: String = t.get("type")?; let mime_type: Option = 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, "")?; - 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, "")?; 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::() { + 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> for State { @@ -150,10 +148,15 @@ impl Service> 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, + headers: Vec<(String, String)>, ) -> Result>> { 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, -- cgit v1.3.1