use std::collections::{HashMap, HashSet}; use std::fs; use std::io; use std::io::{Read}; use std::path::PathBuf; use rand::Rng; use crc::{crc32, Hasher32}; use libpso::{PacketParseError, PSOPacket}; use libpso::packet::patch::*; use libpso::crypto::pc::PSOPCCipher; use ron::de::from_str; use serde::Deserialize; use crate::common::network::{PacketNetworkError}; use crate::common::serverstate::{RecvServerPacket, SendServerPacket, ServerState, OnConnect, ClientId}; #[derive(Debug)] pub enum PatchError { PacketNetworkError(PacketNetworkError), IOError(std::io::Error), } 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, Clone)] pub struct PatchFile { path: PathBuf, checksum: u32, size: u32, } pub enum PatchTreeIterItem { Directory(PathBuf), File(PathBuf, u32), UpDirectory, } #[derive(Debug, Clone)] pub 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) } } #[derive(Debug)] pub enum RecvPatchPacket { PatchWelcomeReply(PatchWelcomeReply), LoginReply(LoginReply), FileInfoReply(FileInfoReply), FileInfoListEnd(FileInfoListEnd), } impl RecvServerPacket for RecvPatchPacket { fn from_bytes(data: &[u8]) -> Result { match data[2] { 0x02 => Ok(RecvPatchPacket::PatchWelcomeReply(PatchWelcomeReply::from_bytes(data)?)), 0x04 => Ok(RecvPatchPacket::LoginReply(LoginReply::from_bytes(data)?)), 0x0F => Ok(RecvPatchPacket::FileInfoReply(FileInfoReply::from_bytes(data)?)), 0x10 => Ok(RecvPatchPacket::FileInfoListEnd(FileInfoListEnd::from_bytes(data)?)), _ => Err(PacketParseError::WrongPacketForServerType(u16::from_le_bytes([data[2], data[3]]), data.to_vec())) } } } #[derive(Debug)] pub enum SendPatchPacket { ChangeDirectory(ChangeDirectory), EndFileSend(EndFileSend), FileInfo(FileInfo), FileSend(FileSend), FilesToPatchMetadata(FilesToPatchMetadata), FinalizePatching(FinalizePatching), Message(Message), PatchEndList(PatchEndList), PatchStartList(PatchStartList), PatchWelcome(PatchWelcome), RequestLogin(RequestLogin), StartFileSend(StartFileSend), UpOneDirectory(UpOneDirectory), } impl SendServerPacket for SendPatchPacket { fn as_bytes(&self) -> Vec { match self { SendPatchPacket::ChangeDirectory(pkt) => pkt.as_bytes(), SendPatchPacket::EndFileSend(pkt) => pkt.as_bytes(), SendPatchPacket::FileInfo(pkt) => pkt.as_bytes(), SendPatchPacket::FileSend(pkt) => pkt.as_bytes(), SendPatchPacket::FilesToPatchMetadata(pkt) => pkt.as_bytes(), SendPatchPacket::FinalizePatching(pkt) => pkt.as_bytes(), SendPatchPacket::Message(pkt) => pkt.as_bytes(), SendPatchPacket::PatchEndList(pkt) => pkt.as_bytes(), SendPatchPacket::PatchStartList(pkt) => pkt.as_bytes(), SendPatchPacket::PatchWelcome(pkt) => pkt.as_bytes(), SendPatchPacket::RequestLogin(pkt) => pkt.as_bytes(), SendPatchPacket::StartFileSend(pkt) => pkt.as_bytes(), SendPatchPacket::UpOneDirectory(pkt) => pkt.as_bytes(), } } } pub struct PatchServerState { patch_file_tree: PatchFileTree, patch_file_lookup: HashMap, patch_file_info: Vec, } impl PatchServerState { pub fn new(patch_file_tree: PatchFileTree, patch_file_lookup: HashMap) -> PatchServerState { PatchServerState { patch_file_tree: patch_file_tree, patch_file_lookup: patch_file_lookup, patch_file_info: Vec::new(), } } } impl ServerState for PatchServerState { type SendPacket = SendPatchPacket; type RecvPacket = RecvPatchPacket; type PacketError = PatchError; fn on_connect(&mut self, _id: ClientId) -> Vec> { let mut rng = rand::thread_rng(); let key_in: u32 = rng.gen(); let key_out: u32 = rng.gen(); vec![OnConnect::Packet(SendPatchPacket::PatchWelcome(PatchWelcome::new(key_out, key_in))), OnConnect::Cipher((Box::new(PSOPCCipher::new(key_in)), Box::new(PSOPCCipher::new(key_out)))) ] } fn handle(&mut self, id: ClientId, pkt: &RecvPatchPacket) -> Result>, PatchError> { Ok(match pkt { RecvPatchPacket::PatchWelcomeReply(_pkt) => { Box::new(vec![SendPatchPacket::RequestLogin(RequestLogin {})].into_iter().map(move |pkt| (id, pkt))) }, RecvPatchPacket::LoginReply(_pkt) => { let mut p = vec![SendPatchPacket::Message(Message::new("hello player".to_string()))]; p.append(&mut get_file_list_packets(&self.patch_file_tree)); p.push(SendPatchPacket::PatchEndList(PatchEndList {})); Box::new(p.into_iter().map(move |pkt| (id, pkt))) }, RecvPatchPacket::FileInfoReply(pkt) => { self.patch_file_info.push(pkt.clone()); Box::new(None.into_iter().map(move |pkt| (id, pkt))) }, RecvPatchPacket::FileInfoListEnd(_pkt) => { let need_update = self.patch_file_info.iter() .filter(|file_info| does_file_need_updating(file_info, &self.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 p = vec![SendPatchPacket::FilesToPatchMetadata(FilesToPatchMetadata::new(total_size, total_files)), SendPatchPacket::PatchStartList(PatchStartList {}) ]; Box::new(p.into_iter().chain(SendFileIterator::new(&self)).map(move |pkt| (id, pkt))) } }) } } 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.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)); } else { let patch_path = path.strip_prefix(basedir).unwrap(); files.push(PatchFileTree::File(patch_path.to_path_buf(), file_ids.len() as u32)); let (checksum, size) = get_checksum_and_size(&path).unwrap(); file_ids.insert(file_ids.len() as u32, PatchFile { path: path, checksum: checksum, size: size, }); } } files.append(&mut dirs); PatchFileTree::Directory(PathBuf::from(patchbase), files) } pub 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) } fn get_file_list_packets(patch_file_tree: &PatchFileTree) -> Vec { let mut pkts = Vec::new(); for item in patch_file_tree.flatten() { match item { PatchTreeIterItem::Directory(path) => { pkts.push(SendPatchPacket::ChangeDirectory(ChangeDirectory::new(path.to_str().unwrap()))); }, PatchTreeIterItem::File(path, id) => { pkts.push(SendPatchPacket::FileInfo(FileInfo::new(path.to_str().unwrap(), id))); }, PatchTreeIterItem::UpDirectory => { pkts.push(SendPatchPacket::UpOneDirectory(UpOneDirectory {})); } } } pkts } 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 } struct SendFileIterator { done: bool, file_iter: Box>, file_ids_to_update: HashSet, patch_file_lookup: HashMap, current_file: Option>, chunk_num: u32, } impl SendFileIterator { fn new(state: &PatchServerState) -> SendFileIterator { let file_ids_to_update = state.patch_file_info.iter() .filter(|file_info| does_file_need_updating(file_info, &state.patch_file_lookup)) .map(|k| k.id) .collect::>(); SendFileIterator { done: false, file_ids_to_update: file_ids_to_update, patch_file_lookup: state.patch_file_lookup.clone(), file_iter: Box::new(state.patch_file_tree.flatten().into_iter()), current_file: None, chunk_num: 0, } } } impl Iterator for SendFileIterator { type Item = SendPatchPacket; fn next(&mut self) -> Option { if self.done { return None; } match self.current_file { Some(ref mut file) => { let mut buf = [0u8; PATCH_FILE_CHUNK_SIZE as usize]; let len = file.read(&mut buf).unwrap(); if len == 0 { self.current_file = None; self.chunk_num = 0; Some(SendPatchPacket::EndFileSend(EndFileSend::new())) } else { let mut crc = crc32::Digest::new(crc32::IEEE); crc.write(&buf[0..len]); let pkt = SendPatchPacket::FileSend(FileSend { chunk_num: self.chunk_num, checksum: crc.sum32(), chunk_size: len as u32, buffer: buf, }); self.chunk_num += 1; Some(pkt) } }, None => { match self.file_iter.next() { Some(next_file) => { match next_file { PatchTreeIterItem::Directory(path) => { Some(SendPatchPacket::ChangeDirectory(ChangeDirectory::new(path.to_str().unwrap()))) }, PatchTreeIterItem::File(path, id) => { if self.file_ids_to_update.contains(&id) { let patch_file = self.patch_file_lookup.get(&id).unwrap(); let file = fs::File::open(&patch_file.path).unwrap(); let size = file.metadata().unwrap().len(); self.current_file = Some(io::BufReader::new(file)); Some(SendPatchPacket::StartFileSend(StartFileSend::new(path.to_str().unwrap(), size as u32, id))) } else { self.next() } }, PatchTreeIterItem::UpDirectory => { Some(SendPatchPacket::UpOneDirectory(UpOneDirectory {})) }, } }, None => { self.current_file = None; self.done = true; Some(SendPatchPacket::FinalizePatching(FinalizePatching {})) } } } } } } #[derive(Debug, Deserialize)] pub struct PatchConfig { pub path: String, pub ip: String, pub port: u16, } pub fn load_config() -> PatchConfig { let ini_file = match fs::File::open(std::path::Path::new("patch.ron")) { Err(err) => panic!("Failed to open patch.ron config file. \n{}", err), Ok(ini_file) => ini_file, }; let mut s = String::new(); if let Err(err) = (&ini_file).read_to_string(&mut s) { panic!("Failed to read patch.ron config file. \n{}", err); } let config: PatchConfig = match from_str(s.as_str()) { Ok(config) => config, Err(err) => panic!("Failed to load values from patch.ron \n{}",err), }; config }