<commit_before>//! Extra streaming decompression functionality.
//!
//! As of now this is mainly inteded for use to build a higher-level wrapper.
use std::io::Cursor;
use std::{cmp, mem};

use crate::inflate::core::{decompress, inflate_flags, DecompressorOxide, TINFL_LZ_DICT_SIZE};
use crate::inflate::TINFLStatus;
use crate::{DataFormat, MZError, MZFlush, MZResult, MZStatus, StreamResult};

/// A struct that compbines a decompressor with extra data for streaming decompression.
///
pub struct InflateState {
    /// Inner decompressor struct
    decomp: DecompressorOxide,

    /// Buffer of input bytes for matches.
    /// TODO: Could probably do this a bit cleaner with some
    /// Cursor-like class.
    /// We may also look into whether we need to keep a buffer here, or just one in the
    /// decompressor struct.
    dict: [u8; TINFL_LZ_DICT_SIZE],
    /// Where in the buffer are we currently at?
    dict_ofs: usize,
    /// How many bytes of data to be flushed is there currently in the buffer?
    dict_avail: usize,

    first_call: bool,
    has_flushed: bool,

    /// Whether the input data is wrapped in a zlib header and checksum.
    /// TODO: This should be stored in the decompressor.
    data_format: DataFormat,
    last_status: TINFLStatus,
}

impl Default for InflateState {
    fn default() -> Self {
        InflateState {
            decomp: DecompressorOxide::default(),
            dict: [0; TINFL_LZ_DICT_SIZE],
            dict_ofs: 0,
            dict_avail: 0,
            first_call: true,
            has_flushed: false,
            data_format: DataFormat::Raw,
            last_status: TINFLStatus::NeedsMoreInput,
        }
    }
}
impl InflateState {
    /// Create a new state.
    ///
    /// Note that this struct is quite large due to internal buffers, and as such storing it on
    /// the stack is not recommended.
    ///
    /// # Parameters
    /// `data_format`: Determines whether the compressed data is assumed to wrapped with zlib
    /// metadata.
    pub fn new(data_format: DataFormat) -> InflateState {
        let mut b = InflateState::default();
        b.data_format = data_format;
        b
    }

    /// Create a new state on the heap.
    ///
    /// # Parameters
    /// `data_format`: Determines whether the compressed data is assumed to wrapped with zlib
    /// metadata.
    pub fn new_boxed(data_format: DataFormat) -> Box<InflateState> {
        let mut b: Box<InflateState> = Box::default();
        b.data_format = data_format;
        b
    }

    /// Access the innner decompressor.
    pub fn decompressor(&mut self) -> &mut DecompressorOxide {
        &mut self.decomp
    }

    /// Return the status of the last call to `inflate` with this `InflateState`.
    pub fn last_status(&self) -> TINFLStatus {
        self.last_status
    }

    /// Create a new state using miniz/zlib style window bits parameter.
    ///
    /// The decompressor does not support different window sizes. As such,
    /// any positive (>0) value will set the zlib header flag, while a negative one
    /// will not.
    pub fn new_boxed_with_window_bits(window_bits: i32) -> Box<InflateState> {
        let mut b: Box<InflateState> = Box::default();
        b.data_format = DataFormat::from_window_bits(window_bits);
        b
    }

    /// Reset the decompressor without re-allocating memory, using the given
    /// data format.
    pub fn reset(&mut self, data_format: DataFormat) {
        self.decompressor().init();
        self.dict = [0; TINFL_LZ_DICT_SIZE];
        self.dict_ofs = 0;
        self.dict_avail = 0;
        self.first_call = true;
        self.has_flushed = false;
        self.data_format = data_format;
        self.last_status = TINFLStatus::NeedsMoreInput;
    }
}

/// Try to decompress from `input` to `output` with the given `InflateState`
///
/// # Errors
///
/// Returns `MZError::Buf` If the size of the `output` slice is empty or no progress was made due to
/// lack of expected input data or called after the decompression was
/// finished without MZFlush::Finish.
///
/// Returns `MZError::Param` if the compressor parameters are set wrong.
pub fn inflate(
    state: &mut InflateState,
    input: &[u8],
    output: &mut [u8],
    flush: MZFlush,
) -> StreamResult {
    let mut bytes_consumed = 0;
    let mut bytes_written = 0;
    let mut next_in = input;
    let mut next_out = output;

    if flush == MZFlush::Full {
        return StreamResult::error(MZError::Stream);
    }

    let mut decomp_flags = inflate_flags::TINFL_FLAG_COMPUTE_ADLER32;
    if state.data_format == DataFormat::Zlib {
        decomp_flags |= inflate_flags::TINFL_FLAG_PARSE_ZLIB_HEADER;
    }

    let first_call = state.first_call;
    state.first_call = false;
    if (state.last_status as i32) < 0 {
        return StreamResult::error(MZError::Data);
    }

    if state.has_flushed && (flush != MZFlush::Finish) {
        return StreamResult::error(MZError::Stream);
    }
    state.has_flushed |= flush == MZFlush::Finish;

    if (flush == MZFlush::Finish) && first_call {
        decomp_flags |= inflate_flags::TINFL_FLAG_USING_NON_WRAPPING_OUTPUT_BUF;

        let status = decompress(
            &mut state.decomp,
            next_in,
            &mut Cursor::new(next_out),
            decomp_flags,
        );
        let in_bytes = status.1;
        let out_bytes = status.2;
        let status = status.0;

        state.last_status = status;

        bytes_consumed += in_bytes;
        bytes_written += out_bytes;

        let ret_status = {
            if (status as i32) < 0 {
                println!("Status: {:?}", status);
                Err(MZError::Data)
            } else if status != TINFLStatus::Done {
                state.last_status = TINFLStatus::Failed;
                Err(MZError::Buf)
            } else {
                Ok(MZStatus::StreamEnd)
            }
        };
        return StreamResult {
            bytes_consumed,
            bytes_written,
            status: ret_status,
        };
    }

    if flush != MZFlush::Finish {
        decomp_flags |= inflate_flags::TINFL_FLAG_HAS_MORE_INPUT;
    }

    if state.dict_avail != 0 {
        bytes_written += push_dict_out(state, &mut next_out);
        return StreamResult {
            bytes_consumed,
            bytes_written,
            status: Ok(
                if (state.last_status == TINFLStatus::Done) && (state.dict_avail == 0) {
                    MZStatus::StreamEnd
                } else {
                    MZStatus::Ok
                },
            ),
        };
    }

    let status = inflate_loop(
        state,
        &mut next_in,
        &mut next_out,
        &mut bytes_consumed,
        &mut bytes_written,
        decomp_flags,
        flush,
    );
    StreamResult {
        bytes_consumed,
        bytes_written,
        status,
    }
}

fn inflate_loop(
    state: &mut InflateState,
    next_in: &mut &[u8],
    next_out: &mut &mut [u8],
    total_in: &mut usize,
    total_out: &mut usize,
    decomp_flags: u32,
    flush: MZFlush,
) -> MZResult {
    let orig_in_len = next_in.len();
    loop {
        let status = {
            let mut cursor = Cursor::new(&mut state.dict[..]);
            cursor.set_position(state.dict_ofs as u64);
            decompress(&mut state.decomp, *next_in, &mut cursor, decomp_flags)
        };

        let in_bytes = status.1;
        let out_bytes = status.2;
        let status = status.0;

        state.last_status = status;

        *next_in = &next_in[in_bytes..];
        *total_in += in_bytes;

        state.dict_avail = out_bytes;
        *total_out += push_dict_out(state, next_out);

        // The stream was corrupted, and decompression failed.
        if (status as i32) < 0 {
            return Err(MZError::Data);
        }

        // The decompressor has flushed all it's data and is waiting for more input, but
        // there was no more input provided.
        if (status == TINFLStatus::NeedsMoreInput) && orig_in_len == 0 {
            return Err(MZError::Buf);
        }

        if flush == MZFlush::Finish {
            if status == TINFLStatus::Done {
                // There is not enough space in the output buffer to flush the remaining
                // decompressed data in the internal buffer.
                return if state.dict_avail != 0 {
                    Err(MZError::Buf)
                } else {
                    Ok(MZStatus::StreamEnd)
                };
            // No more space in the output buffer, but we're not done.
            } else if next_out.is_empty() {
                return Err(MZError::Buf);
            }
        } else {
            // We're not expected to finish, so it's fine if we can't flush everything yet.
            let empty_buf = next_in.is_empty() || next_out.is_empty();
            if (status == TINFLStatus::Done) || empty_buf || (state.dict_avail != 0) {
                return if (status == TINFLStatus::Done) && (state.dict_avail == 0) {
                    // No more data left, we're done.
                    Ok(MZStatus::StreamEnd)
                } else {
                    // Ok for now, still waiting for more input data or output space.
                    Ok(MZStatus::Ok)
                };
            }
        }
    }
}

fn push_dict_out(state: &mut InflateState, next_out: &mut &mut [u8]) -> usize {
    let n = cmp::min(state.dict_avail as usize, next_out.len());
    (next_out[..n]).copy_from_slice(&state.dict[state.dict_ofs..state.dict_ofs + n]);
    *next_out = &mut mem::replace(next_out, &mut [])[n..];
    state.dict_avail -= n;
    state.dict_ofs = (state.dict_ofs + (n)) & (TINFL_LZ_DICT_SIZE - 1);
    n
}

#[cfg(test)]
mod test {
    use super::{inflate, InflateState};
    use crate::{DataFormat, MZFlush, MZStatus};
    #[test]
    fn test_state() {
        let encoded = [
            120u8, 156, 243, 72, 205, 201, 201, 215, 81, 168, 202, 201, 76, 82, 4, 0, 27, 101, 4,
            19,
        ];
        let mut out = vec![0; 50];
        let mut state = InflateState::new_boxed(DataFormat::Zlib);
        let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
        let status = res.status.expect("Failed to decompress!");
        assert_eq!(status, MZStatus::StreamEnd);
        assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
        assert_eq!(res.bytes_consumed, encoded.len());

        state.reset(DataFormat::Zlib);
        let status = res.status.expect("Failed to decompress!");
        assert_eq!(status, MZStatus::StreamEnd);
        assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
        assert_eq!(res.bytes_consumed, encoded.len());
    }
}
<commit_msg>Remove stray println. (#59)<commit_after>//! Extra streaming decompression functionality.
//!
//! As of now this is mainly inteded for use to build a higher-level wrapper.
use std::io::Cursor;
use std::{cmp, mem};

use crate::inflate::core::{decompress, inflate_flags, DecompressorOxide, TINFL_LZ_DICT_SIZE};
use crate::inflate::TINFLStatus;
use crate::{DataFormat, MZError, MZFlush, MZResult, MZStatus, StreamResult};

/// A struct that compbines a decompressor with extra data for streaming decompression.
///
pub struct InflateState {
    /// Inner decompressor struct
    decomp: DecompressorOxide,

    /// Buffer of input bytes for matches.
    /// TODO: Could probably do this a bit cleaner with some
    /// Cursor-like class.
    /// We may also look into whether we need to keep a buffer here, or just one in the
    /// decompressor struct.
    dict: [u8; TINFL_LZ_DICT_SIZE],
    /// Where in the buffer are we currently at?
    dict_ofs: usize,
    /// How many bytes of data to be flushed is there currently in the buffer?
    dict_avail: usize,

    first_call: bool,
    has_flushed: bool,

    /// Whether the input data is wrapped in a zlib header and checksum.
    /// TODO: This should be stored in the decompressor.
    data_format: DataFormat,
    last_status: TINFLStatus,
}

impl Default for InflateState {
    fn default() -> Self {
        InflateState {
            decomp: DecompressorOxide::default(),
            dict: [0; TINFL_LZ_DICT_SIZE],
            dict_ofs: 0,
            dict_avail: 0,
            first_call: true,
            has_flushed: false,
            data_format: DataFormat::Raw,
            last_status: TINFLStatus::NeedsMoreInput,
        }
    }
}
impl InflateState {
    /// Create a new state.
    ///
    /// Note that this struct is quite large due to internal buffers, and as such storing it on
    /// the stack is not recommended.
    ///
    /// # Parameters
    /// `data_format`: Determines whether the compressed data is assumed to wrapped with zlib
    /// metadata.
    pub fn new(data_format: DataFormat) -> InflateState {
        let mut b = InflateState::default();
        b.data_format = data_format;
        b
    }

    /// Create a new state on the heap.
    ///
    /// # Parameters
    /// `data_format`: Determines whether the compressed data is assumed to wrapped with zlib
    /// metadata.
    pub fn new_boxed(data_format: DataFormat) -> Box<InflateState> {
        let mut b: Box<InflateState> = Box::default();
        b.data_format = data_format;
        b
    }

    /// Access the innner decompressor.
    pub fn decompressor(&mut self) -> &mut DecompressorOxide {
        &mut self.decomp
    }

    /// Return the status of the last call to `inflate` with this `InflateState`.
    pub fn last_status(&self) -> TINFLStatus {
        self.last_status
    }

    /// Create a new state using miniz/zlib style window bits parameter.
    ///
    /// The decompressor does not support different window sizes. As such,
    /// any positive (>0) value will set the zlib header flag, while a negative one
    /// will not.
    pub fn new_boxed_with_window_bits(window_bits: i32) -> Box<InflateState> {
        let mut b: Box<InflateState> = Box::default();
        b.data_format = DataFormat::from_window_bits(window_bits);
        b
    }

    /// Reset the decompressor without re-allocating memory, using the given
    /// data format.
    pub fn reset(&mut self, data_format: DataFormat) {
        self.decompressor().init();
        self.dict = [0; TINFL_LZ_DICT_SIZE];
        self.dict_ofs = 0;
        self.dict_avail = 0;
        self.first_call = true;
        self.has_flushed = false;
        self.data_format = data_format;
        self.last_status = TINFLStatus::NeedsMoreInput;
    }
}

/// Try to decompress from `input` to `output` with the given `InflateState`
///
/// # Errors
///
/// Returns `MZError::Buf` If the size of the `output` slice is empty or no progress was made due to
/// lack of expected input data or called after the decompression was
/// finished without MZFlush::Finish.
///
/// Returns `MZError::Param` if the compressor parameters are set wrong.
pub fn inflate(
    state: &mut InflateState,
    input: &[u8],
    output: &mut [u8],
    flush: MZFlush,
) -> StreamResult {
    let mut bytes_consumed = 0;
    let mut bytes_written = 0;
    let mut next_in = input;
    let mut next_out = output;

    if flush == MZFlush::Full {
        return StreamResult::error(MZError::Stream);
    }

    let mut decomp_flags = inflate_flags::TINFL_FLAG_COMPUTE_ADLER32;
    if state.data_format == DataFormat::Zlib {
        decomp_flags |= inflate_flags::TINFL_FLAG_PARSE_ZLIB_HEADER;
    }

    let first_call = state.first_call;
    state.first_call = false;
    if (state.last_status as i32) < 0 {
        return StreamResult::error(MZError::Data);
    }

    if state.has_flushed && (flush != MZFlush::Finish) {
        return StreamResult::error(MZError::Stream);
    }
    state.has_flushed |= flush == MZFlush::Finish;

    if (flush == MZFlush::Finish) && first_call {
        decomp_flags |= inflate_flags::TINFL_FLAG_USING_NON_WRAPPING_OUTPUT_BUF;

        let status = decompress(
            &mut state.decomp,
            next_in,
            &mut Cursor::new(next_out),
            decomp_flags,
        );
        let in_bytes = status.1;
        let out_bytes = status.2;
        let status = status.0;

        state.last_status = status;

        bytes_consumed += in_bytes;
        bytes_written += out_bytes;

        let ret_status = {
            if (status as i32) < 0 {
                Err(MZError::Data)
            } else if status != TINFLStatus::Done {
                state.last_status = TINFLStatus::Failed;
                Err(MZError::Buf)
            } else {
                Ok(MZStatus::StreamEnd)
            }
        };
        return StreamResult {
            bytes_consumed,
            bytes_written,
            status: ret_status,
        };
    }

    if flush != MZFlush::Finish {
        decomp_flags |= inflate_flags::TINFL_FLAG_HAS_MORE_INPUT;
    }

    if state.dict_avail != 0 {
        bytes_written += push_dict_out(state, &mut next_out);
        return StreamResult {
            bytes_consumed,
            bytes_written,
            status: Ok(
                if (state.last_status == TINFLStatus::Done) && (state.dict_avail == 0) {
                    MZStatus::StreamEnd
                } else {
                    MZStatus::Ok
                },
            ),
        };
    }

    let status = inflate_loop(
        state,
        &mut next_in,
        &mut next_out,
        &mut bytes_consumed,
        &mut bytes_written,
        decomp_flags,
        flush,
    );
    StreamResult {
        bytes_consumed,
        bytes_written,
        status,
    }
}

fn inflate_loop(
    state: &mut InflateState,
    next_in: &mut &[u8],
    next_out: &mut &mut [u8],
    total_in: &mut usize,
    total_out: &mut usize,
    decomp_flags: u32,
    flush: MZFlush,
) -> MZResult {
    let orig_in_len = next_in.len();
    loop {
        let status = {
            let mut cursor = Cursor::new(&mut state.dict[..]);
            cursor.set_position(state.dict_ofs as u64);
            decompress(&mut state.decomp, *next_in, &mut cursor, decomp_flags)
        };

        let in_bytes = status.1;
        let out_bytes = status.2;
        let status = status.0;

        state.last_status = status;

        *next_in = &next_in[in_bytes..];
        *total_in += in_bytes;

        state.dict_avail = out_bytes;
        *total_out += push_dict_out(state, next_out);

        // The stream was corrupted, and decompression failed.
        if (status as i32) < 0 {
            return Err(MZError::Data);
        }

        // The decompressor has flushed all it's data and is waiting for more input, but
        // there was no more input provided.
        if (status == TINFLStatus::NeedsMoreInput) && orig_in_len == 0 {
            return Err(MZError::Buf);
        }

        if flush == MZFlush::Finish {
            if status == TINFLStatus::Done {
                // There is not enough space in the output buffer to flush the remaining
                // decompressed data in the internal buffer.
                return if state.dict_avail != 0 {
                    Err(MZError::Buf)
                } else {
                    Ok(MZStatus::StreamEnd)
                };
            // No more space in the output buffer, but we're not done.
            } else if next_out.is_empty() {
                return Err(MZError::Buf);
            }
        } else {
            // We're not expected to finish, so it's fine if we can't flush everything yet.
            let empty_buf = next_in.is_empty() || next_out.is_empty();
            if (status == TINFLStatus::Done) || empty_buf || (state.dict_avail != 0) {
                return if (status == TINFLStatus::Done) && (state.dict_avail == 0) {
                    // No more data left, we're done.
                    Ok(MZStatus::StreamEnd)
                } else {
                    // Ok for now, still waiting for more input data or output space.
                    Ok(MZStatus::Ok)
                };
            }
        }
    }
}

fn push_dict_out(state: &mut InflateState, next_out: &mut &mut [u8]) -> usize {
    let n = cmp::min(state.dict_avail as usize, next_out.len());
    (next_out[..n]).copy_from_slice(&state.dict[state.dict_ofs..state.dict_ofs + n]);
    *next_out = &mut mem::replace(next_out, &mut [])[n..];
    state.dict_avail -= n;
    state.dict_ofs = (state.dict_ofs + (n)) & (TINFL_LZ_DICT_SIZE - 1);
    n
}

#[cfg(test)]
mod test {
    use super::{inflate, InflateState};
    use crate::{DataFormat, MZFlush, MZStatus};
    #[test]
    fn test_state() {
        let encoded = [
            120u8, 156, 243, 72, 205, 201, 201, 215, 81, 168, 202, 201, 76, 82, 4, 0, 27, 101, 4,
            19,
        ];
        let mut out = vec![0; 50];
        let mut state = InflateState::new_boxed(DataFormat::Zlib);
        let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
        let status = res.status.expect("Failed to decompress!");
        assert_eq!(status, MZStatus::StreamEnd);
        assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
        assert_eq!(res.bytes_consumed, encoded.len());

        state.reset(DataFormat::Zlib);
        let status = res.status.expect("Failed to decompress!");
        assert_eq!(status, MZStatus::StreamEnd);
        assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
        assert_eq!(res.bytes_consumed, encoded.len());
    }
}
<|endoftext|>


<commit_before>//! Methods to connect to a WebSocket as a client.

use std::{
    io::{Read, Write},
    net::{SocketAddr, TcpStream, ToSocketAddrs},
    result::Result as StdResult,
};

use http::{request::Parts, Uri};
use log::*;

use url::Url;

use crate::{
    handshake::client::{Request, Response},
    protocol::WebSocketConfig,
};

#[cfg(feature = "tls")]
mod encryption {
    pub use native_tls::TlsStream;
    use native_tls::{HandshakeError as TlsHandshakeError, TlsConnector};
    use std::net::TcpStream;

    pub use crate::stream::Stream as StreamSwitcher;
    /// TCP stream switcher (plain/TLS).
    pub type AutoStream = StreamSwitcher<TcpStream, TlsStream<TcpStream>>;

    use crate::{error::Result, stream::Mode};

    pub fn wrap_stream(stream: TcpStream, domain: &str, mode: Mode) -> Result<AutoStream> {
        match mode {
            Mode::Plain => Ok(StreamSwitcher::Plain(stream)),
            Mode::Tls => {
                let connector = TlsConnector::builder().build()?;
                connector
                    .connect(domain, stream)
                    .map_err(|e| match e {
                        TlsHandshakeError::Failure(f) => f.into(),
                        TlsHandshakeError::WouldBlock(_) => {
                            panic!("Bug: TLS handshake not blocked")
                        }
                    })
                    .map(StreamSwitcher::Tls)
            }
        }
    }
}

#[cfg(not(feature = "tls"))]
mod encryption {
    use std::net::TcpStream;

    use crate::{
        error::{Error, Result},
        stream::Mode,
    };

    /// TLS support is nod compiled in, this is just standard `TcpStream`.
    pub type AutoStream = TcpStream;

    pub fn wrap_stream(stream: TcpStream, _domain: &str, mode: Mode) -> Result<AutoStream> {
        match mode {
            Mode::Plain => Ok(stream),
            Mode::Tls => Err(Error::Url("TLS support not compiled in.".into())),
        }
    }
}

use self::encryption::wrap_stream;
pub use self::encryption::AutoStream;

use crate::{
    error::{Error, Result},
    handshake::{client::ClientHandshake, HandshakeError},
    protocol::WebSocket,
    stream::{Mode, NoDelay},
};

/// Connect to the given WebSocket in blocking mode.
///
/// Uses a websocket configuration passed as an argument to the function. Calling it with `None` is
/// equal to calling `connect()` function.
///
/// The URL may be either ws:// or wss://.
/// To support wss:// URLs, feature "tls" must be turned on.
///
/// This function "just works" for those who wants a simple blocking solution
/// similar to `std::net::TcpStream`. If you want a non-blocking or other
/// custom stream, call `client` instead.
///
/// This function uses `native_tls` to do TLS. If you want to use other TLS libraries,
/// use `client` instead. There is no need to enable the "tls" feature if you don't call
/// `connect` since it's the only function that uses native_tls.
pub fn connect_with_config<Req: IntoClientRequest>(
    request: Req,
    config: Option<WebSocketConfig>,
    max_redirects: u8,
) -> Result<(WebSocket<AutoStream>, Response)> {
    fn try_client_handshake(
        request: Request,
        config: Option<WebSocketConfig>,
    ) -> Result<(WebSocket<AutoStream>, Response)> {
        let uri = request.uri();
        let mode = uri_mode(uri)?;
        let host =
            request.uri().host().ok_or_else(|| Error::Url("No host name in the URL".into()))?;
        let port = uri.port_u16().unwrap_or(match mode {
            Mode::Plain => 80,
            Mode::Tls => 443,
        });
        let addrs = (host, port).to_socket_addrs()?;
        let mut stream = connect_to_some(addrs.as_slice(), &request.uri(), mode)?;
        NoDelay::set_nodelay(&mut stream, true)?;
        client_with_config(request, stream, config).map_err(|e| match e {
            HandshakeError::Failure(f) => f,
            HandshakeError::Interrupted(_) => panic!("Bug: blocking handshake not blocked"),
        })
    }

    fn create_request(parts: &Parts, uri: &Uri) -> Request {
        let mut builder =
            Request::builder().uri(uri.clone()).method(parts.method.clone()).version(parts.version);
        *builder.headers_mut().expect("Failed to create `Request`") = parts.headers.clone();
        builder.body(()).expect("Failed to create `Request`")
    }

    let (parts, _) = request.into_client_request()?.into_parts();
    let mut uri = parts.uri.clone();

    for attempt in 0..(max_redirects + 1) {
        let request = create_request(&parts, &uri);

        match try_client_handshake(request, config) {
            Err(Error::Http(res)) if res.status().is_redirection() && attempt < max_redirects => {
                if let Some(location) = res.headers().get("Location") {
                    uri = location.to_str()?.parse::<Uri>()?;
                    debug!("Redirecting to {:?}", uri);
                    continue;
                } else {
                    warn!("No `Location` found in redirect");
                    return Err(Error::Http(res));
                }
            }
            other => return other,
        }
    }

    unreachable!("Bug in a redirect handling logic")
}

/// Connect to the given WebSocket in blocking mode.
///
/// The URL may be either ws:// or wss://.
/// To support wss:// URLs, feature "tls" must be turned on.
///
/// This function "just works" for those who wants a simple blocking solution
/// similar to `std::net::TcpStream`. If you want a non-blocking or other
/// custom stream, call `client` instead.
///
/// This function uses `native_tls` to do TLS. If you want to use other TLS libraries,
/// use `client` instead. There is no need to enable the "tls" feature if you don't call
/// `connect` since it's the only function that uses native_tls.
pub fn connect<Req: IntoClientRequest>(request: Req) -> Result<(WebSocket<AutoStream>, Response)> {
    connect_with_config(request, None, 3)
}

fn connect_to_some(addrs: &[SocketAddr], uri: &Uri, mode: Mode) -> Result<AutoStream> {
    let domain = uri.host().ok_or_else(|| Error::Url("No host name in the URL".into()))?;
    for addr in addrs {
        debug!("Trying to contact {} at {}...", uri, addr);
        if let Ok(raw_stream) = TcpStream::connect(addr) {
            if let Ok(stream) = wrap_stream(raw_stream, domain, mode) {
                return Ok(stream);
            }
        }
    }
    Err(Error::Url(format!("Unable to connect to {}", uri).into()))
}

/// Get the mode of the given URL.
///
/// This function may be used to ease the creation of custom TLS streams
/// in non-blocking algorithmss or for use with TLS libraries other than `native_tls`.
pub fn uri_mode(uri: &Uri) -> Result<Mode> {
    match uri.scheme_str() {
        Some("ws") => Ok(Mode::Plain),
        Some("wss") => Ok(Mode::Tls),
        _ => Err(Error::Url("URL scheme not supported".into())),
    }
}

/// Do the client handshake over the given stream given a web socket configuration. Passing `None`
/// as configuration is equal to calling `client()` function.
///
/// Use this function if you need a nonblocking handshake support or if you
/// want to use a custom stream like `mio::tcp::TcpStream` or `openssl::ssl::SslStream`.
/// Any stream supporting `Read + Write` will do.
pub fn client_with_config<Stream, Req>(
    request: Req,
    stream: Stream,
    config: Option<WebSocketConfig>,
) -> StdResult<(WebSocket<Stream>, Response), HandshakeError<ClientHandshake<Stream>>>
where
    Stream: Read + Write,
    Req: IntoClientRequest,
{
    ClientHandshake::start(stream, request.into_client_request()?, config)?.handshake()
}

/// Do the client handshake over the given stream.
///
/// Use this function if you need a nonblocking handshake support or if you
/// want to use a custom stream like `mio::tcp::TcpStream` or `openssl::ssl::SslStream`.
/// Any stream supporting `Read + Write` will do.
pub fn client<Stream, Req>(
    request: Req,
    stream: Stream,
) -> StdResult<(WebSocket<Stream>, Response), HandshakeError<ClientHandshake<Stream>>>
where
    Stream: Read + Write,
    Req: IntoClientRequest,
{
    client_with_config(request, stream, None)
}

/// Trait for converting various types into HTTP requests used for a client connection.
///
/// This trait is implemented by default for string slices, strings, `url::Url`, `http::Uri` and
/// `http::Request<()>`.
pub trait IntoClientRequest {
    /// Convert into a `Request` that can be used for a client connection.
    fn into_client_request(self) -> Result<Request>;
}

impl<'a> IntoClientRequest for &'a str {
    fn into_client_request(self) -> Result<Request> {
        self.parse::<Uri>()?.into_client_request()
    }
}

impl<'a> IntoClientRequest for &'a String {
    fn into_client_request(self) -> Result<Request> {
        <&str as IntoClientRequest>::into_client_request(self)
    }
}

impl IntoClientRequest for String {
    fn into_client_request(self) -> Result<Request> {
        <&str as IntoClientRequest>::into_client_request(&self)
    }
}

impl<'a> IntoClientRequest for &'a Uri {
    fn into_client_request(self) -> Result<Request> {
        self.clone().into_client_request()
    }
}

impl IntoClientRequest for Uri {
    fn into_client_request(self) -> Result<Request> {
        Ok(Request::get(self).body(())?)
    }
}

impl<'a> IntoClientRequest for &'a Url {
    fn into_client_request(self) -> Result<Request> {
        self.as_str().into_client_request()
    }
}

impl IntoClientRequest for Url {
    fn into_client_request(self) -> Result<Request> {
        self.as_str().into_client_request()
    }
}

impl IntoClientRequest for Request {
    fn into_client_request(self) -> Result<Request> {
        Ok(self)
    }
}

impl<'h, 'b> IntoClientRequest for httparse::Request<'h, 'b> {
    fn into_client_request(self) -> Result<Request> {
        use crate::handshake::headers::FromHttparse;
        Request::from_httparse(self)
    }
}
<commit_msg>Docfix: mio::tcp has been renamed to mio::net in its latest version (#164)<commit_after>//! Methods to connect to a WebSocket as a client.

use std::{
    io::{Read, Write},
    net::{SocketAddr, TcpStream, ToSocketAddrs},
    result::Result as StdResult,
};

use http::{request::Parts, Uri};
use log::*;

use url::Url;

use crate::{
    handshake::client::{Request, Response},
    protocol::WebSocketConfig,
};

#[cfg(feature = "tls")]
mod encryption {
    pub use native_tls::TlsStream;
    use native_tls::{HandshakeError as TlsHandshakeError, TlsConnector};
    use std::net::TcpStream;

    pub use crate::stream::Stream as StreamSwitcher;
    /// TCP stream switcher (plain/TLS).
    pub type AutoStream = StreamSwitcher<TcpStream, TlsStream<TcpStream>>;

    use crate::{error::Result, stream::Mode};

    pub fn wrap_stream(stream: TcpStream, domain: &str, mode: Mode) -> Result<AutoStream> {
        match mode {
            Mode::Plain => Ok(StreamSwitcher::Plain(stream)),
            Mode::Tls => {
                let connector = TlsConnector::builder().build()?;
                connector
                    .connect(domain, stream)
                    .map_err(|e| match e {
                        TlsHandshakeError::Failure(f) => f.into(),
                        TlsHandshakeError::WouldBlock(_) => {
                            panic!("Bug: TLS handshake not blocked")
                        }
                    })
                    .map(StreamSwitcher::Tls)
            }
        }
    }
}

#[cfg(not(feature = "tls"))]
mod encryption {
    use std::net::TcpStream;

    use crate::{
        error::{Error, Result},
        stream::Mode,
    };

    /// TLS support is nod compiled in, this is just standard `TcpStream`.
    pub type AutoStream = TcpStream;

    pub fn wrap_stream(stream: TcpStream, _domain: &str, mode: Mode) -> Result<AutoStream> {
        match mode {
            Mode::Plain => Ok(stream),
            Mode::Tls => Err(Error::Url("TLS support not compiled in.".into())),
        }
    }
}

use self::encryption::wrap_stream;
pub use self::encryption::AutoStream;

use crate::{
    error::{Error, Result},
    handshake::{client::ClientHandshake, HandshakeError},
    protocol::WebSocket,
    stream::{Mode, NoDelay},
};

/// Connect to the given WebSocket in blocking mode.
///
/// Uses a websocket configuration passed as an argument to the function. Calling it with `None` is
/// equal to calling `connect()` function.
///
/// The URL may be either ws:// or wss://.
/// To support wss:// URLs, feature "tls" must be turned on.
///
/// This function "just works" for those who wants a simple blocking solution
/// similar to `std::net::TcpStream`. If you want a non-blocking or other
/// custom stream, call `client` instead.
///
/// This function uses `native_tls` to do TLS. If you want to use other TLS libraries,
/// use `client` instead. There is no need to enable the "tls" feature if you don't call
/// `connect` since it's the only function that uses native_tls.
pub fn connect_with_config<Req: IntoClientRequest>(
    request: Req,
    config: Option<WebSocketConfig>,
    max_redirects: u8,
) -> Result<(WebSocket<AutoStream>, Response)> {
    fn try_client_handshake(
        request: Request,
        config: Option<WebSocketConfig>,
    ) -> Result<(WebSocket<AutoStream>, Response)> {
        let uri = request.uri();
        let mode = uri_mode(uri)?;
        let host =
            request.uri().host().ok_or_else(|| Error::Url("No host name in the URL".into()))?;
        let port = uri.port_u16().unwrap_or(match mode {
            Mode::Plain => 80,
            Mode::Tls => 443,
        });
        let addrs = (host, port).to_socket_addrs()?;
        let mut stream = connect_to_some(addrs.as_slice(), &request.uri(), mode)?;
        NoDelay::set_nodelay(&mut stream, true)?;
        client_with_config(request, stream, config).map_err(|e| match e {
            HandshakeError::Failure(f) => f,
            HandshakeError::Interrupted(_) => panic!("Bug: blocking handshake not blocked"),
        })
    }

    fn create_request(parts: &Parts, uri: &Uri) -> Request {
        let mut builder =
            Request::builder().uri(uri.clone()).method(parts.method.clone()).version(parts.version);
        *builder.headers_mut().expect("Failed to create `Request`") = parts.headers.clone();
        builder.body(()).expect("Failed to create `Request`")
    }

    let (parts, _) = request.into_client_request()?.into_parts();
    let mut uri = parts.uri.clone();

    for attempt in 0..(max_redirects + 1) {
        let request = create_request(&parts, &uri);

        match try_client_handshake(request, config) {
            Err(Error::Http(res)) if res.status().is_redirection() && attempt < max_redirects => {
                if let Some(location) = res.headers().get("Location") {
                    uri = location.to_str()?.parse::<Uri>()?;
                    debug!("Redirecting to {:?}", uri);
                    continue;
                } else {
                    warn!("No `Location` found in redirect");
                    return Err(Error::Http(res));
                }
            }
            other => return other,
        }
    }

    unreachable!("Bug in a redirect handling logic")
}

/// Connect to the given WebSocket in blocking mode.
///
/// The URL may be either ws:// or wss://.
/// To support wss:// URLs, feature "tls" must be turned on.
///
/// This function "just works" for those who wants a simple blocking solution
/// similar to `std::net::TcpStream`. If you want a non-blocking or other
/// custom stream, call `client` instead.
///
/// This function uses `native_tls` to do TLS. If you want to use other TLS libraries,
/// use `client` instead. There is no need to enable the "tls" feature if you don't call
/// `connect` since it's the only function that uses native_tls.
pub fn connect<Req: IntoClientRequest>(request: Req) -> Result<(WebSocket<AutoStream>, Response)> {
    connect_with_config(request, None, 3)
}

fn connect_to_some(addrs: &[SocketAddr], uri: &Uri, mode: Mode) -> Result<AutoStream> {
    let domain = uri.host().ok_or_else(|| Error::Url("No host name in the URL".into()))?;
    for addr in addrs {
        debug!("Trying to contact {} at {}...", uri, addr);
        if let Ok(raw_stream) = TcpStream::connect(addr) {
            if let Ok(stream) = wrap_stream(raw_stream, domain, mode) {
                return Ok(stream);
            }
        }
    }
    Err(Error::Url(format!("Unable to connect to {}", uri).into()))
}

/// Get the mode of the given URL.
///
/// This function may be used to ease the creation of custom TLS streams
/// in non-blocking algorithmss or for use with TLS libraries other than `native_tls`.
pub fn uri_mode(uri: &Uri) -> Result<Mode> {
    match uri.scheme_str() {
        Some("ws") => Ok(Mode::Plain),
        Some("wss") => Ok(Mode::Tls),
        _ => Err(Error::Url("URL scheme not supported".into())),
    }
}

/// Do the client handshake over the given stream given a web socket configuration. Passing `None`
/// as configuration is equal to calling `client()` function.
///
/// Use this function if you need a nonblocking handshake support or if you
/// want to use a custom stream like `mio::net::TcpStream` or `openssl::ssl::SslStream`.
/// Any stream supporting `Read + Write` will do.
pub fn client_with_config<Stream, Req>(
    request: Req,
    stream: Stream,
    config: Option<WebSocketConfig>,
) -> StdResult<(WebSocket<Stream>, Response), HandshakeError<ClientHandshake<Stream>>>
where
    Stream: Read + Write,
    Req: IntoClientRequest,
{
    ClientHandshake::start(stream, request.into_client_request()?, config)?.handshake()
}

/// Do the client handshake over the given stream.
///
/// Use this function if you need a nonblocking handshake support or if you
/// want to use a custom stream like `mio::net::TcpStream` or `openssl::ssl::SslStream`.
/// Any stream supporting `Read + Write` will do.
pub fn client<Stream, Req>(
    request: Req,
    stream: Stream,
) -> StdResult<(WebSocket<Stream>, Response), HandshakeError<ClientHandshake<Stream>>>
where
    Stream: Read + Write,
    Req: IntoClientRequest,
{
    client_with_config(request, stream, None)
}

/// Trait for converting various types into HTTP requests used for a client connection.
///
/// This trait is implemented by default for string slices, strings, `url::Url`, `http::Uri` and
/// `http::Request<()>`.
pub trait IntoClientRequest {
    /// Convert into a `Request` that can be used for a client connection.
    fn into_client_request(self) -> Result<Request>;
}

impl<'a> IntoClientRequest for &'a str {
    fn into_client_request(self) -> Result<Request> {
        self.parse::<Uri>()?.into_client_request()
    }
}

impl<'a> IntoClientRequest for &'a String {
    fn into_client_request(self) -> Result<Request> {
        <&str as IntoClientRequest>::into_client_request(self)
    }
}

impl IntoClientRequest for String {
    fn into_client_request(self) -> Result<Request> {
        <&str as IntoClientRequest>::into_client_request(&self)
    }
}

impl<'a> IntoClientRequest for &'a Uri {
    fn into_client_request(self) -> Result<Request> {
        self.clone().into_client_request()
    }
}

impl IntoClientRequest for Uri {
    fn into_client_request(self) -> Result<Request> {
        Ok(Request::get(self).body(())?)
    }
}

impl<'a> IntoClientRequest for &'a Url {
    fn into_client_request(self) -> Result<Request> {
        self.as_str().into_client_request()
    }
}

impl IntoClientRequest for Url {
    fn into_client_request(self) -> Result<Request> {
        self.as_str().into_client_request()
    }
}

impl IntoClientRequest for Request {
    fn into_client_request(self) -> Result<Request> {
        Ok(self)
    }
}

impl<'h, 'b> IntoClientRequest for httparse::Request<'h, 'b> {
    fn into_client_request(self) -> Result<Request> {
        use crate::handshake::headers::FromHttparse;
        Request::from_httparse(self)
    }
}
<|endoftext|>


<commit_before>use std::collections::VecDeque;
use std::io;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};

use futures::{Future, Sink, Stream};
use futures::sync::{mpsc, oneshot};

use tokio_core::net::TcpStream;
use tokio_core::reactor::Core;
use tokio_io::AsyncRead;

use super::{error, resp};

/// TODO: comeback and optimise this number
const DEFAULT_BUFFER_SIZE:usize = 100;

/// Connect to a Redis server and return paired Sink and Stream for reading and writing
/// asynchronously.
fn connect(addr: &SocketAddr, core: &Core) -> Box<Future<Item=ClientConnection, Error=io::Error>> {
    let handle = core.handle();
    let client_con = TcpStream::connect(addr, &handle).map(move |socket| {
        let framed = socket.framed(resp::RespCodec);
        let (write_f, read_f) = framed.split();
        let write_b = write_f.buffer(DEFAULT_BUFFER_SIZE);
        ClientConnection {
            sender: ClientSink(Box::new(write_b)),
            receiver: ClientStream(Box::new(read_f))
        }
    });
    Box::new(client_con)
}

/// TODO - is the boxing necessary?  It makes the type signature much simpler
struct ClientSink(Box<Sink<SinkItem=resp::RespValue, SinkError=io::Error>>);
struct ClientStream(Box<Stream<Item=resp::RespValue, Error=error::Error>>);

/// A low-level client connection representing a sender and a receiver.
///
/// The two halves operate independently from one another
struct ClientConnection {
    sender: ClientSink,
    receiver: ClientStream
}

fn paired_connect(addr: &SocketAddr, core: &Core) -> Box<Future<Item=PairedConnection, Error=error::Error>> {
    let handle = core.handle();
    let paired_con = connect(addr, core).map(move |connection| {
        let ClientConnection { sender, receiver } = connection;
        let (out_tx, out_rx) = mpsc::unbounded();
        let sender = out_rx.fold(sender.0, |sender, msg| {
            sender.send(msg).map_err(|_| ())
        });
        let resp_queue:Arc<Mutex<VecDeque<oneshot::Sender<resp::RespValue>>>> = Arc::new(Mutex::new(VecDeque::new()));
        let receiver_queue = resp_queue.clone();
        let receiver = receiver.0.for_each(move |msg| {
            let mut queue = receiver_queue.lock().expect("Lock is tainted");
            let dest = queue.pop_front().expect("Queue is empty");
            match dest.send(msg) {
                Ok(()) => Ok(()),
                // Ignore error as the channel may have been legitimately closed in the meantime
                Err(_) => Ok(())
            }
        });
        handle.spawn(sender.map(|_| ()));
        handle.spawn(receiver.map_err(|_| ()));
        PairedConnection {
            out_tx: out_tx,
            resp_queue: resp_queue
        }
    }).map_err(|e| e.into());
    Box::new(paired_con)
}

struct PairedConnection {
    out_tx: mpsc::UnboundedSender<resp::RespValue>,
    resp_queue: Arc<Mutex<VecDeque<oneshot::Sender<resp::RespValue>>>>
}

impl PairedConnection {
    fn send<R>(&self, msg: R) -> Box<Future<Item=resp::RespValue, Error=error::Error>>
    where R: Into<resp::RespValue> {
        let (tx, rx) = oneshot::channel();
        let mut queue = self.resp_queue.lock().expect("Tainted queue");
        queue.push_back(tx);
        mpsc::UnboundedSender::send(&self.out_tx, msg.into()).expect("Failed to send");
        Box::new(rx.map_err(|e| e.into()))
    }
}

#[cfg(test)]
mod test {
    use std::io;
    use std::rc::Rc;

    use futures::{Future, Sink, Stream};
    use futures::{future, stream};

    use tokio_core::reactor::Core;

    use super::resp;

    #[test]
    fn can_connect() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connection = super::connect(&addr, &core).map_err(|e| e.into()).and_then(|connection| {
            let a = connection.sender.0.send(["PING", "TEST"].as_ref().into()).map_err(|e| e.into());
            let b = connection.receiver.0.take(1).collect();
            a.join(b)
        });

        let (_, values) = core.run(connection).unwrap();
        assert_eq!(values.len(), 1);
        assert_eq!(values[0], "TEST".into());
    }

    #[test]
    fn complex_test() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();
        let connection = super::connect(&addr, &core).map_err(|e| e.into()).and_then(|connection| {
            let mut ops = Vec::<resp::RespValue>::new();
            ops.push(["FLUSH"].as_ref().into());
            ops.extend((0..1000).map(|i| ["SADD", "test_set", &format!("VALUE: {}", i)].as_ref().into()));
            ops.push(["SMEMBERS", "test_set"].as_ref().into());
            let ops_r:Vec<Result<resp::RespValue, io::Error>> = ops.into_iter().map(Result::Ok).collect();
            let send = connection.sender.0.send_all(stream::iter(ops_r)).map_err(|e| e.into());
            let receive = connection.receiver.0.skip(1001).take(1).collect();
            send.join(receive)
        });
        let (_, values) = core.run(connection).unwrap();
        assert_eq!(values.len(), 1);
        let values = match &values[0] {
            &resp::RespValue::Array(ref values) => values.clone(),
            _ => panic!("Not an array")
        };
        assert_eq!(values.len(), 1000);
    }

    #[test]
    fn can_paired_connect() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connect_f = super::paired_connect(&addr, &core).and_then(|connection| {
            let res_f = connection.send(vec!["PING", "TEST"]);
            connection.send(vec!["SET", "X", "123"]);
            let wait_f = connection.send(vec!["GET", "X"]);
            res_f.join(wait_f)
        });
        let (result_1, result_2) = core.run(connect_f).unwrap();
        assert_eq!(result_1, "TEST".into());
        assert_eq!(result_2, "123".into());
    }

    #[test]
    fn complex_paired_connect() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connect_f = super::paired_connect(&addr, &core).and_then(|connection| {
            connection.send(vec!["INCR", "CTR"]).and_then(move |value| {
                let value_str = value.into_string().expect("A string");
                connection.send(vec!["SET", "LASTCTR", &value_str])
            })
        });
        let result = core.run(connect_f).unwrap();
        assert_eq!(result, resp::RespValue::SimpleString("OK".into()));
    }

    #[test]
    fn sending_a_lot_of_data_test() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();

        let test_f = super::paired_connect(&addr, &core);
        let send_data = test_f.and_then(|connection| {
            let mut futures = Vec::with_capacity(1000);
            for i in 0..1000 {
                let key = format!("X_{}", i);
                connection.send(vec!["SET", &key, &i.to_string()]);
                futures.push(connection.send(vec!["GET", &key]));
            }
            futures.remove(999)
        });
        let result = core.run(send_data).unwrap();
        assert_eq!(result.into_string().unwrap(), "999");
    }

    #[test]
    fn realistic_test() {
        let test_data_size = 100;
        let test_data:Vec<_> = (0..test_data_size).map(|x| (x, x.to_string())).collect();
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();
        let test_f = super::paired_connect(&addr, &core);
        let send_data = test_f.and_then(|connection| {
            let connection = Rc::new(connection);
            let futures:Vec<_> = test_data.into_iter().map(move |data| {
                let connection_inner = connection.clone();
                connection.send(vec!["INCR", "realistic_test_ctr"]).and_then(move |ctr| {
                    let key = format!("rt_{}", ctr.into_string().unwrap());
                    let d_val = data.0.to_string();
                    connection_inner.send(vec!["SET", &key, &d_val]);
                    connection_inner.send(vec!["SET", &data.1, &key])
                })
            }).collect();
            future::join_all(futures)
        });
        let result = core.run(send_data).unwrap();
        assert_eq!(result.len(), 100);
    }
}<commit_msg>Handling Handle correctly (vs. Core)<commit_after>use std::collections::VecDeque;
use std::io;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};

use futures::{Future, Sink, Stream};
use futures::sync::{mpsc, oneshot};

use tokio_core::net::TcpStream;
use tokio_core::reactor::Handle;
use tokio_io::AsyncRead;

use super::{error, resp};

/// TODO: comeback and optimise this number
const DEFAULT_BUFFER_SIZE:usize = 100;

/// Connect to a Redis server and return paired Sink and Stream for reading and writing
/// asynchronously.
fn connect(addr: &SocketAddr, handle: &Handle) -> Box<Future<Item=ClientConnection, Error=io::Error>> {
    let client_con = TcpStream::connect(addr, handle).map(move |socket| {
        let framed = socket.framed(resp::RespCodec);
        let (write_f, read_f) = framed.split();
        let write_b = write_f.buffer(DEFAULT_BUFFER_SIZE);
        ClientConnection {
            sender: ClientSink(Box::new(write_b)),
            receiver: ClientStream(Box::new(read_f))
        }
    });
    Box::new(client_con)
}

/// TODO - is the boxing necessary?  It makes the type signature much simpler
struct ClientSink(Box<Sink<SinkItem=resp::RespValue, SinkError=io::Error>>);
struct ClientStream(Box<Stream<Item=resp::RespValue, Error=error::Error>>);

/// A low-level client connection representing a sender and a receiver.
///
/// The two halves operate independently from one another
struct ClientConnection {
    sender: ClientSink,
    receiver: ClientStream
}

fn paired_connect(addr: &SocketAddr, handle: &Handle) -> Box<Future<Item=PairedConnection, Error=error::Error>> {
    let handle = handle.clone();
    let paired_con = connect(addr, &handle).map(move |connection| {
        let ClientConnection { sender, receiver } = connection;
        let (out_tx, out_rx) = mpsc::unbounded();
        let sender = out_rx.fold(sender.0, |sender, msg| {
            sender.send(msg).map_err(|_| ())
        });
        let resp_queue:Arc<Mutex<VecDeque<oneshot::Sender<resp::RespValue>>>> = Arc::new(Mutex::new(VecDeque::new()));
        let receiver_queue = resp_queue.clone();
        let receiver = receiver.0.for_each(move |msg| {
            let mut queue = receiver_queue.lock().expect("Lock is tainted");
            let dest = queue.pop_front().expect("Queue is empty");
            match dest.send(msg) {
                Ok(()) => Ok(()),
                // Ignore error as the channel may have been legitimately closed in the meantime
                Err(_) => Ok(())
            }
        });
        handle.spawn(sender.map(|_| ()));
        handle.spawn(receiver.map_err(|_| ()));
        PairedConnection {
            out_tx: out_tx,
            resp_queue: resp_queue
        }
    }).map_err(|e| e.into());
    Box::new(paired_con)
}

struct PairedConnection {
    out_tx: mpsc::UnboundedSender<resp::RespValue>,
    resp_queue: Arc<Mutex<VecDeque<oneshot::Sender<resp::RespValue>>>>
}

impl PairedConnection {
    fn send<R>(&self, msg: R) -> Box<Future<Item=resp::RespValue, Error=error::Error>>
    where R: Into<resp::RespValue> {
        let (tx, rx) = oneshot::channel();
        let mut queue = self.resp_queue.lock().expect("Tainted queue");
        queue.push_back(tx);
        mpsc::UnboundedSender::send(&self.out_tx, msg.into()).expect("Failed to send");
        Box::new(rx.map_err(|e| e.into()))
    }
}

#[cfg(test)]
mod test {
    use std::io;
    use std::rc::Rc;

    use futures::{Future, Sink, Stream};
    use futures::{future, stream};

    use tokio_core::reactor::Core;

    use super::resp;

    #[test]
    fn can_connect() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connection = super::connect(&addr, &core.handle()).map_err(|e| e.into()).and_then(|connection| {
            let a = connection.sender.0.send(["PING", "TEST"].as_ref().into()).map_err(|e| e.into());
            let b = connection.receiver.0.take(1).collect();
            a.join(b)
        });

        let (_, values) = core.run(connection).unwrap();
        assert_eq!(values.len(), 1);
        assert_eq!(values[0], "TEST".into());
    }

    #[test]
    fn complex_test() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();
        let connection = super::connect(&addr, &core.handle()).map_err(|e| e.into()).and_then(|connection| {
            let mut ops = Vec::<resp::RespValue>::new();
            ops.push(["FLUSH"].as_ref().into());
            ops.extend((0..1000).map(|i| ["SADD", "test_set", &format!("VALUE: {}", i)].as_ref().into()));
            ops.push(["SMEMBERS", "test_set"].as_ref().into());
            let ops_r:Vec<Result<resp::RespValue, io::Error>> = ops.into_iter().map(Result::Ok).collect();
            let send = connection.sender.0.send_all(stream::iter(ops_r)).map_err(|e| e.into());
            let receive = connection.receiver.0.skip(1001).take(1).collect();
            send.join(receive)
        });
        let (_, values) = core.run(connection).unwrap();
        assert_eq!(values.len(), 1);
        let values = match &values[0] {
            &resp::RespValue::Array(ref values) => values.clone(),
            _ => panic!("Not an array")
        };
        assert_eq!(values.len(), 1000);
    }

    #[test]
    fn can_paired_connect() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connect_f = super::paired_connect(&addr, &core.handle()).and_then(|connection| {
            let res_f = connection.send(vec!["PING", "TEST"]);
            connection.send(vec!["SET", "X", "123"]);
            let wait_f = connection.send(vec!["GET", "X"]);
            res_f.join(wait_f)
        });
        let (result_1, result_2) = core.run(connect_f).unwrap();
        assert_eq!(result_1, "TEST".into());
        assert_eq!(result_2, "123".into());
    }

    #[test]
    fn complex_paired_connect() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connect_f = super::paired_connect(&addr, &core.handle()).and_then(|connection| {
            connection.send(vec!["INCR", "CTR"]).and_then(move |value| {
                let value_str = value.into_string().expect("A string");
                connection.send(vec!["SET", "LASTCTR", &value_str])
            })
        });
        let result = core.run(connect_f).unwrap();
        assert_eq!(result, resp::RespValue::SimpleString("OK".into()));
    }

    #[test]
    fn sending_a_lot_of_data_test() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();

        let test_f = super::paired_connect(&addr, &core.handle());
        let send_data = test_f.and_then(|connection| {
            let mut futures = Vec::with_capacity(1000);
            for i in 0..1000 {
                let key = format!("X_{}", i);
                connection.send(vec!["SET", &key, &i.to_string()]);
                futures.push(connection.send(vec!["GET", &key]));
            }
            futures.remove(999)
        });
        let result = core.run(send_data).unwrap();
        assert_eq!(result.into_string().unwrap(), "999");
    }

    #[test]
    fn realistic_test() {
        let test_data_size = 100;
        let test_data:Vec<_> = (0..test_data_size).map(|x| (x, x.to_string())).collect();
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();
        let test_f = super::paired_connect(&addr, &core.handle());
        let send_data = test_f.and_then(|connection| {
            let connection = Rc::new(connection);
            let futures:Vec<_> = test_data.into_iter().map(move |data| {
                let connection_inner = connection.clone();
                connection.send(vec!["INCR", "realistic_test_ctr"]).and_then(move |ctr| {
                    let key = format!("rt_{}", ctr.into_string().unwrap());
                    let d_val = data.0.to_string();
                    connection_inner.send(vec!["SET", &key, &d_val]);
                    connection_inner.send(vec!["SET", &data.1, &key])
                })
            }).collect();
            future::join_all(futures)
        });
        let result = core.run(send_data).unwrap();
        assert_eq!(result.len(), 100);
    }
}<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use fleetspeak_proto::common::{Message, Address};
use prost;
use prost_types;
use std::io::{Read, Write, Result, Error, ErrorKind};

pub struct Connection<'r, 'w, R, W> {
    pub input: &'r mut R,
    pub output: &'w mut W,
    buf: Vec<u8>,
}

impl<'r, 'w, R: Read, W: Write> Connection<'r, 'w, R, W> {

    pub fn new(input: &'r mut R, output: &'w mut W) -> Self {
        Connection {
            input: input,
            output: output,
            buf: vec!(0; 2 * 1024 * 1024), // TODO: Magic.
        }
    }

    pub fn send<M>(&mut self, service: &str, kind: &str, data: M) -> Result<()>
    where
        M: prost::Message,
    {
        self.encode(data)?;

        let msg = Message {
            message_type: kind.to_string(),
            destination: Some(Address {
                service_name: service.to_string(),
                ..Default::default()
            }),
            data: Some(prost_types::Any {
                value: self.buf.to_vec(),
                ..Default::default()
            }),
            ..Default::default()
        };

        self.emit(msg)
    }

    pub fn receive<M>(&mut self) -> Result<M>
    where
        M: prost::Message + Default,
    {
        let msg = self.collect()?;

        // It is not clear what is the best approach here. If there is no data,
        // should we error-out or return a default value? For the time being we
        // stick to the default approach, but if this proves to be not working
        // well in practice, it might be reconsidered.
        let data = match msg.data {
            Some(data) => data,
            None => return Ok(Default::default()),
        };

        match prost::Message::decode(&data.value[..]) {
            Ok(data) => Ok(data),
            Err(err) => {
                let err = Error::new(ErrorKind::InvalidData, Box::new(err));
                Err(err)
            },
        }
    }

    fn emit(&mut self, msg: Message) -> Result<()> {
        let len = prost::Message::encoded_len(&msg);
        self.encode(msg)?;

        self.output.write_u32::<LittleEndian>(len as u32)?;
        self.output.write(&self.buf[..len])?;
        self.output.write_u32::<LittleEndian>(0xf1ee1001)?; // TODO: Magic.

        Ok(())
    }

    fn collect(&mut self) -> Result<Message> {
        let len = self.input.read_i32::<LittleEndian>()? as usize;
        self.input.read_exact(&mut self.buf[..len])?;

        match prost::Message::decode(&self.buf[..len]) {
            Ok(msg) => Ok(msg),
            Err(err) => {
                let err = Error::new(ErrorKind::InvalidData, Box::new(err));
                Err(err)
            },
        }
    }

    fn encode<M>(&mut self, msg: M) -> Result<()>
    where
        M: prost::Message,
    {
        if let Err(err) = msg.encode(&mut self.buf) {
            let err = Error::new(ErrorKind::InvalidData, Box::new(err));
            return Err(err);
        };

        Ok(())
    }
}
<commit_msg>Remove code duplication with error handling.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use fleetspeak_proto::common::{Message, Address};
use prost;
use prost_types;
use std::error::Error;
use std::io::{Read, Write, Result};
use std::marker::{Send, Sync};

pub struct Connection<'r, 'w, R, W> {
    pub input: &'r mut R,
    pub output: &'w mut W,
    buf: Vec<u8>,
}

impl<'r, 'w, R: Read, W: Write> Connection<'r, 'w, R, W> {

    pub fn new(input: &'r mut R, output: &'w mut W) -> Self {
        Connection {
            input: input,
            output: output,
            buf: vec!(0; 2 * 1024 * 1024), // TODO: Magic.
        }
    }

    pub fn send<M>(&mut self, service: &str, kind: &str, data: M) -> Result<()>
    where
        M: prost::Message,
    {
        self.encode(data)?;

        let msg = Message {
            message_type: kind.to_string(),
            destination: Some(Address {
                service_name: service.to_string(),
                ..Default::default()
            }),
            data: Some(prost_types::Any {
                value: self.buf.to_vec(),
                ..Default::default()
            }),
            ..Default::default()
        };

        self.emit(msg)
    }

    pub fn receive<M>(&mut self) -> Result<M>
    where
        M: prost::Message + Default,
    {
        let msg = self.collect()?;

        // It is not clear what is the best approach here. If there is no data,
        // should we error-out or return a default value? For the time being we
        // stick to the default approach, but if this proves to be not working
        // well in practice, it might be reconsidered.
        let data = match msg.data {
            Some(data) => data,
            None => return Ok(Default::default()),
        };

        prost::Message::decode(&data.value[..]).map_err(invalid_data_error)
    }

    fn emit(&mut self, msg: Message) -> Result<()> {
        let len = prost::Message::encoded_len(&msg);
        self.encode(msg)?;

        self.output.write_u32::<LittleEndian>(len as u32)?;
        self.output.write(&self.buf[..len])?;
        self.output.write_u32::<LittleEndian>(0xf1ee1001)?; // TODO: Magic.

        Ok(())
    }

    fn collect(&mut self) -> Result<Message> {
        let len = self.input.read_i32::<LittleEndian>()? as usize;
        self.input.read_exact(&mut self.buf[..len])?;

        prost::Message::decode(&self.buf[..len]).map_err(invalid_data_error)
    }

    fn encode<M>(&mut self, msg: M) -> Result<()>
    where
        M: prost::Message,
    {
        msg.encode(&mut self.buf).map_err(invalid_data_error)
    }
}

fn invalid_data_error<E>(err: E) -> std::io::Error
where
    E: Into<Box<dyn std::error::Error + Send + Sync>>,
{
    return std::io::Error::new(std::io::ErrorKind::InvalidData, err);
}
<|endoftext|>


<commit_before>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Compiles the profiler part of the `compiler-rt` library.
//!
//! See the build.rs for libcompiler_builtins crate for details.

extern crate cc;

use std::env;
use std::path::Path;

fn main() {
    let target = env::var("TARGET").expect("TARGET was not set");
    let cfg = &mut cc::Build::new();

    let mut profile_sources = vec!["GCDAProfiling.c",
                                   "InstrProfiling.c",
                                   "InstrProfilingBuffer.c",
                                   "InstrProfilingFile.c",
                                   "InstrProfilingMerge.c",
                                   "InstrProfilingMergeFile.c",
                                   "InstrProfilingNameVar.c",
                                   "InstrProfilingPlatformDarwin.c",
                                   "InstrProfilingPlatformLinux.c",
                                   "InstrProfilingPlatformOther.c",
                                   "InstrProfilingRuntime.cc",
                                   "InstrProfilingUtil.c",
                                   "InstrProfilingValue.c",
                                   "InstrProfilingWriter.c"];

    if target.contains("msvc") {
        // Don't pull in extra libraries on MSVC
        cfg.flag("/Zl");
        profile_sources.push("WindowsMMap.c");
        cfg.define("strdup", Some("_strdup"));
        cfg.define("open", Some("_open"));
        cfg.define("fdopen", Some("_fdopen"));
    } else {
        // Turn off various features of gcc and such, mostly copying
        // compiler-rt's build system already
        cfg.flag("-fno-builtin");
        cfg.flag("-fvisibility=hidden");
        cfg.flag("-fomit-frame-pointer");
        cfg.flag("-ffreestanding");
        cfg.define("VISIBILITY_HIDDEN", None);
        cfg.define("COMPILER_RT_HAS_UNAME", Some("1"));
    }

    for src in profile_sources {
        cfg.file(Path::new("../libcompiler_builtins/compiler-rt/lib/profile").join(src));
    }

    cfg.compile("profiler-rt");
}
<commit_msg>pgo: Blindly try to fix Windows build.<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Compiles the profiler part of the `compiler-rt` library.
//!
//! See the build.rs for libcompiler_builtins crate for details.

extern crate cc;

use std::env;
use std::path::Path;

fn main() {
    let target = env::var("TARGET").expect("TARGET was not set");
    let cfg = &mut cc::Build::new();

    let mut profile_sources = vec!["GCDAProfiling.c",
                                   "InstrProfiling.c",
                                   "InstrProfilingBuffer.c",
                                   "InstrProfilingFile.c",
                                   "InstrProfilingMerge.c",
                                   "InstrProfilingMergeFile.c",
                                   "InstrProfilingNameVar.c",
                                   "InstrProfilingPlatformDarwin.c",
                                   "InstrProfilingPlatformLinux.c",
                                   "InstrProfilingPlatformOther.c",
                                   "InstrProfilingRuntime.cc",
                                   "InstrProfilingUtil.c",
                                   "InstrProfilingValue.c",
                                   "InstrProfilingWriter.c"];

    if target.contains("msvc") {
        // Don't pull in extra libraries on MSVC
        cfg.flag("/Zl");
        profile_sources.push("WindowsMMap.c");
        cfg.define("strdup", Some("_strdup"));
        cfg.define("open", Some("_open"));
        cfg.define("fdopen", Some("_fdopen"));
        cfg.define("getpid", Some("_getpid"));
        cfg.define("fileno", Some("_fileno"));
    } else {
        // Turn off various features of gcc and such, mostly copying
        // compiler-rt's build system already
        cfg.flag("-fno-builtin");
        cfg.flag("-fvisibility=hidden");
        cfg.flag("-fomit-frame-pointer");
        cfg.flag("-ffreestanding");
        cfg.define("VISIBILITY_HIDDEN", None);
        cfg.define("COMPILER_RT_HAS_UNAME", Some("1"));
    }

    for src in profile_sources {
        cfg.file(Path::new("../libcompiler_builtins/compiler-rt/lib/profile").join(src));
    }

    cfg.compile("profiler-rt");
}
<|endoftext|>


<commit_before>use std::kinds::marker;
use std::str;
use libc;

use {raw, Oid, Repository, Error, Signature, Tree};

/// A structure to represent a git [commit][1]
///
/// [1]: http://git-scm.com/book/en/Git-Internals-Git-Objects
pub struct Commit<'a> {
    raw: *mut raw::git_commit,
    marker1: marker::ContravariantLifetime<'a>,
    marker2: marker::NoSend,
    marker3: marker::NoSync,
}

/// An iterator over the parent commits of a commit.
pub struct Parents<'a, 'b:'a> {
    cur: uint,
    max: uint,
    commit: &'a Commit<'b>,
}

impl<'a> Commit<'a> {
    /// Create a new object from its raw component.
    ///
    /// This method is unsafe as there is no guarantee that `raw` is a valid
    /// pointer.
    pub unsafe fn from_raw(_repo: &Repository,
                           raw: *mut raw::git_commit) -> Commit {
        Commit {
            raw: raw,
            marker1: marker::ContravariantLifetime,
            marker2: marker::NoSend,
            marker3: marker::NoSync,
        }
    }

    /// Get the id (SHA1) of a repository commit
    pub fn id(&self) -> Oid {
        unsafe { Oid::from_raw(raw::git_commit_id(&*self.raw)) }
    }

    /// Get the id of the tree pointed to by this commit.
    ///
    /// No attempts are made to fetch an object from the
    pub fn tree_id(&self) -> Oid {
        unsafe { Oid::from_raw(raw::git_commit_tree_id(&*self.raw)) }
    }

    /// Get access to the underlying raw pointer.
    pub fn raw(&self) -> *mut raw::git_commit { self.raw }

    /// Get the full message of a commit.
    ///
    /// The returned message will be slightly prettified by removing any
    /// potential leading newlines.
    ///
    /// `None` will be returned if the message is not valid utf-8
    pub fn message(&self) -> Option<&str> {
        str::from_utf8(self.message_bytes()).ok()
    }

    /// Get the full message of a commit as a byte slice.
    ///
    /// The returned message will be slightly prettified by removing any
    /// potential leading newlines.
    pub fn message_bytes(&self) -> &[u8] {
        unsafe {
            ::opt_bytes(self, raw::git_commit_message(&*self.raw)).unwrap()
        }
    }

    /// Get the encoding for the message of a commit, as a string representing a
    /// standard encoding name.
    ///
    /// `None` will be returned if the encoding is not known
    pub fn message_encoding(&self) -> Option<&str> {
        let bytes = unsafe {
            ::opt_bytes(self, raw::git_commit_message(&*self.raw))
        };
        bytes.map(|b| str::from_utf8(b).unwrap())
    }

    /// Get the full raw message of a commit.
    ///
    /// `None` will be returned if the message is not valid utf-8
    pub fn message_raw(&self) -> Option<&str> {
        str::from_utf8(self.message_raw_bytes()).ok()
    }

    /// Get the full raw message of a commit.
    pub fn message_raw_bytes(&self) -> &[u8] {
        unsafe {
            ::opt_bytes(self, raw::git_commit_message_raw(&*self.raw)).unwrap()
        }
    }

    /// Get the full raw text of the commit header.
    ///
    /// `None` will be returned if the message is not valid utf-8
    pub fn raw_header(&self) -> Option<&str> {
        str::from_utf8(self.raw_header_bytes()).ok()
    }

    /// Get the full raw text of the commit header.
    pub fn raw_header_bytes(&self) -> &[u8] {
        unsafe {
            ::opt_bytes(self, raw::git_commit_raw_header(&*self.raw)).unwrap()
        }
    }

    /// Get the short "summary" of the git commit message.
    ///
    /// The returned message is the summary of the commit, comprising the first
    /// paragraph of the message with whitespace trimmed and squashed.
    ///
    /// `None` may be returned if an error occurs or if the summary is not valid
    /// utf-8.
    pub fn summary(&mut self) -> Option<&str> {
        self.summary_bytes().and_then(|s| str::from_utf8(s).ok())
    }

    /// Get the short "summary" of the git commit message.
    ///
    /// The returned message is the summary of the commit, comprising the first
    /// paragraph of the message with whitespace trimmed and squashed.
    ///
    /// `None` may be returned if an error occurs
    pub fn summary_bytes(&mut self) -> Option<&[u8]> {
        unsafe { ::opt_bytes(self, raw::git_commit_summary(self.raw)) }
    }

    /// Get the commit time (i.e. committer time) of a commit.
    ///
    /// The first element of the tuple is the time, in seconds, since the epoch.
    /// The second element is the offset, in minutes, of the time zone of the
    /// committer's preferred time zone.
    pub fn time(&self) -> (u64, int) {
        unsafe {
            (raw::git_commit_time(&*self.raw) as u64,
             raw::git_commit_time_offset(&*self.raw) as int)
        }
    }

    /// Creates a new iterator over the parents of this commit.
    pub fn parents<'b>(&'b self) -> Parents<'b, 'a> {
        Parents {
            cur: 0,
            max: unsafe { raw::git_commit_parentcount(&*self.raw) as uint },
            commit: self,
        }
    }

    /// Get the author of this commit.
    pub fn author(&self) -> Signature {
        unsafe {
            let ptr = raw::git_commit_author(&*self.raw);
            Signature::from_raw_const(self, ptr)
        }
    }

    /// Get the committer of this commit.
    pub fn committer(&self) -> Signature {
        unsafe {
            let ptr = raw::git_commit_committer(&*self.raw);
            Signature::from_raw_const(self, ptr)
        }
    }

    /// Amend this existing commit with all non-`None` values
    ///
    /// This creates a new commit that is exactly the same as the old commit,
    /// except that any non-`None` values will be updated. The new commit has
    /// the same parents as the old commit.
    ///
    /// For information about `update_ref`, see `new`.
    pub fn amend(&self,
                 update_ref: Option<&str>,
                 author: Option<&Signature>,
                 committer: Option<&Signature>,
                 message_encoding: Option<&str>,
                 message: Option<&str>,
                 tree: Option<&Tree<'a>>) -> Result<Oid, Error> {
        let mut raw = raw::git_oid { id: [0, ..raw::GIT_OID_RAWSZ] };
        unsafe {
            try_call!(raw::git_commit_amend(&mut raw,
                                            &*self.raw(),
                                            update_ref.map(|s| s.to_c_str()),
                                            author.map(|s| &*s.raw()),
                                            committer.map(|s| &*s.raw()),
                                            message_encoding.map(|s| s.to_c_str()),
                                            message.map(|s| s.to_c_str()),
                                            tree.map(|t| &*t.raw())));
            Ok(Oid::from_raw(&raw))
        }
    }
}

impl<'a, 'b> Iterator<Commit<'a>> for Parents<'b, 'a> {
    fn next(&mut self) -> Option<Commit<'a>> {
        if self.cur == self.max { return None }
        self.cur += 1;
        let mut raw = 0 as *mut raw::git_commit;
        assert_eq!(unsafe {
            raw::git_commit_parent(&mut raw, &*self.commit.raw,
                                   (self.cur - 1) as libc::c_uint)
        }, 0);
        Some(Commit {
            raw: raw,
            marker1: marker::ContravariantLifetime,
            marker2: marker::NoSend,
            marker3: marker::NoSync,
        })
    }
}

#[unsafe_destructor]
impl<'a> Drop for Commit<'a> {
    fn drop(&mut self) {
        unsafe { raw::git_commit_free(self.raw) }
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn smoke() {
        let (_td, repo) = ::test::repo_init();
        let head = repo.head().unwrap();
        let target = head.target().unwrap();
        let mut commit = repo.find_commit(target).unwrap();
        assert_eq!(commit.message(), Some("initial"));
        assert_eq!(commit.id(), target);
        commit.message_raw().unwrap();
        commit.raw_header().unwrap();
        commit.message_encoding();
        commit.summary().unwrap();
        commit.tree_id();
        assert_eq!(commit.parents().count(), 0);

        assert_eq!(commit.author().name(), Some("name"));
        assert_eq!(commit.author().email(), Some("email"));
        assert_eq!(commit.committer().name(), Some("name"));
        assert_eq!(commit.committer().email(), Some("email"));

        let sig = repo.signature().unwrap();
        let tree = repo.find_tree(commit.tree_id()).unwrap();
        let id = repo.commit(Some("HEAD"), &sig, &sig, "bar", &tree,
                             &[&commit]).unwrap();
        let head = repo.find_commit(id).unwrap();

        let new_head = head.amend(Some("HEAD"), None, None, None,
                                  Some("new message"), None).unwrap();
        let new_head = repo.find_commit(new_head).unwrap();
        assert_eq!(new_head.message(), Some("new message"));
    }
}

<commit_msg>Improve parent iterators on commits<commit_after>use std::kinds::marker;
use std::iter::Range;
use std::str;
use libc;

use {raw, Oid, Repository, Error, Signature, Tree};

/// A structure to represent a git [commit][1]
///
/// [1]: http://git-scm.com/book/en/Git-Internals-Git-Objects
pub struct Commit<'a> {
    raw: *mut raw::git_commit,
    marker1: marker::ContravariantLifetime<'a>,
    marker2: marker::NoSend,
    marker3: marker::NoSync,
}

/// An iterator over the parent commits of a commit.
pub struct Parents<'a, 'b:'a> {
    range: Range<uint>,
    commit: &'a Commit<'b>,
}

/// An iterator over the parent commits' ids of a commit.
pub struct ParentIds<'a> {
    range: Range<uint>,
    commit: &'a Commit<'a>,
}

impl<'a> Commit<'a> {
    /// Create a new object from its raw component.
    ///
    /// This method is unsafe as there is no guarantee that `raw` is a valid
    /// pointer.
    pub unsafe fn from_raw(_repo: &Repository,
                           raw: *mut raw::git_commit) -> Commit {
        Commit {
            raw: raw,
            marker1: marker::ContravariantLifetime,
            marker2: marker::NoSend,
            marker3: marker::NoSync,
        }
    }

    /// Get the id (SHA1) of a repository commit
    pub fn id(&self) -> Oid {
        unsafe { Oid::from_raw(raw::git_commit_id(&*self.raw)) }
    }

    /// Get the id of the tree pointed to by this commit.
    ///
    /// No attempts are made to fetch an object from the
    pub fn tree_id(&self) -> Oid {
        unsafe { Oid::from_raw(raw::git_commit_tree_id(&*self.raw)) }
    }

    /// Get access to the underlying raw pointer.
    pub fn raw(&self) -> *mut raw::git_commit { self.raw }

    /// Get the full message of a commit.
    ///
    /// The returned message will be slightly prettified by removing any
    /// potential leading newlines.
    ///
    /// `None` will be returned if the message is not valid utf-8
    pub fn message(&self) -> Option<&str> {
        str::from_utf8(self.message_bytes()).ok()
    }

    /// Get the full message of a commit as a byte slice.
    ///
    /// The returned message will be slightly prettified by removing any
    /// potential leading newlines.
    pub fn message_bytes(&self) -> &[u8] {
        unsafe {
            ::opt_bytes(self, raw::git_commit_message(&*self.raw)).unwrap()
        }
    }

    /// Get the encoding for the message of a commit, as a string representing a
    /// standard encoding name.
    ///
    /// `None` will be returned if the encoding is not known
    pub fn message_encoding(&self) -> Option<&str> {
        let bytes = unsafe {
            ::opt_bytes(self, raw::git_commit_message(&*self.raw))
        };
        bytes.map(|b| str::from_utf8(b).unwrap())
    }

    /// Get the full raw message of a commit.
    ///
    /// `None` will be returned if the message is not valid utf-8
    pub fn message_raw(&self) -> Option<&str> {
        str::from_utf8(self.message_raw_bytes()).ok()
    }

    /// Get the full raw message of a commit.
    pub fn message_raw_bytes(&self) -> &[u8] {
        unsafe {
            ::opt_bytes(self, raw::git_commit_message_raw(&*self.raw)).unwrap()
        }
    }

    /// Get the full raw text of the commit header.
    ///
    /// `None` will be returned if the message is not valid utf-8
    pub fn raw_header(&self) -> Option<&str> {
        str::from_utf8(self.raw_header_bytes()).ok()
    }

    /// Get the full raw text of the commit header.
    pub fn raw_header_bytes(&self) -> &[u8] {
        unsafe {
            ::opt_bytes(self, raw::git_commit_raw_header(&*self.raw)).unwrap()
        }
    }

    /// Get the short "summary" of the git commit message.
    ///
    /// The returned message is the summary of the commit, comprising the first
    /// paragraph of the message with whitespace trimmed and squashed.
    ///
    /// `None` may be returned if an error occurs or if the summary is not valid
    /// utf-8.
    pub fn summary(&mut self) -> Option<&str> {
        self.summary_bytes().and_then(|s| str::from_utf8(s).ok())
    }

    /// Get the short "summary" of the git commit message.
    ///
    /// The returned message is the summary of the commit, comprising the first
    /// paragraph of the message with whitespace trimmed and squashed.
    ///
    /// `None` may be returned if an error occurs
    pub fn summary_bytes(&mut self) -> Option<&[u8]> {
        unsafe { ::opt_bytes(self, raw::git_commit_summary(self.raw)) }
    }

    /// Get the commit time (i.e. committer time) of a commit.
    ///
    /// The first element of the tuple is the time, in seconds, since the epoch.
    /// The second element is the offset, in minutes, of the time zone of the
    /// committer's preferred time zone.
    pub fn time(&self) -> (u64, int) {
        unsafe {
            (raw::git_commit_time(&*self.raw) as u64,
             raw::git_commit_time_offset(&*self.raw) as int)
        }
    }

    /// Creates a new iterator over the parents of this commit.
    pub fn parents<'b>(&'b self) -> Parents<'b, 'a> {
        let max = unsafe { raw::git_commit_parentcount(&*self.raw) as uint };
        Parents {
            range: range(0, max),
            commit: self,
        }
    }

    /// Creates a new iterator over the parents of this commit.
    pub fn parent_ids(&self) -> ParentIds {
        let max = unsafe { raw::git_commit_parentcount(&*self.raw) as uint };
        ParentIds {
            range: range(0, max),
            commit: self,
        }
    }

    /// Get the author of this commit.
    pub fn author(&self) -> Signature {
        unsafe {
            let ptr = raw::git_commit_author(&*self.raw);
            Signature::from_raw_const(self, ptr)
        }
    }

    /// Get the committer of this commit.
    pub fn committer(&self) -> Signature {
        unsafe {
            let ptr = raw::git_commit_committer(&*self.raw);
            Signature::from_raw_const(self, ptr)
        }
    }

    /// Amend this existing commit with all non-`None` values
    ///
    /// This creates a new commit that is exactly the same as the old commit,
    /// except that any non-`None` values will be updated. The new commit has
    /// the same parents as the old commit.
    ///
    /// For information about `update_ref`, see `new`.
    pub fn amend(&self,
                 update_ref: Option<&str>,
                 author: Option<&Signature>,
                 committer: Option<&Signature>,
                 message_encoding: Option<&str>,
                 message: Option<&str>,
                 tree: Option<&Tree<'a>>) -> Result<Oid, Error> {
        let mut raw = raw::git_oid { id: [0, ..raw::GIT_OID_RAWSZ] };
        unsafe {
            try_call!(raw::git_commit_amend(&mut raw,
                                            &*self.raw(),
                                            update_ref.map(|s| s.to_c_str()),
                                            author.map(|s| &*s.raw()),
                                            committer.map(|s| &*s.raw()),
                                            message_encoding.map(|s| s.to_c_str()),
                                            message.map(|s| s.to_c_str()),
                                            tree.map(|t| &*t.raw())));
            Ok(Oid::from_raw(&raw))
        }
    }

    /// Get the specified parent of the commit.
    ///
    /// Use the `parents` iterator to return an iterator over all parents.
    pub fn parent(&self, i: uint) -> Result<Commit<'a>, Error> {
        unsafe {
            let mut raw = 0 as *mut raw::git_commit;
            try_call!(raw::git_commit_parent(&mut raw, &*self.raw,
                                             i as libc::c_uint));
            Ok(Commit {
                raw: raw,
                marker1: marker::ContravariantLifetime,
                marker2: marker::NoSend,
                marker3: marker::NoSync,
            })
        }
    }

    /// Get the specified parent id of the commit.
    ///
    /// This is different from `parent`, which will attemptstempt to load the
    /// parent commit from the ODB.
    ///
    /// Use the `parent_ids` iterator to return an iterator over all parents.
    pub fn parent_id(&self, i: uint) -> Result<Oid, Error> {
        unsafe {
            let id = raw::git_commit_parent_id(self.raw, i as libc::c_uint);
            if id.is_null() {
                Err(Error::from_str("parent index out of bounds"))
            } else {
                Ok(Oid::from_raw(id))
            }
        }
    }
}

impl<'a, 'b> Iterator<Commit<'a>> for Parents<'b, 'a> {
    fn next(&mut self) -> Option<Commit<'a>> {
        self.range.next().map(|i| self.commit.parent(i).unwrap())
    }
    fn size_hint(&self) -> (uint, Option<uint>) { self.range.size_hint() }
}

impl<'a, 'b> DoubleEndedIterator<Commit<'a>> for Parents<'b, 'a> {
    fn next_back(&mut self) -> Option<Commit<'a>> {
        self.range.next_back().map(|i| self.commit.parent(i).unwrap())
    }
}

impl<'a, 'b> ExactSizeIterator<Commit<'a>> for Parents<'b, 'a> {}

impl<'a> Iterator<Oid> for ParentIds<'a> {
    fn next(&mut self) -> Option<Oid> {
        self.range.next().map(|i| self.commit.parent_id(i).unwrap())
    }
    fn size_hint(&self) -> (uint, Option<uint>) { self.range.size_hint() }
}

impl<'a> DoubleEndedIterator<Oid> for ParentIds<'a> {
    fn next_back(&mut self) -> Option<Oid> {
        self.range.next_back().map(|i| self.commit.parent_id(i).unwrap())
    }
}

impl<'a> ExactSizeIterator<Oid> for ParentIds<'a> {}

#[unsafe_destructor]
impl<'a> Drop for Commit<'a> {
    fn drop(&mut self) {
        unsafe { raw::git_commit_free(self.raw) }
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn smoke() {
        let (_td, repo) = ::test::repo_init();
        let head = repo.head().unwrap();
        let target = head.target().unwrap();
        let mut commit = repo.find_commit(target).unwrap();
        assert_eq!(commit.message(), Some("initial"));
        assert_eq!(commit.id(), target);
        commit.message_raw().unwrap();
        commit.raw_header().unwrap();
        commit.message_encoding();
        commit.summary().unwrap();
        commit.tree_id();
        assert_eq!(commit.parents().count(), 0);

        assert_eq!(commit.author().name(), Some("name"));
        assert_eq!(commit.author().email(), Some("email"));
        assert_eq!(commit.committer().name(), Some("name"));
        assert_eq!(commit.committer().email(), Some("email"));

        let sig = repo.signature().unwrap();
        let tree = repo.find_tree(commit.tree_id()).unwrap();
        let id = repo.commit(Some("HEAD"), &sig, &sig, "bar", &tree,
                             &[&commit]).unwrap();
        let head = repo.find_commit(id).unwrap();

        let new_head = head.amend(Some("HEAD"), None, None, None,
                                  Some("new message"), None).unwrap();
        let new_head = repo.find_commit(new_head).unwrap();
        assert_eq!(new_head.message(), Some("new message"));
    }
}

<|endoftext|>


<commit_before>/* The info hash is stored in mongo, which is not binary string safe apparently,
so it needs to be parsed to a hexadecimal string rather than a binary one. */
pub fn parse_info_hash(input: &str) -> Result<String, String> {
    let mut output = String::new();
    let mut input_iterator = input.as_bytes().into_iter();
    while let Some(i) = input_iterator.next() {
        match i {
            &37u8 => { //% in ASCII
                let hexdigits = (input_iterator.next(), input_iterator.next());
                match hexdigits {
                    (Some(hex1), Some(hex2)) => {
                        output.push_str(&String::from_utf8_lossy(&[hex1.clone()]).to_lowercase());
                        output.push_str(&String::from_utf8_lossy(&[hex2.clone()]).to_lowercase());
                    },
                    _ => {},
                }
            }
            _ => {
                output.push_str(&format!("{:X}", i));
            }
        }
    }

    match output.len() {
        40 => return Ok(output),
        _ => return Err("Hash is invalid (too short).".to_string()),
    }
}

#[test]
fn parse_info_hash_test() {
    let output = parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx%9A").unwrap();
    assert_eq!(&output, "123456789abcdef123456789abcdef123456789a");
}
<commit_msg>parse_info_hash: Add failure test cases<commit_after>/* The info hash is stored in mongo, which is not binary string safe apparently,
so it needs to be parsed to a hexadecimal string rather than a binary one. */
pub fn parse_info_hash(input: &str) -> Result<String, String> {
    let mut output = String::new();
    let mut input_iterator = input.as_bytes().into_iter();
    while let Some(i) = input_iterator.next() {
        match i {
            &37u8 => { //% in ASCII
                let hexdigits = (input_iterator.next(), input_iterator.next());
                match hexdigits {
                    (Some(hex1), Some(hex2)) => {
                        output.push_str(&String::from_utf8_lossy(&[hex1.clone()]).to_lowercase());
                        output.push_str(&String::from_utf8_lossy(&[hex2.clone()]).to_lowercase());
                    },
                    _ => {},
                }
            }
            _ => {
                output.push_str(&format!("{:X}", i));
            }
        }
    }

    match output.len() {
        40 => return Ok(output),
        _ => return Err("Hash is invalid (too short).".to_string()),
    }
}

#[test]
fn parse_info_hash_test() {
    // Success
    assert_eq!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx%9A").unwrap(), "123456789abcdef123456789abcdef123456789a");

    // Failures
    assert!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx").is_err()); // too short
    assert!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vxab").is_err()); // too long
    assert!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx%ZA").is_err()); // invalid percent encoding
}
<|endoftext|>