You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

397 lines
14 KiB

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<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, 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<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)
}
}
#[derive(Debug)]
pub enum RecvPatchPacket {
PatchWelcomeReply(PatchWelcomeReply),
LoginReply(LoginReply),
FileInfoReply(FileInfoReply),
FileInfoListEnd(FileInfoListEnd),
}
impl RecvServerPacket for RecvPatchPacket {
fn from_bytes(data: &[u8]) -> Result<RecvPatchPacket, PacketParseError> {
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<u8> {
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<u32, PatchFile>,
patch_file_info: Vec<FileInfoReply>,
}
impl PatchServerState {
pub fn new(patch_file_tree: PatchFileTree, patch_file_lookup: HashMap<u32, PatchFile>) -> 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<OnConnect<Self::SendPacket>> {
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<Box<dyn Iterator<Item = (ClientId, SendPatchPacket)>>, 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::<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 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<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.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<u32, PatchFile>) {
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<SendPatchPacket> {
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<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
}
struct SendFileIterator {
done: bool,
file_iter: Box<dyn Iterator<Item = PatchTreeIterItem>>,
file_ids_to_update: HashSet<u32>,
patch_file_lookup: HashMap<u32, PatchFile>,
current_file: Option<io::BufReader<fs::File>>,
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::<HashSet<_>>();
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<Self::Item> {
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
}