<commit_before>// Copyright (c) 2017 Jason White
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use std::ffi::OsStr;
use std::fmt;
use std::fs;
use std::io;
use std::io::Read;
use std::path::PathBuf;
use std::str::FromStr;

use sha2::{Digest, Sha256};

use serde::{de, Deserialize, Deserializer};

use super::traits::{Error, Resource, ResourceState};

use util::PathExt;

/// A file resource. This can actually be a file *or* directory.
///
/// TODO: Split the directory portion out into a "glob" resource whose state
/// changes when the list of matched files changes.
#[derive(Serialize, Ord, PartialOrd, Eq, PartialEq, Hash, Clone)]
pub struct FilePath {
    path: PathBuf,
}

impl FilePath {
    pub fn new(path: PathBuf) -> FilePath {
        FilePath { path: path.normalize() }
    }

    /// Assumes this resource is a regular file and returns its checksum.
    fn file_state(&self) -> Result<ResourceState, Error> {
        let mut hasher = Sha256::default();

        let mut f = match fs::File::open(&self.path) {
            Ok(f) => Ok(f),
            Err(err) => {
                match err.kind() {
                    io::ErrorKind::NotFound => {
                        return Ok(ResourceState::Missing);
                    }
                    _ => Err(err),
                }
            }
        }?;

        const BUF_SIZE: usize = 16384;

        let mut buf = [0u8; BUF_SIZE];

        loop {
            let n = f.read(&mut buf)?;

            if n == 0 || n < BUF_SIZE {
                break;
            }

            hasher.input(&buf[0..n]);
        }

        Ok(ResourceState::Checksum(hasher.result()))
    }

    /// Assumes this resource is a directory and returns the checksum of its
    /// file contents.
    fn dir_state(&self) -> Result<ResourceState, Error> {
        let mut hasher = Sha256::default();

        let mut names = vec![];

        for entry in fs::read_dir(&self.path)? {
            names.push(entry?.file_name());
        }

        // The order in which files are listed is not guaranteed to be sorted.
        // Whether or not it is sorted depends on the file system
        // implementation. Thus, we sort them to eliminate that potential
        // source of non-determinism.
        names.sort();

        for name in names {
            if let Some(name) = name.to_str() {
                hasher.input(name.as_bytes());
            }
        }

        Ok(ResourceState::Checksum(hasher.result()))
    }
}

impl<'a, T: ?Sized + AsRef<OsStr>> From<&'a T> for FilePath {
    fn from(s: &'a T) -> FilePath {
        FilePath::new(PathBuf::from(s.as_ref()))
    }
}

impl FromStr for FilePath {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(FilePath::new(PathBuf::from(s)))
    }
}

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

impl fmt::Debug for FilePath {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.path)
    }
}

// Derserialize a `FilePath` from a string.
impl<'de> Deserialize<'de> for FilePath {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where D: Deserializer<'de>
    {
        let s = String::deserialize(deserializer)?;
        FromStr::from_str(&s).map_err(de::Error::custom)
    }
}

impl Resource for FilePath {
    /// If a file, the checksum is of the contents of the file. If a directory,
    /// the checksum is of the sorted list of directory entries. Thus, if a file
    /// is added or removed from a directory, the checksum changes.
    fn state(&self) -> Result<ResourceState, Error> {
        if self.path.is_dir() {
            self.dir_state()
        } else {
            // Assume its a file even if its not. It'll error out if there are
            // problems reading it.
            self.file_state()
        }
    }

    /// If a file, simply deletes the file. If a directory, deletes the
    /// directory if it is empty.
    fn delete(&self) -> Result<(), Error> {
        let metadata = match fs::metadata(&self.path) {
            Ok(metadata) => Ok(metadata),
            Err(err) => {
                match err.kind() {
                    io::ErrorKind::NotFound => {
                        return Ok(());
                    }
                    _ => Err(err),
                }
            }
        }?;

        if metadata.is_dir() {
            // Ignore errors for directory deletion. It's not usually a problem
            // if a directory fails to get deleted. Directory deletion can fail
            // for a number of reasons:
            //  - A user may have created a file inside of it.
            //  - An untracked output may have been created inside of it.
            //  - On Windows, someone may have a lock on the directory.
            let _ = fs::remove_dir(&self.path);
            Ok(())
        } else {
            // Assume its a file even if its not. It'll error out if there are
            // problems deleting it.
            fs::remove_file(&self.path)
        }
    }
}

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

    #[test]
    fn test_normalized() {
        assert_eq!(FilePath::new(PathBuf::from("./foo/..//bar/")),
                   FilePath::new(PathBuf::from("bar")));
    }
}
<commit_msg>windows: Add case insensitive FilePath comparison<commit_after>// Copyright (c) 2017 Jason White
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use std::cmp::Ordering;
use std::ffi::OsStr;
use std::fmt;
use std::fs;
use std::io;
use std::io::Read;
use std::path::PathBuf;
use std::str::FromStr;
use std::hash::{Hash, Hasher};

use sha2::{Digest, Sha256};

use serde::{de, Deserialize, Deserializer};

use super::traits::{Error, Resource, ResourceState};

use util::PathExt;

/// A file resource. This can actually be a file *or* directory.
///
/// TODO: Split the directory portion out into a "glob" resource whose state
/// changes when the list of matched files changes.
#[derive(Serialize, Eq, Clone)]
pub struct FilePath {
    path: PathBuf,
}

impl FilePath {
    pub fn new(path: PathBuf) -> FilePath {
        FilePath { path: path.normalize() }
    }

    /// Assumes this resource is a regular file and returns its checksum.
    fn file_state(&self) -> Result<ResourceState, Error> {
        let mut hasher = Sha256::default();

        let mut f = match fs::File::open(&self.path) {
            Ok(f) => Ok(f),
            Err(err) => {
                match err.kind() {
                    io::ErrorKind::NotFound => {
                        return Ok(ResourceState::Missing);
                    }
                    _ => Err(err),
                }
            }
        }?;

        const BUF_SIZE: usize = 16384;

        let mut buf = [0u8; BUF_SIZE];

        loop {
            let n = f.read(&mut buf)?;

            if n == 0 || n < BUF_SIZE {
                break;
            }

            hasher.input(&buf[0..n]);
        }

        Ok(ResourceState::Checksum(hasher.result()))
    }

    /// Assumes this resource is a directory and returns the checksum of its
    /// file contents.
    fn dir_state(&self) -> Result<ResourceState, Error> {
        let mut hasher = Sha256::default();

        let mut names = vec![];

        for entry in fs::read_dir(&self.path)? {
            names.push(entry?.file_name());
        }

        // The order in which files are listed is not guaranteed to be sorted.
        // Whether or not it is sorted depends on the file system
        // implementation. Thus, we sort them to eliminate that potential
        // source of non-determinism.
        names.sort();

        for name in names {
            if let Some(name) = name.to_str() {
                hasher.input(name.as_bytes());
            }
        }

        Ok(ResourceState::Checksum(hasher.result()))
    }
}

impl<'a, T: ?Sized + AsRef<OsStr>> From<&'a T> for FilePath {
    fn from(s: &'a T) -> FilePath {
        FilePath::new(PathBuf::from(s.as_ref()))
    }
}

impl FromStr for FilePath {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(FilePath::new(PathBuf::from(s)))
    }
}

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

impl fmt::Debug for FilePath {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.path)
    }
}

// Derserialize a `FilePath` from a string.
impl<'de> Deserialize<'de> for FilePath {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where D: Deserializer<'de>
    {
        let s = String::deserialize(deserializer)?;
        FromStr::from_str(&s).map_err(de::Error::custom)
    }
}

impl Resource for FilePath {
    /// If a file, the checksum is of the contents of the file. If a directory,
    /// the checksum is of the sorted list of directory entries. Thus, if a file
    /// is added or removed from a directory, the checksum changes.
    fn state(&self) -> Result<ResourceState, Error> {
        if self.path.is_dir() {
            self.dir_state()
        } else {
            // Assume its a file even if its not. It'll error out if there are
            // problems reading it.
            self.file_state()
        }
    }

    /// If a file, simply deletes the file. If a directory, deletes the
    /// directory if it is empty.
    fn delete(&self) -> Result<(), Error> {
        let metadata = match fs::metadata(&self.path) {
            Ok(metadata) => Ok(metadata),
            Err(err) => {
                match err.kind() {
                    io::ErrorKind::NotFound => {
                        return Ok(());
                    }
                    _ => Err(err),
                }
            }
        }?;

        if metadata.is_dir() {
            // Ignore errors for directory deletion. It's not usually a problem
            // if a directory fails to get deleted. Directory deletion can fail
            // for a number of reasons:
            //  - A user may have created a file inside of it.
            //  - An untracked output may have been created inside of it.
            //  - On Windows, someone may have a lock on the directory.
            let _ = fs::remove_dir(&self.path);
            Ok(())
        } else {
            // Assume its a file even if its not. It'll error out if there are
            // problems deleting it.
            fs::remove_file(&self.path)
        }
    }
}

impl Hash for FilePath {
    #[cfg(windows)]
    fn hash<H>(&self, state: &mut H)
        where H: Hasher
    {
        let iter = self.path
                       .to_str()
                       .unwrap_or("")
                       .chars()
                       .flat_map(char::to_lowercase);
        for c in iter {
            c.hash(state);
        }

        '\0'.hash(state);
    }

    #[cfg(unix)]
    fn hash<H>(&self, state: &mut H)
        where H: Hasher
    {
        self.path.hash(state)
    }
}

impl Ord for FilePath {
    #[cfg(windows)]
    fn cmp(&self, other: &FilePath) -> Ordering {
        // The Ord implementation for `std::path::Path` is case-sensitive. It's
        // important that path comparisons on Windows are case-insensitive. The
        // nodes in the build graph don't link up correctly when file paths only
        // differ by case. Thus, we implement our own path comparison here.
        let a = self.path.to_str().unwrap_or("").chars().flat_map(char::to_lowercase);
        let mut b = other.path.to_str().unwrap_or("").chars().flat_map(char::to_lowercase);
        a.cmp(&mut b)
    }

    #[cfg(unix)]
    fn cmp(&self, other: &FilePath) -> Ordering {
        // File paths are case sensitive on non-Windows platforms.
        self.path.cmp(other)
    }
}

impl PartialOrd for FilePath {
    fn partial_cmp(&self, other: &FilePath) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for FilePath {
    fn eq(&self, other: &FilePath) -> bool {
        self.cmp(other) == Ordering::Equal
    }
}

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

    #[test]
    fn test_normalized() {
        assert_eq!(FilePath::new(PathBuf::from("./foo/..//bar/")),
                   FilePath::new(PathBuf::from("bar")));
    }

    #[test]
    #[cfg(windows)]
    fn test_comparison() {
        assert_eq!(FilePath::from("foobar/baz"), FilePath::from("foobar/baz"));
        assert_ne!(FilePath::from("doobar/baz"), FilePath::from("foobar/baz"));

        assert!(FilePath::from("abc") < FilePath::from("abd"));
        assert!(FilePath::from("abc") < FilePath::from("abcd"));

        // Case insensitive comparison
        assert_eq!(FilePath::from("foobar/baz"), FilePath::from("FooBar/Baz"));
        assert_eq!(FilePath::from("foobar/baz"), FilePath::from(r"FooBar\Baz"));
        assert_ne!(FilePath::from("foobar/baz"), FilePath::from("FooBar/Bazz"));
        assert!(FilePath::from("foobar/baz") < FilePath::from("FooBar/Bazz"));
    }
}
<|endoftext|>


<commit_before>use crate::{EarlyContext, EarlyLintPass, LintContext};
use rustc_ast::ast;
use rustc_data_structures::fx::FxHashMap;
use rustc_span::symbol::SymbolStr;

declare_lint! {
    pub NON_ASCII_IDENTS,
    Allow,
    "detects non-ASCII identifiers",
    crate_level_only
}

declare_lint! {
    pub UNCOMMON_CODEPOINTS,
    Warn,
    "detects uncommon Unicode codepoints in identifiers",
    crate_level_only
}

declare_lint! {
    pub CONFUSABLE_IDENTS,
    Warn,
    "detects visually confusable pairs between identifiers",
    crate_level_only
}

declare_lint! {
    pub MIXED_SCRIPT_CONFUSABLES,
    Warn,
    "detects Unicode scripts whose mixed script confusables codepoints are solely used",
    crate_level_only
}

declare_lint_pass!(NonAsciiIdents => [NON_ASCII_IDENTS, UNCOMMON_CODEPOINTS, CONFUSABLE_IDENTS, MIXED_SCRIPT_CONFUSABLES]);

impl EarlyLintPass for NonAsciiIdents {
    fn check_crate(&mut self, cx: &EarlyContext<'_>, _: &ast::Crate) {
        use rustc_session::lint::Level;
        use rustc_span::Span;
        use std::collections::BTreeMap;
        use unicode_security::GeneralSecurityProfile;
        use utils::CowBoxSymStr;

        let check_non_ascii_idents = cx.builder.lint_level(NON_ASCII_IDENTS).0 != Level::Allow;
        let check_uncommon_codepoints =
            cx.builder.lint_level(UNCOMMON_CODEPOINTS).0 != Level::Allow;
        let check_confusable_idents = cx.builder.lint_level(CONFUSABLE_IDENTS).0 != Level::Allow;
        let check_mixed_script_confusables =
            cx.builder.lint_level(MIXED_SCRIPT_CONFUSABLES).0 != Level::Allow;

        if !check_non_ascii_idents
            && !check_uncommon_codepoints
            && !check_confusable_idents
            && !check_mixed_script_confusables
        {
            return;
        }

        let mut has_non_ascii_idents = false;
        let symbols = cx.sess.parse_sess.symbol_gallery.symbols.lock();

        // Sort by `Span` so that error messages make sense with respect to the
        // order of identifier locations in the code.
        let mut symbols: Vec<_> = symbols.iter().collect();
        symbols.sort_by_key(|k| k.1);

        for (symbol, &sp) in symbols.iter() {
            let symbol_str = symbol.as_str();
            if symbol_str.is_ascii() {
                continue;
            }
            has_non_ascii_idents = true;
            cx.struct_span_lint(NON_ASCII_IDENTS, sp, |lint| {
                lint.build("identifier contains non-ASCII characters").emit()
            });
            if check_uncommon_codepoints
                && !symbol_str.chars().all(GeneralSecurityProfile::identifier_allowed)
            {
                cx.struct_span_lint(UNCOMMON_CODEPOINTS, sp, |lint| {
                    lint.build("identifier contains uncommon Unicode codepoints").emit()
                })
            }
        }

        if has_non_ascii_idents && check_confusable_idents {
            let mut skeleton_map: FxHashMap<CowBoxSymStr, (SymbolStr, Span, bool)> =
                FxHashMap::with_capacity_and_hasher(symbols.len(), Default::default());
            let mut str_buf = String::new();
            for (symbol, &sp) in symbols.iter() {
                fn calc_skeleton(symbol_str: &SymbolStr, buffer: &mut String) -> CowBoxSymStr {
                    use std::mem::replace;
                    use unicode_security::confusable_detection::skeleton;
                    buffer.clear();
                    buffer.extend(skeleton(symbol_str));
                    if *symbol_str == *buffer {
                        CowBoxSymStr::Interned(symbol_str.clone())
                    } else {
                        let owned = replace(buffer, String::new());
                        CowBoxSymStr::Owned(owned.into_boxed_str())
                    }
                }
                let symbol_str = symbol.as_str();
                let is_ascii = symbol_str.is_ascii();
                let skeleton = calc_skeleton(&symbol_str, &mut str_buf);
                skeleton_map
                    .entry(skeleton)
                    .and_modify(|(existing_symbolstr, existing_span, existing_is_ascii)| {
                        if !*existing_is_ascii || !is_ascii {
                            cx.struct_span_lint(CONFUSABLE_IDENTS, sp, |lint| {
                                lint.build(&format!(
                                    "identifier pair considered confusable between `{}` and `{}`",
                                    existing_symbolstr, symbol_str
                                ))
                                .span_label(
                                    *existing_span,
                                    "this is where the previous identifier occurred",
                                )
                                .emit();
                            });
                        }
                        if *existing_is_ascii && !is_ascii {
                            *existing_symbolstr = symbol_str.clone();
                            *existing_span = sp;
                            *existing_is_ascii = is_ascii;
                        }
                    })
                    .or_insert((symbol_str, sp, is_ascii));
            }
        }

        if has_non_ascii_idents && check_mixed_script_confusables {
            use unicode_security::is_potential_mixed_script_confusable_char;
            use unicode_security::mixed_script::AugmentedScriptSet;

            #[derive(Clone)]
            enum ScriptSetUsage {
                Suspicious(Vec<char>, Span),
                Verified,
            }

            let mut script_states: FxHashMap<AugmentedScriptSet, ScriptSetUsage> =
                FxHashMap::default();
            let latin_augmented_script_set = AugmentedScriptSet::for_char('A');
            script_states.insert(latin_augmented_script_set, ScriptSetUsage::Verified);

            let mut has_suspicous = false;
            for (symbol, &sp) in symbols.iter() {
                let symbol_str = symbol.as_str();
                for ch in symbol_str.chars() {
                    if ch.is_ascii() {
                        // all ascii characters are covered by exception.
                        continue;
                    }
                    if !GeneralSecurityProfile::identifier_allowed(ch) {
                        // this character is covered by `uncommon_codepoints` lint.
                        continue;
                    }
                    let augmented_script_set = AugmentedScriptSet::for_char(ch);
                    script_states
                        .entry(augmented_script_set)
                        .and_modify(|existing_state| {
                            if let ScriptSetUsage::Suspicious(ch_list, _) = existing_state {
                                if is_potential_mixed_script_confusable_char(ch) {
                                    ch_list.push(ch);
                                } else {
                                    *existing_state = ScriptSetUsage::Verified;
                                }
                            }
                        })
                        .or_insert_with(|| {
                            if !is_potential_mixed_script_confusable_char(ch) {
                                ScriptSetUsage::Verified
                            } else {
                                has_suspicous = true;
                                ScriptSetUsage::Suspicious(vec![ch], sp)
                            }
                        });
                }
            }

            if has_suspicous {
                let verified_augmented_script_sets = script_states
                    .iter()
                    .flat_map(|(k, v)| match v {
                        ScriptSetUsage::Verified => Some(*k),
                        _ => None,
                    })
                    .collect::<Vec<_>>();

                // we're sorting the output here.
                let mut lint_reports: BTreeMap<(Span, Vec<char>), AugmentedScriptSet> =
                    BTreeMap::new();

                'outerloop: for (augment_script_set, usage) in script_states {
                    let (mut ch_list, sp) = match usage {
                        ScriptSetUsage::Verified => continue,
                        ScriptSetUsage::Suspicious(ch_list, sp) => (ch_list, sp),
                    };

                    if augment_script_set.is_all() {
                        continue;
                    }

                    for existing in verified_augmented_script_sets.iter() {
                        if existing.is_all() {
                            continue;
                        }
                        let mut intersect = *existing;
                        intersect.intersect_with(augment_script_set);
                        if !intersect.is_empty() && !intersect.is_all() {
                            continue 'outerloop;
                        }
                    }

                    ch_list.sort();
                    ch_list.dedup();
                    lint_reports.insert((sp, ch_list), augment_script_set);
                }

                for ((sp, ch_list), script_set) in lint_reports {
                    cx.struct_span_lint(MIXED_SCRIPT_CONFUSABLES, sp, |lint| {
                        let message = format!(
                            "The usage of Script Group `{}` in this crate consists solely of mixed script confusables",
                            script_set);
                        let mut note = "The usage includes ".to_string();
                        for (idx, ch) in ch_list.into_iter().enumerate() {
                            if idx != 0 {
                                note += ", ";
                            }
                            let char_info = format!("'{}' (U+{:04X})", ch, ch as u32);
                            note += &char_info;
                        }
                        note += ".";
                        lint.build(&message).note(&note).note("Please recheck to make sure their usages are indeed what you want.").emit()
                    });
                }
            }
        }
    }
}

mod utils {
    use rustc_span::symbol::SymbolStr;
    use std::hash::{Hash, Hasher};
    use std::ops::Deref;

    pub(super) enum CowBoxSymStr {
        Interned(SymbolStr),
        Owned(Box<str>),
    }

    impl Deref for CowBoxSymStr {
        type Target = str;

        fn deref(&self) -> &str {
            match self {
                CowBoxSymStr::Interned(interned) => interned,
                CowBoxSymStr::Owned(ref owned) => owned,
            }
        }
    }

    impl Hash for CowBoxSymStr {
        #[inline]
        fn hash<H: Hasher>(&self, state: &mut H) {
            Hash::hash(&**self, state)
        }
    }

    impl PartialEq<CowBoxSymStr> for CowBoxSymStr {
        #[inline]
        fn eq(&self, other: &CowBoxSymStr) -> bool {
            PartialEq::eq(&**self, &**other)
        }
    }

    impl Eq for CowBoxSymStr {}
}
<commit_msg>Remove `CowBoxSymStr`.<commit_after>use crate::{EarlyContext, EarlyLintPass, LintContext};
use rustc_ast::ast;
use rustc_data_structures::fx::FxHashMap;
use rustc_span::symbol::Symbol;

declare_lint! {
    pub NON_ASCII_IDENTS,
    Allow,
    "detects non-ASCII identifiers",
    crate_level_only
}

declare_lint! {
    pub UNCOMMON_CODEPOINTS,
    Warn,
    "detects uncommon Unicode codepoints in identifiers",
    crate_level_only
}

declare_lint! {
    pub CONFUSABLE_IDENTS,
    Warn,
    "detects visually confusable pairs between identifiers",
    crate_level_only
}

declare_lint! {
    pub MIXED_SCRIPT_CONFUSABLES,
    Warn,
    "detects Unicode scripts whose mixed script confusables codepoints are solely used",
    crate_level_only
}

declare_lint_pass!(NonAsciiIdents => [NON_ASCII_IDENTS, UNCOMMON_CODEPOINTS, CONFUSABLE_IDENTS, MIXED_SCRIPT_CONFUSABLES]);

impl EarlyLintPass for NonAsciiIdents {
    fn check_crate(&mut self, cx: &EarlyContext<'_>, _: &ast::Crate) {
        use rustc_session::lint::Level;
        use rustc_span::Span;
        use std::collections::BTreeMap;
        use unicode_security::GeneralSecurityProfile;

        let check_non_ascii_idents = cx.builder.lint_level(NON_ASCII_IDENTS).0 != Level::Allow;
        let check_uncommon_codepoints =
            cx.builder.lint_level(UNCOMMON_CODEPOINTS).0 != Level::Allow;
        let check_confusable_idents = cx.builder.lint_level(CONFUSABLE_IDENTS).0 != Level::Allow;
        let check_mixed_script_confusables =
            cx.builder.lint_level(MIXED_SCRIPT_CONFUSABLES).0 != Level::Allow;

        if !check_non_ascii_idents
            && !check_uncommon_codepoints
            && !check_confusable_idents
            && !check_mixed_script_confusables
        {
            return;
        }

        let mut has_non_ascii_idents = false;
        let symbols = cx.sess.parse_sess.symbol_gallery.symbols.lock();

        // Sort by `Span` so that error messages make sense with respect to the
        // order of identifier locations in the code.
        let mut symbols: Vec<_> = symbols.iter().collect();
        symbols.sort_by_key(|k| k.1);

        for (symbol, &sp) in symbols.iter() {
            let symbol_str = symbol.as_str();
            if symbol_str.is_ascii() {
                continue;
            }
            has_non_ascii_idents = true;
            cx.struct_span_lint(NON_ASCII_IDENTS, sp, |lint| {
                lint.build("identifier contains non-ASCII characters").emit()
            });
            if check_uncommon_codepoints
                && !symbol_str.chars().all(GeneralSecurityProfile::identifier_allowed)
            {
                cx.struct_span_lint(UNCOMMON_CODEPOINTS, sp, |lint| {
                    lint.build("identifier contains uncommon Unicode codepoints").emit()
                })
            }
        }

        if has_non_ascii_idents && check_confusable_idents {
            let mut skeleton_map: FxHashMap<Symbol, (Symbol, Span, bool)> =
                FxHashMap::with_capacity_and_hasher(symbols.len(), Default::default());
            let mut skeleton_buf = String::new();

            for (&symbol, &sp) in symbols.iter() {
                use unicode_security::confusable_detection::skeleton;

                let symbol_str = symbol.as_str();
                let is_ascii = symbol_str.is_ascii();

                // Get the skeleton as a `Symbol`.
                skeleton_buf.clear();
                skeleton_buf.extend(skeleton(&symbol_str));
                let skeleton_sym = if *symbol_str == *skeleton_buf {
                    symbol
                } else {
                    Symbol::intern(&skeleton_buf)
                };

                skeleton_map
                    .entry(skeleton_sym)
                    .and_modify(|(existing_symbol, existing_span, existing_is_ascii)| {
                        if !*existing_is_ascii || !is_ascii {
                            cx.struct_span_lint(CONFUSABLE_IDENTS, sp, |lint| {
                                lint.build(&format!(
                                    "identifier pair considered confusable between `{}` and `{}`",
                                    existing_symbol.as_str(),
                                    symbol.as_str()
                                ))
                                .span_label(
                                    *existing_span,
                                    "this is where the previous identifier occurred",
                                )
                                .emit();
                            });
                        }
                        if *existing_is_ascii && !is_ascii {
                            *existing_symbol = symbol;
                            *existing_span = sp;
                            *existing_is_ascii = is_ascii;
                        }
                    })
                    .or_insert((symbol, sp, is_ascii));
            }
        }

        if has_non_ascii_idents && check_mixed_script_confusables {
            use unicode_security::is_potential_mixed_script_confusable_char;
            use unicode_security::mixed_script::AugmentedScriptSet;

            #[derive(Clone)]
            enum ScriptSetUsage {
                Suspicious(Vec<char>, Span),
                Verified,
            }

            let mut script_states: FxHashMap<AugmentedScriptSet, ScriptSetUsage> =
                FxHashMap::default();
            let latin_augmented_script_set = AugmentedScriptSet::for_char('A');
            script_states.insert(latin_augmented_script_set, ScriptSetUsage::Verified);

            let mut has_suspicous = false;
            for (symbol, &sp) in symbols.iter() {
                let symbol_str = symbol.as_str();
                for ch in symbol_str.chars() {
                    if ch.is_ascii() {
                        // all ascii characters are covered by exception.
                        continue;
                    }
                    if !GeneralSecurityProfile::identifier_allowed(ch) {
                        // this character is covered by `uncommon_codepoints` lint.
                        continue;
                    }
                    let augmented_script_set = AugmentedScriptSet::for_char(ch);
                    script_states
                        .entry(augmented_script_set)
                        .and_modify(|existing_state| {
                            if let ScriptSetUsage::Suspicious(ch_list, _) = existing_state {
                                if is_potential_mixed_script_confusable_char(ch) {
                                    ch_list.push(ch);
                                } else {
                                    *existing_state = ScriptSetUsage::Verified;
                                }
                            }
                        })
                        .or_insert_with(|| {
                            if !is_potential_mixed_script_confusable_char(ch) {
                                ScriptSetUsage::Verified
                            } else {
                                has_suspicous = true;
                                ScriptSetUsage::Suspicious(vec![ch], sp)
                            }
                        });
                }
            }

            if has_suspicous {
                let verified_augmented_script_sets = script_states
                    .iter()
                    .flat_map(|(k, v)| match v {
                        ScriptSetUsage::Verified => Some(*k),
                        _ => None,
                    })
                    .collect::<Vec<_>>();

                // we're sorting the output here.
                let mut lint_reports: BTreeMap<(Span, Vec<char>), AugmentedScriptSet> =
                    BTreeMap::new();

                'outerloop: for (augment_script_set, usage) in script_states {
                    let (mut ch_list, sp) = match usage {
                        ScriptSetUsage::Verified => continue,
                        ScriptSetUsage::Suspicious(ch_list, sp) => (ch_list, sp),
                    };

                    if augment_script_set.is_all() {
                        continue;
                    }

                    for existing in verified_augmented_script_sets.iter() {
                        if existing.is_all() {
                            continue;
                        }
                        let mut intersect = *existing;
                        intersect.intersect_with(augment_script_set);
                        if !intersect.is_empty() && !intersect.is_all() {
                            continue 'outerloop;
                        }
                    }

                    ch_list.sort();
                    ch_list.dedup();
                    lint_reports.insert((sp, ch_list), augment_script_set);
                }

                for ((sp, ch_list), script_set) in lint_reports {
                    cx.struct_span_lint(MIXED_SCRIPT_CONFUSABLES, sp, |lint| {
                        let message = format!(
                            "The usage of Script Group `{}` in this crate consists solely of mixed script confusables",
                            script_set);
                        let mut note = "The usage includes ".to_string();
                        for (idx, ch) in ch_list.into_iter().enumerate() {
                            if idx != 0 {
                                note += ", ";
                            }
                            let char_info = format!("'{}' (U+{:04X})", ch, ch as u32);
                            note += &char_info;
                        }
                        note += ".";
                        lint.build(&message).note(&note).note("Please recheck to make sure their usages are indeed what you want.").emit()
                    });
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use anyhow::{anyhow, Context, Result};
use async_std::sync::{Mutex, RwLock};
use async_std::task;
use async_std::{stream, task::JoinHandle};
use async_trait::async_trait;
use futures::StreamExt;
use homectl_types::device::CorrelatedColorTemperature;
use homectl_types::{
    device::{Device, DeviceId, DeviceState, Light},
    event::{Message, TxEventChannel},
    integration::{Integration, IntegrationActionPayload, IntegrationId},
};
use palette::Hsv;
use rust_tuyapi::{tuyadevice::TuyaDevice, Payload, PayloadStruct};
use serde::Deserialize;
use serde_json::{json, Value};
use std::ops::Range;
use std::{
    collections::HashMap,
    net::IpAddr,
    str::FromStr,
    sync::Arc,
    time::{Duration, SystemTime},
};

#[derive(Clone, Debug, Deserialize)]
pub struct TuyaDeviceConfig {
    name: String,
    local_key: String,
    ip: String,
    power_on_field: String,
    brightness_field: String,
    color_temp_field: String,
}

#[derive(Clone, Debug, Deserialize)]
pub struct TuyaConfig {
    devices: HashMap<DeviceId, TuyaDeviceConfig>,
}

type DeviceMutex = Arc<Mutex<()>>;

pub struct Tuya {
    id: IntegrationId,
    event_tx: TxEventChannel,
    config: TuyaConfig,
    device_mutexes: HashMap<DeviceId, DeviceMutex>,
    device_expected_states: Arc<RwLock<HashMap<DeviceId, Device>>>,
    poll_handle: Option<JoinHandle<()>>,
}

#[async_trait]
impl Integration for Tuya {
    fn new(id: &IntegrationId, config: &config::Value, event_tx: TxEventChannel) -> Result<Self> {
        let config = config
            .clone()
            .try_into()
            .context("Failed to deserialize config of Dummy integration")?;

        Ok(Tuya {
            id: id.clone(),
            config,
            event_tx,
            device_mutexes: HashMap::new(),
            device_expected_states: Arc::new(RwLock::new(HashMap::new())),
            poll_handle: None,
        })
    }

    async fn register(&mut self) -> Result<()> {
        self.device_mutexes = self
            .config
            .devices
            .keys()
            .map(|device_id| (device_id.clone(), Arc::new(Mutex::new(()))))
            .collect();

        for (device_id, device_config) in &self.config.devices {
            let device_mutex = self
                .device_mutexes
                .entry(device_id.clone())
                .or_insert_with(|| Arc::new(Mutex::new(())));

            // TODO: don't crash if tuya device not found here
            let device = get_tuya_state(device_id, &self.id, device_config, device_mutex).await?;

            {
                let mut device_expected_states = self.device_expected_states.write().await;
                device_expected_states.insert(device_id.clone(), device.clone());
            }

            self.event_tx
                .send(Message::IntegrationDeviceRefresh { device });
        }

        println!("registered tuya integration {}", self.id);

        Ok(())
    }

    async fn start(&mut self) -> Result<()> {
        println!("starting tuya integration {}", self.id);

        let config = self.config.clone();
        let integration_id = self.id.clone();
        let sender = self.event_tx.clone();
        let device_mutexes = self.device_mutexes.clone();
        let device_expected_states = self.device_expected_states.clone();

        self.poll_handle = Some(task::spawn(async {
            poll_lights(
                config,
                integration_id,
                sender,
                device_mutexes,
                device_expected_states,
            )
            .await
        }));

        println!("started tuya integration {}", self.id);

        Ok(())
    }

    async fn set_integration_device_state(&mut self, device: &Device) -> Result<()> {
        {
            let mut device_expected_states = self.device_expected_states.write().await;
            device_expected_states.insert(device.id.clone(), device.clone());
        }

        if let Some(poll_handle) = self.poll_handle.take() {
            poll_handle.cancel().await;
        }

        let config = self.config.clone();
        let integration_id = self.id.clone();
        let sender = self.event_tx.clone();
        let device_mutexes = self.device_mutexes.clone();
        let device_expected_states = self.device_expected_states.clone();

        self.poll_handle = Some(task::spawn(async {
            poll_lights(
                config,
                integration_id,
                sender,
                device_mutexes,
                device_expected_states,
            )
            .await
        }));

        let device_config = self.config.devices.get(&device.id).context(format!(
            "Could not find TuyaDeviceConfig for device with id {}",
            device.id,
        ))?;
        let device_mutex = self
            .device_mutexes
            .entry(device.id.clone())
            .or_insert_with(|| Arc::new(Mutex::new(())));
        set_tuya_state(device, device_config, device_mutex).await?;
        Ok(())
    }

    async fn run_integration_action(&mut self, _: &IntegrationActionPayload) -> Result<()> {
        Ok(())
    }
}

#[derive(Debug)]
struct TuyaState {
    power_on: bool,
    brightness: u32,
    color_temperature: u32,
}

fn to_tuya_state(device: &Device) -> Result<TuyaState> {
    let light_state = match device.state {
        DeviceState::Light(Light {
            brightness,
            color,
            power,
            transition_ms,
            ref cct,
        }) => Ok(Light {
            power,
            brightness,
            color,
            transition_ms,
            cct: cct.clone(),
        }),
        _ => Err(anyhow!("Unsupported device state")),
    }?;

    let color = light_state.color.unwrap_or_else(|| Hsv::new(0.0, 1.0, 1.0));

    let brightness = light_state.brightness.unwrap_or(1.0) * color.value;

    // Brightness goes from 10 to 1000 ¯\_(ツ)_/¯
    let brightness = f32::floor(brightness * 990.0) as u32 + 10;

    let cct = light_state.cct.map(|cct| cct.get_cct()).unwrap_or(4000.0);

    // Range of my bulbs is from 2700K - 4100K (and they express this as a
    // 0-1000 range), this is very likely not true for all Tuya bulbs
    let min_supported_temp = 2700.0;
    let max_supported_temp = 4100.0;

    // Scale the value into 0.0 - 1.0 range
    let q = (cct - min_supported_temp) / (max_supported_temp - min_supported_temp);
    let q = q.clamp(0.0, 1.0);

    // Scale the value into 0 - 1000 range
    let color_temperature = f32::floor(q * 1000.0) as u32;

    let state = TuyaState {
        power_on: light_state.power,
        brightness,
        color_temperature,
    };

    Ok(state)
}

async fn set_tuya_state(
    device: &Device,
    device_config: &TuyaDeviceConfig,
    device_mutex: &DeviceMutex,
) -> Result<()> {
    let device_mutex = device_mutex.lock().await;
    let tuya_state = to_tuya_state(device)?;

    let current_time = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_secs() as u32;

    {
        // Create a TuyaDevice, this is the type used to set/get status to/from a Tuya compatible
        // device.
        let tuya_device = TuyaDevice::create(
            "3.3",
            Some(&device_config.local_key),
            IpAddr::from_str(&device_config.ip).unwrap(),
        )?;

        let mut dps = HashMap::new();
        dps.insert(
            device_config.power_on_field.clone(),
            json!(tuya_state.power_on),
        );
        dps.insert(
            device_config.brightness_field.clone(),
            json!(tuya_state.brightness),
        );
        dps.insert(
            device_config.color_temp_field.clone(),
            json!(tuya_state.color_temperature),
        );

        // Create the payload to be sent, this will be serialized to the JSON format
        let payload = Payload::Struct(PayloadStruct {
            dev_id: device.id.to_string(),
            // gw_id: None,
            gw_id: Some(device.id.to_string()),
            uid: Some(device.id.to_string()),
            t: Some(current_time),
            dp_id: None,
            dps: Some(dps),
        });

        tuya_device.set(payload, 0)?;
    }

    drop(device_mutex);

    Ok(())
}

async fn get_tuya_state(
    device_id: &DeviceId,
    integration_id: &IntegrationId,
    device_config: &TuyaDeviceConfig,
    device_mutex: &DeviceMutex,
) -> Result<Device> {
    let device_mutex = device_mutex.lock().await;
    let current_time = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_secs() as u32;

    let response = {
        // Create a TuyaDevice, this is the type used to set/get status to/from a Tuya compatible
        // device.
        let tuya_device = TuyaDevice::create(
            "3.3",
            Some(&device_config.local_key),
            IpAddr::from_str(&device_config.ip).unwrap(),
        )?;

        // Create the payload to be sent, this will be serialized to the JSON format
        let payload = Payload::Struct(PayloadStruct {
            dev_id: device_id.to_string(),
            gw_id: Some(device_id.to_string()),
            uid: Some(device_id.to_string()),
            t: Some(current_time),
            dp_id: None,
            dps: Some(HashMap::new()),
        });

        tuya_device
            .get(payload, 0)
            .context(format!("Error while polling {}", device_config.name))?
    };

    drop(device_mutex);

    let first = response
        .first()
        .context("Could not find valid Tuya Message in response")?;

    if let Payload::Struct(PayloadStruct { dps: Some(dps), .. }) = &first.payload {
        let power = if let Some(Value::Bool(value)) = dps.get(&device_config.power_on_field) {
            *value
        } else {
            true
        };

        let brightness =
            if let Some(Value::Number(value)) = dps.get(&device_config.brightness_field) {
                // Brightness goes from 10 to 1000 ¯\_(ツ)_/¯
                let bri = value.as_u64().unwrap_or(1000);

                // Scale range to 0-990
                let bri = bri - 10;

                // Scale range to 0-1
                Some(bri as f32 / 990.0)
            } else {
                None
            };

        let cct = if let Some(Value::Number(value)) = dps.get(&device_config.color_temp_field) {
            // Range of my bulbs is from 2700K - 4100K (and they express this as a
            // 0-1000 range), this is very likely not true for all Tuya bulbs
            let min_supported_temp = 2700.0;
            let max_supported_temp = 4100.0;

            let ct = value.as_u64().unwrap_or(1000);

            // Scale range to 0-1
            let q = ct as f32 / 1000.0;

            let cct = q * (max_supported_temp - min_supported_temp) + min_supported_temp;

            Some(CorrelatedColorTemperature::new(
                cct,
                Range {
                    start: min_supported_temp,
                    end: max_supported_temp,
                },
            ))
        } else {
            None
        };

        let state = DeviceState::Light(Light::new_with_cct(power, brightness, cct, Some(1000)));

        let device = Device {
            id: device_id.clone(),
            name: device_config.name.clone(),
            integration_id: integration_id.clone(),
            scene: None,
            state,
            locked: false,
        };

        Ok(device)
    } else {
        Err(anyhow!("Unsupported Tuya response"))
    }
}

pub async fn do_refresh_lights(
    config: TuyaConfig,
    _integration_id: IntegrationId,
    _sender: TxEventChannel,
    device_mutexes: &HashMap<DeviceId, DeviceMutex>,
    device_expected_states: &Arc<RwLock<HashMap<DeviceId, Device>>>,
) -> Result<()> {
    // for (device_id, device_config) in &config.devices {
    // let device_mutex = device_mutexes.get(device_id).unwrap();
    // let device =
    //     get_tuya_state(device_id, &integration_id, device_config, device_mutex).await?;

    // Tuya devices seem to only be able to handle one TCP connection at once.
    // Keeping track of this using Mutexes seems to not be enough
    // async_std::task::sleep(Duration::from_millis(100)).await;

    // HACK:
    // Instead of polling for current state which seems to cause the devices
    // to lock up (no idea why), let's spam them with expected state once
    // every poll interval by lying to homectl about their current state
    // let impossible_cct = Some(CorrelatedColorTemperature::new(
    //     100000.0,
    //     Range {
    //         start: 0.0,
    //         end: 0.0,
    //     },
    // ));
    // let device = Device {
    //     id: device_id.clone(),
    //     name: device_config.name.clone(),
    //     integration_id: integration_id.clone(),
    //     scene: None,
    //     state: DeviceState::Light(Light::new_with_cct(true, None, impossible_cct, None)),
    //     locked: false,
    // };
    // sender.send(Message::IntegrationDeviceRefresh { device });

    // }

    let device_expected_states = { device_expected_states.read().await.clone() };

    for device in device_expected_states.values() {
        let device_config = config.devices.get(&device.id).context(format!(
            "Could not find TuyaDeviceConfig for device with id {}",
            device.id,
        ))?;
        let device_mutex = device_mutexes.get(&device.id).unwrap();
        set_tuya_state(device, device_config, device_mutex)
            .await
            .ok();
    }

    Ok(())
}

pub async fn poll_lights(
    config: TuyaConfig,
    integration_id: IntegrationId,
    sender: TxEventChannel,
    device_mutexes: HashMap<DeviceId, DeviceMutex>,
    device_expected_states: Arc<RwLock<HashMap<DeviceId, Device>>>,
) {
    let poll_rate = Duration::from_millis(1000);
    let mut interval = stream::interval(poll_rate);

    loop {
        interval.next().await;

        println!("{:?}: polling tuya devices", SystemTime::now());
        let sender = sender.clone();
        let result = do_refresh_lights(
            config.clone(),
            integration_id.clone(),
            sender,
            &device_mutexes,
            &device_expected_states,
        )
        .await;

        match result {
            Ok(()) => {}
            Err(e) => println!("Error while polling lights: {:?}", e),
        }
    }
}
<commit_msg>Fix possible Tuya crash on startup, fix Tuya device refresh<commit_after>use anyhow::{anyhow, Context, Result};
use async_std::sync::{Mutex, RwLock};
use async_std::task;
use async_std::{stream, task::JoinHandle};
use async_trait::async_trait;
use futures::StreamExt;
use homectl_types::device::CorrelatedColorTemperature;
use homectl_types::{
    device::{Device, DeviceId, DeviceState, Light},
    event::{Message, TxEventChannel},
    integration::{Integration, IntegrationActionPayload, IntegrationId},
};
use palette::Hsv;
use rust_tuyapi::{tuyadevice::TuyaDevice, Payload, PayloadStruct};
use serde::Deserialize;
use serde_json::{json, Value};
use std::ops::Range;
use std::{
    collections::HashMap,
    net::IpAddr,
    str::FromStr,
    sync::Arc,
    time::{Duration, SystemTime},
};

#[derive(Clone, Debug, Deserialize)]
pub struct TuyaDeviceConfig {
    name: String,
    local_key: String,
    ip: String,
    power_on_field: String,
    brightness_field: String,
    color_temp_field: String,
}

#[derive(Clone, Debug, Deserialize)]
pub struct TuyaConfig {
    devices: HashMap<DeviceId, TuyaDeviceConfig>,
}

type DeviceMutex = Arc<Mutex<()>>;

pub struct Tuya {
    id: IntegrationId,
    event_tx: TxEventChannel,
    config: TuyaConfig,
    device_mutexes: HashMap<DeviceId, DeviceMutex>,
    device_expected_states: Arc<RwLock<HashMap<DeviceId, Device>>>,
    poll_handle: Option<JoinHandle<()>>,
}

#[async_trait]
impl Integration for Tuya {
    fn new(id: &IntegrationId, config: &config::Value, event_tx: TxEventChannel) -> Result<Self> {
        let config = config
            .clone()
            .try_into()
            .context("Failed to deserialize config of Dummy integration")?;

        Ok(Tuya {
            id: id.clone(),
            config,
            event_tx,
            device_mutexes: HashMap::new(),
            device_expected_states: Arc::new(RwLock::new(HashMap::new())),
            poll_handle: None,
        })
    }

    async fn register(&mut self) -> Result<()> {
        self.device_mutexes = self
            .config
            .devices
            .keys()
            .map(|device_id| (device_id.clone(), Arc::new(Mutex::new(()))))
            .collect();

        let integration_id = self.id.clone();

        for (device_id, device_config) in &self.config.devices {
            let device_mutex = self
                .device_mutexes
                .entry(device_id.clone())
                .or_insert_with(|| Arc::new(Mutex::new(())));

            let integration_id = integration_id.clone();
            // TODO: don't crash if tuya device not found here
            let device = get_tuya_state(device_id, &self.id, device_config, device_mutex).await;
            let device = device.unwrap_or_else(|_| {
                println!("Failed to get initial state of Tuya device {}, creating Device with default state", device_config.name);

                Device {
                    id: device_id.clone(),
                    name: device_config.name.clone(),
                    integration_id,
                    scene: None,
                    state: DeviceState::Light(Light { power: false, brightness: None, color: None, cct: None, transition_ms: None }),
                    locked: false
                }
            });

            {
                let mut device_expected_states = self.device_expected_states.write().await;
                device_expected_states.insert(device_id.clone(), device.clone());
            }

            self.event_tx
                .send(Message::IntegrationDeviceRefresh { device });
        }

        println!("registered tuya integration {}", self.id);

        Ok(())
    }

    async fn start(&mut self) -> Result<()> {
        println!("starting tuya integration {}", self.id);

        let config = self.config.clone();
        let integration_id = self.id.clone();
        let sender = self.event_tx.clone();
        let device_mutexes = self.device_mutexes.clone();
        let device_expected_states = self.device_expected_states.clone();

        self.poll_handle = Some(task::spawn(async {
            poll_lights(
                config,
                integration_id,
                sender,
                device_mutexes,
                device_expected_states,
            )
            .await
        }));

        println!("started tuya integration {}", self.id);

        Ok(())
    }

    async fn set_integration_device_state(&mut self, device: &Device) -> Result<()> {
        {
            let mut device_expected_states = self.device_expected_states.write().await;
            device_expected_states.insert(device.id.clone(), device.clone());
        }

        if let Some(poll_handle) = self.poll_handle.take() {
            poll_handle.cancel().await;
        }

        let config = self.config.clone();
        let integration_id = self.id.clone();
        let sender = self.event_tx.clone();
        let device_mutexes = self.device_mutexes.clone();
        let device_expected_states = self.device_expected_states.clone();

        self.poll_handle = Some(task::spawn(async {
            poll_lights(
                config,
                integration_id,
                sender,
                device_mutexes,
                device_expected_states,
            )
            .await
        }));

        let device_config = self.config.devices.get(&device.id).context(format!(
            "Could not find TuyaDeviceConfig for device with id {}",
            device.id,
        ))?;
        let device_mutex = self
            .device_mutexes
            .entry(device.id.clone())
            .or_insert_with(|| Arc::new(Mutex::new(())));
        set_tuya_state(device, device_config, device_mutex).await?;
        Ok(())
    }

    async fn run_integration_action(&mut self, _: &IntegrationActionPayload) -> Result<()> {
        Ok(())
    }
}

#[derive(Debug)]
struct TuyaState {
    power_on: bool,
    brightness: u32,
    color_temperature: u32,
}

fn to_tuya_state(device: &Device) -> Result<TuyaState> {
    let light_state = match device.state {
        DeviceState::Light(Light {
            brightness,
            color,
            power,
            transition_ms,
            ref cct,
        }) => Ok(Light {
            power,
            brightness,
            color,
            transition_ms,
            cct: cct.clone(),
        }),
        _ => Err(anyhow!("Unsupported device state")),
    }?;

    let color = light_state.color.unwrap_or_else(|| Hsv::new(0.0, 1.0, 1.0));

    let brightness = light_state.brightness.unwrap_or(1.0) * color.value;

    // Brightness goes from 10 to 1000 ¯\_(ツ)_/¯
    let brightness = f32::floor(brightness * 990.0) as u32 + 10;

    let cct = light_state.cct.map(|cct| cct.get_cct()).unwrap_or(4000.0);

    // Range of my bulbs is from 2700K - 4100K (and they express this as a
    // 0-1000 range), this is very likely not true for all Tuya bulbs
    let min_supported_temp = 2700.0;
    let max_supported_temp = 4100.0;

    // Scale the value into 0.0 - 1.0 range
    let q = (cct - min_supported_temp) / (max_supported_temp - min_supported_temp);
    let q = q.clamp(0.0, 1.0);

    // Scale the value into 0 - 1000 range
    let color_temperature = f32::floor(q * 1000.0) as u32;

    let state = TuyaState {
        power_on: light_state.power,
        brightness,
        color_temperature,
    };

    Ok(state)
}

async fn set_tuya_state(
    device: &Device,
    device_config: &TuyaDeviceConfig,
    device_mutex: &DeviceMutex,
) -> Result<()> {
    let device_mutex = device_mutex.lock().await;
    let tuya_state = to_tuya_state(device)?;

    let current_time = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_secs() as u32;

    {
        // Create a TuyaDevice, this is the type used to set/get status to/from a Tuya compatible
        // device.
        let tuya_device = TuyaDevice::create(
            "3.3",
            Some(&device_config.local_key),
            IpAddr::from_str(&device_config.ip).unwrap(),
        )?;

        let mut dps = HashMap::new();
        dps.insert(
            device_config.power_on_field.clone(),
            json!(tuya_state.power_on),
        );
        dps.insert(
            device_config.brightness_field.clone(),
            json!(tuya_state.brightness),
        );
        dps.insert(
            device_config.color_temp_field.clone(),
            json!(tuya_state.color_temperature),
        );

        // Create the payload to be sent, this will be serialized to the JSON format
        let payload = Payload::Struct(PayloadStruct {
            dev_id: device.id.to_string(),
            // gw_id: None,
            gw_id: Some(device.id.to_string()),
            uid: Some(device.id.to_string()),
            t: Some(current_time),
            dp_id: None,
            dps: Some(dps),
        });

        tuya_device.set(payload, 0)?;
    }

    drop(device_mutex);

    Ok(())
}

async fn get_tuya_state(
    device_id: &DeviceId,
    integration_id: &IntegrationId,
    device_config: &TuyaDeviceConfig,
    device_mutex: &DeviceMutex,
) -> Result<Device> {
    let device_mutex = device_mutex.lock().await;
    let current_time = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_secs() as u32;

    let response = {
        // Create a TuyaDevice, this is the type used to set/get status to/from a Tuya compatible
        // device.
        let tuya_device = TuyaDevice::create(
            "3.3",
            Some(&device_config.local_key),
            IpAddr::from_str(&device_config.ip).unwrap(),
        )?;

        // Create the payload to be sent, this will be serialized to the JSON format
        let payload = Payload::Struct(PayloadStruct {
            dev_id: device_id.to_string(),
            gw_id: Some(device_id.to_string()),
            uid: Some(device_id.to_string()),
            t: Some(current_time),
            dp_id: None,
            dps: Some(HashMap::new()),
        });

        tuya_device
            .get(payload, 0)
            .context(format!("Error while polling {}", device_config.name))?
    };

    drop(device_mutex);

    let first = response
        .first()
        .context("Could not find valid Tuya Message in response")?;

    if let Payload::Struct(PayloadStruct { dps: Some(dps), .. }) = &first.payload {
        let power = if let Some(Value::Bool(value)) = dps.get(&device_config.power_on_field) {
            *value
        } else {
            true
        };

        let brightness =
            if let Some(Value::Number(value)) = dps.get(&device_config.brightness_field) {
                // Brightness goes from 10 to 1000 ¯\_(ツ)_/¯
                let bri = value.as_u64().unwrap_or(1000);

                // Scale range to 0-990
                let bri = bri - 10;

                // Scale range to 0-1
                Some(bri as f32 / 990.0)
            } else {
                None
            };

        let cct = if let Some(Value::Number(value)) = dps.get(&device_config.color_temp_field) {
            // Range of my bulbs is from 2700K - 4100K (and they express this as a
            // 0-1000 range), this is very likely not true for all Tuya bulbs
            let min_supported_temp = 2700.0;
            let max_supported_temp = 4100.0;

            let ct = value.as_u64().unwrap_or(1000);

            // Scale range to 0-1
            let q = ct as f32 / 1000.0;

            let cct = q * (max_supported_temp - min_supported_temp) + min_supported_temp;

            Some(CorrelatedColorTemperature::new(
                cct,
                Range {
                    start: min_supported_temp,
                    end: max_supported_temp,
                },
            ))
        } else {
            None
        };

        let state = DeviceState::Light(Light::new_with_cct(power, brightness, cct, Some(1000)));

        let device = Device {
            id: device_id.clone(),
            name: device_config.name.clone(),
            integration_id: integration_id.clone(),
            scene: None,
            state,
            locked: false,
        };

        Ok(device)
    } else {
        Err(anyhow!("Unsupported Tuya response"))
    }
}

pub async fn do_refresh_lights(
    config: TuyaConfig,
    _integration_id: IntegrationId,
    sender: TxEventChannel,
    device_mutexes: &HashMap<DeviceId, DeviceMutex>,
    device_expected_states: &Arc<RwLock<HashMap<DeviceId, Device>>>,
) -> Result<()> {
    // for (device_id, device_config) in &config.devices {
    //     let device_mutex = device_mutexes.get(device_id).unwrap();
    //     let device =
    //         get_tuya_state(device_id, &integration_id, device_config, device_mutex).await?;

    //     sender.send(Message::IntegrationDeviceRefresh { device });
    // }

    // Tuya devices seem to only be able to handle one TCP connection at once.
    // Keeping track of this using Mutexes seems to not be enough
    // async_std::task::sleep(Duration::from_millis(100)).await;

    // HACK:
    // Polling state seems to cause Tuya devices to lose WiFi connection more
    // often than if we just spam them with expected state, so let's do that.
    let device_expected_states = { device_expected_states.read().await.clone() };

    for device in device_expected_states.values() {
        let device_config = config.devices.get(&device.id).context(format!(
            "Could not find TuyaDeviceConfig for device with id {}",
            device.id,
        ))?;
        let device_mutex = device_mutexes.get(&device.id).unwrap();
        set_tuya_state(device, device_config, device_mutex)
            .await
            .ok();
    }

    // We still need to send our version of the device state to homectl core, in
    // case it has gone stale.
    for (_, device) in device_expected_states.into_iter() {
        sender.send(Message::IntegrationDeviceRefresh { device });
    }

    Ok(())
}

pub async fn poll_lights(
    config: TuyaConfig,
    integration_id: IntegrationId,
    sender: TxEventChannel,
    device_mutexes: HashMap<DeviceId, DeviceMutex>,
    device_expected_states: Arc<RwLock<HashMap<DeviceId, Device>>>,
) {
    let poll_rate = Duration::from_millis(1000);
    let mut interval = stream::interval(poll_rate);

    loop {
        interval.next().await;

        // println!("{:?}: polling tuya devices", SystemTime::now());
        let sender = sender.clone();
        let result = do_refresh_lights(
            config.clone(),
            integration_id.clone(),
            sender,
            &device_mutexes,
            &device_expected_states,
        )
        .await;

        match result {
            Ok(()) => {}
            Err(e) => println!("Error while polling lights: {:?}", e),
        }
    }
}
<|endoftext|>


<commit_before>//! Base interfaces defining a game whose state space can be searched with MCTS.

use std::cmp::Eq;
use std::fmt::Debug;
use std::hash::Hash;
use std::marker::Reflect;
use std::ops::{Add, AddAssign};

pub trait State: Debug + Hash + Eq + Clone + Reflect {
    type Action: Clone + Debug + Reflect;
    type Payoff: Payoff<State=Self>;
    type PlayerId: Debug + Reflect;

    fn active_player(&self) -> &Self::PlayerId;
    fn for_actions<F>(&self, f: F) where F: FnMut(Self::Action) -> bool;
    fn do_action(&mut self, action: &<Self as State>::Action);
    fn terminated(&self) -> bool;
}

pub trait Payoff: Debug + Add + AddAssign + Default {
    type State: State;
    type PlayerId: Debug + Reflect;

    fn from_state(state: &Self::State) -> Option<Self>;
    fn visits(&self) -> u32;
    fn score(&self, player: &Self::PlayerId) -> f32;
}

pub trait Statistics: Clone + Debug + Default + Reflect {
    type Payoff: Payoff;

    fn as_payoff(&self) -> Self::Payoff;
    fn increment(&self, payoff: &Self::Payoff);
}

pub trait Game: Debug + Clone + Reflect {
    type Action: Clone + Debug + Reflect;
    type PlayerId: Debug + Reflect;
    type Payoff: Payoff<PlayerId=Self::PlayerId, State=Self::State>;
    type State: State<Action=Self::Action, Payoff=Self::Payoff, PlayerId=Self::PlayerId>;
    type Statistics: Statistics<Payoff=Self::Payoff>;
}
<commit_msg>Don't use Reflect marker because it's been gotten rid of.<commit_after>//! Base interfaces defining a game whose state space can be searched with MCTS.

use std::cmp::Eq;
use std::fmt::Debug;
use std::hash::Hash;
use std::ops::{Add, AddAssign};

pub trait State: Debug + Hash + Eq + Clone {
    type Action: Clone + Debug;
    type Payoff: Payoff<State=Self>;
    type PlayerId: Debug;

    fn active_player(&self) -> &Self::PlayerId;
    fn for_actions<F>(&self, f: F) where F: FnMut(Self::Action) -> bool;
    fn do_action(&mut self, action: &<Self as State>::Action);
    fn terminated(&self) -> bool;
}

pub trait Payoff: Debug + Add + AddAssign + Default {
    type State: State;
    type PlayerId: Debug;

    fn from_state(state: &Self::State) -> Option<Self>;
    fn visits(&self) -> u32;
    fn score(&self, player: &Self::PlayerId) -> f32;
}

pub trait Statistics: Clone + Debug + Default {
    type Payoff: Payoff;

    fn as_payoff(&self) -> Self::Payoff;
    fn increment(&self, payoff: &Self::Payoff);
}

pub trait Game: Debug + Clone {
    type Action: Clone + Debug;
    type PlayerId: Debug;
    type Payoff: Payoff<PlayerId=Self::PlayerId, State=Self::State>;
    type State: State<Action=Self::Action, Payoff=Self::Payoff, PlayerId=Self::PlayerId>;
    type Statistics: Statistics<Payoff=Self::Payoff>;
}
<|endoftext|>


<commit_before>use std::io::{Result, Write};

use super::super::general::write_vec;

pub fn begin<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"#![allow(non_camel_case_types)]",
"#![allow(dead_code)]",
"",
"extern crate libc;",
"#[macro_use] extern crate bitflags;",
"extern crate glib_sys as glib_ffi;",
"extern crate gdk_sys as gdk_ffi;",
"extern crate pango_sys as pango_ffi;",
"",
"pub mod enums;",
"",
"use libc::{c_int, c_char, c_float, c_uint, c_double, c_long, c_void, time_t};",
"",
"pub use glib_ffi::{",
"    gboolean, GFALSE, GTRUE, gsize, gpointer, GType, GObject, GPermission,",
"    GList, GSList, GError, GValue};",
    ];

    write_vec(w, &v)
}

pub fn before_func<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"#[repr(C)]",
"pub struct GtkPaperSize; //boxed",
"#[repr(C)]",
"pub struct GtkTextIter; //boxed",
"#[repr(C)]",
"pub struct GtkTreePath; //boxed",
"#[repr(C)]",
"pub struct GtkTextAttributes;",
"#[repr(C)]",
"pub struct GtkRecentInfo;",
"#[repr(C)]",
"pub struct GtkRecentData {",
"    pub display_name: *mut c_char,",
"    pub description: *mut c_char,",
"    pub mime_type: *mut c_char,",
"    pub app_name: *mut c_char,",
"    pub app_exec: *mut c_char,",
"    pub groups: *mut *mut c_char,",
"    pub is_private: gboolean",
"}",
"#[repr(C)]",
"pub struct GtkRecentFilterInfo {",
"    pub contains: enums::RecentFilterFlags,",
"    pub uri: *const c_char,",
"    pub display_name: *const c_char,",
"    pub mime_type: *const c_char,",
"    pub applications: *const *const c_char,",
"    pub groups: *const *const c_char,",
"    pub age: c_int",
"}",
"#[repr(C)]",
"pub struct GtkTreeIter {",
"    pub stamp: c_int,",
"    pub user_data: *mut c_void,",
"    pub user_data2: *mut c_void,",
"    pub user_data3: *mut c_void",
"}",
"",
"pub const GTK_ENTRY_BUFFER_MAX_SIZE: u16 = ::std::u16::MAX;",
"",
"//pub type GtkTreeModelForeachFunc = fn(model: *mut GtkTreeModel, path: *mut GtkTreePath, iter: *mut GtkTreeIter, data: gpointer) -> gboolean;",
"",
"pub const GTK_STYLE_PROVIDER_PRIORITY_FALLBACK: u32 = 1;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_THEME: u32 = 200;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_SETTINGS: u32 = 400;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_APPLICATION: u32 = 600;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_USER: u32 = 800;",
    ];

    write_vec(w, &v)
}
<commit_msg>sys: Fix the libc imports in the predefined code<commit_after>use std::io::{Result, Write};

use super::super::general::write_vec;

pub fn begin<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"#![allow(non_camel_case_types)]",
"#![allow(dead_code)]",
"",
"extern crate libc;",
"#[macro_use] extern crate bitflags;",
"extern crate glib_sys as glib_ffi;",
"extern crate gdk_sys as gdk_ffi;",
"extern crate pango_sys as pango_ffi;",
"",
"pub mod enums;",
"",
"#[allow(unused_imports)]",
"use libc::{c_int, c_char, c_float, c_uint, c_double, c_long, c_void, size_t, ssize_t, time_t};",
"",
"pub use glib_ffi::{",
"    gboolean, GFALSE, GTRUE, gsize, gpointer, GType, GObject, GPermission,",
"    GList, GSList, GError, GValue};",
    ];

    write_vec(w, &v)
}

pub fn before_func<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"#[repr(C)]",
"pub struct GtkPaperSize; //boxed",
"#[repr(C)]",
"pub struct GtkTextIter; //boxed",
"#[repr(C)]",
"pub struct GtkTreePath; //boxed",
"#[repr(C)]",
"pub struct GtkTextAttributes;",
"#[repr(C)]",
"pub struct GtkRecentInfo;",
"#[repr(C)]",
"pub struct GtkRecentData {",
"    pub display_name: *mut c_char,",
"    pub description: *mut c_char,",
"    pub mime_type: *mut c_char,",
"    pub app_name: *mut c_char,",
"    pub app_exec: *mut c_char,",
"    pub groups: *mut *mut c_char,",
"    pub is_private: gboolean",
"}",
"#[repr(C)]",
"pub struct GtkRecentFilterInfo {",
"    pub contains: enums::RecentFilterFlags,",
"    pub uri: *const c_char,",
"    pub display_name: *const c_char,",
"    pub mime_type: *const c_char,",
"    pub applications: *const *const c_char,",
"    pub groups: *const *const c_char,",
"    pub age: c_int",
"}",
"#[repr(C)]",
"pub struct GtkTreeIter {",
"    pub stamp: c_int,",
"    pub user_data: *mut c_void,",
"    pub user_data2: *mut c_void,",
"    pub user_data3: *mut c_void",
"}",
"",
"pub const GTK_ENTRY_BUFFER_MAX_SIZE: u16 = ::std::u16::MAX;",
"",
"//pub type GtkTreeModelForeachFunc = fn(model: *mut GtkTreeModel, path: *mut GtkTreePath, iter: *mut GtkTreeIter, data: gpointer) -> gboolean;",
"",
"pub const GTK_STYLE_PROVIDER_PRIORITY_FALLBACK: u32 = 1;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_THEME: u32 = 200;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_SETTINGS: u32 = 400;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_APPLICATION: u32 = 600;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_USER: u32 = 800;",
    ];

    write_vec(w, &v)
}
<|endoftext|>


<commit_before>// main.rs
// Main entry point for a dining philosophers,
// a concurrency problem.
//
// taken from the rust-lang tutorial:
// https://doc.rust-lang.org/stable/book/dining-philosophers.html
//
// vim: ft=rust sw=4 ts=4

use std::thread;
use std::sync::{Mutex, Arc};

// a philosopher struct
struct Philosopher {
    name: String,
    left: usize,
    right: usize,
}

impl Philosopher {
    fn new(name: &str, left: usize, right: usize) -> Philosopher {
        Philosopher {
            name: name.to_string(),
            left: left,
            right: right,
        }
    }

    fn eat(&self, table: &Table) {
        let _left = table.forks[self.left].lock().unwrap();
        let _right = table.forks[self.right].lock().unwrap();

        println!("{} is eating.", self.name);

        thread::sleep_ms(1000);

        println!("{} is done eating.", self.name);
    }
}

struct Table {
    forks: Vec<Mutex<()>>,
}

fn main() {
    let table = Arc::new(Table { forks: vec![
        Mutex::new(()),
        Mutex::new(()),
        Mutex::new(()),
        Mutex::new(()),
        Mutex::new(()),
        ]});

    let philosophers = vec![
        Philosopher::new("Baruch Spinoza", 0, 1),
        Philosopher::new("Gilles Deleuze", 1, 2),
        Philosopher::new("Karl Marx", 2, 3),
        Philosopher::new("Friedich Nietzsche", 3, 4),
        Philosopher::new("Michel Foucault", 0, 4), // swap "hands" to prevent deadlock
        ];

    let handles: Vec<_> = philosophers.into_iter().map(|p| {
        let table = table.clone();

        thread::spawn(move || {
            p.eat(&table);
        })
    }).collect();

    for h in handles {
        h.join().unwrap();
    }
}

<commit_msg>+ multiple eating cycles<commit_after>// main.rs
// Main entry point for a dining philosophers,
// a concurrency problem.
//
// taken from the rust-lang tutorial:
// https://doc.rust-lang.org/stable/book/dining-philosophers.html
//
// vim: ft=rust sw=4 ts=4

use std::thread;
use std::sync::{Mutex, Arc};

// a philosopher struct
struct Philosopher {
    name: String,
    left: usize,
    right: usize,
}

impl Philosopher {
    fn new(name: &str, left: usize, right: usize) -> Philosopher {
        Philosopher {
            name: name.to_string(),
            left: left,
            right: right,
        }
    }

    fn eat(&self, table: &Table) {
        let _left = table.forks[self.left].lock().unwrap();
        let _right = table.forks[self.right].lock().unwrap();

        println!("{} has started eating.", self.name);

        thread::sleep_ms(1000);

        println!("{} has finished eating.", self.name);
    }

    fn mult_eat(&self, table: &Table, pause: u32, repeat: u32) {
        for x in 0..repeat {
            thread::sleep_ms(pause);
            println!("{} went to table for {}. time.", self.name, (x+1));
            self.eat(table);
        }
    }
}

struct Table {
    forks: Vec<Mutex<()>>,
}

fn main() {
    let table = Arc::new(Table { forks: vec![
        Mutex::new(()),
        Mutex::new(()),
        Mutex::new(()),
        Mutex::new(()),
        Mutex::new(()),
        ]});

    let philosophers = vec![
        Philosopher::new("Baruch Spinoza", 0, 1),
        Philosopher::new("Gilles Deleuze", 1, 2),
        Philosopher::new("Karl Marx", 2, 3),
        Philosopher::new("Friedich Nietzsche", 3, 4),
        Philosopher::new("Michel Foucault", 0, 4), // swap "hands" to prevent deadlock
        ];

    let handles: Vec<_> = philosophers.into_iter().map(|p| {
        let table = table.clone();

        thread::spawn(move || {
            //p.eat(&table);
            p.mult_eat(&table, 2000, 3);
        })
    }).collect();

    for h in handles {
        h.join().unwrap();
    }
}

<|endoftext|>


<commit_before>use core::mem;
use core::sync::atomic::AtomicBool;

use crate::syscall::FloatRegisters;

/// This must be used by the kernel to ensure that context switches are done atomically
/// Compare and exchange this to true when beginning a context switch on any CPU
/// The `Context::switch_to` function will set it back to false, allowing other CPU's to switch
/// This must be done, as no locks can be held on the stack during switch
pub static CONTEXT_SWITCH_LOCK: AtomicBool = AtomicBool::new(false);

const ST_RESERVED: u128 = 0xFFFF_FFFF_FFFF_0000_0000_0000_0000_0000;

#[derive(Clone, Debug)]
#[repr(C)]
pub struct Context {
    /// FX location
    fx: usize,
    /// Page table pointer
    cr3: usize,
    /// RFLAGS register
    rflags: usize,
    /// RBX register
    rbx: usize,
    /// R12 register
    r12: usize,
    /// R13 register
    r13: usize,
    /// R14 register
    r14: usize,
    /// R15 register
    r15: usize,
    /// Base pointer
    rbp: usize,
    /// Stack pointer
    rsp: usize,
    /// FX valid?
    loadable: AbiCompatBool,
}

#[repr(u8)]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum AbiCompatBool {
    False,
    True,
}

impl Context {
    pub fn new() -> Context {
        Context {
            loadable: AbiCompatBool::False,
            fx: 0,
            cr3: 0,
            rflags: 0,
            rbx: 0,
            r12: 0,
            r13: 0,
            r14: 0,
            r15: 0,
            rbp: 0,
            rsp: 0
        }
    }

    pub fn get_page_table(&mut self) -> usize {
        self.cr3
    }

    pub fn get_fx_regs(&self) -> Option<FloatRegisters> {
        if self.loadable == AbiCompatBool::False {
            return None;
        }
        let mut regs = unsafe { *(self.fx as *const FloatRegisters) };
        regs._reserved = 0;
        let mut new_st = regs.st_space;
        for st in &mut new_st {
            // Only allow access to the 80 lowest bits
            *st &= !ST_RESERVED;
        }
        regs.st_space = new_st;
        Some(regs)
    }

    pub fn set_fx_regs(&mut self, mut new: FloatRegisters) -> bool {
        if self.loadable == AbiCompatBool::False {
            return false;
        }

        {
            let old = unsafe { &*(self.fx as *const FloatRegisters) };
            new._reserved = old._reserved;
            let old_st = new.st_space;
            let mut new_st = new.st_space;
            for (new_st, old_st) in new_st.iter_mut().zip(&old_st) {
                *new_st &= !ST_RESERVED;
                *new_st |= old_st & ST_RESERVED;
            }
            new.st_space = new_st;

            // Make sure we don't use `old` from now on
        }

        unsafe {
            *(self.fx as *mut FloatRegisters) = new;
        }
        true
    }

    pub fn set_fx(&mut self, address: usize) {
        self.fx = address;
    }

    pub fn set_page_table(&mut self, address: usize) {
        self.cr3 = address;
    }

    pub fn set_stack(&mut self, address: usize) {
        self.rsp = address;
    }

    pub unsafe fn signal_stack(&mut self, handler: extern fn(usize), sig: u8) {
        self.push_stack(sig as usize);
        self.push_stack(handler as usize);
        self.push_stack(signal_handler_wrapper as usize);
    }

    pub unsafe fn push_stack(&mut self, value: usize) {
        self.rsp -= mem::size_of::<usize>();
        *(self.rsp as *mut usize) = value;
    }

    pub unsafe fn pop_stack(&mut self) -> usize {
        let value = *(self.rsp as *const usize);
        self.rsp += mem::size_of::<usize>();
        value
    }
}

/// Switch to the next context by restoring its stack and registers
/// Check disassembly!
#[cold]
#[inline(never)]
#[naked]
pub unsafe extern "C" fn switch_to(_prev: &mut Context, _next: &mut Context) {
    asm!(
        // As a quick reminder for those who are unfamiliar with the System V ABI (extern "C"):
        //
        // - the current parameters are passed in the registers `rdi`, `rsi`,
        // - we can modify scratch registers, e.g. rax
        // - we cannot change callee-preserved registers arbitrarily, e.g. rbx, which is why we
        //   store them here in the first place.
        "
        // load `prev.fx`
        mov rax, [rdi + 0x00]

        // save processor SSE/FPU/AVX state in `prev.fx` pointee
        fxsave64 [rax]

        // set `prev.loadable` to true
        mov BYTE PTR [rdi + 0x50], {true}
        // compare `next.loadable` with true
        cmp BYTE PTR [rsi + 0x50], {true}
        je switch_to.next_is_loadable

        fninit
        jmp switch_to.after_fx

        switch_to.next_is_loadable:
        mov rax, [rsi + 0x00]
        fxrstor64 [rax]

        switch_to.after_fx:
        // Save the current CR3, and load the next CR3 if not identical
        mov rcx, cr3
        mov [rdi + 0x08], rcx
        mov rax, [rsi + 0x08]
        cmp rax, rcx

        je switch_to.same_cr3
        mov cr3, rax

        switch_to.same_cr3:
        // Save old registers, and load new ones
        mov [rdi + 0x18], rbx
        mov rbx, [rsi + 0x18]

        mov [rdi + 0x20], r12
        mov r12, [rsi + 0x20]

        mov [rdi + 0x28], r13
        mov r13, [rsi + 0x28]

        mov [rdi + 0x30], r14
        mov r14, [rsi + 0x30]

        mov [rdi + 0x38], r15
        mov r15, [rsi + 0x38]

        mov [rdi + 0x40], rbp
        mov rbp, [rsi + 0x40]

        mov [rdi + 0x48], rsp
        mov rsp, [rsi + 0x48]

        // push RFLAGS (can only be modified via stack)
        pushfq
        // pop RFLAGS into `self.rflags`
        pop QWORD PTR [rdi + 0x10]

        // push `next.rflags`
        push QWORD PTR [rsi + 0x10]
        // pop into RFLAGS
        popfq

        // When we return, we cannot even guarantee that the return address on the stack, points to
        // the calling function, `context::switch`. Thus, we have to execute this Rust hook by
        // ourselves, which will unlock the contexts before the later switch.

        call {switch_hook}

        ",

        true = const(AbiCompatBool::True as u8),
        switch_hook = sym crate::context::switch_finish_hook,
    );
}

#[allow(dead_code)]
#[repr(packed)]
pub struct SignalHandlerStack {
    r11: usize,
    r10: usize,
    r9: usize,
    r8: usize,
    rsi: usize,
    rdi: usize,
    rdx: usize,
    rcx: usize,
    rax: usize,
    handler: extern fn(usize),
    sig: usize,
    rip: usize,
}

#[naked]
unsafe extern fn signal_handler_wrapper() {
    #[inline(never)]
    unsafe fn inner(stack: &SignalHandlerStack) {
        (stack.handler)(stack.sig);
    }

    // Push scratch registers
    asm!("push rax
        push rcx
        push rdx
        push rdi
        push rsi
        push r8
        push r9
        push r10
        push r11");

    // Get reference to stack variables
    let rsp: usize;
    asm!("mov {}, rsp", out(reg) rsp);

    // Call inner rust function
    inner(&*(rsp as *const SignalHandlerStack));

    // Pop scratch registers, error code, and return
    asm!("pop r11
        pop r10
        pop r9
        pop r8
        pop rsi
        pop rdi
        pop rdx
        pop rcx
        pop rax
        add rsp, 16");
}
<commit_msg>Rewrite signal_handler_wrapper as single asm block.<commit_after>use core::mem;
use core::sync::atomic::AtomicBool;

use crate::syscall::FloatRegisters;

/// This must be used by the kernel to ensure that context switches are done atomically
/// Compare and exchange this to true when beginning a context switch on any CPU
/// The `Context::switch_to` function will set it back to false, allowing other CPU's to switch
/// This must be done, as no locks can be held on the stack during switch
pub static CONTEXT_SWITCH_LOCK: AtomicBool = AtomicBool::new(false);

const ST_RESERVED: u128 = 0xFFFF_FFFF_FFFF_0000_0000_0000_0000_0000;

#[derive(Clone, Debug)]
#[repr(C)]
pub struct Context {
    /// FX location
    fx: usize,
    /// Page table pointer
    cr3: usize,
    /// RFLAGS register
    rflags: usize,
    /// RBX register
    rbx: usize,
    /// R12 register
    r12: usize,
    /// R13 register
    r13: usize,
    /// R14 register
    r14: usize,
    /// R15 register
    r15: usize,
    /// Base pointer
    rbp: usize,
    /// Stack pointer
    rsp: usize,
    /// FX valid?
    loadable: AbiCompatBool,
}

#[repr(u8)]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum AbiCompatBool {
    False,
    True,
}

impl Context {
    pub fn new() -> Context {
        Context {
            loadable: AbiCompatBool::False,
            fx: 0,
            cr3: 0,
            rflags: 0,
            rbx: 0,
            r12: 0,
            r13: 0,
            r14: 0,
            r15: 0,
            rbp: 0,
            rsp: 0
        }
    }

    pub fn get_page_table(&mut self) -> usize {
        self.cr3
    }

    pub fn get_fx_regs(&self) -> Option<FloatRegisters> {
        if self.loadable == AbiCompatBool::False {
            return None;
        }
        let mut regs = unsafe { *(self.fx as *const FloatRegisters) };
        regs._reserved = 0;
        let mut new_st = regs.st_space;
        for st in &mut new_st {
            // Only allow access to the 80 lowest bits
            *st &= !ST_RESERVED;
        }
        regs.st_space = new_st;
        Some(regs)
    }

    pub fn set_fx_regs(&mut self, mut new: FloatRegisters) -> bool {
        if self.loadable == AbiCompatBool::False {
            return false;
        }

        {
            let old = unsafe { &*(self.fx as *const FloatRegisters) };
            new._reserved = old._reserved;
            let old_st = new.st_space;
            let mut new_st = new.st_space;
            for (new_st, old_st) in new_st.iter_mut().zip(&old_st) {
                *new_st &= !ST_RESERVED;
                *new_st |= old_st & ST_RESERVED;
            }
            new.st_space = new_st;

            // Make sure we don't use `old` from now on
        }

        unsafe {
            *(self.fx as *mut FloatRegisters) = new;
        }
        true
    }

    pub fn set_fx(&mut self, address: usize) {
        self.fx = address;
    }

    pub fn set_page_table(&mut self, address: usize) {
        self.cr3 = address;
    }

    pub fn set_stack(&mut self, address: usize) {
        self.rsp = address;
    }

    pub unsafe fn signal_stack(&mut self, handler: extern fn(usize), sig: u8) {
        self.push_stack(sig as usize);
        self.push_stack(handler as usize);
        self.push_stack(signal_handler_wrapper as usize);
    }

    pub unsafe fn push_stack(&mut self, value: usize) {
        self.rsp -= mem::size_of::<usize>();
        *(self.rsp as *mut usize) = value;
    }

    pub unsafe fn pop_stack(&mut self) -> usize {
        let value = *(self.rsp as *const usize);
        self.rsp += mem::size_of::<usize>();
        value
    }
}

/// Switch to the next context by restoring its stack and registers
/// Check disassembly!
#[cold]
#[inline(never)]
#[naked]
pub unsafe extern "C" fn switch_to(_prev: &mut Context, _next: &mut Context) {
    asm!(
        // As a quick reminder for those who are unfamiliar with the System V ABI (extern "C"):
        //
        // - the current parameters are passed in the registers `rdi`, `rsi`,
        // - we can modify scratch registers, e.g. rax
        // - we cannot change callee-preserved registers arbitrarily, e.g. rbx, which is why we
        //   store them here in the first place.
        "
        // load `prev.fx`
        mov rax, [rdi + 0x00]

        // save processor SSE/FPU/AVX state in `prev.fx` pointee
        fxsave64 [rax]

        // set `prev.loadable` to true
        mov BYTE PTR [rdi + 0x50], {true}
        // compare `next.loadable` with true
        cmp BYTE PTR [rsi + 0x50], {true}
        je switch_to.next_is_loadable

        fninit
        jmp switch_to.after_fx

        switch_to.next_is_loadable:
        mov rax, [rsi + 0x00]
        fxrstor64 [rax]

        switch_to.after_fx:
        // Save the current CR3, and load the next CR3 if not identical
        mov rcx, cr3
        mov [rdi + 0x08], rcx
        mov rax, [rsi + 0x08]
        cmp rax, rcx

        je switch_to.same_cr3
        mov cr3, rax

        switch_to.same_cr3:
        // Save old registers, and load new ones
        mov [rdi + 0x18], rbx
        mov rbx, [rsi + 0x18]

        mov [rdi + 0x20], r12
        mov r12, [rsi + 0x20]

        mov [rdi + 0x28], r13
        mov r13, [rsi + 0x28]

        mov [rdi + 0x30], r14
        mov r14, [rsi + 0x30]

        mov [rdi + 0x38], r15
        mov r15, [rsi + 0x38]

        mov [rdi + 0x40], rbp
        mov rbp, [rsi + 0x40]

        mov [rdi + 0x48], rsp
        mov rsp, [rsi + 0x48]

        // push RFLAGS (can only be modified via stack)
        pushfq
        // pop RFLAGS into `self.rflags`
        pop QWORD PTR [rdi + 0x10]

        // push `next.rflags`
        push QWORD PTR [rsi + 0x10]
        // pop into RFLAGS
        popfq

        // When we return, we cannot even guarantee that the return address on the stack, points to
        // the calling function, `context::switch`. Thus, we have to execute this Rust hook by
        // ourselves, which will unlock the contexts before the later switch.

        call {switch_hook}

        ",

        true = const(AbiCompatBool::True as u8),
        switch_hook = sym crate::context::switch_finish_hook,
    );
}

#[allow(dead_code)]
#[repr(packed)]
pub struct SignalHandlerStack {
    r11: usize,
    r10: usize,
    r9: usize,
    r8: usize,
    rsi: usize,
    rdi: usize,
    rdx: usize,
    rcx: usize,
    rax: usize,
    handler: extern fn(usize),
    sig: usize,
    rip: usize,
}

#[naked]
unsafe extern fn signal_handler_wrapper() {
    #[inline(never)]
    unsafe extern "C" fn inner(stack: &SignalHandlerStack) {
        (stack.handler)(stack.sig);
    }

    // Push scratch registers
    asm!(
        "
            push rax
            push rcx
            push rdx
            push rdi
            push rsi
            push r8
            push r9
            push r10
            push r11

            mov rdi, rsp
            call {inner}

            pop r11
            pop r10
            pop r9
            pop r8
            pop rsi
            pop rdi
            pop rdx
            pop rcx
            pop rax
            add rsp, 16
        ",

        inner = sym inner,
    );
}
<|endoftext|>


<commit_before>use nix::sys::wait::*;
use nix::sys::signal;
use nix::unistd::Pid;
use tracer::*;
use breakpoint::*;
use ptrace_control::*;
use config::Config;
///
/// So we are either:
///     Waiting for an signal
///     Handling a stop
///     Handling an exit
///     Handling being signalled
///     Handling ptrace events
///     Handling an error
///     Continuing execution
///     Handling ptrace error
///
/// So we wait -> other thing -> continuing
///
/// Cannot do wait -> wait (unless non-polling, timeout?)
///
/// Or wait -> other thing -> other thing -> continue
///
/// Or wait -> other thing -> continue -> continue

/// Possible states when executing an inferior process. This is an attempt at
/// a platform agnostic abstracting to provide the potential of future
/// implementations for other operating systems and provides the implementation
/// of the test running state machine
/// T is data used to store the necessary process information to enable tracing
enum State<T> {
    /// Start state. Wait for test to appear and track time to enable timeout
    Start {
        hnd: T,
        start_time: u64,
    },
    /// Initialise: once test process appears instrument
    Initialise {
        hnd: T,
    },
    /// Waiting for breakpoint to be hit or test to end
    Waiting {
        hnd: T,
        start_time: u64,
    },
    /// Test process stopped, check coverage
    Stopped {
        hnd: T,
        stop_state:T,
    },
    /// Unrecoverable error occurred
    Unrecoverable,
    /// Test exited normally
    End,
}

/// Trait for state machines to implement
trait StateMachine<T> {
    /// Update the states
    fn step(self, traces: &mut TracerData, config: &Config) -> State<T>;
}

/// Handle to linux process state
#[derive(Debug, Clone, Copy)]
struct LinuxHnd {
    wait: Option<WaitStatus>,
    parent: Pid,
}

impl StateMachine<LinuxHnd> for State<LinuxHnd> {

    fn step(self, traces: &mut TracerData, config: &Config) -> State<LinuxHnd> {
        match self {
            State::Start{hnd, start_time} => {
                if let Some(res) = self.start(hnd) {
                    res
                } else if start_time > 1000 {
                    State::Unrecoverable::<_>
                } else {
                    State::Start{hnd: hnd, start_time:start_time+1}
                }
            },
            _ => State::Unrecoverable::<_>
        }
    }
}

impl State<LinuxHnd> {

    fn start(self, hnd: LinuxHnd) -> Option<State<LinuxHnd>> {
        match waitpid(hnd.parent, Some(WNOHANG)) {
            Ok(WaitStatus::Stopped(_, signal::SIGTRAP)) => {
                Some(State::Initialise::<LinuxHnd>{hnd})
            },
            Ok(_) => {
                println!("Unexpected signal from test");
                None
            },
            Err(e) => {
                println!("Error on start: {}", e);
                Some(State::Unrecoverable::<_>)
            },
        }
    }

}

<commit_msg>First version of initialisation and start states<commit_after>use std::collections::HashMap;
use nix::sys::wait::*;
use nix::sys::signal;
use nix::unistd::Pid;
use tracer::*;
use breakpoint::*;
use ptrace_control::*;
use config::Config;
///
/// So we are either:
///     Waiting for an signal
///     Handling a stop
///     Handling an exit
///     Handling being signalled
///     Handling ptrace events
///     Handling an error
///     Continuing execution
///     Handling ptrace error
///
/// So we wait -> other thing -> continuing
///
/// Cannot do wait -> wait (unless non-polling, timeout?)
///
/// Or wait -> other thing -> other thing -> continue
///
/// Or wait -> other thing -> continue -> continue

/// Possible states when executing an inferior process. This is an attempt at
/// a platform agnostic abstracting to provide the potential of future
/// implementations for other operating systems and provides the implementation
/// of the test running state machine
/// T is data used to store the necessary process information to enable tracing
pub enum TestState {
    /// Start state. Wait for test to appear and track time to enable timeout
    Start {
        start_time: u64,
    },
    /// Initialise: once test process appears instrument
    Initialise ,
    /// Waiting for breakpoint to be hit or test to end
    Waiting {
        start_time: u64,
    },
    /// Test process stopped, check coverage
    Stopped,
    /// Process timed out
    Timeout,
    /// Unrecoverable error occurred
    Unrecoverable,
    /// Test exited normally
    End,
}

/// Trait for state machines to implement
trait StateMachine<T> {
    /// Update the states
    fn step(self, data: &mut T, config: &Config) -> TestState;
}

/// Handle to linux process state
#[derive(Debug)]
struct LinuxData<'a> {
    wait: WaitStatus,
    current: Pid,
    parent: Pid,
    breakpoints: HashMap<u64, Breakpoint>,
    traces: &'a mut Vec<TracerData>,
}

impl <'a>LinuxData<'a> {
    fn new(traces: &'a mut Vec<TracerData>) -> LinuxData {
        LinuxData {
            wait: WaitStatus::StillAlive,
            current: Pid::from_raw(0),
            parent: Pid::from_raw(0),
            breakpoints: HashMap::new(),
            traces: traces
        }
    }

    fn check_for_tracee(&mut self) -> Option<TestState> {
        match waitpid(self.current, Some(WNOHANG)) {
            Ok(WaitStatus::StillAlive) => None,
            Ok(sig @ WaitStatus::Stopped(_, signal::SIGTRAP)) => {
                if let WaitStatus::Stopped(child, _) = sig {
                    self.current = child;
                }
                self.wait = sig;
                Some(TestState::Initialise)
            },
            Ok(s) => {
                println!("Unexpected signal when starting test");
                None
            },
            Err(e) => {
                println!("Error when starting test: {}", e);
                Some(TestState::Unrecoverable)
            },
        }
    }
}

impl <'a> StateMachine<LinuxData<'a>> for TestState {

    fn step(self, data: &mut LinuxData<'a>, config: &Config) -> TestState {
        match self {
            TestState::Start{start_time} => {
                if let Some(state) = data.check_for_tracee() {
                    state
                } else if start_time < 1000 {
                    TestState::Start{start_time: start_time + 1}
                } else {
                    TestState::Timeout
                }
            },
            TestState::Initialise => {
                if let Err(c) = trace_children(data.current) {
                    println!("Failed to trace child threads");
                }
                for trace in data.traces.iter() {
                    if let Some(addr) = trace.address {
                        match Breakpoint::new(data.current, addr) {
                            Ok(bp) => {
                                let _ = data.breakpoints.insert(addr, bp);
                            },
                            Err(e) => {
                                println!("Failed to instrument");
                            },
                        }
                    }
                }
                if let Ok(_) = continue_exec(data.parent, None) {
                    TestState::Waiting{start_time:0}
                } else {
                    TestState::Unrecoverable
                }
            },
            _ => TestState::Unrecoverable
        }
    }
}

<|endoftext|>