diff --git a/Cargo.toml b/Cargo.toml index 27a352e..ab361ef 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,3 +15,4 @@ libpso = { path = "../libpso" } rand = "0.6.5" walkdir = "2" mio = "0.6" +crc = "^1.0.0" \ No newline at end of file diff --git a/src/patch/main.rs b/src/patch/main.rs index 64ba962..05be267 100644 --- a/src/patch/main.rs +++ b/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), + 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 for PatchError { fn from(err: PacketNetworkError) -> PatchError { PatchError::PacketNetworkError(err) } } +impl From 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), + File(PathBuf, u32), // file_id +} + +impl PatchFileTree { + fn iter_dir(tree: &PatchFileTree) -> Vec { + 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 { + PatchFileTree::iter_dir(self) + } +} + + + struct Client { running: bool, key_in: u32, @@ -56,11 +120,12 @@ struct Client { cipher_in: Box, cipher_out: Box, patch_file_tree: PatchFileTree, - patch_file_lookup: HashMap, + patch_file_lookup: HashMap, + patch_file_info: Vec, } impl Client { - fn new(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap) -> Client { + fn new(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap) -> 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), - File(PathBuf, u32), // file_id -} - -fn load_patch_dir(basedir: &str, patchbase: &str, file_ids: &mut HashMap) -> PatchFileTree { - let paths = fs::read_dir(basedir).unwrap(); +fn load_patch_dir(basedir: &str, patchbase: &str, file_ids: &mut HashMap) -> 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 (PatchFileTree, HashMap) { +fn generate_patch_tree(basedir: &str) -> (PatchFileTree, HashMap) { 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> { let mut pkts: Vec> = 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 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) -> 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::>(); + + 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::>(); + + 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) { +fn new_client(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap) { 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..."); }