# mpc.rs - source
Source of the Rust file `crates/tlsn/src/prover/client/mpc.rs`.

Skip to main content

#
tlsn/prover/client/
mpc.rs

    1//! Implementation of an MPC-TLS client.
    2
    3use crate::{
    4    error::Error as TlsnError,
    5    mpz::{ProverMpc, ProverZk},
    6    prover::client::{DecryptState, TlsClient, TlsOutput},
    7    tag::verify_tags,
    8};
    9use futures::{Future, FutureExt};
    10use mpc_tls::SessionKeys;
    11use mpz_common::Context;
    12use mpz_vm_core::Execute;
    13use std::{
    14    pin::Pin,
    15    sync::{Arc, atomic::AtomicBool},
    16    task::Poll,
    17};
    18use tls_client::ClientConnection;
    19use tlsn_core::transcript::TlsTranscript;
    20use tlsn_deap::Deap;
    21use tokio::sync::Mutex;
    22use tracing::{Span, debug, instrument, trace, warn};
    23
    24type FinalizeFuture =
    25    Box<dyn Future<Output = Result<(InnerState, Context, TlsTranscript), TlsnError>> + Send>;
    26
    27pub(crate) struct MpcTlsClient {
    28    state: State,
    29    decrypt: Arc<DecryptState>,
    30    client_wants_close: bool,
    31    server_closed: bool,
    32}
    33
    34enum State {
    35    Start {
    36        inner: Box<InnerState>,
    37    },
    38    Active {
    39        inner: Box<InnerState>,
    40    },
    41    Busy {
    42        fut: Pin<Box<dyn Future<Output = Result<Box<InnerState>, TlsnError>> + Send>>,
    43    },
    44    CloseActive {
    45        inner: Box<InnerState>,
    46    },
    47    CloseBusy {
    48        fut: Pin<Box<dyn Future<Output = Result<Box<InnerState>, TlsnError>> + Send>>,
    49    },
    50    Finalizing {
    51        fut: Pin<FinalizeFuture>,
    52    },
    53    Finished,
    54    Error,
    55}
    56
    57impl MpcTlsClient {
    58    pub(crate) fn new(
    59        keys: SessionKeys,
    60        vm: Arc<Mutex<Deap<ProverMpc, ProverZk>>>,
    61        span: Span,
    62        tls: ClientConnection,
    63        decrypt: bool,
    64    ) -> Self {
    65        let inner = InnerState {
    66            span,
    67            tls,
    68            vm,
    69            keys,
    70            decrypt,
    71            client_closed: false,
    72        };
    73
    74        let decrypt = DecryptState {
    75            decrypt: AtomicBool::new(decrypt),
    76        };
    77
    78        Self {
    79            decrypt: Arc::new(decrypt),
    80            client_wants_close: false,
    81            server_closed: false,
    82            state: State::Start {
    83                inner: Box::new(inner),
    84            },
    85        }
    86    }
    87
    88    fn inner_client_mut(&mut self) -> Option<&mut ClientConnection> {
    89        if let State::Active { inner } | State::CloseActive { inner } = &mut self.state {
    90            Some(&mut inner.tls)
    91        } else {
    92            None
    93        }
    94    }
    95
    96    fn inner_client(&self) -> Option<&ClientConnection> {
    97        if let State::Active { inner } | State::CloseActive { inner } = &self.state {
    98            Some(&inner.tls)
    99        } else {
    100            None
    101        }
    102    }
    103}
    104
    105impl TlsClient for MpcTlsClient {
    106    type Error = TlsnError;
    107
    108    fn wants_read_tls(&self) -> bool {
    109        if let Some(client) = self.inner_client() {
    110            client.wants_read()
    111        } else {
    112            false
    113        }
    114    }
    115
    116    fn wants_write_tls(&self) -> bool {
    117        if let Some(client) = self.inner_client() {
    118            client.wants_write()
    119        } else {
    120            false
    121        }
    122    }
    123
    124    fn read_tls(&mut self, mut buf: &[u8]) -> Result<usize, Self::Error> {
    125        if let Some(client) = self.inner_client_mut()
    126            && client.wants_read()
    127        {
    128            client.read_tls(&mut buf).map_err(TlsnError::from)
    129        } else {
    130            Ok(0)
    131        }
    132    }
    133
    134    fn write_tls(&mut self, mut buf: &mut [u8]) -> Result<usize, Self::Error> {
    135        if let Some(client) = self.inner_client_mut()
    136            && client.wants_write()
    137        {
    138            client.write_tls(&mut buf).map_err(TlsnError::from)
    139        } else {
    140            Ok(0)
    141        }
    142    }
    143
    144    fn wants_read(&self) -> bool {
    145        if let Some(client) = self.inner_client() {
    146            !client.plaintext_is_empty()
    147        } else {
    148            false
    149        }
    150    }
    151
    152    fn wants_write(&self) -> bool {
    153        if let Some(client) = self.inner_client() {
    154            !client.sendable_plaintext_is_full()
    155        } else {
    156            false
    157        }
    158    }
    159
    160    fn read(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error> {
    161        if let Some(client) = self.inner_client_mut()
    162            && !client.plaintext_is_empty()
    163        {
    164            client.read_plaintext(buf).map_err(TlsnError::from)
    165        } else {
    166            Ok(0)
    167        }
    168    }
    169
    170    fn write(&mut self, buf: &[u8]) -> Result<usize, Self::Error> {
    171        if let Some(client) = self.inner_client_mut()
    172            && !client.sendable_plaintext_is_full()
    173        {
    174            client
    175                .write_plaintext(buf)
    176                .map_err(|err| TlsnError::internal().with_source(err))
    177        } else {
    178            Ok(0)
    179        }
    180    }
    181
    182    fn client_close(&mut self) {
    183        self.client_wants_close = true;
    184    }
    185
    186    fn server_close(&mut self) {
    187        self.server_closed = true;
    188    }
    189
    190    fn decrypt(&self) -> Arc<DecryptState> {
    191        self.decrypt.clone()
    192    }
    193
    194    fn poll(&mut self, cx: &mut std::task::Context) -> Poll<Result<TlsOutput, Self::Error>> {
    195        match std::mem::replace(&mut self.state, State::Error) {
    196            State::Start { inner } => {
    197                trace!("inner client is starting");
    198                self.state = State::Busy {
    199                    fut: Box::pin(inner.start()),
    200                };
    201                self.poll(cx)
    202            }
    203            State::Active { mut inner } => {
    204                trace!("inner client is active");
    205                let decrypt = self.decrypt.is_decrypting();
    206
    207                if !inner.tls.is_handshaking() {
    208                    if self.server_closed {
    209                        self.state = State::CloseBusy {
    210                            fut: Box::pin(inner.server_close()),
    211                        };
    212                    } else if self.client_wants_close {
    213                        self.state = State::Busy {
    214                            fut: Box::pin(inner.client_close()),
    215                        };
    216                    } else if decrypt != inner.decrypt {
    217                        inner
    218                            .tls
    219                            .backend_mut()
    220                            .enable_decryption(decrypt)
    221                            .map_err(|err| TlsnError::internal().with_source(err))?;
    222
    223                        inner.decrypt = decrypt;
    224                        self.state = State::Busy {
    225                            fut: Box::pin(inner.run()),
    226                        };
    227                    } else {
    228                        self.state = State::Busy {
    229                            fut: Box::pin(inner.run()),
    230                        };
    231                    }
    232                    return self.poll(cx);
    233                }
    234                self.state = State::Busy {
    235                    fut: Box::pin(inner.run()),
    236                };
    237                self.poll(cx)
    238            }
    239            State::Busy { mut fut } => {
    240                trace!("inner client is busy");
    241
    242                match fut.as_mut().poll(cx)? {
    243                    Poll::Ready(inner) => {
    244                        self.state = State::Active { inner };
    245                        // re-poll immediately if a close transition is pending
    246                        if self.server_closed {
    247                            return self.poll(cx);
    248                        }
    249                    }
    250                    Poll::Pending => self.state = State::Busy { fut },
    251                }
    252                Poll::Pending
    253            }
    254            State::CloseActive { inner } => {
    255                trace!("inner client is active closing");
    256                self.state = State::CloseBusy {
    257                    fut: Box::pin(inner.run()),
    258                };
    259                self.poll(cx)
    260            }
    261            State::CloseBusy { mut fut } => {
    262                trace!("inner client is busy closing");
    263
    264                match fut.as_mut().poll(cx)? {
    265                    Poll::Ready(mut inner) if inner.is_record_layer_empty() => {
    266                        let (ctx, transcript) = inner
    267                            .tls
    268                            .backend_mut()
    269                            .finish()
    270                            .expect("connection should be closed");
    271                        self.state = State::Finalizing {
    272                            fut: Box::pin(inner.finalize(ctx, transcript)),
    273                        };
    274                        self.poll(cx)
    275                    }
    276                    Poll::Ready(inner) => {
    277                        cx.waker().wake_by_ref();
    278                        self.state = State::CloseActive { inner };
    279                        Poll::Pending
    280                    }
    281                    Poll::Pending => {
    282                        self.state = State::CloseBusy { fut };
    283                        Poll::Pending
    284                    }
    285                }
    286            }
    287            State::Finalizing { mut fut } => match fut.poll_unpin(cx) {
    288                Poll::Ready(output) => {
    289                    let (inner, ctx, tls_transcript) = output?;
    290                    let InnerState { vm, keys, .. } = inner;
    291
    292                    let transcript = tls_transcript
    293                        .to_transcript()
    294                        .expect("transcript is complete");
    295
    296                    let (_, vm) = Arc::into_inner(vm)
    297                        .expect("vm should have only 1 reference")
    298                        .into_inner()
    299                        .into_inner();
    300
    301                    let output = TlsOutput {
    302                        ctx,
    303                        vm,
    304                        keys,
    305                        tls_transcript,
    306                        transcript,
    307                    };
    308
    309                    self.state = State::Finished;
    310                    Poll::Ready(Ok(output))
    311                }
    312                Poll::Pending => {
    313                    self.state = State::Finalizing { fut };
    314                    Poll::Pending
    315                }
    316            },
    317            State::Finished => Poll::Ready(Err(
    318                TlsnError::internal().with_msg("mpc tls client polled again in finished state")
    319            )),
    320            State::Error => Poll::Ready(Err(
    321                TlsnError::internal().with_msg("mpc tls client is in error state")
    322            )),
    323        }
    324    }
    325}
    326
    327struct InnerState {
    328    span: Span,
    329    tls: ClientConnection,
    330    vm: Arc<Mutex<Deap<ProverMpc, ProverZk>>>,
    331    keys: SessionKeys,
    332    decrypt: bool,
    333    client_closed: bool,
    334}
    335
    336impl InnerState {
    337    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
    338    async fn start(mut self: Box<Self>) -> Result<Box<Self>, TlsnError> {
    339        self.tls
    340            .start()
    341            .await
    342            .map_err(|err| TlsnError::internal().with_source(err))?;
    343        Ok(self)
    344    }
    345
    346    #[instrument(parent = &self.span, level = "trace", skip_all, err)]
    347    async fn run(mut self: Box<Self>) -> Result<Box<Self>, TlsnError> {
    348        let mut state = self
    349            .tls
    350            .process_new_packets()
    351            .await
    352            .map_err(|err| TlsnError::internal().with_source(err))?;
    353        loop {
    354            let new_state = self
    355                .tls
    356                .process_new_packets()
    357                .await
    358                .map_err(|err| TlsnError::internal().with_source(err))?;
    359
    360            if new_state.plaintext_bytes_to_read() == state.plaintext_bytes_to_read()
    361                && new_state.tls_bytes_to_write() == state.tls_bytes_to_write()
    362            {
    363                break;
    364            }
    365            state = new_state;
    366        }
    367
    368        Ok(self)
    369    }
    370
    371    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
    372    async fn client_close(mut self: Box<Self>) -> Result<Box<Self>, TlsnError> {
    373        self = self.run().await?;
    374
    375        if !self.client_closed {
    376            debug!("sending close notify");
    377            if let Err(e) = self.tls.send_close_notify().await {
    378                warn!("failed to send close_notify to server: {}", e);
    379            }
    380            self.client_closed = true;
    381        }
    382
    383        Ok(self)
    384    }
    385
    386    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
    387    async fn server_close(mut self: Box<Self>) -> Result<Box<Self>, TlsnError> {
    388        self = self.run().await?;
    389
    390        self.tls
    391            .server_closed()
    392            .await
    393            .map_err(|err| TlsnError::internal().with_source(err))?;
    394        debug!("closed connection serverside");
    395
    396        Ok(self)
    397    }
    398
    399    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
    400    async fn finalize(
    401        self,
    402        mut ctx: Context,
    403        transcript: TlsTranscript,
    404    ) -> Result<(Self, Context, TlsTranscript), TlsnError> {
    405        {
    406            let mut vm = self.vm.try_lock().expect("VM should not be locked");
    407
    408            // Finalize DEAP.
    409            vm.finalize(&mut ctx)
    410                .await
    411                .map_err(|err| TlsnError::internal().with_source(err))?;
    412
    413            debug!("mpc finalized");
    414
    415            // Pull out ZK VM.
    416            let mut zk = vm.zk();
    417
    418            // Prove tag verification of received records.
    419            // The prover drops the proof output.
    420            let _ = verify_tags(
    421                &mut *zk,
    422                (self.keys.server_write_key, self.keys.server_write_iv),
    423                self.keys.server_write_mac_key,
    424                *transcript.version(),
    425                transcript.recv().to_vec(),
    426            )
    427            .map_err(|err| TlsnError::internal().with_source(err))?;
    428            debug!("verified tags from server");
    429
    430            zk.execute_all(&mut ctx)
    431                .await
    432                .map_err(|err| TlsnError::internal().with_source(err))?;
    433        }
    434
    435        debug!("MPC-TLS done");
    436        Ok((self, ctx, transcript))
    437    }
    438
    439    fn is_record_layer_empty(&self) -> bool {
    440        self.tls.plaintext_is_empty() && self.tls.is_empty().unwrap_or(false)
    441    }
    442}