vrl/stdlib/
encode_zlib.rs

1use 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    // TODO consider removal options
27    #[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    // We can safely ignore the error here because the value being read from, `Bytes`, never fails a `read()` call and the value being written to, a `Vec`, never fails a `write()` call
39    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}