elseware/src/patch_server/src/lib.rs

use std::collections::{HashMap, HashSet};
use std::fs;
use std::io;
use std::io::{Read};
use std::path::{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 networking::mainloop::{NetworkError};
use networking::serverstate::{RecvServerPacket, SendServerPacket, ServerState, OnConnect, ClientId};

#[derive(Debug)]
pub enum PatchError {
    NetworkError(NetworkError),
    IOError(std::io::Error),
}

impl From<NetworkError> for PatchError {
    fn from(err: NetworkError) -> PatchError {
        PatchError::NetworkError(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(Box<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(),
        }
    }
}


#[derive(Clone)]
pub struct PatchServerState {
    patch_file_tree: PatchFileTree,
    patch_file_lookup: HashMap<u32, PatchFile>,
    patch_file_info: Vec<FileInfoReply>,
    patch_motd: String,
}

impl PatchServerState {
    pub fn new(patch_file_tree: PatchFileTree, patch_file_lookup: HashMap<u32, PatchFile>, patch_motd: String) -> PatchServerState {
        PatchServerState {
            patch_file_tree,
            patch_file_lookup,
            patch_file_info: Vec::new(),
            patch_motd,
        }
    }
}

#[async_trait::async_trait]
impl ServerState for PatchServerState {
    type SendPacket = SendPatchPacket;
    type RecvPacket = RecvPatchPacket;
    type Cipher = PSOPCCipher;
    type PacketError = PatchError;

    async fn on_connect(&mut self, _id: ClientId) -> Result<Vec<OnConnect<Self::SendPacket, Self::Cipher>>, PatchError> {
        let mut rng = rand::thread_rng();
        let key_in: u32 = rng.gen();
        let key_out: u32 = rng.gen();

        Ok(vec![OnConnect::Packet(SendPatchPacket::PatchWelcome(PatchWelcome::new(key_out, key_in))),
                OnConnect::Cipher(PSOPCCipher::new(key_in), PSOPCCipher::new(key_out))
        ])
    }

    async fn handle(&mut self, id: ClientId, pkt: RecvPatchPacket) -> Result<Vec<(ClientId, SendPatchPacket)>, PatchError> {
        Ok(match pkt {
            RecvPatchPacket::PatchWelcomeReply(_pkt) => {
                vec![SendPatchPacket::RequestLogin(RequestLogin {})]
                    .into_iter()
                    .map(move |pkt| (id, pkt))
                    .collect()
            },
            RecvPatchPacket::LoginReply(_pkt) => {
                let mut pkts = vec![SendPatchPacket::Message(Message::new(self.patch_motd.clone()))];
                pkts.append(&mut get_file_list_packets(&self.patch_file_tree));
                pkts.push(SendPatchPacket::PatchEndList(PatchEndList {}));
                pkts
                    .into_iter()
                    .map(move |pkt| (id, pkt))
                    .collect()
            },
            RecvPatchPacket::FileInfoReply(pkt) => {
                self.patch_file_info.push(pkt);
                Vec::new()
            },
            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;

                vec![SendPatchPacket::FilesToPatchMetadata(FilesToPatchMetadata::new(total_size, total_files)),
                     SendPatchPacket::PatchStartList(PatchStartList {})]
                    .into_iter()
                    .chain(SendFileIterator::new(self))
                    .map(move |pkt| (id, pkt))
                    .collect()
            }
        })
    }

    async fn on_disconnect(&mut self, _id: ClientId) -> Result<Vec<(ClientId, SendPatchPacket)>, PatchError> {
        Ok(Vec::new())
    }
}

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();
        let patch_path = path.strip_prefix(basedir).unwrap();
        if path.is_dir() {
            dirs.push(load_patch_dir(path.to_str().unwrap(), patch_path.to_str().unwrap(), file_ids));
        }
        else {
            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,
                checksum,
                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: &Path) -> 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> + Send>,
    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,
            patch_file_lookup: state.patch_file_lookup.clone(),
            file_iter: Box::new(state.patch_file_tree.flatten().into_iter().filter(move |file| {
                match file {
                    PatchTreeIterItem::File(_path, id) => {
                        file_ids_to_update.contains(id)
                    },
                    _ => true,
                }
            })),
            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(Box::new(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) => {
                                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)))
                            },
                            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, // TODO: this does nothing
    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
}

pub fn load_config_env() -> PatchConfig {
    let patch_path = std::env::var("PATCHFILE_DIR").unwrap();
    let patch_port = std::env::var("PATCH_PORT").unwrap().parse().unwrap();

    PatchConfig {
        path: patch_path,
        ip: "127.0.0.1".into(),
        port: patch_port,
    }
}

pub fn load_motd() -> String {
    if let Ok(m) = fs::read_to_string("patch.motd") {
        m
    }
    else {
        "Welcome to Elseware!".to_string()
    }
}