<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

use {Reactor, WireError};
use bincode::serde::DeserializeError;
use futures::{self, Future};
use protocol::Proto;
use serde::{Deserialize, Serialize};
use std::fmt;
use std::io;
use tokio_core::net::TcpStream;
use tokio_core::reactor;
use tokio_proto::BindClient as ProtoBindClient;
use tokio_proto::multiplex::Multiplex;
use tokio_service::Service;

type WireResponse<Resp, E> = Result<Result<Resp, WireError<E>>, DeserializeError>;
type ResponseFuture<Req, Resp, E> = futures::Map<<BindClient<Req, Resp, E> as Service>::Future,
                                            fn(WireResponse<Resp, E>) -> Result<Resp, ::Error<E>>>;
type BindClient<Req, Resp, E> =
    <Proto<Req, Result<Resp, WireError<E>>> as ProtoBindClient<Multiplex, TcpStream>>::BindClient;

/// A client that impls `tokio_service::Service` that writes and reads bytes.
///
/// Typically, this would be combined with a serialization pre-processing step
/// and a deserialization post-processing step.
#[doc(hidden)]
pub struct Client<Req, Resp, E>
    where Req: Serialize + 'static,
          Resp: Deserialize + 'static,
          E: Deserialize + 'static
{
    inner: BindClient<Req, Resp, E>,
}

impl<Req, Resp, E> Clone for Client<Req, Resp, E>
    where Req: Serialize + 'static,
          Resp: Deserialize + 'static,
          E: Deserialize + 'static
{
    fn clone(&self) -> Self {
        Client { inner: self.inner.clone() }
    }
}

impl<Req, Resp, E> Service for Client<Req, Resp, E>
    where Req: Serialize + Sync + Send + 'static,
          Resp: Deserialize + Sync + Send + 'static,
          E: Deserialize + Sync + Send + 'static
{
    type Request = Req;
    type Response = Result<Resp, ::Error<E>>;
    type Error = io::Error;
    type Future = ResponseFuture<Req, Resp, E>;

    fn call(&self, request: Self::Request) -> Self::Future {
        self.inner.call(request).map(Self::map_err)
    }
}

impl<Req, Resp, E> Client<Req, Resp, E>
    where Req: Serialize + 'static,
          Resp: Deserialize + 'static,
          E: Deserialize + 'static
{
    fn new(inner: BindClient<Req, Resp, E>) -> Self
        where Req: Serialize + Sync + Send + 'static,
              Resp: Deserialize + Sync + Send + 'static,
              E: Deserialize + Sync + Send + 'static
    {
        Client { inner: inner }
    }

    fn map_err(resp: WireResponse<Resp, E>) -> Result<Resp, ::Error<E>> {
        resp.map(|r| r.map_err(::Error::from))
            .map_err(::Error::ClientDeserialize)
            .and_then(|r| r)
    }
}

impl<Req, Resp, E> fmt::Debug for Client<Req, Resp, E>
    where Req: Serialize + 'static,
          Resp: Deserialize + 'static,
          E: Deserialize + 'static
{
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "Client {{ .. }}")
    }
}

/// Additional options to configure how the client connects.
#[derive(Clone, Default)]
pub struct Options {
    reactor: Option<Reactor>,
}

impl Options {
    /// Connect using the given reactor handle.
    pub fn handle(mut self, handle: reactor::Handle) -> Self {
        self.reactor = Some(Reactor::Handle(handle));
        self
    }

    /// Connect using the given reactor remote.
    pub fn remote(mut self, remote: reactor::Remote) -> Self {
        self.reactor = Some(Reactor::Remote(remote));
        self
    }
}

/// Exposes a trait for connecting asynchronously to servers.
pub mod future {
    use {REMOTE, Reactor};
    use futures::{self, Async, Future, future};
    use protocol::Proto;
    use serde::{Deserialize, Serialize};
    use std::io;
    use std::marker::PhantomData;
    use std::net::SocketAddr;
    use super::{Client, Options};
    use tokio_core::{self, reactor};
    use tokio_core::net::TcpStream;
    use tokio_proto::BindClient;

    /// Types that can connect to a server asynchronously.
    pub trait Connect: Sized {
        /// The type of the future returned when calling `connect`.
        type ConnectFut: Future<Item = Self, Error = io::Error>;

        /// Connects to a server located at the given address, using the given options.
        fn connect(addr: SocketAddr, options: Options) -> Self::ConnectFut;
    }

    /// A future that resolves to a `Client` or an `io::Error`.
    #[doc(hidden)]
    pub struct ConnectFuture<Req, Resp, E>
        where Req: Serialize + 'static,
              Resp: Deserialize + 'static,
              E: Deserialize + 'static,
    {
        inner:
            future::Either<
                futures::Map<tokio_core::net::TcpStreamNew, MultiplexConnect<Req, Resp, E>>,
                futures::Flatten<
                    futures::MapErr<
                        futures::Oneshot<io::Result<Client<Req, Resp, E>>>,
                        fn(futures::Canceled) -> io::Error>>>,
    }

    impl<Req, Resp, E> Future for ConnectFuture<Req, Resp, E>
        where Req: Serialize + Sync + Send + 'static,
              Resp: Deserialize + Sync + Send + 'static,
              E: Deserialize + Sync + Send + 'static
    {
        type Item = Client<Req, Resp, E>;
        type Error = io::Error;

        fn poll(&mut self) -> futures::Poll<Self::Item, Self::Error> {
            // Ok to unwrap because we ensure the oneshot is always completed.
            match Future::poll(&mut self.inner)? {
                Async::Ready(client) => Ok(Async::Ready(client)),
                Async::NotReady => Ok(Async::NotReady),
            }
        }
    }

    struct MultiplexConnect<Req, Resp, E>(reactor::Handle, PhantomData<(Req, Resp, E)>);

    impl<Req, Resp, E> MultiplexConnect<Req, Resp, E> {
        fn new(handle: reactor::Handle) -> Self {
            MultiplexConnect(handle, PhantomData)
        }
    }

    impl<Req, Resp, E> FnOnce<(TcpStream,)> for MultiplexConnect<Req, Resp, E>
        where Req: Serialize + Sync + Send + 'static,
              Resp: Deserialize + Sync + Send + 'static,
              E: Deserialize + Sync + Send + 'static
    {
        type Output = Client<Req, Resp, E>;

        extern "rust-call" fn call_once(self, (tcp,): (TcpStream,)) -> Client<Req, Resp, E> {
            Client::new(Proto::new().bind_client(&self.0, tcp))
        }
    }

    impl<Req, Resp, E> Connect for Client<Req, Resp, E>
        where Req: Serialize + Sync + Send + 'static,
              Resp: Deserialize + Sync + Send + 'static,
              E: Deserialize + Sync + Send + 'static
    {
        type ConnectFut = ConnectFuture<Req, Resp, E>;

        fn connect(addr: SocketAddr, options: Options) -> Self::ConnectFut {
            let setup = move |tx: futures::sync::oneshot::Sender<_>| {
                move |handle: &reactor::Handle| {
                    let handle2 = handle.clone();
                    TcpStream::connect(&addr, handle)
                        .map(move |tcp| Client::new(Proto::new().bind_client(&handle2, tcp)))
                        .then(move |result| {
                            tx.complete(result);
                            Ok(())
                        })
                }
            };
            let rx = match options.reactor {
                Some(Reactor::Handle(handle)) => {
                    let tcp = TcpStream::connect(&addr, &handle).map(MultiplexConnect::new(handle));
                    return ConnectFuture { inner: future::Either::A(tcp) };
                }
                Some(Reactor::Remote(remote)) => {
                    let (tx, rx) = futures::oneshot();
                    remote.spawn(setup(tx));
                    rx
                }
                None => {
                    let (tx, rx) = futures::oneshot();
                    REMOTE.spawn(setup(tx));
                    rx
                }
            };
            fn panic(canceled: futures::Canceled) -> io::Error {
                unreachable!(canceled)
            }
            ConnectFuture { inner: future::Either::B(rx.map_err(panic as fn(_) -> _).flatten()) }
        }
    }
}

/// Exposes a trait for connecting synchronously to servers.
pub mod sync {
    use futures::Future;
    use serde::{Deserialize, Serialize};
    use std::io;
    use std::net::ToSocketAddrs;
    use super::{Client, Options};

    /// Types that can connect to a server synchronously.
    pub trait Connect: Sized {
        /// Connects to a server located at the given address.
        fn connect<A>(addr: A, options: Options) -> Result<Self, io::Error> where A: ToSocketAddrs;
    }

    impl<Req, Resp, E> Connect for Client<Req, Resp, E>
        where Req: Serialize + Sync + Send + 'static,
              Resp: Deserialize + Sync + Send + 'static,
              E: Deserialize + Sync + Send + 'static
    {
        fn connect<A>(addr: A, options: Options) -> Result<Self, io::Error>
            where A: ToSocketAddrs
        {
            let addr = if let Some(a) = addr.to_socket_addrs()?.next() {
                a
            } else {
                return Err(io::Error::new(io::ErrorKind::AddrNotAvailable,
                                          "`ToSocketAddrs::to_socket_addrs` returned an empty \
                                           iterator."));
            };
            <Self as super::future::Connect>::connect(addr, options).wait()
        }
    }
}
<commit_msg>Minor refactor<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

use {Reactor, WireError};
use bincode::serde::DeserializeError;
use futures::{self, Future};
use protocol::Proto;
use serde::{Deserialize, Serialize};
use std::fmt;
use std::io;
use tokio_core::net::TcpStream;
use tokio_core::reactor;
use tokio_proto::BindClient as ProtoBindClient;
use tokio_proto::multiplex::Multiplex;
use tokio_service::Service;

type WireResponse<Resp, E> = Result<Result<Resp, WireError<E>>, DeserializeError>;
type ResponseFuture<Req, Resp, E> = futures::Map<<BindClient<Req, Resp, E> as Service>::Future,
                                            fn(WireResponse<Resp, E>) -> Result<Resp, ::Error<E>>>;
type BindClient<Req, Resp, E> =
    <Proto<Req, Result<Resp, WireError<E>>> as ProtoBindClient<Multiplex, TcpStream>>::BindClient;

/// A client that impls `tokio_service::Service` that writes and reads bytes.
///
/// Typically, this would be combined with a serialization pre-processing step
/// and a deserialization post-processing step.
#[doc(hidden)]
pub struct Client<Req, Resp, E>
    where Req: Serialize + 'static,
          Resp: Deserialize + 'static,
          E: Deserialize + 'static
{
    inner: BindClient<Req, Resp, E>,
}

impl<Req, Resp, E> Clone for Client<Req, Resp, E>
    where Req: Serialize + 'static,
          Resp: Deserialize + 'static,
          E: Deserialize + 'static
{
    fn clone(&self) -> Self {
        Client { inner: self.inner.clone() }
    }
}

impl<Req, Resp, E> Service for Client<Req, Resp, E>
    where Req: Serialize + Sync + Send + 'static,
          Resp: Deserialize + Sync + Send + 'static,
          E: Deserialize + Sync + Send + 'static
{
    type Request = Req;
    type Response = Result<Resp, ::Error<E>>;
    type Error = io::Error;
    type Future = ResponseFuture<Req, Resp, E>;

    fn call(&self, request: Self::Request) -> Self::Future {
        self.inner.call(request).map(Self::map_err)
    }
}

impl<Req, Resp, E> Client<Req, Resp, E>
    where Req: Serialize + 'static,
          Resp: Deserialize + 'static,
          E: Deserialize + 'static
{
    fn new(inner: BindClient<Req, Resp, E>) -> Self
        where Req: Serialize + Sync + Send + 'static,
              Resp: Deserialize + Sync + Send + 'static,
              E: Deserialize + Sync + Send + 'static
    {
        Client { inner: inner }
    }

    fn map_err(resp: WireResponse<Resp, E>) -> Result<Resp, ::Error<E>> {
        resp.map(|r| r.map_err(::Error::from))
            .map_err(::Error::ClientDeserialize)
            .and_then(|r| r)
    }
}

impl<Req, Resp, E> fmt::Debug for Client<Req, Resp, E>
    where Req: Serialize + 'static,
          Resp: Deserialize + 'static,
          E: Deserialize + 'static
{
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "Client {{ .. }}")
    }
}

/// Additional options to configure how the client connects.
#[derive(Clone, Default)]
pub struct Options {
    reactor: Option<Reactor>,
}

impl Options {
    /// Connect using the given reactor handle.
    pub fn handle(mut self, handle: reactor::Handle) -> Self {
        self.reactor = Some(Reactor::Handle(handle));
        self
    }

    /// Connect using the given reactor remote.
    pub fn remote(mut self, remote: reactor::Remote) -> Self {
        self.reactor = Some(Reactor::Remote(remote));
        self
    }
}

/// Exposes a trait for connecting asynchronously to servers.
pub mod future {
    use {REMOTE, Reactor};
    use futures::{self, Async, Future, future};
    use protocol::Proto;
    use serde::{Deserialize, Serialize};
    use std::io;
    use std::marker::PhantomData;
    use std::net::SocketAddr;
    use super::{Client, Options};
    use tokio_core::{self, reactor};
    use tokio_core::net::TcpStream;
    use tokio_proto::BindClient;

    /// Types that can connect to a server asynchronously.
    pub trait Connect: Sized {
        /// The type of the future returned when calling `connect`.
        type ConnectFut: Future<Item = Self, Error = io::Error>;

        /// Connects to a server located at the given address, using the given options.
        fn connect(addr: SocketAddr, options: Options) -> Self::ConnectFut;
    }

    /// A future that resolves to a `Client` or an `io::Error`.
    #[doc(hidden)]
    pub struct ConnectFuture<Req, Resp, E>
        where Req: Serialize + 'static,
              Resp: Deserialize + 'static,
              E: Deserialize + 'static,
    {
        inner:
            future::Either<
                futures::Map<tokio_core::net::TcpStreamNew, MultiplexConnect<Req, Resp, E>>,
                futures::Flatten<
                    futures::MapErr<
                        futures::Oneshot<io::Result<Client<Req, Resp, E>>>,
                        fn(futures::Canceled) -> io::Error>>>,
    }

    impl<Req, Resp, E> Future for ConnectFuture<Req, Resp, E>
        where Req: Serialize + Sync + Send + 'static,
              Resp: Deserialize + Sync + Send + 'static,
              E: Deserialize + Sync + Send + 'static
    {
        type Item = Client<Req, Resp, E>;
        type Error = io::Error;

        fn poll(&mut self) -> futures::Poll<Self::Item, Self::Error> {
            // Ok to unwrap because we ensure the oneshot is always completed.
            match Future::poll(&mut self.inner)? {
                Async::Ready(client) => Ok(Async::Ready(client)),
                Async::NotReady => Ok(Async::NotReady),
            }
        }
    }

    struct MultiplexConnect<Req, Resp, E>(reactor::Handle, PhantomData<(Req, Resp, E)>);

    impl<Req, Resp, E> MultiplexConnect<Req, Resp, E> {
        fn new(handle: reactor::Handle) -> Self {
            MultiplexConnect(handle, PhantomData)
        }
    }

    impl<Req, Resp, E> FnOnce<(TcpStream,)> for MultiplexConnect<Req, Resp, E>
        where Req: Serialize + Sync + Send + 'static,
              Resp: Deserialize + Sync + Send + 'static,
              E: Deserialize + Sync + Send + 'static
    {
        type Output = Client<Req, Resp, E>;

        extern "rust-call" fn call_once(self, (tcp,): (TcpStream,)) -> Client<Req, Resp, E> {
            Client::new(Proto::new().bind_client(&self.0, tcp))
        }
    }

    impl<Req, Resp, E> Connect for Client<Req, Resp, E>
        where Req: Serialize + Sync + Send + 'static,
              Resp: Deserialize + Sync + Send + 'static,
              E: Deserialize + Sync + Send + 'static
    {
        type ConnectFut = ConnectFuture<Req, Resp, E>;

        fn connect(addr: SocketAddr, options: Options) -> Self::ConnectFut {
            let setup = move |tx: futures::sync::oneshot::Sender<_>| {
                move |handle: &reactor::Handle| {
                    let handle2 = handle.clone();
                    TcpStream::connect(&addr, handle)
                        .map(move |tcp| Client::new(Proto::new().bind_client(&handle2, tcp)))
                        .then(move |result| {
                            tx.complete(result);
                            Ok(())
                        })
                }
            };
            let rx = match options.reactor {
                Some(Reactor::Handle(handle)) => {
                    let tcp = TcpStream::connect(&addr, &handle).map(MultiplexConnect::new(handle));
                    return ConnectFuture { inner: future::Either::A(tcp) };
                }
                Some(Reactor::Remote(remote)) => {
                    let (tx, rx) = futures::oneshot();
                    remote.spawn(setup(tx));
                    rx
                }
                None => {
                    let (tx, rx) = futures::oneshot();
                    REMOTE.spawn(setup(tx));
                    rx
                }
            };
            fn panic(canceled: futures::Canceled) -> io::Error {
                unreachable!(canceled)
            }
            ConnectFuture { inner: future::Either::B(rx.map_err(panic as fn(_) -> _).flatten()) }
        }
    }
}

/// Exposes a trait for connecting synchronously to servers.
pub mod sync {
    use client::future;
    use futures::Future;
    use serde::{Deserialize, Serialize};
    use std::io;
    use std::net::ToSocketAddrs;
    use super::{Client, Options};
    use util::FirstSocketAddr;

    /// Types that can connect to a server synchronously.
    pub trait Connect: Sized {
        /// Connects to a server located at the given address.
        fn connect<A>(addr: A, options: Options) -> Result<Self, io::Error> where A: ToSocketAddrs;
    }

    impl<Req, Resp, E> Connect for Client<Req, Resp, E>
        where Req: Serialize + Sync + Send + 'static,
              Resp: Deserialize + Sync + Send + 'static,
              E: Deserialize + Sync + Send + 'static
    {
        fn connect<A>(addr: A, options: Options) -> Result<Self, io::Error>
            where A: ToSocketAddrs
        {
            <Self as future::Connect>::connect(addr.try_first_socket_addr()?, options).wait()
        }
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::net::tcp;
use std::io::net::ip::SocketAddr;
use std::io::BufferedReader;
use common::{Header, HTTPVersion, Method};
use Request;
use url::Path;
use sequential::{SequentialReader, SequentialReaderBuilder, SequentialWriterBuilder};

/// A ClientConnection is an object that will store a socket to a client
/// and return Request objects.
pub struct ClientConnection {
    remote_addr: io::IoResult<SocketAddr>,

    source: SequentialReaderBuilder<tcp::TcpStream>,
    sink: SequentialWriterBuilder<tcp::TcpStream>,

    // where to read the next header
	next_header_source: BufferedReader<SequentialReader<tcp::TcpStream>>,

    // set to true if the client sent a "Connection: close" in the previous request
    connection_must_close: bool,
}

impl ClientConnection {
    pub fn new(mut socket: tcp::TcpStream) -> ClientConnection {
        let remote_addr = socket.peer_name();

        let mut source = SequentialReaderBuilder::new(socket.clone());
        let first_header = source.next().unwrap();

        ClientConnection {
            source: source,
            sink: SequentialWriterBuilder::new(socket),
            remote_addr: remote_addr,
            next_header_source: BufferedReader::new(first_header),
            connection_must_close: false,
        }
    }

    /// Generates an IoError with invalid input.
    /// This function is here because it is incredibly annoying to create this error.
    fn gen_invalid_input(desc: &'static str) -> io::IoError {
        io::IoError {
            kind: io::InvalidInput,
            desc: desc,
            detail: None
        }
    }

    /// Parses a "HTTP/1.1" string.
    fn parse_http_version(version: &str) -> io::IoResult<HTTPVersion> {
        let elems = version.splitn('/', 2).map(|e| e.to_string()).collect::<Vec<String>>();
        if elems.len() != 2 {
            return Err(ClientConnection::gen_invalid_input("Wrong HTTP version format"))
        }

        let elems = elems.get(1).as_slice().splitn('.', 2)
            .map(|e| e.to_string()).collect::<Vec<String>>();
        if elems.len() != 2 {
            return Err(ClientConnection::gen_invalid_input("Wrong HTTP version format"))
        }

        match (from_str(elems.get(0).as_slice()), from_str(elems.get(1).as_slice())) {
            (Some(major), Some(minor)) =>
                Ok(HTTPVersion(major, minor)),
            _ => Err(ClientConnection::gen_invalid_input("Wrong HTTP version format"))
        }
    }

    /// Parses the request line of the request.
    /// eg. GET / HTTP/1.1
    fn parse_request_line(line: &str) -> io::IoResult<(Method, Path, HTTPVersion)> {
        let mut words = line.words();

        let method = words.next();
        let path = words.next();
        let version = words.next();

        let (method, path, version) = match (method, path, version) {
            (Some(m), Some(p), Some(v)) => (m, p, v),
            _ => return Err(ClientConnection::gen_invalid_input("Missing element in request line"))
        };

        let method = match from_str(method) {
            Some(method) => method,
            None => return Err(ClientConnection::gen_invalid_input("Could not parse method"))
        };

        let path = match Path::parse(path) {
            Ok(p) => p,
            Err(_) => return Err(ClientConnection::gen_invalid_input("Wrong requested URL"))
        };

        let version = try!(ClientConnection::parse_http_version(version));

        Ok((method, path, version))
    }

    /// Parses a header line.
    /// eg. Host: example.com
    fn parse_header(line: &str) -> io::IoResult<Header> {
        let elems = line.splitn(':', 2).map(|e| e.to_string()).collect::<Vec<String>>();

        if elems.len() <= 1 {
            return Err(ClientConnection::gen_invalid_input(
                "Wrong header format (no ':')"))
        }
        if elems.get(1).as_slice().chars().next() != Some(' ') {
            return Err(ClientConnection::gen_invalid_input(
                "Wrong header format (missing space after ':')"))
        }

        let field = match from_str(elems.get(0).as_slice()) {
            None => return Err(ClientConnection::gen_invalid_input("Could not parse header")),
            Some(f) => f
        };

        Ok(Header {
            field: field,
            value: elems.get(1).as_slice().slice_from(1).to_string()
        })
    }

    /// Reads a request from the stream.
    /// Blocks until the header has been read.
    fn read(&mut self) -> io::IoResult<Request> {
        let (method, path, version, headers) = {
            let mut lines = self.next_header_source.lines();

            // reading the request line
            let (method, path, version) =
                try!(ClientConnection::parse_request_line(
                    match lines.next() {
                        Some(line) => try!(line),
                        None => return Err(ClientConnection::gen_invalid_input(
                                    "Missing request line"))
                    }.as_slice().trim()
                ));

            // getting all headers
            let headers = {
                let mut headers = Vec::new();
                loop {
                    match lines.next() {
                        Some(line) => {
                            let line = try!(line);
                            if line.as_slice().trim().len() == 0 { break };
                            headers.push(
                                try!(ClientConnection::parse_header(line.as_slice().trim()))
                            )
                        },
                        None => break
                    }
                }
                headers
            };

            (method, path, version, headers)
        };

        // finding length of body
        let body_length = headers.iter()
            .find(|h| h.field.equiv(&"Content-Length"))
            .and_then(|h| from_str::<uint>(h.value.as_slice()))
            .unwrap_or(0u);

        // building the next reader
        let data_reader = self.source.next().unwrap();
        self.next_header_source = BufferedReader::new(self.source.next().unwrap());

        // building the request
        Ok(Request {
            data_reader: box data_reader,
            response_writer: box self.sink.next().unwrap(),
            remote_addr: self.remote_addr.clone().unwrap(),     // TODO: could fail
            method: method,
            path: path,
            http_version: version,
            headers: headers,
            body_length: body_length,
        })
    }
}

impl Iterator<Request> for ClientConnection {
    /// Blocks until the next Request is available.
    /// Returns None when no new Requests will come from the client.
    fn next(&mut self) -> Option<Request> {
        // the client sent a "connection: close" header in this previous request
        //  or is using HTTP 1.0, meaning that no new request will come
        if self.connection_must_close {
            return None
        }

        // TODO: send back message to client in case of parsing error
        loop {
            let rq = match self.read() {
                Err(_) => return None,
                Ok(rq) => rq
            };

            // updating the status of the connection
            {
                let connection_header = rq.headers.iter()
                    .find(|h| h.field.equiv(&"Connection")).map(|h| h.value.as_slice());

                if connection_header == Some("close") {
                    self.connection_must_close = true;
                } else if rq.http_version == HTTPVersion(1, 0) &&
                        connection_header != Some("keep-alive")
                {
                    self.connection_must_close = true;
                }
            }

            // returning the request
            return Some(rq);
        }
    }
}

#[cfg(test)]
mod test {
    use super::ClientConnection;

    #[test]
    fn test_parse_request_line() {
        let (method, path, ver) =
            ClientConnection::parse_request_line("GET /hello HTTP/1.1").unwrap();

        assert!(method.equiv(&"get"));
        assert!(path == from_str("/hello").unwrap());
        assert!(ver == ::common::HTTPVersion(1, 1));

        assert!(ClientConnection::parse_request_line("GET /hello").is_err());
        assert!(ClientConnection::parse_request_line("qsd qsd qsd").is_err());
    }

    #[test]
    fn test_parse_header() {
        let header = ClientConnection::parse_header("Content-Type: text/html").unwrap();

        assert!(header.field.equiv(&"content-type"));
        assert!(header.value.as_slice() == "text/html");

        assert!(ClientConnection::parse_header("hello world").is_err());
    }
}
<commit_msg>Add timeout to client connection<commit_after>use std::io;
use std::io::net::tcp;
use std::io::net::ip::SocketAddr;
use std::io::BufferedReader;
use common::{Header, HTTPVersion, Method};
use Request;
use url::Path;
use sequential::{SequentialReader, SequentialReaderBuilder, SequentialWriterBuilder};

/// A ClientConnection is an object that will store a socket to a client
/// and return Request objects.
pub struct ClientConnection {
    remote_addr: io::IoResult<SocketAddr>,

    source: SequentialReaderBuilder<tcp::TcpStream>,
    sink: SequentialWriterBuilder<tcp::TcpStream>,

    // where to read the next header
	next_header_source: BufferedReader<SequentialReader<tcp::TcpStream>>,

    // set to true if the client sent a "Connection: close" in the previous request
    connection_must_close: bool,
}

impl ClientConnection {
    pub fn new(mut socket: tcp::TcpStream) -> ClientConnection {
        socket.set_keepalive(Some(10));
        socket.set_timeout(Some(10000));

        let remote_addr = socket.peer_name();

        let mut source = SequentialReaderBuilder::new(socket.clone());
        let first_header = source.next().unwrap();

        ClientConnection {
            source: source,
            sink: SequentialWriterBuilder::new(socket),
            remote_addr: remote_addr,
            next_header_source: BufferedReader::new(first_header),
            connection_must_close: false,
        }
    }

    /// Generates an IoError with invalid input.
    /// This function is here because it is incredibly annoying to create this error.
    fn gen_invalid_input(desc: &'static str) -> io::IoError {
        io::IoError {
            kind: io::InvalidInput,
            desc: desc,
            detail: None
        }
    }

    /// Parses a "HTTP/1.1" string.
    fn parse_http_version(version: &str) -> io::IoResult<HTTPVersion> {
        let elems = version.splitn('/', 2).map(|e| e.to_string()).collect::<Vec<String>>();
        if elems.len() != 2 {
            return Err(ClientConnection::gen_invalid_input("Wrong HTTP version format"))
        }

        let elems = elems.get(1).as_slice().splitn('.', 2)
            .map(|e| e.to_string()).collect::<Vec<String>>();
        if elems.len() != 2 {
            return Err(ClientConnection::gen_invalid_input("Wrong HTTP version format"))
        }

        match (from_str(elems.get(0).as_slice()), from_str(elems.get(1).as_slice())) {
            (Some(major), Some(minor)) =>
                Ok(HTTPVersion(major, minor)),
            _ => Err(ClientConnection::gen_invalid_input("Wrong HTTP version format"))
        }
    }

    /// Parses the request line of the request.
    /// eg. GET / HTTP/1.1
    fn parse_request_line(line: &str) -> io::IoResult<(Method, Path, HTTPVersion)> {
        let mut words = line.words();

        let method = words.next();
        let path = words.next();
        let version = words.next();

        let (method, path, version) = match (method, path, version) {
            (Some(m), Some(p), Some(v)) => (m, p, v),
            _ => return Err(ClientConnection::gen_invalid_input("Missing element in request line"))
        };

        let method = match from_str(method) {
            Some(method) => method,
            None => return Err(ClientConnection::gen_invalid_input("Could not parse method"))
        };

        let path = match Path::parse(path) {
            Ok(p) => p,
            Err(_) => return Err(ClientConnection::gen_invalid_input("Wrong requested URL"))
        };

        let version = try!(ClientConnection::parse_http_version(version));

        Ok((method, path, version))
    }

    /// Parses a header line.
    /// eg. Host: example.com
    fn parse_header(line: &str) -> io::IoResult<Header> {
        let elems = line.splitn(':', 2).map(|e| e.to_string()).collect::<Vec<String>>();

        if elems.len() <= 1 {
            return Err(ClientConnection::gen_invalid_input(
                "Wrong header format (no ':')"))
        }
        if elems.get(1).as_slice().chars().next() != Some(' ') {
            return Err(ClientConnection::gen_invalid_input(
                "Wrong header format (missing space after ':')"))
        }

        let field = match from_str(elems.get(0).as_slice()) {
            None => return Err(ClientConnection::gen_invalid_input("Could not parse header")),
            Some(f) => f
        };

        Ok(Header {
            field: field,
            value: elems.get(1).as_slice().slice_from(1).to_string()
        })
    }

    /// Reads a request from the stream.
    /// Blocks until the header has been read.
    fn read(&mut self) -> io::IoResult<Request> {
        let (method, path, version, headers) = {
            let mut lines = self.next_header_source.lines();

            // reading the request line
            let (method, path, version) =
                try!(ClientConnection::parse_request_line(
                    match lines.next() {
                        Some(line) => try!(line),
                        None => return Err(ClientConnection::gen_invalid_input(
                                    "Missing request line"))
                    }.as_slice().trim()
                ));

            // getting all headers
            let headers = {
                let mut headers = Vec::new();
                loop {
                    match lines.next() {
                        Some(line) => {
                            let line = try!(line);
                            if line.as_slice().trim().len() == 0 { break };
                            headers.push(
                                try!(ClientConnection::parse_header(line.as_slice().trim()))
                            )
                        },
                        None => break
                    }
                }
                headers
            };

            (method, path, version, headers)
        };

        // finding length of body
        let body_length = headers.iter()
            .find(|h| h.field.equiv(&"Content-Length"))
            .and_then(|h| from_str::<uint>(h.value.as_slice()))
            .unwrap_or(0u);

        // building the next reader
        let data_reader = self.source.next().unwrap();
        self.next_header_source = BufferedReader::new(self.source.next().unwrap());

        // building the request
        Ok(Request {
            data_reader: box data_reader,
            response_writer: box self.sink.next().unwrap(),
            remote_addr: self.remote_addr.clone().unwrap(),     // TODO: could fail
            method: method,
            path: path,
            http_version: version,
            headers: headers,
            body_length: body_length,
        })
    }
}

impl Iterator<Request> for ClientConnection {
    /// Blocks until the next Request is available.
    /// Returns None when no new Requests will come from the client.
    fn next(&mut self) -> Option<Request> {
        // the client sent a "connection: close" header in this previous request
        //  or is using HTTP 1.0, meaning that no new request will come
        if self.connection_must_close {
            return None
        }

        // TODO: send back message to client in case of parsing error
        loop {
            let rq = match self.read() {
                Err(_) => return None,
                Ok(rq) => rq
            };

            // updating the status of the connection
            {
                let connection_header = rq.headers.iter()
                    .find(|h| h.field.equiv(&"Connection")).map(|h| h.value.as_slice());

                if connection_header == Some("close") {
                    self.connection_must_close = true;
                } else if rq.http_version == HTTPVersion(1, 0) &&
                        connection_header != Some("keep-alive")
                {
                    self.connection_must_close = true;
                }
            }

            // returning the request
            return Some(rq);
        }
    }
}

#[cfg(test)]
mod test {
    use super::ClientConnection;

    #[test]
    fn test_parse_request_line() {
        let (method, path, ver) =
            ClientConnection::parse_request_line("GET /hello HTTP/1.1").unwrap();

        assert!(method.equiv(&"get"));
        assert!(path == from_str("/hello").unwrap());
        assert!(ver == ::common::HTTPVersion(1, 1));

        assert!(ClientConnection::parse_request_line("GET /hello").is_err());
        assert!(ClientConnection::parse_request_line("qsd qsd qsd").is_err());
    }

    #[test]
    fn test_parse_header() {
        let header = ClientConnection::parse_header("Content-Type: text/html").unwrap();

        assert!(header.field.equiv(&"content-type"));
        assert!(header.value.as_slice() == "text/html");

        assert!(ClientConnection::parse_header("hello world").is_err());
    }
}
<|endoftext|>


<commit_before>extern crate html2text;
extern crate termion;
use std::io::{self, Write};
use std::fmt::Display;
use html2text::render::text_renderer::{RichAnnotation};

fn to_style(tag: &Vec<RichAnnotation>) -> Box<Display> {
    if tag.len() > 0 {
        Box::new(termion::style::Underline)
    } else {
        Box::new(termion::style::Reset)
    }
}

fn main() {
    let stdin = io::stdin();

    let annotated = html2text::from_read_rich(&mut stdin.lock(), 80);
    for line in annotated {
        for (s, tag) in line.iter() {
            let style = to_style(tag);
            write!(io::stdout(), "{}{}{}", style, s, termion::style::Reset);
        }
        write!(io::stdout(), "\n");
    }
}<commit_msg>Style image titles too.<commit_after>extern crate html2text;
extern crate termion;
use std::io::{self, Write};
use std::fmt::Display;
use html2text::render::text_renderer::{RichAnnotation};

fn to_style(tag: &Vec<RichAnnotation>) -> String {
    let mut style = String::new();

    for ann in tag {
        match *ann {
            RichAnnotation::Default => (),
            RichAnnotation::Link(_) => {
                style.push_str(&format!("{}", termion::style::Underline));
            },
            RichAnnotation::Image => {
                style.push_str(&format!("{}", termion::color::Fg(termion::color::LightBlue)));
            },
        }
    }
    style
}

fn main() {
    let stdin = io::stdin();

    let annotated = html2text::from_read_rich(&mut stdin.lock(), 80);
    for line in annotated {
        for (s, tag) in line.iter() {
            let style = to_style(tag);
            write!(io::stdout(), "{}{}{}", style, s, termion::style::Reset);
        }
        write!(io::stdout(), "\n");
    }
}<|endoftext|>


<commit_before>
use std::io::{Read, Write};
use std::default::Default;

use num::traits::ToPrimitive;
use byteorder::{ByteOrder, BigEndian, ReadBytesExt, WriteBytesExt};

use error::{Result, Error};

pub trait ToByte {
    fn encode<T: Write>(&self, buffer: &mut T) -> Result<()>;
    fn encode_nolen<T: Write>(&self, buffer: &mut T)  -> Result<()> {
        self.encode(buffer)
    }
}

pub trait FromByte {
    type R: Default + FromByte;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()>;
    fn decode_new<T: Read>(buffer: &mut T) -> Result<Self::R> {
        let mut temp: Self::R = Default::default();
        match temp.decode(buffer) {
            Ok(_) => Ok(temp),
            Err(e) => Err(e)
        }
    }
}

impl ToByte for i8 {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        buffer.write_i8(*self).or_else(|e| Err(From::from(e)))
    }
}

impl ToByte for i16 {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        buffer.write_i16::<BigEndian>(*self).or_else(|e| Err(From::from(e)))
    }
}
impl ToByte for i32 {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        buffer.write_i32::<BigEndian>(*self).or_else(|e| Err(From::from(e)))
    }
}
impl ToByte for i64 {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        buffer.write_i64::<BigEndian>(*self).or_else(|e| Err(From::from(e)))
    }
}
impl ToByte for String {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        let l = try!(self.len()
                        .to_i16()
                        .ok_or(Error::CodecError));
        try!(buffer.write_i16::<BigEndian>(l));
        buffer.write_all(self.as_bytes())
                             .or_else(|e| Err(From::from(e)))
    }
}

impl <V:ToByte> ToByte for Vec<V> {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        let l = try!(self.len()
                        .to_i32()
                        .ok_or(Error::CodecError));
        try!(buffer.write_i32::<BigEndian>(l));
        for e in self {
            try!(e.encode(buffer));
        }
        Ok(())
    }
    fn encode_nolen<T:Write>(&self, buffer: &mut T) -> Result<()> {
        for e in self {
            try!(e.encode(buffer));
        }
        Ok(())
    }
}

impl ToByte for Vec<u8>{
    fn encode<T: Write>(&self, buffer: &mut T) -> Result<()> {
        let l = try!(self.len()
                        .to_i32()
                        .ok_or(Error::CodecError));
        try!(buffer.write_i32::<BigEndian>(l));
        buffer.write_all(self).or_else(|e| Err(From::from(e)))
    }
    fn encode_nolen<T: Write>(&self, buffer: &mut T) -> Result<()> {
        buffer.write_all(self).or_else(|e| Err(From::from(e)))
    }
}

macro_rules! dec_helper {
    ($val: expr, $dest:expr) => ({
        match $val {
            Ok(val) => {
                *$dest = val;
                Ok(())
                },
            Err(e) => Err(From::from(e))
        }
    })
}
macro_rules! decode {
    ($src:expr, $dest:expr) => ({
        dec_helper!($src.read_i8(), $dest)

    });
    ($src:expr, $method:ident, $dest:expr) => ({
        dec_helper!($src.$method::<BigEndian>(), $dest)

    });
}

impl FromByte for i8 {
    type R = i8;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        decode!(buffer, self)
    }
}

impl FromByte for i16 {
    type R = i16;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        decode!(buffer, read_i16, self)
    }
}
impl FromByte for i32 {
    type R = i32;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        decode!(buffer, read_i32, self)
    }
}
impl FromByte for i64 {
    type R = i64;
    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        decode!(buffer, read_i64, self)
    }
}
impl FromByte for String {
    type R = String;
    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        let mut length: i16 = 0;
        match decode!(buffer, read_i16, &mut length) {
            Ok(_) => {},
            Err(e) => return Err(e)
        }
        let _ = buffer.take(length as u64).read_to_string(self);
        if self.len() != length as usize {
            return Err(Error::UnexpectedEOF);
        }
        Ok(())
    }
}

impl <V: FromByte + Default> FromByte for Vec<V>{
    type R = Vec<V>;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        let mut length: i32 = 0;
        match decode!(buffer, read_i32, &mut length) {
            Ok(_) => {},
            Err(e) => return Err(e)
        }
        for _ in 0..length {
            let mut e: V = Default::default();
            try!(e.decode(buffer));
            self.push(e);
        }
        Ok(())
    }
}

impl FromByte for Vec<u8>{
    type R = Vec<u8>;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        let mut length: i32 = 0;
        match decode!(buffer, read_i32, &mut length) {
            Ok(_) => {},
            Err(e) => return Err(e)
        }
        if length <= 0 {return Ok(());}
        match buffer.take(length as u64).read_to_end(self) {
            Ok(size) => {
                if size < length as usize {
                    return Err(Error::UnexpectedEOF)
                } else {
                    Ok(())
                }
            },
            Err(e) => Err(From::from(e))
        }
    }
}
<commit_msg>codecs encode decode<commit_after>
use std::io::{Read, Write};
use std::default::Default;

use num::traits::ToPrimitive;
use byteorder::{ByteOrder, BigEndian, ReadBytesExt, WriteBytesExt};

use error::{Result, Error};

pub trait ToByte {
    fn encode<T: Write>(&self, buffer: &mut T) -> Result<()>;
    fn encode_nolen<T: Write>(&self, buffer: &mut T)  -> Result<()> {
        self.encode(buffer)
    }
}

pub trait FromByte {
    type R: Default + FromByte;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()>;
    fn decode_new<T: Read>(buffer: &mut T) -> Result<Self::R> {
        let mut temp: Self::R = Default::default();
        match temp.decode(buffer) {
            Ok(_) => Ok(temp),
            Err(e) => Err(e)
        }
    }
}

impl ToByte for i8 {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        buffer.write_i8(*self).or_else(|e| Err(From::from(e)))
    }
}

impl ToByte for i16 {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        buffer.write_i16::<BigEndian>(*self).or_else(|e| Err(From::from(e)))
    }
}
impl ToByte for i32 {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        buffer.write_i32::<BigEndian>(*self).or_else(|e| Err(From::from(e)))
    }
}
impl ToByte for i64 {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        buffer.write_i64::<BigEndian>(*self).or_else(|e| Err(From::from(e)))
    }
}
impl ToByte for String {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        let l = try!(self.len()
                        .to_i16()
                        .ok_or(Error::CodecError));
        try!(buffer.write_i16::<BigEndian>(l));
        buffer.write_all(self.as_bytes())
                             .or_else(|e| Err(From::from(e)))
    }
}

impl <V:ToByte> ToByte for Vec<V> {
    fn encode<T:Write>(&self, buffer: &mut T) -> Result<()> {
        let l = try!(self.len()
                        .to_i32()
                        .ok_or(Error::CodecError));
        try!(buffer.write_i32::<BigEndian>(l));
        for e in self {
            try!(e.encode(buffer));
        }
        Ok(())
    }
    fn encode_nolen<T:Write>(&self, buffer: &mut T) -> Result<()> {
        for e in self {
            try!(e.encode(buffer));
        }
        Ok(())
    }
}

impl ToByte for Vec<u8>{
    fn encode<T: Write>(&self, buffer: &mut T) -> Result<()> {
        let l = try!(self.len()
                        .to_i32()
                        .ok_or(Error::CodecError));
        try!(buffer.write_i32::<BigEndian>(l));
        buffer.write_all(self).or_else(|e| Err(From::from(e)))
    }
    fn encode_nolen<T: Write>(&self, buffer: &mut T) -> Result<()> {
        buffer.write_all(self).or_else(|e| Err(From::from(e)))
    }
}

macro_rules! dec_helper {
    ($val: expr, $dest:expr) => ({
        match $val {
            Ok(val) => {
                *$dest = val;
                Ok(())
                },
            Err(e) => Err(From::from(e))
        }
    })
}
macro_rules! decode {
    ($src:expr, $dest:expr) => ({
        dec_helper!($src.read_i8(), $dest)

    });
    ($src:expr, $method:ident, $dest:expr) => ({
        dec_helper!($src.$method::<BigEndian>(), $dest)

    });
}

impl FromByte for i8 {
    type R = i8;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        decode!(buffer, self)
    }
}

impl FromByte for i16 {
    type R = i16;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        decode!(buffer, read_i16, self)
    }
}
impl FromByte for i32 {
    type R = i32;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        decode!(buffer, read_i32, self)
    }
}
impl FromByte for i64 {
    type R = i64;
    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        decode!(buffer, read_i64, self)
    }
}
impl FromByte for String {
    type R = String;
    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        let mut length: i16 = 0;
        match decode!(buffer, read_i16, &mut length) {
            Ok(_) => {},
            Err(e) => return Err(e)
        }
        let _ = buffer.take(length as u64).read_to_string(self);
        if self.len() != length as usize {
            return Err(Error::UnexpectedEOF);
        }
        Ok(())
    }
}

impl <V: FromByte + Default> FromByte for Vec<V>{
    type R = Vec<V>;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        let mut length: i32 = 0;
        match decode!(buffer, read_i32, &mut length) {
            Ok(_) => {},
            Err(e) => return Err(e)
        }
        for _ in 0..length {
            let mut e: V = Default::default();
            try!(e.decode(buffer));
            self.push(e);
        }
        Ok(())
    }
}

impl FromByte for Vec<u8>{
    type R = Vec<u8>;

    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
        let mut length: i32 = 0;
        match decode!(buffer, read_i32, &mut length) {
            Ok(_) => {},
            Err(e) => return Err(e)
        }
        if length <= 0 {return Ok(());}
        match buffer.take(length as u64).read_to_end(self) {
            Ok(size) => {
                if size < length as usize {
                    return Err(Error::UnexpectedEOF)
                } else {
                    Ok(())
                }
            },
            Err(e) => Err(From::from(e))
        }
    }
}

#[test]
fn codec_i8() {
    use std::io::Cursor;
    let mut buf = vec!();
    let orig: i8 = 5;

    // Encode into buffer
    orig.encode(&mut buf).unwrap();
    assert_eq!(buf, [5]);

    // Read from buffer into existing variable
    let mut dec1: i8 = 0;
    dec1.decode(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec1, orig);

    // Read from buffer into new variable
    let dec2 = i8::decode_new(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec2, orig);
}

#[test]
fn codec_i16() {
    use std::io::Cursor;
    let mut buf = vec!();
    let orig: i16 = 5;

    // Encode into buffer
    orig.encode(&mut buf).unwrap();
    assert_eq!(buf, [0, 5]);

    // Read from buffer into existing variable
    let mut dec1: i16 = 0;
    dec1.decode(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec1, orig);

    // Read from buffer into new variable
    let dec2 = i16::decode_new(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec2, orig);
}

#[test]
fn codec_32() {
    use std::io::Cursor;
    let mut buf = vec!();
    let orig: i32 = 5;

    // Encode into buffer
    orig.encode(&mut buf).unwrap();
    assert_eq!(buf, [0, 0, 0, 5]);

    // Read from buffer into existing variable
    let mut dec1: i32 = 0;
    dec1.decode(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec1, orig);

    // Read from buffer into new variable
    let dec2 = i32::decode_new(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec2, orig);
}

#[test]
fn codec_i64() {
    use std::io::Cursor;
    let mut buf = vec!();
    let orig: i64 = 5;

    // Encode into buffer
    orig.encode(&mut buf).unwrap();
    assert_eq!(buf, [0, 0, 0, 0, 0, 0, 0, 5]);

    // Read from buffer into existing variable
    let mut dec1: i64 = 0;
    dec1.decode(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec1, orig);

    // Read from buffer into new variable
    let dec2 = i64::decode_new(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec2, orig);
}

#[test]
fn codec_string() {
    use std::io::Cursor;
    let mut buf = vec!();
    let orig = "test".to_string();

    // Encode into buffer
    orig.encode(&mut buf).unwrap();
    assert_eq!(buf, [0, 4, 116, 101, 115, 116]);

    // Read from buffer into existing variable
    let mut dec1 = String::new();
    dec1.decode(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec1, orig);

    // Read from buffer into new variable
    let dec2 = String::decode_new(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec2, orig);
}

#[test]
fn codec_vec_u8() {
    use std::io::Cursor;
    let mut buf = vec!();
    let orig: Vec<u8> = vec!(1, 2, 3);

    // Encode into buffer
    orig.encode(&mut buf).unwrap();
    assert_eq!(buf, [0, 0, 0, 3, 1, 2, 3]);

    // Read from buffer into existing variable
    let mut dec1: Vec<u8> = vec!();
    dec1.decode(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec1, orig);

    // Read from buffer into new variable
    let dec2 = Vec::<u8>::decode_new(&mut Cursor::new(buf.clone())).unwrap();
    assert_eq!(dec2, orig);
}
<|endoftext|>


<commit_before>
pub fn divisible(number: u64, divisor: u64) -> bool {
    number % divisor == 0
}

pub fn even(number: u64) -> bool {
    divisible(number, 2)
}
<commit_msg>quotient-reminder function<commit_after>
pub fn divisible(number: u64, divisor: u64) -> bool {
    number % divisor == 0
}

pub fn even(number: u64) -> bool {
    divisible(number, 2)
}

#[allow(dead_code)]
pub fn quot_rem(divident: u64, divisor: u64) -> (u64, u64) {
    let quotient = divident / divisor;
    let reminder = divident % divisor;
    (quotient, reminder)
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn quot_rem_1() {
        let (q, r) = quot_rem(13, 5);
        assert_eq!(q, 2);
        assert_eq!(r, 3);
    }
}
<|endoftext|>


<commit_before>use std::{
    env, fmt, fs,
    io::{Read, Write},
    path::{Path, PathBuf},
    time::Duration,
};

use anyhow::{bail, ensure, Context, Result};
use app_dirs::{get_app_root, AppDataType};
use log::debug;
use serde_derive::{Deserialize, Serialize};
use yansi::{Color, Style};

use crate::types::PathSource;

pub const CONFIG_FILE_NAME: &str = "config.toml";
pub const MAX_CACHE_AGE: Duration = Duration::from_secs(2_592_000); // 30 days
const DEFAULT_UPDATE_INTERVAL_HOURS: u64 = MAX_CACHE_AGE.as_secs() / 3600; // 30 days

fn default_underline() -> bool {
    false
}

fn default_bold() -> bool {
    false
}

fn default_italic() -> bool {
    false
}

#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "lowercase")]
pub enum RawColor {
    Black,
    Red,
    Green,
    Yellow,
    Blue,
    Magenta,
    Purple, // Backwards compatibility with ansi_term (until tealdeer 1.5.0)
    Cyan,
    White,
    Ansi(u8),
    Rgb { r: u8, g: u8, b: u8 },
}

impl From<RawColor> for Color {
    fn from(raw_color: RawColor) -> Self {
        match raw_color {
            RawColor::Black => Self::Black,
            RawColor::Red => Self::Red,
            RawColor::Green => Self::Green,
            RawColor::Yellow => Self::Yellow,
            RawColor::Blue => Self::Blue,
            RawColor::Magenta | RawColor::Purple => Self::Magenta,
            RawColor::Cyan => Self::Cyan,
            RawColor::White => Self::White,
            RawColor::Ansi(num) => Self::Fixed(num),
            RawColor::Rgb { r, g, b } => Self::RGB(r, g, b),
        }
    }
}

#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
struct RawStyle {
    pub foreground: Option<RawColor>,
    pub background: Option<RawColor>,
    #[serde(default = "default_underline")]
    pub underline: bool,
    #[serde(default = "default_bold")]
    pub bold: bool,
    #[serde(default = "default_italic")]
    pub italic: bool,
}

#[allow(clippy::derivable_impls)] // Explicitly control defaults
impl Default for RawStyle {
    fn default() -> Self {
        Self {
            foreground: None,
            background: None,
            underline: false,
            bold: false,
            italic: false,
        }
    }
}

impl From<RawStyle> for Style {
    fn from(raw_style: RawStyle) -> Self {
        let mut style = Self::default();

        if let Some(foreground) = raw_style.foreground {
            style = style.fg(Color::from(foreground));
        }

        if let Some(background) = raw_style.background {
            style = style.bg(Color::from(background));
        }

        if raw_style.underline {
            style = style.underline();
        }

        if raw_style.bold {
            style = style.bold();
        }

        if raw_style.italic {
            style = style.italic();
        }

        style
    }
}

#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
struct RawStyleConfig {
    #[serde(default)]
    pub description: RawStyle,
    #[serde(default)]
    pub command_name: RawStyle,
    #[serde(default)]
    pub example_text: RawStyle,
    #[serde(default)]
    pub example_code: RawStyle,
    #[serde(default)]
    pub example_variable: RawStyle,
}

#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
struct RawDisplayConfig {
    #[serde(default)]
    pub compact: bool,
    #[serde(default)]
    pub use_pager: bool,
}

/// Serde doesn't support default values yet (tracking issue:
/// <https://github.com/serde-rs/serde/issues/368>), so we need to wrap
/// `DEFAULT_UPDATE_INTERVAL_HOURS` in a function to be able to use
/// `#[serde(default = ...)]`
const fn default_auto_update_interval_hours() -> u64 {
    DEFAULT_UPDATE_INTERVAL_HOURS
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct RawUpdatesConfig {
    #[serde(default)]
    pub auto_update: bool,
    #[serde(default = "default_auto_update_interval_hours")]
    pub auto_update_interval_hours: u64,
}

impl Default for RawUpdatesConfig {
    fn default() -> Self {
        Self {
            auto_update: false,
            auto_update_interval_hours: DEFAULT_UPDATE_INTERVAL_HOURS,
        }
    }
}

#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
struct RawDirectoriesConfig {
    #[serde(default)]
    pub cache_dir: Option<PathBuf>,
    #[serde(default)]
    pub custom_pages_dir: Option<PathBuf>,
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
struct RawConfig {
    style: RawStyleConfig,
    display: RawDisplayConfig,
    updates: RawUpdatesConfig,
    directories: RawDirectoriesConfig,
}

impl RawConfig {
    fn new() -> Self {
        Self::default()
    }
}

impl Default for RawConfig {
    fn default() -> Self {
        let mut raw_config = RawConfig {
            style: RawStyleConfig::default(),
            display: RawDisplayConfig::default(),
            updates: RawUpdatesConfig::default(),
            directories: RawDirectoriesConfig::default(),
        };

        // Set default config
        raw_config.style.example_text.foreground = Some(RawColor::Green);
        raw_config.style.command_name.foreground = Some(RawColor::Cyan);
        raw_config.style.example_code.foreground = Some(RawColor::Cyan);
        raw_config.style.example_variable.foreground = Some(RawColor::Cyan);
        raw_config.style.example_variable.underline = true;

        raw_config
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct StyleConfig {
    pub description: Style,
    pub command_name: Style,
    pub example_text: Style,
    pub example_code: Style,
    pub example_variable: Style,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct DisplayConfig {
    pub compact: bool,
    pub use_pager: bool,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct UpdatesConfig {
    pub auto_update: bool,
    pub auto_update_interval: Duration,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PathWithSource {
    pub path: PathBuf,
    pub source: PathSource,
}

impl PathWithSource {
    pub fn path(&self) -> &Path {
        &self.path
    }
}

impl fmt::Display for PathWithSource {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} ({})", self.path.display(), self.source)
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DirectoriesConfig {
    pub cache_dir: PathWithSource,
    pub custom_pages_dir: Option<PathWithSource>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Config {
    pub style: StyleConfig,
    pub display: DisplayConfig,
    pub updates: UpdatesConfig,
    pub directories: DirectoriesConfig,
}

impl Config {
    /// Convert a `RawConfig` to a high-level `Config`.
    ///
    /// For this, some values need to be converted to other types and some
    /// defaults need to be set (sometimes based on env variables).
    fn from_raw(raw_config: RawConfig) -> Result<Self> {
        // Style config
        let style = StyleConfig {
            command_name: raw_config.style.command_name.into(),
            description: raw_config.style.description.into(),
            example_text: raw_config.style.example_text.into(),
            example_code: raw_config.style.example_code.into(),
            example_variable: raw_config.style.example_variable.into(),
        };

        // Display config
        let display = DisplayConfig {
            compact: raw_config.display.compact,
            use_pager: raw_config.display.use_pager,
        };

        // Updates config
        let updates = UpdatesConfig {
            auto_update: raw_config.updates.auto_update,
            auto_update_interval: Duration::from_secs(
                raw_config.updates.auto_update_interval_hours * 3600,
            ),
        };

        // Directories config
        let cache_dir_env_var = "TEALDEER_CACHE_DIR";
        let cache_dir = if let Ok(env_var) = env::var(cache_dir_env_var) {
            // For backwards compatibility reasons, the cache directory can be
            // overridden using an env variable. This is deprecated and will be
            // phased out in the future.
            eprintln!("Warning: The ${} env variable is deprecated, use the `cache_dir` option in the config file instead.", cache_dir_env_var);
            PathWithSource {
                path: PathBuf::from(env_var),
                source: PathSource::EnvVar,
            }
        } else if let Some(config_value) = raw_config.directories.cache_dir {
            // If the user explicitly configured a cache directory, use that.
            PathWithSource {
                path: config_value,
                source: PathSource::ConfigFile,
            }
        } else if let Ok(default_dir) = get_app_root(AppDataType::UserCache, &crate::APP_INFO) {
            // Otherwise, fall back to the default user cache directory.
            PathWithSource {
                path: default_dir,
                source: PathSource::OsConvention,
            }
        } else {
            // If everything fails, give up
            bail!("Could not determine user cache directory");
        };
        let custom_pages_dir = raw_config
            .directories
            .custom_pages_dir
            .map(|path| PathWithSource {
                path,
                source: PathSource::OsConvention,
            })
            .or_else(|| {
                get_app_root(AppDataType::UserData, &crate::APP_INFO)
                    .map(|path| {
                        // Note: The `join("")` call ensures that there's a trailing slash
                        PathWithSource {
                            path: path.join("pages").join(""),
                            source: PathSource::ConfigFile,
                        }
                    })
                    .ok()
            });
        let directories = DirectoriesConfig {
            cache_dir,
            custom_pages_dir,
        };

        Ok(Self {
            style,
            display,
            updates,
            directories,
        })
    }

    pub fn load(enable_styles: bool) -> Result<Self> {
        debug!("Loading config");

        // Determine path
        let (config_file_path, _) = get_config_path().context("Could not determine config path")?;

        // Load raw config
        let raw_config: RawConfig = if config_file_path.exists() && config_file_path.is_file() {
            let mut config_file = fs::File::open(&config_file_path).with_context(|| {
                format!("Failed to open config file path at {:?}", &config_file_path)
            })?;
            let mut contents = String::new();
            config_file.read_to_string(&mut contents).with_context(|| {
                format!("Failed to read from config file at {:?}", &config_file_path)
            })?;
            toml::from_str(&contents).with_context(|| {
                format!("Failed to parse TOML config file at {:?}", config_file_path)
            })?
        } else {
            RawConfig::new()
        };

        // Convert to config
        let mut config = Self::from_raw(raw_config).context("Could not process raw config")?;

        // Potentially override styles
        if !enable_styles {
            config.style = StyleConfig {
                command_name: Style::default(),
                description: Style::default(),
                example_text: Style::default(),
                example_code: Style::default(),
                example_variable: Style::default(),
            };
        }

        Ok(config)
    }
}

/// Return the path to the config directory.
///
/// The config dir path can be overridden using the `TEALDEER_CONFIG_DIR` env
/// variable. Otherwise, the user config directory is returned.
///
/// Note that this function does not verify whether the directory at that
/// location exists, or is a directory.
pub fn get_config_dir() -> Result<(PathBuf, PathSource)> {
    // Allow overriding the config directory by setting the
    // $TEALDEER_CONFIG_DIR env variable.
    if let Ok(value) = env::var("TEALDEER_CONFIG_DIR") {
        return Ok((PathBuf::from(value), PathSource::EnvVar));
    };

    // Otherwise, fall back to the user config directory.
    let dirs = get_app_root(AppDataType::UserConfig, &crate::APP_INFO)
        .context("Failed to determine the user config directory")?;
    Ok((dirs, PathSource::OsConvention))
}

/// Return the path to the config file.
///
/// Note that this function does not verify whether the file at that location
/// exists, or is a file.
pub fn get_config_path() -> Result<(PathBuf, PathSource)> {
    let (config_dir, source) = get_config_dir()?;
    let config_file_path = config_dir.join(CONFIG_FILE_NAME);
    Ok((config_file_path, source))
}

/// Create default config file.
pub fn make_default_config() -> Result<PathBuf> {
    let (config_dir, _) = get_config_dir()?;

    // Ensure that config directory exists
    if config_dir.exists() {
        ensure!(
            config_dir.is_dir(),
            "Config directory could not be created: {} already exists but is not a directory",
            config_dir.to_string_lossy(),
        );
    } else {
        fs::create_dir_all(&config_dir).context("Could not create config directory")?;
    }

    // Ensure that a config file doesn't get overwritten
    let config_file_path = config_dir.join(CONFIG_FILE_NAME);
    ensure!(
        !config_file_path.is_file(),
        "A configuration file already exists at {}, no action was taken.",
        config_file_path.to_str().unwrap()
    );

    // Create default config
    let serialized_config =
        toml::to_string(&RawConfig::new()).context("Failed to serialize default config")?;

    // Write default config
    let mut config_file =
        fs::File::create(&config_file_path).context("Could not create config file")?;
    let _wc = config_file
        .write(serialized_config.as_bytes())
        .context("Could not write to config file")?;

    Ok(config_file_path)
}

#[test]
fn test_serialize_deserialize() {
    let raw_config = RawConfig::new();
    let serialized = toml::to_string(&raw_config).unwrap();
    let deserialized: RawConfig = toml::from_str(&serialized).unwrap();
    assert_eq!(raw_config, deserialized);
}
<commit_msg>Impl From<RawX> for more config types<commit_after>use std::{
    env, fmt, fs,
    io::{Read, Write},
    path::{Path, PathBuf},
    time::Duration,
};

use anyhow::{bail, ensure, Context, Result};
use app_dirs::{get_app_root, AppDataType};
use log::debug;
use serde_derive::{Deserialize, Serialize};
use yansi::{Color, Style};

use crate::types::PathSource;

pub const CONFIG_FILE_NAME: &str = "config.toml";
pub const MAX_CACHE_AGE: Duration = Duration::from_secs(2_592_000); // 30 days
const DEFAULT_UPDATE_INTERVAL_HOURS: u64 = MAX_CACHE_AGE.as_secs() / 3600; // 30 days

fn default_underline() -> bool {
    false
}

fn default_bold() -> bool {
    false
}

fn default_italic() -> bool {
    false
}

#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "lowercase")]
pub enum RawColor {
    Black,
    Red,
    Green,
    Yellow,
    Blue,
    Magenta,
    Purple, // Backwards compatibility with ansi_term (until tealdeer 1.5.0)
    Cyan,
    White,
    Ansi(u8),
    Rgb { r: u8, g: u8, b: u8 },
}

impl From<RawColor> for Color {
    fn from(raw_color: RawColor) -> Self {
        match raw_color {
            RawColor::Black => Self::Black,
            RawColor::Red => Self::Red,
            RawColor::Green => Self::Green,
            RawColor::Yellow => Self::Yellow,
            RawColor::Blue => Self::Blue,
            RawColor::Magenta | RawColor::Purple => Self::Magenta,
            RawColor::Cyan => Self::Cyan,
            RawColor::White => Self::White,
            RawColor::Ansi(num) => Self::Fixed(num),
            RawColor::Rgb { r, g, b } => Self::RGB(r, g, b),
        }
    }
}

#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
struct RawStyle {
    pub foreground: Option<RawColor>,
    pub background: Option<RawColor>,
    #[serde(default = "default_underline")]
    pub underline: bool,
    #[serde(default = "default_bold")]
    pub bold: bool,
    #[serde(default = "default_italic")]
    pub italic: bool,
}

#[allow(clippy::derivable_impls)] // Explicitly control defaults
impl Default for RawStyle {
    fn default() -> Self {
        Self {
            foreground: None,
            background: None,
            underline: false,
            bold: false,
            italic: false,
        }
    }
}

impl From<RawStyle> for Style {
    fn from(raw_style: RawStyle) -> Self {
        let mut style = Self::default();

        if let Some(foreground) = raw_style.foreground {
            style = style.fg(Color::from(foreground));
        }

        if let Some(background) = raw_style.background {
            style = style.bg(Color::from(background));
        }

        if raw_style.underline {
            style = style.underline();
        }

        if raw_style.bold {
            style = style.bold();
        }

        if raw_style.italic {
            style = style.italic();
        }

        style
    }
}

#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
struct RawStyleConfig {
    #[serde(default)]
    pub description: RawStyle,
    #[serde(default)]
    pub command_name: RawStyle,
    #[serde(default)]
    pub example_text: RawStyle,
    #[serde(default)]
    pub example_code: RawStyle,
    #[serde(default)]
    pub example_variable: RawStyle,
}

impl From<RawStyleConfig> for StyleConfig {
    fn from(raw_style_config: RawStyleConfig) -> Self {
        Self {
            command_name: raw_style_config.command_name.into(),
            description: raw_style_config.description.into(),
            example_text: raw_style_config.example_text.into(),
            example_code: raw_style_config.example_code.into(),
            example_variable: raw_style_config.example_variable.into(),
        }
    }
}

#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
struct RawDisplayConfig {
    #[serde(default)]
    pub compact: bool,
    #[serde(default)]
    pub use_pager: bool,
}

impl From<RawDisplayConfig> for DisplayConfig {
    fn from(raw_display_config: RawDisplayConfig) -> Self {
        Self {
            compact: raw_display_config.compact,
            use_pager: raw_display_config.use_pager,
        }
    }
}

/// Serde doesn't support default values yet (tracking issue:
/// <https://github.com/serde-rs/serde/issues/368>), so we need to wrap
/// `DEFAULT_UPDATE_INTERVAL_HOURS` in a function to be able to use
/// `#[serde(default = ...)]`
const fn default_auto_update_interval_hours() -> u64 {
    DEFAULT_UPDATE_INTERVAL_HOURS
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct RawUpdatesConfig {
    #[serde(default)]
    pub auto_update: bool,
    #[serde(default = "default_auto_update_interval_hours")]
    pub auto_update_interval_hours: u64,
}

impl Default for RawUpdatesConfig {
    fn default() -> Self {
        Self {
            auto_update: false,
            auto_update_interval_hours: DEFAULT_UPDATE_INTERVAL_HOURS,
        }
    }
}

impl From<RawUpdatesConfig> for UpdatesConfig {
    fn from(raw_updates_config: RawUpdatesConfig) -> Self {
        Self {
            auto_update: raw_updates_config.auto_update,
            auto_update_interval: Duration::from_secs(
                raw_updates_config.auto_update_interval_hours * 3600,
            ),
        }
    }
}

#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
struct RawDirectoriesConfig {
    #[serde(default)]
    pub cache_dir: Option<PathBuf>,
    #[serde(default)]
    pub custom_pages_dir: Option<PathBuf>,
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
struct RawConfig {
    style: RawStyleConfig,
    display: RawDisplayConfig,
    updates: RawUpdatesConfig,
    directories: RawDirectoriesConfig,
}

impl RawConfig {
    fn new() -> Self {
        Self::default()
    }
}

impl Default for RawConfig {
    fn default() -> Self {
        let mut raw_config = RawConfig {
            style: RawStyleConfig::default(),
            display: RawDisplayConfig::default(),
            updates: RawUpdatesConfig::default(),
            directories: RawDirectoriesConfig::default(),
        };

        // Set default config
        raw_config.style.example_text.foreground = Some(RawColor::Green);
        raw_config.style.command_name.foreground = Some(RawColor::Cyan);
        raw_config.style.example_code.foreground = Some(RawColor::Cyan);
        raw_config.style.example_variable.foreground = Some(RawColor::Cyan);
        raw_config.style.example_variable.underline = true;

        raw_config
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct StyleConfig {
    pub description: Style,
    pub command_name: Style,
    pub example_text: Style,
    pub example_code: Style,
    pub example_variable: Style,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct DisplayConfig {
    pub compact: bool,
    pub use_pager: bool,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct UpdatesConfig {
    pub auto_update: bool,
    pub auto_update_interval: Duration,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PathWithSource {
    pub path: PathBuf,
    pub source: PathSource,
}

impl PathWithSource {
    pub fn path(&self) -> &Path {
        &self.path
    }
}

impl fmt::Display for PathWithSource {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} ({})", self.path.display(), self.source)
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DirectoriesConfig {
    pub cache_dir: PathWithSource,
    pub custom_pages_dir: Option<PathWithSource>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Config {
    pub style: StyleConfig,
    pub display: DisplayConfig,
    pub updates: UpdatesConfig,
    pub directories: DirectoriesConfig,
}

impl Config {
    /// Convert a `RawConfig` to a high-level `Config`.
    ///
    /// For this, some values need to be converted to other types and some
    /// defaults need to be set (sometimes based on env variables).
    fn from_raw(raw_config: RawConfig) -> Result<Self> {
        let style = raw_config.style.into();
        let display = raw_config.display.into();
        let updates = raw_config.updates.into();

        // Determine directories config. For this, we need to take some
        // additional factory into account, like env variables, or the
        // user config.
        let cache_dir_env_var = "TEALDEER_CACHE_DIR";
        let cache_dir = if let Ok(env_var) = env::var(cache_dir_env_var) {
            // For backwards compatibility reasons, the cache directory can be
            // overridden using an env variable. This is deprecated and will be
            // phased out in the future.
            eprintln!("Warning: The ${} env variable is deprecated, use the `cache_dir` option in the config file instead.", cache_dir_env_var);
            PathWithSource {
                path: PathBuf::from(env_var),
                source: PathSource::EnvVar,
            }
        } else if let Some(config_value) = raw_config.directories.cache_dir {
            // If the user explicitly configured a cache directory, use that.
            PathWithSource {
                path: config_value,
                source: PathSource::ConfigFile,
            }
        } else if let Ok(default_dir) = get_app_root(AppDataType::UserCache, &crate::APP_INFO) {
            // Otherwise, fall back to the default user cache directory.
            PathWithSource {
                path: default_dir,
                source: PathSource::OsConvention,
            }
        } else {
            // If everything fails, give up
            bail!("Could not determine user cache directory");
        };
        let custom_pages_dir = raw_config
            .directories
            .custom_pages_dir
            .map(|path| PathWithSource {
                path,
                source: PathSource::OsConvention,
            })
            .or_else(|| {
                get_app_root(AppDataType::UserData, &crate::APP_INFO)
                    .map(|path| {
                        // Note: The `join("")` call ensures that there's a trailing slash
                        PathWithSource {
                            path: path.join("pages").join(""),
                            source: PathSource::ConfigFile,
                        }
                    })
                    .ok()
            });
        let directories = DirectoriesConfig {
            cache_dir,
            custom_pages_dir,
        };

        Ok(Self {
            style,
            display,
            updates,
            directories,
        })
    }

    pub fn load(enable_styles: bool) -> Result<Self> {
        debug!("Loading config");

        // Determine path
        let (config_file_path, _) = get_config_path().context("Could not determine config path")?;

        // Load raw config
        let raw_config: RawConfig = if config_file_path.exists() && config_file_path.is_file() {
            let mut config_file = fs::File::open(&config_file_path).with_context(|| {
                format!("Failed to open config file path at {:?}", &config_file_path)
            })?;
            let mut contents = String::new();
            config_file.read_to_string(&mut contents).with_context(|| {
                format!("Failed to read from config file at {:?}", &config_file_path)
            })?;
            toml::from_str(&contents).with_context(|| {
                format!("Failed to parse TOML config file at {:?}", config_file_path)
            })?
        } else {
            RawConfig::new()
        };

        // Convert to config
        let mut config = Self::from_raw(raw_config).context("Could not process raw config")?;

        // Potentially override styles
        if !enable_styles {
            config.style = StyleConfig {
                command_name: Style::default(),
                description: Style::default(),
                example_text: Style::default(),
                example_code: Style::default(),
                example_variable: Style::default(),
            };
        }

        Ok(config)
    }
}

/// Return the path to the config directory.
///
/// The config dir path can be overridden using the `TEALDEER_CONFIG_DIR` env
/// variable. Otherwise, the user config directory is returned.
///
/// Note that this function does not verify whether the directory at that
/// location exists, or is a directory.
pub fn get_config_dir() -> Result<(PathBuf, PathSource)> {
    // Allow overriding the config directory by setting the
    // $TEALDEER_CONFIG_DIR env variable.
    if let Ok(value) = env::var("TEALDEER_CONFIG_DIR") {
        return Ok((PathBuf::from(value), PathSource::EnvVar));
    };

    // Otherwise, fall back to the user config directory.
    let dirs = get_app_root(AppDataType::UserConfig, &crate::APP_INFO)
        .context("Failed to determine the user config directory")?;
    Ok((dirs, PathSource::OsConvention))
}

/// Return the path to the config file.
///
/// Note that this function does not verify whether the file at that location
/// exists, or is a file.
pub fn get_config_path() -> Result<(PathBuf, PathSource)> {
    let (config_dir, source) = get_config_dir()?;
    let config_file_path = config_dir.join(CONFIG_FILE_NAME);
    Ok((config_file_path, source))
}

/// Create default config file.
pub fn make_default_config() -> Result<PathBuf> {
    let (config_dir, _) = get_config_dir()?;

    // Ensure that config directory exists
    if config_dir.exists() {
        ensure!(
            config_dir.is_dir(),
            "Config directory could not be created: {} already exists but is not a directory",
            config_dir.to_string_lossy(),
        );
    } else {
        fs::create_dir_all(&config_dir).context("Could not create config directory")?;
    }

    // Ensure that a config file doesn't get overwritten
    let config_file_path = config_dir.join(CONFIG_FILE_NAME);
    ensure!(
        !config_file_path.is_file(),
        "A configuration file already exists at {}, no action was taken.",
        config_file_path.to_str().unwrap()
    );

    // Create default config
    let serialized_config =
        toml::to_string(&RawConfig::new()).context("Failed to serialize default config")?;

    // Write default config
    let mut config_file =
        fs::File::create(&config_file_path).context("Could not create config file")?;
    let _wc = config_file
        .write(serialized_config.as_bytes())
        .context("Could not write to config file")?;

    Ok(config_file_path)
}

#[test]
fn test_serialize_deserialize() {
    let raw_config = RawConfig::new();
    let serialized = toml::to_string(&raw_config).unwrap();
    let deserialized: RawConfig = toml::from_str(&serialized).unwrap();
    assert_eq!(raw_config, deserialized);
}
<|endoftext|>