<commit_before>use std::collections::HashMap;
use std::collections::HashSet;
use std::borrow::BorrowMut;
use std::fmt;
use std::fmt::Display;
use std::cmp;

type Id = String;

/// Returns an `Id` with prefix `x` that is not in the set `vars`
fn fresh_var_like(x: &Id, vars: HashSet<Id>) -> Id {
    let mut suffix = 0;
    let mut result = x.clone();
    let len = result.len();
    loop {
        result.push_str(&suffix.to_string());

        if !vars.contains(&result) {
            return result;
        }

        suffix += 1;
        result.truncate(len);
    }
}

#[derive(Debug)]
pub enum EvaluationError {
    UnboundVariable(String),
    IllformedExpression,
}

impl fmt::Display for EvaluationError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::EvaluationError::*;
        match *self {
            UnboundVariable(ref id) => write!(f, "Encountered unbound variable: {}", id),
            IllformedExpression => write!(f, "Illformed expression"),
        }
    }
}

#[derive(Clone, Debug)]
pub enum Type {
    Var(Id),
    Fun(Box<Type>, Box<Type>),
    Forall(Id, Box<Type>),
}

impl PartialEq for Type {
    /// Equality is defined as alpha-equivalence of types
    fn eq(&self, other: &Type) -> bool {
        use self::Type::*;

        match (self,other) {
            (&Var(ref x),&Var(ref y)) => {
                *x == *y
            },
            (&Fun(ref t1, ref t2),&Fun(ref t3, ref t4)) => {
                *t1 == *t3 && *t2 == *t4
            },
            (&Forall(ref A, ref t1),&Forall(ref B, ref t2)) => {
                t1.subst(&Var(B.clone()), &A) == **t2
            },
            _ => false,
        }
    }
}

impl Eq for Type { }

impl Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Type::Var(ref x) => write!(f,"{}",x),
            Type::Fun(ref a, ref b) => write!(f,"({}) -> {}",a,b),
            Type::Forall(ref x, ref t) => write!(f,"(forall {}. {})", x, t),
        }
    }
}

impl Type {
    /// Returns the result of the capture avoiding substitution `self` {`t`/`X`}
    fn subst(&self, t: &Type, X: &str) -> Type {
        match *self {
            Type::Var(ref Y) => {
                if X == *Y {
                    t.clone()
                } else {
                    self.clone()
                }
            },
            Type::Fun(ref t1, ref t2) => {
                Type::Fun(Box::new(t1.subst(t,X)), Box::new(t2.subst(t,X)))
            },
            Type::Forall(ref Y, ref tau) => {
                if X == *Y {
                    self.clone()
                } else {
                    let mut fvs = t.free_vars();
                    if fvs.contains(Y) {
                        let mut vars = self.vars();
                        vars.extend(fvs.drain());
                        vars.insert(String::from(X));
                        let fresh_id = fresh_var_like(Y, vars);
                        let renamed = self.rename(Y, &fresh_id);
                        renamed.subst(t, X)
                    } else {
                        Type::Forall(Y.clone(), Box::new(tau.subst(t,X)))
                    }
                }
            },
        }
    }

    /// Returns an identical type with `var` renamed into `into`
    fn rename(&self, var: &str, into: &str) -> Type {
        match *self {
            Type::Var(ref X) => {
                if *X == var {
                    Type::Var(String::from(into))
                } else {
                    self.clone()
                }
            },
            Type::Fun(ref t1, ref t2) => {
                Type::Fun(Box::new(t1.rename(var, into)),
                        Box::new(t2.rename(var,into)))
            },
            Type::Forall(ref X, ref tau) => {
                let binder = {
                    if *X == var {
                        String::from(into)
                    } else {
                        X.clone()
                    }
                };
                Type::Forall(binder, Box::new(tau.rename(var, into)))
            },
        }
    }

    /// Returns a set the free variables contained in `self`
    fn free_vars(&self) -> HashSet<Id> {
        match *self {
            Type::Var(ref X) => {
                let mut fvs = HashSet::new();
                fvs.insert(X.clone());
                fvs
            },
            Type::Fun(ref t1, ref t2) => {
                let mut fvs = t1.free_vars();
                let mut t2_fvs = t2.free_vars();
                fvs.extend(t2_fvs.drain());
                fvs
            }
            Type::Forall(ref X, ref t) => {
                let mut fvs = t.free_vars();
                fvs.remove(X);
                fvs
            }
        }
    }

    /// Returns a set containing all of the variables contained in `self`
    fn vars(&self) -> HashSet<Id> {
        match *self {
            Type::Var(ref X) => {
                let mut vars = HashSet::new();
                vars.insert(X.clone());
                vars
            },
            Type::Fun(ref t1, ref t2) => {
                let mut vars = t1.vars();
                let mut t2_vars = t2.vars();
                vars.extend(t2_vars.drain());
                vars
            }
            Type::Forall(ref X, ref t) => {
                let mut vars = t.vars();
                vars.insert(X.clone());
                vars
            }
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Expr {
    Var(Id),
    Lam(Id, Box<Type>, Box<Expr>),
    App(Box<Expr>, Box<Expr>),
    TLam(Id, Box<Expr>),
    TApp(Box<Expr>, Box<Type>),
    Let(Id, Box<Type>, Box<Expr>, Box<Expr>),
    TLet(Id, Box<Type>, Box<Expr>),
}

impl Display for Expr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Expr::Var(ref x) => write!(f,"{}",x),
            Expr::Lam(ref x, ref t, ref e) => write!(f,"(\\{}:{}.{})",x,t,e),
            Expr::App(ref e1, ref e2) => write!(f,"({} {})",e1,e2),
            Expr::TLam(ref X, ref e) => write!(f,"(/\\{}.{})",X,e),
            Expr::TApp(ref e, ref t) => write!(f,"({} [{}])",e,t),
            Expr::Let(ref x, ref t, ref e1, ref e2) => write!(f,"(let {}:{} = {} in {})",x,t,e1,e2),
            Expr::TLet(ref X, ref t, ref e) => write!(f,"(Let {} = {} in {})",X,t,e),
        }
    }
}

impl Expr {
    /// Returns the normal form of `self` under call-by-value semantics
    pub fn eval(&self) -> Result<Expr,EvaluationError> {
        match *self {
            Expr::Var(ref x) => {
                Err(EvaluationError::UnboundVariable(x.clone()))
            },
            Expr::Lam(_,_,_) | Expr::TLam(_,_) => Ok(self.clone()),
            Expr::App(ref e1, ref e2) => {
                let v1 = e1.eval()?;
                let v2 = e2.eval()?;
                if let Expr::Lam(x,t,e3) = v1 {
                    let new_expr = e3.subst(&v2, &x);
                    new_expr.eval()
                } else {
                    Err(EvaluationError::IllformedExpression)
                }
            },
            Expr::TApp(ref e, ref t) => {
                let v = e.eval()?;
                if let Expr::TLam(X,e) = v {
                    let new_expr = e.subst_type(t, &X);
                    new_expr.eval()
                } else {
                    Err(EvaluationError::IllformedExpression)
                }
            },
            Expr::Let(ref x, ref t, ref e1, ref e2) => {
                Expr::App(
                    Box::new(Expr::Lam(
                        x.clone(),
                        t.clone(),
                        e2.clone())),
                    e1.clone()).eval()
            },
            Expr::TLet(ref X, ref t, ref e) => {
                Expr::TApp(
                    Box::new(Expr::TLam(
                        X.clone(),
                        e.clone())),
                    t.clone()).eval()
            },
        }
    }

    /// Returns the result of the capture avoiding substitution `self` {`e`/`x`}
    fn subst(&self, e: &Expr, x: &str) -> Expr {
        match *self {
            Expr::Var(ref y) => {
                if x == *y {
                    e.clone()
                } else {
                    self.clone()
                }
            },
            Expr::Lam(ref y,ref t,ref e1) => {
                if x == *y {
                    self.clone()
                } else {
                    // TODO: check freshness
                    unimplemented!()
                }
            },
            Expr::App(ref e1,ref e2) => {
                Expr::App(Box::new(e1.subst(e,x)), Box::new(e2.subst(e,x)))
            },
            Expr::TLam(ref X,ref e1) => {
                Expr::TLam(X.clone(), Box::new(e1.subst(e,x)))
            },
            Expr::TApp(ref e1,ref t) => {
                Expr::TApp(Box::new(e1.subst(e,x)), t.clone())
            },
            Expr::Let(ref y,ref t,ref e1,ref e2) => {
                let new_e1 = e1.subst(e, x);
                // Unclear
                unimplemented!()
            },
            Expr::TLet(ref X,ref t,ref e1) => {
                Expr::TLet(X.clone(), t.clone(), Box::new(e1.subst(e,x)))
            },
        }
    }

    fn subst_type(&self, to: &Type, from: &str) -> Expr {
        match *self {
            Expr::Var(ref y) => {
                self.clone()
            },
            Expr::Lam(ref y,ref t,ref e1) => {
                let new_t = t.subst(to, from);
                Expr::Lam(y.clone(), Box::new(new_t), Box::new(e1.subst_type(to,from)))
            },
            Expr::App(ref e1,ref e2) => {
                Expr::App(Box::new(e1.subst_type(to,from)), Box::new(e2.subst_type(to,from)))
            },
            Expr::TLam(ref X,ref e1) => {
                if *X == from {
                    self.clone()
                } else {
                    unimplemented!()
                }
            },
            Expr::TApp(ref e1,ref t) => {
                let new_t = t.subst(to, from);
                let new_e1 = e1.subst_type(to, from);
                Expr::TApp(Box::new(new_e1), Box::new(new_t))
            },
            Expr::Let(ref y,ref t,ref e1,ref e2) => {
                let new_t = t.subst(to, from);
                let new_e1 = e1.subst_type(to, from);
                let new_e2 = e2.subst_type(to, from);
                Expr::Let(y.clone(), Box::new(new_t), Box::new(new_e1), Box::new(new_e2))
            },
            Expr::TLet(ref X,ref t,ref e1) => {
                let new_t = t.subst(to, from);
                if *X == from {
                    Expr::TLet(X.clone(), Box::new(new_t), e1.clone())
                }
                else {
                    unimplemented!()
                }
            },
        }
    }

    /// Returns an identical expression with expression variable `var` renamed into `into`
    fn rename(&self, var: &str, into: &str) -> Expr {
        unimplemented!()
    }

    /// Returns an identical expression with type variable `var` renamed into `into`
    fn rename_type(&self, var: &str, into: &str) -> Expr {
        unimplemented!()
    }

    /// Returns a set the free variables contained in `self`
    fn free_vars(&self) -> HashSet<Id> {
        unimplemented!()
    }

    /// Returns a set containing all of the variables contained in `self`
    fn vars(&self) -> HashSet<Id> {
        unimplemented!()
    }

    /// Returns a set the free type variables contained in `self`
    fn free_type_vars(&self) -> HashSet<Id> {
        unimplemented!()
    }

    /// Returns a set containing all of the type variables contained in `self`
    fn type_vars(&self) -> HashSet<Id> {
        unimplemented!()
    }
}<commit_msg>Implement type vars of expressions<commit_after>use std::collections::HashMap;
use std::collections::HashSet;
use std::borrow::BorrowMut;
use std::fmt;
use std::fmt::Display;
use std::cmp;

type Id = String;

/// Returns an `Id` with prefix `x` that is not in the set `vars`
fn fresh_var_like(x: &Id, vars: HashSet<Id>) -> Id {
    let mut suffix = 0;
    let mut result = x.clone();
    let len = result.len();
    loop {
        result.push_str(&suffix.to_string());

        if !vars.contains(&result) {
            return result;
        }

        suffix += 1;
        result.truncate(len);
    }
}

#[derive(Debug)]
pub enum EvaluationError {
    UnboundVariable(String),
    IllformedExpression,
}

impl fmt::Display for EvaluationError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::EvaluationError::*;
        match *self {
            UnboundVariable(ref id) => write!(f, "Encountered unbound variable: {}", id),
            IllformedExpression => write!(f, "Illformed expression"),
        }
    }
}

#[derive(Clone, Debug)]
pub enum Type {
    Var(Id),
    Fun(Box<Type>, Box<Type>),
    Forall(Id, Box<Type>),
}

impl PartialEq for Type {
    /// Equality is defined as alpha-equivalence of types
    fn eq(&self, other: &Type) -> bool {
        use self::Type::*;

        match (self,other) {
            (&Var(ref x),&Var(ref y)) => {
                *x == *y
            },
            (&Fun(ref t1, ref t2),&Fun(ref t3, ref t4)) => {
                *t1 == *t3 && *t2 == *t4
            },
            (&Forall(ref A, ref t1),&Forall(ref B, ref t2)) => {
                t1.subst(&Var(B.clone()), &A) == **t2
            },
            _ => false,
        }
    }
}

impl Eq for Type { }

impl Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Type::Var(ref x) => write!(f,"{}",x),
            Type::Fun(ref a, ref b) => write!(f,"({}) -> {}",a,b),
            Type::Forall(ref x, ref t) => write!(f,"(forall {}. {})", x, t),
        }
    }
}

impl Type {
    /// Returns the result of the capture avoiding substitution `self` {`t`/`X`}
    fn subst(&self, t: &Type, X: &str) -> Type {
        match *self {
            Type::Var(ref Y) => {
                if X == *Y {
                    t.clone()
                } else {
                    self.clone()
                }
            },
            Type::Fun(ref t1, ref t2) => {
                Type::Fun(Box::new(t1.subst(t,X)), Box::new(t2.subst(t,X)))
            },
            Type::Forall(ref Y, ref tau) => {
                if X == *Y {
                    self.clone()
                } else {
                    let mut fvs = t.free_vars();
                    if fvs.contains(Y) {
                        let mut vars = self.vars();
                        vars.extend(fvs.drain());
                        vars.insert(String::from(X));
                        let fresh_id = fresh_var_like(Y, vars);
                        let renamed = self.rename(Y, &fresh_id);
                        renamed.subst(t, X)
                    } else {
                        Type::Forall(Y.clone(), Box::new(tau.subst(t,X)))
                    }
                }
            },
        }
    }

    /// Returns an identical type with `var` renamed into `into`
    fn rename(&self, var: &str, into: &str) -> Type {
        match *self {
            Type::Var(ref X) => {
                if *X == var {
                    Type::Var(String::from(into))
                } else {
                    self.clone()
                }
            },
            Type::Fun(ref t1, ref t2) => {
                Type::Fun(Box::new(t1.rename(var, into)),
                        Box::new(t2.rename(var,into)))
            },
            Type::Forall(ref X, ref tau) => {
                let binder = {
                    if *X == var {
                        String::from(into)
                    } else {
                        X.clone()
                    }
                };
                Type::Forall(binder, Box::new(tau.rename(var, into)))
            },
        }
    }

    /// Returns a set the free variables contained in `self`
    fn free_vars(&self) -> HashSet<Id> {
        match *self {
            Type::Var(ref X) => {
                let mut fvs = HashSet::new();
                fvs.insert(X.clone());
                fvs
            },
            Type::Fun(ref t1, ref t2) => {
                let mut fvs = t1.free_vars();
                let mut t2_fvs = t2.free_vars();
                fvs.extend(t2_fvs.drain());
                fvs
            }
            Type::Forall(ref X, ref t) => {
                let mut fvs = t.free_vars();
                fvs.remove(X);
                fvs
            }
        }
    }

    /// Returns a set containing all of the variables contained in `self`
    fn vars(&self) -> HashSet<Id> {
        match *self {
            Type::Var(ref X) => {
                let mut vars = HashSet::new();
                vars.insert(X.clone());
                vars
            },
            Type::Fun(ref t1, ref t2) => {
                let mut vars = t1.vars();
                let mut t2_vars = t2.vars();
                vars.extend(t2_vars.drain());
                vars
            }
            Type::Forall(ref X, ref t) => {
                let mut vars = t.vars();
                vars.insert(X.clone());
                vars
            }
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Expr {
    Var(Id),
    Lam(Id, Box<Type>, Box<Expr>),
    App(Box<Expr>, Box<Expr>),
    TLam(Id, Box<Expr>),
    TApp(Box<Expr>, Box<Type>),
    Let(Id, Box<Type>, Box<Expr>, Box<Expr>),
    TLet(Id, Box<Type>, Box<Expr>),
}

impl Display for Expr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Expr::Var(ref x) => write!(f,"{}",x),
            Expr::Lam(ref x, ref t, ref e) => write!(f,"(\\{}:{}.{})",x,t,e),
            Expr::App(ref e1, ref e2) => write!(f,"({} {})",e1,e2),
            Expr::TLam(ref X, ref e) => write!(f,"(/\\{}.{})",X,e),
            Expr::TApp(ref e, ref t) => write!(f,"({} [{}])",e,t),
            Expr::Let(ref x, ref t, ref e1, ref e2) => write!(f,"(let {}:{} = {} in {})",x,t,e1,e2),
            Expr::TLet(ref X, ref t, ref e) => write!(f,"(Let {} = {} in {})",X,t,e),
        }
    }
}

impl Expr {
    /// Returns the normal form of `self` under call-by-value semantics
    pub fn eval(&self) -> Result<Expr,EvaluationError> {
        match *self {
            Expr::Var(ref x) => {
                Err(EvaluationError::UnboundVariable(x.clone()))
            },
            Expr::Lam(_,_,_) | Expr::TLam(_,_) => Ok(self.clone()),
            Expr::App(ref e1, ref e2) => {
                let v1 = e1.eval()?;
                let v2 = e2.eval()?;
                if let Expr::Lam(x,t,e3) = v1 {
                    let new_expr = e3.subst(&v2, &x);
                    new_expr.eval()
                } else {
                    Err(EvaluationError::IllformedExpression)
                }
            },
            Expr::TApp(ref e, ref t) => {
                let v = e.eval()?;
                if let Expr::TLam(X,e) = v {
                    let new_expr = e.subst_type(t, &X);
                    new_expr.eval()
                } else {
                    Err(EvaluationError::IllformedExpression)
                }
            },
            Expr::Let(ref x, ref t, ref e1, ref e2) => {
                Expr::App(
                    Box::new(Expr::Lam(
                        x.clone(),
                        t.clone(),
                        e2.clone())),
                    e1.clone()).eval()
            },
            Expr::TLet(ref X, ref t, ref e) => {
                Expr::TApp(
                    Box::new(Expr::TLam(
                        X.clone(),
                        e.clone())),
                    t.clone()).eval()
            },
        }
    }

    /// Returns the result of the capture avoiding substitution `self` {`e`/`x`}
    fn subst(&self, e: &Expr, x: &str) -> Expr {
        match *self {
            Expr::Var(ref y) => {
                if x == *y {
                    e.clone()
                } else {
                    self.clone()
                }
            },
            Expr::Lam(ref y,ref t,ref e1) => {
                if x == *y {
                    self.clone()
                } else {
                    // TODO: check freshness
                    unimplemented!()
                }
            },
            Expr::App(ref e1,ref e2) => {
                Expr::App(Box::new(e1.subst(e,x)), Box::new(e2.subst(e,x)))
            },
            Expr::TLam(ref X,ref e1) => {
                Expr::TLam(X.clone(), Box::new(e1.subst(e,x)))
            },
            Expr::TApp(ref e1,ref t) => {
                Expr::TApp(Box::new(e1.subst(e,x)), t.clone())
            },
            Expr::Let(ref y,ref t,ref e1,ref e2) => {
                let new_e1 = e1.subst(e, x);
                // Unclear
                unimplemented!()
            },
            Expr::TLet(ref X,ref t,ref e1) => {
                Expr::TLet(X.clone(), t.clone(), Box::new(e1.subst(e,x)))
            },
        }
    }

    fn subst_type(&self, to: &Type, from: &str) -> Expr {
        match *self {
            Expr::Var(ref y) => {
                self.clone()
            },
            Expr::Lam(ref y,ref t,ref e1) => {
                let new_t = t.subst(to, from);
                Expr::Lam(y.clone(), Box::new(new_t), Box::new(e1.subst_type(to,from)))
            },
            Expr::App(ref e1,ref e2) => {
                Expr::App(Box::new(e1.subst_type(to,from)), Box::new(e2.subst_type(to,from)))
            },
            Expr::TLam(ref X,ref e1) => {
                if *X == from {
                    self.clone()
                } else {
                    unimplemented!()
                }
            },
            Expr::TApp(ref e1,ref t) => {
                let new_t = t.subst(to, from);
                let new_e1 = e1.subst_type(to, from);
                Expr::TApp(Box::new(new_e1), Box::new(new_t))
            },
            Expr::Let(ref y,ref t,ref e1,ref e2) => {
                let new_t = t.subst(to, from);
                let new_e1 = e1.subst_type(to, from);
                let new_e2 = e2.subst_type(to, from);
                Expr::Let(y.clone(), Box::new(new_t), Box::new(new_e1), Box::new(new_e2))
            },
            Expr::TLet(ref X,ref t,ref e1) => {
                let new_t = t.subst(to, from);
                if *X == from {
                    Expr::TLet(X.clone(), Box::new(new_t), e1.clone())
                }
                else {
                    unimplemented!()
                }
            },
        }
    }

    /// Returns an identical expression with expression variable `var` renamed into `into`
    fn rename(&self, var: &str, into: &str) -> Expr {
        unimplemented!()
    }

    /// Returns an identical expression with type variable `var` renamed into `into`
    fn rename_type(&self, var: &str, into: &str) -> Expr {
        unimplemented!()
    }

    /// Returns a set the free variables contained in `self`
    fn free_vars(&self) -> HashSet<Id> {
        unimplemented!()
    }

    /// Returns a set containing all of the variables contained in `self`
    fn vars(&self) -> HashSet<Id> {
        match *self {
            Expr::Var(ref x) => {
                let mut vars = HashSet::new();
                vars.insert(x.clone());
                vars
            },
            Expr::Lam(ref x, _, ref e) => {
                let mut vars = e.vars();
                vars.insert(x.clone());
                vars
            },
            Expr::App(ref e1, ref e2) => {
                let mut vars = e1.vars();
                let mut e2_vars = e2.vars();
                vars.extend(e2_vars.drain());
                vars
            },
            Expr::TLam(_, ref e) => {
                e.vars()
            },
            Expr::TApp(ref e, _) => {
                e.vars()
            },
            Expr::Let(ref x, _, ref e1, ref e2) => {
                let mut vars = e1.vars();
                let mut e2_vars = e2.vars();
                vars.extend(e2_vars.drain());
                vars.insert(x.clone());
                vars
            },
            Expr::TLet(_, _, ref e) => {
                e.vars()
            }
        }
    }

    /// Returns a set the free type variables contained in `self`
    fn free_type_vars(&self) -> HashSet<Id> {
        unimplemented!()
    }

    /// Returns a set containing all of the type variables contained in `self`
    fn type_vars(&self) -> HashSet<Id> {
        match *self {
            Expr::Var(_) => {
                HashSet::new()
            },
            Expr::Lam(_, ref t, ref e) => {
                let mut vars = t.vars();
                let mut e_vars = e.type_vars();
                vars.extend(e_vars.drain());
                vars
            },
            Expr::App(ref e1, ref e2) => {
                let mut vars = e1.type_vars();
                let mut e2_vars = e2.type_vars();
                vars.extend(e2_vars.drain());
                vars
            },
            Expr::TLam(ref X, ref e) => {
                let mut vars = e.type_vars();
                vars.insert(X.clone());
                vars
            },
            Expr::TApp(ref e, ref t) => {
                let mut vars = e.type_vars();
                let mut t_vars = e.vars();
                vars.extend(t_vars.drain());
                vars
            },
            Expr::Let(_, ref t, ref e1, ref e2) => {
                let mut vars = t.vars();
                let mut e1_vars = e1.type_vars();
                let mut e2_vars = e2.type_vars();
                vars.extend(e1_vars.drain());
                vars.extend(e2_vars.drain());
                vars
            },
            Expr::TLet(ref X, ref t, ref e) => {
                let mut vars = t.vars();
                let mut e_vars = e.type_vars();
                vars.extend(e_vars.drain());
                vars.insert(X.clone());
                vars
            }
        }
    }
}<|endoftext|>


<commit_before>use crate::intrinsics;
use crate::masks::*;
use crate::vector::ptr::{SimdConstPtr, SimdMutPtr};
use crate::vector::*;

/// A representation of a vector as an "array" with indices, implementing
/// operations applicable to any vector type based solely on "having lanes",
/// and describing relationships between vector and scalar types.
pub trait SimdArray<const LANES: usize>: crate::LanesAtMost32
where
    SimdUsize<LANES>: crate::LanesAtMost32,
    SimdIsize<LANES>: crate::LanesAtMost32,
    MaskSize<LANES>: crate::Mask,
    Self: Sized,
{
    /// The scalar type in every lane of this vector type.
    type Scalar: Copy + Sized;
    /// The number of lanes for this vector.
    const LANES: usize = LANES;

    /// Generates a SIMD vector with the same value in every lane.
    #[must_use]
    fn splat(val: Self::Scalar) -> Self;

    /// SIMD gather: construct a SIMD vector by reading from a slice, using potentially discontiguous indices.
    /// If an index is out of bounds, that lane instead selects the value from the "or" vector.
    /// ```
    /// # use core_simd::*;
    /// let vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
    /// let idxs = SimdUsize::<4>::from_array([9, 3, 0, 5]);
    /// let alt = SimdI32::from_array([-5, -4, -3, -2]);
    ///
    /// let result = SimdI32::<4>::gather_or(&vec, idxs, alt); // Note the lane that is out-of-bounds.
    /// assert_eq!(result, SimdI32::from_array([-5, 13, 10, 15]));
    /// ```
    #[must_use]
    #[inline]
    fn gather_or(slice: &[Self::Scalar], idxs: SimdUsize<LANES>, or: Self) -> Self {
        Self::gather_select(slice, MaskSize::splat(true), idxs, or)
    }

    /// SIMD gather: construct a SIMD vector by reading from a slice, using potentially discontiguous indices.
    /// Out-of-bounds indices instead use the default value for that lane (0).
    /// ```
    /// # use core_simd::*;
    /// let vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
    /// let idxs = SimdUsize::<4>::from_array([9, 3, 0, 5]);
    ///
    /// let result = SimdI32::<4>::gather_or_default(&vec, idxs); // Note the lane that is out-of-bounds.
    /// assert_eq!(result, SimdI32::from_array([0, 13, 10, 15]));
    /// ```
    #[must_use]
    #[inline]
    fn gather_or_default(slice: &[Self::Scalar], idxs: SimdUsize<LANES>) -> Self
    where
        Self::Scalar: Default,
    {
        Self::gather_or(slice, idxs, Self::splat(Self::Scalar::default()))
    }

    /// SIMD gather: construct a SIMD vector by reading from a slice, using potentially discontiguous indices.
    /// Out-of-bounds or masked indices instead select the value from the "or" vector.
    /// ```
    /// # use core_simd::*;
    /// let vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
    /// let idxs = SimdUsize::<4>::from_array([9, 3, 0, 5]);
    /// let alt = SimdI32::from_array([-5, -4, -3, -2]);
    /// let mask = MaskSize::from_array([true, true, true, false]); // Note the mask of the last lane.
    ///
    /// let result = SimdI32::<4>::gather_select(&vec, mask, idxs, alt); // Note the lane that is out-of-bounds.
    /// assert_eq!(result, SimdI32::from_array([-5, 13, 10, -2]));
    /// ```
    #[must_use]
    #[inline]
    fn gather_select(
        slice: &[Self::Scalar],
        mask: MaskSize<LANES>,
        idxs: SimdUsize<LANES>,
        or: Self,
    ) -> Self {
        let mask = (mask & idxs.lanes_lt(SimdUsize::splat(slice.len()))).to_int();
        let base_ptr = SimdConstPtr::splat(slice.as_ptr());
        // Ferris forgive me, I have done pointer arithmetic here.
        let ptrs = base_ptr.wrapping_add(idxs);
        // SAFETY: The ptrs have been bounds-masked to prevent memory-unsafe reads insha'allah
        unsafe { intrinsics::simd_gather(or, ptrs, mask) }
    }

    /// SIMD scatter: write a SIMD vector's values into a slice, using potentially discontiguous indices.
    /// Out-of-bounds indices are not written.
    /// ```
    /// # use core_simd::*;
    /// let mut vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
    /// let idxs = SimdUsize::<4>::from_array([9, 3, 0, 5]);
    /// let vals = SimdI32::from_array([-5, -4, -3, -2]);
    ///
    /// vals.scatter(&mut vec, idxs);
    /// assert_eq!(vec, vec![-3, 11, 12, -4, 14, -2, 16, 17, 18]);
    /// ```
    #[inline]
    fn scatter(self, slice: &mut [Self::Scalar], idxs: SimdUsize<LANES>) {
        self.scatter_select(slice, MaskSize::splat(true), idxs)
    }

    /// SIMD scatter: write a SIMD vector's values into a slice, using potentially discontiguous indices.
    /// Out-of-bounds or masked indices are not written.
    /// ```
    /// # use core_simd::*;
    /// let mut vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
    /// let idxs = SimdUsize::<4>::from_array([9, 3, 0, 5]);
    /// let vals = SimdI32::from_array([-5, -4, -3, -2]);
    /// let mask = MaskSize::from_array([true, true, true, false]); // Note the mask of the last lane.
    ///
    /// vals.scatter_select(&mut vec, mask, idxs);
    /// assert_eq!(vec, vec![-3, 11, 12, -4, 14, 15, 16, 17, 18]);
    /// ```
    #[inline]
    fn scatter_select(
        self,
        slice: &mut [Self::Scalar],
        mask: MaskSize<LANES>,
        idxs: SimdUsize<LANES>,
    ) {
        // We must construct our scatter mask before we derive a pointer!
        let mask = (mask & idxs.lanes_lt(SimdUsize::splat(slice.len()))).to_int();
        // SAFETY: This block works with *mut T derived from &mut 'a [T],
        // which means it is delicate in Rust's borrowing model, circa 2021:
        // &mut 'a [T] asserts uniqueness, so deriving &'a [T] invalidates live *mut Ts!
        // Even though this block is largely safe methods, it must be almost exactly this way
        // to prevent invalidating the raw ptrs while they're live.
        // Thus, entering this block requires all values to use being already ready:
        // 0. idxs we want to write to, which are used to construct the mask.
        // 1. mask, which depends on an initial &'a [T] and the idxs.
        // 2. actual values to scatter (self).
        // 3. &mut [T] which will become our base ptr.
        unsafe {
            // Now Entering ☢ *mut T Zone
            let base_ptr = SimdMutPtr::splat(slice.as_mut_ptr());
            // Ferris forgive me, I have done pointer arithmetic here.
            let ptrs = base_ptr.wrapping_add(idxs);
            // The ptrs have been bounds-masked to prevent memory-unsafe writes insha'allah
            intrinsics::simd_scatter(self, ptrs, mask)
            // Cleared ☢ *mut T Zone
        }
    }
}

macro_rules! impl_simdarray_for {
    ($simd:ident {type Scalar = $scalar:ident;}) => {
        impl<const LANES: usize> SimdArray<LANES> for $simd<LANES>
            where SimdUsize<LANES>: crate::LanesAtMost32,
            SimdIsize<LANES>: crate::LanesAtMost32,
            MaskSize<LANES>: crate::Mask,
            Self: crate::LanesAtMost32,
        {
            type Scalar = $scalar;

            #[must_use]
            #[inline]
            fn splat(val: Self::Scalar) -> Self {
                [val; LANES].into()
            }
        }
    };

    ($simd:ident $impl:tt) => {
        impl<const LANES: usize> SimdArray<LANES> for $simd<LANES>
            where SimdUsize<LANES>: crate::LanesAtMost32,
            SimdIsize<LANES>: crate::LanesAtMost32,
            MaskSize<LANES>: crate::Mask,
            Self: crate::LanesAtMost32,
        $impl
    }
}

impl_simdarray_for! {
    SimdUsize {
        type Scalar = usize;
    }
}

impl_simdarray_for! {
    SimdIsize {
        type Scalar = isize;
    }
}

impl_simdarray_for! {
    SimdI8 {
        type Scalar = i8;
    }
}

impl_simdarray_for! {
    SimdI16 {
        type Scalar = i16;
    }
}

impl_simdarray_for! {
    SimdI32 {
        type Scalar = i32;
    }
}

impl_simdarray_for! {
    SimdI64 {
        type Scalar = i64;
    }
}

impl_simdarray_for! {
    SimdU8 {
        type Scalar = u8;
    }
}

impl_simdarray_for! {
    SimdU16 {
        type Scalar = u16;
    }
}

impl_simdarray_for! {
    SimdU32 {
        type Scalar = u32;
    }
}

impl_simdarray_for! {
    SimdU64 {
        type Scalar = u64;
    }
}

impl_simdarray_for! {
    SimdF32 {
        type Scalar = f32;
    }
}

impl_simdarray_for! {
    SimdF64 {
        type Scalar = f64;
    }
}
<commit_msg>Document and test doubled writes in scatter<commit_after>use crate::intrinsics;
use crate::masks::*;
use crate::vector::ptr::{SimdConstPtr, SimdMutPtr};
use crate::vector::*;

/// A representation of a vector as an "array" with indices, implementing
/// operations applicable to any vector type based solely on "having lanes",
/// and describing relationships between vector and scalar types.
pub trait SimdArray<const LANES: usize>: crate::LanesAtMost32
where
    SimdUsize<LANES>: crate::LanesAtMost32,
    SimdIsize<LANES>: crate::LanesAtMost32,
    MaskSize<LANES>: crate::Mask,
    Self: Sized,
{
    /// The scalar type in every lane of this vector type.
    type Scalar: Copy + Sized;
    /// The number of lanes for this vector.
    const LANES: usize = LANES;

    /// Generates a SIMD vector with the same value in every lane.
    #[must_use]
    fn splat(val: Self::Scalar) -> Self;

    /// SIMD gather: construct a SIMD vector by reading from a slice, using potentially discontiguous indices.
    /// If an index is out of bounds, that lane instead selects the value from the "or" vector.
    /// ```
    /// # use core_simd::*;
    /// let vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
    /// let idxs = SimdUsize::<4>::from_array([9, 3, 0, 5]);
    /// let alt = SimdI32::from_array([-5, -4, -3, -2]);
    ///
    /// let result = SimdI32::<4>::gather_or(&vec, idxs, alt); // Note the lane that is out-of-bounds.
    /// assert_eq!(result, SimdI32::from_array([-5, 13, 10, 15]));
    /// ```
    #[must_use]
    #[inline]
    fn gather_or(slice: &[Self::Scalar], idxs: SimdUsize<LANES>, or: Self) -> Self {
        Self::gather_select(slice, MaskSize::splat(true), idxs, or)
    }

    /// SIMD gather: construct a SIMD vector by reading from a slice, using potentially discontiguous indices.
    /// Out-of-bounds indices instead use the default value for that lane (0).
    /// ```
    /// # use core_simd::*;
    /// let vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
    /// let idxs = SimdUsize::<4>::from_array([9, 3, 0, 5]);
    ///
    /// let result = SimdI32::<4>::gather_or_default(&vec, idxs); // Note the lane that is out-of-bounds.
    /// assert_eq!(result, SimdI32::from_array([0, 13, 10, 15]));
    /// ```
    #[must_use]
    #[inline]
    fn gather_or_default(slice: &[Self::Scalar], idxs: SimdUsize<LANES>) -> Self
    where
        Self::Scalar: Default,
    {
        Self::gather_or(slice, idxs, Self::splat(Self::Scalar::default()))
    }

    /// SIMD gather: construct a SIMD vector by reading from a slice, using potentially discontiguous indices.
    /// Out-of-bounds or masked indices instead select the value from the "or" vector.
    /// ```
    /// # use core_simd::*;
    /// let vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
    /// let idxs = SimdUsize::<4>::from_array([9, 3, 0, 5]);
    /// let alt = SimdI32::from_array([-5, -4, -3, -2]);
    /// let mask = MaskSize::from_array([true, true, true, false]); // Note the mask of the last lane.
    ///
    /// let result = SimdI32::<4>::gather_select(&vec, mask, idxs, alt); // Note the lane that is out-of-bounds.
    /// assert_eq!(result, SimdI32::from_array([-5, 13, 10, -2]));
    /// ```
    #[must_use]
    #[inline]
    fn gather_select(
        slice: &[Self::Scalar],
        mask: MaskSize<LANES>,
        idxs: SimdUsize<LANES>,
        or: Self,
    ) -> Self {
        let mask = (mask & idxs.lanes_lt(SimdUsize::splat(slice.len()))).to_int();
        let base_ptr = SimdConstPtr::splat(slice.as_ptr());
        // Ferris forgive me, I have done pointer arithmetic here.
        let ptrs = base_ptr.wrapping_add(idxs);
        // SAFETY: The ptrs have been bounds-masked to prevent memory-unsafe reads insha'allah
        unsafe { intrinsics::simd_gather(or, ptrs, mask) }
    }

    /// SIMD scatter: write a SIMD vector's values into a slice, using potentially discontiguous indices.
    /// Out-of-bounds indices are not written.
    /// `scatter` writes "in order", so if an index receives two writes, only the last is guaranteed.
    /// ```
    /// # use core_simd::*;
    /// let mut vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
    /// let idxs = SimdUsize::<4>::from_array([9, 3, 0, 0]);
    /// let vals = SimdI32::from_array([-27, 82, -41, 124]);
    ///
    /// vals.scatter(&mut vec, idxs); // index 0 receives two writes.
    /// assert_eq!(vec, vec![124, 11, 12, 82, 14, 15, 16, 17, 18]);
    /// ```
    #[inline]
    fn scatter(self, slice: &mut [Self::Scalar], idxs: SimdUsize<LANES>) {
        self.scatter_select(slice, MaskSize::splat(true), idxs)
    }

    /// SIMD scatter: write a SIMD vector's values into a slice, using potentially discontiguous indices.
    /// Out-of-bounds or masked indices are not written.
    /// `scatter_select` writes "in order", so if an index receives two writes, only the last is guaranteed.
    /// ```
    /// # use core_simd::*;
    /// let mut vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
    /// let idxs = SimdUsize::<4>::from_array([9, 3, 0, 0]);
    /// let vals = SimdI32::from_array([-27, 82, -41, 124]);
    /// let mask = MaskSize::from_array([true, true, true, false]); // Note the mask of the last lane.
    ///
    /// vals.scatter_select(&mut vec, mask, idxs); // index 0's second write is masked, thus omitted.
    /// assert_eq!(vec, vec![-41, 11, 12, 82, 14, 15, 16, 17, 18]);
    /// ```
    #[inline]
    fn scatter_select(
        self,
        slice: &mut [Self::Scalar],
        mask: MaskSize<LANES>,
        idxs: SimdUsize<LANES>,
    ) {
        // We must construct our scatter mask before we derive a pointer!
        let mask = (mask & idxs.lanes_lt(SimdUsize::splat(slice.len()))).to_int();
        // SAFETY: This block works with *mut T derived from &mut 'a [T],
        // which means it is delicate in Rust's borrowing model, circa 2021:
        // &mut 'a [T] asserts uniqueness, so deriving &'a [T] invalidates live *mut Ts!
        // Even though this block is largely safe methods, it must be almost exactly this way
        // to prevent invalidating the raw ptrs while they're live.
        // Thus, entering this block requires all values to use being already ready:
        // 0. idxs we want to write to, which are used to construct the mask.
        // 1. mask, which depends on an initial &'a [T] and the idxs.
        // 2. actual values to scatter (self).
        // 3. &mut [T] which will become our base ptr.
        unsafe {
            // Now Entering ☢ *mut T Zone
            let base_ptr = SimdMutPtr::splat(slice.as_mut_ptr());
            // Ferris forgive me, I have done pointer arithmetic here.
            let ptrs = base_ptr.wrapping_add(idxs);
            // The ptrs have been bounds-masked to prevent memory-unsafe writes insha'allah
            intrinsics::simd_scatter(self, ptrs, mask)
            // Cleared ☢ *mut T Zone
        }
    }
}

macro_rules! impl_simdarray_for {
    ($simd:ident {type Scalar = $scalar:ident;}) => {
        impl<const LANES: usize> SimdArray<LANES> for $simd<LANES>
            where SimdUsize<LANES>: crate::LanesAtMost32,
            SimdIsize<LANES>: crate::LanesAtMost32,
            MaskSize<LANES>: crate::Mask,
            Self: crate::LanesAtMost32,
        {
            type Scalar = $scalar;

            #[must_use]
            #[inline]
            fn splat(val: Self::Scalar) -> Self {
                [val; LANES].into()
            }
        }
    };

    ($simd:ident $impl:tt) => {
        impl<const LANES: usize> SimdArray<LANES> for $simd<LANES>
            where SimdUsize<LANES>: crate::LanesAtMost32,
            SimdIsize<LANES>: crate::LanesAtMost32,
            MaskSize<LANES>: crate::Mask,
            Self: crate::LanesAtMost32,
        $impl
    }
}

impl_simdarray_for! {
    SimdUsize {
        type Scalar = usize;
    }
}

impl_simdarray_for! {
    SimdIsize {
        type Scalar = isize;
    }
}

impl_simdarray_for! {
    SimdI8 {
        type Scalar = i8;
    }
}

impl_simdarray_for! {
    SimdI16 {
        type Scalar = i16;
    }
}

impl_simdarray_for! {
    SimdI32 {
        type Scalar = i32;
    }
}

impl_simdarray_for! {
    SimdI64 {
        type Scalar = i64;
    }
}

impl_simdarray_for! {
    SimdU8 {
        type Scalar = u8;
    }
}

impl_simdarray_for! {
    SimdU16 {
        type Scalar = u16;
    }
}

impl_simdarray_for! {
    SimdU32 {
        type Scalar = u32;
    }
}

impl_simdarray_for! {
    SimdU64 {
        type Scalar = u64;
    }
}

impl_simdarray_for! {
    SimdF32 {
        type Scalar = f32;
    }
}

impl_simdarray_for! {
    SimdF64 {
        type Scalar = f64;
    }
}
<|endoftext|>


<commit_before>extern crate devices;
extern crate sys_util;
extern crate x86_64;

use std::result;

use data_model::vm::MachineConfiguration;
use kvm::*;
use memory_model::{GuestAddress, GuestMemory, GuestMemoryError};
use sys_util::EventFd;
use x86_64::{cpuid, interrupts, regs};

pub const KVM_TSS_ADDRESS: usize = 0xfffbd000;

#[derive(Debug)]
pub enum Error {
    AlreadyRunning,
    CpuId(x86_64::cpuid::Error),
    GuestMemory(GuestMemoryError),
    Kvm(sys_util::Error),
    VmFd(sys_util::Error),
    VcpuFd(sys_util::Error),
    VmSetup(sys_util::Error),
    VcpuRun(sys_util::Error),
    GetSupportedCpusFailed(sys_util::Error),
    SetSupportedCpusFailed(sys_util::Error),
    NotEnoughMemory,
    NoMemoryEntry,
    LocalIntConfiguration(interrupts::Error),
    SetUserMemoryRegion(sys_util::Error),
    /// The kernel extends past the end of RAM
    KernelOffsetPastEnd,
    /// Error configuring the MSR registers
    MSRSConfiguration(regs::Error),
    /// Error configuring the general purpose registers
    REGSConfiguration(regs::Error),
    /// Error configuring the special registers
    SREGSConfiguration(regs::Error),
    /// Error configuring the floating point related registers
    FPUConfiguration(regs::Error),
    EventFd(sys_util::Error),
    Irq(sys_util::Error),
}
pub type Result<T> = result::Result<T, Error>;

impl ::std::convert::From<sys_util::Error> for Error {
    fn from(e: sys_util::Error) -> Error {
        Error::SetUserMemoryRegion(e)
    }
}

/// A wrapper around creating and using a VM.
pub struct Vm {
    fd: VmFd,
    guest_mem: Option<GuestMemory>,
}

impl Vm {
    /// Constructs a new `Vm` using the given `Kvm` instance.
    pub fn new(kvm: &Kvm) -> Result<Self> {
        //create fd for interacting with kvm-vm specific functions
        let vm_fd = VmFd::new(&kvm).map_err(Error::VmFd)?;

        Ok(Vm {
            fd: vm_fd,
            guest_mem: None,
        })
    }

    /// Currently this is x86 specific (because of the TSS address setup)
    pub fn memory_init(&mut self, guest_mem: GuestMemory) -> Result<()> {
        guest_mem.with_regions(|index, guest_addr, size, host_addr| {
            // Safe because the guest regions are guaranteed not to overlap.
            self.fd.set_user_memory_region(
                index as u32,
                guest_addr.offset() as u64,
                size as u64,
                host_addr as u64,
                0,
            )
        })?;
        self.guest_mem = Some(guest_mem);

        let tss_addr = GuestAddress(KVM_TSS_ADDRESS);
        self.fd
            .set_tss_address(tss_addr.offset())
            .map_err(Error::VmSetup)?;

        Ok(())
    }

    /// This function creates the irq chip and adds 2 interrupt events to the IRQ
    pub fn setup_irqchip(&self, com_evt_1_3: &EventFd, com_evt_2_4: &EventFd) -> Result<()> {
        self.fd.create_irq_chip().map_err(Error::VmSetup)?;

        self.fd.register_irqfd(com_evt_1_3, 4).map_err(Error::Irq)?;
        self.fd.register_irqfd(com_evt_2_4, 3).map_err(Error::Irq)?;

        Ok(())
    }

    pub fn create_pit(&self) -> Result<()> {
        self.fd.create_pit2().map_err(Error::VmSetup)?;
        Ok(())
    }

    /// Gets a reference to the guest memory owned by this VM.
    ///
    /// Note that `GuestMemory` does not include any device memory that may have been added after
    /// this VM was constructed.
    pub fn get_memory(&self) -> Option<&GuestMemory> {
        self.guest_mem.as_ref()
    }

    /// Gets a reference to the kvm file descriptor owned by this VM.
    ///
    pub fn get_fd(&self) -> &VmFd {
        &self.fd
    }
}

/// A wrapper around creating and using a kvm-based VCPU
pub struct Vcpu {
    cpuid: CpuId,
    fd: VcpuFd,
    id: u8,
}

impl Vcpu {
    /// Constructs a new VCPU for `vm`.
    ///
    /// The `id` argument is the CPU number between [0, max vcpus).
    pub fn new(id: u8, vm: &Vm) -> Result<Self> {
        let kvm_vcpu = VcpuFd::new(id, &vm.fd).map_err(Error::VcpuFd)?;
        // Initially the cpuid per vCPU is the one supported by this VM
        Ok(Vcpu {
            fd: kvm_vcpu,
            cpuid: vm.fd.get_supported_cpuid(),
            id,
        })
    }

    /// Returns a clone of the CPUID entries of this vCPU
    /// For now this function is only used for testing; the cfg(test) should be removed when
    /// this function will be used for configuring the cpu features
    #[cfg(test)]
    pub fn get_cpuid(&self) -> CpuId {
        return self.cpuid.clone();
    }

    /// /// Configures the vcpu and should be called once per vcpu from the vcpu's thread.
    ///
    /// # Arguments
    ///
    /// * `kernel_load_offset` - Offset from `guest_mem` at which the kernel starts.
    /// nr cpus is required for checking populating the kvm_cpuid2 entry for ebx and edx registers
    pub fn configure(
        &mut self,
        machine_config: &MachineConfiguration,
        kernel_start_addr: GuestAddress,
        vm: &Vm,
    ) -> Result<()> {
        // the MachineConfiguration has defaults for ht_enabled and vcpu_count hence it is safe to unwrap
        cpuid::filter_cpuid(
            self.id,
            machine_config.vcpu_count.unwrap(),
            machine_config.ht_enabled.unwrap(),
            &mut self.cpuid,
        ).map_err(|e| Error::CpuId(e))?;
        match machine_config.cpu_template {
            Some(template) => cpuid::set_cpuid_template(template, &mut self.cpuid),
            None => (),
        }

        self.fd
            .set_cpuid2(&self.cpuid)
            .map_err(Error::SetSupportedCpusFailed)?;

        regs::setup_msrs(&self.fd).map_err(Error::MSRSConfiguration)?;
        // Safe to unwrap because this method is called after the VM is configured
        let vm_memory = vm.get_memory().unwrap();
        regs::setup_regs(
            &self.fd,
            kernel_start_addr.offset() as u64,
            x86_64::layout::BOOT_STACK_POINTER as u64,
            x86_64::layout::ZERO_PAGE_START as u64,
        ).map_err(Error::REGSConfiguration)?;
        regs::setup_fpu(&self.fd).map_err(Error::FPUConfiguration)?;
        regs::setup_sregs(vm_memory, &self.fd).map_err(Error::SREGSConfiguration)?;
        interrupts::set_lint(&self.fd).map_err(Error::LocalIntConfiguration)?;
        Ok(())
    }

    /// Runs the VCPU until it exits, returning the reason.
    ///
    /// Note that the state of the VCPU and associated VM must be setup first for this to do
    /// anything useful.
    pub fn run(&self) -> Result<VcpuExit> {
        match self.fd.run() {
            Ok(v) => Ok(v),
            Err(e) => return Err(Error::VcpuRun(<sys_util::Error>::new(e.errno()))),
        }
    }
}

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

    #[test]
    fn create_vm() {
        let kvm = Kvm::new().unwrap();
        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x10000)]).unwrap();
        let mut vm = Vm::new(&kvm).expect("new vm failed");
        assert!(vm.memory_init(gm).is_ok());
    }

    #[test]
    fn get_memory() {
        let kvm = Kvm::new().unwrap();
        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x1000)]).unwrap();
        let mut vm = Vm::new(&kvm).expect("new vm failed");
        assert!(vm.memory_init(gm).is_ok());
        let obj_addr = GuestAddress(0xf0);
        vm.get_memory()
            .unwrap()
            .write_obj_at_addr(67u8, obj_addr)
            .unwrap();
        let read_val: u8 = vm
            .get_memory()
            .unwrap()
            .read_obj_from_addr(obj_addr)
            .unwrap();
        assert_eq!(read_val, 67u8);
    }

    #[test]
    fn create_vcpu() {
        let kvm = Kvm::new().unwrap();
        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x10000)]).unwrap();
        let mut vm = Vm::new(&kvm).expect("new vm failed");
        assert!(vm.memory_init(gm).is_ok());
        Vcpu::new(0, &mut vm).unwrap();
    }

    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    #[test]
    fn test_cpuid() {
        let kvm = Kvm::new().unwrap();
        let mut vm = Vm::new(&kvm).unwrap();
        let mut vcpu = Vcpu::new(0, &mut vm).unwrap();
        assert_eq!(vcpu.get_cpuid(), vm.fd.get_supported_cpuid());
        assert!(cpuid::filter_cpuid(0, 1, true, &mut vcpu.cpuid).is_ok());
        // Test using the T2 template
        cpuid::set_cpuid_template(CpuFeaturesTemplate::T2, &mut vcpu.cpuid);
        assert!(vcpu.fd.set_cpuid2(&vcpu.cpuid).is_ok());
        // Test using the C3 template
        cpuid::set_cpuid_template(CpuFeaturesTemplate::C3, &mut vcpu.cpuid);
        assert!(vcpu.fd.set_cpuid2(&vcpu.cpuid).is_ok());
    }

    #[test]
    fn run_code() {
        use std::io::{self, Write};
        // This example based on https://lwn.net/Articles/658511/
        let code = [
            0xba, 0xf8, 0x03, /* mov $0x3f8, %dx */
            0x00, 0xd8, /* add %bl, %al */
            0x04, '0' as u8, /* add $'0', %al */
            0xee,      /* out %al, (%dx) */
            0xb0, '\n' as u8, /* mov $'\n', %al */
            0xee,       /* out %al, (%dx) */
            0xf4,       /* hlt */
        ];

        let mem_size = 0x1000;
        let load_addr = GuestAddress(0x1000);
        let mem = GuestMemory::new(&vec![(load_addr, mem_size)]).unwrap();

        let kvm = Kvm::new().expect("new kvm failed");
        let mut vm = Vm::new(&kvm).expect("new vm failed");
        assert!(vm.memory_init(mem).is_ok());
        vm.get_memory()
            .unwrap()
            .write_slice_at_addr(&code, load_addr)
            .expect("Writing code to memory failed.");

        let vcpu = Vcpu::new(0, &mut vm).expect("new vcpu failed");

        let mut vcpu_sregs = vcpu.fd.get_sregs().expect("get sregs failed");
        assert_ne!(vcpu_sregs.cs.base, 0);
        assert_ne!(vcpu_sregs.cs.selector, 0);
        vcpu_sregs.cs.base = 0;
        vcpu_sregs.cs.selector = 0;
        vcpu.fd.set_sregs(&vcpu_sregs).expect("set sregs failed");

        let mut vcpu_regs = vcpu.fd.get_regs().expect("get regs failed");
        vcpu_regs.rip = 0x1000;
        vcpu_regs.rax = 2;
        vcpu_regs.rbx = 3;
        vcpu_regs.rflags = 2;
        vcpu.fd.set_regs(&vcpu_regs).expect("set regs failed");

        loop {
            match vcpu.run().expect("run failed") {
                VcpuExit::IoOut(0x3f8, data) => {
                    assert_eq!(data.len(), 1);
                    io::stdout().write(data).unwrap();
                }
                VcpuExit::Hlt => {
                    io::stdout().write(b"KVM_EXIT_HLT\n").unwrap();
                    break;
                }
                r => panic!("unexpected exit reason: {:?}", r),
            }
        }
    }
}
<commit_msg>vmm: log guest memory mapping address<commit_after>extern crate devices;
extern crate sys_util;
extern crate x86_64;

use std::result;

use data_model::vm::MachineConfiguration;
use kvm::*;
use memory_model::{GuestAddress, GuestMemory, GuestMemoryError};
use sys_util::EventFd;
use x86_64::{cpuid, interrupts, regs};

pub const KVM_TSS_ADDRESS: usize = 0xfffbd000;

#[derive(Debug)]
pub enum Error {
    AlreadyRunning,
    CpuId(x86_64::cpuid::Error),
    GuestMemory(GuestMemoryError),
    Kvm(sys_util::Error),
    VmFd(sys_util::Error),
    VcpuFd(sys_util::Error),
    VmSetup(sys_util::Error),
    VcpuRun(sys_util::Error),
    GetSupportedCpusFailed(sys_util::Error),
    SetSupportedCpusFailed(sys_util::Error),
    NotEnoughMemory,
    NoMemoryEntry,
    LocalIntConfiguration(interrupts::Error),
    SetUserMemoryRegion(sys_util::Error),
    /// The kernel extends past the end of RAM
    KernelOffsetPastEnd,
    /// Error configuring the MSR registers
    MSRSConfiguration(regs::Error),
    /// Error configuring the general purpose registers
    REGSConfiguration(regs::Error),
    /// Error configuring the special registers
    SREGSConfiguration(regs::Error),
    /// Error configuring the floating point related registers
    FPUConfiguration(regs::Error),
    EventFd(sys_util::Error),
    Irq(sys_util::Error),
}
pub type Result<T> = result::Result<T, Error>;

impl ::std::convert::From<sys_util::Error> for Error {
    fn from(e: sys_util::Error) -> Error {
        Error::SetUserMemoryRegion(e)
    }
}

/// A wrapper around creating and using a VM.
pub struct Vm {
    fd: VmFd,
    guest_mem: Option<GuestMemory>,
}

impl Vm {
    /// Constructs a new `Vm` using the given `Kvm` instance.
    pub fn new(kvm: &Kvm) -> Result<Self> {
        //create fd for interacting with kvm-vm specific functions
        let vm_fd = VmFd::new(&kvm).map_err(Error::VmFd)?;

        Ok(Vm {
            fd: vm_fd,
            guest_mem: None,
        })
    }

    /// Currently this is x86 specific (because of the TSS address setup)
    pub fn memory_init(&mut self, guest_mem: GuestMemory) -> Result<()> {
        guest_mem.with_regions(|index, guest_addr, size, host_addr| {
            info!("Guest memory starts at {:x?}", host_addr);
            // Safe because the guest regions are guaranteed not to overlap.
            self.fd.set_user_memory_region(
                index as u32,
                guest_addr.offset() as u64,
                size as u64,
                host_addr as u64,
                0,
            )
        })?;
        self.guest_mem = Some(guest_mem);

        let tss_addr = GuestAddress(KVM_TSS_ADDRESS);
        self.fd
            .set_tss_address(tss_addr.offset())
            .map_err(Error::VmSetup)?;

        Ok(())
    }

    /// This function creates the irq chip and adds 2 interrupt events to the IRQ
    pub fn setup_irqchip(&self, com_evt_1_3: &EventFd, com_evt_2_4: &EventFd) -> Result<()> {
        self.fd.create_irq_chip().map_err(Error::VmSetup)?;

        self.fd.register_irqfd(com_evt_1_3, 4).map_err(Error::Irq)?;
        self.fd.register_irqfd(com_evt_2_4, 3).map_err(Error::Irq)?;

        Ok(())
    }

    pub fn create_pit(&self) -> Result<()> {
        self.fd.create_pit2().map_err(Error::VmSetup)?;
        Ok(())
    }

    /// Gets a reference to the guest memory owned by this VM.
    ///
    /// Note that `GuestMemory` does not include any device memory that may have been added after
    /// this VM was constructed.
    pub fn get_memory(&self) -> Option<&GuestMemory> {
        self.guest_mem.as_ref()
    }

    /// Gets a reference to the kvm file descriptor owned by this VM.
    ///
    pub fn get_fd(&self) -> &VmFd {
        &self.fd
    }
}

/// A wrapper around creating and using a kvm-based VCPU
pub struct Vcpu {
    cpuid: CpuId,
    fd: VcpuFd,
    id: u8,
}

impl Vcpu {
    /// Constructs a new VCPU for `vm`.
    ///
    /// The `id` argument is the CPU number between [0, max vcpus).
    pub fn new(id: u8, vm: &Vm) -> Result<Self> {
        let kvm_vcpu = VcpuFd::new(id, &vm.fd).map_err(Error::VcpuFd)?;
        // Initially the cpuid per vCPU is the one supported by this VM
        Ok(Vcpu {
            fd: kvm_vcpu,
            cpuid: vm.fd.get_supported_cpuid(),
            id,
        })
    }

    /// Returns a clone of the CPUID entries of this vCPU
    /// For now this function is only used for testing; the cfg(test) should be removed when
    /// this function will be used for configuring the cpu features
    #[cfg(test)]
    pub fn get_cpuid(&self) -> CpuId {
        return self.cpuid.clone();
    }

    /// /// Configures the vcpu and should be called once per vcpu from the vcpu's thread.
    ///
    /// # Arguments
    ///
    /// * `kernel_load_offset` - Offset from `guest_mem` at which the kernel starts.
    /// nr cpus is required for checking populating the kvm_cpuid2 entry for ebx and edx registers
    pub fn configure(
        &mut self,
        machine_config: &MachineConfiguration,
        kernel_start_addr: GuestAddress,
        vm: &Vm,
    ) -> Result<()> {
        // the MachineConfiguration has defaults for ht_enabled and vcpu_count hence it is safe to unwrap
        cpuid::filter_cpuid(
            self.id,
            machine_config.vcpu_count.unwrap(),
            machine_config.ht_enabled.unwrap(),
            &mut self.cpuid,
        ).map_err(|e| Error::CpuId(e))?;
        match machine_config.cpu_template {
            Some(template) => cpuid::set_cpuid_template(template, &mut self.cpuid),
            None => (),
        }

        self.fd
            .set_cpuid2(&self.cpuid)
            .map_err(Error::SetSupportedCpusFailed)?;

        regs::setup_msrs(&self.fd).map_err(Error::MSRSConfiguration)?;
        // Safe to unwrap because this method is called after the VM is configured
        let vm_memory = vm.get_memory().unwrap();
        regs::setup_regs(
            &self.fd,
            kernel_start_addr.offset() as u64,
            x86_64::layout::BOOT_STACK_POINTER as u64,
            x86_64::layout::ZERO_PAGE_START as u64,
        ).map_err(Error::REGSConfiguration)?;
        regs::setup_fpu(&self.fd).map_err(Error::FPUConfiguration)?;
        regs::setup_sregs(vm_memory, &self.fd).map_err(Error::SREGSConfiguration)?;
        interrupts::set_lint(&self.fd).map_err(Error::LocalIntConfiguration)?;
        Ok(())
    }

    /// Runs the VCPU until it exits, returning the reason.
    ///
    /// Note that the state of the VCPU and associated VM must be setup first for this to do
    /// anything useful.
    pub fn run(&self) -> Result<VcpuExit> {
        match self.fd.run() {
            Ok(v) => Ok(v),
            Err(e) => return Err(Error::VcpuRun(<sys_util::Error>::new(e.errno()))),
        }
    }
}

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

    #[test]
    fn create_vm() {
        let kvm = Kvm::new().unwrap();
        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x10000)]).unwrap();
        let mut vm = Vm::new(&kvm).expect("new vm failed");
        assert!(vm.memory_init(gm).is_ok());
    }

    #[test]
    fn get_memory() {
        let kvm = Kvm::new().unwrap();
        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x1000)]).unwrap();
        let mut vm = Vm::new(&kvm).expect("new vm failed");
        assert!(vm.memory_init(gm).is_ok());
        let obj_addr = GuestAddress(0xf0);
        vm.get_memory()
            .unwrap()
            .write_obj_at_addr(67u8, obj_addr)
            .unwrap();
        let read_val: u8 = vm
            .get_memory()
            .unwrap()
            .read_obj_from_addr(obj_addr)
            .unwrap();
        assert_eq!(read_val, 67u8);
    }

    #[test]
    fn create_vcpu() {
        let kvm = Kvm::new().unwrap();
        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x10000)]).unwrap();
        let mut vm = Vm::new(&kvm).expect("new vm failed");
        assert!(vm.memory_init(gm).is_ok());
        Vcpu::new(0, &mut vm).unwrap();
    }

    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    #[test]
    fn test_cpuid() {
        let kvm = Kvm::new().unwrap();
        let mut vm = Vm::new(&kvm).unwrap();
        let mut vcpu = Vcpu::new(0, &mut vm).unwrap();
        assert_eq!(vcpu.get_cpuid(), vm.fd.get_supported_cpuid());
        assert!(cpuid::filter_cpuid(0, 1, true, &mut vcpu.cpuid).is_ok());
        // Test using the T2 template
        cpuid::set_cpuid_template(CpuFeaturesTemplate::T2, &mut vcpu.cpuid);
        assert!(vcpu.fd.set_cpuid2(&vcpu.cpuid).is_ok());
        // Test using the C3 template
        cpuid::set_cpuid_template(CpuFeaturesTemplate::C3, &mut vcpu.cpuid);
        assert!(vcpu.fd.set_cpuid2(&vcpu.cpuid).is_ok());
    }

    #[test]
    fn run_code() {
        use std::io::{self, Write};
        // This example based on https://lwn.net/Articles/658511/
        let code = [
            0xba, 0xf8, 0x03, /* mov $0x3f8, %dx */
            0x00, 0xd8, /* add %bl, %al */
            0x04, '0' as u8, /* add $'0', %al */
            0xee,      /* out %al, (%dx) */
            0xb0, '\n' as u8, /* mov $'\n', %al */
            0xee,       /* out %al, (%dx) */
            0xf4,       /* hlt */
        ];

        let mem_size = 0x1000;
        let load_addr = GuestAddress(0x1000);
        let mem = GuestMemory::new(&vec![(load_addr, mem_size)]).unwrap();

        let kvm = Kvm::new().expect("new kvm failed");
        let mut vm = Vm::new(&kvm).expect("new vm failed");
        assert!(vm.memory_init(mem).is_ok());
        vm.get_memory()
            .unwrap()
            .write_slice_at_addr(&code, load_addr)
            .expect("Writing code to memory failed.");

        let vcpu = Vcpu::new(0, &mut vm).expect("new vcpu failed");

        let mut vcpu_sregs = vcpu.fd.get_sregs().expect("get sregs failed");
        assert_ne!(vcpu_sregs.cs.base, 0);
        assert_ne!(vcpu_sregs.cs.selector, 0);
        vcpu_sregs.cs.base = 0;
        vcpu_sregs.cs.selector = 0;
        vcpu.fd.set_sregs(&vcpu_sregs).expect("set sregs failed");

        let mut vcpu_regs = vcpu.fd.get_regs().expect("get regs failed");
        vcpu_regs.rip = 0x1000;
        vcpu_regs.rax = 2;
        vcpu_regs.rbx = 3;
        vcpu_regs.rflags = 2;
        vcpu.fd.set_regs(&vcpu_regs).expect("set regs failed");

        loop {
            match vcpu.run().expect("run failed") {
                VcpuExit::IoOut(0x3f8, data) => {
                    assert_eq!(data.len(), 1);
                    io::stdout().write(data).unwrap();
                }
                VcpuExit::Hlt => {
                    io::stdout().write(b"KVM_EXIT_HLT\n").unwrap();
                    break;
                }
                r => panic!("unexpected exit reason: {:?}", r),
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::result;
use std::fmt;
use std::str::FromStr;
use std::net::{SocketAddr, Ipv4Addr, Ipv6Addr, IpAddr};
use pnet::util::MacAddr as PnetMacAddr;
use pnet::packet::PrimitiveValues;
use netinfo::error::*;

/// Udp, Tcp or other packet type on transport layer?
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub enum TransportType {
    /// Using Tcp on transport layer for packet
    Tcp,

    /// Using Udp on transport layer for packet
    Udp,

    // others might get added
}

/// A MAC address (6-bytes)
#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
pub struct MacAddr(pub u8, pub u8, pub u8, pub u8, pub u8, pub u8);
impl From<(u8, u8, u8, u8, u8, u8)> for MacAddr {
    fn from(d: (u8, u8, u8, u8, u8, u8)) -> MacAddr { MacAddr(d.0, d.1, d.2, d.3, d.4, d.5) }
}
impl From<[u8; 6]> for MacAddr {
    fn from(d: [u8; 6]) -> MacAddr { MacAddr(d[0], d[1], d[2], d[3], d[4], d[5]) }
}
impl From<PnetMacAddr> for MacAddr {
    fn from(d: PnetMacAddr) -> MacAddr { d.to_primitive_values().into() }
}
impl From<MacAddr> for PnetMacAddr {
    fn from(d: MacAddr) -> PnetMacAddr { let MacAddr(a, b, c, d, e, f) = d; PnetMacAddr::new(a, b, c, d, e, f) }
}
impl From<MacAddr> for (u8, u8, u8, u8, u8, u8) {
    fn from(d: MacAddr) -> (u8, u8, u8, u8, u8, u8) { let MacAddr(a, b, c, d, e, f) = d; (a, b, c, d, e, f) }
}
impl From<MacAddr> for [u8; 6] {
  fn from(f: MacAddr) -> [u8; 6] { [f.0, f.1, f.2, f.3, f.4, f.5] }
}
impl fmt::Display for MacAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { PnetMacAddr::from(*self).fmt(f) }
}
impl FromStr for MacAddr {
    type Err = Error;
    fn from_str(s: &str) -> result::Result<MacAddr, Error> {
        match PnetMacAddr::from_str(s) {
            Ok(pnet_mac) => Ok(MacAddr::from(pnet_mac)),
            Err(_) => Err(ErrorKind::MacAddrParseError.into()), // TODO: chain_err() would be much nicer
        }
    }
}

/// Describes a network packet.
#[allow(missing_copy_implementations)]
#[derive(Debug, Clone)]
pub struct PacketInfo {

    /// Source IP
    pub sip: SocketAddr,

    /// Destination IP
    pub dip: SocketAddr,

    /// Number of nanoseconds (currently unused and set to 0)
    pub time: u64,

    /// Number of bytes in package (on datalink level)
    pub datalen: u64,

    /// Transport layer type
    pub transport_type: TransportType,

    /// Is Some(Incoming) or Some(Outgoing) if inout type can be determined, None if not
    pub inout_type: Option<InoutType>
}

impl PacketInfo {
    /// Constructor for `PacketInfo` type.
    pub fn new(sip: SocketAddr, dip: SocketAddr, datalen: u64, transport_type: TransportType, inout_type: Option<InoutType>) -> PacketInfo {
        PacketInfo { sip: sip, dip: dip, time: 0, datalen: datalen, transport_type: transport_type, inout_type: inout_type }
    }
}

/// Set IP to zero but leave type and port untouched.
pub fn reset_socket_addr_ip(s: SocketAddr) -> SocketAddr {
    let port = s.port();
    let mut new = reset_socket_addr(s);
    new.set_port(port);
    new
}

/// Set IP and port to zero but leave type untouched.
pub fn reset_socket_addr(s: SocketAddr) -> SocketAddr {
    match s {
        SocketAddr::V4(_) => { SocketAddr::new(IpAddr::V4(Ipv4Addr::from([0u8; 4])), 0) }
        SocketAddr::V6(_) => { SocketAddr::new(IpAddr::V6(Ipv6Addr::from([0u8; 16])), 0) }
    }
}

#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
#[allow(missing_copy_implementations)]
/// Represents a network connection from one "ip-adress:port" to another. It can be
/// Ipv4 or Ipv6.
pub struct Connection {
    /// Local adress.
    pub local: SocketAddr,

    /// Remote adress.
    pub remote: SocketAddr,
}

impl Connection {
    /// Generate new `Connection` from Ipv4 or Ipv6 addressses.
    pub fn new(local: SocketAddr, remote: SocketAddr) -> Connection {
        Connection { local: local, remote: remote }
    }

    /*
    /// Generate new `Connection` from Ipv4 addressses.
    pub fn new_ip4(local: SocketAddrV4, remote: SocketAddrV4) -> Connection {
        Self::new(SocketAddr::V4(local), SocketAddr::V4(remote))
    }

    /// Generate new `Connection` from Ipv6 addressses.
    pub fn new_ip6(local: SocketAddrV6, remote: SocketAddrV6) -> Connection {
        Self::new(SocketAddr::V6(local), SocketAddr::V6(remote))
    }
    */

    /// Get the reverse connection.
    pub fn get_reverse(&self) -> Connection {
        Connection::new(self.remote, self.local)
    }

    /// Get connection with port but ip set to zero.
    pub fn get_resetted_ip(&self) -> Connection {
        Connection::new(reset_socket_addr_ip(self.local), reset_socket_addr_ip(self.remote))
    }

    /// Get connection with port and ip set to zero for remote (but same type).
    pub fn get_resetted_remote(&self) -> Connection {
        Connection::new(self.local, reset_socket_addr(self.remote))
    }
}

impl<'a> From<&'a PacketInfo> for Connection {
    fn from(p: &'a PacketInfo) -> Connection {
        Connection::new(p.sip, p.dip)
    }
}

impl From<PacketInfo> for Connection {
    fn from(p: PacketInfo) -> Connection {
        Connection::new(p.sip, p.dip)
    }
}

#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
/// InoutType can be Incoming or Outgoing (direction of the traffic).
pub enum InoutType {
    /// Local address -> Remote adress
    Incoming,

    /// Remote adress -> Local address
    Outgoing,
}
<commit_msg>Implement NetworkInteface newtype<commit_after>use std::result;
use std::fmt;
use std::str::FromStr;
use std::net::{SocketAddr, Ipv4Addr, Ipv6Addr, IpAddr};
use pnet::util::MacAddr as PnetMacAddr;
use pnet::datalink::NetworkInterface as PnetNetworkInterface;
use pnet::packet::PrimitiveValues;
use netinfo::error::*;

/// Udp, Tcp or other packet type on transport layer?
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub enum TransportType {
    /// Using Tcp on transport layer for packet
    Tcp,

    /// Using Udp on transport layer for packet
    Udp,

    // others might get added
}

/// A MAC address (6-bytes)
#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
pub struct MacAddr(pub u8, pub u8, pub u8, pub u8, pub u8, pub u8);
impl From<(u8, u8, u8, u8, u8, u8)> for MacAddr {
    fn from(d: (u8, u8, u8, u8, u8, u8)) -> MacAddr { MacAddr(d.0, d.1, d.2, d.3, d.4, d.5) }
}
impl From<[u8; 6]> for MacAddr {
    fn from(d: [u8; 6]) -> MacAddr { MacAddr(d[0], d[1], d[2], d[3], d[4], d[5]) }
}
impl From<PnetMacAddr> for MacAddr {
    fn from(d: PnetMacAddr) -> MacAddr { d.to_primitive_values().into() }
}
impl From<MacAddr> for PnetMacAddr {
    fn from(d: MacAddr) -> PnetMacAddr { let MacAddr(a, b, c, d, e, f) = d; PnetMacAddr::new(a, b, c, d, e, f) }
}
impl From<MacAddr> for (u8, u8, u8, u8, u8, u8) {
    fn from(d: MacAddr) -> (u8, u8, u8, u8, u8, u8) { let MacAddr(a, b, c, d, e, f) = d; (a, b, c, d, e, f) }
}
impl From<MacAddr> for [u8; 6] {
  fn from(f: MacAddr) -> [u8; 6] { [f.0, f.1, f.2, f.3, f.4, f.5] }
}
impl fmt::Display for MacAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { PnetMacAddr::from(*self).fmt(f) }
}
impl FromStr for MacAddr {
    type Err = Error;
    fn from_str(s: &str) -> result::Result<MacAddr, Error> {
        match PnetMacAddr::from_str(s) {
            Ok(pnet_mac) => Ok(MacAddr::from(pnet_mac)),
            Err(_) => Err(ErrorKind::MacAddrParseError.into()), // TODO: chain_err() would be much nicer
        }
    }
}

/// This is used to identify the network interface you want to count the traffic on.
pub struct NetworkInterface {
    inner: PnetNetworkInterface,
}
impl From<PnetNetworkInterface> for NetworkInterface {
  fn from(f: PnetNetworkInterface) -> NetworkInterface {
      NetworkInterface { inner: f }
  }
}
impl From<NetworkInterface> for PnetNetworkInterface {
  fn from(f: NetworkInterface) -> PnetNetworkInterface {
      f.inner
  }
}
impl NetworkInterface {
    /// Returns name of network interface (like "enp2s0" or "wlp3s0").
    pub fn get_name(&self) -> String { self.inner.name.clone() }

    /// Returns name of network interface (like "enp2s0" or "wlp3s0").
    pub fn get_name_as_str(&self) -> &str { &self.inner.name }

    /// The interface index (operating system specific)
    pub fn get_index(&self) -> u32 { self.inner.index }

    /// A MAC address for the interface
    pub fn get_mac(&self) -> Option<MacAddr> { self.inner.mac.map(|m| m.into()) }

    /// An IP addresses for the interface
    pub fn get_ips(&self) -> Option<Vec<IpAddr>> { self.inner.ips.clone() }

    /// An IP addresses for the interface
    pub fn get_ips_as_slice(&self) -> Option<&[IpAddr]> { self.inner.ips.as_ref().map(|ips| &ips[..]) }

    /// Operating system specific flags for the interface
    pub fn get_flags(&self) -> u32 { self.inner.flags }
}

/// Describes a network packet.
#[allow(missing_copy_implementations)]
#[derive(Debug, Clone)]
pub struct PacketInfo {

    /// Source IP
    pub sip: SocketAddr,

    /// Destination IP
    pub dip: SocketAddr,

    /// Number of nanoseconds (currently unused and set to 0)
    pub time: u64,

    /// Number of bytes in package (on datalink level)
    pub datalen: u64,

    /// Transport layer type
    pub transport_type: TransportType,

    /// Is Some(Incoming) or Some(Outgoing) if inout type can be determined, None if not
    pub inout_type: Option<InoutType>
}

impl PacketInfo {
    /// Constructor for `PacketInfo` type.
    pub fn new(sip: SocketAddr, dip: SocketAddr, datalen: u64, transport_type: TransportType, inout_type: Option<InoutType>) -> PacketInfo {
        PacketInfo { sip: sip, dip: dip, time: 0, datalen: datalen, transport_type: transport_type, inout_type: inout_type }
    }
}

/// Set IP to zero but leave type and port untouched.
pub fn reset_socket_addr_ip(s: SocketAddr) -> SocketAddr {
    let port = s.port();
    let mut new = reset_socket_addr(s);
    new.set_port(port);
    new
}

/// Set IP and port to zero but leave type untouched.
pub fn reset_socket_addr(s: SocketAddr) -> SocketAddr {
    match s {
        SocketAddr::V4(_) => { SocketAddr::new(IpAddr::V4(Ipv4Addr::from([0u8; 4])), 0) }
        SocketAddr::V6(_) => { SocketAddr::new(IpAddr::V6(Ipv6Addr::from([0u8; 16])), 0) }
    }
}

#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
#[allow(missing_copy_implementations)]
/// Represents a network connection from one "ip-adress:port" to another. It can be
/// Ipv4 or Ipv6.
pub struct Connection {
    /// Local adress.
    pub local: SocketAddr,

    /// Remote adress.
    pub remote: SocketAddr,
}

impl Connection {
    /// Generate new `Connection` from Ipv4 or Ipv6 addressses.
    pub fn new(local: SocketAddr, remote: SocketAddr) -> Connection {
        Connection { local: local, remote: remote }
    }

    /*
    /// Generate new `Connection` from Ipv4 addressses.
    pub fn new_ip4(local: SocketAddrV4, remote: SocketAddrV4) -> Connection {
        Self::new(SocketAddr::V4(local), SocketAddr::V4(remote))
    }

    /// Generate new `Connection` from Ipv6 addressses.
    pub fn new_ip6(local: SocketAddrV6, remote: SocketAddrV6) -> Connection {
        Self::new(SocketAddr::V6(local), SocketAddr::V6(remote))
    }
    */

    /// Get the reverse connection.
    pub fn get_reverse(&self) -> Connection {
        Connection::new(self.remote, self.local)
    }

    /// Get connection with port but ip set to zero.
    pub fn get_resetted_ip(&self) -> Connection {
        Connection::new(reset_socket_addr_ip(self.local), reset_socket_addr_ip(self.remote))
    }

    /// Get connection with port and ip set to zero for remote (but same type).
    pub fn get_resetted_remote(&self) -> Connection {
        Connection::new(self.local, reset_socket_addr(self.remote))
    }
}

impl<'a> From<&'a PacketInfo> for Connection {
    fn from(p: &'a PacketInfo) -> Connection {
        Connection::new(p.sip, p.dip)
    }
}

impl From<PacketInfo> for Connection {
    fn from(p: PacketInfo) -> Connection {
        Connection::new(p.sip, p.dip)
    }
}

#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
/// InoutType can be Incoming or Outgoing (direction of the traffic).
pub enum InoutType {
    /// Local address -> Remote adress
    Incoming,

    /// Remote adress -> Local address
    Outgoing,
}
<|endoftext|>


<commit_before>// The MIT License (MIT)

// Copyright (c) 2014 Y. T. CHUNG <zonyitoo@gmail.com>

// 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.

//! Memcached version

use std::fmt::{Display, Formatter, self};
use std::str::FromStr;

/// Memcached version
///
/// Version(major, minor, patch)
#[derive(Copy, Debug)]
pub struct Version(u32, u32, u32);

impl Version {
    pub fn new(major: u32, minor: u32, patch: u32) -> Version {
        Version(major, minor, patch)
    }
}

impl Display for Version {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let &Version(major, minor, patch) = self;
        write!(f, "{}:{}:{}", major, minor, patch)
    }
}

macro_rules! try_option(
    ($inp:expr) => (
        match $inp {
            Some(v) => { v },
            None => { return None; },
        }
    );
);

impl FromStr for Version {
    fn from_str(s: &str) -> Option<Version> {
        let sp: Vec<&str> = s.split('.').collect();
        Some(Version::new(try_option!(sp[0].parse::<u32>()),
                          try_option!(sp[1].parse::<u32>()),
                          try_option!(sp[2].parse::<u32>())))
    }
}
<commit_msg>updated version parsing<commit_after>// The MIT License (MIT)

// Copyright (c) 2014 Y. T. CHUNG <zonyitoo@gmail.com>

// 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.

//! Memcached version

use std::fmt::{Display, Formatter, self};
use std::str::FromStr;

/// Memcached version
///
/// Version(major, minor, patch)
#[derive(Copy, Debug)]
pub struct Version(u32, u32, u32);

impl Version {
    pub fn new(major: u32, minor: u32, patch: u32) -> Version {
        Version(major, minor, patch)
    }
}

impl Display for Version {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let &Version(major, minor, patch) = self;
        write!(f, "{}:{}:{}", major, minor, patch)
    }
}

macro_rules! try_option(
    ($inp:expr) => (
        match $inp {
            Some(v) => { v },
            None => { return None; },
        }
    );
);

impl FromStr for Version {
    fn from_str(s: &str) -> Option<Version> {
        let mut sp = s.split('.');
        let major = match sp.next() {
            Some(s) => try_option!(s.parse()),
            None => return None,
        };
        let minor = match sp.next() {
            Some(s) => try_option!(s.parse()),
            None => 0,
        };
        let patch = match sp.next() {
            Some(s) => try_option!(s.parse()),
            None => 0,
        };

        Some(Version::new(major, minor, patch))
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2015-2017 Florian Wilkens
//
// 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.

//! Wrapped FFI bindings to Linux-PAM
//!
//! This module provides wrapped versions of some of the functions from
//! the [`raw`](../raw/index.html) module which use the appropriate enums
//! instead of `c_int`. These wrappers should always be preferred as one
//! can safely match on their return types and prevent illegal arguments
//! from beeing passed to the native library.
//!
//! Note: These wrappers get added as I need them. Feel free to open an issue
//! or PR for the ones that you require which haven't been added yet.

use libc::{c_char, c_int, c_void};
use std::ffi::{CStr, CString};
use std::ptr::null;

use types::*;
use raw::*;

/* ------------------------ pam_appl.h -------------------------- */
#[inline]
pub fn start(service: &str,
             user: Option<&str>,
             conversation: &PamConversation,
             handle: *mut *mut PamHandle)
             -> PamReturnCode {
    if let Ok(service) = CString::new(service) {
        if let Some(usr) = user {
            if let Ok(user) = CString::new(usr) {
                unsafe {
                    From::from(pam_start(service.as_ptr(),
                                         user.as_ptr(),
                                         conversation,
                                         handle as *mut *const PamHandle))
                }
            } else {
                PamReturnCode::BUF_ERR
            }
        } else {
            unsafe {
                From::from(pam_start(service.as_ptr(),
                                     null(),
                                     conversation,
                                     handle as *mut *const PamHandle))
            }
        }
    } else {
        PamReturnCode::SERVICE_ERR
    }
}

#[inline]
pub fn end(handle: &mut PamHandle, status: PamReturnCode) -> PamReturnCode {
    From::from(unsafe { pam_end(handle, status as c_int) })
}

#[inline]
pub fn authenticate(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_authenticate(handle, flags as c_int) })
}

#[inline]
pub fn setcred(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_setcred(handle, flags as c_int) })
}

#[inline]
pub fn acct_mgmt(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_acct_mgmt(handle, flags as c_int) })
}

#[inline]
pub fn open_session(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_open_session(handle, flags as c_int) })
}

#[inline]
pub fn close_session(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_close_session(handle, flags as c_int) })
}

#[inline]
pub fn chauthtok(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_chauthtok(handle, flags as c_int) })
}
/* ------------------------ pam_appl.h -------------------------- */

/* ----------------------- _pam_types.h ------------------------- */
#[inline]
pub fn set_item(handle: &mut PamHandle, item_type: PamItemType, item: &c_void) -> PamReturnCode {
    From::from(unsafe { pam_set_item(handle, item_type as c_int, item) })
}

#[inline]
pub fn get_item(handle: &PamHandle,
                item_type: PamItemType,
                item: &mut *const c_void)
                -> PamReturnCode {
    From::from(unsafe { pam_get_item(handle, item_type as c_int, item) })
}

#[inline]
pub fn strerror(handle: &mut PamHandle, errnum: PamReturnCode) -> Option<&str> {
    unsafe { CStr::from_ptr(pam_strerror(handle, errnum as c_int)) }.to_str().ok()
}

#[inline]
pub fn putenv(handle: &mut PamHandle, name_value: &str) -> PamReturnCode {
    if let Ok(name_value) = CString::new(name_value) {
        From::from(unsafe { pam_putenv(handle, name_value.as_ptr()) })
    } else {
        // Not sure whether this is the correct return value
        PamReturnCode::BUF_ERR
    }
}

#[inline]
pub fn getenv<'a>(handle: &'a mut PamHandle, name: &str) -> Option<&'a str> {
    use std::ptr;
	use std::io;
    if let Ok(name) = CString::new(name) {
		let env = unsafe{pam_getenv(handle, name.as_ptr())};
		if env != ptr::null(){
			unsafe { CStr::from_ptr(env) }.to_str().ok()
		}
		else{
			None
		}
	} else {
        None
    }
}

#[inline]
pub fn getenvlist(handle: &mut PamHandle) -> *const *const c_char {
    //TODO: find a convenient way to handle this with Rust types
    unsafe { pam_getenvlist(handle) }
}
/* ----------------------- _pam_types.h ------------------------- */

/* ----------------------- pam_misc.h --------------------------- */
#[inline]
#[cfg(target_os = "linux")]
pub fn misc_paste_env(handle: &mut PamHandle, user_env: &[&str]) -> PamReturnCode {
    // Taken from: https://github.com/rust-lang/rust/issues/9564#issuecomment-95354558
    let user_env: Vec<_> = user_env.iter()
        .map(|&env| CString::new(env).unwrap())
        .collect();
    let env_ptrs: Vec<_> = user_env.iter()
        .map(|env| env.as_ptr())
        .chain(Some(null()))
        .collect();
    From::from(unsafe { pam_misc_paste_env(handle, env_ptrs.as_ptr()) })
}

#[inline]
#[cfg(target_os = "linux")]
pub fn misc_drop_env(env: &mut *mut c_char) -> PamReturnCode {
    From::from(unsafe { pam_misc_drop_env(env) })
}

#[inline]
#[cfg(target_os = "linux")]
pub fn misc_setenv(handle: &mut PamHandle,
                   name: &str,
                   value: &str,
                   readonly: bool)
                   -> PamReturnCode {
    if let (Ok(name), Ok(value)) = (CString::new(name), CString::new(value)) {
        From::from(unsafe {
            pam_misc_setenv(handle,
                            name.as_ptr(),
                            value.as_ptr(),
                            if readonly { 0 } else { 1 })
        })
    } else {
        PamReturnCode::BUF_ERR
    }
}
/* ----------------------- pam_misc.h --------------------------- */

/* ----------------------- pam_modules.h ------------------------ */
#[inline]
pub fn set_data(handle: &mut PamHandle,
                module_data_name: &str,
                data: &mut c_void,
                cleanup: Option<extern "C" fn(*mut PamHandle, *mut c_void, c_int)>)
                -> PamReturnCode {
    if let Ok(module_data_name) = CString::new(module_data_name) {
        From::from(unsafe { pam_set_data(handle, module_data_name.as_ptr(), data, cleanup) })

    } else {
        PamReturnCode::BUF_ERR
    }
}

//pub fn get_data(handle: *const PamHandle, module_data_name: *const c_char, data: *const *const c_void);

pub fn get_user(handle: &PamHandle,
                user: &mut *const c_char,
                prompt: *const c_char)
                -> PamReturnCode {

    From::from(
        unsafe {
            pam_get_user(
                handle,
                user,
                prompt
            )
        }
    )
}

/* ----------------------- pam_modules.h ------------------------ */
<commit_msg>Fixed NULL pointer usage when pam_getenv fail<commit_after>// Copyright (C) 2015-2017 Florian Wilkens
//
// 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.

//! Wrapped FFI bindings to Linux-PAM
//!
//! This module provides wrapped versions of some of the functions from
//! the [`raw`](../raw/index.html) module which use the appropriate enums
//! instead of `c_int`. These wrappers should always be preferred as one
//! can safely match on their return types and prevent illegal arguments
//! from beeing passed to the native library.
//!
//! Note: These wrappers get added as I need them. Feel free to open an issue
//! or PR for the ones that you require which haven't been added yet.

use libc::{c_char, c_int, c_void};
use std::ffi::{CStr, CString};
use std::ptr::null;

use types::*;
use raw::*;

/* ------------------------ pam_appl.h -------------------------- */
#[inline]
pub fn start(service: &str,
             user: Option<&str>,
             conversation: &PamConversation,
             handle: *mut *mut PamHandle)
             -> PamReturnCode {
    if let Ok(service) = CString::new(service) {
        if let Some(usr) = user {
            if let Ok(user) = CString::new(usr) {
                unsafe {
                    From::from(pam_start(service.as_ptr(),
                                         user.as_ptr(),
                                         conversation,
                                         handle as *mut *const PamHandle))
                }
            } else {
                PamReturnCode::BUF_ERR
            }
        } else {
            unsafe {
                From::from(pam_start(service.as_ptr(),
                                     null(),
                                     conversation,
                                     handle as *mut *const PamHandle))
            }
        }
    } else {
        PamReturnCode::SERVICE_ERR
    }
}

#[inline]
pub fn end(handle: &mut PamHandle, status: PamReturnCode) -> PamReturnCode {
    From::from(unsafe { pam_end(handle, status as c_int) })
}

#[inline]
pub fn authenticate(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_authenticate(handle, flags as c_int) })
}

#[inline]
pub fn setcred(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_setcred(handle, flags as c_int) })
}

#[inline]
pub fn acct_mgmt(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_acct_mgmt(handle, flags as c_int) })
}

#[inline]
pub fn open_session(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_open_session(handle, flags as c_int) })
}

#[inline]
pub fn close_session(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_close_session(handle, flags as c_int) })
}

#[inline]
pub fn chauthtok(handle: &mut PamHandle, flags: PamFlag) -> PamReturnCode {
    From::from(unsafe { pam_chauthtok(handle, flags as c_int) })
}
/* ------------------------ pam_appl.h -------------------------- */

/* ----------------------- _pam_types.h ------------------------- */
#[inline]
pub fn set_item(handle: &mut PamHandle, item_type: PamItemType, item: &c_void) -> PamReturnCode {
    From::from(unsafe { pam_set_item(handle, item_type as c_int, item) })
}

#[inline]
pub fn get_item(handle: &PamHandle,
                item_type: PamItemType,
                item: &mut *const c_void)
                -> PamReturnCode {
    From::from(unsafe { pam_get_item(handle, item_type as c_int, item) })
}

#[inline]
pub fn strerror(handle: &mut PamHandle, errnum: PamReturnCode) -> Option<&str> {
    unsafe { CStr::from_ptr(pam_strerror(handle, errnum as c_int)) }.to_str().ok()
}

#[inline]
pub fn putenv(handle: &mut PamHandle, name_value: &str) -> PamReturnCode {
    if let Ok(name_value) = CString::new(name_value) {
        From::from(unsafe { pam_putenv(handle, name_value.as_ptr()) })
    } else {
        // Not sure whether this is the correct return value
        PamReturnCode::BUF_ERR
    }
}

#[inline]
pub fn getenv<'a>(handle: &'a mut PamHandle, name: &str) -> Option<&'a str> {
    use std::ptr;
    if let Ok(name) = CString::new(name) {
		let env = unsafe{pam_getenv(handle, name.as_ptr())};
		if env != ptr::null(){
			unsafe { CStr::from_ptr(env) }.to_str().ok()
		}
		else{
			None
		}
	} else {
        None
    }
}

#[inline]
pub fn getenvlist(handle: &mut PamHandle) -> *const *const c_char {
    //TODO: find a convenient way to handle this with Rust types
    unsafe { pam_getenvlist(handle) }
}
/* ----------------------- _pam_types.h ------------------------- */

/* ----------------------- pam_misc.h --------------------------- */
#[inline]
#[cfg(target_os = "linux")]
pub fn misc_paste_env(handle: &mut PamHandle, user_env: &[&str]) -> PamReturnCode {
    // Taken from: https://github.com/rust-lang/rust/issues/9564#issuecomment-95354558
    let user_env: Vec<_> = user_env.iter()
        .map(|&env| CString::new(env).unwrap())
        .collect();
    let env_ptrs: Vec<_> = user_env.iter()
        .map(|env| env.as_ptr())
        .chain(Some(null()))
        .collect();
    From::from(unsafe { pam_misc_paste_env(handle, env_ptrs.as_ptr()) })
}

#[inline]
#[cfg(target_os = "linux")]
pub fn misc_drop_env(env: &mut *mut c_char) -> PamReturnCode {
    From::from(unsafe { pam_misc_drop_env(env) })
}

#[inline]
#[cfg(target_os = "linux")]
pub fn misc_setenv(handle: &mut PamHandle,
                   name: &str,
                   value: &str,
                   readonly: bool)
                   -> PamReturnCode {
    if let (Ok(name), Ok(value)) = (CString::new(name), CString::new(value)) {
        From::from(unsafe {
            pam_misc_setenv(handle,
                            name.as_ptr(),
                            value.as_ptr(),
                            if readonly { 0 } else { 1 })
        })
    } else {
        PamReturnCode::BUF_ERR
    }
}
/* ----------------------- pam_misc.h --------------------------- */

/* ----------------------- pam_modules.h ------------------------ */
#[inline]
pub fn set_data(handle: &mut PamHandle,
                module_data_name: &str,
                data: &mut c_void,
                cleanup: Option<extern "C" fn(*mut PamHandle, *mut c_void, c_int)>)
                -> PamReturnCode {
    if let Ok(module_data_name) = CString::new(module_data_name) {
        From::from(unsafe { pam_set_data(handle, module_data_name.as_ptr(), data, cleanup) })

    } else {
        PamReturnCode::BUF_ERR
    }
}

//pub fn get_data(handle: *const PamHandle, module_data_name: *const c_char, data: *const *const c_void);

pub fn get_user(handle: &PamHandle,
                user: &mut *const c_char,
                prompt: *const c_char)
                -> PamReturnCode {

    From::from(
        unsafe {
            pam_get_user(
                handle,
                user,
                prompt
            )
        }
    )
}

/* ----------------------- pam_modules.h ------------------------ */
<|endoftext|>