ddtu/src/protocol/packets.rs

/*
ddtu - digital data transferring utility
Copyright (C) 2022  Kasyanov Nikolay Alexeyevich (Unbewohnte)

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.
*/

use std::net;
use std::io::Write;
use std::io::Read;

use crate::util::error::Error;
use crate::util::buf_read::{read_u128_slice_be, read_utf8_string_slice};
use crate::protocol::specs::*;
use crate::fsys::file::ChunkedFile;


pub fn send_packet(conn: &mut net::TcpStream, packet: &dyn Packet) -> Result<(), Error> {
    let mut payload: Vec<u8> = Vec::<u8>::new();

    let packet_bytes: Vec<u8> = packet.as_bytes();
    let packet_len: usize = packet_bytes.len();
    
    payload.extend_from_slice(&(packet_len as u128).to_be_bytes());
    payload.extend_from_slice(&packet_bytes);

    match conn.write_all(&payload) {
        Ok(()) => {
            return Ok(());
        }

        Err(error) => {
            return Err(Error::new(format!("could not write packet to the connection: {}", error).as_str()))
        }
    }
}

pub fn read_next_packet(conn: &mut net::TcpStream) -> Result<Box<dyn Packet>, Error> {
    let mut u128_buf: [u8; 16] = [0; 16];

    match conn.read_exact(&mut u128_buf) {
        Ok(()) => {}
        Err(error) => {
            return Err(Error::new(format!("error while reading initial length bytes: {}", error).as_str()));
        }
    }
    let packet_len: u128 = u128::from_be_bytes(u128_buf);

    let mut packet_bytes: Vec<u8> = vec!(0; packet_len as usize);
    match conn.read_exact(&mut packet_bytes) {
        Ok(()) => {}
        Err(error) => {
            return Err(Error::new(format!("error while reading packet contents: {}", error).as_str()));
        }
    }

    if packet_bytes.len() < packet_len as usize {
        return Err(
            Error::new(
                format!("read {} packet bytes instead of specified {}", packet_bytes.len(), packet_len).as_str())
            );
    }

    let packet_type_byte: u8 = u8::from_be_bytes([packet_bytes[0]]);
    match packet_type_byte {
        CONNECTION_SHUTDOWN_PACKET_ID => {
            return Ok(Box::new(ConnectionShutdown{}));
        }

        TEXT_PACKET_ID => {
            let mut new_text_packet: Text = Text::empty();
            let result = new_text_packet.from_bytes(packet_bytes);
            match result {
                Some(error) => {
                    return Err(Error::new(format!("could not construct new text packet: {}", error.text).as_str()));
                }

                None => {
                    return Ok(Box::new(new_text_packet));
                }
            }
        }

        FILEINFO_PACKET_ID => {
            let mut new_fileinfo_packet: FileInfo = FileInfo::empty();
            let result = new_fileinfo_packet.from_bytes(packet_bytes);
            match result {
                Some(error) => {
                    return Err(Error::new(format!("could not construct new fileinfo packet: {}", error.text).as_str()));
                }

                None => {
                    return Ok(Box::new(new_fileinfo_packet));
                }
            }
        }

        FILEDATA_PACKET_ID => {
            let mut new_filedata_packet: FileData = FileData::empty();
            let result = new_filedata_packet.from_bytes(packet_bytes);
            match result {
                Some(error) => {
                    return Err(Error::new(format!("could not construct new filedata packet: {}", error.text).as_str()));
                }

                None => {
                    return Ok(Box::new(new_filedata_packet));
                }
            }
        }

        _ => {
            return Err(Error::new(format!("invalid packet type ID \"{}\"", packet_type_byte).as_str()));
        }
    }
}


pub struct Text {
    pub text: String,
}

impl Text {
    pub fn empty() -> Text {
        return Text{
            text: "".to_string(),
        };
    }

    pub fn new(txt: &str) -> Text {
        return Text {
            text: txt.to_string(),
        }
    }
}

impl Packet for Text {
    fn get_type(&self) -> PacketType {
        return PacketType::TextData;
    } 

    fn as_bytes(&self) -> Vec<u8> {
        let mut packet_bytes: Vec<u8> = Vec::<u8>::new();

        packet_bytes.extend_from_slice(&TEXT_PACKET_ID.to_be_bytes());

        let text_length: u128 = self.text.len() as u128;
        packet_bytes.extend_from_slice(&text_length.to_be_bytes());
        packet_bytes.extend_from_slice(&self.text.as_bytes());

        return packet_bytes;
    }

    fn from_bytes(&mut self, packet_bytes: Vec<u8>) -> Option<Error> {
        if packet_bytes.len() < 18 {
            return Some(
                Error::new(format!("{} bytes is too small for a text packet to be valid", packet_bytes.len()).as_str()
            ));
        }

        // retrieve and check for packet type byte
        match packet_bytes[0].to_be() {
            TEXT_PACKET_ID => {}
            _ => {
                return Some(Error::new("packet type is not of a text packet"));
            }
        }

        // get text length
        let text_length: u128 = read_u128_slice_be(&packet_bytes[1..17]);
        
        if text_length as usize > packet_bytes[17..].len() {
            return Some(
                Error::new(
                    format!(
                        "text length ({}) is bigger than provided packet bytes length ({})",
                         text_length,
                         packet_bytes[16..].len()
                        ).as_str()
                    )
                );
        } 

        // extract text
        self.text = read_utf8_string_slice(&packet_bytes[17..]);

        return None;
    }
}


pub struct FileInfo {
    pub file_id: u128,
    pub filename: String,
    pub filesize: u128,
    pub relative_path: String,
}

impl FileInfo {
    pub fn empty() -> FileInfo {
        return FileInfo{
            file_id: 0,
            filename: String::new(),
            filesize: 0,
            relative_path: String::new(),
        };
    }

    pub fn new(file: &ChunkedFile, file_id: u128, rel_path: String) -> Result<FileInfo, Error> {
        let mut new_fileinfo_packet: FileInfo = FileInfo::empty();
        
        // id
        new_fileinfo_packet.file_id = file_id;
        
        //filename
        match file.path.file_name() {
            Some(filename) => {
                match filename.to_owned().into_string() {
                    Ok(filename_string) => {
                        new_fileinfo_packet.filename = filename_string;
                    }
                    Err(filename_os_string) => {
                        return Err(
                            Error::new(format!("could not convert \'{:?}\' (OsString) to String", filename_os_string).as_str())
                        );
                    }
                }
            }
            None => {
                return Err(
                    Error::new(format!("could not get filename from \'{}\'", file.path.display()).as_str())
                );
            }
        }

        // size
        new_fileinfo_packet.filesize =  file.size;

        // rel path
        new_fileinfo_packet.relative_path = rel_path;

        return Ok(new_fileinfo_packet);
    }
}

impl Packet for FileInfo {
    fn get_type(&self) -> PacketType {
        return PacketType::FileInfo;
    }

    fn as_bytes(&self) -> Vec<u8> {
        let mut packet_bytes: Vec<u8> = Vec::<u8>::new();

        packet_bytes.extend_from_slice(&FILEINFO_PACKET_ID.to_be_bytes());
        
        // file id
        packet_bytes.extend_from_slice(&self.file_id.to_be_bytes());

        // filename
        let filename_bytes_length: u128 = self.filename.len() as u128;
        packet_bytes.extend_from_slice(&filename_bytes_length.to_be_bytes());
        packet_bytes.extend_from_slice(&self.filename.as_bytes());

        // filesize
        packet_bytes.extend_from_slice(&self.filesize.to_be_bytes());

        // relative path
        let relative_path_bytes_length: u128 = self.relative_path.len() as u128;
        packet_bytes.extend_from_slice(&relative_path_bytes_length.to_be_bytes());
        packet_bytes.extend_from_slice(&self.relative_path.as_bytes());

        return packet_bytes;
    }

    fn from_bytes(&mut self, packet_bytes: Vec<u8>) -> Option<Error> {
        if packet_bytes.len() < 16+16+16+16+1+1 {
            return Some(
                Error::new(format!("{} bytes is too small for a fileinfo packet to be valid", packet_bytes.len()).as_str()
            ));
        }

        // file id
        self.file_id = read_u128_slice_be(&packet_bytes[0..16]);

        // get filename
        let filename_length: u128 = read_u128_slice_be(&packet_bytes[16..32]);
        self.filename = read_utf8_string_slice(&packet_bytes[32..(filename_length+32) as usize]);

        // filesize
        self.filesize = read_u128_slice_be(&packet_bytes[(filename_length+32) as usize..(filename_length+32+16) as usize]);

        // relative path
        let rel_path_length: u128 = read_u128_slice_be(
            &packet_bytes[(filename_length+32+16) as usize..(filename_length+32+16+16)as usize]
        );
        self.relative_path = read_utf8_string_slice(
            &packet_bytes[(filename_length+32+16+16) as usize..(filename_length+32+16+16+rel_path_length) as usize]
        );

        return None;
    }
}


pub struct FileData {
    pub file_id: u128,
    pub chunk_no: u128,
    pub chunk: [u8; CHUNK_SIZE],
}

impl FileData {
    pub fn empty() -> FileData {
        return FileData {
            file_id: 0,
            chunk_no: 0,
            chunk: [0; CHUNK_SIZE],
        }
    }
}

impl Packet for FileData {
    fn get_type(&self) -> PacketType {
        return PacketType::FileData;
    }

    fn as_bytes(&self) -> Vec<u8> {
        let mut packet_bytes: Vec<u8> = Vec::<u8>::new();
        
        packet_bytes.extend_from_slice(&FILEDATA_PACKET_ID.to_be_bytes());

        // file id
        packet_bytes.extend_from_slice(&self.file_id.to_be_bytes());

        // chunk number
        packet_bytes.extend_from_slice(&self.chunk_no.to_be_bytes());

        // chunk
        packet_bytes.extend_from_slice(&self.chunk.len().to_be_bytes());
        packet_bytes.extend_from_slice(&self.chunk);

        return packet_bytes;
    }

    fn from_bytes(&mut self, packet_bytes: Vec<u8>) -> Option<Error> {
        if packet_bytes.len() < 16+16+16+1 {
            return Some(
                Error::new(format!("{} bytes is too small for a fileinfo packet to be valid", packet_bytes.len()).as_str()
            ));
        }

        // file id
        self.file_id = read_u128_slice_be(&packet_bytes[0..16]);
        
        // chunk number
        self.chunk_no = read_u128_slice_be(&packet_bytes[16..32]);

        // chunk bytes
        let chunk_length = read_u128_slice_be(&packet_bytes[32..48]);
        for i in 48..(48+chunk_length) as usize {
            self.chunk[i-48] = packet_bytes[i].to_be();
        } 
        
        return None;
    }
}

pub struct ConnectionShutdown {}

impl Packet for ConnectionShutdown {
    fn get_type(&self) -> PacketType {
        return PacketType::ConnectionShutdown;
    }

    fn as_bytes(&self) -> Vec<u8> {
        return vec!(CONNECTION_SHUTDOWN_PACKET_ID.to_be_bytes()[0]);
    }

    fn from_bytes(&mut self, _packet_bytes: Vec<u8>) -> Option<Error> {
        return None;
    }
}

pub struct Handshake {
    pub protocol_version: u8,
    pub encryption_type: u8,
    pub encryption_key: Vec<u8>,
}

impl Handshake {
    pub fn empty() -> Handshake {
        return Handshake{
            protocol_version: PROTOCOL_LATEST_VERSION,
            encryption_type: ENCRYPTION_NO_ENCRYPTION,
            encryption_key: vec!(0, 0, 0, 0, 0, 0, 0, 0),
        };
    }
}

impl Packet for Handshake {
    fn get_type(&self) -> PacketType {
        return PacketType::Handshake;
    }

    fn as_bytes(&self) -> Vec<u8> {
        let mut packet_bytes: Vec<u8> = Vec::<u8>::new();

        packet_bytes.extend_from_slice(&HANDSHAKE_PACKET_ID.to_be_bytes());
        
        // protocol version
        packet_bytes.extend_from_slice(&PROTOCOL_LATEST_VERSION.to_be_bytes());

        // encryption type
        packet_bytes.extend_from_slice(&self.encryption_type.to_be_bytes());

        // encryption key
        packet_bytes.extend_from_slice(&(self.encryption_key.len() as u128).to_be_bytes());
        packet_bytes.extend_from_slice(&self.encryption_key);

        return packet_bytes;
    }

    fn from_bytes(&mut self, packet_bytes: Vec<u8>) -> Option<Error> {
        if packet_bytes.len() < 1+1+1+16+1 {
            return Some(
                Error::new(format!("{} bytes is too small for a handshake packet to be valid", packet_bytes.len()).as_str())
            );
        }

        // protocol version
        self.protocol_version = packet_bytes[0].to_be();

        // encryption type
        self.encryption_type = packet_bytes[1].to_be();

        // key
        let key_len: u128 = read_u128_slice_be(&packet_bytes[1..17]);
        for i in 17..(17+key_len) as usize {
            self.encryption_key[i-17] = packet_bytes[i].to_be();
        }

        return None;
    }
}