vrl/stdlib/
encode_zlib.rs1use std::io::Read;
2
3use flate2::read::ZlibEncoder;
4use nom::AsBytes;
5
6use crate::compiler::prelude::*;
7use std::sync::LazyLock;
8
9static DEFAULT_COMPRESSION_LEVEL: LazyLock<Value> = LazyLock::new(|| Value::Integer(6));
10
11const MAX_COMPRESSION_LEVEL: u32 = 10;
12
13static PARAMETERS: LazyLock<Vec<Parameter>> = LazyLock::new(|| {
14 vec![
15 Parameter::required("value", kind::BYTES, "The string to encode."),
16 Parameter::optional(
17 "compression_level",
18 kind::INTEGER,
19 "The default compression level.",
20 )
21 .default(&DEFAULT_COMPRESSION_LEVEL),
22 ]
23});
24
25fn encode_zlib(value: Value, compression_level: Value) -> Resolved {
26 #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
28 let level = compression_level.try_integer()? as u32;
29 let compression_level = if level > MAX_COMPRESSION_LEVEL {
30 return Err(format!("compression level must be <= {MAX_COMPRESSION_LEVEL}").into());
31 } else {
32 flate2::Compression::new(level)
33 };
34
35 let value = value.try_bytes()?;
36 let mut buf = Vec::new();
37
38 ZlibEncoder::new(value.as_bytes(), compression_level)
40 .read_to_end(&mut buf)
41 .expect("zlib compression failed, please report");
42
43 Ok(Value::Bytes(buf.into()))
44}
45
46#[derive(Clone, Copy, Debug)]
47pub struct EncodeZlib;
48
49impl Function for EncodeZlib {
50 fn identifier(&self) -> &'static str {
51 "encode_zlib"
52 }
53
54 fn usage(&self) -> &'static str {
55 "Encodes the `value` to [Zlib](https://www.zlib.net)."
56 }
57
58 fn category(&self) -> &'static str {
59 Category::Codec.as_ref()
60 }
61
62 fn return_kind(&self) -> u16 {
63 kind::BYTES
64 }
65
66 fn examples(&self) -> &'static [Example] {
67 &[example! {
68 title: "Encode to Zlib",
69 source: r#"encode_base64(encode_zlib("please encode me"))"#,
70 result: Ok("eJwryElNLE5VSM1Lzk9JVchNBQA0RQX7"),
71 }]
72 }
73
74 fn compile(
75 &self,
76 _state: &state::TypeState,
77 _ctx: &mut FunctionCompileContext,
78 arguments: ArgumentList,
79 ) -> Compiled {
80 let value = arguments.required("value");
81 let compression_level = arguments.optional("compression_level");
82
83 Ok(EncodeZlibFn {
84 value,
85 compression_level,
86 }
87 .as_expr())
88 }
89
90 fn parameters(&self) -> &'static [Parameter] {
91 PARAMETERS.as_slice()
92 }
93}
94
95#[derive(Clone, Debug)]
96struct EncodeZlibFn {
97 value: Box<dyn Expression>,
98 compression_level: Option<Box<dyn Expression>>,
99}
100
101impl FunctionExpression for EncodeZlibFn {
102 fn resolve(&self, ctx: &mut Context) -> Resolved {
103 let value = self.value.resolve(ctx)?;
104
105 let compression_level = self
106 .compression_level
107 .map_resolve_with_default(ctx, || DEFAULT_COMPRESSION_LEVEL.clone())?;
108
109 encode_zlib(value, compression_level)
110 }
111
112 fn type_def(&self, state: &state::TypeState) -> TypeDef {
113 let is_compression_level_valid_constant = if let Some(level) = &self.compression_level {
114 match level.resolve_constant(state) {
115 Some(Value::Integer(level)) => level <= i64::from(MAX_COMPRESSION_LEVEL),
116 _ => false,
117 }
118 } else {
119 true
120 };
121
122 TypeDef::bytes().maybe_fallible(!is_compression_level_valid_constant)
123 }
124}
125
126#[cfg(test)]
127mod test {
128 use crate::value;
129
130 use super::*;
131
132 fn encode(text: &str, level: flate2::Compression) -> Vec<u8> {
133 let mut encoder = ZlibEncoder::new(text.as_bytes(), level);
134 let mut output = vec![];
135 encoder.read_to_end(&mut output).unwrap();
136 output
137 }
138
139 test_function![
140 encode_zlib => EncodeZlib;
141
142 with_defaults {
143 args: func_args![value: value!("you_have_successfully_decoded_me.congratulations.you_are_breathtaking.")],
144 want: Ok(value!(encode("you_have_successfully_decoded_me.congratulations.you_are_breathtaking.", flate2::Compression::default()).as_bytes())),
145 tdef: TypeDef::bytes().infallible(),
146 }
147
148 with_custom_compression_level {
149 args: func_args![value: value!("you_have_successfully_decoded_me.congratulations.you_are_breathtaking."), compression_level: 9],
150 want: Ok(value!(encode("you_have_successfully_decoded_me.congratulations.you_are_breathtaking.", flate2::Compression::new(9)).as_bytes())),
151 tdef: TypeDef::bytes().infallible(),
152 }
153
154 invalid_constant_compression {
155 args: func_args![value: value!("d"), compression_level: 11],
156 want: Err("compression level must be <= 10"),
157 tdef: TypeDef::bytes().fallible(),
158 }
159 ];
160}