vec_router 0.0.3

use std::net::TcpStream;
use std::io::{Read, Write};


/// Reads network traffic sent by the send_network_traffic function. First,
/// a byte is sent denoting the number of byte sequences that will be read.
/// Next, the incoming byte count of the first section is received, followed
/// by the bytes being received. This happens for each section, until all
/// byte sequences are received. 
pub fn read_network_traffic(mut stream: TcpStream, 
                            max_section_count: u8,
                            max_byte_count: usize) -> Result< Vec<Vec<u8>>, &'static str > {
    let mut byte_sequences = Vec::<Vec<u8>>::new();

    let mut section_count: [u8; 1] = [0; 1];

    match stream.read_exact(&mut section_count) {
        Ok(()) => {}
        Err(e) => panic!("Error: {:#?}", e),
    }
    
    if section_count > [max_section_count] {
        return Err("maximum section count exceeded");
    }

    while section_count[0] > 0 {
        section_count[0] -= 1;

        let mut incoming_byte_count: [u8; 8] = [0; 8];

        match stream.read_exact(&mut incoming_byte_count) {
            Ok(()) => {}
            Err(e) => panic!("Error: {:#?}", e),
        }

        let incoming = u64::from_be_bytes(incoming_byte_count);

        let mut incoming = usize::try_from(incoming).unwrap();
        
        if incoming > max_byte_count {
            return Err("Maximum byte count exceeded");
        }

        let mut vec = vec![0; incoming];

        while incoming > 0 {
            incoming -= stream.read(&mut vec).unwrap();
        }

        byte_sequences.push(vec);
    }
    Ok( byte_sequences )
}

/// Sends network traffic to the read_network_traffic function.
pub fn send_network_traffic(
    mut socket: TcpStream,
    byte_sequences: Vec<Vec<u8>>,
) -> TcpStream {
    // The section count is the length of the vector
    let section_count = byte_sequences.len();

    //transmit section count
        socket
        .write(&[section_count as u8])
        .unwrap();

    // transmit byte_sequences byte count
    for n in 0..section_count {
        let byte_sequences_len: [u8; 8] = byte_sequences[n as usize].len().to_be_bytes();
        
        socket.write(&byte_sequences_len).unwrap();
        
             socket
            .write(byte_sequences[n as usize].as_slice())
            .unwrap();
    }

    // return the connected stream
    socket
}