fix:escape keywords in path resolver (#211)

Co-authored-by: $wangjie.wjdew <[email protected]>
cloudwego · Nov 10, 2023 · 2c0ecc7 · 2c0ecc7
1 parent 9606336
commit 2c0ecc7
Show file tree

Hide file tree

Showing 5 changed files with 202 additions and 2 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/pilota-build/Cargo.toml b/pilota-build/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pilota-build"
-version = "0.9.5"
+version = "0.9.6"
 edition = "2021"
 description = "Compile thrift and protobuf idl into rust code at compile-time."
 documentation = "https://docs.rs/pilota-build"

diff --git a/pilota-build/src/middle/resolver.rs b/pilota-build/src/middle/resolver.rs
@@ -23,6 +23,9 @@ pub trait PathResolver: Sync + Send {
                 NodeKind::Variant(_) => cx.rust_name(def_id),
                 _ => panic!(),
             };
+
+            //  keyword escape
+            let name = name.to_string().into();
             segs.push(name);
         }
 

diff --git a/pilota-build/test_data/thrift/path_keyword.rs b/pilota-build/test_data/thrift/path_keyword.rs
@@ -0,0 +1,192 @@
+pub mod path_keyword {
+    #![allow(warnings, clippy::all)]
+
+    pub mod r#enum {
+        #[derive(PartialOrd, Hash, Eq, Ord, Debug, Default, Clone, PartialEq)]
+        pub struct A {
+            pub bytes: ::pilota::Bytes,
+
+            pub vec: ::std::vec::Vec<u8>,
+        }
+        impl ::pilota::thrift::Message for A {
+            fn encode<T: ::pilota::thrift::TOutputProtocol>(
+                &self,
+                protocol: &mut T,
+            ) -> ::std::result::Result<(), ::pilota::thrift::EncodeError> {
+                #[allow(unused_imports)]
+                use ::pilota::thrift::TOutputProtocolExt;
+                let struct_ident = ::pilota::thrift::TStructIdentifier { name: "A" };
+
+                protocol.write_struct_begin(&struct_ident)?;
+                protocol.write_bytes_field(1, (&self.bytes).clone())?;
+                protocol.write_bytes_vec_field(2, &self.vec)?;
+                protocol.write_field_stop()?;
+                protocol.write_struct_end()?;
+                Ok(())
+            }
+
+            fn decode<T: ::pilota::thrift::TInputProtocol>(
+                protocol: &mut T,
+            ) -> ::std::result::Result<Self, ::pilota::thrift::DecodeError> {
+                #[allow(unused_imports)]
+                use ::pilota::{thrift::TLengthProtocolExt, Buf};
+
+                let mut bytes = None;
+                let mut vec = None;
+
+                let mut __pilota_decoding_field_id = None;
+
+                protocol.read_struct_begin()?;
+                if let Err(err) = (|| {
+                    loop {
+                        let field_ident = protocol.read_field_begin()?;
+                        if field_ident.field_type == ::pilota::thrift::TType::Stop {
+                            protocol.field_stop_len();
+                            break;
+                        } else {
+                            protocol.field_begin_len(field_ident.field_type, field_ident.id);
+                        }
+                        __pilota_decoding_field_id = field_ident.id;
+                        match field_ident.id {
+                            Some(1)
+                                if field_ident.field_type == ::pilota::thrift::TType::Binary =>
+                            {
+                                bytes = Some(protocol.read_bytes()?);
+                            }
+                            Some(2)
+                                if field_ident.field_type == ::pilota::thrift::TType::Binary =>
+                            {
+                                vec = Some(protocol.read_bytes_vec()?);
+                            }
+                            _ => {
+                                protocol.skip(field_ident.field_type)?;
+                            }
+                        }
+
+                        protocol.read_field_end()?;
+                        protocol.field_end_len();
+                    }
+                    Ok::<_, ::pilota::thrift::DecodeError>(())
+                })() {
+                    if let Some(field_id) = __pilota_decoding_field_id {
+                        return Err(::pilota::thrift::DecodeError::new(
+                            ::pilota::thrift::DecodeErrorKind::WithContext(::std::boxed::Box::new(
+                                err,
+                            )),
+                            format!("decode struct `A` field(#{}) failed", field_id),
+                        ));
+                    } else {
+                        return Err(err);
+                    }
+                };
+                protocol.read_struct_end()?;
+
+                let Some(bytes) = bytes else {
+                    return Err(::pilota::thrift::DecodeError::new(
+                        ::pilota::thrift::DecodeErrorKind::InvalidData,
+                        "field bytes is required".to_string(),
+                    ));
+                };
+                let Some(vec) = vec else {
+                    return Err(::pilota::thrift::DecodeError::new(
+                        ::pilota::thrift::DecodeErrorKind::InvalidData,
+                        "field vec is required".to_string(),
+                    ));
+                };
+
+                let data = Self { bytes, vec };
+                Ok(data)
+            }
+
+            fn decode_async<'a, T: ::pilota::thrift::TAsyncInputProtocol>(
+                protocol: &'a mut T,
+            ) -> ::std::pin::Pin<
+                ::std::boxed::Box<
+                    dyn ::std::future::Future<
+                            Output = ::std::result::Result<Self, ::pilota::thrift::DecodeError>,
+                        > + Send
+                        + 'a,
+                >,
+            > {
+                ::std::boxed::Box::pin(async move {
+                    let mut bytes = None;
+                    let mut vec = None;
+
+                    let mut __pilota_decoding_field_id = None;
+
+                    protocol.read_struct_begin().await?;
+                    if let Err(err) = async {
+                        loop {
+                            let field_ident = protocol.read_field_begin().await?;
+                            if field_ident.field_type == ::pilota::thrift::TType::Stop {
+                                break;
+                            } else {
+                            }
+                            __pilota_decoding_field_id = field_ident.id;
+                            match field_ident.id {
+                                Some(1)
+                                    if field_ident.field_type
+                                        == ::pilota::thrift::TType::Binary =>
+                                {
+                                    bytes = Some(protocol.read_bytes().await?);
+                                }
+                                Some(2)
+                                    if field_ident.field_type
+                                        == ::pilota::thrift::TType::Binary =>
+                                {
+                                    vec = Some(protocol.read_bytes_vec().await?);
+                                }
+                                _ => {
+                                    protocol.skip(field_ident.field_type).await?;
+                                }
+                            }
+
+                            protocol.read_field_end().await?;
+                        }
+                        Ok::<_, ::pilota::thrift::DecodeError>(())
+                    }
+                    .await
+                    {
+                        if let Some(field_id) = __pilota_decoding_field_id {
+                            return Err(::pilota::thrift::DecodeError::new(
+                                ::pilota::thrift::DecodeErrorKind::WithContext(
+                                    ::std::boxed::Box::new(err),
+                                ),
+                                format!("decode struct `A` field(#{}) failed", field_id),
+                            ));
+                        } else {
+                            return Err(err);
+                        }
+                    };
+                    protocol.read_struct_end().await?;
+
+                    let Some(bytes) = bytes else {
+                        return Err(::pilota::thrift::DecodeError::new(
+                            ::pilota::thrift::DecodeErrorKind::InvalidData,
+                            "field bytes is required".to_string(),
+                        ));
+                    };
+                    let Some(vec) = vec else {
+                        return Err(::pilota::thrift::DecodeError::new(
+                            ::pilota::thrift::DecodeErrorKind::InvalidData,
+                            "field vec is required".to_string(),
+                        ));
+                    };
+
+                    let data = Self { bytes, vec };
+                    Ok(data)
+                })
+            }
+
+            fn size<T: ::pilota::thrift::TLengthProtocol>(&self, protocol: &mut T) -> usize {
+                #[allow(unused_imports)]
+                use ::pilota::thrift::TLengthProtocolExt;
+                protocol.struct_begin_len(&::pilota::thrift::TStructIdentifier { name: "A" })
+                    + protocol.bytes_field_len(Some(1), &self.bytes)
+                    + protocol.bytes_vec_field_len(Some(2), &self.vec)
+                    + protocol.field_stop_len()
+                    + protocol.struct_end_len()
+            }
+        }
+    }
+}
diff --git a/pilota-build/test_data/thrift/path_keyword.thrift b/pilota-build/test_data/thrift/path_keyword.thrift
@@ -0,0 +1,5 @@
+namespace rs enum;
+struct A {
+    1: required binary bytes,
+    2: required binary vec(pilota.rust_type="vec"),
+}