<commit_before>/* A very simple front end to test agb-core (https://github.com/aHeraud/agb)*/

extern crate agb_core;
extern crate sdl2;
extern crate image;
extern crate clap;

mod debugger;

use std::sync::mpsc::sync_channel;
use std::thread;
use std::time::Duration;
use std::thread::sleep;
use std::collections::HashMap;
use std::fs::File;
use std::io::{stdin, stdout, Read, Write, Error};
use std::path::Path;
use std::num::ParseIntError;
use std::sync::{Arc, Mutex};
use std::ops::DerefMut;
use std::net::{TcpListener, TcpStream, SocketAddr, IpAddr, Ipv4Addr};

use agb_core::gameboy::Gameboy;
use agb_core::gameboy::debugger::DebuggerInterface;

use sdl2::pixels::Color;
use sdl2::pixels::PixelFormatEnum;
use sdl2::rect::Rect;
use sdl2::event::Event;
use sdl2::keyboard::Keycode;

use clap::{Arg, App};

const DEFAULT_SCALE: usize = 2;

fn main() {
	/* Create and initialize gameboy */

	//Parse command line arguments
	let matches = App::new("agb")
		.version("0.1")
		.author("Achille Heraud <achille@heraud.xyz>")
		.about("A GameBoy Emulator")
		.arg(Arg::with_name("rom")
			.long("rom")
			.takes_value(true)
			.value_name("FILE")
			.required(true))
		.arg(Arg::with_name("ram")
			.long("ram")
			.takes_value(true)
			.value_name("FILE")
			.required(false))
		.arg(Arg::with_name("paused")
			.long("pause")
			.short("p")
			.required(false))
		.arg(Arg::with_name("print_serial")
			.help("print data written to the serial port")
			.long("print_serial")
			.conflicts_with_all(&["listen", "connect"])
			.required(false))
		.arg(Arg::with_name("listen")
			.help("listen for remote connections (for serial communication over tcp/ip)")
			.long("listen")
			.takes_value(true)
			.value_name("PORT")
			.conflicts_with("connect")
			.required(false))
		.arg(Arg::with_name("connect")
			.help("provide to ip:port to connect to a remote host (for serial communication over tcp/ip")
			.long("connect")
			.takes_value(true)
			.value_name("IP:PORT")
			.required(false))
		.get_matches();

	let rom = read_file(matches.value_of("rom").unwrap()).expect("Could not open rom file.");
	let ram: Option<Box<[u8]>> = if let Some(ram_path) = matches.value_of("ram") {
		Some(read_file(ram_path).expect("failed to read ram file"))
	}
	else {
		None
	};

	let start_paused: bool = matches.occurrences_of("paused") > 0;

	let mut gameboy = agb_core::init(rom, ram).expect("Failed to initialize gameboy");
	let paused: Arc<Mutex<bool>> = Arc::new(Mutex::new(start_paused));
	gameboy.debugger.enable();
	{
		let paused = paused.clone();
		gameboy.register_breakpoint_callback(move |breakpoint| {
			println!("triggered breakpoint access_type: {:?}, address: 0x{:x}", breakpoint.access_type, breakpoint.address);
			let mut paused = paused.lock().unwrap();
			*paused = true;
		});
	}

	if let Some(ref port_str) = matches.value_of("listen") {
		// set up a tcp socket to accept incoming connections
		if let Ok(port) = u16::from_str_radix(port_str, 10) {
			let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127,0,0,1)), port);
			let listener = TcpListener::bind(addr).unwrap();
			let (local_sender, local_receiver) = gameboy.create_serial_channels();

			println!("waiting for connection...");
			let (mut stream, remote_addr) = listener.accept().unwrap();
			println!("connected to {}", remote_addr);

			stream.set_nonblocking(false).expect("failed to set tcp socket to nonblocking mode");
			stream.set_nodelay(true).expect("failed to set tcp socket nodelay mode");

			//create sender thread
			{
				let mut stream = stream.try_clone().unwrap();
				thread::spawn(move || {
					let mut buffer: [u8; 1] = [0;1];
					loop {
						let data = local_receiver.recv().unwrap();
						buffer[0] = data;
						if stream.write(&buffer[..]).is_err() {
							return;
						}
					}
				});
			}

			//create receiver thread
			thread::spawn(move || {
				let mut buffer: [u8; 1] = [0;1];
				loop {
					if stream.read(&mut buffer[..]).is_err() {
						return;
					}
					local_sender.send(buffer[0]).unwrap();
				}
			});
		}
		else {
			println!("Invalid port: port must be an integer in the range 0 to 65535.");
			return;
		}
	}
	else if let Some(ref addr) = matches.value_of("connect") {
		let (local_sender, local_receiver) = gameboy.create_serial_channels();
		let addr: SocketAddr = addr.parse().unwrap();
		println!("attempting to connect...");
		let mut stream = TcpStream::connect(addr).unwrap();
		println!("connected!");

		stream.set_nonblocking(false).expect("failed to set tcp socket to nonblocking mode");
		stream.set_nodelay(true).expect("failed to set tcp socket nodelay mode");

		//create sender thread
		{
			let mut stream = stream.try_clone().unwrap();
			thread::spawn(move || {
				let mut buffer: [u8; 1] = [0;1];
				loop {
					let data = local_receiver.recv().unwrap();
					buffer[0] = data;
					if stream.write(&buffer[..]).is_err() {
						return;
					}
				}
			});
		}

		//create receiver thread
		thread::spawn(move || {
			let mut buffer: [u8; 1] = [0;1];
			loop {
				if stream.read(&mut buffer[..]).is_err() {
					return;
				}
				local_sender.send(buffer[0]).unwrap();
			}
		});
	}
	else if matches.occurrences_of("print_serial") > 0 {
		let (sender, reciever) = gameboy.create_serial_channels();
		thread::spawn(move || {
			loop {
				let data = reciever.recv().unwrap();
				sender.send(0xFF).unwrap();
				print!("{}", data as char);
				stdout().flush().ok().expect("could not flush stdout");
			}
		});
	}

	//debugger text input
	let (tx, rx) = sync_channel(0);
	let main_handle = thread::current();
	thread::spawn(move || {
		println!("debugger enabled - you can type debugger commands here! Type 'help' for more info.");
		loop {
			let mut buf: String = String::new();
			stdin().read_line(&mut buf).unwrap();
			tx.send(buf).unwrap();
			main_handle.unpark();
		}
	});

	//Keys
	let mut keymap: HashMap<Keycode, agb_core::gameboy::Key> = HashMap::new();
	keymap.insert(Keycode::Up, agb_core::gameboy::Key::Up);
	keymap.insert(Keycode::Down, agb_core::gameboy::Key::Down);
	keymap.insert(Keycode::Left, agb_core::gameboy::Key::Left);
	keymap.insert(Keycode::Right, agb_core::gameboy::Key::Right);
	keymap.insert(Keycode::X, agb_core::gameboy::Key::A);
	keymap.insert(Keycode::Z, agb_core::gameboy::Key::B);
	keymap.insert(Keycode::C, agb_core::gameboy::Key::Select);
	keymap.insert(Keycode::V, agb_core::gameboy::Key::Start);

	/* Setup for sdl */
	let sdl_context = sdl2::init().expect("Failed to initialize sdl2");
	let video_subsystem = sdl_context.video().expect("Failed to initialize sdl2 video subsystem");
	let timer_subsystem = sdl_context.timer().expect("Failed to initialize sdl2 timer subsystem");

	//Set resolution
	let width: u32 = (agb_core::WIDTH * DEFAULT_SCALE) as u32;
	let height: u32 = (agb_core::HEIGHT * DEFAULT_SCALE) as u32;

	let window = video_subsystem.window("agb", width, height)
		.position_centered()
		.opengl()
		.build()
		.expect("Failed to create window");

	let mut renderer = window.renderer().build().unwrap();
	let mut event_pump = sdl_context.event_pump().unwrap();
	let mut draw = |gameboy: &mut Gameboy| {
		renderer.set_draw_color(Color::RGB(80, 120, 120));
		renderer.clear();
		const WIDTH: usize = agb_core::WIDTH;
		const HEIGHT: usize = agb_core::HEIGHT;

		let data: &mut[u8] = unsafe {
			//#[allow(mutable_transmutes)]
			//std::mem::transmute::<&[u32], &mut[u8]>(gbc.get_framebuffer())
			let temp: &mut [u32] = agb_core::get_framebuffer_mut(gameboy);
			std::mem::transmute::<&mut[u32], &mut[u8]>(temp)
		};

		let mut texture = renderer.create_texture_streaming(PixelFormatEnum::RGBA8888, WIDTH as u32, HEIGHT as u32).unwrap();
		let _ = texture.update(None, data, WIDTH * 4);
		let rect: Rect = renderer.viewport();
		let _ = renderer.copy(&texture, None, Some(rect));
		renderer.present();
	};

	//Get timer frequency
	let frequency: u64 = timer_subsystem.performance_frequency();

	'running: loop {
		//wait for input from the debugger, but don't wait forever since
		//we don't want the block the gui thread forever
		{
			let paused = paused.lock().unwrap();
			if *paused {
				std::thread::park_timeout(Duration::new(0, 100000000));
			}
		}

		let frame_start: u64 = timer_subsystem.performance_counter();

		//handle debugger input
		if let Ok(input) = rx.try_recv() {
			if input.trim() == "quit" || input.trim() == "exit" {
				break 'running;
			}
			else {
				let mut paused = paused.lock().unwrap();
				debugger::debug(input, &mut gameboy, paused.deref_mut());
			}
		}

		for event in event_pump.poll_iter() {
			match event {
				Event::KeyDown { keycode, .. } => {
					if keycode.is_some() {
						let key = keymap.get(&keycode.unwrap());
						if key.is_some() {
							agb_core::keydown(&mut gameboy, *key.unwrap());
						}
					}
				},
				Event::KeyUp { keycode, .. } => {
					if keycode.is_some() {
						let key = keymap.get(&keycode.unwrap());
						if key.is_some() {
							agb_core::keyup(&mut gameboy, *key.unwrap());
						}
					}
				},
				Event::Quit {..} => {
					break 'running;
				},
				_ => {},
			};
		}

		{
			let paused = paused.lock().unwrap();
			if !*paused {
				gameboy.emulate(Duration::from_millis(1000 / 60));
				draw(&mut gameboy);

				//60hz
				let frame_end: u64 = timer_subsystem.performance_counter();
				let frame_duration: u64 = frame_end - frame_start;
				let ms: u64 = (frame_duration * 1000) / frequency;
				if ms < 1000/60 {
					let duration = Duration::from_millis((1000/60) - ms);
					sleep(duration);
				}
			}
		}
	}
}

pub fn read_file<P: AsRef<Path>>(path: P) -> Result<Box<[u8]>, Error> {
	let mut file = try!(File::open(path));
	let mut buffer = Vec::new();
	try!(file.read_to_end(&mut buffer));
	Ok(buffer.into_boxed_slice())
}

///accepts prefixes (none for base 10, 0x for hex, 0b for binary)
///from_str_radix isn't part of a trait so it can't be generic
pub fn parse_u8(s: &str) -> Result<u8, ParseIntError> {
	if s.starts_with("0x") {
		u8::from_str_radix(&s[2..], 16)
	}
	else if s.starts_with("0b") {
		u8::from_str_radix(&s[2..], 2)
	}
	else {
		u8::from_str_radix(s, 10)
	}
}

///accepts prefixes (none for base 10, 0x for hex, 0b for binary)
pub fn parse_u16(s: &str) -> Result<u16, ParseIntError> {
	if s.starts_with("0x") {
		u16::from_str_radix(&s[2..], 16)
	}
	else if s.starts_with("0b") {
		u16::from_str_radix(&s[2..], 2)
	}
	else {
		u16::from_str_radix(s, 10)
	}
}

///accepts prefixes (none for base 10, 0x for hex, 0b for binary)
pub fn parse_usize(s: &str) -> Result<usize, ParseIntError> {
	if s.starts_with("0x") {
		usize::from_str_radix(&s[2..], 16)
	}
	else if s.starts_with("0b") {
		usize::from_str_radix(&s[2..], 2)
	}
	else {
		usize::from_str_radix(s, 10)
	}
}
<commit_msg>Use the right keyup/keydown/get_framebuffer functions.<commit_after>/* A very simple front end to test agb-core (https://github.com/aHeraud/agb)*/

extern crate agb_core;
extern crate sdl2;
extern crate image;
extern crate clap;

mod debugger;

use std::sync::mpsc::sync_channel;
use std::thread;
use std::time::Duration;
use std::thread::sleep;
use std::collections::HashMap;
use std::fs::File;
use std::io::{stdin, stdout, Read, Write, Error};
use std::path::Path;
use std::num::ParseIntError;
use std::sync::{Arc, Mutex};
use std::ops::DerefMut;
use std::net::{TcpListener, TcpStream, SocketAddr, IpAddr, Ipv4Addr};

use agb_core::gameboy::Gameboy;
use agb_core::gameboy::debugger::DebuggerInterface;

use sdl2::pixels::Color;
use sdl2::pixels::PixelFormatEnum;
use sdl2::rect::Rect;
use sdl2::event::Event;
use sdl2::keyboard::Keycode;

use clap::{Arg, App};

const DEFAULT_SCALE: usize = 2;

fn main() {
	/* Create and initialize gameboy */

	//Parse command line arguments
	let matches = App::new("agb")
		.version("0.1")
		.author("Achille Heraud <achille@heraud.xyz>")
		.about("A GameBoy Emulator")
		.arg(Arg::with_name("rom")
			.long("rom")
			.takes_value(true)
			.value_name("FILE")
			.required(true))
		.arg(Arg::with_name("ram")
			.long("ram")
			.takes_value(true)
			.value_name("FILE")
			.required(false))
		.arg(Arg::with_name("paused")
			.long("pause")
			.short("p")
			.required(false))
		.arg(Arg::with_name("print_serial")
			.help("print data written to the serial port")
			.long("print_serial")
			.conflicts_with_all(&["listen", "connect"])
			.required(false))
		.arg(Arg::with_name("listen")
			.help("listen for remote connections (for serial communication over tcp/ip)")
			.long("listen")
			.takes_value(true)
			.value_name("PORT")
			.conflicts_with("connect")
			.required(false))
		.arg(Arg::with_name("connect")
			.help("provide to ip:port to connect to a remote host (for serial communication over tcp/ip")
			.long("connect")
			.takes_value(true)
			.value_name("IP:PORT")
			.required(false))
		.get_matches();

	let rom = read_file(matches.value_of("rom").unwrap()).expect("Could not open rom file.");
	let ram: Option<Box<[u8]>> = if let Some(ram_path) = matches.value_of("ram") {
		Some(read_file(ram_path).expect("failed to read ram file"))
	}
	else {
		None
	};

	let start_paused: bool = matches.occurrences_of("paused") > 0;

	let mut gameboy = Gameboy::new(rom, ram).expect("Failed to initialize gameboy");
	let paused: Arc<Mutex<bool>> = Arc::new(Mutex::new(start_paused));
	gameboy.debugger.enable();
	{
		let paused = paused.clone();
		gameboy.register_breakpoint_callback(move |breakpoint| {
			println!("triggered breakpoint access_type: {:?}, address: 0x{:x}", breakpoint.access_type, breakpoint.address);
			let mut paused = paused.lock().unwrap();
			*paused = true;
		});
	}

	if let Some(ref port_str) = matches.value_of("listen") {
		// set up a tcp socket to accept incoming connections
		if let Ok(port) = u16::from_str_radix(port_str, 10) {
			let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127,0,0,1)), port);
			let listener = TcpListener::bind(addr).unwrap();
			let (local_sender, local_receiver) = gameboy.create_serial_channels();

			println!("waiting for connection...");
			let (mut stream, remote_addr) = listener.accept().unwrap();
			println!("connected to {}", remote_addr);

			stream.set_nonblocking(false).expect("failed to set tcp socket to nonblocking mode");
			stream.set_nodelay(true).expect("failed to set tcp socket nodelay mode");

			//create sender thread
			{
				let mut stream = stream.try_clone().unwrap();
				thread::spawn(move || {
					let mut buffer: [u8; 1] = [0;1];
					loop {
						let data = local_receiver.recv().unwrap();
						buffer[0] = data;
						if stream.write(&buffer[..]).is_err() {
							return;
						}
					}
				});
			}

			//create receiver thread
			thread::spawn(move || {
				let mut buffer: [u8; 1] = [0;1];
				loop {
					if stream.read(&mut buffer[..]).is_err() {
						return;
					}
					local_sender.send(buffer[0]).unwrap();
				}
			});
		}
		else {
			println!("Invalid port: port must be an integer in the range 0 to 65535.");
			return;
		}
	}
	else if let Some(ref addr) = matches.value_of("connect") {
		let (local_sender, local_receiver) = gameboy.create_serial_channels();
		let addr: SocketAddr = addr.parse().unwrap();
		println!("attempting to connect...");
		let mut stream = TcpStream::connect(addr).unwrap();
		println!("connected!");

		stream.set_nonblocking(false).expect("failed to set tcp socket to nonblocking mode");
		stream.set_nodelay(true).expect("failed to set tcp socket nodelay mode");

		//create sender thread
		{
			let mut stream = stream.try_clone().unwrap();
			thread::spawn(move || {
				let mut buffer: [u8; 1] = [0;1];
				loop {
					let data = local_receiver.recv().unwrap();
					buffer[0] = data;
					if stream.write(&buffer[..]).is_err() {
						return;
					}
				}
			});
		}

		//create receiver thread
		thread::spawn(move || {
			let mut buffer: [u8; 1] = [0;1];
			loop {
				if stream.read(&mut buffer[..]).is_err() {
					return;
				}
				local_sender.send(buffer[0]).unwrap();
			}
		});
	}
	else if matches.occurrences_of("print_serial") > 0 {
		let (sender, reciever) = gameboy.create_serial_channels();
		thread::spawn(move || {
			loop {
				let data = reciever.recv().unwrap();
				sender.send(0xFF).unwrap();
				print!("{}", data as char);
				stdout().flush().ok().expect("could not flush stdout");
			}
		});
	}

	//debugger text input
	let (tx, rx) = sync_channel(0);
	let main_handle = thread::current();
	thread::spawn(move || {
		println!("debugger enabled - you can type debugger commands here! Type 'help' for more info.");
		loop {
			let mut buf: String = String::new();
			stdin().read_line(&mut buf).unwrap();
			tx.send(buf).unwrap();
			main_handle.unpark();
		}
	});

	//Keys
	let mut keymap: HashMap<Keycode, agb_core::gameboy::Key> = HashMap::new();
	keymap.insert(Keycode::Up, agb_core::gameboy::Key::Up);
	keymap.insert(Keycode::Down, agb_core::gameboy::Key::Down);
	keymap.insert(Keycode::Left, agb_core::gameboy::Key::Left);
	keymap.insert(Keycode::Right, agb_core::gameboy::Key::Right);
	keymap.insert(Keycode::X, agb_core::gameboy::Key::A);
	keymap.insert(Keycode::Z, agb_core::gameboy::Key::B);
	keymap.insert(Keycode::C, agb_core::gameboy::Key::Select);
	keymap.insert(Keycode::V, agb_core::gameboy::Key::Start);

	/* Setup for sdl */
	let sdl_context = sdl2::init().expect("Failed to initialize sdl2");
	let video_subsystem = sdl_context.video().expect("Failed to initialize sdl2 video subsystem");
	let timer_subsystem = sdl_context.timer().expect("Failed to initialize sdl2 timer subsystem");

	//Set resolution
	let width: u32 = (agb_core::WIDTH * DEFAULT_SCALE) as u32;
	let height: u32 = (agb_core::HEIGHT * DEFAULT_SCALE) as u32;

	let window = video_subsystem.window("agb", width, height)
		.position_centered()
		.opengl()
		.build()
		.expect("Failed to create window");

	let mut renderer = window.renderer().build().unwrap();
	let mut event_pump = sdl_context.event_pump().unwrap();
	let mut draw = |gameboy: &mut Gameboy| {
		renderer.set_draw_color(Color::RGB(80, 120, 120));
		renderer.clear();
		const WIDTH: usize = agb_core::WIDTH;
		const HEIGHT: usize = agb_core::HEIGHT;

		let data: &mut[u8] = unsafe {
			//#[allow(mutable_transmutes)]
			//std::mem::transmute::<&[u32], &mut[u8]>(gbc.get_framebuffer())
			let temp: &mut [u32] = gameboy.get_framebuffer_mut();
			std::mem::transmute::<&mut[u32], &mut[u8]>(temp)
		};

		let mut texture = renderer.create_texture_streaming(PixelFormatEnum::RGBA8888, WIDTH as u32, HEIGHT as u32).unwrap();
		let _ = texture.update(None, data, WIDTH * 4);
		let rect: Rect = renderer.viewport();
		let _ = renderer.copy(&texture, None, Some(rect));
		renderer.present();
	};

	//Get timer frequency
	let frequency: u64 = timer_subsystem.performance_frequency();

	'running: loop {
		//wait for input from the debugger, but don't wait forever since
		//we don't want the block the gui thread forever
		{
			let paused = paused.lock().unwrap();
			if *paused {
				std::thread::park_timeout(Duration::new(0, 100000000));
			}
		}

		let frame_start: u64 = timer_subsystem.performance_counter();

		//handle debugger input
		if let Ok(input) = rx.try_recv() {
			if input.trim() == "quit" || input.trim() == "exit" {
				break 'running;
			}
			else {
				let mut paused = paused.lock().unwrap();
				debugger::debug(input, &mut gameboy, paused.deref_mut());
			}
		}

		for event in event_pump.poll_iter() {
			match event {
				Event::KeyDown { keycode, .. } => {
					if keycode.is_some() {
						let key = keymap.get(&keycode.unwrap());
						if key.is_some() {
							gameboy.keydown(*key.unwrap());
						}
					}
				},
				Event::KeyUp { keycode, .. } => {
					if keycode.is_some() {
						let key = keymap.get(&keycode.unwrap());
						if key.is_some() {
							gameboy.keyup(*key.unwrap());
						}
					}
				},
				Event::Quit {..} => {
					break 'running;
				},
				_ => {},
			};
		}

		{
			let paused = paused.lock().unwrap();
			if !*paused {
				gameboy.emulate(Duration::from_millis(1000 / 60));
				draw(&mut gameboy);

				//60hz
				let frame_end: u64 = timer_subsystem.performance_counter();
				let frame_duration: u64 = frame_end - frame_start;
				let ms: u64 = (frame_duration * 1000) / frequency;
				if ms < 1000/60 {
					let duration = Duration::from_millis((1000/60) - ms);
					sleep(duration);
				}
			}
		}
	}
}

pub fn read_file<P: AsRef<Path>>(path: P) -> Result<Box<[u8]>, Error> {
	let mut file = try!(File::open(path));
	let mut buffer = Vec::new();
	try!(file.read_to_end(&mut buffer));
	Ok(buffer.into_boxed_slice())
}

///accepts prefixes (none for base 10, 0x for hex, 0b for binary)
///from_str_radix isn't part of a trait so it can't be generic
pub fn parse_u8(s: &str) -> Result<u8, ParseIntError> {
	if s.starts_with("0x") {
		u8::from_str_radix(&s[2..], 16)
	}
	else if s.starts_with("0b") {
		u8::from_str_radix(&s[2..], 2)
	}
	else {
		u8::from_str_radix(s, 10)
	}
}

///accepts prefixes (none for base 10, 0x for hex, 0b for binary)
pub fn parse_u16(s: &str) -> Result<u16, ParseIntError> {
	if s.starts_with("0x") {
		u16::from_str_radix(&s[2..], 16)
	}
	else if s.starts_with("0b") {
		u16::from_str_radix(&s[2..], 2)
	}
	else {
		u16::from_str_radix(s, 10)
	}
}

///accepts prefixes (none for base 10, 0x for hex, 0b for binary)
pub fn parse_usize(s: &str) -> Result<usize, ParseIntError> {
	if s.starts_with("0x") {
		usize::from_str_radix(&s[2..], 16)
	}
	else if s.starts_with("0b") {
		usize::from_str_radix(&s[2..], 2)
	}
	else {
		usize::from_str_radix(s, 10)
	}
}
<|endoftext|>


<commit_before><commit_msg>fix(archive_viewer): correct logic<commit_after><|endoftext|>


<commit_before>#[macro_use] extern crate clap;
extern crate secstr;
extern crate colorhash256;
extern crate interactor;
extern crate rusterpassword;
extern crate ansi_term;
extern crate rustc_serialize;
extern crate cbor;
extern crate freepass_core;

mod util;

use std::{fs,env,io};
use std::io::prelude::*;
use std::collections::btree_map::BTreeMap;
use rustc_serialize::base64::{ToBase64, STANDARD};
use rustc_serialize::hex::ToHex;
use secstr::SecStr;
use interactor::*;
use rusterpassword::*;
use freepass_core::data::*;
use freepass_core::output::*;

fn main() {
    let matches = clap_app!(freepass =>
        (version: env!("CARGO_PKG_VERSION"))
        (author: "Greg V <greg@unrelenting.technology>")
        (about: "The free password manager for power users")
        (@arg FILE: -f --file +takes_value "Sets the vault file to use, by default: $FREEPASS_FILE")
        (@arg NAME: -n --name +takes_value "Sets the user name to use (must be always the same for a vault file!), by default: $FREEPASS_NAME")
        (@arg DEBUG: --debug "Enables logging of data structures for debugging (DO NOT USE ON YOUR REAL DATA)")
        (@subcommand interact =>
            (about: "Launches interactive mode")
        )
    ).get_matches();

    let file_path = opt_or_env(&matches, "FILE", "FREEPASS_FILE");
    let user_name = opt_or_env(&matches, "NAME", "FREEPASS_NAME");
    let debug = matches.is_present("DEBUG");

    freepass_core::init();

    // Do this early because we don't want to ask for the password when we get permission denied or something.
    // Ensure we can write! Maybe someone somewhere would want to open the vault in read-only mode...
    // But the frustration of trying to save the vault while only having read permissions would be worse.
    let file = match fs::OpenOptions::new().read(true).write(true).open(&file_path) {
        Ok(file) => Some(file),
        Err(ref err) if err.kind() == io::ErrorKind::NotFound => None,
        Err(ref err) => panic!("Could not open file {}: {}", &file_path, err),
    };

    let master_key = gen_master_key(util::read_password(), &user_name).unwrap();
    let outer_key = gen_outer_key(&master_key);

    let mut vault = match file {
        Some(f) => Vault::open(&outer_key, f).unwrap(),
        None => Vault::new(),
    };

    if debug {
        debug_output(&vault, "Vault");
    }

    match matches.subcommand() {
        ("interact", _) | _  => interact_entries(&mut vault, &file_path, &outer_key, &master_key, debug),
    }
}

fn opt_or_env(matches: &clap::ArgMatches, opt_name: &str, env_name: &str) -> String {
    match matches.value_of(opt_name).map(|x| x.to_owned()).or(env::var_os(env_name).and_then(|s| s.into_string().ok())) {
        Some(s) => s,
        None => panic!("Option {} or environment variable {} not found", opt_name, env_name)
    }
}

fn debug_output<T: rustc_serialize::Encodable>(data: &T, description: &str) {
    let mut e = cbor::Encoder::from_memory();
    e.encode(&[data]).unwrap();
    println!("--- CBOR debug output (http://cbor.me to decode) of {} ---\n{}\n", description, e.into_bytes().to_hex());
}

macro_rules! interaction {
    ( { $($action_name:expr => $action_fn:expr),+ }, $data:expr, $data_fn:expr ) => {
        {
            let mut items = vec![$(">> ".to_string() + $action_name),+];
            let data_items : Vec<String> = $data.clone().map(|x| " | ".to_string() + x).collect();
            items.extend(data_items.iter().cloned());
            match pick_from_list(util::menu_cmd().as_mut(), &items[..], "Selection: ").unwrap() {
                $(ref x if *x == ">> ".to_string() + $action_name => $action_fn),+
                ref x if data_items.contains(x) => ($data_fn)(&x[3..]),
                ref x => panic!("Unknown selection: {}", x),
            }
        }
    };
    ( { $($action_name:expr => $action_fn:expr),+ }) => {
        {
            let items = vec![$(">> ".to_string() + $action_name),+];
            match pick_from_list(util::menu_cmd().as_mut(), &items[..], "Selection: ").unwrap() {
                $(ref x if *x == ">> ".to_string() + $action_name => $action_fn),+
                ref x => panic!("Unknown selection: {}", x),
            }
        }
    }
}

fn interact_entries(vault: &mut Vault, file_path: &str, outer_key: &SecStr, master_key: &SecStr, debug: bool) {
    let entries_key = gen_entries_key(&master_key);
    loop {
        interaction!({
            "Quit" => {
                return ();
            },
            "Add new entry" => {
                if let Some(entry_name) = util::read_text("Entry name") {
                    interact_entry_edit(vault, file_path, outer_key, master_key, &entries_key, &entry_name, Entry::new(), EntryMetadata::new());
                }
            }
        }, vault.entry_names(), |name| {
            let (entry, meta) = vault.get_entry(&entries_key, name).unwrap();
            if debug {
                debug_output(&entry, &format!("Entry: {}", name));
            }
            interact_entry(vault, file_path, outer_key, master_key, &entries_key, name, entry, meta);
        });
    }
}

fn interact_entry(vault: &mut Vault, file_path: &str, outer_key: &SecStr, master_key: &SecStr, entries_key: &SecStr, entry_name: &str, entry: Entry, meta: EntryMetadata) {
    loop {
        interaction!({
            "Go back" => {
                return ();
            },
            &format!("Name:          {}", entry_name) => {},
            &format!("Last modified: {}", meta.updated_at.to_rfc2822()) => {},
            &format!("Created:       {}", meta.created_at.to_rfc2822()) => {},
            "Edit" => {
                return interact_entry_edit(vault, file_path, outer_key, master_key, entries_key, entry_name, entry, meta);
            }
        }, entry.fields.keys(), |name: &str| {
            let output = process_output(entry_name, master_key, entry.fields.get(name).unwrap()).unwrap();
            match output {
                Output::PrivateText(s) => println!("{}", String::from_utf8(Vec::from(s.unsecure())).unwrap()),
                Output::OpenText(s) => println!("{}", s),
                Output::PrivateBinary(s) => {
                    interaction!({
                        "Go back" => {},
                        "Print as Base64" => { println!("{}", s.unsecure().to_base64(STANDARD)) },
                        "Print as hex" => { println!("{}", s.unsecure().to_hex()) }
                    })
                },
                Output::Ed25519Keypair(usage, _, _) => match usage {
                    Ed25519Usage::SSH => {
                        interaction!({
                            "Go back" => {},
                            "Print public key" => { println!("{}", ssh_public_key_output(&output, entry_name).unwrap()) },
                            "Add private key to ssh-agent" => {
                                ssh_agent_send_message(ssh_private_key_agent_message(&output, entry_name).unwrap()).unwrap()
                            }
                        })
                    }
                    _ => panic!("Unsupported key usage"),
                }
            }
        });
    }
}

fn interact_entry_edit(vault: &mut Vault, file_path: &str, outer_key: &SecStr, master_key: &SecStr, entries_key: &SecStr, entry_name: &str, mut entry: Entry, mut meta: EntryMetadata) {
    interaction!({
        &format!("  Save entry [{}]", entry_name) => {
            vault.put_entry(entries_key, entry_name, &entry, &mut meta).unwrap();
            save_vault(vault, file_path, outer_key);
            return interact_entry(vault, file_path, outer_key, master_key, entries_key, entry_name, entry, meta);
        },
        &format!("Delete entry [{}]", entry_name) => {
            interaction!({
                "Cancel" => {
                    return interact_entry_edit(vault, file_path, outer_key, master_key, entries_key, entry_name, entry, meta);
                },
                &format!("DELETE THE ENTRY '{}'!", entry_name) => {
                    vault.remove_entry(entry_name);
                    save_vault(vault, file_path, outer_key);
                    return ();
                }
            })
        },
        &format!("Rename entry [{}]", entry_name) => {
            let new_entry_name = util::read_text(&format!("New entry name [{}]", entry_name)).unwrap_or(entry_name.to_string());
            vault.remove_entry(entry_name);
            vault.put_entry(entries_key, &new_entry_name, &entry, &mut meta).unwrap();
            return interact_entry_edit(vault, file_path, outer_key, master_key, entries_key, &new_entry_name, entry, meta);
        },
        "Add field" => {
            if let Some(field_name) = util::read_text("Field name") {
                entry = interact_field_edit(vault, entry, field_name);
            }
            return interact_entry_edit(vault, file_path, outer_key, master_key, entries_key, entry_name, entry, meta);
        }
    }, entry.fields.keys(), |name: &str| {
        entry = interact_field_edit(vault, entry, name.to_string());
        return interact_entry_edit(vault, file_path, outer_key, master_key, entries_key, entry_name, entry, meta);
    });
}

fn interact_field_edit(vault: &mut Vault, mut entry: Entry, field_name: String) -> Entry {
    let default_derived = Field::Derived { counter: 0, site_name: None, usage: DerivedUsage::Password(PasswordTemplate::Maximum) };
    let default_stored  = Field::Stored { data: SecStr::new(Vec::new()), usage: StoredUsage::Password };
    let mut field = entry.fields.remove(&field_name).unwrap_or_else(|| default_derived.clone());
    let mut field_actions : BTreeMap<String, Box<Fn(Field) -> Field>> = BTreeMap::new();
    let other_type;
    match field.clone() {
        Field::Derived { counter, site_name, usage } => {
            other_type = "stored";
            field_actions.insert(format!("Counter: {}", counter), Box::new(|f| {
                if let Field::Derived { counter, site_name, usage } = f {
                    let new_counter = util::read_text(&format!("Counter [{}]", counter)).and_then(|c| c.parse::<u32>().ok()).unwrap_or(counter);
                    Field::Derived { counter: new_counter, site_name: site_name, usage: usage }
                } else { unreachable!(); }
            }));
            field_actions.insert(match site_name {
                Some(ref sn) => format!("Site name: {}", sn),
                None => format!("Site name: <same as entry name>"),
            }, Box::new(|f| {
                if let Field::Derived { counter, usage, .. } = f {
                    let new_site_name = util::read_text("Site name");
                    Field::Derived { counter: counter, site_name: new_site_name, usage: usage }
                } else { unreachable!(); }
            }));
            field_actions.insert(format!("Usage: {:?}", usage), Box::new(|f| {
                if let Field::Derived { counter, site_name, .. } = f {
                    let new_usage = interaction!({
                        "Password(Maximum)"   => { DerivedUsage::Password(PasswordTemplate::Maximum) },
                        "Password(Long)"      => { DerivedUsage::Password(PasswordTemplate::Long) },
                        "Password(Medium)"    => { DerivedUsage::Password(PasswordTemplate::Medium) },
                        "Password(Short)"     => { DerivedUsage::Password(PasswordTemplate::Short) },
                        "Password(Basic)"     => { DerivedUsage::Password(PasswordTemplate::Basic) },
                        "Password(Pin)"       => { DerivedUsage::Password(PasswordTemplate::Pin) },
                        "Ed25519Key(SSH)"     => { DerivedUsage::Ed25519Key(Ed25519Usage::SSH) },
                        "Ed25519Key(Signify)" => { DerivedUsage::Ed25519Key(Ed25519Usage::Signify) },
                        "Ed25519Key(SQRL)"    => { DerivedUsage::Ed25519Key(Ed25519Usage::SQRL) },
                        "RawKey"              => { DerivedUsage::RawKey }
                    });
                    Field::Derived { counter: counter, site_name: site_name, usage: new_usage }
                } else { unreachable!(); }
            }));
        },
        Field::Stored { usage, data, .. } => {
            other_type = "derived";
            let txt = String::from_utf8(data.unsecure().to_vec()).unwrap_or("<invalid UTF-8>".to_string());
            field_actions.insert(format!("Change text [{}]", txt), Box::new(|f| {
                if let Field::Stored { usage, data, .. } = f {
                    let txt = String::from_utf8(data.unsecure().to_vec()).unwrap_or("<invalid UTF-8>".to_string());
                    let new_data = util::read_text(&format!("New text [{}]", txt)).unwrap_or(txt);
                    Field::Stored { data: SecStr::from(new_data), usage: usage }
                } else { unreachable!(); }
            }));
            field_actions.insert(format!("Usage: {:?}", usage), Box::new(|f| {
                if let Field::Stored { data, .. } = f {
                    let new_usage = interaction!({
                        "Password"            => { StoredUsage::Password },
                        "Text"                => { StoredUsage::Text }
                    });
                    Field::Stored { data: data, usage: new_usage }
                } else { unreachable!(); }
            }));
        }
    };
    interaction!({
        "Go back" => {
            entry.fields.insert(field_name, field);
            return entry;
        },
        &format!("Delete field [{}]", field_name) => {
            interaction!({
                "Cancel" => {
                    return interact_field_edit(vault, entry, field_name);
                },
                &format!("DELETE THE FIELD '{}'!", field_name) => {
                    entry.fields.remove(&field_name);
                    return entry;
                }
            })
        },
        &format!("Rename field [{}]", field_name) => {
            let new_field_name = util::read_text(&format!("New field name [{}]", field_name)).unwrap_or(field_name.to_string());
            entry.fields.insert(new_field_name.clone(), field);
            return interact_field_edit(vault, entry, new_field_name);
        },
        &format!("Change type to {}", other_type) => {
            entry.fields.remove(&field_name);
            entry.fields.insert(field_name.clone(), match field {
                Field::Derived { .. } => default_stored.clone(),
                Field::Stored { .. } => default_derived.clone(),
            });
            return interact_field_edit(vault, entry, field_name);
        }
    }, field_actions.keys(), |key| {
        field = field_actions.get(key).unwrap()(field);
        entry.fields.insert(field_name.clone(), field);
        return interact_field_edit(vault, entry, field_name);
    })
}

fn save_vault(vault: &mut Vault, file_path: &str, outer_key: &SecStr) {
    // Atomic save!
    vault.save(outer_key, fs::File::create(format!("{}.tmp", file_path)).unwrap()).unwrap();
    fs::rename(format!("{}.tmp", file_path), file_path).unwrap();
}
<commit_msg>Counter must be 1 to follow spec.<commit_after>#[macro_use] extern crate clap;
extern crate secstr;
extern crate colorhash256;
extern crate interactor;
extern crate rusterpassword;
extern crate ansi_term;
extern crate rustc_serialize;
extern crate cbor;
extern crate freepass_core;

mod util;

use std::{fs,env,io};
use std::io::prelude::*;
use std::collections::btree_map::BTreeMap;
use rustc_serialize::base64::{ToBase64, STANDARD};
use rustc_serialize::hex::ToHex;
use secstr::SecStr;
use interactor::*;
use rusterpassword::*;
use freepass_core::data::*;
use freepass_core::output::*;

fn main() {
    let matches = clap_app!(freepass =>
        (version: env!("CARGO_PKG_VERSION"))
        (author: "Greg V <greg@unrelenting.technology>")
        (about: "The free password manager for power users")
        (@arg FILE: -f --file +takes_value "Sets the vault file to use, by default: $FREEPASS_FILE")
        (@arg NAME: -n --name +takes_value "Sets the user name to use (must be always the same for a vault file!), by default: $FREEPASS_NAME")
        (@arg DEBUG: --debug "Enables logging of data structures for debugging (DO NOT USE ON YOUR REAL DATA)")
        (@subcommand interact =>
            (about: "Launches interactive mode")
        )
    ).get_matches();

    let file_path = opt_or_env(&matches, "FILE", "FREEPASS_FILE");
    let user_name = opt_or_env(&matches, "NAME", "FREEPASS_NAME");
    let debug = matches.is_present("DEBUG");

    freepass_core::init();

    // Do this early because we don't want to ask for the password when we get permission denied or something.
    // Ensure we can write! Maybe someone somewhere would want to open the vault in read-only mode...
    // But the frustration of trying to save the vault while only having read permissions would be worse.
    let file = match fs::OpenOptions::new().read(true).write(true).open(&file_path) {
        Ok(file) => Some(file),
        Err(ref err) if err.kind() == io::ErrorKind::NotFound => None,
        Err(ref err) => panic!("Could not open file {}: {}", &file_path, err),
    };

    let master_key = gen_master_key(util::read_password(), &user_name).unwrap();
    let outer_key = gen_outer_key(&master_key);

    let mut vault = match file {
        Some(f) => Vault::open(&outer_key, f).unwrap(),
        None => Vault::new(),
    };

    if debug {
        debug_output(&vault, "Vault");
    }

    match matches.subcommand() {
        ("interact", _) | _  => interact_entries(&mut vault, &file_path, &outer_key, &master_key, debug),
    }
}

fn opt_or_env(matches: &clap::ArgMatches, opt_name: &str, env_name: &str) -> String {
    match matches.value_of(opt_name).map(|x| x.to_owned()).or(env::var_os(env_name).and_then(|s| s.into_string().ok())) {
        Some(s) => s,
        None => panic!("Option {} or environment variable {} not found", opt_name, env_name)
    }
}

fn debug_output<T: rustc_serialize::Encodable>(data: &T, description: &str) {
    let mut e = cbor::Encoder::from_memory();
    e.encode(&[data]).unwrap();
    println!("--- CBOR debug output (http://cbor.me to decode) of {} ---\n{}\n", description, e.into_bytes().to_hex());
}

macro_rules! interaction {
    ( { $($action_name:expr => $action_fn:expr),+ }, $data:expr, $data_fn:expr ) => {
        {
            let mut items = vec![$(">> ".to_string() + $action_name),+];
            let data_items : Vec<String> = $data.clone().map(|x| " | ".to_string() + x).collect();
            items.extend(data_items.iter().cloned());
            match pick_from_list(util::menu_cmd().as_mut(), &items[..], "Selection: ").unwrap() {
                $(ref x if *x == ">> ".to_string() + $action_name => $action_fn),+
                ref x if data_items.contains(x) => ($data_fn)(&x[3..]),
                ref x => panic!("Unknown selection: {}", x),
            }
        }
    };
    ( { $($action_name:expr => $action_fn:expr),+ }) => {
        {
            let items = vec![$(">> ".to_string() + $action_name),+];
            match pick_from_list(util::menu_cmd().as_mut(), &items[..], "Selection: ").unwrap() {
                $(ref x if *x == ">> ".to_string() + $action_name => $action_fn),+
                ref x => panic!("Unknown selection: {}", x),
            }
        }
    }
}

fn interact_entries(vault: &mut Vault, file_path: &str, outer_key: &SecStr, master_key: &SecStr, debug: bool) {
    let entries_key = gen_entries_key(&master_key);
    loop {
        interaction!({
            "Quit" => {
                return ();
            },
            "Add new entry" => {
                if let Some(entry_name) = util::read_text("Entry name") {
                    interact_entry_edit(vault, file_path, outer_key, master_key, &entries_key, &entry_name, Entry::new(), EntryMetadata::new());
                }
            }
        }, vault.entry_names(), |name| {
            let (entry, meta) = vault.get_entry(&entries_key, name).unwrap();
            if debug {
                debug_output(&entry, &format!("Entry: {}", name));
            }
            interact_entry(vault, file_path, outer_key, master_key, &entries_key, name, entry, meta);
        });
    }
}

fn interact_entry(vault: &mut Vault, file_path: &str, outer_key: &SecStr, master_key: &SecStr, entries_key: &SecStr, entry_name: &str, entry: Entry, meta: EntryMetadata) {
    loop {
        interaction!({
            "Go back" => {
                return ();
            },
            &format!("Name:          {}", entry_name) => {},
            &format!("Last modified: {}", meta.updated_at.to_rfc2822()) => {},
            &format!("Created:       {}", meta.created_at.to_rfc2822()) => {},
            "Edit" => {
                return interact_entry_edit(vault, file_path, outer_key, master_key, entries_key, entry_name, entry, meta);
            }
        }, entry.fields.keys(), |name: &str| {
            let output = process_output(entry_name, master_key, entry.fields.get(name).unwrap()).unwrap();
            match output {
                Output::PrivateText(s) => println!("{}", String::from_utf8(Vec::from(s.unsecure())).unwrap()),
                Output::OpenText(s) => println!("{}", s),
                Output::PrivateBinary(s) => {
                    interaction!({
                        "Go back" => {},
                        "Print as Base64" => { println!("{}", s.unsecure().to_base64(STANDARD)) },
                        "Print as hex" => { println!("{}", s.unsecure().to_hex()) }
                    })
                },
                Output::Ed25519Keypair(usage, _, _) => match usage {
                    Ed25519Usage::SSH => {
                        interaction!({
                            "Go back" => {},
                            "Print public key" => { println!("{}", ssh_public_key_output(&output, entry_name).unwrap()) },
                            "Add private key to ssh-agent" => {
                                ssh_agent_send_message(ssh_private_key_agent_message(&output, entry_name).unwrap()).unwrap()
                            }
                        })
                    }
                    _ => panic!("Unsupported key usage"),
                }
            }
        });
    }
}

fn interact_entry_edit(vault: &mut Vault, file_path: &str, outer_key: &SecStr, master_key: &SecStr, entries_key: &SecStr, entry_name: &str, mut entry: Entry, mut meta: EntryMetadata) {
    interaction!({
        &format!("  Save entry [{}]", entry_name) => {
            vault.put_entry(entries_key, entry_name, &entry, &mut meta).unwrap();
            save_vault(vault, file_path, outer_key);
            return interact_entry(vault, file_path, outer_key, master_key, entries_key, entry_name, entry, meta);
        },
        &format!("Delete entry [{}]", entry_name) => {
            interaction!({
                "Cancel" => {
                    return interact_entry_edit(vault, file_path, outer_key, master_key, entries_key, entry_name, entry, meta);
                },
                &format!("DELETE THE ENTRY '{}'!", entry_name) => {
                    vault.remove_entry(entry_name);
                    save_vault(vault, file_path, outer_key);
                    return ();
                }
            })
        },
        &format!("Rename entry [{}]", entry_name) => {
            let new_entry_name = util::read_text(&format!("New entry name [{}]", entry_name)).unwrap_or(entry_name.to_string());
            vault.remove_entry(entry_name);
            vault.put_entry(entries_key, &new_entry_name, &entry, &mut meta).unwrap();
            return interact_entry_edit(vault, file_path, outer_key, master_key, entries_key, &new_entry_name, entry, meta);
        },
        "Add field" => {
            if let Some(field_name) = util::read_text("Field name") {
                entry = interact_field_edit(vault, entry, field_name);
            }
            return interact_entry_edit(vault, file_path, outer_key, master_key, entries_key, entry_name, entry, meta);
        }
    }, entry.fields.keys(), |name: &str| {
        entry = interact_field_edit(vault, entry, name.to_string());
        return interact_entry_edit(vault, file_path, outer_key, master_key, entries_key, entry_name, entry, meta);
    });
}

fn interact_field_edit(vault: &mut Vault, mut entry: Entry, field_name: String) -> Entry {
    let default_derived = Field::Derived { counter: 1, site_name: None, usage: DerivedUsage::Password(PasswordTemplate::Maximum) };
    let default_stored  = Field::Stored { data: SecStr::new(Vec::new()), usage: StoredUsage::Password };
    let mut field = entry.fields.remove(&field_name).unwrap_or_else(|| default_derived.clone());
    let mut field_actions : BTreeMap<String, Box<Fn(Field) -> Field>> = BTreeMap::new();
    let other_type;
    match field.clone() {
        Field::Derived { counter, site_name, usage } => {
            other_type = "stored";
            field_actions.insert(format!("Counter: {}", counter), Box::new(|f| {
                if let Field::Derived { counter, site_name, usage } = f {
                    let new_counter = util::read_text(&format!("Counter [{}]", counter)).and_then(|c| c.parse::<u32>().ok()).unwrap_or(counter);
                    Field::Derived { counter: new_counter, site_name: site_name, usage: usage }
                } else { unreachable!(); }
            }));
            field_actions.insert(match site_name {
                Some(ref sn) => format!("Site name: {}", sn),
                None => format!("Site name: <same as entry name>"),
            }, Box::new(|f| {
                if let Field::Derived { counter, usage, .. } = f {
                    let new_site_name = util::read_text("Site name");
                    Field::Derived { counter: counter, site_name: new_site_name, usage: usage }
                } else { unreachable!(); }
            }));
            field_actions.insert(format!("Usage: {:?}", usage), Box::new(|f| {
                if let Field::Derived { counter, site_name, .. } = f {
                    let new_usage = interaction!({
                        "Password(Maximum)"   => { DerivedUsage::Password(PasswordTemplate::Maximum) },
                        "Password(Long)"      => { DerivedUsage::Password(PasswordTemplate::Long) },
                        "Password(Medium)"    => { DerivedUsage::Password(PasswordTemplate::Medium) },
                        "Password(Short)"     => { DerivedUsage::Password(PasswordTemplate::Short) },
                        "Password(Basic)"     => { DerivedUsage::Password(PasswordTemplate::Basic) },
                        "Password(Pin)"       => { DerivedUsage::Password(PasswordTemplate::Pin) },
                        "Ed25519Key(SSH)"     => { DerivedUsage::Ed25519Key(Ed25519Usage::SSH) },
                        "Ed25519Key(Signify)" => { DerivedUsage::Ed25519Key(Ed25519Usage::Signify) },
                        "Ed25519Key(SQRL)"    => { DerivedUsage::Ed25519Key(Ed25519Usage::SQRL) },
                        "RawKey"              => { DerivedUsage::RawKey }
                    });
                    Field::Derived { counter: counter, site_name: site_name, usage: new_usage }
                } else { unreachable!(); }
            }));
        },
        Field::Stored { usage, data, .. } => {
            other_type = "derived";
            let txt = String::from_utf8(data.unsecure().to_vec()).unwrap_or("<invalid UTF-8>".to_string());
            field_actions.insert(format!("Change text [{}]", txt), Box::new(|f| {
                if let Field::Stored { usage, data, .. } = f {
                    let txt = String::from_utf8(data.unsecure().to_vec()).unwrap_or("<invalid UTF-8>".to_string());
                    let new_data = util::read_text(&format!("New text [{}]", txt)).unwrap_or(txt);
                    Field::Stored { data: SecStr::from(new_data), usage: usage }
                } else { unreachable!(); }
            }));
            field_actions.insert(format!("Usage: {:?}", usage), Box::new(|f| {
                if let Field::Stored { data, .. } = f {
                    let new_usage = interaction!({
                        "Password"            => { StoredUsage::Password },
                        "Text"                => { StoredUsage::Text }
                    });
                    Field::Stored { data: data, usage: new_usage }
                } else { unreachable!(); }
            }));
        }
    };
    interaction!({
        "Go back" => {
            entry.fields.insert(field_name, field);
            return entry;
        },
        &format!("Delete field [{}]", field_name) => {
            interaction!({
                "Cancel" => {
                    return interact_field_edit(vault, entry, field_name);
                },
                &format!("DELETE THE FIELD '{}'!", field_name) => {
                    entry.fields.remove(&field_name);
                    return entry;
                }
            })
        },
        &format!("Rename field [{}]", field_name) => {
            let new_field_name = util::read_text(&format!("New field name [{}]", field_name)).unwrap_or(field_name.to_string());
            entry.fields.insert(new_field_name.clone(), field);
            return interact_field_edit(vault, entry, new_field_name);
        },
        &format!("Change type to {}", other_type) => {
            entry.fields.remove(&field_name);
            entry.fields.insert(field_name.clone(), match field {
                Field::Derived { .. } => default_stored.clone(),
                Field::Stored { .. } => default_derived.clone(),
            });
            return interact_field_edit(vault, entry, field_name);
        }
    }, field_actions.keys(), |key| {
        field = field_actions.get(key).unwrap()(field);
        entry.fields.insert(field_name.clone(), field);
        return interact_field_edit(vault, entry, field_name);
    })
}

fn save_vault(vault: &mut Vault, file_path: &str, outer_key: &SecStr) {
    // Atomic save!
    vault.save(outer_key, fs::File::create(format!("{}.tmp", file_path)).unwrap()).unwrap();
    fs::rename(format!("{}.tmp", file_path), file_path).unwrap();
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

type FreeVariable = u64;
type VariableName<'a> = &'a str;

/// AST node expression
#[derive(Copy, Clone, Debug, PartialEq)]
enum Expression<'a> {
    Number(f64),
    Boolean(bool),

    /// simple variable
    Variable(VariableName<'a>),

    /// fn (x, y, z) -> expr
    Lambda(&'a [VariableName<'a>], &'a Expression<'a>),

    /// let x = expr in expr
    Let(VariableName<'a>, &'a Expression<'a>, &'a Expression<'a>),

    /// some_fn (x, y, z)
    FnCall(&'a Expression<'a>, &'a [ &'a Expression<'a>])
}

#[derive(Clone, Debug, PartialEq)]
enum Type {
    Number,
    Boolean,
    Free(FreeVariable),
    // FIXME: can I avoid having this be boxed and vec'd?
    Lambda(Vec<Type>, Box<Type>),
}

type TypeEnv<'a> = HashMap<VariableName<'a>, Type>;

/// Bind previously unbound type variables to a specified type.
fn bind_type<'a>(env: &mut TypeEnv<'a>, var: FreeVariable, ty: Type) -> Type {
    println!("binding free({}) = {:?}", var, ty);

    for val in env.values_mut() {
        if *val == Type::Free(var) {
            *val = ty.clone();
        }
    }

    ty
}

#[derive(Clone, Copy)]
struct VariableGen { current: FreeVariable }
impl VariableGen {
    fn next(&mut self) -> FreeVariable { self.current += 1; self.current }
}

type UnificationError = &'static str;
type TypeError = &'static str;

/// modifies the environment to unify the given types (or returns false if not possible)
fn unify<'a>(a: Type, b: Type, env: &mut TypeEnv<'a>, inst: &mut VariableGen) -> Result<Type, UnificationError> {
    let res = match (a.clone(), b.clone()) {
        // Easy. Literals match with themselves
        (Type::Number, Type::Number) => Ok(Type::Number),
        (Type::Boolean, Type::Boolean) => Ok(Type::Boolean),

        // FIXME: is this correct?
        (Type::Free(a), Type::Free(b)) => {
            let t1 = bind_type(env, b, Type::Free(inst.next()));

            Ok(bind_type(env, a, t1))
        }

        (Type::Free(n), ty) | (ty, Type::Free(n)) =>
            Ok(bind_type(env, n, ty)),

        (Type::Lambda(args1, body1), Type::Lambda(args2, body2)) => {
            if args1.len() != args2.len() {
                return Err("lambdas take differing numbers of arguments")
            }

            let mut unified_args = Vec::with_capacity(args1.len());

            for iter in args1.into_iter().zip(args2.into_iter()) {
                match unify(iter.0, iter.1, env, inst) {
                    Ok(ty) => unified_args.push(ty),
                    err => return err
                }
            }

            match unify(*body1, *body2, env, inst) {
                Ok(body_ty) => Ok(Type::Lambda(unified_args, Box::new(body_ty))),
                err => err
            }
        }

        _ => Err("could not unify types")
    };

    println!("unify: {:?} <-> {:?} == {:?}", &a, &b, res);

    res
}

/// Run algorithm W on expr with given type env, yielding the resolved
/// type of the expression, or the type error encountered.
fn w<'a>(expr: &Expression<'a>, env: &mut TypeEnv<'a>, inst: &mut VariableGen) -> Type {
    match *expr {
        // simplest case: literals have immediately defined types.
        Expression::Number(_) => Type::Number,
        Expression::Boolean(_) => Type::Boolean,

        // variables we need to look up inside our type environment
        Expression::Variable(name) => match env.get(name) {
            Some(typ) => typ.clone(),
            None => panic!("undefined variable, {}", name)
        },

        Expression::Lambda(args, body) => {
            // Any parameters or variables defined inside the lambda are scoped
            // only to the lambda.
            let mut child_env = env.clone();

            let args_ty = args
                .into_iter()
                .map(|arg| {
                    let ty = inst.next();
                    child_env.insert(arg, Type::Free(ty));
                    Type::Free(ty)
                })
                .collect::<Vec<Type>>();

            let body_ty = w(body, &mut child_env, inst);

            Type::Lambda(args_ty, Box::new(body_ty))
        },

        // TODO: need an occurs check here.
        Expression::Let(name, expr, body) => {
            let expr_ty = w(expr, env, inst);
            let mut child_env = env.clone();

            child_env.insert(name, expr_ty);

            w(body, &mut child_env, inst)
        },

        // 1. ensure that the callable in this function call is truly callable.
        // 2. ensure that the number and type of args match
        // 3. return type of body of lambda
        Expression::FnCall(callable, args) => {
            let mut child_env = env.clone();

            let type_args = args
                .iter()
                .map(|arg| { w(arg, &mut child_env, inst) })
                .collect::<Vec<Type>>();

            // Try to unify the lambda we are calling with the way we are calling it.
            unify(w(callable, &mut child_env, inst),
                  Type::Lambda(type_args, Box::new(Type::Free(inst.next()))),
                  &mut child_env, inst)
                .map(move |ty|{
                    match ty {
                        Type::Lambda(_, return_type) => {
                            *return_type
                        },
                        _ => panic!("BUG: expected unify to return lambda, not {:?}", ty)
                    }
                })
                .or(Err("could not unify lambda types"))
                // FIXME: this should eventually return a result anyway.
                .unwrap()
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    fn infer(exp: Expression) -> Type {
        let mut inst = VariableGen { current: 0 };
        let mut env = HashMap::new();

        w(&exp, &mut env, &mut inst)
    }

    fn unify_(t1: Type, t2: Type) -> Result<Type, UnificationError> {
        let mut inst = VariableGen { current: 0 };
        let mut env = HashMap::new();

        unify(t1, t2, &mut env, &mut inst)
    }

    #[test]
    #[should_panic]
    fn undefined_variable() {
        let var = Expression::Variable("x");

        // FIXME: until we return a result type this will panic.
        infer(var);
    }

    #[test]
    fn literal_inference() {
        let num = Expression::Number(123.0);
        let boolean = Expression::Boolean(true);

        assert_eq!(infer(num), Type::Number);
        assert_eq!(infer(boolean), Type::Boolean);
    }

    #[test]
    fn lambda_expression_inference() {
        let args = ["x"];
        let num = Expression::Number(123.0);

        // \x -> 123.0
        let lambda = Expression::Lambda(&args, &num);

        assert_eq!(infer(lambda), Type::Lambda(vec![Type::Free(1)], Box::new(Type::Number)));

        // \x -> \y -> 123.0
        let lambdalambda = Expression::Lambda(&args, &lambda);
        assert_eq!(infer(lambdalambda),
                   Type::Lambda(vec![Type::Free(1)], Box::new(
                       Type::Lambda(vec![Type::Free(2)], Box::new(Type::Number)))));
    }

    #[test]
    fn let_inference() {
        let boolean = Expression::Boolean(false);
        let var_x = Expression::Variable("x");
        let var_y = Expression::Variable("y");

        // let x = false in x
        let let_ = Expression::Let("x", &boolean, &var_x);

        assert_eq!(infer(let_), Type::Boolean);

        // let y = (let x = false in x) in y
        let nested_let = Expression::Let("y", &let_, &var_y);
        assert_eq!(infer(nested_let), Type::Boolean);
    }

    #[test]
    fn simple_fn_call_inference() {
        let args = ["x"];
        let num = Expression::Number(123.0);
        let boolean = Expression::Boolean(false);

        // \x -> 123.0
        let lambda = Expression::Lambda(&args, &num);

        let app_args = [&boolean];
        let application = Expression::FnCall(&lambda, &app_args);

        assert_eq!(infer(application), Type::Number);
    }

    #[test]
    #[should_panic]
    fn call_non_lambda() {
        let num = Expression::Number(123.0);
        let boolean = Expression::Boolean(false);

        let app_args = [&boolean];
        let application = Expression::FnCall(&num, &app_args);

        infer(application);
    }

    #[test]
    fn fn_returning_lambda() {
        // let
        //    x = (\a -> a)
        // in
        //    let y = (\a -> a)
        //    in
        //       x (y)
        // ===> \a -> a
        let var_a = Expression::Variable("a");
        let var_x = Expression::Variable("x");
        let var_y = Expression::Variable("y");
        let args_a = ["a"];
        let args_y = [&var_y];
        let lambda_x = Expression::Lambda(&args_a, &var_a);
        let lambda_y = Expression::Lambda(&args_a, &var_a);

        let fncall = Expression::FnCall(&var_x, &args_y);

        let let_y = Expression::Let("y", &lambda_y, &fncall);
        let let_x = Expression::Let("x", &lambda_x, &let_y);

        assert_eq!(infer(let_x), Type::Lambda(vec![Type::Free(3)], Box::new(Type::Free(3))))
    }

    #[test]
    fn test_unify() {
        assert!(unify_(Type::Number, Type::Boolean).is_err());

        assert_eq!(unify_(Type::Number, Type::Number).unwrap(), Type::Number);
        assert_eq!(unify_(Type::Free(10), Type::Number).unwrap(), Type::Number);

        assert_eq!(unify_(Type::Lambda(vec![Type::Free(10), Type::Free(20)], Box::new(Type::Number)),
                          Type::Lambda(vec![Type::Free(30), Type::Free(40)], Box::new(Type::Number))).unwrap(),
                   Type::Lambda(vec![Type::Free(1), Type::Free(2)], Box::new(Type::Number)));
    }
}

fn main() {
    println!("patience.");
}
<commit_msg>Change defintion of TypeEnv<commit_after>use std::collections::{HashMap, HashSet};

type FreeVariable = u64;
type VariableName<'a> = &'a str;

/// AST node expression
#[derive(Copy, Clone, Debug, PartialEq)]
enum Expression<'a> {
    Number(f64),
    Boolean(bool),

    /// simple variable
    Variable(VariableName<'a>),

    /// fn (x, y, z) -> expr
    Lambda(&'a [VariableName<'a>], &'a Expression<'a>),

    /// let x = expr in expr
    Let(VariableName<'a>, &'a Expression<'a>, &'a Expression<'a>),

    /// some_fn (x, y, z)
    FnCall(&'a Expression<'a>, &'a [ &'a Expression<'a>])
}

#[derive(Clone, Debug, PartialEq)]
enum Type {
    Number,
    Boolean,
    Free(FreeVariable),
    // FIXME: can I avoid having this be boxed and vec'd?
    Lambda(Vec<Type>, Box<Type>),
}

#[derive(Clone, Debug)]
struct TypeEnv<'a> {
    bindings: HashMap<VariableName<'a>, Type>,
    // FIXME: This isn't the correct naming. Perhaps generic?
    instances: HashSet<FreeVariable>
}

impl <'a> TypeEnv <'a> {
    fn new() -> TypeEnv<'a> {
        TypeEnv {
            bindings: HashMap::new(),
            instances: HashSet::new()
        }
    }

    fn instantiate(&mut self, var: FreeVariable) {
        self.instances.insert(var);
    }

    /// Bind previously unbound type variables to a specified type.
    fn bind_type(&mut self, var: FreeVariable, ty: Type) -> Type {
        println!("binding free({}) = {:?}", var, ty);

        for val in self.bindings.values_mut() {
            if *val == Type::Free(var) {
                *val = ty.clone();
            }
        }

        ty
    }

}

#[derive(Clone, Copy)]
struct VariableGen { current: FreeVariable }
impl VariableGen {
    fn next(&mut self) -> FreeVariable { self.current += 1; self.current }
}

type UnificationError = &'static str;
type TypeError = &'static str;

/// modifies the environment to unify the given types (or returns false if not possible)
fn unify<'a>(a: Type, b: Type, env: &mut TypeEnv<'a>, inst: &mut VariableGen) -> Result<Type, UnificationError> {
    let res = match (a.clone(), b.clone()) {
        // Easy. Literals match with themselves
        (Type::Number, Type::Number) => Ok(Type::Number),
        (Type::Boolean, Type::Boolean) => Ok(Type::Boolean),

        // FIXME: is this correct?
        (Type::Free(a), Type::Free(b)) => {
            let t1 = env.bind_type(b, Type::Free(inst.next()));

            Ok(env.bind_type(a, t1))
        }

        (Type::Free(n), ty) | (ty, Type::Free(n)) =>
            Ok(env.bind_type(n, ty)),

        (Type::Lambda(args1, body1), Type::Lambda(args2, body2)) => {
            if args1.len() != args2.len() {
                return Err("lambdas take differing numbers of arguments")
            }

            let mut unified_args = Vec::with_capacity(args1.len());

            for iter in args1.into_iter().zip(args2.into_iter()) {
                match unify(iter.0, iter.1, env, inst) {
                    Ok(ty) => unified_args.push(ty),
                    err => return err
                }
            }

            match unify(*body1, *body2, env, inst) {
                Ok(body_ty) => Ok(Type::Lambda(unified_args, Box::new(body_ty))),
                err => err
            }
        }

        _ => Err("could not unify types")
    };

    println!("unify: {:?} <-> {:?} == {:?}", &a, &b, res);

    res
}

/// Run algorithm W on expr with given type env, yielding the resolved
/// type of the expression, or the type error encountered.
fn w<'a>(expr: &Expression<'a>, env: &mut TypeEnv<'a>, inst: &mut VariableGen) -> Type {
    match *expr {
        // simplest case: literals have immediately defined types.
        Expression::Number(_) => Type::Number,
        Expression::Boolean(_) => Type::Boolean,

        // variables we need to look up inside our type environment
        Expression::Variable(name) => match env.bindings.get(name) {
            Some(typ) => typ.clone(),
            None => panic!("undefined variable, {}", name)
        },

        Expression::Lambda(args, body) => {
            // Any parameters or variables defined inside the lambda are scoped
            // only to the lambda.
            let mut child_env = env.clone();

            let args_ty = args
                .into_iter()
                .map(|arg| {
                    let ty = inst.next();
                    child_env.bindings.insert(arg, Type::Free(ty));
                    child_env.instantiate(ty);
                    Type::Free(ty)
                })
                .collect::<Vec<Type>>();

            let body_ty = w(body, &mut child_env, inst);

            Type::Lambda(args_ty, Box::new(body_ty))
        },

        // TODO: need an occurs check here.
        Expression::Let(name, expr, body) => {
            let expr_ty = w(expr, env, inst);
            let mut child_env = env.clone();

            child_env.bindings.insert(name, expr_ty);

            w(body, &mut child_env, inst)
        },

        // 1. ensure that the callable in this function call is truly callable.
        // 2. ensure that the number and type of args match
        // 3. return type of body of lambda
        Expression::FnCall(callable, args) => {
            let mut child_env = env.clone();

            let type_args = args
                .iter()
                .map(|arg| {
                    w(arg, &mut child_env, inst)
                })
                .collect::<Vec<Type>>();

            // Try to unify the lambda we are calling with the way we are calling it.
            unify(w(callable, &mut child_env, inst),
                  Type::Lambda(type_args, Box::new(Type::Free(inst.next()))),
                  &mut child_env, inst)
                .map(move |ty|{
                    match ty {
                        Type::Lambda(_, return_type) => {
                            *return_type
                        },
                        _ => panic!("BUG: expected unify to return lambda, not {:?}", ty)
                    }
                })
                .or(Err("could not unify lambda types"))
                // FIXME: this should eventually return a result anyway.
                .unwrap()
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    fn infer(exp: Expression) -> Type {
        let mut inst = VariableGen { current: 0 };
        let mut env = TypeEnv::new();

        w(&exp, &mut env, &mut inst)
    }

    fn unify_(t1: Type, t2: Type) -> Result<Type, UnificationError> {
        let mut inst = VariableGen { current: 0 };
        let mut env = TypeEnv::new();

        unify(t1, t2, &mut env, &mut inst)
    }

    #[test]
    #[should_panic]
    fn undefined_variable() {
        let var = Expression::Variable("x");

        // FIXME: until we return a result type this will panic.
        infer(var);
    }

    #[test]
    fn literal_inference() {
        let num = Expression::Number(123.0);
        let boolean = Expression::Boolean(true);

        assert_eq!(infer(num), Type::Number);
        assert_eq!(infer(boolean), Type::Boolean);
    }

    #[test]
    fn lambda_expression_inference() {
        let args = ["x"];
        let num = Expression::Number(123.0);

        // \x -> 123.0
        let lambda = Expression::Lambda(&args, &num);

        assert_eq!(infer(lambda), Type::Lambda(vec![Type::Free(1)], Box::new(Type::Number)));

        // \x -> \y -> 123.0
        let lambdalambda = Expression::Lambda(&args, &lambda);
        assert_eq!(infer(lambdalambda),
                   Type::Lambda(vec![Type::Free(1)], Box::new(
                       Type::Lambda(vec![Type::Free(2)], Box::new(Type::Number)))));
    }

    #[test]
    fn let_inference() {
        let boolean = Expression::Boolean(false);
        let var_x = Expression::Variable("x");
        let var_y = Expression::Variable("y");

        // let x = false in x
        let let_ = Expression::Let("x", &boolean, &var_x);

        assert_eq!(infer(let_), Type::Boolean);

        // let y = (let x = false in x) in y
        let nested_let = Expression::Let("y", &let_, &var_y);
        assert_eq!(infer(nested_let), Type::Boolean);
    }

    #[test]
    fn simple_lambda_call_inference() {
        let args = ["x"];
        let num = Expression::Number(123.0);
        let boolean = Expression::Boolean(false);

        // \x -> 123.0
        let lambda = Expression::Lambda(&args, &num);

        let app_args = [&boolean];
        let application = Expression::FnCall(&lambda, &app_args);

        assert_eq!(infer(application), Type::Number);
    }

    #[test]
    #[should_panic]
    fn call_non_lambda() {
        let num = Expression::Number(123.0);
        let boolean = Expression::Boolean(false);

        let app_args = [&boolean];
        let application = Expression::FnCall(&num, &app_args);

        infer(application);
    }

    #[test]
    fn lambda_returning_lambda() {
        // let
        //    x = (\a -> a)
        // in
        //    let y = (\a -> a)
        //    in
        //       x (y)
        // ===> \a -> a
        let var_a = Expression::Variable("a");
        let var_x = Expression::Variable("x");
        let var_y = Expression::Variable("y");
        let args_a = ["a"];
        let args_y = [&var_y];
        let lambda_x = Expression::Lambda(&args_a, &var_a);
        let lambda_y = Expression::Lambda(&args_a, &var_a);

        let fncall = Expression::FnCall(&var_x, &args_y);

        let let_y = Expression::Let("y", &lambda_y, &fncall);
        let let_x = Expression::Let("x", &lambda_x, &let_y);

        assert_eq!(infer(let_x), Type::Lambda(vec![Type::Free(3)], Box::new(Type::Free(3))))
    }

    #[test]
    fn test_unify() {
        assert!(unify_(Type::Number, Type::Boolean).is_err());

        assert_eq!(unify_(Type::Number, Type::Number).unwrap(), Type::Number);
        assert_eq!(unify_(Type::Free(10), Type::Number).unwrap(), Type::Number);

        assert_eq!(unify_(Type::Lambda(vec![Type::Free(10), Type::Free(20)], Box::new(Type::Number)),
                          Type::Lambda(vec![Type::Free(30), Type::Free(40)], Box::new(Type::Number))).unwrap(),
                   Type::Lambda(vec![Type::Free(1), Type::Free(2)], Box::new(Type::Number)));
    }
}

fn main() {
    println!("patience.");
}
<|endoftext|>


<commit_before>// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! A wrapper for structures that contain flexible arrays.

use std::marker::PhantomData;
use std::mem::size_of;

// Returns a `Vec<T>` with a size in bytes at least as large as `size_in_bytes`.
fn vec_with_size_in_bytes<T: Default>(size_in_bytes: usize) -> Vec<T> {
    let rounded_size = (size_in_bytes + size_of::<T>() - 1) / size_of::<T>();
    let mut v = Vec::with_capacity(rounded_size);
    for _ in 0..rounded_size {
        v.push(T::default())
    }
    v
}

/// The kernel API has many structs that resemble the following `Foo` structure:
///
/// ```ignore
/// #[repr(C)]
/// struct Foo {
///    some_data: u32,
///    entries: __IncompleteArrayField<__u32>,
/// }
/// ```
///
/// In order to allocate such a structure, `size_of::<Foo>()` would be too small because it would
/// not include any space for `entries`. To make the allocation large enough while still being
/// aligned for `Foo`, a `Vec<Foo>` is created. Only the first element of `Vec<Foo>` would actually
/// be used as a `Foo`. The remaining memory in the `Vec<Foo>` is for `entries`, which must be
/// contiguous with `Foo`. This function is used to make the `Vec<Foo>` with enough space for
/// `count` entries.
pub fn vec_with_array_field<T: Default, F>(count: usize) -> Vec<T> {
    let element_space = count * size_of::<F>();
    let vec_size_bytes = size_of::<T>() + element_space;
    vec_with_size_in_bytes(vec_size_bytes)
}

/// The following code provides generic helpers for creating and accessing flexible array structs.
/// A complete definition of flexible array structs is found in the ISO 9899 specification
/// (http://www.iso-9899.info/n1570.html). A flexible array struct is of the form:
///
/// ```ignore
/// #[repr(C)]
/// struct T {
///    some_data: u32,
///    nents: u32,
///    entries: __IncompleteArrayField<S>,
/// }
/// ```
/// where:
///
/// - `T` is the flexible array struct type
/// - `S` is the flexible array type
/// - `nents` is the flexible array length
/// - `entries` is the flexible array member
///
/// These structures are used by the kernel API.

/// A collection of methods that are required by the FlexibleArrayWrapper type.
///
/// When implemented for `T`, this trait allows the caller to set number of `S` entries and
/// retrieve a slice of `S` entries.  Trait methods must only be called by the FlexibleArrayWrapper
/// type.  Don't implement this trait directly, use the flexible_array! macro to avoid duplication.
pub trait FlexibleArray<S> {
    /// Implementations must set flexible array length in the flexible array struct to the value
    /// specified by `len`. Appropriate conversions (i.e, usize to u32) are allowed so long as
    /// they don't overflow or underflow.
    fn set_len(&mut self, len: usize);
    /// Implementations must return the length of the flexible array member.  Appropriate
    /// conversions (i.e, usize to u32) are allowed so long as they don't overflow or underflow.
    fn get_len(&self) -> usize;
    /// Implementations must return a slice of flexible array member of length `len`.
    unsafe fn get_slice(&self, len: usize) -> &[S];
    /// Implementations must return a mutable slice of flexible array member of length `len`.
    unsafe fn get_mut_slice(&mut self, len: usize) -> &mut [S];
}

/// Always use this macro for implementing the FlexibleArray<`S`> trait for a given `T`.  There
/// exists an 1:1 mapping of macro identifiers to the definitions in the FlexibleArray<`S`>
/// documentation, so refer to that for more information.
#[macro_export]
macro_rules! flexible_array_impl {
    ($T:ident, $S:ident, $nents:ident, $entries:ident) => {
        impl FlexibleArray<$S> for $T {
            fn set_len(&mut self, len: usize) {
                self.$nents = ::std::convert::TryInto::try_into(len).unwrap();
            }

            fn get_len(&self) -> usize {
                self.$nents as usize
            }

            unsafe fn get_slice(&self, len: usize) -> &[$S] {
                self.$entries.as_slice(len)
            }

            unsafe fn get_mut_slice(&mut self, len: usize) -> &mut [$S] {
                self.$entries.as_mut_slice(len)
            }
        }
    };
}

pub struct FlexibleArrayWrapper<T, S> {
    entries: Vec<T>,
    phantom: PhantomData<S>,
    allocated_len: usize,
}

/// Convenience wrapper for flexible array structs.
///
/// The FlexibleArray trait must be implemented for the flexible array struct before using this
/// wrapper.
impl<T, S> FlexibleArrayWrapper<T, S>
where
    T: FlexibleArray<S> + Default,
{
    /// Creates a new FlexibleArrayWrapper for the given flexible array struct type and flexible
    /// array type. The flexible array length is set to `array_len`. vec_with_array_field is used
    /// to make sure the resultant wrapper is appropriately sized.
    pub fn new(array_len: usize) -> FlexibleArrayWrapper<T, S> {
        let mut entries = vec_with_array_field::<T, S>(array_len);
        entries[0].set_len(array_len);

        FlexibleArrayWrapper {
            entries,
            phantom: PhantomData,
            allocated_len: array_len,
        }
    }

    /// Mapping the unsized array to a slice is unsafe because the length isn't known.  Using
    /// the length we originally allocated with eliminates the possibility of overflow.
    fn get_valid_len(&self) -> usize {
        if self.entries[0].get_len() > self.allocated_len {
            self.allocated_len
        } else {
            self.entries[0].get_len()
        }
    }

    /// Returns a slice of the flexible array member, for inspecting. To modify, use
    /// mut_entries_slice instead.
    pub fn entries_slice(&self) -> &[S] {
        let valid_length = self.get_valid_len();
        // Safe because the length has been validated.
        unsafe { self.entries[0].get_slice(valid_length) }
    }

    /// Returns a mutable slice of the flexible array member, for modifying.
    pub fn mut_entries_slice(&mut self) -> &mut [S] {
        let valid_length = self.get_valid_len();
        self.entries[0].set_len(valid_length);
        // Safe because the length has been validated.
        unsafe { self.entries[0].get_mut_slice(valid_length) }
    }

    /// Get a pointer so it can be passed to the kernel. Callers must not access the flexible
    /// array member.  Using this pointer is unsafe.
    pub fn as_ptr(&self) -> *const T {
        &self.entries[0]
    }

    /// Get a mutable pointer so it can be passed to the kernel. Callers must not access the
    /// flexible array member.  Using this pointer is unsafe.
    pub fn as_mut_ptr(&mut self) -> *mut T {
        &mut self.entries[0]
    }
}
<commit_msg>data_model: Add safety section for unsafe functions<commit_after>// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! A wrapper for structures that contain flexible arrays.

use std::marker::PhantomData;
use std::mem::size_of;

// Returns a `Vec<T>` with a size in bytes at least as large as `size_in_bytes`.
fn vec_with_size_in_bytes<T: Default>(size_in_bytes: usize) -> Vec<T> {
    let rounded_size = (size_in_bytes + size_of::<T>() - 1) / size_of::<T>();
    let mut v = Vec::with_capacity(rounded_size);
    for _ in 0..rounded_size {
        v.push(T::default())
    }
    v
}

/// The kernel API has many structs that resemble the following `Foo` structure:
///
/// ```ignore
/// #[repr(C)]
/// struct Foo {
///    some_data: u32,
///    entries: __IncompleteArrayField<__u32>,
/// }
/// ```
///
/// In order to allocate such a structure, `size_of::<Foo>()` would be too small because it would
/// not include any space for `entries`. To make the allocation large enough while still being
/// aligned for `Foo`, a `Vec<Foo>` is created. Only the first element of `Vec<Foo>` would actually
/// be used as a `Foo`. The remaining memory in the `Vec<Foo>` is for `entries`, which must be
/// contiguous with `Foo`. This function is used to make the `Vec<Foo>` with enough space for
/// `count` entries.
pub fn vec_with_array_field<T: Default, F>(count: usize) -> Vec<T> {
    let element_space = count * size_of::<F>();
    let vec_size_bytes = size_of::<T>() + element_space;
    vec_with_size_in_bytes(vec_size_bytes)
}

/// The following code provides generic helpers for creating and accessing flexible array structs.
/// A complete definition of flexible array structs is found in the ISO 9899 specification
/// (http://www.iso-9899.info/n1570.html). A flexible array struct is of the form:
///
/// ```ignore
/// #[repr(C)]
/// struct T {
///    some_data: u32,
///    nents: u32,
///    entries: __IncompleteArrayField<S>,
/// }
/// ```
/// where:
///
/// - `T` is the flexible array struct type
/// - `S` is the flexible array type
/// - `nents` is the flexible array length
/// - `entries` is the flexible array member
///
/// These structures are used by the kernel API.

/// A collection of methods that are required by the FlexibleArrayWrapper type.
///
/// When implemented for `T`, this trait allows the caller to set number of `S` entries and
/// retrieve a slice of `S` entries.  Trait methods must only be called by the FlexibleArrayWrapper
/// type.  Don't implement this trait directly, use the flexible_array! macro to avoid duplication.
pub trait FlexibleArray<S> {
    /// Implementations must set flexible array length in the flexible array struct to the value
    /// specified by `len`. Appropriate conversions (i.e, usize to u32) are allowed so long as
    /// they don't overflow or underflow.
    fn set_len(&mut self, len: usize);
    /// Implementations must return the length of the flexible array member.  Appropriate
    /// conversions (i.e, usize to u32) are allowed so long as they don't overflow or underflow.
    fn get_len(&self) -> usize;
    /// Implementations must return a slice of flexible array member of length `len`.
    /// # Safety
    /// Do not use this function directly, as the FlexibleArrayWrapper will guarantee safety.
    unsafe fn get_slice(&self, len: usize) -> &[S];
    /// Implementations must return a mutable slice of flexible array member of length `len`.
    /// # Safety
    /// Do not use this function directly, as the FlexibleArrayWrapper will guarantee safety.
    unsafe fn get_mut_slice(&mut self, len: usize) -> &mut [S];
}

/// Always use this macro for implementing the FlexibleArray<`S`> trait for a given `T`.  There
/// exists an 1:1 mapping of macro identifiers to the definitions in the FlexibleArray<`S`>
/// documentation, so refer to that for more information.
#[macro_export]
macro_rules! flexible_array_impl {
    ($T:ident, $S:ident, $nents:ident, $entries:ident) => {
        impl FlexibleArray<$S> for $T {
            fn set_len(&mut self, len: usize) {
                self.$nents = ::std::convert::TryInto::try_into(len).unwrap();
            }

            fn get_len(&self) -> usize {
                self.$nents as usize
            }

            unsafe fn get_slice(&self, len: usize) -> &[$S] {
                self.$entries.as_slice(len)
            }

            unsafe fn get_mut_slice(&mut self, len: usize) -> &mut [$S] {
                self.$entries.as_mut_slice(len)
            }
        }
    };
}

pub struct FlexibleArrayWrapper<T, S> {
    entries: Vec<T>,
    phantom: PhantomData<S>,
    allocated_len: usize,
}

/// Convenience wrapper for flexible array structs.
///
/// The FlexibleArray trait must be implemented for the flexible array struct before using this
/// wrapper.
impl<T, S> FlexibleArrayWrapper<T, S>
where
    T: FlexibleArray<S> + Default,
{
    /// Creates a new FlexibleArrayWrapper for the given flexible array struct type and flexible
    /// array type. The flexible array length is set to `array_len`. vec_with_array_field is used
    /// to make sure the resultant wrapper is appropriately sized.
    pub fn new(array_len: usize) -> FlexibleArrayWrapper<T, S> {
        let mut entries = vec_with_array_field::<T, S>(array_len);
        entries[0].set_len(array_len);

        FlexibleArrayWrapper {
            entries,
            phantom: PhantomData,
            allocated_len: array_len,
        }
    }

    /// Mapping the unsized array to a slice is unsafe because the length isn't known.  Using
    /// the length we originally allocated with eliminates the possibility of overflow.
    fn get_valid_len(&self) -> usize {
        if self.entries[0].get_len() > self.allocated_len {
            self.allocated_len
        } else {
            self.entries[0].get_len()
        }
    }

    /// Returns a slice of the flexible array member, for inspecting. To modify, use
    /// mut_entries_slice instead.
    pub fn entries_slice(&self) -> &[S] {
        let valid_length = self.get_valid_len();
        // Safe because the length has been validated.
        unsafe { self.entries[0].get_slice(valid_length) }
    }

    /// Returns a mutable slice of the flexible array member, for modifying.
    pub fn mut_entries_slice(&mut self) -> &mut [S] {
        let valid_length = self.get_valid_len();
        self.entries[0].set_len(valid_length);
        // Safe because the length has been validated.
        unsafe { self.entries[0].get_mut_slice(valid_length) }
    }

    /// Get a pointer so it can be passed to the kernel. Callers must not access the flexible
    /// array member.  Using this pointer is unsafe.
    pub fn as_ptr(&self) -> *const T {
        &self.entries[0]
    }

    /// Get a mutable pointer so it can be passed to the kernel. Callers must not access the
    /// flexible array member.  Using this pointer is unsafe.
    pub fn as_mut_ptr(&mut self) -> *mut T {
        &mut self.entries[0]
    }
}
<|endoftext|>