use crate::error::DbError;
use mlua::{Lua, Value as LuaValue};
fn fetch_max_bytes() -> usize {
static MAX: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
*MAX.get_or_init(|| {
std::env::var("SOLIDB_LUA_FETCH_MAX_BYTES")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.filter(|v| *v > 0)
.unwrap_or(10 * 1024 * 1024)
})
}
async fn validate_fetch_target(
url: &str,
) -> Result<(url::Url, crate::server::ssrf::ValidatedTarget), String> {
if !url.starts_with("http://") && !url.starts_with("https://") {
return Err("only http and https schemes are allowed".into());
}
let parsed = url::Url::parse(url).map_err(|e| format!("invalid URL: {}", e))?;
let for_check = parsed.clone();
let target = tokio::task::spawn_blocking(move || {
crate::server::ssrf::validate_public_url_target(&for_check)
})
.await
.map_err(|e| format!("validation task failed: {}", e))??;
Ok((parsed, target))
}
async fn read_capped_body(mut res: reqwest::Response, max: usize) -> Result<Vec<u8>, String> {
if let Some(len) = res.content_length() {
if len > max as u64 {
return Err(format!("response body exceeds {} bytes", max));
}
}
let mut body = Vec::new();
while let Some(chunk) = res
.chunk()
.await
.map_err(|e| format!("reading body: {}", e))?
{
if body.len() + chunk.len() > max {
return Err(format!("response body exceeds {} bytes", max));
}
body.extend_from_slice(&chunk);
}
Ok(body)
}
pub fn create_fetch_function(lua: &Lua) -> Result<mlua::Function, DbError> {
lua.create_async_function(
|lua, (url, options): (String, Option<LuaValue>)| async move {
let (parsed, target) = validate_fetch_target(&url)
.await
.map_err(|e| mlua::Error::RuntimeError(format!("SSRF protection: {}", e)))?;
let client = reqwest::Client::builder()
.resolve(&target.host, target.addr)
.redirect(reqwest::redirect::Policy::none())
.timeout(std::time::Duration::from_secs(30))
.build()
.map_err(|e| mlua::Error::RuntimeError(format!("HTTP client: {}", e)))?;
let url = parsed.as_str().to_string();
let mut req_builder = client.get(&url);
if let Some(LuaValue::Table(t)) = options {
if let Ok(method) = t.get::<String>("method") {
match method.to_uppercase().as_str() {
"POST" => req_builder = client.post(&url),
"PUT" => req_builder = client.put(&url),
"DELETE" => req_builder = client.delete(&url),
"PATCH" => req_builder = client.patch(&url),
"HEAD" => req_builder = client.head(&url),
_ => {} }
}
if let Ok(LuaValue::Table(h)) = t.get::<LuaValue>("headers") {
for (k, v) in h.pairs::<String, String>().flatten() {
req_builder = req_builder.header(k, v);
}
}
if let Ok(body) = t.get::<String>("body") {
req_builder = req_builder.body(body);
}
}
match req_builder.send().await {
Ok(res) => {
let status = res.status().as_u16();
let headers_map = res.headers().clone();
let bytes = read_capped_body(res, fetch_max_bytes())
.await
.map_err(|e| mlua::Error::RuntimeError(format!("Fetch error: {}", e)))?;
let text = lua.create_string(&bytes)?;
let response_table = lua.create_table()?;
response_table.set("status", status)?;
response_table.set("body", text)?;
response_table.set("ok", (200..300).contains(&status))?;
let resp_headers = lua.create_table()?;
for (k, v) in headers_map.iter() {
if let Ok(val_str) = v.to_str() {
resp_headers.set(k.as_str(), val_str)?;
}
}
response_table.set("headers", resp_headers)?;
Ok(response_table)
}
Err(e) => Err(mlua::Error::RuntimeError(format!("Fetch error: {}", e))),
}
},
)
.map_err(|e| DbError::InternalError(format!("Failed to create fetch function: {}", e)))
}