elseware/src/patch/main.rs

#![allow(unused_imports)]

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};
use libpso::crypto::pc::PSOPCCipher;
use elseware::common::{send_packet, recv_packet, PacketNetworkError};

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)]
pub enum PatchPacket {
    PatchWelcomeReply(PatchWelcomeReply),
    LoginReply(LoginReply),
    FileInfoReply(FileInfoReply),
    FileInfoListEnd(FileInfoListEnd),
}

impl PatchPacket {
    pub fn from_bytes(data: &Vec<u8>) -> Result<PatchPacket, PacketParseError> {
        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<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)
    }
}



struct Client {
    running: bool,
    key_in: u32,
    key_out: u32,
    socket: mio::tcp::TcpStream,
    cipher_in: Box<dyn PSOCipher>,
    cipher_out: Box<dyn PSOCipher>,
    patch_file_tree: PatchFileTree,
    patch_file_lookup: HashMap<u32, PatchFile>,
    patch_file_info: Vec<FileInfoReply>,
}

impl Client {
    fn new(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap<u32, PatchFile>) -> Client {
        let mut rng = rand::thread_rng();
        let key_in: u32 = rng.gen();
        let key_out: u32 = rng.gen();

        Client {
            running: true,
            key_in: key_in,
            key_out: key_out,
            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()
        }
    }

    fn send(&mut self, pkt: &dyn PSOPacket) {
        match send_packet(&mut self.socket, &mut *self.cipher_out, pkt) {
            Ok(_) => {
                println!("[patch] send ({:?}): {:?}", self.socket, pkt);
            },
            Err(err) => {
                println!("[patch] error sending packet to {:?}: {:?}", self.socket, err);
                self.running = false;
            }
        }
    }

}

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


// 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<Box<dyn PSOPacket>> {
    let mut pkts: Vec<Box<dyn PSOPacket>> = Vec::new();

    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
}


fn send_file_list(client: &mut Client) -> Result<(), PatchError> {
    client.send(&PatchStartList {});

    let pkts = get_file_list_packets(&client.patch_file_tree);
    for pkt in pkts {
        client.send(&*pkt);
    }

    client.send(&PatchEndList {});

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

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::<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 file_ids_to_update = need_update.iter().map(|k| k.id).collect::<HashSet<_>>();

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


fn handle_packet(client: &mut Client, pkt: PatchPacket) -> Result<(), PatchError> {
    println!("[patch] recv({:?}): {:?}", client.socket, pkt);

    match pkt {
        PatchPacket::PatchWelcomeReply(_pkt) => {
            client.send(&RequestLogin {});
        },
        PatchPacket::LoginReply(_pkt) => {
            client.send(&Message::new("hello player".to_string()));
            send_file_list(client)?;
        },
        PatchPacket::FileInfoReply(pkt) => {
            client.patch_file_info.push(pkt);
        },
        PatchPacket::FileInfoListEnd(_pkt) => {
            send_file_data(client)?;
        }
    }

    Ok(())
}

fn client_loop(mut client: Client) {
    let poll = mio::Poll::new().unwrap();
    poll.register(&client.socket, Token(0), Ready::readable(), PollOpt::edge()).unwrap();

    let mut events = Events::with_capacity(1024);
    loop {
        poll.poll(&mut events, None).unwrap();

        for event in &events{
            println!("event! {:?}", event);
            if event.token() == Token(0) {
                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;
                        }
                    }
                }
            }
        }
    }
}

fn new_client(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap<u32, PatchFile>) {
    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);
    client.cipher_in = Box::new(PSOPCCipher::new(client.key_in));
    client.cipher_out = Box::new(PSOPCCipher::new(client.key_out));
    client_loop(client);
}


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

    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();
        thread::spawn(move || {
            println!("[patch] accepted connection: {}", addr);
            new_client(socket, local_patch_file_tree, local_patch_file_lookup);
        });
    }


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