<commit_before>#![allow(clippy::too_many_arguments)]
use futures_util::FutureExt;
use libc::c_void;
use std::{
	io,
	ptr::null_mut,
	sync::{
		Arc,
		Mutex,
	},
	task::{
		Context,
		Poll,
	},
};

use crate::{
	cstr,
	dns_consts::{
		Class,
		Type,
	},
	error::Error,
	ffi,
	notify::{
		Notified,
		Notify,
	},
};

// More typesafe than raw "ffi", but still not quite done

struct ManagedService(ffi::DNSServiceRef);

unsafe impl Send for ManagedService {}
unsafe impl Sync for ManagedService {}

impl Drop for ManagedService {
	fn drop(&mut self) {
		unsafe {
			ffi::DNSServiceRefDeallocate(self.0);
		}
	}
}

/// Keeps the service alive
#[derive(Clone)]
pub struct ServiceHandle {
	managed_service: Arc<ManagedService>,
}

impl ServiceHandle {
	fn new(raw: ffi::DNSServiceRef) -> Self {
		Self {
			managed_service: Arc::new(ManagedService(raw)),
		}
	}

	fn as_raw(&self) -> ffi::DNSServiceRef {
		self.managed_service.0
	}
}

pub(crate) trait EventedService: Unpin {
	// in various places we need to drive the underlying (possibly shared)
	// state machine, which will set other readiness events we then check.
	//
	// this underlying state machine will never complete, so there is
	// no need to return a Poll<..> result; but we do need a context to
	// drive the underlying service.
	fn poll_service(&mut self, cx: &mut Context<'_>) -> io::Result<()>;
}

/// Many places can keep the service alive, but a single active user
pub(crate) struct OwnedService {
	handle: ServiceHandle,
	processing: crate::evented::ReadProcessor,
}

impl OwnedService {
	fn new(raw: ffi::DNSServiceRef) -> io::Result<Self> {
		let fd = unsafe { ffi::DNSServiceRefSockFD(raw) };
		let handle = ServiceHandle::new(raw);
		let processing = crate::evented::ReadProcessor::new(fd)?;
		Ok(Self { handle, processing })
	}

	pub(crate) fn share(self) -> SharedService {
		let bg_fail_notify = Notify::new();
		let bg_fail_notified = bg_fail_notify.notified();
		let inner = Arc::new(Mutex::new(SharedInner {
			handle: self.handle,
			bg_error_buf: None,
			bg_failed: false,
			bg_fail_notify,
		}));
		let bg_inner = inner.clone();
		let mut processing = self.processing;

		let bg_task = futures_util::future::poll_fn(move |cx| {
			let mut inner = bg_inner.lock().unwrap();
			let raw = inner.handle.as_raw();
			let r = processing.process(cx, || {
				Error::from(unsafe { ffi::DNSServiceProcessResult(raw) })?;
				Ok(())
			});
			match r {
				Ok(()) => Poll::Pending, // run "forever"
				Err(e) => {
					inner.bg_error_buf = Some(e);
					inner.bg_failed = true;
					inner.bg_fail_notify.notify_waiters();
					Poll::Ready(()) // stop on errors
				},
			}
		});
		SharedService {
			inner,
			bg_task_handle: Arc::new(AbortHandle(tokio::spawn(bg_task))),
			bg_fail_notified,
		}
	}

	// -----

	pub(crate) fn enumerate_domains(
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		callback: ffi::DNSServiceDomainEnumReply,
		context: *mut c_void,
	) -> Result<Self, Error> {
		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceEnumerateDomains(&mut sd_ref, flags, interface_index, callback, context)
		})?;
		Ok(Self::new(sd_ref)?)
	}

	pub(crate) fn register(
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		name: &cstr::NullableCStr<'_>,
		reg_type: &cstr::CStr<'_>,
		domain: &cstr::NullableCStr<'_>,
		host: &cstr::NullableCStr<'_>,
		port: u16,
		txt: &[u8],
		callback: ffi::DNSServiceRegisterReply,
		context: *mut c_void,
	) -> Result<Self, Error> {
		let txt_len = txt.len();
		assert!(txt_len < (1 << 16));
		let txt_len = txt_len as u16;
		let txt_record = txt.as_ptr();

		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceRegister(
				&mut sd_ref,
				flags,
				interface_index,
				name.as_ptr(),
				reg_type.as_ptr(),
				domain.as_ptr(),
				host.as_ptr(),
				port,
				txt_len,
				txt_record,
				callback,
				context,
			)
		})?;
		Ok(Self::new(sd_ref)?)
	}

	pub(crate) fn browse(
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		reg_type: &cstr::CStr<'_>,
		domain: &cstr::NullableCStr<'_>,
		callback: ffi::DNSServiceBrowseReply,
		context: *mut c_void,
	) -> Result<Self, Error> {
		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceBrowse(
				&mut sd_ref,
				flags,
				interface_index,
				reg_type.as_ptr(),
				domain.as_ptr(),
				callback,
				context,
			)
		})?;
		Ok(Self::new(sd_ref)?)
	}

	pub(crate) fn resolve(
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		name: &cstr::CStr<'_>,
		reg_type: &cstr::CStr<'_>,
		domain: &cstr::CStr<'_>,
		callback: ffi::DNSServiceResolveReply,
		context: *mut c_void,
	) -> Result<Self, Error> {
		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceResolve(
				&mut sd_ref,
				flags,
				interface_index,
				name.as_ptr(),
				reg_type.as_ptr(),
				domain.as_ptr(),
				callback,
				context,
			)
		})?;
		Ok(Self::new(sd_ref)?)
	}

	pub(crate) fn query_record(
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		fullname: &cstr::CStr<'_>,
		rr_type: Type,
		rr_class: Class,
		callback: ffi::DNSServiceQueryRecordReply,
		context: *mut c_void,
	) -> Result<Self, Error> {
		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceQueryRecord(
				&mut sd_ref,
				flags,
				interface_index,
				fullname.as_ptr(),
				rr_type.0,
				rr_class.0,
				callback,
				context,
			)
		})?;
		Ok(Self::new(sd_ref)?)
	}
}

impl EventedService for OwnedService {
	fn poll_service(&mut self, cx: &mut Context<'_>) -> io::Result<()> {
		let raw = self.handle.as_raw();
		self.processing.process(cx, || {
			Error::from(unsafe { ffi::DNSServiceProcessResult(raw) })?;
			Ok(())
		})
	}
}

struct AbortHandle(tokio::task::JoinHandle<()>);

impl Drop for AbortHandle {
	fn drop(&mut self) {
		self.0.abort();
	}
}

struct SharedInner {
	// protect ffi calls
	handle: ServiceHandle,
	// forward error from background task
	bg_error_buf: Option<io::Error>,
	// but we can extract error only once, so remember it failed
	bg_failed: bool,
	//
	bg_fail_notify: Notify,
}

#[derive(Clone)]
pub(crate) struct SharedService {
	inner: Arc<Mutex<SharedInner>>,
	// make sure we kill the background task once all users are gone
	bg_task_handle: Arc<AbortHandle>,
	bg_fail_notified: Notified,
}

impl EventedService for SharedService {
	fn poll_service(&mut self, cx: &mut Context<'_>) -> io::Result<()> {
		// service is run in background task; just make sure there wasn't
		// an error yet and to get notified of future errors.
		let mut inner = self.inner.lock().unwrap();
		if let Some(e) = inner.bg_error_buf.take() {
			return Err(e);
		}
		if inner.bg_failed {
			return Err(io::Error::new(io::ErrorKind::NotConnected, "service gone"));
		}
		// should be pending, because we just checked for errors:
		let _ = self.bg_fail_notified.poll_unpin(cx);
		Ok(())
	}
}

impl SharedService {
	pub(crate) fn get_default_txt_record(self) -> DNSRecord {
		DNSRecord {
			service: self,
			raw: DNSRecordRef(null_mut()),
			rr_type: Type::TXT,
		}
	}

	// only valid when `service` was created through "register"
	//
	// still needs a SharedService to synchronize ffi calls
	pub(crate) fn add_record(
		self,
		flags: ffi::DNSServiceFlags,
		rr_type: Type,
		rdata: &[u8],
		ttl: u32,
	) -> Result<DNSRecord, Error> {
		let rd_len = rdata.len();
		assert!(rd_len < (1 << 16));
		let rd_len = rd_len as u16;
		let rdata = rdata.as_ptr();

		let inner = self.inner.lock().unwrap();

		let mut record_ref: ffi::DNSRecordRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceAddRecord(
				inner.handle.as_raw(),
				&mut record_ref,
				flags,
				rr_type.0,
				rd_len,
				rdata,
				ttl,
			)
		})?;

		drop(inner);

		Ok(DNSRecord {
			service: self,
			raw: DNSRecordRef(record_ref),
			rr_type,
		})
	}

	pub(crate) fn create_connection() -> Result<Self, Error> {
		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe { ffi::DNSServiceCreateConnection(&mut sd_ref) })?;
		Ok(OwnedService::new(sd_ref)?.share())
	}

	// only valid when `service` was created through "create_connection"
	pub(crate) fn register_record(
		self,
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		fullname: &cstr::CStr<'_>,
		rr_type: Type,
		rr_class: Class,
		rdata: &[u8],
		ttl: u32,
		callback: ffi::DNSServiceRegisterRecordReply,
		context: *mut c_void,
	) -> Result<DNSRecord, Error> {
		let rd_len = rdata.len();
		assert!(rd_len < (1 << 16));
		let rd_len = rd_len as u16;
		let rdata = rdata.as_ptr();

		let inner = self.inner.lock().unwrap();

		let mut record_ref: ffi::DNSRecordRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceRegisterRecord(
				inner.handle.as_raw(),
				&mut record_ref,
				flags,
				interface_index,
				fullname.as_ptr(),
				rr_type.0,
				rr_class.0,
				rd_len,
				rdata,
				ttl,
				callback,
				context,
			)
		})?;

		drop(inner);

		Ok(DNSRecord {
			service: self,
			raw: DNSRecordRef(record_ref),
			rr_type,
		})
	}
}

// so we don't have to unsafe impl for whole `DNSRecord`
//
// can only be used in combination with service handle, which is protected by mutex
struct DNSRecordRef(ffi::DNSRecordRef);

unsafe impl Sync for DNSRecordRef {}
unsafe impl Send for DNSRecordRef {}

pub(crate) struct DNSRecord {
	service: SharedService,
	raw: DNSRecordRef,
	rr_type: Type,
}

impl Drop for DNSRecord {
	fn drop(&mut self) {
		if !self.raw.0.is_null() {
			let inner = self.service.inner.lock().unwrap();
			unsafe {
				ffi::DNSServiceRemoveRecord(
					inner.handle.as_raw(),
					self.raw.0,
					0, // no flags
				);
			}
		}
	}
}

impl DNSRecord {
	pub(crate) fn update_record(
		&self,
		flags: ffi::DNSServiceFlags,
		rdata: &[u8],
		ttl: u32,
	) -> Result<(), Error> {
		let rd_len = rdata.len();
		assert!(rd_len < (1 << 16));
		let rd_len = rd_len as u16;
		let rdata = rdata.as_ptr();

		let inner = self.service.inner.lock().unwrap();

		Error::from(unsafe {
			ffi::DNSServiceUpdateRecord(
				inner.handle.as_raw(),
				self.raw.0,
				flags,
				rd_len,
				rdata,
				ttl,
			)
		})
	}

	pub(crate) fn rr_type(&self) -> Type {
		self.rr_type
	}

	// keep "forever" (until service is dropped)
	pub(crate) fn keep(mut self) {
		self.raw.0 = null_mut();
	}
}

pub fn reconfirm_record(
	flags: ffi::DNSServiceFlags,
	interface_index: u32,
	fullname: &cstr::CStr<'_>,
	rr_type: Type,
	rr_class: Class,
	rdata: &[u8],
) {
	let rd_len = rdata.len();
	assert!(rd_len < (1 << 16));
	let rd_len = rd_len as u16;
	let rdata = rdata.as_ptr();

	unsafe {
		ffi::DNSServiceReconfirmRecord(
			flags,
			interface_index,
			fullname.as_ptr(),
			rr_type.0,
			rr_class.0,
			rd_len,
			rdata,
		);
	}
}
<commit_msg>Silence build warning (field never read)<commit_after>#![allow(clippy::too_many_arguments)]
use futures_util::FutureExt;
use libc::c_void;
use std::{
	io,
	ptr::null_mut,
	sync::{
		Arc,
		Mutex,
	},
	task::{
		Context,
		Poll,
	},
};

use crate::{
	cstr,
	dns_consts::{
		Class,
		Type,
	},
	error::Error,
	ffi,
	notify::{
		Notified,
		Notify,
	},
};

// More typesafe than raw "ffi", but still not quite done

struct ManagedService(ffi::DNSServiceRef);

unsafe impl Send for ManagedService {}
unsafe impl Sync for ManagedService {}

impl Drop for ManagedService {
	fn drop(&mut self) {
		unsafe {
			ffi::DNSServiceRefDeallocate(self.0);
		}
	}
}

/// Keeps the service alive
#[derive(Clone)]
pub struct ServiceHandle {
	managed_service: Arc<ManagedService>,
}

impl ServiceHandle {
	fn new(raw: ffi::DNSServiceRef) -> Self {
		Self {
			managed_service: Arc::new(ManagedService(raw)),
		}
	}

	fn as_raw(&self) -> ffi::DNSServiceRef {
		self.managed_service.0
	}
}

pub(crate) trait EventedService: Unpin {
	// in various places we need to drive the underlying (possibly shared)
	// state machine, which will set other readiness events we then check.
	//
	// this underlying state machine will never complete, so there is
	// no need to return a Poll<..> result; but we do need a context to
	// drive the underlying service.
	fn poll_service(&mut self, cx: &mut Context<'_>) -> io::Result<()>;
}

/// Many places can keep the service alive, but a single active user
pub(crate) struct OwnedService {
	handle: ServiceHandle,
	processing: crate::evented::ReadProcessor,
}

impl OwnedService {
	fn new(raw: ffi::DNSServiceRef) -> io::Result<Self> {
		let fd = unsafe { ffi::DNSServiceRefSockFD(raw) };
		let handle = ServiceHandle::new(raw);
		let processing = crate::evented::ReadProcessor::new(fd)?;
		Ok(Self { handle, processing })
	}

	pub(crate) fn share(self) -> SharedService {
		let bg_fail_notify = Notify::new();
		let bg_fail_notified = bg_fail_notify.notified();
		let inner = Arc::new(Mutex::new(SharedInner {
			handle: self.handle,
			bg_error_buf: None,
			bg_failed: false,
			bg_fail_notify,
		}));
		let bg_inner = inner.clone();
		let mut processing = self.processing;

		let bg_task = futures_util::future::poll_fn(move |cx| {
			let mut inner = bg_inner.lock().unwrap();
			let raw = inner.handle.as_raw();
			let r = processing.process(cx, || {
				Error::from(unsafe { ffi::DNSServiceProcessResult(raw) })?;
				Ok(())
			});
			match r {
				Ok(()) => Poll::Pending, // run "forever"
				Err(e) => {
					inner.bg_error_buf = Some(e);
					inner.bg_failed = true;
					inner.bg_fail_notify.notify_waiters();
					Poll::Ready(()) // stop on errors
				},
			}
		});
		SharedService {
			inner,
			_bg_task_handle: Arc::new(AbortHandle(tokio::spawn(bg_task))),
			bg_fail_notified,
		}
	}

	// -----

	pub(crate) fn enumerate_domains(
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		callback: ffi::DNSServiceDomainEnumReply,
		context: *mut c_void,
	) -> Result<Self, Error> {
		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceEnumerateDomains(&mut sd_ref, flags, interface_index, callback, context)
		})?;
		Ok(Self::new(sd_ref)?)
	}

	pub(crate) fn register(
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		name: &cstr::NullableCStr<'_>,
		reg_type: &cstr::CStr<'_>,
		domain: &cstr::NullableCStr<'_>,
		host: &cstr::NullableCStr<'_>,
		port: u16,
		txt: &[u8],
		callback: ffi::DNSServiceRegisterReply,
		context: *mut c_void,
	) -> Result<Self, Error> {
		let txt_len = txt.len();
		assert!(txt_len < (1 << 16));
		let txt_len = txt_len as u16;
		let txt_record = txt.as_ptr();

		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceRegister(
				&mut sd_ref,
				flags,
				interface_index,
				name.as_ptr(),
				reg_type.as_ptr(),
				domain.as_ptr(),
				host.as_ptr(),
				port,
				txt_len,
				txt_record,
				callback,
				context,
			)
		})?;
		Ok(Self::new(sd_ref)?)
	}

	pub(crate) fn browse(
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		reg_type: &cstr::CStr<'_>,
		domain: &cstr::NullableCStr<'_>,
		callback: ffi::DNSServiceBrowseReply,
		context: *mut c_void,
	) -> Result<Self, Error> {
		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceBrowse(
				&mut sd_ref,
				flags,
				interface_index,
				reg_type.as_ptr(),
				domain.as_ptr(),
				callback,
				context,
			)
		})?;
		Ok(Self::new(sd_ref)?)
	}

	pub(crate) fn resolve(
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		name: &cstr::CStr<'_>,
		reg_type: &cstr::CStr<'_>,
		domain: &cstr::CStr<'_>,
		callback: ffi::DNSServiceResolveReply,
		context: *mut c_void,
	) -> Result<Self, Error> {
		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceResolve(
				&mut sd_ref,
				flags,
				interface_index,
				name.as_ptr(),
				reg_type.as_ptr(),
				domain.as_ptr(),
				callback,
				context,
			)
		})?;
		Ok(Self::new(sd_ref)?)
	}

	pub(crate) fn query_record(
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		fullname: &cstr::CStr<'_>,
		rr_type: Type,
		rr_class: Class,
		callback: ffi::DNSServiceQueryRecordReply,
		context: *mut c_void,
	) -> Result<Self, Error> {
		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceQueryRecord(
				&mut sd_ref,
				flags,
				interface_index,
				fullname.as_ptr(),
				rr_type.0,
				rr_class.0,
				callback,
				context,
			)
		})?;
		Ok(Self::new(sd_ref)?)
	}
}

impl EventedService for OwnedService {
	fn poll_service(&mut self, cx: &mut Context<'_>) -> io::Result<()> {
		let raw = self.handle.as_raw();
		self.processing.process(cx, || {
			Error::from(unsafe { ffi::DNSServiceProcessResult(raw) })?;
			Ok(())
		})
	}
}

struct AbortHandle(tokio::task::JoinHandle<()>);

impl Drop for AbortHandle {
	fn drop(&mut self) {
		self.0.abort();
	}
}

struct SharedInner {
	// protect ffi calls
	handle: ServiceHandle,
	// forward error from background task
	bg_error_buf: Option<io::Error>,
	// but we can extract error only once, so remember it failed
	bg_failed: bool,
	//
	bg_fail_notify: Notify,
}

#[derive(Clone)]
pub(crate) struct SharedService {
	inner: Arc<Mutex<SharedInner>>,
	// make sure we kill the background task once all users are gone
	_bg_task_handle: Arc<AbortHandle>,
	bg_fail_notified: Notified,
}

impl EventedService for SharedService {
	fn poll_service(&mut self, cx: &mut Context<'_>) -> io::Result<()> {
		// service is run in background task; just make sure there wasn't
		// an error yet and to get notified of future errors.
		let mut inner = self.inner.lock().unwrap();
		if let Some(e) = inner.bg_error_buf.take() {
			return Err(e);
		}
		if inner.bg_failed {
			return Err(io::Error::new(io::ErrorKind::NotConnected, "service gone"));
		}
		// should be pending, because we just checked for errors:
		let _ = self.bg_fail_notified.poll_unpin(cx);
		Ok(())
	}
}

impl SharedService {
	pub(crate) fn get_default_txt_record(self) -> DNSRecord {
		DNSRecord {
			service: self,
			raw: DNSRecordRef(null_mut()),
			rr_type: Type::TXT,
		}
	}

	// only valid when `service` was created through "register"
	//
	// still needs a SharedService to synchronize ffi calls
	pub(crate) fn add_record(
		self,
		flags: ffi::DNSServiceFlags,
		rr_type: Type,
		rdata: &[u8],
		ttl: u32,
	) -> Result<DNSRecord, Error> {
		let rd_len = rdata.len();
		assert!(rd_len < (1 << 16));
		let rd_len = rd_len as u16;
		let rdata = rdata.as_ptr();

		let inner = self.inner.lock().unwrap();

		let mut record_ref: ffi::DNSRecordRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceAddRecord(
				inner.handle.as_raw(),
				&mut record_ref,
				flags,
				rr_type.0,
				rd_len,
				rdata,
				ttl,
			)
		})?;

		drop(inner);

		Ok(DNSRecord {
			service: self,
			raw: DNSRecordRef(record_ref),
			rr_type,
		})
	}

	pub(crate) fn create_connection() -> Result<Self, Error> {
		let mut sd_ref: ffi::DNSServiceRef = null_mut();
		Error::from(unsafe { ffi::DNSServiceCreateConnection(&mut sd_ref) })?;
		Ok(OwnedService::new(sd_ref)?.share())
	}

	// only valid when `service` was created through "create_connection"
	pub(crate) fn register_record(
		self,
		flags: ffi::DNSServiceFlags,
		interface_index: u32,
		fullname: &cstr::CStr<'_>,
		rr_type: Type,
		rr_class: Class,
		rdata: &[u8],
		ttl: u32,
		callback: ffi::DNSServiceRegisterRecordReply,
		context: *mut c_void,
	) -> Result<DNSRecord, Error> {
		let rd_len = rdata.len();
		assert!(rd_len < (1 << 16));
		let rd_len = rd_len as u16;
		let rdata = rdata.as_ptr();

		let inner = self.inner.lock().unwrap();

		let mut record_ref: ffi::DNSRecordRef = null_mut();
		Error::from(unsafe {
			ffi::DNSServiceRegisterRecord(
				inner.handle.as_raw(),
				&mut record_ref,
				flags,
				interface_index,
				fullname.as_ptr(),
				rr_type.0,
				rr_class.0,
				rd_len,
				rdata,
				ttl,
				callback,
				context,
			)
		})?;

		drop(inner);

		Ok(DNSRecord {
			service: self,
			raw: DNSRecordRef(record_ref),
			rr_type,
		})
	}
}

// so we don't have to unsafe impl for whole `DNSRecord`
//
// can only be used in combination with service handle, which is protected by mutex
struct DNSRecordRef(ffi::DNSRecordRef);

unsafe impl Sync for DNSRecordRef {}
unsafe impl Send for DNSRecordRef {}

pub(crate) struct DNSRecord {
	service: SharedService,
	raw: DNSRecordRef,
	rr_type: Type,
}

impl Drop for DNSRecord {
	fn drop(&mut self) {
		if !self.raw.0.is_null() {
			let inner = self.service.inner.lock().unwrap();
			unsafe {
				ffi::DNSServiceRemoveRecord(
					inner.handle.as_raw(),
					self.raw.0,
					0, // no flags
				);
			}
		}
	}
}

impl DNSRecord {
	pub(crate) fn update_record(
		&self,
		flags: ffi::DNSServiceFlags,
		rdata: &[u8],
		ttl: u32,
	) -> Result<(), Error> {
		let rd_len = rdata.len();
		assert!(rd_len < (1 << 16));
		let rd_len = rd_len as u16;
		let rdata = rdata.as_ptr();

		let inner = self.service.inner.lock().unwrap();

		Error::from(unsafe {
			ffi::DNSServiceUpdateRecord(
				inner.handle.as_raw(),
				self.raw.0,
				flags,
				rd_len,
				rdata,
				ttl,
			)
		})
	}

	pub(crate) fn rr_type(&self) -> Type {
		self.rr_type
	}

	// keep "forever" (until service is dropped)
	pub(crate) fn keep(mut self) {
		self.raw.0 = null_mut();
	}
}

pub fn reconfirm_record(
	flags: ffi::DNSServiceFlags,
	interface_index: u32,
	fullname: &cstr::CStr<'_>,
	rr_type: Type,
	rr_class: Class,
	rdata: &[u8],
) {
	let rd_len = rdata.len();
	assert!(rd_len < (1 << 16));
	let rd_len = rd_len as u16;
	let rdata = rdata.as_ptr();

	unsafe {
		ffi::DNSServiceReconfirmRecord(
			flags,
			interface_index,
			fullname.as_ptr(),
			rr_type.0,
			rr_class.0,
			rd_len,
			rdata,
		);
	}
}
<|endoftext|>


<commit_before>use crate::rparser::*;
use crate::{gen_get_variants, Variant};

use ipsec_parser::*;

impl<'a> From<IkeTransformDHType> for Variant<'a> {
    fn from(input: IkeTransformDHType) -> Self {
        Variant::U16(input.0)
    }
}

pub struct IPsecBuilder {}
impl RBuilder for IPsecBuilder {
    fn build(&self) -> Box<dyn RParser> { Box::new(IPsecParser::new(b"IKEv2")) }
    fn probe(&self, i:&[u8]) -> bool { ipsec_probe(i) }
}

pub struct IPsecParser<'a> {
    _name: Option<&'a[u8]>,

    /// The transforms proposed by the initiator
    pub client_proposals : Vec<Vec<IkeV2Transform>>,

    /// The transforms selected by the responder
    pub server_proposals : Vec<Vec<IkeV2Transform>>,

    /// The Diffie-Hellman group from the server KE message, if present.
    pub dh_group: IkeTransformDHType,
}

impl<'a> RParser for IPsecParser<'a> {
    fn parse(&mut self, i: &[u8], direction: u8) -> u32 {
        match parse_ikev2_header(i) {
            Ok((rem,ref hdr)) => {
                debug!("parse_ikev2_header: {:?}",hdr);
                if rem.is_empty() && hdr.length == 28 {
                    return R_STATUS_OK;
                }
                // Rule 0: check version
                if hdr.maj_ver != 2 || hdr.min_ver != 0 {
                    warn!("Unknown header version: {}.{}", hdr.maj_ver, hdr.min_ver);
                }
                match parse_ikev2_payload_list(rem,hdr.next_payload) {
                    Ok((_,Ok(ref p))) => {
                        debug!("parse_ikev2_payload_with_type: {:?}",p);
                        for payload in p {
                            match payload.content {
                                IkeV2PayloadContent::SA(ref prop) => {
                                    // if hdr.flags & IKEV2_FLAG_INITIATOR != 0 {
                                        self.add_proposals(prop, direction);
                                    // }
                                },
                                IkeV2PayloadContent::KE(ref kex) => {
                                    debug!("KEX {:?}", kex.dh_group);
                                    if direction == STREAM_TOCLIENT {
                                        self.dh_group = kex.dh_group;
                                    }
                                },
                                IkeV2PayloadContent::Nonce(ref n) => {
                                    debug!("Nonce: {:?}", n);
                                },
                                IkeV2PayloadContent::Notify(ref n) => {
                                    debug!("Notify: {:?}", n);
                                },
                                _ => {
                                    debug!("Unknown payload content {:?}", payload.content);
                                },
                            }
                        }
                    },
                    e => warn!("parse_ikev2_payload_with_type: {:?}",e),
                };
            },
            e => warn!("parse_ikev2_header: {:?}",e),
        };
        R_STATUS_OK
    }

    gen_get_variants!{IPsecParser, "ikev2.",
        dh_group   => into,
        enc        => |s| { s.get_server_proposal_enc().map(|x| x.into()) },
        prf        => |s| { s.get_server_proposal_prf().map(|x| x.into()) },
        auth       => |s| { s.get_server_proposal_auth().map(|x| x.into()) },
        dh         => |s| { s.get_server_proposal_dh().map(|x| x.into()) },
        esn        => |s| { s.get_server_proposal_esn().map(|x| x.into()) },
    }
}

macro_rules! get_server_proposal {
    ( $t:tt, $n:ident ) => {
        pub fn $n(&self) -> Option<u16> {
            if let Some(xform) = self.server_proposals.first() {
                xform.iter()
                    .find_map(|x| if let IkeV2Transform::$t(e) = x {
                        Some(e.0)
                    } else {
                        None
                    })
            } else {
                None
            }
        }
    };
}

impl<'a> IPsecParser<'a> {
    pub fn new(name: &'a[u8]) -> IPsecParser<'a> {
        IPsecParser{
            _name: Some(name),
            client_proposals: Vec::new(),
            server_proposals: Vec::new(),
            dh_group: IkeTransformDHType::None,
        }
    }

    get_server_proposal!{Encryption, get_server_proposal_enc}
    get_server_proposal!{PRF, get_server_proposal_prf}
    get_server_proposal!{Auth, get_server_proposal_auth}
    get_server_proposal!{DH, get_server_proposal_dh}
    get_server_proposal!{ESN, get_server_proposal_esn}

    #[allow(clippy::cognitive_complexity)]
    fn add_proposals(&mut self, prop: &[IkeV2Proposal], direction: u8) {
        debug!("num_proposals: {}",prop.len());
        for p in prop {
            debug!("proposal: {:?}",p);
            debug!("num_transforms: {}",p.num_transforms);
            for xform in &p.transforms {
                debug!("transform: {:?}", xform);
                debug!("\ttype: {:?}", xform.transform_type);
                match xform.transform_type {
                    IkeTransformType::EncryptionAlgorithm => {
                        debug!("\tEncryptionAlgorithm: {:?}",IkeTransformEncType(xform.transform_id));
                    },
                    IkeTransformType::PseudoRandomFunction => {
                        debug!("\tPseudoRandomFunction: {:?}",IkeTransformPRFType(xform.transform_id));
                    },
                    IkeTransformType::IntegrityAlgorithm => {
                        debug!("\tIntegrityAlgorithm: {:?}",IkeTransformAuthType(xform.transform_id));
                    },
                    IkeTransformType::DiffieHellmanGroup => {
                        debug!("\tDiffieHellmanGroup: {:?}",IkeTransformDHType(xform.transform_id));
                    },
                    IkeTransformType::ExtendedSequenceNumbers => {
                        debug!("\tExtendedSequenceNumbers: {:?}",IkeTransformESNType(xform.transform_id));
                    },
                    _ => warn!("\tUnknown transform type {:?}", xform.transform_type),
                }
                if xform.transform_id == 0 {
                    warn!("\tTransform ID == 0 (choice left to responder)");
                };
            }
            let proposals : Vec<IkeV2Transform> = p.transforms.iter().map(|x| x.into()).collect();
            debug!("Proposals\n{:?}",proposals);
            // Rule 1: warn on weak or unknown transforms
            for prop in &proposals {
                match *prop {
                    IkeV2Transform::Encryption(ref enc) => {
                        match *enc {
                            IkeTransformEncType::ENCR_DES_IV64 |
                            IkeTransformEncType::ENCR_DES |
                            IkeTransformEncType::ENCR_3DES |
                            IkeTransformEncType::ENCR_RC5 |
                            IkeTransformEncType::ENCR_IDEA |
                            IkeTransformEncType::ENCR_CAST |
                            IkeTransformEncType::ENCR_BLOWFISH |
                            IkeTransformEncType::ENCR_3IDEA |
                            IkeTransformEncType::ENCR_DES_IV32 |
                            IkeTransformEncType::ENCR_NULL => {
                                warn!("Weak Encryption: {:?}", enc);
                            },
                            _ => (),
                        }
                    },
                    IkeV2Transform::Auth(ref auth) => {
                        match *auth {
                            IkeTransformAuthType::NONE => {
                                // Note: this could be expected with an AEAD encryption alg.
                                // See rule 4
                            },
                            IkeTransformAuthType::AUTH_HMAC_MD5_96 |
                            IkeTransformAuthType::AUTH_HMAC_SHA1_96 |
                            IkeTransformAuthType::AUTH_DES_MAC |
                            IkeTransformAuthType::AUTH_KPDK_MD5 |
                            IkeTransformAuthType::AUTH_AES_XCBC_96 |
                            IkeTransformAuthType::AUTH_HMAC_MD5_128 |
                            IkeTransformAuthType::AUTH_HMAC_SHA1_160 => {
                                warn!("Weak auth: {:?}", auth);
                            },
                            _ => (),
                        }
                    },
                    IkeV2Transform::DH(ref dh) => {
                        match *dh {
                            IkeTransformDHType::None => {
                                warn!("'None' DH transform proposed");
                            },
                            IkeTransformDHType::Modp768 |
                            IkeTransformDHType::Modp1024 |
                            IkeTransformDHType::Modp1024s160 |
                            IkeTransformDHType::Modp1536 => {
                                warn!("Weak DH: {:?}", dh);
                            },
                            _ => (),
                        }
                    },
                    IkeV2Transform::Unknown(tx_type,tx_id) => {
                        warn!("Unknown proposal: type={}, id={}", tx_type.0, tx_id);
                    },
                    _ => (),
                }
            }
            // Rule 2: check if no DH was proposed
            fn has_dh(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| {
                    if let IkeV2Transform::DH(_) = x { true } else { false }
                })
            }
            if !has_dh(&proposals)
            {
                warn!("No DH transform found");
            }
            // Rule 3: check if proposing AH ([RFC7296] section 3.3.1)
            if p.protocol_id == ProtocolID::AH {
                warn!("Proposal uses protocol AH - no confidentiality");
            }
            // Rule 4: lack of integrity is accepted only if using an AEAD proposal
            // Look if no auth was proposed, including if proposal is Auth::None
            fn has_auth(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| {
                    match *x {
                        IkeV2Transform::Auth(IkeTransformAuthType::NONE) => false,
                        IkeV2Transform::Auth(_)                          => true,
                        _                                                => false,
                    }
                })
            }
            fn has_gcm(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| {
                    if let IkeV2Transform::Encryption(ref enc) = x {
                        enc.is_aead()
                    } else { false }
                })
            }
            if !has_auth(&proposals) && !has_gcm(&proposals)
            {
                warn!("No integrity transform found");
            }
            // Rule 5: Check if an integrity and no integrity are part of the same proposal ?
            // XXX
            // Finally
            if direction == STREAM_TOSERVER {
                self.client_proposals.push(proposals);
            } else {
                self.server_proposals.push(proposals);
            }
        }
        // The server must accept one proposal, or reject them all (RFC 5996 2.7)
        if self.server_proposals.len() > 1 {
            warn!("more than one server proposals")
        }
        debug!("client_proposals: {:?}", self.client_proposals);
        debug!("server_proposals: {:?}", self.server_proposals);
    }
}

pub fn ipsec_probe(i: &[u8]) -> bool {
    if i.len() <= 20 { return false; }
    match parse_ikev2_header(i) {
        Ok((_,ref hdr)) => {
            if hdr.maj_ver != 2 || hdr.min_ver != 0 {
                debug!("ipsec_probe: could be ipsec, but with unsupported/invalid version {}.{}",
                      hdr.maj_ver, hdr.min_ver);
                return false;
            }
            if hdr.exch_type.0 < 34 || hdr.exch_type.0 > 37 {
                debug!("ipsec_probe: could be ipsec, but with unsupported/invalid exchange type {}",
                      hdr.exch_type.0);
                return false;
            }
            true
        }
        _ => false,
    }
}
<commit_msg>IKEv2: trust the initiator flag more than the direction to extract proposals<commit_after>use crate::rparser::*;
use crate::{gen_get_variants, Variant};

use ipsec_parser::*;

impl<'a> From<IkeTransformDHType> for Variant<'a> {
    fn from(input: IkeTransformDHType) -> Self {
        Variant::U16(input.0)
    }
}

pub struct IPsecBuilder {}
impl RBuilder for IPsecBuilder {
    fn build(&self) -> Box<dyn RParser> { Box::new(IPsecParser::new(b"IKEv2")) }
    fn probe(&self, i:&[u8]) -> bool { ipsec_probe(i) }
}

pub struct IPsecParser<'a> {
    _name: Option<&'a[u8]>,

    /// The transforms proposed by the initiator
    pub client_proposals : Vec<Vec<IkeV2Transform>>,

    /// The transforms selected by the responder
    pub server_proposals : Vec<Vec<IkeV2Transform>>,

    /// The Diffie-Hellman group from the server KE message, if present.
    pub dh_group: IkeTransformDHType,
}

impl<'a> RParser for IPsecParser<'a> {
    fn parse(&mut self, i: &[u8], direction: u8) -> u32 {
        match parse_ikev2_header(i) {
            Ok((rem,ref hdr)) => {
                debug!("parse_ikev2_header: {:?}",hdr);
                if rem.is_empty() && hdr.length == 28 {
                    return R_STATUS_OK;
                }
                // Rule 0: check version
                if hdr.maj_ver != 2 || hdr.min_ver != 0 {
                    warn!("Unknown header version: {}.{}", hdr.maj_ver, hdr.min_ver);
                }
                match parse_ikev2_payload_list(rem,hdr.next_payload) {
                    Ok((_,Ok(ref p))) => {
                        debug!("parse_ikev2_payload_with_type: {:?}",p);
                        for payload in p {
                            match payload.content {
                                IkeV2PayloadContent::SA(ref prop) => {
                                    let is_initiator = hdr.flags & IKEV2_FLAG_INITIATOR != 0;
                                    self.add_proposals(prop, is_initiator);
                                },
                                IkeV2PayloadContent::KE(ref kex) => {
                                    debug!("KEX {:?}", kex.dh_group);
                                    if direction == STREAM_TOCLIENT {
                                        self.dh_group = kex.dh_group;
                                    }
                                },
                                IkeV2PayloadContent::Nonce(ref n) => {
                                    debug!("Nonce: {:?}", n);
                                },
                                IkeV2PayloadContent::Notify(ref n) => {
                                    debug!("Notify: {:?}", n);
                                },
                                _ => {
                                    debug!("Unknown payload content {:?}", payload.content);
                                },
                            }
                        }
                    },
                    e => warn!("parse_ikev2_payload_with_type: {:?}",e),
                };
            },
            e => warn!("parse_ikev2_header: {:?}",e),
        };
        R_STATUS_OK
    }

    gen_get_variants!{IPsecParser, "ikev2.",
        dh_group   => into,
        enc        => |s| { s.get_server_proposal_enc().map(|x| x.into()) },
        prf        => |s| { s.get_server_proposal_prf().map(|x| x.into()) },
        auth       => |s| { s.get_server_proposal_auth().map(|x| x.into()) },
        dh         => |s| { s.get_server_proposal_dh().map(|x| x.into()) },
        esn        => |s| { s.get_server_proposal_esn().map(|x| x.into()) },
    }
}

macro_rules! get_server_proposal {
    ( $t:tt, $n:ident ) => {
        pub fn $n(&self) -> Option<u16> {
            if let Some(xform) = self.server_proposals.first() {
                xform.iter()
                    .find_map(|x| if let IkeV2Transform::$t(e) = x {
                        Some(e.0)
                    } else {
                        None
                    })
            } else {
                None
            }
        }
    };
}

impl<'a> IPsecParser<'a> {
    pub fn new(name: &'a[u8]) -> IPsecParser<'a> {
        IPsecParser{
            _name: Some(name),
            client_proposals: Vec::new(),
            server_proposals: Vec::new(),
            dh_group: IkeTransformDHType::None,
        }
    }

    get_server_proposal!{Encryption, get_server_proposal_enc}
    get_server_proposal!{PRF, get_server_proposal_prf}
    get_server_proposal!{Auth, get_server_proposal_auth}
    get_server_proposal!{DH, get_server_proposal_dh}
    get_server_proposal!{ESN, get_server_proposal_esn}

    #[allow(clippy::cognitive_complexity)]
    fn add_proposals(&mut self, prop: &[IkeV2Proposal], is_initiator: bool) {
        debug!("num_proposals: {}",prop.len());
        for p in prop {
            debug!("proposal: {:?}",p);
            debug!("num_transforms: {}",p.num_transforms);
            for xform in &p.transforms {
                debug!("transform: {:?}", xform);
                debug!("\ttype: {:?}", xform.transform_type);
                match xform.transform_type {
                    IkeTransformType::EncryptionAlgorithm => {
                        debug!("\tEncryptionAlgorithm: {:?}",IkeTransformEncType(xform.transform_id));
                    },
                    IkeTransformType::PseudoRandomFunction => {
                        debug!("\tPseudoRandomFunction: {:?}",IkeTransformPRFType(xform.transform_id));
                    },
                    IkeTransformType::IntegrityAlgorithm => {
                        debug!("\tIntegrityAlgorithm: {:?}",IkeTransformAuthType(xform.transform_id));
                    },
                    IkeTransformType::DiffieHellmanGroup => {
                        debug!("\tDiffieHellmanGroup: {:?}",IkeTransformDHType(xform.transform_id));
                    },
                    IkeTransformType::ExtendedSequenceNumbers => {
                        debug!("\tExtendedSequenceNumbers: {:?}",IkeTransformESNType(xform.transform_id));
                    },
                    _ => warn!("\tUnknown transform type {:?}", xform.transform_type),
                }
                if xform.transform_id == 0 {
                    warn!("\tTransform ID == 0 (choice left to responder)");
                };
            }
            let proposals : Vec<IkeV2Transform> = p.transforms.iter().map(|x| x.into()).collect();
            debug!("Proposals\n{:?}",proposals);
            // Rule 1: warn on weak or unknown transforms
            for prop in &proposals {
                match *prop {
                    IkeV2Transform::Encryption(ref enc) => {
                        match *enc {
                            IkeTransformEncType::ENCR_DES_IV64 |
                            IkeTransformEncType::ENCR_DES |
                            IkeTransformEncType::ENCR_3DES |
                            IkeTransformEncType::ENCR_RC5 |
                            IkeTransformEncType::ENCR_IDEA |
                            IkeTransformEncType::ENCR_CAST |
                            IkeTransformEncType::ENCR_BLOWFISH |
                            IkeTransformEncType::ENCR_3IDEA |
                            IkeTransformEncType::ENCR_DES_IV32 |
                            IkeTransformEncType::ENCR_NULL => {
                                warn!("Weak Encryption: {:?}", enc);
                            },
                            _ => (),
                        }
                    },
                    IkeV2Transform::Auth(ref auth) => {
                        match *auth {
                            IkeTransformAuthType::NONE => {
                                // Note: this could be expected with an AEAD encryption alg.
                                // See rule 4
                            },
                            IkeTransformAuthType::AUTH_HMAC_MD5_96 |
                            IkeTransformAuthType::AUTH_HMAC_SHA1_96 |
                            IkeTransformAuthType::AUTH_DES_MAC |
                            IkeTransformAuthType::AUTH_KPDK_MD5 |
                            IkeTransformAuthType::AUTH_AES_XCBC_96 |
                            IkeTransformAuthType::AUTH_HMAC_MD5_128 |
                            IkeTransformAuthType::AUTH_HMAC_SHA1_160 => {
                                warn!("Weak auth: {:?}", auth);
                            },
                            _ => (),
                        }
                    },
                    IkeV2Transform::DH(ref dh) => {
                        match *dh {
                            IkeTransformDHType::None => {
                                warn!("'None' DH transform proposed");
                            },
                            IkeTransformDHType::Modp768 |
                            IkeTransformDHType::Modp1024 |
                            IkeTransformDHType::Modp1024s160 |
                            IkeTransformDHType::Modp1536 => {
                                warn!("Weak DH: {:?}", dh);
                            },
                            _ => (),
                        }
                    },
                    IkeV2Transform::Unknown(tx_type,tx_id) => {
                        warn!("Unknown proposal: type={}, id={}", tx_type.0, tx_id);
                    },
                    _ => (),
                }
            }
            // Rule 2: check if no DH was proposed
            fn has_dh(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| {
                    if let IkeV2Transform::DH(_) = x { true } else { false }
                })
            }
            if !has_dh(&proposals)
            {
                warn!("No DH transform found");
            }
            // Rule 3: check if proposing AH ([RFC7296] section 3.3.1)
            if p.protocol_id == ProtocolID::AH {
                warn!("Proposal uses protocol AH - no confidentiality");
            }
            // Rule 4: lack of integrity is accepted only if using an AEAD proposal
            // Look if no auth was proposed, including if proposal is Auth::None
            fn has_auth(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| {
                    match *x {
                        IkeV2Transform::Auth(IkeTransformAuthType::NONE) => false,
                        IkeV2Transform::Auth(_)                          => true,
                        _                                                => false,
                    }
                })
            }
            fn has_gcm(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| {
                    if let IkeV2Transform::Encryption(ref enc) = x {
                        enc.is_aead()
                    } else { false }
                })
            }
            if !has_auth(&proposals) && !has_gcm(&proposals)
            {
                warn!("No integrity transform found");
            }
            // Rule 5: Check if an integrity and no integrity are part of the same proposal ?
            // XXX
            // Finally
            if is_initiator {
                self.client_proposals.push(proposals);
            } else {
                self.server_proposals.push(proposals);
            }
        }
        // The server must accept one proposal, or reject them all (RFC 5996 2.7)
        if self.server_proposals.len() > 1 {
            warn!("more than one server proposals")
        }
        debug!("client_proposals: {:?}", self.client_proposals);
        debug!("server_proposals: {:?}", self.server_proposals);
    }
}

pub fn ipsec_probe(i: &[u8]) -> bool {
    if i.len() <= 20 { return false; }
    match parse_ikev2_header(i) {
        Ok((_,ref hdr)) => {
            if hdr.maj_ver != 2 || hdr.min_ver != 0 {
                debug!("ipsec_probe: could be ipsec, but with unsupported/invalid version {}.{}",
                      hdr.maj_ver, hdr.min_ver);
                return false;
            }
            if hdr.exch_type.0 < 34 || hdr.exch_type.0 > 37 {
                debug!("ipsec_probe: could be ipsec, but with unsupported/invalid exchange type {}",
                      hdr.exch_type.0);
                return false;
            }
            true
        }
        _ => false,
    }
}
<|endoftext|>


<commit_before>use std::path::PathBuf;
use std::fs::File;
use std::error::Error;
use std::fs;
use std::fmt;
use std::io::{Seek, Read, SeekFrom, copy};
use std::collections::HashMap;

#[derive(Debug)]
pub enum FsError {
    FileNotFound,
    InvalidDirectory,
    NoFileHandle,
    MalformedDataSequence
}
impl Error for FsError {
    fn description(&self) -> &str {
        match *self {
            FsError::FileNotFound => "the folder does not exist or cannot be read from",
            FsError::InvalidDirectory => "the specified directory is not a valid directory",
            FsError::NoFileHandle => "the filesystem did not load a file yet",
            FsError::MalformedDataSequence => "the data sequence did not complete correctly"
        }
    }
}
impl fmt::Display for FsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FsError::FileNotFound => write!(f, "the folder specified could not be found or read from"),
            FsError::InvalidDirectory => write!(f, "the specified directory is not a valid directory"),
            FsError::NoFileHandle => write!(f, "the filesystem did not load a file yet"),
            FsError::MalformedDataSequence => write!(f, "the data sequence did not complete correctly")
        }
    }
}

#[derive(Debug,Clone)]
pub enum CompressionType {
    /// The archive is not compressed and the raw data is the real data.
    None,
    /// The archive is compressed with the Bzip2 compression algorithm.
    Bzip2,
    /// The archive is compressed with (a slightly modified, headerless) Gzip codec.
    Gzip,
    /// The archive is compressed with a modified LZMA2 variant (without size field).
    Lzma
}

#[derive(Debug)]
pub struct FileSystem {
    path: PathBuf,
    mainfile: MainFile,
    indices: HashMap<u32, IndexFile>
}

#[derive(Debug)]
pub struct MainFile {
    file: Option<File>
}

#[derive(Debug)]
pub struct IndexFile {
    id: u32,
    file: File
}

#[derive(Debug,Clone)]
pub struct IndexEntry {
    index: u8,
    id: u32,
    size: u32,
    offset: u64
}

#[derive(Debug,Clone)]
pub struct EntryHeader {
    entry: IndexEntry,
    raw_size: u32,
    real_size: u32,
    compression: CompressionType
}

#[derive(Debug,Clone)]
pub struct BlockHeader {
    big: bool,
    entry_id: u32,
    index_id: u8,

    next_seq: i32,
    next_block: u32
}

impl CompressionType {
    /// Fetches the appropriate type of compression based on the header field
    /// in the archive header.
    pub fn from_code(code: u8) -> CompressionType {
        match code {
            1 => CompressionType::Bzip2,
            2 => CompressionType::Gzip,
            3 => CompressionType::Lzma,
            _ => CompressionType::None
        }
    }
}

impl BlockHeader {
    pub fn from_block(big: bool, data: [u8; 520]) -> BlockHeader {
        match big {
            true => {
                BlockHeader {
                    big: true,
                    entry_id: ((data[0] as u32) << 24) | ((data[1] as u32) << 16) | ((data[2] as u32) << 8) | (data[3] as u32),
                    next_seq: (((data[4] as u32) << 8) | (data[5] as u32)) as i32,
                    next_block: ((data[6] as u32) << 16) | ((data[7] as u32) << 8) | (data[8] as u32),
                    index_id: data[9] as u8
                }
            },
            false => {
                BlockHeader {
                    big: false,
                    entry_id: ((data[0] as u32) << 8) | (data[1] as u32),
                    next_seq: (((data[2] as u32) << 8) | (data[3] as u32)) as i32,
                    next_block: ((data[4] as u32) << 16) | ((data[5] as u32) << 8) | (data[6] as u32),
                    index_id: data[7] as u8
                }
            }
        }
    }
}

impl IndexEntry {
    pub fn index(&self) -> u8 {
        self.index
    }

    /// Gets the id of this block.
    pub fn id(&self) -> u32 {
        self.id
    }

    /// Gets the absolute size of the entry data, not counting the 8-10
    /// byte header in the blocks.
    pub fn size(&self) -> u32 {
        self.size
    }

    /// Gets the absolute offset of the very first 520-byte block of this
    /// entry in the main data file.
    pub fn offset(&self) -> u64 {
        self.offset
    }

    pub fn block(&self) -> u32 {
        (self.offset / 520u64) as u32
    }
}

impl IndexFile {
    pub fn last_entry(&self) -> u64 {
         self.file.metadata().unwrap().len() / 6u64
    }

    pub fn entry(&mut self, id: u32) -> Option<IndexEntry> {
        let ref mut file = self.file;
        let mut tmp: [u8; 6] = [0; 6];

        // Seek to the proper position and read into the temp buffer
        let seek_offset = id as u64 * 6u64;
        let res1 = file.seek(SeekFrom::Start(seek_offset));
        let res2 = file.read(&mut tmp);

        // Check if the seek and read operation succeeded
        if res1.is_err() || res2.is_err() {
            return None;
        }

        // Check if the operations returned the expected results
        if res1.unwrap() != seek_offset || res2.unwrap() != 6 {
            return None;
        }

        // Decode the size and offset from the temp buffer
        let size: u32 = ((tmp[0] as u32) << 16) | ((tmp[1] as u32) << 8) | (tmp[2] as u32);
        let offset: u64 = ((tmp[3] as u64) << 16) | ((tmp[4] as u64) << 8) | (tmp[5] as u64);

        Some(IndexEntry {index: self.id as u8, id: id, size: size, offset: offset * 520u64})
    }
}

impl FileSystem {
    pub fn new(string: &'static str) -> Result<FileSystem, FsError> {
        // Declare some nice variables!!!
        let path = PathBuf::from(string);
        let metadata = fs::metadata(&path);

        // Make sure the folder exists
        if !metadata.is_ok() {
            return Err(FsError::FileNotFound);
        }

        // Make sure it is a directory
        if !metadata.unwrap().is_dir() {
            return Err(FsError::InvalidDirectory);
        }

        // Create mainfile path
        let mut mainfile_path = PathBuf::from(string);
        mainfile_path.push("main_file_cache.dat2");

        // Find all valid index files
        let mut indices: HashMap<u32, IndexFile> = HashMap::new();
        let entries = fs::read_dir(&path).unwrap();
        for entry in entries {
            let e = entry.unwrap();
            let fname = e.file_name().into_string().unwrap();

            // Is this an index?
            if fname.starts_with("main_file_cache.idx") {
                // Parse the index id into an integer
                let idx = fname[19..].parse::<u32>().unwrap();

                // Add the index file to our map with indices
                indices.insert(idx, IndexFile {id: idx, file: File::open(e.path()).unwrap()});
            }
        }

        // Create the filesystem object and return it
        let file = File::open(mainfile_path).ok();
        let mainfile = MainFile{file: file};

        Ok(FileSystem {path: path, mainfile: mainfile, indices: indices})
    }

    /// Gets the mainfile, that is, the main_file_cache.dat2 entry in the folder
    /// that holds the actual binary data of the filesystem entries.
    pub fn mainfile(&mut self) -> &mut MainFile {
        &mut self.mainfile
    }

    /// Gets an index with a specific id if it exists. The index can only exist if the file exists
    /// on the file system.
    pub fn index(&mut self, index: u32) -> Option<&mut IndexFile> {
        self.indices.get_mut(&index)
    }
}

impl MainFile {
    /// Checks if the file exists.
    pub fn exists(&self) -> bool {
        self.file.is_some()
    }

    /// Gets the backing file, if existant. Returns a new instance with a fresh seek pointer.
    pub fn file(&mut self) -> Option<&mut File> {
        self.file.as_mut()
    }

    /// Calculates the number of data blocks in the mainfile (if existant). This is done by
    /// taking the file size and dividing that by 520 (rouding up), because each block
    /// takes up 520 bytes of data.
    pub fn num_blocks(&self) -> Option<u64> {
        match self.file {
            Some(ref x) => Some((x.metadata().unwrap().len() + 519u64) / 520u64),
            None => None
        }
    }

    /// Reads a block of data, specified by the block id. The data is read at 520 * block_id
    /// and is exactly 520 bytes big. It is not guaranteed all 520 bytes are occupied if the
    /// block is the last one, thus possible to be trimmed.
    pub fn read_block(&mut self, block: u32) -> Option<[u8; 520]> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut data: [u8; 520] = [0; 520];
        let file = self.file().unwrap();

        // Seek to the right position and read the data
        file.seek(SeekFrom::Start(block as u64 * 520u64));
        file.read(&mut data);

        return Some(data);
    }

    pub fn read_header(&mut self, entry: IndexEntry) -> Option<EntryHeader> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut hdr: [u8; 9] = [0; 9];
        let file = self.file().unwrap();

        // Seek to the right position and read the data, skipping the block header at start
        let block_header_len = if entry.id() > 0xFFFF { 10 } else { 8 };
        file.seek(SeekFrom::Start(entry.offset() + block_header_len));
        file.read(&mut hdr);

        // Parse the 9 bytes of important info
        let compression_type = hdr[0];
        let raw_size: u32 = ((hdr[1] as u32) << 24) | ((hdr[2] as u32) << 16) | ((hdr[3] as u32) << 8) | (hdr[4] as u32);
        let real_size: u32 = ((hdr[5] as u32) << 24) | ((hdr[6] as u32) << 16) | ((hdr[7] as u32) << 8) | (hdr[8] as u32);

        // Return the new entry header
        return Some(EntryHeader {
            entry: entry,
            raw_size: raw_size,
            real_size: real_size,
            compression: CompressionType::from_code(compression_type)
        });
    }

    pub fn read_entry(&mut self, entry: IndexEntry) -> Result<Vec<i8>, FsError> {
        // Do we have a valid file?
        if self.file.is_none() {
            return Err(FsError::NoFileHandle);
        }

        // Create a vec with what we assume is the size. If not, the vec will
        // perfectly resize itself, so it's only an estimation to help us speed up.
        let mut data: Vec<i8> = Vec::with_capacity(entry.size() as usize);

        let mut current_block = entry.block();
        let mut remaining = entry.size();
        let mut current_seq = 0; // We expect a next part to be '1'

        while remaining > 0 {
            let block_data = self.read_block(current_block).unwrap();
            let block_info = BlockHeader::from_block(entry.id() > 65535, block_data);

            let header_size = if block_info.big {10} else {8};
            let available_data = 520 - header_size;
            let consumable = if remaining > available_data {available_data} else {remaining};
            remaining -= consumable;

            // Do some checks to validate this block.
            if remaining > 0 {
                if block_info.index_id != entry.index() || block_info.next_seq != current_seq {
                    return Err(FsError::MalformedDataSequence);
                } else {
                    // TODO this is so inefficient I should probably feel bad. Really bad. Terribly bad.
                    if block_info.big {
                        for i in block_data[10..520].iter() {
                            data.push(*i as i8);
                        }
                    } else {
                        for i in block_data[8..520].iter() {
                            data.push(*i as i8);
                        }
                    }

                    current_block += 1;
                    current_seq += 1;
                }
            }
        }

        Ok(data)
    }

}
<commit_msg>Ensure errors don't go unnoticed<commit_after>use std::path::PathBuf;
use std::fs::File;
use std::error::Error;
use std::fs;
use std::fmt;
use std::io::{Seek, Read, SeekFrom};
use std::collections::HashMap;

#[derive(Debug)]
pub enum FsError {
    FileNotFound,
    InvalidDirectory,
    NoFileHandle,
    MalformedDataSequence
}
impl Error for FsError {
    fn description(&self) -> &str {
        match *self {
            FsError::FileNotFound => "the folder does not exist or cannot be read from",
            FsError::InvalidDirectory => "the specified directory is not a valid directory",
            FsError::NoFileHandle => "the filesystem did not load a file yet",
            FsError::MalformedDataSequence => "the data sequence did not complete correctly"
        }
    }
}
impl fmt::Display for FsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FsError::FileNotFound => write!(f, "the folder specified could not be found or read from"),
            FsError::InvalidDirectory => write!(f, "the specified directory is not a valid directory"),
            FsError::NoFileHandle => write!(f, "the filesystem did not load a file yet"),
            FsError::MalformedDataSequence => write!(f, "the data sequence did not complete correctly")
        }
    }
}

#[derive(Debug,Clone)]
pub enum CompressionType {
    /// The archive is not compressed and the raw data is the real data.
    None,
    /// The archive is compressed with the Bzip2 compression algorithm.
    Bzip2,
    /// The archive is compressed with (a slightly modified, headerless) Gzip codec.
    Gzip,
    /// The archive is compressed with a modified LZMA2 variant (without size field).
    Lzma
}

#[derive(Debug)]
pub struct FileSystem {
    path: PathBuf,
    mainfile: MainFile,
    indices: HashMap<u32, IndexFile>
}

#[derive(Debug)]
pub struct MainFile {
    file: Option<File>
}

#[derive(Debug)]
pub struct IndexFile {
    id: u32,
    file: File
}

#[derive(Debug,Clone)]
pub struct IndexEntry {
    index: u8,
    id: u32,
    size: u32,
    offset: u64
}

#[derive(Debug,Clone)]
pub struct EntryHeader {
    entry: IndexEntry,
    raw_size: u32,
    real_size: u32,
    compression: CompressionType
}

#[derive(Debug,Clone)]
pub struct BlockHeader {
    big: bool,
    entry_id: u32,
    index_id: u8,

    next_seq: i32,
    next_block: u32
}

impl CompressionType {
    /// Fetches the appropriate type of compression based on the header field
    /// in the archive header.
    pub fn from_code(code: u8) -> CompressionType {
        match code {
            1 => CompressionType::Bzip2,
            2 => CompressionType::Gzip,
            3 => CompressionType::Lzma,
            _ => CompressionType::None
        }
    }
}

impl BlockHeader {
    pub fn from_block(big: bool, data: [u8; 520]) -> BlockHeader {
        match big {
            true => {
                BlockHeader {
                    big: true,
                    entry_id: ((data[0] as u32) << 24) | ((data[1] as u32) << 16) | ((data[2] as u32) << 8) | (data[3] as u32),
                    next_seq: (((data[4] as u32) << 8) | (data[5] as u32)) as i32,
                    next_block: ((data[6] as u32) << 16) | ((data[7] as u32) << 8) | (data[8] as u32),
                    index_id: data[9] as u8
                }
            },
            false => {
                BlockHeader {
                    big: false,
                    entry_id: ((data[0] as u32) << 8) | (data[1] as u32),
                    next_seq: (((data[2] as u32) << 8) | (data[3] as u32)) as i32,
                    next_block: ((data[4] as u32) << 16) | ((data[5] as u32) << 8) | (data[6] as u32),
                    index_id: data[7] as u8
                }
            }
        }
    }
}

impl IndexEntry {
    pub fn index(&self) -> u8 {
        self.index
    }

    /// Gets the id of this block.
    pub fn id(&self) -> u32 {
        self.id
    }

    /// Gets the absolute size of the entry data, not counting the 8-10
    /// byte header in the blocks.
    pub fn size(&self) -> u32 {
        self.size
    }

    /// Gets the absolute offset of the very first 520-byte block of this
    /// entry in the main data file.
    pub fn offset(&self) -> u64 {
        self.offset
    }

    pub fn block(&self) -> u32 {
        (self.offset / 520u64) as u32
    }
}

impl IndexFile {
    pub fn last_entry(&self) -> u64 {
         self.file.metadata().unwrap().len() / 6u64
    }

    pub fn entry(&mut self, id: u32) -> Option<IndexEntry> {
        let ref mut file = self.file;
        let mut tmp: [u8; 6] = [0; 6];

        // Seek to the proper position and read into the temp buffer
        let seek_offset = id as u64 * 6u64;
        let res1 = file.seek(SeekFrom::Start(seek_offset));
        let res2 = file.read(&mut tmp);

        // Check if the seek and read operation succeeded
        if res1.is_err() || res2.is_err() {
            return None;
        }

        // Check if the operations returned the expected results
        if res1.unwrap() != seek_offset || res2.unwrap() != 6 {
            return None;
        }

        // Decode the size and offset from the temp buffer
        let size: u32 = ((tmp[0] as u32) << 16) | ((tmp[1] as u32) << 8) | (tmp[2] as u32);
        let offset: u64 = ((tmp[3] as u64) << 16) | ((tmp[4] as u64) << 8) | (tmp[5] as u64);

        Some(IndexEntry {index: self.id as u8, id: id, size: size, offset: offset * 520u64})
    }
}

impl FileSystem {
    pub fn new(string: &'static str) -> Result<FileSystem, FsError> {
        // Declare some nice variables!!!
        let path = PathBuf::from(string);
        let metadata = fs::metadata(&path);

        // Make sure the folder exists
        if !metadata.is_ok() {
            return Err(FsError::FileNotFound);
        }

        // Make sure it is a directory
        if !metadata.unwrap().is_dir() {
            return Err(FsError::InvalidDirectory);
        }

        // Create mainfile path
        let mut mainfile_path = PathBuf::from(string);
        mainfile_path.push("main_file_cache.dat2");

        // Find all valid index files
        let mut indices: HashMap<u32, IndexFile> = HashMap::new();
        let entries = fs::read_dir(&path).unwrap();
        for entry in entries {
            let e = entry.unwrap();
            let fname = e.file_name().into_string().unwrap();

            // Is this an index?
            if fname.starts_with("main_file_cache.idx") {
                // Parse the index id into an integer
                let idx = fname[19..].parse::<u32>().unwrap();

                // Add the index file to our map with indices
                indices.insert(idx, IndexFile {id: idx, file: File::open(e.path()).unwrap()});
            }
        }

        // Create the filesystem object and return it
        let file = File::open(mainfile_path).ok();
        let mainfile = MainFile{file: file};

        Ok(FileSystem {path: path, mainfile: mainfile, indices: indices})
    }

    /// Gets the mainfile, that is, the main_file_cache.dat2 entry in the folder
    /// that holds the actual binary data of the filesystem entries.
    pub fn mainfile(&mut self) -> &mut MainFile {
        &mut self.mainfile
    }

    /// Gets an index with a specific id if it exists. The index can only exist if the file exists
    /// on the file system.
    pub fn index(&mut self, index: u32) -> Option<&mut IndexFile> {
        self.indices.get_mut(&index)
    }
}

impl MainFile {
    /// Checks if the file exists.
    pub fn exists(&self) -> bool {
        self.file.is_some()
    }

    /// Gets the backing file, if existant. Returns a new instance with a fresh seek pointer.
    pub fn file(&mut self) -> Option<&mut File> {
        self.file.as_mut()
    }

    /// Calculates the number of data blocks in the mainfile (if existant). This is done by
    /// taking the file size and dividing that by 520 (rouding up), because each block
    /// takes up 520 bytes of data.
    pub fn num_blocks(&self) -> Option<u64> {
        match self.file {
            Some(ref x) => Some((x.metadata().unwrap().len() + 519u64) / 520u64),
            None => None
        }
    }

    /// Reads a block of data, specified by the block id. The data is read at 520 * block_id
    /// and is exactly 520 bytes big. It is not guaranteed all 520 bytes are occupied if the
    /// block is the last one, thus possible to be trimmed.
    pub fn read_block(&mut self, block: u32) -> Option<[u8; 520]> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut data: [u8; 520] = [0; 520];
        let file = self.file().unwrap();

        // Seek to the right position and read the data
        file.seek(SeekFrom::Start(block as u64 * 520u64)).unwrap();
        file.read(&mut data).unwrap();

        return Some(data);
    }

    pub fn read_header(&mut self, entry: IndexEntry) -> Option<EntryHeader> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut hdr: [u8; 9] = [0; 9];
        let file = self.file().unwrap();

        // Seek to the right position and read the data, skipping the block header at start
        let block_header_len = if entry.id() > 0xFFFF { 10 } else { 8 };
        file.seek(SeekFrom::Start(entry.offset() + block_header_len)).unwrap();
        file.read(&mut hdr).unwrap();

        // Parse the 9 bytes of important info
        let compression_type = hdr[0];
        let raw_size: u32 = ((hdr[1] as u32) << 24) | ((hdr[2] as u32) << 16) | ((hdr[3] as u32) << 8) | (hdr[4] as u32);
        let real_size: u32 = ((hdr[5] as u32) << 24) | ((hdr[6] as u32) << 16) | ((hdr[7] as u32) << 8) | (hdr[8] as u32);

        // Return the new entry header
        return Some(EntryHeader {
            entry: entry,
            raw_size: raw_size,
            real_size: real_size,
            compression: CompressionType::from_code(compression_type)
        });
    }

    pub fn read_entry(&mut self, entry: IndexEntry) -> Result<Vec<i8>, FsError> {
        // Do we have a valid file?
        if self.file.is_none() {
            return Err(FsError::NoFileHandle);
        }

        // Create a vec with what we assume is the size. If not, the vec will
        // perfectly resize itself, so it's only an estimation to help us speed up.
        let mut data: Vec<i8> = Vec::with_capacity(entry.size() as usize);

        let mut current_block = entry.block();
        let mut remaining = entry.size();
        let mut current_seq = 0; // We expect a next part to be '1'

        while remaining > 0 {
            let block_data = self.read_block(current_block).unwrap();
            let block_info = BlockHeader::from_block(entry.id() > 65535, block_data);

            let header_size = if block_info.big {10} else {8};
            let available_data = 520 - header_size;
            let consumable = if remaining > available_data {available_data} else {remaining};
            remaining -= consumable;

            // Do some checks to validate this block.
            if remaining > 0 {
                if block_info.index_id != entry.index() || block_info.next_seq != current_seq {
                    return Err(FsError::MalformedDataSequence);
                } else {
                    // TODO this is so inefficient I should probably feel bad. Really bad. Terribly bad.
                    if block_info.big {
                        for i in block_data[10..520].iter() {
                            data.push(*i as i8);
                        }
                    } else {
                        for i in block_data[8..520].iter() {
                            data.push(*i as i8);
                        }
                    }

                    current_block += 1;
                    current_seq += 1;
                }
            }
        }

        Ok(data)
    }

}
<|endoftext|>


<commit_before>use nom;
use nom::{alpha, alphanumeric, digit, multispace};
use std::str::from_utf8;

// Identifier
// IDENT := /[a-zA-Z_]\w*/g
named!(ident<&str, &str>, re_match!(r"[a-zA-Z_](?:\w|[\?\!])*"));

// Whitespace
named!(whitespace, chain!(many0!(multispace), || { &b""[..] }));
<commit_msg>Add tests<commit_after>use nom;
use nom::{IResult, alpha, alphanumeric, digit, multispace};

use std::str;
use std::str::{from_utf8, FromStr};

// Identifier
// IDENT := /[a-zA-Z_]\w*/g
named!(ident<&str, &str>, re_match!(r"[a-zA-Z_](?:\w|[?!])*"));

// Whitespace
named!(whitespace, chain!(many0!(multispace), || { &b""[..] }));

// Float literal e.g. `.123` `0.123`
named!(
    float<f32>,
    map_res!(
        map_res!(
            recognize!(
                alt!(
                    delimited!(digit, tag!("."), opt!(complete!(digit))) |
                    delimited!(opt!(digit), tag!("."), digit)
                )
            ),
            str::from_utf8
        ),
        FromStr::from_str
    )
);

#[test]
fn float_test() {
    assert_eq!(float(&b"13.37"[..]), IResult::Done(&b""[..], 13.37));
    assert_eq!(float(&b".1337"[..]), IResult::Done(&b""[..], 0.1337));
}

#[test]
fn ident_test() {
    assert_eq!(ident(&"test?"[..]), IResult::Done("", "test?"));
    assert_eq!(ident(&"test!"[..]), IResult::Done("", "test!"));
    assert_eq!(ident(&"test!?"[..]), IResult::Done("", "test!?"));
    assert_eq!(ident(&"_test!"[..]), IResult::Done("", "_test!"));
}<|endoftext|>


<commit_before>use std::io::{Read, Chars, CharsError};
use std::borrow::Cow;
use std::fmt;
use std::char;

use num::Num;

use error::{ParserError, ParserErrorKind, StreamError};

/// Token types
#[derive(PartialEq)]
pub enum Token {
    /// `(`
    OpenParen,
    /// `)`
    CloseParen,
    /// `#(`
    OpenVectorParen,
    /// `#vu8(`
    OpenBytesParen,
    /// `.`
    Dot,
    /// `'`
    Quote,
    /// `\``
    QuasiQuote,
    /// `,`
    Comma,
    /// `,@`
    UnquoteSplicing,
    Identifier(Cow<'static, str>),
    /// `#t`
    True,
    /// `#f`
    False,
    /// `#\<String>`
    Character(String),
    String(String),
    Numeric(String),
    /// End of character stream
    EOF
}

impl fmt::Debug for Token {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Token::OpenParen => write!(f, "OpenParen"),
            Token::OpenVectorParen => write!(f, "OpenVectorParen"),
            Token::OpenBytesParen => write!(f, "OpenBytesParen"),
            Token::CloseParen => write!(f, "CloseParen"),
            Token::Dot => write!(f, "Dot"),
            Token::Quote => write!(f, "Quote"),
            Token::QuasiQuote => write!(f, "QuasiQuote"),
            Token::Comma => write!(f, "Comma"),
            Token::UnquoteSplicing => write!(f, "UnquoteSplicing"),
            Token::Identifier(ref name) => write!(f, "Identifier({})", name),
            Token::True => write!(f, "#t"),
            Token::False => write!(f, "#f"),
            Token::Character(ref name) => write!(f, "#\\{}", name),
            Token::Numeric(ref rep) => rep.fmt(f),
            Token::String(ref rep) => write!(f, "{:?}", rep),
            Token::EOF => write!(f, "EOF"),
        }
    }
}

/// TokenWrapper provides positional information to each token
pub struct TokenWrapper {
    pub line: usize,
    pub column: usize,
    pub token: Token
}

fn wrap(line: usize, column: usize, t: Token) -> TokenWrapper {
    TokenWrapper {
        line: line,
        column: column,
        token: t
    }
}

fn is_whitespace(c: char) -> bool {
    match c {
        '\t' | '\n' | '\x0b' | '\x0c' | '\r' | ' ' => true,
        _ => false
    }
}

fn is_initial(c: char) -> bool {
    match c {
        'a'...'z' | 'A'...'Z' | '!' | '$' | '%' | '&' | '*' | '/' | ':' | '<' | '=' | '>' | '?' | '^' | '_' | '~' => true,
        _ => false
    }
}

fn is_subsequent(c: char) -> bool {
    if is_initial(c) {
        true
    } else {
        match c {
            '0'...'9' | '+' | '-' | '.' | '@' => true,
            _ => false
        }
    }
}

fn is_delim(c: char) -> bool {
    match c {
        '(' | ')' | '[' | ']' | '"' | ';' => true,
        _ => c.is_whitespace()
    }
}

/// Lexer transforms character stream into a token stream
pub struct Lexer<R> {
    line: usize,
    column: usize,
    stream: Chars<R>,
    lookahead_buf: Option<char>,
}

macro_rules! try_consume {
    ($this:ident) => (
        match $this.consume() {
            Some(Ok(c)) => c,
            Some(Err(e)) => return Err($this.make_error(ParserErrorKind::UnderlyingError(StreamError(e)))),
            None => return Err($this.make_error(ParserErrorKind::UnexpectedEOF))
        }
    )
}

impl <R: Read + Sized> Lexer<R> {
    /// Creates new Lexer from io::Buffer
    pub fn new(stream: R) -> Lexer<R> {
        Lexer {
            line: 1,
            column: 1,
            stream: stream.chars(),
            lookahead_buf: None,
        }
    }

    pub fn lex_token(&mut self) -> Result<TokenWrapper, ParserError> {
        try!(self.consume_whitespace());

        let line = self.line;
        let col = self.column;

        let token = try!(self.lex());

        try!(self.consume_whitespace());

        Ok(wrap(line, col, token))
    }

    /// return next token
    fn lex(&mut self) -> Result<Token, ParserError> {
        let c = match self.consume() {
            Some(Ok(c)) => c,
            None => return Ok(Token::EOF),
            Some(Err(e)) =>
                return Err(self.make_error(ParserErrorKind::UnderlyingError(StreamError(e))))
        };

        let end_of_token = try!(self.is_end_of_token());

        if is_initial(c) {
            let mut init = String::new();
            init.push(c);
            self.lex_ident(init).map(|s| Token::Identifier(Cow::Owned(s)))
        } else if c == '+' {
            if end_of_token {
                Ok(Token::Identifier(Cow::Borrowed("+")))
            } else {
                self.lex_numeric("+".to_string()).map(Token::Numeric)
            }
        } else if c == '-' {
            if end_of_token {
                Ok(Token::Identifier(Cow::Borrowed("-")))
            } else {
                match try!(self.lookahead()) {
                    Some('>') => self.lex_ident("-".to_string()).map(|s| Token::Identifier(Cow::Owned(s))),
                    Some(_) => self.lex_numeric("-".to_string()).map(Token::Numeric),
                    None => Ok(Token::Identifier(Cow::Borrowed("-")))
                }
            }
        } else if c == ',' {
            match try!(self.lookahead()) {
                Some('@') => {
                    self.consume();
                    Ok(Token::UnquoteSplicing)
                },
                _ => Ok(Token::Comma)
            }
        } else if c == '(' {
            Ok(Token::OpenParen)
        } else if c == ')' {
            Ok(Token::CloseParen)
        } else if c == '.' {
            if end_of_token {
                Ok(Token::Dot)
            } else {
                try!(self.expect('.'));
                try!(self.expect('.'));
                if try!(self.is_end_of_token()) {
                    Ok(Token::Identifier(Cow::Borrowed("...")))
                } else {
                    Err(ParserError {
                        line: self.line,
                        column: self.column,
                        kind: ParserErrorKind::InvalidCharacter(try_consume!(self))
                    })
                }
            }
        } else if c == '\'' {
            Ok(Token::Quote)
        } else if c == '`' {
            Ok(Token::QuasiQuote)
        } else if c == '#' {
            let c0 = try_consume!(self);
            match c0 {
                't' | 'T' => Ok(Token::True),
                'f' | 'F' => Ok(Token::False),
                'b' | 'B' | 'o' | 'O' | 'd' | 'D' | 'x' | 'X' | 'i' | 'I' | 'e' | 'E' => {
                    let s = format!("{}{}", c, c0);
                    self.lex_numeric(s).map(Token::Numeric)
                },
                'v' | 'u' => {
                    let rest_prefix = try!(self.read_while(|c| !is_delim(c)));
                    let delim = try_consume!(self);
                    let prefix = format!("{}{}{}", c0, rest_prefix, delim);
                    match prefix.as_ref() {
                        "vu8(" | "u8(" =>
                            Ok(Token::OpenBytesParen),
                        _ =>
                            Err(self.make_error(ParserErrorKind::InvalidToken(prefix)))
                    }
                },
                '\\' => self.lex_char().map(Token::Character),
                '(' => Ok(Token::OpenVectorParen),
                _ => Err(self.make_error(ParserErrorKind::InvalidCharacter(c)))
            }
        } else if c == '"' {
            self.lex_string().map(Token::String)
        } else if c.is_numeric() {
            let s = format!("{}", c);
            self.lex_numeric(s).map(Token::Numeric)
        } else {
            Err(self.make_error(ParserErrorKind::InvalidCharacter(c)))
        }
    }

    fn is_end_of_token(&mut self) -> Result<bool, ParserError> {
        match try!(self.lookahead()) {
            Some(c) => Ok(is_whitespace(c) || is_delim(c)),
            None => Ok(true)
        }
    }

    fn lex_ident(&mut self, initial: String) -> Result<String, ParserError> {
        let mut s = initial;
        let sub = try!(self.read_while(is_subsequent));
        s.push_str(sub.as_ref());
        return Ok(s);
    }

    fn lex_char(&mut self) -> Result<String, ParserError> {
        let c = try_consume!(self);

        let mut s = String::new();
        s.push(c);
        let sub = try!(self.read_while(|c| c.is_alphanumeric()));
        s.push_str(sub.as_ref());
        return Ok(s);
    }

    fn lex_string(&mut self) -> Result<String, ParserError> {
        let mut s = String::new();
        loop {
            match try_consume!(self) {
                '"' => return Ok(s),
                '\\' => match try_consume!(self) {
                    'a' => s.push('\x07'),
                    'b' => s.push('\x08'),
                    't' => s.push('\t'),
                    'n' => s.push('\n'),
                    'v' => s.push('\x0b'),
                    'f' => s.push('\x0c'),
                    'r' => s.push('\r'),
                    '"' => s.push('"'),
                    '\\' => s.push('\\'),
                    'x' => {
                        let mut hex_str = String::new();
                        loop {
                            match try_consume!(self) {
                                ';' => break,
                                c => hex_str.push(c)
                            }
                        }
                        if hex_str.len() == 0 {
                            return Err(self.make_error(ParserErrorKind::InvalidStringEscape(hex_str)))
                        }
                        let code = match Num::from_str_radix(hex_str.as_ref(), 16) {
                            Ok(n) => n,
                            Err(_) => return Err(self.make_error(ParserErrorKind::InvalidStringEscape(hex_str)))
                        };
                        match char::from_u32(code) {
                            Some(c) => s.push(c),
                            None => return Err(self.make_error(ParserErrorKind::InvalidUnicodeRange(code)))
                        }
                    },
                    c =>
                        if c.is_whitespace() {
                            try!(self.read_while(|c| c != '\n' && c.is_whitespace()));
                            if c != '\n' {
                                if let Some(Ok('\n')) = self.consume() {
                                    try!(self.read_while(|c| c != '\n' && c.is_whitespace()));
                                } else {
                                    return Err(self.make_error(ParserErrorKind::InvalidStringLiteral))
                                }
                            }
                        } else {
                            let mut s = String::new();
                            s.push(c);
                            return Err(self.make_error(ParserErrorKind::InvalidStringEscape(s)))
                        }
                },
                c => s.push(c),
            }
        }
    }

    fn lex_numeric(&mut self, init: String) -> Result<String, ParserError> {
        let mut s = init;
        let sub = try!(self.read_while(|c| !is_delim(c)));
        s.push_str(sub.as_ref());
        return Ok(s);
    }

    fn make_error(&self, kind: ParserErrorKind) -> ParserError {
        ParserError {
            line: self.line,
            column: self.column,
            kind: kind
        }
    }

    fn lookahead(&mut self) -> Result<Option<char>, CharsError> {
        match self.lookahead_buf {
            Some(c) => Ok(Some(c)),
            None => {
                match self.stream.next() {
                    Some(Ok(ch)) => {
                        self.lookahead_buf = Some(ch);
                        Ok(Some(ch))
                    },
                    Some(Err(e)) => {
                        Err(e)
                    },
                    None => {
                        Ok(None)
                    }
                }
            }
        }
    }

    fn expect(&mut self, expected: char) -> Result<(), ParserError> {
        let c = try_consume!(self);
        if c == expected {
            Ok(())
        } else {
            Err(ParserError {
                line: self.line,
                column: self.column,
                kind: ParserErrorKind::InvalidCharacter(c)
            })
        }
    }

    fn advance(&mut self, c: char) {
        if c == '\n' {
            self.line += 1;
            self.column = 1;
        } else {
            self.column += 1;
        }
    }

    fn read_while<F>(&mut self, f: F) -> Result<String, ParserError> where
        F: Fn(char) -> bool
    {
        let mut s = match self.lookahead_buf {
            None => String::new(),
            Some(c) => if f(c) {
                self.lookahead_buf = None;
                self.advance(c);
                let mut s = String::new();
                s.push(c);
                s
            } else {
                return Ok(String::new());
            }
        };

        loop {
            match self.stream.next() {
                Some(Ok(c)) => if f(c) {
                    self.advance(c);
                    s.push(c);
                } else {
                    self.lookahead_buf = Some(c);
                    return Ok(s);
                },
                Some(Err(e)) =>
                    return Err(self.make_error(ParserErrorKind::UnderlyingError(StreamError(e)))),
                None => return Ok(s)
            }
        }
    }

    fn consume(&mut self) -> Option<Result<char, CharsError>> {
        let c = match self.lookahead_buf {
            Some(c) => {
                self.lookahead_buf = None;
                Some(Ok(c))
            },
            None => self.stream.next()
        };

        if let Some(Ok(ch)) = c {
            self.advance(ch);
        }

        c
    }

    fn consume_whitespace(&mut self) -> Result<bool, ParserError> {
        let mut consumed = false;
        loop {
            let whitespace = try!(self.read_while(is_whitespace));
            consumed = consumed || whitespace.len() > 0;
            if let Some(';') = try!(self.lookahead()) {
                consumed = true;
                try!(self.read_while(|c| c != '\n'));
                if self.lookahead_buf.is_some() {
                    self.lookahead_buf = None;
                }
            } else {
                return Ok(consumed);
            }
        }
    }
}
<commit_msg>Simpler lexer for `...`<commit_after>use std::io::{Read, Chars, CharsError};
use std::borrow::Cow;
use std::fmt;
use std::char;

use num::Num;

use error::{ParserError, ParserErrorKind, StreamError};

/// Token types
#[derive(PartialEq)]
pub enum Token {
    /// `(`
    OpenParen,
    /// `)`
    CloseParen,
    /// `#(`
    OpenVectorParen,
    /// `#vu8(`
    OpenBytesParen,
    /// `.`
    Dot,
    /// `'`
    Quote,
    /// `\``
    QuasiQuote,
    /// `,`
    Comma,
    /// `,@`
    UnquoteSplicing,
    Identifier(Cow<'static, str>),
    /// `#t`
    True,
    /// `#f`
    False,
    /// `#\<String>`
    Character(String),
    String(String),
    Numeric(String),
    /// End of character stream
    EOF
}

impl fmt::Debug for Token {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Token::OpenParen => write!(f, "OpenParen"),
            Token::OpenVectorParen => write!(f, "OpenVectorParen"),
            Token::OpenBytesParen => write!(f, "OpenBytesParen"),
            Token::CloseParen => write!(f, "CloseParen"),
            Token::Dot => write!(f, "Dot"),
            Token::Quote => write!(f, "Quote"),
            Token::QuasiQuote => write!(f, "QuasiQuote"),
            Token::Comma => write!(f, "Comma"),
            Token::UnquoteSplicing => write!(f, "UnquoteSplicing"),
            Token::Identifier(ref name) => write!(f, "Identifier({})", name),
            Token::True => write!(f, "#t"),
            Token::False => write!(f, "#f"),
            Token::Character(ref name) => write!(f, "#\\{}", name),
            Token::Numeric(ref rep) => rep.fmt(f),
            Token::String(ref rep) => write!(f, "{:?}", rep),
            Token::EOF => write!(f, "EOF"),
        }
    }
}

/// TokenWrapper provides positional information to each token
pub struct TokenWrapper {
    pub line: usize,
    pub column: usize,
    pub token: Token
}

fn wrap(line: usize, column: usize, t: Token) -> TokenWrapper {
    TokenWrapper {
        line: line,
        column: column,
        token: t
    }
}

fn is_whitespace(c: char) -> bool {
    match c {
        '\t' | '\n' | '\x0b' | '\x0c' | '\r' | ' ' => true,
        _ => false
    }
}

fn is_initial(c: char) -> bool {
    match c {
        'a'...'z' | 'A'...'Z' | '!' | '$' | '%' | '&' | '*' | '/' | ':' | '<' | '=' | '>' | '?' | '^' | '_' | '~' => true,
        _ => false
    }
}

fn is_subsequent(c: char) -> bool {
    if is_initial(c) {
        true
    } else {
        match c {
            '0'...'9' | '+' | '-' | '.' | '@' => true,
            _ => false
        }
    }
}

fn is_delim(c: char) -> bool {
    match c {
        '(' | ')' | '[' | ']' | '"' | ';' => true,
        _ => c.is_whitespace()
    }
}

/// Lexer transforms character stream into a token stream
pub struct Lexer<R> {
    line: usize,
    column: usize,
    stream: Chars<R>,
    lookahead_buf: Option<char>,
}

macro_rules! try_consume {
    ($this:ident) => (
        match $this.consume() {
            Some(Ok(c)) => c,
            Some(Err(e)) => return Err($this.make_error(ParserErrorKind::UnderlyingError(StreamError(e)))),
            None => return Err($this.make_error(ParserErrorKind::UnexpectedEOF))
        }
    )
}

impl <R: Read + Sized> Lexer<R> {
    /// Creates new Lexer from io::Buffer
    pub fn new(stream: R) -> Lexer<R> {
        Lexer {
            line: 1,
            column: 1,
            stream: stream.chars(),
            lookahead_buf: None,
        }
    }

    pub fn lex_token(&mut self) -> Result<TokenWrapper, ParserError> {
        try!(self.consume_whitespace());

        let line = self.line;
        let col = self.column;

        let token = try!(self.lex());

        try!(self.consume_whitespace());

        Ok(wrap(line, col, token))
    }

    /// return next token
    fn lex(&mut self) -> Result<Token, ParserError> {
        let c = match self.consume() {
            Some(Ok(c)) => c,
            None => return Ok(Token::EOF),
            Some(Err(e)) =>
                return Err(self.make_error(ParserErrorKind::UnderlyingError(StreamError(e))))
        };

        let end_of_token = try!(self.is_end_of_token());

        if is_initial(c) {
            let mut init = String::new();
            init.push(c);
            self.lex_ident(init).map(|s| Token::Identifier(Cow::Owned(s)))
        } else if c == '+' {
            if end_of_token {
                Ok(Token::Identifier(Cow::Borrowed("+")))
            } else {
                self.lex_numeric("+".to_string()).map(Token::Numeric)
            }
        } else if c == '-' {
            if end_of_token {
                Ok(Token::Identifier(Cow::Borrowed("-")))
            } else {
                match try!(self.lookahead()) {
                    Some('>') => self.lex_ident("-".to_string()).map(|s| Token::Identifier(Cow::Owned(s))),
                    Some(_) => self.lex_numeric("-".to_string()).map(Token::Numeric),
                    None => Ok(Token::Identifier(Cow::Borrowed("-")))
                }
            }
        } else if c == ',' {
            match try!(self.lookahead()) {
                Some('@') => {
                    self.consume();
                    Ok(Token::UnquoteSplicing)
                },
                _ => Ok(Token::Comma)
            }
        } else if c == '(' {
            Ok(Token::OpenParen)
        } else if c == ')' {
            Ok(Token::CloseParen)
        } else if c == '.' {
            if end_of_token {
                Ok(Token::Dot)
            } else {
                let rest_prefix = try!(self.read_while(|c| !is_delim(c)));
                if rest_prefix == ".." {
                    Ok(Token::Identifier(Cow::Borrowed("...")))
                } else {
                    Err(self.make_error(ParserErrorKind::InvalidToken(format!("{}{}", c, rest_prefix))))
                }
            }
        } else if c == '\'' {
            Ok(Token::Quote)
        } else if c == '`' {
            Ok(Token::QuasiQuote)
        } else if c == '#' {
            let c0 = try_consume!(self);
            match c0 {
                't' | 'T' => Ok(Token::True),
                'f' | 'F' => Ok(Token::False),
                'b' | 'B' | 'o' | 'O' | 'd' | 'D' | 'x' | 'X' | 'i' | 'I' | 'e' | 'E' => {
                    let s = format!("{}{}", c, c0);
                    self.lex_numeric(s).map(Token::Numeric)
                },
                'v' | 'u' => {
                    let rest_prefix = try!(self.read_while(|c| !is_delim(c)));
                    let delim = try_consume!(self);
                    let prefix = format!("{}{}{}", c0, rest_prefix, delim);
                    match prefix.as_ref() {
                        "vu8(" | "u8(" =>
                            Ok(Token::OpenBytesParen),
                        _ =>
                            Err(self.make_error(ParserErrorKind::InvalidToken(prefix)))
                    }
                },
                '\\' => self.lex_char().map(Token::Character),
                '(' => Ok(Token::OpenVectorParen),
                _ => Err(self.make_error(ParserErrorKind::InvalidCharacter(c)))
            }
        } else if c == '"' {
            self.lex_string().map(Token::String)
        } else if c.is_numeric() {
            let s = format!("{}", c);
            self.lex_numeric(s).map(Token::Numeric)
        } else {
            Err(self.make_error(ParserErrorKind::InvalidCharacter(c)))
        }
    }

    fn is_end_of_token(&mut self) -> Result<bool, ParserError> {
        match try!(self.lookahead()) {
            Some(c) => Ok(is_whitespace(c) || is_delim(c)),
            None => Ok(true)
        }
    }

    fn lex_ident(&mut self, initial: String) -> Result<String, ParserError> {
        let mut s = initial;
        let sub = try!(self.read_while(is_subsequent));
        s.push_str(sub.as_ref());
        return Ok(s);
    }

    fn lex_char(&mut self) -> Result<String, ParserError> {
        let c = try_consume!(self);

        let mut s = String::new();
        s.push(c);
        let sub = try!(self.read_while(|c| c.is_alphanumeric()));
        s.push_str(sub.as_ref());
        return Ok(s);
    }

    fn lex_string(&mut self) -> Result<String, ParserError> {
        let mut s = String::new();
        loop {
            match try_consume!(self) {
                '"' => return Ok(s),
                '\\' => match try_consume!(self) {
                    'a' => s.push('\x07'),
                    'b' => s.push('\x08'),
                    't' => s.push('\t'),
                    'n' => s.push('\n'),
                    'v' => s.push('\x0b'),
                    'f' => s.push('\x0c'),
                    'r' => s.push('\r'),
                    '"' => s.push('"'),
                    '\\' => s.push('\\'),
                    'x' => {
                        let mut hex_str = String::new();
                        loop {
                            match try_consume!(self) {
                                ';' => break,
                                c => hex_str.push(c)
                            }
                        }
                        if hex_str.len() == 0 {
                            return Err(self.make_error(ParserErrorKind::InvalidStringEscape(hex_str)))
                        }
                        let code = match Num::from_str_radix(hex_str.as_ref(), 16) {
                            Ok(n) => n,
                            Err(_) => return Err(self.make_error(ParserErrorKind::InvalidStringEscape(hex_str)))
                        };
                        match char::from_u32(code) {
                            Some(c) => s.push(c),
                            None => return Err(self.make_error(ParserErrorKind::InvalidUnicodeRange(code)))
                        }
                    },
                    c =>
                        if c.is_whitespace() {
                            try!(self.read_while(|c| c != '\n' && c.is_whitespace()));
                            if c != '\n' {
                                if let Some(Ok('\n')) = self.consume() {
                                    try!(self.read_while(|c| c != '\n' && c.is_whitespace()));
                                } else {
                                    return Err(self.make_error(ParserErrorKind::InvalidStringLiteral))
                                }
                            }
                        } else {
                            let mut s = String::new();
                            s.push(c);
                            return Err(self.make_error(ParserErrorKind::InvalidStringEscape(s)))
                        }
                },
                c => s.push(c),
            }
        }
    }

    fn lex_numeric(&mut self, init: String) -> Result<String, ParserError> {
        let mut s = init;
        let sub = try!(self.read_while(|c| !is_delim(c)));
        s.push_str(sub.as_ref());
        return Ok(s);
    }

    fn make_error(&self, kind: ParserErrorKind) -> ParserError {
        ParserError {
            line: self.line,
            column: self.column,
            kind: kind
        }
    }

    fn lookahead(&mut self) -> Result<Option<char>, CharsError> {
        match self.lookahead_buf {
            Some(c) => Ok(Some(c)),
            None => {
                match self.stream.next() {
                    Some(Ok(ch)) => {
                        self.lookahead_buf = Some(ch);
                        Ok(Some(ch))
                    },
                    Some(Err(e)) => {
                        Err(e)
                    },
                    None => {
                        Ok(None)
                    }
                }
            }
        }
    }

    fn advance(&mut self, c: char) {
        if c == '\n' {
            self.line += 1;
            self.column = 1;
        } else {
            self.column += 1;
        }
    }

    fn read_while<F>(&mut self, f: F) -> Result<String, ParserError> where
        F: Fn(char) -> bool
    {
        let mut s = match self.lookahead_buf {
            None => String::new(),
            Some(c) => if f(c) {
                self.lookahead_buf = None;
                self.advance(c);
                let mut s = String::new();
                s.push(c);
                s
            } else {
                return Ok(String::new());
            }
        };

        loop {
            match self.stream.next() {
                Some(Ok(c)) => if f(c) {
                    self.advance(c);
                    s.push(c);
                } else {
                    self.lookahead_buf = Some(c);
                    return Ok(s);
                },
                Some(Err(e)) =>
                    return Err(self.make_error(ParserErrorKind::UnderlyingError(StreamError(e)))),
                None => return Ok(s)
            }
        }
    }

    fn consume(&mut self) -> Option<Result<char, CharsError>> {
        let c = match self.lookahead_buf {
            Some(c) => {
                self.lookahead_buf = None;
                Some(Ok(c))
            },
            None => self.stream.next()
        };

        if let Some(Ok(ch)) = c {
            self.advance(ch);
        }

        c
    }

    fn consume_whitespace(&mut self) -> Result<bool, ParserError> {
        let mut consumed = false;
        loop {
            let whitespace = try!(self.read_while(is_whitespace));
            consumed = consumed || whitespace.len() > 0;
            if let Some(';') = try!(self.lookahead()) {
                consumed = true;
                try!(self.read_while(|c| c != '\n'));
                if self.lookahead_buf.is_some() {
                    self.lookahead_buf = None;
                }
            } else {
                return Ok(consumed);
            }
        }
    }
}
<|endoftext|>