<commit_before>use std::{
    collections::BTreeMap,
    env,
    io::Write as _,
    path::Path,
    time::{Instant, SystemTime, UNIX_EPOCH},
};

use anyhow::{bail, format_err, Result};
use xshell::{cmd, mkdir_p, pushd, pushenv, read_file, rm_rf};

use crate::flags;

type Unit = String;

impl flags::Metrics {
    pub(crate) fn run(self) -> Result<()> {
        let mut metrics = Metrics::new()?;
        if !self.dry_run {
            rm_rf("./target/release")?;
        }
        if !Path::new("./target/rustc-perf").exists() {
            mkdir_p("./target/rustc-perf")?;
            cmd!("git clone https://github.com/rust-lang/rustc-perf.git ./target/rustc-perf")
                .run()?;
        }
        {
            let _d = pushd("./target/rustc-perf")?;
            cmd!("git reset --hard 1d9288b0da7febf2599917da1b57dc241a1af033").run()?;
        }

        let _env = pushenv("RA_METRICS", "1");

        metrics.measure_build()?;
        metrics.measure_analysis_stats_self()?;
        metrics.measure_analysis_stats("ripgrep")?;
        metrics.measure_analysis_stats("webrender")?;

        if !self.dry_run {
            let _d = pushd("target")?;
            let metrics_token = env::var("METRICS_TOKEN").unwrap();
            cmd!(
                "git clone --depth 1 https://{metrics_token}@github.com/rust-analyzer/metrics.git"
            )
            .run()?;
            let _d = pushd("metrics")?;

            let mut file = std::fs::OpenOptions::new().append(true).open("metrics.json")?;
            writeln!(file, "{}", metrics.json())?;
            cmd!("git add .").run()?;
            cmd!("git -c user.name=Bot -c user.email=dummy@example.com commit --message 📈")
                .run()?;
            cmd!("git push origin master").run()?;
        }
        eprintln!("{:#?}", metrics);
        Ok(())
    }
}

impl Metrics {
    fn measure_build(&mut self) -> Result<()> {
        eprintln!("\nMeasuring build");
        cmd!("cargo fetch").run()?;

        let time = Instant::now();
        cmd!("cargo build --release --package rust-analyzer --bin rust-analyzer").run()?;
        let time = time.elapsed();
        self.report("build", time.as_millis() as u64, "ms".into());
        Ok(())
    }
    fn measure_analysis_stats_self(&mut self) -> Result<()> {
        self.measure_analysis_stats_path("self", &".")
    }
    fn measure_analysis_stats(&mut self, bench: &str) -> Result<()> {
        self.measure_analysis_stats_path(
            bench,
            &format!("./target/rustc-perf/collector/benchmarks/{}", bench),
        )
    }
    fn measure_analysis_stats_path(&mut self, name: &str, path: &str) -> Result<()> {
        eprintln!("\nMeasuring analysis-stats/{}", name);
        let output =
            cmd!("./target/release/rust-analyzer analysis-stats --quiet --memory-usage {path}")
                .read()?;
        for (metric, value, unit) in parse_metrics(&output) {
            self.report(&format!("analysis-stats/{}/{}", name, metric), value, unit.into());
        }
        Ok(())
    }
}

fn parse_metrics(output: &str) -> Vec<(&str, u64, &str)> {
    output
        .lines()
        .filter_map(|it| {
            let entry = it.split(':').collect::<Vec<_>>();
            match entry.as_slice() {
                ["METRIC", name, value, unit] => Some((*name, value.parse().unwrap(), *unit)),
                _ => None,
            }
        })
        .collect()
}

#[derive(Debug)]
struct Metrics {
    host: Host,
    timestamp: SystemTime,
    revision: String,
    metrics: BTreeMap<String, (u64, Unit)>,
}

#[derive(Debug)]
struct Host {
    os: String,
    cpu: String,
    mem: String,
}

impl Metrics {
    fn new() -> Result<Metrics> {
        let host = Host::new()?;
        let timestamp = SystemTime::now();
        let revision = cmd!("git rev-parse HEAD").read()?;
        Ok(Metrics { host, timestamp, revision, metrics: BTreeMap::new() })
    }

    fn report(&mut self, name: &str, value: u64, unit: Unit) {
        self.metrics.insert(name.into(), (value, unit));
    }

    fn json(&self) -> String {
        let mut buf = String::new();
        self.to_json(write_json::object(&mut buf));
        buf
    }

    fn to_json(&self, mut obj: write_json::Object<'_>) {
        self.host.to_json(obj.object("host"));
        let timestamp = self.timestamp.duration_since(UNIX_EPOCH).unwrap();
        obj.number("timestamp", timestamp.as_secs() as f64);
        obj.string("revision", &self.revision);
        let mut metrics = obj.object("metrics");
        for (k, (value, unit)) in &self.metrics {
            metrics.array(k).number(*value as f64).string(unit);
        }
    }
}

impl Host {
    fn new() -> Result<Host> {
        if cfg!(not(target_os = "linux")) {
            bail!("can only collect metrics on Linux ");
        }

        let os = read_field("/etc/os-release", "PRETTY_NAME=")?.trim_matches('"').to_string();

        let cpu =
            read_field("/proc/cpuinfo", "model name")?.trim_start_matches(':').trim().to_string();

        let mem = read_field("/proc/meminfo", "MemTotal:")?;

        return Ok(Host { os, cpu, mem });

        fn read_field<'a>(path: &str, field: &str) -> Result<String> {
            let text = read_file(path)?;

            let line = text
                .lines()
                .find(|it| it.starts_with(field))
                .ok_or_else(|| format_err!("can't parse {}", path))?;
            Ok(line[field.len()..].trim().to_string())
        }
    }
    fn to_json(&self, mut obj: write_json::Object<'_>) {
        obj.string("os", &self.os).string("cpu", &self.cpu).string("mem", &self.mem);
    }
}
<commit_msg>Fix xflags fallout<commit_after>use std::{
    collections::BTreeMap,
    env,
    io::Write as _,
    path::Path,
    time::{Instant, SystemTime, UNIX_EPOCH},
};

use anyhow::{bail, format_err, Result};
use xshell::{cmd, mkdir_p, pushd, pushenv, read_file, rm_rf};

use crate::flags;

type Unit = String;

impl flags::Metrics {
    pub(crate) fn run(self) -> Result<()> {
        let mut metrics = Metrics::new()?;
        if !self.dry_run {
            rm_rf("./target/release")?;
        }
        if !Path::new("./target/rustc-perf").exists() {
            mkdir_p("./target/rustc-perf")?;
            cmd!("git clone https://github.com/rust-lang/rustc-perf.git ./target/rustc-perf")
                .run()?;
        }
        {
            let _d = pushd("./target/rustc-perf")?;
            cmd!("git reset --hard 1d9288b0da7febf2599917da1b57dc241a1af033").run()?;
        }

        let _env = pushenv("RA_METRICS", "1");

        metrics.measure_build()?;
        metrics.measure_analysis_stats_self()?;
        metrics.measure_analysis_stats("ripgrep")?;
        metrics.measure_analysis_stats("webrender")?;

        if !self.dry_run {
            let _d = pushd("target")?;
            let metrics_token = env::var("METRICS_TOKEN").unwrap();
            cmd!(
                "git clone --depth 1 https://{metrics_token}@github.com/rust-analyzer/metrics.git"
            )
            .run()?;
            let _d = pushd("metrics")?;

            let mut file = std::fs::OpenOptions::new().append(true).open("metrics.json")?;
            writeln!(file, "{}", metrics.json())?;
            cmd!("git add .").run()?;
            cmd!("git -c user.name=Bot -c user.email=dummy@example.com commit --message 📈")
                .run()?;
            cmd!("git push origin master").run()?;
        }
        eprintln!("{:#?}", metrics);
        Ok(())
    }
}

impl Metrics {
    fn measure_build(&mut self) -> Result<()> {
        eprintln!("\nMeasuring build");
        cmd!("cargo fetch").run()?;

        let time = Instant::now();
        cmd!("cargo build --release --package rust-analyzer --bin rust-analyzer").run()?;
        let time = time.elapsed();
        self.report("build", time.as_millis() as u64, "ms".into());
        Ok(())
    }
    fn measure_analysis_stats_self(&mut self) -> Result<()> {
        self.measure_analysis_stats_path("self", &".")
    }
    fn measure_analysis_stats(&mut self, bench: &str) -> Result<()> {
        self.measure_analysis_stats_path(
            bench,
            &format!("./target/rustc-perf/collector/benchmarks/{}", bench),
        )
    }
    fn measure_analysis_stats_path(&mut self, name: &str, path: &str) -> Result<()> {
        eprintln!("\nMeasuring analysis-stats/{}", name);
        let output =
            cmd!("./target/release/rust-analyzer --quiet analysis-stats --memory-usage {path}")
                .read()?;
        for (metric, value, unit) in parse_metrics(&output) {
            self.report(&format!("analysis-stats/{}/{}", name, metric), value, unit.into());
        }
        Ok(())
    }
}

fn parse_metrics(output: &str) -> Vec<(&str, u64, &str)> {
    output
        .lines()
        .filter_map(|it| {
            let entry = it.split(':').collect::<Vec<_>>();
            match entry.as_slice() {
                ["METRIC", name, value, unit] => Some((*name, value.parse().unwrap(), *unit)),
                _ => None,
            }
        })
        .collect()
}

#[derive(Debug)]
struct Metrics {
    host: Host,
    timestamp: SystemTime,
    revision: String,
    metrics: BTreeMap<String, (u64, Unit)>,
}

#[derive(Debug)]
struct Host {
    os: String,
    cpu: String,
    mem: String,
}

impl Metrics {
    fn new() -> Result<Metrics> {
        let host = Host::new()?;
        let timestamp = SystemTime::now();
        let revision = cmd!("git rev-parse HEAD").read()?;
        Ok(Metrics { host, timestamp, revision, metrics: BTreeMap::new() })
    }

    fn report(&mut self, name: &str, value: u64, unit: Unit) {
        self.metrics.insert(name.into(), (value, unit));
    }

    fn json(&self) -> String {
        let mut buf = String::new();
        self.to_json(write_json::object(&mut buf));
        buf
    }

    fn to_json(&self, mut obj: write_json::Object<'_>) {
        self.host.to_json(obj.object("host"));
        let timestamp = self.timestamp.duration_since(UNIX_EPOCH).unwrap();
        obj.number("timestamp", timestamp.as_secs() as f64);
        obj.string("revision", &self.revision);
        let mut metrics = obj.object("metrics");
        for (k, (value, unit)) in &self.metrics {
            metrics.array(k).number(*value as f64).string(unit);
        }
    }
}

impl Host {
    fn new() -> Result<Host> {
        if cfg!(not(target_os = "linux")) {
            bail!("can only collect metrics on Linux ");
        }

        let os = read_field("/etc/os-release", "PRETTY_NAME=")?.trim_matches('"').to_string();

        let cpu =
            read_field("/proc/cpuinfo", "model name")?.trim_start_matches(':').trim().to_string();

        let mem = read_field("/proc/meminfo", "MemTotal:")?;

        return Ok(Host { os, cpu, mem });

        fn read_field<'a>(path: &str, field: &str) -> Result<String> {
            let text = read_file(path)?;

            let line = text
                .lines()
                .find(|it| it.starts_with(field))
                .ok_or_else(|| format_err!("can't parse {}", path))?;
            Ok(line[field.len()..].trim().to_string())
        }
    }
    fn to_json(&self, mut obj: write_json::Object<'_>) {
        obj.string("os", &self.os).string("cpu", &self.cpu).string("mem", &self.mem);
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::fs::File;
use std::net::{Shutdown, SocketAddr, TcpListener, TcpStream};
use std::sync::Arc;
use std::time::Duration;
use std::{io, thread};

use crate::lmdb;

use crate::core::core;
use crate::core::core::hash::Hash;
use crate::core::global;
use crate::core::pow::Difficulty;
use crate::handshake::Handshake;
use crate::peer::Peer;
use crate::peers::Peers;
use crate::store::PeerStore;
use crate::types::{
	Capabilities, ChainAdapter, Error, NetAdapter, P2PConfig, ReasonForBan, TxHashSetRead,
};
use crate::util::{Mutex, StopState};
use chrono::prelude::{DateTime, Utc};

/// P2P server implementation, handling bootstrapping to find and connect to
/// peers, receiving connections from other peers and keep track of all of them.
pub struct Server {
	pub config: P2PConfig,
	capabilities: Capabilities,
	handshake: Arc<Handshake>,
	pub peers: Arc<Peers>,
	stop_state: Arc<Mutex<StopState>>,
}

// TODO TLS
impl Server {
	/// Creates a new idle p2p server with no peers
	pub fn new(
		db_env: Arc<lmdb::Environment>,
		capab: Capabilities,
		config: P2PConfig,
		adapter: Arc<dyn ChainAdapter>,
		genesis: Hash,
		stop_state: Arc<Mutex<StopState>>,
	) -> Result<Server, Error> {
		Ok(Server {
			config: config.clone(),
			capabilities: capab,
			handshake: Arc::new(Handshake::new(genesis, config.clone())),
			peers: Arc::new(Peers::new(PeerStore::new(db_env)?, adapter, config)),
			stop_state,
		})
	}

	/// Starts a new TCP server and listen to incoming connections. This is a
	/// blocking call until the TCP server stops.
	pub fn listen(&self) -> Result<(), Error> {
		// start TCP listener and handle incoming connections
		let addr = SocketAddr::new(self.config.host, self.config.port);
		let listener = TcpListener::bind(addr)?;
		listener.set_nonblocking(true)?;

		let mut connected_sockets: HashMap<SocketAddr, TcpStream> = HashMap::new();

		let sleep_time = Duration::from_millis(1);
		loop {
			// Pause peer ingress connection request. Only for tests.
			if self.stop_state.lock().is_paused() {
				thread::sleep(Duration::from_secs(1));
				continue;
			}

			match listener.accept() {
				Ok((stream, peer_addr)) => {
					if !self.check_banned(&stream) {
						let sc = stream.try_clone();
						if let Err(e) = self.handle_new_peer(stream) {
							warn!("Error accepting peer {}: {:?}", peer_addr.to_string(), e);
							let _ = self.peers.add_banned(peer_addr, ReasonForBan::BadHandshake);
						} else if let Ok(s) = sc {
							connected_sockets.insert(peer_addr, s);
						}
					}
					// if any active socket not in our peers list, close it
					self.clean_lost_sockets(&mut connected_sockets);
				}
				Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
					// nothing to do, will retry in next iteration
				}
				Err(e) => {
					warn!("Couldn't establish new client connection: {:?}", e);
				}
			}
			if self.stop_state.lock().is_stopped() {
				break;
			}
			thread::sleep(sleep_time);
		}
		Ok(())
	}

	/// Asks the server to connect to a new peer. Directly returns the peer if
	/// we're already connected to the provided address.
	pub fn connect(&self, addr: &SocketAddr) -> Result<Arc<Peer>, Error> {
		if Peer::is_denied(&self.config, &addr) {
			debug!("connect_peer: peer {} denied, not connecting.", addr);
			return Err(Error::ConnectionClose);
		}

		if global::is_production_mode() {
			let hs = self.handshake.clone();
			let addrs = hs.addrs.read();
			if addrs.contains(&addr) {
				debug!("connect: ignore connecting to PeerWithSelf, addr: {}", addr);
				return Err(Error::PeerWithSelf);
			}
		}

		if let Some(p) = self.peers.get_connected_peer(addr) {
			// if we're already connected to the addr, just return the peer
			trace!("connect_peer: already connected {}", addr);
			return Ok(p);
		}

		trace!(
			"connect_peer: on {}:{}. connecting to {}",
			self.config.host,
			self.config.port,
			addr
		);
		match TcpStream::connect_timeout(addr, Duration::from_secs(10)) {
			Ok(mut stream) => {
				let addr = SocketAddr::new(self.config.host, self.config.port);
				let total_diff = self.peers.total_difficulty();

				let mut peer = Peer::connect(
					&mut stream,
					self.capabilities,
					total_diff,
					addr,
					&self.handshake,
					self.peers.clone(),
				)?;
				peer.start(stream);
				let peer = Arc::new(peer);
				self.peers.add_connected(peer.clone())?;
				Ok(peer)
			}
			Err(e) => {
				trace!(
					"connect_peer: on {}:{}. Could not connect to {}: {:?}",
					self.config.host,
					self.config.port,
					addr,
					e
				);
				Err(Error::Connection(e))
			}
		}
	}

	fn handle_new_peer(&self, mut stream: TcpStream) -> Result<(), Error> {
		let total_diff = self.peers.total_difficulty();

		// accept the peer and add it to the server map
		let mut peer = Peer::accept(
			&mut stream,
			self.capabilities,
			total_diff,
			&self.handshake,
			self.peers.clone(),
		)?;
		peer.start(stream);
		self.peers.add_connected(Arc::new(peer))?;
		Ok(())
	}

	fn check_banned(&self, stream: &TcpStream) -> bool {
		// peer has been banned, go away!
		if let Ok(peer_addr) = stream.peer_addr() {
			if self.peers.is_banned(peer_addr) {
				debug!("Peer {} banned, refusing connection.", peer_addr);
				if let Err(e) = stream.shutdown(Shutdown::Both) {
					debug!("Error shutting down conn: {:?}", e);
				}
				return true;
			}
		}
		false
	}

	/// For all kinds of exception cases, the node could accepted / initiated a peer connection successfully but
	/// failed on the Handshake protocol communication, or a connected peer was closed but without a successful
	/// clean-up on its socket, that will cause this connected (on TcpStream) peer becomes so-called "invisible" peer!
	/// i.e. a peer not included in the 'self.peers.peers' hashmap. This "invisible" peer will cause some security
	/// concern because it still can send something to this node, but without enough visibility as other connected peers.
	/// Another impact is these connections could never be closed, which make the node fully occupied by all such
	/// kind of connections and become un-connectable.
	/// This function can help to clean the peer connections which is "invisible" for this node.
	fn clean_lost_sockets(&self, sockets: &mut HashMap<SocketAddr, TcpStream>) {
		let mut lost_sockets: Vec<SocketAddr> = vec![];
		for (socket, stream) in sockets.iter() {
			if !self.peers.is_known_ip(&socket) {
				if let Ok(_) = stream.shutdown(Shutdown::Both) {
					debug!(
						"clean_lost_sockets: {} cleaned which's not in peers list",
						socket
					);
				}
				lost_sockets.push(socket.clone());
			}
		}

		for socket in lost_sockets {
			sockets.remove(&socket);
		}
	}

	pub fn stop(&self) {
		self.stop_state.lock().stop();
		self.peers.stop();
	}

	/// Pause means: stop all the current peers connection, only for tests.
	/// Note:
	/// 1. must pause the 'seed' thread also, to avoid the new egress peer connection
	/// 2. must pause the 'p2p-server' thread also, to avoid the new ingress peer connection.
	pub fn pause(&self) {
		self.peers.stop();
	}
}

/// A no-op network adapter used for testing.
pub struct DummyAdapter {}

impl ChainAdapter for DummyAdapter {
	fn total_difficulty(&self) -> Difficulty {
		Difficulty::min()
	}
	fn total_height(&self) -> u64 {
		0
	}
	fn get_transaction(&self, _h: Hash) -> Option<core::Transaction> {
		None
	}
	fn tx_kernel_received(&self, _h: Hash, _addr: SocketAddr) {}
	fn transaction_received(&self, _: core::Transaction, _stem: bool) {}
	fn compact_block_received(&self, _cb: core::CompactBlock, _addr: SocketAddr) -> bool {
		true
	}
	fn header_received(&self, _bh: core::BlockHeader, _addr: SocketAddr) -> bool {
		true
	}
	fn block_received(&self, _: core::Block, _: SocketAddr) -> bool {
		true
	}
	fn headers_received(&self, _: &[core::BlockHeader], _: SocketAddr) -> bool {
		true
	}
	fn locate_headers(&self, _: &[Hash]) -> Vec<core::BlockHeader> {
		vec![]
	}
	fn get_block(&self, _: Hash) -> Option<core::Block> {
		None
	}
	fn txhashset_read(&self, _h: Hash) -> Option<TxHashSetRead> {
		unimplemented!()
	}

	fn txhashset_receive_ready(&self) -> bool {
		false
	}

	fn txhashset_write(&self, _h: Hash, _txhashset_data: File, _peer_addr: SocketAddr) -> bool {
		false
	}

	fn txhashset_download_update(
		&self,
		_start_time: DateTime<Utc>,
		_downloaded_size: u64,
		_total_size: u64,
	) -> bool {
		false
	}
}

impl NetAdapter for DummyAdapter {
	fn find_peer_addrs(&self, _: Capabilities) -> Vec<SocketAddr> {
		vec![]
	}
	fn peer_addrs_received(&self, _: Vec<SocketAddr>) {}
	fn peer_difficulty(&self, _: SocketAddr, _: Difficulty, _: u64) {}
	fn is_banned(&self, _: SocketAddr) -> bool {
		false
	}
}
<commit_msg>Refuse duplicate peer IPs (#2347)<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fs::File;
use std::net::{Shutdown, SocketAddr, TcpListener, TcpStream};
use std::sync::Arc;
use std::time::Duration;
use std::{io, thread};

use crate::lmdb;

use crate::core::core;
use crate::core::core::hash::Hash;
use crate::core::global;
use crate::core::pow::Difficulty;
use crate::handshake::Handshake;
use crate::peer::Peer;
use crate::peers::Peers;
use crate::store::PeerStore;
use crate::types::{
	Capabilities, ChainAdapter, Error, NetAdapter, P2PConfig, ReasonForBan, Seeding, TxHashSetRead,
};
use crate::util::{Mutex, StopState};
use chrono::prelude::{DateTime, Utc};

/// P2P server implementation, handling bootstrapping to find and connect to
/// peers, receiving connections from other peers and keep track of all of them.
pub struct Server {
	pub config: P2PConfig,
	capabilities: Capabilities,
	handshake: Arc<Handshake>,
	pub peers: Arc<Peers>,
	stop_state: Arc<Mutex<StopState>>,
}

// TODO TLS
impl Server {
	/// Creates a new idle p2p server with no peers
	pub fn new(
		db_env: Arc<lmdb::Environment>,
		capab: Capabilities,
		config: P2PConfig,
		adapter: Arc<dyn ChainAdapter>,
		genesis: Hash,
		stop_state: Arc<Mutex<StopState>>,
	) -> Result<Server, Error> {
		Ok(Server {
			config: config.clone(),
			capabilities: capab,
			handshake: Arc::new(Handshake::new(genesis, config.clone())),
			peers: Arc::new(Peers::new(PeerStore::new(db_env)?, adapter, config)),
			stop_state,
		})
	}

	/// Starts a new TCP server and listen to incoming connections. This is a
	/// blocking call until the TCP server stops.
	pub fn listen(&self) -> Result<(), Error> {
		// start TCP listener and handle incoming connections
		let addr = SocketAddr::new(self.config.host, self.config.port);
		let listener = TcpListener::bind(addr)?;
		listener.set_nonblocking(true)?;

		let sleep_time = Duration::from_millis(1);
		loop {
			// Pause peer ingress connection request. Only for tests.
			if self.stop_state.lock().is_paused() {
				thread::sleep(Duration::from_secs(1));
				continue;
			}

			match listener.accept() {
				Ok((stream, peer_addr)) => {
					if self.check_undesirable(&stream) {
						continue;
					}
					if let Err(e) = self.handle_new_peer(stream) {
						debug!("Error accepting peer {}: {:?}", peer_addr.to_string(), e);
						let _ = self.peers.add_banned(peer_addr, ReasonForBan::BadHandshake);
					}
				}
				Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
					// nothing to do, will retry in next iteration
				}
				Err(e) => {
					debug!("Couldn't establish new client connection: {:?}", e);
				}
			}
			if self.stop_state.lock().is_stopped() {
				break;
			}
			thread::sleep(sleep_time);
		}
		Ok(())
	}

	/// Asks the server to connect to a new peer. Directly returns the peer if
	/// we're already connected to the provided address.
	pub fn connect(&self, addr: &SocketAddr) -> Result<Arc<Peer>, Error> {
		if Peer::is_denied(&self.config, &addr) {
			debug!("connect_peer: peer {} denied, not connecting.", addr);
			return Err(Error::ConnectionClose);
		}

		if global::is_production_mode() {
			let hs = self.handshake.clone();
			let addrs = hs.addrs.read();
			if addrs.contains(&addr) {
				debug!("connect: ignore connecting to PeerWithSelf, addr: {}", addr);
				return Err(Error::PeerWithSelf);
			}
		}

		if let Some(p) = self.peers.get_connected_peer(addr) {
			// if we're already connected to the addr, just return the peer
			trace!("connect_peer: already connected {}", addr);
			return Ok(p);
		}

		trace!(
			"connect_peer: on {}:{}. connecting to {}",
			self.config.host,
			self.config.port,
			addr
		);
		match TcpStream::connect_timeout(addr, Duration::from_secs(10)) {
			Ok(mut stream) => {
				let addr = SocketAddr::new(self.config.host, self.config.port);
				let total_diff = self.peers.total_difficulty();

				let mut peer = Peer::connect(
					&mut stream,
					self.capabilities,
					total_diff,
					addr,
					&self.handshake,
					self.peers.clone(),
				)?;
				peer.start(stream);
				let peer = Arc::new(peer);
				self.peers.add_connected(peer.clone())?;
				Ok(peer)
			}
			Err(e) => {
				trace!(
					"connect_peer: on {}:{}. Could not connect to {}: {:?}",
					self.config.host,
					self.config.port,
					addr,
					e
				);
				Err(Error::Connection(e))
			}
		}
	}

	fn handle_new_peer(&self, mut stream: TcpStream) -> Result<(), Error> {
		let total_diff = self.peers.total_difficulty();

		// accept the peer and add it to the server map
		let mut peer = Peer::accept(
			&mut stream,
			self.capabilities,
			total_diff,
			&self.handshake,
			self.peers.clone(),
		)?;
		peer.start(stream);
		self.peers.add_connected(Arc::new(peer))?;
		Ok(())
	}

	/// Checks whether there's any reason we don't want to accept a peer
	/// connection. There can be a couple of them:
	/// 1. The peer has been previously banned and the ban period hasn't
	/// expired yet.
	/// 2. We're already connected to a peer at the same IP. While there are
	/// many reasons multiple peers can legitimately share identical IP
	/// addresses (NAT), network distribution is improved if they choose
	/// different sets of peers themselves. In addition, it prevent potential
	/// duplicate connections, malicious or not.
	fn check_undesirable(&self, stream: &TcpStream) -> bool {
		// peer has been banned, go away!
		if let Ok(peer_addr) = stream.peer_addr() {
			let banned = self.peers.is_banned(peer_addr);
			let known_ip =
				self.peers.is_known_ip(&peer_addr) && self.config.seeding_type == Seeding::DNSSeed;
			if banned || known_ip {
				debug!("Peer {} banned or known, refusing connection.", peer_addr);
				if let Err(e) = stream.shutdown(Shutdown::Both) {
					debug!("Error shutting down conn: {:?}", e);
				}
				return true;
			}
		}
		false
	}

	pub fn stop(&self) {
		self.stop_state.lock().stop();
		self.peers.stop();
	}

	/// Pause means: stop all the current peers connection, only for tests.
	/// Note:
	/// 1. must pause the 'seed' thread also, to avoid the new egress peer connection
	/// 2. must pause the 'p2p-server' thread also, to avoid the new ingress peer connection.
	pub fn pause(&self) {
		self.peers.stop();
	}
}

/// A no-op network adapter used for testing.
pub struct DummyAdapter {}

impl ChainAdapter for DummyAdapter {
	fn total_difficulty(&self) -> Difficulty {
		Difficulty::min()
	}
	fn total_height(&self) -> u64 {
		0
	}
	fn get_transaction(&self, _h: Hash) -> Option<core::Transaction> {
		None
	}
	fn tx_kernel_received(&self, _h: Hash, _addr: SocketAddr) {}
	fn transaction_received(&self, _: core::Transaction, _stem: bool) {}
	fn compact_block_received(&self, _cb: core::CompactBlock, _addr: SocketAddr) -> bool {
		true
	}
	fn header_received(&self, _bh: core::BlockHeader, _addr: SocketAddr) -> bool {
		true
	}
	fn block_received(&self, _: core::Block, _: SocketAddr) -> bool {
		true
	}
	fn headers_received(&self, _: &[core::BlockHeader], _: SocketAddr) -> bool {
		true
	}
	fn locate_headers(&self, _: &[Hash]) -> Vec<core::BlockHeader> {
		vec![]
	}
	fn get_block(&self, _: Hash) -> Option<core::Block> {
		None
	}
	fn txhashset_read(&self, _h: Hash) -> Option<TxHashSetRead> {
		unimplemented!()
	}

	fn txhashset_receive_ready(&self) -> bool {
		false
	}

	fn txhashset_write(&self, _h: Hash, _txhashset_data: File, _peer_addr: SocketAddr) -> bool {
		false
	}

	fn txhashset_download_update(
		&self,
		_start_time: DateTime<Utc>,
		_downloaded_size: u64,
		_total_size: u64,
	) -> bool {
		false
	}
}

impl NetAdapter for DummyAdapter {
	fn find_peer_addrs(&self, _: Capabilities) -> Vec<SocketAddr> {
		vec![]
	}
	fn peer_addrs_received(&self, _: Vec<SocketAddr>) {}
	fn peer_difficulty(&self, _: SocketAddr, _: Difficulty, _: u64) {}
	fn is_banned(&self, _: SocketAddr) -> bool {
		false
	}
}
<|endoftext|>


<commit_before>use std::io;

use wayland_sys::client::*;

use event_queue::{create_event_queue, EventQueue};
use generated::client::wl_display::WlDisplay;
use Proxy;

/// Enum representing the possible reasons why connecting to the wayland server failed
#[derive(Debug)]
pub enum ConnectError {
    /// The library was compiled with the `dlopen` feature, and the `libwayland-client.so`
    /// library could not be found at runtime
    NoWaylandLib,
    /// Any needed library was found, but the listening socket of the server could not be
    /// found.
    ///
    /// Most of the time, this means that the program was not started from a wayland session.
    NoCompositorListening
}

/// Connect to the compositor socket
///
/// Attempt to connect to a Wayland compositor according to the environment variables.
///
/// On success, returns the display object, as well as the default event iterator associated with it.
pub fn default_connect() -> Result<(WlDisplay, EventQueue), ConnectError> {
    if !::wayland_sys::client::is_lib_available() { return Err(ConnectError::NoWaylandLib) }
    let ptr = unsafe { ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_connect, ::std::ptr::null()) };
    if ptr.is_null() {
        Err(ConnectError::NoCompositorListening)
    } else {
        let display = unsafe { WlDisplay::from_ptr(ptr as *mut _) };
        let eventiter = unsafe { create_event_queue(display.ptr() as *mut wl_display, None) };
        Ok((display, eventiter))
    }
}

impl WlDisplay {
    /// Non-blocking write to the server
    ///
    /// Will write as many pending requests as possible to the server socket. Never blocks: if not all
    /// requests coul be written, will return an io error `WouldBlock`.
    ///
    /// On success returns the number of written requests.
    pub fn flush(&self) -> io::Result<i32> {
        let ret = unsafe { ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_flush, self.ptr() as *mut _) };
        if ret >= 0 { Ok(ret) } else { Err(io::Error::last_os_error()) }
    }

    /// Create a new EventQueue
    ///
    /// No object is by default attached to it.
    pub fn create_event_queue(&self) -> EventQueue {
        let evq = unsafe { ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_create_queue, self.ptr() as *mut _) };
        unsafe { create_event_queue(self.ptr() as *mut _, Some(evq)) }
    }
}

impl Drop for WlDisplay {
    fn drop(&mut self) {
        unsafe { ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_disconnect, self.ptr() as *mut _) }
    }
}
<commit_msg>client: method to query protocol errors.<commit_after>use std::io;
use std::ffi::CStr;

use wayland_sys::client::*;

use event_queue::{create_event_queue, EventQueue};
use generated::client::wl_display::{WlDisplay, Error as DisplayError};
use Proxy;

/// Enum representing the possible reasons why connecting to the wayland server failed
#[derive(Debug)]
pub enum ConnectError {
    /// The library was compiled with the `dlopen` feature, and the `libwayland-client.so`
    /// library could not be found at runtime
    NoWaylandLib,
    /// Any needed library was found, but the listening socket of the server could not be
    /// found.
    ///
    /// Most of the time, this means that the program was not started from a wayland session.
    NoCompositorListening
}

/// Enum representing possible errors fatal to a wayland session
#[derive(Debug)]
pub enum FatalError {
    /// Session aborted after an I/O error
    Io(io::Error),
    /// Session aborted after a protocol error
    Protocol { kind: DisplayError, interface: String, id: u32 }
}

/// Connect to the compositor socket
///
/// Attempt to connect to a Wayland compositor according to the environment variables.
///
/// On success, returns the display object, as well as the default event iterator associated with it.
pub fn default_connect() -> Result<(WlDisplay, EventQueue), ConnectError> {
    if !::wayland_sys::client::is_lib_available() { return Err(ConnectError::NoWaylandLib) }
    let ptr = unsafe { ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_connect, ::std::ptr::null()) };
    if ptr.is_null() {
        Err(ConnectError::NoCompositorListening)
    } else {
        let display = unsafe { WlDisplay::from_ptr(ptr as *mut _) };
        let eventiter = unsafe { create_event_queue(display.ptr() as *mut wl_display, None) };
        Ok((display, eventiter))
    }
}

impl WlDisplay {
    /// Non-blocking write to the server
    ///
    /// Will write as many pending requests as possible to the server socket. Never blocks: if not all
    /// requests coul be written, will return an io error `WouldBlock`.
    ///
    /// On success returns the number of written requests.
    pub fn flush(&self) -> io::Result<i32> {
        let ret = unsafe { ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_flush, self.ptr() as *mut _) };
        if ret >= 0 { Ok(ret) } else { Err(io::Error::last_os_error()) }
    }

    /// Create a new EventQueue
    ///
    /// No object is by default attached to it.
    pub fn create_event_queue(&self) -> EventQueue {
        let evq = unsafe { ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_create_queue, self.ptr() as *mut _) };
        unsafe { create_event_queue(self.ptr() as *mut _, Some(evq)) }
    }

    /// Get the last error that occured on the session
    ///
    /// Such errors are *fatal*, meaning that if this function does not
    /// return `None`, your session is not usable any longer.
    ///
    /// As such, this function mostly provide diagnistics information. You can have a hint
    /// an error might have been generated if I/O methods of EventQueue start returning errors.
    pub fn last_error(&self) -> Option<FatalError> {
        let err = unsafe { ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_get_error, self.ptr() as *mut _) };
        if err == 0 {
            None
        } else if err == ::libc::EPROTO {
            let mut interface = ::std::ptr::null_mut();
            let mut id = 0;
            let code = unsafe { ffi_dispatch!(
                WAYLAND_CLIENT_HANDLE, wl_display_get_protocol_error,
                self.ptr() as *mut _, &mut interface, &mut id
            ) };
            let interface = if interface.is_null() {
                "<unkown interface>".to_owned()
            } else {
                unsafe { CStr::from_ptr((*interface).name) }.to_string_lossy().into_owned()
            };
            Some(FatalError::Protocol {
                kind: DisplayError::from_raw(code).expect("Unknown protocol error"),
                interface: interface,
                id: id
            })
        } else {
            Some(FatalError::Io(io::Error::from_raw_os_error(err)))
        }
    }
}

impl Drop for WlDisplay {
    fn drop(&mut self) {
        unsafe { ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_disconnect, self.ptr() as *mut _) }
    }
}
<|endoftext|>


<commit_before>use cookie_factory::*;
use elements::{Info, Seek, SeekHead, SegmentElement, Cluster};
use super::ebml::{vint_size, gen_vint, gen_vid, gen_uint};
use serializer::ebml::{gen_u64,gen_f64_ref};

pub fn seek_size(s: &Seek) -> u8 {
    // byte size of id (vid+size)+ data and position vid+size+int
    // FIXME: arbitrarily bad value
    vint_size(vint_size((s.id.len() + 10) as u64) as u64)
}

pub fn gen_segment<'a>(input: (&'a mut [u8], usize),
                       s: &SegmentElement)
                       -> Result<(&'a mut [u8], usize), GenError> {
    unimplemented!();
    /*do_gen!(input,
    gen_call!(gen_vid, 0x18538067) >>
    gen_call!(gen_vint, 4)
  )*/
}

pub fn gen_seek<'a>(input: (&'a mut [u8], usize),
                    s: &Seek)
                    -> Result<(&'a mut [u8], usize), GenError> {
    let capacity = 8 + 2 + vint_size(s.id.len() as u64) as u64 + s.id.len() as u64;

    gen_ebml_master!(input,
    0x4DBB, vint_size(capacity),
    gen_ebml_binary!(0x53AB, s.id) >>
    gen_ebml_uint!(0x53AC, s.position, vint_size(s.position))
  )
}

pub fn gen_seek_head<'a>(input: (&'a mut [u8], usize),
                         s: &SeekHead)
                         -> Result<(&'a mut [u8], usize), GenError> {
    let capacity = s.positions.iter().fold(0u64, |acc, seek| {
        acc + 4 + 8 + 2 + vint_size(seek.id.len() as u64) as u64 + seek.id.len() as u64
    });
    println!("gen_seek_head: calculated capacity: {} -> {} bytes", capacity, vint_size(capacity));

    let byte_capacity = vint_size(capacity as u64);
    gen_ebml_master!(input,
    0x114D9B74, byte_capacity,
    gen_many_ref!(&s.positions, gen_seek)
  )
}

pub fn gen_info<'a>(input: (&'a mut [u8], usize),
                         i: &Info)
                         -> Result<(&'a mut [u8], usize), GenError> {
    //FIXME: probably not large enough
    let capacity = 2 + i.segment_uid.as_ref().map(|s| s.len()).unwrap_or(0)
      + 2 + i.segment_filename.as_ref().map(|s| s.len()).unwrap_or(0)
      + 3 + i.prev_uid.as_ref().map(|s| s.len()).unwrap_or(0)
      + 3 + i.prev_filename.as_ref().map(|s| s.len()).unwrap_or(0)
      + 3 + i.next_uid.as_ref().map(|s| s.len()).unwrap_or(0)
      + 3 + i.next_filename.as_ref().map(|s| s.len()).unwrap_or(0)
      + 2 + i.segment_family.as_ref().map(|s| s.len()).unwrap_or(0)
      //FIXME;
      // chapter translate
      + 8
      + 8
      + 2 + i.date_utc.as_ref().map(|s| s.len()).unwrap_or(0)
      + 2 + i.title.as_ref().map(|s| s.len()).unwrap_or(0)
      + 2 + i.muxing_app.len()
      + 2 + i.writing_app.len();

    let byte_capacity = vint_size(capacity as u64);
    gen_ebml_master!(input,
      0x1549A966, byte_capacity,
      do_gen!(
           gen_opt!( i.segment_uid, gen_ebml_binary!(0x73A4) )
        >> gen_opt!( i.segment_filename, gen_ebml_str!(0x7384) )
        >> gen_opt!( i.prev_uid, gen_ebml_binary!(0x3CB923) )
        >> gen_opt!( i.prev_filename, gen_ebml_str!(0x3C83AB) )
        >> gen_opt!( i.next_uid, gen_ebml_binary!(0x3EB923) )
        >> gen_opt!( i.next_filename, gen_ebml_str!(0x3E83BB) )
        >> gen_opt!( i.segment_family, gen_ebml_binary!(0x4444) )
        //>> gen_opt!( i.chapter_translate, gen_chapter_translate )
        >> gen_call!(gen_u64, 0x2AD7B1, i.timecode_scale)
        >> gen_opt!( i.duration, gen_call!(gen_f64_ref, 0x4489) )
        >> gen_opt!( i.date_utc, gen_ebml_binary!(0x4461) )
        >> gen_opt!( i.title, gen_ebml_str!(0x7BA9) )
        >> gen_ebml_str!(0x4D80, i.muxing_app)
        >> gen_ebml_str!(0x5741, i.writing_app)
      )
    )
}

#[macro_export]
macro_rules! my_gen_many (
    (($i:expr, $idx:expr), $l:expr, $f:ident) => (
        $l.into_iter().fold(
            Ok(($i,$idx)),
            |r,v| {
                match r {
                    Err(e) => Err(e),
                    Ok(x) => { $f(x, v) },
                }
            }
        )
    );
    (($i:expr, $idx:expr), $l:expr, $f:ident!( $($args:tt)* )) => (
        $l.into_iter().fold(
            Ok(($i,$idx)),
            |r,v| {
                match r {
                    Err(e) => Err(e),
                    Ok(x) => {
                      let (i, idx) = x;
                      $f!((i, idx), $($args)*, v)
                    },
                }
            }
        )
    );
);

pub fn gen_cluster<'a>(input: (&'a mut [u8], usize),
                         c: &Cluster)
                         -> Result<(&'a mut [u8], usize), GenError> {
    let capacity = 2 + 8
      // FIXME: serialize SilentTracks
      + 2 + 8
      + 2 + 8
      + 2 + c.simple_block.iter().fold(0, |acc, data| acc+ data.len())
      // FIXME serialize BlockGRoups
      // FIXME: serialize encrypted block
      ;

    let byte_capacity = vint_size(capacity as u64);
    gen_ebml_master!(input,
      0x1F43B675, byte_capacity,
      do_gen!(
           gen_ebml_uint!(0xE7, c.timecode)
        //>> gen_opt!( c.position, gen_ebml_uint!(0xA7) )
        //>> gen_opt!( c.prev_size, gen_ebml_uint!(0xAB) )
        >> my_gen_many!( &c.simple_block, gen_ebml_binary!( 0xA3 ) )
      )
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use nom::*;
    use std::iter::repeat;

    fn test_seek_head_serializer(mut seeks: Vec<(u64, Vec<u8>)>) -> bool {
        println!("testing for {:?}", seeks);

        let mut should_fail = false;
        if seeks.len() == 0 {
            should_fail = true;
        }

        for &(_, ref id) in seeks.iter() {
            println!("id: {}", id.to_hex(16));
            if id.len() == 0 {
                println!("id is empty, returning");
                return true;
                //should_fail = true;
            }
        }

        if should_fail {
            println!("the parser should fail");
        }

        let capacity = seeks.iter().fold(0, |acc, &(_, ref v)| acc + 8 + v.len() + 100);
        println!("defining capacity as {}", capacity);

        let mut data = Vec::with_capacity(capacity);
        data.extend(repeat(0).take(capacity));

        let seek_head = SeekHead {
            positions: seeks.iter()
                .cloned()
                .map(|(position, id)| {
                    Seek {
                        id: id,
                        position: position,
                    }
                })
                .collect(),
        };

        let ser_res = {
            let gen_res = gen_seek_head((&mut data[..], 0), &seek_head);
            println!("gen_res: {:?}", gen_res);
            if let Err(e) = gen_res {
                println!("gen_res is error: {:?}", e);
                println!("should fail: {:?}", should_fail);
                return should_fail;
                /*if should_fail {
          println!("should fail");
          return true;
        }*/
            }
        };

        println!("ser_res: {:?}", ser_res);

        let parse_res = ::elements::segment_element(&data[..]);
        println!("parse_res: {:?}", parse_res);
        match parse_res {
            Ok((rest, SegmentElement::SeekHead(o))) => {
                if should_fail {
                    println!("parser should have failed on input for {:?}", seek_head);
                    println!("{}", (&data[..]).to_hex(16));
                    return false;
                }

                assert_eq!(seek_head, o);
                return true;
            }
            e => {
                if should_fail {
                    return true;
                }

                panic!(format!("parse error: {:?} for input: {:?}", e, seeks))
            }
        }

        false
    }

    quickcheck! {
    fn test_seek_head(seeks: Vec<(u64, Vec<u8>)>) -> bool {
      test_seek_head_serializer(seeks)
    }
  }
}
<commit_msg>some missing members<commit_after>use cookie_factory::*;
use elements::{Info, Seek, SeekHead, SegmentElement, Cluster};
use super::ebml::{vint_size, gen_vint, gen_vid, gen_uint};
use serializer::ebml::{gen_u64,gen_f64_ref};

pub fn seek_size(s: &Seek) -> u8 {
    // byte size of id (vid+size)+ data and position vid+size+int
    // FIXME: arbitrarily bad value
    vint_size(vint_size((s.id.len() + 10) as u64) as u64)
}

pub fn gen_segment<'a>(input: (&'a mut [u8], usize),
                       s: &SegmentElement)
                       -> Result<(&'a mut [u8], usize), GenError> {
    unimplemented!();
    /*do_gen!(input,
    gen_call!(gen_vid, 0x18538067) >>
    gen_call!(gen_vint, 4)
  )*/
}

pub fn gen_seek<'a>(input: (&'a mut [u8], usize),
                    s: &Seek)
                    -> Result<(&'a mut [u8], usize), GenError> {
    let capacity = 8 + 2 + vint_size(s.id.len() as u64) as u64 + s.id.len() as u64;

    gen_ebml_master!(input,
    0x4DBB, vint_size(capacity),
    gen_ebml_binary!(0x53AB, s.id) >>
    gen_ebml_uint!(0x53AC, s.position, vint_size(s.position))
  )
}

pub fn gen_seek_head<'a>(input: (&'a mut [u8], usize),
                         s: &SeekHead)
                         -> Result<(&'a mut [u8], usize), GenError> {
    let capacity = s.positions.iter().fold(0u64, |acc, seek| {
        acc + 4 + 8 + 2 + vint_size(seek.id.len() as u64) as u64 + seek.id.len() as u64
    });
    println!("gen_seek_head: calculated capacity: {} -> {} bytes", capacity, vint_size(capacity));

    let byte_capacity = vint_size(capacity as u64);
    gen_ebml_master!(input,
    0x114D9B74, byte_capacity,
    gen_many_ref!(&s.positions, gen_seek)
  )
}

pub fn gen_info<'a>(input: (&'a mut [u8], usize),
                         i: &Info)
                         -> Result<(&'a mut [u8], usize), GenError> {
    //FIXME: probably not large enough
    let capacity = 2 + i.segment_uid.as_ref().map(|s| s.len()).unwrap_or(0)
      + 2 + i.segment_filename.as_ref().map(|s| s.len()).unwrap_or(0)
      + 3 + i.prev_uid.as_ref().map(|s| s.len()).unwrap_or(0)
      + 3 + i.prev_filename.as_ref().map(|s| s.len()).unwrap_or(0)
      + 3 + i.next_uid.as_ref().map(|s| s.len()).unwrap_or(0)
      + 3 + i.next_filename.as_ref().map(|s| s.len()).unwrap_or(0)
      + 2 + i.segment_family.as_ref().map(|s| s.len()).unwrap_or(0)
      //FIXME;
      // chapter translate
      + 8
      + 8
      + 2 + i.date_utc.as_ref().map(|s| s.len()).unwrap_or(0)
      + 2 + i.title.as_ref().map(|s| s.len()).unwrap_or(0)
      + 2 + i.muxing_app.len()
      + 2 + i.writing_app.len();

    let byte_capacity = vint_size(capacity as u64);
    gen_ebml_master!(input,
      0x1549A966, byte_capacity,
      do_gen!(
           gen_opt!( i.segment_uid, gen_ebml_binary!(0x73A4) )
        >> gen_opt!( i.segment_filename, gen_ebml_str!(0x7384) )
        >> gen_opt!( i.prev_uid, gen_ebml_binary!(0x3CB923) )
        >> gen_opt!( i.prev_filename, gen_ebml_str!(0x3C83AB) )
        >> gen_opt!( i.next_uid, gen_ebml_binary!(0x3EB923) )
        >> gen_opt!( i.next_filename, gen_ebml_str!(0x3E83BB) )
        >> gen_opt!( i.segment_family, gen_ebml_binary!(0x4444) )
        //>> gen_opt!( i.chapter_translate, gen_chapter_translate )
        >> gen_call!(gen_u64, 0x2AD7B1, i.timecode_scale)
        >> gen_opt!( i.duration, gen_call!(gen_f64_ref, 0x4489) )
        >> gen_opt!( i.date_utc, gen_ebml_binary!(0x4461) )
        >> gen_opt!( i.title, gen_ebml_str!(0x7BA9) )
        >> gen_ebml_str!(0x4D80, i.muxing_app)
        >> gen_ebml_str!(0x5741, i.writing_app)
      )
    )
}

#[macro_export]
macro_rules! my_gen_many (
    (($i:expr, $idx:expr), $l:expr, $f:ident) => (
        $l.into_iter().fold(
            Ok(($i,$idx)),
            |r,v| {
                match r {
                    Err(e) => Err(e),
                    Ok(x) => { $f(x, v) },
                }
            }
        )
    );
    (($i:expr, $idx:expr), $l:expr, $f:ident!( $($args:tt)* )) => (
        $l.into_iter().fold(
            Ok(($i,$idx)),
            |r,v| {
                match r {
                    Err(e) => Err(e),
                    Ok(x) => {
                      let (i, idx) = x;
                      $f!((i, idx), $($args)*, v)
                    },
                }
            }
        )
    );
);

pub fn gen_cluster<'a>(input: (&'a mut [u8], usize),
                         c: &Cluster)
                         -> Result<(&'a mut [u8], usize), GenError> {
    let capacity = 2 + 8
      // FIXME: serialize SilentTracks
      + 2 + 8
      + 2 + 8
      + 2 + c.simple_block.iter().fold(0, |acc, data| acc+ data.len())
      // FIXME serialize BlockGRoups
      // FIXME: serialize encrypted block
      ;

    let byte_capacity = vint_size(capacity as u64);
    gen_ebml_master!(input,
      0x1F43B675, byte_capacity,
      do_gen!(
           gen_ebml_uint!(0xE7, c.timecode)
        >> gen_opt!( c.position, gen_ebml_uint!(0xA7) )
        >> gen_opt!( c.prev_size, gen_ebml_uint!(0xAB) )
        >> my_gen_many!( &c.simple_block, gen_ebml_binary!( 0xA3 ) )
      )
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use nom::*;
    use std::iter::repeat;

    fn test_seek_head_serializer(mut seeks: Vec<(u64, Vec<u8>)>) -> bool {
        println!("testing for {:?}", seeks);

        let mut should_fail = false;
        if seeks.len() == 0 {
            should_fail = true;
        }

        for &(_, ref id) in seeks.iter() {
            println!("id: {}", id.to_hex(16));
            if id.len() == 0 {
                println!("id is empty, returning");
                return true;
                //should_fail = true;
            }
        }

        if should_fail {
            println!("the parser should fail");
        }

        let capacity = seeks.iter().fold(0, |acc, &(_, ref v)| acc + 8 + v.len() + 100);
        println!("defining capacity as {}", capacity);

        let mut data = Vec::with_capacity(capacity);
        data.extend(repeat(0).take(capacity));

        let seek_head = SeekHead {
            positions: seeks.iter()
                .cloned()
                .map(|(position, id)| {
                    Seek {
                        id: id,
                        position: position,
                    }
                })
                .collect(),
        };

        let ser_res = {
            let gen_res = gen_seek_head((&mut data[..], 0), &seek_head);
            println!("gen_res: {:?}", gen_res);
            if let Err(e) = gen_res {
                println!("gen_res is error: {:?}", e);
                println!("should fail: {:?}", should_fail);
                return should_fail;
                /*if should_fail {
          println!("should fail");
          return true;
        }*/
            }
        };

        println!("ser_res: {:?}", ser_res);

        let parse_res = ::elements::segment_element(&data[..]);
        println!("parse_res: {:?}", parse_res);
        match parse_res {
            Ok((rest, SegmentElement::SeekHead(o))) => {
                if should_fail {
                    println!("parser should have failed on input for {:?}", seek_head);
                    println!("{}", (&data[..]).to_hex(16));
                    return false;
                }

                assert_eq!(seek_head, o);
                return true;
            }
            e => {
                if should_fail {
                    return true;
                }

                panic!(format!("parse error: {:?} for input: {:?}", e, seeks))
            }
        }

        false
    }

    quickcheck! {
    fn test_seek_head(seeks: Vec<(u64, Vec<u8>)>) -> bool {
      test_seek_head_serializer(seeks)
    }
  }
}
<|endoftext|>


<commit_before>trait StrParseExt {
    fn end_of_start_rest<F1, F2>(&self, is_first: F1, is_rest: F2) -> Option<usize>
        where F1: Fn(char) -> bool,
              F2: Fn(char) -> bool;
}

impl<'a> StrParseExt for &'a str {
    fn end_of_start_rest<F1, F2>(&self, is_first: F1, is_rest: F2) -> Option<usize>
        where F1: Fn(char) -> bool,
              F2: Fn(char) -> bool,
    {
        let mut positions = self.char_indices();

        match positions.next() {
            Some((_, c)) if is_first(c) => (),
            Some((_, _)) => return None,
            None => return None,
        };

        let mut positions = positions.skip_while(|&(_, c)| is_rest(c));
        match positions.next() {
            Some((offset, _)) => Some(offset),
            None => Some(self.len()),
        }
    }
}

pub trait XmlStr {
    /// Find the end of the quoted attribute value, not including the quote
    fn end_of_attribute(&self, quote: &str) -> Option<usize>;
    /// Find the end of the direct character data
    fn end_of_char_data(&self) -> Option<usize>;
    /// Find the end of the CData section, not including the ]]>
    fn end_of_cdata(&self) -> Option<usize>;
    /// Find the end of a run of decimal characters
    fn end_of_decimal_chars(&self) -> Option<usize>;
    /// Find the end of a run of hexidecimal characters
    fn end_of_hex_chars(&self) -> Option<usize>;
    /// Find the end of the comment, not including the -->
    fn end_of_comment(&self) -> Option<usize>;
    /// Find the end of the processing instruction, not including the ?>
    fn end_of_pi_value(&self) -> Option<usize>;
    /// Find the end of the [Name](http://www.w3.org/TR/xml/#NT-Name)
    fn end_of_name(&self) -> Option<usize>;
    /// Find the end of the [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)
    fn end_of_ncname(&self) -> Option<usize>;
    /// Find the end of a run of space characters
    fn end_of_space(&self) -> Option<usize>;
    /// Find the end of the starting tag
    fn end_of_start_tag(&self) -> Option<usize>;
}

impl<'a> XmlStr for &'a str {
    fn end_of_attribute(&self, quote: &str) -> Option<usize> {
        if self.len() == 0 ||
           self.starts_with("&") ||
           self.starts_with("<") ||
           self.starts_with(quote)
        {
            return None;
        }

        let quote_char = quote.chars().next().expect("Cant have null quote");

        let mut positions = self.char_indices().skip_while(|&(_, c)| c != '&' && c != '<' && c != quote_char);

        match positions.next() {
            Some((offset, _)) => Some(offset),
            None => Some(self.len())
        }
    }

    fn end_of_char_data(&self) -> Option<usize> {
        if self.starts_with("<") ||
           self.starts_with("&") ||
           self.starts_with("]]>")
        {
            return None
        }

        // Using a hex literal because emacs' rust-mode doesn't
        // understand ] in a char literal. :-(
        let mut positions = self.char_indices().skip_while(|&(_, c)| c != '<' && c != '&' && c != '\x5d');

        loop {
            match positions.next() {
                None => return Some(self.len()),
                Some((offset, c)) if c == '<' || c == '&' => return Some(offset),
                Some((offset, _)) => {
                    let tail = &self[offset..];
                    if tail.starts_with("]]>") {
                        return Some(offset)
                    } else {
                        // False alarm, resume scanning
                        continue;
                    }
                },
            }
        }
    }

    fn end_of_cdata(&self) -> Option<usize> {
        match self.find("]]>") {
            None => None,
            Some(offset) => Some(offset),
        }
    }

    fn end_of_decimal_chars(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_decimal_char(),
                               |c| c.is_decimal_char())
    }

    fn end_of_hex_chars(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_hex_char(),
                               |c| c.is_hex_char())
    }

    fn end_of_comment(&self) -> Option<usize> {
        // This deliberately does not include the >. -- is not allowed
        // in a comment, so we can just test the end if it matches the
        // complete close delimiter.
        match self.find("--") {
            None => None,
            Some(offset) => Some(offset),
        }
    }

    fn end_of_pi_value(&self) -> Option<usize> {
        match self.find("?>") {
            None => None,
            Some(offset) => Some(offset),
        }
    }

    fn end_of_name(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_name_start_char(), |c| c.is_name_char())
    }

    fn end_of_ncname(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_ncname_start_char(), |c| c.is_ncname_char())
    }

    fn end_of_space(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_space_char(), |c| c.is_space_char())
    }

    fn end_of_start_tag(&self) -> Option<usize> {
        let mut positions = self.char_indices();

        match positions.next() {
            Some((_, c)) if '<' == c => (),
            _ => return None,
        };

        match positions.next() {
            Some((offset, c)) =>
                match c {
                    '?' | '!' | '/' => None,
                    _ => Some(offset),
                },
            None => Some(self.len()),
        }
    }
}

/// Predicates used when parsing an characters in an XML document.
pub trait XmlChar {
    /// Is this a [NameStartChar](http://www.w3.org/TR/xml/#NT-NameStartChar)?
    fn is_name_start_char(self) -> bool;
    /// Is this a [NameChar](http://www.w3.org/TR/xml/#NT-NameChar)?
    fn is_name_char(self) -> bool;
    /// Does this start a [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)?
    fn is_ncname_start_char(self) -> bool;
    /// Is this a component of a [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)?
    fn is_ncname_char(self) -> bool;
    /// Is this an [XML space](http://www.w3.org/TR/xml/#NT-S)?
    fn is_space_char(self) -> bool;
    fn is_decimal_char(self) -> bool;
    fn is_hex_char(self) -> bool;
}

impl XmlChar for char {
    fn is_name_start_char(self) -> bool {
        self == ':' || self.is_ncname_start_char()
    }

    fn is_name_char(self) -> bool {
        self.is_name_start_char() || self.is_ncname_char()
    }

    fn is_ncname_start_char(self) -> bool {
        match self {
            'A'...'Z'                   |
            '_'                         |
            'a'...'z'                   |
            '\u{0000C0}'...'\u{0000D6}' |
            '\u{0000D8}'...'\u{0000F6}' |
            '\u{0000F8}'...'\u{0002FF}' |
            '\u{000370}'...'\u{00037D}' |
            '\u{00037F}'...'\u{001FFF}' |
            '\u{00200C}'...'\u{00200D}' |
            '\u{002070}'...'\u{00218F}' |
            '\u{002C00}'...'\u{002FEF}' |
            '\u{003001}'...'\u{00D7FF}' |
            '\u{00F900}'...'\u{00FDCF}' |
            '\u{00FDF0}'...'\u{00FFFD}' |
            '\u{010000}'...'\u{0EFFFF}' => true,
            _ => false,
        }
    }

    fn is_ncname_char(self) -> bool {
        if self.is_ncname_start_char() { return true; }
        match self {
            '-'                     |
            '.'                     |
            '0'...'9'               |
            '\u{00B7}'              |
            '\u{0300}'...'\u{036F}' |
            '\u{203F}'...'\u{2040}' => true,
            _ => false
        }
    }

    fn is_space_char(self) -> bool {
        match self {
            '\x20' |
            '\x09' |
            '\x0D' |
            '\x0A' => true,
            _ => false,
        }
    }

    fn is_decimal_char(self) -> bool {
        match self {
            '0'...'9' => true,
            _ => false,
        }
    }

    fn is_hex_char(self) -> bool {
        match self {
            '0'...'9' |
            'a'...'f' |
            'A'...'F' => true,
            _ => false,
        }
    }
}

#[cfg(test)]
mod test {

use super::XmlStr;

#[test]
fn end_of_char_data_leading_ampersand() {
    assert_eq!("&".end_of_char_data(), None);
}

#[test]
fn end_of_char_data_leading_less_than() {
    assert_eq!("<".end_of_char_data(), None);
}

#[test]
fn end_of_char_data_leading_cdata_end() {
    assert_eq!("]]>".end_of_char_data(), None);
}

#[test]
fn end_of_char_data_until_ampersand() {
    assert_eq!("hello&world".end_of_char_data(), Some("hello".len()));
}

#[test]
fn end_of_char_data_until_less_than() {
    assert_eq!("hello<world".end_of_char_data(), Some("hello".len()));
}

#[test]
fn end_of_char_data_until_cdata_end() {
    assert_eq!("hello]]>world".end_of_char_data(), Some("hello".len()));
}

#[test]
fn end_of_char_data_includes_right_square() {
    assert_eq!("hello]world".end_of_char_data(), Some("hello]world".len()));
}

}
<commit_msg>No longer need to use hex literal<commit_after>trait StrParseExt {
    fn end_of_start_rest<F1, F2>(&self, is_first: F1, is_rest: F2) -> Option<usize>
        where F1: Fn(char) -> bool,
              F2: Fn(char) -> bool;
}

impl<'a> StrParseExt for &'a str {
    fn end_of_start_rest<F1, F2>(&self, is_first: F1, is_rest: F2) -> Option<usize>
        where F1: Fn(char) -> bool,
              F2: Fn(char) -> bool,
    {
        let mut positions = self.char_indices();

        match positions.next() {
            Some((_, c)) if is_first(c) => (),
            Some((_, _)) => return None,
            None => return None,
        };

        let mut positions = positions.skip_while(|&(_, c)| is_rest(c));
        match positions.next() {
            Some((offset, _)) => Some(offset),
            None => Some(self.len()),
        }
    }
}

pub trait XmlStr {
    /// Find the end of the quoted attribute value, not including the quote
    fn end_of_attribute(&self, quote: &str) -> Option<usize>;
    /// Find the end of the direct character data
    fn end_of_char_data(&self) -> Option<usize>;
    /// Find the end of the CData section, not including the ]]>
    fn end_of_cdata(&self) -> Option<usize>;
    /// Find the end of a run of decimal characters
    fn end_of_decimal_chars(&self) -> Option<usize>;
    /// Find the end of a run of hexidecimal characters
    fn end_of_hex_chars(&self) -> Option<usize>;
    /// Find the end of the comment, not including the -->
    fn end_of_comment(&self) -> Option<usize>;
    /// Find the end of the processing instruction, not including the ?>
    fn end_of_pi_value(&self) -> Option<usize>;
    /// Find the end of the [Name](http://www.w3.org/TR/xml/#NT-Name)
    fn end_of_name(&self) -> Option<usize>;
    /// Find the end of the [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)
    fn end_of_ncname(&self) -> Option<usize>;
    /// Find the end of a run of space characters
    fn end_of_space(&self) -> Option<usize>;
    /// Find the end of the starting tag
    fn end_of_start_tag(&self) -> Option<usize>;
}

impl<'a> XmlStr for &'a str {
    fn end_of_attribute(&self, quote: &str) -> Option<usize> {
        if self.len() == 0 ||
           self.starts_with("&") ||
           self.starts_with("<") ||
           self.starts_with(quote)
        {
            return None;
        }

        let quote_char = quote.chars().next().expect("Cant have null quote");

        let mut positions = self.char_indices().skip_while(|&(_, c)| c != '&' && c != '<' && c != quote_char);

        match positions.next() {
            Some((offset, _)) => Some(offset),
            None => Some(self.len())
        }
    }

    fn end_of_char_data(&self) -> Option<usize> {
        if self.starts_with("<") ||
           self.starts_with("&") ||
           self.starts_with("]]>")
        {
            return None
        }

        let mut positions = self.char_indices().skip_while(|&(_, c)| {
            c != '<' && c != '&' && c != ']'
        });

        loop {
            match positions.next() {
                None => return Some(self.len()),
                Some((offset, c)) if c == '<' || c == '&' => return Some(offset),
                Some((offset, _)) => {
                    let tail = &self[offset..];
                    if tail.starts_with("]]>") {
                        return Some(offset)
                    } else {
                        // False alarm, resume scanning
                        continue;
                    }
                },
            }
        }
    }

    fn end_of_cdata(&self) -> Option<usize> {
        match self.find("]]>") {
            None => None,
            Some(offset) => Some(offset),
        }
    }

    fn end_of_decimal_chars(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_decimal_char(),
                               |c| c.is_decimal_char())
    }

    fn end_of_hex_chars(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_hex_char(),
                               |c| c.is_hex_char())
    }

    fn end_of_comment(&self) -> Option<usize> {
        // This deliberately does not include the >. -- is not allowed
        // in a comment, so we can just test the end if it matches the
        // complete close delimiter.
        match self.find("--") {
            None => None,
            Some(offset) => Some(offset),
        }
    }

    fn end_of_pi_value(&self) -> Option<usize> {
        match self.find("?>") {
            None => None,
            Some(offset) => Some(offset),
        }
    }

    fn end_of_name(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_name_start_char(), |c| c.is_name_char())
    }

    fn end_of_ncname(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_ncname_start_char(), |c| c.is_ncname_char())
    }

    fn end_of_space(&self) -> Option<usize> {
        self.end_of_start_rest(|c| c.is_space_char(), |c| c.is_space_char())
    }

    fn end_of_start_tag(&self) -> Option<usize> {
        let mut positions = self.char_indices();

        match positions.next() {
            Some((_, c)) if '<' == c => (),
            _ => return None,
        };

        match positions.next() {
            Some((offset, c)) =>
                match c {
                    '?' | '!' | '/' => None,
                    _ => Some(offset),
                },
            None => Some(self.len()),
        }
    }
}

/// Predicates used when parsing an characters in an XML document.
pub trait XmlChar {
    /// Is this a [NameStartChar](http://www.w3.org/TR/xml/#NT-NameStartChar)?
    fn is_name_start_char(self) -> bool;
    /// Is this a [NameChar](http://www.w3.org/TR/xml/#NT-NameChar)?
    fn is_name_char(self) -> bool;
    /// Does this start a [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)?
    fn is_ncname_start_char(self) -> bool;
    /// Is this a component of a [NCName](http://www.w3.org/TR/REC-xml-names/#NT-NCName)?
    fn is_ncname_char(self) -> bool;
    /// Is this an [XML space](http://www.w3.org/TR/xml/#NT-S)?
    fn is_space_char(self) -> bool;
    fn is_decimal_char(self) -> bool;
    fn is_hex_char(self) -> bool;
}

impl XmlChar for char {
    fn is_name_start_char(self) -> bool {
        self == ':' || self.is_ncname_start_char()
    }

    fn is_name_char(self) -> bool {
        self.is_name_start_char() || self.is_ncname_char()
    }

    fn is_ncname_start_char(self) -> bool {
        match self {
            'A'...'Z'                   |
            '_'                         |
            'a'...'z'                   |
            '\u{0000C0}'...'\u{0000D6}' |
            '\u{0000D8}'...'\u{0000F6}' |
            '\u{0000F8}'...'\u{0002FF}' |
            '\u{000370}'...'\u{00037D}' |
            '\u{00037F}'...'\u{001FFF}' |
            '\u{00200C}'...'\u{00200D}' |
            '\u{002070}'...'\u{00218F}' |
            '\u{002C00}'...'\u{002FEF}' |
            '\u{003001}'...'\u{00D7FF}' |
            '\u{00F900}'...'\u{00FDCF}' |
            '\u{00FDF0}'...'\u{00FFFD}' |
            '\u{010000}'...'\u{0EFFFF}' => true,
            _ => false,
        }
    }

    fn is_ncname_char(self) -> bool {
        if self.is_ncname_start_char() { return true; }
        match self {
            '-'                     |
            '.'                     |
            '0'...'9'               |
            '\u{00B7}'              |
            '\u{0300}'...'\u{036F}' |
            '\u{203F}'...'\u{2040}' => true,
            _ => false
        }
    }

    fn is_space_char(self) -> bool {
        match self {
            '\x20' |
            '\x09' |
            '\x0D' |
            '\x0A' => true,
            _ => false,
        }
    }

    fn is_decimal_char(self) -> bool {
        match self {
            '0'...'9' => true,
            _ => false,
        }
    }

    fn is_hex_char(self) -> bool {
        match self {
            '0'...'9' |
            'a'...'f' |
            'A'...'F' => true,
            _ => false,
        }
    }
}

#[cfg(test)]
mod test {

use super::XmlStr;

#[test]
fn end_of_char_data_leading_ampersand() {
    assert_eq!("&".end_of_char_data(), None);
}

#[test]
fn end_of_char_data_leading_less_than() {
    assert_eq!("<".end_of_char_data(), None);
}

#[test]
fn end_of_char_data_leading_cdata_end() {
    assert_eq!("]]>".end_of_char_data(), None);
}

#[test]
fn end_of_char_data_until_ampersand() {
    assert_eq!("hello&world".end_of_char_data(), Some("hello".len()));
}

#[test]
fn end_of_char_data_until_less_than() {
    assert_eq!("hello<world".end_of_char_data(), Some("hello".len()));
}

#[test]
fn end_of_char_data_until_cdata_end() {
    assert_eq!("hello]]>world".end_of_char_data(), Some("hello".len()));
}

#[test]
fn end_of_char_data_includes_right_square() {
    assert_eq!("hello]world".end_of_char_data(), Some("hello]world".len()));
}

}
<|endoftext|>


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use chrono::naive::datetime::NaiveDateTime;
use toml_query::delete::TomlValueDeleteExt;
use toml_query::insert::TomlValueInsertExt;
use toml_query::read::TomlValueReadExt;
use toml::Value;

use libimagstore::store::Entry;
use libimagerror::into::IntoError;

use error::DateErrorKind as DEK;
use error::*;
use result::Result;
use range::DateTimeRange;

pub trait EntryDate {

    fn delete_date(&mut self) -> Result<()>;
    fn read_date(&self) -> Result<NaiveDateTime>;
    fn set_date(&mut self, d: NaiveDateTime) -> Result<Option<Result<NaiveDateTime>>>;

    fn delete_date_range(&mut self) -> Result<()>;
    fn read_date_range(&self) -> Result<DateTimeRange>;
    fn set_date_range(&mut self, start: NaiveDateTime, end: NaiveDateTime) -> Result<Option<Result<DateTimeRange>>>;

}

lazy_static! {
    static ref DATE_HEADER_LOCATION : String             = String::from("date.value");
    static ref DATE_RANGE_START_HEADER_LOCATION : String = String::from("date.range.start");
    static ref DATE_RANGE_END_HEADER_LOCATION : String   = String::from("date.range.end");
    static ref DATE_FMT : &'static str                   = "%Y-%m-%d %H:%M:%S";
}

impl EntryDate for Entry {

    fn delete_date(&mut self) -> Result<()> {
        self.get_header_mut()
            .delete(&DATE_HEADER_LOCATION)
            .map(|_| ())
            .map_err_into(DEK::DeleteDateError)
    }

    fn read_date(&self) -> Result<NaiveDateTime> {
        self.get_header()
            .read(&DATE_HEADER_LOCATION)
            .map_err_into(DEK::ReadDateError)
            .and_then(|v| {
                match v {
                    &Value::String(ref s) => s.parse::<NaiveDateTime>()
                        .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            })
    }

    /// Set a Date for this entry
    ///
    /// # Return value
    ///
    /// This function returns funny things, I know. But I find it more attractive to be explicit
    /// what failed when here, instead of beeing nice to the user here.
    ///
    /// So here's a list how things are returned:
    ///
    /// - Err(_) if the inserting failed
    /// - Ok(None) if the inserting succeeded and _did not replace an existing value_.
    /// - Ok(Some(Ok(_))) if the inserting succeeded, but replaced an existing value which then got
    /// parsed into a NaiveDateTime object
    /// - Ok(Some(Err(_))) if the inserting succeeded, but replaced an existing value which then
    /// got parsed into a NaiveDateTime object, where the parsing failed for some reason.
    ///
    fn set_date(&mut self, d: NaiveDateTime) -> Result<Option<Result<NaiveDateTime>>> {
        let date = d.format(&DATE_FMT).to_string();

        self.get_header_mut()
            .insert(&DATE_HEADER_LOCATION, Value::String(date))
            .map(|opt| opt.map(|stri| {
                match stri {
                    Value::String(ref s) => s.parse::<NaiveDateTime>()
                                             .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            }))
            .map_err_into(DEK::SetDateError)
    }

    /// Deletes the date range
    ///
    /// # Warning
    ///
    /// First deletes the start, then the end. If the first operation fails, this might leave the
    /// header in an inconsistent state.
    ///
    fn delete_date_range(&mut self) -> Result<()> {
        let _ = try!(self
             .get_header_mut()
            .delete(&DATE_RANGE_START_HEADER_LOCATION)
            .map(|_| ())
            .map_err_into(DEK::DeleteDateTimeRangeError));

        self.get_header_mut()
            .delete(&DATE_RANGE_END_HEADER_LOCATION)
            .map(|_| ())
            .map_err_into(DEK::DeleteDateTimeRangeError)
    }

    fn read_date_range(&self) -> Result<DateTimeRange> {
        let start = try!(self
            .get_header()
            .read(&DATE_RANGE_START_HEADER_LOCATION)
            .map_err_into(DEK::ReadDateTimeRangeError)
            .and_then(|v| {
                match v {
                    &Value::String(ref s) => s.parse::<NaiveDateTime>()
                        .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            }));

        let end = try!(self
            .get_header()
            .read(&DATE_RANGE_START_HEADER_LOCATION)
            .map_err_into(DEK::ReadDateTimeRangeError)
            .and_then(|v| {
                match v {
                    &Value::String(ref s) => s.parse::<NaiveDateTime>()
                        .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            }));

        DateTimeRange::new(start, end)
            .map_err_into(DEK::DateTimeRangeError)
    }

    /// Set the date range
    ///
    /// # Warning
    ///
    /// This first sets the start, then the end. If the first operation fails, this might leave the
    /// header in an inconsistent state.
    ///
    fn set_date_range(&mut self, start: NaiveDateTime, end: NaiveDateTime)
        -> Result<Option<Result<DateTimeRange>>>
    {
        let start = start.format(&DATE_FMT).to_string();
        let end   = end.format(&DATE_FMT).to_string();

        let opt_old_start = try!(self
            .get_header_mut()
            .insert(&DATE_RANGE_START_HEADER_LOCATION, Value::String(start))
            .map(|opt| opt.map(|stri| {
                match stri {
                    Value::String(ref s) => s.parse::<NaiveDateTime>()
                                             .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            }))
            .map_err_into(DEK::SetDateTimeRangeError));

        let opt_old_end = try!(self
            .get_header_mut()
            .insert(&DATE_RANGE_END_HEADER_LOCATION, Value::String(end))
            .map(|opt| opt.map(|stri| {
                match stri {
                    Value::String(ref s) => s.parse::<NaiveDateTime>()
                                             .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            }))
            .map_err_into(DEK::SetDateTimeRangeError));

        match (opt_old_start, opt_old_end) {
            (Some(Ok(old_start)), Some(Ok(old_end))) => {
                let dr = DateTimeRange::new(old_start, old_end)
                    .map_err_into(DEK::DateTimeRangeError);

                Ok(Some(dr))
            },

            (Some(Err(e)), _) => Err(e),
            (_, Some(Err(e))) => Err(e),
            _ => {
                Ok(None)
            },
        }
    }

}

<commit_msg>Add test for testing Entry::set_date()<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use chrono::naive::datetime::NaiveDateTime;
use toml_query::delete::TomlValueDeleteExt;
use toml_query::insert::TomlValueInsertExt;
use toml_query::read::TomlValueReadExt;
use toml::Value;

use libimagstore::store::Entry;
use libimagerror::into::IntoError;

use error::DateErrorKind as DEK;
use error::*;
use result::Result;
use range::DateTimeRange;

pub trait EntryDate {

    fn delete_date(&mut self) -> Result<()>;
    fn read_date(&self) -> Result<NaiveDateTime>;
    fn set_date(&mut self, d: NaiveDateTime) -> Result<Option<Result<NaiveDateTime>>>;

    fn delete_date_range(&mut self) -> Result<()>;
    fn read_date_range(&self) -> Result<DateTimeRange>;
    fn set_date_range(&mut self, start: NaiveDateTime, end: NaiveDateTime) -> Result<Option<Result<DateTimeRange>>>;

}

lazy_static! {
    static ref DATE_HEADER_LOCATION : &'static str              = "date.value";
    static ref DATE_RANGE_START_HEADER_LOCATION : &'static str  = "date.range.start";
    static ref DATE_RANGE_END_HEADER_LOCATION : &'static str    = "date.range.end";
    static ref DATE_FMT : &'static str                          = "%Y-%m-%dT%H:%M:%S";
}

impl EntryDate for Entry {

    fn delete_date(&mut self) -> Result<()> {
        self.get_header_mut()
            .delete(&DATE_HEADER_LOCATION)
            .map(|_| ())
            .map_err_into(DEK::DeleteDateError)
    }

    fn read_date(&self) -> Result<NaiveDateTime> {
        self.get_header()
            .read(&DATE_HEADER_LOCATION)
            .map_err_into(DEK::ReadDateError)
            .and_then(|v| {
                match v {
                    &Value::String(ref s) => s.parse::<NaiveDateTime>()
                        .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            })
    }

    /// Set a Date for this entry
    ///
    /// # Return value
    ///
    /// This function returns funny things, I know. But I find it more attractive to be explicit
    /// what failed when here, instead of beeing nice to the user here.
    ///
    /// So here's a list how things are returned:
    ///
    /// - Err(_) if the inserting failed
    /// - Ok(None) if the inserting succeeded and _did not replace an existing value_.
    /// - Ok(Some(Ok(_))) if the inserting succeeded, but replaced an existing value which then got
    /// parsed into a NaiveDateTime object
    /// - Ok(Some(Err(_))) if the inserting succeeded, but replaced an existing value which then
    /// got parsed into a NaiveDateTime object, where the parsing failed for some reason.
    ///
    fn set_date(&mut self, d: NaiveDateTime) -> Result<Option<Result<NaiveDateTime>>> {
        let date = d.format(&DATE_FMT).to_string();

        self.get_header_mut()
            .insert(&DATE_HEADER_LOCATION, Value::String(date))
            .map(|opt| opt.map(|stri| {
                match stri {
                    Value::String(ref s) => s.parse::<NaiveDateTime>()
                                             .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            }))
            .map_err_into(DEK::SetDateError)
    }

    /// Deletes the date range
    ///
    /// # Warning
    ///
    /// First deletes the start, then the end. If the first operation fails, this might leave the
    /// header in an inconsistent state.
    ///
    fn delete_date_range(&mut self) -> Result<()> {
        let _ = try!(self
             .get_header_mut()
            .delete(&DATE_RANGE_START_HEADER_LOCATION)
            .map(|_| ())
            .map_err_into(DEK::DeleteDateTimeRangeError));

        self.get_header_mut()
            .delete(&DATE_RANGE_END_HEADER_LOCATION)
            .map(|_| ())
            .map_err_into(DEK::DeleteDateTimeRangeError)
    }

    fn read_date_range(&self) -> Result<DateTimeRange> {
        let start = try!(self
            .get_header()
            .read(&DATE_RANGE_START_HEADER_LOCATION)
            .map_err_into(DEK::ReadDateTimeRangeError)
            .and_then(|v| {
                match v {
                    &Value::String(ref s) => s.parse::<NaiveDateTime>()
                        .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            }));

        let end = try!(self
            .get_header()
            .read(&DATE_RANGE_START_HEADER_LOCATION)
            .map_err_into(DEK::ReadDateTimeRangeError)
            .and_then(|v| {
                match v {
                    &Value::String(ref s) => s.parse::<NaiveDateTime>()
                        .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            }));

        DateTimeRange::new(start, end)
            .map_err_into(DEK::DateTimeRangeError)
    }

    /// Set the date range
    ///
    /// # Warning
    ///
    /// This first sets the start, then the end. If the first operation fails, this might leave the
    /// header in an inconsistent state.
    ///
    fn set_date_range(&mut self, start: NaiveDateTime, end: NaiveDateTime)
        -> Result<Option<Result<DateTimeRange>>>
    {
        let start = start.format(&DATE_FMT).to_string();
        let end   = end.format(&DATE_FMT).to_string();

        let opt_old_start = try!(self
            .get_header_mut()
            .insert(&DATE_RANGE_START_HEADER_LOCATION, Value::String(start))
            .map(|opt| opt.map(|stri| {
                match stri {
                    Value::String(ref s) => s.parse::<NaiveDateTime>()
                                             .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            }))
            .map_err_into(DEK::SetDateTimeRangeError));

        let opt_old_end = try!(self
            .get_header_mut()
            .insert(&DATE_RANGE_END_HEADER_LOCATION, Value::String(end))
            .map(|opt| opt.map(|stri| {
                match stri {
                    Value::String(ref s) => s.parse::<NaiveDateTime>()
                                             .map_err_into(DEK::DateTimeParsingError),
                    _ => Err(DEK::DateHeaderFieldTypeError.into_error()),
                }
            }))
            .map_err_into(DEK::SetDateTimeRangeError));

        match (opt_old_start, opt_old_end) {
            (Some(Ok(old_start)), Some(Ok(old_end))) => {
                let dr = DateTimeRange::new(old_start, old_end)
                    .map_err_into(DEK::DateTimeRangeError);

                Ok(Some(dr))
            },

            (Some(Err(e)), _) => Err(e),
            (_, Some(Err(e))) => Err(e),
            _ => {
                Ok(None)
            },
        }
    }

}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;

    use super::*;

    use libimagstore::store::Store;

    use chrono::naive::datetime::NaiveDateTime;
    use chrono::naive::date::NaiveDate;
    use chrono::naive::time::NaiveTime;

    pub fn get_store() -> Store {
        Store::new(PathBuf::from("/"), None).unwrap()
    }

    #[test]
    fn test_set_date() {
        let store = get_store();

        let date = {
            let date = NaiveDate::from_ymd(2000, 01, 02);
            let time = NaiveTime::from_hms(03, 04, 05);

            NaiveDateTime::new(date, time)
        };

        let mut entry   = store.create(PathBuf::from("test")).unwrap();
        let res         = entry.set_date(date);

        assert!(res.is_ok(), format!("Error: {:?}", res));
        let res = res.unwrap();

        assert!(res.is_none()); // There shouldn't be an existing value

        // Check whether the header is set correctly

        let hdr_field = entry.get_header().read(&DATE_HEADER_LOCATION);

        assert!(hdr_field.is_ok());
        let hdr_field = hdr_field.unwrap();

        match *hdr_field {
            Value::String(ref s) => assert_eq!("2000-01-02T03:04:05", s),
            _ => assert!(false, "Wrong header type"),
        }
    }
}

<|endoftext|>