exoplanet/src/state.rs

use crate::pronouns::Pronouns;
use crate::protocol::*;
use protocol::*;
use std::collections::VecDeque;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Result as IoResult};
use std::net::{Ipv4Addr, SocketAddr, UdpSocket};
use std::sync::Arc;

pub struct Identity {
    pub username: String,
    pub about: String,
    pub pronouns: Option<Pronouns>,
}

pub struct LoginInfo {
    pub identity: Identity,
    pub listen_port: u16,
}

pub struct LocalSocket {
    pub listen_socket: UdpSocket,
}

impl LocalSocket {
    pub fn new(port: u16) -> Result<Self, IoError> {
        let ip = Ipv4Addr::UNSPECIFIED.into();
        let addr = SocketAddr::new(ip, port);
        let listen_socket = UdpSocket::bind(addr)?;
        listen_socket.set_nonblocking(true)?;
        Ok(Self { listen_socket })
    }
}

pub struct State {
    pub local_socket: Arc<LocalSocket>,
    pub connections: Vec<Connection>,
    pub identity: Identity,
    pub next_uid: usize,
}

impl State {
    pub fn new(info: LoginInfo) -> Result<Self, IoError> {
        let LoginInfo {
            identity,
            listen_port,
        } = info;

        let local_socket = Arc::new(LocalSocket::new(listen_port)?);
        let connections = Vec::new();
        let identity = identity;
        let next_uid = 0;

        Ok(Self {
            local_socket,
            connections,
            identity,
            next_uid,
        })
    }

    pub fn poll(&mut self) -> Result<(), IoError> {
        let mut buf = [0u8; 65507];

        loop {
            match self.local_socket.listen_socket.recv_from(&mut buf) {
                Ok((len, from)) => {
                    let packet = &buf[0..len];
                    self.on_packet(from, packet)?;
                }
                Err(e) if e.kind() == IoErrorKind::WouldBlock => break,
                Err(e) => return Err(e),
            }
        }

        Ok(())
    }

    pub fn get_packet_kind(reader: &mut impl std::io::Read) -> IoResult<PacketKind> {
        let kind = Var::<u16>::decode(reader)?.0;
        kind.try_into()
            .map_err(|_| IoErrorKind::InvalidInput.into())
    }

    pub fn on_packet(&mut self, from: SocketAddr, mut packet: &[u8]) -> IoResult<()> {
        let kind = Self::get_packet_kind(&mut packet)?;
        let connection = self.connections.iter_mut().find(|c| c.other_addr == from);

        if let Some(connection) = connection {
            connection.on_packet(kind, &mut packet)?;
        } else if kind == PacketKind::HandshakeRequest {
            let sock = self.local_socket.to_owned();
            let uid = self.next_uid;
            let connection = Connection::new_requested(sock, uid, from);
            self.connections.push(connection);
            self.next_uid += 1;
        }

        Ok(())
    }

    pub fn connect(&mut self, to: SocketAddr) -> IoResult<()> {
        let sock = self.local_socket.to_owned();
        let uid = self.next_uid;
        let connection = Connection::connect(sock, uid, to)?;
        self.connections.push(connection);
        self.next_uid += 1;
        Ok(())
    }
}

pub struct Connection {
    /// A handle to the local socket this connection is using.
    pub local_socket: Arc<LocalSocket>,

    /// Local unique identifier for this connection.
    pub uid: usize,

    /// The remote address of the other end of this connection.
    pub other_addr: SocketAddr,

    /// The current state.
    pub state: ConnectionState,
}

impl Connection {
    pub fn connect(
        local_socket: Arc<LocalSocket>,
        uid: usize,
        other_addr: SocketAddr,
    ) -> IoResult<Self> {
        let connection = Self {
            local_socket,
            uid,
            other_addr,
            state: ConnectionState::Connecting,
        };

        connection.send_empty_packet(PacketKind::HandshakeRequest)?;
        Ok(connection)
    }

    pub fn new_requested(
        local_socket: Arc<LocalSocket>,
        uid: usize,
        other_addr: SocketAddr,
    ) -> Self {
        Self {
            local_socket,
            uid,
            other_addr,
            state: ConnectionState::Pending,
        }
    }

    pub fn on_packet(&mut self, kind: PacketKind, mut packet: &[u8]) -> Result<(), IoError> {
        match self.state {
            ConnectionState::Connecting => match kind {
                PacketKind::HandshakeAccept => {
                    self.state = ConnectionState::Connected;
                    self.send_empty_packet(PacketKind::HandshakeFinal)?;
                }
                _ => {}
            },
            ConnectionState::Pending => {}
            ConnectionState::Accepted => match kind {
                PacketKind::HandshakeFinal => {
                    self.state = ConnectionState::Connected;
                }
                _ => {}
            },
            ConnectionState::Connected => match kind {
                PacketKind::Ping => self.send_empty_packet(PacketKind::Pong)?,
                PacketKind::Pong => {}
                PacketKind::Message => {
                    let message = Message::decode(&mut packet)?;
                    eprintln!("message: {:?}", message);
                }
                _ => {}
            },
        }

        Ok(())
    }

    pub fn send_packet(
        &self,
        kind: PacketKind,
        encode: impl FnOnce(&mut Vec<u8>) -> std::io::Result<()>,
    ) -> IoResult<()> {
        let mut buf = Vec::new();
        Var(kind as u16).encode(&mut buf)?;
        encode(&mut buf)?;
        self.local_socket
            .listen_socket
            .send_to(&buf, self.other_addr)?;
        Ok(())
    }

    pub fn send_empty_packet(&self, kind: PacketKind) -> IoResult<()> {
        self.send_packet(kind, |_| Ok(()))
    }
}

#[derive(PartialEq, Eq)]
pub enum ConnectionState {
    Connecting,
    Pending,
    Accepted,
    Connected,
}