1 files changed, 36 insertions, 19 deletions

diff --git a/src/codec.rs b/src/codec.rs
index e0214b7..35dfb2e 100644
--- a/src/codec.rs
+++ b/src/codec.rs
@@ -5,7 +5,11 @@ use crate::call::MessageReader;
 use crate::error::Result;
 
 pub type DeserializeFn<T> = fn(MessageReader) -> Result<T>;
-pub type SerializeFn<T> = fn(&T, &mut GrpcSlice);
+pub type SerializeFn<T> = fn(&T, &mut GrpcSlice) -> Result<()>;
+
+/// According to https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md, grpc uses
+/// a four bytes to describe the length of a message, so it should not exceed u32::MAX.
+pub const MAX_MESSAGE_SIZE: usize = std::u32::MAX as usize;
 
 /// Defines how to serialize and deserialize between the specialized type and byte slice.
 pub struct Marshaller<T> {
@@ -29,18 +33,25 @@ pub struct Marshaller<T> {
 pub mod pb_codec {
     use protobuf::{CodedInputStream, CodedOutputStream, Message};
 
-    use super::MessageReader;
+    use super::{MessageReader, MAX_MESSAGE_SIZE};
     use crate::buf::GrpcSlice;
-    use crate::error::Result;
+    use crate::error::{Error, Result};
 
     #[inline]
-    pub fn ser<T: Message>(t: &T, buf: &mut GrpcSlice) {
-        let cap = t.compute_size();
-        unsafe {
-            let bytes = buf.realloc(cap as usize);
-            let raw_bytes = &mut *(bytes as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]);
-            let mut s = CodedOutputStream::bytes(raw_bytes);
-            t.write_to_with_cached_sizes(&mut s).unwrap();
+    pub fn ser<T: Message>(t: &T, buf: &mut GrpcSlice) -> Result<()> {
+        let cap = t.compute_size() as usize;
+        // FIXME: This is not a practical fix until stepancheg/rust-protobuf#530 is fixed.
+        if cap <= MAX_MESSAGE_SIZE {
+            unsafe {
+                let bytes = buf.realloc(cap);
+                let raw_bytes = &mut *(bytes as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]);
+                let mut s = CodedOutputStream::bytes(raw_bytes);
+                t.write_to_with_cached_sizes(&mut s).map_err(Into::into)
+            }
+        } else {
+            Err(Error::Codec(
+                format!("message is too large: {} > {}", cap, MAX_MESSAGE_SIZE).into(),
+            ))
         }
     }
 
@@ -57,19 +68,25 @@ pub mod pb_codec {
 pub mod pr_codec {
     use prost::Message;
 
-    use super::MessageReader;
+    use super::{MessageReader, MAX_MESSAGE_SIZE};
     use crate::buf::GrpcSlice;
-    use crate::error::Result;
+    use crate::error::{Error, Result};
 
     #[inline]
-    pub fn ser<M: Message>(msg: &M, buf: &mut GrpcSlice) {
+    pub fn ser<M: Message>(msg: &M, buf: &mut GrpcSlice) -> Result<()> {
         let size = msg.encoded_len();
-        unsafe {
-            let bytes = buf.realloc(size);
-            let mut b = &mut *(bytes as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]);
-            msg.encode(&mut b)
-                .expect("Writing message to buffer failed");
-            debug_assert!(b.is_empty());
+        if size <= MAX_MESSAGE_SIZE {
+            unsafe {
+                let bytes = buf.realloc(size);
+                let mut b = &mut *(bytes as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]);
+                msg.encode(&mut b)?;
+                debug_assert!(b.is_empty());
+            }
+            Ok(())
+        } else {
+            Err(Error::Codec(
+                format!("message is too large: {} > {}", size, MAX_MESSAGE_SIZE).into(),
+            ))
         }
     }