vrl/stdlib/
encode_proto.rs

1use super::util::example_path_or_basename;
2use crate::compiler::prelude::*;
3use crate::protobuf::descriptor::get_message_descriptor;
4use crate::protobuf::encode::encode_proto;
5use prost_reflect::MessageDescriptor;
6use std::ffi::OsString;
7use std::path::Path;
8use std::path::PathBuf;
9use std::sync::LazyLock;
10
11#[derive(Clone, Copy, Debug)]
12pub struct EncodeProto;
13
14// This needs to be static because parse_proto needs to read a file
15// and the file path needs to be a literal.
16static EXAMPLE_ENCODE_PROTO_EXPR: LazyLock<&str> = LazyLock::new(|| {
17    let path = example_path_or_basename("protobuf/test_protobuf/v1/test_protobuf.desc");
18
19    Box::leak(
20        format!(
21            r#"encode_base64(encode_proto!({{ "name": "someone", "phones": [{{"number": "123456"}}]}}, "{path}", "test_protobuf.v1.Person"))"#
22        )
23        .into_boxed_str(),
24    )
25});
26
27static EXAMPLES: LazyLock<Vec<Example>> = LazyLock::new(|| {
28    vec![example! {
29        title: "Encode to proto",
30        source: &EXAMPLE_ENCODE_PROTO_EXPR,
31        result: Ok("Cgdzb21lb25lIggKBjEyMzQ1Ng=="),
32    }]
33});
34
35impl Function for EncodeProto {
36    fn identifier(&self) -> &'static str {
37        "encode_proto"
38    }
39
40    fn summary(&self) -> &'static str {
41        "Encodes a value into a protobuf"
42    }
43
44    fn usage(&self) -> &'static str {
45        "Encodes the `value` into a protocol buffer payload."
46    }
47
48    fn category(&self) -> &'static str {
49        Category::Codec.as_ref()
50    }
51
52    fn internal_failure_reasons(&self) -> &'static [&'static str] {
53        &[
54            "`desc_file` file does not exist.",
55            "`message_type` message type does not exist in the descriptor file.",
56        ]
57    }
58
59    fn return_kind(&self) -> u16 {
60        kind::BYTES
61    }
62
63    fn parameters(&self) -> &'static [Parameter] {
64        const PARAMETERS: &[Parameter] = &[
65            Parameter::required(
66                "value",
67                kind::ANY,
68                "The object to convert to a protocol buffer payload.",
69            ),
70            Parameter::required(
71                "desc_file",
72                kind::BYTES,
73                "The path to the protobuf descriptor set file. Must be a literal string.
74
75This file is the output of protoc -o <path> ...",
76            ),
77            Parameter::required(
78                "message_type",
79                kind::BYTES,
80                "The name of the message type to use for serializing.
81
82Must be a literal string.",
83            ),
84        ];
85        PARAMETERS
86    }
87
88    fn examples(&self) -> &'static [Example] {
89        EXAMPLES.as_slice()
90    }
91
92    fn compile(
93        &self,
94        state: &state::TypeState,
95        _ctx: &mut FunctionCompileContext,
96        arguments: ArgumentList,
97    ) -> Compiled {
98        let value = arguments.required("value");
99        let desc_file = arguments.required_literal("desc_file", state)?;
100        let desc_file_str = desc_file
101            .try_bytes_utf8_lossy()
102            .expect("descriptor file must be a string");
103        let message_type = arguments.required_literal("message_type", state)?;
104        let message_type_str = message_type
105            .try_bytes_utf8_lossy()
106            .expect("message_type must be a string");
107        let os_string: OsString = desc_file_str.into_owned().into();
108        let path_buf = PathBuf::from(os_string);
109        let path = Path::new(&path_buf);
110        let descriptor =
111            get_message_descriptor(path, &message_type_str).expect("message type not found");
112
113        Ok(EncodeProtoFn { descriptor, value }.as_expr())
114    }
115}
116
117#[derive(Debug, Clone)]
118struct EncodeProtoFn {
119    descriptor: MessageDescriptor,
120    value: Box<dyn Expression>,
121}
122
123impl FunctionExpression for EncodeProtoFn {
124    fn resolve(&self, ctx: &mut Context) -> Resolved {
125        let value = self.value.resolve(ctx)?;
126        encode_proto(&self.descriptor, value)
127    }
128
129    fn type_def(&self, _: &state::TypeState) -> TypeDef {
130        TypeDef::bytes().fallible()
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use crate::value;
138    use std::{env, fs};
139
140    fn test_data_dir() -> PathBuf {
141        PathBuf::from(env::var_os("CARGO_MANIFEST_DIR").unwrap()).join("tests/data/protobuf")
142    }
143
144    fn read_pb_file(protobuf_bin_message_path: &str) -> String {
145        fs::read_to_string(test_data_dir().join(protobuf_bin_message_path)).unwrap()
146    }
147
148    test_function![
149        encode_proto => EncodeProto;
150
151        encodes {
152            args: func_args![ value: value!({ name: "Someone", phones: [{number: "123-456"}] }),
153                desc_file: test_data_dir().join("test_protobuf/v1/test_protobuf.desc").to_str().unwrap().to_owned(),
154                message_type: "test_protobuf.v1.Person"],
155            want: Ok(value!(read_pb_file("test_protobuf/v1/input/person_someone.pb"))),
156            tdef: TypeDef::bytes().fallible(),
157        }
158
159        encodes_proto3 {
160            args: func_args![
161                value: value!({ name: "Someone", phones: [{number: "123-456", type: "PHONE_TYPE_MOBILE"}] }),
162                desc_file: test_data_dir().join("test_protobuf3/v1/test_protobuf3.desc").to_str().unwrap().to_owned(),
163                message_type: "test_protobuf3.v1.Person"],
164            want: Ok(value!(read_pb_file("test_protobuf3/v1/input/person_someone.pb"))),
165            tdef: TypeDef::bytes().fallible(),
166        }
167    ];
168}