PATCHSERVER!

Cargo.toml

@@ -15,3 +15,4 @@ libpso = { path = "../libpso" }
 rand = "0.6.5"
 walkdir = "2"
 mio = "0.6"
+crc = "^1.0.0"

src/patch/main.rs

@@ -1,14 +1,17 @@
 #![allow(unused_imports)]
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::net::{TcpListener, SocketAddr, Ipv4Addr};
 use std::net;
 use std::thread;
 use std::fs;
+use std::io;
+use std::io::{Read, Write};
 use std::path::{Path, PathBuf, Components};
 use std::convert::AsRef;
 use mio::{Events, Poll, Token, Ready, PollOpt};
 use rand::{Rng, RngCore};
+use crc::{crc32, Hasher32};
 use libpso::{PacketParseError, PSOPacket};
 use libpso::patch::packet::*;
 use libpso::crypto::{CipherError, PSOCipher, NullCipher};
@@ -16,26 +19,38 @@ use libpso::crypto::pc::PSOPCCipher;
 use elseware::common::{send_packet, recv_packet, PacketNetworkError};
 
 const PATCH_PORT: u16 = 11000;
-const DATA_PORT: u16 = 11001;
-
 
 #[derive(Debug)]
 enum PatchError {
     PacketNetworkError(PacketNetworkError),
     UnexpectedPacket(Box<dyn PSOPacket>),
+    IOError(std::io::Error),
     CloseConnection,
 }
 
+// TODO: something like
+// convert_error!(PacketNetworkError into PatchError)
+// or
+// convert_error!(std::io::Error into PatchError as IOError)
+
 impl From<PacketNetworkError> for PatchError {
     fn from(err: PacketNetworkError) -> PatchError {
         PatchError::PacketNetworkError(err)
     }
 }
 
+impl From<std::io::Error> for PatchError {
+    fn from(err: std::io::Error) -> PatchError {
+        PatchError::IOError(err)
+    }
+}
+
 #[derive(Debug)]
 pub enum PatchPacket {
     PatchWelcomeReply(PatchWelcomeReply),
     LoginReply(LoginReply),
+    FileInfoReply(FileInfoReply),
+    FileInfoListEnd(FileInfoListEnd),
 }
 
 impl PatchPacket {
@@ -43,11 +58,60 @@ impl PatchPacket {
         match data[2] {
             0x02 => Ok(PatchPacket::PatchWelcomeReply(PatchWelcomeReply::from_bytes(data)?)),
             0x04 => Ok(PatchPacket::LoginReply(LoginReply::from_bytes(data)?)),
+            0x0F => Ok(PatchPacket::FileInfoReply(FileInfoReply::from_bytes(data)?)),
+            0x10 => Ok(PatchPacket::FileInfoListEnd(FileInfoListEnd::from_bytes(data)?)),
             _ => Err(PacketParseError::WrongPacketForServerType)
         }
     }
 }
 
+#[derive(Debug, Clone)]
+struct PatchFile {
+    path: PathBuf,
+    checksum: u32,
+    size: u32,
+}
+
+
+enum PatchTreeIterItem {
+    Directory(PathBuf),
+    File(PathBuf, u32),
+    UpDirectory,
+}
+
+#[derive(Debug, Clone)]
+enum PatchFileTree {
+    Directory(PathBuf, Vec<PatchFileTree>),
+    File(PathBuf, u32), // file_id
+}
+
+impl PatchFileTree {
+    fn iter_dir(tree: &PatchFileTree) -> Vec<PatchTreeIterItem> {
+        let mut v = Vec::new();
+
+        match tree {
+            PatchFileTree::Directory(dir, files) => {
+                v.push(PatchTreeIterItem::Directory(dir.clone()));
+                for file in files {
+                    v.append(&mut PatchFileTree::iter_dir(&file));
+                }
+                v.push(PatchTreeIterItem::UpDirectory);
+            },
+            PatchFileTree::File(path, id) => {
+                v.push(PatchTreeIterItem::File(path.clone(), *id));
+            }
+        }
+
+        v
+    }
+
+    pub fn flatten(&self) -> Vec<PatchTreeIterItem> {
+        PatchFileTree::iter_dir(self)
+    }
+}
+
+
+
 struct Client {
     running: bool,
     key_in: u32,
@@ -56,11 +120,12 @@ struct Client {
     cipher_in: Box<dyn PSOCipher>,
     cipher_out: Box<dyn PSOCipher>,
     patch_file_tree: PatchFileTree,
-    patch_file_lookup: HashMap<u32, PathBuf>,
+    patch_file_lookup: HashMap<u32, PatchFile>,
+    patch_file_info: Vec<FileInfoReply>,
 }
 
 impl Client {
-    fn new(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap<u32, PathBuf>) -> Client {
+    fn new(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap<u32, PatchFile>) -> Client {
         let mut rng = rand::thread_rng();
         let key_in: u32 = rng.gen();
         let key_out: u32 = rng.gen();
@@ -69,11 +134,12 @@ impl Client {
             running: true,
             key_in: key_in,
             key_out: key_out,
-            socket: mio::tcp::TcpStream::from_stream(socket).unwrap(),
+            socket: mio::tcp::TcpStream::from_stream(socket).expect("could not convert socket to nonblocking"),
             cipher_in: Box::new(NullCipher {}),
             cipher_out: Box::new(NullCipher {}),
             patch_file_tree: patch_file_tree,
             patch_file_lookup: patch_file_lookup,
+            patch_file_info: Vec::new()
         }
     }
 
@@ -91,19 +157,13 @@ impl Client {
 
 }
 
-#[derive(Debug, Clone)]
-enum PatchFileTree {
-    Directory(PathBuf, Vec<PatchFileTree>),
-    File(PathBuf, u32), // file_id
-}
-
-fn load_patch_dir(basedir: &str, patchbase: &str, file_ids: &mut HashMap<u32, PathBuf>) -> PatchFileTree {
-    let paths = fs::read_dir(basedir).unwrap();
+fn load_patch_dir(basedir: &str, patchbase: &str, file_ids: &mut HashMap<u32, PatchFile>) -> PatchFileTree {
+    let paths = fs::read_dir(basedir).expect("could not read directory");
 
     let mut files = Vec::new();
     let mut dirs = Vec::new();
     for p in paths {
-        let path = p.unwrap().path();
+        let path = p.expect("not a real path").path();
         if path.is_dir() {
             let patch_path = path.strip_prefix(basedir).unwrap();
             dirs.push(load_patch_dir(path.to_str().unwrap(), patch_path.to_str().unwrap(), file_ids));
@@ -111,7 +171,12 @@ fn load_patch_dir(basedir: &str, patchbase: &str, file_ids: &mut HashMap<u32, Pa
         else {
             let patch_path = path.strip_prefix(basedir).unwrap();
             files.push(PatchFileTree::File(patch_path.to_path_buf(), file_ids.len() as u32));
-            file_ids.insert(file_ids.len() as u32, path);
+            let (checksum, size) = get_checksum_and_size(&path).unwrap();
+            file_ids.insert(file_ids.len() as u32, PatchFile {
+                path: path,
+                checksum: checksum,
+                size: size,
+            });
         }
     }
 
@@ -120,28 +185,30 @@ fn load_patch_dir(basedir: &str, patchbase: &str, file_ids: &mut HashMap<u32, Pa
     PatchFileTree::Directory(PathBuf::from(patchbase), files)
 }
 
-fn generate_patch_tree(basedir: &str) -> (PatchFileTree, HashMap<u32, PathBuf>) {
+fn generate_patch_tree(basedir: &str) -> (PatchFileTree, HashMap<u32, PatchFile>) {
     let mut file_ids = HashMap::new();
 
     let patch_tree = load_patch_dir(basedir, "", &mut file_ids);
-    
+
     (patch_tree, file_ids)
 }
 
 
+// TODO: this should be a function that takes client and sends packets instead of returning a vec of packets
 fn get_file_list_packets(patch_file_tree: &PatchFileTree) -> Vec<Box<dyn PSOPacket>> {
     let mut pkts: Vec<Box<dyn PSOPacket>> = Vec::new();
 
-    match patch_file_tree {
-        PatchFileTree::Directory(dir, files) => {
-            pkts.push(Box::new(ChangeDirectory::new(dir.to_str().unwrap())));
-            for file in files {
-                pkts.append(&mut get_file_list_packets(&file));
+    for item in patch_file_tree.flatten() {
+        match item {
+            PatchTreeIterItem::Directory(path) => {
+                pkts.push(Box::new(ChangeDirectory::new(path.to_str().unwrap())));
+            },
+            PatchTreeIterItem::File(path, id) => {
+                pkts.push(Box::new(FileInfo::new(path.to_str().unwrap(), id)));
+            },
+            PatchTreeIterItem::UpDirectory => {
+                pkts.push(Box::new(UpOneDirectory {}));
             }
-            pkts.push(Box::new(UpOneDirectory {}));
-        },
-        PatchFileTree::File(path, id) => {
-            pkts.push(Box::new(FileInfo::new(path.to_str().unwrap(), *id)));
         }
     }
 
@@ -149,19 +216,105 @@ fn get_file_list_packets(patch_file_tree: &PatchFileTree) -> Vec<Box<dyn PSOPack
 }
 
 
-// so I can not figure out why this doesnt work!
 fn send_file_list(client: &mut Client) -> Result<(), PatchError> {
     client.send(&PatchStartList {});
 
     let pkts = get_file_list_packets(&client.patch_file_tree);
-
-    /*for pkt in pkts {
+    for pkt in pkts {
         client.send(&*pkt);
-    }*/
-    
-    client.send(&ChangeDirectory::new(""));
+    }
+
     client.send(&PatchEndList {});
-    client.send(&EndIt {});
+
+    Ok(())
+}
+
+fn get_checksum_and_size(path: &PathBuf) -> Result<(u32, u32), PatchError> {
+    let file = fs::File::open(path)?;
+    let size = file.metadata().unwrap().len();
+    let mut crc = crc32::Digest::new(crc32::IEEE);
+    let mut buf = [0u8; 1024 * 32];
+    let mut reader = io::BufReader::new(file);
+    while let Ok(len) = reader.read(&mut buf) {
+        if len == 0 {
+            break;
+        }
+        crc.write(&buf[0..len]);
+    }
+
+    Ok((crc.sum32(), size as u32))
+}
+
+fn does_file_need_updating(file_info: &FileInfoReply, patch_file_lookup: &HashMap<u32, PatchFile>) -> bool {
+    let patch_file = patch_file_lookup.get(&file_info.id).unwrap();
+    patch_file.checksum != file_info.checksum || patch_file.size != file_info.size
+}
+
+fn send_file(client: &mut Client, actual_path: &PathBuf, patch_path: &PathBuf, id: u32) -> Result<(), PatchError>{
+    let file = fs::File::open(actual_path)?;
+    let size = file.metadata().unwrap().len();
+
+    client.send(&StartFileSend::new(patch_path.to_str().unwrap(), size as u32, id));
+
+    let mut buf = [0u8; PATCH_FILE_CHUNK_SIZE as usize];
+    let mut reader = io::BufReader::new(file);
+    let mut chunk_num = 0;
+    while let Ok(len) = reader.read(&mut buf) {
+        if len == 0 {
+            break;
+        }
+
+        let mut crc = crc32::Digest::new(crc32::IEEE);
+        crc.write(&buf[0..len]);
+
+        let pkt = FileSend {
+            chunk_num: chunk_num,
+            checksum: crc.sum32(),
+            chunk_size: len as u32,
+            buffer: buf
+        };
+
+        client.send(&pkt);
+
+        chunk_num += 1;
+    }
+
+    client.send(&EndFileSend::new());
+    Ok(())
+}
+
+
+fn send_file_data(client: &mut Client) -> Result<(), PatchError> {
+    let need_update = client.patch_file_info.iter()
+        .filter(|file_info| does_file_need_updating(file_info, &client.patch_file_lookup))
+        .collect::<Vec<_>>();
+
+    let total_size = need_update.iter().fold(0, |a, file_info| a + file_info.size);
+    let total_files = need_update.len() as u32;
+
+    let file_ids_to_update = need_update.iter().map(|k| k.id).collect::<HashSet<_>>();
+
+    client.send(&FilesToPatchMetadata::new(total_size, total_files));
+    client.send(&PatchStartList {});
+
+    for file in client.patch_file_tree.flatten() {
+        match file {
+            PatchTreeIterItem::Directory(path) => {
+                client.send(&ChangeDirectory::new(path.to_str().unwrap()));
+            },
+            PatchTreeIterItem::File(path, id) => {
+                if file_ids_to_update.contains(&id) {
+                    let patch_file = client.patch_file_lookup.get(&id).unwrap().clone();
+                    send_file(client, &patch_file.path, &path, id)?;
+                }
+            },
+            PatchTreeIterItem::UpDirectory => {
+                client.send(&UpOneDirectory {});
+            }
+        }
+
+    }
+    client.send(&FinalizePatching {});
 
     Ok(())
 }
@@ -171,13 +324,19 @@ fn handle_packet(client: &mut Client, pkt: PatchPacket) -> Result<(), PatchError
     println!("[patch] recv({:?}): {:?}", client.socket, pkt);
 
     match pkt {
-        PatchPacket::PatchWelcomeReply(pkt) => {
+        PatchPacket::PatchWelcomeReply(_pkt) => {
             client.send(&RequestLogin {});
         },
-        PatchPacket::LoginReply(pkt) => {
+        PatchPacket::LoginReply(_pkt) => {
             client.send(&Message::new("hello player".to_string()));
-            send_file_list(client);
+            send_file_list(client)?;
         },
+        PatchPacket::FileInfoReply(pkt) => {
+            client.patch_file_info.push(pkt);
+        },
+        PatchPacket::FileInfoListEnd(_pkt) => {
+            send_file_data(client)?;
+        }
     }
 
     Ok(())
@@ -185,29 +344,30 @@ fn handle_packet(client: &mut Client, pkt: PatchPacket) -> Result<(), PatchError
 
 fn client_loop(mut client: Client) {
     let poll = mio::Poll::new().unwrap();
-    poll.register(&client.socket, Token(0), Ready::readable(), PollOpt::edge());
+    poll.register(&client.socket, Token(0), Ready::readable(), PollOpt::edge()).unwrap();
 
     let mut events = Events::with_capacity(1024);
     loop {
-        println!("about to poll {:?}", client.socket);
         poll.poll(&mut events, None).unwrap();
-        println!("polled!");
-        
+
         for event in &events{
             println!("event! {:?}", event);
             if event.token() == Token(0) {
-                let pkt = recv_packet(&mut client.socket, &mut *client.cipher_in)
-                    .and_then(|pkt| {
-                        PatchPacket::from_bytes(&pkt)
-                            .map_err(|err| err.into())
-                    });
-
-                match pkt {
-                    Ok(pkt) => {
-                        handle_packet(&mut client, pkt);
-                    },
-                    Err(err) => {
-                        println!("[patch] error recv-ing packet with {:?}: {:?}", client.socket, err);
+                loop {
+                    let pkt = recv_packet(&mut client.socket, &mut *client.cipher_in)
+                        .and_then(|pkt| {
+                            PatchPacket::from_bytes(&pkt)
+                                .map_err(|err| err.into())
+                        });
+
+                    match pkt {
+                        Ok(pkt) => {
+                            handle_packet(&mut client, pkt).expect("could not handle packet");
+                        },
+                        Err(err) => {
+                            println!("[patch] error recv-ing packet with {:?}: {:?}", client.socket, err);
+                            break;
+                        }
                     }
                 }
             }
@@ -215,7 +375,7 @@ fn client_loop(mut client: Client) {
     }
 }
 
-fn new_client(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap<u32, PathBuf>) {
+fn new_client(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap<u32, PatchFile>) {
     let mut client = Client::new(socket, patch_file_tree, patch_file_lookup);
     let welcome_pkt = PatchWelcome::new(client.key_out, client.key_in);
     client.send(&welcome_pkt);
@@ -231,7 +391,7 @@ fn main() {
     let listener = TcpListener::bind(&SocketAddr::from((Ipv4Addr::new(0,0,0,0), PATCH_PORT))).unwrap();
 
     let (patch_file_tree, patch_file_lookup) = generate_patch_tree("patchfiles/");
-    
+
     while let Ok((socket, addr)) = listener.accept() {
         let local_patch_file_tree = patch_file_tree.clone();
         let local_patch_file_lookup = patch_file_lookup.clone();
@@ -241,6 +401,6 @@ fn main() {
         });
     }
 
-    
+
     println!("[patch] exiting...");
 }