<commit_before>use std::collections::BTreeSet;
use std::fmt::{Display, Error, Formatter};
use action::{POSSIBLE_DIRECTIONS,POSSIBLE_MOVE_TYPES,Coordinate,Move};

pub const NUM_COLS: usize = 5;
pub const NUM_ROWS: usize = 4;

// Hacks for the three states of the lower-left of the board in JunSeok vs YeonSeung game
#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum BoardType {
    Left,
    LeftOrMiddle,
    Middle,
    RightOrMiddle,
    Right,
}

impl BoardType {
    fn is_final(&self) -> bool {
        match *self {
            BoardType::Left => true,
            BoardType::Middle => true,
            BoardType::Right => true,
            BoardType::LeftOrMiddle => false,
            BoardType::RightOrMiddle => false,
        }
    }

    // Can a board of type `current` become this type?
    fn applies_to(&self, current: Option<BoardType>) -> bool {
        match (current, *self) {
            // a none board type can change to anything.
            (None, _) => true,
            // LeftOrMiddle can change to itself, left, or middle.
            (Some(BoardType::LeftOrMiddle), BoardType::LeftOrMiddle) => true,
            (Some(BoardType::LeftOrMiddle), BoardType::Left) => true,
            (Some(BoardType::LeftOrMiddle), BoardType::Middle) => true,
            (Some(BoardType::LeftOrMiddle), _) => false,
            // RightOrMiddle can change to itself, right, or middle.
            (Some(BoardType::RightOrMiddle), BoardType::RightOrMiddle) => true,
            (Some(BoardType::RightOrMiddle), BoardType::Right) => true,
            (Some(BoardType::RightOrMiddle), BoardType::Middle) => true,
            (Some(BoardType::RightOrMiddle), _) => false,
            // Left, Middle, Right can change to themselves only.
            (Some(BoardType::Left), BoardType::Left) => true,
            (Some(BoardType::Left), _) => false,
            (Some(BoardType::Middle), BoardType::Middle) => true,
            (Some(BoardType::Middle), _) => false,
            (Some(BoardType::Right), BoardType::Right) => true,
            (Some(BoardType::Right), _) => false,
        }
    }

    fn induced_by(&self, c: Coordinate) -> bool {
        // Not in the lower left, so it's a free pass.
        if !c.induces_board_type() {
            return true;
        }

        match *self {
            BoardType::Left          => c != Coordinate{row: 2, col: 1} && c != Coordinate{row: 1, col: 1},
            BoardType::LeftOrMiddle  => c == Coordinate{row: 1, col: 0},
            BoardType::Middle        => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 1, col: 1},
            BoardType::RightOrMiddle => c == Coordinate{row: 3, col: 1},
            BoardType::Right         => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 2, col: 0},
        }
    }
}

#[derive(Copy, Clone)]
enum Orientation {
    UpDown,
    LeftRight,
    UpLeft,
    UpRight,
    DownLeft,
    DownRight,
}

impl Orientation {
    fn to_str(&self) -> &'static str {
        match *self {
            Orientation::UpDown => "║",
            Orientation::LeftRight => "═",
            Orientation::UpLeft => "╝",
            Orientation::UpRight => "╚",
            Orientation::DownLeft => "╗",
            Orientation::DownRight => "╔",
        }
    }
}

#[derive(Copy, Clone)]
enum OrientationOption {
    Fixed(Orientation),
    IfRight(Orientation, Orientation),
    IfLeft(Orientation, Orientation),
    LeftOrMiddle(Orientation, Orientation),
    RightOrMiddle(Orientation, Orientation),
}

impl OrientationOption {
    fn for_board(&self, board_type: Option<BoardType>) -> &'static str {
        let orientation = match (*self, board_type) {
            (OrientationOption::Fixed(orientation), _) => orientation,
            (OrientationOption::IfRight(ifright, _), Some(BoardType::Right)) => ifright,
            (OrientationOption::IfRight(_, notright), _) => notright,
            (OrientationOption::IfLeft(ifleft, _), Some(BoardType::Left)) => ifleft,
            (OrientationOption::IfLeft(_, notleft), _) => notleft,
            (OrientationOption::LeftOrMiddle(ifleft, _), Some(BoardType::Left)) => ifleft,
            (OrientationOption::LeftOrMiddle(_, ifmiddle), Some(BoardType::Middle)) => ifmiddle,
            (OrientationOption::LeftOrMiddle(_, _), _) => unreachable!(),
            (OrientationOption::RightOrMiddle(ifright, _), Some(BoardType::Right)) => ifright,
            (OrientationOption::RightOrMiddle(_, ifmiddle), Some(BoardType::Middle)) => ifmiddle,
            (OrientationOption::RightOrMiddle(_, _), _) => unreachable!(),
        };
        orientation.to_str()
    }
}

const ORIENTATIONS: [[OrientationOption; NUM_COLS]; NUM_ROWS] = [
    [
        OrientationOption::Fixed(Orientation::DownRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::DownLeft)
    ],
    [
        OrientationOption::IfRight(Orientation::UpRight, Orientation::UpDown),
        OrientationOption::Fixed(Orientation::DownLeft),
        OrientationOption::Fixed(Orientation::DownRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::UpLeft)
    ],
    [
        OrientationOption::LeftOrMiddle(Orientation::UpDown, Orientation::UpRight),
        OrientationOption::RightOrMiddle(Orientation::UpDown, Orientation::DownLeft),
        OrientationOption::Fixed(Orientation::UpRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::DownLeft)
    ],
    [
        OrientationOption::Fixed(Orientation::UpRight),
        OrientationOption::IfLeft(Orientation::LeftRight, Orientation::UpRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::UpLeft)
    ],
];

const POSSIBLE_BOARD_TYPES: [BoardType; 5] = [
    BoardType::Left,
    BoardType::LeftOrMiddle,
    BoardType::Middle,
    BoardType::RightOrMiddle,
    BoardType::Right,
];

pub type BoardArray = [[bool; NUM_COLS]; NUM_ROWS];
pub struct Board {
    board: BoardArray,
    board_type: Option<BoardType>,
}

impl Board {
    pub fn new(array: BoardArray, board_type: Option<BoardType>) -> Board {
        Board{board: array, board_type: board_type}
    }

    pub fn make_move(&mut self, m: Move) {
        if let Some(bt) = m.new_board_type {
            if !bt.applies_to(self.board_type) {
                panic!("Board type is {:?}, not compatible with {:?}", self.board_type, bt);
            }
            self.board_type = m.new_board_type
        }
        self.set_squares(m, true)
    }

    pub fn undo_move(&mut self, m: Move) {
        self.board_type = m.old_board_type;
        self.set_squares(m, false)
    }

    fn set_squares(&mut self, m: Move, mode: bool) {
        self.board[m.coord.row][m.coord.col] = mode;
        for other_space in m.extensions().iter() {
            self.board[other_space.row][other_space.col] = mode;
        }
    }

    fn occupied(&self, c: Coordinate) -> bool {
        self.board[c.row][c.col]
    }

    // This assesses whether a coordinate can be placed on the board,
    // given the current type of the board.
    fn compatible(&self, c: Coordinate) -> bool {
        // Not in the lower left, so it's a free pass.
        if !c.induces_board_type() {
            return true;
        }

        match self.board_type {
            Some(BoardType::Left)          => c != Coordinate{row: 2, col: 1} && c != Coordinate{row: 1, col: 1},
            Some(BoardType::LeftOrMiddle)  => c != Coordinate{row: 1, col: 1},
            Some(BoardType::Middle)        => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 1, col: 1},
            Some(BoardType::RightOrMiddle) => c != Coordinate{row: 3, col: 0},
            Some(BoardType::Right)         => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 2, col: 0},
            None => true,
        }
    }

    fn frontier(&self) -> Vec<Coordinate> {
        let mut results = Vec::new();
        for row in 0..NUM_ROWS {
            for col in 0..NUM_COLS {
                let coord = Coordinate{row: row, col: col};
                if self.occupied(coord) || !self.compatible(coord) {
                    continue;
                }
                let have_neighbor = POSSIBLE_DIRECTIONS.iter().any(|dir| {
                    if let Some(dest) = coord.move_in(*dir, 1) { self.occupied(dest) } else { false }
                });
                if have_neighbor {
                    results.push(coord);
                }
            }
        }
        results
    }

    fn board_type_final(&self) -> bool {
        if let Some(x) = self.board_type { x.is_final() } else { false }
    }

    pub fn legal_moves(&self) -> Vec<Move> {
        let mut results = Vec::new();
        for frontier_space in self.frontier().iter() {
            for move_type in POSSIBLE_MOVE_TYPES.iter() {
                let mov = match Move::new(*frontier_space, *move_type, self.board_type) {
                    Some(x) => x,
                    None => continue,
                };
                let mut other_space_taken = false;
                let mut induces_board_type = frontier_space.induces_board_type();
                for other_space in mov.extensions().iter() {
                    if other_space.induces_board_type() {
                        induces_board_type = true;
                    }
                    if self.occupied(*other_space) || !self.compatible(*other_space) {
                        other_space_taken = true;
                        break;
                    }
                }
                if other_space_taken {
                    continue;
                }

                if !induces_board_type || self.board_type_final() {
                    results.push(mov);
                    continue;
                }

                let mut ok_board_types: BTreeSet<_> = POSSIBLE_BOARD_TYPES.iter().cloned().filter(|board_type| {
                    board_type.applies_to(self.board_type) &&
                        board_type.induced_by(*frontier_space) &&
                        mov.extensions().iter().all(|coord| board_type.induced_by(*coord))
                }).collect();

                // Dominated board types...
                if ok_board_types.contains(&BoardType::LeftOrMiddle) {
                    ok_board_types.remove(&BoardType::Left);
                    ok_board_types.remove(&BoardType::Middle);
                }
                if ok_board_types.contains(&BoardType::RightOrMiddle) {
                    ok_board_types.remove(&BoardType::Right);
                    ok_board_types.remove(&BoardType::Middle);
                }

                for board_type in ok_board_types.iter() {
                    results.push(mov.with_board_type(*board_type));
                }
            }
        }
        results
    }
}

impl Display for Board {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), Error> {
        // Print header row
        try!(formatter.write_str("    "));
        for i in 0..NUM_COLS {
            try!(write!(formatter, " {}", i));
        }
        try!(formatter.write_str("\n"));

        // Print top box border
        try!(formatter.write_str("    ┌─"));
        for _ in 0..NUM_COLS - 1 {
            try!(formatter.write_str("┬─"));
        }
        try!(formatter.write_str("┐\n"));

        for (i, row) in self.board.iter().enumerate() {
            // Print cell content
            try!(write!(formatter, "{: >2}  │", i));
            for (j, col) in row.iter().enumerate() {
                let chr = if *col { ORIENTATIONS[i][j].for_board(self.board_type) } else { " " };
                try!(write!(formatter, "{}│", chr));
            }
            try!(formatter.write_str("\n"));

            // Print box border between rows
            if i != NUM_ROWS - 1 {
                try!(formatter.write_str("    ├─"));
                for _ in 0..NUM_COLS - 1 {
                    try!(formatter.write_str("┼─"));
                }
                try!(formatter.write_str("┤\n"));
            }
        }

        // Print bottom box border
        try!(formatter.write_str("    └─"));
        for _ in 0..NUM_COLS - 1 {
            try!(formatter.write_str("┴─"));
        }
        try!(formatter.write_str("┘\n"));

        write!(formatter, "{:?}", self.board_type)
    }
}
<commit_msg>Board#display: No need to show board type anymore<commit_after>use std::collections::BTreeSet;
use std::fmt::{Display, Error, Formatter};
use action::{POSSIBLE_DIRECTIONS,POSSIBLE_MOVE_TYPES,Coordinate,Move};

pub const NUM_COLS: usize = 5;
pub const NUM_ROWS: usize = 4;

// Hacks for the three states of the lower-left of the board in JunSeok vs YeonSeung game
#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum BoardType {
    Left,
    LeftOrMiddle,
    Middle,
    RightOrMiddle,
    Right,
}

impl BoardType {
    fn is_final(&self) -> bool {
        match *self {
            BoardType::Left => true,
            BoardType::Middle => true,
            BoardType::Right => true,
            BoardType::LeftOrMiddle => false,
            BoardType::RightOrMiddle => false,
        }
    }

    // Can a board of type `current` become this type?
    fn applies_to(&self, current: Option<BoardType>) -> bool {
        match (current, *self) {
            // a none board type can change to anything.
            (None, _) => true,
            // LeftOrMiddle can change to itself, left, or middle.
            (Some(BoardType::LeftOrMiddle), BoardType::LeftOrMiddle) => true,
            (Some(BoardType::LeftOrMiddle), BoardType::Left) => true,
            (Some(BoardType::LeftOrMiddle), BoardType::Middle) => true,
            (Some(BoardType::LeftOrMiddle), _) => false,
            // RightOrMiddle can change to itself, right, or middle.
            (Some(BoardType::RightOrMiddle), BoardType::RightOrMiddle) => true,
            (Some(BoardType::RightOrMiddle), BoardType::Right) => true,
            (Some(BoardType::RightOrMiddle), BoardType::Middle) => true,
            (Some(BoardType::RightOrMiddle), _) => false,
            // Left, Middle, Right can change to themselves only.
            (Some(BoardType::Left), BoardType::Left) => true,
            (Some(BoardType::Left), _) => false,
            (Some(BoardType::Middle), BoardType::Middle) => true,
            (Some(BoardType::Middle), _) => false,
            (Some(BoardType::Right), BoardType::Right) => true,
            (Some(BoardType::Right), _) => false,
        }
    }

    fn induced_by(&self, c: Coordinate) -> bool {
        // Not in the lower left, so it's a free pass.
        if !c.induces_board_type() {
            return true;
        }

        match *self {
            BoardType::Left          => c != Coordinate{row: 2, col: 1} && c != Coordinate{row: 1, col: 1},
            BoardType::LeftOrMiddle  => c == Coordinate{row: 1, col: 0},
            BoardType::Middle        => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 1, col: 1},
            BoardType::RightOrMiddle => c == Coordinate{row: 3, col: 1},
            BoardType::Right         => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 2, col: 0},
        }
    }
}

#[derive(Copy, Clone)]
enum Orientation {
    UpDown,
    LeftRight,
    UpLeft,
    UpRight,
    DownLeft,
    DownRight,
}

impl Orientation {
    fn to_str(&self) -> &'static str {
        match *self {
            Orientation::UpDown => "║",
            Orientation::LeftRight => "═",
            Orientation::UpLeft => "╝",
            Orientation::UpRight => "╚",
            Orientation::DownLeft => "╗",
            Orientation::DownRight => "╔",
        }
    }
}

#[derive(Copy, Clone)]
enum OrientationOption {
    Fixed(Orientation),
    IfRight(Orientation, Orientation),
    IfLeft(Orientation, Orientation),
    LeftOrMiddle(Orientation, Orientation),
    RightOrMiddle(Orientation, Orientation),
}

impl OrientationOption {
    fn for_board(&self, board_type: Option<BoardType>) -> &'static str {
        let orientation = match (*self, board_type) {
            (OrientationOption::Fixed(orientation), _) => orientation,
            (OrientationOption::IfRight(ifright, _), Some(BoardType::Right)) => ifright,
            (OrientationOption::IfRight(_, notright), _) => notright,
            (OrientationOption::IfLeft(ifleft, _), Some(BoardType::Left)) => ifleft,
            (OrientationOption::IfLeft(_, notleft), _) => notleft,
            (OrientationOption::LeftOrMiddle(ifleft, _), Some(BoardType::Left)) => ifleft,
            (OrientationOption::LeftOrMiddle(_, ifmiddle), Some(BoardType::Middle)) => ifmiddle,
            (OrientationOption::LeftOrMiddle(_, _), _) => unreachable!(),
            (OrientationOption::RightOrMiddle(ifright, _), Some(BoardType::Right)) => ifright,
            (OrientationOption::RightOrMiddle(_, ifmiddle), Some(BoardType::Middle)) => ifmiddle,
            (OrientationOption::RightOrMiddle(_, _), _) => unreachable!(),
        };
        orientation.to_str()
    }
}

const ORIENTATIONS: [[OrientationOption; NUM_COLS]; NUM_ROWS] = [
    [
        OrientationOption::Fixed(Orientation::DownRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::DownLeft)
    ],
    [
        OrientationOption::IfRight(Orientation::UpRight, Orientation::UpDown),
        OrientationOption::Fixed(Orientation::DownLeft),
        OrientationOption::Fixed(Orientation::DownRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::UpLeft)
    ],
    [
        OrientationOption::LeftOrMiddle(Orientation::UpDown, Orientation::UpRight),
        OrientationOption::RightOrMiddle(Orientation::UpDown, Orientation::DownLeft),
        OrientationOption::Fixed(Orientation::UpRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::DownLeft)
    ],
    [
        OrientationOption::Fixed(Orientation::UpRight),
        OrientationOption::IfLeft(Orientation::LeftRight, Orientation::UpRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::UpLeft)
    ],
];

const POSSIBLE_BOARD_TYPES: [BoardType; 5] = [
    BoardType::Left,
    BoardType::LeftOrMiddle,
    BoardType::Middle,
    BoardType::RightOrMiddle,
    BoardType::Right,
];

pub type BoardArray = [[bool; NUM_COLS]; NUM_ROWS];
pub struct Board {
    board: BoardArray,
    board_type: Option<BoardType>,
}

impl Board {
    pub fn new(array: BoardArray, board_type: Option<BoardType>) -> Board {
        Board{board: array, board_type: board_type}
    }

    pub fn make_move(&mut self, m: Move) {
        if let Some(bt) = m.new_board_type {
            if !bt.applies_to(self.board_type) {
                panic!("Board type is {:?}, not compatible with {:?}", self.board_type, bt);
            }
            self.board_type = m.new_board_type
        }
        self.set_squares(m, true)
    }

    pub fn undo_move(&mut self, m: Move) {
        self.board_type = m.old_board_type;
        self.set_squares(m, false)
    }

    fn set_squares(&mut self, m: Move, mode: bool) {
        self.board[m.coord.row][m.coord.col] = mode;
        for other_space in m.extensions().iter() {
            self.board[other_space.row][other_space.col] = mode;
        }
    }

    fn occupied(&self, c: Coordinate) -> bool {
        self.board[c.row][c.col]
    }

    // This assesses whether a coordinate can be placed on the board,
    // given the current type of the board.
    fn compatible(&self, c: Coordinate) -> bool {
        // Not in the lower left, so it's a free pass.
        if !c.induces_board_type() {
            return true;
        }

        match self.board_type {
            Some(BoardType::Left)          => c != Coordinate{row: 2, col: 1} && c != Coordinate{row: 1, col: 1},
            Some(BoardType::LeftOrMiddle)  => c != Coordinate{row: 1, col: 1},
            Some(BoardType::Middle)        => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 1, col: 1},
            Some(BoardType::RightOrMiddle) => c != Coordinate{row: 3, col: 0},
            Some(BoardType::Right)         => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 2, col: 0},
            None => true,
        }
    }

    fn frontier(&self) -> Vec<Coordinate> {
        let mut results = Vec::new();
        for row in 0..NUM_ROWS {
            for col in 0..NUM_COLS {
                let coord = Coordinate{row: row, col: col};
                if self.occupied(coord) || !self.compatible(coord) {
                    continue;
                }
                let have_neighbor = POSSIBLE_DIRECTIONS.iter().any(|dir| {
                    if let Some(dest) = coord.move_in(*dir, 1) { self.occupied(dest) } else { false }
                });
                if have_neighbor {
                    results.push(coord);
                }
            }
        }
        results
    }

    fn board_type_final(&self) -> bool {
        if let Some(x) = self.board_type { x.is_final() } else { false }
    }

    pub fn legal_moves(&self) -> Vec<Move> {
        let mut results = Vec::new();
        for frontier_space in self.frontier().iter() {
            for move_type in POSSIBLE_MOVE_TYPES.iter() {
                let mov = match Move::new(*frontier_space, *move_type, self.board_type) {
                    Some(x) => x,
                    None => continue,
                };
                let mut other_space_taken = false;
                let mut induces_board_type = frontier_space.induces_board_type();
                for other_space in mov.extensions().iter() {
                    if other_space.induces_board_type() {
                        induces_board_type = true;
                    }
                    if self.occupied(*other_space) || !self.compatible(*other_space) {
                        other_space_taken = true;
                        break;
                    }
                }
                if other_space_taken {
                    continue;
                }

                if !induces_board_type || self.board_type_final() {
                    results.push(mov);
                    continue;
                }

                let mut ok_board_types: BTreeSet<_> = POSSIBLE_BOARD_TYPES.iter().cloned().filter(|board_type| {
                    board_type.applies_to(self.board_type) &&
                        board_type.induced_by(*frontier_space) &&
                        mov.extensions().iter().all(|coord| board_type.induced_by(*coord))
                }).collect();

                // Dominated board types...
                if ok_board_types.contains(&BoardType::LeftOrMiddle) {
                    ok_board_types.remove(&BoardType::Left);
                    ok_board_types.remove(&BoardType::Middle);
                }
                if ok_board_types.contains(&BoardType::RightOrMiddle) {
                    ok_board_types.remove(&BoardType::Right);
                    ok_board_types.remove(&BoardType::Middle);
                }

                for board_type in ok_board_types.iter() {
                    results.push(mov.with_board_type(*board_type));
                }
            }
        }
        results
    }
}

impl Display for Board {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), Error> {
        // Print header row
        try!(formatter.write_str("    "));
        for i in 0..NUM_COLS {
            try!(write!(formatter, " {}", i));
        }
        try!(formatter.write_str("\n"));

        // Print top box border
        try!(formatter.write_str("    ┌─"));
        for _ in 0..NUM_COLS - 1 {
            try!(formatter.write_str("┬─"));
        }
        try!(formatter.write_str("┐\n"));

        for (i, row) in self.board.iter().enumerate() {
            // Print cell content
            try!(write!(formatter, "{: >2}  │", i));
            for (j, col) in row.iter().enumerate() {
                let chr = if *col { ORIENTATIONS[i][j].for_board(self.board_type) } else { " " };
                try!(write!(formatter, "{}│", chr));
            }
            try!(formatter.write_str("\n"));

            // Print box border between rows
            if i != NUM_ROWS - 1 {
                try!(formatter.write_str("    ├─"));
                for _ in 0..NUM_COLS - 1 {
                    try!(formatter.write_str("┼─"));
                }
                try!(formatter.write_str("┤\n"));
            }
        }

        // Print bottom box border
        try!(formatter.write_str("    └─"));
        for _ in 0..NUM_COLS - 1 {
            try!(formatter.write_str("┴─"));
        }
        formatter.write_str("┘\n")
    }
}
<|endoftext|>


<commit_before>use crate::algorithm::{AlgorithmType, SigningAlgorithm, VerifyingAlgorithm};
use crate::error::Error;
use crate::SEPARATOR;
use openssl::bn::BigNum;
use openssl::ecdsa::EcdsaSig;
use openssl::hash::MessageDigest;
use openssl::nid::Nid;
use openssl::pkey::{Id, PKey, Private, Public};
use openssl::sign::{Signer, Verifier};

pub struct PKeyWithDigest<T> {
    pub digest: MessageDigest,
    pub key: PKey<T>,
}

impl<T> PKeyWithDigest<T> {
    fn algorithm_type(&self) -> AlgorithmType {
        match (self.key.id(), self.digest.type_()) {
            (Id::RSA, Nid::SHA256) => AlgorithmType::Rs256,
            (Id::EC, Nid::SHA256) => AlgorithmType::Es256,
            _ => panic!("Invalid algorithm type"),
        }
    }
}

impl SigningAlgorithm for PKeyWithDigest<Private> {
    fn algorithm_type(&self) -> AlgorithmType {
        PKeyWithDigest::algorithm_type(self)
    }

    fn sign(&self, header: &str, claims: &str) -> Result<String, Error> {
        let mut signer = Signer::new(self.digest.clone(), &self.key)?;
        signer.update(header.as_bytes())?;
        signer.update(SEPARATOR.as_bytes())?;
        signer.update(claims.as_bytes())?;
        let signer_signature = signer.sign_to_vec()?;

        let signature = if self.key.id() == Id::EC {
            let der_signature = signer_signature;

            let signature = EcdsaSig::from_der(&der_signature)?;
            let r = signature.r().to_vec();
            let s = signature.s().to_vec();
            [r, s].concat()
        } else {
            signer_signature
        };

        Ok(base64::encode_config(&signature, base64::URL_SAFE_NO_PAD))
    }
}

impl VerifyingAlgorithm for PKeyWithDigest<Public> {
    fn algorithm_type(&self) -> AlgorithmType {
        PKeyWithDigest::algorithm_type(self)
    }

    fn verify_bytes(&self, header: &str, claims: &str, signature: &[u8]) -> Result<bool, Error> {
        let mut verifier = Verifier::new(self.digest.clone(), &self.key)?;
        verifier.update(header.as_bytes())?;
        verifier.update(SEPARATOR.as_bytes())?;
        verifier.update(claims.as_bytes())?;

        let verified = if self.key.id() == Id::EC {
            let (r, s) = signature.split_at(signature.len() / 2);
            let ecdsa_signature =
                EcdsaSig::from_private_components(BigNum::from_slice(r)?, BigNum::from_slice(s)?)?;
            let der_signature = ecdsa_signature.to_der()?;
            verifier.verify(&der_signature)?
        } else {
            verifier.verify(signature)?
        };

        Ok(verified)
    }
}

#[cfg(test)]
mod tests {
    use crate::algorithm::openssl::PKeyWithDigest;
    use crate::algorithm::{SigningAlgorithm, VerifyingAlgorithm};
    use openssl::hash::MessageDigest;
    use openssl::pkey::PKey;

    #[test]
    fn rs256_sign() {
        let pem = include_bytes!("../../test/private_rsa.pem");

        let algorithm = PKeyWithDigest {
            digest: MessageDigest::sha256(),
            key: PKey::private_key_from_pem(pem).unwrap(),
        };

        let result = algorithm
            .sign(
                "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9",
                "eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9",
            )
            .unwrap();
        assert_eq!(result, "EkN-DOsnsuRjRO6BxXemmJDm3HbxrbRzXglbN2S4sOkopdU4IsDxTI8jO19W_A4K8ZPJijNLis4EZsHeY559a4DFOd50_OqgHGuERTqYZyuhtF39yxJPAjUESwxk2J5k_4zM3O-vtd1Ghyo4IbqKKSy6J9mTniYJPenn5-HIirE");
    }

    #[test]
    fn rs256_verify() {
        let pem = include_bytes!("../../test/public_rsa.pem");

        let algorithm = PKeyWithDigest {
            digest: MessageDigest::sha256(),
            key: PKey::public_key_from_pem(pem).unwrap(),
        };

        assert!(algorithm.verify(
            "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9",
            "eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9",
            "EkN-DOsnsuRjRO6BxXemmJDm3HbxrbRzXglbN2S4sOkopdU4IsDxTI8jO19W_A4K8ZPJijNLis4EZsHeY559a4DFOd50_OqgHGuERTqYZyuhtF39yxJPAjUESwxk2J5k_4zM3O-vtd1Ghyo4IbqKKSy6J9mTniYJPenn5-HIirE",
        ).unwrap_or(false));
    }

    #[test]
    fn es256() {
        let header = "eyJhbGciOiJFUzI1NiIsInR5cCI6IkpXVCJ9";
        let claims = "eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9";

        let private_pem = include_bytes!("../../test/es256-private.pem");
        let private_key = PKeyWithDigest {
            digest: MessageDigest::sha256(),
            key: PKey::private_key_from_pem(private_pem).unwrap(),
        };

        let signature = private_key.sign(header, claims).unwrap();

        let public_pem = include_bytes!("../../test/es256-public.pem");

        let public_key = PKeyWithDigest {
            digest: MessageDigest::sha256(),
            key: PKey::public_key_from_pem(public_pem).unwrap(),
        };

        assert!(public_key
            .verify(header, claims, &*signature)
            .unwrap_or(false));
    }
}
<commit_msg>Refactor tests<commit_after>use crate::algorithm::{AlgorithmType, SigningAlgorithm, VerifyingAlgorithm};
use crate::error::Error;
use crate::SEPARATOR;
use openssl::bn::BigNum;
use openssl::ecdsa::EcdsaSig;
use openssl::hash::MessageDigest;
use openssl::nid::Nid;
use openssl::pkey::{Id, PKey, Private, Public};
use openssl::sign::{Signer, Verifier};

pub struct PKeyWithDigest<T> {
    pub digest: MessageDigest,
    pub key: PKey<T>,
}

impl<T> PKeyWithDigest<T> {
    fn algorithm_type(&self) -> AlgorithmType {
        match (self.key.id(), self.digest.type_()) {
            (Id::RSA, Nid::SHA256) => AlgorithmType::Rs256,
            (Id::EC, Nid::SHA256) => AlgorithmType::Es256,
            _ => panic!("Invalid algorithm type"),
        }
    }
}

impl SigningAlgorithm for PKeyWithDigest<Private> {
    fn algorithm_type(&self) -> AlgorithmType {
        PKeyWithDigest::algorithm_type(self)
    }

    fn sign(&self, header: &str, claims: &str) -> Result<String, Error> {
        let mut signer = Signer::new(self.digest.clone(), &self.key)?;
        signer.update(header.as_bytes())?;
        signer.update(SEPARATOR.as_bytes())?;
        signer.update(claims.as_bytes())?;
        let signer_signature = signer.sign_to_vec()?;

        let signature = if self.key.id() == Id::EC {
            let der_signature = signer_signature;

            let signature = EcdsaSig::from_der(&der_signature)?;
            let r = signature.r().to_vec();
            let s = signature.s().to_vec();
            [r, s].concat()
        } else {
            signer_signature
        };

        Ok(base64::encode_config(&signature, base64::URL_SAFE_NO_PAD))
    }
}

impl VerifyingAlgorithm for PKeyWithDigest<Public> {
    fn algorithm_type(&self) -> AlgorithmType {
        PKeyWithDigest::algorithm_type(self)
    }

    fn verify_bytes(&self, header: &str, claims: &str, signature: &[u8]) -> Result<bool, Error> {
        let mut verifier = Verifier::new(self.digest.clone(), &self.key)?;
        verifier.update(header.as_bytes())?;
        verifier.update(SEPARATOR.as_bytes())?;
        verifier.update(claims.as_bytes())?;

        let verified = if self.key.id() == Id::EC {
            let (r, s) = signature.split_at(signature.len() / 2);
            let ecdsa_signature =
                EcdsaSig::from_private_components(BigNum::from_slice(r)?, BigNum::from_slice(s)?)?;
            let der_signature = ecdsa_signature.to_der()?;
            verifier.verify(&der_signature)?
        } else {
            verifier.verify(signature)?
        };

        Ok(verified)
    }
}

#[cfg(test)]
mod tests {
    use crate::algorithm::openssl::PKeyWithDigest;
    use crate::algorithm::{SigningAlgorithm, VerifyingAlgorithm};
    use openssl::hash::MessageDigest;
    use openssl::pkey::PKey;

    // {"alg":"RS256","typ":"JWT"}
    const RS256_HEADER: &'static str = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9";
    // {"alg":"ES256","typ":"JWT"}
    const ES256_HEADER: &'static str = "eyJhbGciOiJFUzI1NiIsInR5cCI6IkpXVCJ9";
    // {"sub":"1234567890","name":"John Doe","admin":true}
    const CLAIMS: &'static str =
        "eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9";

    const RS256_SIGNATURE: &'static str =
    "EkN-DOsnsuRjRO6BxXemmJDm3HbxrbRzXglbN2S4sOkopdU4IsDxTI8jO19W_A4K8ZPJijNLis4EZsHeY559a4DFOd50_OqgHGuERTqYZyuhtF39y\
    xJPAjUESwxk2J5k_4zM3O-vtd1Ghyo4IbqKKSy6J9mTniYJPenn5-HIirE";

    #[test]
    fn rs256_sign() {
        let pem = include_bytes!("../../test/private_rsa.pem");

        let algorithm = PKeyWithDigest {
            digest: MessageDigest::sha256(),
            key: PKey::private_key_from_pem(pem).unwrap(),
        };

        let result = algorithm.sign(RS256_HEADER, CLAIMS).unwrap();
        assert_eq!(result, RS256_SIGNATURE);
    }

    #[test]
    fn rs256_verify() {
        let pem = include_bytes!("../../test/public_rsa.pem");

        let algorithm = PKeyWithDigest {
            digest: MessageDigest::sha256(),
            key: PKey::public_key_from_pem(pem).unwrap(),
        };

        assert!(algorithm
            .verify(RS256_HEADER, CLAIMS, RS256_SIGNATURE)
            .unwrap_or(false));
    }

    #[test]
    fn es256() {
        let private_pem = include_bytes!("../../test/es256-private.pem");
        let private_key = PKeyWithDigest {
            digest: MessageDigest::sha256(),
            key: PKey::private_key_from_pem(private_pem).unwrap(),
        };

        let signature = private_key.sign(ES256_HEADER, CLAIMS).unwrap();

        let public_pem = include_bytes!("../../test/es256-public.pem");

        let public_key = PKeyWithDigest {
            digest: MessageDigest::sha256(),
            key: PKey::public_key_from_pem(public_pem).unwrap(),
        };

        assert!(public_key
            .verify(ES256_HEADER, CLAIMS, &*signature)
            .unwrap_or(false));
    }
}
<|endoftext|>


<commit_before>use std::time::{SystemTime, Duration};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::cmp::{min, max};

use mioco::sync::{Mutex};
use mioco::sync::mpsc::{Sender};
use trust_dns::rr::{Name, RecordType, Record};
use trust_dns::op::{Message, OpCode, MessageType, ResponseCode, Query};

use utils::{CloneExt, MessageExt, AsDisplay};
use resolver::{RcResolver, RcResolverWeak};

pub type Key = Name;

#[derive(Debug, Copy, Clone)]
pub enum TtlMode {
    Original,
    Fixed(u32),
    Relative(SystemTime),
}
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum RecordSource {
    Forwarder,
    Recursive,
    Pinned,
}
pub struct RecordTypeEntry {
    message: Message,
    expiration: Option<SystemTime>,
    ttl: TtlMode,
    expiration_notifier: Option<Sender<Query>>,
    expiration_notified: AtomicBool,
    resolver: RcResolverWeak,
    source: RecordSource,
    global_expiration_counter: Arc<AtomicUsize>,
    record_expiration_counter: usize,
}
impl RecordTypeEntry {
    pub fn maybe_notify_expiration(&self) {
        let sender = match self.expiration_notifier {
            None => return,
            Some(ref s) => s,
        };
        if !self.is_expired() {
            return;
        }
        if self.expiration_notified.load(Ordering::Relaxed) {
            // Already notified
            return;
        }
        if self.expiration_notified.swap(true, Ordering::Relaxed) {
            // Lost race
            return;
        }
        let q = self.message.get_queries()[0].clone();
        debug!("Entry expired: {}", q.as_disp());
        sender.send(q.into()).is_ok();
    }
    pub fn is_expired(&self) -> bool {
        match self.expiration {
            None => false,
            Some(ref t) => {
                if self.record_expiration_counter !=
                    self.global_expiration_counter.load(Ordering::Relaxed)
                {
                    true
                } else {
                    t <= &SystemTime::now()
                }
            },
        }
    }
    pub fn adjust_ttl(&self, rec: &mut Record) {
        match self.ttl {
            TtlMode::Original => {},
            TtlMode::Fixed(secs) => { rec.ttl(secs); },
            TtlMode::Relative(relto) => {
                match relto.elapsed() {
                    Err(e) => {
                        warn!("Failed to get elapsed time for setting TTL: {:?}", e);
                    },
                    Ok(dur) => {
                        let diff = min(dur.as_secs(), i32::max_value() as u64) as u32;
                        let mut new_ttl = rec.get_ttl().saturating_sub(diff);
                        if let Some(res) = self.resolver.upgrade() {
                            new_ttl = max(new_ttl, res.config.cache.min_response_ttl);
                        }
                        if self.is_expired() {
                            new_ttl = 1
                        }
                        rec.ttl(new_ttl);
                    },
                };
            },
        };
    }
}
pub struct Entry {
    records: HashMap<RecordType, RecordTypeEntry>,
    expiration_notifier: Option<Sender<Query>>,
    global_expiration_counter: Arc<AtomicUsize>,
}
impl Entry {
    pub fn lookup_entry(&self, t: RecordType) -> Option<&RecordTypeEntry> {
        self.records.get(&t)
    }
    pub fn purge(&mut self, t: RecordType) -> Option<Message> {
        self.records.remove(&t).map(|entry| entry.message)
    }
    pub fn lookup_adjusted(&self, t: RecordType) -> Option<Message> {
        self.records.get(&t).map(|entry| {
            entry.maybe_notify_expiration();
            let mut ret = Message::new();
            ret.copy_resp_with(&entry.message, |rec| {
                let mut new_rec = rec.clone();
                entry.adjust_ttl(&mut new_rec);
                new_rec
            });
            ret
        })
    }
    fn get_ttl_mode(&self, source: RecordSource) -> TtlMode {
        match source {
            RecordSource::Pinned => TtlMode::Original,
            RecordSource::Forwarder => TtlMode::Fixed(1),
            RecordSource::Recursive => TtlMode::Relative(SystemTime::now()),
        }
    }
    pub fn update(&mut self, msg: &Message, source: RecordSource) {
        // TODO: Validate message before updating
        // TODO: Confirm that the new message is more preferable than existing one
        assert!(msg.get_queries().len() == 1);
        let t = msg.get_queries()[0].get_query_type();
        let accepted_responses = [ResponseCode::NoError, ResponseCode::NXDomain];
        if !accepted_responses.contains(&msg.get_response_code()) {
            return;
        }
        if let Some(existing) = self.records.get(&t) {
            if !existing.is_expired() && existing.source >= source {
                // Existing record is more preferable
                return;
            }
        }
        let all_records = msg.get_answers().iter()
        .chain(msg.get_name_servers())
        .chain(msg.get_additional())
        .filter(|x| x.get_rr_type() != RecordType::OPT);
        // TODO: NXDomain should be cached for shorter time
        let resolver = RcResolver::current();
        let min_cache_ttl = resolver.config.cache.min_cache_ttl;
        let cache_ttl = max(
            min_cache_ttl,
            all_records.map(|x| x.get_ttl()).min().unwrap_or(min_cache_ttl),
        ) as u64;
        debug!("Updating cache: {} -> {} (TTL: {})",
               msg.get_queries()[0].as_disp(), msg.as_disp(), cache_ttl);
        let mut cache_msg = msg.clone_resp();
        cache_msg.add_query(msg.get_queries()[0].clone());
        let ttl_mode = self.get_ttl_mode(source);
        self.records.insert(t, RecordTypeEntry {
            message: cache_msg,
            expiration: Some(SystemTime::now() + Duration::from_secs(cache_ttl)),
            ttl: ttl_mode,
            expiration_notifier: self.expiration_notifier.clone(),
            expiration_notified: AtomicBool::new(false),
            resolver: resolver.to_weak(),
            source: source,
            global_expiration_counter: self.global_expiration_counter.clone(),
            record_expiration_counter: self.global_expiration_counter.load(Ordering::Relaxed),
        });
    }
}

#[derive(Default)]
pub struct CachePlain {
    entries: HashMap<Key, Entry>,
    expiration_notifier: Option<Sender<Query>>,
    global_expiration_counter: Arc<AtomicUsize>,
}
pub type CacheInst = Mutex<CachePlain>;

impl CachePlain {
    pub fn lookup(&self, key: &Key) -> Option<&Entry> {
        self.entries.get(key)
    }
    fn lookup_mut_internal(&mut self, key: &Key) -> &mut Entry {
        self.entries.get_mut(key).unwrap()
    }
    fn create_entry(&mut self, key: &Key) -> &mut Entry {
        let notifier = self.expiration_notifier.clone();
        let global_expiration_counter = self.global_expiration_counter.clone();
        self.entries.entry(key.clone()).or_insert_with(move || Entry {
            records: HashMap::default(),
            expiration_notifier: notifier,
            global_expiration_counter: global_expiration_counter,
        })
    }
    pub fn lookup_or_insert(&mut self, key: &Key) -> &mut Entry {
        // We can't inline these functions due to errors from borrow checker
        if self.entries.contains_key(key) {
            self.lookup_mut_internal(key)
        } else {
            self.create_entry(key)
        }
    }
}
pub struct Cache {
    inst: CacheInst,
    global_expiration_counter: Arc<AtomicUsize>,
}
pub type RcCache = Arc<Cache>;
impl Cache {
    pub fn with_expiration_notifier(notifier: Sender<Query>) -> Self {
        let counter = Arc::new(AtomicUsize::default());
        Cache {
            global_expiration_counter: counter.clone(),
            inst: Mutex::new(CachePlain {
                entries: HashMap::new(),
                expiration_notifier: Some(notifier),
                global_expiration_counter: counter.clone(),
            }),
        }
    }
    pub fn lookup_with_type<F, R>(&self, key: &Key, t: RecordType, op: F) -> Option<R>
        where F: FnOnce(Message) -> R,
    {
        let guard = self.inst.lock().expect("The mutex shouldn't be poisoned");
        guard.lookup(key).and_then(|entry| entry.lookup_adjusted(t)).map(op)
    }
    pub fn operate<F, R>(&self, key: &Key, op: F) -> R
        where F: FnOnce(&mut Entry) -> R,
    {
        let mut guard = self.inst.lock().expect("The mutex shouldn't be poisoned");
        op(guard.lookup_or_insert(key))
    }
    pub fn update_from_message(&self, msg: &Message, source: RecordSource) {
        if cfg!(debug_assertions) {
            assert!(msg.get_op_code() == OpCode::Query);
            assert!(msg.get_message_type() == MessageType::Response);
            assert!(msg.get_queries().len() == 1);
            if msg.get_answers().is_empty() {
                assert!(msg.get_response_code() != ResponseCode::NoError);
            } else {
                let name = msg.get_queries()[0].get_name();
                assert!(msg.get_answers().iter().any(|x| x.get_name() == name));
            }
        }
        let name = msg.get_queries()[0].get_name();
        self.operate(name, |entry| entry.update(msg, source));
    }
    pub fn expire_all(&self) {
        self.global_expiration_counter.fetch_add(1, Ordering::Relaxed);
    }
    pub fn clear(&self) {
        let mut guard = self.inst.lock().expect("The mutex shouldn't be poisoned");
        guard.entries.clear();
    }
}
impl Default for Cache {
    fn default() -> Self {
        let counter = Arc::new(AtomicUsize::default());
        Cache {
            global_expiration_counter: counter.clone(),
            inst: Mutex::new(CachePlain {
                entries: HashMap::new(),
                expiration_notifier: None,
                global_expiration_counter: counter.clone(),
            }),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::*;
    use std::time::{SystemTime, Duration};
    use trust_dns::rr::*;
    use trust_dns::op::*;
    use super::*;
    use resolver::*;
    use utils::*;
    use ::mioco_config_start;

    fn name(raw: &str) -> Name {
        Name::parse(raw, Some(&Name::root())).unwrap()
    }
    #[test]
    fn test_record_preference() {
        mioco_config_start(move || {
            let resolver = RcResolver::default();
            let mut msg1 = Message::new();
            msg1.message_type(MessageType::Response);
            msg1.op_code(OpCode::Query);
            let name = Name::parse("www.google.com", Some(&Name::root())).unwrap();
            let t = RecordType::A;
            let mut q = Query::new();
            q.name(name.clone());
            q.query_type(t);
            msg1.add_query(q.clone());
            let mut msg2 = msg1.clone_resp_for(&q);
            let mut rec = Record::new();
            rec.ttl(3600);
            rec.name(name.clone());
            msg1.add_answer(rec.clone());
            msg2.add_answer(rec.clone());
            msg2.add_answer(rec.clone());

            let cache = &resolver.cache;
            assert!(cache.lookup_with_type(&name, t, |x| x).is_none());
            cache.update_from_message(&msg1, RecordSource::Forwarder);
            assert_eq!(cache.lookup_with_type(&name, t, |x| x).unwrap().get_answers().len(), 1);

            // Should replace cache with better message
            cache.update_from_message(&msg2, RecordSource::Recursive);
            assert_eq!(cache.lookup_with_type(&name, t, |x| x).unwrap().get_answers().len(), 2);

            // Should not replace cache with worse message
            cache.update_from_message(&msg1, RecordSource::Forwarder);
            assert_eq!(cache.lookup_with_type(&name, t, |x| x).unwrap().get_answers().len(), 2);

            // Should not replace cache until expiration
            cache.update_from_message(&msg1, RecordSource::Recursive);
            assert_eq!(cache.lookup_with_type(&name, t, |x| x).unwrap().get_answers().len(), 2);

            // Even better message
            cache.update_from_message(&msg1, RecordSource::Pinned);
            assert_eq!(cache.lookup_with_type(&name, t, |x| x).unwrap().get_answers().len(), 1);
        }).unwrap();
    }
    #[test]
    fn test_cache_case_insensitivity() {
        let mut cache = CachePlain::default();
        cache.lookup_or_insert(&name("www.google.com"));
        assert!(cache.lookup(&name("www.google.com")).is_some());
        assert!(cache.lookup(&name("www.baidu.com")).is_none());
        assert!(cache.lookup(&name("wWw.goOgle.coM")).is_some());
    }
    #[test]
    fn test_ttl_adjust() {
        mioco_config_start(move || {
            let resolver = RcResolver::default();
            let min_ttl = resolver.config.cache.min_response_ttl;
            let mut entry = RecordTypeEntry {
                message: Message::new(),
                expiration: None,
                ttl: TtlMode::Original,
                expiration_notifier: None,
                expiration_notified: AtomicBool::new(false),
                resolver: resolver.to_weak(),
                source: RecordSource::Pinned,
            };
            let mut rec = Record::new();
            macro_rules! adjust {
                ($e:expr) => {{
                    entry.ttl = $e;
                    entry.adjust_ttl(&mut rec);
                }};
            };
            rec.ttl(120);
            adjust!(TtlMode::Original);
            assert_eq!(rec.get_ttl(), 120);
            adjust!(TtlMode::Fixed(240));
            assert_eq!(rec.get_ttl(), 240);

            adjust!(TtlMode::Relative(SystemTime::now()));
            assert_eq!(rec.get_ttl(), 240);
            adjust!(TtlMode::Relative(SystemTime::now() - Duration::from_secs(10)));
            assert_eq!(rec.get_ttl(), 230);
            adjust!(TtlMode::Relative(SystemTime::now() - Duration::new(10, 999999999)));
            assert_eq!(rec.get_ttl(), 219);
            adjust!(TtlMode::Relative(SystemTime::now() - Duration::from_secs(500)));
            assert_eq!(rec.get_ttl(), min_ttl);

            rec.ttl(240);
            adjust!(TtlMode::Relative(SystemTime::now() - Duration::from_secs(u32::max_value() as u64)));
            assert_eq!(rec.get_ttl(), min_ttl);
            rec.ttl(240);

            // Using `Duration::from_secs(u64::max_value())` will silently overflow and give us
            // incorrect result, so we test with `u64::max_value() >> 1` here.
            // Ref: https://github.com/rust-lang/rust/issues/32070
            adjust!(TtlMode::Relative(SystemTime::now() - Duration::from_secs(u64::max_value() >> 1)));
            assert_eq!(rec.get_ttl(), min_ttl);
        }).unwrap();
    }
}
<commit_msg>Fix test<commit_after>use std::time::{SystemTime, Duration};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::cmp::{min, max};

use mioco::sync::{Mutex};
use mioco::sync::mpsc::{Sender};
use trust_dns::rr::{Name, RecordType, Record};
use trust_dns::op::{Message, OpCode, MessageType, ResponseCode, Query};

use utils::{CloneExt, MessageExt, AsDisplay};
use resolver::{RcResolver, RcResolverWeak};

pub type Key = Name;

#[derive(Debug, Copy, Clone)]
pub enum TtlMode {
    Original,
    Fixed(u32),
    Relative(SystemTime),
}
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum RecordSource {
    Forwarder,
    Recursive,
    Pinned,
}
pub struct RecordTypeEntry {
    message: Message,
    expiration: Option<SystemTime>,
    ttl: TtlMode,
    expiration_notifier: Option<Sender<Query>>,
    expiration_notified: AtomicBool,
    resolver: RcResolverWeak,
    source: RecordSource,
    global_expiration_counter: Arc<AtomicUsize>,
    record_expiration_counter: usize,
}
impl RecordTypeEntry {
    pub fn maybe_notify_expiration(&self) {
        let sender = match self.expiration_notifier {
            None => return,
            Some(ref s) => s,
        };
        if !self.is_expired() {
            return;
        }
        if self.expiration_notified.load(Ordering::Relaxed) {
            // Already notified
            return;
        }
        if self.expiration_notified.swap(true, Ordering::Relaxed) {
            // Lost race
            return;
        }
        let q = self.message.get_queries()[0].clone();
        debug!("Entry expired: {}", q.as_disp());
        sender.send(q.into()).is_ok();
    }
    pub fn is_expired(&self) -> bool {
        match self.expiration {
            None => false,
            Some(ref t) => {
                if self.record_expiration_counter !=
                    self.global_expiration_counter.load(Ordering::Relaxed)
                {
                    true
                } else {
                    t <= &SystemTime::now()
                }
            },
        }
    }
    pub fn adjust_ttl(&self, rec: &mut Record) {
        match self.ttl {
            TtlMode::Original => {},
            TtlMode::Fixed(secs) => { rec.ttl(secs); },
            TtlMode::Relative(relto) => {
                match relto.elapsed() {
                    Err(e) => {
                        warn!("Failed to get elapsed time for setting TTL: {:?}", e);
                    },
                    Ok(dur) => {
                        let diff = min(dur.as_secs(), i32::max_value() as u64) as u32;
                        let mut new_ttl = rec.get_ttl().saturating_sub(diff);
                        if let Some(res) = self.resolver.upgrade() {
                            new_ttl = max(new_ttl, res.config.cache.min_response_ttl);
                        }
                        if self.is_expired() {
                            new_ttl = 1
                        }
                        rec.ttl(new_ttl);
                    },
                };
            },
        };
    }
}
pub struct Entry {
    records: HashMap<RecordType, RecordTypeEntry>,
    expiration_notifier: Option<Sender<Query>>,
    global_expiration_counter: Arc<AtomicUsize>,
}
impl Entry {
    pub fn lookup_entry(&self, t: RecordType) -> Option<&RecordTypeEntry> {
        self.records.get(&t)
    }
    pub fn purge(&mut self, t: RecordType) -> Option<Message> {
        self.records.remove(&t).map(|entry| entry.message)
    }
    pub fn lookup_adjusted(&self, t: RecordType) -> Option<Message> {
        self.records.get(&t).map(|entry| {
            entry.maybe_notify_expiration();
            let mut ret = Message::new();
            ret.copy_resp_with(&entry.message, |rec| {
                let mut new_rec = rec.clone();
                entry.adjust_ttl(&mut new_rec);
                new_rec
            });
            ret
        })
    }
    fn get_ttl_mode(&self, source: RecordSource) -> TtlMode {
        match source {
            RecordSource::Pinned => TtlMode::Original,
            RecordSource::Forwarder => TtlMode::Fixed(1),
            RecordSource::Recursive => TtlMode::Relative(SystemTime::now()),
        }
    }
    pub fn update(&mut self, msg: &Message, source: RecordSource) {
        // TODO: Validate message before updating
        // TODO: Confirm that the new message is more preferable than existing one
        assert!(msg.get_queries().len() == 1);
        let t = msg.get_queries()[0].get_query_type();
        let accepted_responses = [ResponseCode::NoError, ResponseCode::NXDomain];
        if !accepted_responses.contains(&msg.get_response_code()) {
            return;
        }
        if let Some(existing) = self.records.get(&t) {
            if !existing.is_expired() && existing.source >= source {
                // Existing record is more preferable
                return;
            }
        }
        let all_records = msg.get_answers().iter()
        .chain(msg.get_name_servers())
        .chain(msg.get_additional())
        .filter(|x| x.get_rr_type() != RecordType::OPT);
        // TODO: NXDomain should be cached for shorter time
        let resolver = RcResolver::current();
        let min_cache_ttl = resolver.config.cache.min_cache_ttl;
        let cache_ttl = max(
            min_cache_ttl,
            all_records.map(|x| x.get_ttl()).min().unwrap_or(min_cache_ttl),
        ) as u64;
        debug!("Updating cache: {} -> {} (TTL: {})",
               msg.get_queries()[0].as_disp(), msg.as_disp(), cache_ttl);
        let mut cache_msg = msg.clone_resp();
        cache_msg.add_query(msg.get_queries()[0].clone());
        let ttl_mode = self.get_ttl_mode(source);
        self.records.insert(t, RecordTypeEntry {
            message: cache_msg,
            expiration: Some(SystemTime::now() + Duration::from_secs(cache_ttl)),
            ttl: ttl_mode,
            expiration_notifier: self.expiration_notifier.clone(),
            expiration_notified: AtomicBool::new(false),
            resolver: resolver.to_weak(),
            source: source,
            global_expiration_counter: self.global_expiration_counter.clone(),
            record_expiration_counter: self.global_expiration_counter.load(Ordering::Relaxed),
        });
    }
}

#[derive(Default)]
pub struct CachePlain {
    entries: HashMap<Key, Entry>,
    expiration_notifier: Option<Sender<Query>>,
    global_expiration_counter: Arc<AtomicUsize>,
}
pub type CacheInst = Mutex<CachePlain>;

impl CachePlain {
    pub fn lookup(&self, key: &Key) -> Option<&Entry> {
        self.entries.get(key)
    }
    fn lookup_mut_internal(&mut self, key: &Key) -> &mut Entry {
        self.entries.get_mut(key).unwrap()
    }
    fn create_entry(&mut self, key: &Key) -> &mut Entry {
        let notifier = self.expiration_notifier.clone();
        let global_expiration_counter = self.global_expiration_counter.clone();
        self.entries.entry(key.clone()).or_insert_with(move || Entry {
            records: HashMap::default(),
            expiration_notifier: notifier,
            global_expiration_counter: global_expiration_counter,
        })
    }
    pub fn lookup_or_insert(&mut self, key: &Key) -> &mut Entry {
        // We can't inline these functions due to errors from borrow checker
        if self.entries.contains_key(key) {
            self.lookup_mut_internal(key)
        } else {
            self.create_entry(key)
        }
    }
}
pub struct Cache {
    inst: CacheInst,
    global_expiration_counter: Arc<AtomicUsize>,
}
pub type RcCache = Arc<Cache>;
impl Cache {
    pub fn with_expiration_notifier(notifier: Sender<Query>) -> Self {
        let counter = Arc::new(AtomicUsize::default());
        Cache {
            global_expiration_counter: counter.clone(),
            inst: Mutex::new(CachePlain {
                entries: HashMap::new(),
                expiration_notifier: Some(notifier),
                global_expiration_counter: counter.clone(),
            }),
        }
    }
    pub fn lookup_with_type<F, R>(&self, key: &Key, t: RecordType, op: F) -> Option<R>
        where F: FnOnce(Message) -> R,
    {
        let guard = self.inst.lock().expect("The mutex shouldn't be poisoned");
        guard.lookup(key).and_then(|entry| entry.lookup_adjusted(t)).map(op)
    }
    pub fn operate<F, R>(&self, key: &Key, op: F) -> R
        where F: FnOnce(&mut Entry) -> R,
    {
        let mut guard = self.inst.lock().expect("The mutex shouldn't be poisoned");
        op(guard.lookup_or_insert(key))
    }
    pub fn update_from_message(&self, msg: &Message, source: RecordSource) {
        if cfg!(debug_assertions) {
            assert!(msg.get_op_code() == OpCode::Query);
            assert!(msg.get_message_type() == MessageType::Response);
            assert!(msg.get_queries().len() == 1);
            if msg.get_answers().is_empty() {
                assert!(msg.get_response_code() != ResponseCode::NoError);
            } else {
                let name = msg.get_queries()[0].get_name();
                assert!(msg.get_answers().iter().any(|x| x.get_name() == name));
            }
        }
        let name = msg.get_queries()[0].get_name();
        self.operate(name, |entry| entry.update(msg, source));
    }
    pub fn expire_all(&self) {
        self.global_expiration_counter.fetch_add(1, Ordering::Relaxed);
    }
    pub fn clear(&self) {
        let mut guard = self.inst.lock().expect("The mutex shouldn't be poisoned");
        guard.entries.clear();
    }
}
impl Default for Cache {
    fn default() -> Self {
        let counter = Arc::new(AtomicUsize::default());
        Cache {
            global_expiration_counter: counter.clone(),
            inst: Mutex::new(CachePlain {
                entries: HashMap::new(),
                expiration_notifier: None,
                global_expiration_counter: counter.clone(),
            }),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::*;
    use std::time::{SystemTime, Duration};
    use trust_dns::rr::*;
    use trust_dns::op::*;
    use super::*;
    use resolver::*;
    use utils::*;
    use ::mioco_config_start;

    fn name(raw: &str) -> Name {
        Name::parse(raw, Some(&Name::root())).unwrap()
    }
    #[test]
    fn test_record_preference() {
        mioco_config_start(move || {
            let resolver = RcResolver::default();
            let mut msg1 = Message::new();
            msg1.message_type(MessageType::Response);
            msg1.op_code(OpCode::Query);
            let name = Name::parse("www.google.com", Some(&Name::root())).unwrap();
            let t = RecordType::A;
            let mut q = Query::new();
            q.name(name.clone());
            q.query_type(t);
            msg1.add_query(q.clone());
            let mut msg2 = msg1.clone_resp_for(&q);
            let mut rec = Record::new();
            rec.ttl(3600);
            rec.name(name.clone());
            msg1.add_answer(rec.clone());
            msg2.add_answer(rec.clone());
            msg2.add_answer(rec.clone());

            let cache = &resolver.cache;
            assert!(cache.lookup_with_type(&name, t, |x| x).is_none());
            cache.update_from_message(&msg1, RecordSource::Forwarder);
            assert_eq!(cache.lookup_with_type(&name, t, |x| x).unwrap().get_answers().len(), 1);

            // Should replace cache with better message
            cache.update_from_message(&msg2, RecordSource::Recursive);
            assert_eq!(cache.lookup_with_type(&name, t, |x| x).unwrap().get_answers().len(), 2);

            // Should not replace cache with worse message
            cache.update_from_message(&msg1, RecordSource::Forwarder);
            assert_eq!(cache.lookup_with_type(&name, t, |x| x).unwrap().get_answers().len(), 2);

            // Should not replace cache until expiration
            cache.update_from_message(&msg1, RecordSource::Recursive);
            assert_eq!(cache.lookup_with_type(&name, t, |x| x).unwrap().get_answers().len(), 2);

            // Even better message
            cache.update_from_message(&msg1, RecordSource::Pinned);
            assert_eq!(cache.lookup_with_type(&name, t, |x| x).unwrap().get_answers().len(), 1);
        }).unwrap();
    }
    #[test]
    fn test_cache_case_insensitivity() {
        let mut cache = CachePlain::default();
        cache.lookup_or_insert(&name("www.google.com"));
        assert!(cache.lookup(&name("www.google.com")).is_some());
        assert!(cache.lookup(&name("www.baidu.com")).is_none());
        assert!(cache.lookup(&name("wWw.goOgle.coM")).is_some());
    }
    #[test]
    fn test_ttl_adjust() {
        mioco_config_start(move || {
            let resolver = RcResolver::default();
            let min_ttl = resolver.config.cache.min_response_ttl;
            let mut entry = RecordTypeEntry {
                message: Message::new(),
                expiration: None,
                ttl: TtlMode::Original,
                expiration_notifier: None,
                expiration_notified: AtomicBool::new(false),
                resolver: resolver.to_weak(),
                source: RecordSource::Pinned,
                record_expiration_counter: 0,
                global_expiration_counter: Default::default(),
            };
            let mut rec = Record::new();
            macro_rules! adjust {
                ($e:expr) => {{
                    entry.ttl = $e;
                    entry.adjust_ttl(&mut rec);
                }};
            };
            rec.ttl(120);
            adjust!(TtlMode::Original);
            assert_eq!(rec.get_ttl(), 120);
            adjust!(TtlMode::Fixed(240));
            assert_eq!(rec.get_ttl(), 240);

            adjust!(TtlMode::Relative(SystemTime::now()));
            assert_eq!(rec.get_ttl(), 240);
            adjust!(TtlMode::Relative(SystemTime::now() - Duration::from_secs(10)));
            assert_eq!(rec.get_ttl(), 230);
            adjust!(TtlMode::Relative(SystemTime::now() - Duration::new(10, 999999999)));
            assert_eq!(rec.get_ttl(), 219);
            adjust!(TtlMode::Relative(SystemTime::now() - Duration::from_secs(500)));
            assert_eq!(rec.get_ttl(), min_ttl);

            rec.ttl(240);
            adjust!(TtlMode::Relative(SystemTime::now() - Duration::from_secs(u32::max_value() as u64)));
            assert_eq!(rec.get_ttl(), min_ttl);
            rec.ttl(240);

            // Using `Duration::from_secs(u64::max_value())` will silently overflow and give us
            // incorrect result, so we test with `u64::max_value() >> 1` here.
            // Ref: https://github.com/rust-lang/rust/issues/32070
            adjust!(TtlMode::Relative(SystemTime::now() - Duration::from_secs(u64::max_value() >> 1)));
            assert_eq!(rec.get_ttl(), min_ttl);
        }).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::ffi::CString;
use std::string::String;
use std::sync::{Arc, Mutex};

use url::{UrlParser,SchemeType};

use dfs::HdfsFS;
use native::*;

pub static LOCALFS_SCHEME: &'static str = "file";

/// for HDFS URL scheme (i.e., hdfs://)
pub fn hdfs_scheme_handler(scheme: &str) -> SchemeType
{
  match scheme {
    "file" => SchemeType::FileLike,
    "hdfs" => SchemeType::Relative(50070),
    _ => panic!("Unsupported scheme: {}", scheme)
  }
}

/// HdfsFs Cache. Basically, It is a cache as well as a way to guarantees
/// thread-safe when you get HdfsFs.
pub struct HdfsFsCache<'a> {
  fs_map: HashMap<String, HdfsFS<'a>>,
  lock: Arc<Mutex<i32>>,
  url_parser: UrlParser<'a>
}

impl<'a> HdfsFsCache<'a> {

  pub fn new() -> HdfsFsCache<'a> {
    let mut url_parser = UrlParser::new();
    url_parser.scheme_type_mapper(hdfs_scheme_handler);

    HdfsFsCache {
      fs_map: HashMap::new(),
      lock: Arc::new(Mutex::new(0)),
      url_parser: url_parser
    }
  }

  pub fn get(&mut self, path: &str) -> Option<&HdfsFS> {
    let mut namenode_uri = String::new();

    match self.url_parser.parse(path) {
      Ok(url) => {

        if &url.scheme == LOCALFS_SCHEME {
          namenode_uri.push_str("file:///");
        } else {

          if url.host().is_some() {
            namenode_uri.push_str(&(
              format!("{}://{}", &url.scheme, url.host().unwrap())));

            if url.port().is_some() {
              namenode_uri.push_str(&(format!(":{}", url.port().unwrap())));
            }

          } else {
            panic!("No hostname");
          }
        }
      },
      Err(_) => {
        error!("Path parsing failed: {}", path);
      }
    };

    self.lock.lock();

    info!("Connect to Namenode ({})", &namenode_uri);

    if !self.fs_map.contains_key(&namenode_uri) {
      unsafe {
        let hdfs_builder = hdfsNewBuilder();
        let namenode_uri_bytes: Vec<u8> = namenode_uri.bytes().collect();
        let namenode_cstr = CString::new(namenode_uri_bytes).unwrap();
        hdfsBuilderSetNameNode(hdfs_builder, namenode_cstr.as_ptr());
        let hdfs_fs = hdfsBuilderConnect(hdfs_builder);

        if hdfs_fs.is_null() {
          return None
        }

        self.fs_map.insert(
          namenode_uri.clone(),
          HdfsFS::new(namenode_uri.clone(), hdfs_fs));
      }
    }

    Some(self.fs_map.get(&namenode_uri).unwrap())
  }
}<commit_msg>Reformat cache.rs.<commit_after>use std::collections::HashMap;
use std::ffi::CString;
use std::string::String;
use std::sync::{Arc, Mutex};

use url::{UrlParser,SchemeType};

use dfs::HdfsFS;
use native::*;

pub static LOCAL_FS_SCHEME: &'static str = "file";

/// for HDFS URL scheme (i.e., hdfs://)
pub fn hdfs_scheme_handler(scheme: &str) -> SchemeType
{
  match scheme {
    "file" => SchemeType::FileLike,
    "hdfs" => SchemeType::Relative(50070),
    _ => panic!("Unsupported scheme: {}", scheme)
  }
}

/// HdfsFs Cache. Basically, It is a cache as well as a way to guarantees
/// thread-safe when you get HdfsFs.
pub struct HdfsFsCache<'a>
{
  fs_map: HashMap<String, HdfsFS<'a>>,
  lock: Arc<Mutex<i32>>,
  url_parser: UrlParser<'a>
}

impl<'a> HdfsFsCache<'a>
{
  pub fn new() -> HdfsFsCache<'a>
  {
    let mut url_parser = UrlParser::new();
    url_parser.scheme_type_mapper(hdfs_scheme_handler);

    HdfsFsCache {
      fs_map: HashMap::new(),
      lock: Arc::new(Mutex::new(0)),
      url_parser: url_parser
    }
  }

  pub fn get(&mut self, path: &str) -> Option<&HdfsFS>
  {
    let mut namenode_uri = String::new();

    match self.url_parser.parse(path) {
      Ok(url) => {

        if &url.scheme == LOCAL_FS_SCHEME {
          namenode_uri.push_str("file:///");
        } else {

          if url.host().is_some() {
            namenode_uri.push_str(&(
              format!("{}://{}", &url.scheme, url.host().unwrap())));

            if url.port().is_some() {
              namenode_uri.push_str(&(format!(":{}", url.port().unwrap())));
            }

          } else {
            panic!("No hostname");
          }
        }
      },
      Err(_) => {
        error!("Path parsing failed: {}", path);
      }
    };

    self.lock.lock();

    info!("Connect to Namenode ({})", &namenode_uri);

    if !self.fs_map.contains_key(&namenode_uri) {
      unsafe {
        let hdfs_builder = hdfsNewBuilder();
        let namenode_uri_bytes: Vec<u8> = namenode_uri.bytes().collect();
        let namenode_cstr = CString::new(namenode_uri_bytes).unwrap();
        hdfsBuilderSetNameNode(hdfs_builder, namenode_cstr.as_ptr());
        let hdfs_fs = hdfsBuilderConnect(hdfs_builder);

        if hdfs_fs.is_null() {
          return None
        }

        self.fs_map.insert(
          namenode_uri.clone(),
          HdfsFS::new(namenode_uri.clone(), hdfs_fs));
      }
    }

    Some(self.fs_map.get(&namenode_uri).unwrap())
  }
}<|endoftext|>


<commit_before>
fn gf_multiply(a_in: u8, b_in: u8) -> u8 {
    let mut a = a_in;
    let mut b = b_in;
    let mut result = 0;
    for _ in 0..8 {
        if b & 1 != 0 {
            result ^= a;
        }
        let high_bit = a & 0x80;
        a <<= 1;
        b >>= 1;
        if high_bit != 0 {
            a ^= 0x1B;
        }
    }
    result
}

#[test]
fn gf_multiply_test() {
    assert_eq!(0x01, gf_multiply(0x03, 0xf6));
    assert_eq!(0x04, gf_multiply(0x02, 0x02));
}

fn generate_sbox() -> [u8; 256] {
    let inv3: u8 = 0xf6;
    let mut up: u8 = 1;
    let mut down: u8 = 1;
    let mut sbox = [0; 256];

    while up != 1 || sbox[1] == 0 {
        up = gf_multiply(up, 0x03);
        down = gf_multiply(down, inv3);
        let mut value = down;
        value ^= down.rotate_left(1);
        value ^= down.rotate_left(2);
        value ^= down.rotate_left(3);
        value ^= down.rotate_left(4);
        value ^= 0x63;
        sbox[up as usize] = value;
    }
    sbox[0] = 0x63;
    sbox
}

fn invert_sbox(inbox: [u8; 256]) -> [u8; 256] {
    let mut outbox = [0; 256];
    for (i, v) in inbox.iter().enumerate() {
        outbox[*v as usize] = i as u8;
    }
    outbox
}

#[test]
fn generate_sbox_test() {
    for (i, v) in expand_key_128([0;16], &generate_sbox()).iter().enumerate() {
        print!("{:02x} ", v);
        if i % 16 == 15 {
            println!("");
        }
    }
}

fn schedule_core(bytes: &mut [u8; 4], iterations: u8, sbox: &[u8; 256]) {
    let fst = bytes[0];
    for i in 0..3 {
        bytes[i] = sbox[bytes[i+1] as usize];
    }
    bytes[3] = sbox[fst as usize];
    let mut rcon = 1;
    for i in 1..iterations {
        rcon = gf_multiply(rcon, 2);
    }
    bytes[0] ^= rcon;
}

fn expand_key_128(key: [u8; 16], sbox: &[u8; 256]) -> [u8; 176] {
    let mut tmp = [0; 4];
    let mut count = 16;
    let mut iterations = 1;
    let mut expanded: [u8; 176] = [0; 176];
    for (index, byte) in key.iter().enumerate() {
        expanded[index] = *byte;
    }

    while count < 176 {
        for i in 0..4 {
            tmp[i] = expanded[count + i - 4];
        }
        if count % 16 == 0 {
            schedule_core(&mut tmp, iterations, sbox);
            iterations += 1;
        }
        for i in 0..4 {
            expanded[count] = expanded[count-16] ^ tmp[i];
            count += 1;
        }
    }
    expanded
}

fn encrypt_128(key: [u8; 16], block: [u8; 16], sbox: &[u8; 256]) -> [u8; 16] {
    let mut state = block;
    let ekey = expand_key_128(key, sbox);
    add_round_key(&mut state, &ekey[0..16]);
    for round in 1..10 {
        sub_bytes(&mut state, sbox);
        shift_rows(&mut state);
        mix_columns(&mut state);
        add_round_key(&mut state, &ekey[round*16 .. (round+1)*16]);
    }
    sub_bytes(&mut state, sbox);
    shift_rows(&mut state);
    add_round_key(&mut state, &ekey[10*16 .. 11 * 16]);
    state
}

fn add_round_key(state: &mut [u8; 16], round_key: &[u8]) {
    for (i, byte) in round_key.iter().enumerate() {
        state[i] ^= *byte;
    }
}

fn sub_bytes(state: &mut [u8; 16], sbox: &[u8; 256]) {
    for i in 0..16 {
        state[i] = sbox[state[i] as usize];
    }
}

fn shift_rows(state: &mut [u8; 16]) {
    let copy = state.clone();
    for i in 0..16 {
        let row = i % 4;
        let col = i / 4;
        let target = (row + col) % 4;
        state[i] = copy[target * 4 + row];
    }
}

fn mix_columns(state: &mut [u8; 16]) {
    for i in 0..4 {
        let b0 = state[4*i+0];
        let b1 = state[4*i+1];
        let b2 = state[4*i+2];
        let b3 = state[4*i+3];
        state[4*i+0] = gf_multiply(2, b0) ^ gf_multiply(3, b1) ^ b2 ^ b3;
        state[4*i+1] = b0 ^ gf_multiply(2, b1) ^ gf_multiply(3, b2) ^ b3;
        state[4*i+2] = b0 ^ b1 ^ gf_multiply(2, b2) ^ gf_multiply(3, b3);
        state[4*i+3] = gf_multiply(3, b0) ^ b1 ^ b2 ^ gf_multiply(2, b3);
    }
}
#[test]
fn encrypt_test() {
    let plain = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff];
    let key = [0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f];
    for (i, v) in encrypt_128(key, plain, &generate_sbox()).iter().enumerate() {
        print!("{:02x} ", v);
        if i % 16 == 15 {
            println!("");
        }
    }
}

<commit_msg>Improved tests<commit_after>
fn gf_multiply(a_in: u8, b_in: u8) -> u8 {
    let mut a = a_in;
    let mut b = b_in;
    let mut result = 0;
    for _ in 0..8 {
        if b & 1 != 0 {
            result ^= a;
        }
        let high_bit = a & 0x80;
        a <<= 1;
        b >>= 1;
        if high_bit != 0 {
            a ^= 0x1B;
        }
    }
    result
}

#[test]
fn gf_multiply_test() {
    assert_eq!(0x01, gf_multiply(0x03, 0xf6));
    assert_eq!(0x04, gf_multiply(0x02, 0x02));
}

fn generate_sbox() -> [u8; 256] {
    let inv3: u8 = 0xf6;
    let mut up: u8 = 1;
    let mut down: u8 = 1;
    let mut sbox = [0; 256];

    while up != 1 || sbox[1] == 0 {
        up = gf_multiply(up, 0x03);
        down = gf_multiply(down, inv3);
        let mut value = down;
        value ^= down.rotate_left(1);
        value ^= down.rotate_left(2);
        value ^= down.rotate_left(3);
        value ^= down.rotate_left(4);
        value ^= 0x63;
        sbox[up as usize] = value;
    }
    sbox[0] = 0x63;
    sbox
}

fn invert_sbox(inbox: [u8; 256]) -> [u8; 256] {
    let mut outbox = [0; 256];
    for (i, v) in inbox.iter().enumerate() {
        outbox[*v as usize] = i as u8;
    }
    outbox
}

#[test]
fn generate_sbox_test() {
    let expected: [u8; 256] =
        [0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76
        ,0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0
        ,0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15
        ,0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75
        ,0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84
        ,0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF
        ,0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8
        ,0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2
        ,0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73
        ,0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB
        ,0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79
        ,0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08
        ,0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A
        ,0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E
        ,0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF
        ,0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16];
    assert_eq!(expected[..], generate_sbox()[..]);
}

fn schedule_core(bytes: &mut [u8; 4], iterations: u8, sbox: &[u8; 256]) {
    let fst = bytes[0];
    for i in 0..3 {
        bytes[i] = sbox[bytes[i+1] as usize];
    }
    bytes[3] = sbox[fst as usize];
    let mut rcon = 1;
    for i in _..iterations {
        rcon = gf_multiply(rcon, 2);
    }
    bytes[0] ^= rcon;
}

fn expand_key_128(key: [u8; 16], sbox: &[u8; 256]) -> [u8; 176] {
    let mut tmp = [0; 4];
    let mut count = 16;
    let mut iterations = 1;
    let mut expanded: [u8; 176] = [0; 176];
    for (index, byte) in key.iter().enumerate() {
        expanded[index] = *byte;
    }

    while count < 176 {
        for i in 0..4 {
            tmp[i] = expanded[count + i - 4];
        }
        if count % 16 == 0 {
            schedule_core(&mut tmp, iterations, sbox);
            iterations += 1;
        }
        for i in 0..4 {
            expanded[count] = expanded[count-16] ^ tmp[i];
            count += 1;
        }
    }
    expanded
}

fn encrypt_128(key: [u8; 16], block: [u8; 16], sbox: &[u8; 256]) -> [u8; 16] {
    let mut state = block;
    let ekey = expand_key_128(key, sbox);
    add_round_key(&mut state, &ekey[0..16]);
    for round in 1..10 {
        sub_bytes(&mut state, sbox);
        shift_rows(&mut state);
        mix_columns(&mut state);
        add_round_key(&mut state, &ekey[round*16 .. (round+1)*16]);
    }
    sub_bytes(&mut state, sbox);
    shift_rows(&mut state);
    add_round_key(&mut state, &ekey[10*16 .. 11 * 16]);
    state
}

fn add_round_key(state: &mut [u8; 16], round_key: &[u8]) {
    for (i, byte) in round_key.iter().enumerate() {
        state[i] ^= *byte;
    }
}

fn sub_bytes(state: &mut [u8; 16], sbox: &[u8; 256]) {
    for i in 0..16 {
        state[i] = sbox[state[i] as usize];
    }
}

fn shift_rows(state: &mut [u8; 16]) {
    let copy = state.clone();
    for i in 0..16 {
        let row = i % 4;
        let col = i / 4;
        let target = (row + col) % 4;
        state[i] = copy[target * 4 + row];
    }
}

fn mix_columns(state: &mut [u8; 16]) {
    for i in 0..4 {
        let b0 = state[4*i+0];
        let b1 = state[4*i+1];
        let b2 = state[4*i+2];
        let b3 = state[4*i+3];
        state[4*i+0] = gf_multiply(2, b0) ^ gf_multiply(3, b1) ^ b2 ^ b3;
        state[4*i+1] = b0 ^ gf_multiply(2, b1) ^ gf_multiply(3, b2) ^ b3;
        state[4*i+2] = b0 ^ b1 ^ gf_multiply(2, b2) ^ gf_multiply(3, b3);
        state[4*i+3] = gf_multiply(3, b0) ^ b1 ^ b2 ^ gf_multiply(2, b3);
    }
}
#[test]
fn encrypt_test() {
    let plain = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff];
    let key = [0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f];
    let expected = [0x69, 0xc4, 0xe0, 0xd8, 0x6a, 0x7b, 0x04, 0x30, 0xd8, 0xcd, 0xb7, 0x80, 0x70, 0xb4, 0xc5, 0x5a];
    let actual = encrypt_128(key, plain, &generate_sbox());
    assert_eq!(expected, actual);
}

<|endoftext|>


<commit_before>// rust
// http://rosettacode.org/wiki/foo
// Accepts two command line arguments
// cargo run --name agm arg1 arg2

use std::os;
use std::num;

fn main () {
    let args = os::args();
    let args = args.as_slice();

    let x = from_str::<f32>(args[1].as_slice()).unwrap() ;
    let y = from_str::<f32>(args[2].as_slice()).unwrap() ;

    let result = agm(x,y);
    println!("The arithmetic-geometric mean is {}", result);
}

fn agm (x: f32, y: f32) -> (f32) {
    let e: f32 = 0.000001;
    let mut a = x;
    let mut g = y;
    let mut a1: f32;
    let mut g1: f32;

    loop {
            a1 = (a + g) / 2f32;
            g1 = (a * g).sqrt();
            a = a1;
            g = g1;
            if num::abs( a - g) < e {  return a; }
        }

}

#[test]
fn_test_agm_12_6() {
    assert! ( agm(12,6) == 8.740746 );
}
<commit_msg>fixed test<commit_after>// rust
// http://rosettacode.org/wiki/foo
// Accepts two command line arguments
// cargo run --name agm arg1 arg2

use std::os;
use std::num;

fn main () {
    let args = os::args();
    let args = args.as_slice();

    let x = from_str::<f32>(args[1].as_slice()).unwrap() ;
    let y = from_str::<f32>(args[2].as_slice()).unwrap() ;

    let result = agm(x,y);
    println!("The arithmetic-geometric mean is {}", result);
}

fn agm (x: f32, y: f32) -> (f32) {
    let e: f32 = 0.000001;
    let mut a = x;
    let mut g = y;
    let mut a1: f32;
    let mut g1: f32;

    loop {
            a1 = (a + g) / 2f32;
            g1 = (a * g).sqrt();
            a = a1;
            g = g1;
            if num::abs( a - g) < e {  return a; }
        }

}

#[test]
fn test_agm_12_6() {
    assert! ( agm(12,6) == 8.740746 );
}
<|endoftext|>