codec: check serialization result length

benchmark/src/bench.rs

@@ -100,8 +100,9 @@ impl Generic {
 
 #[inline]
 #[allow(clippy::ptr_arg)]
-pub fn bin_ser(t: &Vec<u8>, buf: &mut Vec<u8>) {
-    buf.extend_from_slice(t)
+pub fn bin_ser(t: &Vec<u8>, buf: &mut Vec<u8>) -> grpc::Result<()> {
+    buf.extend_from_slice(t);
+    Ok(())
 }
 
 #[inline]

src/call/client.rs

@@ -103,7 +103,7 @@ impl Call {
     ) -> Result<ClientUnaryReceiver<Resp>> {
         let call = channel.create_call(method, &opt)?;
         let mut payload = vec![];
-        (method.req_ser())(req, &mut payload);
+        (method.req_ser())(req, &mut payload)?;
         let cq_f = check_run(BatchType::CheckRead, |ctx, tag| unsafe {
             grpc_sys::grpcwrap_call_start_unary(
                 call.call,
@@ -157,7 +157,7 @@ impl Call {
     ) -> Result<ClientSStreamReceiver<Resp>> {
         let call = channel.create_call(method, &opt)?;
         let mut payload = vec![];
-        (method.req_ser())(req, &mut payload);
+        (method.req_ser())(req, &mut payload)?;
         let cq_f = check_run(BatchType::Finish, |ctx, tag| unsafe {
             grpc_sys::grpcwrap_call_start_server_streaming(
                 call.call,

src/call/mod.rs

@@ -657,7 +657,7 @@ impl SinkBase {
         }
 
         self.buf.clear();
-        ser(t, &mut self.buf);
+        ser(t, &mut self.buf)?;
         if flags.get_buffer_hint() && self.send_metadata {
             // temporary fix: buffer hint with send meta will not send out any metadata.
             flags = flags.buffer_hint(false);

src/call/server.rs

@@ -330,11 +330,20 @@ macro_rules! impl_unary_sink {
             }
 
             fn complete(mut self, status: RpcStatus, t: Option<T>) -> $rt {
-                let data = t.as_ref().map(|t| {
-                    let mut buf = vec![];
-                    (self.ser)(t, &mut buf);
-                    buf
-                });
+                let data = match t {
+                    Some(t) => {
+                        let mut buf = vec![];
+                        match (self.ser)(&t, &mut buf) {
+                            Ok(()) => Some(buf),
+                            Err(e) => return $rt {
+                                call: self.call.take().unwrap(),
+                                cq_f: None,
+                                err: Some(e),
+                            }
+                        }
+                    },
+                    None => None,
+                };
 
                 let write_flags = self.write_flags;
                 let res = self.call.as_mut().unwrap().call(|c| {

src/codec.rs

@@ -4,7 +4,7 @@ use crate::call::MessageReader;
 use crate::error::Result;
 
 pub type DeserializeFn<T> = fn(MessageReader) -> Result<T>;
-pub type SerializeFn<T> = fn(&T, &mut Vec<u8>);
+pub type SerializeFn<T> = fn(&T, &mut Vec<u8>) -> Result<()>;
 
 /// Defines how to serialize and deserialize between the specialized type and byte slice.
 pub struct Marshaller<T> {
@@ -29,11 +29,18 @@ pub mod pb_codec {
     use protobuf::{CodedInputStream, Message};
 
     use super::MessageReader;
-    use crate::error::Result;
+    use crate::error::{Error, Result};
 
     #[inline]
-    pub fn ser<T: Message>(t: &T, buf: &mut Vec<u8>) {
-        t.write_to_vec(buf).unwrap()
+    pub fn ser<T: Message>(t: &T, buf: &mut Vec<u8>) -> Result<()> {
+        t.write_to_vec(buf)?;
+        if buf.len() <= u32::MAX as usize {
+            Ok(())
+        } else {
+            Err(Error::Codec(
+                format!("message is too large: {} > u32::MAX", buf.len()).into(),
+            ))
+        }
     }
 
     #[inline]
@@ -47,15 +54,21 @@ pub mod pb_codec {
 
 #[cfg(feature = "prost-codec")]
 pub mod pr_codec {
-    use bytes::buf::BufMut;
     use prost::Message;
 
     use super::MessageReader;
-    use crate::error::Result;
+    use crate::error::{Error, Result};
 
     #[inline]
-    pub fn ser<M: Message, B: BufMut>(msg: &M, buf: &mut B) {
-        msg.encode(buf).expect("Writing message to buffer failed");
+    pub fn ser<M: Message>(msg: &M, buf: &mut Vec<u8>) -> Result<()> {
+        msg.encode(buf)?;
+        if buf.len() <= u32::MAX as usize {
+            Ok(())
+        } else {
+            Err(Error::Codec(
+                format!("message is too large: {} > u32::MAX", buf.len()).into(),
+            ))
+        }
     }
 
     #[inline]

src/error.rs

@@ -5,8 +5,6 @@ use std::{error, fmt, result};
 use crate::call::RpcStatus;
 use crate::grpc_sys::grpc_call_error;
 
-#[cfg(feature = "prost-codec")]
-use prost::DecodeError;
 #[cfg(feature = "protobuf-codec")]
 use protobuf::ProtobufError;
 
@@ -64,8 +62,15 @@ impl From<ProtobufError> for Error {
 }
 
 #[cfg(feature = "prost-codec")]
-impl From<DecodeError> for Error {
-    fn from(e: DecodeError) -> Error {
+impl From<prost::DecodeError> for Error {
+    fn from(e: prost::DecodeError) -> Error {
+        Error::Codec(Box::new(e))
+    }
+}
+
+#[cfg(feature = "prost-codec")]
+impl From<prost::EncodeError> for Error {
+    fn from(e: prost::EncodeError) -> Error {
         Error::Codec(Box::new(e))
     }
 }