vrl/stdlib/
encode_gzip.rs

1use std::io::Read;
2
3use flate2::read::GzEncoder;
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_gzip(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    // 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
38    GzEncoder::new(value.as_bytes(), compression_level)
39        .read_to_end(&mut buf)
40        .expect("gzip compression failed, please report");
41
42    Ok(Value::Bytes(buf.into()))
43}
44
45#[derive(Clone, Copy, Debug)]
46pub struct EncodeGzip;
47
48impl Function for EncodeGzip {
49    fn identifier(&self) -> &'static str {
50        "encode_gzip"
51    }
52
53    fn usage(&self) -> &'static str {
54        "Encodes the `value` to [Gzip](https://www.gzip.org/)."
55    }
56
57    fn category(&self) -> &'static str {
58        Category::Codec.as_ref()
59    }
60
61    fn return_kind(&self) -> u16 {
62        kind::BYTES
63    }
64
65    fn examples(&self) -> &'static [Example] {
66        &[example! {
67            title: "Encode to Gzip",
68            source: r#"encode_base64(encode_gzip("please encode me"))"#,
69            result: Ok("H4sIAAAAAAAA/yvISU0sTlVIzUvOT0lVyE0FAI4R4vcQAAAA"),
70        }]
71    }
72
73    fn compile(
74        &self,
75        _state: &state::TypeState,
76        _ctx: &mut FunctionCompileContext,
77        arguments: ArgumentList,
78    ) -> Compiled {
79        let value = arguments.required("value");
80        let compression_level = arguments.optional("compression_level");
81
82        Ok(EncodeGzipFn {
83            value,
84            compression_level,
85        }
86        .as_expr())
87    }
88
89    fn parameters(&self) -> &'static [Parameter] {
90        PARAMETERS.as_slice()
91    }
92}
93
94#[derive(Clone, Debug)]
95struct EncodeGzipFn {
96    value: Box<dyn Expression>,
97    compression_level: Option<Box<dyn Expression>>,
98}
99
100impl FunctionExpression for EncodeGzipFn {
101    fn resolve(&self, ctx: &mut Context) -> Resolved {
102        let value = self.value.resolve(ctx)?;
103
104        let compression_level = self
105            .compression_level
106            .map_resolve_with_default(ctx, || DEFAULT_COMPRESSION_LEVEL.clone())?;
107
108        encode_gzip(value, compression_level)
109    }
110
111    fn type_def(&self, state: &state::TypeState) -> TypeDef {
112        let is_compression_level_valid_constant = if let Some(level) = &self.compression_level {
113            match level.resolve_constant(state) {
114                Some(Value::Integer(level)) => level <= i64::from(MAX_COMPRESSION_LEVEL),
115                _ => false,
116            }
117        } else {
118            true
119        };
120
121        TypeDef::bytes().maybe_fallible(!is_compression_level_valid_constant)
122    }
123}
124
125#[cfg(test)]
126mod test {
127    use crate::value;
128
129    use super::*;
130
131    fn encode(text: &str, level: flate2::Compression) -> Vec<u8> {
132        let mut encoder = GzEncoder::new(text.as_bytes(), level);
133        let mut output = vec![];
134        encoder.read_to_end(&mut output).unwrap();
135        output
136    }
137
138    test_function![
139        encode_gzip => EncodeGzip;
140
141        with_defaults {
142            args: func_args![value: value!("you_have_successfully_decoded_me.congratulations.you_are_breathtaking.")],
143            want: Ok(value!(encode("you_have_successfully_decoded_me.congratulations.you_are_breathtaking.", flate2::Compression::default()).as_bytes())),
144            tdef: TypeDef::bytes().infallible(),
145        }
146
147        with_custom_compression_level {
148            args: func_args![value: value!("you_have_successfully_decoded_me.congratulations.you_are_breathtaking."), compression_level: 9],
149            want: Ok(value!(encode("you_have_successfully_decoded_me.congratulations.you_are_breathtaking.", flate2::Compression::new(9)).as_bytes())),
150            tdef: TypeDef::bytes().infallible(),
151        }
152
153        invalid_constant_compression {
154            args: func_args![value: value!("test"), compression_level: 11],
155            want: Err("compression level must be <= 10"),
156            tdef: TypeDef::bytes().fallible(),
157        }
158    ];
159}