<commit_before>#[feature(globs)];

extern mod native;
extern mod libusb;

use libusb::*;
use std::unstable::intrinsics;
use std::libc::{c_int, c_void, size_t, malloc, free};
use std::vec;
use std::ptr::{to_mut_unsafe_ptr};
use std::result::Result;
use std::comm::{Port, Chan, SharedChan};
use std::cast::transmute;
use std::mem::size_of;

use std::sync::arc::UnsafeArc;
use std::sync::atomics::{AtomicInt, SeqCst};

use native::task;

pub struct ContextData {
	priv ctx: *mut libusb_context,
	priv open_device_count: AtomicInt
}

impl Drop for ContextData {
	fn drop(&mut self) {
		unsafe {
			assert!(self.open_device_count.load(SeqCst) == 0);
			// TODO: make sure backend thread is dead if the last device just closed
			libusb_exit(self.ctx);
		}
	}
}

pub struct Context {
	priv bx: UnsafeArc<ContextData>
}

impl Context {
	pub fn new() -> Context {
		unsafe{
			let mut ctx: *mut libusb_context = intrinsics::init();
			let r = libusb_init(&mut ctx);

			Context{
				bx: UnsafeArc::new(ContextData{
					ctx: ctx,
					open_device_count: AtomicInt::new(0)
				})
			}
		}
	}

	pub fn ptr(&self) -> *mut libusb_context {
		unsafe{
			(*self.bx.get()).ctx
		}
	}

	pub fn setDebug(&self, level: int) {
		unsafe{
			libusb_set_debug(self.ptr(), level as c_int);
		}
	}

	pub fn listDevices(&self) -> ~[Device] {
		unsafe{
			let mut list: *mut *mut libusb_device = intrinsics::init();
			let num_devices = libusb_get_device_list(self.ptr(), &mut list);
			let r = vec::raw::mut_buf_as_slice(list, num_devices as uint, |l|{
				l.iter().map(|i| Device{dev: *i, ctx: (*self).clone()}).collect()
			});

			libusb_free_device_list(list, 0);
			r
		}
	}

	pub fn find_by_vid_pid(&self, vid: uint, pid: uint) -> Option<Device> {
		self.listDevices().move_iter().find(|d| {
			let desc = d.descriptor();
			desc.idVendor as uint == vid && desc.idProduct as uint == pid
		})
	}

	fn device_opened(&self) {
		let count = unsafe { &mut (*self.bx.get()).open_device_count };
		let old_count = count.fetch_add(1, SeqCst);

		if old_count == 0 {
			let bx = self.bx.clone();

			fn threadfn(tbx: &UnsafeArc<ContextData>) {
				unsafe {
					let ctx = (*tbx.get()).ctx;
					let count = &(*tbx.get()).open_device_count;

					while (count.load(SeqCst) > 0) {
						libusb_handle_events(ctx);
					}
				}
			}

			do task::spawn {
				threadfn(&bx);
			}
		}
	}

	fn device_closed(&self) {
		let count = unsafe { &mut (*self.bx.get()).open_device_count };
		count.fetch_sub(1, SeqCst);
	}
}

extern fn rust_usb_callback(transfer: *mut libusb_transfer) {
	unsafe {
		let chan: ~Chan<()> = transmute((*transfer).user_data);
		chan.send(());
	}
}

struct TH {
	t: *mut libusb_transfer,
	c: SharedChan<*mut libusb_transfer>,
}

impl Drop for TH{
	fn drop(&mut self) {
		unsafe {
			free((*self.t).buffer as *c_void);
			libusb_free_transfer(self.t);
		}
	}
}

extern fn rust_usb_stream_callback(transfer: *mut libusb_transfer) {
	unsafe {
		let chan: &SharedChan<*mut libusb_transfer> = transmute((*transfer).user_data);
		chan.send(transfer);
	}
}

impl Clone for Context{
	fn clone(&self) -> Context{
		Context{bx: self.bx.clone()}
	}
}

pub struct Device {
	priv dev: *mut libusb_device,
	ctx: Context
}

impl Device {
	pub fn descriptor(&self) -> ~libusb_device_descriptor {
		unsafe{
			let mut d: ~libusb_device_descriptor = ~intrinsics::uninit();
			libusb_get_device_descriptor(self.dev, to_mut_unsafe_ptr(d));
			d
		}
	}

	pub fn bus(&self) -> int {
		unsafe {
			libusb_get_bus_number(self.dev) as int
		}
	}

	pub fn address(&self) -> int {
		unsafe {
			libusb_get_device_address(self.dev) as int
		}
	}

	pub fn open(&self) -> Result<DeviceHandle, int> {
		unsafe {
			let mut handle: *mut libusb_device_handle = intrinsics::uninit();
			let r = libusb_open(self.dev, &mut handle);
			if (r == 0){
				self.ctx.device_opened();
				Ok(DeviceHandle {
					bx: UnsafeArc::new(DeviceHandleData {
						dev: handle,
						ctx: self.ctx.clone()
					})
				})
			}else{
				Err(r as int)
			}
		}
	}
}

impl Drop for Device {
	fn drop(&mut self) {
		unsafe {
			libusb_unref_device(self.dev);
		}
	}
}

impl Clone for Device {
	fn clone(&self) -> Device {
		unsafe {
			libusb_ref_device(self.dev);
		}
		Device{dev: self.dev, ctx: self.ctx.clone()}
	}
}

struct DeviceHandleData{
	priv dev: *mut libusb_device_handle,
	priv ctx: Context
}

impl Drop for DeviceHandleData {
	fn drop(&mut self) {
		unsafe {
			self.ctx.device_closed();
			libusb_close(self.dev);
		}
	}
}

pub struct DeviceHandle {
	priv bx: UnsafeArc<DeviceHandleData>
}

impl DeviceHandle {
	pub fn ptr(&self) -> *mut libusb_device_handle {
		unsafe {
			(*self.bx.get()).dev
		}
	}

	pub fn claim_interface(&self,
		iface_num: uint) {
		unsafe {
		libusb_claim_interface(self.ptr(), iface_num as c_int);
		}
	}

	pub unsafe fn submit_transfer_sync(&self,
		endpoint: u8,
		transfer_type: libusb_transfer_type,
		length: uint,
		buffer: *mut u8) -> (libusb_transfer_status, uint) {

		let (port, chan): (Port<()>, Chan<()>) = Chan::new();

		let t = libusb_alloc_transfer(0);
		(*t).dev_handle = self.ptr();
		(*t).endpoint = endpoint;
		(*t).transfer_type = transfer_type as u8;
		(*t).timeout = 0;
		(*t).length = length as c_int;
		(*t).callback = rust_usb_callback;
		(*t).user_data = transmute(~chan);
		(*t).buffer = buffer;

		libusb_submit_transfer(t);
		port.recv();
		let r = ((*t).get_status(), (*t).actual_length as uint);
		libusb_free_transfer(t);
		return r;
	}

	pub fn read(&self,
			endpoint: u8,
			transfer_type: libusb_transfer_type,
			size: uint
			) -> Result<~[u8], libusb_transfer_status> {
		let mut buf: ~[u8] = vec::from_elem(size, 0u8);
		unsafe {
			let ptr = buf.as_mut_ptr();
			let (status, actual_length) = self.submit_transfer_sync(
				endpoint, transfer_type, size, ptr);

			if status == LIBUSB_TRANSFER_COMPLETED {
				buf.truncate(actual_length);
				Ok(buf)
			} else {
				Err(status)
			}
		}
	}

	pub fn write(&self, endpoint: u8, transfer_type: libusb_transfer_type, buf: &[u8]) -> Result<(), libusb_transfer_status> {
		unsafe {
			let ptr = buf.as_ptr() as *mut u8;

			let (status, _) = self.submit_transfer_sync(
				endpoint, transfer_type, buf.len(), ptr);

			if status == LIBUSB_TRANSFER_COMPLETED {
				Ok(())
			} else {
				Err(status)
			}
		}
	}

	pub fn ctrl_read(&self, bmRequestType: u8, bRequest: u8,
		wValue:u16, wIndex: u16, length: uint) -> Result<~[u8], libusb_transfer_status> {

		let setup_length = size_of::<libusb_control_setup>();
		let total_length = setup_length + length as uint;
		let mut buf: ~[u8] = vec::from_elem(total_length, 0u8);

		unsafe{
			let ptr = fill_setup_buf(buf, bmRequestType, bRequest, wValue, wIndex, length);

			let (status, actual_length) = self.submit_transfer_sync(
				0, LIBUSB_TRANSFER_TYPE_CONTROL, total_length, ptr);

			if status == LIBUSB_TRANSFER_COMPLETED {
				Ok(buf.slice(setup_length, setup_length+actual_length).to_owned())
			} else {
				Err(status)
			}
		}
	}

	pub fn ctrl_write(&self, bmRequestType: u8, bRequest: u8,
		wValue:u16, wIndex: u16, buf: &[u8]) -> Result<(), libusb_transfer_status> {
		unsafe {
			let mut setup_buf = vec::from_elem(size_of::<libusb_control_setup>(), 0u8);
			fill_setup_buf(setup_buf, bmRequestType, bRequest, wValue, wIndex, buf.len());
			self.write(0, LIBUSB_TRANSFER_TYPE_CONTROL, setup_buf+buf)
		}
	}

	unsafe fn stream_transfers(&self, endpoint: u8,
			transfer_type: libusb_transfer_type, size: uint,
			num_transfers: uint) -> (Port<*mut libusb_transfer>, ~[TH]) {

		let (port, chan) = SharedChan::new();

		let transfers = vec::from_fn(num_transfers, |_| { TH {
			t: libusb_alloc_transfer(0),
			c: chan.clone(),
		}});

		for th in transfers.iter() {
			(*th.t).dev_handle = self.ptr();
			(*th.t).endpoint = endpoint;
			(*th.t).transfer_type = transfer_type as u8;
			(*th.t).timeout = 0;
			(*th.t).length = size as i32;
			(*th.t).callback = rust_usb_stream_callback;
			(*th.t).buffer = malloc(size as size_t) as *mut u8;
			(*th.t).user_data = transmute(&th.c);
		}

		return (port, transfers);
	}

	pub fn read_stream(&self, endpoint: u8,
			transfer_type: libusb_transfer_type,
			size: uint, mut num_transfers: uint, cb: |Result<&[u8], libusb_transfer_status>| -> bool) {

		unsafe {
			let mut running = true;
			let (port, transfers) = self.stream_transfers(
				endpoint, transfer_type, size, num_transfers);

			for th in transfers.iter() {
				libusb_submit_transfer(th.t);
			}

			while (num_transfers > 0) {
				let transfer: *mut libusb_transfer = port.recv();

				if ((*transfer).get_status() == LIBUSB_TRANSFER_COMPLETED) {
					vec::raw::buf_as_slice((*transfer).buffer as *u8, size, |b| {
						running &= cb(Ok(b))
					});
				} else {
					running = false;
					running &= cb(Err((*transfer).get_status()));
				}

				if (running) {
					let r = libusb_submit_transfer(transfer);
					assert!(r == 0);
				} else {
					num_transfers -= 1;
				}
			}
		}
	}

	pub fn write_stream(&self, endpoint: u8,
			transfer_type: libusb_transfer_type,
			size: uint, num_transfers: uint, cb: |Result<(&mut[u8]), libusb_transfer_status>| -> bool) {

		unsafe {
			let mut running = true;
			let mut running_transfers = 0;
			let (port, transfers) = self.stream_transfers(
				endpoint, transfer_type, size, num_transfers);

			for th in transfers.iter() {
				vec::raw::mut_buf_as_slice((*th.t).buffer, size, |b| {
					running &= cb(Ok(b));
				});
				if running {
					libusb_submit_transfer(th.t);
					running_transfers += 1;
				} else {
					break;
				}
			}

			while (running_transfers > 0) {
				let transfer: *mut libusb_transfer = port.recv();

				if ((*transfer).get_status() == LIBUSB_TRANSFER_COMPLETED) {
					vec::raw::mut_buf_as_slice((*transfer).buffer, size, |b| {
						running &= cb(Ok(b))
					});
				} else {
					running = false;
					running &= cb(Err((*transfer).get_status()));
				}

				if (running) {
					let r = libusb_submit_transfer(transfer);
					assert!(r == 0);
				} else {
					running_transfers -= 1;
				}
			}
		}
	}
}

unsafe fn fill_setup_buf(buf: &mut [u8], bmRequestType: u8,
	bRequest: u8, wValue:u16, wIndex: u16, length: uint) -> *mut u8 {
	let ptr = buf.as_mut_ptr();
	let setup = ptr as *mut libusb_control_setup;

	// TODO: these are always little-endian
	(*setup).bmRequestType = bmRequestType;
	(*setup).bRequest = bRequest;
	(*setup).wValue = wValue;
	(*setup).wIndex = wIndex;
	(*setup).wLength = length as u16;

	return ptr;
}

impl Clone for DeviceHandle {
	fn clone(&self) -> DeviceHandle {
		DeviceHandle{bx: self.bx.clone()}
	}
}
<commit_msg>updates for latest rust<commit_after>#[feature(globs)];

extern crate native;
extern crate libusb;

use libusb::*;
use std::unstable::intrinsics;
use std::libc::{c_int, c_void, size_t, malloc, free};
use std::vec;
//use std::ptr::{to_mut_unsafe_ptr};
use std::result::Result;
use std::comm::{Port, Chan};
use std::cast::transmute;
use std::mem::size_of;

use std::sync::arc::UnsafeArc;
use std::sync::atomics::{AtomicInt, SeqCst};

use native::task;

pub struct ContextData {
	ctx: *mut libusb_context,
	open_device_count: AtomicInt
}

impl Drop for ContextData {
	fn drop(&mut self) {
		unsafe {
			assert!(self.open_device_count.load(SeqCst) == 0);
			// TODO: make sure backend thread is dead if the last device just closed
			libusb_exit(self.ctx);
		}
	}
}

pub struct Context {
	bx: UnsafeArc<ContextData>
}

impl Context {
	pub fn new() -> Context {
		unsafe{
			let mut ctx: *mut libusb_context = intrinsics::init();
			let r = libusb_init(&mut ctx);

			Context{
				bx: UnsafeArc::new(ContextData{
					ctx: ctx,
					open_device_count: AtomicInt::new(0)
				})
			}
		}
	}

	pub fn ptr(&self) -> *mut libusb_context {
		unsafe{
			(*self.bx.get()).ctx
		}
	}

	pub fn setDebug(&self, level: int) {
		unsafe{
			libusb_set_debug(self.ptr(), level as c_int);
		}
	}

	pub fn listDevices(&self) -> ~[Device] {
		unsafe{
			let mut list: *mut *mut libusb_device = intrinsics::init();
			let num_devices = libusb_get_device_list(self.ptr(), &mut list);
			let r = vec::raw::mut_buf_as_slice(list, num_devices as uint, |l|{
				l.iter().map(|i| Device{dev: *i, ctx: (*self).clone()}).collect()
			});

			libusb_free_device_list(list, 0);
			r
		}
	}

	pub fn find_by_vid_pid(&self, vid: uint, pid: uint) -> Option<Device> {
		self.listDevices().move_iter().find(|d| {
			let desc = d.descriptor();
			desc.idVendor as uint == vid && desc.idProduct as uint == pid
		})
	}

	fn device_opened(&self) {
		let count = unsafe { &mut (*self.bx.get()).open_device_count };
		let old_count = count.fetch_add(1, SeqCst);

		if old_count == 0 {
			let bx = self.bx.clone();

			fn threadfn(tbx: &UnsafeArc<ContextData>) {
				unsafe {
					let ctx = (*tbx.get()).ctx;
					let count = &(*tbx.get()).open_device_count;

					while (count.load(SeqCst) > 0) {
						libusb_handle_events(ctx);
					}
				}
			}

			task::spawn(proc() {threadfn(&bx);});
		}
	}

	fn device_closed(&self) {
		let count = unsafe { &mut (*self.bx.get()).open_device_count };
		count.fetch_sub(1, SeqCst);
	}
}

extern fn rust_usb_callback(transfer: *mut libusb_transfer) {
	unsafe {
		let chan: ~Chan<()> = transmute((*transfer).user_data);
		chan.send(());
	}
}

struct TH {
	t: *mut libusb_transfer,
	c: Chan<*mut libusb_transfer>,
}

impl Drop for TH{
	fn drop(&mut self) {
		unsafe {
			free((*self.t).buffer as *mut c_void);
			libusb_free_transfer(self.t);
		}
	}
}

extern fn rust_usb_stream_callback(transfer: *mut libusb_transfer) {
	unsafe {
		let chan: &Chan<*mut libusb_transfer> = transmute((*transfer).user_data);
		chan.send(transfer);
	}
}

impl Clone for Context{
	fn clone(&self) -> Context{
		Context{bx: self.bx.clone()}
	}
}

pub struct Device {
	dev: *mut libusb_device,
	ctx: Context
}

impl Device {
	pub fn descriptor(&self) -> ~libusb_device_descriptor {
		unsafe{
			let mut d: libusb_device_descriptor = intrinsics::uninit();
			libusb_get_device_descriptor(self.dev, transmute(&d));
			~d
		}
	}

	pub fn bus(&self) -> int {
		unsafe {
			libusb_get_bus_number(self.dev) as int
		}
	}

	pub fn address(&self) -> int {
		unsafe {
			libusb_get_device_address(self.dev) as int
		}
	}

	pub fn open(&self) -> Result<DeviceHandle, int> {
		unsafe {
			let mut handle: *mut libusb_device_handle = intrinsics::uninit();
			let r = libusb_open(self.dev, &mut handle);
			if (r == 0){
				self.ctx.device_opened();
				Ok(DeviceHandle {
					bx: UnsafeArc::new(DeviceHandleData {
						dev: handle,
						ctx: self.ctx.clone()
					})
				})
			}else{
				Err(r as int)
			}
		}
	}
}

impl Drop for Device {
	fn drop(&mut self) {
		unsafe {
			libusb_unref_device(self.dev);
		}
	}
}

impl Clone for Device {
	fn clone(&self) -> Device {
		unsafe {
			libusb_ref_device(self.dev);
		}
		Device{dev: self.dev, ctx: self.ctx.clone()}
	}
}

struct DeviceHandleData{
	dev: *mut libusb_device_handle,
	ctx: Context
}

impl Drop for DeviceHandleData {
	fn drop(&mut self) {
		unsafe {
			self.ctx.device_closed();
			libusb_close(self.dev);
		}
	}
}

pub struct DeviceHandle {
	bx: UnsafeArc<DeviceHandleData>
}

impl DeviceHandle {
	pub fn ptr(&self) -> *mut libusb_device_handle {
		unsafe {
			(*self.bx.get()).dev
		}
	}

	pub fn claim_interface(&self,
		iface_num: uint) {
		unsafe {
		libusb_claim_interface(self.ptr(), iface_num as c_int);
		}
	}

	pub unsafe fn submit_transfer_sync(&self,
		endpoint: u8,
		transfer_type: libusb_transfer_type,
		length: uint,
		buffer: *mut u8) -> (libusb_transfer_status, uint) {

		let (port, chan): (Port<()>, Chan<()>) = Chan::new();

		let t = libusb_alloc_transfer(0);
		(*t).dev_handle = self.ptr();
		(*t).endpoint = endpoint;
		(*t).transfer_type = transfer_type as u8;
		(*t).timeout = 0;
		(*t).length = length as c_int;
		(*t).callback = rust_usb_callback;
		(*t).user_data = transmute(~chan);
		(*t).buffer = buffer;

		libusb_submit_transfer(t);
		port.recv();
		let r = ((*t).get_status(), (*t).actual_length as uint);
		libusb_free_transfer(t);
		return r;
	}

	pub fn read(&self,
			endpoint: u8,
			transfer_type: libusb_transfer_type,
			size: uint
			) -> Result<~[u8], libusb_transfer_status> {
		let mut buf: ~[u8] = vec::from_elem(size, 0u8);
		unsafe {
			let ptr = buf.as_mut_ptr();
			let (status, actual_length) = self.submit_transfer_sync(
				endpoint, transfer_type, size, ptr);

			if status == LIBUSB_TRANSFER_COMPLETED {
				buf.truncate(actual_length);
				Ok(buf)
			} else {
				Err(status)
			}
		}
	}

	pub fn write(&self, endpoint: u8, transfer_type: libusb_transfer_type, buf: &[u8]) -> Result<(), libusb_transfer_status> {
		unsafe {
			let ptr = buf.as_ptr() as *mut u8;

			let (status, _) = self.submit_transfer_sync(
				endpoint, transfer_type, buf.len(), ptr);

			if status == LIBUSB_TRANSFER_COMPLETED {
				Ok(())
			} else {
				Err(status)
			}
		}
	}

	pub fn ctrl_read(&self, bmRequestType: u8, bRequest: u8,
		wValue:u16, wIndex: u16, length: uint) -> Result<~[u8], libusb_transfer_status> {

		let setup_length = size_of::<libusb_control_setup>();
		let total_length = setup_length + length as uint;
		let mut buf: ~[u8] = vec::from_elem(total_length, 0u8);

		unsafe{
			let ptr = fill_setup_buf(buf, bmRequestType, bRequest, wValue, wIndex, length);

			let (status, actual_length) = self.submit_transfer_sync(
				0, LIBUSB_TRANSFER_TYPE_CONTROL, total_length, ptr);

			if status == LIBUSB_TRANSFER_COMPLETED {
				Ok(buf.slice(setup_length, setup_length+actual_length).to_owned())
			} else {
				Err(status)
			}
		}
	}

	pub fn ctrl_write(&self, bmRequestType: u8, bRequest: u8,
		wValue:u16, wIndex: u16, buf: &[u8]) -> Result<(), libusb_transfer_status> {
		unsafe {
			let mut setup_buf = vec::from_elem(size_of::<libusb_control_setup>(), 0u8);
			fill_setup_buf(setup_buf, bmRequestType, bRequest, wValue, wIndex, buf.len());
			self.write(0, LIBUSB_TRANSFER_TYPE_CONTROL, setup_buf+buf)
		}
	}

	unsafe fn stream_transfers(&self, endpoint: u8,
			transfer_type: libusb_transfer_type, size: uint,
			num_transfers: uint) -> (Port<*mut libusb_transfer>, ~[TH]) {

		let (port, chan) = Chan::new();

		let transfers = vec::from_fn(num_transfers, |_| { TH {
			t: libusb_alloc_transfer(0),
			c: chan.clone(),
		}});

		for th in transfers.iter() {
			(*th.t).dev_handle = self.ptr();
			(*th.t).endpoint = endpoint;
			(*th.t).transfer_type = transfer_type as u8;
			(*th.t).timeout = 0;
			(*th.t).length = size as i32;
			(*th.t).callback = rust_usb_stream_callback;
			(*th.t).buffer = malloc(size as size_t) as *mut u8;
			(*th.t).user_data = transmute(&th.c);
		}

		return (port, transfers);
	}

	pub fn read_stream(&self, endpoint: u8,
			transfer_type: libusb_transfer_type,
			size: uint, mut num_transfers: uint, cb: |Result<&[u8], libusb_transfer_status>| -> bool) {

		unsafe {
			let mut running = true;
			let (port, transfers) = self.stream_transfers(
				endpoint, transfer_type, size, num_transfers);

			for th in transfers.iter() {
				libusb_submit_transfer(th.t);
			}

			while (num_transfers > 0) {
				let transfer: *mut libusb_transfer = port.recv();

				if ((*transfer).get_status() == LIBUSB_TRANSFER_COMPLETED) {
					vec::raw::buf_as_slice((*transfer).buffer as *u8, size, |b| {
						running &= cb(Ok(b))
					});
				} else {
					running = false;
					running &= cb(Err((*transfer).get_status()));
				}

				if (running) {
					let r = libusb_submit_transfer(transfer);
					assert!(r == 0);
				} else {
					num_transfers -= 1;
				}
			}
		}
	}

	pub fn write_stream(&self, endpoint: u8,
			transfer_type: libusb_transfer_type,
			size: uint, num_transfers: uint, cb: |Result<(&mut[u8]), libusb_transfer_status>| -> bool) {

		unsafe {
			let mut running = true;
			let mut running_transfers = 0;
			let (port, transfers) = self.stream_transfers(
				endpoint, transfer_type, size, num_transfers);

			for th in transfers.iter() {
				vec::raw::mut_buf_as_slice((*th.t).buffer, size, |b| {
					running &= cb(Ok(b));
				});
				if running {
					libusb_submit_transfer(th.t);
					running_transfers += 1;
				} else {
					break;
				}
			}

			while (running_transfers > 0) {
				let transfer: *mut libusb_transfer = port.recv();

				if ((*transfer).get_status() == LIBUSB_TRANSFER_COMPLETED) {
					vec::raw::mut_buf_as_slice((*transfer).buffer, size, |b| {
						running &= cb(Ok(b))
					});
				} else {
					running = false;
					running &= cb(Err((*transfer).get_status()));
				}

				if (running) {
					let r = libusb_submit_transfer(transfer);
					assert!(r == 0);
				} else {
					running_transfers -= 1;
				}
			}
		}
	}
}

unsafe fn fill_setup_buf(buf: &mut [u8], bmRequestType: u8,
	bRequest: u8, wValue:u16, wIndex: u16, length: uint) -> *mut u8 {
	let ptr = buf.as_mut_ptr();
	let setup = ptr as *mut libusb_control_setup;

	// TODO: these are always little-endian
	(*setup).bmRequestType = bmRequestType;
	(*setup).bRequest = bRequest;
	(*setup).wValue = wValue;
	(*setup).wIndex = wIndex;
	(*setup).wLength = length as u16;

	return ptr;
}

impl Clone for DeviceHandle {
	fn clone(&self) -> DeviceHandle {
		DeviceHandle{bx: self.bx.clone()}
	}
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;

extern crate serde;
extern crate serde_json;

use base64_vlq;

static SOURCE_MAP_VERSION: u32 = 3;

#[allow(non_snake_case)]
#[derive(Deserialize, Debug)]
struct SourceMap {
  version: u32,
  sources: Vec<String>,
  names: Vec<String>,
  sourceRoot: Option<String>,
  mappings: String,
  file: Option<String>

  // We skip this. Keeping megabytes of data that we do not care about
  // in memory seems reckless to caches.
  //sourcesContent: Option<vec<String>>,
}

#[derive(Clone, Eq, PartialEq, Debug)]
pub struct CodePosition {
  /** Line number in a code file, starting from 1 */
  pub line: u32,
  /** Column number in a code file, starting from 0 */
  pub column: u32
}

#[derive(Clone, Eq, PartialEq, Debug)]
pub struct Mapping {
  /** The position in the generated file */
  pub generated: CodePosition,
  /** The position in the corresponding original source file */
  pub original: CodePosition,
  /** The original source file */
  pub source: String,
  /** The original source name of the function/class, if applicable */
  pub name: String
}

pub struct Cache {
  generated_mappings: Vec<Mapping>,
  /** The path prefix of mapping source paths */
  pub source_root: String
}

/**
 * consume parses a SourceMap into a cache that can be queried for mappings
 *
 * The only parameter is the raw source map as a JSON string.
 * According to the [source map spec][source-map-spec], source maps have the following attributes:
 *
 *   - version: Which version of the source map spec this map is following.
 *   - sources: An array of URLs to the original source files.
 *   - names: An array of identifiers which can be referrenced by individual mappings.
 *   - sourceRoot: Optional. The URL root from which all sources are relative.
 *   - sourcesContent: Optional. An array of contents of the original source files.
 *   - mappings: A string of base64 VLQs which contain the actual mappings.
 *   - file: Optional. The generated file this source map is associated with.
 *
 * Here is an example source map:
 *
 * ```json
 *     {
 *       "version": 3,
 *       "file": "out.js",
 *       "sourceRoot" : "",
 *       "sources": ["foo.js", "bar.js"],
 *       "names": ["src", "maps", "are", "fun"],
 *       "mappings": "AA,AB;;ABCDE;"
 *     }
 * ```
 *
 * [source-map-spec]: https://docs.google.com/document/d/1U1RGAehQwRypUTovF1KRlpiOFze0b-_2gc6fAH0KY0k/edit?pli=1#
 */
pub fn consume(source_map_json: &str) -> Result<Cache, String> {
  let source_map: SourceMap = match serde_json::from_str(source_map_json) {
    Ok(x) => x,
    Err(err) => return Err(format!("{}", err))
  };

  parse_mappings(&source_map)
}

fn parse_mappings(source_map: &SourceMap) -> Result<Cache, String>{
  if source_map.version != SOURCE_MAP_VERSION {
    return Err("Only Source Map version 3 is implemented".into())
  }

  let sources_length = source_map.sources.len() as u32;
  let names_length = source_map.names.len() as u32;

  let mut generated_mappings: Vec<Mapping> = Vec::new();

  let mut generated_line: u32 = 0;
  let mut previous_original_line: u32 = 0;
  let mut previous_original_column: u32 = 0;
  let mut previous_source: u32 = 0;
  let mut previous_name: u32 = 0;

  for line in source_map.mappings.as_bytes().split(|&x| x == (';' as u8)) {
    generated_line += 1;
    let mut previous_generated_column: u32 = 0;

    for segment in line.split(|&x| x == (',' as u8)) {
      let segment_length = segment.len();
      let mut fields: Vec<i32> = Vec::new();
      let mut character_index = 0;
      while character_index < segment_length {
        match base64_vlq::decode(&segment[character_index..segment_length]) {
          Some((value, field_length)) => {
            fields.push(value);
            character_index += field_length;
          },
          None => return Err("Invalid VLQ mapping field".into())
        };
      }

      if fields.len() < 1 {
        continue;
      }

      if fields.len() == 2 {
        return Err("Found a source, but no line and column".into());
      }

      if fields.len() == 3 {
        return Err("Found a source and line, but no column".into());
      }

      let mut mapping = Mapping {
        generated: CodePosition {
          line: generated_line,
          column: ((previous_generated_column as i32) + fields[0]) as u32
        },
        original: CodePosition {
          line: 0,
          column: 0
        },
        source: "".into(),
        name: "".into()
      };

      previous_generated_column = mapping.generated.column;

      if fields.len() > 1 {
        // Original source.
        previous_source = ((previous_source as i32) + fields[1]) as u32;
        if previous_source < sources_length {
          mapping.source = source_map.sources[previous_source as usize].to_owned();
        } else {
          return Err(format!("Invalid source map: reference to source index {} when source list length is {}", previous_source, sources_length));
        }

        // Original line.
        previous_original_line = ((previous_original_line as i32) + fields[2]) as u32;
        // Lines are stored 0-based
        mapping.original.line = previous_original_line + 1;

        // Original column.
        previous_original_column = ((previous_original_column as i32) + fields[3]) as u32;
        mapping.original.column = previous_original_column;

        if fields.len() > 4 {
          // Original name.
          previous_name = ((previous_name as i32) + fields[4]) as u32;
          if previous_name < names_length {
            mapping.name = source_map.names[previous_name as usize].to_owned();
          } else {
            return Err(format!("Invalid source map: reference to name index {} when name list length is {}", previous_name, names_length));
          }
        }
      }

      generated_mappings.push(mapping);
    }
  }

  if generated_mappings.len() < 1 {
    return Err("Source Map contains no mappings".to_owned());
  }

  fn sort_key(mapping: &Mapping) -> (u32, u32) {
    (mapping.generated.line, mapping.generated.column)
  }
  generated_mappings.sort_by(|a, b| sort_key(a).cmp(&sort_key(b)));

  Ok(Cache {
    generated_mappings: generated_mappings,
    source_root: match &source_map.sourceRoot {
      &Some(ref x) => x.to_owned(),
      &None => "".into()
    }
  })
}

impl Cache {
  /**
   * Returns the original source, line, column and name information for the generated
   * source's line and column positions provided.
   *
   * # Arguments
   *
   * * line: The line number in the generated source.
   * * column: The column number in the generated source.
   *
   * # Examples
   *
   * ```
   * use js_source_mapper::consume;
   *
   * let cache = consume(r#"{ "version": 3, "file": "foo.js", "sources": ["source.js"], "names": ["name1", "name1", "name3"], "mappings": ";EAACA;;IAEEA;;MAEEE", "sourceRoot": "http://example.com" }"#).unwrap();
   *
   * println!("{:?}", cache.mapping_for_generated_position(2, 2));
   * // => Mapping {
   * //   generated: CodePosition { line: 2, column: 2 },
   * //   original: CodePosition { line: 1, column: 1 },
   * //   source: "source.js"
   * //   name: "name1"
   * // }
   * ```
   *
   */

  pub fn mapping_for_generated_position(&self, line: u32, column: u32) -> Mapping {
    let matcher = |mapping: &Mapping| -> Ordering {
      (mapping.generated.line, mapping.generated.column).cmp(&(line, column))
    };
    let mappings = &self.generated_mappings;
    match mappings.binary_search_by(matcher) {
      Ok(index) => &self.generated_mappings[index],
      Err(index) => &self.generated_mappings[if index > mappings.len() { mappings.len() - 1 } else { index }]
    }.clone()
  }
}

macro_rules! assert_equal_mappings(
  ($a:expr, $b:expr) => (
    if $a != $b {
      panic!(format!("\n\n{:?}\n\n!=\n\n{:?}\n\n", $a, $b));
    }
  );
);

#[test]
fn test_source_map_issue_64() {
  let cache = consume(r#"{
    "version": 3,
    "file": "foo.js",
    "sourceRoot": "http://example.com/",
    "sources": ["/a"],
    "names": [],
    "mappings": "AACA",
    "sourcesContent": ["foo"]
  }"#).unwrap();

  let expected = Mapping {
    generated: CodePosition { line: 1, column: 0 },
    original: CodePosition { line: 2, column: 0 },
    source: "/a".into(),
    name: "".into()
  };
  let actual = cache.mapping_for_generated_position(1, 0);
  assert_equal_mappings!(actual, expected);
}

#[test]
fn test_source_map_issue_72_duplicate_sources() {
  let cache = consume(r#"{
    "version": 3,
    "file": "foo.js",
    "sources": ["source1.js", "source1.js", "source3.js"],
    "names": [],
    "mappings": ";EAAC;;IAEE;;MEEE",
    "sourceRoot": "http://example.com"
  }"#).unwrap();

  {
    let expected = Mapping {
      generated: CodePosition { line: 2, column: 2 },
      original: CodePosition { line: 1, column: 1 },
      source: "source1.js".into(),
      name: "".into()
    };
    let actual = cache.mapping_for_generated_position(2, 2);
    assert_equal_mappings!(actual, expected);
  }

  {
    let expected = Mapping {
      generated: CodePosition { line: 4, column: 4 },
      original: CodePosition { line: 3, column: 3 },
      source: "source1.js".into(),
      name: "".into()
    };
    let actual = cache.mapping_for_generated_position(4, 4);
    assert_equal_mappings!(actual, expected);
  }

  {
    let expected = Mapping {
      generated: CodePosition { line: 6, column: 6 },
      original: CodePosition { line: 5, column: 5 },
      source: "source3.js".into(),
      name: "".into()
    };
    let actual = cache.mapping_for_generated_position(6, 6);
    assert_equal_mappings!(actual, expected);
  }
}

#[test]
fn test_source_map_issue_72_duplicate_names() {
  let cache = consume(r#"{
    "version": 3,
    "file": "foo.js",
    "sources": ["source.js"],
    "names": ["name1", "name1", "name3"],
    "mappings": ";EAACA;;IAEEA;;MAEEE",
    "sourceRoot": "http://example.com"
  }"#).unwrap();

  {
    let expected = Mapping {
      generated: CodePosition { line: 2, column: 2 },
      original: CodePosition { line: 1, column: 1 },
      source: "source.js".into(),
      name: "name1".into()
    };
    let actual = cache.mapping_for_generated_position(2, 2);
    assert_equal_mappings!(actual, expected);
  }

  {
    let expected = Mapping {
      generated: CodePosition { line: 4, column: 4 },
      original: CodePosition { line: 3, column: 3 },
      source: "source.js".into(),
      name: "name1".into()
    };
    let actual = cache.mapping_for_generated_position(4, 4);
    assert_equal_mappings!(actual, expected);
  }

  {
    let expected = Mapping {
      generated: CodePosition { line: 6, column: 6 },
      original: CodePosition { line: 5, column: 5 },
      source: "source.js".into(),
      name: "name3".into()
    };
    let actual = cache.mapping_for_generated_position(6, 6);
    assert_equal_mappings!(actual, expected);
  }
}

#[test]
fn it_allows_omitting_source_root() {
  let cache_result: Result<Cache, String> = consume(r#"{
    "version": 3,
    "file": "foo.js",
    "sources": ["source.js"],
    "names": ["name1", "name1", "name3"],
    "mappings": ";EAACA;;IAEEA;;MAEEE"
  }"#);
  match cache_result {
    Ok(_) => {},
    Err(s) => panic!(format!("Error due to omitting: '{}'", s))
  }
}

#[test]
fn it_rejects_older_source_map_revisions() {
  let cache_result = consume(r#"{
    "version": 2,
    "file": "",
    "sources": ["source.js"],
    "names": ["name1", "name1", "name3"],
    "mappings": ";EAACA;;IAEEA;;MAEEE",
    "sourceRoot": "http://example.com"
  }"#);
  match cache_result {
    Ok(_) => panic!("Source Map revision < 3 should be rejected"),
    Err(_) => {}
  }
}

#[test]
fn it_does_not_panic_due_to_malformed_source_maps() {
  let cache_result = consume(r#"{
    "version": 3,
    "file": "",
    "sources": [],
    "names": [],
    "mappings": ";EAACA;;IAEEA;;MAEEE"
  }"#);
  match cache_result {
    Ok(_) => panic!("Invalid source maps should be rejected"),
    Err(_) => {}
  }
}

#[test]
fn it_returns_error_when_there_are_no_mappings() {
  let cache_result = consume(r#"{
    "version": 3,
    "file": "foo.js",
    "sources": ["source.js"],
    "names": ["name1", "name1", "name3"],
    "mappings": ";;;"
  }"#);
  match cache_result {
    Ok(_) => panic!("Source maps with no mappings should be rejected"),
    Err(_) => {}
  }
}
<commit_msg>Fix bounds check (fuzz)<commit_after>use std::cmp::Ordering;

extern crate serde;
extern crate serde_json;

use base64_vlq;

static SOURCE_MAP_VERSION: u32 = 3;

#[allow(non_snake_case)]
#[derive(Deserialize, Debug)]
struct SourceMap {
  version: u32,
  sources: Vec<String>,
  names: Vec<String>,
  sourceRoot: Option<String>,
  mappings: String,
  file: Option<String>

  // We skip this. Keeping megabytes of data that we do not care about
  // in memory seems reckless to caches.
  //sourcesContent: Option<vec<String>>,
}

#[derive(Clone, Eq, PartialEq, Debug)]
pub struct CodePosition {
  /** Line number in a code file, starting from 1 */
  pub line: u32,
  /** Column number in a code file, starting from 0 */
  pub column: u32
}

#[derive(Clone, Eq, PartialEq, Debug)]
pub struct Mapping {
  /** The position in the generated file */
  pub generated: CodePosition,
  /** The position in the corresponding original source file */
  pub original: CodePosition,
  /** The original source file */
  pub source: String,
  /** The original source name of the function/class, if applicable */
  pub name: String
}

pub struct Cache {
  generated_mappings: Vec<Mapping>,
  /** The path prefix of mapping source paths */
  pub source_root: String
}

/**
 * consume parses a SourceMap into a cache that can be queried for mappings
 *
 * The only parameter is the raw source map as a JSON string.
 * According to the [source map spec][source-map-spec], source maps have the following attributes:
 *
 *   - version: Which version of the source map spec this map is following.
 *   - sources: An array of URLs to the original source files.
 *   - names: An array of identifiers which can be referrenced by individual mappings.
 *   - sourceRoot: Optional. The URL root from which all sources are relative.
 *   - sourcesContent: Optional. An array of contents of the original source files.
 *   - mappings: A string of base64 VLQs which contain the actual mappings.
 *   - file: Optional. The generated file this source map is associated with.
 *
 * Here is an example source map:
 *
 * ```json
 *     {
 *       "version": 3,
 *       "file": "out.js",
 *       "sourceRoot" : "",
 *       "sources": ["foo.js", "bar.js"],
 *       "names": ["src", "maps", "are", "fun"],
 *       "mappings": "AA,AB;;ABCDE;"
 *     }
 * ```
 *
 * [source-map-spec]: https://docs.google.com/document/d/1U1RGAehQwRypUTovF1KRlpiOFze0b-_2gc6fAH0KY0k/edit?pli=1#
 */
pub fn consume(source_map_json: &str) -> Result<Cache, String> {
  let source_map: SourceMap = match serde_json::from_str(source_map_json) {
    Ok(x) => x,
    Err(err) => return Err(format!("{}", err))
  };

  parse_mappings(&source_map)
}

fn parse_mappings(source_map: &SourceMap) -> Result<Cache, String>{
  if source_map.version != SOURCE_MAP_VERSION {
    return Err("Only Source Map version 3 is implemented".into())
  }

  let sources_length = source_map.sources.len() as u32;
  let names_length = source_map.names.len() as u32;

  let mut generated_mappings: Vec<Mapping> = Vec::new();

  let mut generated_line: u32 = 0;
  let mut previous_original_line: u32 = 0;
  let mut previous_original_column: u32 = 0;
  let mut previous_source: u32 = 0;
  let mut previous_name: u32 = 0;

  for line in source_map.mappings.as_bytes().split(|&x| x == (';' as u8)) {
    generated_line += 1;
    let mut previous_generated_column: u32 = 0;

    for segment in line.split(|&x| x == (',' as u8)) {
      let segment_length = segment.len();
      let mut fields: Vec<i32> = Vec::new();
      let mut character_index = 0;
      while character_index < segment_length {
        match base64_vlq::decode(&segment[character_index..segment_length]) {
          Some((value, field_length)) => {
            fields.push(value);
            character_index += field_length;
          },
          None => return Err("Invalid VLQ mapping field".into())
        };
      }

      if fields.len() < 1 {
        continue;
      }

      if fields.len() == 2 {
        return Err("Found a source, but no line and column".into());
      }

      if fields.len() == 3 {
        return Err("Found a source and line, but no column".into());
      }

      let mut mapping = Mapping {
        generated: CodePosition {
          line: generated_line,
          column: ((previous_generated_column as i32) + fields[0]) as u32
        },
        original: CodePosition {
          line: 0,
          column: 0
        },
        source: "".into(),
        name: "".into()
      };

      previous_generated_column = mapping.generated.column;

      if fields.len() > 1 {
        // Original source.
        previous_source = ((previous_source as i32) + fields[1]) as u32;
        if previous_source < sources_length {
          mapping.source = source_map.sources[previous_source as usize].to_owned();
        } else {
          return Err(format!("Invalid source map: reference to source index {} when source list length is {}", previous_source, sources_length));
        }

        // Original line.
        previous_original_line = ((previous_original_line as i32) + fields[2]) as u32;
        // Lines are stored 0-based
        mapping.original.line = previous_original_line + 1;

        // Original column.
        previous_original_column = ((previous_original_column as i32) + fields[3]) as u32;
        mapping.original.column = previous_original_column;

        if fields.len() > 4 {
          // Original name.
          previous_name = ((previous_name as i32) + fields[4]) as u32;
          if previous_name < names_length {
            mapping.name = source_map.names[previous_name as usize].to_owned();
          } else {
            return Err(format!("Invalid source map: reference to name index {} when name list length is {}", previous_name, names_length));
          }
        }
      }

      generated_mappings.push(mapping);
    }
  }

  if generated_mappings.len() < 1 {
    return Err("Source Map contains no mappings".to_owned());
  }

  fn sort_key(mapping: &Mapping) -> (u32, u32) {
    (mapping.generated.line, mapping.generated.column)
  }
  generated_mappings.sort_by(|a, b| sort_key(a).cmp(&sort_key(b)));

  Ok(Cache {
    generated_mappings: generated_mappings,
    source_root: match &source_map.sourceRoot {
      &Some(ref x) => x.to_owned(),
      &None => "".into()
    }
  })
}

impl Cache {
  /**
   * Returns the original source, line, column and name information for the generated
   * source's line and column positions provided.
   *
   * # Arguments
   *
   * * line: The line number in the generated source.
   * * column: The column number in the generated source.
   *
   * # Examples
   *
   * ```
   * use js_source_mapper::consume;
   *
   * let cache = consume(r#"{ "version": 3, "file": "foo.js", "sources": ["source.js"], "names": ["name1", "name1", "name3"], "mappings": ";EAACA;;IAEEA;;MAEEE", "sourceRoot": "http://example.com" }"#).unwrap();
   *
   * println!("{:?}", cache.mapping_for_generated_position(2, 2));
   * // => Mapping {
   * //   generated: CodePosition { line: 2, column: 2 },
   * //   original: CodePosition { line: 1, column: 1 },
   * //   source: "source.js"
   * //   name: "name1"
   * // }
   * ```
   *
   */

  pub fn mapping_for_generated_position(&self, line: u32, column: u32) -> Mapping {
    let matcher = |mapping: &Mapping| -> Ordering {
      (mapping.generated.line, mapping.generated.column).cmp(&(line, column))
    };
    let mappings = &self.generated_mappings;
    match mappings.binary_search_by(matcher) {
      Ok(index) => &self.generated_mappings[index],
      Err(index) => &self.generated_mappings[if index >= mappings.len() { mappings.len() - 1 } else { index }]
    }.clone()
  }
}

macro_rules! assert_equal_mappings(
  ($a:expr, $b:expr) => (
    if $a != $b {
      panic!(format!("\n\n{:?}\n\n!=\n\n{:?}\n\n", $a, $b));
    }
  );
);

#[test]
fn test_source_map_issue_64() {
  let cache = consume(r#"{
    "version": 3,
    "file": "foo.js",
    "sourceRoot": "http://example.com/",
    "sources": ["/a"],
    "names": [],
    "mappings": "AACA",
    "sourcesContent": ["foo"]
  }"#).unwrap();

  let expected = Mapping {
    generated: CodePosition { line: 1, column: 0 },
    original: CodePosition { line: 2, column: 0 },
    source: "/a".into(),
    name: "".into()
  };
  let actual = cache.mapping_for_generated_position(1, 0);
  assert_equal_mappings!(actual, expected);
}

#[test]
fn test_source_map_issue_72_duplicate_sources() {
  let cache = consume(r#"{
    "version": 3,
    "file": "foo.js",
    "sources": ["source1.js", "source1.js", "source3.js"],
    "names": [],
    "mappings": ";EAAC;;IAEE;;MEEE",
    "sourceRoot": "http://example.com"
  }"#).unwrap();

  {
    let expected = Mapping {
      generated: CodePosition { line: 2, column: 2 },
      original: CodePosition { line: 1, column: 1 },
      source: "source1.js".into(),
      name: "".into()
    };
    let actual = cache.mapping_for_generated_position(2, 2);
    assert_equal_mappings!(actual, expected);
  }

  {
    let expected = Mapping {
      generated: CodePosition { line: 4, column: 4 },
      original: CodePosition { line: 3, column: 3 },
      source: "source1.js".into(),
      name: "".into()
    };
    let actual = cache.mapping_for_generated_position(4, 4);
    assert_equal_mappings!(actual, expected);
  }

  {
    let expected = Mapping {
      generated: CodePosition { line: 6, column: 6 },
      original: CodePosition { line: 5, column: 5 },
      source: "source3.js".into(),
      name: "".into()
    };
    let actual = cache.mapping_for_generated_position(6, 6);
    assert_equal_mappings!(actual, expected);
  }
}

#[test]
fn test_source_map_issue_72_duplicate_names() {
  let cache = consume(r#"{
    "version": 3,
    "file": "foo.js",
    "sources": ["source.js"],
    "names": ["name1", "name1", "name3"],
    "mappings": ";EAACA;;IAEEA;;MAEEE",
    "sourceRoot": "http://example.com"
  }"#).unwrap();

  {
    let expected = Mapping {
      generated: CodePosition { line: 2, column: 2 },
      original: CodePosition { line: 1, column: 1 },
      source: "source.js".into(),
      name: "name1".into()
    };
    let actual = cache.mapping_for_generated_position(2, 2);
    assert_equal_mappings!(actual, expected);
  }

  {
    let expected = Mapping {
      generated: CodePosition { line: 4, column: 4 },
      original: CodePosition { line: 3, column: 3 },
      source: "source.js".into(),
      name: "name1".into()
    };
    let actual = cache.mapping_for_generated_position(4, 4);
    assert_equal_mappings!(actual, expected);
  }

  {
    let expected = Mapping {
      generated: CodePosition { line: 6, column: 6 },
      original: CodePosition { line: 5, column: 5 },
      source: "source.js".into(),
      name: "name3".into()
    };
    let actual = cache.mapping_for_generated_position(6, 6);
    assert_equal_mappings!(actual, expected);
  }
}

#[test]
fn it_allows_omitting_source_root() {
  let cache_result: Result<Cache, String> = consume(r#"{
    "version": 3,
    "file": "foo.js",
    "sources": ["source.js"],
    "names": ["name1", "name1", "name3"],
    "mappings": ";EAACA;;IAEEA;;MAEEE"
  }"#);
  match cache_result {
    Ok(_) => {},
    Err(s) => panic!(format!("Error due to omitting: '{}'", s))
  }
}

#[test]
fn it_rejects_older_source_map_revisions() {
  let cache_result = consume(r#"{
    "version": 2,
    "file": "",
    "sources": ["source.js"],
    "names": ["name1", "name1", "name3"],
    "mappings": ";EAACA;;IAEEA;;MAEEE",
    "sourceRoot": "http://example.com"
  }"#);
  match cache_result {
    Ok(_) => panic!("Source Map revision < 3 should be rejected"),
    Err(_) => {}
  }
}

#[test]
fn it_does_not_panic_due_to_malformed_source_maps() {
  let cache_result = consume(r#"{
    "version": 3,
    "file": "",
    "sources": [],
    "names": [],
    "mappings": ";EAACA;;IAEEA;;MAEEE"
  }"#);
  match cache_result {
    Ok(_) => panic!("Invalid source maps should be rejected"),
    Err(_) => {}
  }
}

#[test]
fn it_returns_error_when_there_are_no_mappings() {
  let cache_result = consume(r#"{
    "version": 3,
    "file": "foo.js",
    "sources": ["source.js"],
    "names": ["name1", "name1", "name3"],
    "mappings": ";;;"
  }"#);
  match cache_result {
    Ok(_) => panic!("Source maps with no mappings should be rejected"),
    Err(_) => {}
  }
}

#[test]
fn it_does_not_panic_when_querying_for_position_2() {
  let cache = consume(r#"{
    "version": 3,
    "file": "foo.js",
    "sources": ["source.js"],
    "names": ["name1", "name1", "name3"],
    "mappings": "Z",
    "sourceRoot": "http://example.com"
  }"#).unwrap();
  cache.mapping_for_generated_position(2, 2);
}
<|endoftext|>


<commit_before>//! Shadowsocks Server Context

use std::{
    io,
    net::SocketAddr,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

#[cfg(feature = "local-dns-relay")]
use std::net::IpAddr;

use bloomfilter::Bloom;
#[cfg(feature = "local-dns-relay")]
use lru_time_cache::LruCache;
use spin::Mutex;
use tokio::runtime::Handle;
#[cfg(feature = "trust-dns")]
use trust_dns_resolver::TokioAsyncResolver;

#[cfg(feature = "trust-dns")]
use crate::relay::dns_resolver::create_resolver;
#[cfg(feature = "local-flow-stat")]
use crate::relay::flow::ServerFlowStatistic;
use crate::{
    config::{Config, ConfigType, ServerConfig},
    relay::{dns_resolver::resolve, socks5::Address},
};

// Entries for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_SERVER: usize = 1_000_000;

// Entries for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_CLIENT: usize = 10_000;

// Error rate for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_SERVER: f64 = 1e-6;

// Error rate for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_CLIENT: f64 = 1e-15;

// A bloom filter borrowed from shadowsocks-libev's `ppbloom`
//
// It contains 2 bloom filters and each one holds 1/2 entries.
// Use them as a ring buffer.
struct PingPongBloom {
    blooms: [Bloom<[u8]>; 2],
    bloom_count: [usize; 2],
    item_count: usize,
    current: usize,
}

impl PingPongBloom {
    fn new(ty: ConfigType) -> PingPongBloom {
        let (mut item_count, fp_p) = if ty.is_local() {
            (BF_NUM_ENTRIES_FOR_CLIENT, BF_ERROR_RATE_FOR_CLIENT)
        } else {
            (BF_NUM_ENTRIES_FOR_SERVER, BF_ERROR_RATE_FOR_SERVER)
        };

        item_count /= 2;

        PingPongBloom {
            blooms: [
                Bloom::new_for_fp_rate(item_count, fp_p),
                Bloom::new_for_fp_rate(item_count, fp_p),
            ],
            bloom_count: [0, 0],
            item_count,
            current: 0,
        }
    }

    // Check if data in `buf` exist.
    //
    // Set into the current bloom filter if not exist.
    //
    // Return `true` if data exist in bloom filter.
    fn check_and_set(&mut self, buf: &[u8]) -> bool {
        for bloom in &self.blooms {
            if bloom.check(buf) {
                return true;
            }
        }

        if self.bloom_count[self.current] >= self.item_count {
            // Current bloom filter is full,
            // Create a new one and use that one as current.

            self.current = (self.current + 1) % 2;

            self.bloom_count[self.current] = 0;
            self.blooms[self.current].clear();
        }

        // Cannot be optimized by `check_and_set`
        // Because we have to check every filters in `blooms` before `set`
        self.blooms[self.current].set(buf);
        self.bloom_count[self.current] += 1;

        false
    }
}

/// Server's global running status
///
/// Shared between UDP and TCP servers
pub struct ServerState {
    #[cfg(feature = "trust-dns")]
    dns_resolver: Option<TokioAsyncResolver>,
}

#[cfg(feature = "trust-dns")]
impl ServerState {
    /// Create a global shared server state
    pub async fn new_shared(config: &Config, rt: Handle) -> SharedServerState {
        let state = ServerState {
            dns_resolver: match create_resolver(config.get_dns_config(), config.timeout, config.ipv6_first, rt).await {
                Ok(resolver) => Some(resolver),
                Err(..) => None,
            },
        };

        Arc::new(state)
    }

    /// Get the global shared resolver
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.dns_resolver.as_ref()
    }
}

#[cfg(not(feature = "trust-dns"))]
impl ServerState {
    /// Create a global shared server state
    pub async fn new_shared(_config: &Config, _rt: Handle) -> SharedServerState {
        Arc::new(ServerState {})
    }
}

/// `ServerState` wrapped in `Arc`
pub type SharedServerState = Arc<ServerState>;

/// Shared basic configuration for the whole server
pub struct Context {
    config: Config,

    // Shared variables for all servers
    server_state: SharedServerState,

    // Server's running indicator
    // For killing all background jobs
    server_running: AtomicBool,

    // Check for duplicated IV/Nonce, for prevent replay attack
    // https://github.com/shadowsocks/shadowsocks-org/issues/44
    nonce_ppbloom: Mutex<PingPongBloom>,

    // For Android's flow stat report
    #[cfg(feature = "local-flow-stat")]
    local_flow_statistic: ServerFlowStatistic,

    // For DNS relay's ACL domain name reverse lookup -- whether the IP shall be forwarded
    #[cfg(feature = "local-dns-relay")]
    reverse_lookup_cache: Mutex<LruCache<IpAddr, bool>>,
}

/// Unique context thw whole server
pub type SharedContext = Arc<Context>;

impl Context {
    /// Create a non-shared Context
    fn new(config: Config, server_state: SharedServerState) -> Context {
        let nonce_ppbloom = Mutex::new(PingPongBloom::new(config.config_type));
        #[cfg(feature = "local-dns-relay")]
        let reverse_lookup_cache = Mutex::new(LruCache::<IpAddr, bool>::with_capacity(8192));

        Context {
            config,
            server_state,
            server_running: AtomicBool::new(true),
            nonce_ppbloom,
            #[cfg(feature = "local-flow-stat")]
            local_flow_statistic: ServerFlowStatistic::new(),
            #[cfg(feature = "local-dns-relay")]
            reverse_lookup_cache,
        }
    }

    /// Create a shared Context, wrapped in `Arc`
    pub fn new_shared(config: Config, server_state: SharedServerState) -> SharedContext {
        SharedContext::new(Context::new(config, server_state))
    }

    /// Config for TCP server
    pub fn config(&self) -> &Config {
        &self.config
    }

    /// Clone the internal ServerState
    pub fn clone_server_state(&self) -> SharedServerState {
        self.server_state.clone()
    }

    /// Mutable Config for TCP server
    ///
    /// NOTE: Only for launching plugins
    pub fn config_mut(&mut self) -> &mut Config {
        &mut self.config
    }

    /// Get ServerConfig by index
    pub fn server_config(&self, idx: usize) -> &ServerConfig {
        &self.config.server[idx]
    }

    /// Get mutable ServerConfig by index
    pub fn server_config_mut(&mut self, idx: usize) -> &mut ServerConfig {
        &mut self.config.server[idx]
    }

    #[cfg(feature = "trust-dns")]
    /// Get the global shared resolver
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.server_state.dns_resolver()
    }

    /// Perform a DNS resolution
    pub async fn dns_resolve(&self, host: &str, port: u16) -> io::Result<Vec<SocketAddr>> {
        resolve(self, host, port).await
    }

    /// Check if the server is still in running state
    pub fn server_running(&self) -> bool {
        self.server_running.load(Ordering::Acquire)
    }

    /// Stops the server, kills all detached running tasks
    pub fn set_server_stopped(&self) {
        self.server_running.store(false, Ordering::Release)
    }

    /// Check if nonce exist or not
    ///
    /// If not, set into the current bloom filter
    pub fn check_nonce_and_set(&self, nonce: &[u8]) -> bool {
        // Plain cipher doesn't have a nonce
        // Always treated as non-duplicated
        if nonce.is_empty() {
            return false;
        }

        let mut ppbloom = self.nonce_ppbloom.lock();
        ppbloom.check_and_set(nonce)
    }

    /// Check client ACL (for server)
    pub fn check_client_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_client_blocked(addr),
        }
    }

    /// Check outbound address ACL (for server)
    pub fn check_outbound_blocked(&self, addr: &Address) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_outbound_blocked(addr),
        }
    }

    /// Check resolved outbound address ACL (for server)
    pub fn check_resolved_outbound_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_resolved_outbound_blocked(addr),
        }
    }

    /// Add a record to the reverse lookup cache
    #[cfg(feature = "local-dns-relay")]
    pub fn add_to_reverse_lookup_cache(&self, addr: &IpAddr, forward: bool) {
        if self.check_ip_in_proxy_list(addr) == forward {
            return;
        }
        let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
        reverse_lookup_cache.insert(addr.clone(), forward);
    }

    /// Check if domain name is in proxy_list.
    /// If so, it should be resolved from remote (for Android's DNS relay)
    pub fn check_qname_in_proxy_list(&self, qname: &Address) -> Option<bool> {
        match self.config.acl {
            // Proxy everything by default
            None => None,
            Some(ref a) => a.check_qname_in_proxy_list(qname),
        }
    }

    #[cfg(feature = "local-dns-relay")]
    pub fn check_ip_in_proxy_list(&self, ip: &IpAddr) -> bool {
        match self.config.acl {
            // Proxy everything by default
            None => true,
            Some(ref a) => {
                // do the reverse lookup in our local cache
                let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
                // if a qname is found
                if let Some(forward) = reverse_lookup_cache.get(ip) {
                    *forward
                } else {
                    a.check_ip_in_proxy_list(ip)
                }
            }
        }
    }

    /// Check target address ACL (for client)
    pub async fn check_target_bypassed(&self, target: &Address) -> bool {
        match self.config.acl {
            // Proxy everything by default
            None => false,
            Some(ref a) => {
                #[cfg(feature = "local-dns-relay")]
                {
                    if let Address::SocketAddress(ref saddr) = target {
                        // do the reverse lookup in our local cache
                        let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
                        // if a qname is found
                        if let Some(forward) = reverse_lookup_cache.get(&saddr.ip()) {
                            return !*forward
                        }
                    }
                }
                a.check_target_bypassed(self, target).await
            }
        }
    }

    /// Get client flow statistics
    #[cfg(feature = "local-flow-stat")]
    pub fn local_flow_statistic(&self) -> &ServerFlowStatistic {
        &self.local_flow_statistic
    }
}
<commit_msg>Refine reverse_lookup_cache<commit_after>//! Shadowsocks Server Context

use std::{
    io,
    net::SocketAddr,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

#[cfg(feature = "local-dns-relay")]
use std::{
    net::IpAddr,
    time::Duration,
};

use bloomfilter::Bloom;
#[cfg(feature = "local-dns-relay")]
use lru_time_cache::LruCache;
use spin::Mutex;
use tokio::runtime::Handle;
#[cfg(feature = "trust-dns")]
use trust_dns_resolver::TokioAsyncResolver;

#[cfg(feature = "trust-dns")]
use crate::relay::dns_resolver::create_resolver;
#[cfg(feature = "local-flow-stat")]
use crate::relay::flow::ServerFlowStatistic;
use crate::{
    config::{Config, ConfigType, ServerConfig},
    relay::{dns_resolver::resolve, socks5::Address},
};

// Entries for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_SERVER: usize = 1_000_000;

// Entries for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_CLIENT: usize = 10_000;

// Error rate for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_SERVER: f64 = 1e-6;

// Error rate for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_CLIENT: f64 = 1e-15;

// A bloom filter borrowed from shadowsocks-libev's `ppbloom`
//
// It contains 2 bloom filters and each one holds 1/2 entries.
// Use them as a ring buffer.
struct PingPongBloom {
    blooms: [Bloom<[u8]>; 2],
    bloom_count: [usize; 2],
    item_count: usize,
    current: usize,
}

impl PingPongBloom {
    fn new(ty: ConfigType) -> PingPongBloom {
        let (mut item_count, fp_p) = if ty.is_local() {
            (BF_NUM_ENTRIES_FOR_CLIENT, BF_ERROR_RATE_FOR_CLIENT)
        } else {
            (BF_NUM_ENTRIES_FOR_SERVER, BF_ERROR_RATE_FOR_SERVER)
        };

        item_count /= 2;

        PingPongBloom {
            blooms: [
                Bloom::new_for_fp_rate(item_count, fp_p),
                Bloom::new_for_fp_rate(item_count, fp_p),
            ],
            bloom_count: [0, 0],
            item_count,
            current: 0,
        }
    }

    // Check if data in `buf` exist.
    //
    // Set into the current bloom filter if not exist.
    //
    // Return `true` if data exist in bloom filter.
    fn check_and_set(&mut self, buf: &[u8]) -> bool {
        for bloom in &self.blooms {
            if bloom.check(buf) {
                return true;
            }
        }

        if self.bloom_count[self.current] >= self.item_count {
            // Current bloom filter is full,
            // Create a new one and use that one as current.

            self.current = (self.current + 1) % 2;

            self.bloom_count[self.current] = 0;
            self.blooms[self.current].clear();
        }

        // Cannot be optimized by `check_and_set`
        // Because we have to check every filters in `blooms` before `set`
        self.blooms[self.current].set(buf);
        self.bloom_count[self.current] += 1;

        false
    }
}

/// Server's global running status
///
/// Shared between UDP and TCP servers
pub struct ServerState {
    #[cfg(feature = "trust-dns")]
    dns_resolver: Option<TokioAsyncResolver>,
}

#[cfg(feature = "trust-dns")]
impl ServerState {
    /// Create a global shared server state
    pub async fn new_shared(config: &Config, rt: Handle) -> SharedServerState {
        let state = ServerState {
            dns_resolver: match create_resolver(config.get_dns_config(), config.timeout, config.ipv6_first, rt).await {
                Ok(resolver) => Some(resolver),
                Err(..) => None,
            },
        };

        Arc::new(state)
    }

    /// Get the global shared resolver
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.dns_resolver.as_ref()
    }
}

#[cfg(not(feature = "trust-dns"))]
impl ServerState {
    /// Create a global shared server state
    pub async fn new_shared(_config: &Config, _rt: Handle) -> SharedServerState {
        Arc::new(ServerState {})
    }
}

/// `ServerState` wrapped in `Arc`
pub type SharedServerState = Arc<ServerState>;

/// Shared basic configuration for the whole server
pub struct Context {
    config: Config,

    // Shared variables for all servers
    server_state: SharedServerState,

    // Server's running indicator
    // For killing all background jobs
    server_running: AtomicBool,

    // Check for duplicated IV/Nonce, for prevent replay attack
    // https://github.com/shadowsocks/shadowsocks-org/issues/44
    nonce_ppbloom: Mutex<PingPongBloom>,

    // For Android's flow stat report
    #[cfg(feature = "local-flow-stat")]
    local_flow_statistic: ServerFlowStatistic,

    // For DNS relay's ACL domain name reverse lookup -- whether the IP shall be forwarded
    #[cfg(feature = "local-dns-relay")]
    reverse_lookup_cache: Mutex<LruCache<IpAddr, bool>>,
}

/// Unique context thw whole server
pub type SharedContext = Arc<Context>;

impl Context {
    /// Create a non-shared Context
    fn new(config: Config, server_state: SharedServerState) -> Context {
        let nonce_ppbloom = Mutex::new(PingPongBloom::new(config.config_type));
        #[cfg(feature = "local-dns-relay")]
        let reverse_lookup_cache = Mutex::new(LruCache::<IpAddr, bool>::with_expiry_duration(
            Duration::from_secs(3 * 24 * 60 * 60)));

        Context {
            config,
            server_state,
            server_running: AtomicBool::new(true),
            nonce_ppbloom,
            #[cfg(feature = "local-flow-stat")]
            local_flow_statistic: ServerFlowStatistic::new(),
            #[cfg(feature = "local-dns-relay")]
            reverse_lookup_cache,
        }
    }

    /// Create a shared Context, wrapped in `Arc`
    pub fn new_shared(config: Config, server_state: SharedServerState) -> SharedContext {
        SharedContext::new(Context::new(config, server_state))
    }

    /// Config for TCP server
    pub fn config(&self) -> &Config {
        &self.config
    }

    /// Clone the internal ServerState
    pub fn clone_server_state(&self) -> SharedServerState {
        self.server_state.clone()
    }

    /// Mutable Config for TCP server
    ///
    /// NOTE: Only for launching plugins
    pub fn config_mut(&mut self) -> &mut Config {
        &mut self.config
    }

    /// Get ServerConfig by index
    pub fn server_config(&self, idx: usize) -> &ServerConfig {
        &self.config.server[idx]
    }

    /// Get mutable ServerConfig by index
    pub fn server_config_mut(&mut self, idx: usize) -> &mut ServerConfig {
        &mut self.config.server[idx]
    }

    #[cfg(feature = "trust-dns")]
    /// Get the global shared resolver
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.server_state.dns_resolver()
    }

    /// Perform a DNS resolution
    pub async fn dns_resolve(&self, host: &str, port: u16) -> io::Result<Vec<SocketAddr>> {
        resolve(self, host, port).await
    }

    /// Check if the server is still in running state
    pub fn server_running(&self) -> bool {
        self.server_running.load(Ordering::Acquire)
    }

    /// Stops the server, kills all detached running tasks
    pub fn set_server_stopped(&self) {
        self.server_running.store(false, Ordering::Release)
    }

    /// Check if nonce exist or not
    ///
    /// If not, set into the current bloom filter
    pub fn check_nonce_and_set(&self, nonce: &[u8]) -> bool {
        // Plain cipher doesn't have a nonce
        // Always treated as non-duplicated
        if nonce.is_empty() {
            return false;
        }

        let mut ppbloom = self.nonce_ppbloom.lock();
        ppbloom.check_and_set(nonce)
    }

    /// Check client ACL (for server)
    pub fn check_client_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_client_blocked(addr),
        }
    }

    /// Check outbound address ACL (for server)
    pub fn check_outbound_blocked(&self, addr: &Address) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_outbound_blocked(addr),
        }
    }

    /// Check resolved outbound address ACL (for server)
    pub fn check_resolved_outbound_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_resolved_outbound_blocked(addr),
        }
    }

    /// Add a record to the reverse lookup cache
    #[cfg(feature = "local-dns-relay")]
    pub fn add_to_reverse_lookup_cache(&self, addr: &IpAddr, forward: bool) {
        let is_exception = self.check_ip_in_proxy_list(addr) != forward;
        let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
        match reverse_lookup_cache.get_mut(addr) {
            Some(value) => if is_exception {
                *value = forward;
            } else {
                // we do not need to remember the entry if it is already matched correctly
                reverse_lookup_cache.remove(addr);
            }
            None => if is_exception {
                reverse_lookup_cache.insert(addr.clone(), forward);
            }
        }
    }

    /// Check if domain name is in proxy_list.
    /// If so, it should be resolved from remote (for Android's DNS relay)
    pub fn check_qname_in_proxy_list(&self, qname: &Address) -> Option<bool> {
        match self.config.acl {
            // Proxy everything by default
            None => None,
            Some(ref a) => a.check_qname_in_proxy_list(qname),
        }
    }

    #[cfg(feature = "local-dns-relay")]
    pub fn check_ip_in_proxy_list(&self, ip: &IpAddr) -> bool {
        match self.config.acl {
            // Proxy everything by default
            None => true,
            Some(ref a) => {
                // do the reverse lookup in our local cache
                let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
                // if a qname is found
                if let Some(forward) = reverse_lookup_cache.get(ip) {
                    *forward
                } else {
                    a.check_ip_in_proxy_list(ip)
                }
            }
        }
    }

    /// Check target address ACL (for client)
    pub async fn check_target_bypassed(&self, target: &Address) -> bool {
        match self.config.acl {
            // Proxy everything by default
            None => false,
            Some(ref a) => {
                #[cfg(feature = "local-dns-relay")]
                {
                    if let Address::SocketAddress(ref saddr) = target {
                        // do the reverse lookup in our local cache
                        let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
                        // if a qname is found
                        if let Some(forward) = reverse_lookup_cache.get(&saddr.ip()) {
                            return !*forward
                        }
                    }
                }
                a.check_target_bypassed(self, target).await
            }
        }
    }

    /// Get client flow statistics
    #[cfg(feature = "local-flow-stat")]
    pub fn local_flow_statistic(&self) -> &ServerFlowStatistic {
        &self.local_flow_statistic
    }
}
<|endoftext|>


<commit_before>use std::io::{self, Write};
use {Color, Style};

/// Extension to the `Write` trait.
///
/// This extension to the `Write` trait is capable of producing the correct ANSI escape sequences
/// for given commands, effectively controlling the terminal.
pub trait TermWrite {

    /// Print the CSI (control sequence introducer) followed by a byte string.
    fn csi(&mut self, b: &[u8]) -> io::Result<usize>;
    /// Print OSC (operating system command) followed by a byte string.
    fn osc(&mut self, b: &[u8]) -> io::Result<usize>;
    /// Print DSC (device control string) followed by a byte string.
    fn dsc(&mut self, b: &[u8]) -> io::Result<usize>;

    /// Clear the entire screen.
    fn clear(&mut self) -> io::Result<usize> {
        self.csi(b"2J")
    }
    /// Clear everything _after_ the cursor.
    fn clear_after(&mut self) -> io::Result<usize> {
        self.csi(b"J")
    }
    /// Clear everything _before_ the cursor.
    fn clear_before(&mut self) -> io::Result<usize> {
        self.csi(b"1J")
    }
    /// Clear the current line.
    fn clear_line(&mut self) -> io::Result<usize> {
        self.csi(b"2K")
    }
    /// Clear from the cursor until newline.
    fn clear_until_newline(&mut self) -> io::Result<usize> {
        self.csi(b"K")
    }
    /// Show the cursor.
    fn show_cursor(&mut self) -> io::Result<usize> {
        self.csi(b"?25h")
    }
    /// Hide the cursor.
    fn hide_cursor(&mut self) -> io::Result<usize> {
        self.csi(b"?25l")
    }

    // TODO
    // fn mode

    /// Reset the rendition mode.
    ///
    /// This will reset both the current style and color.
    fn reset(&mut self) -> io::Result<usize> {
        self.csi(b"m")
    }
    /// Restore the defaults.
    ///
    /// This will reset color, position, cursor state, and so on. It is recommended that you use
    /// this before you exit your program, to avoid messing up the user's terminal.
    fn restore(&mut self) -> io::Result<usize> {
        Ok(try!(self.reset()) + try!(self.clear()) + try!(self.goto(0, 0)) + try!(self.show_cursor()))
    }

    /// Go to a given position.
    ///
    /// The position is 0-based.
    fn goto(&mut self, mut x: u16, mut y: u16) -> io::Result<usize> {
        x += 1;
        y += 1;

        self.csi(&[
            b'0' + (y / 10000) as u8,
            b'0' + (y / 1000) as u8 % 10,
            b'0' + (y / 100) as u8 % 10,
            b'0' + (y / 10) as u8 % 10,
            b'0' + y as u8 % 10,
            b';',
            b'0' + (x / 10000) as u8,
            b'0' + (x / 1000) as u8 % 10,
            b'0' + (x / 100) as u8 % 10,
            b'0' + (x / 10) as u8 % 10,
            b'0' + x as u8 % 10,
            b'H',
        ])
    }
    /// Set graphic rendition.
    fn rendition(&mut self, r: u8) -> io::Result<usize> {
        self.csi(&[
            b'0' + r / 100,
            b'0' + r / 10 % 10,
            b'0' + r % 10,
            b'm',
        ])
    }
    /// Set foreground color.
    fn color(&mut self, color: Color) -> io::Result<usize> {
        let ansi = color.to_ansi_val();
        self.csi(&[
            b'3',
            b'8',
            b';',
            b'5',
            b';',
            b'0' + ansi / 100,
            b'0' + ansi / 10 % 10,
            b'0' + ansi % 10,
            b'm',
        ])
    }
    /// Set background color.
    fn bg_color(&mut self, color: Color) -> io::Result<usize> {
        let ansi = color.to_ansi_val();
        self.csi(&[
            b'4',
            b'8',
            b';',
            b'5',
            b';',
            b'0' + ansi / 100,
            b'0' + ansi / 10 % 10,
            b'0' + ansi % 10,
            b'm',
        ])
    }
    /// Set rendition mode (SGR).
    fn style(&mut self, mode: Style) -> io::Result<usize> {
        self.rendition(mode as u8)
    }
}

impl<W: Write> TermWrite for W {
    fn csi(&mut self, b: &[u8]) -> io::Result<usize> {
        Ok(try!(self.write(b"\x1B[")) + try!(self.write(b)))
    }
    fn osc(&mut self, b: &[u8]) -> io::Result<usize> {
        Ok(try!(self.write(b"\x1B]")) + try!(self.write(b)))
    }
    fn dsc(&mut self, b: &[u8]) -> io::Result<usize> {
        Ok(try!(self.write(b"\x1BP")) + try!(self.write(b)))
    }
}

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

    #[test]
    fn test_csi() {
        let mut buf = Cursor::new(Vec::new());
        buf.csi(b"bluh").unwrap();

        assert_eq!(buf.get_ref(), b"\x1B[bluh");

        buf.csi(b"blah").unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[bluh\x1B[blah");
    }
    #[test]
    fn test_csi_partial() {
        let mut buf = [0; 3];
        let mut buf = &mut buf[..];
        assert_eq!(buf.csi(b"blu").unwrap(), 3);
        assert_eq!(buf.csi(b"").unwrap(), 0);
        assert_eq!(buf.csi(b"nooooo").unwrap(), 0);
    }
    #[test]
    fn test_osc() {
        let mut buf = Cursor::new(Vec::new());
        buf.osc(b"bluh").unwrap();

        assert_eq!(buf.get_ref(), b"\x1B]bluh");

        buf.osc(b"blah").unwrap();
        assert_eq!(buf.get_ref(), b"\x1B]bluh\x1B]blah");
    }
    #[test]
    fn test_osc_partial() {
        let mut buf = [0; 3];
        let mut buf = &mut buf[..];
        assert_eq!(buf.osc(b"blu").unwrap(), 3);
        assert_eq!(buf.osc(b"").unwrap(), 0);
        assert_eq!(buf.osc(b"nooooo").unwrap(), 0);
    }
    #[test]
    fn test_dsc() {
        let mut buf = Cursor::new(Vec::new());
        buf.dsc(b"bluh").unwrap();

        assert_eq!(buf.get_ref(), b"\x1BPbluh");

        buf.dsc(b"blah").unwrap();
        assert_eq!(buf.get_ref(), b"\x1BPbluh\x1BPblah");
    }
    #[test]
    fn test_dsc_partial() {
        let mut buf = [0; 3];
        let mut buf = &mut buf[..];
        assert_eq!(buf.dsc(b"blu").unwrap(), 3);
        assert_eq!(buf.dsc(b"").unwrap(), 0);
        assert_eq!(buf.dsc(b"nooooo").unwrap(), 0);
    }
    #[test]
    fn test_clear() {
        let mut buf = Cursor::new(Vec::new());
        buf.clear().unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[2J");
        buf.clear().unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[2J\x1B[2J");
    }
    #[test]
    fn test_goto() {
        let mut buf = Cursor::new(Vec::new());
        buf.goto(34, 43).unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[00044;00035H");
        buf.goto(24, 45).unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[00044;00035H\x1B[00046;00025H");
    }
    #[test]
    fn test_style() {
        use Style;

        let mut buf = Cursor::new(Vec::new());
        buf.style(Style::Bold).unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[001m");
        buf.style(Style::Italic).unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[001m\x1B[003m");
    }
}
<commit_msg>Use 8-color mode when possible.<commit_after>use std::io::{self, Write};
use {Color, Style};

/// Extension to the `Write` trait.
///
/// This extension to the `Write` trait is capable of producing the correct ANSI escape sequences
/// for given commands, effectively controlling the terminal.
pub trait TermWrite {

    /// Print the CSI (control sequence introducer) followed by a byte string.
    fn csi(&mut self, b: &[u8]) -> io::Result<usize>;
    /// Print OSC (operating system command) followed by a byte string.
    fn osc(&mut self, b: &[u8]) -> io::Result<usize>;
    /// Print DSC (device control string) followed by a byte string.
    fn dsc(&mut self, b: &[u8]) -> io::Result<usize>;

    /// Clear the entire screen.
    fn clear(&mut self) -> io::Result<usize> {
        self.csi(b"2J")
    }
    /// Clear everything _after_ the cursor.
    fn clear_after(&mut self) -> io::Result<usize> {
        self.csi(b"J")
    }
    /// Clear everything _before_ the cursor.
    fn clear_before(&mut self) -> io::Result<usize> {
        self.csi(b"1J")
    }
    /// Clear the current line.
    fn clear_line(&mut self) -> io::Result<usize> {
        self.csi(b"2K")
    }
    /// Clear from the cursor until newline.
    fn clear_until_newline(&mut self) -> io::Result<usize> {
        self.csi(b"K")
    }
    /// Show the cursor.
    fn show_cursor(&mut self) -> io::Result<usize> {
        self.csi(b"?25h")
    }
    /// Hide the cursor.
    fn hide_cursor(&mut self) -> io::Result<usize> {
        self.csi(b"?25l")
    }

    // TODO
    // fn mode

    /// Reset the rendition mode.
    ///
    /// This will reset both the current style and color.
    fn reset(&mut self) -> io::Result<usize> {
        self.csi(b"m")
    }
    /// Restore the defaults.
    ///
    /// This will reset color, position, cursor state, and so on. It is recommended that you use
    /// this before you exit your program, to avoid messing up the user's terminal.
    fn restore(&mut self) -> io::Result<usize> {
        Ok(try!(self.reset()) + try!(self.clear()) + try!(self.goto(0, 0)) + try!(self.show_cursor()))
    }

    /// Go to a given position.
    ///
    /// The position is 0-based.
    fn goto(&mut self, mut x: u16, mut y: u16) -> io::Result<usize> {
        x += 1;
        y += 1;

        self.csi(&[
            b'0' + (y / 10000) as u8,
            b'0' + (y / 1000) as u8 % 10,
            b'0' + (y / 100) as u8 % 10,
            b'0' + (y / 10) as u8 % 10,
            b'0' + y as u8 % 10,
            b';',
            b'0' + (x / 10000) as u8,
            b'0' + (x / 1000) as u8 % 10,
            b'0' + (x / 100) as u8 % 10,
            b'0' + (x / 10) as u8 % 10,
            b'0' + x as u8 % 10,
            b'H',
        ])
    }
    /// Set graphic rendition.
    fn rendition(&mut self, r: u8) -> io::Result<usize> {
        self.csi(&[
            b'0' + r / 100,
            b'0' + r / 10 % 10,
            b'0' + r % 10,
            b'm',
        ])
    }
    /// Set foreground color.
    fn color(&mut self, color: Color) -> io::Result<usize> {
        let ansi = color.to_ansi_val();
        if ansi <= 7 {
            self.csi(&[
                b'3',
                b'0' + ansi,
                b'm',
            ])
        } else {
            self.csi(&[
                b'3',
                b'8',
                b';',
                b'5',
                b';',
                b'0' + ansi / 100,
                b'0' + ansi / 10 % 10,
                b'0' + ansi % 10,
                b'm',
            ])
        }
    }
    /// Set background color.
    fn bg_color(&mut self, color: Color) -> io::Result<usize> {
        let ansi = color.to_ansi_val();

        if ansi <= 7 {
            self.csi(&[
                b'4',
                b'0' + ansi,
                b'm',
            ])
        } else {
            self.csi(&[
                b'4',
                b'8',
                b';',
                b'5',
                b';',
                b'0' + ansi / 100,
                b'0' + ansi / 10 % 10,
                b'0' + ansi % 10,
                b'm',
            ])
        }
    }
    /// Set rendition mode (SGR).
    fn style(&mut self, mode: Style) -> io::Result<usize> {
        self.rendition(mode as u8)
    }
}

impl<W: Write> TermWrite for W {
    fn csi(&mut self, b: &[u8]) -> io::Result<usize> {
        Ok(try!(self.write(b"\x1B[")) + try!(self.write(b)))
    }
    fn osc(&mut self, b: &[u8]) -> io::Result<usize> {
        Ok(try!(self.write(b"\x1B]")) + try!(self.write(b)))
    }
    fn dsc(&mut self, b: &[u8]) -> io::Result<usize> {
        Ok(try!(self.write(b"\x1BP")) + try!(self.write(b)))
    }
}

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

    #[test]
    fn test_csi() {
        let mut buf = Cursor::new(Vec::new());
        buf.csi(b"bluh").unwrap();

        assert_eq!(buf.get_ref(), b"\x1B[bluh");

        buf.csi(b"blah").unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[bluh\x1B[blah");
    }
    #[test]
    fn test_csi_partial() {
        let mut buf = [0; 3];
        let mut buf = &mut buf[..];
        assert_eq!(buf.csi(b"blu").unwrap(), 3);
        assert_eq!(buf.csi(b"").unwrap(), 0);
        assert_eq!(buf.csi(b"nooooo").unwrap(), 0);
    }
    #[test]
    fn test_osc() {
        let mut buf = Cursor::new(Vec::new());
        buf.osc(b"bluh").unwrap();

        assert_eq!(buf.get_ref(), b"\x1B]bluh");

        buf.osc(b"blah").unwrap();
        assert_eq!(buf.get_ref(), b"\x1B]bluh\x1B]blah");
    }
    #[test]
    fn test_osc_partial() {
        let mut buf = [0; 3];
        let mut buf = &mut buf[..];
        assert_eq!(buf.osc(b"blu").unwrap(), 3);
        assert_eq!(buf.osc(b"").unwrap(), 0);
        assert_eq!(buf.osc(b"nooooo").unwrap(), 0);
    }
    #[test]
    fn test_dsc() {
        let mut buf = Cursor::new(Vec::new());
        buf.dsc(b"bluh").unwrap();

        assert_eq!(buf.get_ref(), b"\x1BPbluh");

        buf.dsc(b"blah").unwrap();
        assert_eq!(buf.get_ref(), b"\x1BPbluh\x1BPblah");
    }
    #[test]
    fn test_dsc_partial() {
        let mut buf = [0; 3];
        let mut buf = &mut buf[..];
        assert_eq!(buf.dsc(b"blu").unwrap(), 3);
        assert_eq!(buf.dsc(b"").unwrap(), 0);
        assert_eq!(buf.dsc(b"nooooo").unwrap(), 0);
    }
    #[test]
    fn test_clear() {
        let mut buf = Cursor::new(Vec::new());
        buf.clear().unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[2J");
        buf.clear().unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[2J\x1B[2J");
    }
    #[test]
    fn test_goto() {
        let mut buf = Cursor::new(Vec::new());
        buf.goto(34, 43).unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[00044;00035H");
        buf.goto(24, 45).unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[00044;00035H\x1B[00046;00025H");
    }
    #[test]
    fn test_style() {
        use Style;

        let mut buf = Cursor::new(Vec::new());
        buf.style(Style::Bold).unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[001m");
        buf.style(Style::Italic).unwrap();
        assert_eq!(buf.get_ref(), b"\x1B[001m\x1B[003m");
    }
}
<|endoftext|>