assay/lua/policy/
apply.rs1use std::sync::Arc;
7
8use mlua::{Lua, MultiValue, Table, Value};
9
10use super::{
11 active, credential, guard_http, is_redacted_header, redact_json_text, redact_keys,
12 response_limit,
13};
14
15const VERBS: &[&str] = &["get", "post", "put", "patch", "delete"];
16
17type Target = Arc<dyn Fn(&MultiValue) -> Option<(String, String, usize)>>;
20
21pub fn apply(lua: &Lua) -> mlua::Result<()> {
22 let Some(http) = lua.globals().get::<Option<Table>>("http")? else {
23 return Ok(());
24 };
25 for verb in VERBS {
26 wrap(lua, &http, verb, verb_target(verb))?;
27 }
28 wrap(lua, &http, "_client_request", client_request_target())?;
31 wrap(lua, &http, "download", verb_target("get"))?;
32 Ok(())
33}
34
35fn verb_target(verb: &'static str) -> Target {
36 Arc::new(move |args| arg_string(args, 0).map(|url| (verb.to_string(), url, 0)))
37}
38
39fn client_request_target() -> Target {
40 Arc::new(|args| Some((arg_string(args, 1)?, arg_string(args, 2)?, 2)))
41}
42
43fn arg_string(args: &MultiValue, at: usize) -> Option<String> {
44 match args.iter().nth(at) {
45 Some(Value::String(s)) => s.to_str().ok().map(|s| s.to_string()),
46 _ => None,
47 }
48}
49
50fn wrap(lua: &Lua, http: &Table, name: &str, target: Target) -> mlua::Result<()> {
51 let Value::Function(inner) = http.get::<Value>(name)? else {
52 return Ok(());
53 };
54 let wrapper = lua.create_async_function(move |lua, args: MultiValue| {
55 let inner = inner.clone();
56 let target = Arc::clone(&target);
57 async move {
58 let mut args = args;
59 if let Some((method, url, url_index)) = target(&args) {
60 credential::reject_in_url(&url)?;
61 guard_http(&lua, &method, &url)?;
62 args = fill_credentials(&lua, args, url_index)?;
63 }
64 let result = inner.call_async::<Value>(args).await?;
65 sanitize(&lua, result)
66 }
67 })?;
68 http.set(name, wrapper)
69}
70
71fn fill_credentials(lua: &Lua, args: MultiValue, url_index: usize) -> mlua::Result<MultiValue> {
75 let Some(policy) = active(lua) else {
76 return Ok(args);
77 };
78 if policy.credentials.is_empty() {
79 return Ok(args);
80 }
81 let mut out = Vec::with_capacity(args.len());
82 for (i, value) in args.into_iter().enumerate() {
83 out.push(if i == url_index {
84 value
85 } else {
86 credential::substitute(lua, &policy, value)?
87 });
88 }
89 Ok(MultiValue::from_iter(out))
90}
91
92fn sanitize(lua: &Lua, result: Value) -> mlua::Result<Value> {
96 let Value::Table(table) = &result else {
97 return Ok(result);
98 };
99 if let Some(limit) = response_limit(lua)
100 && let Ok(body) = table.get::<mlua::String>("body")
101 && body.as_bytes().len() > limit
102 {
103 return Err(mlua::Error::runtime(format!(
104 "policy: response body exceeds max_response_bytes ({limit})"
105 )));
106 }
107
108 let keys = redact_keys(lua);
109 if keys.is_empty() {
110 return Ok(result);
111 }
112 redact_body(lua, table, &keys)?;
113 redact_headers(table, &keys)?;
114 Ok(result)
115}
116
117fn redact_body(lua: &Lua, table: &Table, keys: &[String]) -> mlua::Result<()> {
118 let Ok(body) = table.get::<mlua::String>("body") else {
119 return Ok(());
120 };
121 let bytes = body.as_bytes();
122 let Some(redacted) = std::str::from_utf8(&bytes)
123 .ok()
124 .and_then(|text| redact_json_text(text, keys))
125 else {
126 return Ok(());
127 };
128 table.set("body", lua.create_string(redacted.as_bytes())?)
129}
130
131fn redact_headers(table: &Table, keys: &[String]) -> mlua::Result<()> {
132 let Ok(headers) = table.get::<Table>("headers") else {
133 return Ok(());
134 };
135 let names: Vec<String> = headers
136 .clone()
137 .pairs::<String, Value>()
138 .filter_map(|pair| pair.ok().map(|(name, _)| name))
139 .filter(|name| is_redacted_header(name, keys))
140 .collect();
141 for name in names {
142 headers.set(name, "[redacted]")?;
143 }
144 Ok(())
145}