<commit_before>//! Utility functions to write nifti images.

use std::fs::File;
use std::io::BufWriter;
use std::ops::{Div, Sub};
use std::path::{Path, PathBuf};

use byteorder::{LittleEndian, WriteBytesExt};
use flate2::write::GzEncoder;
use flate2::Compression;
use ndarray::{ArrayBase, ArrayView, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;
use safe_transmute::{guarded_transmute_to_bytes_pod_many, PodTransmutable};

use {
    header::{MAGIC_CODE_NI1, MAGIC_CODE_NIP1},
    util::{is_gz_file, is_hdr_file},
    volume::element::DataElement,
    NiftiHeader, NiftiType, Result,
};

// TODO make this configurable. The Nifti standard does not specify a specific field for endianness,
// but it is encoded in `dim[0]`. "if dim[0] is outside range 1..7, then swap".
type B = LittleEndian;

/// Write a nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields. The voxels
/// intensity will be subtracted by `scl_slope` and divided by `scl_inter`. If `reference` is not
/// given, a default `NiftiHeader` will be built and written.
///
/// In all cases, the `dim`, `datatype` and `bitpix` fields will depend only on `data`, not on the
/// header. In other words, the `datatype` defined in `reference` will be ignored.
pub fn write_nifti<P, A, S, D>(
    header_path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = A>,
    A: Copy,
    A: DataElement,
    A: Div<Output = A>,
    A: FromPrimitive,
    A: PodTransmutable,
    A: ScalarOperand,
    A: Sub<Output = A>,
    D: Dimension + RemoveAxis,
{
    let compression_level = Compression::fast();
    let is_gz = is_gz_file(&header_path);
    let (header, data_path) = prepare_header_and_paths(&header_path, data, reference, A::DATA_TYPE);

    // Need the transpose for fortran ordering used in nifti file format.
    let data = data.t();

    let header_file = File::create(header_path)?;
    let mut header_writer = BufWriter::new(header_file);
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            write_data(&mut e, &header, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_data(&mut header_writer, &header, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        let mut data_writer = BufWriter::new(data_file);
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            let _ = e.finish()?;

            let mut e = GzEncoder::new(data_writer, compression_level);
            write_data(&mut e, &header, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_data(&mut data_writer, &header, data)?;
        }
    }

    Ok(())
}

/// Write a RGB nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields, except those
/// necessary to be recognized as a RGB image. `scl_slope` will be set to 1.0 and `scl_inter` to
/// 0.0.  If `reference` is not given, a default `NiftiHeader` will be built and written.
pub fn write_rgb_nifti<P, S, D>(
    header_path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = [u8; 3]>,
    D: Dimension + RemoveAxis,
{
    let compression_level = Compression::fast();
    let is_gz = is_gz_file(&header_path);
    let (mut header, data_path) =
        prepare_header_and_paths(&header_path, data, reference, NiftiType::Rgb24);

    // The `scl_slope` and `scl_inter` fields are ignored on the Rgb24 type.
    header.scl_slope = 1.0;
    header.scl_inter = 0.0;

    // Need the transpose for fortran used in nifti file format.
    let data = data.t();

    let header_file = File::create(header_path)?;
    let mut header_writer = BufWriter::new(header_file);
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            write_slices(&mut e, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_slices(&mut header_writer, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        let mut data_writer = BufWriter::new(data_file);
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            let _ = e.finish()?;

            let mut e = GzEncoder::new(data_writer, compression_level);
            write_slices(&mut e, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_slices(&mut data_writer, data)?;
        }
    }

    Ok(())
}

fn prepare_header_and_paths<P, T, D>(
    header_path: P,
    data: &ArrayBase<T, D>,
    reference: Option<&NiftiHeader>,
    datatype: NiftiType,
) -> (NiftiHeader, PathBuf)
where
    P: AsRef<Path>,
    T: Data,
    D: Dimension,
{
    let mut dim = [1; 8];
    dim[0] = data.ndim() as u16;
    for (i, s) in data.shape().iter().enumerate() {
        dim[i + 1] = *s as u16;
    }

    // If no reference header is given, use the default.
    let reference = match reference {
        Some(r) => r.clone(),
        None => {
            let mut header = NiftiHeader::default();
            header.pixdim = [1.0; 8];
            header.sform_code = 2;
            header.srow_x = [1.0, 0.0, 0.0, 0.0];
            header.srow_y = [0.0, 1.0, 0.0, 0.0];
            header.srow_z = [0.0, 0.0, 1.0, 0.0];
            header
        }
    };

    let mut header = NiftiHeader {
        dim,
        sizeof_hdr: 348,
        datatype: datatype as i16,
        bitpix: (datatype.size_of() * 8) as i16,
        vox_offset: 352.0,
        magic: *MAGIC_CODE_NIP1,
        // All other fields are copied from reference header
        ..reference
    };

    let mut path_buf = PathBuf::from(header_path.as_ref());
    let data_path = if is_hdr_file(&header_path) {
        header.vox_offset = 0.0;
        header.magic = *MAGIC_CODE_NI1;
        let data_path = if is_gz_file(&header_path) {
            let _ = path_buf.set_extension("");
            path_buf.with_extension("img.gz")
        } else {
            path_buf.with_extension("img")
        };
        data_path
    } else {
        path_buf
    };

    (header, data_path)
}

fn write_header<W>(writer: &mut W, header: &NiftiHeader) -> Result<()>
where
    W: WriteBytesExt,
{
    writer.write_i32::<B>(header.sizeof_hdr)?;
    writer.write_all(&header.data_type)?;
    writer.write_all(&header.db_name)?;
    writer.write_i32::<B>(header.extents)?;
    writer.write_i16::<B>(header.session_error)?;
    writer.write_u8(header.regular)?;
    writer.write_u8(header.dim_info)?;
    for s in &header.dim {
        writer.write_u16::<B>(*s)?;
    }
    writer.write_f32::<B>(header.intent_p1)?;
    writer.write_f32::<B>(header.intent_p2)?;
    writer.write_f32::<B>(header.intent_p3)?;
    writer.write_i16::<B>(header.intent_code)?;
    writer.write_i16::<B>(header.datatype)?;
    writer.write_i16::<B>(header.bitpix)?;
    writer.write_i16::<B>(header.slice_start)?;
    for f in &header.pixdim {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_f32::<B>(header.vox_offset)?;
    writer.write_f32::<B>(header.scl_slope)?;
    writer.write_f32::<B>(header.scl_inter)?;
    writer.write_i16::<B>(header.slice_end)?;
    writer.write_u8(header.slice_code)?;
    writer.write_u8(header.xyzt_units)?;
    writer.write_f32::<B>(header.cal_max)?;
    writer.write_f32::<B>(header.cal_min)?;
    writer.write_f32::<B>(header.slice_duration)?;
    writer.write_f32::<B>(header.toffset)?;
    writer.write_i32::<B>(header.glmax)?;
    writer.write_i32::<B>(header.glmin)?;

    writer.write_all(&header.descrip)?;
    writer.write_all(&header.aux_file)?;
    writer.write_i16::<B>(header.qform_code)?;
    writer.write_i16::<B>(header.sform_code)?;
    for f in &[
        header.quatern_b,
        header.quatern_c,
        header.quatern_d,
        header.quatern_x,
        header.quatern_y,
        header.quatern_z,
    ] {
        writer.write_f32::<B>(*f)?;
    }
    for f in header
        .srow_x
        .iter()
        .chain(&header.srow_y)
        .chain(&header.srow_z)
    {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_all(&header.intent_name)?;
    writer.write_all(&header.magic)?;

    // Empty 4 bytes after the header
    // TODO Support writing extension data.
    writer.write_u32::<B>(0)?;

    Ok(())
}

/// Write the data in 'f' order.
///
/// Like NiBabel, we iterate by "slice" to improve speed and use less memory.
fn write_data<T, D, W>(writer: &mut W, header: &NiftiHeader, data: ArrayView<T, D>) -> Result<()>
where
    T: Clone + PodTransmutable,
    T: Div<Output = T>,
    T: FromPrimitive,
    T: PodTransmutable,
    T: ScalarOperand,
    T: Sub<Output = T>,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    // `1.0x + 0.0` would give the same results, but we avoid a lot of divisions
    let slope = if header.scl_slope == 0.0 {
        1.0
    } else {
        header.scl_slope
    };
    if slope != 1.0 || header.scl_inter != 0.0 {
        // TODO Use linear transformation like when reading. An scl_slope of 0.5 would turn all
        // voxel values to 0 if we pass an ndarray of integers.
        let slope = T::from_f32(slope).unwrap();
        let inter = T::from_f32(header.scl_inter).unwrap();
        for arr_data in data.axis_iter(Axis(0)) {
            write_slice(writer, arr_data.sub(inter).div(slope))?;
        }
    } else {
        write_slices(writer, data)?;
    }
    Ok(())
}

fn write_slices<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    let mut iter = data.axis_iter(Axis(0));
    if let Some(arr_data) = iter.next() {
        // Keep slice voxels in a separate array to ensure `C` ordering even after `into_shape`.
        let mut slice = arr_data.to_owned();
        write_slice(writer, slice.view())?;
        for arr_data in iter {
            slice.assign(&arr_data);
            write_slice(writer, slice.view())?;
        }
    }
    Ok(())
}

fn write_slice<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension,
    W: WriteBytesExt,
{
    let len = data.len();
    let arr_data = data.into_shape(len).unwrap();
    let slice = arr_data.as_slice().unwrap();
    writer.write_all(guarded_transmute_to_bytes_pod_many(slice))?;
    Ok(())
}
<commit_msg>Do not give BufWriter to GzEncoder<commit_after>//! Utility functions to write nifti images.

use std::fs::File;
use std::io::BufWriter;
use std::ops::{Div, Sub};
use std::path::{Path, PathBuf};

use byteorder::{LittleEndian, WriteBytesExt};
use flate2::write::GzEncoder;
use flate2::Compression;
use ndarray::{ArrayBase, ArrayView, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;
use safe_transmute::{guarded_transmute_to_bytes_pod_many, PodTransmutable};

use {
    header::{MAGIC_CODE_NI1, MAGIC_CODE_NIP1},
    util::{is_gz_file, is_hdr_file},
    volume::element::DataElement,
    NiftiHeader, NiftiType, Result,
};

// TODO make this configurable. The Nifti standard does not specify a specific field for endianness,
// but it is encoded in `dim[0]`. "if dim[0] is outside range 1..7, then swap".
type B = LittleEndian;

/// Write a nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields. The voxels
/// intensity will be subtracted by `scl_slope` and divided by `scl_inter`. If `reference` is not
/// given, a default `NiftiHeader` will be built and written.
///
/// In all cases, the `dim`, `datatype` and `bitpix` fields will depend only on `data`, not on the
/// header. In other words, the `datatype` defined in `reference` will be ignored.
pub fn write_nifti<P, A, S, D>(
    header_path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = A>,
    A: Copy,
    A: DataElement,
    A: Div<Output = A>,
    A: FromPrimitive,
    A: PodTransmutable,
    A: ScalarOperand,
    A: Sub<Output = A>,
    D: Dimension + RemoveAxis,
{
    let compression_level = Compression::fast();
    let is_gz = is_gz_file(&header_path);
    let (header, data_path) = prepare_header_and_paths(&header_path, data, reference, A::DATA_TYPE);

    // Need the transpose for fortran ordering used in nifti file format.
    let data = data.t();

    let header_file = File::create(header_path)?;
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut e = GzEncoder::new(header_file, compression_level);
            write_header(&mut e, &header)?;
            write_data(&mut e, &header, data)?;
            let _ = e.finish()?;
        } else {
            let mut header_writer = BufWriter::new(header_file);
            write_header(&mut header_writer, &header)?;
            write_data(&mut header_writer, &header, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        if is_gz {
            let mut e = GzEncoder::new(header_file, compression_level);
            write_header(&mut e, &header)?;
            let _ = e.finish()?;

            let mut e = GzEncoder::new(data_file, compression_level);
            write_data(&mut e, &header, data)?;
            let _ = e.finish()?;
        } else {
            let mut header_writer = BufWriter::new(header_file);
            let mut data_writer = BufWriter::new(data_file);
            write_header(&mut header_writer, &header)?;
            write_data(&mut data_writer, &header, data)?;
        }
    }

    Ok(())
}

/// Write a RGB nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields, except those
/// necessary to be recognized as a RGB image. `scl_slope` will be set to 1.0 and `scl_inter` to
/// 0.0.  If `reference` is not given, a default `NiftiHeader` will be built and written.
pub fn write_rgb_nifti<P, S, D>(
    header_path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = [u8; 3]>,
    D: Dimension + RemoveAxis,
{
    let compression_level = Compression::fast();
    let is_gz = is_gz_file(&header_path);
    let (mut header, data_path) =
        prepare_header_and_paths(&header_path, data, reference, NiftiType::Rgb24);

    // The `scl_slope` and `scl_inter` fields are ignored on the Rgb24 type.
    header.scl_slope = 1.0;
    header.scl_inter = 0.0;

    // Need the transpose for fortran used in nifti file format.
    let data = data.t();

    let header_file = File::create(header_path)?;
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut e = GzEncoder::new(header_file, compression_level);
            write_header(&mut e, &header)?;
            write_slices(&mut e, data)?;
            let _ = e.finish()?;
        } else {
            let mut header_writer = BufWriter::new(header_file);
            write_header(&mut header_writer, &header)?;
            write_slices(&mut header_writer, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        if is_gz {
            let mut e = GzEncoder::new(header_file, compression_level);
            write_header(&mut e, &header)?;
            let _ = e.finish()?;

            let mut e = GzEncoder::new(data_file, compression_level);
            write_slices(&mut e, data)?;
            let _ = e.finish()?;
        } else {
            let mut header_writer = BufWriter::new(header_file);
            let mut data_writer = BufWriter::new(data_file);
            write_header(&mut header_writer, &header)?;
            write_slices(&mut data_writer, data)?;
        }
    }

    Ok(())
}

fn prepare_header_and_paths<P, T, D>(
    header_path: P,
    data: &ArrayBase<T, D>,
    reference: Option<&NiftiHeader>,
    datatype: NiftiType,
) -> (NiftiHeader, PathBuf)
where
    P: AsRef<Path>,
    T: Data,
    D: Dimension,
{
    let mut dim = [1; 8];
    dim[0] = data.ndim() as u16;
    for (i, s) in data.shape().iter().enumerate() {
        dim[i + 1] = *s as u16;
    }

    // If no reference header is given, use the default.
    let reference = match reference {
        Some(r) => r.clone(),
        None => {
            let mut header = NiftiHeader::default();
            header.pixdim = [1.0; 8];
            header.sform_code = 2;
            header.srow_x = [1.0, 0.0, 0.0, 0.0];
            header.srow_y = [0.0, 1.0, 0.0, 0.0];
            header.srow_z = [0.0, 0.0, 1.0, 0.0];
            header
        }
    };

    let mut header = NiftiHeader {
        dim,
        sizeof_hdr: 348,
        datatype: datatype as i16,
        bitpix: (datatype.size_of() * 8) as i16,
        vox_offset: 352.0,
        magic: *MAGIC_CODE_NIP1,
        // All other fields are copied from reference header
        ..reference
    };

    let mut path_buf = PathBuf::from(header_path.as_ref());
    let data_path = if is_hdr_file(&header_path) {
        header.vox_offset = 0.0;
        header.magic = *MAGIC_CODE_NI1;
        let data_path = if is_gz_file(&header_path) {
            let _ = path_buf.set_extension("");
            path_buf.with_extension("img.gz")
        } else {
            path_buf.with_extension("img")
        };
        data_path
    } else {
        path_buf
    };

    (header, data_path)
}

fn write_header<W>(writer: &mut W, header: &NiftiHeader) -> Result<()>
where
    W: WriteBytesExt,
{
    writer.write_i32::<B>(header.sizeof_hdr)?;
    writer.write_all(&header.data_type)?;
    writer.write_all(&header.db_name)?;
    writer.write_i32::<B>(header.extents)?;
    writer.write_i16::<B>(header.session_error)?;
    writer.write_u8(header.regular)?;
    writer.write_u8(header.dim_info)?;
    for s in &header.dim {
        writer.write_u16::<B>(*s)?;
    }
    writer.write_f32::<B>(header.intent_p1)?;
    writer.write_f32::<B>(header.intent_p2)?;
    writer.write_f32::<B>(header.intent_p3)?;
    writer.write_i16::<B>(header.intent_code)?;
    writer.write_i16::<B>(header.datatype)?;
    writer.write_i16::<B>(header.bitpix)?;
    writer.write_i16::<B>(header.slice_start)?;
    for f in &header.pixdim {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_f32::<B>(header.vox_offset)?;
    writer.write_f32::<B>(header.scl_slope)?;
    writer.write_f32::<B>(header.scl_inter)?;
    writer.write_i16::<B>(header.slice_end)?;
    writer.write_u8(header.slice_code)?;
    writer.write_u8(header.xyzt_units)?;
    writer.write_f32::<B>(header.cal_max)?;
    writer.write_f32::<B>(header.cal_min)?;
    writer.write_f32::<B>(header.slice_duration)?;
    writer.write_f32::<B>(header.toffset)?;
    writer.write_i32::<B>(header.glmax)?;
    writer.write_i32::<B>(header.glmin)?;

    writer.write_all(&header.descrip)?;
    writer.write_all(&header.aux_file)?;
    writer.write_i16::<B>(header.qform_code)?;
    writer.write_i16::<B>(header.sform_code)?;
    for f in &[
        header.quatern_b,
        header.quatern_c,
        header.quatern_d,
        header.quatern_x,
        header.quatern_y,
        header.quatern_z,
    ] {
        writer.write_f32::<B>(*f)?;
    }
    for f in header
        .srow_x
        .iter()
        .chain(&header.srow_y)
        .chain(&header.srow_z)
    {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_all(&header.intent_name)?;
    writer.write_all(&header.magic)?;

    // Empty 4 bytes after the header
    // TODO Support writing extension data.
    writer.write_u32::<B>(0)?;

    Ok(())
}

/// Write the data in 'f' order.
///
/// Like NiBabel, we iterate by "slice" to improve speed and use less memory.
fn write_data<T, D, W>(writer: &mut W, header: &NiftiHeader, data: ArrayView<T, D>) -> Result<()>
where
    T: Clone + PodTransmutable,
    T: Div<Output = T>,
    T: FromPrimitive,
    T: PodTransmutable,
    T: ScalarOperand,
    T: Sub<Output = T>,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    // `1.0x + 0.0` would give the same results, but we avoid a lot of divisions
    let slope = if header.scl_slope == 0.0 {
        1.0
    } else {
        header.scl_slope
    };
    if slope != 1.0 || header.scl_inter != 0.0 {
        // TODO Use linear transformation like when reading. An scl_slope of 0.5 would turn all
        // voxel values to 0 if we pass an ndarray of integers.
        let slope = T::from_f32(slope).unwrap();
        let inter = T::from_f32(header.scl_inter).unwrap();
        for arr_data in data.axis_iter(Axis(0)) {
            write_slice(writer, arr_data.sub(inter).div(slope))?;
        }
    } else {
        write_slices(writer, data)?;
    }
    Ok(())
}

fn write_slices<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    let mut iter = data.axis_iter(Axis(0));
    if let Some(arr_data) = iter.next() {
        // Keep slice voxels in a separate array to ensure `C` ordering even after `into_shape`.
        let mut slice = arr_data.to_owned();
        write_slice(writer, slice.view())?;
        for arr_data in iter {
            slice.assign(&arr_data);
            write_slice(writer, slice.view())?;
        }
    }
    Ok(())
}

fn write_slice<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension,
    W: WriteBytesExt,
{
    let len = data.len();
    let arr_data = data.into_shape(len).unwrap();
    let slice = arr_data.as_slice().unwrap();
    writer.write_all(guarded_transmute_to_bytes_pod_many(slice))?;
    Ok(())
}
<|endoftext|>


<commit_before>//! A module to manage protobuf serialization

use std::io::Write;

use errors::Result;
use message::MessageWrite;

use byteorder::WriteBytesExt;
use byteorder::LittleEndian as LE;

/// A struct to write protobuf messages
pub struct Writer<W> {
    inner: W,
}

impl<W: Write> Writer<W> {

    /// Creates a new `ProtobufWriter`
    pub fn new(w: W) -> Writer<W> {
        Writer { inner: w }
    }

    /// Writes a `varint` (compacted `u64`)
    pub fn write_varint(&mut self, mut v: u64) -> Result<()> {
        while v > 0x7F {
            self.inner.write_u8(((v as u8) & 0x7F) | 0x80)?;
            v >>= 7;
        }
        self.inner.write_u8(v as u8).map_err(|e| e.into())
    }

    /// Writes a tag, which represents both the field number and the wire type
    pub fn write_tag(&mut self, tag: u32) -> Result<()> {
        self.write_varint(tag as u64)
    }

    /// Writes a `int32` which is internally coded as a `varint`
    pub fn write_int32(&mut self, v: i32) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `int64` which is internally coded as a `varint`
    pub fn write_int64(&mut self, v: i64) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `uint32` which is internally coded as a `varint`
    pub fn write_uint32(&mut self, v: u32) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `uint64` which is internally coded as a `varint`
    pub fn write_uint64(&mut self, v: u64) -> Result<()> {
        self.write_varint(v)
    }

    /// Writes a `sint32` which is internally coded as a `varint`
    pub fn write_sint32(&mut self, v: i32) -> Result<()> {
        self.write_varint(((v << 1) ^ (v >> 31)) as u64)
    }

    /// Writes a `sint64` which is internally coded as a `varint`
    pub fn write_sint64(&mut self, v: i64) -> Result<()> {
        self.write_varint(((v << 1) ^ (v >> 63)) as u64)
    }

    /// Writes a `fixed64` which is little endian coded `u64`
    pub fn write_fixed64(&mut self, v: u64) -> Result<()> {
        self.inner.write_u64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `fixed32` which is little endian coded `u32`
    pub fn write_fixed32(&mut self, v: u32) -> Result<()> {
        self.inner.write_u32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `sfixed64` which is little endian coded `i64`
    pub fn write_sfixed64(&mut self, v: i64) -> Result<()> {
        self.inner.write_i64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `sfixed32` which is little endian coded `i32`
    pub fn write_sfixed32(&mut self, v: i32) -> Result<()> {
        self.inner.write_i32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `float`
    pub fn write_float(&mut self, v: f32) -> Result<()> {
        self.inner.write_f32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `double`
    pub fn write_double(&mut self, v: f64) -> Result<()> {
        self.inner.write_f64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `bool` 1 = true, 0 = false
    pub fn write_bool(&mut self, v: bool) -> Result<()> {
        self.write_varint(if v { 1 } else { 0 })
    }

    /// Writes an `enum` converting it to a `i32` first
    pub fn write_enum(&mut self, v: i32) -> Result<()> {
        self.write_int32(v)
    }

    /// Writes `bytes`: length first then the chunk of data
    pub fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> {
        self.write_varint(bytes.len() as u64)?;
        self.inner.write_all(bytes).map_err(|e| e.into())
    }

    /// Writes `string`: length first then the chunk of data
    pub fn write_string(&mut self, s: &str) -> Result<()> {
        self.write_bytes(s.as_bytes())
    }

    /// Writes packed repeated field: length first then the chunk of data
    pub fn write_packed_repeated_field<M, F, S>(&mut self, v: &[M], mut write: F, size: &S) -> Result<()>
        where F: FnMut(&mut Self, &M) -> Result<()>,
              S: Fn(&M) -> usize,
    {
        if v.is_empty() {
            return Ok(());
        }
        let len: usize = v.iter().map(|m| size(m)).sum();
        self.write_varint(len as u64)?;
        for m in v {
            write(self, m)?;
        }
        Ok(())
    }

    /// Writes packed repeated field when we know the size of items
    ///
    /// `item_size` is internally used to compute the total length
    /// As the length is fixed (and the same as rust internal representation, we can directly dump
    /// all data at once
    pub fn write_packed_fixed_size<M>(&mut self, v: &[M], item_size: usize) -> Result<()> {
        let len = v.len() * item_size;
        let bytes = unsafe { ::std::slice::from_raw_parts(v as *const [M] as *const M as *const u8, len) };
        self.write_bytes(bytes)
    }

    /// Writes a message which implements `MessageWrite`
    pub fn write_message<M: MessageWrite>(&mut self, m: &M) -> Result<()> {
        let len = m.get_size();
        self.write_varint(len as u64)?;
        m.write_message(self)
    }

    /// Writes tag then `int32`
    pub fn write_int32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `int64`
    pub fn write_int64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `uint32`
    pub fn write_uint32_with_tag(&mut self, tag: u32, v: u32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `uint64`
    pub fn write_uint64_with_tag(&mut self, tag: u32, v: u64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v)
    }

    /// Writes tag then `sint32`
    pub fn write_sint32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_sint32(v)
    }

    /// Writes tag then `sint64`
    pub fn write_sint64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_sint64(v)
    }

    /// Writes tag then `fixed64`
    pub fn write_fixed64_with_tag(&mut self, tag: u32, v: u64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_u64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `fixed32`
    pub fn write_fixed32_with_tag(&mut self, tag: u32, v: u32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_u32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `sfixed64`
    pub fn write_sfixed64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_i64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `sfixed32`
    pub fn write_sfixed32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_i32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `float`
    pub fn write_float_with_tag(&mut self, tag: u32, v: f32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_f32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `double`
    pub fn write_double_with_tag(&mut self, tag: u32, v: f64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_f64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `bool`
    pub fn write_bool_with_tag(&mut self, tag: u32, v: bool) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(if v { 1 } else { 0 })
    }

    /// Writes tag then `bytes`
    pub fn write_bytes_with_tag(&mut self, tag: u32, bytes: &[u8]) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(bytes.len() as u64)?;
        self.inner.write_all(bytes).map_err(|e| e.into())
    }

    /// Writes tag then `string`
    pub fn write_string_with_tag(&mut self, tag: u32, s: &str) -> Result<()> {
        self.write_tag(tag)?;
        self.write_bytes(s.as_bytes())
    }

    /// Writes tag then repeated field
    ///
    /// If array is empty, then do nothing (do not even write the tag)
    pub fn write_packed_repeated_field_with_tag<M, F, S>(&mut self,
                                                         tag: u32,
                                                         v: &[M],
                                                         mut write: F,
                                                         size: &S) -> Result<()>
        where F: FnMut(&mut Self, &M) -> Result<()>,
              S: Fn(&M) -> usize,
    {
        if v.is_empty() {
            return Ok(());
        }

        self.write_tag(tag)?;
        let len: usize = v.iter().map(|m| size(m)).sum();
        self.write_varint(len as u64)?;
        for m in v {
            write(self, m)?;
        }
        Ok(())
    }

    /// Writes tag then repeated field with fixed length item size
    ///
    /// If array is empty, then do nothing (do not even write the tag)
    pub fn write_packed_fixed_size_with_tag<M>(&mut self,
                                               tag: u32,
                                               v: &[M],
                                               item_size: usize) -> Result<()> {
        if v.is_empty() {
            return Ok(());
        }
        self.write_tag(tag)?;
        let len = v.len() * item_size;
        let bytes = unsafe { ::std::slice::from_raw_parts(v as *const [M] as *const M as *const u8, len) };
        self.write_bytes(bytes)
    }

    /// Writes tag then message
    pub fn write_message_with_tag<M: MessageWrite>(&mut self, tag: u32, m: &M) -> Result<()> {
        self.write_tag(tag)?;
        self.write_message(m)
    }

    /// Writes tag then enum
    pub fn write_enum_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_int32(v)
    }
}
<commit_msg>add example in writer documentation<commit_after>//! A module to manage protobuf serialization

use std::io::Write;

use errors::Result;
use message::MessageWrite;

use byteorder::WriteBytesExt;
use byteorder::LittleEndian as LE;

/// A struct to write protobuf messages
///
/// # Examples
///
/// ```rust
/// // an automatically generated module which is in a separate file in general
/// mod foo_bar {
///     # use quick_protobuf::{MessageWrite, Writer, Result};
///     # use std::borrow::Cow;
///     # use std::io::Write;
///     pub struct Foo<'a> { pub name: Option<Cow<'a, str>>, }
///     pub struct Bar { pub id: Option<u32> }
///     pub struct FooBar<'a> { pub foos: Vec<Foo<'a>>, pub bars: Vec<Bar>, }
///     impl<'a> MessageWrite for FooBar<'a> {
///         // implements
///         // fn get_size(&self) -> usize { ... }
///         // fn write_message<W: Write>(&self, r: &mut Writer<W>) -> Result<()> { ... }
///         # fn get_size(&self) -> usize { 0 }
///         # fn write_message<W: Write>(&self, _: &mut Writer<W>) -> Result<()> { Ok(()) }
///     }
/// }
///
/// // FooBar is a message generated from a proto file
/// // in parcicular it contains a `write_message` function
/// use foo_bar::{FooBar, Foo, Bar};
/// use std::fs::File;
/// use std::borrow::Cow;
/// use quick_protobuf::{MessageWrite, Writer};
///
/// fn main() {
///     // let's say we want to write to a file
///     let mut file: File;
///     # let mut file = Vec::new();
///     let mut writer = Writer::new(&mut file);
///
///     // manually generates a FooBar for the example
///     let foobar = FooBar {
///         foos: vec![Foo { name: Some(Cow::Borrowed("test!")) }, Foo { name: None }],
///         bars: vec![Bar { id: Some(43) }, Bar { id: None }],
///     };
///
///     // now using the generated module
///     writer.write_message(&foobar).expect("Cannot write FooBar");
/// }
/// ```
pub struct Writer<W> {
    inner: W,
}

impl<W: Write> Writer<W> {

    /// Creates a new `ProtobufWriter`
    pub fn new(w: W) -> Writer<W> {
        Writer { inner: w }
    }

    /// Writes a `varint` (compacted `u64`)
    pub fn write_varint(&mut self, mut v: u64) -> Result<()> {
        while v > 0x7F {
            self.inner.write_u8(((v as u8) & 0x7F) | 0x80)?;
            v >>= 7;
        }
        self.inner.write_u8(v as u8).map_err(|e| e.into())
    }

    /// Writes a tag, which represents both the field number and the wire type
    pub fn write_tag(&mut self, tag: u32) -> Result<()> {
        self.write_varint(tag as u64)
    }

    /// Writes a `int32` which is internally coded as a `varint`
    pub fn write_int32(&mut self, v: i32) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `int64` which is internally coded as a `varint`
    pub fn write_int64(&mut self, v: i64) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `uint32` which is internally coded as a `varint`
    pub fn write_uint32(&mut self, v: u32) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `uint64` which is internally coded as a `varint`
    pub fn write_uint64(&mut self, v: u64) -> Result<()> {
        self.write_varint(v)
    }

    /// Writes a `sint32` which is internally coded as a `varint`
    pub fn write_sint32(&mut self, v: i32) -> Result<()> {
        self.write_varint(((v << 1) ^ (v >> 31)) as u64)
    }

    /// Writes a `sint64` which is internally coded as a `varint`
    pub fn write_sint64(&mut self, v: i64) -> Result<()> {
        self.write_varint(((v << 1) ^ (v >> 63)) as u64)
    }

    /// Writes a `fixed64` which is little endian coded `u64`
    pub fn write_fixed64(&mut self, v: u64) -> Result<()> {
        self.inner.write_u64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `fixed32` which is little endian coded `u32`
    pub fn write_fixed32(&mut self, v: u32) -> Result<()> {
        self.inner.write_u32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `sfixed64` which is little endian coded `i64`
    pub fn write_sfixed64(&mut self, v: i64) -> Result<()> {
        self.inner.write_i64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `sfixed32` which is little endian coded `i32`
    pub fn write_sfixed32(&mut self, v: i32) -> Result<()> {
        self.inner.write_i32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `float`
    pub fn write_float(&mut self, v: f32) -> Result<()> {
        self.inner.write_f32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `double`
    pub fn write_double(&mut self, v: f64) -> Result<()> {
        self.inner.write_f64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `bool` 1 = true, 0 = false
    pub fn write_bool(&mut self, v: bool) -> Result<()> {
        self.write_varint(if v { 1 } else { 0 })
    }

    /// Writes an `enum` converting it to a `i32` first
    pub fn write_enum(&mut self, v: i32) -> Result<()> {
        self.write_int32(v)
    }

    /// Writes `bytes`: length first then the chunk of data
    pub fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> {
        self.write_varint(bytes.len() as u64)?;
        self.inner.write_all(bytes).map_err(|e| e.into())
    }

    /// Writes `string`: length first then the chunk of data
    pub fn write_string(&mut self, s: &str) -> Result<()> {
        self.write_bytes(s.as_bytes())
    }

    /// Writes packed repeated field: length first then the chunk of data
    pub fn write_packed_repeated_field<M, F, S>(&mut self, v: &[M], mut write: F, size: &S) -> Result<()>
        where F: FnMut(&mut Self, &M) -> Result<()>,
              S: Fn(&M) -> usize,
    {
        if v.is_empty() {
            return Ok(());
        }
        let len: usize = v.iter().map(|m| size(m)).sum();
        self.write_varint(len as u64)?;
        for m in v {
            write(self, m)?;
        }
        Ok(())
    }

    /// Writes packed repeated field when we know the size of items
    ///
    /// `item_size` is internally used to compute the total length
    /// As the length is fixed (and the same as rust internal representation, we can directly dump
    /// all data at once
    pub fn write_packed_fixed_size<M>(&mut self, v: &[M], item_size: usize) -> Result<()> {
        let len = v.len() * item_size;
        let bytes = unsafe { ::std::slice::from_raw_parts(v as *const [M] as *const M as *const u8, len) };
        self.write_bytes(bytes)
    }

    /// Writes a message which implements `MessageWrite`
    pub fn write_message<M: MessageWrite>(&mut self, m: &M) -> Result<()> {
        let len = m.get_size();
        self.write_varint(len as u64)?;
        m.write_message(self)
    }

    /// Writes tag then `int32`
    pub fn write_int32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `int64`
    pub fn write_int64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `uint32`
    pub fn write_uint32_with_tag(&mut self, tag: u32, v: u32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `uint64`
    pub fn write_uint64_with_tag(&mut self, tag: u32, v: u64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v)
    }

    /// Writes tag then `sint32`
    pub fn write_sint32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_sint32(v)
    }

    /// Writes tag then `sint64`
    pub fn write_sint64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_sint64(v)
    }

    /// Writes tag then `fixed64`
    pub fn write_fixed64_with_tag(&mut self, tag: u32, v: u64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_u64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `fixed32`
    pub fn write_fixed32_with_tag(&mut self, tag: u32, v: u32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_u32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `sfixed64`
    pub fn write_sfixed64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_i64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `sfixed32`
    pub fn write_sfixed32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_i32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `float`
    pub fn write_float_with_tag(&mut self, tag: u32, v: f32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_f32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `double`
    pub fn write_double_with_tag(&mut self, tag: u32, v: f64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_f64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `bool`
    pub fn write_bool_with_tag(&mut self, tag: u32, v: bool) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(if v { 1 } else { 0 })
    }

    /// Writes tag then `bytes`
    pub fn write_bytes_with_tag(&mut self, tag: u32, bytes: &[u8]) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(bytes.len() as u64)?;
        self.inner.write_all(bytes).map_err(|e| e.into())
    }

    /// Writes tag then `string`
    pub fn write_string_with_tag(&mut self, tag: u32, s: &str) -> Result<()> {
        self.write_tag(tag)?;
        self.write_bytes(s.as_bytes())
    }

    /// Writes tag then repeated field
    ///
    /// If array is empty, then do nothing (do not even write the tag)
    pub fn write_packed_repeated_field_with_tag<M, F, S>(&mut self,
                                                         tag: u32,
                                                         v: &[M],
                                                         mut write: F,
                                                         size: &S) -> Result<()>
        where F: FnMut(&mut Self, &M) -> Result<()>,
              S: Fn(&M) -> usize,
    {
        if v.is_empty() {
            return Ok(());
        }

        self.write_tag(tag)?;
        let len: usize = v.iter().map(|m| size(m)).sum();
        self.write_varint(len as u64)?;
        for m in v {
            write(self, m)?;
        }
        Ok(())
    }

    /// Writes tag then repeated field with fixed length item size
    ///
    /// If array is empty, then do nothing (do not even write the tag)
    pub fn write_packed_fixed_size_with_tag<M>(&mut self,
                                               tag: u32,
                                               v: &[M],
                                               item_size: usize) -> Result<()> {
        if v.is_empty() {
            return Ok(());
        }
        self.write_tag(tag)?;
        let len = v.len() * item_size;
        let bytes = unsafe { ::std::slice::from_raw_parts(v as *const [M] as *const M as *const u8, len) };
        self.write_bytes(bytes)
    }

    /// Writes tag then message
    pub fn write_message_with_tag<M: MessageWrite>(&mut self, tag: u32, m: &M) -> Result<()> {
        self.write_tag(tag)?;
        self.write_message(m)
    }

    /// Writes tag then enum
    pub fn write_enum_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_int32(v)
    }
}
<|endoftext|>


<commit_before>extern crate bip_handshake;
extern crate bip_peer;
extern crate bip_util;
extern crate chan;

use std::io::Write;
use std::sync::mpsc::{self, Sender, Receiver};

use bip_handshake::{Handshaker};
use bip_peer::disk::{ODiskMessage, IDiskMessage, DiskManagerAccess};
use bip_peer::protocol::{self, OProtocolMessage, OProtocolMessageKind};
use bip_peer::selector::{OSelectorMessage, OSelectorMessageKind};
use bip_peer::LayerRegistration;
use bip_peer::token::{Token};
use bip_util::send::TrySender;

struct MockDiskManager;
impl DiskManagerAccess for MockDiskManager {
    fn write_block(&self, _token: Token, _read_bytes: &[u8]) {
        unimplemented!()
    }

    fn read_block(&self, _token: Token, _write_bytes: &mut Write) {
        unimplemented!()
    }

    fn new_request_token(&mut self) -> Token {
        unimplemented!()
    }
}
impl TrySender<IDiskMessage> for MockDiskManager {
    fn try_send(&self, _msg: IDiskMessage) -> Option<IDiskMessage> {
        unimplemented!()
    }
}

struct MockDiskRegistration;
impl LayerRegistration<ODiskMessage, IDiskMessage> for MockDiskRegistration {
    type SS2 = MockDiskManager;

    fn register(&mut self, _send: Box<TrySender<ODiskMessage>>) -> MockDiskManager {
        MockDiskManager
    }
}

struct MockSelectionRegistration {
    send: Sender<OProtocolMessage>,
}
impl LayerRegistration<OSelectorMessage, OProtocolMessage> for MockSelectionRegistration {
    type SS2 = Sender<OProtocolMessage>;

    fn register(&mut self, _send: Box<TrySender<OSelectorMessage>>) -> Sender<OProtocolMessage> {
        self.send.clone()
    }
}

fn main() {
    let (send, _recv): (Sender<()>, Receiver<()>) = mpsc::channel();

    let (protocol_send, protocol_recv) = mpsc::channel();
    let selection_registration = MockSelectionRegistration{ send: protocol_send };
    let disk_registration = MockDiskRegistration;

    let pid = ['-' as u8, 'U' as u8, 'T' as u8, '5' as u8, '5' as u8, '5' as u8, '5' as u8, '-' as u8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0].into();

    let mut handshaker: () = protocol::spawn_tcp_handshaker(send, "0.0.0.0:0".parse().unwrap(), pid, disk_registration, selection_registration).unwrap();

    let hash = [0xf8, 0x8e, 0xd0, 0xc1, 0x6c, 0xf7, 0xf4, 0x52, 0xa5, 0xc7, 0x37, 0xa0, 0xb7, 0x50, 0x3f, 0x92, 0x5e, 0x11, 0xfe, 0x00].into();
    handshaker.register(hash);

    handshaker.connect(None, hash, "10.0.0.18:55194".parse().unwrap());

    let (id, kind) = protocol_recv.recv().unwrap().destroy();
    let protocol_send = match kind {
        OProtocolMessageKind::PeerConnect(peer_send, _) => peer_send,
        _ => panic!("1")
    };

    //thread::sleep(Duration::from_millis(2000));

    protocol_send.try_send(OSelectorMessage::new(id, OSelectorMessageKind::PeerKeepAlive));

    loop {
        match protocol_recv.recv().unwrap().destroy().1 {
            OProtocolMessageKind::PeerChoke => println!("PeerChoke"),
            OProtocolMessageKind::PeerUnChoke => println!("PeerUnChoke"),
            OProtocolMessageKind::PeerInterested => println!("PeerInterested"),
            OProtocolMessageKind::PeerUnInterested => println!("PeerUnInterested"),
            OProtocolMessageKind::PeerHave(..) => println!("PeerHave"),
            OProtocolMessageKind::PeerBitField(b) => println!("PeerBitField {:?}", b),
            OProtocolMessageKind::PeerRequest(..) => println!("PeerRequest"),
            OProtocolMessageKind::PeerPiece(..) => println!("PeerPiece"),
            OProtocolMessageKind::PeerCancel(..) => println!("PeerCancel"),
            OProtocolMessageKind::PeerDisconnect => panic!("PeerDisconnect"),
            _ => panic!("ASD")
        }
    }
}<commit_msg>Fix Type Check<commit_after>extern crate bip_handshake;
extern crate bip_peer;
extern crate bip_util;
extern crate chan;

use std::io::Write;
use std::sync::mpsc::{self, Sender, Receiver};

use bip_handshake::{Handshaker};
use bip_peer::disk::{ODiskMessage, IDiskMessage, DiskManagerAccess};
use bip_peer::protocol::{self, OProtocolMessage, OProtocolMessageKind};
use bip_peer::selector::{OSelectorMessage, OSelectorMessageKind};
use bip_peer::LayerRegistration;
use bip_peer::token::{Token};
use bip_util::send::TrySender;

struct MockDiskManager;
impl DiskManagerAccess for MockDiskManager {
    fn write_block(&self, _token: Token, _read_bytes: &[u8]) {
        unimplemented!()
    }

    fn read_block(&self, _token: Token, _write_bytes: &mut Write) {
        unimplemented!()
    }

    fn new_request_token(&mut self) -> Token {
        unimplemented!()
    }
}
impl TrySender<IDiskMessage> for MockDiskManager {
    fn try_send(&self, _msg: IDiskMessage) -> Option<IDiskMessage> {
        unimplemented!()
    }
}

struct MockDiskRegistration;
impl LayerRegistration<ODiskMessage, IDiskMessage> for MockDiskRegistration {
    type SS2 = MockDiskManager;

    fn register(&mut self, _send: Box<TrySender<ODiskMessage>>) -> MockDiskManager {
        MockDiskManager
    }
}

struct MockSelectionRegistration {
    send: Sender<OProtocolMessage>,
}
impl LayerRegistration<OSelectorMessage, OProtocolMessage> for MockSelectionRegistration {
    type SS2 = Sender<OProtocolMessage>;

    fn register(&mut self, _send: Box<TrySender<OSelectorMessage>>) -> Sender<OProtocolMessage> {
        self.send.clone()
    }
}

fn main() {
    let (send, _recv): (Sender<()>, Receiver<()>) = mpsc::channel();

    let (protocol_send, protocol_recv) = mpsc::channel();
    let selection_registration = MockSelectionRegistration{ send: protocol_send };
    let disk_registration = MockDiskRegistration;

    let pid = ['-' as u8, 'U' as u8, 'T' as u8, '5' as u8, '5' as u8, '5' as u8, '5' as u8, '-' as u8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0].into();

    let mut handshaker = protocol::spawn_tcp_handshaker(send, "0.0.0.0:0".parse().unwrap(), pid, disk_registration, selection_registration).unwrap();

    let hash = [0xf8, 0x8e, 0xd0, 0xc1, 0x6c, 0xf7, 0xf4, 0x52, 0xa5, 0xc7, 0x37, 0xa0, 0xb7, 0x50, 0x3f, 0x92, 0x5e, 0x11, 0xfe, 0x00].into();
    handshaker.register(hash);

    handshaker.connect(None, hash, "10.0.0.18:55194".parse().unwrap());

    let (id, kind) = protocol_recv.recv().unwrap().destroy();
    let protocol_send = match kind {
        OProtocolMessageKind::PeerConnect(peer_send, _) => peer_send,
        _ => panic!("1")
    };

    //thread::sleep(Duration::from_millis(2000));

    protocol_send.try_send(OSelectorMessage::new(id, OSelectorMessageKind::PeerKeepAlive));

    loop {
        match protocol_recv.recv().unwrap().destroy().1 {
            OProtocolMessageKind::PeerChoke => println!("PeerChoke"),
            OProtocolMessageKind::PeerUnChoke => println!("PeerUnChoke"),
            OProtocolMessageKind::PeerInterested => println!("PeerInterested"),
            OProtocolMessageKind::PeerUnInterested => println!("PeerUnInterested"),
            OProtocolMessageKind::PeerHave(..) => println!("PeerHave"),
            OProtocolMessageKind::PeerBitField(b) => println!("PeerBitField {:?}", b),
            OProtocolMessageKind::PeerRequest(..) => println!("PeerRequest"),
            OProtocolMessageKind::PeerPiece(..) => println!("PeerPiece"),
            OProtocolMessageKind::PeerCancel(..) => println!("PeerCancel"),
            OProtocolMessageKind::PeerDisconnect => panic!("PeerDisconnect"),
            _ => panic!("ASD")
        }
    }
}<|endoftext|>


<commit_before>use crc32;
use constants::{CR, DEFAULT_LINE_SIZE, ESCAPE, LF, NUL};

use std::path::Path;
use std::fs::File;
use std::io;
use std::io::{Read, Seek, SeekFrom, Write};

/// Options for encoding
#[derive(Debug)]
pub struct EncodeOptions {
    line_length: u8,
    parts: u32,
    part: u32,
    begin: u64,
    end: u64,
}

impl Default for EncodeOptions {
    /// Constructs a new EncodeOptions instance, with the following defaults:
    /// line_length = 128.
    /// parts = 1,
    /// part = begin = end = 0
    fn default() -> Self {
        EncodeOptions {
            line_length: DEFAULT_LINE_SIZE,
            parts: 1,
            part: 0,
            begin: 0,
            end: 0,
        }
    }
}

impl EncodeOptions {
    /// Sets the maximum line length.
    pub fn line_length(mut self, line_length: u8) -> EncodeOptions {
        self.line_length = line_length;
        self
    }

    /// Sets the number of parts (default=1).
    /// When the number of parts equals 1, no '=ypart' line will be written
    /// in the ouput.
    pub fn parts(mut self, parts: u32) -> EncodeOptions {
        self.parts = parts;
        self
    }

    /// Sets the part number.
    /// Only used when `parts > 1`.
    /// The part number count starts at 1.
    pub fn part(mut self, part: u32) -> EncodeOptions {
        self.part = part;
        self
    }

    /// Sets the begin (which is the file offset + 1).
    /// Only used when `parts > 1`.
    /// The size of the part is `end - begin + 1`.
    pub fn begin(mut self, begin: u64) -> EncodeOptions {
        self.begin = begin;
        self
    }

    /// Sets the end.
    /// Only used when `parts > 1`.
    /// The size of the part is `end - begin + 1`.
    /// `end` should be larger than `begin`, otherwise an overflow error occurrs.
    pub fn end(mut self, end: u64) -> EncodeOptions {
        self.end = end;
        self
    }
}

/// Encodes the input file in a new output file.
/// # Example
/// ```rust,no_run
/// let encode_options = yenc::EncodeOptions::default().parts(1);
/// let mut output_file = std::fs::File::create("test1.bin.yenc").unwrap();
/// yenc::encode_file("test1.bin", &encode_options, &mut output_file);
/// ```
/// # Errors
/// - when the output file already exists
///
pub fn encode_file<P: AsRef<Path>>(
    input_path: P,
    encode_options: &EncodeOptions,
    output: &mut Write,
) -> Result<(), io::Error> {
    let mut checksum = crc32::Crc32::new();
    let mut buffer = [0u8; 8192];
    let mut col = 0;
    let input_filename = input_path.as_ref().file_name();
    let input_filename = match input_filename {
        Some(s) => s.to_str().unwrap_or(""),
        None => "",
    };
    let mut input_file = File::open(&input_path)?;

    if encode_options.parts == 1 {
        output.write_all(
            format!(
                "=ybegin line={} size={} name={}\r\n",
                encode_options.line_length,
                input_file.metadata()?.len(),
                input_filename
            ).as_bytes(),
        )?;
    } else {
        output.write_all(
            format!(
                "=ybegin part={} line={} size={} name={}\r\n",
                encode_options.part,
                encode_options.line_length,
                input_file.metadata()?.len(),
                input_filename
            ).as_bytes(),
        )?;
    }

    if encode_options.parts > 1 {
        output.write_all(
            format!(
                "=ypart begin={} end={}\r\n",
                encode_options.begin, encode_options.end
            ).as_bytes(),
        )?;
    }

    input_file.seek(SeekFrom::Start(encode_options.begin - 1))?;

    let mut remainder = (encode_options.end - encode_options.begin + 1) as usize;
    while remainder > 0 {
        let bytes_to_read = if remainder > buffer.len() {
            buffer.len()
        } else {
            remainder
        };
        input_file.read_exact(&mut buffer[0..bytes_to_read])?;
        checksum.update_with_slice(&buffer[0..bytes_to_read]);
        encode_buffer(
            &buffer[0..bytes_to_read],
            &mut col,
            encode_options.line_length,
            output,
        );
        remainder -= bytes_to_read;
    }

    if encode_options.parts > 1 {
        output.write_all(
            format!(
                "\r\n=yend size={} part={} pcrc32={:08x}\r\n",
                checksum.num_bytes, encode_options.part, checksum.crc
            ).as_bytes(),
        )?;
    } else {
        output.write_all(
            format!(
                "\r\n=yend size={} crc32={:08x}\r\n",
                checksum.num_bytes, checksum.crc
            ).as_bytes(),
        )?;
    }
    Ok(())
}

/// Encode the byte slice into a vector of yEncoded bytes.
///
/// Lines are wrapped with a maximum of `line_length` characters per line.
/// Does not include the header and footer lines. These are only produced
/// by `encode_stream` and `encode_file`.
pub fn encode_buffer(input: &[u8], col: &mut u8, line_length: u8, writer: &mut Write) {
    for &b in input {
        let (encoded, encoded_len) = encode_byte(b);
        writer.write_all(&encoded[0..encoded_len]).unwrap();
        *col += encoded_len as u8;
        if *col >= line_length {
            writer.write_all(&[CR, LF]).unwrap();
            *col = 0;
        }
    }
}

#[inline]
fn encode_byte(input_byte: u8) -> ([u8; 2], usize) {
    let mut output = [0u8; 2];
    let mut idx = 0;
    let mut output_byte = input_byte.overflowing_add(42).0;
    match output_byte {
        NUL | CR | LF | ESCAPE => {
            output[idx] = ESCAPE;
            idx += 1;
            output_byte = output_byte.overflowing_add(64).0;
        }
        _ => {}
    };
    output[idx] = output_byte;
    (output, idx + 1)
}

#[cfg(test)]
mod tests {
    use super::{encode_buffer, encode_byte};
    use super::super::constants::{CR, ESCAPE, LF};

    #[test]
    fn escape_null() {
        assert_eq!(([ESCAPE, 0x40], 2), encode_byte(214));
    }

    /*
    #[test]
    fn escape_tab() {
        let mut output = [0u8; 2];
        assert_eq!(2, encode_byte(214 + TAB, &mut output));
        assert_eq!(vec![ESCAPE, 0x49], output);
    }
    */

    #[test]
    fn escape_lf() {
        assert_eq!(([ESCAPE, 0x4A], 2), encode_byte(214 + LF));
    }

    #[test]
    fn escape_cr() {
        assert_eq!(([ESCAPE, 0x4D], 2), encode_byte(214 + CR));
    }

    /*
    #[test]
    fn escape_space() {
        let mut output = [0u8; 2];
        assert_eq!(2, encode_byte(214 + SPACE, &mut output));
        assert_eq!(vec![ESCAPE, 0x60], output);
    }
    */

    #[test]
    fn escape_equal_sign() {
        assert_eq!(([ESCAPE, 0x7D], 2), encode_byte(ESCAPE - 42));
    }

    #[test]
    fn non_escaped() {
        assert_eq!(([42, 0], 1), encode_byte(0));
    }
}
<commit_msg>Improve encoding performance<commit_after>use constants::{CR, DEFAULT_LINE_SIZE, ESCAPE, LF, NUL};
use crc32;

use std::fs::File;
use std::io;
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::Path;

/// Options for encoding
#[derive(Debug)]
pub struct EncodeOptions {
    line_length: u8,
    parts: u32,
    part: u32,
    begin: u64,
    end: u64,
}

impl Default for EncodeOptions {
    /// Constructs a new EncodeOptions instance, with the following defaults:
    /// line_length = 128.
    /// parts = 1,
    /// part = begin = end = 0
    fn default() -> Self {
        EncodeOptions {
            line_length: DEFAULT_LINE_SIZE,
            parts: 1,
            part: 0,
            begin: 0,
            end: 0,
        }
    }
}

impl EncodeOptions {
    /// Sets the maximum line length.
    pub fn line_length(mut self, line_length: u8) -> EncodeOptions {
        self.line_length = line_length;
        self
    }

    /// Sets the number of parts (default=1).
    /// When the number of parts equals 1, no '=ypart' line will be written
    /// in the ouput.
    pub fn parts(mut self, parts: u32) -> EncodeOptions {
        self.parts = parts;
        self
    }

    /// Sets the part number.
    /// Only used when `parts > 1`.
    /// The part number count starts at 1.
    pub fn part(mut self, part: u32) -> EncodeOptions {
        self.part = part;
        self
    }

    /// Sets the begin (which is the file offset + 1).
    /// Only used when `parts > 1`.
    /// The size of the part is `end - begin + 1`.
    pub fn begin(mut self, begin: u64) -> EncodeOptions {
        self.begin = begin;
        self
    }

    /// Sets the end.
    /// Only used when `parts > 1`.
    /// The size of the part is `end - begin + 1`.
    /// `end` should be larger than `begin`, otherwise an overflow error occurrs.
    pub fn end(mut self, end: u64) -> EncodeOptions {
        self.end = end;
        self
    }
}

/// Encodes the input file in a new output file.
/// # Example
/// ```rust,no_run
/// let encode_options = yenc::EncodeOptions::default().parts(1);
/// let mut output_file = std::fs::File::create("test1.bin.yenc").unwrap();
/// yenc::encode_file("test1.bin", &encode_options, &mut output_file);
/// ```
/// # Errors
/// - when the output file already exists
///
pub fn encode_file<P: AsRef<Path>>(
    input_path: P,
    encode_options: &EncodeOptions,
    output: &mut Write,
) -> Result<(), io::Error> {
    let mut checksum = crc32::Crc32::new();
    let mut buffer = [0u8; 8192];
    let mut col = 0;
    let input_filename = input_path.as_ref().file_name();
    let input_filename = match input_filename {
        Some(s) => s.to_str().unwrap_or(""),
        None => "",
    };
    let mut input_file = File::open(&input_path)?;

    if encode_options.parts == 1 {
        output.write_all(
            format!(
                "=ybegin line={} size={} name={}\r\n",
                encode_options.line_length,
                input_file.metadata()?.len(),
                input_filename
            ).as_bytes(),
        )?;
    } else {
        output.write_all(
            format!(
                "=ybegin part={} line={} size={} name={}\r\n",
                encode_options.part,
                encode_options.line_length,
                input_file.metadata()?.len(),
                input_filename
            ).as_bytes(),
        )?;
    }

    if encode_options.parts > 1 {
        output.write_all(
            format!(
                "=ypart begin={} end={}\r\n",
                encode_options.begin, encode_options.end
            ).as_bytes(),
        )?;
    }

    input_file.seek(SeekFrom::Start(encode_options.begin - 1))?;

    let mut remainder = (encode_options.end - encode_options.begin + 1) as usize;
    while remainder > 0 {
        let bytes_to_read = if remainder > buffer.len() {
            buffer.len()
        } else {
            remainder
        };
        input_file.read_exact(&mut buffer[0..bytes_to_read])?;
        checksum.update_with_slice(&buffer[0..bytes_to_read]);
        encode_buffer(
            &buffer[0..bytes_to_read],
            &mut col,
            encode_options.line_length,
            output,
        );
        remainder -= bytes_to_read;
    }

    if encode_options.parts > 1 {
        output.write_all(
            format!(
                "\r\n=yend size={} part={} pcrc32={:08x}\r\n",
                checksum.num_bytes, encode_options.part, checksum.crc
            ).as_bytes(),
        )?;
    } else {
        output.write_all(
            format!(
                "\r\n=yend size={} crc32={:08x}\r\n",
                checksum.num_bytes, checksum.crc
            ).as_bytes(),
        )?;
    }
    Ok(())
}

/// Encode the byte slice into a vector of yEncoded bytes.
///
/// Lines are wrapped with a maximum of `line_length` characters per line.
/// Does not include the header and footer lines. These are only produced
/// by `encode_stream` and `encode_file`.
pub fn encode_buffer(input: &[u8], col: &mut u8, line_length: u8, writer: &mut Write) {
    for &b in input {
        let (encoded, encoded_len) = encode_byte(b);
        writer.write_all(&encoded[0..encoded_len]).unwrap();
        *col += encoded_len as u8;
        if *col >= line_length {
            writer.write_all(&[CR, LF]).unwrap();
            *col = 0;
        }
    }
}

#[inline]
fn encode_byte(input_byte: u8) -> ([u8; 2], usize) {
    let mut output = (0, 0);

    let output_byte = input_byte.overflowing_add(42).0;
    let len = match output_byte {
        NUL | CR | LF | ESCAPE => {
            output.0 = ESCAPE;
            output.1 = output_byte.overflowing_add(64).0;
            2
        }
        _ => {
            output.0 = output_byte;
            1
        }
    };
    let output_array = [output.0, output.1];
    (output_array, len)
}

#[cfg(test)]
mod tests {
    use super::super::constants::{CR, ESCAPE, LF};
    use super::{encode_buffer, encode_byte};

    #[test]
    fn escape_null() {
        assert_eq!(([ESCAPE, 0x40], 2), encode_byte(214));
    }

    /*
    #[test]
    fn escape_tab() {
        let mut output = [0u8; 2];
        assert_eq!(2, encode_byte(214 + TAB, &mut output));
        assert_eq!(vec![ESCAPE, 0x49], output);
    }
    */

    #[test]
    fn escape_lf() {
        assert_eq!(([ESCAPE, 0x4A], 2), encode_byte(214 + LF));
    }

    #[test]
    fn escape_cr() {
        assert_eq!(([ESCAPE, 0x4D], 2), encode_byte(214 + CR));
    }

    /*
    #[test]
    fn escape_space() {
        let mut output = [0u8; 2];
        assert_eq!(2, encode_byte(214 + SPACE, &mut output));
        assert_eq!(vec![ESCAPE, 0x60], output);
    }
    */

    #[test]
    fn escape_equal_sign() {
        assert_eq!(([ESCAPE, 0x7D], 2), encode_byte(ESCAPE - 42));
    }

    #[test]
    fn non_escaped() {
        assert_eq!(([42, 0], 1), encode_byte(0));
    }
}
<|endoftext|>


<commit_before>// Third party
use reqwest::Error as HttpError;
use reqwest::StatusCode;
use serde_json::error::Error as SerdeError;
use std::io::Error as IoError;

// Ours
use crate::Errors;

/// an enumeration over potential errors
/// that may happen when sending a request to jira
#[derive(Debug)]
pub enum Error {
    /// error associated with http request
    Http(HttpError),
    /// error associated IO
    IO(IoError),
    /// error associated with parsing or serializing
    Serde(SerdeError),
    /// client request errors
    Fault { code: StatusCode, errors: Errors },
    /// invalid credentials
    Unauthorized,
    /// HTTP method is not allowed
    MethodNotAllowed,
    /// Page not found
    NotFound,
}

impl From<SerdeError> for Error {
    fn from(error: SerdeError) -> Error {
        Error::Serde(error)
    }
}

impl From<HttpError> for Error {
    fn from(error: HttpError) -> Error {
        Error::Http(error)
    }
}

impl From<IoError> for Error {
    fn from(error: IoError) -> Error {
        Error::IO(error)
    }
}

impl ::std::fmt::Display for Error {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        use crate::Error::*;

        match *self {
            Http(ref e) => writeln!(f, "Http Error: {}", e),
            IO(ref e) => writeln!(f, "IO Error: {}", e),
            Serde(ref e) => writeln!(f, "Serialization Error: {}", e),
            Fault {
                ref code,
                ref errors,
            } => writeln!(f, "Jira Client Error ({}):\n{:#?}", code, errors),
            _ => writeln!(f, "Could not connect to Jira: {}!", self),
        }
    }
}

impl ::std::error::Error for Error {
    fn description(&self) -> &str {
        use crate::Error::*;

        match *self {
            Http(ref e) => e.description(),
            IO(ref e) => e.description(),
            Serde(ref e) => e.description(),
            Fault { .. } => "Jira client error",
            Unauthorized => "Unauthorized",
            MethodNotAllowed => "MethodNotAllowed",
            NotFound => "NotFound",
        }
    }

    fn cause(&self) -> Option<&::std::error::Error> {
        use crate::Error::*;

        match *self {
            Http(ref e) => Some(e),
            IO(ref e) => Some(e),
            Serde(ref e) => Some(e),
            Fault { .. } => None,
            _ => None,
        }
    }
}
<commit_msg>Fixed errors in stable<commit_after>// Third party
use reqwest::Error as HttpError;
use reqwest::StatusCode;
use serde_json::error::Error as SerdeError;
use std::io::Error as IoError;

// Ours
use crate::Errors;

/// an enumeration over potential errors
/// that may happen when sending a request to jira
#[derive(Debug)]
pub enum Error {
    /// error associated with http request
    Http(HttpError),
    /// error associated IO
    IO(IoError),
    /// error associated with parsing or serializing
    Serde(SerdeError),
    /// client request errors
    Fault { code: StatusCode, errors: Errors },
    /// invalid credentials
    Unauthorized,
    /// HTTP method is not allowed
    MethodNotAllowed,
    /// Page not found
    NotFound,
}

impl From<SerdeError> for Error {
    fn from(error: SerdeError) -> Error {
        Error::Serde(error)
    }
}

impl From<HttpError> for Error {
    fn from(error: HttpError) -> Error {
        Error::Http(error)
    }
}

impl From<IoError> for Error {
    fn from(error: IoError) -> Error {
        Error::IO(error)
    }
}

impl ::std::fmt::Display for Error {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        use crate::Error::*;

        match *self {
            Http(ref e) => writeln!(f, "Http Error: {}", e),
            IO(ref e) => writeln!(f, "IO Error: {}", e),
            Serde(ref e) => writeln!(f, "Serialization Error: {}", e),
            Fault {
                ref code,
                ref errors,
            } => writeln!(f, "Jira Client Error ({}):\n{:#?}", code, errors),
            _ => writeln!(f, "Could not connect to Jira: {}!", self),
        }
    }
}

impl ::std::error::Error for Error {
    fn description(&self) -> &str {
        use crate::Error::*;

        match *self {
            Http(ref e) => e.description(),
            IO(ref e) => e.description(),
            Serde(ref e) => e.description(),
            Fault { .. } => "Jira client error",
            Unauthorized => "Unauthorized",
            MethodNotAllowed => "MethodNotAllowed",
            NotFound => "NotFound",
        }
    }

    fn cause(&self) -> Option<&dyn (::std::error::Error)> {
        use crate::Error::*;

        match *self {
            Http(ref e) => Some(e),
            IO(ref e) => Some(e),
            Serde(ref e) => Some(e),
            Fault { .. } => None,
            _ => None,
        }
    }
}
<|endoftext|>


<commit_before>//! String escaping.
//!
//! Although filenames in Linux are commonly represented as UTF-8
//! sequences, there is no system requirement that this be the case.  As a
//! consequence, this means that it is possible for filenames in Linux to
//! not be valid UTF-8, and therefore not representable as strings.
//!
//! To prevent encoding problems, as well as to allow certain characters,
//! such as space, to separate tokens in the sure file format, we escape
//! some bytes in strings by replacing them with "=xx" where "xx" is the
//! lower-cased hex version of the string.  The range of valid characters
//! is fairly straightforward, including all of the printable characters
//! from '!' to '~' except for the '=', which is always escaped.  This
//! means, for example, that a 2-byte encoded UTF-8 sequence will expand to
//! take 6 bytes.

use std::{io::prelude::*, result};
use thiserror::Error;

pub trait Escape {
    fn escaped(&self) -> String;
}

pub trait Unescape {
    fn unescape(&self) -> EscapeResult<Vec<u8>>;
}

pub type EscapeResult<T> = result::Result<T, EscapeError>;

#[derive(Error, Debug)]
pub enum EscapeError {
    #[error("Invalid hex character: {0:?}")]
    InvalidHexCharacter(u8),
    #[error("Invalid hex length")]
    InvalidHexLength,
}

// The basic encoding converts a sequence of bytes into a string.
impl Escape for [u8] {
    fn escaped(&self) -> String {
        let mut result = vec![];
        for &ch in self.iter() {
            // TODO: Can be made more efficient.
            if (b'!'..=b'~').contains(&ch) && ch != b'=' && ch != b'[' && ch != b']' {
                result.push(ch);
            } else {
                write!(&mut result, "={:02x}", ch).unwrap();
            }
        }

        // TODO: String::from_utf8_unchecked(result)
        String::from_utf8(result).unwrap()
    }
}

impl Unescape for str {
    fn unescape(&self) -> EscapeResult<Vec<u8>> {
        // Will overestimate.
        let mut buf = Vec::with_capacity(self.len() / 2);
        let mut phase = 0;
        let mut tmp = 0;

        for byte in self.bytes() {
            if phase == 0 {
                if byte == b'=' {
                    phase = 1;
                } else {
                    buf.push(byte);
                }
            } else {
                tmp <<= 4;
                match byte {
                    b'A'..=b'F' => tmp |= byte - b'A' + 10,
                    b'a'..=b'f' => tmp |= byte - b'a' + 10,
                    b'0'..=b'f' => tmp |= byte - b'0',
                    _ => return Err(EscapeError::InvalidHexCharacter(byte)),
                }
                phase += 1;
                if phase == 3 {
                    buf.push(tmp);
                    phase = 0;
                    tmp = 0;
                }
            }
        }

        if phase != 0 {
            return Err(EscapeError::InvalidHexLength);
        }

        Ok(buf)
    }
}

#[test]
fn test_unescape() {
    macro_rules! assert_error_kind {
        ( $expr:expr, $kind:pat ) => {
            match $expr {
                Err($kind) => (),
                Err(e) => panic!(
                    "Unexpected error kind: {:?} (want {})",
                    e,
                    stringify!($kind)
                ),
                Ok(_) => panic!("Unexpected success"),
            }
        };
    }

    assert_eq!("=00".unescape().unwrap(), vec![0]);
    assert_error_kind!("=00=0".unescape(), EscapeError::InvalidHexLength);
    assert_error_kind!("=00=".unescape(), EscapeError::InvalidHexLength);
    assert_error_kind!("=4g".unescape(), EscapeError::InvalidHexCharacter(b'g'));
}

#[test]
fn test_escape() {
    let buf: Vec<u8> = (0u32..256).map(|i| i as u8).collect();
    let text = (&buf[..]).escaped();
    assert_eq!(text.unescape().unwrap(), buf);
}
<commit_msg>clippy: Fix an overlapping match range<commit_after>//! String escaping.
//!
//! Although filenames in Linux are commonly represented as UTF-8
//! sequences, there is no system requirement that this be the case.  As a
//! consequence, this means that it is possible for filenames in Linux to
//! not be valid UTF-8, and therefore not representable as strings.
//!
//! To prevent encoding problems, as well as to allow certain characters,
//! such as space, to separate tokens in the sure file format, we escape
//! some bytes in strings by replacing them with "=xx" where "xx" is the
//! lower-cased hex version of the string.  The range of valid characters
//! is fairly straightforward, including all of the printable characters
//! from '!' to '~' except for the '=', which is always escaped.  This
//! means, for example, that a 2-byte encoded UTF-8 sequence will expand to
//! take 6 bytes.

use std::{io::prelude::*, result};
use thiserror::Error;

pub trait Escape {
    fn escaped(&self) -> String;
}

pub trait Unescape {
    fn unescape(&self) -> EscapeResult<Vec<u8>>;
}

pub type EscapeResult<T> = result::Result<T, EscapeError>;

#[derive(Error, Debug)]
pub enum EscapeError {
    #[error("Invalid hex character: {0:?}")]
    InvalidHexCharacter(u8),
    #[error("Invalid hex length")]
    InvalidHexLength,
}

// The basic encoding converts a sequence of bytes into a string.
impl Escape for [u8] {
    fn escaped(&self) -> String {
        let mut result = vec![];
        for &ch in self.iter() {
            // TODO: Can be made more efficient.
            if (b'!'..=b'~').contains(&ch) && ch != b'=' && ch != b'[' && ch != b']' {
                result.push(ch);
            } else {
                write!(&mut result, "={:02x}", ch).unwrap();
            }
        }

        // TODO: String::from_utf8_unchecked(result)
        String::from_utf8(result).unwrap()
    }
}

impl Unescape for str {
    fn unescape(&self) -> EscapeResult<Vec<u8>> {
        // Will overestimate.
        let mut buf = Vec::with_capacity(self.len() / 2);
        let mut phase = 0;
        let mut tmp = 0;

        for byte in self.bytes() {
            if phase == 0 {
                if byte == b'=' {
                    phase = 1;
                } else {
                    buf.push(byte);
                }
            } else {
                tmp <<= 4;
                match byte {
                    b'A'..=b'F' => tmp |= byte - b'A' + 10,
                    b'a'..=b'f' => tmp |= byte - b'a' + 10,
                    b'0'..=b'9' => tmp |= byte - b'0',
                    _ => return Err(EscapeError::InvalidHexCharacter(byte)),
                }
                phase += 1;
                if phase == 3 {
                    buf.push(tmp);
                    phase = 0;
                    tmp = 0;
                }
            }
        }

        if phase != 0 {
            return Err(EscapeError::InvalidHexLength);
        }

        Ok(buf)
    }
}

#[test]
fn test_unescape() {
    macro_rules! assert_error_kind {
        ( $expr:expr, $kind:pat ) => {
            match $expr {
                Err($kind) => (),
                Err(e) => panic!(
                    "Unexpected error kind: {:?} (want {})",
                    e,
                    stringify!($kind)
                ),
                Ok(_) => panic!("Unexpected success"),
            }
        };
    }

    assert_eq!("=00".unescape().unwrap(), vec![0]);
    assert_error_kind!("=00=0".unescape(), EscapeError::InvalidHexLength);
    assert_error_kind!("=00=".unescape(), EscapeError::InvalidHexLength);
    assert_error_kind!("=4g".unescape(), EscapeError::InvalidHexCharacter(b'g'));
}

#[test]
fn test_escape() {
    let buf: Vec<u8> = (0u32..256).map(|i| i as u8).collect();
    let text = (&buf[..]).escaped();
    assert_eq!(text.unescape().unwrap(), buf);
}
<|endoftext|>


<commit_before>extern crate crypto;
extern crate reqwest;

use std::fs::{self, File};
use std::io::prelude::*;
use std::path::Path;

use crypto::md5::Md5;
use crypto::digest::Digest;

use reqwest::{Client, StatusCode};

#[derive(Debug)]
struct LoLVersion {
    region: &'static str,
    version: &'static str,
    desc: &'static str,
}

static VERSIONS: &'static [LoLVersion] = &[
    LoLVersion { region: "EUW", version: "2016_11_10", desc: "EU West" },
    LoLVersion { region: "NA", version: "2016_05_13", desc: "North America" },
    LoLVersion { region: "EUNE", version: "2016_11_10", desc: "EU Nordic & East" },
    LoLVersion { region: "OC1", version: "2016_05_13", desc: "Oceania" },
    LoLVersion { region: "RU", version: "2016_05_13", desc: "Russia" },
    LoLVersion { region: "LA1", version: "2016_05_26", desc: "Latin America North" },
    LoLVersion { region: "LA2", version: "2016_05_27", desc: "Latin America South" },
    LoLVersion { region: "BR", version: "2016_05_13", desc: "Brazil" },
];

fn main() {
    let cache_path = Path::new("filecache");

    if !cache_path.exists() {
        fs::create_dir(cache_path).expect("Could not create cache directory");
    }

    let mut found = vec![];
    let mut hash = Md5::new();
    let client = Client::new().expect("Could not create http client");

    for v in VERSIONS {
        let url = format!("https://riotgamespatcher-a.akamaihd.net/ShellInstaller/{region}/LeagueofLegends_{region}_Installer_{version}.exe",
            region = v.region, version = v.version);

        println!("Calculating hash for url: {}", url);

        let exec = cache_path.join(format!("LeagueofLegends_{}_Installer_{}.exe", v.region, v.version));
        let mut buf = vec![];

        if exec.exists() {
            if File::open(&exec).and_then(|mut f| f.read_to_end(&mut buf)).is_err() {
                println!("Warning: Could not read from installer: {:?}", exec);
                continue;
            }
        } else {
            match client.get(&url).send() {
                // We are only interested in a successful response.
                Ok(ref mut r) if *r.status() == StatusCode::Ok => {
                    if r.read_to_end(&mut buf).is_err() {
                        println!("Warning: Could not read from downloaded installer: {:?}", exec);
                        continue;
                    }

                    if File::create(&exec).and_then(|mut f| f.write_all(&buf)).is_err() {
                        println!("Warning: Could not save downloaded installer: {:?}", exec);
                        continue;
                    }
                }
                _ => {
                    println!("Warning: Could not download installer: {}", url);
                    continue;
                }
            }
        }

        hash.input(&buf);
        let hash_str = hash.result_str();
        hash.reset();

        println!("{}", hash_str);

        found.push((hash_str, url, v));
    }

    // Create string outputs for hashes, urls, regions, and descriptions
    let (mut hashes, mut urls, mut regions, mut desc) = found.iter()
        .fold((String::new(), String::new(), String::new(), String::new()), |s, &(ref hash, ref url, v)| {
            ((s.0 + "[\"" + v.desc + "\"]=\"" + hash + "\" "),
             (s.1 + "[\"" + v.desc + "\"]=\"" + url + "\" "),
             (s.2 + "[\"" + v.desc + "\"]=\"" + v.region + "\" "),
             (s.3 + v.desc + "|"))
        });

    // For now we just pop the extra last character off of each string because of the fold.
    hashes.pop(); urls.pop(); regions.pop(); desc.pop();

    println!("declare -A HASHES=( {} )", hashes);
    println!("declare -A URLS=( {} )", urls);
    println!("declare -A REGIONS=( {} )", regions);
    println!("REGION_DESC=\"{}\"", desc);
}<commit_msg>Go directly to strings.<commit_after>extern crate crypto;
extern crate reqwest;

use std::fs::{self, File};
use std::io::prelude::*;
use std::path::Path;

use crypto::md5::Md5;
use crypto::digest::Digest;

use reqwest::{Client, StatusCode};

#[derive(Debug)]
struct LoLVersion {
    region: &'static str,
    version: &'static str,
    desc: &'static str,
}

static VERSIONS: &'static [LoLVersion] = &[
    LoLVersion { region: "EUW", version: "2016_11_10", desc: "EU West" },
    LoLVersion { region: "NA", version: "2016_05_13", desc: "North America" },
    LoLVersion { region: "EUNE", version: "2016_11_10", desc: "EU Nordic & East" },
    LoLVersion { region: "OC1", version: "2016_05_13", desc: "Oceania" },
    LoLVersion { region: "RU", version: "2016_05_13", desc: "Russia" },
    LoLVersion { region: "LA1", version: "2016_05_26", desc: "Latin America North" },
    LoLVersion { region: "LA2", version: "2016_05_27", desc: "Latin America South" },
    LoLVersion { region: "BR", version: "2016_05_13", desc: "Brazil" },
];

fn main() {
    let cache_path = Path::new("filecache");

    if !cache_path.exists() {
        fs::create_dir(cache_path).expect("Could not create cache directory");
    }

    let client = Client::new().expect("Could not create http client");
    let mut hash = Md5::new();

    let (mut hashes, mut urls, mut regions, mut desc) = VERSIONS.iter()
        .filter_map(|v| {
            let url = format!("https://riotgamespatcher-a.akamaihd.net/ShellInstaller/{region}/LeagueofLegends_{region}_Installer_{version}.exe",
                region = v.region, version = v.version);

            println!("Calculating hash for: {}", v.desc);

            let exec = cache_path.join(format!("LeagueofLegends_{}_Installer_{}.exe", v.region, v.version));
            let mut buf = vec![];

            if exec.exists() {
                // If the installer already exists then we just read from it.
                if File::open(&exec).and_then(|mut f| f.read_to_end(&mut buf)).is_err() {
                    println!("Warning: Could not read from installer: {:?}", exec);
                    return None;
                }
            } else {
                // If not we attempt to download and cache the installer.
                if match client.get(&url).send() {
                    Ok(ref mut r) if *r.status() == StatusCode::Ok => {
                        r.read_to_end(&mut buf)
                            .and_then(|_| File::create(&exec))
                            .and_then(|mut f| f.write_all(&buf))
                            .ok()
                    }
                    _ => Some(()),
                }.is_none() {
                    println!("Warning: Could not download or save installer: {:?}", exec);
                    return None;
                }
            }

            hash.input(&buf);
            let hash_str = hash.result_str();
            println!("{}", &hash_str);
            hash.reset();

            Some((hash_str, url, v))
        })
        .fold((String::new(), String::new(), String::new(), String::new()), |s, (ref hash, ref url, v)| {
            ((s.0 + "[\"" + v.desc + "\"]=\"" + hash + "\" "),
             (s.1 + "[\"" + v.desc + "\"]=\"" + url + "\" "),
             (s.2 + "[\"" + v.desc + "\"]=\"" + v.region + "\" "),
             (s.3 + v.desc + "|"))
        });

    // For now we just pop the extra last character off of each string because of the fold.
    hashes.pop(); urls.pop(); regions.pop(); desc.pop();

    println!("declare -A HASHES=( {} )", hashes);
    println!("declare -A URLS=( {} )", urls);
    println!("declare -A REGIONS=( {} )", regions);
    println!("REGION_DESC=\"{}\"", desc);
}<|endoftext|>