Skip to content

Commit dee7a80

Browse files
th4sth4s
committed
cleanup
1 parent 2eee88c commit dee7a80

File tree

3 files changed

+22
-45
lines changed

3 files changed

+22
-45
lines changed

‎crates/tlsn/src/prover.rs‎

Lines changed: 10 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -226,7 +226,7 @@ impl Prover<state::CommitAccepted> {
226226
}
227227

228228
impl Prover<state::Connected> {
229-
/// Returns `true` if the prover can read TLS data from the server.
229+
/// Returns `true` if the prover wants to read TLS data from the server.
230230
pub fn wants_read_tls(&self) -> bool {
231231
self.state.tls_client.wants_read_tls()
232232
}
@@ -242,10 +242,7 @@ impl Prover<state::Connected> {
242242
///
243243
/// * `buf` - The buffer to read the TLS data from.
244244
pub fn read_tls(&mut self, buf: &[u8]) -> Result<usize, ProverError> {
245-
self.state
246-
.tls_client
247-
.read_tls(buf)
248-
.map_err(ProverError::from)
245+
self.state.tls_client.read_tls(buf)
249246
}
250247

251248
/// Writes TLS data for the server into the provided buffer.
@@ -254,10 +251,7 @@ impl Prover<state::Connected> {
254251
///
255252
/// * `buf` - The buffer to write the TLS data to.
256253
pub fn write_tls(&mut self, buf: &mut [u8]) -> Result<usize, ProverError> {
257-
self.state
258-
.tls_client
259-
.write_tls(buf)
260-
.map_err(ProverError::from)
254+
self.state.tls_client.write_tls(buf)
261255
}
262256

263257
/// Returns `true` if the prover wants to read plaintext data.
@@ -276,7 +270,7 @@ impl Prover<state::Connected> {
276270
///
277271
/// * `buf` - The buffer where the plaintext data gets written to.
278272
pub fn read(&mut self, buf: &mut [u8]) -> Result<usize, ProverError> {
279-
self.state.tls_client.read(buf).map_err(ProverError::from)
273+
self.state.tls_client.read(buf)
280274
}
281275

282276
/// Writes plaintext data to be sent to the server.
@@ -285,23 +279,17 @@ impl Prover<state::Connected> {
285279
///
286280
/// * `buf` - The buffer to read the plaintext data from.
287281
pub fn write(&mut self, buf: &[u8]) -> Result<usize, ProverError> {
288-
self.state.tls_client.write(buf).map_err(ProverError::from)
282+
self.state.tls_client.write(buf)
289283
}
290284

291285
/// Closes the connection from the client side.
292286
pub fn client_close(&mut self) -> Result<(), ProverError> {
293-
self.state
294-
.tls_client
295-
.client_close()
296-
.map_err(ProverError::from)
287+
self.state.tls_client.client_close()
297288
}
298289

299290
/// Closes the connection from the server side.
300291
pub fn server_close(&mut self) -> Result<(), ProverError> {
301-
self.state
302-
.tls_client
303-
.server_close()
304-
.map_err(ProverError::from)
292+
self.state.tls_client.server_close()
305293
}
306294

307295
/// Enables or disables the decryption of data from the server until the
@@ -311,10 +299,7 @@ impl Prover<state::Connected> {
311299
///
312300
/// * `enable` - Whether to enable or disable decryption.
313301
pub fn enable_decryption(&mut self, enable: bool) -> Result<(), ProverError> {
314-
self.state
315-
.tls_client
316-
.enable_decryption(enable)
317-
.map_err(ProverError::from)
302+
self.state.tls_client.enable_decryption(enable)
318303
}
319304

320305
/// Returns `true` if decryption of TLS traffic from the server is active.
@@ -330,9 +315,9 @@ impl Prover<state::Connected> {
330315
pub fn poll(&mut self, cx: &mut Context) -> Poll<Result<(), ProverError>> {
331316
let _ = self.state.mux_fut.poll_unpin(cx)?;
332317

333-
match self.state.tls_client.poll(cx) {
318+
match self.state.tls_client.poll(cx)? {
334319
Poll::Ready(output) => {
335-
self.state.output = Some(output?);
320+
self.state.output = Some(output);
336321
Poll::Ready(Ok(()))
337322
}
338323
Poll::Pending => Poll::Pending,

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ pub(crate) use mpc::MpcTlsClient;
1313
pub(crate) trait TlsClient {
1414
type Error: std::error::Error + Send + Sync + Unpin + 'static;
1515

16-
/// Returns `true` if the client can read TLS data from the server.
16+
/// Returns `true` if the client wants to read TLS data from the server.
1717
fn wants_read_tls(&self) -> bool;
1818

1919
/// Returns `true` if the client wants to write TLS data to the server.
@@ -25,7 +25,7 @@ pub(crate) trait TlsClient {
2525
/// Writes TLS data for the server into the provided buffer.
2626
fn write_tls(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error>;
2727

28-
/// Returns `true` if the client can read plaintext data.
28+
/// Returns `true` if the client wants to read plaintext data.
2929
fn wants_read(&self) -> bool;
3030

3131
/// Returns `true` if the client wants to write plaintext data.
@@ -46,7 +46,7 @@ pub(crate) trait TlsClient {
4646
/// Enables or disables decryption of TLS traffic sent by the server.
4747
fn enable_decryption(&mut self, enable: bool) -> Result<(), Self::Error>;
4848

49-
/// Returns if decryption of TLS traffic from the server is active.
49+
/// Returns `true` if decryption of TLS traffic from the server is active.
5050
fn is_decrypting(&self) -> bool;
5151

5252
/// Polls the client to make progress.

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

Lines changed: 9 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -126,9 +126,7 @@ impl TlsClient for MpcTlsClient {
126126
if let Some(client) = self.inner_client_mut()
127127
&& client.wants_read()
128128
{
129-
let res = client.read_tls(&mut buf).map_err(ProverError::from);
130-
131-
res
129+
client.read_tls(&mut buf).map_err(ProverError::from)
132130
} else {
133131
Ok(0)
134132
}
@@ -138,9 +136,7 @@ impl TlsClient for MpcTlsClient {
138136
if let Some(client) = self.inner_client_mut()
139137
&& client.wants_write()
140138
{
141-
let res = client.write_tls(&mut buf).map_err(ProverError::from);
142-
143-
res
139+
client.write_tls(&mut buf).map_err(ProverError::from)
144140
} else {
145141
Ok(0)
146142
}
@@ -166,9 +162,7 @@ impl TlsClient for MpcTlsClient {
166162
if let Some(client) = self.inner_client_mut()
167163
&& !client.sendable_plaintext_is_full()
168164
{
169-
let res = client.read_plaintext(buf).map_err(ProverError::from);
170-
171-
res
165+
client.read_plaintext(buf).map_err(ProverError::from)
172166
} else {
173167
Ok(0)
174168
}
@@ -178,9 +172,7 @@ impl TlsClient for MpcTlsClient {
178172
if let Some(client) = self.inner_client_mut()
179173
&& !client.plaintext_is_empty()
180174
{
181-
let res = client.write_plaintext(buf).map_err(ProverError::from);
182-
183-
res
175+
client.write_plaintext(buf).map_err(ProverError::from)
184176
} else {
185177
Ok(0)
186178
}
@@ -225,6 +217,7 @@ impl TlsClient for MpcTlsClient {
225217
fn enable_decryption(&mut self, enable: bool) -> Result<(), Self::Error> {
226218
match std::mem::replace(&mut self.state, State::Error) {
227219
State::Active { inner, mpc } => {
220+
self.decrypt = enable;
228221
self.state = State::Busy {
229222
mpc,
230223
fut: Box::pin(inner.set_decrypt(enable)),
@@ -262,9 +255,8 @@ impl TlsClient for MpcTlsClient {
262255
// yet.
263256
let _ = mpc.as_mut().poll(cx)?;
264257

265-
match fut.as_mut().poll(cx) {
266-
Poll::Ready(res) => {
267-
let inner = res?;
258+
match fut.as_mut().poll(cx)? {
259+
Poll::Ready(inner) => {
268260
self.state = State::Active { mpc, inner };
269261
}
270262
Poll::Pending => self.state = State::Busy { mpc, fut },
@@ -339,7 +331,7 @@ impl TlsClient for MpcTlsClient {
339331
}
340332
self.poll(cx)
341333
}
342-
State::Finalizing { mut fut } => match fut.as_mut().poll(cx) {
334+
State::Finalizing { mut fut } => match fut.poll_unpin(cx) {
343335
Poll::Ready(output) => {
344336
let (inner, ctx, tls_transcript) = output?;
345337
let InnerState { vm, keys, .. } = inner;
@@ -366,7 +358,7 @@ impl TlsClient for MpcTlsClient {
366358
}
367359
Poll::Pending => {
368360
self.state = State::Finalizing { fut };
369-
Poll::Pending
361+
self.poll(cx)
370362
}
371363
},
372364
State::Finished => Poll::Ready(Err(ProverError::state(

0 commit comments

Comments
 (0)