rustls-mio: fix tlsserver for mio 0.7

This commit is contained in:
Emmanuel Gil Peyrot 2020-04-28 19:34:42 +02:00 committed by ctz
parent 83c6495770
commit d5ecfb4b75
1 changed files with 32 additions and 38 deletions

View File

@ -1,7 +1,7 @@
use std::sync::Arc;
use mio;
use mio::tcp::{TcpListener, TcpStream, Shutdown};
use mio::net::{TcpListener, TcpStream};
#[macro_use]
extern crate log;
@ -75,8 +75,9 @@ impl TlsServer {
let token = mio::Token(self.next_id);
self.next_id += 1;
self.connections.insert(token, Connection::new(socket, token, mode, tls_session));
self.connections[&token].register(poll);
let mut connection = Connection::new(socket, token, mode, tls_session);
connection.register(poll);
self.connections.insert(token, connection);
true
}
Err(e) => {
@ -86,7 +87,7 @@ impl TlsServer {
}
}
fn conn_event(&mut self, poll: &mut mio::Poll, event: &mio::Event) {
fn conn_event(&mut self, poll: &mut mio::Poll, event: &mio::event::Event) {
let token = event.token();
if self.connections.contains_key(&token) {
@ -123,7 +124,7 @@ fn open_back(mode: &ServerMode) -> Option<TcpStream> {
match *mode {
ServerMode::Forward(ref port) => {
let addr = net::SocketAddrV4::new(net::Ipv4Addr::new(127, 0, 0, 1), *port);
let conn = TcpStream::connect(&net::SocketAddr::V4(addr)).unwrap();
let conn = TcpStream::connect(net::SocketAddr::V4(addr)).unwrap();
Some(conn)
}
_ => None,
@ -165,26 +166,25 @@ impl Connection {
}
/// We're a connection, and we have something to do.
fn ready(&mut self, poll: &mut mio::Poll, ev: &mio::Event) {
fn ready(&mut self, poll: &mut mio::Poll, ev: &mio::event::Event) {
// If we're readable: read some TLS. Then
// see if that yielded new plaintext. Then
// see if the backend is readable too.
if ev.readiness().is_readable() {
if ev.is_readable() {
self.do_tls_read();
self.try_plain_read();
self.try_back_read();
}
if ev.readiness().is_writable() {
if ev.is_writable() {
self.do_tls_write_and_handle_error();
}
if self.closing {
let _ = self.socket.shutdown(Shutdown::Both);
let _ = self.socket.shutdown(net::Shutdown::Both);
self.close_back();
self.closed = true;
} else {
self.reregister(poll);
self.deregister(poll);
}
}
@ -192,7 +192,7 @@ impl Connection {
fn close_back(&mut self) {
if self.back.is_some() {
let back = self.back.as_mut().unwrap();
back.shutdown(Shutdown::Both).unwrap();
back.shutdown(net::Shutdown::Both).unwrap();
}
self.back = None;
}
@ -326,50 +326,45 @@ impl Connection {
}
}
fn register(&self, poll: &mut mio::Poll) {
poll.register(&self.socket,
self.token,
self.event_set(),
mio::PollOpt::level() | mio::PollOpt::oneshot())
fn register(&mut self, poll: &mut mio::Poll) {
let registry = poll.registry();
let event_set = self.event_set();
registry.register(&mut self.socket,
self.token,
event_set)
.unwrap();
if self.back.is_some() {
poll.register(self.back.as_ref().unwrap(),
self.token,
mio::Ready::readable(),
mio::PollOpt::level() | mio::PollOpt::oneshot())
registry.register(self.back.as_mut().unwrap(),
self.token,
mio::Interest::READABLE)
.unwrap();
}
}
fn reregister(&self, poll: &mut mio::Poll) {
poll.reregister(&self.socket,
self.token,
self.event_set(),
mio::PollOpt::level() | mio::PollOpt::oneshot())
fn deregister(&mut self, poll: &mut mio::Poll) {
let registry = poll.registry();
registry.deregister(&mut self.socket)
.unwrap();
if self.back.is_some() {
poll.reregister(self.back.as_ref().unwrap(),
self.token,
mio::Ready::readable(),
mio::PollOpt::level() | mio::PollOpt::oneshot())
registry.deregister(self.back.as_mut().unwrap())
.unwrap();
}
}
/// What IO events we're currently waiting for,
/// based on wants_read/wants_write.
fn event_set(&self) -> mio::Ready {
fn event_set(&self) -> mio::Interest {
let rd = self.tls_session.wants_read();
let wr = self.tls_session.wants_write();
if rd && wr {
mio::Ready::readable() | mio::Ready::writable()
mio::Interest::READABLE | mio::Interest::WRITABLE
} else if wr {
mio::Ready::writable()
mio::Interest::WRITABLE
} else {
mio::Ready::readable()
mio::Interest::READABLE
}
}
@ -606,13 +601,12 @@ fn main() {
let config = make_config(&args);
let listener = TcpListener::bind(&addr).expect("cannot listen on port");
let mut listener = TcpListener::bind(addr).expect("cannot listen on port");
let mut poll = mio::Poll::new()
.unwrap();
poll.register(&listener,
poll.registry().register(&mut listener,
LISTENER,
mio::Ready::readable(),
mio::PollOpt::level())
mio::Interest::READABLE)
.unwrap();
let mode = if args.cmd_echo {