feat: add new socket api

Author: th4s

crates/tlsn/src/prover.rs

@@ -1,10 +1,14 @@
 //! Prover.
 
 mod client;
+mod conn;
+mod control;
 mod error;
 mod prove;
 pub mod state;
 
+pub use conn::{ConnectionFuture, TlsConnection};
+pub use control::ProverControl;
 pub use error::ProverError;
 pub use tlsn_core::ProverOutput;
 
@@ -21,7 +25,7 @@ use futures::{AsyncRead, AsyncWrite, FutureExt, TryFutureExt};
 use rustls_pki_types::CertificateDer;
 use serio::{SinkExt, stream::IoStreamExt};
 use std::{
-    sync::Arc,
+    sync::{Arc, Mutex},
     task::{Context, Poll},
 };
 use tls_client::{ClientConnection, ServerName as TlsServerName};
@@ -223,6 +227,35 @@ impl Prover<state::CommitAccepted> {
         };
         Ok(prover)
     }
+
+    /// Connects the prover and attaches a socket.
+    ///
+    /// This is a convenience function which returns
+    ///   - [`TlsConnection`] for reading and writing traffic as well as other
+    ///     connection-specific settings.
+    ///   - [`ConnectionFuture`] which has to be polled for driving the
+    ///     connection forward.
+    ///
+    /// # Arguments
+    ///
+    /// * `config` - The TLS client configuration.
+    #[instrument(parent = &self.span, level = "debug", skip_all, err)]
+    pub async fn connect_with<S>(
+        self,
+        config: TlsClientConfig,
+        socket: S,
+    ) -> Result<(TlsConnection, ConnectionFuture<S>), ProverError>
+    where
+        S: AsyncRead + AsyncWrite + Send,
+    {
+        let prover = self.connect(config).await?;
+        let prover = Arc::new(Mutex::new(prover));
+
+        let conn = TlsConnection::new(prover.clone());
+        let fut = ConnectionFuture::new(socket, prover.clone());
+
+        Ok((conn, fut))
+    }
 }
 
 impl Prover<state::Connected> {

crates/tlsn/src/prover/conn.rs

@@ -0,0 +1,94 @@
+use std::sync::{Arc, Mutex};
+
+use crate::prover::{Prover, ProverError, control::ProverControl, state};
+use futures::{AsyncRead, AsyncWrite};
+
+pub(crate) type ProverConnected = Arc<Mutex<Prover<state::Connected>>>;
+
+/// A TLS connection to a server.
+///
+/// This type implements [`AsyncRead`] and [`AsyncWrite`] and can be used to
+/// communicate with a server using TLS.
+///
+/// # Note
+///
+/// This connection is closed on a best-effort basis if this is dropped. To
+/// ensure a clean close, you should call
+/// [`AsyncWriteExt::close`](futures::io::AsyncWriteExt::close) to close the
+/// connection.
+pub struct TlsConnection {
+    prover: ProverConnected,
+}
+
+impl TlsConnection {
+    pub(crate) fn new(prover: ProverConnected) -> Self {
+        Self { prover }
+    }
+}
+
+impl AsyncRead for TlsConnection {
+    fn poll_read(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &mut [u8],
+    ) -> std::task::Poll<std::io::Result<usize>> {
+        todo!()
+    }
+}
+
+impl AsyncWrite for TlsConnection {
+    fn poll_write(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &[u8],
+    ) -> std::task::Poll<std::io::Result<usize>> {
+        todo!()
+    }
+
+    fn poll_flush(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<std::io::Result<()>> {
+        todo!()
+    }
+
+    fn poll_close(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<std::io::Result<()>> {
+        todo!()
+    }
+}
+
+/// A future to drive the connection. Must be polled to make progress.
+pub struct ConnectionFuture<S> {
+    socket: S,
+    prover: ProverConnected,
+}
+
+impl<S> ConnectionFuture<S> {
+    pub(crate) fn new(socket: S, prover: ProverConnected) -> Self {
+        Self { socket, prover }
+    }
+
+    /// Returns a handle to control the prover.
+    pub fn handle(&self) -> ProverControl {
+        ProverControl {
+            prover: self.prover.clone(),
+        }
+    }
+}
+
+impl<S> Future for ConnectionFuture<S>
+where
+    S: AsyncRead + AsyncWrite + Send,
+{
+    type Output = Result<Prover<state::Committed>, ProverError>;
+
+    fn poll(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Self::Output> {
+        todo!()
+    }
+}

crates/tlsn/src/prover/control.rs

@@ -0,0 +1,31 @@
+use crate::prover::{ProverError, conn::ProverConnected};
+
+/// A controller for the prover.
+#[derive(Clone)]
+pub struct ProverControl {
+    pub(crate) prover: ProverConnected,
+}
+
+impl ProverControl {
+    /// Returns whether the prover is decrypting the server traffic.
+    pub fn is_decrypting(&self) -> bool {
+        let prover = self
+            .prover
+            .lock()
+            .expect("should be able to acquire prover handle");
+        prover.is_decrypting()
+    }
+
+    /// Enables or disables the decryption of server traffic.
+    ///
+    /// # Arguments
+    ///
+    /// * `enable` - If decryption should be enabled or disabled.
+    pub fn enable_decryption(&self, enable: bool) -> Result<(), ProverError> {
+        let mut prover = self
+            .prover
+            .lock()
+            .expect("should be able to acquire prover handle");
+        prover.enable_decryption(enable)
+    }
+}