1use anyhow::Result;
2use base64::{engine::general_purpose, Engine as _};
3use hcl::eval::{Context as HclContext, FuncArgs, FuncDef, ParamType};
4use hcl::Value;
5
6pub fn register_stdlib(ctx: &mut HclContext) {
7 ctx.declare_func(
9 "upper",
10 FuncDef::builder().param(ParamType::String).build(upper),
11 );
12 ctx.declare_func(
13 "lower",
14 FuncDef::builder().param(ParamType::String).build(lower),
15 );
16 ctx.declare_func(
17 "trim",
18 FuncDef::builder()
19 .param(ParamType::String)
20 .param(ParamType::String)
21 .build(trim),
22 );
23 ctx.declare_func(
24 "trimspace",
25 FuncDef::builder().param(ParamType::String).build(trimspace),
26 );
27 ctx.declare_func(
28 "split",
29 FuncDef::builder()
30 .param(ParamType::String)
31 .param(ParamType::String)
32 .build(split),
33 );
34 ctx.declare_func(
35 "join",
36 FuncDef::builder()
37 .param(ParamType::String)
38 .param(ParamType::Array(Box::new(ParamType::Any)))
39 .build(join),
40 );
41 ctx.declare_func(
42 "replace",
43 FuncDef::builder()
44 .param(ParamType::String)
45 .param(ParamType::String)
46 .param(ParamType::String)
47 .build(replace),
48 );
49
50 ctx.declare_func(
52 "jsonencode",
53 FuncDef::builder().param(ParamType::Any).build(jsonencode),
54 );
55 ctx.declare_func(
56 "jsondecode",
57 FuncDef::builder()
58 .param(ParamType::String)
59 .build(jsondecode),
60 );
61 ctx.declare_func(
62 "base64encode",
63 FuncDef::builder()
64 .param(ParamType::String)
65 .build(base64encode),
66 );
67 ctx.declare_func(
68 "base64decode",
69 FuncDef::builder()
70 .param(ParamType::String)
71 .build(base64decode),
72 );
73
74 ctx.declare_func(
76 "length",
77 FuncDef::builder().param(ParamType::Any).build(length),
78 );
79 ctx.declare_func(
80 "keys",
81 FuncDef::builder()
82 .param(ParamType::Object(Box::new(ParamType::Any)))
83 .build(keys),
84 );
85 ctx.declare_func(
86 "values",
87 FuncDef::builder()
88 .param(ParamType::Object(Box::new(ParamType::Any)))
89 .build(values),
90 );
91 ctx.declare_func(
92 "contains",
93 FuncDef::builder()
94 .param(ParamType::Array(Box::new(ParamType::Any)))
95 .param(ParamType::Any)
96 .build(contains),
97 );
98
99 ctx.declare_func(
101 "coalesce",
102 FuncDef::builder()
103 .variadic_param(ParamType::Any)
104 .build(coalesce),
105 );
106}
107
108fn upper(args: FuncArgs) -> Result<Value, String> {
109 if let Some(Value::String(s)) = args.first() {
110 Ok(Value::String(s.to_uppercase()))
111 } else {
112 Err("upper() expects a string argument".to_string())
113 }
114}
115
116fn lower(args: FuncArgs) -> Result<Value, String> {
117 if let Some(Value::String(s)) = args.first() {
118 Ok(Value::String(s.to_lowercase()))
119 } else {
120 Err("lower() expects a string argument".to_string())
121 }
122}
123
124fn trim(args: FuncArgs) -> Result<Value, String> {
125 if args.len() != 2 {
126 return Err("trim() expects exactly 2 arguments".to_string());
127 }
128 match (&args[0], &args[1]) {
129 (Value::String(s), Value::String(cutset)) => {
130 let cutset_chars: Vec<char> = cutset.chars().collect();
131 Ok(Value::String(
132 s.trim_matches(|c| cutset_chars.contains(&c)).to_string(),
133 ))
134 }
135 _ => Err("trim() expects string arguments".to_string()),
136 }
137}
138
139fn trimspace(args: FuncArgs) -> Result<Value, String> {
140 if let Some(Value::String(s)) = args.first() {
141 Ok(Value::String(s.trim().to_string()))
142 } else {
143 Err("trimspace() expects a string argument".to_string())
144 }
145}
146
147fn split(args: FuncArgs) -> Result<Value, String> {
148 if args.len() != 2 {
149 return Err("split() expects exactly 2 arguments".to_string());
150 }
151 match (&args[0], &args[1]) {
152 (Value::String(sep), Value::String(s)) => {
153 let parts: Vec<Value> = s
154 .split(sep)
155 .map(|part| Value::String(part.to_string()))
156 .collect();
157 Ok(Value::Array(parts))
158 }
159 _ => Err("split() expects string arguments".to_string()),
160 }
161}
162
163fn join(args: FuncArgs) -> Result<Value, String> {
164 if args.len() != 2 {
165 return Err("join() expects exactly 2 arguments".to_string());
166 }
167 match (&args[0], &args[1]) {
168 (Value::String(sep), Value::Array(arr)) => {
169 let mut strings = Vec::new();
170 for item in arr {
171 match item {
172 Value::String(s) => strings.push(s.clone()),
173 _ => strings.push(item.to_string()),
174 }
175 }
176 Ok(Value::String(strings.join(sep)))
177 }
178 _ => Err("join() expects a string and an array".to_string()),
179 }
180}
181
182fn replace(args: FuncArgs) -> Result<Value, String> {
183 if args.len() != 3 {
184 return Err("replace() expects exactly 3 arguments".to_string());
185 }
186 match (&args[0], &args[1], &args[2]) {
187 (Value::String(s), Value::String(old), Value::String(new)) => {
188 Ok(Value::String(s.replace(old, new)))
189 }
190 _ => Err("replace() expects string arguments".to_string()),
191 }
192}
193
194fn jsonencode(args: FuncArgs) -> Result<Value, String> {
195 if args.len() != 1 {
196 return Err("jsonencode() expects exactly 1 argument".to_string());
197 }
198 match serde_json::to_string(&args[0]) {
199 Ok(s) => Ok(Value::String(s)),
200 Err(e) => Err(format!("jsonencode() failed: {}", e)),
201 }
202}
203
204fn jsondecode(args: FuncArgs) -> Result<Value, String> {
205 if let Some(Value::String(s)) = args.first() {
206 match serde_json::from_str::<Value>(s) {
207 Ok(v) => Ok(v),
208 Err(e) => Err(format!("jsondecode() failed: {}", e)),
209 }
210 } else {
211 Err("jsondecode() expects a string argument".to_string())
212 }
213}
214
215fn base64encode(args: FuncArgs) -> Result<Value, String> {
216 if let Some(Value::String(s)) = args.first() {
217 Ok(Value::String(general_purpose::STANDARD.encode(s)))
218 } else {
219 Err("base64encode() expects a string argument".to_string())
220 }
221}
222
223fn base64decode(args: FuncArgs) -> Result<Value, String> {
224 if let Some(Value::String(s)) = args.first() {
225 match general_purpose::STANDARD.decode(s) {
226 Ok(bytes) => match String::from_utf8(bytes) {
227 Ok(decoded) => Ok(Value::String(decoded)),
228 Err(_) => Err("base64decode() resulted in invalid UTF-8".to_string()),
229 },
230 Err(e) => Err(format!("base64decode() failed: {}", e)),
231 }
232 } else {
233 Err("base64decode() expects a string argument".to_string())
234 }
235}
236
237fn length(args: FuncArgs) -> Result<Value, String> {
238 if args.len() != 1 {
239 return Err("length() expects exactly 1 argument".to_string());
240 }
241 match &args[0] {
242 Value::String(s) => Ok(Value::Number(hcl::Number::from(s.len() as u64))),
243 Value::Array(a) => Ok(Value::Number(hcl::Number::from(a.len() as u64))),
244 Value::Object(o) => Ok(Value::Number(hcl::Number::from(o.len() as u64))),
245 _ => Err("length() expects a string, array, or object".to_string()),
246 }
247}
248
249fn keys(args: FuncArgs) -> Result<Value, String> {
250 if let Some(Value::Object(o)) = args.first() {
251 let keys_arr: Vec<Value> = o.keys().map(|k| Value::String(k.to_string())).collect();
252 Ok(Value::Array(keys_arr))
253 } else {
254 Err("keys() expects an object argument".to_string())
255 }
256}
257
258fn values(args: FuncArgs) -> Result<Value, String> {
259 if let Some(Value::Object(o)) = args.first() {
260 let vals_arr: Vec<Value> = o.values().cloned().collect();
261 Ok(Value::Array(vals_arr))
262 } else {
263 Err("values() expects an object argument".to_string())
264 }
265}
266
267fn contains(args: FuncArgs) -> Result<Value, String> {
268 if args.len() != 2 {
269 return Err("contains() expects exactly 2 arguments".to_string());
270 }
271 if let Value::Array(arr) = &args[0] {
272 let target = &args[1];
273 Ok(Value::Bool(arr.contains(target)))
274 } else {
275 Err("contains() expects an array as its first argument".to_string())
276 }
277}
278
279fn coalesce(args: FuncArgs) -> Result<Value, String> {
285 if args.is_empty() {
286 return Err("coalesce() expects at least 1 argument".to_string());
287 }
288 for arg in args.iter() {
289 if arg.is_null() {
290 continue;
291 }
292 if let Value::String(s) = arg {
293 if s.is_empty() {
294 continue;
295 }
296 }
297 return Ok(arg.clone());
298 }
299 Err("coalesce(): no non-null, non-empty-string arguments provided".to_string())
300}
301
302#[cfg(test)]
303mod tests {
304 use super::*;
305 use crate::hcl_eval::evaluate_raw_expr;
306 use hcl::eval::Context;
307 use serde_json::json;
308
309 fn eval_with_stdlib(expr: &str) -> serde_json::Value {
310 let mut ctx = Context::new();
311 register_stdlib(&mut ctx);
312 evaluate_raw_expr(expr, &ctx).expect("evaluation failed")
313 }
314
315 #[test]
316 fn test_string_functions() {
317 assert_eq!(eval_with_stdlib("upper(\"hello\")"), json!("HELLO"));
318 assert_eq!(eval_with_stdlib("lower(\"HELLO\")"), json!("hello"));
319 assert_eq!(
320 eval_with_stdlib("trim(\"?!hello?!\", \"?!\")"),
321 json!("hello")
322 );
323 assert_eq!(
324 eval_with_stdlib("trimspace(\" hello \\n\")"),
325 json!("hello")
326 );
327
328 let split_res = eval_with_stdlib("split(\",\", \"a,b,c\")");
329 assert_eq!(split_res, json!(["a", "b", "c"]));
330
331 assert_eq!(
332 eval_with_stdlib("join(\"-\", [\"a\", \"b\", \"c\"])"),
333 json!("a-b-c")
334 );
335 assert_eq!(
336 eval_with_stdlib("replace(\"hello world\", \"world\", \"there\")"),
337 json!("hello there")
338 );
339 }
340
341 #[test]
342 fn test_encoding_functions() {
343 let encode_res = eval_with_stdlib("jsonencode({\"a\" = 1})");
344 assert_eq!(encode_res, json!("{\"a\":1}"));
345
346 let decode_res = eval_with_stdlib("jsondecode(\"{\\\"a\\\": 1}\")");
347 assert_eq!(decode_res, json!({"a": 1}));
348
349 assert_eq!(
350 eval_with_stdlib("base64encode(\"hello\")"),
351 json!("aGVsbG8=")
352 );
353 assert_eq!(
354 eval_with_stdlib("base64decode(\"aGVsbG8=\")"),
355 json!("hello")
356 );
357 }
358
359 #[test]
360 fn test_collection_functions() {
361 assert_eq!(eval_with_stdlib("length(\"hello\")"), json!(5));
362 assert_eq!(eval_with_stdlib("length([1, 2, 3])"), json!(3));
363
364 assert_eq!(
365 eval_with_stdlib("contains([\"a\", \"b\"], \"b\")"),
366 json!(true)
367 );
368 assert_eq!(
369 eval_with_stdlib("contains([\"a\", \"b\"], \"c\")"),
370 json!(false)
371 );
372 }
373
374 #[test]
375 fn test_coalesce_skips_null_and_empty_string() {
376 assert_eq!(
378 eval_with_stdlib("coalesce(\"\", null, \"first\", \"second\")"),
379 json!("first")
380 );
381 }
382
383 #[test]
384 fn test_coalesce_returns_zero_number() {
385 assert_eq!(eval_with_stdlib("coalesce(0, 1)"), json!(0));
387 }
388
389 #[test]
390 fn test_coalesce_returns_false_bool() {
391 assert_eq!(eval_with_stdlib("coalesce(false, true)"), json!(false));
393 }
394
395 #[test]
396 fn test_coalesce_returns_empty_array() {
397 assert_eq!(eval_with_stdlib("coalesce([], [1, 2])"), json!([]));
399 }
400}