Skip to main content

rustls_util/
lib.rs

1use std::io;
2
3mod key_log_file;
4pub use key_log_file::KeyLogFile;
5use rustls::{Connection, VecInput};
6
7mod stream;
8pub use crate::stream::{Stream, StreamOwned};
9
10/// This function uses `io` to complete any outstanding IO for
11/// the connection.
12///
13/// This is a convenience function which solely uses other parts
14/// of the public API.
15///
16/// What this means depends on the connection  state:
17///
18/// - If the connection [`is_handshaking()`], then IO is performed until
19///   the handshake is complete.
20/// - Otherwise, if [`wants_write()`] is true, [`write_tls()`] is invoked
21///   until it is all written.
22/// - Otherwise, if [`wants_read()`] is true, [`VecInput::read()`] is invoked
23///   once.
24///
25/// The return value is the number of bytes read from and written
26/// to `io`, respectively. Once both `read()` and `write()` yield `WouldBlock`,
27/// this function will propagate the error.
28///
29/// Errors from TLS record handling (i.e., from [`process_new_packets()`])
30/// are wrapped in an `io::ErrorKind::InvalidData`-kind error.
31///
32/// [`is_handshaking()`]: rustls::CommonState::is_handshaking
33/// [`wants_read()`]: rustls::Connection::wants_read
34/// [`wants_write()`]: rustls::Connection::wants_write
35/// [`write_tls()`]: rustls::Connection::write_tls
36/// [`process_new_packets()`]: rustls::Connection::process_new_packets
37pub fn complete_io(
38    io: &mut (impl io::Read + io::Write),
39    input: &mut VecInput,
40    conn: &mut dyn Connection,
41) -> Result<(usize, usize), io::Error> {
42    let mut eof = false;
43    let mut wrlen = 0;
44    let mut rdlen = 0;
45    loop {
46        let (mut blocked_write, mut blocked_read) = (None, None);
47        let until_handshaked = conn.is_handshaking();
48
49        if !conn.wants_write() && !conn.wants_read() {
50            // We will make no further progress.
51            return Ok((rdlen, wrlen));
52        }
53
54        while conn.wants_write() {
55            match conn.write_tls(io) {
56                Ok(0) => {
57                    io.flush()?;
58                    return Ok((rdlen, wrlen)); // EOF.
59                }
60                Ok(n) => wrlen += n,
61                Err(err) if err.kind() == io::ErrorKind::WouldBlock => {
62                    blocked_write = Some(err);
63                    break;
64                }
65                Err(err) => return Err(err),
66            }
67        }
68        if wrlen > 0 {
69            io.flush()?;
70        }
71
72        if !until_handshaked && wrlen > 0 {
73            return Ok((rdlen, wrlen));
74        }
75
76        // If we want to write, but are WouldBlocked by the underlying IO, *and*
77        // have no desire to read; that is everything.
78        if let (Some(_), false) = (&blocked_write, conn.wants_read()) {
79            return match wrlen {
80                0 => Err(blocked_write.unwrap()),
81                _ => Ok((rdlen, wrlen)),
82            };
83        }
84
85        while !eof && conn.wants_read() {
86            let read_size = match input.read(io) {
87                Ok(0) => {
88                    eof = true;
89                    Some(0)
90                }
91                Ok(n) => {
92                    rdlen += n;
93                    Some(n)
94                }
95                Err(err) if err.kind() == io::ErrorKind::WouldBlock => {
96                    blocked_read = Some(err);
97                    break;
98                }
99                Err(err) if err.kind() == io::ErrorKind::Interrupted => None, // nothing to do
100                Err(err) => return Err(err),
101            };
102            if read_size.is_some() {
103                break;
104            }
105        }
106
107        if let Err(e) = conn.process_new_packets(input) {
108            // In case we have an alert to send describing this error, try a last-gasp
109            // write -- but don't predate the primary error.
110            let _ignored = conn.write_tls(io);
111            let _ignored = io.flush();
112            return Err(io::Error::new(io::ErrorKind::InvalidData, e));
113        };
114
115        // If we want to read, but are WouldBlocked by the underlying IO, *and*
116        // have no desire to write; that is everything.
117        if let (Some(_), false) = (&blocked_read, conn.wants_write()) {
118            return match rdlen {
119                0 => Err(blocked_read.unwrap()),
120                _ => Ok((rdlen, wrlen)),
121            };
122        }
123
124        // if we're doing IO until handshaked, and we believe we've finished handshaking,
125        // but process_new_packets() has queued TLS data to send, loop around again to write
126        // the queued messages.
127        if until_handshaked && !conn.is_handshaking() && conn.wants_write() {
128            continue;
129        }
130
131        let blocked = blocked_write.zip(blocked_read);
132        match (eof, until_handshaked, conn.is_handshaking(), blocked) {
133            (_, true, false, _) => return Ok((rdlen, wrlen)),
134            (_, _, _, Some((e, _))) if rdlen == 0 && wrlen == 0 => return Err(e),
135            (_, false, _, _) => return Ok((rdlen, wrlen)),
136            (true, true, true, _) => return Err(io::Error::from(io::ErrorKind::UnexpectedEof)),
137            _ => {}
138        }
139    }
140}