elseware/src/patch/main.rs

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};
use std::path::PathBuf;
use rand::{Rng, RngCore};
use crc::{crc32, Hasher32};
use mio::tcp::TcpStream;
use libpso::{PacketParseError, PSOPacket};
use libpso::packet::patch::*;
use libpso::crypto::pc::PSOPCCipher;
use elseware::common::network::{PacketNetworkError};
use elseware::common::client::Client;
use elseware::common::serverstate::{RecvServerPacket, SendServerPacket, ServerState, OnConnect};

const PATCH_PORT: u16 = 11000;

#[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, 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)
    }
}


#[derive(Debug)]
pub enum RecvPatchPacket {
    PatchWelcomeReply(PatchWelcomeReply),
    LoginReply(LoginReply),
    FileInfoReply(FileInfoReply),
    FileInfoListEnd(FileInfoListEnd),
}

impl RecvServerPacket for RecvPatchPacket {
    fn from_bytes(data: &Vec<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)
        }
    }
}

#[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(),
        }
    }
}


struct PatchServerState {
    patch_file_tree: PatchFileTree,
    patch_file_lookup: HashMap<u32, PatchFile>,
    patch_file_info: Vec<FileInfoReply>,
}

impl PatchServerState {
    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) -> 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, pkt: &RecvPatchPacket) -> Box<dyn Iterator<Item = SendPatchPacket>> {
        match pkt {
            RecvPatchPacket::PatchWelcomeReply(_pkt) => {
                Box::new(vec![SendPatchPacket::RequestLogin(RequestLogin {})].into_iter())
            },
            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())
            },
            RecvPatchPacket::FileInfoReply(pkt) => {
                self.patch_file_info.push(pkt.clone());
                Box::new(None.into_iter())
            },
            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)))
            }
        }
    }
}

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)
}

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 {}))
                    }
                }
            }
        }
    }
}



fn new_client(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap<u32, PatchFile>) {
    let state = PatchServerState::new(patch_file_tree, patch_file_lookup);
    let client = Client::new(TcpStream::from_stream(socket).unwrap(), Box::new(state));
    client.io_loop();
}

fn main() {
    println!("[patch] starting server");

    let (patch_file_tree, patch_file_lookup) = generate_patch_tree("patchfiles/");
    println!("[patch] files to patch:");
    let mut indent = 0;
    for item in patch_file_tree.flatten() {
        match item {
            PatchTreeIterItem::Directory(path) => {
                let s = path.to_str().unwrap();
                println!("{: >2$}\u{2517}\u{2500}\u{2500} {}", "", s, indent * 4);
                indent += 1;
            },
            PatchTreeIterItem::File(path, id) => {
                let s = path.to_str().unwrap();
                println!("{: >3$}\u{2520}\u{2500}\u{2500} {} ({})", "", s, id, indent * 4);
            },
            PatchTreeIterItem::UpDirectory => {
                indent -= 1;
            }
        }
    }

    let listener = TcpListener::bind(&SocketAddr::from((Ipv4Addr::new(0,0,0,0), PATCH_PORT))).unwrap();
    
    println!("[patch] waiting for connections");
    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();
        thread::spawn(move || {
            println!("[patch] accepted connection: {}", addr);
            new_client(socket, local_patch_file_tree, local_patch_file_lookup);
        });
    }

    println!("[patch] exiting...");
}