First of all, I am here to explore unsafe features of Rust, if you are the kind of guy who likes "safe" Rust and will start nagging about using safe features of Rust I will still happily read your comment and learn about the safe alternatives.
That said, I am looking to enhance my knowledge about unsafe Rust. The documentation is very scarce and not very well documented. Little by little, thanks to help from Stack Overflow users (especially @kmdreko) and my perseverance I was able to make it work. My program is able to send large files (4Gb) over UDP.
As I always look to improve the quality of my code I wonder if someone could take a look at my source code and let me know how I could improve it.
server.rs
const UDP_HEADER: usize = 8;
const IP_HEADER: usize = 20;
const AG_HEADER: usize = 4;
const MAX_CHUNK_SIZE: usize = (64 * 1024 - 1) - UDP_HEADER - IP_HEADER - AG_HEADER;
use std::net::UdpSocket;
use std::io;
use std::fs::File;
use std::io::prelude::*;
use std::alloc::{alloc, dealloc, Layout};
use std::mem;
use std::{mem::MaybeUninit};
// cmp -l 1.jpg 2.jpg
#[inline(always)]
fn memcpy(dst_ptr:*mut u8, src_ptr:*const u8, len:usize) {
unsafe {
std::ptr::copy_nonoverlapping(src_ptr, dst_ptr, len);
}
}
#[inline(always)]
fn next_power_of_two(n:u32) -> u32 {
return 32 - (n - 1).leading_zeros();
}
#[inline(always)]
fn write_chunks_to_file(filename: &str, bytes:&[u8]) -> Result<bool, io::Error> {
let mut file = File::create(filename)?;
file.write_all(bytes)?;
Ok(true)
}
use std::thread;
fn main() {
let socket = UdpSocket::bind("0.0.0.0:8888").expect("Could not bind socket");
let filename = "2.jpg";
let mut count = 0;
let mut chunks_cnt:u32 = 0xffff;
let mut total_size:usize = 0;
let mut layout;
unsafe { layout = MaybeUninit::<Layout>::uninit().assume_init(); };
let mut bytes_buf;
unsafe { bytes_buf = MaybeUninit::<*mut u8>::uninit().assume_init(); };
loop {
let mut buf = [0u8; MAX_CHUNK_SIZE + AG_HEADER];
let sock = socket.try_clone().expect("Failed to clone socket");
match socket.recv_from(&mut buf) {
Ok((size, src)) => { // thanks https://doc.rust-lang.org/beta/std/net/struct.UdpSocket.html#method.recv_from
total_size += size;
let packet_index:usize = (buf[0] as usize) << 8 | buf[1] as usize;
if count == 0 {
chunks_cnt = (buf[2] as u32) << 8 | buf[3] as u32;
let n:usize = 0x10000 << next_power_of_two(chunks_cnt);
// assert_eq!(n.count_ones(), 1); // can check with this function that n is aligned on power of 2
unsafe {
layout = Layout::from_size_align_unchecked(n, mem::align_of::<u8>());
bytes_buf = alloc(layout);
}
}
unsafe {
let dst_ptr = bytes_buf.offset((packet_index*MAX_CHUNK_SIZE) as isize);
memcpy(dst_ptr, &buf[AG_HEADER], size-AG_HEADER);
};
thread::spawn(move || {
//let s = String::from_utf8_lossy(&buf);
println!("receiving packet {} from: {}", packet_index, src);
sock.send_to(&buf, &src).expect("Failed to send a response");
});
println!("count: {}", count);
count+=1;
}
Err(e) => {
eprintln!("couldn't recieve a datagram: {}", e);
}
}
if count == chunks_cnt { // all chunks have been collected, write bytes to file
let bytes = unsafe { std::slice::from_raw_parts(bytes_buf, total_size) };
let result = write_chunks_to_file(filename, &bytes);
match result {
Ok(true) => println!("Succesfully created file: {}", true),
Ok(false) => println!("Could not create file: {}", false),
Err(e) => println!("Error: {}", e),
}
count = 0;
total_size = 0;
unsafe { dealloc(bytes_buf, layout); }
}
}
}
client.rs
const UDP_HEADER: usize = 8;
const IP_HEADER: usize = 20;
const AG_HEADER: usize = 4;
const MAX_DATA_LENGTH: usize = (64 * 1024 - 1) - UDP_HEADER - IP_HEADER;
use std::io::Read;
use std::net::UdpSocket;
use std::io;
pub fn get_chunks_from_file(mut filename: String,total_size: &mut usize) -> Result<Vec<Vec<u8>>, io::Error> {
filename.pop(); // get read of the trailing '\n' in user input.
let mut file = std::fs::File::open(filename)?;
let mut list_of_chunks = Vec::new();
let chunk_size = MAX_DATA_LENGTH - AG_HEADER;
loop {
let mut chunk = Vec::with_capacity(chunk_size);
let n = file
.by_ref()
.take(chunk_size as u64)
.read_to_end(&mut chunk)?;
*total_size += n;
if n == 0 {
break;
}
list_of_chunks.push(chunk);
if n < chunk_size {
break;
}
}
Ok(list_of_chunks)
}
fn main() {
let socket = UdpSocket::bind("127.0.0.1:8000").expect("Could not bind client socket");
let mut buffer = [0u8; MAX_DATA_LENGTH];
socket.connect("127.0.0.1:8888").expect("Could not connect to server");
loop {
let mut input = String::new();
io::stdin()
.read_line(&mut input)
.expect("Failed to read from stdin");
println!("{}", input);
let mut total_size: usize = 0;
let result: Result<Vec<Vec<u8>>, io::Error> = get_chunks_from_file(input, &mut total_size); // : Result<u8:u8>
match result {
Ok(chunks) => {
let nb: u16 = chunks.len() as u16;
let mut index: u16 = 0;
let header: &mut[u8;4] = &mut[
(index >> 8) as u8,
(index & 0xff) as u8,
(nb >> 8) as u8,
(nb & 0xff) as u8,
]; //input.as_bytes();
for chunk in chunks.iter() {
header[1] = (index & 0xff) as u8;
header[0] = (index >> 8) as u8;
let data:Vec<u8> = [header.as_ref(), chunk].concat();
//println!("FILE {} BYTES\n {:?}", index, chunk);
println!(
"size: {} FILE {:?} of {} BYTES\n {:?}",
total_size,
(header[0] as u16) << 8 | header[1] as u16,
nb - 1,
[0]
);
println!("{}", index);
socket.send(&data).expect("Failed to write to server");
socket.recv_from(&mut buffer).expect("Could not read into buffer");
index += 1;
}
}
Err(e) => println!("Error: {}", e),
}
//print!( "{}",str::from_utf8(&buffer).expect("Could not write buffer as string"));
// println!( "Chunk received by server {:?}", &buffer);
}
}
The ones I am aware of
There are a few things I am especially concerned and already know that could be improved like this function:
fn write_chunks_to_file(filename: &str, bytes:&[u8]) -> Result<bool, io::Error> {
let mut file = File::create(filename)?;
file.write_all(bytes)?;
Ok(true)
}
I wrote it this way in order to make the compilation works but even I, the creator, can acknowledge that it does not make any sense! What would be a good Result to return?
And also the following warning when I compile the server file:
warning: the type `std::alloc::Layout` does not permit being left uninitialized
--> udp-server.rs:43:23
|
43 | unsafe { layout = MaybeUninit::<Layout>::uninit().assume_init(); };
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
| |
| this code causes undefined behavior when executed
| help: use `MaybeUninit<T>` instead, and only call `assume_init` after initialization is done
|
= note: `#[warn(invalid_value)]` on by default
note: `std::num::NonZeroUsize` must be non-null (in this struct field)
Kindly appreciating any peer review.
