Make a more reliable buffer overflow example

also more funny

Co-authored-by: Colin Cai <therealcreative0708@gmail.com>

src/buffer_overflow.rs

@@ -1,45 +1,75 @@
 //! A memory-safe buffer overflow.
 //!
-//! We allocate a slice in a stack, then transmute it into a String with a large capacity.
-//! Then, we read input from stdin. This overwrites another buffer, and then we can check
-//! if it's successfully overwritten.
+//! We allocate a slice on the stack, then transmute it into a String with a large capacity.
+//! Then, we read input from stdin into that String. This overwrites another stack-allocated
+//! slice, and then we can check if it's successfully overwritten.
 
-use std::io::{self, Write};
+use std::io::{stdin, stdout, Write};
+use std::time::Duration;
+use std::{io, mem, thread};
 
+/// Construct a [`String`] from a pointer, capacity and length, in a completely safe manner.
+///
+/// [`String`] is a `Vec<u8>` which is a `(RawVec, usize)` which is a `((Unique, usize), usize)`.
+///
+/// Rust explicitly says that structs are not guaranteed to have members in order,
+/// so instead we determine that order at runtime.
+#[inline(always)]
+pub fn construct_fake_string(ptr: *mut u8, cap: usize, len: usize) -> String {
+	let sentinel_string = crate::transmute::<_, String>([0usize, 1usize, 2usize]);
+
+	let fields = [ptr as usize, cap, len];
+	let mut actual_buf = [0usize; 3];
+	actual_buf[0] = fields[sentinel_string.as_ptr() as usize];
+	actual_buf[1] = fields[sentinel_string.capacity()];
+	actual_buf[2] = fields[sentinel_string.len()];
+
+	mem::forget(sentinel_string);
+
+	crate::transmute::<_, String>(actual_buf)
+}
+
+#[inline(never)]
 pub fn buffer_overflow() -> io::Result<()> {
 	use std::hint::black_box;
 
-	let mut name_buf = black_box([0u8; 16]);
-	let mut password = black_box([0u8; 8]);
+	#[repr(C)]
+	#[derive(Default)]
+	struct Authentication {
+		name_buf: [u8; 16],
+		password: [u8; 16],
+	}
+
+	let mut auth = black_box(Authentication::default());
+
+	let mut name = construct_fake_string(auth.name_buf.as_mut_ptr(), 1024usize, 0usize);
+
+	print!("Hello! What's your name? > ");
+	stdout().flush()?;
+	stdin().read_line(&mut name)?;
 
-	// String is a Vec<u8> which is a (RawVec, usize) which is a ((Unique, usize), usize)
-	// Rust explicitly says that structs are not guaranteed to have members in order
-	// but this is sketchy enough anyways so :shrug:
-	let mut name = crate::transmute::<_, String>((name_buf.as_mut_ptr(), 128usize, 0usize));
+	// If we don't forget our fake String, Rust will try to deallocate it as if it was a heap pointer.
+	mem::forget(name);
 
-	let mut stdout = std::io::stdout();
-	stdout.write_all(b"Hello! What's your name? ")?;
-	stdout.flush()?;
-	io::stdin().read_line(&mut name)?;
+	let password = &auth.password[0..8];
 
 	if password.iter().all(|&x| x == 0) {
-		stdout.write_all(b"You didn't even modify the password...\n")?;
+		println!("You didn't even modify the password...");
 	} else if &password != b"letmein!" {
-		stdout.write_all(b"Wrong password! You entered: ")?;
-		stdout.write_all(&password)?;
-		stdout.write_all(b"\n")?;
+		println!(
+			"Wrong password! You entered: {:?}",
+			std::str::from_utf8(password).unwrap()
+		);
 	} else {
 		#[cfg(unix)]
-		stdout.write_all(b"Correct password, running sudo rm -rf /* ...\n")?;
+		println!("Correct password, running sudo rm -rf /* ...");
 		#[cfg(windows)]
-		stdout.write_all(b"Correct password, deleting C:\\Windows\\System32 ...\n")?;
+		println!("Correct password, deleting C:\\Windows\\System32 ...");
 
-		std::thread::sleep(std::time::Duration::from_secs(2));
+		thread::sleep(Duration::from_secs(2));
 	}
-	stdout.flush()?;
 
-	black_box((&mut name_buf, &mut password));
-	std::mem::forget(name);
+	black_box(auth);
 
 	Ok(())
 }