From 61264d9bb13f9343b35fcdc587a72db1fdf7010f Mon Sep 17 00:00:00 2001
From: Jake Probst <jake.probst@gmail.com>
Date: Thu, 18 Jul 2019 22:25:33 -0700
Subject: [PATCH] properly handle sending packets when send buffer gets filled

---
 src/common/client.rs  | 79 ++++++++++++++++++++++++++++---------------
 src/common/network.rs | 16 ++-------
 2 files changed, 54 insertions(+), 41 deletions(-)

diff --git a/src/common/client.rs b/src/common/client.rs
index ad3b378..727fdc9 100644
--- a/src/common/client.rs
+++ b/src/common/client.rs
@@ -2,9 +2,10 @@ use libpso::crypto::{PSOCipher, NullCipher};
 use libpso::{PSOPacket, PacketParseError};
 
 use crate::common::serverstate::{ServerState, ServerPacket, OnConnect};
-use crate::common::network::{send_packet, recv_packet};
+use crate::common::network::{recv_packet, PacketNetworkError};
 
 use std::net;
+use std::io::Write;
 use mio::tcp::TcpStream;
 use mio::{Poll, Events, Token, Ready, PollOpt};
 
@@ -14,6 +15,7 @@ pub struct Client<P, E> {
     cipher_in: Box<dyn PSOCipher>,
     cipher_out: Box<dyn PSOCipher>,
     state: Box<dyn ServerState<Packet = P, PacketError = E>>,
+    send_buffer: Vec<u8>,
 }
 
 impl<P: ServerPacket + std::fmt::Debug, E: std::fmt::Debug> Client<P, E> {
@@ -24,6 +26,7 @@ impl<P: ServerPacket + std::fmt::Debug, E: std::fmt::Debug> Client<P, E> {
             cipher_in: Box::new(NullCipher {}),
             cipher_out: Box::new(NullCipher {}),
             state: state,
+            send_buffer: Vec::new(),
         };
 
         for task in client.state.on_connect() {
@@ -39,53 +42,73 @@ impl<P: ServerPacket + std::fmt::Debug, E: std::fmt::Debug> Client<P, E> {
         client
     }
 
-    fn send(&mut self, pkt: &dyn PSOPacket) {
-        match send_packet(&mut self.socket, &mut *self.cipher_out, pkt) {
-            Ok(_) => {
-                println!("[patch] send ({:?}): {:?}", self.socket, pkt);
+    fn send_data(&mut self) {
+        if self.send_buffer.len() == 0 {
+            return;
+        }
+        match self.socket.write(&self.send_buffer) {
+            Ok(len) => {
+                if len == 0 {
+                    self.running = false;
+                }
+                self.send_buffer.drain(..len);
             },
             Err(err) => {
-                println!("[patch] error sending packet to {:?}: {:?}", self.socket, err);
-                self.running = false;
+                println!("[client] error sending data to {:?}: {:?}", self.socket, err);
             }
         }
     }
 
-    pub fn io_loop(mut self) {
-        let poll = mio::Poll::new().unwrap();
-        poll.register(&self.socket, Token(0), Ready::readable(), PollOpt::edge()).unwrap();
+    // TODO: this may need to pad to 8 bytes for bb cipher
+    fn send(&mut self, pkt: &dyn PSOPacket) {
+        let buf = pkt.as_bytes();
+        println!("[send]: {:X?}", buf);
+        let mut cbuf = self.cipher_out.encrypt(&buf).unwrap();
+        self.send_buffer.append(&mut cbuf);
+        self.send_data();
+    }
 
+    pub fn io_loop(mut self) {
+        let poll = Poll::new().unwrap();
+        poll.register(&self.socket, Token(0), Ready::readable() | Ready::writable(), PollOpt::edge()).unwrap();
 
         let mut events = Events::with_capacity(1024);
-        loop {
+        while self.running {
             poll.poll(&mut events, None).unwrap();
 
             for event in &events{
-                println!("event! {:?}", event);
                 if event.token() == Token(0) {
-                    loop {
-                        let pkt = recv_packet(&mut self.socket, &mut *self.cipher_in)
-                            .and_then(|pkt| {
-                                P::from_bytes(&pkt)
-                                    .map_err(|err| err.into())
-                            });
+                    if event.readiness().is_writable() {
+                        self.send_data();
+                    }
+                    if event.readiness().is_readable() {
+                        loop {
+                            let pkt = recv_packet(&mut self.socket, &mut *self.cipher_in)
+                                .and_then(|pkt| {
+                                    P::from_bytes(&pkt)
+                                        .map_err(|err| err.into())
+                                });
 
-                        match pkt {
-                            Ok(pkt) => {
-                                let response = self.state.handle(&pkt);
-                                for r in response {
-                                    self.send(&*r);
+                            match pkt {
+                                Ok(pkt) => {
+                                    let response = self.state.handle(&pkt);
+                                    for r in response {
+                                        self.send(&*r);
+                                    }
+                                },
+                                Err(err) => {
+                                    println!("error recv-ing packet with {:?}: {:?}", self.socket, err);
+                                    match err {
+                                        PacketNetworkError::ClientDisconnected => self.running = false,
+                                        _ => println!("error recv-ing packet with {:?}: {:?}", self.socket, err),
+                                    }
+                                    break;
                                 }
-                            },
-                            Err(err) => {
-                                println!("error recv-ing packet with {:?}: {:?}", self.socket, err);
-                                break;
                             }
                         }
                     }
                 }
             }
         }
-        
     }
 }
diff --git a/src/common/network.rs b/src/common/network.rs
index f110462..ba9ce1f 100644
--- a/src/common/network.rs
+++ b/src/common/network.rs
@@ -1,7 +1,7 @@
-use std::net::TcpStream;
-use std::io::{Read, Write};
+use std::io::{Read};
+
 use libpso::crypto::{PSOCipher, CipherError};
-use libpso::{PSOPacket, PacketParseError};
+use libpso::PacketParseError;
 
 #[derive(Debug)]
 pub enum PacketNetworkError {
@@ -71,15 +71,5 @@ pub fn recv_packet<T: Read>(socket: &mut T, cipher: &mut dyn PSOCipher) -> Resul
     full_buf.append(&mut dec_size_buf);
     full_buf.append(&mut dec_data_buf);
 
-    //println!("[recv]: {:X?}", full_buf);
     Ok(full_buf)
 }
-
-// TODO: this fails horribly when the send buffer is full
-// if it didnt send it all make a mio poll
-pub fn send_packet<T: Write>(socket: &mut T, cipher: &mut dyn PSOCipher, pkt: &dyn PSOPacket) -> Result<(), PacketNetworkError> {
-    let buf = pkt.as_bytes();
-    //println!("[send]: {:X?}", buf);
-    let cbuf = cipher.encrypt(&buf)?;
-    Ok(socket.write_all(&cbuf)?)
-}