WIP: tidying up state room...

Author: th4s

crates/tlsn/src/prover/client/mpc.rs

@@ -30,35 +30,29 @@ pub(crate) struct MpcTlsClient {
 }
 
 enum State {
-    ActiveIdle {
-        inner: InnerState,
+    Active {
+        mpc: Pin<MpcFuture>,
+        inner: Box<InnerState>,
     },
-    ActiveBusy {
+    Busy {
+        mpc: Pin<MpcFuture>,
         fut: Pin<Box<dyn Future<Output = Result<Box<InnerState>, ProverError>>>>,
     },
-    SetDecryptIdle {
+    SetDecrypt {
         enable: bool,
-        inner: InnerState,
+        mpc: Pin<MpcFuture>,
+        inner: Box<InnerState>,
     },
-    SetDecryptBusy {
-        fut: Pin<Box<dyn Future<Output = Result<Box<InnerState>, ProverError>>>>,
-    },
-    ClientCloseIdle {
-        inner: InnerState,
-    },
-    ClientCloseBusy {
-        fut: Pin<Box<dyn Future<Output = Result<Box<InnerState>, ProverError>>>>,
-    },
-    ServerCloseIdle {
-        inner: InnerState,
-    },
-    ServerCloseBusy {
-        fut: Pin<Box<dyn Future<Output = Result<Box<InnerState>, ProverError>>>>,
+    ClientClose {
+        mpc: Pin<MpcFuture>,
+        inner: Box<InnerState>,
     },
-    ClosingIdle {
-        inner: InnerState,
+    ServerClose {
+        mpc: Pin<MpcFuture>,
+        inner: Box<InnerState>,
     },
-    ClosingBusy {
+    Closing {
+        mpc: Pin<MpcFuture>,
         fut: Pin<Box<dyn Future<Output = Result<Box<InnerState>, ProverError>>>>,
     },
     Finalizing {
@@ -85,31 +79,33 @@ impl MpcTlsClient {
         tls: ClientConnection,
         decrypt: bool,
     ) -> Self {
+        let inner = InnerState {
+            span,
+            tls,
+            vm,
+            keys,
+            mpc_ctrl,
+        };
+
         Self {
-            state: State::ActiveIdle {
-                inner: InnerState {
-                    span,
-                    tls,
-                    mpc: Box::into_pin(mpc),
-                    vm,
-                    keys,
-                    mpc_ctrl,
-                    decrypt,
-                },
+            decrypt,
+            state: State::Active {
+                mpc: Box::into_pin(mpc),
+                inner: Box::new(inner),
             },
         }
     }
 
     fn inner_client_mut(&mut self) -> Option<&mut ClientConnection> {
-        if let State::ActiveIdle { inner } = &mut self.state {
+        if let State::Active { inner, .. } = &mut self.state {
             Some(&mut inner.tls)
         } else {
             None
         }
     }
 
     fn inner_client(&self) -> Option<&ClientConnection> {
-        if let State::ActiveIdle { inner } = &self.state {
+        if let State::Active { inner, .. } = &self.state {
             Some(&inner.tls)
         } else {
             None
@@ -118,17 +114,17 @@ impl MpcTlsClient {
 
     fn set_status(&mut self, status: Status) -> bool {
         match std::mem::replace(&mut self.state, State::Error) {
-            State::ActiveIdle { inner } => match status {
+            State::Active { inner, mpc } => match status {
                 Status::SetDecrypt(enable) => {
-                    self.state = State::SetDecryptIdle { enable, inner };
+                    self.state = State::SetDecrypt { enable, mpc, inner };
                     true
                 }
                 Status::ClientClose => {
-                    self.state = State::ClientCloseIdle { inner };
+                    self.state = State::ClientClose { mpc, inner };
                     true
                 }
                 Status::ServerClose => {
-                    self.state = State::ServerCloseIdle { inner };
+                    self.state = State::ServerClose { mpc, inner };
                     true
                 }
             },
@@ -256,30 +252,74 @@ impl TlsClient for MpcTlsClient {
 
     fn poll(&mut self, cx: &mut std::task::Context) -> Poll<Result<TlsOutput, Self::Error>> {
         match std::mem::replace(&mut self.state, State::Error) {
-            State::ActiveIdle { inner } => {
-                trace!("inner client is active idle");
-                let _ = inner.mpc.as_mut().poll(cx)?;
-                self.state = State::ActiveBusy { fut: inner.run() }
+            State::Active { mpc, inner } => {
+                trace!("inner client is active");
+
+                self.state = State::Busy {
+                    mpc,
+                    fut: Box::pin(inner.run()),
+                };
+                self.poll(cx)
             }
-            State::ActiveBusy { fut } => match fut.as_mut().poll(cx) {
-                Poll::Ready(res) => {
-                    let inner = res?;
-                    self.state.decrypt = *inner.decrypt;
-                    self.state = State::ActiveIdle { inner };
+            State::Busy { mut mpc, mut fut } => {
+                trace!("inner client is busy");
+                let _ = mpc.as_mut().poll(cx)?;
+
+                match fut.as_mut().poll(cx) {
+                    Poll::Ready(res) => {
+                        let inner = res?;
+                        self.state = State::Active { mpc, inner };
+                    }
+                    Poll::Pending => self.state = State::Busy { mpc, fut },
                 }
-                Poll::Pending => todo!(),
-            },
-            State::SetDecryptIdle {
-                enable,
-                inner: state,
-            } => todo!(),
-            State::SetDecryptBusy { fut } => todo!(),
-            State::ClientCloseIdle { inner: state } => todo!(),
-            State::ClientCloseBusy { fut } => todo!(),
-            State::ServerCloseIdle { inner: state } => todo!(),
-            State::ServerCloseBusy { fut } => todo!(),
-            State::ClosingIdle { inner: state } => todo!(),
-            State::ClosingBusy { fut } => todo!(),
+                Poll::Pending
+            }
+            State::SetDecrypt { enable, mpc, inner } => {
+                self.state = State::Busy {
+                    mpc,
+                    fut: Box::pin(inner.set_decrypt(enable)),
+                };
+                self.decrypt = enable;
+
+                debug!("set decryption to {}", enable);
+                self.poll(cx)
+            }
+            State::ClientClose { mpc, inner } => {
+                debug!("attempting to close connection clientside");
+                self.state = State::Closing {
+                    mpc,
+                    fut: Box::pin(inner.client_close()),
+                };
+                self.poll(cx)
+            }
+            State::ServerClose { mpc, inner } => {
+                debug!("attempting to close connection serverside");
+                self.state = State::Closing {
+                    mpc,
+                    fut: Box::pin(inner.server_close()),
+                };
+                self.poll(cx)
+            }
+            State::Closing { mut mpc, mut fut } => {
+                let Poll::Ready(res) = fut.as_mut().poll(cx) else {
+                    self.state = State::Closing { mpc, fut };
+                    return Poll::Pending;
+                };
+                let inner = res?;
+                let Poll::Ready(res) = mpc.as_mut().poll(cx) else {
+                    self.state = State::Closing {
+                        mpc,
+                        fut: Box::pin(inner.client_close()),
+                    };
+                    return Poll::Pending;
+                };
+
+                let (ctx, transcript) = res?;
+                self.state = State::Finalizing {
+                    fut: Box::pin(inner.finalize(ctx, transcript)),
+                };
+                self.poll(cx)
+            }
             State::Finalizing { mut fut } => match fut.as_mut().poll(cx) {
                 Poll::Ready(output) => {
                     let (state, ctx, tls_transcript) = output?;
@@ -323,18 +363,45 @@ impl TlsClient for MpcTlsClient {
 struct InnerState {
     span: Span,
     tls: ClientConnection,
-    mpc: Pin<MpcFuture>,
     vm: Arc<Mutex<Deap<ProverMpc, ProverZk>>>,
     keys: SessionKeys,
     mpc_ctrl: LeaderCtrl,
 }
 
 impl InnerState {
     #[instrument(parent = &self.span, level = "debug", skip_all, err)]
+    async fn set_decrypt(self: Box<Self>, enable: bool) -> Result<Box<Self>, ProverError> {
+        self.mpc_ctrl.enable_decryption(enable).await?;
+        self.run().await
+    }
+
+    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
+    async fn client_close(mut self: Box<Self>) -> Result<Box<Self>, ProverError> {
+        if self.tls.plaintext_is_empty() && self.tls.is_empty().await? {
+            if let Err(e) = self.tls.send_close_notify().await {
+                warn!("failed to send close_notify to server: {}", e);
+            };
+
+            self.mpc_ctrl.stop().await?;
+            debug!("closed connection");
+        }
+        self.run().await
+    }
+
+    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
+    async fn server_close(mut self: Box<Self>) -> Result<Box<Self>, ProverError> {
+        if self.tls.plaintext_is_empty() && self.tls.is_empty().await? {
+            self.tls.server_closed().await?;
+
+            self.mpc_ctrl.stop().await?;
+            debug!("closed connection");
+        }
+        self.run().await
+    }
+
+    #[instrument(parent = &self.span, level = "trace", skip_all, err)]
     async fn run(mut self: Box<Self>) -> Result<Box<Self>, ProverError> {
-        trace!("processing new packets");
         self.tls.process_new_packets().await?;
-
         Ok(self)
     }