sync ttrpc: Add size checks where packets are created

server: check before send to client
client: check before send to server

Signed-off-by: Tim Zhang <[email protected]>

Author: Tim-Zhang

src/error.rs

@@ -14,7 +14,7 @@
 
 //! Error and Result of ttrpc and relevant functions, macros.
 
-use crate::proto::{Code, Status};
+use crate::proto::{Code, Response, Status};
 use std::result;
 use thiserror::Error;
 
@@ -48,6 +48,20 @@ pub enum Error {
     Others(String),
 }
 
+impl From<Error> for Response {
+    fn from(e: Error) -> Self {
+        let status = if let Error::RpcStatus(stat) = e {
+            stat
+        } else {
+            get_status(Code::UNKNOWN, e)
+        };
+
+        let mut res = Response::new();
+        res.set_status(status);
+        res
+    }
+}
+
 /// A specialized Result type for ttrpc.
 pub type Result<T> = result::Result<T, Error>;
 

src/sync/client.rs

@@ -17,12 +17,14 @@
 #[cfg(unix)]
 use std::os::unix::io::RawFd;
 
+use protobuf::Message;
 use std::collections::HashMap;
 use std::sync::mpsc;
 use std::sync::{Arc, Mutex};
 use std::thread;
 use std::time::Duration;
 
+use crate::common::check_oversize;
 use crate::error::{Error, Result};
 use crate::proto::{Code, Codec, MessageHeader, Request, Response, MESSAGE_TYPE_RESPONSE};
 use crate::sync::channel::{read_message, write_message};
@@ -150,7 +152,10 @@ impl Client {
         })
     }
     pub fn request(&self, req: Request) -> Result<Response> {
+        check_oversize(req.compute_size() as usize, false)?;
+
         let buf = req.encode().map_err(err_to_others_err!(e, ""))?;
+        // Notice: pure client problem can't be rpc error
 
         let (tx, rx) = mpsc::sync_channel(0);
 

src/sync/utils.rs

@@ -3,9 +3,9 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use crate::error::{get_status, Error, Result};
-use crate::proto::{Code, MessageHeader, Request, Response, MESSAGE_TYPE_RESPONSE};
-use protobuf::Message;
+use crate::common::check_oversize;
+use crate::error::{Error, Result};
+use crate::proto::{Codec, MessageHeader, Request, Response, MESSAGE_TYPE_RESPONSE};
 use std::collections::HashMap;
 
 /// Response message through a channel.
@@ -15,19 +15,20 @@ pub fn response_to_channel(
     res: Response,
     tx: std::sync::mpsc::Sender<(MessageHeader, Vec<u8>)>,
 ) -> Result<()> {
-    let mut buf = Vec::with_capacity(res.compute_size() as usize);
-    let mut s = protobuf::CodedOutputStream::vec(&mut buf);
-    res.write_to(&mut s).map_err(err_to_others_err!(e, ""))?;
-    s.flush().map_err(err_to_others_err!(e, ""))?;
+    let mut buf = res.encode().map_err(err_to_others_err!(e, ""))?;
 
-    drop(s);
+    if let Err(e) = check_oversize(buf.len(), true) {
+        let resp: Response = e.into();
+        buf = resp.encode().map_err(err_to_others_err!(e, ""))?;
+    };
 
     let mh = MessageHeader {
         length: buf.len() as u32,
         stream_id,
         type_: MESSAGE_TYPE_RESPONSE,
         flags: 0,
     };
+
     tx.send((mh, buf)).map_err(err_to_others_err!(e, ""))?;
 
     Ok(())
@@ -38,16 +39,7 @@ pub fn response_error_to_channel(
     e: Error,
     tx: std::sync::mpsc::Sender<(MessageHeader, Vec<u8>)>,
 ) -> Result<()> {
-    let status = if let Error::RpcStatus(stat) = e {
-        stat
-    } else {
-        get_status(Code::UNKNOWN, e)
-    };
-
-    let mut res = Response::new();
-    res.set_status(status);
-
-    response_to_channel(stream_id, res, tx)
+    response_to_channel(stream_id, e.into(), tx)
 }
 
 /// Handle request in sync mode.