<commit_before>//! Async WebSocket usage.
//!
//! This library is an implementation of WebSocket handshakes and streams. It
//! is based on the crate which implements all required WebSocket protocol
//! logic. So this crate basically just brings tokio support / tokio integration
//! to it.
//!
//! Each WebSocket stream implements the required `Stream` and `Sink` traits,
//! so the socket is just a stream of messages coming in and going out.

#![deny(missing_docs, unused_must_use, unused_mut, unused_imports, unused_import_braces)]

pub use tungstenite;

mod compat;
#[cfg(feature = "connect")]
mod connect;
mod handshake;
#[cfg(feature = "stream")]
mod stream;
#[cfg(any(feature = "native-tls", feature = "__rustls-tls", feature = "connect"))]
mod tls;

use std::io::{Read, Write};

use compat::{cvt, AllowStd, ContextWaker};
use futures_util::{
    sink::{Sink, SinkExt},
    stream::Stream,
};
use log::*;
use std::{
    pin::Pin,
    task::{Context, Poll},
};
use tokio::io::{AsyncRead, AsyncWrite};

use tungstenite::{
    client::IntoClientRequest,
    error::Error as WsError,
    handshake::{
        client::{ClientHandshake, Response},
        server::{Callback, NoCallback},
        HandshakeError,
    },
    protocol::{Message, Role, WebSocket, WebSocketConfig},
};

#[cfg(any(feature = "native-tls", feature = "__rustls-tls", feature = "connect"))]
pub use tls::Connector;
#[cfg(any(feature = "native-tls", feature = "__rustls-tls"))]
pub use tls::{client_async_tls, client_async_tls_with_config};

#[cfg(feature = "connect")]
pub use connect::{connect_async, connect_async_with_config};

#[cfg(all(any(feature = "native-tls", feature = "__rustls-tls"), feature = "connect"))]
pub use connect::connect_async_tls_with_config;

#[cfg(feature = "stream")]
pub use stream::MaybeTlsStream;

use tungstenite::protocol::CloseFrame;

/// Creates a WebSocket handshake from a request and a stream.
/// For convenience, the user may call this with a url string, a URL,
/// or a `Request`. Calling with `Request` allows the user to add
/// a WebSocket protocol or other custom headers.
///
/// Internally, this custom creates a handshake representation and returns
/// a future representing the resolution of the WebSocket handshake. The
/// returned future will resolve to either `WebSocketStream<S>` or `Error`
/// depending on whether the handshake is successful.
///
/// This is typically used for clients who have already established, for
/// example, a TCP connection to the remote server.
pub async fn client_async<'a, R, S>(
    request: R,
    stream: S,
) -> Result<(WebSocketStream<S>, Response), WsError>
where
    R: IntoClientRequest + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    client_async_with_config(request, stream, None).await
}

/// The same as `client_async()` but the one can specify a websocket configuration.
/// Please refer to `client_async()` for more details.
pub async fn client_async_with_config<'a, R, S>(
    request: R,
    stream: S,
    config: Option<WebSocketConfig>,
) -> Result<(WebSocketStream<S>, Response), WsError>
where
    R: IntoClientRequest + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    let f = handshake::client_handshake(stream, move |allow_std| {
        let request = request.into_client_request()?;
        let cli_handshake = ClientHandshake::start(allow_std, request, config)?;
        cli_handshake.handshake()
    });
    f.await.map_err(|e| match e {
        HandshakeError::Failure(e) => e,
        e => WsError::Io(std::io::Error::new(std::io::ErrorKind::Other, e.to_string())),
    })
}

/// Accepts a new WebSocket connection with the provided stream.
///
/// This function will internally call `server::accept` to create a
/// handshake representation and returns a future representing the
/// resolution of the WebSocket handshake. The returned future will resolve
/// to either `WebSocketStream<S>` or `Error` depending if it's successful
/// or not.
///
/// This is typically used after a socket has been accepted from a
/// `TcpListener`. That socket is then passed to this function to perform
/// the server half of the accepting a client's websocket connection.
pub async fn accept_async<S>(stream: S) -> Result<WebSocketStream<S>, WsError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    accept_hdr_async(stream, NoCallback).await
}

/// The same as `accept_async()` but the one can specify a websocket configuration.
/// Please refer to `accept_async()` for more details.
pub async fn accept_async_with_config<S>(
    stream: S,
    config: Option<WebSocketConfig>,
) -> Result<WebSocketStream<S>, WsError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    accept_hdr_async_with_config(stream, NoCallback, config).await
}

/// Accepts a new WebSocket connection with the provided stream.
///
/// This function does the same as `accept_async()` but accepts an extra callback
/// for header processing. The callback receives headers of the incoming
/// requests and is able to add extra headers to the reply.
pub async fn accept_hdr_async<S, C>(stream: S, callback: C) -> Result<WebSocketStream<S>, WsError>
where
    S: AsyncRead + AsyncWrite + Unpin,
    C: Callback + Unpin,
{
    accept_hdr_async_with_config(stream, callback, None).await
}

/// The same as `accept_hdr_async()` but the one can specify a websocket configuration.
/// Please refer to `accept_hdr_async()` for more details.
pub async fn accept_hdr_async_with_config<S, C>(
    stream: S,
    callback: C,
    config: Option<WebSocketConfig>,
) -> Result<WebSocketStream<S>, WsError>
where
    S: AsyncRead + AsyncWrite + Unpin,
    C: Callback + Unpin,
{
    let f = handshake::server_handshake(stream, move |allow_std| {
        tungstenite::accept_hdr_with_config(allow_std, callback, config)
    });
    f.await.map_err(|e| match e {
        HandshakeError::Failure(e) => e,
        e => WsError::Io(std::io::Error::new(std::io::ErrorKind::Other, e.to_string())),
    })
}

/// A wrapper around an underlying raw stream which implements the WebSocket
/// protocol.
///
/// A `WebSocketStream<S>` represents a handshake that has been completed
/// successfully and both the server and the client are ready for receiving
/// and sending data. Message from a `WebSocketStream<S>` are accessible
/// through the respective `Stream` and `Sink`. Check more information about
/// them in `futures-rs` crate documentation or have a look on the examples
/// and unit tests for this crate.
#[derive(Debug)]
pub struct WebSocketStream<S> {
    inner: WebSocket<AllowStd<S>>,
    closing: bool,
    ended: bool,
}

impl<S> WebSocketStream<S> {
    /// Convert a raw socket into a WebSocketStream without performing a
    /// handshake.
    pub async fn from_raw_socket(stream: S, role: Role, config: Option<WebSocketConfig>) -> Self
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        handshake::without_handshake(stream, move |allow_std| {
            WebSocket::from_raw_socket(allow_std, role, config)
        })
        .await
    }

    /// Convert a raw socket into a WebSocketStream without performing a
    /// handshake.
    pub async fn from_partially_read(
        stream: S,
        part: Vec<u8>,
        role: Role,
        config: Option<WebSocketConfig>,
    ) -> Self
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        handshake::without_handshake(stream, move |allow_std| {
            WebSocket::from_partially_read(allow_std, part, role, config)
        })
        .await
    }

    pub(crate) fn new(ws: WebSocket<AllowStd<S>>) -> Self {
        WebSocketStream { inner: ws, closing: false, ended: false }
    }

    fn with_context<F, R>(&mut self, ctx: Option<(ContextWaker, &mut Context<'_>)>, f: F) -> R
    where
        S: Unpin,
        F: FnOnce(&mut WebSocket<AllowStd<S>>) -> R,
        AllowStd<S>: Read + Write,
    {
        trace!("{}:{} WebSocketStream.with_context", file!(), line!());
        if let Some((kind, ctx)) = ctx {
            self.inner.get_mut().set_waker(kind, ctx.waker());
        }
        f(&mut self.inner)
    }

    /// Returns a shared reference to the inner stream.
    pub fn get_ref(&self) -> &S
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        self.inner.get_ref().get_ref()
    }

    /// Returns a mutable reference to the inner stream.
    pub fn get_mut(&mut self) -> &mut S
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        self.inner.get_mut().get_mut()
    }

    /// Returns a reference to the configuration of the tungstenite stream.
    pub fn get_config(&self) -> &WebSocketConfig {
        self.inner.get_config()
    }

    /// Close the underlying web socket
    pub async fn close(&mut self, msg: Option<CloseFrame<'_>>) -> Result<(), WsError>
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        let msg = msg.map(|msg| msg.into_owned());
        self.send(Message::Close(msg)).await
    }
}

impl<T> Stream for WebSocketStream<T>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    type Item = Result<Message, WsError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        trace!("{}:{} Stream.poll_next", file!(), line!());

        // The connection has been closed or a critical error has occurred.
        // We have already returned the error to the user, the `Stream` is unusable,
        // so we assume that the stream has been "fused".
        if self.ended {
            return Poll::Ready(None);
        }

        match futures_util::ready!(self.with_context(Some((ContextWaker::Read, cx)), |s| {
            trace!("{}:{} Stream.with_context poll_next -> read_message()", file!(), line!());
            cvt(s.read_message())
        })) {
            Ok(v) => Poll::Ready(Some(Ok(v))),
            Err(e) => {
                self.ended = true;
                if matches!(e, WsError::AlreadyClosed | WsError::ConnectionClosed) {
                    Poll::Ready(None)
                } else {
                    Poll::Ready(Some(Err(e)))
                }
            }
        }
    }
}

impl<T> Sink<Message> for WebSocketStream<T>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    type Error = WsError;

    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        (*self).with_context(Some((ContextWaker::Write, cx)), |s| cvt(s.write_pending()))
    }

    fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
        match (*self).with_context(None, |s| s.write_message(item)) {
            Ok(()) => Ok(()),
            Err(::tungstenite::Error::Io(err)) if err.kind() == std::io::ErrorKind::WouldBlock => {
                // the message was accepted and queued
                // isn't an error.
                Ok(())
            }
            Err(e) => {
                debug!("websocket start_send error: {}", e);
                Err(e)
            }
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        (*self).with_context(Some((ContextWaker::Write, cx)), |s| cvt(s.write_pending()))
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let res = if self.closing {
            // After queueing it, we call `write_pending` to drive the close handshake to completion.
            (*self).with_context(Some((ContextWaker::Write, cx)), |s| s.write_pending())
        } else {
            (*self).with_context(Some((ContextWaker::Write, cx)), |s| s.close(None))
        };

        match res {
            Ok(()) => Poll::Ready(Ok(())),
            Err(::tungstenite::Error::ConnectionClosed) => Poll::Ready(Ok(())),
            Err(::tungstenite::Error::Io(err)) if err.kind() == std::io::ErrorKind::WouldBlock => {
                trace!("WouldBlock");
                self.closing = true;
                Poll::Pending
            }
            Err(err) => {
                debug!("websocket close error: {}", err);
                Poll::Ready(Err(err))
            }
        }
    }
}

/// Get a domain from an URL.
#[cfg(any(feature = "connect", feature = "native-tls", feature = "__rustls-tls"))]
#[inline]
fn domain(request: &tungstenite::handshake::client::Request) -> Result<String, WsError> {
    match request.uri().host() {
        Some(d) => Ok(d.to_string()),
        None => Err(WsError::Url(tungstenite::error::UrlError::NoHostName)),
    }
}

#[cfg(test)]
mod tests {
    #[cfg(feature = "connect")]
    use crate::stream::MaybeTlsStream;
    use crate::{compat::AllowStd, WebSocketStream};
    use std::io::{Read, Write};
    #[cfg(feature = "connect")]
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    fn is_read<T: Read>() {}
    fn is_write<T: Write>() {}
    #[cfg(feature = "connect")]
    fn is_async_read<T: AsyncReadExt>() {}
    #[cfg(feature = "connect")]
    fn is_async_write<T: AsyncWriteExt>() {}
    fn is_unpin<T: Unpin>() {}

    #[test]
    fn web_socket_stream_has_traits() {
        is_read::<AllowStd<tokio::net::TcpStream>>();
        is_write::<AllowStd<tokio::net::TcpStream>>();

        #[cfg(feature = "connect")]
        is_async_read::<MaybeTlsStream<tokio::net::TcpStream>>();
        #[cfg(feature = "connect")]
        is_async_write::<MaybeTlsStream<tokio::net::TcpStream>>();

        is_unpin::<WebSocketStream<tokio::net::TcpStream>>();
        #[cfg(feature = "connect")]
        is_unpin::<WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>>();
    }
}
<commit_msg>Implement the `FusedStream` along with `Stream`<commit_after>//! Async WebSocket usage.
//!
//! This library is an implementation of WebSocket handshakes and streams. It
//! is based on the crate which implements all required WebSocket protocol
//! logic. So this crate basically just brings tokio support / tokio integration
//! to it.
//!
//! Each WebSocket stream implements the required `Stream` and `Sink` traits,
//! so the socket is just a stream of messages coming in and going out.

#![deny(missing_docs, unused_must_use, unused_mut, unused_imports, unused_import_braces)]

pub use tungstenite;

mod compat;
#[cfg(feature = "connect")]
mod connect;
mod handshake;
#[cfg(feature = "stream")]
mod stream;
#[cfg(any(feature = "native-tls", feature = "__rustls-tls", feature = "connect"))]
mod tls;

use std::io::{Read, Write};

use compat::{cvt, AllowStd, ContextWaker};
use futures_util::{
    sink::{Sink, SinkExt},
    stream::{FusedStream, Stream},
};
use log::*;
use std::{
    pin::Pin,
    task::{Context, Poll},
};
use tokio::io::{AsyncRead, AsyncWrite};

use tungstenite::{
    client::IntoClientRequest,
    error::Error as WsError,
    handshake::{
        client::{ClientHandshake, Response},
        server::{Callback, NoCallback},
        HandshakeError,
    },
    protocol::{Message, Role, WebSocket, WebSocketConfig},
};

#[cfg(any(feature = "native-tls", feature = "__rustls-tls", feature = "connect"))]
pub use tls::Connector;
#[cfg(any(feature = "native-tls", feature = "__rustls-tls"))]
pub use tls::{client_async_tls, client_async_tls_with_config};

#[cfg(feature = "connect")]
pub use connect::{connect_async, connect_async_with_config};

#[cfg(all(any(feature = "native-tls", feature = "__rustls-tls"), feature = "connect"))]
pub use connect::connect_async_tls_with_config;

#[cfg(feature = "stream")]
pub use stream::MaybeTlsStream;

use tungstenite::protocol::CloseFrame;

/// Creates a WebSocket handshake from a request and a stream.
/// For convenience, the user may call this with a url string, a URL,
/// or a `Request`. Calling with `Request` allows the user to add
/// a WebSocket protocol or other custom headers.
///
/// Internally, this custom creates a handshake representation and returns
/// a future representing the resolution of the WebSocket handshake. The
/// returned future will resolve to either `WebSocketStream<S>` or `Error`
/// depending on whether the handshake is successful.
///
/// This is typically used for clients who have already established, for
/// example, a TCP connection to the remote server.
pub async fn client_async<'a, R, S>(
    request: R,
    stream: S,
) -> Result<(WebSocketStream<S>, Response), WsError>
where
    R: IntoClientRequest + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    client_async_with_config(request, stream, None).await
}

/// The same as `client_async()` but the one can specify a websocket configuration.
/// Please refer to `client_async()` for more details.
pub async fn client_async_with_config<'a, R, S>(
    request: R,
    stream: S,
    config: Option<WebSocketConfig>,
) -> Result<(WebSocketStream<S>, Response), WsError>
where
    R: IntoClientRequest + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    let f = handshake::client_handshake(stream, move |allow_std| {
        let request = request.into_client_request()?;
        let cli_handshake = ClientHandshake::start(allow_std, request, config)?;
        cli_handshake.handshake()
    });
    f.await.map_err(|e| match e {
        HandshakeError::Failure(e) => e,
        e => WsError::Io(std::io::Error::new(std::io::ErrorKind::Other, e.to_string())),
    })
}

/// Accepts a new WebSocket connection with the provided stream.
///
/// This function will internally call `server::accept` to create a
/// handshake representation and returns a future representing the
/// resolution of the WebSocket handshake. The returned future will resolve
/// to either `WebSocketStream<S>` or `Error` depending if it's successful
/// or not.
///
/// This is typically used after a socket has been accepted from a
/// `TcpListener`. That socket is then passed to this function to perform
/// the server half of the accepting a client's websocket connection.
pub async fn accept_async<S>(stream: S) -> Result<WebSocketStream<S>, WsError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    accept_hdr_async(stream, NoCallback).await
}

/// The same as `accept_async()` but the one can specify a websocket configuration.
/// Please refer to `accept_async()` for more details.
pub async fn accept_async_with_config<S>(
    stream: S,
    config: Option<WebSocketConfig>,
) -> Result<WebSocketStream<S>, WsError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    accept_hdr_async_with_config(stream, NoCallback, config).await
}

/// Accepts a new WebSocket connection with the provided stream.
///
/// This function does the same as `accept_async()` but accepts an extra callback
/// for header processing. The callback receives headers of the incoming
/// requests and is able to add extra headers to the reply.
pub async fn accept_hdr_async<S, C>(stream: S, callback: C) -> Result<WebSocketStream<S>, WsError>
where
    S: AsyncRead + AsyncWrite + Unpin,
    C: Callback + Unpin,
{
    accept_hdr_async_with_config(stream, callback, None).await
}

/// The same as `accept_hdr_async()` but the one can specify a websocket configuration.
/// Please refer to `accept_hdr_async()` for more details.
pub async fn accept_hdr_async_with_config<S, C>(
    stream: S,
    callback: C,
    config: Option<WebSocketConfig>,
) -> Result<WebSocketStream<S>, WsError>
where
    S: AsyncRead + AsyncWrite + Unpin,
    C: Callback + Unpin,
{
    let f = handshake::server_handshake(stream, move |allow_std| {
        tungstenite::accept_hdr_with_config(allow_std, callback, config)
    });
    f.await.map_err(|e| match e {
        HandshakeError::Failure(e) => e,
        e => WsError::Io(std::io::Error::new(std::io::ErrorKind::Other, e.to_string())),
    })
}

/// A wrapper around an underlying raw stream which implements the WebSocket
/// protocol.
///
/// A `WebSocketStream<S>` represents a handshake that has been completed
/// successfully and both the server and the client are ready for receiving
/// and sending data. Message from a `WebSocketStream<S>` are accessible
/// through the respective `Stream` and `Sink`. Check more information about
/// them in `futures-rs` crate documentation or have a look on the examples
/// and unit tests for this crate.
#[derive(Debug)]
pub struct WebSocketStream<S> {
    inner: WebSocket<AllowStd<S>>,
    closing: bool,
    ended: bool,
}

impl<S> WebSocketStream<S> {
    /// Convert a raw socket into a WebSocketStream without performing a
    /// handshake.
    pub async fn from_raw_socket(stream: S, role: Role, config: Option<WebSocketConfig>) -> Self
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        handshake::without_handshake(stream, move |allow_std| {
            WebSocket::from_raw_socket(allow_std, role, config)
        })
        .await
    }

    /// Convert a raw socket into a WebSocketStream without performing a
    /// handshake.
    pub async fn from_partially_read(
        stream: S,
        part: Vec<u8>,
        role: Role,
        config: Option<WebSocketConfig>,
    ) -> Self
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        handshake::without_handshake(stream, move |allow_std| {
            WebSocket::from_partially_read(allow_std, part, role, config)
        })
        .await
    }

    pub(crate) fn new(ws: WebSocket<AllowStd<S>>) -> Self {
        WebSocketStream { inner: ws, closing: false, ended: false }
    }

    fn with_context<F, R>(&mut self, ctx: Option<(ContextWaker, &mut Context<'_>)>, f: F) -> R
    where
        S: Unpin,
        F: FnOnce(&mut WebSocket<AllowStd<S>>) -> R,
        AllowStd<S>: Read + Write,
    {
        trace!("{}:{} WebSocketStream.with_context", file!(), line!());
        if let Some((kind, ctx)) = ctx {
            self.inner.get_mut().set_waker(kind, ctx.waker());
        }
        f(&mut self.inner)
    }

    /// Returns a shared reference to the inner stream.
    pub fn get_ref(&self) -> &S
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        self.inner.get_ref().get_ref()
    }

    /// Returns a mutable reference to the inner stream.
    pub fn get_mut(&mut self) -> &mut S
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        self.inner.get_mut().get_mut()
    }

    /// Returns a reference to the configuration of the tungstenite stream.
    pub fn get_config(&self) -> &WebSocketConfig {
        self.inner.get_config()
    }

    /// Close the underlying web socket
    pub async fn close(&mut self, msg: Option<CloseFrame<'_>>) -> Result<(), WsError>
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        let msg = msg.map(|msg| msg.into_owned());
        self.send(Message::Close(msg)).await
    }
}

impl<T> Stream for WebSocketStream<T>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    type Item = Result<Message, WsError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        trace!("{}:{} Stream.poll_next", file!(), line!());

        // The connection has been closed or a critical error has occurred.
        // We have already returned the error to the user, the `Stream` is unusable,
        // so we assume that the stream has been "fused".
        if self.ended {
            return Poll::Ready(None);
        }

        match futures_util::ready!(self.with_context(Some((ContextWaker::Read, cx)), |s| {
            trace!("{}:{} Stream.with_context poll_next -> read_message()", file!(), line!());
            cvt(s.read_message())
        })) {
            Ok(v) => Poll::Ready(Some(Ok(v))),
            Err(e) => {
                self.ended = true;
                if matches!(e, WsError::AlreadyClosed | WsError::ConnectionClosed) {
                    Poll::Ready(None)
                } else {
                    Poll::Ready(Some(Err(e)))
                }
            }
        }
    }
}

impl<T> FusedStream for WebSocketStream<T>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    fn is_terminated(&self) -> bool {
        self.ended
    }
}

impl<T> Sink<Message> for WebSocketStream<T>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    type Error = WsError;

    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        (*self).with_context(Some((ContextWaker::Write, cx)), |s| cvt(s.write_pending()))
    }

    fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
        match (*self).with_context(None, |s| s.write_message(item)) {
            Ok(()) => Ok(()),
            Err(::tungstenite::Error::Io(err)) if err.kind() == std::io::ErrorKind::WouldBlock => {
                // the message was accepted and queued
                // isn't an error.
                Ok(())
            }
            Err(e) => {
                debug!("websocket start_send error: {}", e);
                Err(e)
            }
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        (*self).with_context(Some((ContextWaker::Write, cx)), |s| cvt(s.write_pending()))
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let res = if self.closing {
            // After queueing it, we call `write_pending` to drive the close handshake to completion.
            (*self).with_context(Some((ContextWaker::Write, cx)), |s| s.write_pending())
        } else {
            (*self).with_context(Some((ContextWaker::Write, cx)), |s| s.close(None))
        };

        match res {
            Ok(()) => Poll::Ready(Ok(())),
            Err(::tungstenite::Error::ConnectionClosed) => Poll::Ready(Ok(())),
            Err(::tungstenite::Error::Io(err)) if err.kind() == std::io::ErrorKind::WouldBlock => {
                trace!("WouldBlock");
                self.closing = true;
                Poll::Pending
            }
            Err(err) => {
                debug!("websocket close error: {}", err);
                Poll::Ready(Err(err))
            }
        }
    }
}

/// Get a domain from an URL.
#[cfg(any(feature = "connect", feature = "native-tls", feature = "__rustls-tls"))]
#[inline]
fn domain(request: &tungstenite::handshake::client::Request) -> Result<String, WsError> {
    match request.uri().host() {
        Some(d) => Ok(d.to_string()),
        None => Err(WsError::Url(tungstenite::error::UrlError::NoHostName)),
    }
}

#[cfg(test)]
mod tests {
    #[cfg(feature = "connect")]
    use crate::stream::MaybeTlsStream;
    use crate::{compat::AllowStd, WebSocketStream};
    use std::io::{Read, Write};
    #[cfg(feature = "connect")]
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    fn is_read<T: Read>() {}
    fn is_write<T: Write>() {}
    #[cfg(feature = "connect")]
    fn is_async_read<T: AsyncReadExt>() {}
    #[cfg(feature = "connect")]
    fn is_async_write<T: AsyncWriteExt>() {}
    fn is_unpin<T: Unpin>() {}

    #[test]
    fn web_socket_stream_has_traits() {
        is_read::<AllowStd<tokio::net::TcpStream>>();
        is_write::<AllowStd<tokio::net::TcpStream>>();

        #[cfg(feature = "connect")]
        is_async_read::<MaybeTlsStream<tokio::net::TcpStream>>();
        #[cfg(feature = "connect")]
        is_async_write::<MaybeTlsStream<tokio::net::TcpStream>>();

        is_unpin::<WebSocketStream<tokio::net::TcpStream>>();
        #[cfg(feature = "connect")]
        is_unpin::<WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>>();
    }
}
<|endoftext|>


<commit_before>//! # mdBook
//!
//! **mdBook** is similar to Gitbook but implemented in Rust.
//! It offers a command line interface, but can also be used as a regular crate.
//!
//! This is the API doc, but you can find a [less "low-level" documentation here](../index.html)
//! that contains information about the command line tool, format, structure etc.
//! It is also rendered with mdBook to showcase the features and default theme.
//!
//! Some reasons why you would want to use the crate (over the cli):
//!
//! - Integrate mdbook in a current project
//! - Extend the capabilities of mdBook
//! - Do some processing or test before building your book
//! - Write a new Renderer
//! - ...
//!
//! ## Example
//!
//! ```no_run
//! extern crate mdbook;
//!
//! use mdbook::MDBook;
//!
//! # #[allow(unused_variables)]
//! fn main() {
//!     let mut book =  MDBook::new("my-book")        // Path to root
//!                         .with_source("src")       // Path from root to source directory
//!                         .with_destination("book") // Path from root to output directory
//!                         .read_config()            // Parse book.toml configuration file
//!                         .expect("I don't handle configuration file errors, but you should!");
//!
//!     book.build().unwrap();                        // Render the book
//! }
//! ```
//!
//! ## Implementing a new Renderer
//!
//! If you want to create a new renderer for mdBook, the only thing you have to do is to implement
//! the [Renderer trait](renderer/renderer/trait.Renderer.html)
//!
//! And then you can swap in your renderer like this:
//!
//! ```no_run
//! # extern crate mdbook;
//! #
//! # use mdbook::MDBook;
//! # use mdbook::renderer::HtmlHandlebars;
//! #
//! # #[allow(unused_variables)]
//! # fn main() {
//! #   let your_renderer = HtmlHandlebars::new();
//! #
//! let book = MDBook::new("my-book").set_renderer(Box::new(your_renderer));
//! # }
//! ```
//! If you make a renderer, you get the book constructed in form of `Vec<BookItems>` and you get
//! the book config in a `BookConfig` struct.
//!
//! It's your responsability to create the necessary files in the correct directories.
//!
//! ## utils
//!
//! I have regrouped some useful functions in the [utils](utils/index.html) module, like the
//! following function [`utils::fs::create_file(path:
//! &Path)`](utils/fs/fn.create_file.html)
//!
//! This function creates a file and returns it. But before creating the file
//! it checks every directory in the path to see if it exists, and if it does
//! not it will be created.
//!
//! Make sure to take a look at it.

#[macro_use]
extern crate error_chain;
extern crate handlebars;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
extern crate pulldown_cmark;
extern crate regex;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate serde_json;
extern crate tempdir;
extern crate toml;

mod parse;
mod preprocess;
pub mod book;
pub mod config;
pub mod renderer;
pub mod theme;
pub mod utils;

pub use book::MDBook;
pub use book::BookItem;
pub use renderer::Renderer;

/// The error types used through out this crate.
pub mod errors {
    error_chain!{
        foreign_links {
            Io(::std::io::Error);
            HandlebarsRender(::handlebars::RenderError);
            HandlebarsTemplate(Box<::handlebars::TemplateError>);
            Utf8(::std::string::FromUtf8Error);
        }

        errors {
            Subprocess(message: String, output: ::std::process::Output) {
                description("A subprocess failed")
                display("{}: {}", message, String::from_utf8_lossy(&output.stdout))
            }
        }
    }

    // Box to halve the size of Error
    impl From<::handlebars::TemplateError> for Error {
        fn from(e: ::handlebars::TemplateError) -> Error {
            From::from(Box::new(e))
        }
    }
}
<commit_msg>All tests pass again :)<commit_after>//! # mdBook
//!
//! **mdBook** is similar to Gitbook but implemented in Rust.
//! It offers a command line interface, but can also be used as a regular crate.
//!
//! This is the API doc, but you can find a [less "low-level" documentation here](../index.html)
//! that contains information about the command line tool, format, structure etc.
//! It is also rendered with mdBook to showcase the features and default theme.
//!
//! Some reasons why you would want to use the crate (over the cli):
//!
//! - Integrate mdbook in a current project
//! - Extend the capabilities of mdBook
//! - Do some processing or test before building your book
//! - Write a new Renderer
//! - ...
//!
//! ## Example
//!
//! ```no_run
//! extern crate mdbook;
//!
//! use mdbook::MDBook;
//! use std::path::PathBuf;
//!
//! fn main() {
//!     let mut md = MDBook::new("my-book");
//!
//!     // tweak the book configuration a bit
//!     md.config.book.src = PathBuf::from("source");
//!     md.config.book.build_dir = PathBuf::from("book");
//!
//!     // Render the book
//!     md.build().unwrap();
//! }
//! ```
//!
//! ## Implementing a new Renderer
//!
//! If you want to create a new renderer for mdBook, the only thing you have to do is to implement
//! the [Renderer trait](renderer/renderer/trait.Renderer.html)
//!
//! And then you can swap in your renderer like this:
//!
//! ```no_run
//! # extern crate mdbook;
//! #
//! # use mdbook::MDBook;
//! # use mdbook::renderer::HtmlHandlebars;
//! #
//! # #[allow(unused_variables)]
//! # fn main() {
//! #   let your_renderer = HtmlHandlebars::new();
//! #
//! let book = MDBook::new("my-book").set_renderer(Box::new(your_renderer));
//! # }
//! ```
//! If you make a renderer, you get the book constructed in form of `Vec<BookItems>` and you get
//! the book config in a `BookConfig` struct.
//!
//! It's your responsability to create the necessary files in the correct directories.
//!
//! ## utils
//!
//! I have regrouped some useful functions in the [utils](utils/index.html) module, like the
//! following function [`utils::fs::create_file(path:
//! &Path)`](utils/fs/fn.create_file.html)
//!
//! This function creates a file and returns it. But before creating the file
//! it checks every directory in the path to see if it exists, and if it does
//! not it will be created.
//!
//! Make sure to take a look at it.

#[macro_use]
extern crate error_chain;
extern crate handlebars;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
extern crate pulldown_cmark;
extern crate regex;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate serde_json;
extern crate tempdir;
extern crate toml;

mod parse;
mod preprocess;
pub mod book;
pub mod config;
pub mod renderer;
pub mod theme;
pub mod utils;

pub use book::MDBook;
pub use book::BookItem;
pub use renderer::Renderer;

/// The error types used through out this crate.
pub mod errors {
    error_chain!{
        foreign_links {
            Io(::std::io::Error);
            HandlebarsRender(::handlebars::RenderError);
            HandlebarsTemplate(Box<::handlebars::TemplateError>);
            Utf8(::std::string::FromUtf8Error);
        }

        errors {
            Subprocess(message: String, output: ::std::process::Output) {
                description("A subprocess failed")
                display("{}: {}", message, String::from_utf8_lossy(&output.stdout))
            }
        }
    }

    // Box to halve the size of Error
    impl From<::handlebars::TemplateError> for Error {
        fn from(e: ::handlebars::TemplateError) -> Error {
            From::from(Box::new(e))
        }
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(not(feature = "std"), no_std)]
#![forbid(unsafe_code)]

//! This crate provides you with the ability to generate and parse ASN.1
//! encoded data. More precisely, it provides you with the ability to generate
//! and parse data encoded with ASN.1's DER (Distinguished Encoding Rules)
//! encoding. It does not support BER (Basic Encoding Rules), CER (Canonical
//! Encoding Rules), XER (XML Encoding Rules), CXER (Canonical XML Encoding
//! Rules), or any other alphabet soup encodings -- and it never will.
//!
//! If you wanted to parse an ASN.1 structure like this:
//! ```text
//! Signature ::= SEQUENCE {
//!     r INTEGER,
//!     s INTEGER
//! }
//! ```
//!
//! Then you'd write the following code:
//! ```
//! # let data = b"";
//! let result: asn1::ParseResult<_> = asn1::parse(data, |d| {
//!     return d.read_element::<asn1::Sequence>()?.parse(|d| {
//!         let r = d.read_element::<u64>()?;
//!         let s = d.read_element::<u64>()?;
//!         return Ok((r, s));
//!     })
//! });
//! ```
//!
//! In general everything about parsing is driven by providing different type
//! parameters to `Parser.read_element`. Some types directly implement the
//! `Asn1Element` trait, as seen with `u64` or `&[u8]` (`OCTET STRING`), while
//! others use placeholder types which differ from the return type
//! (`PrintableString` or `UtcTime`). There are also types such as `Implicit`
//! and `Explicit` for handling tagged values, `Choice1`, `Choice2`, and
//! `Choice3` available for choices, and `Option<T>` for handling `OPTIONAL`
//!  values.
//!
//! To serialize DER for the `Sequence` structure, you'd write the following:
//! ```
//! # let r = 0u64;
//! # let s = 0u64;
//! let result = asn1::write(|w| {
//!     w.write_element(&asn1::SequenceWriter::new(&|w| {
//!         w.write_element(&r);
//!         w.write_element(&s);
//!     }));
//! });
//! ```
//!
//! # Derive
//!
//! When built with the `derive` feature (enabled by default), these can also
//! be expressed as Rust structs:
//! ```text
//! #[derive(asn1::Asn1Read, asn1::Asn1Write)]
//! struct Signature {
//!     r: u64,
//!     s: u64,
//! }
//!
//! # let data = b"";
//! # let r = 0u64;
//! # let s = 0u64;
//! let sig = asn1::parse_single::<Signature>(data);
//! let result = asn1::write_single(&Signature{r, s});
//! ```
//!
//! On Rust >= 1.51.0, [`Explicit`] and [`Implicit`] tagging may be specified
//! with struct members of those types. However on Rust < 1.51.0, this is not
//! possible, since they require const generics. Instead, the `#[implicit]`
//! and `#[explicit]` attributes may be used:
//! ```text
//! #[derive(asn1::Asn1Read, asn1::Asn1Write)]
//! struct SomeSequence<'a> {
//!     #[implicit(0)]
//!     a: Option<&'a [u8]>,
//!     #[explicit(1)]
//!     b: Option<u64>,
//! }
//! ```

extern crate alloc;

mod bit_string;
mod object_identitifer;
mod parser;
mod types;
mod writer;

pub use crate::bit_string::BitString;
pub use crate::object_identitifer::ObjectIdentifier;
pub use crate::parser::{parse, parse_single, ParseError, ParseResult, Parser};
pub use crate::types::{
    Asn1Readable, Asn1Writable, BigUint, Choice1, Choice2, Choice3, PrintableString, Sequence,
    SequenceOf, SequenceOfWriter, SequenceWriter, SetOf, SetOfWriter, SimpleAsn1Readable,
    SimpleAsn1Writable, Tlv, UtcTime,
};
#[cfg(feature = "const-generics")]
pub use crate::types::{Explicit, Implicit};
pub use crate::writer::{write, write_single, Writer};

#[cfg(feature = "derive")]
pub use asn1_derive::{Asn1Read, Asn1Write};
<commit_msg>updated docs to match reality (#128)<commit_after>#![cfg_attr(not(feature = "std"), no_std)]
#![forbid(unsafe_code)]

//! This crate provides you with the ability to generate and parse ASN.1
//! encoded data. More precisely, it provides you with the ability to generate
//! and parse data encoded with ASN.1's DER (Distinguished Encoding Rules)
//! encoding. It does not support BER (Basic Encoding Rules), CER (Canonical
//! Encoding Rules), XER (XML Encoding Rules), CXER (Canonical XML Encoding
//! Rules), or any other alphabet soup encodings -- and it never will.
//!
//! If you wanted to parse an ASN.1 structure like this:
//! ```text
//! Signature ::= SEQUENCE {
//!     r INTEGER,
//!     s INTEGER
//! }
//! ```
//!
//! Then you'd write the following code:
//! ```
//! # let data = b"";
//! let result: asn1::ParseResult<_> = asn1::parse(data, |d| {
//!     return d.read_element::<asn1::Sequence>()?.parse(|d| {
//!         let r = d.read_element::<u64>()?;
//!         let s = d.read_element::<u64>()?;
//!         return Ok((r, s));
//!     })
//! });
//! ```
//!
//! In general everything about parsing is driven by providing different type
//! parameters to `Parser.read_element`. Some types implement the
//! `Asn1Readable` trait directly on a basic type, as seen with `u64` or
//! `&[u8]` (`OCTET STRING`), while others use wrapper types which simply
//! provide ASN.1 encoding and decoding for some other type (`PrintableString`
//! or `UtcTime`). There are also types such as `Implicit` and `Explicit` for
//! handling tagged values, `Choice1`, `Choice2`, and `Choice3` available for
//! choices, and `Option<T>` for handling `OPTIONAL` values.
//!
//! To serialize DER for the `Sequence` structure, you'd write the following:
//! ```
//! # let r = 0u64;
//! # let s = 0u64;
//! let result = asn1::write(|w| {
//!     w.write_element(&asn1::SequenceWriter::new(&|w| {
//!         w.write_element(&r);
//!         w.write_element(&s);
//!     }));
//! });
//! ```
//!
//! # Derive
//!
//! When built with the `derive` feature (enabled by default), these can also
//! be expressed as Rust structs:
//! ```text
//! #[derive(asn1::Asn1Read, asn1::Asn1Write)]
//! struct Signature {
//!     r: u64,
//!     s: u64,
//! }
//!
//! # let data = b"";
//! # let r = 0u64;
//! # let s = 0u64;
//! let sig = asn1::parse_single::<Signature>(data);
//! let result = asn1::write_single(&Signature{r, s});
//! ```
//!
//! On Rust >= 1.51.0, [`Explicit`] and [`Implicit`] tagging may be specified
//! with struct members of those types. However on Rust < 1.51.0, this is not
//! possible, since they require const generics. Instead, the `#[implicit]`
//! and `#[explicit]` attributes may be used:
//! ```text
//! #[derive(asn1::Asn1Read, asn1::Asn1Write)]
//! struct SomeSequence<'a> {
//!     #[implicit(0)]
//!     a: Option<&'a [u8]>,
//!     #[explicit(1)]
//!     b: Option<u64>,
//! }
//! ```

extern crate alloc;

mod bit_string;
mod object_identitifer;
mod parser;
mod types;
mod writer;

pub use crate::bit_string::BitString;
pub use crate::object_identitifer::ObjectIdentifier;
pub use crate::parser::{parse, parse_single, ParseError, ParseResult, Parser};
pub use crate::types::{
    Asn1Readable, Asn1Writable, BigUint, Choice1, Choice2, Choice3, PrintableString, Sequence,
    SequenceOf, SequenceOfWriter, SequenceWriter, SetOf, SetOfWriter, SimpleAsn1Readable,
    SimpleAsn1Writable, Tlv, UtcTime,
};
#[cfg(feature = "const-generics")]
pub use crate::types::{Explicit, Implicit};
pub use crate::writer::{write, write_single, Writer};

#[cfg(feature = "derive")]
pub use asn1_derive::{Asn1Read, Asn1Write};
<|endoftext|>


<commit_before>// Module for parsing ISO Base Media Format aka video/mp4 streams.

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

/// Basic ISO box structure.
pub struct BoxHeader {
    /// Four character box type
pub name: u32,
    /// Size of the box in bytes
pub size: u64,
    /// Offset to the start of the contained data (or header size).
pub offset: u64,
}

/// File type box 'ftyp'.
pub struct FileTypeBox {
    name: u32,
    size: u64,
    major_brand: u32,
    minor_version: u32,
    compatible_brands: Vec<u32>,
}

/// Movie header box 'mvhd'.
pub struct MovieHeaderBox {
    pub name: u32,
    pub size: u64,
    pub timescale: u32,
    pub duration: u64,
    // Ignore other fields.
}

pub struct TrackHeaderBox {
    pub name: u32,
    pub size: u64,
    pub track_id: u32,
    pub duration: u64,
    pub width: u32,
    pub height: u32,
}

extern crate byteorder;
use byteorder::{BigEndian, ReadBytesExt};
use std::io::{Read, Result, Seek, SeekFrom, Take};
use std::io::Cursor;

/// Parse a box out of a data buffer.
pub fn read_box_header<T: ReadBytesExt>(src: &mut T) -> Option<BoxHeader> {
    let tmp_size = src.read_u32::<BigEndian>().unwrap();
    let name = src.read_u32::<BigEndian>().unwrap();
    let size = match tmp_size {
        1 => src.read_u64::<BigEndian>().unwrap(),
        _ => tmp_size as u64,
    };
    assert!(size >= 8);
    if tmp_size == 1 {
        assert!(size >= 16);
    }
    let offset = match tmp_size {
        1 => 4 + 4 + 8,
        _ => 4 + 4,
    };
    assert!(offset <= size);
    Some(BoxHeader{
        name: name,
        size: size,
        offset: offset,
    })
}

/// Parse the extra header fields for a full box.
fn read_fullbox_extra<T: ReadBytesExt>(src: &mut T) -> (u8, u32) {
    let version = src.read_u8().unwrap();
    let flags_a = src.read_u8().unwrap();
    let flags_b = src.read_u8().unwrap();
    let flags_c = src.read_u8().unwrap();
    (version, (flags_a as u32) << 16 |
              (flags_b as u32) <<  8 |
              (flags_c as u32))
}

/// Skip over the contents of a box.
pub fn skip_box_content<T: ReadBytesExt + Seek>
  (src: &mut T, header: &BoxHeader)
  -> std::io::Result<u64>
{
    src.seek(SeekFrom::Current((header.size - header.offset) as i64))
}

/// Helper to construct a Take over the contents of a box.
fn limit<'a, T: Read>(f: &'a mut T, h: &BoxHeader) -> Take<&'a mut T> {
    f.take(h.size - h.offset)
}

/// Helper to construct a Cursor over the contents of a box.
fn recurse<T: Read>(f: &mut T, h: &BoxHeader) {
    println!("{} -- recursing", h);
    // FIXME: I couldn't figure out how to do this without copying.
    // We use Seek on the Read we return in skip_box_content, but
    // that trait isn't implemented for a Take like our limit()
    // returns. Slurping the buffer and wrapping it in a Cursor
    // functions as a work around.
    let buf: Vec<u8> = limit(f, &h)
        .bytes()
        .map(|u| u.unwrap())
        .collect();
    let mut content = Cursor::new(buf);
    loop {
        read_box(&mut content);
    }
    println!("{} -- end", h);
}

/// Read the contents of a box, including sub boxes.
/// Right now it just prints the box value rather than
/// returning anything.
pub fn read_box<T: Read + Seek>(f: &mut T) {
    read_box_header(f).and_then(|h| {
        match &(fourcc_to_string(h.name))[..] {
            "ftyp" => {
                let mut content = limit(f, &h);
                let ftyp = read_ftyp(&mut content, &h).unwrap();
                println!("{}", ftyp);
            },
            "moov" => recurse(f, &h),
            "mvhd" => {
                let mut content = limit(f, &h);
                let mvhd = read_mvhd(&mut content, &h).unwrap();
                println!("  {}", mvhd);
            },
            "trak" => recurse(f, &h),
            "tkhd" => {
                let mut content = limit(f, &h);
                let tkhd = read_tkhd(&mut content, &h).unwrap();
                println!("  {}", tkhd);
            },
            _ => {
                // Skip the contents of unknown chunks.
                println!("{}", h);
                skip_box_content(f, &h).unwrap();
            },
        };
        Some(()) // and_then needs an Option to return.
    });
}

/// Entry point for C language callers.
/// Take a ptr, size and return a cursor over it.
fn cursor_from_cbuf(buffer: *mut u8, size: usize) -> Cursor<Vec<u8>> {
    unsafe {
        let b = Vec::from_raw_parts(buffer, size, size);
        Cursor::new(b)
    }
}

/// Take a buffer and call read_box() on it.
#[no_mangle]
pub extern fn read_box_from_buffer(buffer: *mut u8, size: usize) -> bool {
    use std::thread;

    // Validate arguments from C.
    if buffer.is_null() || size < 8 {
        return false;
    }

    // Parse in a subthread.
    let mut c = cursor_from_cbuf(buffer, size);
    let task = thread::spawn(move || {
        read_box(&mut c);
        drop(c);
    });
    // Catch any panics.
    task.join().is_ok()
}

/// Parse an ftype box.
pub fn read_ftyp<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<FileTypeBox> {
    let major = src.read_u32::<BigEndian>().unwrap();
    let minor = src.read_u32::<BigEndian>().unwrap();
    let brand_count = (head.size - 8 - 8) /4;
    let mut brands = Vec::new();
    for _ in 0..brand_count {
        brands.push(src.read_u32::<BigEndian>().unwrap());
    }
    Some(FileTypeBox{
        name: head.name,
        size: head.size,
        major_brand: major,
        minor_version: minor,
        compatible_brands: brands,
    })
}

/// Parse an mvhd box.
pub fn read_mvhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<MovieHeaderBox> {
    let (version, _) = read_fullbox_extra(src);
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        _ => panic!("invalid mhdr version"),
    }
    let timescale = src.read_u32::<BigEndian>().unwrap();
    let duration = match version {
        1 => src.read_u64::<BigEndian>().unwrap(),
        0 => src.read_u32::<BigEndian>().unwrap() as u64,
        _ => panic!("invalid mhdr version"),
    };
    // Skip remaining fields.
    let mut skip: Vec<u8> = vec![0; 80];
    let r = src.read(&mut skip).unwrap();
    assert!(r == skip.len());
    Some(MovieHeaderBox {
        name: head.name,
        size: head.size,
        timescale: timescale,
        duration: duration,
    })
}

/// Parse a tkhd box.
pub fn read_tkhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<TrackHeaderBox> {
    let (version, flags) = read_fullbox_extra(src);
    if flags & 0x1u32 == 0 || flags & 0x2u32 == 0 {
        // Track is disabled.
        return None;
    }
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        _ => panic!("invalid tkhd version"),
    }
    let track_id = src.read_u32::<BigEndian>().unwrap();
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    assert!(_reserved == 0);
    let duration = match version {
        1 => {
            src.read_u64::<BigEndian>().unwrap()
        },
        0 => src.read_u32::<BigEndian>().unwrap() as u64,
        _ => panic!("invalid tkhd version"),
    };
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    // Skip uninterested fields.
    let mut skip: Vec<u8> = vec![0; 44];
    let r = src.read(&mut skip).unwrap();
    assert!(r == skip.len());
    let width = src.read_u32::<BigEndian>().unwrap();
    let height = src.read_u32::<BigEndian>().unwrap();
    Some(TrackHeaderBox {
        name: head.name,
        size: head.size,
        track_id: track_id,
        duration: duration,
        width: width,
        height: height,
    })
}

/// Convert the iso box type or other 4-character value to a string.
fn fourcc_to_string(name: u32) -> String {
    let u32_to_vec = |u| {
        vec!((u >> 24 & 0xffu32) as u8,
             (u >> 16 & 0xffu32) as u8,
             (u >>  8 & 0xffu32) as u8,
             (u & 0xffu32) as u8)
    };
    let name_bytes = u32_to_vec(name);
    String::from_utf8_lossy(&name_bytes).into_owned()
}

use std::fmt;
impl fmt::Display for BoxHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "'{}' {} bytes", fourcc_to_string(self.name), self.size)
    }
}

impl fmt::Display for FileTypeBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        let brand = fourcc_to_string(self.major_brand);
        write!(f, "'{}' {} bytes '{}' v{}", name, self.size,
            brand, self.minor_version)
    }
}

impl fmt::Display for MovieHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        write!(f, "'{}' {} bytes duration {}s", name, self.size,
            (self.duration as f64)/(self.timescale as f64))
    }
}

use std::u16;
impl fmt::Display for TrackHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        // Dimensions are 16.16 fixed-point.
        let base = u16::MAX as f64 + 1.0;
        let width = (self.width as f64) / base;
        let height = (self.height as f64) / base;
        write!(f, "'{}' {} bytes duration {} id {} {}x{}",
            name, self.size, self.duration, self.track_id,
            width, height)
    }
}

#[test]
fn test_read_box_header() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 8];  // minimal box length
    write!(&mut test, "test").unwrap(); // box type
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1952805748);
    assert_eq!(parsed.size, 8);
    println!("box {}", parsed);
}

#[test]
fn test_read_box_header_long() {
    use std::io::Cursor;
    let mut test: Vec<u8> = vec![0, 0, 0, 1]; // long box extension code
    test.extend("long".to_string().into_bytes()); // box type
    test.extend(vec![0, 0, 0, 0, 0, 0, 16, 0]); // 64 bit size
    // Skip generating box content.
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1819242087);
    assert_eq!(parsed.size, 4096);
    println!("box {}", parsed);
}

#[test]
fn test_read_ftyp() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 24]; // size
    write!(&mut test, "ftyp").unwrap(); // type
    write!(&mut test, "mp42").unwrap(); // major brand
    test.extend(vec![0, 0, 0, 0]);      // minor version
    write!(&mut test, "isom").unwrap(); // compatible brands...
    write!(&mut test, "mp42").unwrap();
    assert_eq!(test.len(), 24);

    let mut stream = Cursor::new(test);
    let header = read_box_header(&mut stream).unwrap();
    let parsed = read_ftyp(&mut stream, &header).unwrap();
    assert_eq!(parsed.name, 1718909296);
    assert_eq!(parsed.size, 24);
    assert_eq!(parsed.major_brand, 1836069938);
    assert_eq!(parsed.minor_version, 0);
    assert_eq!(parsed.compatible_brands.len(), 2);
    println!("box {}", parsed);
}
<commit_msg>Return a Result from read_box.<commit_after>// Module for parsing ISO Base Media Format aka video/mp4 streams.

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

/// Basic ISO box structure.
pub struct BoxHeader {
    /// Four character box type
pub name: u32,
    /// Size of the box in bytes
pub size: u64,
    /// Offset to the start of the contained data (or header size).
pub offset: u64,
}

/// File type box 'ftyp'.
pub struct FileTypeBox {
    name: u32,
    size: u64,
    major_brand: u32,
    minor_version: u32,
    compatible_brands: Vec<u32>,
}

/// Movie header box 'mvhd'.
pub struct MovieHeaderBox {
    pub name: u32,
    pub size: u64,
    pub timescale: u32,
    pub duration: u64,
    // Ignore other fields.
}

pub struct TrackHeaderBox {
    pub name: u32,
    pub size: u64,
    pub track_id: u32,
    pub duration: u64,
    pub width: u32,
    pub height: u32,
}

extern crate byteorder;
use byteorder::{BigEndian, ReadBytesExt};
use std::io::{Read, Result, Seek, SeekFrom, Take};
use std::io::Cursor;

/// Parse a box out of a data buffer.
pub fn read_box_header<T: ReadBytesExt>(src: &mut T) -> Result<BoxHeader> {
    let tmp_size = try!(src.read_u32::<BigEndian>());
    let name = try!(src.read_u32::<BigEndian>());
    let size = match tmp_size {
        1 => try!(src.read_u64::<BigEndian>()),
        _ => tmp_size as u64,
    };
    assert!(size >= 8);
    if tmp_size == 1 {
        assert!(size >= 16);
    }
    let offset = match tmp_size {
        1 => 4 + 4 + 8,
        _ => 4 + 4,
    };
    assert!(offset <= size);
    Ok(BoxHeader{
      name: name,
      size: size,
      offset: offset,
    })
}

/// Parse the extra header fields for a full box.
fn read_fullbox_extra<T: ReadBytesExt>(src: &mut T) -> (u8, u32) {
    let version = src.read_u8().unwrap();
    let flags_a = src.read_u8().unwrap();
    let flags_b = src.read_u8().unwrap();
    let flags_c = src.read_u8().unwrap();
    (version, (flags_a as u32) << 16 |
              (flags_b as u32) <<  8 |
              (flags_c as u32))
}

/// Skip over the contents of a box.
pub fn skip_box_content<T: ReadBytesExt + Seek>
  (src: &mut T, header: &BoxHeader)
  -> std::io::Result<u64>
{
    src.seek(SeekFrom::Current((header.size - header.offset) as i64))
}

/// Helper to construct a Take over the contents of a box.
fn limit<'a, T: Read>(f: &'a mut T, h: &BoxHeader) -> Take<&'a mut T> {
    f.take(h.size - h.offset)
}

/// Helper to construct a Cursor over the contents of a box.
fn recurse<T: Read>(f: &mut T, h: &BoxHeader) {
    println!("{} -- recursing", h);
    // FIXME: I couldn't figure out how to do this without copying.
    // We use Seek on the Read we return in skip_box_content, but
    // that trait isn't implemented for a Take like our limit()
    // returns. Slurping the buffer and wrapping it in a Cursor
    // functions as a work around.
    let buf: Vec<u8> = limit(f, &h)
        .bytes()
        .map(|u| u.unwrap())
        .collect();
    let mut content = Cursor::new(buf);
    loop {
        read_box(&mut content).unwrap();
    }
    println!("{} -- end", h);
}

/// Read the contents of a box, including sub boxes.
/// Right now it just prints the box value rather than
/// returning anything.
pub fn read_box<T: Read + Seek>(f: &mut T) -> Result<()> {
    read_box_header(f).and_then(|h| {
        match &(fourcc_to_string(h.name))[..] {
            "ftyp" => {
                let mut content = limit(f, &h);
                let ftyp = read_ftyp(&mut content, &h).unwrap();
                println!("{}", ftyp);
            },
            "moov" => recurse(f, &h),
            "mvhd" => {
                let mut content = limit(f, &h);
                let mvhd = read_mvhd(&mut content, &h).unwrap();
                println!("  {}", mvhd);
            },
            "trak" => recurse(f, &h),
            "tkhd" => {
                let mut content = limit(f, &h);
                let tkhd = read_tkhd(&mut content, &h).unwrap();
                println!("  {}", tkhd);
            },
            _ => {
                // Skip the contents of unknown chunks.
                println!("{} (skipped)", h);
                skip_box_content(f, &h).unwrap();
            },
        };
        Ok(()) // and_then needs a Result to return.
    })
}

/// Entry point for C language callers.
/// Take a ptr, size and return a cursor over it.
fn cursor_from_cbuf(buffer: *mut u8, size: usize) -> Cursor<Vec<u8>> {
    unsafe {
        let b = Vec::from_raw_parts(buffer, size, size);
        Cursor::new(b)
    }
}

/// Take a buffer and call read_box() on it.
#[no_mangle]
pub extern fn read_box_from_buffer(buffer: *mut u8, size: usize) -> bool {
    use std::thread;

    // Validate arguments from C.
    if buffer.is_null() || size < 8 {
        return false;
    }

    // Parse in a subthread.
    let mut c = cursor_from_cbuf(buffer, size);
    let task = thread::spawn(move || {
        read_box(&mut c).unwrap();
        drop(c);
    });
    // Catch any panics.
    task.join().is_ok()
}

/// Parse an ftype box.
pub fn read_ftyp<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<FileTypeBox> {
    let major = src.read_u32::<BigEndian>().unwrap();
    let minor = src.read_u32::<BigEndian>().unwrap();
    let brand_count = (head.size - 8 - 8) /4;
    let mut brands = Vec::new();
    for _ in 0..brand_count {
        brands.push(src.read_u32::<BigEndian>().unwrap());
    }
    Some(FileTypeBox{
        name: head.name,
        size: head.size,
        major_brand: major,
        minor_version: minor,
        compatible_brands: brands,
    })
}

/// Parse an mvhd box.
pub fn read_mvhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<MovieHeaderBox> {
    let (version, _) = read_fullbox_extra(src);
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        _ => panic!("invalid mhdr version"),
    }
    let timescale = src.read_u32::<BigEndian>().unwrap();
    let duration = match version {
        1 => src.read_u64::<BigEndian>().unwrap(),
        0 => src.read_u32::<BigEndian>().unwrap() as u64,
        _ => panic!("invalid mhdr version"),
    };
    // Skip remaining fields.
    let mut skip: Vec<u8> = vec![0; 80];
    let r = src.read(&mut skip).unwrap();
    assert!(r == skip.len());
    Some(MovieHeaderBox {
        name: head.name,
        size: head.size,
        timescale: timescale,
        duration: duration,
    })
}

/// Parse a tkhd box.
pub fn read_tkhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<TrackHeaderBox> {
    let (version, flags) = read_fullbox_extra(src);
    if flags & 0x1u32 == 0 || flags & 0x2u32 == 0 {
        // Track is disabled.
        return None;
    }
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        _ => panic!("invalid tkhd version"),
    }
    let track_id = src.read_u32::<BigEndian>().unwrap();
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    assert!(_reserved == 0);
    let duration = match version {
        1 => {
            src.read_u64::<BigEndian>().unwrap()
        },
        0 => src.read_u32::<BigEndian>().unwrap() as u64,
        _ => panic!("invalid tkhd version"),
    };
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    // Skip uninterested fields.
    let mut skip: Vec<u8> = vec![0; 44];
    let r = src.read(&mut skip).unwrap();
    assert!(r == skip.len());
    let width = src.read_u32::<BigEndian>().unwrap();
    let height = src.read_u32::<BigEndian>().unwrap();
    Some(TrackHeaderBox {
        name: head.name,
        size: head.size,
        track_id: track_id,
        duration: duration,
        width: width,
        height: height,
    })
}

/// Convert the iso box type or other 4-character value to a string.
fn fourcc_to_string(name: u32) -> String {
    let u32_to_vec = |u| {
        vec!((u >> 24 & 0xffu32) as u8,
             (u >> 16 & 0xffu32) as u8,
             (u >>  8 & 0xffu32) as u8,
             (u & 0xffu32) as u8)
    };
    let name_bytes = u32_to_vec(name);
    String::from_utf8_lossy(&name_bytes).into_owned()
}

use std::fmt;
impl fmt::Display for BoxHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "'{}' {} bytes", fourcc_to_string(self.name), self.size)
    }
}

impl fmt::Display for FileTypeBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        let brand = fourcc_to_string(self.major_brand);
        write!(f, "'{}' {} bytes '{}' v{}", name, self.size,
            brand, self.minor_version)
    }
}

impl fmt::Display for MovieHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        write!(f, "'{}' {} bytes duration {}s", name, self.size,
            (self.duration as f64)/(self.timescale as f64))
    }
}

use std::u16;
impl fmt::Display for TrackHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        // Dimensions are 16.16 fixed-point.
        let base = u16::MAX as f64 + 1.0;
        let width = (self.width as f64) / base;
        let height = (self.height as f64) / base;
        write!(f, "'{}' {} bytes duration {} id {} {}x{}",
            name, self.size, self.duration, self.track_id,
            width, height)
    }
}

#[test]
fn test_read_box_header() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 8];  // minimal box length
    write!(&mut test, "test").unwrap(); // box type
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1952805748);
    assert_eq!(parsed.size, 8);
    println!("box {}", parsed);
}

#[test]
fn test_read_box_header_long() {
    use std::io::Cursor;
    let mut test: Vec<u8> = vec![0, 0, 0, 1]; // long box extension code
    test.extend("long".to_string().into_bytes()); // box type
    test.extend(vec![0, 0, 0, 0, 0, 0, 16, 0]); // 64 bit size
    // Skip generating box content.
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1819242087);
    assert_eq!(parsed.size, 4096);
    println!("box {}", parsed);
}

#[test]
fn test_read_ftyp() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 24]; // size
    write!(&mut test, "ftyp").unwrap(); // type
    write!(&mut test, "mp42").unwrap(); // major brand
    test.extend(vec![0, 0, 0, 0]);      // minor version
    write!(&mut test, "isom").unwrap(); // compatible brands...
    write!(&mut test, "mp42").unwrap();
    assert_eq!(test.len(), 24);

    let mut stream = Cursor::new(test);
    let header = read_box_header(&mut stream).unwrap();
    let parsed = read_ftyp(&mut stream, &header).unwrap();
    assert_eq!(parsed.name, 1718909296);
    assert_eq!(parsed.size, 24);
    assert_eq!(parsed.major_brand, 1836069938);
    assert_eq!(parsed.minor_version, 0);
    assert_eq!(parsed.compatible_brands.len(), 2);
    println!("box {}", parsed);
}
<|endoftext|>


<commit_before>#![feature(std_misc)]

extern crate html5ever;
extern crate string_cache;

use std::path::AsPath;
use std::fs::{File, read_dir};
use std::io::{self, Read, Write};
use std::error::FromError;
use std::default::Default;
use std::borrow::Cow;
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::fmt;
use std::mem::swap;

use html5ever::Attribute;
use html5ever::tokenizer::{Tokenizer, TokenSink, Tag, TagKind};
use html5ever::tokenizer::Token as HtmlToken;

use string_cache::atom::Atom;

pub use codegen::{JavaScript, JsVarState, Rust, RustVisibility};

use codegen::{Token, Content, Codegen, ContentType, Scope};
use fragments::{Fragment, ReturnType};

use parser::ExtensibleMap;

pub mod codegen;
pub mod fragments;
mod parser;

#[derive(Debug)]
pub enum Error {
    Io(io::Error),
    Parse(Vec<Cow<'static, str>>)
}

impl FromError<io::Error> for Error {
    fn from_error(e: io::Error) -> Error {
        Error::Io(e)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::Io(ref e) => e.fmt(f),
            &Error::Parse(ref e) => {
                try!(write!(f, "Parse errors: "));
                for (i, error) in e.iter().enumerate() {
                    if i == 0 {
                        try!(error.fmt(f));
                    } else {
                        try!(write!(f, ", {}", error));
                    }
                }
                Ok(())
            }
        }
    }
}

///A collection of templates.
pub struct TemplateGroup<'a> {
    templates: Vec<ParsedTemplate>,
    fragments: HashMap<&'static str, Box<Fragment + 'a>>
}

impl<'a> TemplateGroup<'a> {
    pub fn new() -> TemplateGroup<'a> {
        TemplateGroup {
            templates: vec![],
            fragments: init_fragments()
        }
    }

    ///Add a template from a string.
    pub fn parse_string(&mut self, name: String, source: String) -> Result<(), Error> {
        let mut template = Template::extending(&self.fragments);
        try!(template.parse_string(source));
        self.templates.push(ParsedTemplate {
            name: name,
            parameters: template.parameters,
            tokens: template.tokens
        });

        Ok(())
    }

    ///Add every `*.html` and `*.htm` file in a directory to the template group.
    pub fn parse_directory<P: AsPath>(&mut self, dir: P) -> Result<(), Error> {
        let directory = try!(read_dir(dir));

        for path in directory {
            let path = try!(path).path();
            let info = (
                path.file_stem().and_then(|n| n.to_str()),
                path.extension().and_then(|e| e.to_str())
            );
            if let (Some(name), Some(ext)) = info {
                if ext == "html" || ext == "htm" {
                    let mut source = String::new();
                    try!(File::open(&path).and_then(|mut f| f.read_to_string(&mut source)));
                    try!(self.parse_string(name.to_string(), source));
                }
            }
        }

        Ok(())
    }

    ///Add an already parsed template to the group.
    pub fn add_template(&mut self, name: String, template: Template) {
        self.templates.push(ParsedTemplate {
            name: name,
            parameters: template.parameters,
            tokens: template.tokens
        });
    }

    ///Register a custom fragment.
    pub fn register_fragment<F: Fragment + 'a>(&mut self, fragment: F) {
        let fragment = Box::new(fragment) as Box<Fragment>;
        self.fragments.insert(fragment.identifier(), fragment);
    }

    ///Write template code.
    pub fn emit_code<W: Write, C: Codegen>(&self, writer: &mut W, codegen: &C) -> io::Result<()> {
        codegen.build_module(writer, |w, indent| {
            for template in &self.templates {
                try!(codegen.build_template(w, &template.name, indent, &template.parameters, &template.tokens));
            }
            Ok(())
        })
    }
}

struct ParsedTemplate {
    name: String,
    parameters: HashMap<String, ContentType>,
    tokens: Vec<codegen::Token>
}

///A single template.
pub struct Template<'a> {
    fragments: ExtensibleMap<'a, &'static str, Box<Fragment + 'a>>,
    parameters: HashMap<String, ContentType>,
    tokens: Vec<codegen::Token>,
    errors: Vec<Cow<'static, str>>
}

impl<'a> Template<'a> {
    ///Create an empty template.
    pub fn new() -> Template<'a> {
        Template {
            fragments: ExtensibleMap::Owned(init_fragments()),
            parameters: HashMap::new(),
            tokens: vec![],
            errors: vec![]
        }
    }

    fn extending(fragments: &'a HashMap<&'static str, Box<Fragment + 'a>>) -> Template<'a> {
        Template {
            fragments: ExtensibleMap::extend(fragments),
            parameters: HashMap::new(),
            tokens: vec![],
            errors: vec![]
        }
    }

    ///Add content from a string to the template.
    pub fn parse_string(&mut self, source: String) -> Result<(), Error> {
        {
            let mut tokenizer = Tokenizer::new(&mut*self, Default::default());
            tokenizer.feed(source);
        }

        if self.errors.len() > 0 {
            let mut e = vec![];
            swap(&mut e, &mut self.errors);
            Err(Error::Parse(e))
        } else {
            Ok(())
        }
    }

    ///Register a custom fragment.
    pub fn register_fragment<F: Fragment + 'a>(&mut self, fragment: F) {
        let fragment = Box::new(fragment) as Box<Fragment>;
        self.fragments.insert(fragment.identifier(), fragment);
    }

    ///Write template code.
    pub fn emit_code<W: Write, C: Codegen>(&self, template_name: &str, writer: &mut W, codegen: &C) -> io::Result<()> {
        codegen.build_template(writer, template_name, 0, &self.parameters, &self.tokens)
    }

    fn open_tag(&mut self, name: Atom, attributes: Vec<Attribute>, self_closing: bool) -> Result<(), Cow<'static, str>> {
        let void = is_void(name.as_slice());
        self.tokens.push(Token::BeginTag(name));

        for attribute in attributes {
            let mut content = try!(parser::parse_content(&attribute.value, &self.fragments)).into_iter();
            match content.next() {
                Some(ReturnType::Content(Content::String(text))) => self.tokens.push(Token::BeginAttribute(attribute.name.local, Content::String(text))),
                Some(ReturnType::Content(Content::Placeholder(name))) => {
                    self.reg_string(name.clone());
                    self.tokens.push(Token::BeginAttribute(attribute.name.local, Content::Placeholder(name)));
                },
                Some(ReturnType::Logic(_)) => return Err(Cow::Borrowed("logic can not be used as text")),
                Some(ReturnType::Scope(scope)) => {
                    self.reg_scope_vars(&scope);
                    self.tokens.push(Token::Scope(scope));
                },
                Some(ReturnType::End) => self.tokens.push(Token::End),
                None => self.tokens.push(Token::BeginAttribute(attribute.name.local, Content::String(String::new())))
            }

            for part in content {
                match part {
                    ReturnType::Content(Content::String(text)) => self.tokens.push(Token::AppendToAttribute(Content::String(text))),
                    ReturnType::Content(Content::Placeholder(name)) => {
                        self.reg_string(name.clone());
                        self.tokens.push(Token::AppendToAttribute(Content::Placeholder(name)));
                    },
                    ReturnType::Logic(_) => return Err(Cow::Borrowed("logic can not be used as text")),
                    ReturnType::Scope(scope) => {
                        self.reg_scope_vars(&scope);
                        self.tokens.push(Token::Scope(scope));
                    },
                    ReturnType::End => self.tokens.push(Token::End),
                }
            }

            self.tokens.push(Token::EndAttribute);
        }

        self.tokens.push(Token::EndTag(void || self_closing));

        Ok(())
    }

    fn add_text(&mut self, text: String) -> Result<(), Cow<'static, str>> {
        let mut content = try!(parser::parse_content(&text, &self.fragments)).into_iter();
        match content.next() {
            Some(ReturnType::Content(Content::String(text))) => self.tokens.push(Token::BeginText(Content::String(text))),
            Some(ReturnType::Content(Content::Placeholder(name))) => {
                self.reg_string(name.clone());
                self.tokens.push(Token::BeginText(Content::Placeholder(name)));
            },
            Some(ReturnType::Logic(_)) => return Err(Cow::Borrowed("logic can not be used as text")),
            Some(ReturnType::Scope(scope)) => {
                self.reg_scope_vars(&scope);
                self.tokens.push(Token::Scope(scope));
            },
            Some(ReturnType::End) => self.tokens.push(Token::End),
            None => return Ok(())
        }

        for part in content {
            match part {
                ReturnType::Content(Content::String(text)) => self.tokens.push(Token::AppendToText(Content::String(text))),
                ReturnType::Content(Content::Placeholder(name)) => {
                    self.reg_string(name.clone());
                    self.tokens.push(Token::AppendToText(Content::Placeholder(name)));
                },
                ReturnType::Logic(_) => return Err(Cow::Borrowed("logic can not be used as text")),
                ReturnType::Scope(scope) => {
                    self.reg_scope_vars(&scope);
                    self.tokens.push(Token::Scope(scope));
                },
                ReturnType::End => self.tokens.push(Token::End),
            }
        }
        self.tokens.push(Token::EndText);

        Ok(())
    }

    fn close_tag(&mut self, tag: Atom) {
        self.tokens.push(Token::CloseTag(tag));
    }

    fn reg_scope_vars(&mut self, scope: &Scope) {
        match scope {
            &Scope::If(ref logic) => logic.parameters(&mut |param| self.reg_bool(param.clone()))
        }
    }

    fn reg_string(&mut self, param: String) {
        match self.parameters.entry(param) {
            Entry::Occupied(mut entry) => match entry.get_mut() {
                &mut ContentType::String => {},
                entry => *entry = ContentType::OptionalString
            },
            Entry::Vacant(entry) => { entry.insert(ContentType::String); },
        }
    }

    fn reg_bool(&mut self, param: String) {
        match self.parameters.entry(param) {
            Entry::Occupied(mut entry) => match entry.get_mut() {
                &mut ContentType::Bool => {},
                entry => *entry = ContentType::OptionalString
            },
            Entry::Vacant(entry) => { entry.insert(ContentType::Bool); },
        }
    }
}

impl<'a, 'b: 'a> TokenSink for &'a mut Template<'b> {
    fn process_token(&mut self, token: HtmlToken) {
        match token {
            HtmlToken::TagToken(Tag {
                kind: TagKind::StartTag,
                name,
                attrs,
                self_closing
            }) => if let Err(e) = self.open_tag(name, attrs, self_closing) {
                self.errors.push(e);
            },
            HtmlToken::TagToken(Tag {
                kind: TagKind::EndTag,
                name,
                ..
            }) => self.close_tag(name),
            HtmlToken::CharacterTokens(text) => if let Err(e) = self.add_text(text) {
                self.errors.push(e);
            },
            _ => {}
        }
    }
}

// http://www.w3.org/TR/html5/syntax.html#void-elements
fn is_void(tag: &str) -> bool {
    match tag {
        "area" |
        "base" |
        "br" |
        "col" |
        "embed" |
        "hr" |
        "img" |
        "input" |
        "keygen" |
        "link" |
        "meta" |
        "param" |
        "source" |
        "track" |
        "wbr" => true,
        _ => false
    }
}

fn init_fragments<'a>() -> HashMap<&'static str, Box<Fragment + 'a>> {
    let mut map = HashMap::new();

    let f = fragments::If;
    map.insert(f.identifier(), Box::new(f) as Box<Fragment>);

    let f = fragments::And;
    map.insert(f.identifier(), Box::new(f) as Box<Fragment>);

    let f = fragments::Or;
    map.insert(f.identifier(), Box::new(f) as Box<Fragment>);

    let f = fragments::Not;
    map.insert(f.identifier(), Box::new(f) as Box<Fragment>);

    map
}<commit_msg>Rust update<commit_after>#![feature(convert)]

extern crate html5ever;
extern crate string_cache;

use std::convert::AsRef;
use std::path::Path;
use std::fs::{File, read_dir};
use std::io::{self, Read, Write};
use std::error::FromError;
use std::default::Default;
use std::borrow::Cow;
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::fmt;
use std::mem::swap;

use html5ever::Attribute;
use html5ever::tokenizer::{Tokenizer, TokenSink, Tag, TagKind};
use html5ever::tokenizer::Token as HtmlToken;

use string_cache::atom::Atom;

pub use codegen::{JavaScript, JsVarState, Rust, RustVisibility};

use codegen::{Token, Content, Codegen, ContentType, Scope};
use fragments::{Fragment, ReturnType};

use parser::ExtensibleMap;

pub mod codegen;
pub mod fragments;
mod parser;

#[derive(Debug)]
pub enum Error {
    Io(io::Error),
    Parse(Vec<Cow<'static, str>>)
}

impl FromError<io::Error> for Error {
    fn from_error(e: io::Error) -> Error {
        Error::Io(e)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::Io(ref e) => e.fmt(f),
            &Error::Parse(ref e) => {
                try!(write!(f, "Parse errors: "));
                for (i, error) in e.iter().enumerate() {
                    if i == 0 {
                        try!(error.fmt(f));
                    } else {
                        try!(write!(f, ", {}", error));
                    }
                }
                Ok(())
            }
        }
    }
}

///A collection of templates.
pub struct TemplateGroup<'a> {
    templates: Vec<ParsedTemplate>,
    fragments: HashMap<&'static str, Box<Fragment + 'a>>
}

impl<'a> TemplateGroup<'a> {
    pub fn new() -> TemplateGroup<'a> {
        TemplateGroup {
            templates: vec![],
            fragments: init_fragments()
        }
    }

    ///Add a template from a string.
    pub fn parse_string(&mut self, name: String, source: String) -> Result<(), Error> {
        let mut template = Template::extending(&self.fragments);
        try!(template.parse_string(source));
        self.templates.push(ParsedTemplate {
            name: name,
            parameters: template.parameters,
            tokens: template.tokens
        });

        Ok(())
    }

    ///Add every `*.html` and `*.htm` file in a directory to the template group.
    pub fn parse_directory<P: AsRef<Path>>(&mut self, dir: P) -> Result<(), Error> {
        let directory = try!(read_dir(dir));

        for path in directory {
            let path = try!(path).path();
            let info = (
                path.file_stem().and_then(|n| n.to_str()),
                path.extension().and_then(|e| e.to_str())
            );
            if let (Some(name), Some(ext)) = info {
                if ext == "html" || ext == "htm" {
                    let mut source = String::new();
                    try!(File::open(&path).and_then(|mut f| f.read_to_string(&mut source)));
                    try!(self.parse_string(name.to_string(), source));
                }
            }
        }

        Ok(())
    }

    ///Add an already parsed template to the group.
    pub fn add_template(&mut self, name: String, template: Template) {
        self.templates.push(ParsedTemplate {
            name: name,
            parameters: template.parameters,
            tokens: template.tokens
        });
    }

    ///Register a custom fragment.
    pub fn register_fragment<F: Fragment + 'a>(&mut self, fragment: F) {
        let fragment = Box::new(fragment);
        self.fragments.insert(fragment.identifier(), fragment);
    }

    ///Write template code.
    pub fn emit_code<W: Write, C: Codegen>(&self, writer: &mut W, codegen: &C) -> io::Result<()> {
        codegen.build_module(writer, |w, indent| {
            for template in &self.templates {
                try!(codegen.build_template(w, &template.name, indent, &template.parameters, &template.tokens));
            }
            Ok(())
        })
    }
}

struct ParsedTemplate {
    name: String,
    parameters: HashMap<String, ContentType>,
    tokens: Vec<codegen::Token>
}

///A single template.
pub struct Template<'a> {
    fragments: ExtensibleMap<'a, &'static str, Box<Fragment + 'a>>,
    parameters: HashMap<String, ContentType>,
    tokens: Vec<codegen::Token>,
    errors: Vec<Cow<'static, str>>
}

impl<'a> Template<'a> {
    ///Create an empty template.
    pub fn new() -> Template<'a> {
        Template {
            fragments: ExtensibleMap::Owned(init_fragments()),
            parameters: HashMap::new(),
            tokens: vec![],
            errors: vec![]
        }
    }

    fn extending(fragments: &'a HashMap<&'static str, Box<Fragment + 'a>>) -> Template<'a> {
        Template {
            fragments: ExtensibleMap::extend(fragments),
            parameters: HashMap::new(),
            tokens: vec![],
            errors: vec![]
        }
    }

    ///Add content from a string to the template.
    pub fn parse_string(&mut self, source: String) -> Result<(), Error> {
        {
            let mut tokenizer = Tokenizer::new(&mut*self, Default::default());
            tokenizer.feed(source);
        }

        if self.errors.len() > 0 {
            let mut e = vec![];
            swap(&mut e, &mut self.errors);
            Err(Error::Parse(e))
        } else {
            Ok(())
        }
    }

    ///Register a custom fragment.
    pub fn register_fragment<F: Fragment + 'a>(&mut self, fragment: F) {
        let fragment = Box::new(fragment);
        self.fragments.insert(fragment.identifier(), fragment);
    }

    ///Write template code.
    pub fn emit_code<W: Write, C: Codegen>(&self, template_name: &str, writer: &mut W, codegen: &C) -> io::Result<()> {
        codegen.build_template(writer, template_name, 0, &self.parameters, &self.tokens)
    }

    fn open_tag(&mut self, name: Atom, attributes: Vec<Attribute>, self_closing: bool) -> Result<(), Cow<'static, str>> {
        let void = is_void(name.as_slice());
        self.tokens.push(Token::BeginTag(name));

        for attribute in attributes {
            let mut content = try!(parser::parse_content(&attribute.value, &self.fragments)).into_iter();
            match content.next() {
                Some(ReturnType::Content(Content::String(text))) => self.tokens.push(Token::BeginAttribute(attribute.name.local, Content::String(text))),
                Some(ReturnType::Content(Content::Placeholder(name))) => {
                    self.reg_string(name.clone());
                    self.tokens.push(Token::BeginAttribute(attribute.name.local, Content::Placeholder(name)));
                },
                Some(ReturnType::Logic(_)) => return Err(Cow::Borrowed("logic can not be used as text")),
                Some(ReturnType::Scope(scope)) => {
                    self.reg_scope_vars(&scope);
                    self.tokens.push(Token::Scope(scope));
                },
                Some(ReturnType::End) => self.tokens.push(Token::End),
                None => self.tokens.push(Token::BeginAttribute(attribute.name.local, Content::String(String::new())))
            }

            for part in content {
                match part {
                    ReturnType::Content(Content::String(text)) => self.tokens.push(Token::AppendToAttribute(Content::String(text))),
                    ReturnType::Content(Content::Placeholder(name)) => {
                        self.reg_string(name.clone());
                        self.tokens.push(Token::AppendToAttribute(Content::Placeholder(name)));
                    },
                    ReturnType::Logic(_) => return Err(Cow::Borrowed("logic can not be used as text")),
                    ReturnType::Scope(scope) => {
                        self.reg_scope_vars(&scope);
                        self.tokens.push(Token::Scope(scope));
                    },
                    ReturnType::End => self.tokens.push(Token::End),
                }
            }

            self.tokens.push(Token::EndAttribute);
        }

        self.tokens.push(Token::EndTag(void || self_closing));

        Ok(())
    }

    fn add_text(&mut self, text: String) -> Result<(), Cow<'static, str>> {
        let mut content = try!(parser::parse_content(&text, &self.fragments)).into_iter();
        match content.next() {
            Some(ReturnType::Content(Content::String(text))) => self.tokens.push(Token::BeginText(Content::String(text))),
            Some(ReturnType::Content(Content::Placeholder(name))) => {
                self.reg_string(name.clone());
                self.tokens.push(Token::BeginText(Content::Placeholder(name)));
            },
            Some(ReturnType::Logic(_)) => return Err(Cow::Borrowed("logic can not be used as text")),
            Some(ReturnType::Scope(scope)) => {
                self.reg_scope_vars(&scope);
                self.tokens.push(Token::Scope(scope));
            },
            Some(ReturnType::End) => self.tokens.push(Token::End),
            None => return Ok(())
        }

        for part in content {
            match part {
                ReturnType::Content(Content::String(text)) => self.tokens.push(Token::AppendToText(Content::String(text))),
                ReturnType::Content(Content::Placeholder(name)) => {
                    self.reg_string(name.clone());
                    self.tokens.push(Token::AppendToText(Content::Placeholder(name)));
                },
                ReturnType::Logic(_) => return Err(Cow::Borrowed("logic can not be used as text")),
                ReturnType::Scope(scope) => {
                    self.reg_scope_vars(&scope);
                    self.tokens.push(Token::Scope(scope));
                },
                ReturnType::End => self.tokens.push(Token::End),
            }
        }
        self.tokens.push(Token::EndText);

        Ok(())
    }

    fn close_tag(&mut self, tag: Atom) {
        self.tokens.push(Token::CloseTag(tag));
    }

    fn reg_scope_vars(&mut self, scope: &Scope) {
        match scope {
            &Scope::If(ref logic) => logic.parameters(&mut |param| self.reg_bool(param.clone()))
        }
    }

    fn reg_string(&mut self, param: String) {
        match self.parameters.entry(param) {
            Entry::Occupied(mut entry) => match entry.get_mut() {
                &mut ContentType::String => {},
                entry => *entry = ContentType::OptionalString
            },
            Entry::Vacant(entry) => { entry.insert(ContentType::String); },
        }
    }

    fn reg_bool(&mut self, param: String) {
        match self.parameters.entry(param) {
            Entry::Occupied(mut entry) => match entry.get_mut() {
                &mut ContentType::Bool => {},
                entry => *entry = ContentType::OptionalString
            },
            Entry::Vacant(entry) => { entry.insert(ContentType::Bool); },
        }
    }
}

impl<'a, 'b: 'a> TokenSink for &'a mut Template<'b> {
    fn process_token(&mut self, token: HtmlToken) {
        match token {
            HtmlToken::TagToken(Tag {
                kind: TagKind::StartTag,
                name,
                attrs,
                self_closing
            }) => if let Err(e) = self.open_tag(name, attrs, self_closing) {
                self.errors.push(e);
            },
            HtmlToken::TagToken(Tag {
                kind: TagKind::EndTag,
                name,
                ..
            }) => self.close_tag(name),
            HtmlToken::CharacterTokens(text) => if let Err(e) = self.add_text(text) {
                self.errors.push(e);
            },
            _ => {}
        }
    }
}

// http://www.w3.org/TR/html5/syntax.html#void-elements
fn is_void(tag: &str) -> bool {
    match tag {
        "area" |
        "base" |
        "br" |
        "col" |
        "embed" |
        "hr" |
        "img" |
        "input" |
        "keygen" |
        "link" |
        "meta" |
        "param" |
        "source" |
        "track" |
        "wbr" => true,
        _ => false
    }
}

fn init_fragments<'a>() -> HashMap<&'static str, Box<Fragment + 'a>> {
    let mut map = HashMap::new();

    let f: Box<Fragment + 'a> = Box::new(fragments::If);
    map.insert(f.identifier(), f);

    let f = fragments::And;
    map.insert(f.identifier(), Box::new(f));

    let f = fragments::Or;
    map.insert(f.identifier(), Box::new(f));

    let f = fragments::Not;
    map.insert(f.identifier(), Box::new(f));

    map
}<|endoftext|>