<commit_before>//! Rust bindings for [libressl](http://libressl.org)'s libtls
//! For the authoritative source on the inner workings of libtls check
//! the [manpage](http://www.openbsd.org/cgi-bin/man.cgi/OpenBSD-current/man3/tls_accept_fds.3?query=tls_init&sec=3).
//! The [raw](raw/index.html) module holds the bindings around libtls. A more idiomatic API is
//! is provided here.
//!
//! ## Client
//!
//! ```no_run
//! use std::io::Write;
//! use std::net::TcpStream;
//! let tcp = TcpStream::connect("google.com:443").unwrap();
//! let mut client = telos::new_client()
//!     .from_socket(&tcp, "google.com")
//!     .unwrap();
//! client.write("GET / HTTP/1.1\n\n".as_bytes()).unwrap();
//! ```
//!
//! ## Server
//!
//! The library does not handle TCP listening and binding, you need to handle the
//! TCP server accept() and then call `TlsServer::accept`
//!
//! ```no_run
//! use std::net::TcpListener;
//! let srv = TcpListener::bind("127.0.0.1:0").unwrap();
//! let addr = srv.local_addr().unwrap();
//! let mut tls_srv = telos::new_server()
//!     .key_file("tests/private_key.key")
//!     .cert_file("tests/certificate.crt")
//!     .bind().unwrap();
//! // Accept TCP connection, and then start TLS over it
//! let tcp_conn = srv.incoming().next().unwrap().unwrap();
//! let mut tls_conn = tls_srv.accept(&tcp_conn).unwrap();
//! ```
//!
//! ## Certificate Verification
//!
//! By default libtls will verify certificates using the system certificate store (usually defined
//! as /etc/ssl/cert.pem). In some Linux flavours and in Windows this file does not exist and you
//! will need to use one of the appropriate methods to load the correct certificates for your
//! system - check the Builder classes for the ca methods.
//!
//! ## Connection Lifetime
//!
//! The `from_socket()` and `accept()` methods build new connections on top of
//! `RawFd`/`RawSocket`. Destroying the TlsStream object WILL NOT cause the underlying
//! sockets to be closed, you need to close them yourself.

extern crate chrono;
extern crate libc;

use std::error::Error;
use std::io;
use std::io::{Read, Write};
#[cfg(unix)]
use std::os::unix::io::AsRawFd;
#[cfg(windows)]
use std::os::windows::io::AsRawSocket;
use chrono::datetime::DateTime;
use chrono::offset::utc::UTC;

mod util;
pub mod raw;
use raw::{TlsConfig, TlsContext};

pub use raw::{TlsResult, TlsError};

pub struct ClientBuilder {
    cfg: Option<TlsConfig>,
    error: Option<TlsError>,
}

impl ClientBuilder {
    /// Load CA certificates from PEM file
    pub fn ca_file(mut self, path: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_ca_file(path).err();
        }
        self
    }
    /// Load CA certificates from folder
    pub fn ca_path(mut self, path: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_ca_path(path).err();
        }
        self
    }
    /// Use CA certificates from PEM string
    pub fn ca(mut self, ca: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_ca_mem(ca).err();
        }
        self
    }
    pub fn verify_depth(mut self, depth: i32) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            cfg.set_verify_depth(depth);
        }
        self
    }
    pub fn protocols(mut self, protocols: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_protocols(protocols).err();
        }
        self
    }
    pub fn ciphers(mut self, ciphers: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_ciphers(ciphers).err();
        }
        self
    }
    /// Disable certificate verification
    pub fn insecure_noverifycert(mut self) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            cfg.insecure_noverifycert();
        }
        self
    }
    /// Disable hostname verification
    pub fn insecure_noverifyname(mut self) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            cfg.insecure_noverifyname();
        }
        self
    }

    /// Create client context from settings
    fn new_ctx(self) -> TlsResult<TlsContext> {
        if let Some(err) = self.error {
            Err(err)
        } else {
            let mut cli = try!(TlsContext::new_client());
            // This unwrap should be safe, we can't have a cfg without an error
            try!(cli.configure(self.cfg.unwrap()));
            Ok(cli)
        }
    }

    #[cfg(unix)]
    /// Establish a TLS connection over the given socket
    pub fn from_socket<F: AsRawFd>(self, ifd: &F, servername: &str) -> TlsResult<TlsStream> {
        let mut ctx = try!(self.new_ctx());
        let fd = ifd.as_raw_fd();
        try!(ctx.connect_socket(fd, servername));
        Ok(TlsStream {
            ctx: ctx,
        })
    }

    #[cfg(windows)]
    /// Establish a TLS connection over the given socket
    pub fn from_socket<F: AsRawSocket>(self, isock: &F, servername: &str) -> TlsResult<TlsStream> {
        let mut ctx = try!(self.new_ctx());
        let sock = isock.as_raw_socket();
        try!(ctx.connect_socket(sock, servername));
        Ok(TlsStream { ctx: ctx })
    }
}

/// Create a new TLS client
pub fn new_client() -> ClientBuilder {
    if !raw::init() {
        return ClientBuilder {
            cfg: None,
            error: Some(TlsError::new("Failed to initialize libtls")),
        }
    }

    match TlsConfig::new() {
        Ok(cfg) => {
            ClientBuilder {
                cfg: Some(cfg),
                error: None,
            }
        }
        Err(err) => {
            ClientBuilder {
                cfg: None,
                error: Some(err),
            }
        }
    }
}

pub struct TlsStream {
    ctx: TlsContext,
}

impl TlsStream {
    /// Executes the TLS handshake. This function is automatically called when reading or writing,
    /// you usually don't need to call it unless you want to force the handshake to finish sooner.
    ///
    /// Calling handshake multiple times, if the other end of the connection is not expecting it
    /// will usually result in an error.
    pub fn handshake(&mut self) -> TlsResult<()> {
        self.ctx.handshake()
    }

    /// Close TLS connection. This will not close the underlying transport.
    pub fn shutdown(&mut self) -> io::Result<()> {
        if let Err(err) = self.ctx.close() {
            if err.wants_more() {
                try!(self.ctx.close());
            } else {
                return Err(io::Error::from(err));
            }
        }
        Ok(())
    }

    /// Calling this method before the handshake is complete causes this method
    /// to return an empty string. See [handshake()](#method.handshake).
    pub fn certificate_issuer(&self) -> String {
        self.ctx.peer_cert_issuer()
    }
    /// Calling this method before the handshake is complete causes this method
    /// to return an empty string. See [handshake()](#method.handshake).
    pub fn certificate_hash(&self) -> String {
        self.ctx.peer_cert_hash()
    }
    /// Calling this method before the handshake is complete causes this method
    /// to return an empty string. See [handshake()](#method.handshake).
    pub fn certificate_subject(&self) -> String {
        self.ctx.peer_cert_subject()
    }
    pub fn peer_cert_provided(&self) -> bool {
        self.ctx.peer_cert_provided()
    }
    pub fn peer_cert_notbefore(&self) -> TlsResult<DateTime<UTC>> {
        self.ctx.peer_cert_notbefore()
    }
    pub fn peer_cert_notafter(&self) -> TlsResult<DateTime<UTC>> {
        self.ctx.peer_cert_notafter()
    }
    pub fn peer_cert_contains_name(&self, name: &str) -> bool {
        self.ctx.peer_cert_contains_name(name)
    }
    /// Calling this method before the handshake is complete causes this method
    /// to return an empty string. See [handshake()](#method.handshake).
    pub fn version(&self) -> String {
        self.ctx.conn_version()
    }
    /// Calling this method before the handshake is complete causes this method
    /// to return an empty string. See [handshake()](#method.handshake).
    pub fn cipher(&self) -> String {
        self.ctx.conn_cipher()
    }
}

impl Read for TlsStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.ctx
            .read(buf)
            .map_err(|err| io::Error::from(err))
    }
}

impl Write for TlsStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.ctx
            .write(buf)
            .map_err(|err| io::Error::from(err))
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Drop for TlsStream {
    fn drop(&mut self) {
        let _ = self.ctx.close();
    }
}

pub struct ServerBuilder {
    cfg: Option<TlsConfig>,
    error: Option<TlsError>,
}

impl ServerBuilder {
    pub fn key_file(mut self, path: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_key_file(path).err();
        }
        self
    }
    pub fn cert_file(mut self, path: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_cert_file(path).err();
        }
        self
    }

    /// Create server context from settings
    fn new_ctx(self) -> TlsResult<TlsContext> {
        if let Some(err) = self.error {
            Err(err)
        } else {
            let mut cli = try!(TlsContext::new_server());
            // This unwrap should be safe, we can't have a cfg without an error
            try!(cli.configure(self.cfg.unwrap()));
            Ok(cli)
        }
    }
    pub fn bind(self) -> TlsResult<TlsServer> {
        let ctx = try!(self.new_ctx());
        Ok(TlsServer { ctx: ctx })
    }
}

/// Create a new TLS server
pub fn new_server() -> ServerBuilder {
    if !raw::init() {
        return ServerBuilder {
            cfg: None,
            error: Some(TlsError::new("Failed to initialize libtls")),
        }
    }

    match TlsConfig::new() {
        Ok(cfg) => {
            ServerBuilder {
                cfg: Some(cfg),
                error: None,
            }
        }
        Err(err) => {
            ServerBuilder {
                cfg: None,
                error: Some(err),
            }
        }
    }
}

/// TLS Server, used to start TLS session over existing sockets.
pub struct TlsServer {
    ctx: TlsContext,
}

impl TlsServer {
    #[cfg(unix)]
    /// Start a new TLS connection over an existing socket (server-side)
    pub fn accept<F: AsRawFd>(&mut self, ifd: &F) -> io::Result<TlsStream> {
        let fd = ifd.as_raw_fd();
        let c = try!(self.ctx.accept_socket(fd));
        Ok(TlsStream {
            ctx: c,
        })
    }

    #[cfg(windows)]
    /// Start a new TLS connection over an existing socket (server-side)
    pub fn accept<F: AsRawSocket>(&mut self, isock: &F) -> TlsResult<TlsStream> {
        let sock = isock.as_raw_socket();
        let c = try!(self.ctx.accept_socket(sock));
        Ok(TlsStream { ctx: c })
    }
}

#[test]
fn test_protocols() {
    let mut cfg = TlsConfig::new().unwrap();

    // The following are all supported
    cfg.set_protocols("all").unwrap();
    cfg.set_protocols("legacy").unwrap();
    cfg.set_protocols("default").unwrap();
    cfg.set_protocols("secure").unwrap();
    cfg.set_protocols("tlsv1").unwrap();
    cfg.set_protocols("tlsv1.0").unwrap();
    cfg.set_protocols("tlsv1.1").unwrap();
    cfg.set_protocols("tlsv1.2").unwrap();

    // This is not valid
    assert!(cfg.set_protocols("unknown-proto").is_err());
}

#[test]
fn client_ctx_defs() {
    let c = TlsContext::new_client().unwrap();

    // These are the defaults before the connection is set
    assert_eq!(c.conn_version(), "");
    assert_eq!(c.conn_cipher(), "");
    assert!(c.peer_cert_notbefore().is_err());
    assert!(c.peer_cert_notafter().is_err());
    assert_eq!(c.peer_cert_issuer(), "");
    assert_eq!(c.peer_cert_subject(), "");
    assert_eq!(c.peer_cert_hash(), "");
    assert_eq!(c.peer_cert_contains_name("some.name"), false);
    assert_eq!(c.peer_cert_provided(), false);
}
<commit_msg>Docs: fix typo<commit_after>//! Rust bindings for [libressl](http://libressl.org)'s libtls
//! For the authoritative source on the inner workings of libtls check
//! the [manpage](http://www.openbsd.org/cgi-bin/man.cgi/OpenBSD-current/man3/tls_accept_fds.3?query=tls_init&sec=3).
//! The [raw](raw/index.html) module holds the bindings around libtls. A more idiomatic API
//! is provided here.
//!
//! ## Client
//!
//! ```no_run
//! use std::io::Write;
//! use std::net::TcpStream;
//! let tcp = TcpStream::connect("google.com:443").unwrap();
//! let mut client = telos::new_client()
//!     .from_socket(&tcp, "google.com")
//!     .unwrap();
//! client.write("GET / HTTP/1.1\n\n".as_bytes()).unwrap();
//! ```
//!
//! ## Server
//!
//! The library does not handle TCP listening and binding, you need to handle the
//! TCP server accept() and then call `TlsServer::accept`
//!
//! ```no_run
//! use std::net::TcpListener;
//! let srv = TcpListener::bind("127.0.0.1:0").unwrap();
//! let addr = srv.local_addr().unwrap();
//! let mut tls_srv = telos::new_server()
//!     .key_file("tests/private_key.key")
//!     .cert_file("tests/certificate.crt")
//!     .bind().unwrap();
//! // Accept TCP connection, and then start TLS over it
//! let tcp_conn = srv.incoming().next().unwrap().unwrap();
//! let mut tls_conn = tls_srv.accept(&tcp_conn).unwrap();
//! ```
//!
//! ## Certificate Verification
//!
//! By default libtls will verify certificates using the system certificate store (usually defined
//! as /etc/ssl/cert.pem). In some Linux flavours and in Windows this file does not exist and you
//! will need to use one of the appropriate methods to load the correct certificates for your
//! system - check the Builder classes for the ca methods.
//!
//! ## Connection Lifetime
//!
//! The `from_socket()` and `accept()` methods build new connections on top of
//! `RawFd`/`RawSocket`. Destroying the TlsStream object WILL NOT cause the underlying
//! sockets to be closed, you need to close them yourself.

extern crate chrono;
extern crate libc;

use std::error::Error;
use std::io;
use std::io::{Read, Write};
#[cfg(unix)]
use std::os::unix::io::AsRawFd;
#[cfg(windows)]
use std::os::windows::io::AsRawSocket;
use chrono::datetime::DateTime;
use chrono::offset::utc::UTC;

mod util;
pub mod raw;
use raw::{TlsConfig, TlsContext};

pub use raw::{TlsResult, TlsError};

pub struct ClientBuilder {
    cfg: Option<TlsConfig>,
    error: Option<TlsError>,
}

impl ClientBuilder {
    /// Load CA certificates from PEM file
    pub fn ca_file(mut self, path: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_ca_file(path).err();
        }
        self
    }
    /// Load CA certificates from folder
    pub fn ca_path(mut self, path: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_ca_path(path).err();
        }
        self
    }
    /// Use CA certificates from PEM string
    pub fn ca(mut self, ca: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_ca_mem(ca).err();
        }
        self
    }
    pub fn verify_depth(mut self, depth: i32) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            cfg.set_verify_depth(depth);
        }
        self
    }
    pub fn protocols(mut self, protocols: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_protocols(protocols).err();
        }
        self
    }
    pub fn ciphers(mut self, ciphers: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_ciphers(ciphers).err();
        }
        self
    }
    /// Disable certificate verification
    pub fn insecure_noverifycert(mut self) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            cfg.insecure_noverifycert();
        }
        self
    }
    /// Disable hostname verification
    pub fn insecure_noverifyname(mut self) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            cfg.insecure_noverifyname();
        }
        self
    }

    /// Create client context from settings
    fn new_ctx(self) -> TlsResult<TlsContext> {
        if let Some(err) = self.error {
            Err(err)
        } else {
            let mut cli = try!(TlsContext::new_client());
            // This unwrap should be safe, we can't have a cfg without an error
            try!(cli.configure(self.cfg.unwrap()));
            Ok(cli)
        }
    }

    #[cfg(unix)]
    /// Establish a TLS connection over the given socket
    pub fn from_socket<F: AsRawFd>(self, ifd: &F, servername: &str) -> TlsResult<TlsStream> {
        let mut ctx = try!(self.new_ctx());
        let fd = ifd.as_raw_fd();
        try!(ctx.connect_socket(fd, servername));
        Ok(TlsStream {
            ctx: ctx,
        })
    }

    #[cfg(windows)]
    /// Establish a TLS connection over the given socket
    pub fn from_socket<F: AsRawSocket>(self, isock: &F, servername: &str) -> TlsResult<TlsStream> {
        let mut ctx = try!(self.new_ctx());
        let sock = isock.as_raw_socket();
        try!(ctx.connect_socket(sock, servername));
        Ok(TlsStream { ctx: ctx })
    }
}

/// Create a new TLS client
pub fn new_client() -> ClientBuilder {
    if !raw::init() {
        return ClientBuilder {
            cfg: None,
            error: Some(TlsError::new("Failed to initialize libtls")),
        }
    }

    match TlsConfig::new() {
        Ok(cfg) => {
            ClientBuilder {
                cfg: Some(cfg),
                error: None,
            }
        }
        Err(err) => {
            ClientBuilder {
                cfg: None,
                error: Some(err),
            }
        }
    }
}

pub struct TlsStream {
    ctx: TlsContext,
}

impl TlsStream {
    /// Executes the TLS handshake. This function is automatically called when reading or writing,
    /// you usually don't need to call it unless you want to force the handshake to finish sooner.
    ///
    /// Calling handshake multiple times, if the other end of the connection is not expecting it
    /// will usually result in an error.
    pub fn handshake(&mut self) -> TlsResult<()> {
        self.ctx.handshake()
    }

    /// Close TLS connection. This will not close the underlying transport.
    pub fn shutdown(&mut self) -> io::Result<()> {
        if let Err(err) = self.ctx.close() {
            if err.wants_more() {
                try!(self.ctx.close());
            } else {
                return Err(io::Error::from(err));
            }
        }
        Ok(())
    }

    /// Calling this method before the handshake is complete causes this method
    /// to return an empty string. See [handshake()](#method.handshake).
    pub fn certificate_issuer(&self) -> String {
        self.ctx.peer_cert_issuer()
    }
    /// Calling this method before the handshake is complete causes this method
    /// to return an empty string. See [handshake()](#method.handshake).
    pub fn certificate_hash(&self) -> String {
        self.ctx.peer_cert_hash()
    }
    /// Calling this method before the handshake is complete causes this method
    /// to return an empty string. See [handshake()](#method.handshake).
    pub fn certificate_subject(&self) -> String {
        self.ctx.peer_cert_subject()
    }
    pub fn peer_cert_provided(&self) -> bool {
        self.ctx.peer_cert_provided()
    }
    pub fn peer_cert_notbefore(&self) -> TlsResult<DateTime<UTC>> {
        self.ctx.peer_cert_notbefore()
    }
    pub fn peer_cert_notafter(&self) -> TlsResult<DateTime<UTC>> {
        self.ctx.peer_cert_notafter()
    }
    pub fn peer_cert_contains_name(&self, name: &str) -> bool {
        self.ctx.peer_cert_contains_name(name)
    }
    /// Calling this method before the handshake is complete causes this method
    /// to return an empty string. See [handshake()](#method.handshake).
    pub fn version(&self) -> String {
        self.ctx.conn_version()
    }
    /// Calling this method before the handshake is complete causes this method
    /// to return an empty string. See [handshake()](#method.handshake).
    pub fn cipher(&self) -> String {
        self.ctx.conn_cipher()
    }
}

impl Read for TlsStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.ctx
            .read(buf)
            .map_err(|err| io::Error::from(err))
    }
}

impl Write for TlsStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.ctx
            .write(buf)
            .map_err(|err| io::Error::from(err))
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Drop for TlsStream {
    fn drop(&mut self) {
        let _ = self.ctx.close();
    }
}

pub struct ServerBuilder {
    cfg: Option<TlsConfig>,
    error: Option<TlsError>,
}

impl ServerBuilder {
    pub fn key_file(mut self, path: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_key_file(path).err();
        }
        self
    }
    pub fn cert_file(mut self, path: &str) -> Self {
        if self.error.is_some() {
            return self;
        }
        if let Some(cfg) = self.cfg.as_mut() {
            self.error = cfg.set_cert_file(path).err();
        }
        self
    }

    /// Create server context from settings
    fn new_ctx(self) -> TlsResult<TlsContext> {
        if let Some(err) = self.error {
            Err(err)
        } else {
            let mut cli = try!(TlsContext::new_server());
            // This unwrap should be safe, we can't have a cfg without an error
            try!(cli.configure(self.cfg.unwrap()));
            Ok(cli)
        }
    }
    pub fn bind(self) -> TlsResult<TlsServer> {
        let ctx = try!(self.new_ctx());
        Ok(TlsServer { ctx: ctx })
    }
}

/// Create a new TLS server
pub fn new_server() -> ServerBuilder {
    if !raw::init() {
        return ServerBuilder {
            cfg: None,
            error: Some(TlsError::new("Failed to initialize libtls")),
        }
    }

    match TlsConfig::new() {
        Ok(cfg) => {
            ServerBuilder {
                cfg: Some(cfg),
                error: None,
            }
        }
        Err(err) => {
            ServerBuilder {
                cfg: None,
                error: Some(err),
            }
        }
    }
}

/// TLS Server, used to start TLS session over existing sockets.
pub struct TlsServer {
    ctx: TlsContext,
}

impl TlsServer {
    #[cfg(unix)]
    /// Start a new TLS connection over an existing socket (server-side)
    pub fn accept<F: AsRawFd>(&mut self, ifd: &F) -> io::Result<TlsStream> {
        let fd = ifd.as_raw_fd();
        let c = try!(self.ctx.accept_socket(fd));
        Ok(TlsStream {
            ctx: c,
        })
    }

    #[cfg(windows)]
    /// Start a new TLS connection over an existing socket (server-side)
    pub fn accept<F: AsRawSocket>(&mut self, isock: &F) -> TlsResult<TlsStream> {
        let sock = isock.as_raw_socket();
        let c = try!(self.ctx.accept_socket(sock));
        Ok(TlsStream { ctx: c })
    }
}

#[test]
fn test_protocols() {
    let mut cfg = TlsConfig::new().unwrap();

    // The following are all supported
    cfg.set_protocols("all").unwrap();
    cfg.set_protocols("legacy").unwrap();
    cfg.set_protocols("default").unwrap();
    cfg.set_protocols("secure").unwrap();
    cfg.set_protocols("tlsv1").unwrap();
    cfg.set_protocols("tlsv1.0").unwrap();
    cfg.set_protocols("tlsv1.1").unwrap();
    cfg.set_protocols("tlsv1.2").unwrap();

    // This is not valid
    assert!(cfg.set_protocols("unknown-proto").is_err());
}

#[test]
fn client_ctx_defs() {
    let c = TlsContext::new_client().unwrap();

    // These are the defaults before the connection is set
    assert_eq!(c.conn_version(), "");
    assert_eq!(c.conn_cipher(), "");
    assert!(c.peer_cert_notbefore().is_err());
    assert!(c.peer_cert_notafter().is_err());
    assert_eq!(c.peer_cert_issuer(), "");
    assert_eq!(c.peer_cert_subject(), "");
    assert_eq!(c.peer_cert_hash(), "");
    assert_eq!(c.peer_cert_contains_name("some.name"), false);
    assert_eq!(c.peer_cert_provided(), false);
}
<|endoftext|>


<commit_before>//! Http client request
#![allow(unused_imports, dead_code)]
use std::{fmt, io, str};
use std::rc::Rc;
use std::time::Duration;
use std::cell::UnsafeCell;

use base64;
use rand;
use cookie::Cookie;
use bytes::BytesMut;
use http::{HttpTryFrom, StatusCode, Error as HttpError};
use http::header::{self, HeaderName, HeaderValue};
use sha1::Sha1;
use futures::{Async, Future, Poll, Stream};
use futures::future::{Either, err as FutErr};
use tokio_core::net::TcpStream;

use actix::prelude::*;

use body::Binary;
use error::UrlParseError;
use server::shared::SharedBytes;

use server::{utils, IoStream};
use client::{ClientRequest, ClientRequestBuilder,
             HttpResponseParser, HttpResponseParserError, HttpClientWriter};
use client::{Connect, Connection, ClientConnector, ClientConnectorError};

use super::Message;
use super::proto::{CloseCode, OpCode};
use super::frame::Frame;

pub type WsClientFuture =
    Future<Item=(WsClientReader, WsClientWriter), Error=WsClientError>;

/// Websockt client error
#[derive(Fail, Debug)]
pub enum WsClientError {
    #[fail(display="Invalid url")]
    InvalidUrl,
    #[fail(display="Invalid response status")]
    InvalidResponseStatus,
    #[fail(display="Invalid upgrade header")]
    InvalidUpgradeHeader,
    #[fail(display="Invalid connection header")]
    InvalidConnectionHeader,
    #[fail(display="Invalid challenge response")]
    InvalidChallengeResponse,
    #[fail(display="Http parsing error")]
    Http(HttpError),
    #[fail(display="Url parsing error")]
    Url(UrlParseError),
    #[fail(display="Response parsing error")]
    ResponseParseError(HttpResponseParserError),
    #[fail(display="{}", _0)]
    Connector(ClientConnectorError),
    #[fail(display="{}", _0)]
    Io(io::Error),
    #[fail(display="Disconnected")]
    Disconnected,
}

impl From<HttpError> for WsClientError {
    fn from(err: HttpError) -> WsClientError {
        WsClientError::Http(err)
    }
}

impl From<UrlParseError> for WsClientError {
    fn from(err: UrlParseError) -> WsClientError {
        WsClientError::Url(err)
    }
}

impl From<ClientConnectorError> for WsClientError {
    fn from(err: ClientConnectorError) -> WsClientError {
        WsClientError::Connector(err)
    }
}

impl From<io::Error> for WsClientError {
    fn from(err: io::Error) -> WsClientError {
        WsClientError::Io(err)
    }
}

impl From<HttpResponseParserError> for WsClientError {
    fn from(err: HttpResponseParserError) -> WsClientError {
        WsClientError::ResponseParseError(err)
    }
}

/// Websockt client
pub struct WsClient {
    request: ClientRequestBuilder,
    err: Option<WsClientError>,
    http_err: Option<HttpError>,
    origin: Option<HeaderValue>,
    protocols: Option<String>,
    conn: Address<ClientConnector>,
}

impl WsClient {

    /// Create new websocket connection
    pub fn new<S: AsRef<str>>(uri: S) -> WsClient {
        WsClient::with_connector(uri, ClientConnector::from_registry())
    }

    /// Create new websocket connection with custom `ClientConnector`
    pub fn with_connector<S: AsRef<str>>(uri: S, conn: Address<ClientConnector>) -> WsClient {
        let mut cl = WsClient {
            request: ClientRequest::build(),
            err: None,
            http_err: None,
            origin: None,
            protocols: None,
            conn: conn,
        };
        cl.request.uri(uri.as_ref());
        cl
    }

    /// Set supported websocket protocols
    pub fn protocols<U, V>(&mut self, protos: U) -> &mut Self
        where U: IntoIterator<Item=V> + 'static,
              V: AsRef<str>
    {
        let mut protos = protos.into_iter()
            .fold(String::new(), |acc, s| {acc + s.as_ref() + ","});
        protos.pop();
        self.protocols = Some(protos);
        self
    }

    /// Set cookie for handshake request
    pub fn cookie<'c>(&mut self, cookie: Cookie<'c>) -> &mut Self {
        self.request.cookie(cookie);
        self
    }

    /// Set request Origin
    pub fn origin<V>(&mut self, origin: V) -> &mut Self
        where HeaderValue: HttpTryFrom<V>
    {
        match HeaderValue::try_from(origin) {
            Ok(value) => self.origin = Some(value),
            Err(e) => self.http_err = Some(e.into()),
        }
        self
    }

    /// Set request header
    pub fn header<K, V>(&mut self, key: K, value: V) -> &mut Self
        where HeaderName: HttpTryFrom<K>, HeaderValue: HttpTryFrom<V>
    {
        self.request.header(key, value);
        self
    }

    /// Connect to websocket server and do ws handshake
    pub fn connect(&mut self) -> Result<Box<WsClientFuture>, WsClientError> {
        if let Some(e) = self.err.take() {
            return Err(e)
        }
        if let Some(e) = self.http_err.take() {
            return Err(e.into())
        }

        // origin
        if let Some(origin) = self.origin.take() {
            self.request.set_header(header::ORIGIN, origin);
        }

        self.request.set_header(header::UPGRADE, "websocket");
        self.request.set_header(header::CONNECTION, "upgrade");
        self.request.set_header("SEC-WEBSOCKET-VERSION", "13");

        if let Some(protocols) = self.protocols.take() {
            self.request.set_header("SEC-WEBSOCKET-PROTOCOL", protocols.as_str());
        }
        let request = self.request.finish()?;

        if request.uri().host().is_none() {
            return Err(WsClientError::InvalidUrl)
        }
        if let Some(scheme) = request.uri().scheme_part() {
            if scheme != "http" && scheme != "https" && scheme != "ws" && scheme != "wss" {
                return Err(WsClientError::InvalidUrl);
            }
        } else {
            return Err(WsClientError::InvalidUrl);
        }

        // get connection and start handshake
        Ok(Box::new(
            self.conn.call_fut(Connect(request.uri().clone()))
                .map_err(|_| WsClientError::Disconnected)
                .and_then(|res| match res {
                    Ok(stream) => Either::A(WsHandshake::new(stream, request)),
                    Err(err) => Either::B(FutErr(err.into())),
                })
        ))
    }
}

struct WsInner {
    conn: Connection,
    writer: HttpClientWriter,
    parser: HttpResponseParser,
    parser_buf: BytesMut,
    closed: bool,
    error_sent: bool,
}

struct WsHandshake {
    inner: Option<WsInner>,
    request: ClientRequest,
    sent: bool,
    key: String,
}

impl WsHandshake {
    fn new(conn: Connection, mut request: ClientRequest) -> WsHandshake {
        // Generate a random key for the `Sec-WebSocket-Key` header.
        // a base64-encoded (see Section 4 of [RFC4648]) value that,
        // when decoded, is 16 bytes in length (RFC 6455)
        let sec_key: [u8; 16] = rand::random();
        let key = base64::encode(&sec_key);

        request.headers_mut().insert(
            HeaderName::try_from("SEC-WEBSOCKET-KEY").unwrap(),
            HeaderValue::try_from(key.as_str()).unwrap());

        let inner = WsInner {
            conn: conn,
            writer: HttpClientWriter::new(SharedBytes::default()),
            parser: HttpResponseParser::default(),
            parser_buf: BytesMut::new(),
            closed: false,
            error_sent: false,
        };

        WsHandshake {
            key: key,
            inner: Some(inner),
            request: request,
            sent: false,
        }
    }
}

impl Future for WsHandshake {
    type Item = (WsClientReader, WsClientWriter);
    type Error = WsClientError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut inner = self.inner.take().unwrap();

        if !self.sent {
            self.sent = true;
            inner.writer.start(&mut self.request);
        }
        if let Err(err) = inner.writer.poll_completed(&mut inner.conn, false) {
            return Err(err.into())
        }

        match inner.parser.parse(&mut inner.conn, &mut inner.parser_buf) {
            Ok(Async::Ready(resp)) => {
                // verify response
                if resp.status() != StatusCode::SWITCHING_PROTOCOLS {
                    return Err(WsClientError::InvalidResponseStatus)
                }
                // Check for "UPGRADE" to websocket header
                let has_hdr = if let Some(hdr) = resp.headers().get(header::UPGRADE) {
                    if let Ok(s) = hdr.to_str() {
                        s.to_lowercase().contains("websocket")
                    } else {
                        false
                    }
                } else {
                    false
                };
                if !has_hdr {
                    return Err(WsClientError::InvalidUpgradeHeader)
                }
                // Check for "CONNECTION" header
                let has_hdr = if let Some(conn) = resp.headers().get(header::CONNECTION) {
                    if let Ok(s) = conn.to_str() {
                        s.to_lowercase().contains("upgrade")
                    } else { false }
                } else { false };
                if !has_hdr {
                    return Err(WsClientError::InvalidConnectionHeader)
                }

                let match_key = if let Some(key) = resp.headers().get(
                    HeaderName::try_from("SEC-WEBSOCKET-ACCEPT").unwrap())
                {
                    // field is constructed by concatenating /key/
                    // with the string "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" (RFC 6455)
                    const WS_GUID: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
                    let mut sha1 = Sha1::new();
                    sha1.update(self.key.as_ref());
                    sha1.update(WS_GUID);
                    key.as_bytes() == base64::encode(&sha1.digest().bytes()).as_bytes()
                } else {
                    false
                };
                if !match_key {
                    return Err(WsClientError::InvalidChallengeResponse)
                }

                let inner = Rc::new(UnsafeCell::new(Inner{inner: inner}));
                Ok(Async::Ready(
                    (WsClientReader{inner: Rc::clone(&inner)},
                     WsClientWriter{inner: inner})))
            },
            Ok(Async::NotReady) => {
                self.inner = Some(inner);
                Ok(Async::NotReady)
            },
            Err(err) => Err(err.into())
        }
    }
}

struct Inner {
    inner: WsInner,
}

pub struct WsClientReader {
    inner: Rc<UnsafeCell<Inner>>
}

impl fmt::Debug for WsClientReader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "WsClientReader()")
    }
}

impl WsClientReader {
    #[inline]
    fn as_mut(&mut self) -> &mut Inner {
        unsafe{ &mut *self.inner.get() }
    }
}

impl Stream for WsClientReader {
    type Item = Message;
    type Error = WsClientError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        let inner = self.as_mut();
        let mut done = false;

        match utils::read_from_io(&mut inner.inner.conn, &mut inner.inner.parser_buf) {
            Ok(Async::Ready(0)) => {
                done = true;
                inner.inner.closed = true;
            },
            Ok(Async::Ready(_)) | Ok(Async::NotReady) => (),
            Err(err) =>
                return Err(err.into())
        }

        // write
        let _ = inner.inner.writer.poll_completed(&mut inner.inner.conn, false);

        // read
        match Frame::parse(&mut inner.inner.parser_buf) {
            Ok(Some(frame)) => {
                // trace!("WsFrame {}", frame);
                let (_finished, opcode, payload) = frame.unpack();

                match opcode {
                    OpCode::Continue => unimplemented!(),
                    OpCode::Bad =>
                        Ok(Async::Ready(Some(Message::Error))),
                    OpCode::Close => {
                        inner.inner.closed = true;
                        inner.inner.error_sent = true;
                        Ok(Async::Ready(Some(Message::Closed)))
                    },
                    OpCode::Ping =>
                        Ok(Async::Ready(Some(
                            Message::Ping(
                                String::from_utf8_lossy(payload.as_ref()).into())))),
                    OpCode::Pong =>
                        Ok(Async::Ready(Some(
                            Message::Pong(
                                String::from_utf8_lossy(payload.as_ref()).into())))),
                    OpCode::Binary =>
                        Ok(Async::Ready(Some(Message::Binary(payload)))),
                    OpCode::Text => {
                        let tmp = Vec::from(payload.as_ref());
                        match String::from_utf8(tmp) {
                            Ok(s) =>
                                Ok(Async::Ready(Some(Message::Text(s)))),
                            Err(_) =>
                                Ok(Async::Ready(Some(Message::Error))),
                        }
                    }
                }
            }
            Ok(None) => {
                if done {
                    Ok(Async::Ready(None))
                } else if inner.inner.closed {
                    if !inner.inner.error_sent {
                        inner.inner.error_sent = true;
                        Ok(Async::Ready(Some(Message::Closed)))
                    } else {
                        Ok(Async::Ready(None))
                    }
                } else {
                    Ok(Async::NotReady)
                }
            },
            Err(err) => {
                inner.inner.closed = true;
                inner.inner.error_sent = true;
                Err(err.into())
            }
        }
    }
}

pub struct WsClientWriter {
    inner: Rc<UnsafeCell<Inner>>
}

impl WsClientWriter {
    #[inline]
    fn as_mut(&mut self) -> &mut Inner {
        unsafe{ &mut *self.inner.get() }
    }
}

impl WsClientWriter {

    /// Write payload
    #[inline]
    fn write<B: Into<Binary>>(&mut self, data: B) {
        if !self.as_mut().inner.closed {
            let _ = self.as_mut().inner.writer.write(&data.into());
        } else {
            warn!("Trying to write to disconnected response");
        }
    }

    /// Send text frame
    pub fn text(&mut self, text: &str) {
        let mut frame = Frame::message(Vec::from(text), OpCode::Text, true);
        let mut buf = Vec::new();
        frame.format(&mut buf).unwrap();

        self.write(buf);
    }

    /// Send binary frame
    pub fn binary<B: Into<Binary>>(&mut self, data: B) {
        let mut frame = Frame::message(data, OpCode::Binary, true);
        let mut buf = Vec::new();
        frame.format(&mut buf).unwrap();

        self.write(buf);
    }

    /// Send ping frame
    pub fn ping(&mut self, message: &str) {
        let mut frame = Frame::message(Vec::from(message), OpCode::Ping, true);
        let mut buf = Vec::new();
        frame.format(&mut buf).unwrap();

        self.write(buf);
    }

    /// Send pong frame
    pub fn pong(&mut self, message: &str) {
        let mut frame = Frame::message(Vec::from(message), OpCode::Pong, true);
        let mut buf = Vec::new();
        frame.format(&mut buf).unwrap();

        self.write(buf);
    }

    /// Send close frame
    pub fn close(&mut self, code: CloseCode, reason: &str) {
        let mut frame = Frame::close(code, reason);
        let mut buf = Vec::new();
        frame.format(&mut buf).unwrap();
        self.write(buf);
    }
}
<commit_msg>link to websocket example<commit_after>//! Http client request
#![allow(unused_imports, dead_code)]
use std::{fmt, io, str};
use std::rc::Rc;
use std::time::Duration;
use std::cell::UnsafeCell;

use base64;
use rand;
use cookie::Cookie;
use bytes::BytesMut;
use http::{HttpTryFrom, StatusCode, Error as HttpError};
use http::header::{self, HeaderName, HeaderValue};
use sha1::Sha1;
use futures::{Async, Future, Poll, Stream};
use futures::future::{Either, err as FutErr};
use tokio_core::net::TcpStream;

use actix::prelude::*;

use body::Binary;
use error::UrlParseError;
use server::shared::SharedBytes;

use server::{utils, IoStream};
use client::{ClientRequest, ClientRequestBuilder,
             HttpResponseParser, HttpResponseParserError, HttpClientWriter};
use client::{Connect, Connection, ClientConnector, ClientConnectorError};

use super::Message;
use super::proto::{CloseCode, OpCode};
use super::frame::Frame;

pub type WsClientFuture =
    Future<Item=(WsClientReader, WsClientWriter), Error=WsClientError>;

/// Websockt client error
#[derive(Fail, Debug)]
pub enum WsClientError {
    #[fail(display="Invalid url")]
    InvalidUrl,
    #[fail(display="Invalid response status")]
    InvalidResponseStatus,
    #[fail(display="Invalid upgrade header")]
    InvalidUpgradeHeader,
    #[fail(display="Invalid connection header")]
    InvalidConnectionHeader,
    #[fail(display="Invalid challenge response")]
    InvalidChallengeResponse,
    #[fail(display="Http parsing error")]
    Http(HttpError),
    #[fail(display="Url parsing error")]
    Url(UrlParseError),
    #[fail(display="Response parsing error")]
    ResponseParseError(HttpResponseParserError),
    #[fail(display="{}", _0)]
    Connector(ClientConnectorError),
    #[fail(display="{}", _0)]
    Io(io::Error),
    #[fail(display="Disconnected")]
    Disconnected,
}

impl From<HttpError> for WsClientError {
    fn from(err: HttpError) -> WsClientError {
        WsClientError::Http(err)
    }
}

impl From<UrlParseError> for WsClientError {
    fn from(err: UrlParseError) -> WsClientError {
        WsClientError::Url(err)
    }
}

impl From<ClientConnectorError> for WsClientError {
    fn from(err: ClientConnectorError) -> WsClientError {
        WsClientError::Connector(err)
    }
}

impl From<io::Error> for WsClientError {
    fn from(err: io::Error) -> WsClientError {
        WsClientError::Io(err)
    }
}

impl From<HttpResponseParserError> for WsClientError {
    fn from(err: HttpResponseParserError) -> WsClientError {
        WsClientError::ResponseParseError(err)
    }
}

/// WebSocket client
///
/// Example of WebSocket client usage is available in
/// [websocket example](
/// https://github.com/actix/actix-web/blob/master/examples/websocket/src/client.rs#L24)
pub struct WsClient {
    request: ClientRequestBuilder,
    err: Option<WsClientError>,
    http_err: Option<HttpError>,
    origin: Option<HeaderValue>,
    protocols: Option<String>,
    conn: Address<ClientConnector>,
}

impl WsClient {

    /// Create new websocket connection
    pub fn new<S: AsRef<str>>(uri: S) -> WsClient {
        WsClient::with_connector(uri, ClientConnector::from_registry())
    }

    /// Create new websocket connection with custom `ClientConnector`
    pub fn with_connector<S: AsRef<str>>(uri: S, conn: Address<ClientConnector>) -> WsClient {
        let mut cl = WsClient {
            request: ClientRequest::build(),
            err: None,
            http_err: None,
            origin: None,
            protocols: None,
            conn: conn,
        };
        cl.request.uri(uri.as_ref());
        cl
    }

    /// Set supported websocket protocols
    pub fn protocols<U, V>(&mut self, protos: U) -> &mut Self
        where U: IntoIterator<Item=V> + 'static,
              V: AsRef<str>
    {
        let mut protos = protos.into_iter()
            .fold(String::new(), |acc, s| {acc + s.as_ref() + ","});
        protos.pop();
        self.protocols = Some(protos);
        self
    }

    /// Set cookie for handshake request
    pub fn cookie<'c>(&mut self, cookie: Cookie<'c>) -> &mut Self {
        self.request.cookie(cookie);
        self
    }

    /// Set request Origin
    pub fn origin<V>(&mut self, origin: V) -> &mut Self
        where HeaderValue: HttpTryFrom<V>
    {
        match HeaderValue::try_from(origin) {
            Ok(value) => self.origin = Some(value),
            Err(e) => self.http_err = Some(e.into()),
        }
        self
    }

    /// Set request header
    pub fn header<K, V>(&mut self, key: K, value: V) -> &mut Self
        where HeaderName: HttpTryFrom<K>, HeaderValue: HttpTryFrom<V>
    {
        self.request.header(key, value);
        self
    }

    /// Connect to websocket server and do ws handshake
    pub fn connect(&mut self) -> Result<Box<WsClientFuture>, WsClientError> {
        if let Some(e) = self.err.take() {
            return Err(e)
        }
        if let Some(e) = self.http_err.take() {
            return Err(e.into())
        }

        // origin
        if let Some(origin) = self.origin.take() {
            self.request.set_header(header::ORIGIN, origin);
        }

        self.request.set_header(header::UPGRADE, "websocket");
        self.request.set_header(header::CONNECTION, "upgrade");
        self.request.set_header("SEC-WEBSOCKET-VERSION", "13");

        if let Some(protocols) = self.protocols.take() {
            self.request.set_header("SEC-WEBSOCKET-PROTOCOL", protocols.as_str());
        }
        let request = self.request.finish()?;

        if request.uri().host().is_none() {
            return Err(WsClientError::InvalidUrl)
        }
        if let Some(scheme) = request.uri().scheme_part() {
            if scheme != "http" && scheme != "https" && scheme != "ws" && scheme != "wss" {
                return Err(WsClientError::InvalidUrl);
            }
        } else {
            return Err(WsClientError::InvalidUrl);
        }

        // get connection and start handshake
        Ok(Box::new(
            self.conn.call_fut(Connect(request.uri().clone()))
                .map_err(|_| WsClientError::Disconnected)
                .and_then(|res| match res {
                    Ok(stream) => Either::A(WsHandshake::new(stream, request)),
                    Err(err) => Either::B(FutErr(err.into())),
                })
        ))
    }
}

struct WsInner {
    conn: Connection,
    writer: HttpClientWriter,
    parser: HttpResponseParser,
    parser_buf: BytesMut,
    closed: bool,
    error_sent: bool,
}

struct WsHandshake {
    inner: Option<WsInner>,
    request: ClientRequest,
    sent: bool,
    key: String,
}

impl WsHandshake {
    fn new(conn: Connection, mut request: ClientRequest) -> WsHandshake {
        // Generate a random key for the `Sec-WebSocket-Key` header.
        // a base64-encoded (see Section 4 of [RFC4648]) value that,
        // when decoded, is 16 bytes in length (RFC 6455)
        let sec_key: [u8; 16] = rand::random();
        let key = base64::encode(&sec_key);

        request.headers_mut().insert(
            HeaderName::try_from("SEC-WEBSOCKET-KEY").unwrap(),
            HeaderValue::try_from(key.as_str()).unwrap());

        let inner = WsInner {
            conn: conn,
            writer: HttpClientWriter::new(SharedBytes::default()),
            parser: HttpResponseParser::default(),
            parser_buf: BytesMut::new(),
            closed: false,
            error_sent: false,
        };

        WsHandshake {
            key: key,
            inner: Some(inner),
            request: request,
            sent: false,
        }
    }
}

impl Future for WsHandshake {
    type Item = (WsClientReader, WsClientWriter);
    type Error = WsClientError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut inner = self.inner.take().unwrap();

        if !self.sent {
            self.sent = true;
            inner.writer.start(&mut self.request);
        }
        if let Err(err) = inner.writer.poll_completed(&mut inner.conn, false) {
            return Err(err.into())
        }

        match inner.parser.parse(&mut inner.conn, &mut inner.parser_buf) {
            Ok(Async::Ready(resp)) => {
                // verify response
                if resp.status() != StatusCode::SWITCHING_PROTOCOLS {
                    return Err(WsClientError::InvalidResponseStatus)
                }
                // Check for "UPGRADE" to websocket header
                let has_hdr = if let Some(hdr) = resp.headers().get(header::UPGRADE) {
                    if let Ok(s) = hdr.to_str() {
                        s.to_lowercase().contains("websocket")
                    } else {
                        false
                    }
                } else {
                    false
                };
                if !has_hdr {
                    return Err(WsClientError::InvalidUpgradeHeader)
                }
                // Check for "CONNECTION" header
                let has_hdr = if let Some(conn) = resp.headers().get(header::CONNECTION) {
                    if let Ok(s) = conn.to_str() {
                        s.to_lowercase().contains("upgrade")
                    } else { false }
                } else { false };
                if !has_hdr {
                    return Err(WsClientError::InvalidConnectionHeader)
                }

                let match_key = if let Some(key) = resp.headers().get(
                    HeaderName::try_from("SEC-WEBSOCKET-ACCEPT").unwrap())
                {
                    // field is constructed by concatenating /key/
                    // with the string "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" (RFC 6455)
                    const WS_GUID: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
                    let mut sha1 = Sha1::new();
                    sha1.update(self.key.as_ref());
                    sha1.update(WS_GUID);
                    key.as_bytes() == base64::encode(&sha1.digest().bytes()).as_bytes()
                } else {
                    false
                };
                if !match_key {
                    return Err(WsClientError::InvalidChallengeResponse)
                }

                let inner = Rc::new(UnsafeCell::new(Inner{inner: inner}));
                Ok(Async::Ready(
                    (WsClientReader{inner: Rc::clone(&inner)},
                     WsClientWriter{inner: inner})))
            },
            Ok(Async::NotReady) => {
                self.inner = Some(inner);
                Ok(Async::NotReady)
            },
            Err(err) => Err(err.into())
        }
    }
}

struct Inner {
    inner: WsInner,
}

pub struct WsClientReader {
    inner: Rc<UnsafeCell<Inner>>
}

impl fmt::Debug for WsClientReader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "WsClientReader()")
    }
}

impl WsClientReader {
    #[inline]
    fn as_mut(&mut self) -> &mut Inner {
        unsafe{ &mut *self.inner.get() }
    }
}

impl Stream for WsClientReader {
    type Item = Message;
    type Error = WsClientError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        let inner = self.as_mut();
        let mut done = false;

        match utils::read_from_io(&mut inner.inner.conn, &mut inner.inner.parser_buf) {
            Ok(Async::Ready(0)) => {
                done = true;
                inner.inner.closed = true;
            },
            Ok(Async::Ready(_)) | Ok(Async::NotReady) => (),
            Err(err) =>
                return Err(err.into())
        }

        // write
        let _ = inner.inner.writer.poll_completed(&mut inner.inner.conn, false);

        // read
        match Frame::parse(&mut inner.inner.parser_buf) {
            Ok(Some(frame)) => {
                // trace!("WsFrame {}", frame);
                let (_finished, opcode, payload) = frame.unpack();

                match opcode {
                    OpCode::Continue => unimplemented!(),
                    OpCode::Bad =>
                        Ok(Async::Ready(Some(Message::Error))),
                    OpCode::Close => {
                        inner.inner.closed = true;
                        inner.inner.error_sent = true;
                        Ok(Async::Ready(Some(Message::Closed)))
                    },
                    OpCode::Ping =>
                        Ok(Async::Ready(Some(
                            Message::Ping(
                                String::from_utf8_lossy(payload.as_ref()).into())))),
                    OpCode::Pong =>
                        Ok(Async::Ready(Some(
                            Message::Pong(
                                String::from_utf8_lossy(payload.as_ref()).into())))),
                    OpCode::Binary =>
                        Ok(Async::Ready(Some(Message::Binary(payload)))),
                    OpCode::Text => {
                        let tmp = Vec::from(payload.as_ref());
                        match String::from_utf8(tmp) {
                            Ok(s) =>
                                Ok(Async::Ready(Some(Message::Text(s)))),
                            Err(_) =>
                                Ok(Async::Ready(Some(Message::Error))),
                        }
                    }
                }
            }
            Ok(None) => {
                if done {
                    Ok(Async::Ready(None))
                } else if inner.inner.closed {
                    if !inner.inner.error_sent {
                        inner.inner.error_sent = true;
                        Ok(Async::Ready(Some(Message::Closed)))
                    } else {
                        Ok(Async::Ready(None))
                    }
                } else {
                    Ok(Async::NotReady)
                }
            },
            Err(err) => {
                inner.inner.closed = true;
                inner.inner.error_sent = true;
                Err(err.into())
            }
        }
    }
}

pub struct WsClientWriter {
    inner: Rc<UnsafeCell<Inner>>
}

impl WsClientWriter {
    #[inline]
    fn as_mut(&mut self) -> &mut Inner {
        unsafe{ &mut *self.inner.get() }
    }
}

impl WsClientWriter {

    /// Write payload
    #[inline]
    fn write<B: Into<Binary>>(&mut self, data: B) {
        if !self.as_mut().inner.closed {
            let _ = self.as_mut().inner.writer.write(&data.into());
        } else {
            warn!("Trying to write to disconnected response");
        }
    }

    /// Send text frame
    pub fn text(&mut self, text: &str) {
        let mut frame = Frame::message(Vec::from(text), OpCode::Text, true);
        let mut buf = Vec::new();
        frame.format(&mut buf).unwrap();

        self.write(buf);
    }

    /// Send binary frame
    pub fn binary<B: Into<Binary>>(&mut self, data: B) {
        let mut frame = Frame::message(data, OpCode::Binary, true);
        let mut buf = Vec::new();
        frame.format(&mut buf).unwrap();

        self.write(buf);
    }

    /// Send ping frame
    pub fn ping(&mut self, message: &str) {
        let mut frame = Frame::message(Vec::from(message), OpCode::Ping, true);
        let mut buf = Vec::new();
        frame.format(&mut buf).unwrap();

        self.write(buf);
    }

    /// Send pong frame
    pub fn pong(&mut self, message: &str) {
        let mut frame = Frame::message(Vec::from(message), OpCode::Pong, true);
        let mut buf = Vec::new();
        frame.format(&mut buf).unwrap();

        self.write(buf);
    }

    /// Send close frame
    pub fn close(&mut self, code: CloseCode, reason: &str) {
        let mut frame = Frame::close(code, reason);
        let mut buf = Vec::new();
        frame.format(&mut buf).unwrap();
        self.write(buf);
    }
}
<|endoftext|>


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use error::TimeTrackError as TTE;
use error::TimeTrackErrorKind as TTEK;
use error::MapErrInto;

use libimagstore::store::FileLockEntry;
use libimagstore::store::Store;
use libimagstore::storeid::StoreIdIterator;
use libimagerror::into::IntoError;

use iter::get::GetTimeTrackIter;
use tag::TimeTrackingTag as TTT;
use timetracking::TimeTracking;

pub struct WithOneOf<'a> {
    iter: GetTimeTrackIter<'a>,
    allowed_tags: &'a Vec<TTT>,
}

impl<'a> WithOneOf<'a> {

    pub fn new(iter: GetTimeTrackIter<'a>, allowed_tags: &'a Vec<TTT>) -> WithOneOf<'a> {
        WithOneOf {
            iter: iter,
            allowed_tags: allowed_tags
        }
    }
}

impl<'a> Iterator for WithOneOf<'a> {
    type Item = Result<FileLockEntry<'a>, TTE>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.iter.next() {
                Some(Ok(fle)) => {
                    match fle.get_timetrack_tag() {
                        Err(e) => return Some(Err(e)),
                        Ok(t) => if self.allowed_tags.contains(&t) {
                            return Some(Ok(fle))
                        } else {
                            // loop
                        },
                    }
                },
                Some(Err(e)) => return Some(Err(e)),
                None => return None,
            }
        }
    }
}

<commit_msg>Add iterator to forbid tags<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use error::TimeTrackError as TTE;
use error::TimeTrackErrorKind as TTEK;
use error::MapErrInto;

use libimagstore::store::FileLockEntry;
use libimagstore::store::Store;
use libimagstore::storeid::StoreIdIterator;
use libimagerror::into::IntoError;

use iter::get::GetTimeTrackIter;
use tag::TimeTrackingTag as TTT;
use timetracking::TimeTracking;

pub struct WithOneOf<'a> {
    iter: GetTimeTrackIter<'a>,
    allowed_tags: &'a Vec<TTT>,
}

impl<'a> WithOneOf<'a> {

    pub fn new(iter: GetTimeTrackIter<'a>, allowed_tags: &'a Vec<TTT>) -> WithOneOf<'a> {
        WithOneOf {
            iter: iter,
            allowed_tags: allowed_tags
        }
    }
}

impl<'a> Iterator for WithOneOf<'a> {
    type Item = Result<FileLockEntry<'a>, TTE>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.iter.next() {
                Some(Ok(fle)) => {
                    match fle.get_timetrack_tag() {
                        Err(e) => return Some(Err(e)),
                        Ok(t) => if self.allowed_tags.contains(&t) {
                            return Some(Ok(fle))
                        } else {
                            // loop
                        },
                    }
                },
                Some(Err(e)) => return Some(Err(e)),
                None => return None,
            }
        }
    }
}

pub struct WithNoneOf<'a> {
    iter: GetTimeTrackIter<'a>,
    disallowed_tags: &'a Vec<TTT>,
}

impl<'a> WithNoneOf<'a> {

    pub fn new(iter: GetTimeTrackIter<'a>, disallowed_tags: &'a Vec<TTT>) -> WithNoneOf<'a> {
        WithNoneOf {
            iter: iter,
            disallowed_tags: disallowed_tags
        }
    }
}

impl<'a> Iterator for WithNoneOf<'a> {
    type Item = Result<FileLockEntry<'a>, TTE>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.iter.next() {
                Some(Ok(fle)) => {
                    match fle.get_timetrack_tag() {
                        Err(e) => return Some(Err(e)),
                        Ok(t) => if !self.disallowed_tags.contains(&t) {
                            return Some(Ok(fle))
                        } else {
                            // loop
                        },
                    }
                },
                Some(Err(e)) => return Some(Err(e)),
                None => return None,
            }
        }
    }
}

<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::path::{PathBuf, Path};
use config::Config;

/// Used to track the state of logical conditions
#[derive(Debug, Clone, Default, Hash, PartialEq, Eq, PartialOrd)]
pub struct LogicState {
    /// Whether the condition has been observed as true
    pub beenTrue: bool,
    /// Whether the condition has been observed as false
    pub beenFalse: bool,
}

/// Shows what type of coverage data is being collected by a given trace
#[derive(Debug, Clone, Default, Hash, PartialEq, Eq, PartialOrd)]
pub enum CoverageStat {
    /// Line coverage data (whether line has been hit)
    Line(u64),
    /// Branch coverage data (whether branch has been true and false
    Branch(LogicState),
    /// Condition coverage data (each boolean subcondition true and false)
    Condition(Vec<LogicState>),
}

#[derive(Debug, Clone, Default, Hash, PartialEq, Eq, PartialOrd)]
pub struct Trace {
    /// Line the trace is on in the file
    pub line: u64,
    /// Optional address showing location in the test artefact
    pub address: Option<u64>,
    /// Length of the instruction (useful to get entire condition/branch)
    pub length: usize,
    /// Coverage stats
    pub stats: CoverageStat,
}

/// Stores all the program traces mapped to files and provides an interface to
/// add, query and change traces.
pub struct TraceMap<'a> {
    /// Traces in the program mapped to the given file
    traces: BTreeMap<PathBuf, Vec<Trace>>,
    /// Tarpaulin config - is this needed?
    config: &'a Config
}

impl<'a> TraceMap<'a> {
    /// Create a new TraceMap
    pub fn new(config: &'a Config) -> TraceMap<'a> {
        TraceMap {
            config: config,
            traces: BTreeMap::new(),
        }
    }

    /// Add a trace to the tracemap for the given file
    pub fn add_trace(&mut self, file: &Path, trace: Trace) {
        if self.traces.contains_key(file) {
            if let Some(trace_vec) = self.traces.get_mut(file) {
                trace_vec.push(trace);
            }
        } else {
            self.traces.insert(file.to_path_buf(), vec![trace]);
        }
    }

    /// Gets an immutable reference to a trace from an address. Returns None if
    /// there is no trace at that address
    pub fn get_trace(&self, address: u64) -> Option<&Trace> {
        self.all_traces()
            .iter()
            .find(|x| x.address == Some(address))
            .map(|x| *x)
    }

    /// Gets a mutable reference to a trace at a given address
    /// Returns None if there is no trace at that address
    pub fn get_trace_mut(&mut self, address: u64) -> Option<&mut Trace> {
        None
    }

    /// Gets all traces below a certain path
    pub fn get_child_traces(&self, root: &Path) -> Vec<&Trace> {
        self.traces.iter()
                   .filter(|&(ref k, _)| k.starts_with(root))
                   .flat_map(|(_, ref v)| v.iter())
                   .collect()
    }

    /// Gets all traces
    fn all_traces(&self) -> Vec<&Trace> {
        self.traces.values().flat_map(|ref x| x.iter()).collect()
    }

}
<commit_msg>Added get_trace_mut<commit_after>use std::collections::BTreeMap;
use std::path::{PathBuf, Path};
use config::Config;

/// Used to track the state of logical conditions
#[derive(Debug, Clone, Default, Hash, PartialEq, Eq, PartialOrd)]
pub struct LogicState {
    /// Whether the condition has been observed as true
    pub beenTrue: bool,
    /// Whether the condition has been observed as false
    pub beenFalse: bool,
}

/// Shows what type of coverage data is being collected by a given trace
#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd)]
pub enum CoverageStat {
    /// Line coverage data (whether line has been hit)
    Line(u64),
    /// Branch coverage data (whether branch has been true and false
    Branch(LogicState),
    /// Condition coverage data (each boolean subcondition true and false)
    Condition(Vec<LogicState>),
}

#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd)]
pub struct Trace {
    /// Line the trace is on in the file
    pub line: u64,
    /// Optional address showing location in the test artefact
    pub address: Option<u64>,
    /// Length of the instruction (useful to get entire condition/branch)
    pub length: usize,
    /// Coverage stats
    pub stats: CoverageStat,
}

/// Stores all the program traces mapped to files and provides an interface to
/// add, query and change traces.
pub struct TraceMap<'a> {
    /// Traces in the program mapped to the given file
    traces: BTreeMap<PathBuf, Vec<Trace>>,
    /// Tarpaulin config - is this needed?
    config: &'a Config
}

impl<'a> TraceMap<'a> {
    /// Create a new TraceMap
    pub fn new(config: &'a Config) -> TraceMap<'a> {
        TraceMap {
            config: config,
            traces: BTreeMap::new(),
        }
    }

    /// Add a trace to the tracemap for the given file
    pub fn add_trace(&mut self, file: &Path, trace: Trace) {
        if self.traces.contains_key(file) {
            if let Some(trace_vec) = self.traces.get_mut(file) {
                trace_vec.push(trace);
            }
        } else {
            self.traces.insert(file.to_path_buf(), vec![trace]);
        }
    }

    /// Gets an immutable reference to a trace from an address. Returns None if
    /// there is no trace at that address
    pub fn get_trace(&self, address: u64) -> Option<&Trace> {
        self.all_traces()
            .iter()
            .find(|x| x.address == Some(address))
            .map(|x| *x)
    }

    /// Gets a mutable reference to a trace at a given address
    /// Returns None if there is no trace at that address
    pub fn get_trace_mut(&mut self, address: u64) -> Option<&mut Trace> {
        for val in self.all_traces_mut() {
            if val.address == Some(address) {
                return Some(val);
            }
        }
        None
    }

    /// Gets all traces below a certain path
    pub fn get_child_traces(&self, root: &Path) -> Vec<&Trace> {
        self.traces.iter()
                   .filter(|&(ref k, _)| k.starts_with(root))
                   .flat_map(|(_, ref v)| v.iter())
                   .collect()
    }

    /// Gets all traces
    fn all_traces(&self) -> Vec<&Trace> {
        self.traces.values().flat_map(|ref x| x.iter()).collect()
    }

    fn all_traces_mut(&mut self) -> Vec<&mut Trace> {
        self.traces.values_mut().flat_map(|x| x.iter_mut()).collect()
    }

}
<|endoftext|>


<commit_before>use std::ptr::null;
use libc::{c_int, c_char, c_double};
use std::ffi::CString;
use std::sync::{StaticMutex, MUTEX_INIT};
use utils::_string;

#[link(name="gdal")]
extern {
    fn OGRRegisterAll();
    fn OGROpen(pszName: *const c_char, bUpdate: c_int, pahDriverList: *const ()) -> *const ();
    fn OGR_DS_GetLayerCount(hDS: *const ()) -> c_int;
    fn OGR_DS_Destroy(hDataSource: *const ());
    fn OGR_DS_GetLayer(hDS: *const (), iLayer: c_int) -> *const ();
    fn OGR_L_GetLayerDefn(hLayer: *const ()) -> *const ();
    fn OGR_L_GetNextFeature(hLayer: *const ()) -> *const ();
    fn OGR_FD_GetFieldCount(hDefn: *const ()) -> c_int;
    fn OGR_FD_GetFieldDefn(hDefn: *const (), iField: c_int) -> *const ();
    fn OGR_F_GetFieldIndex(hFeat: *const (), pszName: *const c_char) -> c_int;
    fn OGR_F_GetFieldDefnRef(hFeat: *const (), i: c_int) -> *const ();
    fn OGR_F_GetFieldAsString(hFeat: *const (), iField: c_int) -> *const c_char;
    fn OGR_F_GetFieldAsDouble(hFeat: *const (), iField: c_int) -> c_double;
    fn OGR_F_GetGeometryRef(hFeat: *const ()) -> *const ();
    fn OGR_F_Destroy(hFeat: *const ());
    fn OGR_G_ExportToWkt(hGeom: *const (), ppszSrcText: &mut *const c_char) -> c_int;
    fn OGR_G_ExportToJson(hGeometry: *const ()) -> *const c_char;
    fn OGR_Fld_GetNameRef(hDefn: *const ()) -> *const c_char;
    fn OGR_Fld_GetType(hDefn: *const ()) -> c_int;
    fn OGRFree(ptr: *const ());
    fn VSIFree(ptr: *const ());
}

const OFT_REAL:             c_int = 2;
const OFT_STRING:           c_int = 4;

static mut LOCK: StaticMutex = MUTEX_INIT;
static mut registered_drivers: bool = false;

fn register_drivers() {
    unsafe {
        let _g = LOCK.lock();
        if ! registered_drivers {
            OGRRegisterAll();
            registered_drivers = true;
        }
    }
}

pub struct VectorDataset {
    c_dataset: *const (),
}

impl VectorDataset {
    pub fn count(&self) -> isize {
        return unsafe { OGR_DS_GetLayerCount(self.c_dataset) } as isize;
    }

    pub fn layer<'a>(&'a self, idx: isize) -> Option<Layer<'a>> {
        let c_layer = unsafe { OGR_DS_GetLayer(self.c_dataset, idx as c_int) };
        return match c_layer.is_null() {
            true  => None,
            false => Some(Layer{_vector_dataset: self, c_layer: c_layer}),
        };
    }
}

impl Drop for VectorDataset {
    fn drop(&mut self) {
        unsafe { OGR_DS_Destroy(self.c_dataset); }
    }
}

pub struct Layer<'a> {
    _vector_dataset: &'a VectorDataset,
    c_layer: *const (),
}

impl<'a> Layer<'a> {
    pub fn fields(&'a self) -> FieldIterator<'a> {
        let c_feature_defn = unsafe { OGR_L_GetLayerDefn(self.c_layer) };
        let total = unsafe { OGR_FD_GetFieldCount(c_feature_defn) } as isize;
        return FieldIterator{
            layer: self,
            c_feature_defn: c_feature_defn,
            next_id: 0,
            total: total
        };
    }

    pub fn features(&'a self) -> FeatureIterator<'a> {
        return FeatureIterator{layer: self};
    }
}

pub struct FieldIterator<'a> {
    layer: &'a Layer<'a>,
    c_feature_defn: *const (),
    next_id: isize,
    total: isize,
}

impl<'a> Iterator for FieldIterator<'a> {
    type Item = Field<'a>;

    #[inline]
    fn next(&mut self) -> Option<Field<'a>> {
        if self.next_id == self.total {
            return None;
        }
        let field = Field{
            _layer: self.layer,
            c_field_defn: unsafe { OGR_FD_GetFieldDefn(
                self.c_feature_defn,
                self.next_id as c_int
            ) }
        };
        self.next_id += 1;
        return Some(field);
    }
}

pub struct Field<'a> {
    _layer: &'a Layer<'a>,
    c_field_defn: *const (),
}

impl<'a> Field<'a> {
    pub fn name(&'a self) -> String {
        let rv = unsafe { OGR_Fld_GetNameRef(self.c_field_defn) };
        return _string(rv);
    }
}

pub struct FeatureIterator<'a> {
    layer: &'a Layer<'a>,
}

impl<'a> Iterator for FeatureIterator<'a> {
    type Item = Feature<'a>;

    #[inline]
    fn next(&mut self) -> Option<Feature<'a>> {
        let c_feature = unsafe { OGR_L_GetNextFeature(self.layer.c_layer) };
        return match c_feature.is_null() {
            true  => None,
            false => Some(Feature{_layer: self.layer, c_feature: c_feature}),
        };
    }
}

pub struct Feature<'a> {
    _layer: &'a Layer<'a>,
    c_feature: *const (),
}

impl<'a> Feature<'a> {
    pub fn field(&self, name: &str) -> Option<FieldValue> {
        let c_name = CString::from_slice(name.as_bytes());
        let field_id = unsafe { OGR_F_GetFieldIndex(self.c_feature, c_name.as_ptr()) };
        if field_id == -1 {
            return None;
        }
        let field_defn = unsafe { OGR_F_GetFieldDefnRef(self.c_feature, field_id) };
        let field_type = unsafe { OGR_Fld_GetType(field_defn) };
        return match field_type {
            OFT_STRING => {
                let rv = unsafe { OGR_F_GetFieldAsString(self.c_feature, field_id) };
                return Some(FieldValue::StringValue(_string(rv)));
            },
            OFT_REAL => {
                let rv = unsafe { OGR_F_GetFieldAsDouble(self.c_feature, field_id) };
                return Some(FieldValue::RealValue(rv as f64));
            },
            _ => panic!("Unknown field type {}", field_type)
        }
    }

    pub fn wkt(&self) -> String {
        unsafe {
            let c_geom = OGR_F_GetGeometryRef(self.c_feature);
            let mut c_wkt: *const c_char = null();
            OGR_G_ExportToWkt(c_geom, &mut c_wkt);
            let wkt = _string(c_wkt);
            OGRFree(c_wkt as *const ());
            return wkt;
        }
    }

    pub fn json(&self) -> String {
        unsafe {
            let c_geom = OGR_F_GetGeometryRef(self.c_feature);
            let c_json = OGR_G_ExportToJson(c_geom);
            let json = _string(c_json);
            VSIFree(c_json as *const ());
            return json;
        }
    }
}

#[unsafe_destructor]
impl<'a> Drop for Feature<'a> {
    fn drop(&mut self) {
        unsafe { OGR_F_Destroy(self.c_feature); }
    }
}

pub fn open(path: &Path) -> Option<VectorDataset> {
    register_drivers();
    let filename = path.as_str().unwrap();
    let c_filename = CString::from_slice(filename.as_bytes());
    let c_dataset = unsafe { OGROpen(c_filename.as_ptr(), 0, null()) };
    return match c_dataset.is_null() {
        true  => None,
        false => Some(VectorDataset{c_dataset: c_dataset}),
    };
}

pub enum FieldValue {
    StringValue(String),
    RealValue(f64),
}

impl FieldValue {
    pub fn as_string(self) -> String {
        match self {
            FieldValue::StringValue(rv) => rv,
            _ => panic!("not a StringValue")
        }
    }

    pub fn as_real(self) -> f64 {
        match self {
            FieldValue::RealValue(rv) => rv,
            _ => panic!("not a RealValue")
        }
    }
}

#[cfg(test)]
mod test {
    use std::path::Path;
    use super::{Feature, FeatureIterator, open};

    fn fixtures() -> Path {
        return Path::new(file!()).dir_path().dir_path().join("fixtures");
    }

    fn assert_almost_eq(a: f64, b: f64) {
        let f: f64 = a / b;
        assert!(f < 1.00001);
        assert!(f > 0.99999);
    }

    #[test]
    fn test_layer_count() {
        let ds = open(&fixtures().join("roads.geojson")).unwrap();
        assert_eq!(ds.count(), 1);
    }

    fn with_features<F>(fixture: &str, f: F) where F: Fn(FeatureIterator) {
        let ds = open(&fixtures().join(fixture)).unwrap();
        let layer = ds.layer(0).unwrap();
        f(layer.features());
    }

    fn with_first_feature<F>(fixture: &str, f: F) where F: Fn(Feature) {
        with_features(fixture, |mut features| f(features.next().unwrap()));
    }

    #[test]
    fn test_iterate_features() {
        with_features("roads.geojson", |features| {
            let feature_vec: Vec<Feature> = features.collect();
            assert_eq!(feature_vec.len(), 21);
        });
    }

    #[test]
    fn test_string_field() {
        with_features("roads.geojson", |mut features| {
            let feature = features.next().unwrap();
            assert_eq!(feature.field("highway")
                              .unwrap()
                              .as_string(),
                       "footway".to_string());
            assert_eq!(
                features.filter(|field| {
                    let highway = field.field("highway")
                                       .unwrap()
                                       .as_string();
                    highway == "residential".to_string() })
                    .count(),
                2);
        });
    }

    #[test]
    fn test_float_field() {
        with_first_feature("roads.geojson", |feature| {
            assert_almost_eq(
                feature.field("sort_key")
                       .unwrap()
                       .as_real(),
                -9.0
            );
        });
    }

    #[test]
    fn test_missing_field() {
        with_first_feature("roads.geojson", |feature| {
            assert!(feature.field("no such field").is_none());
        });
    }

    #[test]
    fn test_wkt() {
        with_first_feature("roads.geojson", |feature| {
            let wkt = feature.wkt();
            let wkt_ok = format!("{}{}",
                "LINESTRING (26.1019276 44.4302748,",
                "26.1019382 44.4303191,26.1020002 44.4304202)"
                ).to_string();
            assert_eq!(wkt, wkt_ok);
        });
    }

    #[test]
    fn test_json() {
        with_first_feature("roads.geojson", |feature| {
            let json = feature.json();
            let json_ok = format!("{}{}{}{}",
                "{ \"type\": \"LineString\", \"coordinates\": [ ",
                "[ 26.1019276, 44.4302748 ], ",
                "[ 26.1019382, 44.4303191 ], ",
                "[ 26.1020002, 44.4304202 ] ] }"
                ).to_string();
            assert_eq!(json, json_ok);
        });
    }

    #[test]
    fn test_schema() {
        let ds = open(&fixtures().join("roads.geojson")).unwrap();
        let layer = ds.layer(0).unwrap();
        let name_list: Vec<String> = layer
            .fields()
            .map(|f| f.name())
            .collect();
        let ok_names: Vec<String> = vec!(
            "kind", "sort_key", "is_link", "is_tunnel",
            "is_bridge", "railway", "highway")
            .iter().map(|s| s.to_string()).collect();
        assert_eq!(name_list, ok_names);
    }
}
<commit_msg>OGRFree and VSIFree expect mutable arguments<commit_after>use std::ptr::null;
use libc::{c_int, c_char, c_double};
use std::ffi::CString;
use std::sync::{StaticMutex, MUTEX_INIT};
use utils::_string;

#[link(name="gdal")]
extern {
    fn OGRRegisterAll();
    fn OGROpen(pszName: *const c_char, bUpdate: c_int, pahDriverList: *const ()) -> *const ();
    fn OGR_DS_GetLayerCount(hDS: *const ()) -> c_int;
    fn OGR_DS_Destroy(hDataSource: *const ());
    fn OGR_DS_GetLayer(hDS: *const (), iLayer: c_int) -> *const ();
    fn OGR_L_GetLayerDefn(hLayer: *const ()) -> *const ();
    fn OGR_L_GetNextFeature(hLayer: *const ()) -> *const ();
    fn OGR_FD_GetFieldCount(hDefn: *const ()) -> c_int;
    fn OGR_FD_GetFieldDefn(hDefn: *const (), iField: c_int) -> *const ();
    fn OGR_F_GetFieldIndex(hFeat: *const (), pszName: *const c_char) -> c_int;
    fn OGR_F_GetFieldDefnRef(hFeat: *const (), i: c_int) -> *const ();
    fn OGR_F_GetFieldAsString(hFeat: *const (), iField: c_int) -> *const c_char;
    fn OGR_F_GetFieldAsDouble(hFeat: *const (), iField: c_int) -> c_double;
    fn OGR_F_GetGeometryRef(hFeat: *const ()) -> *const ();
    fn OGR_F_Destroy(hFeat: *const ());
    fn OGR_G_ExportToWkt(hGeom: *const (), ppszSrcText: &mut *const c_char) -> c_int;
    fn OGR_G_ExportToJson(hGeometry: *const ()) -> *const c_char;
    fn OGR_Fld_GetNameRef(hDefn: *const ()) -> *const c_char;
    fn OGR_Fld_GetType(hDefn: *const ()) -> c_int;
    fn OGRFree(ptr: *mut ());
    fn VSIFree(ptr: *mut ());
}

const OFT_REAL:             c_int = 2;
const OFT_STRING:           c_int = 4;

static mut LOCK: StaticMutex = MUTEX_INIT;
static mut registered_drivers: bool = false;

fn register_drivers() {
    unsafe {
        let _g = LOCK.lock();
        if ! registered_drivers {
            OGRRegisterAll();
            registered_drivers = true;
        }
    }
}

pub struct VectorDataset {
    c_dataset: *const (),
}

impl VectorDataset {
    pub fn count(&self) -> isize {
        return unsafe { OGR_DS_GetLayerCount(self.c_dataset) } as isize;
    }

    pub fn layer<'a>(&'a self, idx: isize) -> Option<Layer<'a>> {
        let c_layer = unsafe { OGR_DS_GetLayer(self.c_dataset, idx as c_int) };
        return match c_layer.is_null() {
            true  => None,
            false => Some(Layer{_vector_dataset: self, c_layer: c_layer}),
        };
    }
}

impl Drop for VectorDataset {
    fn drop(&mut self) {
        unsafe { OGR_DS_Destroy(self.c_dataset); }
    }
}

pub struct Layer<'a> {
    _vector_dataset: &'a VectorDataset,
    c_layer: *const (),
}

impl<'a> Layer<'a> {
    pub fn fields(&'a self) -> FieldIterator<'a> {
        let c_feature_defn = unsafe { OGR_L_GetLayerDefn(self.c_layer) };
        let total = unsafe { OGR_FD_GetFieldCount(c_feature_defn) } as isize;
        return FieldIterator{
            layer: self,
            c_feature_defn: c_feature_defn,
            next_id: 0,
            total: total
        };
    }

    pub fn features(&'a self) -> FeatureIterator<'a> {
        return FeatureIterator{layer: self};
    }
}

pub struct FieldIterator<'a> {
    layer: &'a Layer<'a>,
    c_feature_defn: *const (),
    next_id: isize,
    total: isize,
}

impl<'a> Iterator for FieldIterator<'a> {
    type Item = Field<'a>;

    #[inline]
    fn next(&mut self) -> Option<Field<'a>> {
        if self.next_id == self.total {
            return None;
        }
        let field = Field{
            _layer: self.layer,
            c_field_defn: unsafe { OGR_FD_GetFieldDefn(
                self.c_feature_defn,
                self.next_id as c_int
            ) }
        };
        self.next_id += 1;
        return Some(field);
    }
}

pub struct Field<'a> {
    _layer: &'a Layer<'a>,
    c_field_defn: *const (),
}

impl<'a> Field<'a> {
    pub fn name(&'a self) -> String {
        let rv = unsafe { OGR_Fld_GetNameRef(self.c_field_defn) };
        return _string(rv);
    }
}

pub struct FeatureIterator<'a> {
    layer: &'a Layer<'a>,
}

impl<'a> Iterator for FeatureIterator<'a> {
    type Item = Feature<'a>;

    #[inline]
    fn next(&mut self) -> Option<Feature<'a>> {
        let c_feature = unsafe { OGR_L_GetNextFeature(self.layer.c_layer) };
        return match c_feature.is_null() {
            true  => None,
            false => Some(Feature{_layer: self.layer, c_feature: c_feature}),
        };
    }
}

pub struct Feature<'a> {
    _layer: &'a Layer<'a>,
    c_feature: *const (),
}

impl<'a> Feature<'a> {
    pub fn field(&self, name: &str) -> Option<FieldValue> {
        let c_name = CString::from_slice(name.as_bytes());
        let field_id = unsafe { OGR_F_GetFieldIndex(self.c_feature, c_name.as_ptr()) };
        if field_id == -1 {
            return None;
        }
        let field_defn = unsafe { OGR_F_GetFieldDefnRef(self.c_feature, field_id) };
        let field_type = unsafe { OGR_Fld_GetType(field_defn) };
        return match field_type {
            OFT_STRING => {
                let rv = unsafe { OGR_F_GetFieldAsString(self.c_feature, field_id) };
                return Some(FieldValue::StringValue(_string(rv)));
            },
            OFT_REAL => {
                let rv = unsafe { OGR_F_GetFieldAsDouble(self.c_feature, field_id) };
                return Some(FieldValue::RealValue(rv as f64));
            },
            _ => panic!("Unknown field type {}", field_type)
        }
    }

    pub fn wkt(&self) -> String {
        unsafe {
            let c_geom = OGR_F_GetGeometryRef(self.c_feature);
            let mut c_wkt: *const c_char = null();
            OGR_G_ExportToWkt(c_geom, &mut c_wkt);
            let wkt = _string(c_wkt);
            OGRFree(c_wkt as *mut ());
            return wkt;
        }
    }

    pub fn json(&self) -> String {
        unsafe {
            let c_geom = OGR_F_GetGeometryRef(self.c_feature);
            let c_json = OGR_G_ExportToJson(c_geom);
            let json = _string(c_json);
            VSIFree(c_json as *mut ());
            return json;
        }
    }
}

#[unsafe_destructor]
impl<'a> Drop for Feature<'a> {
    fn drop(&mut self) {
        unsafe { OGR_F_Destroy(self.c_feature); }
    }
}

pub fn open(path: &Path) -> Option<VectorDataset> {
    register_drivers();
    let filename = path.as_str().unwrap();
    let c_filename = CString::from_slice(filename.as_bytes());
    let c_dataset = unsafe { OGROpen(c_filename.as_ptr(), 0, null()) };
    return match c_dataset.is_null() {
        true  => None,
        false => Some(VectorDataset{c_dataset: c_dataset}),
    };
}

pub enum FieldValue {
    StringValue(String),
    RealValue(f64),
}

impl FieldValue {
    pub fn as_string(self) -> String {
        match self {
            FieldValue::StringValue(rv) => rv,
            _ => panic!("not a StringValue")
        }
    }

    pub fn as_real(self) -> f64 {
        match self {
            FieldValue::RealValue(rv) => rv,
            _ => panic!("not a RealValue")
        }
    }
}

#[cfg(test)]
mod test {
    use std::path::Path;
    use super::{Feature, FeatureIterator, open};

    fn fixtures() -> Path {
        return Path::new(file!()).dir_path().dir_path().join("fixtures");
    }

    fn assert_almost_eq(a: f64, b: f64) {
        let f: f64 = a / b;
        assert!(f < 1.00001);
        assert!(f > 0.99999);
    }

    #[test]
    fn test_layer_count() {
        let ds = open(&fixtures().join("roads.geojson")).unwrap();
        assert_eq!(ds.count(), 1);
    }

    fn with_features<F>(fixture: &str, f: F) where F: Fn(FeatureIterator) {
        let ds = open(&fixtures().join(fixture)).unwrap();
        let layer = ds.layer(0).unwrap();
        f(layer.features());
    }

    fn with_first_feature<F>(fixture: &str, f: F) where F: Fn(Feature) {
        with_features(fixture, |mut features| f(features.next().unwrap()));
    }

    #[test]
    fn test_iterate_features() {
        with_features("roads.geojson", |features| {
            let feature_vec: Vec<Feature> = features.collect();
            assert_eq!(feature_vec.len(), 21);
        });
    }

    #[test]
    fn test_string_field() {
        with_features("roads.geojson", |mut features| {
            let feature = features.next().unwrap();
            assert_eq!(feature.field("highway")
                              .unwrap()
                              .as_string(),
                       "footway".to_string());
            assert_eq!(
                features.filter(|field| {
                    let highway = field.field("highway")
                                       .unwrap()
                                       .as_string();
                    highway == "residential".to_string() })
                    .count(),
                2);
        });
    }

    #[test]
    fn test_float_field() {
        with_first_feature("roads.geojson", |feature| {
            assert_almost_eq(
                feature.field("sort_key")
                       .unwrap()
                       .as_real(),
                -9.0
            );
        });
    }

    #[test]
    fn test_missing_field() {
        with_first_feature("roads.geojson", |feature| {
            assert!(feature.field("no such field").is_none());
        });
    }

    #[test]
    fn test_wkt() {
        with_first_feature("roads.geojson", |feature| {
            let wkt = feature.wkt();
            let wkt_ok = format!("{}{}",
                "LINESTRING (26.1019276 44.4302748,",
                "26.1019382 44.4303191,26.1020002 44.4304202)"
                ).to_string();
            assert_eq!(wkt, wkt_ok);
        });
    }

    #[test]
    fn test_json() {
        with_first_feature("roads.geojson", |feature| {
            let json = feature.json();
            let json_ok = format!("{}{}{}{}",
                "{ \"type\": \"LineString\", \"coordinates\": [ ",
                "[ 26.1019276, 44.4302748 ], ",
                "[ 26.1019382, 44.4303191 ], ",
                "[ 26.1020002, 44.4304202 ] ] }"
                ).to_string();
            assert_eq!(json, json_ok);
        });
    }

    #[test]
    fn test_schema() {
        let ds = open(&fixtures().join("roads.geojson")).unwrap();
        let layer = ds.layer(0).unwrap();
        let name_list: Vec<String> = layer
            .fields()
            .map(|f| f.name())
            .collect();
        let ok_names: Vec<String> = vec!(
            "kind", "sort_key", "is_link", "is_tunnel",
            "is_bridge", "railway", "highway")
            .iter().map(|s| s.to_string()).collect();
        assert_eq!(name_list, ok_names);
    }
}
<|endoftext|>