handle disconnected clients, use psocipher.header_size

src/common/mod.rs

@@ -1,3 +1,5 @@
+pub mod cipherkeys;
+
 use std::net::TcpStream;
 use std::io::{Read, Write};
 use libpso::crypto::{PSOCipher, CipherError};
@@ -10,6 +12,7 @@ pub enum PacketNetworkError {
     PacketParseError(PacketParseError),
     IOError(std::io::Error),
     DataNotReady,
+    ClientDisconnected,
 }
 
 impl From<CipherError> for PacketNetworkError {
@@ -30,43 +33,54 @@ impl From<PacketParseError> for PacketNetworkError {
     }
 }
 
-
-
 pub fn recv_packet<T: Read>(socket: &mut T, cipher: &mut dyn PSOCipher) -> Result<Vec<u8>, PacketNetworkError> {
-    let mut size_buf = [0u8; 4];
+    let mut size_buf = vec![0u8; cipher.header_size()];
-    //socket.read_exact(&mut size_buf)?;
+    let mut offset = 0;
-    /*let mut offset = 0;
+    while offset < cipher.header_size() {
-    while offset < 4 {
+        let diff = socket.read(&mut size_buf[offset..])?;
-        offset += socket.read(&mut size_buf[offset..])?;
+        println!("! {} {:?}", diff, size_buf);
-    }*/
+        if diff == 0 {
-    let len = socket.read(&mut size_buf)?;
+            return Err(PacketNetworkError::ClientDisconnected);
-    if len != 4 {
+        }
-        return Err(PacketNetworkError::DataNotReady)
+        offset += diff;
     }
 
-    let mut dec_size_buf = cipher.decrypt(&size_buf.to_vec())?;
+    let mut dec_size_buf = cipher.decrypt(&size_buf)?;
-    let size = u16::from_le_bytes([dec_size_buf[0], dec_size_buf[1]]);
+    let size = u16::from_le_bytes([dec_size_buf[0], dec_size_buf[1]]) as usize;
 
-    let mut data_buf = vec![0u8; (size - 4) as usize];
+    let mut data_buf = vec![0u8; size - cipher.header_size()];
-    //socket.read_exact(&mut data_buf)?;
     let mut offset = 0;
-    while offset < size as usize - 4{
+    while offset < size - cipher.header_size() {
-        offset += socket.read(&mut data_buf[offset..])?;
+        let diff = socket.read(&mut data_buf[offset..])?;
+        if diff == 0 {
+            return Err(PacketNetworkError::ClientDisconnected);
+        }
+        offset += diff;
+    }
+
+    let mut null_data_count = 0;
+    while data_buf.len() % cipher.header_size() != 0 {
+        data_buf.push(0);
+        null_data_count += 1;
     }
     
     let mut dec_data_buf = cipher.decrypt(&data_buf.to_vec())?;
 
+    for _ in  0..null_data_count {
+        dec_data_buf.pop();
+    }
+    
     let mut full_buf = Vec::new();
     full_buf.append(&mut dec_size_buf);
     full_buf.append(&mut dec_data_buf);
 
-    println!("[recv]: {:X?}", full_buf);
+    //println!("[recv]: {:X?}", full_buf);
     Ok(full_buf)
 }
 
 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);
+    //println!("[send]: {:X?}", buf);
     let cbuf = cipher.encrypt(&buf)?;
     Ok(socket.write_all(&cbuf)?)
 }