sync ttrpc: Fix the bug caused by oversized packets

Fixes: #198

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

Author: Tim-Zhang

src/common.rs

@@ -6,7 +6,6 @@
 
 //! Common functions.
 
-use crate::error::{Error, Result};
 #[cfg(any(
     feature = "async",
     not(any(target_os = "linux", target_os = "android"))
@@ -16,6 +15,8 @@ use nix::fcntl::{fcntl, FcntlArg, OFlag};
 use nix::sys::socket::*;
 use std::os::unix::io::RawFd;
 
+use crate::error::{Error, Result};
+
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub(crate) enum Domain {
     Unix,

src/proto.rs

@@ -13,11 +13,11 @@ pub use compiled::ttrpc::*;
 use byteorder::{BigEndian, ByteOrder};
 use protobuf::{CodedInputStream, CodedOutputStream};
 
-#[cfg(feature = "async")]
 use crate::error::{get_rpc_status, Error, Result as TtResult};
 
 pub const MESSAGE_HEADER_LENGTH: usize = 10;
 pub const MESSAGE_LENGTH_MAX: usize = 4 << 20;
+pub const DEFAULT_PAGE_SIZE: usize = 4 << 10;
 
 pub const MESSAGE_TYPE_REQUEST: u8 = 0x1;
 pub const MESSAGE_TYPE_RESPONSE: u8 = 0x2;
@@ -27,6 +27,24 @@ pub const FLAG_REMOTE_CLOSED: u8 = 0x1;
 pub const FLAG_REMOTE_OPEN: u8 = 0x2;
 pub const FLAG_NO_DATA: u8 = 0x4;
 
+pub(crate) fn check_oversize(len: usize, return_rpc_error: bool) -> TtResult<()> {
+    if len > MESSAGE_LENGTH_MAX {
+        let msg = format!(
+            "message length {} exceed maximum message size of {}",
+            len, MESSAGE_LENGTH_MAX
+        );
+        let e = if return_rpc_error {
+            get_rpc_status(Code::INVALID_ARGUMENT, msg)
+        } else {
+            Error::Others(msg)
+        };
+
+        return Err(e);
+    }
+
+    Ok(())
+}
+
 /// Message header of ttrpc.
 #[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
 pub struct MessageHeader {

src/sync/channel.rs

@@ -12,10 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-
-use crate::error::{get_rpc_status, sock_error_msg, Error, Result};
+use crate::error::{sock_error_msg, Error, Result};
+use crate::proto::{check_oversize, MessageHeader, DEFAULT_PAGE_SIZE, MESSAGE_HEADER_LENGTH};
 use crate::sync::sys::PipeConnection;
-use crate::proto::{Code, MessageHeader, MESSAGE_HEADER_LENGTH, MESSAGE_LENGTH_MAX};
 
 fn read_count(conn: &PipeConnection, count: usize) -> Result<Vec<u8>> {
     let mut v: Vec<u8> = vec![0; count];
@@ -51,7 +50,7 @@ fn write_count(conn: &PipeConnection, buf: &[u8], count: usize) -> Result<usize>
     }
 
     loop {
-        match conn.write(&buf[len..]){
+        match conn.write(&buf[len..]) {
             Ok(l) => {
                 len += l;
                 if len == count {
@@ -67,6 +66,18 @@ fn write_count(conn: &PipeConnection, buf: &[u8], count: usize) -> Result<usize>
     Ok(len)
 }
 
+fn discard_count(conn: &PipeConnection, count: usize) -> Result<()> {
+    let mut need_discard = count;
+
+    while need_discard > 0 {
+        let once_discard = std::cmp::min(DEFAULT_PAGE_SIZE, need_discard);
+        read_count(conn, once_discard)?;
+        need_discard -= once_discard;
+    }
+
+    Ok(())
+}
+
 fn read_message_header(conn: &PipeConnection) -> Result<MessageHeader> {
     let buf = read_count(conn, MESSAGE_HEADER_LENGTH)?;
     let size = buf.len();
@@ -82,18 +93,14 @@ fn read_message_header(conn: &PipeConnection) -> Result<MessageHeader> {
     Ok(mh)
 }
 
-pub fn read_message(conn: &PipeConnection) -> Result<(MessageHeader, Vec<u8>)> {
+pub fn read_message(conn: &PipeConnection) -> Result<(MessageHeader, Result<Vec<u8>>)> {
     let mh = read_message_header(conn)?;
     trace!("Got Message header {:?}", mh);
 
-    if mh.length > MESSAGE_LENGTH_MAX as u32 {
-        return Err(get_rpc_status(
-            Code::INVALID_ARGUMENT,
-            format!(
-                "message length {} exceed maximum message size of {}",
-                mh.length, MESSAGE_LENGTH_MAX
-            ),
-        ));
+    let mh_len = mh.length as usize;
+    if let Err(e) = check_oversize(mh_len, true) {
+        discard_count(conn, mh_len)?;
+        return Ok((mh, Err(e)));
     }
 
     let buf = read_count(conn, mh.length as usize)?;
@@ -106,7 +113,7 @@ pub fn read_message(conn: &PipeConnection) -> Result<(MessageHeader, Vec<u8>)> {
     }
     trace!("Got Message body {:?}", buf);
 
-    Ok((mh, buf))
+    Ok((mh, Ok(buf)))
 }
 
 fn write_message_header(conn: &PipeConnection, mh: MessageHeader) -> Result<()> {

src/sync/client.rs

@@ -111,12 +111,9 @@ impl Client {
                     }
                 }
 
-                let mh;
-                let buf;
                 match read_message(&receiver_connection) {
-                    Ok((x, y)) => {
-                        mh = x;
-                        buf = y;
+                    Ok((mh, buf)) => {
+                        trans_resp(recver_map_orig.clone(), mh, buf);
                     }
                     Err(x) => match x {
                         Error::Socket(y) => {
@@ -138,8 +135,6 @@ impl Client {
                         }
                     },
                 };
-
-                trans_resp(recver_map_orig.clone(), mh, Ok(buf));
             }
 
             let _ = receiver_client

src/sync/server.rs

@@ -26,7 +26,7 @@ use std::sync::{Arc, Mutex};
 use std::thread;
 use std::thread::JoinHandle;
 
-use super::utils::response_to_channel;
+use super::utils::{response_error_to_channel, response_to_channel};
 use crate::context;
 use crate::error::{get_status, Error, Result};
 use crate::proto::{Code, MessageHeader, Request, Response, MESSAGE_TYPE_REQUEST};
@@ -43,8 +43,8 @@ const DEFAULT_WAIT_THREAD_COUNT_MAX: usize = 5;
 
 type MessageSender = Sender<(MessageHeader, Vec<u8>)>;
 type MessageReceiver = Receiver<(MessageHeader, Vec<u8>)>;
-type WorkloadSender = crossbeam::channel::Sender<(MessageHeader, Vec<u8>)>;
-type WorkloadReceiver = crossbeam::channel::Receiver<(MessageHeader, Vec<u8>)>;
+type WorkloadSender = crossbeam::channel::Sender<(MessageHeader, Result<Vec<u8>>)>;
+type WorkloadReceiver = crossbeam::channel::Receiver<(MessageHeader, Result<Vec<u8>>)>;
 
 /// A ttrpc Server (sync).
 pub struct Server {
@@ -134,20 +134,22 @@ fn start_method_handler_thread(
             let mh;
             let buf;
             match result {
-                Ok((x, y)) => {
+                Ok((x, Ok(y))) => {
                     mh = x;
                     buf = y;
                 }
+                Ok((mh, Err(e))) => {
+                    if let Err(x) = response_error_to_channel(mh.stream_id, e, res_tx.clone()) {
+                        debug!("response_error_to_channel get error {:?}", x);
+                        quit_connection(quit, control_tx);
+                        break;
+                    }
+                    continue;
+                }
                 Err(x) => match x {
                     crossbeam::channel::RecvError => {
                         trace!("workload_rx recv error");
-                        quit.store(true, Ordering::SeqCst);
-                        // the workload tx would be dropped and
-                        // the connection dealing main thread would
-                        // have exited.
-                        control_tx
-                            .send(())
-                            .unwrap_or_else(|err| trace!("Failed to send {:?}", err));
+                        quit_connection(quit, control_tx);
                         trace!("workload_rx recv error, send control_tx");
                         break;
                     }
@@ -165,13 +167,7 @@ fn start_method_handler_thread(
                 res.set_status(status);
                 if let Err(x) = response_to_channel(mh.stream_id, res, res_tx.clone()) {
                     debug!("response_to_channel get error {:?}", x);
-                    quit.store(true, Ordering::SeqCst);
-                    // the client connection would be closed and
-                    // the connection dealing main thread would have
-                    // exited.
-                    control_tx
-                        .send(())
-                        .unwrap_or_else(|err| trace!("Failed to send {:?}", err));
+                    quit_connection(quit, control_tx);
                     break;
                 }
                 continue;
@@ -187,13 +183,7 @@ fn start_method_handler_thread(
                 res.set_status(status);
                 if let Err(x) = response_to_channel(mh.stream_id, res, res_tx.clone()) {
                     info!("response_to_channel get error {:?}", x);
-                    quit.store(true, Ordering::SeqCst);
-                    // the client connection would be closed and
-                    // the connection dealing main thread would have
-                    // exited.
-                    control_tx
-                        .send(())
-                        .unwrap_or_else(|err| trace!("Failed to send {:?}", err));
+                    quit_connection(quit, control_tx);
                     break;
                 }
                 continue;
@@ -208,13 +198,7 @@ fn start_method_handler_thread(
             };
             if let Err(x) = method.handler(ctx, req) {
                 debug!("method handle {} get error {:?}", path, x);
-                quit.store(true, Ordering::SeqCst);
-                // the client connection would be closed and
-                // the connection dealing main thread would have
-                // exited.
-                control_tx
-                    .send(())
-                    .unwrap_or_else(|err| trace!("Failed to send {:?}", err));
+                quit_connection(quit, control_tx);
                 break;
             }
         }
@@ -595,3 +579,13 @@ impl AsRawFd for Server {
         self.listeners[0].as_raw_fd()
     }
 }
+
+fn quit_connection(quit: Arc<AtomicBool>, control_tx: SyncSender<()>) {
+    quit.store(true, Ordering::SeqCst);
+    // the client connection would be closed and
+    // the connection dealing main thread would
+    // have exited.
+    control_tx
+        .send(())
+        .unwrap_or_else(|err| debug!("Failed to send {:?}", err));
+}

src/sync/utils.rs

@@ -3,8 +3,8 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use crate::error::{Error, Result};
-use crate::proto::{MessageHeader, Request, Response, MESSAGE_TYPE_RESPONSE};
+use crate::error::{get_status, Error, Result};
+use crate::proto::{Code, MessageHeader, Request, Response, MESSAGE_TYPE_RESPONSE};
 use protobuf::Message;
 use std::collections::HashMap;
 
@@ -33,6 +33,23 @@ pub fn response_to_channel(
     Ok(())
 }
 
+pub fn response_error_to_channel(
+    stream_id: u32,
+    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)
+}
+
 /// Handle request in sync mode.
 #[macro_export]
 macro_rules! request_handler {