<commit_before>use crate::AsArray2D;
use russell_lab::Matrix;
use std::collections::HashMap;

pub enum BlockKind {
    Qua4,
    Qua8,
    Hex8,
    Hex20,
}

pub enum Constraint {}

pub struct ConstraintData {
    kind: Constraint,
}

/// Defines a polygon on polyhedron that can be split into smaller shapes
///
/// # Block shape
///
/// ## 2D: Local indices of points and edges
///
/// ```text
///             Points (4 or 8)                        Edges (4)
///   y
///   |        3      6      2                             2
///   +--x      @-----@-----@                        +-----------+
///             |           |                        |           |
///             |           |                        |           |
///           7 @           @ 5                     3|           |1
///             |           |                        |           |
///             |           |                        |           |
///             @-----@-----@                        +-----------+
///            0      4      1                             0
/// ```
///
/// ## 3D: Local indices of points, edges, and faces
///
/// ```text
///                   Points (8 or 20)                   Faces (6)
///    z
///    |           4        15        7
///   ,+--y         @-------@--------@                 +----------------+
/// x'            ,'|              ,'|               ,'|              ,'|
///          12 @'  |         14 ,'  |             ,'  |  ___       ,'  |
///           ,'    |16        ,@    |19         ,'    |,'5,'  [0],'    |
///     5   ,'      @      6 ,'      @         ,'      |~~~     ,'      |
///       @'=======@=======@'        |       +'===============+'  ,'|   |
///       |      13 |      |         |       |   ,'|   |      |   |3|   |
///       |         |      |  11     |       |   |2|   |      |   |,'   |
///    17 |       0 @- - - | @- - - -@       |   |,'   +- - - | +- - - -+
///       @       ,'       @       ,' 3      |       ,'       |       ,'
///       |   8 @'      18 |     ,'          |     ,' [1]  ___|     ,'
///       |   ,'           |   ,@ 10         |   ,'      ,'4,'|   ,'
///       | ,'             | ,'              | ,'        ~~~  | ,'
///       @-------@--------@'                +----------------+'
///     1         9         2
///
///
///                  Edges (12)
///
///                         7
///                 +----------------+
///               ,'|              ,'|
///           4 ,'  |8          6,'  |
///           ,'    |          ,'    |
///         ,'      |   5    ,'      |11
///       +'===============+'        |
///       |         |      |         |
///       |         |      |  3      |
///       |         .- - - | -  - - -+
///      9|       ,'       |       ,'
///       |    0,'         |10   ,'
///       |   ,'           |   ,' 2
///       | ,'             | ,'
///       +----------------+'
///               1
///
pub struct Block {
    group: usize,             // group of all elements in this block
    ndim: usize,              // space dimension
    npoint: usize,            // number of points (4, 8, or 20)
    nedge: usize,             // number of edges (4 or 12)
    nface: usize,             // number of faces (4 or 6)
    coords: Matrix,           // coordinates of points (npoint, ndim)
    point_groups: Vec<usize>, // point groups (npoint)
    edge_groups: Vec<usize>,  // edge groups (nedge)
    face_groups: Vec<usize>,  // face groups (nface)
    num_div: Vec<usize>,      // number of divisions along each dim (ndim)
    weights: Vec<Vec<f64>>,   // weights along each dimension (ndim, nweight)
    sum_weights: Vec<f64>,    // sum of weights along each dimension (ndim)

    // maps side to constraint (num_sides)
    constraints: HashMap<Constraint, ConstraintData>,
}

impl Block {
    /// Creates a new Block with default options
    pub fn new(group: usize, kind: BlockKind) -> Self {
        let (ndim, npoint, nedge, nface) = match kind {
            BlockKind::Qua4 => (2, 4, 4, 0),
            BlockKind::Qua8 => (2, 8, 4, 0),
            BlockKind::Hex8 => (3, 8, 12, 6),
            BlockKind::Hex20 => (3, 20, 12, 6),
        };
        Block {
            group,
            ndim,
            npoint,
            nedge,
            nface,
            coords: Matrix::new(npoint, ndim),
            point_groups: vec![0; npoint],
            edge_groups: vec![0; nedge],
            face_groups: vec![0; nface],
            num_div: vec![0; ndim],
            weights: vec![Vec::new(); ndim],
            sum_weights: vec![0.0; ndim],
            constraints: HashMap::new(),
        }
    }

    /// Sets 2D point
    ///
    /// # Input
    ///
    /// `p` -- index of point in [0, npoint-1]
    /// `group` -- group of point
    /// `x`, `y` -- coordinates of point
    pub fn set_point_2d(&mut self, p: usize, group: usize, x: f64, y: f64) -> &mut Self {
        assert_eq!(self.ndim, 2);
        assert!(p < self.npoint);
        self.point_groups[p] = group;
        self.coords[p][0] = x;
        self.coords[p][1] = y;
        self
    }

    /// Sets 3D point
    ///
    /// # Input
    ///
    /// `p` -- index of point in [0, npoint-1]
    /// `group` -- group of point
    /// `x`, `y`, `z` -- coordinates of point
    pub fn set_point_3d(&mut self, p: usize, group: usize, x: f64, y: f64, z: f64) -> &mut Self {
        assert_eq!(self.ndim, 3);
        assert!(p < self.npoint);
        self.point_groups[p] = group;
        self.coords[p][0] = x;
        self.coords[p][1] = y;
        self.coords[p][2] = z;
        self
    }

    /// Sets the coordinates of all 2D or 3D points
    ///
    /// # Input
    ///
    /// `coords` -- (npoint, ndim) matrix with all coordinates
    pub fn set_coords<'a, T, U>(&mut self, coords: &'a T) -> &mut Self
    where
        T: AsArray2D<'a, U>,
        U: 'a + Into<f64>,
    {
        let (nrow, ncol) = coords.size();
        assert_eq!(nrow, self.npoint);
        assert_eq!(ncol, self.ndim);
        for i in 0..nrow {
            for j in 0..ncol {
                self.coords[i][j] = coords.at(i, j).into();
            }
        }
        self
    }

    /// Sets the group of a point
    ///
    /// # Input
    ///
    /// `p` -- index of point in [0, npoint-1]
    pub fn set_point_group(&mut self, p: usize, group: usize) -> &mut Self {
        assert!(p < self.npoint);
        self.point_groups[p] = group;
        self
    }

    /// Sets the group of an edge
    ///
    /// # Input
    ///
    /// `e` -- index of edge in [0, nedge-1]
    pub fn set_edge_group(&mut self, e: usize, group: usize) -> &mut Self {
        assert!(e < self.nedge);
        self.edge_groups[e] = group;
        self
    }

    /// Sets the group of a face
    ///
    /// # Input
    ///
    /// `f` -- index of edge in [0, nface-1]
    pub fn set_face_group(&mut self, f: usize, group: usize) -> &mut Self {
        assert!(f < self.nface);
        self.face_groups[f] = group;
        self
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

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

    #[test]
    fn new_works() {
        let block = Block::new(1, BlockKind::Qua4);
        assert_eq!(block.group, 1);
        assert_eq!(block.ndim, 2);
        assert_eq!(block.npoint, 4);
        assert_eq!(block.nedge, 4);
        assert_eq!(block.nface, 0);
    }

    #[test]
    fn set_point_2d_works() {
        let mut block = Block::new(1, BlockKind::Qua4);
        block
            .set_point_2d(0, 1, 0.0, 0.0)
            .set_point_2d(1, 2, 2.0, 0.0)
            .set_point_2d(2, 3, 2.0, 2.0)
            .set_point_2d(3, 4, 0.0, 2.0);
        assert_eq!(
            format!("{}", block.coords),
            "┌     ┐\n\
             │ 0 0 │\n\
             │ 2 0 │\n\
             │ 2 2 │\n\
             │ 0 2 │\n\
             └     ┘"
        );
    }

    #[test]
    fn set_point_3d_works() {
        let mut block = Block::new(1, BlockKind::Hex8);
        block
            .set_point_3d(0, 1, 0.0, 0.0, 0.0)
            .set_point_3d(1, 2, 2.0, 0.0, 0.0)
            .set_point_3d(2, 3, 2.0, 2.0, 0.0)
            .set_point_3d(3, 4, 0.0, 2.0, 0.0)
            .set_point_3d(4, 1, 0.0, 0.0, 2.0)
            .set_point_3d(5, 2, 2.0, 0.0, 2.0)
            .set_point_3d(6, 3, 2.0, 2.0, 2.0)
            .set_point_3d(7, 4, 0.0, 2.0, 2.0);
        assert_eq!(
            format!("{}", block.coords),
            "┌       ┐\n\
             │ 0 0 0 │\n\
             │ 2 0 0 │\n\
             │ 2 2 0 │\n\
             │ 0 2 0 │\n\
             │ 0 0 2 │\n\
             │ 2 0 2 │\n\
             │ 2 2 2 │\n\
             │ 0 2 2 │\n\
             └       ┘"
        );
    }

    #[test]
    fn set_coords_works() {
        let mut qua4 = Block::new(1, BlockKind::Qua4);
        #[rustfmt::skip]
        qua4.set_coords(&[
            [0.0, 0.0],
            [2.0, 0.0],
            [2.0, 2.0],
            [0.0, 2.0],
        ]);
        assert_eq!(
            format!("{}", qua4.coords),
            "┌     ┐\n\
             │ 0 0 │\n\
             │ 2 0 │\n\
             │ 2 2 │\n\
             │ 0 2 │\n\
             └     ┘"
        );

        let mut hex8 = Block::new(1, BlockKind::Hex8);
        #[rustfmt::skip]
        hex8.set_coords(&[
            [0.0, 0.0, 0.0],
            [2.0, 0.0, 0.0],
            [2.0, 2.0, 0.0],
            [0.0, 2.0, 0.0],
            [0.0, 0.0, 2.0],
            [2.0, 0.0, 2.0],
            [2.0, 2.0, 2.0],
            [0.0, 2.0, 2.0],
        ]);
        assert_eq!(
            format!("{}", hex8.coords),
            "┌       ┐\n\
             │ 0 0 0 │\n\
             │ 2 0 0 │\n\
             │ 2 2 0 │\n\
             │ 0 2 0 │\n\
             │ 0 0 2 │\n\
             │ 2 0 2 │\n\
             │ 2 2 2 │\n\
             │ 0 2 2 │\n\
             └       ┘"
        );
    }

    #[test]
    fn set_point_group_works() {
        let mut block = Block::new(1, BlockKind::Qua4);
        block.set_point_group(0, 111);
        assert_eq!(block.point_groups, &[111, 0, 0, 0]);
    }

    #[test]
    fn set_edge_group_works() {
        let mut block = Block::new(1, BlockKind::Qua4);
        block.set_edge_group(0, 111);
        assert_eq!(block.edge_groups, &[111, 0, 0, 0]);
    }

    #[test]
    fn set_face_group_works() {
        let mut block = Block::new(1, BlockKind::Hex8);
        block.set_face_group(0, 111);
        assert_eq!(block.face_groups, &[111, 0, 0, 0, 0, 0]);
    }
}
<commit_msg>Improve doc<commit_after>use crate::AsArray2D;
use russell_lab::Matrix;
use std::collections::HashMap;

pub enum BlockKind {
    Qua4,
    Qua8,
    Hex8,
    Hex20,
}

pub enum Constraint {}

pub struct ConstraintData {
    kind: Constraint,
}

/// Defines a polygon on polyhedron that can be split into smaller shapes
///
/// # Block shape
///
/// ## 2D: Local indices of points and edges
///
/// ```text
///             Points (4 or 8)                        Edges (4)
///   y
///   |        3      6      2                             2
///   +--x      @-----@-----@                        +-----------+
///             |           |                        |           |
///             |           |                        |           |
///           7 @           @ 5                     3|           |1
///             |           |                        |           |
///             |           |                        |           |
///             @-----@-----@                        +-----------+
///            0      4      1                             0
/// ```
///
/// ## 3D: Local indices of points, edges, and faces
///
/// ```text
///                   Points (8 or 20)                   Faces (6)
///    z
///    |           4        15        7
///   ,+--y         @-------@--------@                 +----------------+
/// x'            ,'|              ,'|               ,'|              ,'|
///          12 @'  |         14 ,'  |             ,'  |  ___       ,'  |
///           ,'    |16        ,@    |19         ,'    |,'5,'  [0],'    |
///     5   ,'      @      6 ,'      @         ,'      |~~~     ,'      |
///       @'=======@=======@'        |       +'===============+'  ,'|   |
///       |      13 |      |         |       |   ,'|   |      |   |3|   |
///       |         |      |  11     |       |   |2|   |      |   |,'   |
///    17 |       0 @- - - | @- - - -@       |   |,'   +- - - | +- - - -+
///       @       ,'       @       ,' 3      |       ,'       |       ,'
///       |   8 @'      18 |     ,'          |     ,' [1]  ___|     ,'
///       |   ,'           |   ,@ 10         |   ,'      ,'4,'|   ,'
///       | ,'             | ,'              | ,'        ~~~  | ,'
///       @-------@--------@'                +----------------+'
///     1         9         2
///
///
///                  Edges (12)
///
///                         7
///                 +----------------+
///               ,'|              ,'|
///           4 ,'  |8          6,'  |
///           ,'    |          ,'    |
///         ,'      |   5    ,'      |11
///       +'===============+'        |
///       |         |      |         |
///       |         |      |  3      |
///       |         .- - - | -  - - -+
///      9|       ,'       |       ,'
///       |    0,'         |10   ,'
///       |   ,'           |   ,' 2
///       | ,'             | ,'
///       +----------------+'
///               1
///
pub struct Block {
    group: usize,             // group of all elements in this block
    ndim: usize,              // space dimension
    npoint: usize,            // number of points (4, 8, or 20)
    nedge: usize,             // number of edges (4 or 12)
    nface: usize,             // number of faces (4 or 6)
    coords: Matrix,           // coordinates of points (npoint, ndim)
    point_groups: Vec<usize>, // point groups (npoint)
    edge_groups: Vec<usize>,  // edge groups (nedge)
    face_groups: Vec<usize>,  // face groups (nface)
    num_div: Vec<usize>,      // number of divisions along each dim (ndim)
    weights: Vec<Vec<f64>>,   // weights along each dimension (ndim, nweight)
    sum_weights: Vec<f64>,    // sum of weights along each dimension (ndim)

    // maps side to constraint (num_sides)
    constraints: HashMap<Constraint, ConstraintData>,
}

impl Block {
    /// Creates a new Block with default options
    pub fn new(group: usize, kind: BlockKind) -> Self {
        let (ndim, npoint, nedge, nface) = match kind {
            BlockKind::Qua4 => (2, 4, 4, 0),
            BlockKind::Qua8 => (2, 8, 4, 0),
            BlockKind::Hex8 => (3, 8, 12, 6),
            BlockKind::Hex20 => (3, 20, 12, 6),
        };
        Block {
            group,
            ndim,
            npoint,
            nedge,
            nface,
            coords: Matrix::new(npoint, ndim),
            point_groups: vec![0; npoint],
            edge_groups: vec![0; nedge],
            face_groups: vec![0; nface],
            num_div: vec![0; ndim],
            weights: vec![Vec::new(); ndim],
            sum_weights: vec![0.0; ndim],
            constraints: HashMap::new(),
        }
    }

    /// Sets 2D point
    ///
    /// # Input
    ///
    /// * `p` -- index of point in [0, npoint-1]
    /// * `group` -- group of point
    /// * `x`, `y` -- coordinates of point
    pub fn set_point_2d(&mut self, p: usize, group: usize, x: f64, y: f64) -> &mut Self {
        assert_eq!(self.ndim, 2);
        assert!(p < self.npoint);
        self.point_groups[p] = group;
        self.coords[p][0] = x;
        self.coords[p][1] = y;
        self
    }

    /// Sets 3D point
    ///
    /// # Input
    ///
    /// * `p` -- index of point in [0, npoint-1]
    /// * `group` -- group of point
    /// * `x`, `y`, `z` -- coordinates of point
    pub fn set_point_3d(&mut self, p: usize, group: usize, x: f64, y: f64, z: f64) -> &mut Self {
        assert_eq!(self.ndim, 3);
        assert!(p < self.npoint);
        self.point_groups[p] = group;
        self.coords[p][0] = x;
        self.coords[p][1] = y;
        self.coords[p][2] = z;
        self
    }

    /// Sets the coordinates of all 2D or 3D points
    ///
    /// # Input
    ///
    /// * `coords` -- (npoint, ndim) matrix with all coordinates
    pub fn set_coords<'a, T, U>(&mut self, coords: &'a T) -> &mut Self
    where
        T: AsArray2D<'a, U>,
        U: 'a + Into<f64>,
    {
        let (nrow, ncol) = coords.size();
        assert_eq!(nrow, self.npoint);
        assert_eq!(ncol, self.ndim);
        for i in 0..nrow {
            for j in 0..ncol {
                self.coords[i][j] = coords.at(i, j).into();
            }
        }
        self
    }

    /// Sets the group of a point
    ///
    /// # Input
    ///
    /// * `p` -- index of point in [0, npoint-1]
    pub fn set_point_group(&mut self, p: usize, group: usize) -> &mut Self {
        assert!(p < self.npoint);
        self.point_groups[p] = group;
        self
    }

    /// Sets the group of an edge
    ///
    /// # Input
    ///
    /// * `e` -- index of edge in [0, nedge-1]
    pub fn set_edge_group(&mut self, e: usize, group: usize) -> &mut Self {
        assert!(e < self.nedge);
        self.edge_groups[e] = group;
        self
    }

    /// Sets the group of a face
    ///
    /// # Input
    ///
    /// * `f` -- index of edge in [0, nface-1]
    pub fn set_face_group(&mut self, f: usize, group: usize) -> &mut Self {
        assert!(f < self.nface);
        self.face_groups[f] = group;
        self
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

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

    #[test]
    fn new_works() {
        let block = Block::new(1, BlockKind::Qua4);
        assert_eq!(block.group, 1);
        assert_eq!(block.ndim, 2);
        assert_eq!(block.npoint, 4);
        assert_eq!(block.nedge, 4);
        assert_eq!(block.nface, 0);
    }

    #[test]
    fn set_point_2d_works() {
        let mut block = Block::new(1, BlockKind::Qua4);
        block
            .set_point_2d(0, 1, 0.0, 0.0)
            .set_point_2d(1, 2, 2.0, 0.0)
            .set_point_2d(2, 3, 2.0, 2.0)
            .set_point_2d(3, 4, 0.0, 2.0);
        assert_eq!(
            format!("{}", block.coords),
            "┌     ┐\n\
             │ 0 0 │\n\
             │ 2 0 │\n\
             │ 2 2 │\n\
             │ 0 2 │\n\
             └     ┘"
        );
    }

    #[test]
    fn set_point_3d_works() {
        let mut block = Block::new(1, BlockKind::Hex8);
        block
            .set_point_3d(0, 1, 0.0, 0.0, 0.0)
            .set_point_3d(1, 2, 2.0, 0.0, 0.0)
            .set_point_3d(2, 3, 2.0, 2.0, 0.0)
            .set_point_3d(3, 4, 0.0, 2.0, 0.0)
            .set_point_3d(4, 1, 0.0, 0.0, 2.0)
            .set_point_3d(5, 2, 2.0, 0.0, 2.0)
            .set_point_3d(6, 3, 2.0, 2.0, 2.0)
            .set_point_3d(7, 4, 0.0, 2.0, 2.0);
        assert_eq!(
            format!("{}", block.coords),
            "┌       ┐\n\
             │ 0 0 0 │\n\
             │ 2 0 0 │\n\
             │ 2 2 0 │\n\
             │ 0 2 0 │\n\
             │ 0 0 2 │\n\
             │ 2 0 2 │\n\
             │ 2 2 2 │\n\
             │ 0 2 2 │\n\
             └       ┘"
        );
    }

    #[test]
    fn set_coords_works() {
        let mut qua4 = Block::new(1, BlockKind::Qua4);
        #[rustfmt::skip]
        qua4.set_coords(&[
            [0.0, 0.0],
            [2.0, 0.0],
            [2.0, 2.0],
            [0.0, 2.0],
        ]);
        assert_eq!(
            format!("{}", qua4.coords),
            "┌     ┐\n\
             │ 0 0 │\n\
             │ 2 0 │\n\
             │ 2 2 │\n\
             │ 0 2 │\n\
             └     ┘"
        );

        let mut hex8 = Block::new(1, BlockKind::Hex8);
        #[rustfmt::skip]
        hex8.set_coords(&[
            [0.0, 0.0, 0.0],
            [2.0, 0.0, 0.0],
            [2.0, 2.0, 0.0],
            [0.0, 2.0, 0.0],
            [0.0, 0.0, 2.0],
            [2.0, 0.0, 2.0],
            [2.0, 2.0, 2.0],
            [0.0, 2.0, 2.0],
        ]);
        assert_eq!(
            format!("{}", hex8.coords),
            "┌       ┐\n\
             │ 0 0 0 │\n\
             │ 2 0 0 │\n\
             │ 2 2 0 │\n\
             │ 0 2 0 │\n\
             │ 0 0 2 │\n\
             │ 2 0 2 │\n\
             │ 2 2 2 │\n\
             │ 0 2 2 │\n\
             └       ┘"
        );
    }

    #[test]
    fn set_point_group_works() {
        let mut block = Block::new(1, BlockKind::Qua4);
        block.set_point_group(0, 111);
        assert_eq!(block.point_groups, &[111, 0, 0, 0]);
    }

    #[test]
    fn set_edge_group_works() {
        let mut block = Block::new(1, BlockKind::Qua4);
        block.set_edge_group(0, 111);
        assert_eq!(block.edge_groups, &[111, 0, 0, 0]);
    }

    #[test]
    fn set_face_group_works() {
        let mut block = Block::new(1, BlockKind::Hex8);
        block.set_face_group(0, 111);
        assert_eq!(block.face_groups, &[111, 0, 0, 0, 0, 0]);
    }
}
<|endoftext|>


<commit_before>extern crate nom;

use std::str;
use nom::*;

fn is_uchar(c: char) -> bool {
    c == '\\' || c == 'u' || is_hex_digit(c as u8)
}

fn is_echar(c: char) -> bool {
    c == '\\' || c == 't' || c == 'b' || c == 'n' || c == 'r' || c == 'f' || c == '"' || c == '\''
}
/* [26] UCHAR */
named!(uchar<&str, &str>, alt!(
    verify!(
        take_s!(6),
        |val:&str| {
            let len = val.char_indices().count();
            val.char_indices().all(|(idx, c)| {
                if idx == 0 {
                    c == '\\'
                } else if idx == 1 {
                    c == 'u'
                } else {
                    is_hex_digit(c as u8)
                }
            })
        }
    ) |
    verify!(
        take_s!(10),
        |val:&str| {
            let len = val.char_indices().count();
            val.char_indices().all(|(idx, c)| {
                if idx == 0 {
                    c == '\\'
                } else if idx == 1 {
                    c == 'U'
                } else {
                    is_hex_digit(c as u8)
                }
            })
        }
    )
));

/* [159s] ECHAR */
named!(echar<&str, &str>, verify!(
    take_s!(2),
    |val:&str| {
        let len = val.char_indices().count();
        val.char_indices().all(|(idx, c)| {
            if idx == 0 {
                c == '\\'
            } else {
                is_echar(c)
            }
        })
    }
));

named!(string_literal_long_quote_body<&str, Vec<&str>>, many0!(
    alt!(tag!("\"") | echar | uchar)
));
/* [25] STRING_LITERAL_LONG_QUOTE */
named!(string_literal_long_quote<&str, &str>, delimited!(
    tag!("\"\"\""),
    verify!(
        take_until_s!("\"\"\""),
        |val: &str| {
            match string_literal_long_quote_body(val) {
                IResult::Error(_) => false,
                _                 => true,
            }
        }
    ),
    tag!("\"\"\"")
));

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

    #[test]
    fn uchar_test() {
        assert_eq!(uchar("\\u02FFa")    , IResult::Done("a", "\\u02FF")    );
        assert_eq!(uchar("\\U02FFAABBa"), IResult::Done("a", "\\U02FFAABB"));
    }

    #[test]
    fn echar_test() {
        assert_eq!(echar("\\trest") , IResult::Done("rest", "\\t") );
        assert_eq!(echar("\\brest") , IResult::Done("rest", "\\b") );
        assert_eq!(echar("\\nrest") , IResult::Done("rest", "\\n") );
        assert_eq!(echar("\\rrest") , IResult::Done("rest", "\\r") );
        assert_eq!(echar("\\frest") , IResult::Done("rest", "\\f") );
        assert_eq!(echar("\\\"rest"), IResult::Done("rest", "\\\""));
        assert_eq!(echar("\\'rest") , IResult::Done("rest", "\\'") );
        assert_eq!(echar("\\\\rest"), IResult::Done("rest", "\\\\"));
    }

    #[test]
    fn string_literal_long_quote_test() {
        let input    = "\"\"\"\\t\\u02FF\"\"\"rest";
        let expected = "\\t\\u02FF";
        assert_eq!(string_literal_long_quote(input), IResult::Done("rest", expected));
    }
}
<commit_msg>Implement STRING_LITERAL_LONG_SINGLE_QUOTE<commit_after>extern crate nom;

use std::str;
use nom::*;

fn is_uchar(c: char) -> bool {
    c == '\\' || c == 'u' || is_hex_digit(c as u8)
}

fn is_echar(c: char) -> bool {
    c == '\\' || c == 't' || c == 'b' || c == 'n' || c == 'r' || c == 'f' || c == '"' || c == '\''
}
/* [26] UCHAR */
named!(uchar<&str, &str>, alt!(
    verify!(
        take_s!(6),
        |val:&str| {
            let len = val.char_indices().count();
            val.char_indices().all(|(idx, c)| {
                if idx == 0 {
                    c == '\\'
                } else if idx == 1 {
                    c == 'u'
                } else {
                    is_hex_digit(c as u8)
                }
            })
        }
    ) |
    verify!(
        take_s!(10),
        |val:&str| {
            let len = val.char_indices().count();
            val.char_indices().all(|(idx, c)| {
                if idx == 0 {
                    c == '\\'
                } else if idx == 1 {
                    c == 'U'
                } else {
                    is_hex_digit(c as u8)
                }
            })
        }
    )
));

/* [159s] ECHAR */
named!(echar<&str, &str>, verify!(
    take_s!(2),
    |val:&str| {
        let len = val.char_indices().count();
        val.char_indices().all(|(idx, c)| {
            if idx == 0 {
                c == '\\'
            } else {
                is_echar(c)
            }
        })
    }
));

named!(string_literal_long_single_quote_body<&str, Vec<&str>>, many0!(
    alt!(tag!("'") | echar | uchar)
));
/* [24] STRING_LITERAL_LONG_SINGLE_QUOTE */
named!(string_literal_long_single_quote<&str, &str>, delimited!(
    tag!("'''"),
    verify!(
        take_until_s!("'''"),
        |val: &str| {
            match string_literal_long_single_quote_body(val) {
                IResult::Error(_) => false,
                _                 => true,
            }
        }
    ),
    tag!("'''")
));

named!(string_literal_long_quote_body<&str, Vec<&str>>, many0!(
    alt!(tag!("\"") | echar | uchar)
));
/* [25] STRING_LITERAL_LONG_QUOTE */
named!(string_literal_long_quote<&str, &str>, delimited!(
    tag!("\"\"\""),
    verify!(
        take_until_s!("\"\"\""),
        |val: &str| {
            match string_literal_long_quote_body(val) {
                IResult::Error(_) => false,
                _                 => true,
            }
        }
    ),
    tag!("\"\"\"")
));

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

    #[test]
    fn uchar_test() {
        assert_eq!(uchar("\\u02FFa")    , IResult::Done("a", "\\u02FF")    );
        assert_eq!(uchar("\\U02FFAABBa"), IResult::Done("a", "\\U02FFAABB"));
    }

    #[test]
    fn echar_test() {
        assert_eq!(echar("\\trest") , IResult::Done("rest", "\\t") );
        assert_eq!(echar("\\brest") , IResult::Done("rest", "\\b") );
        assert_eq!(echar("\\nrest") , IResult::Done("rest", "\\n") );
        assert_eq!(echar("\\rrest") , IResult::Done("rest", "\\r") );
        assert_eq!(echar("\\frest") , IResult::Done("rest", "\\f") );
        assert_eq!(echar("\\\"rest"), IResult::Done("rest", "\\\""));
        assert_eq!(echar("\\'rest") , IResult::Done("rest", "\\'") );
        assert_eq!(echar("\\\\rest"), IResult::Done("rest", "\\\\"));
    }

    #[test]
    fn string_literal_long_single_quote_test() {
        let input    = "'''\\t\\u02FF'''rest";
        let expected = "\\t\\u02FF";
        assert_eq!(string_literal_long_single_quote(input), IResult::Done("rest", expected));
    }

    #[test]
    fn string_literal_long_quote_test() {
        let input    = "\"\"\"\\t\\u02FF\"\"\"rest";
        let expected = "\\t\\u02FF";
        assert_eq!(string_literal_long_quote(input), IResult::Done("rest", expected));
    }
}
<|endoftext|>


<commit_before>use cairo::{self, MatrixTrait};
use cairo_sys;
use glib::translate::*;
use libc;
use std::cell::Cell;

use attributes::Attribute;
use coord_units::CoordUnits;
use draw;
use drawing_ctx::{self, RsvgDrawingCtx};
use handle::RsvgHandle;
use length::{LengthDir, RsvgLength};
use node::{boxed_node_new, NodeResult, NodeTrait, NodeType, RsvgCNodeImpl, RsvgNode};
use parsers::{parse, Parse};
use property_bag::PropertyBag;
use state::Opacity;

coord_units!(MaskUnits, CoordUnits::ObjectBoundingBox);
coord_units!(MaskContentUnits, CoordUnits::UserSpaceOnUse);

// remove this once cairo-rs has this mask_surface()
fn cairo_mask_surface(cr: &cairo::Context, surface: &cairo::Surface, x: f64, y: f64) {
    unsafe {
        let raw_cr = cr.to_glib_none().0;

        cairo_sys::cairo_mask_surface(raw_cr, surface.to_raw_none(), x, y);
    }
}

pub struct NodeMask {
    x: Cell<RsvgLength>,
    y: Cell<RsvgLength>,
    width: Cell<RsvgLength>,
    height: Cell<RsvgLength>,

    units: Cell<MaskUnits>,
    content_units: Cell<MaskContentUnits>,
}

impl NodeMask {
    fn new() -> NodeMask {
        NodeMask {
            x: Cell::new(NodeMask::get_default_pos(LengthDir::Horizontal)),
            y: Cell::new(NodeMask::get_default_pos(LengthDir::Vertical)),

            width: Cell::new(NodeMask::get_default_size(LengthDir::Horizontal)),
            height: Cell::new(NodeMask::get_default_size(LengthDir::Vertical)),

            units: Cell::new(MaskUnits::default()),
            content_units: Cell::new(MaskContentUnits::default()),
        }
    }

    fn get_default_pos(dir: LengthDir) -> RsvgLength {
        RsvgLength::parse("-10%", dir).unwrap()
    }

    fn get_default_size(dir: LengthDir) -> RsvgLength {
        RsvgLength::parse("120%", dir).unwrap()
    }

    pub fn generate_cairo_mask(&self, node: &RsvgNode, draw_ctx: *mut RsvgDrawingCtx) {
        let values = &node.get_computed_values();

        let width = drawing_ctx::get_width(draw_ctx) as i32;
        let height = drawing_ctx::get_height(draw_ctx) as i32;

        let mut surface = match cairo::ImageSurface::create(cairo::Format::ARgb32, width, height) {
            Ok(surface) => surface,
            Err(_) => return,
        };

        let mask_units = CoordUnits::from(self.units.get());
        let content_units = CoordUnits::from(self.content_units.get());

        if mask_units == CoordUnits::ObjectBoundingBox {
            drawing_ctx::push_view_box(draw_ctx, 1.0, 1.0);
        }

        let x = self.x.get().normalize(draw_ctx);
        let y = self.y.get().normalize(draw_ctx);
        let w = self.width.get().normalize(draw_ctx);
        let h = self.height.get().normalize(draw_ctx);

        if mask_units == CoordUnits::ObjectBoundingBox {
            drawing_ctx::pop_view_box(draw_ctx);
        }

        // Use a scope because mask_cr needs to release the
        // reference to the surface before we access the pixels
        {
            let save_cr = drawing_ctx::get_cairo_context(draw_ctx);
            let mask_cr = cairo::Context::new(&surface);
            drawing_ctx::set_cairo_context(draw_ctx, &mask_cr);

            if mask_units == CoordUnits::ObjectBoundingBox {
                let bbox = drawing_ctx::get_bbox(draw_ctx);
                let rect = bbox.rect.unwrap();

                draw::add_clipping_rect(
                    draw_ctx,
                    &values.affine,
                    x * rect.width + rect.x,
                    y * rect.height + rect.y,
                    w * rect.width,
                    h * rect.height,
                );
            } else {
                draw::add_clipping_rect(draw_ctx, &values.affine, x, y, w, h);
            }

            // Horribly dirty hack to have the bbox premultiplied to everything
            let node_state = node.get_state_mut();
            let affinesave = node_state.affine;
            if content_units == CoordUnits::ObjectBoundingBox {
                let bbox = drawing_ctx::get_bbox(draw_ctx);
                let rect = bbox.rect.unwrap();
                let mut bbtransform =
                    cairo::Matrix::new(rect.width, 0.0, 0.0, rect.height, rect.x, rect.y);

                bbtransform = cairo::Matrix::multiply(&bbtransform, &affinesave);
                node_state.affine = bbtransform;
                drawing_ctx::push_view_box(draw_ctx, 1.0, 1.0);
            }

            drawing_ctx::state_push(draw_ctx);
            node.draw_children(draw_ctx, 0, false);
            drawing_ctx::state_pop(draw_ctx);

            if content_units == CoordUnits::ObjectBoundingBox {
                drawing_ctx::pop_view_box(draw_ctx);
                node_state.affine = affinesave;
            }

            drawing_ctx::set_cairo_context(draw_ctx, &save_cr);
        }

        {
            let rowstride = surface.get_stride() as usize;
            let mut pixels = surface.get_data().unwrap();
            let opacity = {
                let Opacity(o) = values.opacity;
                u8::from(o)
            };

            for row in pixels.chunks_mut(rowstride) {
                for p in row[..width as usize].chunks_mut(4) {
                    //  Assuming, the pixel is linear RGB (not sRGB)
                    //  y = luminance
                    //  Y = 0.2126 R + 0.7152 G + 0.0722 B
                    //  1.0 opacity = 255
                    //
                    //  When Y = 1.0, pixel for mask should be 0xFFFFFFFF
                    //    (you get 1.0 luminance from 255 from R, G and B)
                    //
                    // r_mult = 0xFFFFFFFF / (255.0 * 255.0) * .2126 = 14042.45  ~= 14042
                    // g_mult = 0xFFFFFFFF / (255.0 * 255.0) * .7152 = 47239.69  ~= 47240
                    // b_mult = 0xFFFFFFFF / (255.0 * 255.0) * .0722 =  4768.88  ~= 4769
                    //
                    // This allows for the following expected behaviour:
                    //    (we only care about the most sig byte)
                    // if pixel = 0x00FFFFFF, pixel' = 0xFF......
                    // if pixel = 0x00020202, pixel' = 0x02......
                    // if pixel = 0x00000000, pixel' = 0x00......
                    let (r, g, b, o) = (p[1] as u32, p[2] as u32, p[3] as u32, opacity as u32);
                    p[0] = (((r * 14042 + g * 47240 + b * 4769) * o) >> 24) as u8;
                }
            }
        }

        let cr = drawing_ctx::get_cairo_context(draw_ctx);
        cr.identity_matrix();

        let (xofs, yofs) = drawing_ctx::get_offset(draw_ctx);
        cairo_mask_surface(&cr, &surface, xofs, yofs);
    }
}

impl NodeTrait for NodeMask {
    fn set_atts(&self, _: &RsvgNode, _: *const RsvgHandle, pbag: &PropertyBag) -> NodeResult {
        for (_key, attr, value) in pbag.iter() {
            match attr {
                Attribute::X => self.x.set(parse("x", value, LengthDir::Horizontal, None)?),
                Attribute::Y => self.y.set(parse("y", value, LengthDir::Vertical, None)?),
                Attribute::Width => self.width.set(parse(
                    "width",
                    value,
                    LengthDir::Horizontal,
                    Some(RsvgLength::check_nonnegative),
                )?),
                Attribute::Height => self.height.set(parse(
                    "height",
                    value,
                    LengthDir::Vertical,
                    Some(RsvgLength::check_nonnegative),
                )?),

                Attribute::MaskUnits => self.units.set(parse("maskUnits", value, (), None)?),

                Attribute::MaskContentUnits => {
                    self.content_units
                        .set(parse("maskContentUnits", value, (), None)?)
                }

                _ => (),
            }
        }

        Ok(())
    }

    fn draw(&self, _: &RsvgNode, _: *mut RsvgDrawingCtx, _: i32, _: bool) {
        // nothing; masks are handled specially
    }

    fn get_c_impl(&self) -> *const RsvgCNodeImpl {
        unreachable!();
    }
}

#[no_mangle]
pub extern "C" fn rsvg_node_mask_new(
    _: *const libc::c_char,
    raw_parent: *const RsvgNode,
) -> *const RsvgNode {
    boxed_node_new(NodeType::Mask, raw_parent, Box::new(NodeMask::new()))
}
<commit_msg>mask.rs: Iterate all the pixels in each row<commit_after>use cairo::{self, MatrixTrait};
use cairo_sys;
use glib::translate::*;
use libc;
use std::cell::Cell;

use attributes::Attribute;
use coord_units::CoordUnits;
use draw;
use drawing_ctx::{self, RsvgDrawingCtx};
use handle::RsvgHandle;
use length::{LengthDir, RsvgLength};
use node::{boxed_node_new, NodeResult, NodeTrait, NodeType, RsvgCNodeImpl, RsvgNode};
use parsers::{parse, Parse};
use property_bag::PropertyBag;
use state::Opacity;

coord_units!(MaskUnits, CoordUnits::ObjectBoundingBox);
coord_units!(MaskContentUnits, CoordUnits::UserSpaceOnUse);

// remove this once cairo-rs has this mask_surface()
fn cairo_mask_surface(cr: &cairo::Context, surface: &cairo::Surface, x: f64, y: f64) {
    unsafe {
        let raw_cr = cr.to_glib_none().0;

        cairo_sys::cairo_mask_surface(raw_cr, surface.to_raw_none(), x, y);
    }
}

pub struct NodeMask {
    x: Cell<RsvgLength>,
    y: Cell<RsvgLength>,
    width: Cell<RsvgLength>,
    height: Cell<RsvgLength>,

    units: Cell<MaskUnits>,
    content_units: Cell<MaskContentUnits>,
}

impl NodeMask {
    fn new() -> NodeMask {
        NodeMask {
            x: Cell::new(NodeMask::get_default_pos(LengthDir::Horizontal)),
            y: Cell::new(NodeMask::get_default_pos(LengthDir::Vertical)),

            width: Cell::new(NodeMask::get_default_size(LengthDir::Horizontal)),
            height: Cell::new(NodeMask::get_default_size(LengthDir::Vertical)),

            units: Cell::new(MaskUnits::default()),
            content_units: Cell::new(MaskContentUnits::default()),
        }
    }

    fn get_default_pos(dir: LengthDir) -> RsvgLength {
        RsvgLength::parse("-10%", dir).unwrap()
    }

    fn get_default_size(dir: LengthDir) -> RsvgLength {
        RsvgLength::parse("120%", dir).unwrap()
    }

    pub fn generate_cairo_mask(&self, node: &RsvgNode, draw_ctx: *mut RsvgDrawingCtx) {
        let values = &node.get_computed_values();

        let width = drawing_ctx::get_width(draw_ctx) as i32;
        let height = drawing_ctx::get_height(draw_ctx) as i32;

        let mut surface = match cairo::ImageSurface::create(cairo::Format::ARgb32, width, height) {
            Ok(surface) => surface,
            Err(_) => return,
        };

        let mask_units = CoordUnits::from(self.units.get());
        let content_units = CoordUnits::from(self.content_units.get());

        if mask_units == CoordUnits::ObjectBoundingBox {
            drawing_ctx::push_view_box(draw_ctx, 1.0, 1.0);
        }

        let x = self.x.get().normalize(draw_ctx);
        let y = self.y.get().normalize(draw_ctx);
        let w = self.width.get().normalize(draw_ctx);
        let h = self.height.get().normalize(draw_ctx);

        if mask_units == CoordUnits::ObjectBoundingBox {
            drawing_ctx::pop_view_box(draw_ctx);
        }

        // Use a scope because mask_cr needs to release the
        // reference to the surface before we access the pixels
        {
            let save_cr = drawing_ctx::get_cairo_context(draw_ctx);
            let mask_cr = cairo::Context::new(&surface);
            drawing_ctx::set_cairo_context(draw_ctx, &mask_cr);

            if mask_units == CoordUnits::ObjectBoundingBox {
                let bbox = drawing_ctx::get_bbox(draw_ctx);
                let rect = bbox.rect.unwrap();

                draw::add_clipping_rect(
                    draw_ctx,
                    &values.affine,
                    x * rect.width + rect.x,
                    y * rect.height + rect.y,
                    w * rect.width,
                    h * rect.height,
                );
            } else {
                draw::add_clipping_rect(draw_ctx, &values.affine, x, y, w, h);
            }

            // Horribly dirty hack to have the bbox premultiplied to everything
            let node_state = node.get_state_mut();
            let affinesave = node_state.affine;
            if content_units == CoordUnits::ObjectBoundingBox {
                let bbox = drawing_ctx::get_bbox(draw_ctx);
                let rect = bbox.rect.unwrap();
                let mut bbtransform =
                    cairo::Matrix::new(rect.width, 0.0, 0.0, rect.height, rect.x, rect.y);

                bbtransform = cairo::Matrix::multiply(&bbtransform, &affinesave);
                node_state.affine = bbtransform;
                drawing_ctx::push_view_box(draw_ctx, 1.0, 1.0);
            }

            drawing_ctx::state_push(draw_ctx);
            node.draw_children(draw_ctx, 0, false);
            drawing_ctx::state_pop(draw_ctx);

            if content_units == CoordUnits::ObjectBoundingBox {
                drawing_ctx::pop_view_box(draw_ctx);
                node_state.affine = affinesave;
            }

            drawing_ctx::set_cairo_context(draw_ctx, &save_cr);
        }

        {
            let rowstride = surface.get_stride() as usize;
            let mut pixels = surface.get_data().unwrap();
            let opacity = {
                let Opacity(o) = values.opacity;
                u8::from(o)
            };

            for row in pixels.chunks_mut(rowstride) {
                for p in row[..4 * width as usize].chunks_mut(4) {
                    //  Assuming, the pixel is linear RGB (not sRGB)
                    //  y = luminance
                    //  Y = 0.2126 R + 0.7152 G + 0.0722 B
                    //  1.0 opacity = 255
                    //
                    //  When Y = 1.0, pixel for mask should be 0xFFFFFFFF
                    //    (you get 1.0 luminance from 255 from R, G and B)
                    //
                    // r_mult = 0xFFFFFFFF / (255.0 * 255.0) * .2126 = 14042.45  ~= 14042
                    // g_mult = 0xFFFFFFFF / (255.0 * 255.0) * .7152 = 47239.69  ~= 47240
                    // b_mult = 0xFFFFFFFF / (255.0 * 255.0) * .0722 =  4768.88  ~= 4769
                    //
                    // This allows for the following expected behaviour:
                    //    (we only care about the most sig byte)
                    // if pixel = 0x00FFFFFF, pixel' = 0xFF......
                    // if pixel = 0x00020202, pixel' = 0x02......
                    // if pixel = 0x00000000, pixel' = 0x00......
                    let (r, g, b, o) = (p[1] as u32, p[2] as u32, p[3] as u32, opacity as u32);
                    p[0] = (((r * 14042 + g * 47240 + b * 4769) * o) >> 24) as u8;
                }
            }
        }

        let cr = drawing_ctx::get_cairo_context(draw_ctx);
        cr.identity_matrix();

        let (xofs, yofs) = drawing_ctx::get_offset(draw_ctx);
        cairo_mask_surface(&cr, &surface, xofs, yofs);
    }
}

impl NodeTrait for NodeMask {
    fn set_atts(&self, _: &RsvgNode, _: *const RsvgHandle, pbag: &PropertyBag) -> NodeResult {
        for (_key, attr, value) in pbag.iter() {
            match attr {
                Attribute::X => self.x.set(parse("x", value, LengthDir::Horizontal, None)?),
                Attribute::Y => self.y.set(parse("y", value, LengthDir::Vertical, None)?),
                Attribute::Width => self.width.set(parse(
                    "width",
                    value,
                    LengthDir::Horizontal,
                    Some(RsvgLength::check_nonnegative),
                )?),
                Attribute::Height => self.height.set(parse(
                    "height",
                    value,
                    LengthDir::Vertical,
                    Some(RsvgLength::check_nonnegative),
                )?),

                Attribute::MaskUnits => self.units.set(parse("maskUnits", value, (), None)?),

                Attribute::MaskContentUnits => {
                    self.content_units
                        .set(parse("maskContentUnits", value, (), None)?)
                }

                _ => (),
            }
        }

        Ok(())
    }

    fn draw(&self, _: &RsvgNode, _: *mut RsvgDrawingCtx, _: i32, _: bool) {
        // nothing; masks are handled specially
    }

    fn get_c_impl(&self) -> *const RsvgCNodeImpl {
        unreachable!();
    }
}

#[no_mangle]
pub extern "C" fn rsvg_node_mask_new(
    _: *const libc::c_char,
    raw_parent: *const RsvgNode,
) -> *const RsvgNode {
    boxed_node_new(NodeType::Mask, raw_parent, Box::new(NodeMask::new()))
}
<|endoftext|>


<commit_before>mod reader;
mod writer;

use std::time::{Duration, Instant};
use std::{io, mem};
use std::net::SocketAddr;

use url::percent_encoding::percent_encode_byte;
use url::Url;

use {amy, bencode, PEER_ID};
use self::writer::Writer;
use self::reader::{ReadRes, Reader};
use socket::TSocket;
use tracker::{self, dns, Announce, Error, ErrorKind, Response, Result, ResultExt, TrackerResponse};
use util::{AView, UHashMap};

const TIMEOUT_MS: u64 = 5_000;

pub struct Handler {
    reg: amy::Registrar,
    connections: UHashMap<Tracker>,
}

enum Event {
    DNSResolved(dns::QueryResponse),
    Readable,
    Writable,
}

struct Tracker {
    torrent: usize,
    url: AView<Url>,
    last_updated: Instant,
    redirect: bool,
    state: TrackerState,
}

enum TrackerState {
    Error,
    ResolvingDNS {
        sock: TSocket,
        req: Vec<u8>,
        port: u16,
    },
    Writing {
        sock: TSocket,
        writer: Writer,
    },
    Reading {
        sock: TSocket,
        reader: Reader,
    },
    Redirect(String),
    Complete(TrackerResponse),
}

enum HTTPRes {
    None,
    Redirect(String),
    Complete(TrackerResponse),
}

impl TrackerState {
    fn new(sock: TSocket, req: Vec<u8>, port: u16) -> TrackerState {
        TrackerState::ResolvingDNS { sock, req, port }
    }

    fn handle(&mut self, event: Event) -> Result<HTTPRes> {
        let s = mem::replace(self, TrackerState::Error);
        match s.next(event)? {
            TrackerState::Complete(r) => Ok(HTTPRes::Complete(r)),
            TrackerState::Redirect(l) => Ok(HTTPRes::Redirect(l)),
            n => {
                mem::replace(self, n);
                Ok(HTTPRes::None)
            }
        }
    }

    fn next(self, event: Event) -> Result<TrackerState> {
        match (self, event) {
            (
                TrackerState::ResolvingDNS {
                    mut sock,
                    req,
                    port,
                },
                Event::DNSResolved(r),
            ) => {
                let addr = SocketAddr::new(r.res?, port);
                sock.connect(addr).chain_err(|| ErrorKind::IO)?;
                Ok(TrackerState::Writing {
                    sock,
                    writer: Writer::new(req),
                }.next(Event::Writable)?
                    .next(Event::Readable)?)
            }
            (
                TrackerState::Writing {
                    mut sock,
                    mut writer,
                },
                _,
            ) => match writer.writable(&mut sock)? {
                Some(()) => {
                    debug!("Tracker write completed, beginning read");
                    let r = Reader::new();
                    Ok(TrackerState::Reading { sock, reader: r }.next(Event::Readable)?)
                }
                None => Ok(TrackerState::Writing { sock, writer }),
            },
            (
                TrackerState::Reading {
                    mut sock,
                    mut reader,
                },
                _,
            ) => match reader.readable(&mut sock)? {
                ReadRes::Done(data) => {
                    let content = bencode::decode_buf(&data)
                        .chain_err(|| ErrorKind::InvalidResponse("Invalid BEncoded response!"))?;
                    let resp = TrackerResponse::from_bencode(content)?;
                    Ok(TrackerState::Complete(resp))
                }
                ReadRes::Redirect(l) => Ok(TrackerState::Redirect(l)),
                ReadRes::None => Ok(TrackerState::Reading { sock, reader }),
            },
            (s @ TrackerState::ResolvingDNS { .. }, _) => Ok(s),
            _ => bail!("Unknown state transition encountered!"),
        }
    }
}

impl Handler {
    pub fn new(reg: &amy::Registrar) -> io::Result<Handler> {
        Ok(Handler {
            reg: reg.clone(),
            connections: UHashMap::default(),
        })
    }

    pub fn active_requests(&self) -> usize {
        self.connections.len()
    }

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

    pub fn contains(&self, id: usize) -> bool {
        self.connections.contains_key(&id)
    }

    pub fn dns_resolved(&mut self, resp: dns::QueryResponse) -> Option<Response> {
        let id = resp.id;
        debug!("Received a DNS resp for {:?}", id);
        let resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::DNSResolved(resp)) {
                Ok(_) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };
        if resp.is_some() {
            self.connections.remove(&id);
        }
        resp
    }

    pub fn writable(&mut self, id: usize) -> Option<Response> {
        let resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Writable) {
                Ok(_) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };
        if resp.is_some() {
            self.connections.remove(&id);
        }
        resp
    }

    pub fn readable(&mut self, id: usize, dns: &mut dns::Resolver) -> Option<Response> {
        let mut loc = None;
        let mut resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Readable) {
                Ok(HTTPRes::Complete(r)) => {
                    debug!("Announce response received for {:?} succesfully", id);
                    Some(Response::Tracker {
                        tid: trk.torrent,
                        url: trk.url.clone(),
                        resp: Ok(r),
                    })
                }
                Ok(HTTPRes::Redirect(l)) => {
                    loc = Some(l);
                    None
                }
                Ok(HTTPRes::None) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };

        if resp.is_some() {
            self.connections.remove(&id);
        }

        if let Some(l) = loc {
            let trk = self.connections.remove(&id).unwrap();
            // Disallow 2 levels of redirection
            if trk.redirect {
                resp = Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(ErrorKind::InvalidResponse("Too many redirects").into()),
                });
            }
            if let Err(e) = self.try_redirect(&l, trk.torrent, dns) {
                debug!(
                    "Announce response received for {:?}, redirecting!",
                    trk.torrent
                );
                resp = Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                });
            }
        }
        resp
    }

    fn try_redirect(&mut self, url: &str, torrent: usize, dns: &mut dns::Resolver) -> Result<()> {
        let url = Url::parse(url).chain_err(|| ErrorKind::InvalidResponse("Malformed redirect!"))?;
        let mut http_req = Vec::with_capacity(50);
        http_req.extend_from_slice(b"GET ");
        http_req.extend_from_slice(url.path().as_bytes());
        if let Some(q) = url.query() {
            http_req.extend_from_slice(b"?");
            http_req.extend_from_slice(q.as_bytes());
        }

        http_req.extend_from_slice(b" HTTP/1.1\r\n");
        http_req.extend_from_slice(b"Connection: close\r\n");
        http_req.extend_from_slice(b"Host: ");
        let host = url.host_str()
            .ok_or_else(|| Error::from(ErrorKind::InvalidResponse("Malformed redirect!")))?;
        let port = url.port().unwrap_or(80);
        http_req.extend_from_slice(host.as_bytes());
        http_req.extend_from_slice(b"\r\n\r\n");

        let ohost = if url.scheme() == "https" {
            Some(host.to_owned())
        } else {
            None
        };

        // Setup actual connection and start DNS query
        let sock = TSocket::new_v4(ohost).chain_err(|| ErrorKind::IO)?;
        let id = self.reg
            .register(&sock, amy::Event::Both)
            .chain_err(|| ErrorKind::IO)?;
        self.connections.insert(
            id,
            Tracker {
                last_updated: Instant::now(),
                redirect: true,
                torrent,
                url: AView::value(url.clone()),
                state: TrackerState::new(sock, http_req, port),
            },
        );

        debug!("Dispatching redirect DNS req, id {:?}", id);
        dns.new_query(id, host);
        Ok(())
    }

    pub fn tick(&mut self) -> Vec<Response> {
        let mut resps = Vec::new();
        self.connections.retain(|id, trk| {
            if trk.last_updated.elapsed() > Duration::from_millis(TIMEOUT_MS) {
                debug!("Announce {:?} timed out", id);
                resps.push(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(ErrorKind::Timeout.into()),
                });
                false
            } else {
                true
            }
        });
        resps
    }

    pub fn new_announce(&mut self, req: Announce, dns: &mut dns::Resolver) -> Result<()> {
        debug!("Received a new announce req for {:?}", req.url);
        let mut http_req = Vec::with_capacity(50);
        // Encode GET req
        http_req.extend_from_slice(b"GET ");

        // Encode the URL
        http_req.extend_from_slice(req.url.path().as_bytes());
        http_req.extend_from_slice(b"?");
        append_query_pair(&mut http_req, "info_hash", &encode_param(&req.hash));
        append_query_pair(&mut http_req, "peer_id", &encode_param(&PEER_ID[..]));
        append_query_pair(&mut http_req, "uploaded", &req.uploaded.to_string());
        append_query_pair(&mut http_req, "downloaded", &req.downloaded.to_string());
        append_query_pair(&mut http_req, "left", &req.left.to_string());
        append_query_pair(&mut http_req, "compact", "1");
        append_query_pair(&mut http_req, "port", &req.port.to_string());
        if let Some(nw) = req.num_want {
            append_query_pair(&mut http_req, "numwant", &nw.to_string());
        }
        match req.event {
            Some(tracker::Event::Started) => {
                append_query_pair(&mut http_req, "event", "started");
            }
            Some(tracker::Event::Stopped) => {
                append_query_pair(&mut http_req, "event", "stopped");
            }
            Some(tracker::Event::Completed) => {
                append_query_pair(&mut http_req, "event", "completed");
            }
            None => {}
        }

        // Encode HTTP protocol
        http_req.extend_from_slice(b" HTTP/1.1\r\n");
        // Don't keep alive
        http_req.extend_from_slice(b"Connection: close\r\n");
        // Encode host header
        http_req.extend_from_slice(b"Host: ");
        let host = req.url.host_str().ok_or_else(|| {
            Error::from(ErrorKind::InvalidRequest(
                "Tracker announce url has no host!".to_owned(),
            ))
        })?;
        let port =
            req.url
                .port()
                .unwrap_or_else(|| if req.url.scheme() == "https" { 443 } else { 80 });
        http_req.extend_from_slice(host.as_bytes());
        http_req.extend_from_slice(b"\r\n");
        // Encode empty line to terminate request
        http_req.extend_from_slice(b"\r\n");

        let ohost = if req.url.scheme() == "https" {
            Some(host.to_owned())
        } else {
            None
        };

        // Setup actual connection and start DNS query
        let sock = TSocket::new_v4(ohost).chain_err(|| ErrorKind::IO)?;
        let id = self.reg
            .register(&sock, amy::Event::Both)
            .chain_err(|| ErrorKind::IO)?;
        self.connections.insert(
            id,
            Tracker {
                url: req.url.clone(),
                last_updated: Instant::now(),
                torrent: req.id,
                state: TrackerState::new(sock, http_req, port),
                redirect: false,
            },
        );

        debug!("Dispatching DNS req, id {:?}", id);
        dns.new_query(id, host);

        Ok(())
    }
}

fn append_query_pair(s: &mut Vec<u8>, k: &str, v: &str) {
    s.extend_from_slice(k.as_bytes());
    s.extend_from_slice(b"=");
    s.extend_from_slice(v.as_bytes());
    s.extend_from_slice(b"&");
}

fn encode_param(data: &[u8]) -> String {
    let mut resp = String::new();
    for byte in data {
        let c = char::from(*byte);
        if (*byte > 0x20 && *byte < 0x7E) && (c.is_numeric() || c.is_alphabetic() || c == '-') {
            resp.push(c);
        } else {
            resp += percent_encode_byte(*byte);
        }
    }
    resp
}
<commit_msg>Add User-Agent to HTTP tracker requests<commit_after>mod reader;
mod writer;

use std::time::{Duration, Instant};
use std::{io, mem};
use std::net::SocketAddr;

use url::percent_encoding::percent_encode_byte;
use url::Url;

use {amy, bencode, PEER_ID};
use self::writer::Writer;
use self::reader::{ReadRes, Reader};
use socket::TSocket;
use tracker::{self, dns, Announce, Error, ErrorKind, Response, Result, ResultExt, TrackerResponse};
use util::{AView, UHashMap};

const TIMEOUT_MS: u64 = 5_000;

pub struct Handler {
    reg: amy::Registrar,
    connections: UHashMap<Tracker>,
}

enum Event {
    DNSResolved(dns::QueryResponse),
    Readable,
    Writable,
}

struct Tracker {
    torrent: usize,
    url: AView<Url>,
    last_updated: Instant,
    redirect: bool,
    state: TrackerState,
}

enum TrackerState {
    Error,
    ResolvingDNS {
        sock: TSocket,
        req: Vec<u8>,
        port: u16,
    },
    Writing {
        sock: TSocket,
        writer: Writer,
    },
    Reading {
        sock: TSocket,
        reader: Reader,
    },
    Redirect(String),
    Complete(TrackerResponse),
}

enum HTTPRes {
    None,
    Redirect(String),
    Complete(TrackerResponse),
}

impl TrackerState {
    fn new(sock: TSocket, req: Vec<u8>, port: u16) -> TrackerState {
        TrackerState::ResolvingDNS { sock, req, port }
    }

    fn handle(&mut self, event: Event) -> Result<HTTPRes> {
        let s = mem::replace(self, TrackerState::Error);
        match s.next(event)? {
            TrackerState::Complete(r) => Ok(HTTPRes::Complete(r)),
            TrackerState::Redirect(l) => Ok(HTTPRes::Redirect(l)),
            n => {
                mem::replace(self, n);
                Ok(HTTPRes::None)
            }
        }
    }

    fn next(self, event: Event) -> Result<TrackerState> {
        match (self, event) {
            (
                TrackerState::ResolvingDNS {
                    mut sock,
                    req,
                    port,
                },
                Event::DNSResolved(r),
            ) => {
                let addr = SocketAddr::new(r.res?, port);
                sock.connect(addr).chain_err(|| ErrorKind::IO)?;
                Ok(TrackerState::Writing {
                    sock,
                    writer: Writer::new(req),
                }.next(Event::Writable)?
                    .next(Event::Readable)?)
            }
            (
                TrackerState::Writing {
                    mut sock,
                    mut writer,
                },
                _,
            ) => match writer.writable(&mut sock)? {
                Some(()) => {
                    debug!("Tracker write completed, beginning read");
                    let r = Reader::new();
                    Ok(TrackerState::Reading { sock, reader: r }.next(Event::Readable)?)
                }
                None => Ok(TrackerState::Writing { sock, writer }),
            },
            (
                TrackerState::Reading {
                    mut sock,
                    mut reader,
                },
                _,
            ) => match reader.readable(&mut sock)? {
                ReadRes::Done(data) => {
                    let content = bencode::decode_buf(&data)
                        .chain_err(|| ErrorKind::InvalidResponse("Invalid BEncoded response!"))?;
                    let resp = TrackerResponse::from_bencode(content)?;
                    Ok(TrackerState::Complete(resp))
                }
                ReadRes::Redirect(l) => Ok(TrackerState::Redirect(l)),
                ReadRes::None => Ok(TrackerState::Reading { sock, reader }),
            },
            (s @ TrackerState::ResolvingDNS { .. }, _) => Ok(s),
            _ => bail!("Unknown state transition encountered!"),
        }
    }
}

impl Handler {
    pub fn new(reg: &amy::Registrar) -> io::Result<Handler> {
        Ok(Handler {
            reg: reg.clone(),
            connections: UHashMap::default(),
        })
    }

    pub fn active_requests(&self) -> usize {
        self.connections.len()
    }

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

    pub fn contains(&self, id: usize) -> bool {
        self.connections.contains_key(&id)
    }

    pub fn dns_resolved(&mut self, resp: dns::QueryResponse) -> Option<Response> {
        let id = resp.id;
        debug!("Received a DNS resp for {:?}", id);
        let resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::DNSResolved(resp)) {
                Ok(_) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };
        if resp.is_some() {
            self.connections.remove(&id);
        }
        resp
    }

    pub fn writable(&mut self, id: usize) -> Option<Response> {
        let resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Writable) {
                Ok(_) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };
        if resp.is_some() {
            self.connections.remove(&id);
        }
        resp
    }

    pub fn readable(&mut self, id: usize, dns: &mut dns::Resolver) -> Option<Response> {
        let mut loc = None;
        let mut resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Readable) {
                Ok(HTTPRes::Complete(r)) => {
                    debug!("Announce response received for {:?} succesfully", id);
                    Some(Response::Tracker {
                        tid: trk.torrent,
                        url: trk.url.clone(),
                        resp: Ok(r),
                    })
                }
                Ok(HTTPRes::Redirect(l)) => {
                    loc = Some(l);
                    None
                }
                Ok(HTTPRes::None) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };

        if resp.is_some() {
            self.connections.remove(&id);
        }

        if let Some(l) = loc {
            let trk = self.connections.remove(&id).unwrap();
            // Disallow 2 levels of redirection
            if trk.redirect {
                resp = Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(ErrorKind::InvalidResponse("Too many redirects").into()),
                });
            }
            if let Err(e) = self.try_redirect(&l, trk.torrent, dns) {
                debug!(
                    "Announce response received for {:?}, redirecting!",
                    trk.torrent
                );
                resp = Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                });
            }
        }
        resp
    }

    fn try_redirect(&mut self, url: &str, torrent: usize, dns: &mut dns::Resolver) -> Result<()> {
        let url = Url::parse(url).chain_err(|| ErrorKind::InvalidResponse("Malformed redirect!"))?;
        let mut http_req = Vec::with_capacity(50);
        http_req.extend_from_slice(b"GET ");
        http_req.extend_from_slice(url.path().as_bytes());
        if let Some(q) = url.query() {
            http_req.extend_from_slice(b"?");
            http_req.extend_from_slice(q.as_bytes());
        }

        http_req.extend_from_slice(b" HTTP/1.1\r\n");
        let user_agent = format!(
            "User-Agent: {}/{}",
            "synapse",
            env!("CARGO_PKG_VERSION")
        );
        http_req.extend_from_slice(user_agent.as_bytes());
        http_req.extend_from_slice(b"Connection: close\r\n");
        http_req.extend_from_slice(b"Host: ");
        let host = url.host_str()
            .ok_or_else(|| Error::from(ErrorKind::InvalidResponse("Malformed redirect!")))?;
        let port = url.port().unwrap_or(80);
        http_req.extend_from_slice(host.as_bytes());
        http_req.extend_from_slice(b"\r\n\r\n");

        let ohost = if url.scheme() == "https" {
            Some(host.to_owned())
        } else {
            None
        };

        // Setup actual connection and start DNS query
        let sock = TSocket::new_v4(ohost).chain_err(|| ErrorKind::IO)?;
        let id = self.reg
            .register(&sock, amy::Event::Both)
            .chain_err(|| ErrorKind::IO)?;
        self.connections.insert(
            id,
            Tracker {
                last_updated: Instant::now(),
                redirect: true,
                torrent,
                url: AView::value(url.clone()),
                state: TrackerState::new(sock, http_req, port),
            },
        );

        debug!("Dispatching redirect DNS req, id {:?}", id);
        dns.new_query(id, host);
        Ok(())
    }

    pub fn tick(&mut self) -> Vec<Response> {
        let mut resps = Vec::new();
        self.connections.retain(|id, trk| {
            if trk.last_updated.elapsed() > Duration::from_millis(TIMEOUT_MS) {
                debug!("Announce {:?} timed out", id);
                resps.push(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(ErrorKind::Timeout.into()),
                });
                false
            } else {
                true
            }
        });
        resps
    }

    pub fn new_announce(&mut self, req: Announce, dns: &mut dns::Resolver) -> Result<()> {
        debug!("Received a new announce req for {:?}", req.url);
        let mut http_req = Vec::with_capacity(50);
        // Encode GET req
        http_req.extend_from_slice(b"GET ");

        // Encode the URL
        http_req.extend_from_slice(req.url.path().as_bytes());
        http_req.extend_from_slice(b"?");
        append_query_pair(&mut http_req, "info_hash", &encode_param(&req.hash));
        append_query_pair(&mut http_req, "peer_id", &encode_param(&PEER_ID[..]));
        append_query_pair(&mut http_req, "uploaded", &req.uploaded.to_string());
        append_query_pair(&mut http_req, "downloaded", &req.downloaded.to_string());
        append_query_pair(&mut http_req, "left", &req.left.to_string());
        append_query_pair(&mut http_req, "compact", "1");
        append_query_pair(&mut http_req, "port", &req.port.to_string());
        if let Some(nw) = req.num_want {
            append_query_pair(&mut http_req, "numwant", &nw.to_string());
        }
        match req.event {
            Some(tracker::Event::Started) => {
                append_query_pair(&mut http_req, "event", "started");
            }
            Some(tracker::Event::Stopped) => {
                append_query_pair(&mut http_req, "event", "stopped");
            }
            Some(tracker::Event::Completed) => {
                append_query_pair(&mut http_req, "event", "completed");
            }
            None => {}
        }

        // Encode HTTP protocol
        http_req.extend_from_slice(b" HTTP/1.1\r\n");
        // Don't keep alive
        http_req.extend_from_slice(b"Connection: close\r\n");
        // Encode host header
        http_req.extend_from_slice(b"Host: ");
        let host = req.url.host_str().ok_or_else(|| {
            Error::from(ErrorKind::InvalidRequest(
                "Tracker announce url has no host!".to_owned(),
            ))
        })?;
        let port =
            req.url
                .port()
                .unwrap_or_else(|| if req.url.scheme() == "https" { 443 } else { 80 });
        http_req.extend_from_slice(host.as_bytes());
        http_req.extend_from_slice(b"\r\n");
        // Encode empty line to terminate request
        http_req.extend_from_slice(b"\r\n");

        let ohost = if req.url.scheme() == "https" {
            Some(host.to_owned())
        } else {
            None
        };

        // Setup actual connection and start DNS query
        let sock = TSocket::new_v4(ohost).chain_err(|| ErrorKind::IO)?;
        let id = self.reg
            .register(&sock, amy::Event::Both)
            .chain_err(|| ErrorKind::IO)?;
        self.connections.insert(
            id,
            Tracker {
                url: req.url.clone(),
                last_updated: Instant::now(),
                torrent: req.id,
                state: TrackerState::new(sock, http_req, port),
                redirect: false,
            },
        );

        debug!("Dispatching DNS req, id {:?}", id);
        dns.new_query(id, host);

        Ok(())
    }
}

fn append_query_pair(s: &mut Vec<u8>, k: &str, v: &str) {
    s.extend_from_slice(k.as_bytes());
    s.extend_from_slice(b"=");
    s.extend_from_slice(v.as_bytes());
    s.extend_from_slice(b"&");
}

fn encode_param(data: &[u8]) -> String {
    let mut resp = String::new();
    for byte in data {
        let c = char::from(*byte);
        if (*byte > 0x20 && *byte < 0x7E) && (c.is_numeric() || c.is_alphabetic() || c == '-') {
            resp.push(c);
        } else {
            resp += percent_encode_byte(*byte);
        }
    }
    resp
}
<|endoftext|>


<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Maciej Dziardziel <fiedzia@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (ToDO) signalname pids

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use libc::{c_int, pid_t};
use std::io::Error;
use uucore::signals::ALL_SIGNALS;
use uucore::InvalidEncodingHandling;

static ABOUT: &str = "Send signal to processes or list information about signals.";

static EXIT_OK: i32 = 0;
static EXIT_ERR: i32 = 1;

pub mod options {
    pub static PIDS_OR_SIGNALS: &str = "pids_of_signals";
    pub static LIST: &str = "list";
    pub static TABLE: &str = "table";
    pub static TABLE_OLD: &str = "table_old";
    pub static SIGNAL: &str = "signal";
}

#[derive(Clone, Copy)]
pub enum Mode {
    Kill,
    Table,
    List,
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let args = args
        .collect_str(InvalidEncodingHandling::Ignore)
        .accept_any();
    let (args, obs_signal) = handle_obsolete(args);

    let usage = format!("{} [OPTIONS]... PID...", executable!());
    let matches = App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(options::LIST)
                .short("l")
                .long(options::LIST)
                .help("Lists signals")
                .conflicts_with(options::TABLE)
                .conflicts_with(options::TABLE_OLD),
        )
        .arg(
            Arg::with_name(options::TABLE)
                .short("t")
                .long(options::TABLE)
                .help("Lists table of signals"),
        )
        .arg(Arg::with_name(options::TABLE_OLD).short("L").hidden(true))
        .arg(
            Arg::with_name(options::SIGNAL)
                .short("s")
                .long(options::SIGNAL)
                .help("Sends given signal")
                .takes_value(true),
        )
        .arg(
            Arg::with_name(options::PIDS_OR_SIGNALS)
                .hidden(true)
                .multiple(true),
        )
        .get_matches_from(args);

    let mode = if matches.is_present(options::TABLE) || matches.is_present(options::TABLE_OLD) {
        Mode::Table
    } else if matches.is_present(options::LIST) {
        Mode::List
    } else {
        Mode::Kill
    };

    let pids_or_signals: Vec<String> = matches
        .values_of(options::PIDS_OR_SIGNALS)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    match mode {
        Mode::Kill => {
            let sig = match (obs_signal, matches.value_of(options::SIGNAL)) {
                (Some(s), Some(_)) => s, // -s takes precedence
                (Some(s), None) => s,
                (None, Some(s)) => s.to_owned(),
                (None, None) => "TERM".to_owned(),
            };
            return kill(&sig, &pids_or_signals);
        }
        Mode::Table => table(),
        Mode::List => list(pids_or_signals.get(0).cloned()),
    }

    EXIT_OK
}

fn handle_obsolete(mut args: Vec<String>) -> (Vec<String>, Option<String>) {
    let mut i = 0;
    while i < args.len() {
        // this is safe because slice is valid when it is referenced
        let slice = &args[i].clone();
        if slice.starts_with('-') && slice.chars().nth(1).map_or(false, |c| c.is_digit(10)) {
            let val = &slice[1..];
            match val.parse() {
                Ok(num) => {
                    if uucore::signals::is_signal(num) {
                        args.remove(i);
                        return (args, Some(val.to_owned()));
                    }
                }
                Err(_) => break, /* getopts will error out for us */
            }
        }
        i += 1;
    }
    (args, None)
}

fn table() {
    let mut name_width = 0;
    /* Compute the maximum width of a signal name. */
    for s in &ALL_SIGNALS {
        if s.len() > name_width {
            name_width = s.len()
        }
    }

    for (idx, signal) in ALL_SIGNALS.iter().enumerate() {
        print!("{0: >#2} {1: <#8}", idx + 1, signal);
        //TODO: obtain max signal width here

        if (idx + 1) % 7 == 0 {
            println!();
        }
    }
}

fn print_signal(signal_name_or_value: &str) {
    for (value, &signal) in ALL_SIGNALS.iter().enumerate() {
        if signal == signal_name_or_value || (format!("SIG{}", signal)) == signal_name_or_value {
            println!("{}", value);
            exit!(EXIT_OK as i32)
        } else if signal_name_or_value == value.to_string() {
            println!("{}", signal);
            exit!(EXIT_OK as i32)
        }
    }
    crash!(EXIT_ERR, "unknown signal name {}", signal_name_or_value)
}

fn print_signals() {
    let mut pos = 0;
    for (idx, signal) in ALL_SIGNALS.iter().enumerate() {
        pos += signal.len();
        print!("{}", signal);
        if idx > 0 && pos > 73 {
            println!();
            pos = 0;
        } else {
            pos += 1;
            print!(" ");
        }
    }
}

fn list(arg: Option<String>) {
    match arg {
        Some(ref x) => print_signal(x),
        None => print_signals(),
    };
}

fn kill(signalname: &str, pids: &[String]) -> i32 {
    let mut status = 0;
    let optional_signal_value = uucore::signals::signal_by_name_or_value(signalname);
    let signal_value = match optional_signal_value {
        Some(x) => x,
        None => crash!(EXIT_ERR, "unknown signal name {}", signalname),
    };
    for pid in pids {
        match pid.parse::<usize>() {
            Ok(x) => {
                if unsafe { libc::kill(x as pid_t, signal_value as c_int) } != 0 {
                    show_error!("{}", Error::last_os_error());
                    status = 1;
                }
            }
            Err(e) => crash!(EXIT_ERR, "failed to parse argument {}: {}", pid, e),
        };
    }
    status
}
<commit_msg>kill: fix signal table printing<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Maciej Dziardziel <fiedzia@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (ToDO) signalname pids

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use libc::{c_int, pid_t};
use std::io::Error;
use uucore::signals::ALL_SIGNALS;
use uucore::InvalidEncodingHandling;

static ABOUT: &str = "Send signal to processes or list information about signals.";

static EXIT_OK: i32 = 0;
static EXIT_ERR: i32 = 1;

pub mod options {
    pub static PIDS_OR_SIGNALS: &str = "pids_of_signals";
    pub static LIST: &str = "list";
    pub static TABLE: &str = "table";
    pub static TABLE_OLD: &str = "table_old";
    pub static SIGNAL: &str = "signal";
}

#[derive(Clone, Copy)]
pub enum Mode {
    Kill,
    Table,
    List,
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let args = args
        .collect_str(InvalidEncodingHandling::Ignore)
        .accept_any();
    let (args, obs_signal) = handle_obsolete(args);

    let usage = format!("{} [OPTIONS]... PID...", executable!());
    let matches = App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(options::LIST)
                .short("l")
                .long(options::LIST)
                .help("Lists signals")
                .conflicts_with(options::TABLE)
                .conflicts_with(options::TABLE_OLD),
        )
        .arg(
            Arg::with_name(options::TABLE)
                .short("t")
                .long(options::TABLE)
                .help("Lists table of signals"),
        )
        .arg(Arg::with_name(options::TABLE_OLD).short("L").hidden(true))
        .arg(
            Arg::with_name(options::SIGNAL)
                .short("s")
                .long(options::SIGNAL)
                .help("Sends given signal")
                .takes_value(true),
        )
        .arg(
            Arg::with_name(options::PIDS_OR_SIGNALS)
                .hidden(true)
                .multiple(true),
        )
        .get_matches_from(args);

    let mode = if matches.is_present(options::TABLE) || matches.is_present(options::TABLE_OLD) {
        Mode::Table
    } else if matches.is_present(options::LIST) {
        Mode::List
    } else {
        Mode::Kill
    };

    let pids_or_signals: Vec<String> = matches
        .values_of(options::PIDS_OR_SIGNALS)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    match mode {
        Mode::Kill => {
            let sig = match (obs_signal, matches.value_of(options::SIGNAL)) {
                (Some(s), Some(_)) => s, // -s takes precedence
                (Some(s), None) => s,
                (None, Some(s)) => s.to_owned(),
                (None, None) => "TERM".to_owned(),
            };
            return kill(&sig, &pids_or_signals);
        }
        Mode::Table => table(),
        Mode::List => list(pids_or_signals.get(0).cloned()),
    }

    EXIT_OK
}

fn handle_obsolete(mut args: Vec<String>) -> (Vec<String>, Option<String>) {
    let mut i = 0;
    while i < args.len() {
        // this is safe because slice is valid when it is referenced
        let slice = &args[i].clone();
        if slice.starts_with('-') && slice.chars().nth(1).map_or(false, |c| c.is_digit(10)) {
            let val = &slice[1..];
            match val.parse() {
                Ok(num) => {
                    if uucore::signals::is_signal(num) {
                        args.remove(i);
                        return (args, Some(val.to_owned()));
                    }
                }
                Err(_) => break, /* getopts will error out for us */
            }
        }
        i += 1;
    }
    (args, None)
}

fn table() {
    let name_width = ALL_SIGNALS.iter().map(|n| n.len()).max().unwrap();

    for (idx, signal) in ALL_SIGNALS.iter().enumerate() {
        print!("{0: >#2} {1: <#2$}", idx, signal, name_width + 2);
        if (idx + 1) % 7 == 0 {
            println!();
        }
    }
    println!()
}

fn print_signal(signal_name_or_value: &str) {
    for (value, &signal) in ALL_SIGNALS.iter().enumerate() {
        if signal == signal_name_or_value || (format!("SIG{}", signal)) == signal_name_or_value {
            println!("{}", value);
            exit!(EXIT_OK as i32)
        } else if signal_name_or_value == value.to_string() {
            println!("{}", signal);
            exit!(EXIT_OK as i32)
        }
    }
    crash!(EXIT_ERR, "unknown signal name {}", signal_name_or_value)
}

fn print_signals() {
    let mut pos = 0;
    for (idx, signal) in ALL_SIGNALS.iter().enumerate() {
        pos += signal.len();
        print!("{}", signal);
        if idx > 0 && pos > 73 {
            println!();
            pos = 0;
        } else {
            pos += 1;
            print!(" ");
        }
    }
}

fn list(arg: Option<String>) {
    match arg {
        Some(ref x) => print_signal(x),
        None => print_signals(),
    };
}

fn kill(signalname: &str, pids: &[String]) -> i32 {
    let mut status = 0;
    let optional_signal_value = uucore::signals::signal_by_name_or_value(signalname);
    let signal_value = match optional_signal_value {
        Some(x) => x,
        None => crash!(EXIT_ERR, "unknown signal name {}", signalname),
    };
    for pid in pids {
        match pid.parse::<usize>() {
            Ok(x) => {
                if unsafe { libc::kill(x as pid_t, signal_value as c_int) } != 0 {
                    show_error!("{}", Error::last_os_error());
                    status = 1;
                }
            }
            Err(e) => crash!(EXIT_ERR, "failed to parse argument {}: {}", pid, e),
        };
    }
    status
}
<|endoftext|>