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

Skip to main content

#
tlsn/
prover.rs

    1//! Prover.
    2
    3mod client;
    4mod conn;
    5mod control;
    6mod future;
    7mod prove;
    8pub mod state;
    9
    10pub use conn::TlsConnection;
    11pub use control::ProverControl;
    12pub use future::ProverFuture;
    13pub use tlsn_core::ProverOutput;
    14
    15use crate::{
    16    Error, Result,
    17    mpz::{ProverDeps, build_prover_deps, translate_keys},
    18    msg::{ProveRequestMsg, Response, TlsCommitRequestMsg},
    19    prover::{
    20        client::{MpcTlsClient, TlsOutput},
    21        state::ConnectedProj,
    22    },
    23};
    24
    25use futures::{AsyncRead, AsyncWrite, ready};
    26use mpz_common::Context;
    27use rustls_pki_types::CertificateDer;
    28use serio::{SinkExt, stream::IoStreamExt};
    29use std::{pin::Pin, sync::Arc, task::Poll};
    30use tls_client::{ClientConnection, ServerName as TlsServerName};
    31use tlsn_core::{
    32    config::{
    33        prove::ProveConfig,
    34        prover::ProverConfig,
    35        tls::TlsClientConfig,
    36        tls_commit::{TlsCommitConfig, TlsCommitProtocolConfig},
    37    },
    38    connection::{HandshakeData, ServerName},
    39    transcript::{TlsTranscript, Transcript},
    40};
    41use tracing::{Span, debug, info_span, instrument};
    42use webpki::anchor_from_trusted_cert;
    43
    44const BUF_CAP: usize = 16 * 1024 * 1024;
    45
    46/// A prover instance.
    47#[derive(Debug)]
    48pub struct Prover<T: state::ProverState = state::Initialized> {
    49    config: ProverConfig,
    50    span: Span,
    51    ctx: Option<Context>,
    52    state: T,
    53}
    54
    55impl Prover<state::Initialized> {
    56    /// Creates a new prover.
    57    ///
    58    /// # Arguments
    59    ///
    60    /// * `ctx` - A thread context.
    61    /// * `config` - The configuration for the prover.
    62    pub(crate) fn new(ctx: Context, config: ProverConfig) -> Self {
    63        let span = info_span!("prover");
    64        Self {
    65            config,
    66            span,
    67            ctx: Some(ctx),
    68            state: state::Initialized,
    69        }
    70    }
    71
    72    /// Starts the TLS commitment protocol.
    73    ///
    74    /// This initiates the TLS commitment protocol, including performing any
    75    /// necessary preprocessing operations.
    76    ///
    77    /// # Arguments
    78    ///
    79    /// * `config` - The TLS commitment configuration.
    80    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
    81    pub async fn commit(
    82        mut self,
    83        config: TlsCommitConfig,
    84    ) -> Result<Prover<state::CommitAccepted>> {
    85        let mut ctx = self
    86            .ctx
    87            .take()
    88            .ok_or_else(|| Error::internal().with_msg("commitment protocol context was dropped"))?;
    89
    90        // Sends protocol configuration to verifier for compatibility check.
    91        ctx.io_mut()
    92            .send(TlsCommitRequestMsg {
    93                request: config.to_request(),
    94                version: crate::VERSION.clone(),
    95            })
    96            .await
    97            .map_err(|e| {
    98                Error::io()
    99                    .with_msg("commitment protocol failed to send request")
    100                    .with_source(e)
    101            })?;
    102
    103        ctx.io_mut()
    104            .expect_next::<Response>()
    105            .await
    106            .map_err(|e| {
    107                Error::io()
    108                    .with_msg("commitment protocol failed to receive response")
    109                    .with_source(e)
    110            })?
    111            .result
    112            .map_err(|e| {
    113                Error::user()
    114                    .with_msg("commitment protocol rejected by verifier")
    115                    .with_source(e)
    116            })?;
    117
    118        let TlsCommitProtocolConfig::Mpc(mpc_tls_config) = config.protocol().clone() else {
    119            unreachable!("only MPC TLS is supported");
    120        };
    121
    122        let ProverDeps { vm, mut mpc_tls } = build_prover_deps(mpc_tls_config, ctx);
    123
    124        // Allocate resources for MPC-TLS in the VM.
    125        let mut keys = mpc_tls.alloc().map_err(|e| {
    126            Error::internal()
    127                .with_msg("commitment protocol failed to allocate mpc-tls resources")
    128                .with_source(e)
    129        })?;
    130        let vm_lock = vm.try_lock().expect("VM is not locked");
    131        translate_keys(&mut keys, &vm_lock);
    132        drop(vm_lock);
    133
    134        debug!("setting up mpc-tls");
    135
    136        mpc_tls.preprocess().await.map_err(|e| {
    137            Error::internal()
    138                .with_msg("commitment protocol failed during mpc-tls preprocessing")
    139                .with_source(e)
    140        })?;
    141
    142        debug!("mpc-tls setup complete");
    143
    144        Ok(Prover {
    145            config: self.config,
    146            span: self.span,
    147            ctx: None,
    148            state: state::CommitAccepted { mpc_tls, keys, vm },
    149        })
    150    }
    151}
    152
    153impl Prover<state::CommitAccepted> {
    154    /// Connects to the server using the provided socket.
    155    ///
    156    /// Returns a handle to the TLS connection, a future which returns the
    157    /// prover once the connection is closed and the TLS transcript is
    158    /// committed.
    159    ///
    160    /// # Arguments
    161    ///
    162    /// * `config` - The TLS client configuration.
    163    /// * `socket` - The socket to the server.
    164    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
    165    pub fn connect<S: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
    166        self,
    167        config: TlsClientConfig,
    168        socket: S,
    169    ) -> Result<(TlsConnection, ProverFuture<S>)> {
    170        let state::CommitAccepted {
    171            mpc_tls, keys, vm, ..
    172        } = self.state;
    173
    174        let decrypt = mpc_tls.is_decrypting();
    175
    176        let ServerName::Dns(server_name) = config.server_name();
    177        let server_name =
    178            TlsServerName::try_from(server_name.as_ref()).expect("name was validated");
    179
    180        let root_store = tls_client::RootCertStore {
    181            roots: config
    182                .root_store()
    183                .roots
    184                .iter()
    185                .map(|cert| {
    186                    let der = CertificateDer::from_slice(&cert.0);
    187                    anchor_from_trusted_cert(&der)
    188                        .map(|anchor| anchor.to_owned())
    189                        .map_err(|e| {
    190                            Error::config()
    191                                .with_msg("failed to parse root certificate")
    192                                .with_source(e)
    193                        })
    194                })
    195                .collect::<Result<Vec<_>, _>>()?,
    196        };
    197
    198        let rustls_config = tls_client::ClientConfig::builder()
    199            .with_safe_defaults()
    200            .with_root_certificates(root_store);
    201
    202        let rustls_config = if let Some((cert, key)) = config.client_auth() {
    203            rustls_config
    204                .with_single_cert(
    205                    cert.iter()
    206                        .map(|cert| tls_client::Certificate(cert.0.clone()))
    207                        .collect(),
    208                    tls_client::PrivateKey(key.0.clone()),
    209                )
    210                .map_err(|e| {
    211                    Error::config()
    212                        .with_msg("failed to configure client authentication")
    213                        .with_source(e)
    214                })?
    215        } else {
    216            rustls_config.with_no_client_auth()
    217        };
    218
    219        let client = ClientConnection::new(Arc::new(rustls_config), Box::new(mpc_tls), server_name)
    220            .map_err(|e| {
    221                Error::config()
    222                    .with_msg("failed to create tls client connection")
    223                    .with_source(e)
    224            })?;
    225
    226        let span = self.span.clone();
    227        let mpc_tls = MpcTlsClient::new(keys, vm, span, client, decrypt);
    228
    229        let (client_io, tlsn_conn) = futures_plex::duplex(BUF_CAP);
    230        let (client_to_server, server_to_client) = futures_plex::duplex(BUF_CAP);
    231
    232        let prover = Prover {
    233            ctx: self.ctx,
    234            config: self.config,
    235            span: self.span,
    236            state: state::Connected {
    237                server_name: config.server_name().clone(),
    238                tls_client: Box::new(mpc_tls),
    239                client_io,
    240                output: None,
    241                server_socket: socket,
    242                client_to_server,
    243                server_to_client,
    244                client_closed: false,
    245                server_closed: false,
    246            },
    247        };
    248
    249        let conn = TlsConnection::new(tlsn_conn);
    250        let fut = ProverFuture {
    251            prover: Some(prover),
    252        };
    253
    254        Ok((conn, fut))
    255    }
    256}
    257
    258impl<S> Future for Prover<state::Connected<S>>
    259where
    260    S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
    261{
    262    type Output = Result<(), Error>;
    263
    264    fn poll(
    265        mut self: std::pin::Pin<&mut Self>,
    266        cx: &mut std::task::Context<'_>,
    267    ) -> Poll<Self::Output> {
    268        let mut state = Pin::new(&mut self.state).project();
    269
    270        Self::io_to_tls_client(&mut state, cx)?;
    271
    272        if state.output.is_none()
    273            && let Poll::Ready(output) = state.tls_client.poll(cx)?
    274        {
    275            *state.output = Some(output);
    276        }
    277
    278        Self::io_from_tls_client(&mut state, cx)?;
    279
    280        if *state.server_closed && state.output.is_some() {
    281            ready!(state.client_io.poll_close(cx))?;
    282            ready!(state.server_socket.poll_close(cx))?;
    283
    284            return Poll::Ready(Ok(()));
    285        }
    286
    287        Poll::Pending
    288    }
    289}
    290
    291impl<S> Prover<state::Connected<S>>
    292where
    293    S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
    294{
    295    fn finish(self) -> Result<Prover<state::Committed>, Error> {
    296        let TlsOutput {
    297            ctx,
    298            vm,
    299            keys,
    300            tls_transcript,
    301            transcript,
    302        } = self
    303            .state
    304            .output
    305            .ok_or(Error::internal().with_msg("prover has not yet closed the connection"))?;
    306
    307        let prover = Prover {
    308            config: self.config,
    309            span: self.span,
    310            ctx: Some(ctx),
    311            state: state::Committed {
    312                vm,
    313                server_name: self.state.server_name,
    314                keys,
    315                tls_transcript,
    316                transcript,
    317            },
    318        };
    319
    320        Ok(prover)
    321    }
    322
    323    fn io_to_tls_client(
    324        state: &mut ConnectedProj<S>,
    325        cx: &mut std::task::Context<'_>,
    326    ) -> Result<(), Error> {
    327        // tls_conn -> tls_client
    328        // Always poll to register wakers, then check wants_write()
    329        if let Poll::Ready(mut simplex) = state.client_io.as_mut().poll_lock_read(cx)
    330            && let Poll::Ready(buf) = simplex.poll_get(cx)?
    331        {
    332            if !buf.is_empty() {
    333                if state.tls_client.wants_write() {
    334                    let write = state.tls_client.write(buf)?;
    335                    if write > 0 {
    336                        simplex.advance(write);
    337                    }
    338                } else {
    339                    cx.waker().wake_by_ref();
    340                }
    341            } else if !*state.client_closed && !*state.server_closed {
    342                *state.client_closed = true;
    343                state.tls_client.client_close();
    344            }
    345        }
    346
    347        // server_socket -> buf
    348        if let Poll::Ready(write) = state
    349            .server_to_client
    350            .poll_write_from(cx, state.server_socket.as_mut())?
    351        {
    352            if write == 0 && !*state.server_closed {
    353                *state.server_closed = true;
    354            } else if write > 0 {
    355                cx.waker().wake_by_ref();
    356            }
    357        }
    358
    359        // buf -> tls_client
    360        // Always poll to register wakers, then check wants_read_tls()
    361        if let Poll::Ready(mut simplex) = state.client_to_server.as_mut().poll_lock_read(cx)
    362            && let Poll::Ready(buf) = simplex.poll_get(cx)?
    363        {
    364            if state.tls_client.wants_read_tls() {
    365                let read = state.tls_client.read_tls(buf)?;
    366                if read > 0 {
    367                    simplex.advance(read);
    368                    cx.waker().wake_by_ref();
    369                }
    370            } else if !buf.is_empty() {
    371                cx.waker().wake_by_ref();
    372            }
    373        } else if *state.server_closed {
    374            state.tls_client.server_close();
    375        }
    376
    377        Ok(())
    378    }
    379
    380    fn io_from_tls_client(
    381        state: &mut ConnectedProj<S>,
    382        cx: &mut std::task::Context<'_>,
    383    ) -> Result<(), Error> {
    384        // tls_client -> buf
    385        // Always poll to register wakers, then check wants_write_tls()
    386        if let Poll::Ready(mut simplex) = state.client_to_server.as_mut().poll_lock_write(cx)
    387            && let Poll::Ready(buf) = simplex.poll_mut(cx)?
    388        {
    389            let write = state.tls_client.write_tls(buf)?;
    390            if write > 0 {
    391                simplex.advance_mut(write);
    392            } else if state.tls_client.wants_write_tls() {
    393                cx.waker().wake_by_ref();
    394            }
    395        }
    396
    397        // buf -> server_socket
    398        match state
    399            .server_to_client
    400            .poll_read_to(cx, state.server_socket.as_mut())
    401        {
    402            // do not attempt to write into closed sockets
    403            Poll::Ready(Err(err)) if matches!(err.kind(), std::io::ErrorKind::BrokenPipe) => {}
    404            Poll::Ready(Err(err)) if matches!(err.kind(), std::io::ErrorKind::ConnectionReset) => {}
    405            Poll::Ready(Err(err)) => return Err(Error::from(err)),
    406            _ => {}
    407        }
    408
    409        // tls_client -> tls_conn
    410        // Always poll to register wakers, then check wants_read()
    411        if let Poll::Ready(mut simplex) = state.client_io.as_mut().poll_lock_write(cx)
    412            && let Poll::Ready(buf) = simplex.poll_mut(cx)?
    413        {
    414            let read = state.tls_client.read(buf)?;
    415            if read > 0 {
    416                simplex.advance_mut(read);
    417            } else if state.tls_client.wants_read() {
    418                cx.waker().wake_by_ref();
    419            }
    420        }
    421
    422        Ok(())
    423    }
    424}
    425
    426impl Prover<state::Committed> {
    427    /// Returns the TLS transcript.
    428    pub fn tls_transcript(&self) -> &TlsTranscript {
    429        &self.state.tls_transcript
    430    }
    431
    432    /// Returns the transcript.
    433    pub fn transcript(&self) -> &Transcript {
    434        &self.state.transcript
    435    }
    436
    437    /// Proves information to the verifier.
    438    ///
    439    /// # Arguments
    440    ///
    441    /// * `config` - The disclosure configuration.
    442    #[instrument(parent = &self.span, level = "info", skip_all, err)]
    443    pub async fn prove(&mut self, config: &ProveConfig) -> Result<ProverOutput> {
    444        let ctx = self
    445            .ctx
    446            .as_mut()
    447            .ok_or_else(|| Error::internal().with_msg("proving context was dropped"))?;
    448        let state::Committed {
    449            vm,
    450            keys,
    451            server_name,
    452            tls_transcript,
    453            transcript,
    454            ..
    455        } = &mut self.state;
    456
    457        let handshake = config.server_identity().then(|| {
    458            (
    459                server_name.clone(),
    460                HandshakeData {
    461                    certs: tls_transcript
    462                        .server_cert_chain()
    463                        .expect("server cert chain is present")
    464                        .to_vec(),
    465                    sig: tls_transcript
    466                        .server_signature()
    467                        .expect("server signature is present")
    468                        .clone(),
    469                    binding: tls_transcript.certificate_binding().clone(),
    470                },
    471            )
    472        });
    473
    474        let partial_transcript = config
    475            .reveal()
    476            .map(|(sent, recv)| transcript.to_partial(sent.clone(), recv.clone()));
    477
    478        let msg = ProveRequestMsg {
    479            request: config.to_request(),
    480            handshake,
    481            transcript: partial_transcript,
    482        };
    483
    484        ctx.io_mut().send(msg).await.map_err(|e| {
    485            Error::io()
    486                .with_msg("failed to send prove configuration")
    487                .with_source(e)
    488        })?;
    489        ctx.io_mut()
    490            .expect_next::<Response>()
    491            .await
    492            .map_err(|e| {
    493                Error::io()
    494                    .with_msg("failed to receive prove response from verifier")
    495                    .with_source(e)
    496            })?
    497            .result
    498            .map_err(|e| {
    499                Error::user()
    500                    .with_msg("proving rejected by verifier")
    501                    .with_source(e)
    502            })?;
    503
    504        let output = prove::prove(ctx, vm, keys, transcript, tls_transcript, config).await?;
    505
    506        Ok(output)
    507    }
    508
    509    /// Closes the connection with the verifier.
    510    #[instrument(parent = &self.span, level = "info", skip_all, err)]
    511    pub async fn close(self) -> Result<()> {
    512        Ok(())
    513    }
    514}