26250451051-1

// Copyright (C) 2013-2020 Blockstack PBC, a public benefit corporation
// Copyright (C) 2020-2023 Stacks Open Internet Foundation
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use std::collections::{HashMap, HashSet};
use std::io::ErrorKind;
use std::net::{Shutdown, SocketAddr};
use std::time::Duration;
use std::{io, time};

use mio::{self, net as mio_net, PollOpt, Ready, Token};
use rand::{self, RngCore};
use stacks_common::util::sleep_ms;

use crate::net::Error as net_error;

const SERVER: Token = mio::Token(0);

pub struct NetworkPollState {
    pub new: HashMap<usize, mio_net::TcpStream>,
    pub ready: Vec<usize>,
}

impl NetworkPollState {
    pub fn new() -> NetworkPollState {
        NetworkPollState {
            new: HashMap::new(),
            ready: vec![],
        }
    }
}

// state for a single network server
#[derive(Debug)]
pub struct NetworkServerState {
    addr: SocketAddr,
    server_socket: mio_net::TcpListener,
    server_event: mio::Token,
}

// state for the entire network
#[derive(Debug)]
pub struct NetworkState {
    poll: mio::Poll,
    events: mio::Events,
    event_capacity: usize,
    servers: Vec<NetworkServerState>,
    count: usize,
    event_map: HashMap<usize, usize>, // map socket events to their registered server socket (including server sockets)
}

impl NetworkState {
    pub fn new(event_capacity: usize) -> Result<NetworkState, net_error> {
        let poll = mio::Poll::new().map_err(|e| {
            error!("Failed to initialize poller: {:?}", e);
            net_error::BindError
        })?;

        let events = mio::Events::with_capacity(event_capacity);

        Ok(NetworkState {
            poll,
            events,
            event_capacity,
            servers: vec![],
            count: 1,
            event_map: HashMap::new(),
        })
    }

    #[cfg_attr(test, mutants::skip)]
    pub fn num_events(&self) -> usize {
        self.event_map.len()
    }

    fn bind_address(addr: &SocketAddr) -> Result<mio_net::TcpListener, net_error> {
        if !cfg!(test) && !cfg!(feature = "testing") {
            mio_net::TcpListener::bind(addr).map_err(|e| {
                error!("Failed to bind to {:?}: {:?}", addr, e);
                net_error::BindError
            })
        } else {
            let mut backoff = 1000;
            let mut rng = rand::thread_rng();
            let mut count = 1000;
            loop {
                match mio_net::TcpListener::bind(addr) {
                    Ok(server) => {
                        return Ok(server);
                    }
                    Err(e) => match e.kind() {
                        io::ErrorKind::AddrInUse => {
                            debug!(
                                "Waiting {} millis and trying to bind {:?} again",
                                backoff, addr
                            );
                            sleep_ms(backoff);
                            backoff = count + (rng.next_u64() % count);
                            count += count;
                            continue;
                        }
                        _ => {
                            debug!("Failed to bind {:?}: {:?}", addr, &e);
                            return Err(net_error::BindError);
                        }
                    },
                }
            }
        }
    }

    /// Bind to the given socket address.
    /// Returns the handle to the poll state and the bound address, used to key network poll events.
    pub fn bind(&mut self, addr: &SocketAddr) -> Result<(usize, SocketAddr), net_error> {
        let server = NetworkState::bind_address(addr)?;
        let next_server_event = self.next_event_id()?;

        self.poll
            .register(
                &server,
                mio::Token(next_server_event),
                Ready::all(),
                PollOpt::edge(),
            )
            .map_err(|e| {
                error!("Failed to register server socket: {:?}", &e);
                net_error::BindError
            })?;

        // N.B. the port for `addr` might be 0, in which case, `local_addr` may not be equal to
        // `addr` since the port will be system-assigned.  Use `local_adddr`.
        let local_addr = server.local_addr().map_err(|e| {
            error!("Failed to get local address for server: {:?}", &e);
            net_error::BindError
        })?;

        let network_server = NetworkServerState {
            addr: local_addr,
            server_socket: server,
            server_event: mio::Token(next_server_event),
        };

        assert!(
            !self.event_map.contains_key(&next_server_event),
            "BUG: failed to generate an unused server event ID"
        );

        self.servers.push(network_server);
        self.event_map.insert(next_server_event, 0); // server events always mapped to 0

        Ok((next_server_event, local_addr))
    }

    /// Register a socket for read/write notifications with this poller.
    /// Try to use the given hint_event_id value, but generate a different event ID if it's been
    /// taken.
    /// Return the actual event ID used (it may be different than hint_event_id)
    pub fn register(
        &mut self,
        server_event_id: usize,
        hint_event_id: usize,
        sock: &mio_net::TcpStream,
    ) -> Result<usize, net_error> {
        let hint_event_id = hint_event_id % (self.event_capacity + self.servers.len());
        if let Some(x) = self.event_map.get(&server_event_id) {
            if x != &0 {
                // not a server event
                error!(
                    "Server event ID {} not mapped to a server token, but to {}",
                    &server_event_id, x
                );
                return Err(net_error::RegisterError);
            }
        } else {
            // not a server event
            panic!("Not a server event ID: {}", &server_event_id);
        }

        // if the event ID is in use, then find another one
        let event_id = if self.event_map.contains_key(&hint_event_id) {
            self.next_event_id()?
        } else {
            hint_event_id
        };

        assert!(
            self.event_map.len() <= self.event_capacity + self.servers.len(),
            "BUG: event map exceeded event capacity ({} > {} + {})",
            self.event_map.len(),
            self.event_capacity,
            self.servers.len()
        );

        self.poll
            .register(sock, mio::Token(event_id), Ready::all(), PollOpt::edge())
            .map_err(|e| {
                error!(
                    "Failed to register socket on server {} event ID {} ({}): {:?}",
                    server_event_id, event_id, hint_event_id, &e
                );
                net_error::RegisterError
            })?;

        self.event_map.insert(event_id, server_event_id);

        debug!(
            "Socket registered: {}, hint {}, {:?} on server {} (Events total: {}, max: {})",
            event_id,
            hint_event_id,
            sock,
            server_event_id,
            self.event_map.len(),
            self.event_capacity
        );
        Ok(event_id)
    }

    /// Deregister a socket event
    pub fn deregister(
        &mut self,
        event_id: usize,
        sock: &mio_net::TcpStream,
    ) -> Result<(), net_error> {
        assert!(
            self.event_map.contains_key(&event_id),
            "BUG: no such socket {}",
            event_id
        );
        self.event_map.remove(&event_id);

        if let Err(e) = self.poll.deregister(sock) {
            warn!("Failed to deregister socket {}: {:?}", event_id, &e);
        };

        debug!(
            "Socket deregistered: {}, {:?} (Events total: {}, max: {})",
            event_id,
            sock,
            self.event_map.len(),
            self.event_capacity
        );

        if let Err(e) = sock.shutdown(Shutdown::Both) {
            debug!("Failed to shut down socket {}: {:?}", event_id, &e);
        }

        Ok(())
    }

    fn make_next_event_id(&self, cur_count: usize, in_use: &HashSet<usize>) -> Option<usize> {
        let mut ret = cur_count;

        let mut in_use_count = 0;
        let mut event_map_count = 0;

        for _ in 0..(self.event_capacity + self.servers.len()) {
            if self.event_map.contains_key(&ret) || in_use.contains(&ret) {
                ret = (ret + 1) % (self.event_capacity + self.servers.len());

                if in_use.contains(&ret) {
                    in_use_count += 1;
                } else {
                    event_map_count += 1;
                }
            } else {
                return Some(ret);
            }
        }

        debug!(
            "Too many peers (events: {}, in_use: {}, max: {})",
            event_map_count, in_use_count, self.event_capacity
        );
        None
    }

    /// next event ID
    pub fn next_event_id(&mut self) -> Result<usize, net_error> {
        let ret = self
            .make_next_event_id(self.count, &HashSet::new())
            .ok_or(net_error::TooManyPeers)?;
        self.count = (ret + 1) % (self.event_capacity + self.servers.len());
        Ok(ret)
    }

    /// Connect to a remote peer, but don't register it with the poll handle.
    /// The underlying connect(2) is _asynchronous_, so the caller will need to register it with a
    /// poll handle and wait for it to be connected.
    pub fn connect(
        addr: &SocketAddr,
        socket_send_buffer: u32,
        socket_recv_buffer: u32,
    ) -> Result<mio_net::TcpStream, net_error> {
        let stream = mio_net::TcpStream::connect(addr).map_err(|_e| {
            test_debug!("Failed to convert to mio stream: {:?}", &_e);
            net_error::ConnectionError
        })?;

        // set some helpful defaults
        // Don't go crazy on TIME_WAIT states; have them all die after 5 seconds
        stream
            .set_linger(Some(time::Duration::from_millis(5000)))
            .map_err(|e| {
                warn!("Failed to set SO_LINGER: {:?}", &e);
                net_error::ConnectionError
            })?;

        // Disable Nagle algorithm
        stream.set_nodelay(true).map_err(|_e| {
            warn!("Failed to set TCP_NODELAY: {:?}", &_e);
            net_error::ConnectionError
        })?;

        // Make sure keep-alive is on, since at least in p2p messages, we keep sockets around
        // for a while.  Linux default is 7200 seconds, so make sure we keep it here.
        stream
            .set_keepalive(Some(time::Duration::from_millis(7200 * 1000)))
            .map_err(|e| {
                warn!("Failed to set TCP_KEEPALIVE and/or SO_KEEPALIVE: {:?}", &e);
                net_error::ConnectionError
            })?;

        if cfg!(test) {
            if std::env::var("STACKS_TEST_DISABLE_EDGE_TRIGGER_TEST") != Ok("1".to_string()) {
                // edge-trigger torture test
                stream.set_send_buffer_size(32).unwrap();
                stream.set_recv_buffer_size(32).unwrap();
            }
        } else {
            stream
                .set_send_buffer_size(socket_send_buffer as usize)
                .map_err(|e| {
                    warn!(
                        "Failed to set socket write buffer size to {}: {:?}",
                        socket_send_buffer, &e
                    );
                    net_error::ConnectionError
                })?;

            stream
                .set_recv_buffer_size(socket_recv_buffer as usize)
                .map_err(|e| {
                    warn!(
                        "Failed to set socket read buffer size to {}: {:?}",
                        socket_send_buffer, &e
                    );
                    net_error::ConnectionError
                })?;
        }

        test_debug!("New socket connected to {:?}: {:?}", addr, &stream);
        Ok(stream)
    }

    /// Poll all server sockets.
    /// Returns a map between network server handles (returned by bind()) and their new polling state
    pub fn poll(&mut self, timeout: u64) -> Result<HashMap<usize, NetworkPollState>, net_error> {
        self.events.clear();
        self.poll
            .poll(&mut self.events, Some(Duration::from_millis(timeout)))
            .map_err(|e| {
                error!("Failed to poll: {:?}", &e);
                net_error::PollError
            })?;

        let mut poll_states = HashMap::new();
        for server in self.servers.iter() {
            // pre-populate with server tokens
            let server_event_id = usize::from(server.server_event);
            poll_states.insert(server_event_id, NetworkPollState::new());
        }

        let mut new_events = HashSet::new();

        for event in &self.events {
            let token = event.token();
            let mut is_server_event = false;

            for server in self.servers.iter() {
                // server token?
                if token == server.server_event {
                    // new inbound connection(s)
                    let poll_state =
                        poll_states.get_mut(&usize::from(token)).unwrap_or_else(|| {
                            panic!(
                                "BUG: FATAL: no poll state registered for server {}",
                                usize::from(token)
                            )
                        });

                    loop {
                        let (client_sock, client_addr) = match server.server_socket.accept() {
                            Ok((client_sock, client_addr)) => (client_sock, client_addr),
                            Err(e) => match e.kind() {
                                ErrorKind::WouldBlock => {
                                    break;
                                }
                                _ => {
                                    error!("Network error: {}", e);
                                    return Err(net_error::AcceptError);
                                }
                            },
                        };

                        // this does the same thing as next_event_id(), but we can't borrow self
                        // mutably here (so we'll just do the increment-mod directly).
                        let next_event_id = match self.make_next_event_id(self.count, &new_events) {
                            Some(eid) => eid,
                            None => {
                                // no poll slots available. Close the socket and carry on.
                                info!("Too many peers on {:?}, closing {:?} (events: {}, in-flight: {}, capacity: {})", &server.server_socket, &client_sock, self.event_map.len(), new_events.len(), self.event_capacity);
                                let _ = client_sock.shutdown(Shutdown::Both);
                                continue;
                            }
                        };

                        self.count =
                            (next_event_id + 1) % (self.event_capacity + self.servers.len());

                        new_events.insert(next_event_id);

                        debug!(
                            "New socket event: {}, {:?} addr={:?} (Events total: {}, max: {}) on server {:?}",
                            next_event_id,
                            &client_sock,
                            &client_addr,
                            self.event_map.len(),
                            self.event_capacity,
                            &server.server_socket
                        );

                        poll_state.new.insert(next_event_id, client_sock);
                    }

                    is_server_event = true;
                    break;
                }
            }

            if is_server_event {
                continue;
            }

            // event for a client of one of our servers.  which one?
            let event_id = usize::from(token);
            match self.event_map.get(&event_id) {
                Some(server_event_id) => {
                    if let Some(poll_state) = poll_states.get_mut(server_event_id) {
                        test_debug!(
                            "Wakeup socket event {} on server {}",
                            event_id,
                            server_event_id
                        );
                        poll_state.ready.push(event_id);
                    } else {
                        warn!("Unknown server event ID {}", server_event_id);
                    }
                }
                None => {
                    warn!("Surreptitious readiness event {}", event_id);
                }
            }
        }

        Ok(poll_states)
    }
}

#[cfg(test)]
mod test {
    use std::collections::HashSet;

    use super::*;

    #[test]
    fn test_bind() {
        let mut ns = NetworkState::new(100).unwrap();
        let mut server_events = HashSet::new();
        for _ in 0..10 {
            let addr = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
            let (event_id, _local_addr) = ns.bind(&addr).unwrap();
            assert!(!server_events.contains(&event_id));
            server_events.insert(event_id);
        }
    }

    #[test]
    fn test_register_deregister() {
        let mut ns = NetworkState::new(100).unwrap();
        let mut server_events = vec![];
        let mut event_ids = HashSet::new();
        let mut ports = vec![];
        for _ in 0..10 {
            let addr = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
            let (event_id, local_addr) = ns.bind(&addr).unwrap();
            server_events.push(event_id);
            event_ids.insert(event_id);

            ports.push(local_addr.port());
        }

        let mut client_events = vec![];
        for (i, port) in ports.iter().enumerate() {
            let addr = format!("127.0.0.1:{}", &port)
                .parse::<SocketAddr>()
                .unwrap();
            let sock = NetworkState::connect(&addr, 4096, 4096).unwrap();

            let event_id = ns.register(server_events[i], 1, &sock).unwrap();
            assert!(event_id != 0);
            assert!(!event_ids.contains(&event_id));
            ns.deregister(event_id, &sock).unwrap();

            let event_id = ns.register(server_events[i], 101, &sock).unwrap();
            assert!(event_id != 0);
            assert!(!event_ids.contains(&event_id));
            ns.deregister(event_id, &sock).unwrap();

            let event_id = ns
                .register(server_events[i], server_events[i], &sock)
                .unwrap();
            assert!(event_id != 0);
            assert!(!event_ids.contains(&event_id));
            ns.deregister(event_id, &sock).unwrap();

            let event_id = ns.register(server_events[i], 11, &sock).unwrap();
            assert!(!event_ids.contains(&event_id));

            event_ids.insert(event_id);
            client_events.push(event_id);
        }

        test_debug!("=====");
        for (i, port) in ports.iter().enumerate() {
            let addr = format!("127.0.0.1:{}", &port)
                .parse::<SocketAddr>()
                .unwrap();
            let sock = NetworkState::connect(&addr, 4096, 4096).unwrap();

            // can't use non-server events
            assert_eq!(
                Err(net_error::RegisterError),
                ns.register(client_events[i], i + 1, &sock)
            );
        }
    }

    #[test]
    fn test_register_too_many_peers() {
        let mut ns = NetworkState::new(10).unwrap();
        let mut event_ids = HashSet::new();
        let addr = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
        let (server_event_id, local_addr) = ns.bind(&addr).unwrap();
        let port = local_addr.port();
        for _ in 0..10 {
            let addr = format!("127.0.0.1:{}", &port)
                .parse::<SocketAddr>()
                .unwrap();
            event_ids.insert(server_event_id);

            let sock = NetworkState::connect(&addr, 4096, 4096).unwrap();

            // register 10 client events
            let event_id = ns.register(server_event_id, 11, &sock).unwrap();
            assert!(!event_ids.contains(&event_id));
        }

        // the 11th socket should fail
        let addr = format!("127.0.0.1:{}", port).parse::<SocketAddr>().unwrap();
        let sock = NetworkState::connect(&addr, 4096, 4096).unwrap();
        let res = ns.register(server_event_id, 11, &sock);
        assert_eq!(Err(net_error::TooManyPeers), res);
    }

    #[test]
    fn test_register_deregister_stress() {
        let mut ns = NetworkState::new(20).unwrap();
        let count = 0;
        let mut in_use = HashSet::new();
        let mut events_in = vec![];

        for _ in 0..20 {
            let next_eid = ns.make_next_event_id(count, &in_use).unwrap();
            ns.event_map.insert(next_eid, 0);
            events_in.push(next_eid);
        }

        assert_eq!(ns.event_map.len(), 20);

        for _ in 0..20 {
            assert!(ns.make_next_event_id(count, &in_use).is_none());
        }

        for eid in events_in.iter() {
            ns.event_map.remove(eid);
        }

        events_in.clear();

        for _ in 0..20 {
            let next_eid = ns.make_next_event_id(count, &in_use).unwrap();
            events_in.push(next_eid);
            in_use.insert(next_eid);
        }

        assert!(ns.event_map.is_empty());

        for _ in 0..20 {
            assert!(ns.make_next_event_id(count, &in_use).is_none());
        }
    }
}

1	// Copyright (C) 2013-2020 Blockstack PBC, a public benefit corporation
2	// Copyright (C) 2020-2023 Stacks Open Internet Foundation
3	//
4	// This program is free software: you can redistribute it and/or modify
5	// it under the terms of the GNU General Public License as published by
6	// the Free Software Foundation, either version 3 of the License, or
7	// (at your option) any later version.
8	//
9	// This program is distributed in the hope that it will be useful,
10	// but WITHOUT ANY WARRANTY; without even the implied warranty of
11	// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12	// GNU General Public License for more details.
13	//
14	// You should have received a copy of the GNU General Public License
15	// along with this program. If not, see <http://www.gnu.org/licenses/>.
16
17	use std::collections::{HashMap, HashSet};
18	use std::io::ErrorKind;
19	use std::net::{Shutdown, SocketAddr};
20	use std::time::Duration;
21	use std::{io, time};
22
23	use mio::{self, net as mio_net, PollOpt, Ready, Token};
24	use rand::{self, RngCore};
25	use stacks_common::util::sleep_ms;
26
27	use crate::net::Error as net_error;
28
29	const SERVER: Token = mio::Token(0);
30
31	pub struct NetworkPollState {
32	pub new: HashMap<usize, mio_net::TcpStream>,
33	pub ready: Vec<usize>,
34	}
35
36	impl NetworkPollState {
37	pub fn new() -> NetworkPollState {	20,770,988✔
38	NetworkPollState {	20,770,988✔
39	new: HashMap::new(),	20,770,988✔
40	ready: vec![],	20,770,988✔
41	}	20,770,988✔
42	}	20,770,988✔
43	}
44
45	// state for a single network server
46	#[derive(Debug)]
47	pub struct NetworkServerState {
48	addr: SocketAddr,
49	server_socket: mio_net::TcpListener,
50	server_event: mio::Token,
51	}
52
53	// state for the entire network
54	#[derive(Debug)]
55	pub struct NetworkState {
56	poll: mio::Poll,
57	events: mio::Events,
58	event_capacity: usize,
59	servers: Vec<NetworkServerState>,
60	count: usize,
61	event_map: HashMap<usize, usize>, // map socket events to their registered server socket (including server sockets)
62	}
63
64	impl NetworkState {
65	pub fn new(event_capacity: usize) -> Result<NetworkState, net_error> {	3,313✔
66	let poll = mio::Poll::new().map_err(\|e\| {	3,313✔
67	error!("Failed to initialize poller: {:?}", e);	×
68	net_error::BindError	×
69	})?;	×
70
71	let events = mio::Events::with_capacity(event_capacity);	3,313✔
72
73	Ok(NetworkState {	3,313✔
74	poll,	3,313✔
75	events,	3,313✔
76	event_capacity,	3,313✔
77	servers: vec![],	3,313✔
78	count: 1,	3,313✔
79	event_map: HashMap::new(),	3,313✔
80	})	3,313✔
81	}	3,313✔
82
83	#[cfg_attr(test, mutants::skip)]
84	pub fn num_events(&self) -> usize {	51✔
85	self.event_map.len()	51✔
86	}	51✔
87
88	fn bind_address(addr: &SocketAddr) -> Result<mio_net::TcpListener, net_error> {	6,639✔
89	if !cfg!(test) && !cfg!(feature = "testing") {	6,639✔
90	mio_net::TcpListener::bind(addr).map_err(\|e\| {	×
91	error!("Failed to bind to {:?}: {:?}", addr, e);	×
92	net_error::BindError	×
93	})	×
94	} else {
95	let mut backoff = 1000;	6,639✔
96	let mut rng = rand::thread_rng();	6,639✔
97	let mut count = 1000;	6,639✔
98	loop {
99	match mio_net::TcpListener::bind(addr) {	6,748✔
100	Ok(server) => {	6,639✔
101	return Ok(server);	6,639✔
102	}
103	Err(e) => match e.kind() {	109✔
104	io::ErrorKind::AddrInUse => {
105	debug!(	109✔
106	"Waiting {} millis and trying to bind {:?} again",
107	backoff, addr
108	);
109	sleep_ms(backoff);	109✔
110	backoff = count + (rng.next_u64() % count);	109✔
111	count += count;	109✔
112	continue;	109✔
113	}
114	_ => {
115	debug!("Failed to bind {:?}: {:?}", addr, &e);	×
116	return Err(net_error::BindError);	×
117	}
118	},
119	}
120	}
121	}
122	}	6,639✔
123
124	/// Bind to the given socket address.
125	/// Returns the handle to the poll state and the bound address, used to key network poll events.
126	pub fn bind(&mut self, addr: &SocketAddr) -> Result<(usize, SocketAddr), net_error> {	6,639✔
127	let server = NetworkState::bind_address(addr)?;	6,639✔
128	let next_server_event = self.next_event_id()?;	6,639✔
129
130	self.poll	6,639✔
131	.register(	6,639✔
132	&server,	6,639✔
133	mio::Token(next_server_event),	6,639✔
134	Ready::all(),	6,639✔
135	PollOpt::edge(),	6,639✔
136	)
137	.map_err(\|e\| {	6,639✔
138	error!("Failed to register server socket: {:?}", &e);	×
139	net_error::BindError	×
140	})?;	×
141
142	// N.B. the port for `addr` might be 0, in which case, `local_addr` may not be equal to
143	// `addr` since the port will be system-assigned. Use `local_adddr`.
144	let local_addr = server.local_addr().map_err(\|e\| {	6,639✔
145	error!("Failed to get local address for server: {:?}", &e);	×
146	net_error::BindError	×
147	})?;	×
148
149	let network_server = NetworkServerState {	6,639✔
150	addr: local_addr,	6,639✔
151	server_socket: server,	6,639✔
152	server_event: mio::Token(next_server_event),	6,639✔
153	};	6,639✔
154
155	assert!(	6,639✔
156	!self.event_map.contains_key(&next_server_event),	6,639✔
157	"BUG: failed to generate an unused server event ID"
158	);
159
160	self.servers.push(network_server);	6,639✔
161	self.event_map.insert(next_server_event, 0); // server events always mapped to 0	6,639✔
162
163	Ok((next_server_event, local_addr))	6,639✔
164	}	6,639✔
165
166	/// Register a socket for read/write notifications with this poller.
167	/// Try to use the given hint_event_id value, but generate a different event ID if it's been
168	/// taken.
169	/// Return the actual event ID used (it may be different than hint_event_id)
170	pub fn register(	1,124,913✔
171	&mut self,	1,124,913✔
172	server_event_id: usize,	1,124,913✔
173	hint_event_id: usize,	1,124,913✔
174	sock: &mio_net::TcpStream,	1,124,913✔
175	) -> Result<usize, net_error> {	1,124,913✔
176	let hint_event_id = hint_event_id % (self.event_capacity + self.servers.len());	1,124,913✔
177	if let Some(x) = self.event_map.get(&server_event_id) {	1,124,913✔
178	if x != &0 {	1,124,913✔
179	// not a server event
180	error!(	10✔
181	"Server event ID {} not mapped to a server token, but to {}",
182	&server_event_id, x	10✔
183	);
184	return Err(net_error::RegisterError);	10✔
185	}	1,124,903✔
186	} else {
187	// not a server event
188	panic!("Not a server event ID: {}", &server_event_id);	×
189	}
190
191	// if the event ID is in use, then find another one
192	let event_id = if self.event_map.contains_key(&hint_event_id) {	1,124,903✔
193	self.next_event_id()?	39✔
194	} else {
195	hint_event_id	1,124,864✔
196	};
197
198	assert!(	1,124,902✔
199	self.event_map.len() <= self.event_capacity + self.servers.len(),	1,124,902✔
200	"BUG: event map exceeded event capacity ({} > {} + {})",
201	self.event_map.len(),	×
202	self.event_capacity,
203	self.servers.len()	×
204	);
205
206	self.poll	1,124,902✔
207	.register(sock, mio::Token(event_id), Ready::all(), PollOpt::edge())	1,124,902✔
208	.map_err(\|e\| {	1,124,902✔
209	error!(	×
210	"Failed to register socket on server {} event ID {} ({}): {:?}",
211	server_event_id, event_id, hint_event_id, &e	×
212	);
213	net_error::RegisterError	×
214	})?;	×
215
216	self.event_map.insert(event_id, server_event_id);	1,124,902✔
217
218	debug!(	1,124,902✔
219	"Socket registered: {}, hint {}, {:?} on server {} (Events total: {}, max: {})",
220	event_id,
221	hint_event_id,
222	sock,
223	server_event_id,
224	self.event_map.len(),	×
225	self.event_capacity
226	);
227	Ok(event_id)	1,124,902✔
228	}	1,124,913✔
229
230	/// Deregister a socket event
231	pub fn deregister(	1,111,023✔
232	&mut self,	1,111,023✔
233	event_id: usize,	1,111,023✔
234	sock: &mio_net::TcpStream,	1,111,023✔
235	) -> Result<(), net_error> {	1,111,023✔
236	assert!(	1,111,023✔
237	self.event_map.contains_key(&event_id),	1,111,023✔
238	"BUG: no such socket {}",
239	event_id
240	);
241	self.event_map.remove(&event_id);	1,111,023✔
242
243	if let Err(e) = self.poll.deregister(sock) {	1,111,023✔
244	warn!("Failed to deregister socket {}: {:?}", event_id, &e);	×
245	};	1,111,023✔
246
247	debug!(	1,111,023✔
248	"Socket deregistered: {}, {:?} (Events total: {}, max: {})",
249	event_id,
250	sock,
251	self.event_map.len(),	×
252	self.event_capacity
253	);
254
255	if let Err(e) = sock.shutdown(Shutdown::Both) {	1,111,023✔
256	debug!("Failed to shut down socket {}: {:?}", event_id, &e);	1,322✔
257	}	1,109,701✔
258
259	Ok(())	1,111,023✔
260	}	1,111,023✔
261
262	fn make_next_event_id(&self, cur_count: usize, in_use: &HashSet<usize>) -> Option<usize> {	1,131,610✔
263	let mut ret = cur_count;	1,131,610✔
264
265	let mut in_use_count = 0;	1,131,610✔
266	let mut event_map_count = 0;	1,131,610✔
267
268	for _ in 0..(self.event_capacity + self.servers.len()) {	1,131,610✔
269	if self.event_map.contains_key(&ret) \|\| in_use.contains(&ret) {	1,134,506✔
270	ret = (ret + 1) % (self.event_capacity + self.servers.len());	2,937✔
271
272	if in_use.contains(&ret) {	2,937✔
273	in_use_count += 1;	571✔
274	} else {	2,366✔
275	event_map_count += 1;	2,366✔
276	}	2,366✔
277	} else {
278	return Some(ret);	1,131,569✔
279	}
280	}
281
282	debug!(	41✔
283	"Too many peers (events: {}, in_use: {}, max: {})",
284	event_map_count, in_use_count, self.event_capacity
285	);
286	None	41✔
287	}	1,131,610✔
288
289	/// next event ID
290	pub fn next_event_id(&mut self) -> Result<usize, net_error> {	25,380✔
291	let ret = self	25,380✔
292	.make_next_event_id(self.count, &HashSet::new())	25,380✔
293	.ok_or(net_error::TooManyPeers)?;	25,380✔
294	self.count = (ret + 1) % (self.event_capacity + self.servers.len());	25,379✔
295	Ok(ret)	25,379✔
296	}	25,380✔
297
298	/// Connect to a remote peer, but don't register it with the poll handle.
299	/// The underlying connect(2) is _asynchronous_, so the caller will need to register it with a
300	/// poll handle and wait for it to be connected.
301	pub fn connect(	18,733✔
302	addr: &SocketAddr,	18,733✔
303	socket_send_buffer: u32,	18,733✔
304	socket_recv_buffer: u32,	18,733✔
305	) -> Result<mio_net::TcpStream, net_error> {	18,733✔
306	let stream = mio_net::TcpStream::connect(addr).map_err(\|_e\| {	18,733✔
307	test_debug!("Failed to convert to mio stream: {:?}", &_e);	×
308	net_error::ConnectionError	×
309	})?;	×
310
311	// set some helpful defaults
312	// Don't go crazy on TIME_WAIT states; have them all die after 5 seconds
313	stream	18,733✔
314	.set_linger(Some(time::Duration::from_millis(5000)))	18,733✔
315	.map_err(\|e\| {	18,733✔
316	warn!("Failed to set SO_LINGER: {:?}", &e);	×
317	net_error::ConnectionError	×
318	})?;	×
319
320	// Disable Nagle algorithm
321	stream.set_nodelay(true).map_err(\|_e\| {	18,733✔
322	warn!("Failed to set TCP_NODELAY: {:?}", &_e);	×
323	net_error::ConnectionError	×
324	})?;	×
325
326	// Make sure keep-alive is on, since at least in p2p messages, we keep sockets around
327	// for a while. Linux default is 7200 seconds, so make sure we keep it here.
328	stream	18,733✔
329	.set_keepalive(Some(time::Duration::from_millis(7200 * 1000)))	18,733✔
330	.map_err(\|e\| {	18,733✔
331	warn!("Failed to set TCP_KEEPALIVE and/or SO_KEEPALIVE: {:?}", &e);	×
332	net_error::ConnectionError	×
333	})?;	×
334
335	if cfg!(test) {	18,733✔
336	if std::env::var("STACKS_TEST_DISABLE_EDGE_TRIGGER_TEST") != Ok("1".to_string()) {	10,768✔
337	// edge-trigger torture test	10,743✔
338	stream.set_send_buffer_size(32).unwrap();	10,743✔
339	stream.set_recv_buffer_size(32).unwrap();	10,743✔
340	}	10,743✔
341	} else {
342	stream	7,965✔
343	.set_send_buffer_size(socket_send_buffer as usize)	7,965✔
344	.map_err(\|e\| {	7,965✔
345	warn!(	×
346	"Failed to set socket write buffer size to {}: {:?}",
347	socket_send_buffer, &e	×
348	);
349	net_error::ConnectionError	×
350	})?;	×
351
352	stream	7,965✔
353	.set_recv_buffer_size(socket_recv_buffer as usize)	7,965✔
354	.map_err(\|e\| {	7,965✔
355	warn!(	×
356	"Failed to set socket read buffer size to {}: {:?}",
357	socket_send_buffer, &e	×
358	);
359	net_error::ConnectionError	×
360	})?;	×
361	}
362
363	test_debug!("New socket connected to {:?}: {:?}", addr, &stream);	18,733✔
364	Ok(stream)	18,733✔
365	}	18,733✔
366
367	/// Poll all server sockets.
368	/// Returns a map between network server handles (returned by bind()) and their new polling state
369	pub fn poll(&mut self, timeout: u64) -> Result<HashMap<usize, NetworkPollState>, net_error> {	10,386,016✔
370	self.events.clear();	10,386,016✔
371	self.poll	10,386,016✔
372	.poll(&mut self.events, Some(Duration::from_millis(timeout)))	10,386,016✔
373	.map_err(\|e\| {	10,386,016✔
374	error!("Failed to poll: {:?}", &e);	×
375	net_error::PollError	×
376	})?;	×
377
378	let mut poll_states = HashMap::new();	10,386,016✔
379	for server in self.servers.iter() {	3✔
380	// pre-populate with server tokens	3✔
381	let server_event_id = usize::from(server.server_event);	3✔
382	poll_states.insert(server_event_id, NetworkPollState::new());	3✔
383	}	3✔
384
385	let mut new_events = HashSet::new();	10,386,016✔
386
387	for event in &self.events {	13,093,762✔
388	let token = event.token();	12,042,179✔
389	let mut is_server_event = false;	12,042,179✔
390
391	for server in self.servers.iter() {	24,074,120✔
392	// server token?
393	if token == server.server_event {	24,074,120✔
394	// new inbound connection(s)
395	let poll_state =	801,895✔
396	poll_states.get_mut(&usize::from(token)).unwrap_or_else(\|\| {	801,895✔
397	panic!(	×
398	"BUG: FATAL: no poll state registered for server {}",
399	usize::from(token)	×
400	)
401	});
402
403	loop {
404	let (client_sock, client_addr) = match server.server_socket.accept() {	1,908,045✔
405	Ok((client_sock, client_addr)) => (client_sock, client_addr),	1,106,150✔
406	Err(e) => match e.kind() {	801,895✔
407	ErrorKind::WouldBlock => {
408	break;	1✔
409	}
410	_ => {
411	error!("Network error: {}", e);	522✔
412	return Err(net_error::AcceptError);	522✔
413	}
414	},
415	};
416
417	// this does the same thing as next_event_id(), but we can't borrow self
418	// mutably here (so we'll just do the increment-mod directly).
419	let next_event_id = match self.make_next_event_id(self.count, &new_events) {	1,106,150✔
420	Some(eid) => eid,	1,106,150✔
421	None => {
422	// no poll slots available. Close the socket and carry on.
423	info!("Too many peers on {:?}, closing {:?} (events: {}, in-flight: {}, capacity: {})", &server.server_socket, &client_sock, self.event_map.len(), new_events.len(), self.event_capacity);	×
424	let _ = client_sock.shutdown(Shutdown::Both);	×
425	continue;	×
426	}
427	};
428
429	self.count =	1,106,150✔
430	(next_event_id + 1) % (self.event_capacity + self.servers.len());	1,106,150✔
431
432	new_events.insert(next_event_id);	1,106,150✔
433
434	debug!(	1,106,150✔
435	"New socket event: {}, {:?} addr={:?} (Events total: {}, max: {}) on server {:?}",
436	next_event_id,
437	&client_sock,	×
438	&client_addr,	×
439	self.event_map.len(),	×
440	self.event_capacity,
441	&server.server_socket	×
442	);
443
444	poll_state.new.insert(next_event_id, client_sock);	1,106,150✔
445	}
446
447	is_server_event = true;	1✔
448	break;	1✔
449	}	23,272,225✔
450	}
451
452	if is_server_event {	1✔
453	continue;	1✔
454	}	11,240,284✔
455
456	// event for a client of one of our servers. which one?
457	let event_id = usize::from(token);	11,240,284✔
458	match self.event_map.get(&event_id) {	11,240,284✔
459	Some(server_event_id) => {	11,240,284✔
460	if let Some(poll_state) = poll_states.get_mut(server_event_id) {	11,240,284✔
461	test_debug!(	11,240,284✔
462	"Wakeup socket event {} on server {}",
463	event_id,
464	server_event_id
465	);
466	poll_state.ready.push(event_id);	11,240,284✔
467	} else {
468	warn!("Unknown server event ID {}", server_event_id);	×
469	}
470	}
471	None => {
472	warn!("Surreptitious readiness event {}", event_id);	×
473	}
474	}
475	}
476
477	Ok(poll_states)	10,385,494✔
478	}	10,386,016✔
479	}
480
481	#[cfg(test)]
482	mod test {
483	use std::collections::HashSet;
484
485	use super::*;
486
487	#[test]
488	fn test_bind() {	1✔
489	let mut ns = NetworkState::new(100).unwrap();	1✔
490	let mut server_events = HashSet::new();	1✔
491	for _ in 0..10 {	1✔
492	let addr = "127.0.0.1:0".parse::<SocketAddr>().unwrap();	10✔
493	let (event_id, _local_addr) = ns.bind(&addr).unwrap();	10✔
494	assert!(!server_events.contains(&event_id));	10✔
495	server_events.insert(event_id);	10✔
496	}
497	}	1✔
498
499	#[test]
500	fn test_register_deregister() {	1✔
501	let mut ns = NetworkState::new(100).unwrap();	1✔
502	let mut server_events = vec![];	1✔
503	let mut event_ids = HashSet::new();	1✔
504	let mut ports = vec![];	1✔
505	for _ in 0..10 {	10✔
506	let addr = "127.0.0.1:0".parse::<SocketAddr>().unwrap();	10✔
507	let (event_id, local_addr) = ns.bind(&addr).unwrap();	10✔
508	server_events.push(event_id);	10✔
509	event_ids.insert(event_id);	10✔
510		10✔
511	ports.push(local_addr.port());	10✔
512	}	10✔
513
514	let mut client_events = vec![];	1✔
515	for (i, port) in ports.iter().enumerate() {	10✔
516	let addr = format!("127.0.0.1:{}", &port)	10✔
517	.parse::<SocketAddr>()	10✔
518	.unwrap();	10✔
519	let sock = NetworkState::connect(&addr, 4096, 4096).unwrap();	10✔
520
521	let event_id = ns.register(server_events[i], 1, &sock).unwrap();	10✔
522	assert!(event_id != 0);	10✔
523	assert!(!event_ids.contains(&event_id));	10✔
524	ns.deregister(event_id, &sock).unwrap();	10✔
525
526	let event_id = ns.register(server_events[i], 101, &sock).unwrap();	10✔
527	assert!(event_id != 0);	10✔
528	assert!(!event_ids.contains(&event_id));	10✔
529	ns.deregister(event_id, &sock).unwrap();	10✔
530
531	let event_id = ns	10✔
532	.register(server_events[i], server_events[i], &sock)	10✔
533	.unwrap();	10✔
534	assert!(event_id != 0);	10✔
535	assert!(!event_ids.contains(&event_id));	10✔
536	ns.deregister(event_id, &sock).unwrap();	10✔
537
538	let event_id = ns.register(server_events[i], 11, &sock).unwrap();	10✔
539	assert!(!event_ids.contains(&event_id));	10✔
540
541	event_ids.insert(event_id);	10✔
542	client_events.push(event_id);	10✔
543	}
544
545	test_debug!("=====");	1✔
546	for (i, port) in ports.iter().enumerate() {	10✔
547	let addr = format!("127.0.0.1:{}", &port)	10✔
548	.parse::<SocketAddr>()	10✔
549	.unwrap();	10✔
550	let sock = NetworkState::connect(&addr, 4096, 4096).unwrap();	10✔
551
552	// can't use non-server events
553	assert_eq!(	10✔
554	Err(net_error::RegisterError),
555	ns.register(client_events[i], i + 1, &sock)	10✔
556	);
557	}
558	}	1✔
559
560	#[test]
561	fn test_register_too_many_peers() {	1✔
562	let mut ns = NetworkState::new(10).unwrap();	1✔
563	let mut event_ids = HashSet::new();	1✔
564	let addr = "127.0.0.1:0".parse::<SocketAddr>().unwrap();	1✔
565	let (server_event_id, local_addr) = ns.bind(&addr).unwrap();	1✔
566	let port = local_addr.port();	1✔
567	for _ in 0..10 {	1✔
568	let addr = format!("127.0.0.1:{}", &port)	10✔
569	.parse::<SocketAddr>()	10✔
570	.unwrap();	10✔
571	event_ids.insert(server_event_id);	10✔
572
573	let sock = NetworkState::connect(&addr, 4096, 4096).unwrap();	10✔
574
575	// register 10 client events
576	let event_id = ns.register(server_event_id, 11, &sock).unwrap();	10✔
577	assert!(!event_ids.contains(&event_id));	10✔
578	}
579
580	// the 11th socket should fail
581	let addr = format!("127.0.0.1:{}", port).parse::<SocketAddr>().unwrap();	1✔
582	let sock = NetworkState::connect(&addr, 4096, 4096).unwrap();	1✔
583	let res = ns.register(server_event_id, 11, &sock);	1✔
584	assert_eq!(Err(net_error::TooManyPeers), res);	1✔
585	}	1✔
586
587	#[test]
588	fn test_register_deregister_stress() {	1✔
589	let mut ns = NetworkState::new(20).unwrap();	1✔
590	let count = 0;	1✔
591	let mut in_use = HashSet::new();	1✔
592	let mut events_in = vec![];	1✔
593
594	for _ in 0..20 {	20✔
595	let next_eid = ns.make_next_event_id(count, &in_use).unwrap();	20✔
596	ns.event_map.insert(next_eid, 0);	20✔
597	events_in.push(next_eid);	20✔
598	}	20✔
599
600	assert_eq!(ns.event_map.len(), 20);	1✔
601
602	for _ in 0..20 {	1✔
603	assert!(ns.make_next_event_id(count, &in_use).is_none());	20✔
604	}
605
606	for eid in events_in.iter() {	20✔
607	ns.event_map.remove(eid);	20✔
608	}	20✔
609
610	events_in.clear();	1✔
611
612	for _ in 0..20 {	20✔
613	let next_eid = ns.make_next_event_id(count, &in_use).unwrap();	20✔
614	events_in.push(next_eid);	20✔
615	in_use.insert(next_eid);	20✔
616	}	20✔
617
618	assert!(ns.event_map.is_empty());	1✔
619
620	for _ in 0..20 {	1✔
621	assert!(ns.make_next_event_id(count, &in_use).is_none());	20✔
622	}
623	}	1✔
624	}

stacks-network / stacks-core / 26250451051-1

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous