<commit_before>use std::cell::RefCell;
use std::io;
use std::rc::{Rc, Weak};
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use mio::{Event, PollOpt, Ready, Token};
use mio::net::UdpSocket;

use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::datagram_buffer::DatagramBuffer;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packetizer::Packetizer;
use super::selector::Selector;

const TAG: &'static str = "UDPConnection";

pub struct UDPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    socket: UdpSocket,
    token: Token,
    client_to_network: DatagramBuffer,
    network_to_client: Packetizer,
    closed: bool,
}

impl UDPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, reference_packet: &IPv4Packet) -> io::Result<Rc<RefCell<Self>>> {
        let socket = UDPConnection::create_socket(&id)?;
        let raw: &[u8] = reference_packet.raw();
        let ipv4_header = reference_packet.ipv4_header().clone();
        let transport_header = reference_packet.transport_header().as_ref().unwrap().clone();
        let rc = Rc::new(RefCell::new(Self {
            id: id,
            client: client,
            socket: socket,
            token: Token(0), // default value, will be set afterwards
            client_to_network: DatagramBuffer::new(4 * MAX_PACKET_LENGTH),
            network_to_client: Packetizer::new(raw, ipv4_header, transport_header),
            closed: false,
        }));

        {
            let rc_clone = rc.clone();
            let handler = move |selector: &mut Selector, ready| {
                let mut self_ref = rc_clone.borrow_mut();
                self_ref.on_ready(selector, ready);
            };
            let mut self_ref = rc.borrow_mut();
            let token = selector.register(&self_ref.socket, handler, Ready::readable(), PollOpt::level())?;
            self_ref.token = token;
        }
        Ok(rc)
    }

    fn create_socket(id: &ConnectionId) -> io::Result<UdpSocket> {
        let autobind_addr = SocketAddr::new(Ipv4Addr::new(0, 0, 0, 0).into(), 0);
        let udp_socket = UdpSocket::bind(&autobind_addr)?;
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        udp_socket.connect(rewritten_destination)?;
        Ok(udp_socket)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.disconnect();

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        if let Err(_) = self.write() {
            error!(target: TAG, "Cannot write");
            self.close(selector);
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        // TODO
    }

    fn update_interests(&mut self, selector: &mut Selector) -> io::Result<()> {
        let mut ready = Ready::empty();
        if self.may_read() {
            ready = Ready::readable();
        }
        if self.may_write() {
            ready = ready | Ready::writable();
        }
        selector.reregister(&self.socket, self.token, ready, PollOpt::level())
    }

    fn read(&mut self) -> io::Result<()> {
        let packet = self.network_to_client.packetize(&mut self.socket)?;
        // TODO
        Ok(())
    }

    fn write(&mut self) -> io::Result<()> {
        self.client_to_network.write_to(&mut self.socket)?;
        Ok(())
    }

    fn may_read(&self) -> bool {
        // TODO
        true
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }<commit_msg>Introduce pending packet for UDPConnection<commit_after>use std::cell::RefCell;
use std::io;
use std::rc::{Rc, Weak};
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use mio::{Event, PollOpt, Ready, Token};
use mio::net::UdpSocket;

use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::datagram_buffer::DatagramBuffer;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packetizer::Packetizer;
use super::selector::Selector;

const TAG: &'static str = "UDPConnection";

pub struct UDPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    socket: UdpSocket,
    token: Token,
    client_to_network: DatagramBuffer,
    network_to_client: Packetizer,
    pending_packet_for_client: Option<Box<[u8]>>, // used when the client buffer is full, should be rare
    closed: bool,
}

impl UDPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, reference_packet: &IPv4Packet) -> io::Result<Rc<RefCell<Self>>> {
        let socket = UDPConnection::create_socket(&id)?;
        let raw: &[u8] = reference_packet.raw();
        let ipv4_header = reference_packet.ipv4_header().clone();
        let transport_header = reference_packet.transport_header().as_ref().unwrap().clone();
        let rc = Rc::new(RefCell::new(Self {
            id: id,
            client: client,
            socket: socket,
            token: Token(0), // default value, will be set afterwards
            client_to_network: DatagramBuffer::new(4 * MAX_PACKET_LENGTH),
            network_to_client: Packetizer::new(raw, ipv4_header, transport_header),
            pending_packet_for_client: None,
            closed: false,
        }));

        {
            let rc_clone = rc.clone();
            let handler = move |selector: &mut Selector, ready| {
                let mut self_ref = rc_clone.borrow_mut();
                self_ref.on_ready(selector, ready);
            };
            let mut self_ref = rc.borrow_mut();
            let token = selector.register(&self_ref.socket, handler, Ready::readable(), PollOpt::level())?;
            self_ref.token = token;
        }
        Ok(rc)
    }

    fn create_socket(id: &ConnectionId) -> io::Result<UdpSocket> {
        let autobind_addr = SocketAddr::new(Ipv4Addr::new(0, 0, 0, 0).into(), 0);
        let udp_socket = UdpSocket::bind(&autobind_addr)?;
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        udp_socket.connect(rewritten_destination)?;
        Ok(udp_socket)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.disconnect();

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        if let Err(err) = self.write() {
            error!(target: TAG, "Cannot write: {}", err);
            self.close(selector);
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        if let Err(err) = self.read() {
            error!(target: TAG, "Cannot read: {}", err);
            self.close(selector);
        }
        // TODO push to client
    }

    fn update_interests(&mut self, selector: &mut Selector) -> io::Result<()> {
        let mut ready = Ready::empty();
        if self.may_read() {
            ready = Ready::readable();
        }
        if self.may_write() {
            ready = ready | Ready::writable();
        }
        selector.reregister(&self.socket, self.token, ready, PollOpt::level())
    }

    fn read(&mut self) -> io::Result<()> {
        assert!(self.pending_packet_for_client.is_none());
        let packet = self.network_to_client.packetize(&mut self.socket)?;
        // TODO
        Ok(())
    }

    fn write(&mut self) -> io::Result<()> {
        self.client_to_network.write_to(&mut self.socket)?;
        Ok(())
    }

    fn may_read(&self) -> bool {
        // TODO
        true
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
{'fraction_non_alphanumeric': 0.13769187038089825, 'fraction_numerical': 0.0037521318931210915, 'mean_word_length': 2.6589018302828618, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '394714', 'n_tokens_mistral': 2709, 'n_tokens_neox': 2651, 'n_words': 596}