<commit_before>//! This module specifies the input to rust-analyzer. In some sense, this is
//! **the** most important module, because all other fancy stuff is strictly
//! derived from this input.
//!
//! Note that neither this module, nor any other part of the analyzer's core do
//! actual IO. See `vfs` and `project_model` in the `rust-analyzer` crate for how
//! actual IO is done and lowered to input.

use std::{fmt, ops, str::FromStr, sync::Arc};

use ra_cfg::CfgOptions;
use ra_syntax::SmolStr;
use ra_tt::TokenExpander;
use rustc_hash::{FxHashMap, FxHashSet};
use vfs::file_set::FileSet;

pub use vfs::FileId;

/// Files are grouped into source roots. A source root is a directory on the
/// file systems which is watched for changes. Typically it corresponds to a
/// Rust crate. Source roots *might* be nested: in this case, a file belongs to
/// the nearest enclosing source root. Paths to files are always relative to a
/// source root, and the analyzer does not know the root path of the source root at
/// all. So, a file from one source root can't refer to a file in another source
/// root by path.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct SourceRootId(pub u32);

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SourceRoot {
    /// Sysroot or crates.io library.
    ///
    /// Libraries are considered mostly immutable, this assumption is used to
    /// optimize salsa's query structure
    pub is_library: bool,
    pub(crate) file_set: FileSet,
}

impl SourceRoot {
    pub fn new_local(file_set: FileSet) -> SourceRoot {
        SourceRoot { is_library: false, file_set }
    }
    pub fn new_library(file_set: FileSet) -> SourceRoot {
        SourceRoot { is_library: true, file_set }
    }
    pub fn iter(&self) -> impl Iterator<Item = FileId> + '_ {
        self.file_set.iter()
    }
}

/// `CrateGraph` is a bit of information which turns a set of text files into a
/// number of Rust crates. Each crate is defined by the `FileId` of its root module,
/// the set of cfg flags (not yet implemented) and the set of dependencies. Note
/// that, due to cfg's, there might be several crates for a single `FileId`! As
/// in the rust-lang proper, a crate does not have a name. Instead, names are
/// specified on dependency edges. That is, a crate might be known under
/// different names in different dependent crates.
///
/// Note that `CrateGraph` is build-system agnostic: it's a concept of the Rust
/// language proper, not a concept of the build system. In practice, we get
/// `CrateGraph` by lowering `cargo metadata` output.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CrateGraph {
    arena: FxHashMap<CrateId, CrateData>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CrateId(pub u32);

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CrateName(SmolStr);

impl CrateName {
    /// Creates a crate name, checking for dashes in the string provided.
    /// Dashes are not allowed in the crate names,
    /// hence the input string is returned as `Err` for those cases.
    pub fn new(name: &str) -> Result<CrateName, &str> {
        if name.contains('-') {
            Err(name)
        } else {
            Ok(Self(SmolStr::new(name)))
        }
    }

    /// Creates a crate name, unconditionally replacing the dashes with underscores.
    pub fn normalize_dashes(name: &str) -> CrateName {
        Self(SmolStr::new(name.replace('-', "_")))
    }
}

impl fmt::Display for CrateName {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl ops::Deref for CrateName {
    type Target = str;
    fn deref(&self) -> &Self::Target {
        &*self.0
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct ProcMacroId(pub u32);

#[derive(Debug, Clone)]
pub struct ProcMacro {
    pub name: SmolStr,
    pub expander: Arc<dyn TokenExpander>,
}

impl Eq for ProcMacro {}
impl PartialEq for ProcMacro {
    fn eq(&self, other: &ProcMacro) -> bool {
        self.name == other.name && Arc::ptr_eq(&self.expander, &other.expander)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CrateData {
    pub root_file_id: FileId,
    pub edition: Edition,
    /// The name to display to the end user.
    /// This actual crate name can be different in a particular dependent crate
    /// or may even be missing for some cases, such as a dummy crate for the code snippet.
    pub display_name: Option<String>,
    pub cfg_options: CfgOptions,
    pub env: Env,
    pub dependencies: Vec<Dependency>,
    pub proc_macro: Vec<ProcMacro>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Edition {
    Edition2018,
    Edition2015,
}

#[derive(Default, Debug, Clone, PartialEq, Eq)]
pub struct Env {
    entries: FxHashMap<String, String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Dependency {
    pub crate_id: CrateId,
    pub name: CrateName,
}

impl CrateGraph {
    pub fn add_crate_root(
        &mut self,
        file_id: FileId,
        edition: Edition,
        display_name: Option<String>,
        cfg_options: CfgOptions,
        env: Env,
        proc_macro: Vec<(SmolStr, Arc<dyn ra_tt::TokenExpander>)>,
    ) -> CrateId {
        let proc_macro =
            proc_macro.into_iter().map(|(name, it)| ProcMacro { name, expander: it }).collect();

        let data = CrateData {
            root_file_id: file_id,
            edition,
            display_name,
            cfg_options,
            env,
            proc_macro,
            dependencies: Vec::new(),
        };
        let crate_id = CrateId(self.arena.len() as u32);
        let prev = self.arena.insert(crate_id, data);
        assert!(prev.is_none());
        crate_id
    }

    pub fn add_dep(
        &mut self,
        from: CrateId,
        name: CrateName,
        to: CrateId,
    ) -> Result<(), CyclicDependenciesError> {
        if self.dfs_find(from, to, &mut FxHashSet::default()) {
            return Err(CyclicDependenciesError);
        }
        self.arena.get_mut(&from).unwrap().add_dep(name, to);
        Ok(())
    }

    pub fn is_empty(&self) -> bool {
        self.arena.is_empty()
    }

    pub fn iter(&self) -> impl Iterator<Item = CrateId> + '_ {
        self.arena.keys().copied()
    }

    // FIXME: this only finds one crate with the given root; we could have multiple
    pub fn crate_id_for_crate_root(&self, file_id: FileId) -> Option<CrateId> {
        let (&crate_id, _) =
            self.arena.iter().find(|(_crate_id, data)| data.root_file_id == file_id)?;
        Some(crate_id)
    }

    /// Extends this crate graph by adding a complete disjoint second crate
    /// graph.
    ///
    /// The ids of the crates in the `other` graph are shifted by the return
    /// amount.
    pub fn extend(&mut self, other: CrateGraph) -> u32 {
        let start = self.arena.len() as u32;
        self.arena.extend(other.arena.into_iter().map(|(id, mut data)| {
            let new_id = id.shift(start);
            for dep in &mut data.dependencies {
                dep.crate_id = dep.crate_id.shift(start);
            }
            (new_id, data)
        }));
        start
    }

    fn dfs_find(&self, target: CrateId, from: CrateId, visited: &mut FxHashSet<CrateId>) -> bool {
        if !visited.insert(from) {
            return false;
        }

        if target == from {
            return true;
        }

        for dep in &self[from].dependencies {
            let crate_id = dep.crate_id;
            if self.dfs_find(target, crate_id, visited) {
                return true;
            }
        }
        false
    }
}

impl ops::Index<CrateId> for CrateGraph {
    type Output = CrateData;
    fn index(&self, crate_id: CrateId) -> &CrateData {
        &self.arena[&crate_id]
    }
}

impl CrateId {
    pub fn shift(self, amount: u32) -> CrateId {
        CrateId(self.0 + amount)
    }
}

impl CrateData {
    fn add_dep(&mut self, name: CrateName, crate_id: CrateId) {
        self.dependencies.push(Dependency { name, crate_id })
    }
}

impl FromStr for Edition {
    type Err = ParseEditionError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let res = match s {
            "2015" => Edition::Edition2015,
            "2018" => Edition::Edition2018,
            _ => return Err(ParseEditionError { invalid_input: s.to_string() }),
        };
        Ok(res)
    }
}

impl fmt::Display for Edition {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Edition::Edition2015 => "2015",
            Edition::Edition2018 => "2018",
        })
    }
}

impl<'a, T> From<T> for Env
where
    T: Iterator<Item = (&'a String, &'a String)>,
{
    fn from(iter: T) -> Self {
        let mut result = Self::default();

        for (k, v) in iter {
            result.entries.insert(k.to_owned(), v.to_owned());
        }

        result
    }
}

impl Env {
    pub fn set(&mut self, env: &str, value: String) {
        self.entries.insert(env.to_owned(), value);
    }

    pub fn get(&self, env: &str) -> Option<String> {
        self.entries.get(env).cloned()
    }
}

#[derive(Debug)]
pub struct ParseEditionError {
    invalid_input: String,
}

impl fmt::Display for ParseEditionError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "invalid edition: {:?}", self.invalid_input)
    }
}

impl std::error::Error for ParseEditionError {}

#[derive(Debug)]
pub struct CyclicDependenciesError;

#[cfg(test)]
mod tests {
    use super::{CfgOptions, CrateGraph, CrateName, Dependency, Edition::Edition2018, Env, FileId};

    #[test]
    fn detect_cyclic_dependency_indirect() {
        let mut graph = CrateGraph::default();
        let crate1 = graph.add_crate_root(
            FileId(1u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate2 = graph.add_crate_root(
            FileId(2u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate3 = graph.add_crate_root(
            FileId(3u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        assert!(graph.add_dep(crate1, CrateName::new("crate2").unwrap(), crate2).is_ok());
        assert!(graph.add_dep(crate2, CrateName::new("crate3").unwrap(), crate3).is_ok());
        assert!(graph.add_dep(crate3, CrateName::new("crate1").unwrap(), crate1).is_err());
    }

    #[test]
    fn detect_cyclic_dependency_direct() {
        let mut graph = CrateGraph::default();
        let crate1 = graph.add_crate_root(
            FileId(1u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate2 = graph.add_crate_root(
            FileId(2u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        assert!(graph.add_dep(crate1, CrateName::new("crate2").unwrap(), crate2).is_ok());
        assert!(graph.add_dep(crate2, CrateName::new("crate2").unwrap(), crate2).is_err());
    }

    #[test]
    fn it_works() {
        let mut graph = CrateGraph::default();
        let crate1 = graph.add_crate_root(
            FileId(1u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate2 = graph.add_crate_root(
            FileId(2u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate3 = graph.add_crate_root(
            FileId(3u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        assert!(graph.add_dep(crate1, CrateName::new("crate2").unwrap(), crate2).is_ok());
        assert!(graph.add_dep(crate2, CrateName::new("crate3").unwrap(), crate3).is_ok());
    }

    #[test]
    fn dashes_are_normalized() {
        let mut graph = CrateGraph::default();
        let crate1 = graph.add_crate_root(
            FileId(1u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate2 = graph.add_crate_root(
            FileId(2u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        assert!(graph
            .add_dep(crate1, CrateName::normalize_dashes("crate-name-with-dashes"), crate2)
            .is_ok());
        assert_eq!(
            graph[crate1].dependencies,
            vec![Dependency {
                crate_id: crate2,
                name: CrateName::new("crate_name_with_dashes").unwrap()
            }]
        );
    }
}
<commit_msg>Add a transitive deps iterator to `CrateGraph`<commit_after>//! This module specifies the input to rust-analyzer. In some sense, this is
//! **the** most important module, because all other fancy stuff is strictly
//! derived from this input.
//!
//! Note that neither this module, nor any other part of the analyzer's core do
//! actual IO. See `vfs` and `project_model` in the `rust-analyzer` crate for how
//! actual IO is done and lowered to input.

use std::{fmt, ops, str::FromStr, sync::Arc};

use ra_cfg::CfgOptions;
use ra_syntax::SmolStr;
use ra_tt::TokenExpander;
use rustc_hash::{FxHashMap, FxHashSet};
use vfs::file_set::FileSet;

pub use vfs::FileId;

/// Files are grouped into source roots. A source root is a directory on the
/// file systems which is watched for changes. Typically it corresponds to a
/// Rust crate. Source roots *might* be nested: in this case, a file belongs to
/// the nearest enclosing source root. Paths to files are always relative to a
/// source root, and the analyzer does not know the root path of the source root at
/// all. So, a file from one source root can't refer to a file in another source
/// root by path.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct SourceRootId(pub u32);

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SourceRoot {
    /// Sysroot or crates.io library.
    ///
    /// Libraries are considered mostly immutable, this assumption is used to
    /// optimize salsa's query structure
    pub is_library: bool,
    pub(crate) file_set: FileSet,
}

impl SourceRoot {
    pub fn new_local(file_set: FileSet) -> SourceRoot {
        SourceRoot { is_library: false, file_set }
    }
    pub fn new_library(file_set: FileSet) -> SourceRoot {
        SourceRoot { is_library: true, file_set }
    }
    pub fn iter(&self) -> impl Iterator<Item = FileId> + '_ {
        self.file_set.iter()
    }
}

/// `CrateGraph` is a bit of information which turns a set of text files into a
/// number of Rust crates. Each crate is defined by the `FileId` of its root module,
/// the set of cfg flags (not yet implemented) and the set of dependencies. Note
/// that, due to cfg's, there might be several crates for a single `FileId`! As
/// in the rust-lang proper, a crate does not have a name. Instead, names are
/// specified on dependency edges. That is, a crate might be known under
/// different names in different dependent crates.
///
/// Note that `CrateGraph` is build-system agnostic: it's a concept of the Rust
/// language proper, not a concept of the build system. In practice, we get
/// `CrateGraph` by lowering `cargo metadata` output.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CrateGraph {
    arena: FxHashMap<CrateId, CrateData>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CrateId(pub u32);

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CrateName(SmolStr);

impl CrateName {
    /// Creates a crate name, checking for dashes in the string provided.
    /// Dashes are not allowed in the crate names,
    /// hence the input string is returned as `Err` for those cases.
    pub fn new(name: &str) -> Result<CrateName, &str> {
        if name.contains('-') {
            Err(name)
        } else {
            Ok(Self(SmolStr::new(name)))
        }
    }

    /// Creates a crate name, unconditionally replacing the dashes with underscores.
    pub fn normalize_dashes(name: &str) -> CrateName {
        Self(SmolStr::new(name.replace('-', "_")))
    }
}

impl fmt::Display for CrateName {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl ops::Deref for CrateName {
    type Target = str;
    fn deref(&self) -> &Self::Target {
        &*self.0
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct ProcMacroId(pub u32);

#[derive(Debug, Clone)]
pub struct ProcMacro {
    pub name: SmolStr,
    pub expander: Arc<dyn TokenExpander>,
}

impl Eq for ProcMacro {}
impl PartialEq for ProcMacro {
    fn eq(&self, other: &ProcMacro) -> bool {
        self.name == other.name && Arc::ptr_eq(&self.expander, &other.expander)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CrateData {
    pub root_file_id: FileId,
    pub edition: Edition,
    /// The name to display to the end user.
    /// This actual crate name can be different in a particular dependent crate
    /// or may even be missing for some cases, such as a dummy crate for the code snippet.
    pub display_name: Option<String>,
    pub cfg_options: CfgOptions,
    pub env: Env,
    pub dependencies: Vec<Dependency>,
    pub proc_macro: Vec<ProcMacro>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Edition {
    Edition2018,
    Edition2015,
}

#[derive(Default, Debug, Clone, PartialEq, Eq)]
pub struct Env {
    entries: FxHashMap<String, String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Dependency {
    pub crate_id: CrateId,
    pub name: CrateName,
}

impl CrateGraph {
    pub fn add_crate_root(
        &mut self,
        file_id: FileId,
        edition: Edition,
        display_name: Option<String>,
        cfg_options: CfgOptions,
        env: Env,
        proc_macro: Vec<(SmolStr, Arc<dyn ra_tt::TokenExpander>)>,
    ) -> CrateId {
        let proc_macro =
            proc_macro.into_iter().map(|(name, it)| ProcMacro { name, expander: it }).collect();

        let data = CrateData {
            root_file_id: file_id,
            edition,
            display_name,
            cfg_options,
            env,
            proc_macro,
            dependencies: Vec::new(),
        };
        let crate_id = CrateId(self.arena.len() as u32);
        let prev = self.arena.insert(crate_id, data);
        assert!(prev.is_none());
        crate_id
    }

    pub fn add_dep(
        &mut self,
        from: CrateId,
        name: CrateName,
        to: CrateId,
    ) -> Result<(), CyclicDependenciesError> {
        if self.dfs_find(from, to, &mut FxHashSet::default()) {
            return Err(CyclicDependenciesError);
        }
        self.arena.get_mut(&from).unwrap().add_dep(name, to);
        Ok(())
    }

    pub fn is_empty(&self) -> bool {
        self.arena.is_empty()
    }

    pub fn iter(&self) -> impl Iterator<Item = CrateId> + '_ {
        self.arena.keys().copied()
    }

    /// Returns an iterator over all transitive dependencies of the given crate.
    pub fn transitive_deps(&self, of: CrateId) -> impl Iterator<Item = CrateId> + '_ {
        let mut worklist = vec![of];
        let mut deps = FxHashSet::default();

        while let Some(krate) = worklist.pop() {
            if !deps.insert(krate) {
                continue;
            }

            worklist.extend(self[krate].dependencies.iter().map(|dep| dep.crate_id));
        }

        deps.remove(&of);
        deps.into_iter()
    }

    // FIXME: this only finds one crate with the given root; we could have multiple
    pub fn crate_id_for_crate_root(&self, file_id: FileId) -> Option<CrateId> {
        let (&crate_id, _) =
            self.arena.iter().find(|(_crate_id, data)| data.root_file_id == file_id)?;
        Some(crate_id)
    }

    /// Extends this crate graph by adding a complete disjoint second crate
    /// graph.
    ///
    /// The ids of the crates in the `other` graph are shifted by the return
    /// amount.
    pub fn extend(&mut self, other: CrateGraph) -> u32 {
        let start = self.arena.len() as u32;
        self.arena.extend(other.arena.into_iter().map(|(id, mut data)| {
            let new_id = id.shift(start);
            for dep in &mut data.dependencies {
                dep.crate_id = dep.crate_id.shift(start);
            }
            (new_id, data)
        }));
        start
    }

    fn dfs_find(&self, target: CrateId, from: CrateId, visited: &mut FxHashSet<CrateId>) -> bool {
        if !visited.insert(from) {
            return false;
        }

        if target == from {
            return true;
        }

        for dep in &self[from].dependencies {
            let crate_id = dep.crate_id;
            if self.dfs_find(target, crate_id, visited) {
                return true;
            }
        }
        false
    }
}

impl ops::Index<CrateId> for CrateGraph {
    type Output = CrateData;
    fn index(&self, crate_id: CrateId) -> &CrateData {
        &self.arena[&crate_id]
    }
}

impl CrateId {
    pub fn shift(self, amount: u32) -> CrateId {
        CrateId(self.0 + amount)
    }
}

impl CrateData {
    fn add_dep(&mut self, name: CrateName, crate_id: CrateId) {
        self.dependencies.push(Dependency { name, crate_id })
    }
}

impl FromStr for Edition {
    type Err = ParseEditionError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let res = match s {
            "2015" => Edition::Edition2015,
            "2018" => Edition::Edition2018,
            _ => return Err(ParseEditionError { invalid_input: s.to_string() }),
        };
        Ok(res)
    }
}

impl fmt::Display for Edition {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Edition::Edition2015 => "2015",
            Edition::Edition2018 => "2018",
        })
    }
}

impl<'a, T> From<T> for Env
where
    T: Iterator<Item = (&'a String, &'a String)>,
{
    fn from(iter: T) -> Self {
        let mut result = Self::default();

        for (k, v) in iter {
            result.entries.insert(k.to_owned(), v.to_owned());
        }

        result
    }
}

impl Env {
    pub fn set(&mut self, env: &str, value: String) {
        self.entries.insert(env.to_owned(), value);
    }

    pub fn get(&self, env: &str) -> Option<String> {
        self.entries.get(env).cloned()
    }
}

#[derive(Debug)]
pub struct ParseEditionError {
    invalid_input: String,
}

impl fmt::Display for ParseEditionError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "invalid edition: {:?}", self.invalid_input)
    }
}

impl std::error::Error for ParseEditionError {}

#[derive(Debug)]
pub struct CyclicDependenciesError;

#[cfg(test)]
mod tests {
    use super::{CfgOptions, CrateGraph, CrateName, Dependency, Edition::Edition2018, Env, FileId};

    #[test]
    fn detect_cyclic_dependency_indirect() {
        let mut graph = CrateGraph::default();
        let crate1 = graph.add_crate_root(
            FileId(1u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate2 = graph.add_crate_root(
            FileId(2u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate3 = graph.add_crate_root(
            FileId(3u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        assert!(graph.add_dep(crate1, CrateName::new("crate2").unwrap(), crate2).is_ok());
        assert!(graph.add_dep(crate2, CrateName::new("crate3").unwrap(), crate3).is_ok());
        assert!(graph.add_dep(crate3, CrateName::new("crate1").unwrap(), crate1).is_err());
    }

    #[test]
    fn detect_cyclic_dependency_direct() {
        let mut graph = CrateGraph::default();
        let crate1 = graph.add_crate_root(
            FileId(1u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate2 = graph.add_crate_root(
            FileId(2u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        assert!(graph.add_dep(crate1, CrateName::new("crate2").unwrap(), crate2).is_ok());
        assert!(graph.add_dep(crate2, CrateName::new("crate2").unwrap(), crate2).is_err());
    }

    #[test]
    fn it_works() {
        let mut graph = CrateGraph::default();
        let crate1 = graph.add_crate_root(
            FileId(1u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate2 = graph.add_crate_root(
            FileId(2u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate3 = graph.add_crate_root(
            FileId(3u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        assert!(graph.add_dep(crate1, CrateName::new("crate2").unwrap(), crate2).is_ok());
        assert!(graph.add_dep(crate2, CrateName::new("crate3").unwrap(), crate3).is_ok());
    }

    #[test]
    fn dashes_are_normalized() {
        let mut graph = CrateGraph::default();
        let crate1 = graph.add_crate_root(
            FileId(1u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        let crate2 = graph.add_crate_root(
            FileId(2u32),
            Edition2018,
            None,
            CfgOptions::default(),
            Env::default(),
            Default::default(),
        );
        assert!(graph
            .add_dep(crate1, CrateName::normalize_dashes("crate-name-with-dashes"), crate2)
            .is_ok());
        assert_eq!(
            graph[crate1].dependencies,
            vec![Dependency {
                crate_id: crate2,
                name: CrateName::new("crate_name_with_dashes").unwrap()
            }]
        );
    }
}
<|endoftext|>


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

use super::*;

pub(crate) fn prefix_encode(prefix: &[u8], buf: &[u8]) -> Vec<u8> {
    assert!(
        prefix <= buf,
        "prefix {:?} must be lexicographically <= to the encoded buf {:?}",
        prefix,
        buf
    );
    let limit = std::cmp::min(std::u8::MAX as usize, buf.len());
    let mut prefix_len = 0usize;
    for (i, c) in prefix.iter().take(limit).enumerate() {
        if buf[i] == *c {
            prefix_len += 1;
        } else {
            break;
        }
    }

    let encoded_len = 1 + buf.len() - prefix_len;
    let mut ret = Vec::with_capacity(encoded_len);
    unsafe {
        ret.set_len(encoded_len);
    }
    ret[1..].copy_from_slice(&buf[prefix_len..]);
    ret[0] = prefix_len as u8;
    ret
}

pub(crate) fn prefix_decode(prefix: &[u8], buf: &[u8]) -> Vec<u8> {
    assert!(buf.len() >= 1);
    let prefix_len = buf[0] as usize;
    let mut ret = Vec::with_capacity(prefix_len + buf.len() - 1);
    unsafe {
        ret.set_len(prefix_len + buf.len() - 1);
    }
    ret[0..prefix_len].copy_from_slice(&prefix[0..prefix_len]);
    ret[prefix_len..].copy_from_slice(&buf[1..]);
    ret
}

// NB: the correctness of this function depends on
// the invariant that the prefix is ALWAYS lexicographically
// Less than or Equal to the keys that have been encoded
// using it. Otherwise this comparison would make no sense.
pub(crate) fn prefix_cmp(a: &[u8], b: &[u8]) -> Ordering {
    if a.is_empty() && b.is_empty() {
        return Ordering::Equal;
    } else if a.is_empty() && !b.is_empty() {
        return Ordering::Less;
    } else if !a.is_empty() && b.is_empty() {
        return Ordering::Greater;
    }

    if a[0] > b[0] {
        return Ordering::Less;
    } else if a[0] < b[0] {
        return Ordering::Greater;
    } else {
        return a[1..].cmp(&b[1..]);
    }
}

/// Compare `a` and `b`, assuming that `a` is prefix encoded and `b` is not.
pub fn prefix_cmp_encoded(a: &[u8], mut b: &[u8], mut prefix: &[u8]) -> Ordering {
    assert!(a.len() >= 1 && a[0] as usize <= prefix.len());

    let mut a_prefix_len = a[0];
    let a_suffix = &a[1..];

    while a_prefix_len > 0 {
        if b.is_empty() {
            return Ordering::Greater;
        } else if prefix[0] > b[0] {
            return Ordering::Greater;
        } else if prefix[0] < b[0] {
            return Ordering::Less;
        }

        a_prefix_len -= 1;
        b = &b[1..];
        prefix = &prefix[1..];
    }

    a_suffix.cmp(b)
}

#[test]
fn test_prefix() {
    let prefix = b"cat";
    assert_eq!(
        prefix_encode(prefix, prefix),
        vec![prefix.len() as u8]
    );
    assert_eq!(prefix_encode(prefix, b"catt"), vec![3, b't']);
    assert_eq!(prefix_encode(prefix, b"cb"), vec![1, b'b']);
    assert_eq!(prefix_encode(prefix, b"caz"), vec![2, b'z']);
    assert_eq!(
        prefix_encode(prefix, b"cvar"),
        vec![1, b'v', b'a', b'r']
    );
    assert_eq!(
        prefix_encode(prefix, b"zig"),
        vec![0, b'z', b'i', b'g']
    );

    let prefix = b"";
    assert_eq!(
        prefix_encode(prefix, prefix),
        vec![prefix.len() as u8]
    );
    assert_eq!(prefix_encode(prefix, b"ca"), vec![0, b'c', b'a']);
    assert_eq!(
        prefix_encode(prefix, b"cat"),
        vec![0, b'c', b'a', b't']
    );
    assert_eq!(
        prefix_encode(prefix, b"cab"),
        vec![0, b'c', b'a', b'b']
    );
    assert_eq!(
        prefix_encode(prefix, b"cvar"),
        vec![0, b'c', b'v', b'a', b'r']
    );
    assert_eq!(
        prefix_encode(prefix, b"zig"),
        vec![0, b'z', b'i', b'g']
    );

    let rtt =
        vec![b"" as &[u8], b"\x00cat", b"\x00", b"oyunwytounw\x00"];
    for item in rtt {
        assert_eq!(
            prefix_decode(prefix, &*prefix_encode(prefix, item)),
            item.to_vec()
        );
    }
}

#[test]
fn test_prefix_cmp() {
    assert_eq!(prefix_cmp(&[], &[]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[0], &[]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[], &[0]), Ordering::Less);

    assert_eq!(prefix_cmp(&[3], &[4]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[4, 3], &[3, 4]), Ordering::Less);

    assert_eq!(prefix_cmp(&[1], &[1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 3], &[1, 1]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 3]), Ordering::Less);
}

#[test]
fn test_prefix_cmp_encoded() {
    fn assert_pce(a: &[u8], b: &[u8], prefix: &[u8], expected: Ordering) {
        assert_eq!(prefix_cmp_encoded(a, &prefix_decode(prefix, b), prefix), expected);
    }

    let prefix = &vec![3, 3, 3, 3];
    assert_pce(&[0], &[0], prefix, Ordering::Equal);
    assert_pce(&[1], &[0], prefix, Ordering::Greater);
    assert_pce(&[0], &[1], prefix, Ordering::Less);

    assert_pce(&[3], &[2], prefix, Ordering::Greater);
    assert_pce(&[4, 3], &[3, 4], prefix, Ordering::Less);

    assert_pce(&[1], &[1], prefix, Ordering::Equal);
    assert_pce(&[1, 3, 3, 1], &[3, 1], prefix, Ordering::Equal);
    assert_pce(&[1, 3], &[1, 1], prefix, Ordering::Greater);
    assert_pce(&[1, 1], &[3, 3], prefix, Ordering::Less);
}
<commit_msg>Update prefix.rs<commit_after>use std::cmp::Ordering;

use super::*;

pub(crate) fn prefix_encode(prefix: &[u8], buf: &[u8]) -> Vec<u8> {
    assert!(
        prefix <= buf,
        "prefix {:?} must be lexicographically <= to the encoded buf {:?}",
        prefix,
        buf
    );
    let limit = std::cmp::min(std::u8::MAX as usize, buf.len());
    let mut prefix_len = 0usize;
    for (i, c) in prefix.iter().take(limit).enumerate() {
        if buf[i] == *c {
            prefix_len += 1;
        } else {
            break;
        }
    }

    let encoded_len = 1 + buf.len() - prefix_len;
    let mut ret = Vec::with_capacity(encoded_len);
    unsafe {
        ret.set_len(encoded_len);
    }
    ret[1..].copy_from_slice(&buf[prefix_len..]);
    ret[0] = prefix_len as u8;
    ret
}

pub(crate) fn prefix_decode(prefix: &[u8], buf: &[u8]) -> Vec<u8> {
    assert!(buf.len() >= 1);
    let prefix_len = buf[0] as usize;
    let mut ret = Vec::with_capacity(prefix_len + buf.len() - 1);
    unsafe {
        ret.set_len(prefix_len + buf.len() - 1);
    }
    ret[0..prefix_len].copy_from_slice(&prefix[0..prefix_len]);
    ret[prefix_len..].copy_from_slice(&buf[1..]);
    ret
}

// NB: the correctness of this function depends on
// the invariant that the prefix is ALWAYS lexicographically
// Less than or Equal to the keys that have been encoded
// using it. Otherwise this comparison would make no sense.
pub(crate) fn prefix_cmp(a: &[u8], b: &[u8]) -> Ordering {
    if a.is_empty() && b.is_empty() {
        return Ordering::Equal;
    } else if a.is_empty() && !b.is_empty() {
        return Ordering::Less;
    } else if !a.is_empty() && b.is_empty() {
        return Ordering::Greater;
    }

    if a[0] > b[0] {
        return Ordering::Less;
    } else if a[0] < b[0] {
        return Ordering::Greater;
    } else {
        return a[1..].cmp(&b[1..]);
    }
}

/// Compare `a` and `b`, assuming that `a` is prefix encoded and `b` is not.
pub(crate) fn prefix_cmp_encoded(a: &[u8], mut b: &[u8], mut prefix: &[u8]) -> Ordering {
    assert!(a.len() >= 1 && a[0] as usize <= prefix.len());

    let mut a_prefix_len = a[0];
    let a_suffix = &a[1..];

    while a_prefix_len > 0 {
        if b.is_empty() {
            return Ordering::Greater;
        } else if prefix[0] > b[0] {
            return Ordering::Greater;
        } else if prefix[0] < b[0] {
            return Ordering::Less;
        }

        a_prefix_len -= 1;
        b = &b[1..];
        prefix = &prefix[1..];
    }

    a_suffix.cmp(b)
}

#[test]
fn test_prefix() {
    let prefix = b"cat";
    assert_eq!(
        prefix_encode(prefix, prefix),
        vec![prefix.len() as u8]
    );
    assert_eq!(prefix_encode(prefix, b"catt"), vec![3, b't']);
    assert_eq!(prefix_encode(prefix, b"cb"), vec![1, b'b']);
    assert_eq!(prefix_encode(prefix, b"caz"), vec![2, b'z']);
    assert_eq!(
        prefix_encode(prefix, b"cvar"),
        vec![1, b'v', b'a', b'r']
    );
    assert_eq!(
        prefix_encode(prefix, b"zig"),
        vec![0, b'z', b'i', b'g']
    );

    let prefix = b"";
    assert_eq!(
        prefix_encode(prefix, prefix),
        vec![prefix.len() as u8]
    );
    assert_eq!(prefix_encode(prefix, b"ca"), vec![0, b'c', b'a']);
    assert_eq!(
        prefix_encode(prefix, b"cat"),
        vec![0, b'c', b'a', b't']
    );
    assert_eq!(
        prefix_encode(prefix, b"cab"),
        vec![0, b'c', b'a', b'b']
    );
    assert_eq!(
        prefix_encode(prefix, b"cvar"),
        vec![0, b'c', b'v', b'a', b'r']
    );
    assert_eq!(
        prefix_encode(prefix, b"zig"),
        vec![0, b'z', b'i', b'g']
    );

    let rtt =
        vec![b"" as &[u8], b"\x00cat", b"\x00", b"oyunwytounw\x00"];
    for item in rtt {
        assert_eq!(
            prefix_decode(prefix, &*prefix_encode(prefix, item)),
            item.to_vec()
        );
    }
}

#[test]
fn test_prefix_cmp() {
    assert_eq!(prefix_cmp(&[], &[]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[0], &[]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[], &[0]), Ordering::Less);

    assert_eq!(prefix_cmp(&[3], &[4]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[4, 3], &[3, 4]), Ordering::Less);

    assert_eq!(prefix_cmp(&[1], &[1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 3], &[1, 1]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 3]), Ordering::Less);
}

#[test]
fn test_prefix_cmp_encoded() {
    fn assert_pce(a: &[u8], b: &[u8], prefix: &[u8], expected: Ordering) {
        assert_eq!(prefix_cmp_encoded(a, &prefix_decode(prefix, b), prefix), expected);
    }

    let prefix = &vec![3, 3, 3, 3];
    assert_pce(&[0], &[0], prefix, Ordering::Equal);
    assert_pce(&[1], &[0], prefix, Ordering::Greater);
    assert_pce(&[0], &[1], prefix, Ordering::Less);

    assert_pce(&[3], &[2], prefix, Ordering::Greater);
    assert_pce(&[4, 3], &[3, 4], prefix, Ordering::Less);

    assert_pce(&[1], &[1], prefix, Ordering::Equal);
    assert_pce(&[1, 3, 3, 1], &[3, 1], prefix, Ordering::Equal);
    assert_pce(&[1, 3], &[1, 1], prefix, Ordering::Greater);
    assert_pce(&[1, 1], &[3, 3], prefix, Ordering::Less);
}
<|endoftext|>


<commit_before>
use libc::{c_int, c_char};
use std::ffi::{CString, CStr};
use std::borrow::{Cow};
use std::ptr;
use std::fs;
use std::fs::{File, OpenOptions};
use std::io;
use std::io::SeekFrom;
use std::io::prelude::*;
use std::process::{Command, Stdio};
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::fmt;
use openssl;
use bindings::*;
use openssl::sign::Verifier;
use openssl::pkey::PKey;
use openssl::hash::MessageDigest;
use toml;
use toml::Value;
use uuid::Uuid;

#[no_mangle]
pub fn pam_sm_authenticate(
    pamh: *mut pam_handle_t,
    flags: c_int,
    _: c_int,
    _: *mut *const c_char,
) -> c_int {
    authenticate_via_hello(pamh).unwrap_or_else(|err| {
        if (flags & PAM_SILENT) == 0 {
            println!("WSL Hello error: {}", err);
        }
        match err {
            HelloAuthenticationError::PublicKeyFileError(ref err)
                if err.kind() == io::ErrorKind::NotFound  => PAM_USER_UNKNOWN,
            HelloAuthenticationError::AuthenticatorLaunchError(_) => PAM_AUTHINFO_UNAVAIL,
                HelloAuthenticationError::AuthenticatorConnectionError(_) => PAM_AUTHINFO_UNAVAIL,
                HelloAuthenticationError::AuthenticatorSignalled => PAM_AUTHINFO_UNAVAIL,
                _ => PAM_AUTH_ERR,
        }
    })
}

fn get_user(pamh: *mut pam_handle_t, prompt: Option<&str>) -> Result<Cow<str>, i32> {
    let mut c_user: *const c_char = ptr::null();
    let c_prompt = match prompt {
        Some(prompt_str) => CString::new(prompt_str).unwrap().as_ptr(),
        None => ptr::null(),
    };
    let err;
    unsafe {
        err = pam_get_user(pamh, &mut c_user, c_prompt);
    }
    match err {
        PAM_SUCCESS => unsafe {
            let c_user_str = CStr::from_ptr(c_user);
            Ok(c_user_str.to_string_lossy())
        },
        err => Err(err),
    }
}

#[derive(Debug)]
enum ConfigError {
    Io(io::Error),
    Toml(toml::de::Error),
    MissingField(String),
    InvalidValueType(String),
}

impl From<io::Error> for ConfigError {
    fn from(err: io::Error) -> ConfigError {
        ConfigError::Io(err)
    }
}

impl From<toml::de::Error> for ConfigError {
    fn from(err: toml::de::Error) -> ConfigError {
        ConfigError::Toml(err)
    }
}

impl fmt::Display for ConfigError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ConfigError::Io(ref ioerr) => write!(f, "{}", ioerr),
            ConfigError::Toml(_) => write!(f, "TOML format error"),
            ConfigError::MissingField(_) => write!(f, "field: 'authenticator_path' is not found"),
            ConfigError::InvalidValueType(_) => write!(f, "field: 'authenticator_path' has an invalid value type"),
        }
    }
}

fn get_authenticator_path() -> Result<String, ConfigError> {
    let mut config_file = File::open("/etc/pam_wsl_hello/config")?;
    let mut config = String::new();
    config_file.read_to_string(&mut config)?;

    let config_value = config.parse::<Value>()?;
    let authenticator_path = config_value.get("authenticator_path")
        .ok_or(ConfigError::MissingField("authenticator_path".to_owned()))?
        .as_str()
        .ok_or(ConfigError::InvalidValueType("authenticator_path".to_owned()))?;
    Ok(authenticator_path.to_owned())
}

#[derive(Debug)]
enum HelloAuthenticationError {
    GetUserError(i32),
    ConfigError(ConfigError),
    PublicKeyFileError(io::Error),
    Io(io::Error),
    InvalidPublicKey(openssl::error::ErrorStack),
    OpenSSLError(openssl::error::ErrorStack),
    AuthenticatorLaunchError(io::Error),
    AuthenticatorConnectionError(io::Error),
    AuthenticatorSignalled,
    HelloAuthenticationFail(String),
    SignAuthenticationFail,
}

impl From<io::Error> for HelloAuthenticationError {
    fn from(err: io::Error) -> HelloAuthenticationError {
        HelloAuthenticationError::Io(err)
    }
}

impl From<ConfigError> for HelloAuthenticationError {
    fn from(err: ConfigError) -> HelloAuthenticationError {
        HelloAuthenticationError::ConfigError(err)
    }
}

impl fmt::Display for HelloAuthenticationError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            HelloAuthenticationError::ConfigError(ref err) => {
                write!(f, "config error; {}", err)
            },
            HelloAuthenticationError::PublicKeyFileError(ref err) => {
                match err.kind() {
                   io::ErrorKind::NotFound => {
                       write!(f, "cannot find the credential public key for this user")
                   },
                   _ => {
                       write!(f, "{}", err)
                   },
                }
            },
            HelloAuthenticationError::Io(ref err) => {
                write!(f, "{}", err)
            },
            HelloAuthenticationError::InvalidPublicKey(_) => {
                write!(f, "the pem file of the public key is invalid")
            },
            HelloAuthenticationError::AuthenticatorLaunchError(ref err) => {
                write!(f, "cannot launch Windows Hello; {}", err)
            },
            HelloAuthenticationError::AuthenticatorConnectionError(ref err) => {
                write!(f, "cannot communicate with Windows Hello; {}", err)
            },
            HelloAuthenticationError::HelloAuthenticationFail(ref msg) => {
                write!(f, "authentication failed; {}", msg)
            },
            HelloAuthenticationError::SignAuthenticationFail => {
                write!(f, "the result of signature verification of the credential is failure")
            },
            ref err => {
                write!(f, "internal error; {:?}", err)
            }
        }
    }
}

fn authenticate_via_hello(pamh: *mut pam_handle_t) -> Result<i32, HelloAuthenticationError> {

    let user_name = get_user(pamh, None).map_err(|e| HelloAuthenticationError::GetUserError(e))?;
    let credential_key_name = format!("pam_wsl_hello_{}", user_name);

    let mut hello_public_key_file = File::open(format!("/etc/pam_wsl_hello/public_keys/{}.pem", credential_key_name)).map_err(|io| HelloAuthenticationError::PublicKeyFileError(io))?;
    let mut key_str = String::new();
    hello_public_key_file.read_to_string(&mut key_str)?;
    let hello_public_key = PKey::public_key_from_pem(key_str.as_bytes()).map_err(|e| HelloAuthenticationError::InvalidPublicKey(e))?;

    let challenge = format!("pam_wsl_hello:{}:{}", user_name, Uuid::new_v4());

    let auth_res;
    let challenge_tmpfile_path = &format!("/tmp/{}", challenge);
    {
        // Since there seems to be a bug that C# applications cannot read from pipes on WSL,
        // we create a temporary file to redirect
        let mut challenge_tmpfile = OpenOptions::new().write(true)
            .read(true)
            .create_new(true)
            .open(challenge_tmpfile_path)?;
        challenge_tmpfile.write_all(challenge.as_bytes())?;
        challenge_tmpfile.seek(SeekFrom::Start(0))?;
        let challenge_tmpfile_in = unsafe {Stdio::from_raw_fd(challenge_tmpfile.as_raw_fd())};

        let authenticator_path = get_authenticator_path()?;
        let authenticator = Command::new(&authenticator_path)
            .arg(credential_key_name)
            .current_dir("/mnt/c")
            .stdin(challenge_tmpfile_in)
            .stdout(Stdio::piped())
            .spawn()
            .map_err(|e| HelloAuthenticationError::AuthenticatorLaunchError(e))?;

        auth_res = authenticator
            .wait_with_output()
            .map_err(|e| HelloAuthenticationError::AuthenticatorConnectionError(e))?;
    }
    fs::remove_file(challenge_tmpfile_path)?;

    match auth_res.status.code() {
        Some(code) if code == 0 => {/* Success */},
        Some(_) => return Err(HelloAuthenticationError::HelloAuthenticationFail(String::from_utf8(auth_res.stdout)
                                                                                .unwrap_or("invalid utf8 output".to_string()))),
        None => return Err(HelloAuthenticationError::AuthenticatorSignalled),
    }
    let signature = auth_res.stdout;

    let mut verifier = Verifier::new(MessageDigest::sha256(), &hello_public_key).unwrap();
    verifier.update(challenge.as_bytes()).map_err(|e| HelloAuthenticationError::OpenSSLError(e))?;

    match verifier.finish(&signature).map_err(|e| HelloAuthenticationError::OpenSSLError(e))? {
        true => Ok(PAM_SUCCESS),
        false => Err(HelloAuthenticationError::SignAuthenticationFail),
    }
}
<commit_msg>Format files by rstfmt<commit_after>
use libc::{c_char, c_int};
use std::ffi::{CStr, CString};
use std::borrow::Cow;
use std::ptr;
use std::fs;
use std::fs::{File, OpenOptions};
use std::io;
use std::io::SeekFrom;
use std::io::prelude::*;
use std::process::{Command, Stdio};
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::fmt;
use openssl;
use bindings::*;
use openssl::sign::Verifier;
use openssl::pkey::PKey;
use openssl::hash::MessageDigest;
use toml;
use toml::Value;
use uuid::Uuid;

#[no_mangle]
pub fn pam_sm_authenticate(
    pamh: *mut pam_handle_t,
    flags: c_int,
    _: c_int,
    _: *mut *const c_char,
) -> c_int {
    authenticate_via_hello(pamh).unwrap_or_else(|err| {
        if (flags & PAM_SILENT) == 0 {
            println!("WSL Hello error: {}", err);
        }
        match err {
            HelloAuthenticationError::PublicKeyFileError(ref err)
                if err.kind() == io::ErrorKind::NotFound =>
            {
                PAM_USER_UNKNOWN
            }
            HelloAuthenticationError::AuthenticatorLaunchError(_) => PAM_AUTHINFO_UNAVAIL,
            HelloAuthenticationError::AuthenticatorConnectionError(_) => PAM_AUTHINFO_UNAVAIL,
            HelloAuthenticationError::AuthenticatorSignalled => PAM_AUTHINFO_UNAVAIL,
            _ => PAM_AUTH_ERR,
        }
    })
}

fn get_user(pamh: *mut pam_handle_t, prompt: Option<&str>) -> Result<Cow<str>, i32> {
    let mut c_user: *const c_char = ptr::null();
    let c_prompt = match prompt {
        Some(prompt_str) => CString::new(prompt_str).unwrap().as_ptr(),
        None => ptr::null(),
    };
    let err;
    unsafe {
        err = pam_get_user(pamh, &mut c_user, c_prompt);
    }
    match err {
        PAM_SUCCESS => unsafe {
            let c_user_str = CStr::from_ptr(c_user);
            Ok(c_user_str.to_string_lossy())
        },
        err => Err(err),
    }
}

#[derive(Debug)]
enum ConfigError {
    Io(io::Error),
    Toml(toml::de::Error),
    MissingField(String),
    InvalidValueType(String),
}

impl From<io::Error> for ConfigError {
    fn from(err: io::Error) -> ConfigError {
        ConfigError::Io(err)
    }
}

impl From<toml::de::Error> for ConfigError {
    fn from(err: toml::de::Error) -> ConfigError {
        ConfigError::Toml(err)
    }
}

impl fmt::Display for ConfigError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ConfigError::Io(ref ioerr) => write!(f, "{}", ioerr),
            ConfigError::Toml(_) => write!(f, "TOML format error"),
            ConfigError::MissingField(_) => write!(f, "field: 'authenticator_path' is not found"),
            ConfigError::InvalidValueType(_) => {
                write!(f, "field: 'authenticator_path' has an invalid value type")
            }
        }
    }
}

fn get_authenticator_path() -> Result<String, ConfigError> {
    let mut config_file = File::open("/etc/pam_wsl_hello/config")?;
    let mut config = String::new();
    config_file.read_to_string(&mut config)?;

    let config_value = config.parse::<Value>()?;
    let authenticator_path = config_value
        .get("authenticator_path")
        .ok_or(ConfigError::MissingField("authenticator_path".to_owned()))?
        .as_str()
        .ok_or(ConfigError::InvalidValueType(
            "authenticator_path".to_owned(),
        ))?;
    Ok(authenticator_path.to_owned())
}

#[derive(Debug)]
enum HelloAuthenticationError {
    GetUserError(i32),
    ConfigError(ConfigError),
    PublicKeyFileError(io::Error),
    Io(io::Error),
    InvalidPublicKey(openssl::error::ErrorStack),
    OpenSSLError(openssl::error::ErrorStack),
    AuthenticatorLaunchError(io::Error),
    AuthenticatorConnectionError(io::Error),
    AuthenticatorSignalled,
    HelloAuthenticationFail(String),
    SignAuthenticationFail,
}

impl From<io::Error> for HelloAuthenticationError {
    fn from(err: io::Error) -> HelloAuthenticationError {
        HelloAuthenticationError::Io(err)
    }
}

impl From<ConfigError> for HelloAuthenticationError {
    fn from(err: ConfigError) -> HelloAuthenticationError {
        HelloAuthenticationError::ConfigError(err)
    }
}

impl fmt::Display for HelloAuthenticationError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            HelloAuthenticationError::ConfigError(ref err) => write!(f, "config error; {}", err),
            HelloAuthenticationError::PublicKeyFileError(ref err) => match err.kind() {
                io::ErrorKind::NotFound => {
                    write!(f, "cannot find the credential public key for this user")
                }
                _ => write!(f, "{}", err),
            },
            HelloAuthenticationError::Io(ref err) => write!(f, "{}", err),
            HelloAuthenticationError::InvalidPublicKey(_) => {
                write!(f, "the pem file of the public key is invalid")
            }
            HelloAuthenticationError::AuthenticatorLaunchError(ref err) => {
                write!(f, "cannot launch Windows Hello; {}", err)
            }
            HelloAuthenticationError::AuthenticatorConnectionError(ref err) => {
                write!(f, "cannot communicate with Windows Hello; {}", err)
            }
            HelloAuthenticationError::HelloAuthenticationFail(ref msg) => {
                write!(f, "authentication failed; {}", msg)
            }
            HelloAuthenticationError::SignAuthenticationFail => write!(
                f,
                "the result of signature verification of the credential is failure"
            ),
            ref err => write!(f, "internal error; {:?}", err),
        }
    }
}

fn authenticate_via_hello(pamh: *mut pam_handle_t) -> Result<i32, HelloAuthenticationError> {
    let user_name = get_user(pamh, None).map_err(|e| HelloAuthenticationError::GetUserError(e))?;
    let credential_key_name = format!("pam_wsl_hello_{}", user_name);

    let mut hello_public_key_file =
        File::open(format!(
            "/etc/pam_wsl_hello/public_keys/{}.pem",
            credential_key_name
        )).map_err(|io| HelloAuthenticationError::PublicKeyFileError(io))?;
    let mut key_str = String::new();
    hello_public_key_file.read_to_string(&mut key_str)?;
    let hello_public_key = PKey::public_key_from_pem(key_str.as_bytes())
        .map_err(|e| HelloAuthenticationError::InvalidPublicKey(e))?;

    let challenge = format!("pam_wsl_hello:{}:{}", user_name, Uuid::new_v4());

    let auth_res;
    let challenge_tmpfile_path = &format!("/tmp/{}", challenge);
    {
        // Since there seems to be a bug that C# applications cannot read from pipes on WSL,
        // we create a temporary file to redirect
        let mut challenge_tmpfile = OpenOptions::new()
            .write(true)
            .read(true)
            .create_new(true)
            .open(challenge_tmpfile_path)?;
        challenge_tmpfile.write_all(challenge.as_bytes())?;
        challenge_tmpfile.seek(SeekFrom::Start(0))?;
        let challenge_tmpfile_in = unsafe { Stdio::from_raw_fd(challenge_tmpfile.as_raw_fd()) };

        let authenticator_path = get_authenticator_path()?;
        let authenticator = Command::new(&authenticator_path)
            .arg(credential_key_name)
            .current_dir("/mnt/c")
            .stdin(challenge_tmpfile_in)
            .stdout(Stdio::piped())
            .spawn()
            .map_err(|e| HelloAuthenticationError::AuthenticatorLaunchError(e))?;

        auth_res = authenticator.wait_with_output().map_err(|e| {
            HelloAuthenticationError::AuthenticatorConnectionError(e)
        })?;
    }
    fs::remove_file(challenge_tmpfile_path)?;

    match auth_res.status.code() {
        Some(code) if code == 0 => { /* Success */ }
        Some(_) => {
            return Err(HelloAuthenticationError::HelloAuthenticationFail(
                String::from_utf8(auth_res.stdout).unwrap_or("invalid utf8 output".to_string()),
            ))
        }
        None => return Err(HelloAuthenticationError::AuthenticatorSignalled),
    }
    let signature = auth_res.stdout;

    let mut verifier = Verifier::new(MessageDigest::sha256(), &hello_public_key).unwrap();
    verifier
        .update(challenge.as_bytes())
        .map_err(|e| HelloAuthenticationError::OpenSSLError(e))?;

    match verifier
        .finish(&signature)
        .map_err(|e| HelloAuthenticationError::OpenSSLError(e))?
    {
        true => Ok(PAM_SUCCESS),
        false => Err(HelloAuthenticationError::SignAuthenticationFail),
    }
}
<|endoftext|>


<commit_before>//! Emulates the various functionality of the cartridges
use std::slice::bytes::copy_memory;

use cart::MemoryBankController::{NoMbc, Mbc1, Mbc2, Mbc3};
#[derive(Debug, Copy, PartialEq)]
pub enum MemoryBankController {
    NoMbc, // No memory bank controller (32Kbyte ROM only)
    Mbc1,  // Max 2MBbyte ROM and/or 32KByte RAM
    Mbc2,  // Max 256Kbyte ROM and 512x4 bits RAM
    Mbc3,  // Max 2MByte ROM and/or 32KByte RAM and Timer
    // FIXME(minor): add support for Huc1 controller
    // Huc1,  // MBC with Infrared Controller
}

#[derive(Copy, PartialEq)]
pub enum BankingMode {
    Rom,
    Ram,
}

pub trait SaveFile: Send {
    fn load(&mut self, data: &mut [u8]);
    fn save(&mut self, data: &[u8]);
}

pub struct Cartridge {
    mbc: MemoryBankController,
    selected_banking_mode: BankingMode,

    /// Rom banks (maximum of 128 banks = 2MB, mapped to: 0x0000-0x7FFFF)
    pub rom: [[u8; 0x4000]; 128],
    /// The currently mapped rom bank (bank 0 is always mapped)
    pub rom_bank: usize,

    /// External ram banks (maximum of 4 banks = 32KB, mapped to: 0xA000-0xBFFF)
    pub ram: [u8; 0x2000 * 4],
    /// The the currently mapped external ram bank (bank 0 is always mapped)
    pub ram_bank: usize,
    // Indicates if the ram has been enabled
    pub ram_enabled: bool,

    /// The save file attached to the cartridge
    pub save_file: Option<Box<SaveFile>>,
}

impl Cartridge {
    pub fn new() -> Cartridge {
        Cartridge {
            mbc: NoMbc,
            selected_banking_mode: BankingMode::Rom,

            rom: [[0; 0x4000]; 128],
            rom_bank: 1,

            ram: [0; 0x2000 * 4],
            ram_bank: 0,
            ram_enabled: false,

            save_file: None,
        }
    }

    /// Convert the specified address into an index into the ram array, adjusting the index
    /// according to the currently enabled ram bank
    pub fn ram_index(&self, addr: u16) -> usize {
        0x2000 * self.ram_bank + (addr & 0x1FFF) as usize
    }

    pub fn load(&mut self, data: &[u8], save_file: Option<Box<SaveFile>>) {
        for (i, chunk) in data.chunks(0x4000).enumerate() {
            copy_memory(&mut self.rom[i], chunk);
        }

        // Get the type of memory bank controller from the cart
        // TODO: Handle more of the settings
        let mbc_type = self.rom[0][0x0147];
        // 00h  ROM ONLY                 13h  MBC3+RAM+BATTERY
        // 01h  MBC1                     15h  MBC4
        // 02h  MBC1+RAM                 16h  MBC4+RAM
        // 03h  MBC1+RAM+BATTERY         17h  MBC4+RAM+BATTERY
        // 05h  MBC2                     19h  MBC5
        // 06h  MBC2+BATTERY             1Ah  MBC5+RAM
        // 08h  ROM+RAM                  1Bh  MBC5+RAM+BATTERY
        // 09h  ROM+RAM+BATTERY          1Ch  MBC5+RUMBLE
        // 0Bh  MMM01                    1Dh  MBC5+RUMBLE+RAM
        // 0Ch  MMM01+RAM                1Eh  MBC5+RUMBLE+RAM+BATTERY
        // 0Dh  MMM01+RAM+BATTERY        FCh  POCKET CAMERA
        // 0Fh  MBC3+TIMER+BATTERY       FDh  BANDAI TAMA5
        // 10h  MBC3+TIMER+RAM+BATTERY   FEh  HuC3
        // 11h  MBC3                     FFh  HuC1+RAM+BATTERY
        // 12h  MBC3+RAM
        match mbc_type {
            0x00 | 0x08 | 0x09 => { self.mbc = NoMbc; },
            0x01 | 0x02 | 0x03 => { self.mbc = Mbc1; },
            0x05 | 0x06 => { self.mbc = Mbc2; },
            0x11 | 0x12 | 0x0F | 0x10 | 0x13 => { self.mbc = Mbc3; }

            x => { panic!("Unsupported cartridge: {:X}", x); }
        }
        println!("{:?}", self.mbc);

        // Load the save file if there is one
        self.save_file = save_file;
        if let Some(ref mut save) = self.save_file {
            save.load(&mut self.ram);
        }
    }

    pub fn read(&self, addr: u16) -> u8 {
        // FIXME(minor): should probably handle invalid memory accesses depending on the memory bank
        // controller more carefully.
        match addr {
            0x0000 ... 0x3FFF => self.rom[0][(addr & 0x3FFF) as usize],
            0x4000 ... 0x7FFF => self.rom[self.rom_bank][(addr & 0x3FFF) as usize],
            0xA000 ... 0xBFFF => self.ram[self.ram_index(addr)],

            _ => unreachable!(),
        }
    }

    pub fn write(&mut self, addr: u16, value: u8) {
        match addr {
            0x0000 ... 0x1FFF => match self.mbc {
                NoMbc => {},
                Mbc1 => self.ram_enabled = value & 0xF == 0x0A,
                Mbc2 => {
                    if addr & 0x0100 == 0 {
                        self.ram_enabled = !self.ram_enabled;
                    }
                },
                Mbc3 => {
                    self.ram_enabled = value & 0xF == 0x0A;

                    // FIXME(major): uncomment the following line when the real-time-clock is added.
                    // self.rtc.enabled = value == 0x0A;
                },
            },

            // Set ROM bank number
            0x2000 ... 0x3FFF => match self.mbc {
                NoMbc => {},
                Mbc1 => {
                    let high_bits = self.rom_bank as u8 & 0x60;
                    self.rom_bank =
                        if (value & 0x1F) == 0 { (high_bits | 1) as usize }
                        else { (high_bits | (value & 0x1F)) as usize };
                },
                Mbc2 => self.rom_bank = (value & 0x0F) as usize,
                Mbc3 => {
                    self.rom_bank =
                        if value == 0x00 { 1 }
                        else { (value & 0x7F) as usize };
                },
            },

            // Set RAM bank number, upper bits of ROM bank number or real-time-clock register select
            0x4000 ... 0x5FFF => match self.mbc {
                NoMbc | Mbc2 => {}
                Mbc1 => {
                    let bits = value & 0x3;
                    if self.selected_banking_mode == BankingMode::Rom {
                        self.rom_bank = (bits << 5 | (self.rom_bank as u8 & 0x1F)) as usize;
                    }
                    else {
                        self.ram_bank = bits as usize;
                    }
                },
                Mbc3 => {
                    // FIXME(major): handle real time clock here
                    self.ram_bank = (value & 0x3) as usize;
                },
            },

            // Set ROM/RAM mode select or latch clock data
            0x6000 ... 0x7FFF => match self.mbc {
                NoMbc | Mbc2 => {},
                Mbc1 => {
                    self.selected_banking_mode =
                        if value & 0x1 == 0 { BankingMode::Rom }
                        else { BankingMode::Ram };
                },
                Mbc3 => {
                    // FIXME(major): Handle real-time-clock latch clock data
                },
            },

            // Write to external ram
            0xA000 ... 0xBFFF => self.ram[self.ram_index(addr)] = value,

            _ => unreachable!(),
        }
    }
}

impl Drop for Cartridge {
    fn drop(&mut self) {
        if let Some(ref mut save_file) = self.save_file {
            save_file.save(&self.ram);
        }
    }
}
<commit_msg>Added basic RTC support.<commit_after>//! Emulates the various functionality of the cartridges
use std::slice::bytes::copy_memory;

use cart::MemoryBankController::{NoMbc, Mbc1, Mbc2, Mbc3};
#[derive(Debug, Copy, PartialEq)]
pub enum MemoryBankController {
    NoMbc, // No memory bank controller (32Kbyte ROM only)
    Mbc1,  // Max 2MBbyte ROM and/or 32KByte RAM
    Mbc2,  // Max 256Kbyte ROM and 512x4 bits RAM
    Mbc3,  // Max 2MByte ROM and/or 32KByte RAM and Timer
    // FIXME(minor): add support for Huc1 controller
    // Huc1,  // MBC with Infrared Controller
}

#[derive(Copy, PartialEq)]
pub enum BankingMode {
    Rom,
    Ram,
}

pub trait SaveFile: Send {
    fn load(&mut self, data: &mut [u8]);
    fn save(&mut self, data: &[u8]);
}

pub struct RealTimeClock {
    state: u8,
    carry: bool,
    halt: bool,

    seconds: u8,
    minutes: u8,
    hours: u8,
    days: u16,
}

impl RealTimeClock {
    fn new() -> RealTimeClock {
        RealTimeClock {
            state: 0,
            carry: false,
            halt: false,

            seconds: 0,
            minutes: 0,
            hours: 0,
            days: 0,
        }
    }

    fn read(&self) -> u8 {
        match self.state & 0x7 {
            0 => self.seconds,
            1 => self.minutes,
            2 => self.hours,
            3 => self.days as u8,
            4 => {
                ((self.days >> 8) as u8) |
                if self.halt { 0x40 } else { 0 } |
                if self.carry { 0x80 } else { 0 }
            },

            _ => 0x00,
        }
    }

    fn write(&mut self, value: u8) {
        match self.state & 0x7 {
            0 => self.seconds = value % 60,
            1 => self.minutes = value % 60,
            2 => self.hours = value % 24,
            3 => self.days = (self.days & 0x100) | (value as u16),
            4 => {
                self.days = (((value & 1) as u16) << 8) | (self.days & 0xFF);
                self.halt = (value & 0x40) != 0;
                self.carry = (value & 0x80) != 0;
            },
            _ => {},
        }
    }

    fn is_enabled(&self) -> bool {
        self.state & 0x8 != 0
    }

    pub fn tick(&mut self) {
        if self.halt {
            return;
        }

        self.seconds += 1;
        if self.seconds >= 60 {
            self.seconds = 0;

            self.minutes += 1;
            if self.minutes >= 60 {
                self.minutes = 0;

                self.hours += 1;
                if self.hours >= 24 {
                    self.hours = 0;

                    self.days += 1;
                    if self.days >= 356 {
                        self.days = 0;

                        self.carry = true;
                    }
                }
            }
        }
    }

}

pub struct Cartridge {
    mbc: MemoryBankController,
    selected_banking_mode: BankingMode,

    /// Rom banks (maximum of 128 banks = 2MB, mapped to: 0x0000-0x7FFFF)
    pub rom: [[u8; 0x4000]; 128],
    /// The currently mapped rom bank (bank 0 is always mapped)
    pub rom_bank: usize,

    /// External ram banks (maximum of 4 banks = 32KB, mapped to: 0xA000-0xBFFF)
    pub ram: [u8; 0x2000 * 4],
    /// The the currently mapped external ram bank (bank 0 is always mapped)
    pub ram_bank: usize,
    // Indicates if the ram has been enabled
    pub ram_enabled: bool,

    /// The save file attached to the cartridge
    pub save_file: Option<Box<SaveFile>>,

    /// The realtime clock attached to the cartridge
    pub rtc: RealTimeClock,
}

impl Cartridge {
    pub fn new() -> Cartridge {
        Cartridge {
            mbc: NoMbc,
            selected_banking_mode: BankingMode::Rom,

            rom: [[0; 0x4000]; 128],
            rom_bank: 1,

            ram: [0; 0x2000 * 4],
            ram_bank: 0,
            ram_enabled: false,

            save_file: None,
            rtc: RealTimeClock::new(),
        }
    }

    /// Convert the specified address into an index into the ram array, adjusting the index
    /// according to the currently enabled ram bank
    pub fn ram_index(&self, addr: u16) -> usize {
        0x2000 * self.ram_bank + (addr & 0x1FFF) as usize
    }

    pub fn load(&mut self, data: &[u8], save_file: Option<Box<SaveFile>>) {
        for (i, chunk) in data.chunks(0x4000).enumerate() {
            copy_memory(&mut self.rom[i], chunk);
        }

        // Get the type of memory bank controller from the cart
        // TODO: Handle more of the settings
        let mbc_type = self.rom[0][0x0147];
        // 00h  ROM ONLY                 13h  MBC3+RAM+BATTERY
        // 01h  MBC1                     15h  MBC4
        // 02h  MBC1+RAM                 16h  MBC4+RAM
        // 03h  MBC1+RAM+BATTERY         17h  MBC4+RAM+BATTERY
        // 05h  MBC2                     19h  MBC5
        // 06h  MBC2+BATTERY             1Ah  MBC5+RAM
        // 08h  ROM+RAM                  1Bh  MBC5+RAM+BATTERY
        // 09h  ROM+RAM+BATTERY          1Ch  MBC5+RUMBLE
        // 0Bh  MMM01                    1Dh  MBC5+RUMBLE+RAM
        // 0Ch  MMM01+RAM                1Eh  MBC5+RUMBLE+RAM+BATTERY
        // 0Dh  MMM01+RAM+BATTERY        FCh  POCKET CAMERA
        // 0Fh  MBC3+TIMER+BATTERY       FDh  BANDAI TAMA5
        // 10h  MBC3+TIMER+RAM+BATTERY   FEh  HuC3
        // 11h  MBC3                     FFh  HuC1+RAM+BATTERY
        // 12h  MBC3+RAM
        match mbc_type {
            0x00 | 0x08 | 0x09 => { self.mbc = NoMbc; },
            0x01 | 0x02 | 0x03 => { self.mbc = Mbc1; },
            0x05 | 0x06 => { self.mbc = Mbc2; },
            0x11 | 0x12 | 0x0F | 0x10 | 0x13 => { self.mbc = Mbc3; }

            x => { panic!("Unsupported cartridge: {:X}", x); }
        }
        println!("{:?}", self.mbc);

        // Load the save file if there is one
        self.save_file = save_file;
        if let Some(ref mut save) = self.save_file {
            save.load(&mut self.ram);
        }
    }

    pub fn read(&self, addr: u16) -> u8 {
        // FIXME(minor): should probably handle invalid memory accesses depending on the memory bank
        // controller more carefully.
        match addr {
            0x0000 ... 0x3FFF => self.rom[0][(addr & 0x3FFF) as usize],
            0x4000 ... 0x7FFF => self.rom[self.rom_bank][(addr & 0x3FFF) as usize],
            0xA000 ... 0xBFFF => {
                if self.rtc.is_enabled() { self.rtc.read() }
                else { self.ram[self.ram_index(addr)] }
            },

            _ => unreachable!(),
        }
    }

    pub fn write(&mut self, addr: u16, value: u8) {
        match addr {
            0x0000 ... 0x1FFF => match self.mbc {
                NoMbc => {},
                Mbc1 => self.ram_enabled = value & 0xF == 0x0A,
                Mbc2 => {
                    if addr & 0x0100 == 0 {
                        self.ram_enabled = !self.ram_enabled;
                    }
                },
                Mbc3 => {
                    self.ram_enabled = value & 0xF == 0x0A;
                    // CHECKME: Not sure if this also enables the rtc
                    // self.rtc.enabled = value == 0x0A;
                },
            },

            // Set ROM bank number
            0x2000 ... 0x3FFF => match self.mbc {
                NoMbc => {},
                Mbc1 => {
                    let high_bits = self.rom_bank as u8 & 0x60;
                    self.rom_bank =
                        if (value & 0x1F) == 0 { (high_bits | 1) as usize }
                        else { (high_bits | (value & 0x1F)) as usize };
                },
                Mbc2 => self.rom_bank = (value & 0x0F) as usize,
                Mbc3 => {
                    self.rom_bank =
                        if value == 0x00 { 1 }
                        else { (value & 0x7F) as usize };
                },
            },

            // Set RAM bank number, upper bits of ROM bank number or real-time-clock register select
            0x4000 ... 0x5FFF => match self.mbc {
                NoMbc | Mbc2 => {}
                Mbc1 => {
                    let bits = value & 0x3;
                    if self.selected_banking_mode == BankingMode::Rom {
                        self.rom_bank = (bits << 5 | (self.rom_bank as u8 & 0x1F)) as usize;
                    }
                    else {
                        self.ram_bank = bits as usize;
                    }
                },
                Mbc3 => {
                    self.rtc.state = value & 0xF;
                    self.ram_bank = (value & 0x3) as usize;
                },
            },

            // Set ROM/RAM mode select or latch clock data
            0x6000 ... 0x7FFF => match self.mbc {
                NoMbc | Mbc2 => {},
                Mbc1 => {
                    self.selected_banking_mode =
                        if value & 0x1 == 0 { BankingMode::Rom }
                        else { BankingMode::Ram };
                },
                Mbc3 => {
                    // FIXME(major): Handle real-time-clock latch clock data
                },
            },

            // Write to external ram
            0xA000 ... 0xBFFF => self.ram[self.ram_index(addr)] = value,

            _ => unreachable!(),
        }
    }
}

impl Drop for Cartridge {
    fn drop(&mut self) {
        if let Some(ref mut save_file) = self.save_file {
            // TODO: Save the rtc
            save_file.save(&self.ram);
        }
    }
}
<|endoftext|>


<commit_before>//! Thread-safe connections on IrcStreams.
#![experimental]
use std::sync::{Mutex, MutexGuard};
use std::io::{BufferedReader, BufferedWriter, IoResult, TcpStream};
#[cfg(any(feature = "encode", feature = "ssl"))] use std::io::{IoError, IoErrorKind};
#[cfg(feature = "encode")] use encoding::{DecoderTrap, EncoderTrap, Encoding};
#[cfg(feature = "encode")] use encoding::label::encoding_from_whatwg_label;
use data::kinds::{IrcReader, IrcWriter};
use data::message::Message;
#[cfg(feature = "ssl")] use openssl::ssl::{SslContext, SslMethod, SslStream};
#[cfg(feature = "ssl")] use openssl::ssl::error::SslError;

/// A thread-safe connection.
#[experimental]
pub struct Connection<T: IrcReader, U: IrcWriter> {
    reader: Mutex<T>,
    writer: Mutex<U>,
}

/// A Connection over a buffered NetStream.
pub type NetConnection = Connection<BufferedReader<NetStream>, BufferedWriter<NetStream>>;

impl Connection<BufferedReader<NetStream>, BufferedWriter<NetStream>> {
    /// Creates a thread-safe TCP connection to the specified server.
    #[experimental]
    pub fn connect(host: &str, port: u16) -> IoResult<NetConnection> {
        Connection::connect_internal(host, port, None)
    }

    /// Creates a thread-safe TCP connection to the specified server with a given timeout in
    /// milliseconds.
    #[experimental]
    pub fn connect_with_timeout(host: &str, port: u16, timeout_ms: u64)
        -> IoResult<NetConnection> {
        Connection::connect_internal(host, port, Some(timeout_ms))
    }

    /// Creates a thread-safe TCP connection with an optional timeout.
    #[experimental]
    fn connect_internal(host: &str, port: u16, timeout_ms: Option<u64>)
    -> IoResult<NetConnection> {
        let mut socket = try!(TcpStream::connect(format!("{}:{}", host, port)[]));
        socket.set_timeout(timeout_ms);
        Ok(Connection::new(
            BufferedReader::new(NetStream::UnsecuredTcpStream(socket.clone())),
            BufferedWriter::new(NetStream::UnsecuredTcpStream(socket))
        ))
    }

    /// Creates a thread-safe TCP connection to the specified server over SSL.
    /// If the library is compiled without SSL support, this method panics.
    #[experimental]
    pub fn connect_ssl(host: &str, port: u16) -> IoResult<NetConnection> {
        Connection::connect_ssl_internal(host, port, None)
    }

    /// Creates a thread-safe TCP connection to the specificed server over SSL with a given timeout
    /// in milliseconds. If the library is compiled without SSL support, this method panics.
    #[experimental]
    pub fn connect_ssl_with_timeout(host: &str, port: u16, timeout_ms: u64)
        -> IoResult<NetConnection> {
        Connection::connect_ssl_internal(host, port, Some(timeout_ms))
    }

    /// Panics because SSL support was not included at compilation.
    #[experimental]
    #[cfg(not(feature = "ssl"))]
    fn connect_ssl_internal(host: &str, port: u16, _: Option<u64>)
    -> IoResult<NetConnection> {
        panic!("Cannot connect to {}:{} over SSL without compiling with SSL support.", host, port)
    }

    /// Creates a thread-safe TCP connection over SSL with an optional timeout.
    #[experimental]
    #[cfg(feature = "ssl")]
    fn connect_ssl_internal(host: &str, port: u16, timeout_ms: Option<u64>)
    -> IoResult<NetConnection> {
        let mut socket = try!(TcpStream::connect(format!("{}:{}", host, port)[]));
        socket.set_timeout(timeout_ms);
        let ssl = try!(ssl_to_io(SslContext::new(SslMethod::Tlsv1)));
        let ssl_socket = try!(ssl_to_io(SslStream::new(&ssl, socket)));
        Ok(Connection::new(
            BufferedReader::new(NetStream::SslTcpStream(ssl_socket.clone())),
            BufferedWriter::new(NetStream::SslTcpStream(ssl_socket)),
        ))
    }

    /// Sets the keepalive for the network stream.
    #[experimental]
    pub fn set_keepalive(&self, delay_in_seconds: Option<uint>) -> IoResult<()> {
        match self.reader.lock().get_mut() {
            &NetStream::UnsecuredTcpStream(ref mut tcp) => tcp.set_keepalive(delay_in_seconds),
            #[cfg(feature = "ssl")]
            &NetStream::SslTcpStream(ref mut ssl) => ssl.get_mut().set_keepalive(delay_in_seconds),
        }
    }
}

impl<T: IrcReader, U: IrcWriter> Connection<T, U> {
    /// Creates a new connection from an IrcReader and an IrcWriter.
    #[experimental]
    pub fn new(reader: T, writer: U) -> Connection<T, U> {
        Connection {
            reader: Mutex::new(reader),
            writer: Mutex::new(writer),
        }
    }

    /// Sends a Message over this connection.
    #[experimental]
    #[cfg(feature = "encode")]
    pub fn send(&self, message: Message, encoding: &str) -> IoResult<()> {
        let encoding = match encoding_from_whatwg_label(encoding) {
            Some(enc) => enc,
            None => return Err(IoError {
                kind: IoErrorKind::InvalidInput,
                desc: "Failed to find decoder.",
                detail: Some(format!("Invalid decoder: {}", encoding))
            })
        };
        let data = match encoding.encode(message.into_string()[], EncoderTrap::Replace) {
            Ok(data) => data,
            Err(data) => return Err(IoError {
                kind: IoErrorKind::InvalidInput,
                desc: "Failed to decode message.",
                detail: Some(format!("Failed to decode {} as {}.", data, encoding.name())),
            })
        };
        let mut writer = self.writer.lock();
        try!(writer.write(data[]));
        writer.flush()
    }

    /// Sends a message over this connection.
    #[experimental]
    #[cfg(not(feature = "encode"))]
    pub fn send(&self, message: Message) -> IoResult<()> {
        let mut writer = self.writer.lock();
        try!(writer.write_str(message.into_string()[]));
        writer.flush()
    }

    /// Receives a single line from this connection.
    #[experimental]
    #[cfg(feature = "encoding")]
    pub fn recv(&self, encoding: &str) -> IoResult<String> {
        let encoding = match encoding_from_whatwg_label(encoding) {
            Some(enc) => enc,
            None => return Err(IoError {
                kind: IoErrorKind::InvalidInput,
                desc: "Failed to find decoder.",
                detail: Some(format!("Invalid decoder: {}", encoding))
            })
        };
        self.reader.lock().read_until(b'\n').and_then(|line|
            match encoding.decode(line[], DecoderTrap::Replace) {
                Ok(data) => Ok(data),
                Err(data) => Err(IoError {
                    kind: IoErrorKind::InvalidInput,
                    desc: "Failed to decode message.",
                    detail: Some(format!("Failed to decode {} as {}.", data, encoding.name())),
                })
            }
        )
    }

    /// Receives a single line from this connection.
    #[experimental]
    #[cfg(not(feature = "encoding"))]
    pub fn recv(&self) -> IoResult<String> {
        self.reader.lock().read_line()
    }

    /// Acquires the Reader lock.
    #[experimental]
    pub fn reader<'a>(&'a self) -> MutexGuard<'a, T> {
        self.reader.lock()
    }

    /// Acquires the Writer lock.
    #[experimental]
    pub fn writer<'a>(&'a self) -> MutexGuard<'a, U> {
        self.writer.lock()
    }
}

/// Converts a Result<T, SslError> into an IoResult<T>.
#[cfg(feature = "ssl")]
fn ssl_to_io<T>(res: Result<T, SslError>) -> IoResult<T> {
    match res {
        Ok(x) => Ok(x),
        Err(e) => Err(IoError {
            kind: IoErrorKind::OtherIoError,
            desc: "An SSL error occurred.",
            detail: Some(format!("{}", e)),
        }),
    }
}

/// An abstraction over different networked streams.
#[experimental]
pub enum NetStream {
    /// An unsecured TcpStream.
    UnsecuredTcpStream(TcpStream),
    /// An SSL-secured TcpStream.
    /// This is only available when compiled with SSL support.
    #[cfg(feature = "ssl")]
    SslTcpStream(SslStream<TcpStream>),
}

impl Reader for NetStream {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        match self {
            &NetStream::UnsecuredTcpStream(ref mut stream) => stream.read(buf),
            #[cfg(feature = "ssl")]
            &NetStream::SslTcpStream(ref mut stream) => stream.read(buf),
        }
    }
}

impl Writer for NetStream {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        match self {
            &NetStream::UnsecuredTcpStream(ref mut stream) => stream.write(buf),
            #[cfg(feature = "ssl")]
            &NetStream::SslTcpStream(ref mut stream) => stream.write(buf),
        }
    }
}

#[cfg(test)]
mod test {
    use super::Connection;
    use std::io::{MemReader, MemWriter};
    use std::io::util::{NullReader, NullWriter};
    use data::message::Message;
    #[cfg(feature = "encode")] use encoding::{DecoderTrap, Encoding};
    #[cfg(feature = "encode")] use encoding::all::{ISO_8859_15, UTF_8};

    #[test]
    #[cfg(not(feature = "encode"))]
    fn send() {
        let conn = Connection::new(NullReader, MemWriter::new());
        assert!(conn.send(
            Message::new(None, "PRIVMSG", Some(vec!["test"]), Some("Testing!"))
        ).is_ok());
        let data = String::from_utf8(conn.writer().get_ref().to_vec()).unwrap();
        assert_eq!(data[], "PRIVMSG test :Testing!\r\n");
    }

    #[test]
    #[cfg(feature = "encode")]
    fn send_utf8() {
        let conn = Connection::new(NullReader, MemWriter::new());
        assert!(conn.send(
            Message::new(None, "PRIVMSG", Some(vec!["test"]), Some("€ŠšŽžŒœŸ")), "UTF-8"
        ).is_ok());
        let data = UTF_8.decode(conn.writer().get_ref(), DecoderTrap::Strict).unwrap();
        assert_eq!(data[], "PRIVMSG test :€ŠšŽžŒœŸ\r\n");
    }

    #[test]
    #[cfg(feature = "encode")]
    fn send_iso885915() {
        let conn = Connection::new(NullReader, MemWriter::new());
        assert!(conn.send(
            Message::new(None, "PRIVMSG", Some(vec!["test"]), Some("€ŠšŽžŒœŸ")), "l9"
        ).is_ok());
        let data = ISO_8859_15.decode(conn.writer().get_ref(), DecoderTrap::Strict).unwrap();
        assert_eq!(data[], "PRIVMSG test :€ŠšŽžŒœŸ\r\n");
    }

    #[test]
    #[cfg(not(feature = "encode"))]
    fn recv() {
        let conn = Connection::new(
            MemReader::new("PRIVMSG test :Testing!\r\n".as_bytes().to_vec()), NullWriter
        );
        assert_eq!(conn.recv().unwrap()[], "PRIVMSG test :Testing!\r\n");
    }

    #[test]
    #[cfg(feature = "encode")]
    fn recv_utf8() {
        let conn = Connection::new(
            MemReader::new(b"PRIVMSG test :Testing!\r\n".to_vec()), NullWriter
        );
        assert_eq!(conn.recv("UTF-8").unwrap()[], "PRIVMSG test :Testing!\r\n");
    }

    #[test]
    #[cfg(feature = "encode")]
    fn recv_iso885915() {
        let conn = Connection::new(
            MemReader::new({
                let mut vec = Vec::new();
                vec.push_all(b"PRIVMSG test :");
                vec.push_all(&[0xA4, 0xA6, 0xA8, 0xB4, 0xB8, 0xBC, 0xBD, 0xBE]);
                vec.push_all(b"\r\n");
                vec
            }), NullWriter
        );
        assert_eq!(conn.recv("l9").unwrap()[], "PRIVMSG test :€ŠšŽžŒœŸ\r\n");
    }
}
<commit_msg>Added set_timeout to Connection.<commit_after>//! Thread-safe connections on IrcStreams.
#![experimental]
use std::sync::{Mutex, MutexGuard};
use std::io::{BufferedReader, BufferedWriter, IoResult, TcpStream};
#[cfg(any(feature = "encode", feature = "ssl"))] use std::io::{IoError, IoErrorKind};
#[cfg(feature = "encode")] use encoding::{DecoderTrap, EncoderTrap, Encoding};
#[cfg(feature = "encode")] use encoding::label::encoding_from_whatwg_label;
use data::kinds::{IrcReader, IrcWriter};
use data::message::Message;
#[cfg(feature = "ssl")] use openssl::ssl::{SslContext, SslMethod, SslStream};
#[cfg(feature = "ssl")] use openssl::ssl::error::SslError;

/// A thread-safe connection.
#[experimental]
pub struct Connection<T: IrcReader, U: IrcWriter> {
    reader: Mutex<T>,
    writer: Mutex<U>,
}

/// A Connection over a buffered NetStream.
pub type NetConnection = Connection<BufferedReader<NetStream>, BufferedWriter<NetStream>>;

impl Connection<BufferedReader<NetStream>, BufferedWriter<NetStream>> {
    /// Creates a thread-safe TCP connection to the specified server.
    #[experimental]
    pub fn connect(host: &str, port: u16) -> IoResult<NetConnection> {
        Connection::connect_internal(host, port, None)
    }

    /// Creates a thread-safe TCP connection to the specified server with a given timeout in
    /// milliseconds.
    #[experimental]
    pub fn connect_with_timeout(host: &str, port: u16, timeout_ms: u64)
        -> IoResult<NetConnection> {
        Connection::connect_internal(host, port, Some(timeout_ms))
    }

    /// Creates a thread-safe TCP connection with an optional timeout.
    #[experimental]
    fn connect_internal(host: &str, port: u16, timeout_ms: Option<u64>)
    -> IoResult<NetConnection> {
        let mut socket = try!(TcpStream::connect(format!("{}:{}", host, port)[]));
        socket.set_timeout(timeout_ms);
        Ok(Connection::new(
            BufferedReader::new(NetStream::UnsecuredTcpStream(socket.clone())),
            BufferedWriter::new(NetStream::UnsecuredTcpStream(socket))
        ))
    }

    /// Creates a thread-safe TCP connection to the specified server over SSL.
    /// If the library is compiled without SSL support, this method panics.
    #[experimental]
    pub fn connect_ssl(host: &str, port: u16) -> IoResult<NetConnection> {
        Connection::connect_ssl_internal(host, port, None)
    }

    /// Creates a thread-safe TCP connection to the specificed server over SSL with a given timeout
    /// in milliseconds. If the library is compiled without SSL support, this method panics.
    #[experimental]
    pub fn connect_ssl_with_timeout(host: &str, port: u16, timeout_ms: u64)
        -> IoResult<NetConnection> {
        Connection::connect_ssl_internal(host, port, Some(timeout_ms))
    }

    /// Panics because SSL support was not included at compilation.
    #[experimental]
    #[cfg(not(feature = "ssl"))]
    fn connect_ssl_internal(host: &str, port: u16, _: Option<u64>)
    -> IoResult<NetConnection> {
        panic!("Cannot connect to {}:{} over SSL without compiling with SSL support.", host, port)
    }

    /// Creates a thread-safe TCP connection over SSL with an optional timeout.
    #[experimental]
    #[cfg(feature = "ssl")]
    fn connect_ssl_internal(host: &str, port: u16, timeout_ms: Option<u64>)
    -> IoResult<NetConnection> {
        let mut socket = try!(TcpStream::connect(format!("{}:{}", host, port)[]));
        socket.set_timeout(timeout_ms);
        let ssl = try!(ssl_to_io(SslContext::new(SslMethod::Tlsv1)));
        let ssl_socket = try!(ssl_to_io(SslStream::new(&ssl, socket)));
        Ok(Connection::new(
            BufferedReader::new(NetStream::SslTcpStream(ssl_socket.clone())),
            BufferedWriter::new(NetStream::SslTcpStream(ssl_socket)),
        ))
    }

    /// Sets the keepalive for the network stream.
    #[experimental]
    pub fn set_keepalive(&self, delay_in_seconds: Option<uint>) -> IoResult<()> {
        self.mod_stream(|tcp| tcp.set_keepalive(delay_in_seconds))
    }

    /// Sets the timeout for the network stream.
    #[experimental]
    pub fn set_timeout(&self, timeout_ms: Option<u64>) {
        self.mod_stream(|tcp| Ok(tcp.set_timeout(timeout_ms))).unwrap(); // this cannot fail.
    }

    /// Modifies the internal TcpStream using a function.
    fn mod_stream(&self, f: |&mut TcpStream| -> IoResult<()>) -> IoResult<()> {
        match self.reader.lock().get_mut() {
            &NetStream::UnsecuredTcpStream(ref mut tcp) => f(tcp),
            #[cfg(feature = "ssl")]
            &NetStream::SslTcpStream(ref mut ssl) => f(ssl.get_mut()),
        }
    }
}

impl<T: IrcReader, U: IrcWriter> Connection<T, U> {
    /// Creates a new connection from an IrcReader and an IrcWriter.
    #[experimental]
    pub fn new(reader: T, writer: U) -> Connection<T, U> {
        Connection {
            reader: Mutex::new(reader),
            writer: Mutex::new(writer),
        }
    }

    /// Sends a Message over this connection.
    #[experimental]
    #[cfg(feature = "encode")]
    pub fn send(&self, message: Message, encoding: &str) -> IoResult<()> {
        let encoding = match encoding_from_whatwg_label(encoding) {
            Some(enc) => enc,
            None => return Err(IoError {
                kind: IoErrorKind::InvalidInput,
                desc: "Failed to find decoder.",
                detail: Some(format!("Invalid decoder: {}", encoding))
            })
        };
        let data = match encoding.encode(message.into_string()[], EncoderTrap::Replace) {
            Ok(data) => data,
            Err(data) => return Err(IoError {
                kind: IoErrorKind::InvalidInput,
                desc: "Failed to decode message.",
                detail: Some(format!("Failed to decode {} as {}.", data, encoding.name())),
            })
        };
        let mut writer = self.writer.lock();
        try!(writer.write(data[]));
        writer.flush()
    }

    /// Sends a message over this connection.
    #[experimental]
    #[cfg(not(feature = "encode"))]
    pub fn send(&self, message: Message) -> IoResult<()> {
        let mut writer = self.writer.lock();
        try!(writer.write_str(message.into_string()[]));
        writer.flush()
    }

    /// Receives a single line from this connection.
    #[experimental]
    #[cfg(feature = "encoding")]
    pub fn recv(&self, encoding: &str) -> IoResult<String> {
        let encoding = match encoding_from_whatwg_label(encoding) {
            Some(enc) => enc,
            None => return Err(IoError {
                kind: IoErrorKind::InvalidInput,
                desc: "Failed to find decoder.",
                detail: Some(format!("Invalid decoder: {}", encoding))
            })
        };
        self.reader.lock().read_until(b'\n').and_then(|line|
            match encoding.decode(line[], DecoderTrap::Replace) {
                Ok(data) => Ok(data),
                Err(data) => Err(IoError {
                    kind: IoErrorKind::InvalidInput,
                    desc: "Failed to decode message.",
                    detail: Some(format!("Failed to decode {} as {}.", data, encoding.name())),
                })
            }
        )
    }

    /// Receives a single line from this connection.
    #[experimental]
    #[cfg(not(feature = "encoding"))]
    pub fn recv(&self) -> IoResult<String> {
        self.reader.lock().read_line()
    }

    /// Acquires the Reader lock.
    #[experimental]
    pub fn reader<'a>(&'a self) -> MutexGuard<'a, T> {
        self.reader.lock()
    }

    /// Acquires the Writer lock.
    #[experimental]
    pub fn writer<'a>(&'a self) -> MutexGuard<'a, U> {
        self.writer.lock()
    }
}

/// Converts a Result<T, SslError> into an IoResult<T>.
#[cfg(feature = "ssl")]
fn ssl_to_io<T>(res: Result<T, SslError>) -> IoResult<T> {
    match res {
        Ok(x) => Ok(x),
        Err(e) => Err(IoError {
            kind: IoErrorKind::OtherIoError,
            desc: "An SSL error occurred.",
            detail: Some(format!("{}", e)),
        }),
    }
}

/// An abstraction over different networked streams.
#[experimental]
pub enum NetStream {
    /// An unsecured TcpStream.
    UnsecuredTcpStream(TcpStream),
    /// An SSL-secured TcpStream.
    /// This is only available when compiled with SSL support.
    #[cfg(feature = "ssl")]
    SslTcpStream(SslStream<TcpStream>),
}

impl Reader for NetStream {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        match self {
            &NetStream::UnsecuredTcpStream(ref mut stream) => stream.read(buf),
            #[cfg(feature = "ssl")]
            &NetStream::SslTcpStream(ref mut stream) => stream.read(buf),
        }
    }
}

impl Writer for NetStream {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        match self {
            &NetStream::UnsecuredTcpStream(ref mut stream) => stream.write(buf),
            #[cfg(feature = "ssl")]
            &NetStream::SslTcpStream(ref mut stream) => stream.write(buf),
        }
    }
}

#[cfg(test)]
mod test {
    use super::Connection;
    use std::io::{MemReader, MemWriter};
    use std::io::util::{NullReader, NullWriter};
    use data::message::Message;
    #[cfg(feature = "encode")] use encoding::{DecoderTrap, Encoding};
    #[cfg(feature = "encode")] use encoding::all::{ISO_8859_15, UTF_8};

    #[test]
    #[cfg(not(feature = "encode"))]
    fn send() {
        let conn = Connection::new(NullReader, MemWriter::new());
        assert!(conn.send(
            Message::new(None, "PRIVMSG", Some(vec!["test"]), Some("Testing!"))
        ).is_ok());
        let data = String::from_utf8(conn.writer().get_ref().to_vec()).unwrap();
        assert_eq!(data[], "PRIVMSG test :Testing!\r\n");
    }

    #[test]
    #[cfg(feature = "encode")]
    fn send_utf8() {
        let conn = Connection::new(NullReader, MemWriter::new());
        assert!(conn.send(
            Message::new(None, "PRIVMSG", Some(vec!["test"]), Some("€ŠšŽžŒœŸ")), "UTF-8"
        ).is_ok());
        let data = UTF_8.decode(conn.writer().get_ref(), DecoderTrap::Strict).unwrap();
        assert_eq!(data[], "PRIVMSG test :€ŠšŽžŒœŸ\r\n");
    }

    #[test]
    #[cfg(feature = "encode")]
    fn send_iso885915() {
        let conn = Connection::new(NullReader, MemWriter::new());
        assert!(conn.send(
            Message::new(None, "PRIVMSG", Some(vec!["test"]), Some("€ŠšŽžŒœŸ")), "l9"
        ).is_ok());
        let data = ISO_8859_15.decode(conn.writer().get_ref(), DecoderTrap::Strict).unwrap();
        assert_eq!(data[], "PRIVMSG test :€ŠšŽžŒœŸ\r\n");
    }

    #[test]
    #[cfg(not(feature = "encode"))]
    fn recv() {
        let conn = Connection::new(
            MemReader::new("PRIVMSG test :Testing!\r\n".as_bytes().to_vec()), NullWriter
        );
        assert_eq!(conn.recv().unwrap()[], "PRIVMSG test :Testing!\r\n");
    }

    #[test]
    #[cfg(feature = "encode")]
    fn recv_utf8() {
        let conn = Connection::new(
            MemReader::new(b"PRIVMSG test :Testing!\r\n".to_vec()), NullWriter
        );
        assert_eq!(conn.recv("UTF-8").unwrap()[], "PRIVMSG test :Testing!\r\n");
    }

    #[test]
    #[cfg(feature = "encode")]
    fn recv_iso885915() {
        let conn = Connection::new(
            MemReader::new({
                let mut vec = Vec::new();
                vec.push_all(b"PRIVMSG test :");
                vec.push_all(&[0xA4, 0xA6, 0xA8, 0xB4, 0xB8, 0xBC, 0xBD, 0xBE]);
                vec.push_all(b"\r\n");
                vec
            }), NullWriter
        );
        assert_eq!(conn.recv("l9").unwrap()[], "PRIVMSG test :€ŠšŽžŒœŸ\r\n");
    }
}
<|endoftext|>