You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
321 lines
8.5 KiB
321 lines
8.5 KiB
//! Defines a worker model for peer connections.
|
|
|
|
use std::io;
|
|
use std::net::SocketAddr;
|
|
|
|
use thiserror::Error;
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
|
|
use crate::core::{FrameReader, FrameWriter, Worker};
|
|
use crate::peer::{Message, PeerConnectionType, PeerInit};
|
|
|
|
// Peer states:
|
|
//
|
|
// - closed
|
|
// - open
|
|
// - waiting for pierce firewall
|
|
// - cannot connect
|
|
//
|
|
// Transitions:
|
|
//
|
|
// - closed:
|
|
// -> open:
|
|
// connect
|
|
// send peerinit
|
|
// -> waiting for pierce firewall:
|
|
// connect failed, or send peerinit failed
|
|
// send connecttopeer to server
|
|
// -> open:
|
|
// accept connection
|
|
// receive peerinit
|
|
// - open:
|
|
// -> closed:
|
|
// connection closed
|
|
// - waiting for pierce firewall:
|
|
// -> open:
|
|
// accept connection
|
|
// receive pierce firewall
|
|
// ???send peer init???
|
|
// -> cannot connect:
|
|
// receive cannot connect
|
|
//
|
|
|
|
/// A peer to connect to.
|
|
#[derive(Debug)]
|
|
pub struct PeerConnection {
|
|
/// The address of the peer.
|
|
pub address: SocketAddr,
|
|
|
|
/// The type of connection to establish.
|
|
pub connection_type: PeerConnectionType,
|
|
|
|
/// The user name as which to identify ourselves to the peer.
|
|
pub our_user_name: String,
|
|
}
|
|
|
|
/// An error that arose while establishing a connection to a peer.
|
|
#[derive(Debug, Error)]
|
|
pub enum PeerConnectionError {
|
|
#[error("error establishing network connection: {0}")]
|
|
ConnectError(#[source] io::Error),
|
|
|
|
#[error("error sending initial message: {0}")]
|
|
WriteError(#[source] io::Error),
|
|
}
|
|
|
|
/// A `Worker` that handles an open connection to a peer.
|
|
type PeerWorker = Worker<Message, Message>;
|
|
|
|
impl PeerConnection {
|
|
pub async fn connect(self) -> Result<PeerWorker, PeerConnectionError> {
|
|
let stream = TcpStream::connect(self.address)
|
|
.await
|
|
.map_err(PeerConnectionError::ConnectError)?;
|
|
|
|
let (read_half, write_half) = stream.into_split();
|
|
let reader = FrameReader::new(read_half);
|
|
let mut writer = FrameWriter::new(write_half);
|
|
|
|
writer
|
|
.write(&Message::PeerInit(PeerInit {
|
|
user_name: self.our_user_name,
|
|
connection_type: self.connection_type,
|
|
token: 0,
|
|
}))
|
|
.await
|
|
.map_err(PeerConnectionError::WriteError)?;
|
|
|
|
Ok(Worker::from_parts(reader, writer))
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct IncomingPeer {
|
|
pub user_name: String,
|
|
pub address: SocketAddr,
|
|
pub connection_type: PeerConnectionType,
|
|
pub token: u32,
|
|
pub worker: PeerWorker,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct IncomingConnection {
|
|
address: SocketAddr,
|
|
stream: TcpStream,
|
|
}
|
|
|
|
#[derive(Debug, Error)]
|
|
pub enum IncomingHandshakeError {
|
|
#[error("error reading message: {0}")]
|
|
ReadError(#[source] io::Error),
|
|
|
|
#[error("stream closed unexpectedly")]
|
|
StreamClosed,
|
|
|
|
#[error("unexpected message: {0:?}")]
|
|
UnexpectedMessage(Message),
|
|
}
|
|
|
|
impl IncomingConnection {
|
|
pub async fn handshake(self) -> Result<IncomingPeer, IncomingHandshakeError> {
|
|
let (read_half, write_half) = self.stream.into_split();
|
|
let mut reader = FrameReader::new(read_half);
|
|
let writer = FrameWriter::new(write_half);
|
|
|
|
let optional_message = reader
|
|
.read()
|
|
.await
|
|
.map_err(IncomingHandshakeError::ReadError)?;
|
|
|
|
match optional_message {
|
|
Some(Message::PeerInit(peer_init)) => Ok(IncomingPeer {
|
|
user_name: peer_init.user_name,
|
|
address: self.address,
|
|
connection_type: peer_init.connection_type,
|
|
token: peer_init.token,
|
|
worker: Worker::from_parts(reader, writer),
|
|
}),
|
|
Some(message) => Err(IncomingHandshakeError::UnexpectedMessage(message)),
|
|
None => Err(IncomingHandshakeError::StreamClosed),
|
|
}
|
|
}
|
|
}
|
|
|
|
struct PeerListener {
|
|
inner: TcpListener,
|
|
}
|
|
|
|
impl PeerListener {
|
|
pub fn new(listener: TcpListener) -> Self {
|
|
Self { inner: listener }
|
|
}
|
|
|
|
pub async fn accept(&mut self) -> io::Result<IncomingConnection> {
|
|
let (stream, address) = self.inner.accept().await?;
|
|
Ok(IncomingConnection { stream, address })
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
|
|
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
|
|
use crate::core::{FrameReader, FrameWriter};
|
|
use crate::peer::{Message, PeerConnectionType, PeerInit};
|
|
|
|
use super::{
|
|
IncomingHandshakeError, PeerConnection, PeerConnectionError, PeerListener,
|
|
};
|
|
|
|
#[tokio::test]
|
|
async fn peer_connection_connect_error() {
|
|
let connection = PeerConnection {
|
|
address: SocketAddr::V4(SocketAddrV4::new(
|
|
// TODO: use example IP instead, ensuring this fails.
|
|
Ipv4Addr::new(0, 0, 0, 1),
|
|
42,
|
|
)),
|
|
connection_type: PeerConnectionType::File,
|
|
our_user_name: "me".to_string(),
|
|
};
|
|
|
|
let err = connection.connect().await.unwrap_err();
|
|
|
|
match err {
|
|
PeerConnectionError::ConnectError(_) => (),
|
|
_ => panic!("Wrong error: {:?}", err),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn peer_connection_success() {
|
|
let listener = TcpListener::bind("localhost:0")
|
|
.await
|
|
.expect("binding listener");
|
|
let address = listener.local_addr().expect("getting local address");
|
|
|
|
let connection = PeerConnection {
|
|
address,
|
|
connection_type: PeerConnectionType::File,
|
|
our_user_name: "me".to_string(),
|
|
};
|
|
|
|
let (connect_result, accept_result) =
|
|
tokio::join!(connection.connect(), listener.accept());
|
|
|
|
let _worker = connect_result.expect("connecting");
|
|
let (stream, _remote_address) = accept_result.expect("accepting");
|
|
|
|
let message = FrameReader::new(stream).read().await.expect("reading");
|
|
assert_eq!(
|
|
message,
|
|
Some(Message::PeerInit(PeerInit {
|
|
user_name: "me".to_string(),
|
|
connection_type: PeerConnectionType::File,
|
|
token: 0,
|
|
}))
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn peer_accept_stream_closed() {
|
|
let tcp_listener = TcpListener::bind("localhost:0").await.expect("binding");
|
|
let listener_address =
|
|
tcp_listener.local_addr().expect("getting local address");
|
|
|
|
let mut peer_listener = PeerListener::new(tcp_listener);
|
|
|
|
let connect_task = async {
|
|
// Open connection and immediately drop/close it.
|
|
TcpStream::connect(listener_address)
|
|
.await
|
|
.expect("connecting");
|
|
};
|
|
|
|
let ((), accept_result) =
|
|
tokio::join!(connect_task, peer_listener.accept());
|
|
|
|
let incoming_connection = accept_result.expect("accepting");
|
|
let err = incoming_connection
|
|
.handshake()
|
|
.await
|
|
.expect_err("performing handshake");
|
|
|
|
match err {
|
|
IncomingHandshakeError::StreamClosed => (),
|
|
_ => panic!("Wrong error: {:?}", err),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn peer_accept_unexpected_message() {
|
|
let tcp_listener = TcpListener::bind("localhost:0").await.expect("binding");
|
|
let listener_address =
|
|
tcp_listener.local_addr().expect("getting local address");
|
|
|
|
let mut peer_listener = PeerListener::new(tcp_listener);
|
|
|
|
let connect_task = async {
|
|
let stream = TcpStream::connect(listener_address)
|
|
.await
|
|
.expect("connecting");
|
|
|
|
let mut writer = FrameWriter::new(stream);
|
|
writer.write("garbage").await.expect("writing");
|
|
|
|
// Return writer so that the connection is kept alive until we notice the
|
|
// error on the receiving side.
|
|
writer
|
|
};
|
|
|
|
let (_writer, accept_result) =
|
|
tokio::join!(connect_task, peer_listener.accept());
|
|
|
|
let incoming_connection = accept_result.expect("accepting");
|
|
let err = incoming_connection
|
|
.handshake()
|
|
.await
|
|
.expect_err("performing handshake");
|
|
|
|
match err {
|
|
IncomingHandshakeError::ReadError(_) => (),
|
|
_ => panic!("Wrong error: {:?}", err),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn peer_accept_success() {
|
|
let tcp_listener = TcpListener::bind("localhost:0").await.expect("binding");
|
|
let listener_address =
|
|
tcp_listener.local_addr().expect("getting local address");
|
|
|
|
let mut peer_listener = PeerListener::new(tcp_listener);
|
|
|
|
let outgoing_connection = PeerConnection {
|
|
address: listener_address,
|
|
connection_type: PeerConnectionType::File,
|
|
our_user_name: "olabode".to_string(),
|
|
};
|
|
|
|
let accept_task = async {
|
|
let incoming_connection =
|
|
peer_listener.accept().await.expect("accepting");
|
|
|
|
incoming_connection
|
|
.handshake()
|
|
.await
|
|
.expect("performing handshake")
|
|
};
|
|
|
|
let (connect_result, incoming_peer) =
|
|
tokio::join!(outgoing_connection.connect(), accept_task);
|
|
|
|
let _worker = connect_result.expect("connecting");
|
|
|
|
assert_eq!(incoming_peer.user_name, "olabode");
|
|
assert_eq!(incoming_peer.connection_type, PeerConnectionType::File);
|
|
assert_eq!(incoming_peer.token, 0);
|
|
}
|
|
}
|