<commit_before>#[cfg(feature = "serde-serialize")]
use serde::{Deserialize, Serialize};

use num::One;
use alga::general::ComplexField;

use crate::allocator::Allocator;
use crate::base::{DefaultAllocator, Matrix, MatrixMN, MatrixN, SquareMatrix, Vector};
use crate::constraint::{SameNumberOfRows, ShapeConstraint};
use crate::dimension::{Dim, DimAdd, DimSum, DimDiff, DimSub, Dynamic, U1};
use crate::storage::{Storage, StorageMut};

/// The Cholesky decomposition of a symmetric-definite-positive matrix.
#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
#[cfg_attr(
    feature = "serde-serialize",
    serde(bound(serialize = "DefaultAllocator: Allocator<N, D>,
         MatrixN<N, D>: Serialize"))
)]
#[cfg_attr(
    feature = "serde-serialize",
    serde(bound(deserialize = "DefaultAllocator: Allocator<N, D>,
         MatrixN<N, D>: Deserialize<'de>"))
)]
#[derive(Clone, Debug)]
pub struct Cholesky<N: ComplexField, D: Dim>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    chol: MatrixN<N, D>,
}

impl<N: ComplexField, D: Dim> Copy for Cholesky<N, D>
where
    DefaultAllocator: Allocator<N, D, D>,
    MatrixN<N, D>: Copy,
{
}

impl<N: ComplexField, D: DimSub<Dynamic>> Cholesky<N, D>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    /// Attempts to compute the Cholesky decomposition of `matrix`.
    ///
    /// Returns `None` if the input matrix is not definite-positive. The input matrix is assumed
    /// to be symmetric and only the lower-triangular part is read.
    pub fn new(mut matrix: MatrixN<N, D>) -> Option<Self> {
        assert!(matrix.is_square(), "The input matrix must be square.");

        let n = matrix.nrows();

        for j in 0..n {
            for k in 0..j {
                let factor = unsafe { -*matrix.get_unchecked((j, k)) };

                let (mut col_j, col_k) = matrix.columns_range_pair_mut(j, k);
                let mut col_j = col_j.rows_range_mut(j..);
                let col_k = col_k.rows_range(j..);

                col_j.axpy(factor.conjugate(), &col_k, N::one());
            }

            let diag = unsafe { *matrix.get_unchecked((j, j)) };
            if !diag.is_zero() {
                if let Some(denom) = diag.try_sqrt() {
                    unsafe {
                        *matrix.get_unchecked_mut((j, j)) = denom;
                    }

                    let mut col = matrix.slice_range_mut(j + 1.., j);
                    col /= denom;
                    continue;
                }
            }

            // The diagonal element is either zero or its square root could not
            // be taken (e.g. for negative real numbers).
            return None;
        }

        Some(Cholesky { chol: matrix })
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition with its strictly
    /// upper-triangular part filled with zeros.
    pub fn unpack(mut self) -> MatrixN<N, D> {
        self.chol.fill_upper_triangle(N::zero(), 1);
        self.chol
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition, without zeroing-out
    /// its strict upper-triangular part.
    ///
    /// The values of the strict upper-triangular part are garbage and should be ignored by further
    /// computations.
    pub fn unpack_dirty(self) -> MatrixN<N, D> {
        self.chol
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition with its strictly
    /// uppen-triangular part filled with zeros.
    pub fn l(&self) -> MatrixN<N, D> {
        self.chol.lower_triangle()
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition, without zeroing-out
    /// its strict upper-triangular part.
    ///
    /// This is an allocation-less version of `self.l()`. The values of the strict upper-triangular
    /// part are garbage and should be ignored by further computations.
    pub fn l_dirty(&self) -> &MatrixN<N, D> {
        &self.chol
    }

    /// Solves the system `self * x = b` where `self` is the decomposed matrix and `x` the unknown.
    ///
    /// The result is stored on `b`.
    pub fn solve_mut<R2: Dim, C2: Dim, S2>(&self, b: &mut Matrix<N, R2, C2, S2>)
    where
        S2: StorageMut<N, R2, C2>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        let _ = self.chol.solve_lower_triangular_mut(b);
        let _ = self.chol.ad_solve_lower_triangular_mut(b);
    }

    /// Returns the solution of the system `self * x = b` where `self` is the decomposed matrix and
    /// `x` the unknown.
    pub fn solve<R2: Dim, C2: Dim, S2>(&self, b: &Matrix<N, R2, C2, S2>) -> MatrixMN<N, R2, C2>
    where
        S2: Storage<N, R2, C2>,
        DefaultAllocator: Allocator<N, R2, C2>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        let mut res = b.clone_owned();
        self.solve_mut(&mut res);
        res
    }

    /// Computes the inverse of the decomposed matrix.
    pub fn inverse(&self) -> MatrixN<N, D> {
        let shape = self.chol.data.shape();
        let mut res = MatrixN::identity_generic(shape.0, shape.1);

        self.solve_mut(&mut res);
        res
    }

    /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `v`,
    /// performs a rank one update such that we end up with the decomposition of `M + sigma * v*v.adjoint()`.
    #[inline]
    pub fn rank_one_update<R2: Dim, S2>(&mut self, x: &Vector<N, R2, S2>, sigma: N::RealField)
    where
        S2: Storage<N, R2, U1>,
        DefaultAllocator: Allocator<N, R2, U1>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        Self::xx_rank_one_update(&mut self.chol, &mut x.clone_owned(), sigma)
    }

    /// Updates the decomposition such that we get the decomposition of a matrix with the given column `col` in the `j`th position.
    /// Since the matrix is square, an identical row will be added in the `j`th row.
    pub fn insert_column<R2, S2>(
        &self,
        j: usize,
        col: Vector<N, R2, S2>,
    ) -> Cholesky<N, DimSum<D, U1>>
    where
        D: DimAdd<U1>,
        R2: Dim,
        S2: Storage<N, R2, U1>,
        DefaultAllocator: Allocator<N, DimSum<D, U1>, DimSum<D, U1>> + Allocator<N, R2>,
        ShapeConstraint: SameNumberOfRows<R2, DimSum<D, U1>>,
    {
        let mut col = col.into_owned();
        // for an explanation of the formulas, see https://en.wikipedia.org/wiki/Cholesky_decomposition#Updating_the_decomposition
        let n = col.nrows();
        assert_eq!(n, self.chol.nrows() + 1, "The new column must have the size of the factored matrix plus one.");
        assert!(j < n, "j needs to be within the bound of the new matrix.");

        // loads the data into a new matrix with an additional jth row/column
        let mut chol = unsafe { Matrix::new_uninitialized_generic(self.chol.data.shape().0.add(U1), self.chol.data.shape().1.add(U1)) };
        chol.slice_range_mut(..j, ..j).copy_from(&self.chol.slice_range(..j, ..j));
        chol.slice_range_mut(..j, j+1..).copy_from(&self.chol.slice_range(..j, j..));
        chol.slice_range_mut(j+1.., ..j).copy_from(&self.chol.slice_range(j.., ..j));
        chol.slice_range_mut(j+1.., j+1..).copy_from(&self.chol.slice_range(j.., j..));

        // update the jth row
        let top_left_corner = self.chol.slice_range(..j, ..j);

        let col_j = col[j];
        let (mut new_rowj_adjoint, mut new_colj) = col.rows_range_pair_mut(..j, j + 1..);
        assert!(top_left_corner.solve_lower_triangular_mut(&mut new_rowj_adjoint), "Cholesky::insert_column : Unable to solve lower triangular system!");

        new_rowj_adjoint.adjoint_to(&mut chol.slice_range_mut(j, ..j));

        // update the center element
        let center_element = N::sqrt(col_j - N::from_real(new_rowj_adjoint.norm_squared()));
        chol[(j, j)] = center_element;

        // update the jth column
        let bottom_left_corner = self.chol.slice_range(j.., ..j);
        // new_colj = (col_jplus - bottom_left_corner * new_rowj.adjoint()) / center_element;
        new_colj.gemm(-N::one() / center_element, &bottom_left_corner, &new_rowj_adjoint, N::one() / center_element);
        chol.slice_range_mut(j+1.., j).copy_from(&new_colj);

        // update the bottom right corner
        let mut bottom_right_corner = chol.slice_range_mut(j+1.., j+1..);
        Self::xx_rank_one_update(&mut bottom_right_corner, &mut new_colj, -N::RealField::one());

        Cholesky { chol }
    }

    /// Updates the decomposition such that we get the decomposition of the factored matrix with its `j`th column removed.
    /// Since the matrix is square, the `j`th row will also be removed.
    pub fn remove_column(
        &self,
        j: usize,
    ) -> Cholesky<N, DimDiff<D, U1>>
        where
            D: DimSub<U1>,
            DefaultAllocator: Allocator<N, DimDiff<D, U1>, DimDiff<D, U1>> + Allocator<N, D>
    {
        let n = self.chol.nrows();
        assert!(n > 0, "The matrix needs at least one column.");
        assert!(j < n, "j needs to be within the bound of the matrix.");

        // loads the data into a new matrix except for the jth row/column
        let mut chol = unsafe { Matrix::new_uninitialized_generic(self.chol.data.shape().0.sub(U1), self.chol.data.shape().1.sub(U1)) };
        chol.slice_range_mut(..j, ..j).copy_from(&self.chol.slice_range(..j, ..j));
        chol.slice_range_mut(..j, j..).copy_from(&self.chol.slice_range(..j, j+1..));
        chol.slice_range_mut(j.., ..j).copy_from(&self.chol.slice_range(j+1.., ..j));
        chol.slice_range_mut(j.., j..).copy_from(&self.chol.slice_range(j+1.., j+1..));

        // updates the bottom right corner
        let mut bottom_right_corner = chol.slice_range_mut(j.., j..);
        let mut workspace = self.chol.column(j).clone_owned();
        let mut old_colj = workspace.rows_range_mut(j+1..);
        Self::xx_rank_one_update(&mut bottom_right_corner, &mut old_colj, N::RealField::one());

        Cholesky { chol }
    }

    /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `v`,
    /// performs a rank one update such that we end up with the decomposition of `M + sigma * x*x.adjoint()`.
    ///
    /// This helper method is calling for by `rank_one_update` but also `insert_column` and `remove_column`
    /// where it is used on a square slice of the decomposition
    fn xx_rank_one_update<Dm, Sm, Rx, Sx>(chol : &mut Matrix<N, Dm, Dm, Sm>, x: &mut Vector<N, Rx, Sx>, sigma: N::RealField)
        where
            //N: ComplexField,
            Dm: Dim,
            Rx: Dim,
            Sm: StorageMut<N, Dm, Dm>,
            Sx: StorageMut<N, Rx, U1>,
    {
        // heavily inspired by Eigen's `llt_rank_update_lower` implementation https://eigen.tuxfamily.org/dox/LLT_8h_source.html
        let n = x.nrows();
        assert_eq!(
            n,
            chol.nrows(),
            "The input vector must be of the same size as the factorized matrix."
        );

        let mut beta = crate::one::<N::RealField>();

        for j in 0..n {
            // updates the diagonal
            let diag = N::real(unsafe { *chol.get_unchecked((j, j)) });
            let diag2 = diag * diag;
            let xj = unsafe { *x.get_unchecked(j) };
            let sigma_xj2 = sigma * N::modulus_squared(xj);
            let gamma = diag2 * beta + sigma_xj2;
            let new_diag = (diag2 + sigma_xj2 / beta).sqrt();
            unsafe { *chol.get_unchecked_mut((j, j)) = N::from_real(new_diag) };
            beta += sigma_xj2 / diag2;
            // updates the terms of L
            let mut xjplus = x.rows_range_mut(j + 1..);
            let mut col_j = chol.slice_range_mut(j + 1.., j);
            // temp_jplus -= (wj / N::from_real(diag)) * col_j;
            xjplus.axpy(-xj / N::from_real(diag), &col_j, N::one());
            if gamma != crate::zero::<N::RealField>() {
                // col_j = N::from_real(nljj / diag) * col_j  + (N::from_real(nljj * sigma / gamma) * N::conjugate(wj)) * temp_jplus;
                col_j.axpy(
                    N::from_real(new_diag * sigma / gamma) * N::conjugate(xj),
                    &xjplus,
                    N::from_real(new_diag / diag),
                );
            }
        }
    }
}

impl<N: ComplexField, D: DimSub<Dynamic>, S: Storage<N, D, D>> SquareMatrix<N, D, S>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    /// Attempts to compute the Cholesky decomposition of this matrix.
    ///
    /// Returns `None` if the input matrix is not definite-positive. The input matrix is assumed
    /// to be symmetric and only the lower-triangular part is read.
    pub fn cholesky(self) -> Option<Cholesky<N, D>> {
        Cholesky::new(self.into_owned())
    }
}
<commit_msg>Add some missing spaces.<commit_after>#[cfg(feature = "serde-serialize")]
use serde::{Deserialize, Serialize};

use num::One;
use alga::general::ComplexField;

use crate::allocator::Allocator;
use crate::base::{DefaultAllocator, Matrix, MatrixMN, MatrixN, SquareMatrix, Vector};
use crate::constraint::{SameNumberOfRows, ShapeConstraint};
use crate::dimension::{Dim, DimAdd, DimSum, DimDiff, DimSub, Dynamic, U1};
use crate::storage::{Storage, StorageMut};

/// The Cholesky decomposition of a symmetric-definite-positive matrix.
#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
#[cfg_attr(
    feature = "serde-serialize",
    serde(bound(serialize = "DefaultAllocator: Allocator<N, D>,
         MatrixN<N, D>: Serialize"))
)]
#[cfg_attr(
    feature = "serde-serialize",
    serde(bound(deserialize = "DefaultAllocator: Allocator<N, D>,
         MatrixN<N, D>: Deserialize<'de>"))
)]
#[derive(Clone, Debug)]
pub struct Cholesky<N: ComplexField, D: Dim>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    chol: MatrixN<N, D>,
}

impl<N: ComplexField, D: Dim> Copy for Cholesky<N, D>
where
    DefaultAllocator: Allocator<N, D, D>,
    MatrixN<N, D>: Copy,
{
}

impl<N: ComplexField, D: DimSub<Dynamic>> Cholesky<N, D>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    /// Attempts to compute the Cholesky decomposition of `matrix`.
    ///
    /// Returns `None` if the input matrix is not definite-positive. The input matrix is assumed
    /// to be symmetric and only the lower-triangular part is read.
    pub fn new(mut matrix: MatrixN<N, D>) -> Option<Self> {
        assert!(matrix.is_square(), "The input matrix must be square.");

        let n = matrix.nrows();

        for j in 0..n {
            for k in 0..j {
                let factor = unsafe { -*matrix.get_unchecked((j, k)) };

                let (mut col_j, col_k) = matrix.columns_range_pair_mut(j, k);
                let mut col_j = col_j.rows_range_mut(j..);
                let col_k = col_k.rows_range(j..);

                col_j.axpy(factor.conjugate(), &col_k, N::one());
            }

            let diag = unsafe { *matrix.get_unchecked((j, j)) };
            if !diag.is_zero() {
                if let Some(denom) = diag.try_sqrt() {
                    unsafe {
                        *matrix.get_unchecked_mut((j, j)) = denom;
                    }

                    let mut col = matrix.slice_range_mut(j + 1.., j);
                    col /= denom;
                    continue;
                }
            }

            // The diagonal element is either zero or its square root could not
            // be taken (e.g. for negative real numbers).
            return None;
        }

        Some(Cholesky { chol: matrix })
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition with its strictly
    /// upper-triangular part filled with zeros.
    pub fn unpack(mut self) -> MatrixN<N, D> {
        self.chol.fill_upper_triangle(N::zero(), 1);
        self.chol
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition, without zeroing-out
    /// its strict upper-triangular part.
    ///
    /// The values of the strict upper-triangular part are garbage and should be ignored by further
    /// computations.
    pub fn unpack_dirty(self) -> MatrixN<N, D> {
        self.chol
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition with its strictly
    /// uppen-triangular part filled with zeros.
    pub fn l(&self) -> MatrixN<N, D> {
        self.chol.lower_triangle()
    }

    /// Retrieves the lower-triangular factor of the Cholesky decomposition, without zeroing-out
    /// its strict upper-triangular part.
    ///
    /// This is an allocation-less version of `self.l()`. The values of the strict upper-triangular
    /// part are garbage and should be ignored by further computations.
    pub fn l_dirty(&self) -> &MatrixN<N, D> {
        &self.chol
    }

    /// Solves the system `self * x = b` where `self` is the decomposed matrix and `x` the unknown.
    ///
    /// The result is stored on `b`.
    pub fn solve_mut<R2: Dim, C2: Dim, S2>(&self, b: &mut Matrix<N, R2, C2, S2>)
    where
        S2: StorageMut<N, R2, C2>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        let _ = self.chol.solve_lower_triangular_mut(b);
        let _ = self.chol.ad_solve_lower_triangular_mut(b);
    }

    /// Returns the solution of the system `self * x = b` where `self` is the decomposed matrix and
    /// `x` the unknown.
    pub fn solve<R2: Dim, C2: Dim, S2>(&self, b: &Matrix<N, R2, C2, S2>) -> MatrixMN<N, R2, C2>
    where
        S2: Storage<N, R2, C2>,
        DefaultAllocator: Allocator<N, R2, C2>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        let mut res = b.clone_owned();
        self.solve_mut(&mut res);
        res
    }

    /// Computes the inverse of the decomposed matrix.
    pub fn inverse(&self) -> MatrixN<N, D> {
        let shape = self.chol.data.shape();
        let mut res = MatrixN::identity_generic(shape.0, shape.1);

        self.solve_mut(&mut res);
        res
    }

    /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `v`,
    /// performs a rank one update such that we end up with the decomposition of `M + sigma * (v * v.adjoint())`.
    #[inline]
    pub fn rank_one_update<R2: Dim, S2>(&mut self, x: &Vector<N, R2, S2>, sigma: N::RealField)
    where
        S2: Storage<N, R2, U1>,
        DefaultAllocator: Allocator<N, R2, U1>,
        ShapeConstraint: SameNumberOfRows<R2, D>,
    {
        Self::xx_rank_one_update(&mut self.chol, &mut x.clone_owned(), sigma)
    }

    /// Updates the decomposition such that we get the decomposition of a matrix with the given column `col` in the `j`th position.
    /// Since the matrix is square, an identical row will be added in the `j`th row.
    pub fn insert_column<R2, S2>(
        &self,
        j: usize,
        col: Vector<N, R2, S2>,
    ) -> Cholesky<N, DimSum<D, U1>>
    where
        D: DimAdd<U1>,
        R2: Dim,
        S2: Storage<N, R2, U1>,
        DefaultAllocator: Allocator<N, DimSum<D, U1>, DimSum<D, U1>> + Allocator<N, R2>,
        ShapeConstraint: SameNumberOfRows<R2, DimSum<D, U1>>,
    {
        let mut col = col.into_owned();
        // for an explanation of the formulas, see https://en.wikipedia.org/wiki/Cholesky_decomposition#Updating_the_decomposition
        let n = col.nrows();
        assert_eq!(n, self.chol.nrows() + 1, "The new column must have the size of the factored matrix plus one.");
        assert!(j < n, "j needs to be within the bound of the new matrix.");

        // loads the data into a new matrix with an additional jth row/column
        let mut chol = unsafe { Matrix::new_uninitialized_generic(self.chol.data.shape().0.add(U1), self.chol.data.shape().1.add(U1)) };
        chol.slice_range_mut(..j, ..j).copy_from(&self.chol.slice_range(..j, ..j));
        chol.slice_range_mut(..j, j + 1..).copy_from(&self.chol.slice_range(..j, j..));
        chol.slice_range_mut(j + 1.., ..j).copy_from(&self.chol.slice_range(j.., ..j));
        chol.slice_range_mut(j + 1.., j + 1..).copy_from(&self.chol.slice_range(j.., j..));

        // update the jth row
        let top_left_corner = self.chol.slice_range(..j, ..j);

        let col_j = col[j];
        let (mut new_rowj_adjoint, mut new_colj) = col.rows_range_pair_mut(..j, j + 1..);
        assert!(top_left_corner.solve_lower_triangular_mut(&mut new_rowj_adjoint), "Cholesky::insert_column : Unable to solve lower triangular system!");

        new_rowj_adjoint.adjoint_to(&mut chol.slice_range_mut(j, ..j));

        // update the center element
        let center_element = N::sqrt(col_j - N::from_real(new_rowj_adjoint.norm_squared()));
        chol[(j, j)] = center_element;

        // update the jth column
        let bottom_left_corner = self.chol.slice_range(j.., ..j);
        // new_colj = (col_jplus - bottom_left_corner * new_rowj.adjoint()) / center_element;
        new_colj.gemm(-N::one() / center_element, &bottom_left_corner, &new_rowj_adjoint, N::one() / center_element);
        chol.slice_range_mut(j + 1.., j).copy_from(&new_colj);

        // update the bottom right corner
        let mut bottom_right_corner = chol.slice_range_mut(j + 1.., j + 1..);
        Self::xx_rank_one_update(&mut bottom_right_corner, &mut new_colj, -N::RealField::one());

        Cholesky { chol }
    }

    /// Updates the decomposition such that we get the decomposition of the factored matrix with its `j`th column removed.
    /// Since the matrix is square, the `j`th row will also be removed.
    pub fn remove_column(
        &self,
        j: usize,
    ) -> Cholesky<N, DimDiff<D, U1>>
        where
            D: DimSub<U1>,
            DefaultAllocator: Allocator<N, DimDiff<D, U1>, DimDiff<D, U1>> + Allocator<N, D>
    {
        let n = self.chol.nrows();
        assert!(n > 0, "The matrix needs at least one column.");
        assert!(j < n, "j needs to be within the bound of the matrix.");

        // loads the data into a new matrix except for the jth row/column
        let mut chol = unsafe { Matrix::new_uninitialized_generic(self.chol.data.shape().0.sub(U1), self.chol.data.shape().1.sub(U1)) };
        chol.slice_range_mut(..j, ..j).copy_from(&self.chol.slice_range(..j, ..j));
        chol.slice_range_mut(..j, j..).copy_from(&self.chol.slice_range(..j, j + 1..));
        chol.slice_range_mut(j.., ..j).copy_from(&self.chol.slice_range(j + 1.., ..j));
        chol.slice_range_mut(j.., j..).copy_from(&self.chol.slice_range(j + 1.., j + 1..));

        // updates the bottom right corner
        let mut bottom_right_corner = chol.slice_range_mut(j.., j..);
        let mut workspace = self.chol.column(j).clone_owned();
        let mut old_colj = workspace.rows_range_mut(j + 1..);
        Self::xx_rank_one_update(&mut bottom_right_corner, &mut old_colj, N::RealField::one());

        Cholesky { chol }
    }

    /// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `v`,
    /// performs a rank one update such that we end up with the decomposition of `M + sigma * (x * x.adjoint())`.
    ///
    /// This helper method is calling for by `rank_one_update` but also `insert_column` and `remove_column`
    /// where it is used on a square slice of the decomposition
    fn xx_rank_one_update<Dm, Sm, Rx, Sx>(chol : &mut Matrix<N, Dm, Dm, Sm>, x: &mut Vector<N, Rx, Sx>, sigma: N::RealField)
        where
            //N: ComplexField,
            Dm: Dim,
            Rx: Dim,
            Sm: StorageMut<N, Dm, Dm>,
            Sx: StorageMut<N, Rx, U1>,
    {
        // heavily inspired by Eigen's `llt_rank_update_lower` implementation https://eigen.tuxfamily.org/dox/LLT_8h_source.html
        let n = x.nrows();
        assert_eq!(
            n,
            chol.nrows(),
            "The input vector must be of the same size as the factorized matrix."
        );

        let mut beta = crate::one::<N::RealField>();

        for j in 0..n {
            // updates the diagonal
            let diag = N::real(unsafe { *chol.get_unchecked((j, j)) });
            let diag2 = diag * diag;
            let xj = unsafe { *x.get_unchecked(j) };
            let sigma_xj2 = sigma * N::modulus_squared(xj);
            let gamma = diag2 * beta + sigma_xj2;
            let new_diag = (diag2 + sigma_xj2 / beta).sqrt();
            unsafe { *chol.get_unchecked_mut((j, j)) = N::from_real(new_diag) };
            beta += sigma_xj2 / diag2;
            // updates the terms of L
            let mut xjplus = x.rows_range_mut(j + 1..);
            let mut col_j = chol.slice_range_mut(j + 1.., j);
            // temp_jplus -= (wj / N::from_real(diag)) * col_j;
            xjplus.axpy(-xj / N::from_real(diag), &col_j, N::one());
            if gamma != crate::zero::<N::RealField>() {
                // col_j = N::from_real(nljj / diag) * col_j  + (N::from_real(nljj * sigma / gamma) * N::conjugate(wj)) * temp_jplus;
                col_j.axpy(
                    N::from_real(new_diag * sigma / gamma) * N::conjugate(xj),
                    &xjplus,
                    N::from_real(new_diag / diag),
                );
            }
        }
    }
}

impl<N: ComplexField, D: DimSub<Dynamic>, S: Storage<N, D, D>> SquareMatrix<N, D, S>
where
    DefaultAllocator: Allocator<N, D, D>,
{
    /// Attempts to compute the Cholesky decomposition of this matrix.
    ///
    /// Returns `None` if the input matrix is not definite-positive. The input matrix is assumed
    /// to be symmetric and only the lower-triangular part is read.
    pub fn cholesky(self) -> Option<Cholesky<N, D>> {
        Cholesky::new(self.into_owned())
    }
}
<|endoftext|>


<commit_before>
pub struct Task {
        id: int,
        todo: ~str,
        priority: char,
        finished: bool,
        raw_todo: ~str,
        create_date: ~str,
        finish_date: ~str,
        contexts: & 'static [~str],
        projects: & 'static [~str]
}

impl Task {
        pub fn to_str(&self) -> ~str {
                return format!("{:d} {:b} |{:s}| ({:c}) |{:s}| '{:s}' \"{:s}\"",
                                self.id, self.finished, self.finish_date, self.priority,
                                self.create_date, self.todo, self.raw_todo)
        }

        pub fn create(text: ~str, id: int) -> Task {
                let mut task = Task{id: id, raw_todo: text.clone(), todo: text.clone(),
                                        finished: false, priority: '^',
                                        create_date: ~"", finish_date: ~"",
                                        contexts: &[], projects: &[]};

                let mut todo = text.clone();

                if (todo.slice(0, 2) == "x ") {
                        task.finished = true;
                        todo = todo.slice(2, todo.len()).to_owned();
                }

                let mut chars = todo.chars().to_owned_vec();

                chars = todo.chars().to_owned_vec();
                if (chars[10] == ' ' && chars[4] == '-' &&
                    chars[7] == '-' && task.finished) {
                        task.finish_date = todo.slice(0, 10).to_owned();
                        todo = todo.slice(11, todo.len()).to_owned();
                }

                chars = todo.chars().to_owned_vec();
                if (chars[0] == '(' && chars[2] == ')' &&
                    chars[1].is_uppercase() && chars[3] == ' ') {

                        task.priority = chars[1];
                        todo = todo.slice(4, todo.len()).to_owned();
                }

                chars = todo.chars().to_owned_vec();
                if (chars[10] == ' ' && chars[4] == '-' && chars[7] == '-') {
                        task.create_date = todo.slice(0, 10).to_owned();
                        todo = todo.slice(11, todo.len()).to_owned();
                }

                task.todo = todo;

                return task;
        }
}

pub struct Tasks {
        list: ~[Task]
}

impl Tasks {
        pub fn create() -> Tasks {
                let tasks = Tasks{list: ~[]};

                return tasks;
        }

        pub fn add(&mut self, task: Task) {
                self.list.push(task);
        }

        pub fn add_task_from_string(&mut self, text: ~str) {
                let id = self.len();
                self.list.push(Task::create(text, id));
        }

        pub fn len(&self) -> int {
                return self.list.len() as int;
        }
}

#[test]
fn simple_todo_create_test() {
        let t = Task::create(~"some important task", 1);

        assert_eq!(t.id , 1);
        assert_eq!(t.priority , '^');
        assert_eq!(t.finished , false);
        assert_eq!(t.todo , ~"some important task");
        assert_eq!(t.raw_todo , ~"some important task");
}

#[test]
fn simple_todo_with_priority_test() {
        let t = Task::create(~"(A) some important task with priority", 1);

        assert_eq!(t.id , 1);
        assert_eq!(t.priority , 'A');
        assert_eq!(t.finished , false);
        assert_eq!(t.todo , ~"some important task with priority");
        assert_eq!(t.raw_todo , ~"(A) some important task with priority");

        let ts = Task::create(~"(a) some important task with priority", 1);

        assert_eq!(ts.id , 1);
        assert_eq!(ts.priority , '^');
        assert_eq!(ts.finished , false);
        assert_eq!(ts.todo , ~"(a) some important task with priority");
        assert_eq!(ts.raw_todo , ~"(a) some important task with priority");

}

#[test]
fn simple_todo_with_date_test() {
        let t = Task::create(~"(A) 2014-03-02 some important task with priority", 1);

        assert_eq!(t.id , 1);
        assert_eq!(t.priority , 'A');
        assert_eq!(t.finished , false);
        assert_eq!(t.todo , ~"some important task with priority");
        assert_eq!(t.raw_todo , ~"(A) 2014-03-02 some important task with priority");
        assert_eq!(t.create_date, ~"2014-03-02");

        let ft = Task::create(~"x 2014-03-03 (A) 2014-03-02 some finished important task", 1);

        assert_eq!(ft.id , 1);
        assert_eq!(ft.priority , 'A');
        assert_eq!(ft.finished , true);
        assert_eq!(ft.todo , ~"some finished important task");
        assert_eq!(ft.raw_todo , ~"x 2014-03-03 (A) 2014-03-02 some finished important task");
        assert_eq!(ft.create_date, ~"2014-03-02");
        assert_eq!(ft.finish_date, ~"2014-03-03");

}

#[test]
fn todo_to_str_test() {
        let t = Task::create(~"some important task", 1);

        assert_eq!(t.to_str(), ~"1 false || (^) || 'some important task' \"some important task\"")

        let x = Task::create(~"x some important task", 1);
        assert_eq!(x.to_str(), ~"1 true || (^) || 'some important task' \"x some important task\"")
}
<commit_msg>a bit of cleanup<commit_after>
pub struct Task {
        id: int,
        todo: ~str,
        priority: char,
        finished: bool,
        raw_todo: ~str,
        create_date: ~str,
        finish_date: ~str,
        contexts: ~[~str],
        projects: ~[~str]
}

impl Task {
        pub fn to_str(&self) -> ~str {
                return format!("{:d} {:b} |{:s}| ({:c}) |{:s}| '{:s}' \"{:s}\"",
                                self.id, self.finished, self.finish_date, self.priority,
                                self.create_date, self.todo, self.raw_todo)
        }

        pub fn create(text: ~str, id: int) -> Task {
                let mut task = Task{id: id, raw_todo: text.clone(), todo: text.clone(),
                                        finished: false, priority: '^',
                                        create_date: ~"", finish_date: ~"",
                                        contexts: ~[], projects: ~[]};

                let mut todo = text.clone();

                if (todo.slice(0, 2) == "x ") {
                        task.finished = true;
                        todo = todo.slice(2, todo.len()).to_owned();
                }

                let mut chars = todo.chars().to_owned_vec();

                if (chars[10] == ' ' && chars[4] == '-' &&
                    chars[7] == '-' && task.finished) {
                        task.finish_date = todo.slice(0, 10).to_owned();
                        todo = todo.slice(11, todo.len()).to_owned();
                }

                chars = todo.chars().to_owned_vec();
                if (chars[0] == '(' && chars[2] == ')' &&
                    chars[1].is_uppercase() && chars[3] == ' ') {

                        task.priority = chars[1];
                        todo = todo.slice(4, todo.len()).to_owned();
                }

                chars = todo.chars().to_owned_vec();
                if (chars[10] == ' ' && chars[4] == '-' && chars[7] == '-') {
                        task.create_date = todo.slice(0, 10).to_owned();
                        todo = todo.slice(11, todo.len()).to_owned();
                }

                task.todo = todo;

                return task;
        }
}

pub struct Tasks {
        list: ~[Task]
}

impl Tasks {
        pub fn create() -> Tasks {
                let tasks = Tasks{list: ~[]};

                return tasks;
        }

        pub fn add(&mut self, task: Task) {
                self.list.push(task);
        }

        pub fn add_task_from_string(&mut self, text: ~str) {
                let id = self.len();
                self.list.push(Task::create(text, id));
        }

        pub fn len(&self) -> int {
                return self.list.len() as int;
        }
}

#[test]
fn simple_todo_create_test() {
        let t = Task::create(~"some important task", 1);

        assert_eq!(t.id , 1);
        assert_eq!(t.priority , '^');
        assert_eq!(t.finished , false);
        assert_eq!(t.todo , ~"some important task");
        assert_eq!(t.raw_todo , ~"some important task");
}

#[test]
fn simple_todo_with_priority_test() {
        let t = Task::create(~"(A) some important task with priority", 1);

        assert_eq!(t.id , 1);
        assert_eq!(t.priority , 'A');
        assert_eq!(t.finished , false);
        assert_eq!(t.todo , ~"some important task with priority");
        assert_eq!(t.raw_todo , ~"(A) some important task with priority");

        let ts = Task::create(~"(a) some important task with priority", 1);

        assert_eq!(ts.id , 1);
        assert_eq!(ts.priority , '^');
        assert_eq!(ts.finished , false);
        assert_eq!(ts.todo , ~"(a) some important task with priority");
        assert_eq!(ts.raw_todo , ~"(a) some important task with priority");

}

#[test]
fn simple_todo_with_date_test() {
        let t = Task::create(~"(A) 2014-03-02 some important task with priority", 1);

        assert_eq!(t.id , 1);
        assert_eq!(t.priority , 'A');
        assert_eq!(t.finished , false);
        assert_eq!(t.todo , ~"some important task with priority");
        assert_eq!(t.raw_todo , ~"(A) 2014-03-02 some important task with priority");
        assert_eq!(t.create_date, ~"2014-03-02");

        let ft = Task::create(~"x 2014-03-03 (A) 2014-03-02 some finished important task", 1);

        assert_eq!(ft.id , 1);
        assert_eq!(ft.priority , 'A');
        assert_eq!(ft.finished , true);
        assert_eq!(ft.todo , ~"some finished important task");
        assert_eq!(ft.raw_todo , ~"x 2014-03-03 (A) 2014-03-02 some finished important task");
        assert_eq!(ft.create_date, ~"2014-03-02");
        assert_eq!(ft.finish_date, ~"2014-03-03");

}

#[test]
fn todo_to_str_test() {
        let t = Task::create(~"some important task", 1);

        assert_eq!(t.to_str(), ~"1 false || (^) || 'some important task' \"some important task\"")

        let x = Task::create(~"x some important task", 1);
        assert_eq!(x.to_str(), ~"1 true || (^) || 'some important task' \"x some important task\"")
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
extern crate rustc_serialize;

mod download;
mod extract;
mod find_package;
mod package;
mod verify;

use clap::{App, Arg};

arg_enum! {
    #[derive(Debug)]
    enum Method {
        Install,
        Remove
    }
}

fn main() {
    let arg_matches = App::new("wiz")
        .version("0.0.1")

        .arg(Arg::with_name("method")
             .help("Specifies what to do with the specified package")
             .index(1)
             .possible_values(&Method::variants())
             .required(true))

        .arg(Arg::with_name("package")
             .help("Specifies package")
             .index(2)
             .required(true))

        .get_matches();

    let method = value_t!(arg_matches.value_of("method"), Method).unwrap();
    let package = arg_matches.value_of("package").unwrap_or("invalid");

    match method {
        Method::Install => println!("install: {}", package),
        Method::Remove => println!("remove: {}", package)
    }
}
<commit_msg>Add package_version.rs<commit_after>#[macro_use]
extern crate clap;
extern crate rustc_serialize;

mod download;
mod extract;
mod find_package;
mod package;
mod package_version;
mod verify;

use clap::{App, Arg};

arg_enum! {
    #[derive(Debug)]
    enum Method {
        Install,
        Remove
    }
}

fn main() {
    let arg_matches = App::new("wiz")
        .version("0.0.1")

        .arg(Arg::with_name("method")
             .help("Specifies what to do with the specified package")
             .index(1)
             .possible_values(&Method::variants())
             .required(true))

        .arg(Arg::with_name("package")
             .help("Specifies package")
             .index(2)
             .required(true))

        .get_matches();

    let method = value_t!(arg_matches.value_of("method"), Method).unwrap();
    let package = arg_matches.value_of("package").unwrap_or("invalid");

    match method {
        Method::Install => println!("install: {}", package),
        Method::Remove => println!("remove: {}", package)
    }
}
<|endoftext|>


<commit_before>#![feature(async_closure)]

extern crate dirs;
extern crate hyper;
extern crate hyper_rustls;
extern crate net2;
extern crate pem;
#[macro_use]
extern crate prettytable;
extern crate rustls;
extern crate structopt;
extern crate structopt_derive;
#[macro_use]
extern crate tera;
extern crate tokio_io;
extern crate tokio_process;
extern crate tokio_rustls;
extern crate tokio_service;

extern crate ct_tools;

use ct_tools::common::{sha256_hex, Log};
use ct_tools::ct::submit_cert_to_logs;
use ct_tools::google::{fetch_all_ct_logs, fetch_trusted_ct_logs};
use ct_tools::{crtsh, letsencrypt};
use futures::stream::{StreamExt, TryStreamExt};
use net2::unix::UnixTcpBuilderExt;
use rustls::Session;
use std::fs::{self, File};
use std::io::Read;
use std::net::SocketAddr;
use std::process::Stdio;
use std::sync::Arc;
use std::time::Duration;
use structopt::StructOpt;
use tokio_io::AsyncWriteExt;
use tokio_net::driver::Handle;
use tokio_process::Command;

fn pems_to_chain(data: &[u8]) -> Vec<Vec<u8>> {
    pem::parse_many(data)
        .into_iter()
        .filter(|p| p.tag == "CERTIFICATE")
        .map(|p| p.contents)
        .collect()
}

fn new_http_client(
) -> hyper::Client<hyper_rustls::HttpsConnector<hyper::client::connect::HttpConnector>> {
    hyper::Client::builder().build(hyper_rustls::HttpsConnector::new(4))
}

fn compute_paths(paths: &[String]) -> Vec<String> {
    paths
        .iter()
        .flat_map(|p| {
            if fs::metadata(p).unwrap().is_dir() {
                fs::read_dir(p)
                    .unwrap()
                    .map(|d| d.unwrap().path().to_str().unwrap().to_string())
                    .collect()
            } else {
                vec![p.clone()]
            }
        })
        .collect()
}

async fn submit(paths: &[String], all_logs: bool) {
    let http_client = new_http_client();

    let logs = if all_logs {
        fetch_all_ct_logs(&http_client).await
    } else {
        fetch_trusted_ct_logs(&http_client).await
    };

    let all_paths = compute_paths(paths);

    let work = all_paths
        .iter()
        .map(|path| {
            let path = path.to_string();
            let mut contents = Vec::new();
            File::open(&path)
                .unwrap()
                .read_to_end(&mut contents)
                .unwrap();

            let http_client = &http_client;
            let logs = &logs;
            let mut chain = pems_to_chain(&contents);
            async move {
                if chain.len() == 1 {
                    // TODO: There's got to be some way to do this ourselves, instead of using crt.sh
                    // as a glorified AIA chaser.
                    println!(
                    "[{}] Only one certificate in chain, using crt.sh to build a full chain ...",
                    &path
                );
                    let new_chain = crtsh::build_chain_for_cert(&http_client, &chain[0]).await;
                    chain = match new_chain {
                        Ok(c) => c,
                        Err(()) => {
                            println!("[{}] Unable to build a chain", path);
                            return futures::future::ready(());
                        }
                    }
                }
                println!("[{}] Submitting ...", &path);
                let timeout = Duration::from_secs(30);
                let scts = submit_cert_to_logs(&http_client, &logs, &chain, timeout).await;

                if !scts.is_empty() {
                    println!(
                        "[{}] Find the cert on crt.sh: {}",
                        path,
                        crtsh::url_for_cert(&chain[0])
                    );
                    let mut table = prettytable::Table::new();
                    table.add_row(row!["Log"]);
                    for (log_idx, _) in scts {
                        let log = &logs[log_idx];
                        table.add_row(row![log.description]);
                    }
                    table.printstd();
                    println!();
                    println!();
                } else {
                    println!("[{}] No SCTs obtained", &path);
                }

                futures::future::ready(())
            }
        })
        .collect::<futures::stream::FuturesOrdered<_>>()
        .buffered(4)
        .for_each(async move |()| ());
    work.await;
}

async fn check(paths: &[String]) {
    let http_client = new_http_client();

    let all_paths = compute_paths(paths);

    let work = all_paths
        .iter()
        .map(|path| {
            let path = path.to_string();
            let mut contents = Vec::new();
            File::open(&path)
                .unwrap()
                .read_to_end(&mut contents)
                .unwrap();

            let http_client = &http_client;
            let chain = pems_to_chain(&contents);
            async move {
                if !chain.is_empty() && crtsh::is_cert_logged(&http_client, &chain[0]).await {
                    println!("{} was already logged", path);
                } else {
                    println!("{} has not been logged", path);
                }

                futures::future::ready(())
            }
        })
        .collect::<futures::stream::FuturesOrdered<_>>()
        .buffered(16)
        .for_each(async move |()| ());
    work.await;
}

async fn handle_request<C: hyper::client::connect::Connect + 'static>(
    request: hyper::Request<hyper::Body>,
    templates: Arc<tera::Tera>,
    http_client: Arc<hyper::Client<C>>,
    logs: Arc<Vec<Log>>,
    client_certs: Option<Vec<rustls::Certificate>>,
) -> Result<hyper::Response<hyper::Body>, hyper::Error> {
    let mut crtsh_url = None;
    let mut rendered_cert = None;
    if let Some(peer_chain) = client_certs {
        if request.method() == hyper::Method::POST {
            if let Ok(chain) = crtsh::build_chain_for_cert(&http_client, &peer_chain[0].0).await {
                let timeout = Duration::from_secs(5);
                let scts = submit_cert_to_logs(&http_client, &logs, &chain, timeout).await;
                if !scts.is_empty() {
                    crtsh_url = Some(crtsh::url_for_cert(&chain[0]));
                    println!("Successfully submitted: {}", sha256_hex(&chain[0]));
                }
            }
        }

        let mut process = Command::new("openssl")
            .arg("x509")
            .arg("-text")
            .arg("-noout")
            .arg("-inform")
            .arg("der")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .unwrap();
        let cert_bytes = peer_chain[0].0.clone();
        let mut stdin = process.stdin().take().unwrap();
        AsyncWriteExt::write_all(&mut stdin, &cert_bytes)
            .await
            .unwrap();
        let out = process.wait_with_output().await.unwrap();
        rendered_cert = Some(String::from_utf8_lossy(&out.stdout).into_owned());
    }

    let mut context = tera::Context::new();
    context.insert("cert", &rendered_cert);
    context.insert("crtsh_url", &crtsh_url);
    let body = templates.render("home.html", &context).unwrap();
    Ok(hyper::Response::builder()
        .status(hyper::StatusCode::OK)
        .body(body.into())
        .unwrap())
}

struct NoVerificationCertificateVerifier;
impl rustls::ClientCertVerifier for NoVerificationCertificateVerifier {
    fn client_auth_mandatory(&self) -> bool {
        false
    }

    fn verify_client_cert(
        &self,
        _: &[rustls::Certificate],
    ) -> Result<rustls::ClientCertVerified, rustls::TLSError> {
        Ok(rustls::ClientCertVerified::assertion())
    }

    fn client_auth_root_subjects(&self) -> rustls::DistinguishedNames {
        rustls::DistinguishedNames::new()
    }
}

async fn server(local_dev: bool, domain: Option<&str>, letsencrypt_env: Option<&str>) {
    match (local_dev, domain, letsencrypt_env) {
        (true, Some(_), _) | (true, _, Some(_)) => {
            panic!("Can't use both --local-dev and --letsencrypt-env or --domain")
        }
        (_, Some(_), None) | (_, None, Some(_)) => {
            panic!("When using Let's Encrypt, must set both --letsencrypt-env and --domain")
        }
        (false, _, None) => panic!("Must set at least one of --local-dev or --letsencrypt-env"),
        _ => {}
    };

    // Disable certificate verification. In any normal context, this would be horribly insecure!
    // However, all we're doing is taking the certs and then turning around and submitting them to
    // CT logs, so it doesn't matter if they're verified.
    let mut tls_config = rustls::ServerConfig::new(Arc::new(NoVerificationCertificateVerifier));
    if local_dev {
        // TODO: not all the details on the cert are perfect, but it's fine.
        let (cert, pkey) = letsencrypt::generate_temporary_cert("localhost");
        tls_config
            .set_single_cert(
                vec![letsencrypt::openssl_cert_to_rustls(&cert)],
                letsencrypt::openssl_pkey_to_rustls(&pkey),
            )
            .unwrap();
    } else {
        let letsencrypt_url = match letsencrypt_env.unwrap() {
            "prod" => "https://acme-v01.api.letsencrypt.org/directory",
            "dev" => "https://acme-staging.api.letsencrypt.org/directory",
            _ => unreachable!(),
        };
        let cert_cache = letsencrypt::DiskCache::new(
            dirs::home_dir()
                .unwrap()
                .join(".ct-tools")
                .join("certificates"),
        );
        tls_config.cert_resolver = Arc::new(letsencrypt::AutomaticCertResolver::new(
            letsencrypt_url,
            vec![domain.unwrap().to_string()],
            cert_cache,
        ));
    }
    tls_config.set_persistence(rustls::ServerSessionMemoryCache::new(1024));
    tls_config.ticketer = rustls::Ticketer::new();

    let http_client = Arc::new(new_http_client());
    let logs = Arc::new(fetch_trusted_ct_logs(&http_client).await);
    let templates = Arc::new(compile_templates!("templates/*"));

    let addr = if local_dev {
        "0.0.0.0:8000"
    } else {
        "0.0.0.0:443"
    }
    .parse()
    .unwrap();

    // If there aren't at least two threads, the Let's Encrypt integration will deadlock.
    println!("Listening on https://{} ...", addr);
    let listener = match addr {
        SocketAddr::V4(_) => net2::TcpBuilder::new_v4().unwrap(),
        SocketAddr::V6(_) => net2::TcpBuilder::new_v6().unwrap(),
    };
    listener.reuse_port(true).unwrap();
    listener.reuse_address(true).unwrap();
    listener.bind(addr).unwrap();
    let tls_acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(tls_config));
    let connections =
        tokio::net::TcpListener::from_std(listener.listen(1024).unwrap(), &Handle::default())
            .unwrap()
            .incoming()
            .and_then(move |sock| tls_acceptor.accept(sock))
            .filter(move |s| futures::future::ready(s.is_ok()))
            .boxed();
    let server = hyper::Server::builder(connections).serve(hyper::service::make_service_fn(
        move |conn: &tokio_rustls::server::TlsStream<tokio::net::TcpStream>| {
            let http_client = Arc::clone(&http_client);
            let templates = Arc::clone(&templates);
            let logs = Arc::clone(&logs);
            let client_certs = conn.get_ref().1.get_peer_certificates();

            async {
                Ok::<_, hyper::Error>(hyper::service::service_fn(
                    move |r: hyper::Request<hyper::Body>| {
                        handle_request(
                            r,
                            Arc::clone(&templates),
                            Arc::clone(&http_client),
                            Arc::clone(&logs),
                            client_certs.clone(),
                        )
                    },
                ))
            }
        },
    ));

    server.await.unwrap()
}

#[derive(StructOpt)]
#[structopt(name = "ct-tools")]
enum Opt {
    #[structopt(name = "submit", about = "Directly submits certificates to CT logs")]
    Submit {
        #[structopt(
            long = "all-logs",
            help = "Submit to all logs, instead of just ones trusted by Chrome"
        )]
        all_logs: bool,
        #[structopt(help = "Path to certificate or chain")]
        paths: Vec<String>,
    },

    #[structopt(
        name = "check",
        about = "Checks whether a certificate exists in CT logs"
    )]
    Check {
        #[structopt(help = "Path to certificate or chain")]
        paths: Vec<String>,
    },

    #[structopt(
        name = "server",
        about = "Run an HTTPS server that submits client certificates to CT logs"
    )]
    Server {
        #[structopt(
            long = "--local-dev",
            help = "Local development, do not obtain a certificate"
        )]
        local_dev: bool,
        #[structopt(long = "domain", help = "Domain this is running as")]
        domain: Option<String>,
        #[structopt(
            long = "letsencrypt-env",
            raw(possible_values = "&[\"dev\", \"prod\"]"),
            help = "Let's Encrypt environment to use"
        )]
        letsencrypt_env: Option<String>,
    },
}

#[tokio::main]
async fn main() {
    match Opt::from_args() {
        Opt::Submit { paths, all_logs } => {
            submit(&paths, all_logs).await;
        }
        Opt::Check { paths } => {
            check(&paths).await;
        }
        Opt::Server {
            local_dev,
            domain,
            letsencrypt_env,
        } => {
            server(
                local_dev,
                domain.as_ref().map(|s| s.as_ref()),
                letsencrypt_env.as_ref().map(|s| s.as_ref()),
            )
            .await;
        }
    }
}
<commit_msg>Only pass the first client cert in the chain to the http handler<commit_after>#![feature(async_closure)]

extern crate dirs;
extern crate hyper;
extern crate hyper_rustls;
extern crate net2;
extern crate pem;
#[macro_use]
extern crate prettytable;
extern crate rustls;
extern crate structopt;
extern crate structopt_derive;
#[macro_use]
extern crate tera;
extern crate tokio_io;
extern crate tokio_process;
extern crate tokio_rustls;
extern crate tokio_service;

extern crate ct_tools;

use ct_tools::common::{sha256_hex, Log};
use ct_tools::ct::submit_cert_to_logs;
use ct_tools::google::{fetch_all_ct_logs, fetch_trusted_ct_logs};
use ct_tools::{crtsh, letsencrypt};
use futures::stream::{StreamExt, TryStreamExt};
use net2::unix::UnixTcpBuilderExt;
use rustls::Session;
use std::fs::{self, File};
use std::io::Read;
use std::net::SocketAddr;
use std::process::Stdio;
use std::sync::Arc;
use std::time::Duration;
use structopt::StructOpt;
use tokio_io::AsyncWriteExt;
use tokio_net::driver::Handle;
use tokio_process::Command;

fn pems_to_chain(data: &[u8]) -> Vec<Vec<u8>> {
    pem::parse_many(data)
        .into_iter()
        .filter(|p| p.tag == "CERTIFICATE")
        .map(|p| p.contents)
        .collect()
}

fn new_http_client(
) -> hyper::Client<hyper_rustls::HttpsConnector<hyper::client::connect::HttpConnector>> {
    hyper::Client::builder().build(hyper_rustls::HttpsConnector::new(4))
}

fn compute_paths(paths: &[String]) -> Vec<String> {
    paths
        .iter()
        .flat_map(|p| {
            if fs::metadata(p).unwrap().is_dir() {
                fs::read_dir(p)
                    .unwrap()
                    .map(|d| d.unwrap().path().to_str().unwrap().to_string())
                    .collect()
            } else {
                vec![p.clone()]
            }
        })
        .collect()
}

async fn submit(paths: &[String], all_logs: bool) {
    let http_client = new_http_client();

    let logs = if all_logs {
        fetch_all_ct_logs(&http_client).await
    } else {
        fetch_trusted_ct_logs(&http_client).await
    };

    let all_paths = compute_paths(paths);

    let work = all_paths
        .iter()
        .map(|path| {
            let path = path.to_string();
            let mut contents = Vec::new();
            File::open(&path)
                .unwrap()
                .read_to_end(&mut contents)
                .unwrap();

            let http_client = &http_client;
            let logs = &logs;
            let mut chain = pems_to_chain(&contents);
            async move {
                if chain.len() == 1 {
                    // TODO: There's got to be some way to do this ourselves, instead of using crt.sh
                    // as a glorified AIA chaser.
                    println!(
                    "[{}] Only one certificate in chain, using crt.sh to build a full chain ...",
                    &path
                );
                    let new_chain = crtsh::build_chain_for_cert(&http_client, &chain[0]).await;
                    chain = match new_chain {
                        Ok(c) => c,
                        Err(()) => {
                            println!("[{}] Unable to build a chain", path);
                            return futures::future::ready(());
                        }
                    }
                }
                println!("[{}] Submitting ...", &path);
                let timeout = Duration::from_secs(30);
                let scts = submit_cert_to_logs(&http_client, &logs, &chain, timeout).await;

                if !scts.is_empty() {
                    println!(
                        "[{}] Find the cert on crt.sh: {}",
                        path,
                        crtsh::url_for_cert(&chain[0])
                    );
                    let mut table = prettytable::Table::new();
                    table.add_row(row!["Log"]);
                    for (log_idx, _) in scts {
                        let log = &logs[log_idx];
                        table.add_row(row![log.description]);
                    }
                    table.printstd();
                    println!();
                    println!();
                } else {
                    println!("[{}] No SCTs obtained", &path);
                }

                futures::future::ready(())
            }
        })
        .collect::<futures::stream::FuturesOrdered<_>>()
        .buffered(4)
        .for_each(async move |()| ());
    work.await;
}

async fn check(paths: &[String]) {
    let http_client = new_http_client();

    let all_paths = compute_paths(paths);

    let work = all_paths
        .iter()
        .map(|path| {
            let path = path.to_string();
            let mut contents = Vec::new();
            File::open(&path)
                .unwrap()
                .read_to_end(&mut contents)
                .unwrap();

            let http_client = &http_client;
            let chain = pems_to_chain(&contents);
            async move {
                if !chain.is_empty() && crtsh::is_cert_logged(&http_client, &chain[0]).await {
                    println!("{} was already logged", path);
                } else {
                    println!("{} has not been logged", path);
                }

                futures::future::ready(())
            }
        })
        .collect::<futures::stream::FuturesOrdered<_>>()
        .buffered(16)
        .for_each(async move |()| ());
    work.await;
}

async fn handle_request<C: hyper::client::connect::Connect + 'static>(
    request: hyper::Request<hyper::Body>,
    templates: Arc<tera::Tera>,
    http_client: Arc<hyper::Client<C>>,
    logs: Arc<Vec<Log>>,
    client_cert: Option<rustls::Certificate>,
) -> Result<hyper::Response<hyper::Body>, hyper::Error> {
    let mut crtsh_url = None;
    let mut rendered_cert = None;
    if let Some(cert) = client_cert {
        if request.method() == hyper::Method::POST {
            if let Ok(chain) = crtsh::build_chain_for_cert(&http_client, &cert.0).await {
                let timeout = Duration::from_secs(5);
                let scts = submit_cert_to_logs(&http_client, &logs, &chain, timeout).await;
                if !scts.is_empty() {
                    crtsh_url = Some(crtsh::url_for_cert(&chain[0]));
                    println!("Successfully submitted: {}", sha256_hex(&chain[0]));
                }
            }
        }

        let mut process = Command::new("openssl")
            .arg("x509")
            .arg("-text")
            .arg("-noout")
            .arg("-inform")
            .arg("der")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .unwrap();
        let cert_bytes = cert.0.clone();
        let mut stdin = process.stdin().take().unwrap();
        AsyncWriteExt::write_all(&mut stdin, &cert_bytes)
            .await
            .unwrap();
        let out = process.wait_with_output().await.unwrap();
        rendered_cert = Some(String::from_utf8_lossy(&out.stdout).into_owned());
    }

    let mut context = tera::Context::new();
    context.insert("cert", &rendered_cert);
    context.insert("crtsh_url", &crtsh_url);
    let body = templates.render("home.html", &context).unwrap();
    Ok(hyper::Response::builder()
        .status(hyper::StatusCode::OK)
        .body(body.into())
        .unwrap())
}

struct NoVerificationCertificateVerifier;
impl rustls::ClientCertVerifier for NoVerificationCertificateVerifier {
    fn client_auth_mandatory(&self) -> bool {
        false
    }

    fn verify_client_cert(
        &self,
        _: &[rustls::Certificate],
    ) -> Result<rustls::ClientCertVerified, rustls::TLSError> {
        Ok(rustls::ClientCertVerified::assertion())
    }

    fn client_auth_root_subjects(&self) -> rustls::DistinguishedNames {
        rustls::DistinguishedNames::new()
    }
}

async fn server(local_dev: bool, domain: Option<&str>, letsencrypt_env: Option<&str>) {
    match (local_dev, domain, letsencrypt_env) {
        (true, Some(_), _) | (true, _, Some(_)) => {
            panic!("Can't use both --local-dev and --letsencrypt-env or --domain")
        }
        (_, Some(_), None) | (_, None, Some(_)) => {
            panic!("When using Let's Encrypt, must set both --letsencrypt-env and --domain")
        }
        (false, _, None) => panic!("Must set at least one of --local-dev or --letsencrypt-env"),
        _ => {}
    };

    // Disable certificate verification. In any normal context, this would be horribly insecure!
    // However, all we're doing is taking the certs and then turning around and submitting them to
    // CT logs, so it doesn't matter if they're verified.
    let mut tls_config = rustls::ServerConfig::new(Arc::new(NoVerificationCertificateVerifier));
    if local_dev {
        // TODO: not all the details on the cert are perfect, but it's fine.
        let (cert, pkey) = letsencrypt::generate_temporary_cert("localhost");
        tls_config
            .set_single_cert(
                vec![letsencrypt::openssl_cert_to_rustls(&cert)],
                letsencrypt::openssl_pkey_to_rustls(&pkey),
            )
            .unwrap();
    } else {
        let letsencrypt_url = match letsencrypt_env.unwrap() {
            "prod" => "https://acme-v01.api.letsencrypt.org/directory",
            "dev" => "https://acme-staging.api.letsencrypt.org/directory",
            _ => unreachable!(),
        };
        let cert_cache = letsencrypt::DiskCache::new(
            dirs::home_dir()
                .unwrap()
                .join(".ct-tools")
                .join("certificates"),
        );
        tls_config.cert_resolver = Arc::new(letsencrypt::AutomaticCertResolver::new(
            letsencrypt_url,
            vec![domain.unwrap().to_string()],
            cert_cache,
        ));
    }
    tls_config.set_persistence(rustls::ServerSessionMemoryCache::new(1024));
    tls_config.ticketer = rustls::Ticketer::new();

    let http_client = Arc::new(new_http_client());
    let logs = Arc::new(fetch_trusted_ct_logs(&http_client).await);
    let templates = Arc::new(compile_templates!("templates/*"));

    let addr = if local_dev {
        "0.0.0.0:8000"
    } else {
        "0.0.0.0:443"
    }
    .parse()
    .unwrap();

    // If there aren't at least two threads, the Let's Encrypt integration will deadlock.
    println!("Listening on https://{} ...", addr);
    let listener = match addr {
        SocketAddr::V4(_) => net2::TcpBuilder::new_v4().unwrap(),
        SocketAddr::V6(_) => net2::TcpBuilder::new_v6().unwrap(),
    };
    listener.reuse_port(true).unwrap();
    listener.reuse_address(true).unwrap();
    listener.bind(addr).unwrap();
    let tls_acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(tls_config));
    let connections =
        tokio::net::TcpListener::from_std(listener.listen(1024).unwrap(), &Handle::default())
            .unwrap()
            .incoming()
            .and_then(move |sock| tls_acceptor.accept(sock))
            .filter(move |s| futures::future::ready(s.is_ok()))
            .boxed();
    let server = hyper::Server::builder(connections).serve(hyper::service::make_service_fn(
        move |conn: &tokio_rustls::server::TlsStream<tokio::net::TcpStream>| {
            let http_client = Arc::clone(&http_client);
            let templates = Arc::clone(&templates);
            let logs = Arc::clone(&logs);
            let client_cert = conn
                .get_ref()
                .1
                .get_peer_certificates()
                .map(|mut chain| chain.remove(0));

            async {
                Ok::<_, hyper::Error>(hyper::service::service_fn(
                    move |r: hyper::Request<hyper::Body>| {
                        handle_request(
                            r,
                            Arc::clone(&templates),
                            Arc::clone(&http_client),
                            Arc::clone(&logs),
                            client_cert.clone(),
                        )
                    },
                ))
            }
        },
    ));

    server.await.unwrap()
}

#[derive(StructOpt)]
#[structopt(name = "ct-tools")]
enum Opt {
    #[structopt(name = "submit", about = "Directly submits certificates to CT logs")]
    Submit {
        #[structopt(
            long = "all-logs",
            help = "Submit to all logs, instead of just ones trusted by Chrome"
        )]
        all_logs: bool,
        #[structopt(help = "Path to certificate or chain")]
        paths: Vec<String>,
    },

    #[structopt(
        name = "check",
        about = "Checks whether a certificate exists in CT logs"
    )]
    Check {
        #[structopt(help = "Path to certificate or chain")]
        paths: Vec<String>,
    },

    #[structopt(
        name = "server",
        about = "Run an HTTPS server that submits client certificates to CT logs"
    )]
    Server {
        #[structopt(
            long = "--local-dev",
            help = "Local development, do not obtain a certificate"
        )]
        local_dev: bool,
        #[structopt(long = "domain", help = "Domain this is running as")]
        domain: Option<String>,
        #[structopt(
            long = "letsencrypt-env",
            raw(possible_values = "&[\"dev\", \"prod\"]"),
            help = "Let's Encrypt environment to use"
        )]
        letsencrypt_env: Option<String>,
    },
}

#[tokio::main]
async fn main() {
    match Opt::from_args() {
        Opt::Submit { paths, all_logs } => {
            submit(&paths, all_logs).await;
        }
        Opt::Check { paths } => {
            check(&paths).await;
        }
        Opt::Server {
            local_dev,
            domain,
            letsencrypt_env,
        } => {
            server(
                local_dev,
                domain.as_ref().map(|s| s.as_ref()),
                letsencrypt_env.as_ref().map(|s| s.as_ref()),
            )
            .await;
        }
    }
}
<|endoftext|>


<commit_before>fn main() {
    println!("Hello again!");
}
<commit_msg>Read from STDIN and PTTY input and print to PTTY output<commit_after>#![recursion_limit="256"]
// use futures::{select, FutureExt};
use async_std::io;
use async_std::os::unix::io::AsRawFd;
use async_std::prelude::*;
use async_std::task;
use async_std::fs;
use termios::{self, Termios};

type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;

async fn get_tty() -> Result<fs::File> {
    let tty = fs::OpenOptions::new()
        .read(true)
        .write(true)
        .open("/dev/tty")
        .await?;

    Ok(tty)
}

fn main() -> Result<()> {
    task::block_on(async {
        // Get all inputs
        let stdin = io::stdin();
        let tty_in = get_tty().await?;

        // Get all outputs
        let mut tty_out = get_tty().await?;

        // TODO: This raw mode may need to be different than tty_out
        // BEGIN: RAW MODE
        let fd = tty_in.as_raw_fd();
        let tty = Termios::from_fd(fd)?;
        let mut raw_tty = tty;
        raw_tty.c_lflag &= !(termios::ICANON | termios::ECHO);
        termios::tcsetattr(fd, termios::TCSANOW, &raw_tty)?;
        // END:   RAW MODE

        // BEGIN: RAW MODE
        let fd = tty_out.as_raw_fd();
        let tty = Termios::from_fd(fd)?;
        let mut raw_tty = tty;
        raw_tty.c_lflag &= !(termios::ICANON | termios::ECHO);
        termios::tcsetattr(fd, termios::TCSANOW, &raw_tty)?;
        // END:   RAW MODE

        let std_reader = io::BufReader::new(stdin);
        let tty_reader = io::BufReader::new(tty_in);

        let std_lines = std_reader.lines();
        let tty_lines = tty_reader.lines();

        let mut all = futures::stream::select(tty_lines, std_lines);

        while let Some(line) = all.next().await {
            let l = format!("{}\n", line?);
            tty_out.write_all(l.as_bytes()).await?;
            tty_out.flush().await?;
        }

        Ok(())
    })
}
<|endoftext|>


<commit_before>extern crate rand;

use rand::thread_rng;
use rand::distributions::Range;
use individual::Individual;

mod individual;

fn main() {
    let mut rng = thread_rng();
    let range = Range::new(-512.03_f64, 511.97); // range for Schwefel problem

    // initialize population
    let population: Vec<_> = (0..128).map(|_| {
        Individual::new(&range, &mut rng)
    }).collect();

    loop {
        // generate mutated offspring
        let offspring: Vec<_> = population.iter().map(|x| {
            x.mutate(&range, &mut rng)
        }).collect();

        let best = offspring.iter().min().unwrap();
        println!("Fitness: {}", best.fitness);

        if best.fitness < 1000_f64 {
            println!("Solution: {:?}", best.solution);
            break;
        }
    }
}
<commit_msg>Run algorithm for a number of iterations<commit_after>extern crate rand;

use rand::thread_rng;
use rand::distributions::Range;
use individual::Individual;

mod individual;

fn main() {
    let mut rng = thread_rng();
    let range = Range::new(-512.03_f64, 511.97); // range for Schwefel problem

    // initialize population
    let population: Vec<_> = (0..128).map(|_| {
        Individual::new(&range, &mut rng)
    }).collect();

    for i in 0..10000 {
        // generate mutated offspring
        let offspring: Vec<_> = population.iter().map(|x| {
            x.mutate(&range, &mut rng)
        }).collect();

        let best = offspring.iter().min().unwrap();
        println!("{}th fitness: {}", i, best.fitness);

        if best.fitness < 1000_f64 {
            println!("Solution: {:?}", best.solution);
            return;
        }
    }
    println!("Failed to converge.");
}
<|endoftext|>


<commit_before>extern crate mio;
extern crate http_muncher;
extern crate sha1;
extern crate rustc_serialize;

use std::collections::HashMap;
use std::collections::LinkedList;
use std::cell::RefCell;
use std::rc::Rc;
use std::fmt;
use std::sync::mpsc;
use std::thread;
use std::str::FromStr;
use std::io::Read;
use std::io::Write;

use mio::*;
use mio::tcp::*;
use http_muncher::{Parser, ParserHandler};
use rustc_serialize::base64::{ToBase64, STANDARD};

enum InternalMessage {
    NewClient{client_token: Token},
    CloseClient{client_token: Token},
    Data{client_token: Token, format: String, data: String},
}

fn gen_key(key: &String) -> String {
    let mut m = sha1::Sha1::new();
    let mut buf = [0u8; 20];

    m.update(key.as_bytes());
    m.update("258EAFA5-E914-47DA-95CA-C5AB0DC85B11".as_bytes());

    m.output(&mut buf);

    return buf.to_base64(STANDARD);
}

struct HttpParser {
    current_key: Option<String>,
    headers: Rc<RefCell<HashMap<String, String>>>
}

impl ParserHandler for HttpParser {
    fn on_header_field(&mut self, s: &[u8]) -> bool {
        self.current_key = Some(std::str::from_utf8(s).unwrap().to_string());
        true
    }

    fn on_header_value(&mut self, s: &[u8]) -> bool {
        self.headers.borrow_mut()
            .insert(self.current_key.clone().unwrap(),
                    std::str::from_utf8(s).unwrap().to_string());
        true
    }

    fn on_headers_complete(&mut self) -> bool {
        false
    }
}

#[derive(PartialEq)]
enum ClientState {
    AwaitingHandshake,
    HandshakeResponse,
    Open,
}

enum SubState {
    Waiting,
    Doing,
}

struct WebSocketClient {
    socket: TcpStream,
    headers: Rc<RefCell<HashMap<String, String>>>,
    http_parser: Parser<HttpParser>,
    interest: EventSet,
    state: ClientState,
    writing_substate: SubState,
    reading_substate: SubState,
    outgoing_messages: std::collections::LinkedList<String>,
}

impl WebSocketClient {
    fn new(socket: TcpStream) -> WebSocketClient {
        let headers = Rc::new(RefCell::new(HashMap::new()));

        WebSocketClient {
            socket: socket,
            headers: headers.clone(),
            http_parser: Parser::request(HttpParser {
                current_key: None,
                headers: headers.clone()
            }),
            interest: EventSet::readable(),
            state: ClientState::AwaitingHandshake,
            writing_substate: SubState::Waiting,
            reading_substate: SubState::Waiting,
            outgoing_messages: std::collections::LinkedList::new(),
        }
    }

    fn write(&mut self) {
        let headers = self.headers.borrow();
        let response_key = gen_key(&headers.get("Sec-WebSocket-Key").unwrap());
        let response = fmt::format(format_args!("HTTP/1.1 101 Switching Protocols\r\n\
                                                 Connection: Upgrade\r\n\
                                                 Sec-WebSocket-Accept: {}\r\n\
                                                 Upgrade: websocket\r\n\r\n", response_key));
        self.socket.try_write(response.as_bytes()).unwrap();

        // Change the state
        self.state = ClientState::Open;

        self.interest.remove(EventSet::writable());
        self.interest.insert(EventSet::readable());
    }

    fn read(&mut self) {
        loop {
            let mut buf = [0; 2048];
            match self.socket.try_read(&mut buf) {
                Err(e) => {
                    println!("Error while reading socket: {:?}", e);
                    return
                },
                Ok(None) =>
                    // Socket buffer has got no more bytes.
                    break,
                Ok(Some(_len)) => {
                    self.http_parser.parse(&buf);
                    if self.http_parser.is_upgrade() {
                        // Change the current state
                        self.state = ClientState::HandshakeResponse;

                        // Change current interest to `Writable`
                        self.interest.remove(EventSet::readable());
                        self.interest.insert(EventSet::writable());
                        break;
                    }
                }
            }
        }
    }
}

struct Counter {
    value: usize,
}

impl Counter {
    fn new() -> Counter {
        Counter{value:0}
    }

    fn next(&mut self) -> Token {
        self.value += 1;
        return Token(self.value - 1);
    }
}

struct InternalReader {
    event_buffer: LinkedList<InternalMessage>,
    output_rx: mpsc::Receiver<InternalMessage>,
}

impl InternalReader {
    fn new(output_rx: mpsc::Receiver<InternalMessage>) -> InternalReader {
        InternalReader {
            event_buffer: LinkedList::new(),
            output_rx: output_rx,
        }
    }

    fn bread(&mut self) -> InternalMessage {
        if self.event_buffer.len() == 0 {
            match self.output_rx.recv() {
                Ok(m)  => self.event_buffer.push_back(m),
                Err(_) => panic!("whattodo"),
            }
        }
        return self.event_buffer.pop_front().unwrap();
    }

    fn read(&mut self) -> Option<InternalMessage> {
        match self.output_rx.try_recv() {
            Ok(m)  => self.event_buffer.push_back(m),
            Err(_) => {},
        }

        self.event_buffer.pop_front()
    }
}

struct InternalWriter {
    pipe_writer: unix::PipeWriter,
    input_tx: mpsc::Sender<InternalMessage>,
}

impl InternalWriter {
    fn new(pipe_writer: unix::PipeWriter,
           input_tx: mpsc::Sender<InternalMessage>) -> InternalWriter {
        InternalWriter {
            input_tx    : input_tx,
            pipe_writer : pipe_writer,
        }
    }

    fn write(&mut self, msg: InternalMessage) {
        let v = vec![1];
        self.pipe_writer.write(&v).unwrap();

        self.input_tx.send(msg).unwrap();
    }
}

struct WebSocketServer {
    counter: Counter,
    socket: TcpListener,
    clients: HashMap<Token, WebSocketClient>,
    input_rx: mpsc::Receiver<InternalMessage>,
    output_tx: mpsc::Sender<InternalMessage>,
    server_token: Token,
    pipe_token: Token,
    pipe_reader: unix::PipeReader,
}

impl WebSocketServer {
    fn new(_ip: &str, _port: u32) -> (WebSocketServer, InternalReader, InternalWriter) {
        let (p_reader, p_writer) = unix::pipe().unwrap();
        // i/o wrt to the event loop
        let (input_tx,  input_rx)  = mpsc::channel();
        let (output_tx, output_rx) = mpsc::channel();

        // FIXME: use ip + port
        let server_socket = TcpSocket::v4().unwrap();
        let address = FromStr::from_str("0.0.0.0:10000").unwrap();
        server_socket.bind(&address).unwrap();
        let server_socket = server_socket.listen(256).unwrap();

        let mut counter = Counter::new();

        (WebSocketServer {
            clients:        HashMap::new(),
            socket:         server_socket,
            input_rx:       input_rx,
            output_tx:      output_tx,
            server_token:   counter.next(),
            pipe_token:     counter.next(),
            counter:        counter,
            pipe_reader:    p_reader,
        },
        InternalReader::new(output_rx),
        InternalWriter::new(p_writer, input_tx))
    }

    fn start(&mut self) {
        let mut event_loop = EventLoop::new().unwrap();
        event_loop.register_opt(&self.socket,
                                self.server_token,
                                EventSet::readable(),
                                PollOpt::edge()).unwrap();
        event_loop.register_opt(&self.pipe_reader,
                                self.pipe_token,
                                EventSet::readable(),
                                PollOpt::edge()).unwrap();
        event_loop.run(self).unwrap();
    }
}

impl Handler for WebSocketServer {
    type Timeout = usize;
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<WebSocketServer>,
             token: Token, events: EventSet) {
        let (reregister_rtoken, reregister_wtoken) : (Option<Token>, Option<Token>);
        if events.is_readable() {
            println!("SERVER_TOKEN: {:?}", self.server_token);
            if token == self.server_token {
                let client_socket = match self.socket.accept() {
                    Ok(Some(sock)) => sock,
                    Ok(None) => unreachable!(),
                    Err(e) => {
                        println!("Accept error: {}", e);
                        return;
                    }
                };

                let new_token = self.counter.next();
                self.clients.insert(new_token, WebSocketClient::new(client_socket));

                event_loop.register_opt(&self.clients[&new_token].socket,
                                        new_token, EventSet::readable(),
                                        PollOpt::edge() | PollOpt::oneshot()).unwrap();
                self.output_tx.send(InternalMessage::NewClient{client_token: new_token}).unwrap();
                reregister_rtoken = None;
            } else if token == self.pipe_token {
                let mut buff = Vec::new();
                let size = self.pipe_reader.read(&mut buff);
                println!("READ: {:?} {:?}", size, buff);
                println!("PIPE_TOKEN: {:?}", self.pipe_token);
                match self.input_rx.recv().unwrap() {
                    InternalMessage::CloseClient{client_token: _token} => {
                        // FIXME: implement
                        reregister_rtoken = None;
                    },
                    InternalMessage::Data{client_token: token, data, format: _} => {
                        let client = self.clients.get_mut(&token).unwrap();
                        client.outgoing_messages.push_back(data);
                        if client.state == ClientState::Open {
                            client.writing_substate = SubState::Doing;
                            reregister_rtoken = Some(token);
                        } else {
                            reregister_rtoken = None;
                        }
                    },
                    _ => {
                        reregister_rtoken = None;
                    },
                }
            } else {
                let mut client = self.clients.get_mut(&token).unwrap();
                client.read();
                reregister_rtoken = Some(token);
            }
        } else {
            reregister_rtoken = None;
        }

        if events.is_writable() {
            if token == self.server_token || token == self.pipe_token {
                reregister_wtoken = None;
            } else {
                let mut client = self.clients.get_mut(&token).unwrap();
                client.write();
                reregister_wtoken = Some(token);
            }
        } else {
            reregister_wtoken = None;
        }

        let vtoken = vec![&reregister_rtoken, &reregister_wtoken];
        for v in vtoken {
            match v.clone() {
                Some(token) => {
                    let client = self.clients.get(&token).unwrap();
                    let mut interest = EventSet::none();
                    match client.state {
                        ClientState::AwaitingHandshake => unreachable!(),
                        ClientState::HandshakeResponse => interest.insert(EventSet::writable()),
                        ClientState::Open => {
                            interest.insert(EventSet::readable());
                            match client.reading_substate {
                                SubState::Waiting => {},
                                SubState::Doing   => {},
                            }

                            match client.writing_substate {
                                SubState::Waiting => {},
                                SubState::Doing   => interest.insert(EventSet::writable()),
                            }
                        }
                    }

                    event_loop.reregister(&client.socket, token, interest,
                                          PollOpt::edge() | PollOpt::oneshot()).unwrap();
                },
                None => {},
            }
        }
    }
}

fn main() {
    let (mut server, mut reader, mut writer) = WebSocketServer::new("0.0.0.0", 10000);

    // Receiver Process
    thread::spawn(move || {
        loop {
            match reader.bread() {
                InternalMessage::NewClient{client_token: token} => {
                    println!("New Client: {:?}", token);
                },
                InternalMessage::CloseClient{client_token: token} => {
                    println!("Close Client: {:?}", token);
                },
                InternalMessage::Data{client_token: _token, format: _format, data: _data} =>
                    {},
            }
        }
    });

    // Sender Process
    thread::spawn(move || {
        loop {
            loop {}
            let msg = InternalMessage::Data{client_token: Token(2),
                format: String::from("Aaron"),
                data:   String::from("Burrow"),
            };
            writer.write(msg);
        }
    });

    server.start();
}
<commit_msg>cleanup reregister handling code<commit_after>extern crate mio;
extern crate http_muncher;
extern crate sha1;
extern crate rustc_serialize;

use std::collections::HashMap;
use std::collections::LinkedList;
use std::cell::RefCell;
use std::rc::Rc;
use std::fmt;
use std::sync::mpsc;
use std::thread;
use std::str::FromStr;
use std::io::Read;
use std::io::Write;

use mio::*;
use mio::tcp::*;
use http_muncher::{Parser, ParserHandler};
use rustc_serialize::base64::{ToBase64, STANDARD};

enum InternalMessage {
    NewClient{client_token: Token},
    CloseClient{client_token: Token},
    Data{client_token: Token, format: String, data: String},
}

fn gen_key(key: &String) -> String {
    let mut m = sha1::Sha1::new();
    let mut buf = [0u8; 20];

    m.update(key.as_bytes());
    m.update("258EAFA5-E914-47DA-95CA-C5AB0DC85B11".as_bytes());

    m.output(&mut buf);

    return buf.to_base64(STANDARD);
}

struct HttpParser {
    current_key: Option<String>,
    headers: Rc<RefCell<HashMap<String, String>>>
}

impl ParserHandler for HttpParser {
    fn on_header_field(&mut self, s: &[u8]) -> bool {
        self.current_key = Some(std::str::from_utf8(s).unwrap().to_string());
        true
    }

    fn on_header_value(&mut self, s: &[u8]) -> bool {
        self.headers.borrow_mut()
            .insert(self.current_key.clone().unwrap(),
                    std::str::from_utf8(s).unwrap().to_string());
        true
    }

    fn on_headers_complete(&mut self) -> bool {
        false
    }
}

#[derive(PartialEq)]
enum ClientState {
    AwaitingHandshake,
    HandshakeResponse,
    Open,
}

enum SubState {
    Waiting,
    Doing,
}

struct WebSocketClient {
    socket: TcpStream,
    headers: Rc<RefCell<HashMap<String, String>>>,
    http_parser: Parser<HttpParser>,
    interest: EventSet,
    state: ClientState,
    writing_substate: SubState,
    reading_substate: SubState,
    outgoing_messages: std::collections::LinkedList<String>,
}

impl WebSocketClient {
    fn new(socket: TcpStream) -> WebSocketClient {
        let headers = Rc::new(RefCell::new(HashMap::new()));

        WebSocketClient {
            socket: socket,
            headers: headers.clone(),
            http_parser: Parser::request(HttpParser {
                current_key: None,
                headers: headers.clone()
            }),
            interest: EventSet::readable(),
            state: ClientState::AwaitingHandshake,
            writing_substate: SubState::Waiting,
            reading_substate: SubState::Waiting,
            outgoing_messages: std::collections::LinkedList::new(),
        }
    }

    fn handshake_response(&mut self) {
        let headers = self.headers.borrow();
        let response_key = gen_key(&headers.get("Sec-WebSocket-Key").unwrap());
        let response = fmt::format(format_args!("HTTP/1.1 101 Switching Protocols\r\n\
                                                 Connection: Upgrade\r\n\
                                                 Sec-WebSocket-Accept: {}\r\n\
                                                 Upgrade: websocket\r\n\r\n", response_key));
        self.socket.try_write(response.as_bytes()).unwrap();
    }

    fn read(&mut self) {
        loop {
            let mut buf = [0; 2048];
            match self.socket.try_read(&mut buf) {
                Err(e) => {
                    println!("Error while reading socket: {:?}", e);
                    return
                },
                Ok(None) =>
                    // Socket buffer has got no more bytes.
                    break,
                Ok(Some(_len)) => {
                    self.http_parser.parse(&buf);
                    if self.http_parser.is_upgrade() {
                        // Change the current state
                        self.state = ClientState::HandshakeResponse;

                        // Change current interest to `Writable`
                        self.interest.remove(EventSet::readable());
                        self.interest.insert(EventSet::writable());
                        break;
                    }
                }
            }
        }
    }
}

struct Counter {
    value: usize,
}

impl Counter {
    fn new() -> Counter {
        Counter{value:0}
    }

    fn next(&mut self) -> Token {
        self.value += 1;
        return Token(self.value - 1);
    }
}

struct InternalReader {
    event_buffer: LinkedList<InternalMessage>,
    output_rx: mpsc::Receiver<InternalMessage>,
}

impl InternalReader {
    fn new(output_rx: mpsc::Receiver<InternalMessage>) -> InternalReader {
        InternalReader {
            event_buffer: LinkedList::new(),
            output_rx: output_rx,
        }
    }

    fn bread(&mut self) -> InternalMessage {
        if self.event_buffer.len() == 0 {
            match self.output_rx.recv() {
                Ok(m)  => self.event_buffer.push_back(m),
                Err(_) => panic!("whattodo"),
            }
        }
        return self.event_buffer.pop_front().unwrap();
    }

    fn read(&mut self) -> Option<InternalMessage> {
        match self.output_rx.try_recv() {
            Ok(m)  => self.event_buffer.push_back(m),
            Err(_) => {},
        }

        self.event_buffer.pop_front()
    }
}

struct InternalWriter {
    pipe_writer: unix::PipeWriter,
    input_tx: mpsc::Sender<InternalMessage>,
}

impl InternalWriter {
    fn new(pipe_writer: unix::PipeWriter,
           input_tx: mpsc::Sender<InternalMessage>) -> InternalWriter {
        InternalWriter {
            input_tx    : input_tx,
            pipe_writer : pipe_writer,
        }
    }

    fn write(&mut self, msg: InternalMessage) {
        let v = vec![1];
        self.pipe_writer.write(&v).unwrap();

        self.input_tx.send(msg).unwrap();
    }
}

struct WebSocketServer {
    counter: Counter,
    socket: TcpListener,
    clients: HashMap<Token, WebSocketClient>,
    input_rx: mpsc::Receiver<InternalMessage>,
    output_tx: mpsc::Sender<InternalMessage>,
    server_token: Token,
    pipe_token: Token,
    pipe_reader: unix::PipeReader,
}

impl WebSocketServer {
    fn new(_ip: &str, _port: u32) -> (WebSocketServer, InternalReader, InternalWriter) {
        // i/o wrt to the event loop
        let (p_reader, p_writer)   = unix::pipe().unwrap();
        let (input_tx,  input_rx)  = mpsc::channel();
        let (output_tx, output_rx) = mpsc::channel();

        // FIXME: use ip + port
        let server_socket = TcpSocket::v4().unwrap();
        let address = FromStr::from_str("0.0.0.0:10000").unwrap();
        server_socket.bind(&address).unwrap();
        let server_socket = server_socket.listen(256).unwrap();

        let mut counter = Counter::new();

        (WebSocketServer {
            clients:        HashMap::new(),
            socket:         server_socket,
            input_rx:       input_rx,
            output_tx:      output_tx,
            server_token:   counter.next(),
            pipe_token:     counter.next(),
            counter:        counter,
            pipe_reader:    p_reader,
        },
        InternalReader::new(output_rx),
        InternalWriter::new(p_writer, input_tx))
    }

    fn start(&mut self) {
        let mut event_loop = EventLoop::new().unwrap();
        event_loop.register_opt(&self.socket,
                                self.server_token,
                                EventSet::readable(),
                                PollOpt::edge()).unwrap();
        event_loop.register_opt(&self.pipe_reader,
                                self.pipe_token,
                                EventSet::readable(),
                                PollOpt::edge()).unwrap();
        event_loop.run(self).unwrap();
    }
}

impl Handler for WebSocketServer {
    type Timeout = usize;
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<WebSocketServer>,
             token: Token, events: EventSet) {
        let mut reregister_token : Option<Token> = None;
        if events.is_readable() {
            println!("SERVER_TOKEN: {:?}", self.server_token);
            if token == self.server_token {
                let client_socket = match self.socket.accept() {
                    Ok(Some(sock)) => sock,
                    Ok(None) => unreachable!(),
                    Err(e) => {
                        println!("Accept error: {}", e);
                        return;
                    }
                };

                let new_token = self.counter.next();
                self.clients.insert(new_token, WebSocketClient::new(client_socket));

                event_loop.register_opt(&self.clients[&new_token].socket,
                                        new_token, EventSet::readable(),
                                        PollOpt::edge() | PollOpt::oneshot()).unwrap();
                self.output_tx.send(InternalMessage::NewClient{client_token: new_token}).unwrap();
            } else if token == self.pipe_token {
                let mut buff = Vec::new();
                let size = self.pipe_reader.read(&mut buff);
                println!("READ: {:?} {:?}", size, buff);
                println!("PIPE_TOKEN: {:?}", self.pipe_token);
                match self.input_rx.recv().unwrap() {
                    InternalMessage::CloseClient{client_token: _token} => {
                        // FIXME: implement
                    },
                    InternalMessage::Data{client_token: token, data, format: _} => {
                        let client = self.clients.get_mut(&token).unwrap();
                        client.outgoing_messages.push_back(data);
                        if client.state == ClientState::Open {
                            client.writing_substate = SubState::Doing;
                            reregister_token = Some(token);
                        }
                    },
                    _ => {},
                }
            } else {
                let mut client = self.clients.get_mut(&token).unwrap();
                client.read();
                reregister_token = Some(token);
            }
        }

        if events.is_writable() {
            if token != self.server_token && token != self.pipe_token {
                let mut client = self.clients.get_mut(&token).unwrap();
                match client.state {
                    ClientState::AwaitingHandshake => unreachable!(),
                    ClientState::HandshakeResponse => {
                        client.handshake_response();
                        client.state      = ClientState::Open;
                        reregister_token = Some(token);
                    },
                    ClientState::Open => {},
                }
            }
        }

        match reregister_token {
            Some(token) => {
                let client = self.clients.get(&token).unwrap();
                let mut interest = EventSet::none();
                match client.state {
                    ClientState::AwaitingHandshake => unreachable!(),
                    ClientState::HandshakeResponse => interest.insert(EventSet::writable()),
                    ClientState::Open => {
                        interest.insert(EventSet::readable());
                        match client.reading_substate {
                            SubState::Waiting => {},
                            SubState::Doing   => {},
                        }

                        match client.writing_substate {
                            SubState::Waiting => {},
                            SubState::Doing   => interest.insert(EventSet::writable()),
                        }
                    }
                }

                event_loop.reregister(&client.socket, token, interest,
                                      PollOpt::edge() | PollOpt::oneshot()).unwrap();
            },
            None => {},
        }
    }
}

fn main() {
    let (mut server, mut reader, mut writer) = WebSocketServer::new("0.0.0.0", 10000);

    // Receiver Process
    thread::spawn(move || {
        loop {
            match reader.bread() {
                InternalMessage::NewClient{client_token: token} => {
                    println!("New Client: {:?}", token);
                },
                InternalMessage::CloseClient{client_token: token} => {
                    println!("Close Client: {:?}", token);
                },
                InternalMessage::Data{client_token: _token, format: _format, data: _data} =>
                    {},
            }
        }
    });

    // Sender Process
    thread::spawn(move || {
        loop {
            loop {}
            let msg = InternalMessage::Data{client_token: Token(2),
                format: String::from("Aaron"),
                data:   String::from("Burrow"),
            };
            writer.write(msg);
        }
    });

    server.start();
}
<|endoftext|>


<commit_before>extern crate regex;
extern crate irc;
extern crate xdg_basedir;
mod github;

use regex::Regex;
use irc::client::prelude::*;
use irc::client::server::NetIrcServer;
use xdg_basedir::*;

fn parse_post(post: &str) -> Option<String> {
    let postregex = Regex::new(r"(\S+)/(\S+)#(\d+)").unwrap();
    let cap = match postregex.captures(post) {
        Some(x) => x,
        None => { return None; }
    };
    let user = cap.at(1).unwrap();
    let repo = cap.at(2).unwrap();
    let number = cap.at(3).unwrap();
    return github::get_display_text(user, repo, number).ok();
}

fn handle_message(server: &NetIrcServer, from: String, to: String, message: String) {
    let nickname = server.config().nickname.clone().unwrap();

    if to == nickname && from == "ids1024" {
        let mut words = message.split_whitespace();
        let command = words.next().unwrap_or("");
        let parameter = words.next().unwrap_or("");
        match command {
            "join" => {
                server.send_join(parameter).unwrap();
            },
            "part" => {
                server.send(Command::PART(parameter.to_string(), None)).unwrap();
            },
            "quit" => {
                server.send_quit("").unwrap();
            }
            _ => {},
        }
    }

    match parse_post(&message) {
        Some(x) => {
            for line in x.lines() {
                server.send_privmsg(&to, &line).unwrap();
            }
        },
        None => {}
    };
}

fn main() {
    let mut configpath = get_config_home().unwrap();
    configpath.push("idsbot/config.json");

    let config = Config::load(configpath).unwrap();
    let server = IrcServer::from_config(config).unwrap();
    server.identify().unwrap();

    for message in server.iter() {
        let message = message.unwrap();
        print!("{}", message.into_string());
        if &message.command[..] == "PRIVMSG" {
            let from = message.get_source_nickname().unwrap().to_owned();
            let to = message.args[0].to_owned();
            let content = message.suffix.unwrap();
            handle_message(&server, from, to, content);
        }
    }
}
<commit_msg>Use &str instead of String<commit_after>extern crate regex;
extern crate irc;
extern crate xdg_basedir;
mod github;

use regex::Regex;
use irc::client::prelude::*;
use irc::client::server::NetIrcServer;
use xdg_basedir::*;

fn parse_post(post: &str) -> Option<String> {
    let postregex = Regex::new(r"(\S+)/(\S+)#(\d+)").unwrap();
    let cap = match postregex.captures(post) {
        Some(x) => x,
        None => { return None; }
    };
    let user = cap.at(1).unwrap();
    let repo = cap.at(2).unwrap();
    let number = cap.at(3).unwrap();
    return github::get_display_text(user, repo, number).ok();
}

fn handle_message(server: &NetIrcServer, from: &str, to: &str, message: &str) {
    let nickname = server.config().nickname.as_ref().unwrap();

    if to == nickname && from == "ids1024" {
        let mut words = message.split_whitespace();
        let command = words.next().unwrap_or("");
        let parameter = words.next().unwrap_or("");
        match command {
            "join" => {
                server.send_join(parameter).unwrap();
            },
            "part" => {
                server.send(Command::PART(parameter.to_string(), None)).unwrap();
            },
            "quit" => {
                server.send_quit("").unwrap();
            }
            _ => {},
        }
    }

    match parse_post(&message) {
        Some(x) => {
            for line in x.lines() {
                server.send_privmsg(&to, &line).unwrap();
            }
        },
        None => {}
    };
}

fn main() {
    let mut configpath = get_config_home().unwrap();
    configpath.push("idsbot/config.json");

    let config = Config::load(configpath).unwrap();
    let server = IrcServer::from_config(config).unwrap();
    server.identify().unwrap();

    for message in server.iter() {
        let message = message.unwrap();
        print!("{}", message.into_string());
        if &message.command[..] == "PRIVMSG" {
            let from = message.get_source_nickname().unwrap().to_owned();
            let to = message.args[0].to_owned();
            let content = message.suffix.unwrap();
            handle_message(&server, &from, &to, &content);
        }
    }
}
<|endoftext|>