refactor: add tls-client trait

Author: th4s

crates/mpc-tls/src/leader.rs

@@ -378,15 +378,29 @@ impl MpcTlsLeader {
         Ok(())
     }
 
-    /// Defers decryption of any incoming messages.
+    /// Enables or disables the decryption of any incoming messages.
+    ///
+    /// # Arguments
+    ///
+    /// * `enable` - Whether to enable or disable decryption.
     #[instrument(level = "debug", skip_all, err)]
-    pub async fn defer_decryption(&mut self) -> Result<(), MpcTlsError> {
-        self.is_decrypting = false;
-        self.notifier.clear();
+    pub fn enable_decryption(&mut self, enable: bool) -> Result<(), MpcTlsError> {
+        self.is_decrypting = enable;
+
+        if enable {
+            self.notifier.set();
+        } else {
+            self.notifier.clear();
+        }
 
         Ok(())
     }
 
+    /// Returns if incoming messages are decrypted.
+    pub fn is_decrypting(&self) -> bool {
+        self.is_decrypting
+    }
+
     /// Stops the actor.
     pub fn stop(&mut self, ctx: &mut LudiContext<Self>) {
         ctx.stop();

crates/mpc-tls/src/leader/actor.rs

@@ -32,10 +32,14 @@ impl MpcTlsLeaderCtrl {
         Self { address }
     }
 
-    /// Defers decryption of any incoming messages.
-    pub async fn defer_decryption(&self) -> Result<(), MpcTlsError> {
+    /// Enables or disables the decryption of any incoming messages.
+    ///
+    /// # Arguments
+    ///
+    /// * `enable` - Whether to enable or disable decryption.
+    pub async fn enable_decryption(&self, enable: bool) -> Result<(), MpcTlsError> {
         self.address
-            .send(DeferDecryption)
+            .send(EnableDecryption { enable })
             .await
             .map_err(MpcTlsError::actor)?
     }
@@ -981,7 +985,7 @@ impl Handler<BackendMsgServerClosed> for MpcTlsLeader {
     }
 }
 
-impl Dispatch<MpcTlsLeader> for DeferDecryption {
+impl Dispatch<MpcTlsLeader> for EnableDecryption {
     fn dispatch<R: FnOnce(Self::Return) + Send>(
         self,
         actor: &mut MpcTlsLeader,
@@ -992,13 +996,13 @@ impl Dispatch<MpcTlsLeader> for DeferDecryption {
     }
 }
 
-impl Handler<DeferDecryption> for MpcTlsLeader {
+impl Handler<EnableDecryption> for MpcTlsLeader {
     async fn handle(
         &mut self,
-        _msg: DeferDecryption,
+        msg: EnableDecryption,
         _ctx: &mut LudiCtx<Self>,
-    ) -> <DeferDecryption as Message>::Return {
-        self.defer_decryption().await
+    ) -> <EnableDecryption as Message>::Return {
+        self.enable_decryption(msg.enable)
     }
 }
 
@@ -1048,7 +1052,7 @@ pub enum MpcTlsLeaderMsg {
     BackendMsgGetNotify(BackendMsgGetNotify),
     BackendMsgIsEmpty(BackendMsgIsEmpty),
     BackendMsgServerClosed(BackendMsgServerClosed),
-    DeferDecryption(DeferDecryption),
+    DeferDecryption(EnableDecryption),
     Stop(Stop),
 }
 
@@ -1083,7 +1087,7 @@ pub enum MpcTlsLeaderMsgReturn {
     BackendMsgGetNotify(<BackendMsgGetNotify as Message>::Return),
     BackendMsgIsEmpty(<BackendMsgIsEmpty as Message>::Return),
     BackendMsgServerClosed(<BackendMsgServerClosed as Message>::Return),
-    DeferDecryption(<DeferDecryption as Message>::Return),
+    DeferDecryption(<EnableDecryption as Message>::Return),
     Stop(<Stop as Message>::Return),
 }
 
@@ -1732,23 +1736,25 @@ impl Wrap<BackendMsgServerClosed> for MpcTlsLeaderMsg {
     }
 }
 
-/// Message to start deferring the decryption.
+/// Message to enable or disable the decryption of messages.
 #[allow(missing_docs)]
 #[derive(Debug)]
-pub struct DeferDecryption;
+pub struct EnableDecryption {
+    pub enable: bool,
+}
 
-impl Message for DeferDecryption {
+impl Message for EnableDecryption {
     type Return = Result<(), MpcTlsError>;
 }
 
-impl From<DeferDecryption> for MpcTlsLeaderMsg {
-    fn from(value: DeferDecryption) -> Self {
+impl From<EnableDecryption> for MpcTlsLeaderMsg {
+    fn from(value: EnableDecryption) -> Self {
         MpcTlsLeaderMsg::DeferDecryption(value)
     }
 }
 
-impl Wrap<DeferDecryption> for MpcTlsLeaderMsg {
-    fn unwrap_return(ret: Self::Return) -> Result<<DeferDecryption as Message>::Return, Error> {
+impl Wrap<EnableDecryption> for MpcTlsLeaderMsg {
+    fn unwrap_return(ret: Self::Return) -> Result<<EnableDecryption as Message>::Return, Error> {
         match ret {
             Self::Return::DeferDecryption(value) => Ok(value),
             _ => Err(Error::Wrapper),

crates/mpc-tls/tests/test.rs

@@ -88,7 +88,7 @@ async fn leader_task(mut leader: MpcTlsLeader) {
     let mut buf = vec![0u8; 48];
     conn.read_exact(&mut buf).await.unwrap();
 
-    leader_ctrl.defer_decryption().await.unwrap();
+    leader_ctrl.enable_decryption(false).await.unwrap();
 
     let msg = concat!(
         "POST /echo HTTP/1.1\r\n",

crates/tls/client-async/src/lib.rs

@@ -197,8 +197,8 @@ pub fn bind_client<T: AsyncRead + AsyncWrite + Send + Unpin + 'static>(
 
                         sent.extend(&data);
                         client
-                            .write_all_plaintext(&data)
-                            .await?;
+                            .write_all_plaintext(&data)?;
+                        client.process_new_packets().await?;
 
                         tx_recv_fut = tx_receiver.next().fuse();
                     } else {

crates/tls/client/src/conn.rs

@@ -690,6 +690,11 @@ impl CommonState {
         self.received_plaintext.is_empty()
     }
 
+    /// Returns true if the buffer for sendable plaintext is full.
+    pub fn sendable_plaintext_is_full(&self) -> bool {
+        self.sendable_plaintext.is_full()
+    }
+
     /// Returns true if the connection is currently performing the TLS
     /// handshake.
     ///

crates/tls/client/src/vecbuf.rs

@@ -35,6 +35,15 @@ impl ChunkVecBuffer {
         self.chunks.is_empty()
     }
 
+    /// If the buffer has reached limit.
+    pub(crate) fn is_full(&self) -> bool {
+        if let Some(limit) = self.limit {
+            self.len() >= limit
+        } else {
+            false
+        }
+    }
+
     /// How many bytes we're storing
     pub(crate) fn len(&self) -> usize {
         let mut len = 0;