tlsn/
session.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    sync::{
5        Arc, Mutex,
6        atomic::{AtomicBool, Ordering},
7    },
8    task::{Context, Poll, Waker},
9};
10
11use futures::{AsyncRead, AsyncWrite};
12use mpz_common::{Session as MpzSession, ThreadPool, ThreadPoolBuildError, io::Io, mux::Mux};
13use tlsn_core::config::{prover::ProverConfig, verifier::VerifierConfig};
14use tlsn_mux::{Connection, Handle};
15
16use crate::{
17    Error, Result,
18    prover::{Prover, state as prover_state},
19    verifier::{Verifier, state as verifier_state},
20};
21
22/// Configuration for a [`Session`].
23#[derive(Debug, Clone)]
24pub struct SessionConfig {
25    mux: tlsn_mux::Config,
26}
27
28impl SessionConfig {
29    /// Creates a new builder.
30    pub fn builder() -> SessionConfigBuilder {
31        SessionConfigBuilder::default()
32    }
33
34    /// Returns the underlying [`tlsn_mux::Config`].
35    fn to_mux_config(&self) -> tlsn_mux::Config {
36        self.mux.clone()
37    }
38}
39
40impl Default for SessionConfig {
41    fn default() -> Self {
42        Self {
43            mux: session_mux_config(None).expect("default mux config is valid"),
44        }
45    }
46}
47
48/// Builds a [`tlsn_mux::Config`] with the session's fixed settings applied.
49///
50/// All mux-specific invariants (e.g. the receive window / stream limit
51/// coupling) are validated by [`tlsn_mux::ConfigBuilder`].
52fn session_mux_config(
53    max_num_streams: Option<usize>,
54) -> Result<tlsn_mux::Config, tlsn_mux::ConfigError> {
55    let mut builder = tlsn_mux::Config::builder()
56        .keep_alive(true)
57        .close_sync(true);
58
59    if let Some(max_num_streams) = max_num_streams {
60        builder = builder.max_num_streams(max_num_streams);
61    }
62
63    builder.build()
64}
65
66/// Builder for [`SessionConfig`].
67#[derive(Debug, Default)]
68pub struct SessionConfigBuilder {
69    max_num_streams: Option<usize>,
70}
71
72impl SessionConfigBuilder {
73    /// Sets the maximum number of concurrent streams per session.
74    ///
75    /// Defaults to the mux default (512) when unset. The limit is bounded by
76    /// the session's receive window (1 GiB by default); higher values are
77    /// rejected by [`build`](Self::build).
78    pub fn max_num_streams(mut self, max_num_streams: usize) -> Self {
79        self.max_num_streams = Some(max_num_streams);
80        self
81    }
82
83    /// Builds the configuration.
84    ///
85    /// Returns an error if the requested stream limit exceeds what the
86    /// session's receive window supports.
87    pub fn build(self) -> Result<SessionConfig, SessionConfigError> {
88        Ok(SessionConfig {
89            mux: session_mux_config(self.max_num_streams)?,
90        })
91    }
92}
93
94/// Error for [`SessionConfig`].
95#[derive(Debug, thiserror::Error)]
96#[error(transparent)]
97pub struct SessionConfigError(#[from] tlsn_mux::ConfigError);
98
99/// A TLSNotary session over a communication channel.
100///
101/// Wraps an async IO stream and provides multiplexing for the protocol. Use
102/// [`new_prover`](Self::new_prover) or [`new_verifier`](Self::new_verifier) to
103/// create protocol participants.
104///
105/// The session must be polled continuously (either directly or via
106/// [`split`](Self::split)) to drive the underlying connection. After the
107/// session closes, the underlying IO can be reclaimed with
108/// [`try_take`](Self::try_take).
109///
110/// **Important**: The order in which provers and verifiers are created must
111/// match on both sides. For example, if the prover side calls `new_prover`
112/// then `new_verifier`, the verifier side must call `new_verifier` then
113/// `new_prover`.
114#[must_use = "session must be polled continuously to make progress, including during closing."]
115pub struct Session<Io> {
116    conn: Option<Connection<Io>>,
117    executor: MpzSession,
118    handle: Handle,
119}
120
121impl<Io> Session<Io>
122where
123    Io: AsyncRead + AsyncWrite + Unpin,
124{
125    /// Creates a new session over `io`, the prover ↔ verifier channel.
126    ///
127    /// On TCP transports, disable Nagle's algorithm; see
128    /// [Performance](crate#performance).
129    pub fn new(io: Io) -> Self {
130        Self::new_with_config(io, SessionConfig::default())
131    }
132
133    /// Creates a new session over `io` with the given configuration.
134    ///
135    /// On TCP transports, disable Nagle's algorithm; see
136    /// [Performance](crate#performance).
137    pub fn new_with_config(io: Io, config: SessionConfig) -> Self {
138        let mux_config = config.to_mux_config();
139
140        let conn = tlsn_mux::Connection::new(io, mux_config);
141        let handle = conn.handle().expect("handle should be available");
142        let executor = build_executor(MuxHandle {
143            handle: handle.clone(),
144        });
145
146        Self {
147            conn: Some(conn),
148            executor,
149            handle,
150        }
151    }
152
153    /// Creates a new prover.
154    pub fn new_prover(
155        &mut self,
156        config: ProverConfig,
157    ) -> Result<Prover<prover_state::Initialized>> {
158        let ctx = self.executor.new_context().map_err(|e| {
159            Error::internal()
160                .with_msg("failed to create new prover")
161                .with_source(e)
162        })?;
163
164        Ok(Prover::new(ctx, self.handle.clone(), config))
165    }
166
167    /// Creates a new verifier.
168    pub fn new_verifier(
169        &mut self,
170        config: VerifierConfig,
171    ) -> Result<Verifier<verifier_state::Initialized>> {
172        let ctx = self.executor.new_context().map_err(|e| {
173            Error::internal()
174                .with_msg("failed to create new verifier")
175                .with_source(e)
176        })?;
177
178        Ok(Verifier::new(ctx, self.handle.clone(), config))
179    }
180
181    /// Returns `true` if the session is closed.
182    pub fn is_closed(&self) -> bool {
183        self.conn
184            .as_ref()
185            .map(|mux| mux.is_complete())
186            .unwrap_or_default()
187    }
188
189    /// Closes the session.
190    ///
191    /// This will cause the session to begin closing. Session must continue to
192    /// be polled until completion.
193    pub fn close(&mut self) {
194        if let Some(conn) = self.conn.as_mut() {
195            conn.close()
196        }
197    }
198
199    /// Attempts to take the IO, returning an error if it is not available.
200    pub fn try_take(&mut self) -> Result<Io> {
201        let conn = self.conn.take().ok_or_else(|| {
202            Error::io().with_msg("failed to take the session io, it was already taken")
203        })?;
204
205        match conn.try_into_io() {
206            Err(conn) => {
207                self.conn = Some(conn);
208                Err(Error::io()
209                    .with_msg("failed to take the session io, session was not completed yet"))
210            }
211            Ok(conn) => Ok(conn),
212        }
213    }
214
215    /// Polls the session.
216    pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
217        self.conn
218            .as_mut()
219            .ok_or_else(|| {
220                Error::io()
221                    .with_msg("failed to poll the session connection because it has been taken")
222            })?
223            .poll(cx)
224            .map_err(|e| {
225                Error::io()
226                    .with_msg("error occurred while polling the session connection")
227                    .with_source(e)
228            })
229    }
230
231    /// Splits the session into a driver and handle.
232    ///
233    /// The driver must be polled to make progress. The handle is used
234    /// for creating provers/verifiers and closing the session.
235    pub fn split(self) -> (SessionDriver<Io>, SessionHandle) {
236        let should_close = Arc::new(AtomicBool::new(false));
237        let waker = Arc::new(Mutex::new(None::<Waker>));
238
239        (
240            SessionDriver {
241                conn: self.conn,
242                should_close: should_close.clone(),
243                waker: waker.clone(),
244            },
245            SessionHandle {
246                executor: self.executor,
247                should_close,
248                waker,
249                handle: self.handle,
250            },
251        )
252    }
253}
254
255impl<Io> Future for Session<Io>
256where
257    Io: AsyncRead + AsyncWrite + Unpin,
258{
259    type Output = Result<()>;
260
261    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
262        Session::poll(&mut (*self), cx)
263    }
264}
265
266/// The polling half of a split session.
267///
268/// Must be polled continuously to drive the session. Returns the underlying
269/// IO when the session closes.
270#[must_use = "driver must be polled to make progress"]
271pub struct SessionDriver<Io> {
272    conn: Option<Connection<Io>>,
273    should_close: Arc<AtomicBool>,
274    waker: Arc<Mutex<Option<Waker>>>,
275}
276
277impl<Io> SessionDriver<Io>
278where
279    Io: AsyncRead + AsyncWrite + Unpin,
280{
281    /// Polls the driver.
282    pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<Io>> {
283        // Store the waker so the handle can wake us when close() is called.
284        {
285            let mut waker_guard = self.waker.lock().unwrap();
286            *waker_guard = Some(cx.waker().clone());
287        }
288
289        let conn = self
290            .conn
291            .as_mut()
292            .ok_or_else(|| Error::io().with_msg("session driver already completed"))?;
293
294        if self.should_close.load(Ordering::Acquire) {
295            conn.close();
296        }
297
298        match conn.poll(cx) {
299            Poll::Ready(Ok(())) => {}
300            Poll::Ready(Err(e)) => {
301                return Poll::Ready(Err(Error::io()
302                    .with_msg("error polling session connection")
303                    .with_source(e)));
304            }
305            Poll::Pending => return Poll::Pending,
306        }
307
308        let conn = self.conn.take().unwrap();
309        Poll::Ready(
310            conn.try_into_io()
311                .map_err(|_| Error::io().with_msg("failed to take session io")),
312        )
313    }
314}
315
316impl<Io> Future for SessionDriver<Io>
317where
318    Io: AsyncRead + AsyncWrite + Unpin,
319{
320    type Output = Result<Io>;
321
322    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
323        SessionDriver::poll(&mut *self, cx)
324    }
325}
326
327/// The control half of a split session.
328///
329/// Used to create provers/verifiers and control the session lifecycle.
330pub struct SessionHandle {
331    executor: MpzSession,
332    should_close: Arc<AtomicBool>,
333    waker: Arc<Mutex<Option<Waker>>>,
334    handle: Handle,
335}
336
337impl SessionHandle {
338    /// Creates a new prover.
339    pub fn new_prover(
340        &mut self,
341        config: ProverConfig,
342    ) -> Result<Prover<prover_state::Initialized>> {
343        let ctx = self.executor.new_context().map_err(|e| {
344            Error::internal()
345                .with_msg("failed to create new prover")
346                .with_source(e)
347        })?;
348
349        Ok(Prover::new(ctx, self.handle.clone(), config))
350    }
351
352    /// Creates a new verifier.
353    pub fn new_verifier(
354        &mut self,
355        config: VerifierConfig,
356    ) -> Result<Verifier<verifier_state::Initialized>> {
357        let ctx = self.executor.new_context().map_err(|e| {
358            Error::internal()
359                .with_msg("failed to create new verifier")
360                .with_source(e)
361        })?;
362
363        Ok(Verifier::new(ctx, self.handle.clone(), config))
364    }
365
366    /// Signals the session to close.
367    ///
368    /// The driver must continue to be polled until it completes.
369    pub fn close(&self) {
370        self.should_close.store(true, Ordering::Release);
371        if let Some(waker) = self.waker.lock().unwrap().take() {
372            waker.wake();
373        }
374    }
375}
376
377/// Multiplexer controller providing streams.
378struct MuxHandle {
379    handle: Handle,
380}
381
382impl std::fmt::Debug for MuxHandle {
383    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
384        f.debug_struct("MuxHandle").finish_non_exhaustive()
385    }
386}
387
388impl Mux for MuxHandle {
389    fn open(&self, id: &[u8]) -> Result<Io, std::io::Error> {
390        let stream = self.handle.new_stream(id).map_err(std::io::Error::other)?;
391        let io = Io::from_io(stream);
392
393        Ok(io)
394    }
395}
396
397/// Builds a session backed by the process-wide global thread pool, with the
398/// given muxer.
399fn build_executor(mux: MuxHandle) -> MpzSession {
400    #[cfg(all(feature = "web", target_arch = "wasm32"))]
401    let cores = web_spawn::available_parallelism().map(|n| n.get());
402
403    #[cfg(not(all(feature = "web", target_arch = "wasm32")))]
404    let cores = std::thread::available_parallelism().map(|n| n.get());
405
406    let builder = ThreadPool::builder().num_threads(cores.unwrap_or(8));
407
408    #[cfg(all(feature = "web", target_arch = "wasm32"))]
409    let builder = builder.spawn(|f| {
410        let _ = web_spawn::spawn(f);
411        Ok(())
412    });
413
414    // Install our configured pool as the process-wide global pool on first
415    // use.
416    match builder.build_global() {
417        Ok(()) | Err(ThreadPoolBuildError::AlreadyInitialized) => {}
418        Err(e) => panic!("failed to build global thread pool: {e}"),
419    }
420
421    MpzSession::builder()
422        .pool(ThreadPool::global())
423        .build(mux)
424        .expect("session should build")
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430
431    #[test]
432    fn default_config_is_valid() {
433        assert!(SessionConfig::builder().build().is_ok());
434    }
435
436    #[test]
437    fn oversized_max_num_streams_is_rejected() {
438        let err = SessionConfig::builder()
439            .max_num_streams(usize::MAX)
440            .build()
441            .unwrap_err();
442        assert!(matches!(
443            err.0,
444            tlsn_mux::ConfigError::ReceiveWindowTooSmall { .. }
445        ));
446    }
447}
tlsn/session.rs

tlsn/
session.rs
