<commit_before>//! Handles the IPC (Inter-process Communication) protocol used by the renderer client and server.
//!
//! This is the common "language" between client/server. It defines both the protocol for
//! connection and for send messages back and forth.

mod async_ipc_receiver;
mod messages;

pub use messages::*;

use std::io;

use thiserror::Error;
use serde::{Serialize, Deserialize};
use ipc_channel::ipc::{self, IpcOneShotServer, IpcSender};

use crate::renderer_client::RendererServerProcess;

use async_ipc_receiver::AsyncIpcReceiver;

#[derive(Debug, Error)]
#[error(transparent)]
pub enum ConnectionError {
    IpcChannelError(#[from] ipc_channel::Error),
    IOError(#[from] io::Error),
    JoinError(#[from] tokio::task::JoinError),
}

#[derive(Debug, Error)]
#[error(transparent)]
pub enum SendError {
    IpcChannelError(#[from] ipc_channel::Error),
    // IpcError does not implement Display for some reason...
    #[error("IPC Error: {0:?}")]
    IpcError(ipc_channel::ipc::IpcError),
}

impl From<ipc_channel::ipc::IpcError> for SendError {
    fn from(err: ipc_channel::ipc::IpcError) -> Self {
        SendError::IpcError(err)
    }
}

/// Represents either a sender from the connection handshake or a response
#[derive(Debug, Serialize, Deserialize)]
enum HandshakeResponse {
    /// Represents that the handshake has been completed successfully.
    ///
    /// Provides the sender that the client should use to send further requests to the server. This
    /// message is only ever sent **once** for each client at the beginning when the connection is
    /// first established.
    HandshakeFinish(IpcSender<ClientRequest>),

    /// A response from the server sent in response to a request
    Response(ServerResponse),
}

/// Represents the client side of the IPC connection
#[derive(Debug)]
pub struct ClientConnection {
    sender: IpcSender<ClientRequest>,
    receiver: AsyncIpcReceiver<HandshakeResponse>,
}

impl ClientConnection {
    pub async fn new(process: &mut RendererServerProcess) -> Result<Self, ConnectionError> {
        // Send the oneshot token to the server which will then respond with its own oneshot token
        let (server, server_name) = IpcOneShotServer::new()?;
        process.writeln(server_name).await?;

        let (receiver, response): (_, HandshakeResponse) = tokio::task::spawn_blocking(|| {
            server.accept()
        }).await??;

        let sender = match response {
            HandshakeResponse::HandshakeFinish(sender) => sender,
            _ => unreachable!("bug: server did not send back Sender at the end of handshake"),
        };
        let receiver = AsyncIpcReceiver::new(receiver);

        Ok(Self {sender, receiver})
    }

    pub async fn send(&mut self, req: ClientRequest) -> Result<ServerResponse, SendError> {
        self.sender.send(req)?;
        let response = self.receiver.recv().await?;
        match response {
            HandshakeResponse::Response(response) => Ok(response),
            _ => unreachable!("bug: server did not send response after request"),
        }
    }
}

/// Represents the server side of the IPC connection
#[derive(Debug)]
pub struct ServerConnection {
    sender: IpcSender<HandshakeResponse>,
    receiver: AsyncIpcReceiver<ClientRequest>,
}

impl ServerConnection {
    /// Establishes a connection with IPC channel one shot server with the given name
    pub fn connect(one_shot_name: String) -> Result<Self, ConnectionError> {
        let (server_sender, receiver) = ipc::channel()?;
        let sender = IpcSender::connect(one_shot_name)?;

        // Finish handshake by giving client a sender it can use to send messages to the server
        sender.send(HandshakeResponse::HandshakeFinish(server_sender))?;

        let receiver = AsyncIpcReceiver::new(receiver);

        Ok(Self {sender, receiver})
    }
}
<commit_msg>Added recv and send methods to ServerConnection<commit_after>//! Handles the IPC (Inter-process Communication) protocol used by the renderer client and server.
//!
//! This is the common "language" between client/server. It defines both the protocol for
//! connection and for send messages back and forth.

mod async_ipc_receiver;
mod messages;

pub use messages::*;

use std::io;

use thiserror::Error;
use serde::{Serialize, Deserialize};
use ipc_channel::ipc::{self, IpcOneShotServer, IpcSender, IpcError};

use crate::renderer_client::RendererServerProcess;

use async_ipc_receiver::AsyncIpcReceiver;

#[derive(Debug, Error)]
#[error(transparent)]
pub enum ConnectionError {
    IpcChannelError(#[from] ipc_channel::Error),
    IOError(#[from] io::Error),
    JoinError(#[from] tokio::task::JoinError),
}

#[derive(Debug, Error)]
#[error(transparent)]
pub enum SendError {
    IpcChannelError(#[from] ipc_channel::Error),
    // IpcError does not implement Display for some reason...
    #[error("IPC Error: {0:?}")]
    IpcError(IpcError),
}

impl From<ipc_channel::ipc::IpcError> for SendError {
    fn from(err: ipc_channel::ipc::IpcError) -> Self {
        SendError::IpcError(err)
    }
}

/// Represents either a sender from the connection handshake or a response
#[derive(Debug, Serialize, Deserialize)]
enum HandshakeResponse {
    /// Represents that the handshake has been completed successfully.
    ///
    /// Provides the sender that the client should use to send further requests to the server. This
    /// message is only ever sent **once** for each client at the beginning when the connection is
    /// first established.
    HandshakeFinish(IpcSender<ClientRequest>),

    /// A response from the server sent in response to a request
    Response(ServerResponse),
}

/// Represents the client side of the IPC connection
#[derive(Debug)]
pub struct ClientConnection {
    sender: IpcSender<ClientRequest>,
    receiver: AsyncIpcReceiver<HandshakeResponse>,
}

impl ClientConnection {
    pub async fn new(process: &mut RendererServerProcess) -> Result<Self, ConnectionError> {
        // Send the oneshot token to the server which will then respond with its own oneshot token
        let (server, server_name) = IpcOneShotServer::new()?;
        process.writeln(server_name).await?;

        let (receiver, response): (_, HandshakeResponse) = tokio::task::spawn_blocking(|| {
            server.accept()
        }).await??;

        let sender = match response {
            HandshakeResponse::HandshakeFinish(sender) => sender,
            _ => unreachable!("bug: server did not send back Sender at the end of handshake"),
        };
        let receiver = AsyncIpcReceiver::new(receiver);

        Ok(Self {sender, receiver})
    }

    /// Sends a request and awaits the response
    pub async fn send(&mut self, req: ClientRequest) -> Result<ServerResponse, SendError> {
        self.sender.send(req)?;
        let response = self.receiver.recv().await?;
        match response {
            HandshakeResponse::Response(response) => Ok(response),
            _ => unreachable!("bug: server did not send response after request"),
        }
    }
}

/// Represents the server side of the IPC connection
#[derive(Debug)]
pub struct ServerConnection {
    sender: IpcSender<HandshakeResponse>,
    receiver: AsyncIpcReceiver<ClientRequest>,
}

impl ServerConnection {
    /// Establishes a connection with the IPC channel one shot server with the given name
    pub fn connect(one_shot_name: String) -> Result<Self, ConnectionError> {
        let (server_sender, receiver) = ipc::channel()?;
        let sender = IpcSender::connect(one_shot_name)?;

        // Finish handshake by giving client a sender it can use to send messages to the server
        sender.send(HandshakeResponse::HandshakeFinish(server_sender))?;

        let receiver = AsyncIpcReceiver::new(receiver);

        Ok(Self {sender, receiver})
    }

    /// Returns the next request, waiting until one is available
    pub async fn recv(&mut self) -> Result<ClientRequest, IpcError> {
        self.receiver.recv().await
    }

    /// Sends a response to the client
    ///
    /// This should only ever be done in response to a request
    pub fn send(&mut self, res: ServerResponse) -> Result<(), ipc_channel::Error> {
        self.sender.send(HandshakeResponse::Response(res))
    }
}
<|endoftext|>


<commit_before>// The MIT License (MIT)

// Copyright (c) 2014 Y. T. CHUNG <zonyitoo@gmail.com>

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

//! TcpRelay implementation

use std::io::{self, Read, Write};
use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
use std::rc::Rc;
use std::mem;
use std::cmp;

use crypto::cipher;
use crypto::CryptoMode;
use relay::socks5::Address;
use relay::{BoxIoFuture, boxed_future};
use relay::dns_resolver::DnsResolver;
use config::{ServerConfig, ServerAddr};

use tokio_core::net::TcpStream;
use tokio_core::reactor::Handle;
use tokio_core::io::{read_exact, write_all};
use tokio_core::io::{ReadHalf, WriteHalf};
use tokio_core::io::Io;

use futures::{self, Future, Poll};

use net2::TcpBuilder;

use ip::IpAddr;

use self::stream::{EncryptedWriter, DecryptedReader};

pub mod local;
mod socks5_local;
mod http_local;
pub mod server;
mod stream;
pub mod client;

/// Directions in the tunnel
#[derive(Debug, Copy, Clone)]
pub enum TunnelDirection {
    /// Client -> Server
    Client2Server,
    /// Client <- Server
    Server2Client,
}

/// ReadHalf of TcpStream with decryption
pub type DecryptedHalf = DecryptedReader<ReadHalf<TcpStream>>;
/// WriteHalf of TcpStream with encryption
pub type EncryptedHalf = EncryptedWriter<WriteHalf<TcpStream>>;

/// Boxed future of DecryptedHalf
pub type DecryptedHalfFut = BoxIoFuture<DecryptedHalf>;
/// Boxed future of EncryptedHalf
pub type EncryptedHalfFut = BoxIoFuture<EncryptedHalf>;

fn connect_proxy_server(handle: &Handle, svr_cfg: Rc<ServerConfig>) -> BoxIoFuture<TcpStream> {
    match svr_cfg.addr() {
        &ServerAddr::SocketAddr(ref addr) => TcpStream::connect(addr, handle).boxed(),
        &ServerAddr::DomainName(ref domain, port) => {
            let handle = handle.clone();
            let fut = DnsResolver::resolve(&domain[..]).and_then(move |sockaddr| {
                let sockaddr = match sockaddr {
                    IpAddr::V4(v4) => SocketAddr::V4(SocketAddrV4::new(v4, port)),
                    IpAddr::V6(v6) => SocketAddr::V6(SocketAddrV6::new(v6, port, 0, 0)),
                };
                TcpStream::connect(&sockaddr, &handle)
            });
            Box::new(fut)
        }
    }
}

/// Handshake logic for ShadowSocks Client
pub fn proxy_server_handshake(remote_stream: TcpStream,
                              svr_cfg: Rc<ServerConfig>,
                              relay_addr: Address)
                              -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> {
    let fut = proxy_handshake(remote_stream, svr_cfg).and_then(|(r_fut, w_fut)| {;
        let w_fut = w_fut.and_then(move |enc_w| {
            trace!("Got encrypt stream and going to send addr: {:?}",
                   relay_addr);

            // Send relay address to remote
            let local_buf = Vec::new();
            relay_addr.write_to(local_buf)
                .and_then(|buf| enc_w.write_all_encrypted(buf))
                .map(|(w, _)| w)
        });

        Ok((r_fut, boxed_future(w_fut)))
    });
    boxed_future(fut)
}

/// ShadowSocks Client-Server handshake protocol
/// Exchange cipher IV and creates stream wrapper
pub fn proxy_handshake(remote_stream: TcpStream,
                       svr_cfg: Rc<ServerConfig>)
                       -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> {
    let fut = futures::lazy(|| Ok(remote_stream.split())).and_then(|(r, w)| {

        let svr_cfg_cloned = svr_cfg.clone();

        let enc = {
            // Encrypt data to remote server

            // Send initialize vector to remote and create encryptor

            let local_iv = svr_cfg.method().gen_init_vec();
            trace!("Going to send initialize vector: {:?}", local_iv);

            write_all(w, local_iv).and_then(move |(w, local_iv)| {
                let encryptor = cipher::with_type(svr_cfg.method(),
                                                  svr_cfg.key(),
                                                  &local_iv[..],
                                                  CryptoMode::Encrypt);

                Ok(EncryptedWriter::new(w, encryptor))
            })
        };

        let dec = {
            let svr_cfg = svr_cfg_cloned;

            // Decrypt data from remote server
            let iv_len = svr_cfg.method().iv_size();
            read_exact(r, vec![0u8; iv_len]).and_then(move |(r, remote_iv)| {
                trace!("Got initialize vector {:?}", remote_iv);

                let decryptor = cipher::with_type(svr_cfg.method(),
                                                  svr_cfg.key(),
                                                  &remote_iv[..],
                                                  CryptoMode::Decrypt);
                let decrypt_stream = DecryptedReader::new(r, decryptor);

                Ok(decrypt_stream)
            })
        };

        Ok((boxed_future(dec), boxed_future(enc)))
    });

    boxed_future(fut)
}

/// Copy exactly N bytes by encryption
pub enum CopyExactEncrypted<R, W>
    where R: Read,
          W: Write
{
    Pending {
        reader: R,
        writer: EncryptedWriter<W>,
        buf: [u8; 4096],
        remain: usize,
        pos: usize,
        cap: usize,
        enc_buf: Vec<u8>,
    },
    Empty,
}

impl<R, W> CopyExactEncrypted<R, W>
    where R: Read,
          W: Write
{
    pub fn new(r: R, w: EncryptedWriter<W>, amt: usize) -> CopyExactEncrypted<R, W> {
        CopyExactEncrypted::Pending {
            reader: r,
            writer: w,
            buf: [0u8; 4096],
            remain: amt,
            pos: 0,
            cap: 0,
            enc_buf: Vec::new(),
        }
    }
}

impl<R, W> Future for CopyExactEncrypted<R, W>
    where R: Read,
          W: Write
{
    type Item = (R, EncryptedWriter<W>);
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        match self {
            &mut CopyExactEncrypted::Empty => panic!("poll after CopyExactEncrypted is finished"),
            &mut CopyExactEncrypted::Pending { ref mut reader,
                                               ref mut writer,
                                               ref mut buf,
                                               ref mut remain,
                                               ref mut pos,
                                               ref mut cap,
                                               ref mut enc_buf } => {
                loop {
                    // If our buffer is empty, then we need to read some data to
                    // continue.
                    if *pos == *cap && *remain != 0 {
                        let buf_len = cmp::min(*remain, buf.len());
                        let n = try_nb!(reader.read(&mut buf[..buf_len]));
                        if n == 0 {
                            // Unexpected EOF!
                            let err = io::Error::new(io::ErrorKind::UnexpectedEof, "Unexpected Eof");
                            return Err(err);
                        } else {
                            *pos = 0;
                            *remain -= n;

                            enc_buf.clear();
                            try!(writer.cipher_update(&buf[..n], enc_buf));
                            *cap = enc_buf.len();
                        }
                    }

                    // If our buffer has some data, let's write it out!
                    while *pos < *cap {
                        let i = try_nb!(writer.write(&enc_buf[*pos..*cap]).and_then(|x| writer.flush().map(|_| x)));
                        *pos += i;
                    }

                    // If we've written al the data and we've seen EOF, flush out the
                    // data and finish the transfer.
                    // done with the entire transfer.
                    if *pos == *cap && *remain == 0 {
                        try_nb!(writer.flush());
                        break; // The only path to execute the following logic
                    }
                }
            }
        }

        match mem::replace(self, CopyExactEncrypted::Empty) {
            CopyExactEncrypted::Pending { reader, writer, .. } => Ok((reader, writer).into()),
            CopyExactEncrypted::Empty => unreachable!(),
        }
    }
}

/// Copy all bytes from reader and write all encrypted data into writer
pub fn copy_exact_encrypted<R, W>(r: R, w: EncryptedWriter<W>, amt: usize) -> CopyExactEncrypted<R, W>
    where R: Read,
          W: Write
{
    CopyExactEncrypted::new(r, w, amt)
}

/// Establish tunnel between server and client
pub fn tunnel<CF, SF>(addr: Address, c2s: CF, s2c: SF) -> BoxIoFuture<()>
    where CF: Future<Item = u64, Error = io::Error> + 'static,
          SF: Future<Item = u64, Error = io::Error> + 'static
{
    let addr = Rc::new(addr);

    let cloned_addr = addr.clone();
    let c2s = c2s.then(move |res| {
        match res {
            Ok(amt) => {
                // Continue reading response from remote server
                trace!("Relay {} client -> server is finished, relayed {} bytes",
                       cloned_addr,
                       amt);

                Ok(TunnelDirection::Client2Server)
            }
            Err(err) => {
                error!("Relay {} client -> server aborted: {}", cloned_addr, err);
                Err(err)
            }
        }
    });

    let cloned_addr = addr.clone();
    let s2c = s2c.then(move |res| {
        match res {
            Ok(amt) => {
                trace!("Relay {} client <- server is finished, relayed {} bytes",
                       cloned_addr,
                       amt);

                Ok(TunnelDirection::Server2Client)
            }
            Err(err) => {
                error!("Relay {} client <- server aborted: {}", cloned_addr, err);
                Err(err)
            }
        }
    });

    let fut = c2s.select(s2c)
        .and_then(|(dir, _)| {
            match dir {
                TunnelDirection::Server2Client => trace!("client <- server is closed, abort connection"),
                TunnelDirection::Client2Server => trace!("server -> client is closed, abort connection"),
            }

            Ok(())
        })
        .map_err(|(err, _)| err);

    boxed_future(fut)
}

/// Read until EOF, and ignore
pub enum IgnoreUntilEnd<R: Read> {
    Pending { r: R, amt: u64 },
    Empty,
}

impl<R: Read> Future for IgnoreUntilEnd<R> {
    type Item = u64;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        match self {
            &mut IgnoreUntilEnd::Empty => panic!("poll IgnoreUntilEnd after it is finished"),
            &mut IgnoreUntilEnd::Pending { ref mut r, ref mut amt } => {
                let mut buf = [0u8; 4096];
                loop {
                    let n = try_nb!(r.read(&mut buf));
                    *amt += n as u64;

                    if n == 0 {
                        break;
                    }
                }
            }
        }

        match mem::replace(self, IgnoreUntilEnd::Empty) {
            IgnoreUntilEnd::Pending { amt, .. } => Ok(amt.into()),
            IgnoreUntilEnd::Empty => unreachable!(),
        }
    }
}

/// Ignore all data from the reader
pub fn ignore_until_end<R: Read>(r: R) -> IgnoreUntilEnd<R> {
    IgnoreUntilEnd::Pending { r: r, amt: 0 }
}

#[cfg(unix)]
fn reuse_port(builder: &TcpBuilder) -> io::Result<&TcpBuilder> {
    use net2::unix::UnixTcpBuilderExt;
    builder.reuse_port(true)
}

#[cfg(windows)]
fn reuse_port(builder: &TcpBuilder) -> io::Result<&TcpBuilder> {
    Ok(builder)
}<commit_msg>no flush before copy all<commit_after>// The MIT License (MIT)

// Copyright (c) 2014 Y. T. CHUNG <zonyitoo@gmail.com>

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

//! TcpRelay implementation

use std::io::{self, Read, Write};
use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
use std::rc::Rc;
use std::mem;
use std::cmp;

use crypto::cipher;
use crypto::CryptoMode;
use relay::socks5::Address;
use relay::{BoxIoFuture, boxed_future};
use relay::dns_resolver::DnsResolver;
use config::{ServerConfig, ServerAddr};

use tokio_core::net::TcpStream;
use tokio_core::reactor::Handle;
use tokio_core::io::{read_exact, write_all};
use tokio_core::io::{ReadHalf, WriteHalf};
use tokio_core::io::Io;

use futures::{self, Future, Poll};

use net2::TcpBuilder;

use ip::IpAddr;

use self::stream::{EncryptedWriter, DecryptedReader};

pub mod local;
mod socks5_local;
mod http_local;
pub mod server;
mod stream;
pub mod client;

/// Directions in the tunnel
#[derive(Debug, Copy, Clone)]
pub enum TunnelDirection {
    /// Client -> Server
    Client2Server,
    /// Client <- Server
    Server2Client,
}

/// ReadHalf of TcpStream with decryption
pub type DecryptedHalf = DecryptedReader<ReadHalf<TcpStream>>;
/// WriteHalf of TcpStream with encryption
pub type EncryptedHalf = EncryptedWriter<WriteHalf<TcpStream>>;

/// Boxed future of DecryptedHalf
pub type DecryptedHalfFut = BoxIoFuture<DecryptedHalf>;
/// Boxed future of EncryptedHalf
pub type EncryptedHalfFut = BoxIoFuture<EncryptedHalf>;

fn connect_proxy_server(handle: &Handle, svr_cfg: Rc<ServerConfig>) -> BoxIoFuture<TcpStream> {
    match svr_cfg.addr() {
        &ServerAddr::SocketAddr(ref addr) => TcpStream::connect(addr, handle).boxed(),
        &ServerAddr::DomainName(ref domain, port) => {
            let handle = handle.clone();
            let fut = DnsResolver::resolve(&domain[..]).and_then(move |sockaddr| {
                let sockaddr = match sockaddr {
                    IpAddr::V4(v4) => SocketAddr::V4(SocketAddrV4::new(v4, port)),
                    IpAddr::V6(v6) => SocketAddr::V6(SocketAddrV6::new(v6, port, 0, 0)),
                };
                TcpStream::connect(&sockaddr, &handle)
            });
            Box::new(fut)
        }
    }
}

/// Handshake logic for ShadowSocks Client
pub fn proxy_server_handshake(remote_stream: TcpStream,
                              svr_cfg: Rc<ServerConfig>,
                              relay_addr: Address)
                              -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> {
    let fut = proxy_handshake(remote_stream, svr_cfg).and_then(|(r_fut, w_fut)| {;
        let w_fut = w_fut.and_then(move |enc_w| {
            trace!("Got encrypt stream and going to send addr: {:?}",
                   relay_addr);

            // Send relay address to remote
            let local_buf = Vec::new();
            relay_addr.write_to(local_buf)
                .and_then(|buf| enc_w.write_all_encrypted(buf))
                .map(|(w, _)| w)
        });

        Ok((r_fut, boxed_future(w_fut)))
    });
    boxed_future(fut)
}

/// ShadowSocks Client-Server handshake protocol
/// Exchange cipher IV and creates stream wrapper
pub fn proxy_handshake(remote_stream: TcpStream,
                       svr_cfg: Rc<ServerConfig>)
                       -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> {
    let fut = futures::lazy(|| Ok(remote_stream.split())).and_then(|(r, w)| {

        let svr_cfg_cloned = svr_cfg.clone();

        let enc = {
            // Encrypt data to remote server

            // Send initialize vector to remote and create encryptor

            let local_iv = svr_cfg.method().gen_init_vec();
            trace!("Going to send initialize vector: {:?}", local_iv);

            write_all(w, local_iv).and_then(move |(w, local_iv)| {
                let encryptor = cipher::with_type(svr_cfg.method(),
                                                  svr_cfg.key(),
                                                  &local_iv[..],
                                                  CryptoMode::Encrypt);

                Ok(EncryptedWriter::new(w, encryptor))
            })
        };

        let dec = {
            let svr_cfg = svr_cfg_cloned;

            // Decrypt data from remote server
            let iv_len = svr_cfg.method().iv_size();
            read_exact(r, vec![0u8; iv_len]).and_then(move |(r, remote_iv)| {
                trace!("Got initialize vector {:?}", remote_iv);

                let decryptor = cipher::with_type(svr_cfg.method(),
                                                  svr_cfg.key(),
                                                  &remote_iv[..],
                                                  CryptoMode::Decrypt);
                let decrypt_stream = DecryptedReader::new(r, decryptor);

                Ok(decrypt_stream)
            })
        };

        Ok((boxed_future(dec), boxed_future(enc)))
    });

    boxed_future(fut)
}

/// Copy exactly N bytes by encryption
pub enum CopyExactEncrypted<R, W>
    where R: Read,
          W: Write
{
    Pending {
        reader: R,
        writer: EncryptedWriter<W>,
        buf: [u8; 4096],
        remain: usize,
        pos: usize,
        cap: usize,
        enc_buf: Vec<u8>,
    },
    Empty,
}

impl<R, W> CopyExactEncrypted<R, W>
    where R: Read,
          W: Write
{
    pub fn new(r: R, w: EncryptedWriter<W>, amt: usize) -> CopyExactEncrypted<R, W> {
        CopyExactEncrypted::Pending {
            reader: r,
            writer: w,
            buf: [0u8; 4096],
            remain: amt,
            pos: 0,
            cap: 0,
            enc_buf: Vec::new(),
        }
    }
}

impl<R, W> Future for CopyExactEncrypted<R, W>
    where R: Read,
          W: Write
{
    type Item = (R, EncryptedWriter<W>);
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        match self {
            &mut CopyExactEncrypted::Empty => panic!("poll after CopyExactEncrypted is finished"),
            &mut CopyExactEncrypted::Pending { ref mut reader,
                                               ref mut writer,
                                               ref mut buf,
                                               ref mut remain,
                                               ref mut pos,
                                               ref mut cap,
                                               ref mut enc_buf } => {
                loop {
                    // If our buffer is empty, then we need to read some data to
                    // continue.
                    if *pos == *cap && *remain != 0 {
                        let buf_len = cmp::min(*remain, buf.len());
                        let n = try_nb!(reader.read(&mut buf[..buf_len]));
                        if n == 0 {
                            // Unexpected EOF!
                            let err = io::Error::new(io::ErrorKind::UnexpectedEof, "Unexpected Eof");
                            return Err(err);
                        } else {
                            *pos = 0;
                            *remain -= n;

                            enc_buf.clear();
                            try!(writer.cipher_update(&buf[..n], enc_buf));
                            *cap = enc_buf.len();
                        }
                    }

                    // If our buffer has some data, let's write it out!
                    while *pos < *cap {
                        let i = try_nb!(writer.write(&enc_buf[*pos..*cap]));
                        *pos += i;
                    }

                    // If we've written al the data and we've seen EOF, flush out the
                    // data and finish the transfer.
                    // done with the entire transfer.
                    if *pos == *cap && *remain == 0 {
                        try_nb!(writer.flush());
                        break; // The only path to execute the following logic
                    }
                }
            }
        }

        match mem::replace(self, CopyExactEncrypted::Empty) {
            CopyExactEncrypted::Pending { reader, writer, .. } => Ok((reader, writer).into()),
            CopyExactEncrypted::Empty => unreachable!(),
        }
    }
}

/// Copy all bytes from reader and write all encrypted data into writer
pub fn copy_exact_encrypted<R, W>(r: R, w: EncryptedWriter<W>, amt: usize) -> CopyExactEncrypted<R, W>
    where R: Read,
          W: Write
{
    CopyExactEncrypted::new(r, w, amt)
}

/// Establish tunnel between server and client
pub fn tunnel<CF, SF>(addr: Address, c2s: CF, s2c: SF) -> BoxIoFuture<()>
    where CF: Future<Item = u64, Error = io::Error> + 'static,
          SF: Future<Item = u64, Error = io::Error> + 'static
{
    let addr = Rc::new(addr);

    let cloned_addr = addr.clone();
    let c2s = c2s.then(move |res| {
        match res {
            Ok(amt) => {
                // Continue reading response from remote server
                trace!("Relay {} client -> server is finished, relayed {} bytes",
                       cloned_addr,
                       amt);

                Ok(TunnelDirection::Client2Server)
            }
            Err(err) => {
                error!("Relay {} client -> server aborted: {}", cloned_addr, err);
                Err(err)
            }
        }
    });

    let cloned_addr = addr.clone();
    let s2c = s2c.then(move |res| {
        match res {
            Ok(amt) => {
                trace!("Relay {} client <- server is finished, relayed {} bytes",
                       cloned_addr,
                       amt);

                Ok(TunnelDirection::Server2Client)
            }
            Err(err) => {
                error!("Relay {} client <- server aborted: {}", cloned_addr, err);
                Err(err)
            }
        }
    });

    let fut = c2s.select(s2c)
        .and_then(|(dir, _)| {
            match dir {
                TunnelDirection::Server2Client => trace!("client <- server is closed, abort connection"),
                TunnelDirection::Client2Server => trace!("server -> client is closed, abort connection"),
            }

            Ok(())
        })
        .map_err(|(err, _)| err);

    boxed_future(fut)
}

/// Read until EOF, and ignore
pub enum IgnoreUntilEnd<R: Read> {
    Pending { r: R, amt: u64 },
    Empty,
}

impl<R: Read> Future for IgnoreUntilEnd<R> {
    type Item = u64;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        match self {
            &mut IgnoreUntilEnd::Empty => panic!("poll IgnoreUntilEnd after it is finished"),
            &mut IgnoreUntilEnd::Pending { ref mut r, ref mut amt } => {
                let mut buf = [0u8; 4096];
                loop {
                    let n = try_nb!(r.read(&mut buf));
                    *amt += n as u64;

                    if n == 0 {
                        break;
                    }
                }
            }
        }

        match mem::replace(self, IgnoreUntilEnd::Empty) {
            IgnoreUntilEnd::Pending { amt, .. } => Ok(amt.into()),
            IgnoreUntilEnd::Empty => unreachable!(),
        }
    }
}

/// Ignore all data from the reader
pub fn ignore_until_end<R: Read>(r: R) -> IgnoreUntilEnd<R> {
    IgnoreUntilEnd::Pending { r: r, amt: 0 }
}

#[cfg(unix)]
fn reuse_port(builder: &TcpBuilder) -> io::Result<&TcpBuilder> {
    use net2::unix::UnixTcpBuilderExt;
    builder.reuse_port(true)
}

#[cfg(windows)]
fn reuse_port(builder: &TcpBuilder) -> io::Result<&TcpBuilder> {
    Ok(builder)
}<|endoftext|>


<commit_before>iface iterable<A> {
    fn iter(blk: fn(A));
}

impl<A> of iterable<A> for fn@(fn(A)) {
    fn iter(blk: fn(A)) {
        self(blk);
    }
}

// accomodate the fact that int/uint are passed by value by default:
impl of iterable<int> for fn@(fn(int)) {
    fn iter(blk: fn(&&int)) {
        self {|i| blk(i)}
    }
}

impl of iterable<uint> for fn@(fn(uint)) {
    fn iter(blk: fn(&&uint)) {
        self {|i| blk(i)}
    }
}

impl<A> of iterable<A> for [A] {
    fn iter(blk: fn(A)) {
        vec::iter(self, blk)
    }
}

impl<A> of iterable<A> for option<A> {
    fn iter(blk: fn(A)) {
        option::iter(self, blk)
    }
}

impl of iterable<char> for str {
    fn iter(blk: fn(&&char)) {
        str::chars_iter(self) { |ch| blk(ch) }
    }
}

fn enumerate<A,IA:iterable<A>>(self: IA, blk: fn(uint, A)) {
    let mut i = 0u;
    self.iter {|a|
        blk(i, a);
        i += 1u;
    }
}

// Here: we have to use fn@ for predicates and map functions, because
// we will be binding them up into a closure.  Disappointing.  A true
// region type system might be able to do better than this.

fn filter<A,IA:iterable<A>>(self: IA, prd: fn@(A) -> bool, blk: fn(A)) {
    self.iter {|a|
        if prd(a) { blk(a) }
    }
}

fn filter_map<A,B,IA:iterable<A>>(self: IA, cnv: fn@(A) -> option<B>,
                                  blk: fn(B)) {
    self.iter {|a|
        alt cnv(a) {
          some(b) { blk(b) }
          none { }
        }
    }
}

fn map<A,B,IA:iterable<A>>(self: IA, cnv: fn@(A) -> B, blk: fn(B)) {
    self.iter {|a|
        let b = cnv(a);
        blk(b);
    }
}

fn flat_map<A,B,IA:iterable<A>,IB:iterable<B>>(
    self: IA, cnv: fn@(A) -> IB, blk: fn(B)) {
    self.iter {|a|
        cnv(a).iter(blk)
    }
}

fn foldl<A,B,IA:iterable<A>>(self: IA, +b0: B, blk: fn(-B, A) -> B) -> B {
    let mut b <- b0;
    self.iter {|a|
        b = blk(b, a);
    }
    ret b;
}

fn foldr<A:copy,B,IA:iterable<A>>(
    self: IA, +b0: B, blk: fn(A, -B) -> B) -> B {

    let mut b <- b0;
    reversed(self) {|a|
        b = blk(a, b);
    }
    ret b;
}

fn to_vec<A:copy,IA:iterable<A>>(self: IA) -> [A] {
    foldl::<A,[A],IA>(self, [], {|r, a| r + [a]})
}

// FIXME: This could be made more efficient with an riterable interface
// #2005
fn reversed<A:copy,IA:iterable<A>>(self: IA, blk: fn(A)) {
    vec::riter(to_vec(self), blk)
}

fn count<A,IA:iterable<A>>(self: IA, x: A) -> uint {
    foldl(self, 0u) {|count, value|
        if value == x {
            count + 1u
        } else {
            count
        }
    }
}

fn repeat(times: uint, blk: fn()) {
    let mut i = 0u;
    while i < times {
        blk();
        i += 1u;
    }
}

fn min<A:copy,IA:iterable<A>>(self: IA) -> A {
    alt foldl::<A,option<A>,IA>(self, none) {|a, b|
        alt a {
          some(a_) if a_ < b {
            // FIXME: Not sure if this is successfully optimized to a move
            // #2005
            a
          }
          _ { some(b) }
        }
    } {
        some(val) { val }
        none { fail "min called on empty iterator" }
    }
}

fn max<A:copy,IA:iterable<A>>(self: IA) -> A {
    alt foldl::<A,option<A>,IA>(self, none) {|a, b|
        alt a {
          some(a_) if a_ > b {
            // FIXME: Not sure if this is successfully optimized to a move
            // #2005
            a
          }
          _ { some(b) }
        }
    } {
        some(val) { val }
        none { fail "max called on empty iterator" }
    }
}

#[test]
fn test_enumerate() {
    enumerate(["0", "1", "2"]) {|i,j|
        assert #fmt["%u",i] == j;
    }
}

#[test]
fn test_map_and_to_vec() {
    let a = bind vec::iter([0, 1, 2], _);
    let b = bind map(a, {|i| 2*i}, _);
    let c = to_vec(b);
    assert c == [0, 2, 4];
}

#[test]
fn test_map_directly_on_vec() {
    let b = bind map([0, 1, 2], {|i| 2*i}, _);
    let c = to_vec(b);
    assert c == [0, 2, 4];
}

#[test]
fn test_filter_on_int_range() {
    fn is_even(&&i: int) -> bool {
        ret (i % 2) == 0;
    }

    let l = to_vec(bind filter(bind int::range(0, 10, _), is_even, _));
    assert l == [0, 2, 4, 6, 8];
}

#[test]
fn test_filter_on_uint_range() {
    fn is_even(&&i: uint) -> bool {
        ret (i % 2u) == 0u;
    }

    let l = to_vec(bind filter(bind uint::range(0u, 10u, _), is_even, _));
    assert l == [0u, 2u, 4u, 6u, 8u];
}

#[test]
fn test_filter_map() {
    fn negativate_the_evens(&&i: int) -> option<int> {
        if i % 2 == 0 {
            some(-i)
        } else {
            none
        }
    }

    let l = to_vec(bind filter_map(
        bind int::range(0, 5, _), negativate_the_evens, _));
    assert l == [0, -2, -4];
}

#[test]
fn test_flat_map_with_option() {
    fn if_even(&&i: int) -> option<int> {
        if (i % 2) == 0 { some(i) }
        else { none }
    }

    let a = bind vec::iter([0, 1, 2], _);
    let b = bind flat_map(a, if_even, _);
    let c = to_vec(b);
    assert c == [0, 2];
}

#[test]
fn test_flat_map_with_list() {
    fn repeat(&&i: int) -> [int] {
        let mut r = [];
        int::range(0, i) {|_j| r += [i]; }
        r
    }

    let a = bind vec::iter([0, 1, 2, 3], _);
    let b = bind flat_map(a, repeat, _);
    let c = to_vec(b);
    #debug["c = %?", c];
    assert c == [1, 2, 2, 3, 3, 3];
}

#[test]
fn test_repeat() {
    let mut c = [], i = 0u;
    repeat(5u) {||
        c += [(i * i)];
        i += 1u;
    };
    #debug["c = %?", c];
    assert c == [0u, 1u, 4u, 9u, 16u];
}

#[test]
fn test_min() {
    assert min([5, 4, 1, 2, 3]) == 1;
}

#[test]
#[should_fail]
#[ignore(cfg(target_os = "win32"))]
fn test_min_empty() {
    min::<int, [int]>([]);
}

#[test]
fn test_max() {
    assert max([1, 2, 4, 2, 3]) == 4;
}

#[test]
#[should_fail]
#[ignore(cfg(target_os = "win32"))]
fn test_max_empty() {
    max::<int, [int]>([]);
}

#[test]
fn test_reversed() {
    assert to_vec(bind reversed([1, 2, 3], _)) == [3, 2, 1];
}

#[test]
fn test_count() {
    assert count([1, 2, 1, 2, 1], 1) == 3u;
}

#[test]
fn test_foldr() {
    fn sub(&&a: int, -b: int) -> int {
        a - b
    }
    let sum = foldr([1, 2, 3, 4], 0, sub);
    assert sum == -2;
}<commit_msg>change to use && mode (can't move from upvar)<commit_after>iface iterable<A> {
    fn iter(blk: fn(A));
}

impl<A> of iterable<A> for fn@(fn(A)) {
    fn iter(blk: fn(A)) {
        self(blk);
    }
}

// accomodate the fact that int/uint are passed by value by default:
impl of iterable<int> for fn@(fn(int)) {
    fn iter(blk: fn(&&int)) {
        self {|i| blk(i)}
    }
}

impl of iterable<uint> for fn@(fn(uint)) {
    fn iter(blk: fn(&&uint)) {
        self {|i| blk(i)}
    }
}

impl<A> of iterable<A> for [A] {
    fn iter(blk: fn(A)) {
        vec::iter(self, blk)
    }
}

impl<A> of iterable<A> for option<A> {
    fn iter(blk: fn(A)) {
        option::iter(self, blk)
    }
}

impl of iterable<char> for str {
    fn iter(blk: fn(&&char)) {
        str::chars_iter(self) { |ch| blk(ch) }
    }
}

fn enumerate<A,IA:iterable<A>>(self: IA, blk: fn(uint, A)) {
    let mut i = 0u;
    self.iter {|a|
        blk(i, a);
        i += 1u;
    }
}

// Here: we have to use fn@ for predicates and map functions, because
// we will be binding them up into a closure.  Disappointing.  A true
// region type system might be able to do better than this.

fn filter<A,IA:iterable<A>>(self: IA, prd: fn@(A) -> bool, blk: fn(A)) {
    self.iter {|a|
        if prd(a) { blk(a) }
    }
}

fn filter_map<A,B,IA:iterable<A>>(self: IA, cnv: fn@(A) -> option<B>,
                                  blk: fn(B)) {
    self.iter {|a|
        alt cnv(a) {
          some(b) { blk(b) }
          none { }
        }
    }
}

fn map<A,B,IA:iterable<A>>(self: IA, cnv: fn@(A) -> B, blk: fn(B)) {
    self.iter {|a|
        let b = cnv(a);
        blk(b);
    }
}

fn flat_map<A,B,IA:iterable<A>,IB:iterable<B>>(
    self: IA, cnv: fn@(A) -> IB, blk: fn(B)) {
    self.iter {|a|
        cnv(a).iter(blk)
    }
}

fn foldl<A,B,IA:iterable<A>>(self: IA, +b0: B, blk: fn(B, A) -> B) -> B {
    let mut b <- b0;
    self.iter {|a|
        b <- blk(b, a);
    }
    ret b;
}

fn foldr<A:copy,B,IA:iterable<A>>(
    self: IA, +b0: B, blk: fn(A, B) -> B) -> B {

    let mut b <- b0;
    reversed(self) {|a|
        b <- blk(a, b);
    }
    ret b;
}

fn to_vec<A:copy,IA:iterable<A>>(self: IA) -> [A] {
    foldl::<A,[A],IA>(self, [], {|r, a| r + [a]})
}

// FIXME: This could be made more efficient with an riterable interface
// #2005
fn reversed<A:copy,IA:iterable<A>>(self: IA, blk: fn(A)) {
    vec::riter(to_vec(self), blk)
}

fn count<A,IA:iterable<A>>(self: IA, x: A) -> uint {
    foldl(self, 0u) {|count, value|
        if value == x {
            count + 1u
        } else {
            count
        }
    }
}

fn repeat(times: uint, blk: fn()) {
    let mut i = 0u;
    while i < times {
        blk();
        i += 1u;
    }
}

fn min<A:copy,IA:iterable<A>>(self: IA) -> A {
    alt foldl::<A,option<A>,IA>(self, none) {|a, b|
        alt a {
          some(a_) if a_ < b {
            // FIXME: Not sure if this is successfully optimized to a move
            // #2005
            a
          }
          _ { some(b) }
        }
    } {
        some(val) { val }
        none { fail "min called on empty iterator" }
    }
}

fn max<A:copy,IA:iterable<A>>(self: IA) -> A {
    alt foldl::<A,option<A>,IA>(self, none) {|a, b|
        alt a {
          some(a_) if a_ > b {
            // FIXME: Not sure if this is successfully optimized to a move
            // #2005
            a
          }
          _ { some(b) }
        }
    } {
        some(val) { val }
        none { fail "max called on empty iterator" }
    }
}

#[test]
fn test_enumerate() {
    enumerate(["0", "1", "2"]) {|i,j|
        assert #fmt["%u",i] == j;
    }
}

#[test]
fn test_map_and_to_vec() {
    let a = bind vec::iter([0, 1, 2], _);
    let b = bind map(a, {|i| 2*i}, _);
    let c = to_vec(b);
    assert c == [0, 2, 4];
}

#[test]
fn test_map_directly_on_vec() {
    let b = bind map([0, 1, 2], {|i| 2*i}, _);
    let c = to_vec(b);
    assert c == [0, 2, 4];
}

#[test]
fn test_filter_on_int_range() {
    fn is_even(&&i: int) -> bool {
        ret (i % 2) == 0;
    }

    let l = to_vec(bind filter(bind int::range(0, 10, _), is_even, _));
    assert l == [0, 2, 4, 6, 8];
}

#[test]
fn test_filter_on_uint_range() {
    fn is_even(&&i: uint) -> bool {
        ret (i % 2u) == 0u;
    }

    let l = to_vec(bind filter(bind uint::range(0u, 10u, _), is_even, _));
    assert l == [0u, 2u, 4u, 6u, 8u];
}

#[test]
fn test_filter_map() {
    fn negativate_the_evens(&&i: int) -> option<int> {
        if i % 2 == 0 {
            some(-i)
        } else {
            none
        }
    }

    let l = to_vec(bind filter_map(
        bind int::range(0, 5, _), negativate_the_evens, _));
    assert l == [0, -2, -4];
}

#[test]
fn test_flat_map_with_option() {
    fn if_even(&&i: int) -> option<int> {
        if (i % 2) == 0 { some(i) }
        else { none }
    }

    let a = bind vec::iter([0, 1, 2], _);
    let b = bind flat_map(a, if_even, _);
    let c = to_vec(b);
    assert c == [0, 2];
}

#[test]
fn test_flat_map_with_list() {
    fn repeat(&&i: int) -> [int] {
        let mut r = [];
        int::range(0, i) {|_j| r += [i]; }
        r
    }

    let a = bind vec::iter([0, 1, 2, 3], _);
    let b = bind flat_map(a, repeat, _);
    let c = to_vec(b);
    #debug["c = %?", c];
    assert c == [1, 2, 2, 3, 3, 3];
}

#[test]
fn test_repeat() {
    let mut c = [], i = 0u;
    repeat(5u) {||
        c += [(i * i)];
        i += 1u;
    };
    #debug["c = %?", c];
    assert c == [0u, 1u, 4u, 9u, 16u];
}

#[test]
fn test_min() {
    assert min([5, 4, 1, 2, 3]) == 1;
}

#[test]
#[should_fail]
#[ignore(cfg(target_os = "win32"))]
fn test_min_empty() {
    min::<int, [int]>([]);
}

#[test]
fn test_max() {
    assert max([1, 2, 4, 2, 3]) == 4;
}

#[test]
#[should_fail]
#[ignore(cfg(target_os = "win32"))]
fn test_max_empty() {
    max::<int, [int]>([]);
}

#[test]
fn test_reversed() {
    assert to_vec(bind reversed([1, 2, 3], _)) == [3, 2, 1];
}

#[test]
fn test_count() {
    assert count([1, 2, 1, 2, 1], 1) == 3u;
}

#[test]
fn test_foldr() {
    fn sub(&&a: int, &&b: int) -> int {
        a - b
    }
    let sum = foldr([1, 2, 3, 4], 0, sub);
    assert sum == -2;
}<|endoftext|>


<commit_before>// Copyright 2012 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.

use cast;
use iterator::Iterator;
use libc;
use ops::Drop;
use option::{Option, Some, None};
use ptr::RawPtr;
use ptr;
use str::StrSlice;
use vec::ImmutableVector;

/// The representation of a C String.
///
/// This structure wraps a `*libc::c_char`, and will automatically free the
/// memory it is pointing to when it goes out of scope.
pub struct CString {
    priv buf: *libc::c_char,
    priv owns_buffer_: bool,
}

impl CString {
    /// Create a C String from a pointer.
    pub fn new(buf: *libc::c_char, owns_buffer: bool) -> CString {
        CString { buf: buf, owns_buffer_: owns_buffer }
    }

    /// Unwraps the wrapped `*libc::c_char` from the `CString` wrapper.
    pub unsafe fn unwrap(self) -> *libc::c_char {
        let mut c_str = self;
        c_str.owns_buffer_ = false;
        c_str.buf
    }

    /// Calls a closure with a reference to the underlying `*libc::c_char`.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    pub fn with_ref<T>(&self, f: &fn(*libc::c_char) -> T) -> T {
        if self.buf.is_null() { fail!("CString is null!"); }
        f(self.buf)
    }

    /// Calls a closure with a mutable reference to the underlying `*libc::c_char`.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    pub fn with_mut_ref<T>(&mut self, f: &fn(*mut libc::c_char) -> T) -> T {
        if self.buf.is_null() { fail!("CString is null!"); }
        f(unsafe { cast::transmute(self.buf) })
    }

    /// Returns true if the CString is a null.
    pub fn is_null(&self) -> bool {
        self.buf.is_null()
    }

    /// Returns true if the CString is not null.
    pub fn is_not_null(&self) -> bool {
        self.buf.is_not_null()
    }

    /// Returns whether or not the `CString` owns the buffer.
    pub fn owns_buffer(&self) -> bool {
        self.owns_buffer_
    }

    /// Converts the CString into a `&[u8]` without copying.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    pub fn as_bytes<'a>(&'a self) -> &'a [u8] {
        if self.buf.is_null() { fail!("CString is null!"); }
        unsafe {
            let len = libc::strlen(self.buf) as uint;
            cast::transmute((self.buf, len + 1))
        }
    }

    /// Return a CString iterator.
    fn iter<'a>(&'a self) -> CStringIterator<'a> {
        CStringIterator {
            ptr: self.buf,
            lifetime: unsafe { cast::transmute(self.buf) },
        }
    }
}

impl Drop for CString {
    fn drop(&self) {
        if self.owns_buffer_ && self.buf.is_not_null() {
            unsafe {
                libc::free(self.buf as *libc::c_void)
            }
        }
    }
}

/// A generic trait for converting a value to a CString.
pub trait ToCStr {
    /// Create a C String.
    fn to_c_str(&self) -> CString;
}

impl<'self> ToCStr for &'self str {
    #[inline]
    fn to_c_str(&self) -> CString {
        self.as_bytes().to_c_str()
    }
}

impl<'self> ToCStr for &'self [u8] {
    fn to_c_str(&self) -> CString {
        do self.as_imm_buf |self_buf, self_len| {
            unsafe {
                let buf = libc::malloc(self_len as u64 + 1) as *mut u8;
                if buf.is_null() {
                    fail!("failed to allocate memory!");
                }

                ptr::copy_memory(buf, self_buf, self_len);
                *ptr::mut_offset(buf, self_len as int) = 0;

                CString::new(buf as *libc::c_char, true)
            }
        }
    }
}

/// External iterator for a CString's bytes.
///
/// Use with the `std::iterator` module.
pub struct CStringIterator<'self> {
    priv ptr: *libc::c_char,
    priv lifetime: &'self libc::c_char, // FIXME: #5922
}

impl<'self> Iterator<libc::c_char> for CStringIterator<'self> {
    fn next(&mut self) -> Option<libc::c_char> {
        let ch = unsafe { *self.ptr };
        if ch == 0 {
            None
        } else {
            self.ptr = ptr::offset(self.ptr, 1);
            Some(ch)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use libc;
    use ptr;
    use option::{Some, None};

    #[test]
    fn test_to_c_str() {
        do "".to_c_str().with_ref |buf| {
            unsafe {
                assert_eq!(*ptr::offset(buf, 0), 0);
            }
        }

        do "hello".to_c_str().with_ref |buf| {
            unsafe {
                assert_eq!(*ptr::offset(buf, 0), 'h' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 1), 'e' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 2), 'l' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 3), 'l' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 4), 'o' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 5), 0);
            }
        }
    }

    #[test]
    fn test_is_null() {
        let c_str = CString::new(ptr::null(), false);
        assert!(c_str.is_null());
        assert!(!c_str.is_not_null());
    }

    #[test]
    fn test_unwrap() {
        let c_str = "hello".to_c_str();
        unsafe { libc::free(c_str.unwrap() as *libc::c_void) }
    }

    #[test]
    fn test_with_ref() {
        let c_str = "hello".to_c_str();
        let len = unsafe { c_str.with_ref(|buf| libc::strlen(buf)) };
        assert!(!c_str.is_null());
        assert!(c_str.is_not_null());
        assert_eq!(len, 5);
    }

    #[test]
    #[should_fail]
    #[ignore(cfg(windows))]
    fn test_with_ref_empty_fail() {
        let c_str = unsafe { CString::new(ptr::null(), false) };
        c_str.with_ref(|_| ());
    }

    #[test]
    fn test_iterator() {
        let c_str = "".to_c_str();
        let mut iter = c_str.iter();
        assert_eq!(iter.next(), None);

        let c_str = "hello".to_c_str();
        let mut iter = c_str.iter();
        assert_eq!(iter.next(), Some('h' as libc::c_char));
        assert_eq!(iter.next(), Some('e' as libc::c_char));
        assert_eq!(iter.next(), Some('l' as libc::c_char));
        assert_eq!(iter.next(), Some('l' as libc::c_char));
        assert_eq!(iter.next(), Some('o' as libc::c_char));
        assert_eq!(iter.next(), None);
    }
}
<commit_msg>std: remove unnecessary test from c_str.drop and use safer transmute<commit_after>// Copyright 2012 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.

use cast;
use iterator::Iterator;
use libc;
use ops::Drop;
use option::{Option, Some, None};
use ptr::RawPtr;
use ptr;
use str::StrSlice;
use vec::ImmutableVector;

/// The representation of a C String.
///
/// This structure wraps a `*libc::c_char`, and will automatically free the
/// memory it is pointing to when it goes out of scope.
pub struct CString {
    priv buf: *libc::c_char,
    priv owns_buffer_: bool,
}

impl CString {
    /// Create a C String from a pointer.
    pub fn new(buf: *libc::c_char, owns_buffer: bool) -> CString {
        CString { buf: buf, owns_buffer_: owns_buffer }
    }

    /// Unwraps the wrapped `*libc::c_char` from the `CString` wrapper.
    pub unsafe fn unwrap(self) -> *libc::c_char {
        let mut c_str = self;
        c_str.owns_buffer_ = false;
        c_str.buf
    }

    /// Calls a closure with a reference to the underlying `*libc::c_char`.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    pub fn with_ref<T>(&self, f: &fn(*libc::c_char) -> T) -> T {
        if self.buf.is_null() { fail!("CString is null!"); }
        f(self.buf)
    }

    /// Calls a closure with a mutable reference to the underlying `*libc::c_char`.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    pub fn with_mut_ref<T>(&mut self, f: &fn(*mut libc::c_char) -> T) -> T {
        if self.buf.is_null() { fail!("CString is null!"); }
        f(unsafe { cast::transmute_mut_unsafe(self.buf) })
    }

    /// Returns true if the CString is a null.
    pub fn is_null(&self) -> bool {
        self.buf.is_null()
    }

    /// Returns true if the CString is not null.
    pub fn is_not_null(&self) -> bool {
        self.buf.is_not_null()
    }

    /// Returns whether or not the `CString` owns the buffer.
    pub fn owns_buffer(&self) -> bool {
        self.owns_buffer_
    }

    /// Converts the CString into a `&[u8]` without copying.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    pub fn as_bytes<'a>(&'a self) -> &'a [u8] {
        if self.buf.is_null() { fail!("CString is null!"); }
        unsafe {
            let len = libc::strlen(self.buf) as uint;
            cast::transmute((self.buf, len + 1))
        }
    }

    /// Return a CString iterator.
    fn iter<'a>(&'a self) -> CStringIterator<'a> {
        CStringIterator {
            ptr: self.buf,
            lifetime: unsafe { cast::transmute(self.buf) },
        }
    }
}

impl Drop for CString {
    fn drop(&self) {
        if self.owns_buffer_ {
            unsafe {
                libc::free(self.buf as *libc::c_void)
            }
        }
    }
}

/// A generic trait for converting a value to a CString.
pub trait ToCStr {
    /// Create a C String.
    fn to_c_str(&self) -> CString;
}

impl<'self> ToCStr for &'self str {
    #[inline]
    fn to_c_str(&self) -> CString {
        self.as_bytes().to_c_str()
    }
}

impl<'self> ToCStr for &'self [u8] {
    fn to_c_str(&self) -> CString {
        do self.as_imm_buf |self_buf, self_len| {
            unsafe {
                let buf = libc::malloc(self_len as u64 + 1) as *mut u8;
                if buf.is_null() {
                    fail!("failed to allocate memory!");
                }

                ptr::copy_memory(buf, self_buf, self_len);
                *ptr::mut_offset(buf, self_len as int) = 0;

                CString::new(buf as *libc::c_char, true)
            }
        }
    }
}

/// External iterator for a CString's bytes.
///
/// Use with the `std::iterator` module.
pub struct CStringIterator<'self> {
    priv ptr: *libc::c_char,
    priv lifetime: &'self libc::c_char, // FIXME: #5922
}

impl<'self> Iterator<libc::c_char> for CStringIterator<'self> {
    fn next(&mut self) -> Option<libc::c_char> {
        let ch = unsafe { *self.ptr };
        if ch == 0 {
            None
        } else {
            self.ptr = ptr::offset(self.ptr, 1);
            Some(ch)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use libc;
    use ptr;
    use option::{Some, None};

    #[test]
    fn test_to_c_str() {
        do "".to_c_str().with_ref |buf| {
            unsafe {
                assert_eq!(*ptr::offset(buf, 0), 0);
            }
        }

        do "hello".to_c_str().with_ref |buf| {
            unsafe {
                assert_eq!(*ptr::offset(buf, 0), 'h' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 1), 'e' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 2), 'l' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 3), 'l' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 4), 'o' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 5), 0);
            }
        }
    }

    #[test]
    fn test_is_null() {
        let c_str = CString::new(ptr::null(), false);
        assert!(c_str.is_null());
        assert!(!c_str.is_not_null());
    }

    #[test]
    fn test_unwrap() {
        let c_str = "hello".to_c_str();
        unsafe { libc::free(c_str.unwrap() as *libc::c_void) }
    }

    #[test]
    fn test_with_ref() {
        let c_str = "hello".to_c_str();
        let len = unsafe { c_str.with_ref(|buf| libc::strlen(buf)) };
        assert!(!c_str.is_null());
        assert!(c_str.is_not_null());
        assert_eq!(len, 5);
    }

    #[test]
    #[should_fail]
    #[ignore(cfg(windows))]
    fn test_with_ref_empty_fail() {
        let c_str = unsafe { CString::new(ptr::null(), false) };
        c_str.with_ref(|_| ());
    }

    #[test]
    fn test_iterator() {
        let c_str = "".to_c_str();
        let mut iter = c_str.iter();
        assert_eq!(iter.next(), None);

        let c_str = "hello".to_c_str();
        let mut iter = c_str.iter();
        assert_eq!(iter.next(), Some('h' as libc::c_char));
        assert_eq!(iter.next(), Some('e' as libc::c_char));
        assert_eq!(iter.next(), Some('l' as libc::c_char));
        assert_eq!(iter.next(), Some('l' as libc::c_char));
        assert_eq!(iter.next(), Some('o' as libc::c_char));
        assert_eq!(iter.next(), None);
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 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.

//! Traits for working with Errors.
//!
//! # The `Error` trait
//!
//! `Error` is a trait representing the basic expectations for error values,
//! i.e. values of type `E` in `Result<T, E>`. At a minimum, errors must provide
//! a description, but they may optionally provide additional detail and cause
//! chain information:
//!
//! ```
//! trait Error: Send {
//!     fn description(&self) -> &str;
//!
//!     fn detail(&self) -> Option<String> { None }
//!     fn cause(&self) -> Option<&Error> { None }
//! }
//! ```
//!
//! The `cause` method is generally used when errors cross "abstraction
//! boundaries", i.e.  when a one module must report an error that is "caused"
//! by an error from a lower-level module. This setup makes it possible for the
//! high-level module to provide its own errors that do not commit to any
//! particular implementation, but also reveal some of its implementation for
//! debugging via `cause` chains.
//!
//! The trait inherits from `Any` to allow *downcasting*: converting from a
//! trait object to a specific concrete type when applicable.
//!
//! # The `FromError` trait
//!
//! `FromError` is a simple trait that expresses conversions between different
//! error types. To provide maximum flexibility, it does not require either of
//! the types to actually implement the `Error` trait, although this will be the
//! common case.
//!
//! The main use of this trait is in the `try!` macro, which uses it to
//! automatically convert a given error to the error specified in a function's
//! return type.
//!
//! For example,
//!
//! ```
//! use std::error::FromError;
//! use std::io::{File, IoError};
//! use std::os::{MemoryMap, MapError};
//! use std::path::Path;
//!
//! enum MyError {
//!     Io(IoError),
//!     Map(MapError)
//! }
//!
//! impl FromError<IoError> for MyError {
//!     fn from_error(err: IoError) -> MyError {
//!         Io(err)
//!     }
//! }
//!
//! impl FromError<MapError> for MyError {
//!     fn from_error(err: MapError) -> MyError {
//!         Map(err)
//!     }
//! }
//!
//! #[allow(unused_variables)]
//! fn open_and_map() -> Result<(), MyError> {
//!     let f = try!(File::open(&Path::new("foo.txt")));
//!     let m = try!(MemoryMap::new(0, &[]));
//!     // do something interesting here...
//!     Ok(())
//! }
//! ```

use option::{Option, None};
use kinds::Send;
use string::String;

/// Base functionality for all errors in Rust.
pub trait Error: Send {
    /// A short description of the error; usually a static string.
    fn description(&self) -> &str;

    /// A detailed description of the error, usually including dynamic information.
    fn detail(&self) -> Option<String> { None }

    /// The lower-level cause of this error, if any.
    fn cause(&self) -> Option<&Error> { None }
}

/// A trait for types that can be converted from a given error type `E`.
pub trait FromError<E> {
    /// Perform the conversion.
    fn from_error(err: E) -> Self;
}

// Any type is convertable from itself
impl<E> FromError<E> for E {
    fn from_error(err: E) -> E {
        err
    }
}
<commit_msg>std::error: fix stray doc comment<commit_after>// Copyright 2014 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.

//! Traits for working with Errors.
//!
//! # The `Error` trait
//!
//! `Error` is a trait representing the basic expectations for error values,
//! i.e. values of type `E` in `Result<T, E>`. At a minimum, errors must provide
//! a description, but they may optionally provide additional detail and cause
//! chain information:
//!
//! ```
//! trait Error: Send {
//!     fn description(&self) -> &str;
//!
//!     fn detail(&self) -> Option<String> { None }
//!     fn cause(&self) -> Option<&Error> { None }
//! }
//! ```
//!
//! The `cause` method is generally used when errors cross "abstraction
//! boundaries", i.e.  when a one module must report an error that is "caused"
//! by an error from a lower-level module. This setup makes it possible for the
//! high-level module to provide its own errors that do not commit to any
//! particular implementation, but also reveal some of its implementation for
//! debugging via `cause` chains.
//!
//! # The `FromError` trait
//!
//! `FromError` is a simple trait that expresses conversions between different
//! error types. To provide maximum flexibility, it does not require either of
//! the types to actually implement the `Error` trait, although this will be the
//! common case.
//!
//! The main use of this trait is in the `try!` macro, which uses it to
//! automatically convert a given error to the error specified in a function's
//! return type.
//!
//! For example,
//!
//! ```
//! use std::error::FromError;
//! use std::io::{File, IoError};
//! use std::os::{MemoryMap, MapError};
//! use std::path::Path;
//!
//! enum MyError {
//!     Io(IoError),
//!     Map(MapError)
//! }
//!
//! impl FromError<IoError> for MyError {
//!     fn from_error(err: IoError) -> MyError {
//!         Io(err)
//!     }
//! }
//!
//! impl FromError<MapError> for MyError {
//!     fn from_error(err: MapError) -> MyError {
//!         Map(err)
//!     }
//! }
//!
//! #[allow(unused_variables)]
//! fn open_and_map() -> Result<(), MyError> {
//!     let f = try!(File::open(&Path::new("foo.txt")));
//!     let m = try!(MemoryMap::new(0, &[]));
//!     // do something interesting here...
//!     Ok(())
//! }
//! ```

use option::{Option, None};
use kinds::Send;
use string::String;

/// Base functionality for all errors in Rust.
pub trait Error: Send {
    /// A short description of the error; usually a static string.
    fn description(&self) -> &str;

    /// A detailed description of the error, usually including dynamic information.
    fn detail(&self) -> Option<String> { None }

    /// The lower-level cause of this error, if any.
    fn cause(&self) -> Option<&Error> { None }
}

/// A trait for types that can be converted from a given error type `E`.
pub trait FromError<E> {
    /// Perform the conversion.
    fn from_error(err: E) -> Self;
}

// Any type is convertable from itself
impl<E> FromError<E> for E {
    fn from_error(err: E) -> E {
        err
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2013 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.

/*!

Higher level communication abstractions.

*/

#[allow(missing_doc)];

use std::comm;

/// An extension of `pipes::stream` that allows both sending and receiving.
pub struct DuplexStream<T, U> {
    priv tx: Sender<T>,
    priv rx: Receiver<U>,
}

/// Creates a bidirectional stream.
pub fn duplex<T: Send, U: Send>() -> (DuplexStream<T, U>, DuplexStream<U, T>) {
    let (tx1, rx1) = channel();
    let (tx2, rx2) = channel();
    (DuplexStream { tx: tx1, rx: rx2 },
     DuplexStream { tx: tx2, rx: rx1 })
}

// Allow these methods to be used without import:
impl<T:Send,U:Send> DuplexStream<T, U> {
    pub fn send(&self, x: T) {
        self.tx.send(x)
    }
    pub fn try_send(&self, x: T) -> bool {
        self.tx.try_send(x)
    }
    pub fn recv(&self) -> U {
        self.rx.recv()
    }
    pub fn try_recv(&self) -> comm::TryRecvResult<U> {
        self.rx.try_recv()
    }
    pub fn recv_opt(&self) -> Option<U> {
        self.rx.recv_opt()
    }
}

/// An extension of `pipes::stream` that provides synchronous message sending.
pub struct SyncSender<T> { priv duplex_stream: DuplexStream<T, ()> }
/// An extension of `pipes::stream` that acknowledges each message received.
pub struct SyncReceiver<T> { priv duplex_stream: DuplexStream<(), T> }

impl<T: Send> SyncSender<T> {
    pub fn send(&self, val: T) {
        assert!(self.try_send(val), "SyncSender.send: receiving port closed");
    }

    /// Sends a message, or report if the receiver has closed the connection
    /// before receiving.
    pub fn try_send(&self, val: T) -> bool {
        self.duplex_stream.try_send(val) && self.duplex_stream.recv_opt().is_some()
    }
}

impl<T: Send> SyncReceiver<T> {
    pub fn recv(&self) -> T {
        self.recv_opt().expect("SyncReceiver.recv: sending channel closed")
    }

    pub fn recv_opt(&self) -> Option<T> {
        self.duplex_stream.recv_opt().map(|val| {
            self.duplex_stream.try_send(());
            val
        })
    }

    pub fn try_recv(&self) -> comm::TryRecvResult<T> {
        match self.duplex_stream.try_recv() {
            comm::Data(t) => { self.duplex_stream.try_send(()); comm::Data(t) }
            state => state,
        }
    }
}

/// Creates a stream whose channel, upon sending a message, blocks until the
/// message is received.
pub fn rendezvous<T: Send>() -> (SyncReceiver<T>, SyncSender<T>) {
    let (chan_stream, port_stream) = duplex();
    (SyncReceiver { duplex_stream: port_stream },
     SyncSender { duplex_stream: chan_stream })
}

#[cfg(test)]
mod test {
    use comm::{duplex, rendezvous};

    #[test]
    pub fn DuplexStream1() {
        let (left, right) = duplex();

        left.send(~"abc");
        right.send(123);

        assert!(left.recv() == 123);
        assert!(right.recv() == ~"abc");
    }

    #[test]
    pub fn basic_rendezvous_test() {
        let (port, chan) = rendezvous();

        spawn(proc() {
            chan.send("abc");
        });

        assert!(port.recv() == "abc");
    }

    #[test]
    fn recv_a_lot() {
        // Rendezvous streams should be able to handle any number of messages being sent
        let (port, chan) = rendezvous();
        spawn(proc() {
            for _ in range(0, 10000) { chan.send(()); }
        });
        for _ in range(0, 10000) { port.recv(); }
    }

    #[test]
    fn send_and_fail_and_try_recv() {
        let (port, chan) = rendezvous();
        spawn(proc() {
            chan.duplex_stream.send(()); // Can't access this field outside this module
            fail!()
        });
        port.recv()
    }

    #[test]
    fn try_send_and_recv_then_fail_before_ack() {
        let (port, chan) = rendezvous();
        spawn(proc() {
            port.duplex_stream.recv();
            fail!()
        });
        chan.try_send(());
    }

    #[test]
    #[should_fail]
    fn send_and_recv_then_fail_before_ack() {
        let (port, chan) = rendezvous();
        spawn(proc() {
            port.duplex_stream.recv();
            fail!()
        });
        chan.send(());
    }
}
<commit_msg>auto merge of #13093 : Havvy/rust/master, r=sfackler<commit_after>// Copyright 2012-2013 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.

/*!

Higher level communication abstractions.

*/

#[allow(missing_doc)];

use std::comm;

/// An extension of `pipes::stream` that allows both sending and receiving.
pub struct DuplexStream<S, R> {
    priv tx: Sender<S>,
    priv rx: Receiver<R>,
}

/// Creates a bidirectional stream.
pub fn duplex<S: Send, R: Send>() -> (DuplexStream<S, R>, DuplexStream<R, S>) {
    let (tx1, rx1) = channel();
    let (tx2, rx2) = channel();
    (DuplexStream { tx: tx1, rx: rx2 },
     DuplexStream { tx: tx2, rx: rx1 })
}

// Allow these methods to be used without import:
impl<S:Send,R:Send> DuplexStream<S, R> {
    pub fn send(&self, x: S) {
        self.tx.send(x)
    }
    pub fn try_send(&self, x: S) -> bool {
        self.tx.try_send(x)
    }
    pub fn recv(&self) -> R {
        self.rx.recv()
    }
    pub fn try_recv(&self) -> comm::TryRecvResult<R> {
        self.rx.try_recv()
    }
    pub fn recv_opt(&self) -> Option<R> {
        self.rx.recv_opt()
    }
}

/// An extension of `pipes::stream` that provides synchronous message sending.
pub struct SyncSender<S> { priv duplex_stream: DuplexStream<S, ()> }
/// An extension of `pipes::stream` that acknowledges each message received.
pub struct SyncReceiver<R> { priv duplex_stream: DuplexStream<(), R> }

impl<S: Send> SyncSender<S> {
    pub fn send(&self, val: S) {
        assert!(self.try_send(val), "SyncSender.send: receiving port closed");
    }

    /// Sends a message, or report if the receiver has closed the connection
    /// before receiving.
    pub fn try_send(&self, val: S) -> bool {
        self.duplex_stream.try_send(val) && self.duplex_stream.recv_opt().is_some()
    }
}

impl<R: Send> SyncReceiver<R> {
    pub fn recv(&self) -> R {
        self.recv_opt().expect("SyncReceiver.recv: sending channel closed")
    }

    pub fn recv_opt(&self) -> Option<R> {
        self.duplex_stream.recv_opt().map(|val| {
            self.duplex_stream.try_send(());
            val
        })
    }

    pub fn try_recv(&self) -> comm::TryRecvResult<R> {
        match self.duplex_stream.try_recv() {
            comm::Data(t) => { self.duplex_stream.try_send(()); comm::Data(t) }
            state => state,
        }
    }
}

/// Creates a stream whose channel, upon sending a message, blocks until the
/// message is received.
pub fn rendezvous<T: Send>() -> (SyncReceiver<T>, SyncSender<T>) {
    let (chan_stream, port_stream) = duplex();
    (SyncReceiver { duplex_stream: port_stream },
     SyncSender { duplex_stream: chan_stream })
}

#[cfg(test)]
mod test {
    use comm::{duplex, rendezvous};

    #[test]
    pub fn DuplexStream1() {
        let (left, right) = duplex();

        left.send(~"abc");
        right.send(123);

        assert!(left.recv() == 123);
        assert!(right.recv() == ~"abc");
    }

    #[test]
    pub fn basic_rendezvous_test() {
        let (port, chan) = rendezvous();

        spawn(proc() {
            chan.send("abc");
        });

        assert!(port.recv() == "abc");
    }

    #[test]
    fn recv_a_lot() {
        // Rendezvous streams should be able to handle any number of messages being sent
        let (port, chan) = rendezvous();
        spawn(proc() {
            for _ in range(0, 10000) { chan.send(()); }
        });
        for _ in range(0, 10000) { port.recv(); }
    }

    #[test]
    fn send_and_fail_and_try_recv() {
        let (port, chan) = rendezvous();
        spawn(proc() {
            chan.duplex_stream.send(()); // Can't access this field outside this module
            fail!()
        });
        port.recv()
    }

    #[test]
    fn try_send_and_recv_then_fail_before_ack() {
        let (port, chan) = rendezvous();
        spawn(proc() {
            port.duplex_stream.recv();
            fail!()
        });
        chan.try_send(());
    }

    #[test]
    #[should_fail]
    fn send_and_recv_then_fail_before_ack() {
        let (port, chan) = rendezvous();
        spawn(proc() {
            port.duplex_stream.recv();
            fail!()
        });
        chan.send(());
    }
}
<|endoftext|>