<commit_before>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::{BufMut, BytesMut};
use std::cmp;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use futures::{Async,AsyncSink,Poll,Sink,Stream,Future,future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use channel::BasicProperties;
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec {
    pub frame_max: u32,
}

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("amqp decoder; frame={:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      // Ensure we at least allocate room for two complete frames, to make sure that forthcoming
      // `BytesMut::extend` calls do not call the allocator each time.
      let frame_max = cmp::max(self.frame_max, 8192) as usize;
      if buf.remaining_mut() < frame_max / 2 {
        trace!("encoder; reserve={}", frame_max * 2);
        buf.reserve(frame_max * 2);
      }
      trace!("encoder; frame={:?}", frame);
      let offset = buf.len();
      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, offset)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, offset)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, offset), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, offset), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, offset), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            trace!("encoder; frame_size={}", sz - offset);
            return Ok(());
          },
          Err(GenError::BufferTooSmall(sz)) => {
            // BufferTooSmall error variant returns the index the next write would have
            // occured if there was enough space in the buffer. Thus we subtract the
            // buffer's length to know how much bytes we sould make available.
            let length = buf.len();
            trace!("encoder; sz={} length={} extend={}", sz, length, sz - length);
            buf.extend(repeat(0).take(sz - length));
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream: Framed<T,AMQPCodec>,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: Send,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(stream: T, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error> + Send> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_frame_max(options.frame_max);
    conn.set_heartbeat(options.heartbeat);
    if let Err(e) = conn.connect() {
      let err = format!("Failed to connect: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let codec = AMQPCodec {
      frame_max: conn.configuration.frame_max,
    };
    let t = AMQPTransport {
      upstream:     stream.framed(codec),
      conn:         conn,
    };
    let connector = AMQPTransportConnector {
      transport: Some(t),
    };
    Box::new(connector)
  }

  /// Send a frame to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frame to a queue, to actually send the frame you have to
  /// call either `poll` or `poll_send`.
  pub fn send_frame(&mut self, frame: Frame) {
    self.conn.frame_queue.push_back(frame);
  }

  /// Send content frames to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frames to a queue, to actually send the frames you have to
  /// call either `poll` or `poll_send`.
  pub fn send_content_frames(&mut self, channel_id: u16, payload: &[u8], properties: BasicProperties) {
    self.conn.send_content_frames(channel_id, 60, payload, properties);
  }

  /// Poll the network to receive & handle incoming frames.
  ///
  /// # Return value
  ///
  /// This function will always return `Ok(Async::NotReady)` except in two cases:
  ///
  /// * In case of error, it will return `Err(e)`
  /// * If the socket was closed, it will return `Ok(Async::Ready(()))`
  pub fn poll_recv(&mut self) -> Poll<(), io::Error> {
    loop {
      match self.upstream.poll() {
        Ok(Async::Ready(Some(frame))) => {
          trace!("transport poll_recv; frame={:?}", frame);
          if let Err(e) = self.conn.handle_frame(frame) {
            let err = format!("failed to handle frame: {:?}", e);
            return Err(io::Error::new(io::ErrorKind::Other, err));
          }
        },
        Ok(Async::Ready(None)) => {
          trace!("transport poll_recv; status=Ready(None)");
          return Ok(Async::Ready(()));
        },
        Ok(Async::NotReady) => {
          trace!("transport poll_recv; status=NotReady");
          return Ok(Async::NotReady);
        },
        Err(e) => {
          error!("transport poll_recv; status=Err({:?})", e);
          return Err(From::from(e));
        },
      };
    }
  }

  /// Poll the network to send outcoming frames.
  pub fn poll_send(&mut self) -> Poll<(), io::Error> {
    while let Some(frame) = self.conn.next_frame() {
      trace!("transport poll_send; frame={:?}", frame);
      match self.upstream.start_send(frame)? {
        AsyncSink::Ready => {
          trace!("transport poll_send; status=Ready");
        },
        AsyncSink::NotReady(frame) => {
          trace!("transport poll_send; status=NotReady");
          self.conn.frame_queue.push_front(frame);
          return Ok(Async::NotReady);
        }
      }
    }
    self.upstream.poll_complete()
  }
}

impl<T> Future for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<(), io::Error> {
        trace!("transport poll");
        if let Async::Ready(_) = self.poll_recv()? {
            return Ok(Async::Ready(()));
        }
        self.poll_send()?;
        Ok(Async::NotReady)
    }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    trace!("connector poll; has_transport={:?}", !self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    if let Async::Ready(_) = transport.poll()? {
      trace!("connector poll transport; status=Ready");
      return Err(io::Error::new(io::ErrorKind::Other, "The connection was closed during the handshake"));
    }

    trace!("connector poll; state=ConnectionState::{:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      return Ok(Async::Ready(transport))
    }

    self.transport = Some(transport);
    Ok(Async::NotReady)
  }
}

#[macro_export]
macro_rules! try_lock_transport (
    ($t: expr) => ({
        match $t.try_lock() {
            Ok(t) => t,
            Err(_) => if $t.is_poisoned() {
                return Err(io::Error::new(io::ErrorKind::Other, "Transport mutex is poisoned"))
            } else {
                task::current().notify();
                return Ok(Async::NotReady)
            }
        }
    });
);

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

  #[test]
  fn encode_multiple_frames() {
    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::with_capacity(8192);
    let r = codec.encode(Frame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(8, buffer.len());
    let r = codec.encode(Frame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(16, buffer.len());
    let r = codec.encode(Frame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(24, buffer.len());
  }

  #[test]
  fn encode_extend_buffer() {
    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::new();
    let r = codec.encode(Frame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(true, buffer.capacity() >= (codec.frame_max as usize));
    assert_ne!(0, buffer.len());
  }
}
<commit_msg>Bring back Sink impl for AMQPTransport<commit_after>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::{BufMut, BytesMut};
use std::cmp;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use futures::{Async,AsyncSink,Poll,Sink,StartSend,Stream,Future,future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use channel::BasicProperties;
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec {
    pub frame_max: u32,
}

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("amqp decoder; frame={:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      // Ensure we at least allocate room for two complete frames, to make sure that forthcoming
      // `BytesMut::extend` calls do not call the allocator each time.
      let frame_max = cmp::max(self.frame_max, 8192) as usize;
      if buf.remaining_mut() < frame_max / 2 {
        trace!("encoder; reserve={}", frame_max * 2);
        buf.reserve(frame_max * 2);
      }
      trace!("encoder; frame={:?}", frame);
      let offset = buf.len();
      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, offset)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, offset)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, offset), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, offset), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, offset), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            trace!("encoder; frame_size={}", sz - offset);
            return Ok(());
          },
          Err(GenError::BufferTooSmall(sz)) => {
            // BufferTooSmall error variant returns the index the next write would have
            // occured if there was enough space in the buffer. Thus we subtract the
            // buffer's length to know how much bytes we sould make available.
            let length = buf.len();
            trace!("encoder; sz={} length={} extend={}", sz, length, sz - length);
            buf.extend(repeat(0).take(sz - length));
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream: Framed<T,AMQPCodec>,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: Send,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(stream: T, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error> + Send> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_frame_max(options.frame_max);
    conn.set_heartbeat(options.heartbeat);
    if let Err(e) = conn.connect() {
      let err = format!("Failed to connect: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let codec = AMQPCodec {
      frame_max: conn.configuration.frame_max,
    };
    let t = AMQPTransport {
      upstream:     stream.framed(codec),
      conn:         conn,
    };
    let connector = AMQPTransportConnector {
      transport: Some(t),
    };
    Box::new(connector)
  }

  /// Send a frame to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frame to a queue, to actually send the frame you have to
  /// call either `poll` or `poll_send`.
  pub fn send_frame(&mut self, frame: Frame) {
    self.conn.frame_queue.push_back(frame);
  }

  /// Send content frames to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frames to a queue, to actually send the frames you have to
  /// call either `poll` or `poll_send`.
  pub fn send_content_frames(&mut self, channel_id: u16, payload: &[u8], properties: BasicProperties) {
    self.conn.send_content_frames(channel_id, 60, payload, properties);
  }

  /// Poll the network to receive & handle incoming frames.
  ///
  /// # Return value
  ///
  /// This function will always return `Ok(Async::NotReady)` except in two cases:
  ///
  /// * In case of error, it will return `Err(e)`
  /// * If the socket was closed, it will return `Ok(Async::Ready(()))`
  pub fn poll_recv(&mut self) -> Poll<(), io::Error> {
    loop {
      match self.upstream.poll() {
        Ok(Async::Ready(Some(frame))) => {
          trace!("transport poll_recv; frame={:?}", frame);
          if let Err(e) = self.conn.handle_frame(frame) {
            let err = format!("failed to handle frame: {:?}", e);
            return Err(io::Error::new(io::ErrorKind::Other, err));
          }
        },
        Ok(Async::Ready(None)) => {
          trace!("transport poll_recv; status=Ready(None)");
          return Ok(Async::Ready(()));
        },
        Ok(Async::NotReady) => {
          trace!("transport poll_recv; status=NotReady");
          return Ok(Async::NotReady);
        },
        Err(e) => {
          error!("transport poll_recv; status=Err({:?})", e);
          return Err(From::from(e));
        },
      };
    }
  }

  /// Poll the network to send outcoming frames.
  pub fn poll_send(&mut self) -> Poll<(), io::Error> {
    while let Some(frame) = self.conn.next_frame() {
      trace!("transport poll_send; frame={:?}", frame);
      match self.upstream.start_send(frame)? {
        AsyncSink::Ready => {
          trace!("transport poll_send; status=Ready");
        },
        AsyncSink::NotReady(frame) => {
          trace!("transport poll_send; status=NotReady");
          self.conn.frame_queue.push_front(frame);
          return Ok(Async::NotReady);
        }
      }
    }
    self.upstream.poll_complete()
  }
}

impl<T> Future for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<(), io::Error> {
        trace!("transport poll");
        if let Async::Ready(_) = self.poll_recv()? {
            return Ok(Async::Ready(()));
        }
        self.poll_send()?;
        Ok(Async::NotReady)
    }
}

impl <T> Sink for AMQPTransport<T>
    where T: AsyncWrite,
          T: Send,
          T: 'static {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, frame: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        trace!("transport start_send; frame={:?}", frame);
        self.upstream.start_send(frame)
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        trace!("transport poll_complete");
        self.upstream.poll_complete()
    }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    trace!("connector poll; has_transport={:?}", !self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    if let Async::Ready(_) = transport.poll()? {
      trace!("connector poll transport; status=Ready");
      return Err(io::Error::new(io::ErrorKind::Other, "The connection was closed during the handshake"));
    }

    trace!("connector poll; state=ConnectionState::{:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      return Ok(Async::Ready(transport))
    }

    self.transport = Some(transport);
    Ok(Async::NotReady)
  }
}

#[macro_export]
macro_rules! try_lock_transport (
    ($t: expr) => ({
        match $t.try_lock() {
            Ok(t) => t,
            Err(_) => if $t.is_poisoned() {
                return Err(io::Error::new(io::ErrorKind::Other, "Transport mutex is poisoned"))
            } else {
                task::current().notify();
                return Ok(Async::NotReady)
            }
        }
    });
);

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

  #[test]
  fn encode_multiple_frames() {
    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::with_capacity(8192);
    let r = codec.encode(Frame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(8, buffer.len());
    let r = codec.encode(Frame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(16, buffer.len());
    let r = codec.encode(Frame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(24, buffer.len());
  }

  #[test]
  fn encode_extend_buffer() {
    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::new();
    let r = codec.encode(Frame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(true, buffer.capacity() >= (codec.frame_max as usize));
    assert_ne!(0, buffer.len());
  }
}
<|endoftext|>


<commit_before>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::cmp;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use std::time::Duration;
use futures::{Async,Poll,Sink,Stream,StartSend,Future,future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use tokio_timer::{Interval,Timer};
use channel::BasicProperties;
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec {
    pub frame_max: u32,
}

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length    = buf.len();
      // Ensure we at least allocate 8192 so that the buffer is big enough for the frame_max
      // negociation. Afterwards, use frame_max if > 8192.
      let frame_max = cmp::max(self.frame_max, 8192) as usize;
      if length < frame_max {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(frame_max - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream: Framed<T,AMQPCodec>,
  heartbeat: Option<Interval>,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: Sync+Send,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(stream: T, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_frame_max(options.frame_max);
    conn.set_heartbeat(options.heartbeat);
    if let Err(e) = conn.connect() {
      let err = format!("Failed to connect: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let codec = AMQPCodec {
      frame_max: conn.configuration.frame_max,
    };
    let mut t = AMQPTransport {
      upstream:  stream.framed(codec),
      heartbeat: None,
      conn:      conn,
    };

    if let Err(e) = t.send_and_handle_frames() {
      let err = format!("Failed to handle frames: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let mut connector = AMQPTransportConnector {
      transport: Some(t),
    };

    trace!("pre-poll");
    if let Err(e) = connector.poll() {
      let err = format!("Failed to handle frames: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }
    trace!("post-poll");

    Box::new(connector)
  }

  pub fn start_heartbeat(&mut self) {
      let heartbeat = self.conn.configuration.heartbeat as u64;
      if heartbeat > 0 {
          self.heartbeat = Some(Timer::default().interval(Duration::from_secs(heartbeat)));
      }
  }

  fn poll_heartbeat(&mut self) -> Result<(), io::Error> {
    if let Some(Ok(Async::Ready(_))) = self.heartbeat.as_mut().map(Interval::poll) {
      trace!("Sending heartbeat");
      if let Err(e) = self.send_frame(Frame::Heartbeat(0)) {
        error!("Failed to send heartbeat: {:?}", e);
        return Err(e);
      }
    }
    Ok(())
  }

  fn poll_upstream(&mut self) -> Poll<Option<()>, io::Error> {
    trace!("poll upstream");
    let value = match self.upstream.poll() {
      Ok(Async::Ready(t)) => t,
      Ok(Async::NotReady) => {
        trace!("upstream poll gave NotReady");
        return Ok(Async::NotReady);
      },
      Err(e) => {
        error!("upstream poll gave error: {:?}", e);
        return Err(From::from(e));
      },
    };

    if let Some(frame) = value {
      trace!("upstream poll gave frame: {:?}", frame);
      if let Err(e) = self.conn.handle_frame(frame) {
        let err = format!("failed to handle frame: {:?}", e);
        return Err(io::Error::new(io::ErrorKind::Other, err));
      }
      self.send_frames()?;
      Ok(Async::Ready(Some(())))
    } else {
      error!("upstream poll gave Ready(None)");
      Ok(Async::Ready(None))
    }
  }

  pub fn send_and_handle_frames(&mut self) -> Poll<Option<()>, io::Error> {
    self.send_frames()?;
    self.handle_frames()
  }

  pub fn send_and_handle_frames_with_payload(&mut self, channel_id: u16, payload: &[u8], properties: BasicProperties) -> Poll<Option<()>, io::Error> {
      self.send_and_handle_frames()?;
      self.conn.send_content_frames(channel_id, 60, payload, properties);
      self.send_and_handle_frames()
  }

  fn send_frames(&mut self) -> Result<(), io::Error> {
    //FIXME: find a way to use a future here
    trace!("send frames");
    while let Some(f) = self.conn.next_frame() {
      if let Err(e) = self.send_frame(f) {
        error!("Failed to send frame: {:?}", e);
        return Err(e);
      }
    }
    Ok(())
  }

  fn send_frame(&mut self, frame: Frame) -> Poll<(), io::Error> {
      self.start_send(frame).and_then(|_| self.poll_complete())
  }

  fn handle_frames(&mut self) -> Poll<Option<()>, io::Error> {
    trace!("handle frames");
    for _ in 0..30 {
      if try_ready!(self.poll()).is_none() {
        return Ok(Async::Ready(None));
      }
    }
    self.poll()
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: Sync+Send,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.send_and_handle_frames()?;

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    match transport.poll()? {
      Async::Ready(Some(_)) => {
        if transport.conn.state == ConnectionState::Connected {
          // Upstream had frames available and we're connected, the transport is ready
          Ok(Async::Ready(transport))
        } else {
          // Upstream had frames but we're not yet connected, continue polling
          let poll_ret = transport.poll();
          self.transport = Some(transport);
          poll_ret?;
          Ok(Async::NotReady)
        }
      },
      _ => {
        // Upstream had no frames
        self.transport = Some(transport);
        Ok(Async::NotReady)
      },
    }
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: Sync+Send,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
      trace!("stream poll");
      self.poll_heartbeat()?;
      self.poll_upstream()
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite,
          T: Sync+Send {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}
<commit_msg>futures: log heartbeat sending at debug log level<commit_after>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::cmp;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use std::time::Duration;
use futures::{Async,Poll,Sink,Stream,StartSend,Future,future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use tokio_timer::{Interval,Timer};
use channel::BasicProperties;
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec {
    pub frame_max: u32,
}

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length    = buf.len();
      // Ensure we at least allocate 8192 so that the buffer is big enough for the frame_max
      // negociation. Afterwards, use frame_max if > 8192.
      let frame_max = cmp::max(self.frame_max, 8192) as usize;
      if length < frame_max {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(frame_max - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream: Framed<T,AMQPCodec>,
  heartbeat: Option<Interval>,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: Sync+Send,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(stream: T, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_frame_max(options.frame_max);
    conn.set_heartbeat(options.heartbeat);
    if let Err(e) = conn.connect() {
      let err = format!("Failed to connect: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let codec = AMQPCodec {
      frame_max: conn.configuration.frame_max,
    };
    let mut t = AMQPTransport {
      upstream:  stream.framed(codec),
      heartbeat: None,
      conn:      conn,
    };

    if let Err(e) = t.send_and_handle_frames() {
      let err = format!("Failed to handle frames: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let mut connector = AMQPTransportConnector {
      transport: Some(t),
    };

    trace!("pre-poll");
    if let Err(e) = connector.poll() {
      let err = format!("Failed to handle frames: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }
    trace!("post-poll");

    Box::new(connector)
  }

  pub fn start_heartbeat(&mut self) {
      let heartbeat = self.conn.configuration.heartbeat as u64;
      if heartbeat > 0 {
          self.heartbeat = Some(Timer::default().interval(Duration::from_secs(heartbeat)));
      }
  }

  fn poll_heartbeat(&mut self) -> Result<(), io::Error> {
    if let Some(Ok(Async::Ready(_))) = self.heartbeat.as_mut().map(Interval::poll) {
      debug!("Sending heartbeat");
      if let Err(e) = self.send_frame(Frame::Heartbeat(0)) {
        error!("Failed to send heartbeat: {:?}", e);
        return Err(e);
      }
    }
    Ok(())
  }

  fn poll_upstream(&mut self) -> Poll<Option<()>, io::Error> {
    trace!("poll upstream");
    let value = match self.upstream.poll() {
      Ok(Async::Ready(t)) => t,
      Ok(Async::NotReady) => {
        trace!("upstream poll gave NotReady");
        return Ok(Async::NotReady);
      },
      Err(e) => {
        error!("upstream poll gave error: {:?}", e);
        return Err(From::from(e));
      },
    };

    if let Some(frame) = value {
      trace!("upstream poll gave frame: {:?}", frame);
      if let Err(e) = self.conn.handle_frame(frame) {
        let err = format!("failed to handle frame: {:?}", e);
        return Err(io::Error::new(io::ErrorKind::Other, err));
      }
      self.send_frames()?;
      Ok(Async::Ready(Some(())))
    } else {
      error!("upstream poll gave Ready(None)");
      Ok(Async::Ready(None))
    }
  }

  pub fn send_and_handle_frames(&mut self) -> Poll<Option<()>, io::Error> {
    self.send_frames()?;
    self.handle_frames()
  }

  pub fn send_and_handle_frames_with_payload(&mut self, channel_id: u16, payload: &[u8], properties: BasicProperties) -> Poll<Option<()>, io::Error> {
      self.send_and_handle_frames()?;
      self.conn.send_content_frames(channel_id, 60, payload, properties);
      self.send_and_handle_frames()
  }

  fn send_frames(&mut self) -> Result<(), io::Error> {
    //FIXME: find a way to use a future here
    trace!("send frames");
    while let Some(f) = self.conn.next_frame() {
      if let Err(e) = self.send_frame(f) {
        error!("Failed to send frame: {:?}", e);
        return Err(e);
      }
    }
    Ok(())
  }

  fn send_frame(&mut self, frame: Frame) -> Poll<(), io::Error> {
      self.start_send(frame).and_then(|_| self.poll_complete())
  }

  fn handle_frames(&mut self) -> Poll<Option<()>, io::Error> {
    trace!("handle frames");
    for _ in 0..30 {
      if try_ready!(self.poll()).is_none() {
        return Ok(Async::Ready(None));
      }
    }
    self.poll()
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: Sync+Send,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.send_and_handle_frames()?;

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    match transport.poll()? {
      Async::Ready(Some(_)) => {
        if transport.conn.state == ConnectionState::Connected {
          // Upstream had frames available and we're connected, the transport is ready
          Ok(Async::Ready(transport))
        } else {
          // Upstream had frames but we're not yet connected, continue polling
          let poll_ret = transport.poll();
          self.transport = Some(transport);
          poll_ret?;
          Ok(Async::NotReady)
        }
      },
      _ => {
        // Upstream had no frames
        self.transport = Some(transport);
        Ok(Async::NotReady)
      },
    }
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: Sync+Send,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
      trace!("stream poll");
      self.poll_heartbeat()?;
      self.poll_upstream()
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite,
          T: Sync+Send {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}
<|endoftext|>


<commit_before>use std::cmp;

use lz77::LongestMatch;
use util::{ZOPFLI_CACHE_LENGTH, ZOPFLI_MAX_MATCH, ZOPFLI_MIN_MATCH};

// Cache used by ZopfliFindLongestMatch to remember previously found length/dist
// values.
// This is needed because the squeeze runs will ask these values multiple times for
// the same position.
// Uses large amounts of memory, since it has to remember the distance belonging
// to every possible shorter-than-the-best length (the so called "sublen" array).
pub struct ZopfliLongestMatchCache {
    length: Vec<u16>,
    dist: Vec<u16>,
    sublen: Vec<u8>,
}

impl ZopfliLongestMatchCache {
    pub fn new(blocksize: usize) -> ZopfliLongestMatchCache {
        ZopfliLongestMatchCache {
            /* length > 0 and dist 0 is invalid combination, which indicates on purpose
            that this cache value is not filled in yet. */
            length: vec![1; blocksize],
            dist: vec![0; blocksize],
            /* Rather large amount of memory. */
            sublen: vec![0; ZOPFLI_CACHE_LENGTH * blocksize * 3],
        }
    }

    pub fn length_at(&self, pos: usize) -> u16 {
        self.length[pos]
    }

    pub fn dist_at(&self, pos: usize) -> u16 {
        self.dist[pos]
    }

    pub fn store_length_at(&mut self, pos: usize, val: u16) {
        self.length[pos] = val;
    }

    pub fn store_dist_at(&mut self, pos: usize, val: u16) {
        self.dist[pos] = val;
    }

    /// Returns the length up to which could be stored in the cache.
    pub fn max_sublen(&self, pos: usize) -> u32 {
        let start = ZOPFLI_CACHE_LENGTH * pos * 3;
        if self.sublen[start + 1] == 0 && self.sublen[start + 2] == 0 {
            return 0;  // No sublen cached.
        }
        self.sublen[start + ((ZOPFLI_CACHE_LENGTH - 1) * 3)] as u32 + 3
    }

    /// Stores sublen array in the cache.
    pub fn store_sublen(&mut self, sublen: &[u16], pos: usize, length: usize) {
        if length < 3 {
            return;
        }

        let start = ZOPFLI_CACHE_LENGTH * pos * 3;
        let mut i = 3;
        let mut j = 0;
        let mut bestlength = 0;
        while i <= length {
            if i == length || sublen[i] != sublen[i + 1] {
                self.sublen[start + (j * 3)] = (i - 3) as u8;
                self.sublen[start + (j * 3 + 1)] = sublen[i].wrapping_rem(256) as u8;
                self.sublen[start + (j * 3 + 2)] = ((sublen[i] >> 8)).wrapping_rem(256) as u8;
                bestlength = i as u32;
                j += 1;
                if j >= ZOPFLI_CACHE_LENGTH {
                    break;
                }
            }
            i += 1;
        }

        if j < ZOPFLI_CACHE_LENGTH {
            assert!(bestlength == length as u32);
            self.sublen[start + ((ZOPFLI_CACHE_LENGTH - 1) * 3)] = (bestlength - 3) as u8;
        } else {
            assert!(bestlength <= length as u32);
        }
        assert!(bestlength == self.max_sublen(pos));
    }

    /// Extracts sublen array from the cache.
    pub fn fetch_sublen(&self, pos: usize, length: usize, sublen: &mut [u16]) {
        if length < 3 {
            return;
        }

        let start = ZOPFLI_CACHE_LENGTH * pos * 3;
        let maxlength = self.max_sublen(pos) as usize;
        let mut prevlength = 0;

        for j in 0..ZOPFLI_CACHE_LENGTH {
            let length = self.sublen[start + (j * 3)] as usize + 3;
            let dist = self.sublen[start + (j * 3 + 1)] as u16 + 256 * self.sublen[start + (j * 3 + 2)] as u16;

            let mut i = prevlength;
            while i <= length {
                sublen[i] = dist;
                i += 1;
            }
            if length == maxlength {
                break;
            }
            prevlength = length + 1;
        }
    }
}

pub trait Cache {
    fn try_get(&self, pos: usize, limit: usize, sublen: &mut Option<&mut [u16]>, blockstart: usize) -> LongestMatch;
    fn store(&mut self, pos: usize, limit: usize, sublen: &mut Option<&mut [u16]>, distance: u16, length: u16, blockstart: usize);
}

pub struct NoCache;

impl Cache for NoCache {
    fn try_get(&self, _: usize, limit: usize, _: &mut Option<&mut [u16]>, _: usize) -> LongestMatch {
        LongestMatch::new(limit)
    }

    fn store(&mut self, _: usize, _: usize, _: &mut Option<&mut [u16]>, _: u16, _: u16, _: usize) {
        // Nowhere to store
    }
}

impl Cache for ZopfliLongestMatchCache {
    fn try_get(&self, pos: usize, mut limit: usize, sublen: &mut Option<&mut [u16]>, blockstart: usize) -> LongestMatch {
        let mut longest_match = LongestMatch::new(limit);

        /* The LMC cache starts at the beginning of the block rather than the
        beginning of the whole array. */
        let lmcpos = pos - blockstart;

        /* Length > 0 and dist 0 is invalid combination, which indicates on purpose
        that this cache value is not filled in yet. */
        let length_lmcpos = self.length_at(lmcpos);
        let dist_lmcpos = self.dist_at(lmcpos);
        let cache_available = length_lmcpos == 0 || dist_lmcpos != 0;
        let max_sublen = self.max_sublen(lmcpos);
        let limit_ok_for_cache = limit == ZOPFLI_MAX_MATCH || length_lmcpos <= limit as u16 || ( sublen.is_some() && max_sublen >= limit as u32);

        if limit_ok_for_cache && cache_available {
            if sublen.is_none() || length_lmcpos as u32 <= max_sublen {
                let length = cmp::min(length_lmcpos, limit as u16);
                let distance;
                if let &mut Some(ref mut subl) = sublen {
                    self.fetch_sublen(lmcpos, length as usize, subl);
                    distance = subl[length as usize];

                    if limit == ZOPFLI_MAX_MATCH && length >= ZOPFLI_MIN_MATCH as u16 {
                        assert!(subl[length as usize] == dist_lmcpos);
                    }
                } else {
                    distance = dist_lmcpos;
                }
                longest_match.distance = distance;
                longest_match.length = length;
                longest_match.from_cache = true;
                return longest_match;
            }
            /* Can't use much of the cache, since the "sublens" need to be calculated,
            but at least we already know when to stop. */
            limit = length_lmcpos as usize;
            longest_match.limit = limit;
        }

        longest_match
    }

    fn store(&mut self, pos: usize, limit: usize, sublen: &mut Option<&mut [u16]>, distance: u16, length: u16, blockstart: usize) {
        if let &mut Some(ref mut subl) = sublen {
            /* Length > 0 and dist 0 is invalid combination, which indicates on purpose
            that this cache value is not filled in yet. */
            let lmcpos = pos - blockstart;
            let cache_available = self.length_at(lmcpos) == 0 || self.dist_at(lmcpos) != 0;

            if limit == ZOPFLI_MAX_MATCH && !cache_available {
                assert!(self.length_at(lmcpos) == 1 && self.dist_at(lmcpos) == 0);
                if length < ZOPFLI_MIN_MATCH as u16 {
                    self.store_dist_at(lmcpos, 0);
                    self.store_length_at(lmcpos, 0);
                } else {
                    self.store_dist_at(lmcpos, distance);
                    self.store_length_at(lmcpos, length);
                }
                assert!(!(self.length_at(lmcpos) == 1 && self.dist_at(lmcpos) == 0));
                self.store_sublen(subl, lmcpos, length as usize);
            }
        }
    }
}
<commit_msg>Remove unneeded pub<commit_after>use std::cmp;

use lz77::LongestMatch;
use util::{ZOPFLI_CACHE_LENGTH, ZOPFLI_MAX_MATCH, ZOPFLI_MIN_MATCH};

// Cache used by ZopfliFindLongestMatch to remember previously found length/dist
// values.
// This is needed because the squeeze runs will ask these values multiple times for
// the same position.
// Uses large amounts of memory, since it has to remember the distance belonging
// to every possible shorter-than-the-best length (the so called "sublen" array).
pub struct ZopfliLongestMatchCache {
    length: Vec<u16>,
    dist: Vec<u16>,
    sublen: Vec<u8>,
}

impl ZopfliLongestMatchCache {
    pub fn new(blocksize: usize) -> ZopfliLongestMatchCache {
        ZopfliLongestMatchCache {
            /* length > 0 and dist 0 is invalid combination, which indicates on purpose
            that this cache value is not filled in yet. */
            length: vec![1; blocksize],
            dist: vec![0; blocksize],
            /* Rather large amount of memory. */
            sublen: vec![0; ZOPFLI_CACHE_LENGTH * blocksize * 3],
        }
    }

    pub fn length_at(&self, pos: usize) -> u16 {
        self.length[pos]
    }

    pub fn dist_at(&self, pos: usize) -> u16 {
        self.dist[pos]
    }

    pub fn store_length_at(&mut self, pos: usize, val: u16) {
        self.length[pos] = val;
    }

    pub fn store_dist_at(&mut self, pos: usize, val: u16) {
        self.dist[pos] = val;
    }

    /// Returns the length up to which could be stored in the cache.
    pub fn max_sublen(&self, pos: usize) -> u32 {
        let start = ZOPFLI_CACHE_LENGTH * pos * 3;
        if self.sublen[start + 1] == 0 && self.sublen[start + 2] == 0 {
            return 0;  // No sublen cached.
        }
        self.sublen[start + ((ZOPFLI_CACHE_LENGTH - 1) * 3)] as u32 + 3
    }

    /// Stores sublen array in the cache.
    pub fn store_sublen(&mut self, sublen: &[u16], pos: usize, length: usize) {
        if length < 3 {
            return;
        }

        let start = ZOPFLI_CACHE_LENGTH * pos * 3;
        let mut i = 3;
        let mut j = 0;
        let mut bestlength = 0;
        while i <= length {
            if i == length || sublen[i] != sublen[i + 1] {
                self.sublen[start + (j * 3)] = (i - 3) as u8;
                self.sublen[start + (j * 3 + 1)] = sublen[i].wrapping_rem(256) as u8;
                self.sublen[start + (j * 3 + 2)] = ((sublen[i] >> 8)).wrapping_rem(256) as u8;
                bestlength = i as u32;
                j += 1;
                if j >= ZOPFLI_CACHE_LENGTH {
                    break;
                }
            }
            i += 1;
        }

        if j < ZOPFLI_CACHE_LENGTH {
            assert!(bestlength == length as u32);
            self.sublen[start + ((ZOPFLI_CACHE_LENGTH - 1) * 3)] = (bestlength - 3) as u8;
        } else {
            assert!(bestlength <= length as u32);
        }
        assert!(bestlength == self.max_sublen(pos));
    }

    /// Extracts sublen array from the cache.
    fn fetch_sublen(&self, pos: usize, length: usize, sublen: &mut [u16]) {
        if length < 3 {
            return;
        }

        let start = ZOPFLI_CACHE_LENGTH * pos * 3;
        let maxlength = self.max_sublen(pos) as usize;
        let mut prevlength = 0;

        for j in 0..ZOPFLI_CACHE_LENGTH {
            let length = self.sublen[start + (j * 3)] as usize + 3;
            let dist = self.sublen[start + (j * 3 + 1)] as u16 + 256 * self.sublen[start + (j * 3 + 2)] as u16;

            let mut i = prevlength;
            while i <= length {
                sublen[i] = dist;
                i += 1;
            }
            if length == maxlength {
                break;
            }
            prevlength = length + 1;
        }
    }
}

pub trait Cache {
    fn try_get(&self, pos: usize, limit: usize, sublen: &mut Option<&mut [u16]>, blockstart: usize) -> LongestMatch;
    fn store(&mut self, pos: usize, limit: usize, sublen: &mut Option<&mut [u16]>, distance: u16, length: u16, blockstart: usize);
}

pub struct NoCache;

impl Cache for NoCache {
    fn try_get(&self, _: usize, limit: usize, _: &mut Option<&mut [u16]>, _: usize) -> LongestMatch {
        LongestMatch::new(limit)
    }

    fn store(&mut self, _: usize, _: usize, _: &mut Option<&mut [u16]>, _: u16, _: u16, _: usize) {
        // Nowhere to store
    }
}

impl Cache for ZopfliLongestMatchCache {
    fn try_get(&self, pos: usize, mut limit: usize, sublen: &mut Option<&mut [u16]>, blockstart: usize) -> LongestMatch {
        let mut longest_match = LongestMatch::new(limit);

        /* The LMC cache starts at the beginning of the block rather than the
        beginning of the whole array. */
        let lmcpos = pos - blockstart;

        /* Length > 0 and dist 0 is invalid combination, which indicates on purpose
        that this cache value is not filled in yet. */
        let length_lmcpos = self.length_at(lmcpos);
        let dist_lmcpos = self.dist_at(lmcpos);
        let cache_available = length_lmcpos == 0 || dist_lmcpos != 0;
        let max_sublen = self.max_sublen(lmcpos);
        let limit_ok_for_cache = limit == ZOPFLI_MAX_MATCH || length_lmcpos <= limit as u16 || ( sublen.is_some() && max_sublen >= limit as u32);

        if limit_ok_for_cache && cache_available {
            if sublen.is_none() || length_lmcpos as u32 <= max_sublen {
                let length = cmp::min(length_lmcpos, limit as u16);
                let distance;
                if let &mut Some(ref mut subl) = sublen {
                    self.fetch_sublen(lmcpos, length as usize, subl);
                    distance = subl[length as usize];

                    if limit == ZOPFLI_MAX_MATCH && length >= ZOPFLI_MIN_MATCH as u16 {
                        assert!(subl[length as usize] == dist_lmcpos);
                    }
                } else {
                    distance = dist_lmcpos;
                }
                longest_match.distance = distance;
                longest_match.length = length;
                longest_match.from_cache = true;
                return longest_match;
            }
            /* Can't use much of the cache, since the "sublens" need to be calculated,
            but at least we already know when to stop. */
            limit = length_lmcpos as usize;
            longest_match.limit = limit;
        }

        longest_match
    }

    fn store(&mut self, pos: usize, limit: usize, sublen: &mut Option<&mut [u16]>, distance: u16, length: u16, blockstart: usize) {
        if let &mut Some(ref mut subl) = sublen {
            /* Length > 0 and dist 0 is invalid combination, which indicates on purpose
            that this cache value is not filled in yet. */
            let lmcpos = pos - blockstart;
            let cache_available = self.length_at(lmcpos) == 0 || self.dist_at(lmcpos) != 0;

            if limit == ZOPFLI_MAX_MATCH && !cache_available {
                assert!(self.length_at(lmcpos) == 1 && self.dist_at(lmcpos) == 0);
                if length < ZOPFLI_MIN_MATCH as u16 {
                    self.store_dist_at(lmcpos, 0);
                    self.store_length_at(lmcpos, 0);
                } else {
                    self.store_dist_at(lmcpos, distance);
                    self.store_length_at(lmcpos, length);
                }
                assert!(!(self.length_at(lmcpos) == 1 && self.dist_at(lmcpos) == 0));
                self.store_sublen(subl, lmcpos, length as usize);
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    env,
    ffi::OsStr,
    fs,
    io::{Cursor, Read, Seek},
    iter,
    path::{Path, PathBuf},
};

use app_dirs::{get_app_root, AppDataType};
use log::debug;
use reqwest::{blocking::Client, Proxy};
use std::time::{Duration, SystemTime};
use walkdir::{DirEntry, WalkDir};
use zip::ZipArchive;

use crate::{
    error::TealdeerError::{self, CacheError, UpdateError},
    types::{OsType, PathSource},
};

static CACHE_DIR_ENV_VAR: &str = "TEALDEER_CACHE_DIR";

pub static TLDR_PAGES_DIR: &str = "tldr-pages";
static TLDR_OLD_PAGES_DIR: &str = "tldr-master";

#[derive(Debug)]
pub struct Cache {
    url: String,
    os: OsType,
}

#[derive(Debug)]
pub struct PageLookupResult {
    page_path: PathBuf,
    patch_path: Option<PathBuf>,
}

impl PageLookupResult {
    pub fn with_page(page_path: PathBuf) -> Self {
        Self {
            page_path,
            patch_path: None,
        }
    }

    pub fn with_optional_patch(mut self, patch_path: Option<PathBuf>) -> Self {
        self.patch_path = patch_path;
        self
    }

    pub fn paths(&self) -> impl Iterator<Item = &Path> {
        iter::once(self.page_path.as_path()).chain(self.patch_path.as_deref())
    }
}

pub enum CacheFreshness {
    /// The cache is still fresh (less than MAX_CACHE_AGE old)
    Fresh,
    /// The cache is stale and should be updated
    Stale(Duration),
    /// The cache is missing
    Missing,
}

impl Cache {
    pub fn new<S>(url: S, os: OsType) -> Self
    where
        S: Into<String>,
    {
        Self {
            url: url.into(),
            os,
        }
    }

    /// Return the path to the cache directory.
    pub fn get_cache_dir() -> Result<(PathBuf, PathSource), TealdeerError> {
        // Allow overriding the cache directory by setting the env variable.
        if let Ok(value) = env::var(CACHE_DIR_ENV_VAR) {
            let path = PathBuf::from(value);
            let (path_exists, path_is_dir) = path
                .metadata()
                .map_or((false, false), |md| (true, md.is_dir()));
            if path_exists && !path_is_dir {
                return Err(CacheError(format!(
                    "Path specified by ${} is not a directory.",
                    CACHE_DIR_ENV_VAR
                )));
            }
            if !path_exists {
                // Try to create the complete directory path.
                fs::create_dir_all(&path).map_err(|_| {
                    CacheError(format!(
                        "Directory path specified by ${} cannot be created.",
                        CACHE_DIR_ENV_VAR
                    ))
                })?;
                eprintln!(
                    "Successfully created cache directory path `{}`.",
                    path.to_str().unwrap()
                );
            }
            return Ok((path, PathSource::EnvVar));
        };

        // Otherwise, fall back to user cache directory.
        match get_app_root(AppDataType::UserCache, &crate::APP_INFO) {
            Ok(dirs) => Ok((dirs, PathSource::OsConvention)),
            Err(_) => Err(CacheError(
                "Could not determine user cache directory.".into(),
            )),
        }
    }

    /// Download the archive
    fn download(&self) -> Result<Vec<u8>, TealdeerError> {
        let mut builder = Client::builder();
        if let Ok(ref host) = env::var("HTTP_PROXY") {
            if let Ok(proxy) = Proxy::http(host) {
                builder = builder.proxy(proxy);
            }
        }
        if let Ok(ref host) = env::var("HTTPS_PROXY") {
            if let Ok(proxy) = Proxy::https(host) {
                builder = builder.proxy(proxy);
            }
        }
        let client = builder.build().unwrap_or_else(|_| Client::new());
        let mut resp = client.get(&self.url).send()?;
        let mut buf: Vec<u8> = vec![];
        let bytes_downloaded = resp.copy_to(&mut buf)?;
        debug!("{} bytes downloaded", bytes_downloaded);
        Ok(buf)
    }

    /// Decompress and open the archive
    fn decompress<R: Read + Seek>(reader: R) -> ZipArchive<R> {
        ZipArchive::new(reader).unwrap()
    }

    /// Update the pages cache.
    pub fn update(&self) -> Result<(), TealdeerError> {
        // First, download the compressed data
        let bytes: Vec<u8> = self.download()?;

        // Decompress the response body into an `Archive`
        let mut archive = Self::decompress(Cursor::new(bytes));

        // Determine paths
        let (cache_dir, _) = Self::get_cache_dir()?;
        let pages_dir = cache_dir.join(TLDR_PAGES_DIR);

        // Make sure that cache directory exists
        debug!("Ensure cache directory {:?} exists", &cache_dir);
        fs::create_dir_all(&cache_dir)
            .map_err(|e| UpdateError(format!("Could not create cache directory: {}", e)))?;

        // Clear cache directory
        // Note: This is not the best solution. Ideally we would download the
        // archive to a temporary directory and then swap the two directories.
        // But renaming a directory doesn't work across filesystems and Rust
        // does not yet offer a recursive directory copying function. So for
        // now, we'll use this approach.
        Self::clear()?;

        // Extract archive
        archive
            .extract(&pages_dir)
            .map_err(|e| UpdateError(format!("Could not unpack compressed data: {}", e)))?;

        Ok(())
    }

    /// Return the duration since the cache directory was last modified.
    pub fn last_update() -> Option<Duration> {
        if let Ok((cache_dir, _)) = Self::get_cache_dir() {
            if let Ok(metadata) = fs::metadata(cache_dir.join(TLDR_PAGES_DIR)) {
                if let Ok(mtime) = metadata.modified() {
                    let now = SystemTime::now();
                    return now.duration_since(mtime).ok();
                };
            };
        };
        None
    }

    /// Return the freshness of the cache (fresh, stale or missing).
    pub fn freshness() -> CacheFreshness {
        match Cache::last_update() {
            Some(ago) if ago > crate::config::MAX_CACHE_AGE => CacheFreshness::Stale(ago),
            Some(_) => CacheFreshness::Fresh,
            None => CacheFreshness::Missing,
        }
    }

    /// Return the platform directory.
    fn get_platform_dir(&self) -> Option<&'static str> {
        match self.os {
            OsType::Linux => Some("linux"),
            OsType::OsX => Some("osx"),
            OsType::SunOs => Some("sunos"),
            OsType::Windows => Some("windows"),
            OsType::Other => None,
        }
    }

    /// Check for pages for a given platform in one of the given languages.
    fn find_page_for_platform(
        page_name: &str,
        cache_dir: &Path,
        platform: &str,
        language_dirs: &[String],
    ) -> Option<PathBuf> {
        language_dirs
            .iter()
            .map(|lang_dir| cache_dir.join(lang_dir).join(platform).join(page_name))
            .find(|path| path.exists() && path.is_file())
    }

    /// Look up custom patch (<name>.patch). If it exists, store it in a variable.
    fn find_patch(patch_name: &str, custom_pages_dir: Option<&Path>) -> Option<PathBuf> {
        custom_pages_dir
            .map(|custom_dir| custom_dir.join(patch_name))
            .filter(|path| path.exists() && path.is_file())
    }

    /// Search for a page and return the path to it.
    pub fn find_page(
        &self,
        name: &str,
        languages: &[String],
        custom_pages_dir: Option<&Path>,
    ) -> Option<PageLookupResult> {
        let page_filename = format!("{}.md", name);
        let patch_filename = format!("{}.patch", name);
        let custom_filename = format!("{}.page", name);

        // Get cache dir
        let cache_dir = match Self::get_cache_dir() {
            Ok((cache_dir, _)) => cache_dir.join(TLDR_PAGES_DIR),
            Err(e) => {
                log::error!("Could not get cache directory: {}", e);
                return None;
            }
        };

        let lang_dirs: Vec<String> = languages
            .iter()
            .map(|lang| {
                if lang == "en" {
                    String::from("pages")
                } else {
                    format!("pages.{}", lang)
                }
            })
            .collect();

        // Look up custom page (<name>.page). If it exists, return it directly
        if let Some(config_dir) = custom_pages_dir {
            let custom_page = config_dir.join(custom_filename);
            if custom_page.exists() && custom_page.is_file() {
                return Some(PageLookupResult::with_page(custom_page));
            }
        }

        let patch_path = Self::find_patch(&patch_filename, custom_pages_dir);

        // Try to find a platform specific path next, append custom patch to it.
        if let Some(pf) = self.get_platform_dir() {
            if let Some(page) =
                Self::find_page_for_platform(&page_filename, &cache_dir, pf, &lang_dirs)
            {
                return Some(PageLookupResult::with_page(page).with_optional_patch(patch_path));
            }
        }

        // Did not find platform specific results, fall back to "common"
        Self::find_page_for_platform(&page_filename, &cache_dir, "common", &lang_dirs)
            .map(|page| PageLookupResult::with_page(page).with_optional_patch(patch_path))
    }

    /// Return the available pages.
    pub fn list_pages(&self) -> Result<Vec<String>, TealdeerError> {
        // Determine platforms directory and platform
        let (cache_dir, _) = Self::get_cache_dir()?;
        let platforms_dir = cache_dir.join(TLDR_PAGES_DIR).join("pages");
        let platform_dir = self.get_platform_dir();

        // Closure that allows the WalkDir instance to traverse platform
        // specific and common page directories, but not others.
        let should_walk = |entry: &DirEntry| -> bool {
            let file_type = entry.file_type();
            let file_name = match entry.file_name().to_str() {
                Some(name) => name,
                None => return false,
            };
            if file_type.is_dir() {
                if file_name == "common" {
                    return true;
                }
                if let Some(platform) = platform_dir {
                    return file_name == platform;
                }
            } else if file_type.is_file() {
                return true;
            }
            false
        };

        // Recursively walk through common and (if applicable) platform specific directory
        let mut pages = WalkDir::new(platforms_dir)
            .min_depth(1) // Skip root directory
            .into_iter()
            .filter_entry(|e| should_walk(e)) // Filter out pages for other architectures
            .filter_map(Result::ok) // Convert results to options, filter out errors
            .filter_map(|e| {
                let path = e.path();
                let extension = &path.extension().and_then(OsStr::to_str).unwrap_or("");
                if e.file_type().is_file() && extension == &"md" {
                    path.file_stem()
                        .and_then(|stem| stem.to_str().map(|s| s.into()))
                } else {
                    None
                }
            })
            .collect::<Vec<String>>();
        pages.sort();
        pages.dedup();
        Ok(pages)
    }

    /// Delete the cache directory.
    pub fn clear() -> Result<(), TealdeerError> {
        let (path, _) = Self::get_cache_dir()?;
        if path.exists() && path.is_dir() {
            fs::remove_dir_all(&path).map_err(|_| {
                CacheError(format!(
                    "Could not remove cache directory ({}).",
                    path.display()
                ))
            })?;
        } else if path.exists() {
            return Err(CacheError(format!(
                "Cache path ({}) is not a directory.",
                path.display()
            )));
        } else {
            return Err(CacheError(format!(
                "Cache path ({}) does not exist.",
                path.display()
            )));
        };
        Ok(())
    }
}

/// Unit Tests for cache module
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_page_lookup_result_iter_with_patch() {
        let lookup = PageLookupResult::with_page(PathBuf::from("test.page"))
            .with_optional_patch(Some(PathBuf::from("test.patch")));
        let mut iter = lookup.paths();
        assert_eq!(iter.next(), Some(Path::new("test.page")));
        assert_eq!(iter.next(), Some(Path::new("test.patch")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_page_lookup_result_iter_no_patch() {
        let lookup = PageLookupResult::with_page(PathBuf::from("test.page"));
        let mut iter = lookup.paths();
        assert_eq!(iter.next(), Some(Path::new("test.page")));
        assert_eq!(iter.next(), None);
    }
}
<commit_msg>When clearing cache, only delete pages, not entire cache dir<commit_after>use std::{
    env,
    ffi::OsStr,
    fs,
    io::{Cursor, Read, Seek},
    iter,
    path::{Path, PathBuf},
};

use app_dirs::{get_app_root, AppDataType};
use log::debug;
use reqwest::{blocking::Client, Proxy};
use std::time::{Duration, SystemTime};
use walkdir::{DirEntry, WalkDir};
use zip::ZipArchive;

use crate::{
    error::TealdeerError::{self, CacheError, UpdateError},
    types::{OsType, PathSource},
};

static CACHE_DIR_ENV_VAR: &str = "TEALDEER_CACHE_DIR";

pub static TLDR_PAGES_DIR: &str = "tldr-pages";
static TLDR_OLD_PAGES_DIR: &str = "tldr-master";

#[derive(Debug)]
pub struct Cache {
    url: String,
    os: OsType,
}

#[derive(Debug)]
pub struct PageLookupResult {
    page_path: PathBuf,
    patch_path: Option<PathBuf>,
}

impl PageLookupResult {
    pub fn with_page(page_path: PathBuf) -> Self {
        Self {
            page_path,
            patch_path: None,
        }
    }

    pub fn with_optional_patch(mut self, patch_path: Option<PathBuf>) -> Self {
        self.patch_path = patch_path;
        self
    }

    pub fn paths(&self) -> impl Iterator<Item = &Path> {
        iter::once(self.page_path.as_path()).chain(self.patch_path.as_deref())
    }
}

pub enum CacheFreshness {
    /// The cache is still fresh (less than MAX_CACHE_AGE old)
    Fresh,
    /// The cache is stale and should be updated
    Stale(Duration),
    /// The cache is missing
    Missing,
}

impl Cache {
    pub fn new<S>(url: S, os: OsType) -> Self
    where
        S: Into<String>,
    {
        Self {
            url: url.into(),
            os,
        }
    }

    /// Return the path to the cache directory.
    pub fn get_cache_dir() -> Result<(PathBuf, PathSource), TealdeerError> {
        // Allow overriding the cache directory by setting the env variable.
        if let Ok(value) = env::var(CACHE_DIR_ENV_VAR) {
            let path = PathBuf::from(value);
            let (path_exists, path_is_dir) = path
                .metadata()
                .map_or((false, false), |md| (true, md.is_dir()));
            if path_exists && !path_is_dir {
                return Err(CacheError(format!(
                    "Path specified by ${} is not a directory.",
                    CACHE_DIR_ENV_VAR
                )));
            }
            if !path_exists {
                // Try to create the complete directory path.
                fs::create_dir_all(&path).map_err(|_| {
                    CacheError(format!(
                        "Directory path specified by ${} cannot be created.",
                        CACHE_DIR_ENV_VAR
                    ))
                })?;
                eprintln!(
                    "Successfully created cache directory path `{}`.",
                    path.to_str().unwrap()
                );
            }
            return Ok((path, PathSource::EnvVar));
        };

        // Otherwise, fall back to user cache directory.
        match get_app_root(AppDataType::UserCache, &crate::APP_INFO) {
            Ok(dirs) => Ok((dirs, PathSource::OsConvention)),
            Err(_) => Err(CacheError(
                "Could not determine user cache directory.".into(),
            )),
        }
    }

    /// Download the archive
    fn download(&self) -> Result<Vec<u8>, TealdeerError> {
        let mut builder = Client::builder();
        if let Ok(ref host) = env::var("HTTP_PROXY") {
            if let Ok(proxy) = Proxy::http(host) {
                builder = builder.proxy(proxy);
            }
        }
        if let Ok(ref host) = env::var("HTTPS_PROXY") {
            if let Ok(proxy) = Proxy::https(host) {
                builder = builder.proxy(proxy);
            }
        }
        let client = builder.build().unwrap_or_else(|_| Client::new());
        let mut resp = client.get(&self.url).send()?;
        let mut buf: Vec<u8> = vec![];
        let bytes_downloaded = resp.copy_to(&mut buf)?;
        debug!("{} bytes downloaded", bytes_downloaded);
        Ok(buf)
    }

    /// Decompress and open the archive
    fn decompress<R: Read + Seek>(reader: R) -> ZipArchive<R> {
        ZipArchive::new(reader).unwrap()
    }

    /// Update the pages cache.
    pub fn update(&self) -> Result<(), TealdeerError> {
        // First, download the compressed data
        let bytes: Vec<u8> = self.download()?;

        // Decompress the response body into an `Archive`
        let mut archive = Self::decompress(Cursor::new(bytes));

        // Determine paths
        let (cache_dir, _) = Self::get_cache_dir()?;
        let pages_dir = cache_dir.join(TLDR_PAGES_DIR);

        // Make sure that cache directory exists
        debug!("Ensure cache directory {:?} exists", &cache_dir);
        fs::create_dir_all(&cache_dir)
            .map_err(|e| UpdateError(format!("Could not create cache directory: {}", e)))?;

        // Clear cache directory
        // Note: This is not the best solution. Ideally we would download the
        // archive to a temporary directory and then swap the two directories.
        // But renaming a directory doesn't work across filesystems and Rust
        // does not yet offer a recursive directory copying function. So for
        // now, we'll use this approach.
        Self::clear()?;

        // Extract archive
        archive
            .extract(&pages_dir)
            .map_err(|e| UpdateError(format!("Could not unpack compressed data: {}", e)))?;

        Ok(())
    }

    /// Return the duration since the cache directory was last modified.
    pub fn last_update() -> Option<Duration> {
        if let Ok((cache_dir, _)) = Self::get_cache_dir() {
            if let Ok(metadata) = fs::metadata(cache_dir.join(TLDR_PAGES_DIR)) {
                if let Ok(mtime) = metadata.modified() {
                    let now = SystemTime::now();
                    return now.duration_since(mtime).ok();
                };
            };
        };
        None
    }

    /// Return the freshness of the cache (fresh, stale or missing).
    pub fn freshness() -> CacheFreshness {
        match Cache::last_update() {
            Some(ago) if ago > crate::config::MAX_CACHE_AGE => CacheFreshness::Stale(ago),
            Some(_) => CacheFreshness::Fresh,
            None => CacheFreshness::Missing,
        }
    }

    /// Return the platform directory.
    fn get_platform_dir(&self) -> Option<&'static str> {
        match self.os {
            OsType::Linux => Some("linux"),
            OsType::OsX => Some("osx"),
            OsType::SunOs => Some("sunos"),
            OsType::Windows => Some("windows"),
            OsType::Other => None,
        }
    }

    /// Check for pages for a given platform in one of the given languages.
    fn find_page_for_platform(
        page_name: &str,
        cache_dir: &Path,
        platform: &str,
        language_dirs: &[String],
    ) -> Option<PathBuf> {
        language_dirs
            .iter()
            .map(|lang_dir| cache_dir.join(lang_dir).join(platform).join(page_name))
            .find(|path| path.exists() && path.is_file())
    }

    /// Look up custom patch (<name>.patch). If it exists, store it in a variable.
    fn find_patch(patch_name: &str, custom_pages_dir: Option<&Path>) -> Option<PathBuf> {
        custom_pages_dir
            .map(|custom_dir| custom_dir.join(patch_name))
            .filter(|path| path.exists() && path.is_file())
    }

    /// Search for a page and return the path to it.
    pub fn find_page(
        &self,
        name: &str,
        languages: &[String],
        custom_pages_dir: Option<&Path>,
    ) -> Option<PageLookupResult> {
        let page_filename = format!("{}.md", name);
        let patch_filename = format!("{}.patch", name);
        let custom_filename = format!("{}.page", name);

        // Get cache dir
        let cache_dir = match Self::get_cache_dir() {
            Ok((cache_dir, _)) => cache_dir.join(TLDR_PAGES_DIR),
            Err(e) => {
                log::error!("Could not get cache directory: {}", e);
                return None;
            }
        };

        let lang_dirs: Vec<String> = languages
            .iter()
            .map(|lang| {
                if lang == "en" {
                    String::from("pages")
                } else {
                    format!("pages.{}", lang)
                }
            })
            .collect();

        // Look up custom page (<name>.page). If it exists, return it directly
        if let Some(config_dir) = custom_pages_dir {
            let custom_page = config_dir.join(custom_filename);
            if custom_page.exists() && custom_page.is_file() {
                return Some(PageLookupResult::with_page(custom_page));
            }
        }

        let patch_path = Self::find_patch(&patch_filename, custom_pages_dir);

        // Try to find a platform specific path next, append custom patch to it.
        if let Some(pf) = self.get_platform_dir() {
            if let Some(page) =
                Self::find_page_for_platform(&page_filename, &cache_dir, pf, &lang_dirs)
            {
                return Some(PageLookupResult::with_page(page).with_optional_patch(patch_path));
            }
        }

        // Did not find platform specific results, fall back to "common"
        Self::find_page_for_platform(&page_filename, &cache_dir, "common", &lang_dirs)
            .map(|page| PageLookupResult::with_page(page).with_optional_patch(patch_path))
    }

    /// Return the available pages.
    pub fn list_pages(&self) -> Result<Vec<String>, TealdeerError> {
        // Determine platforms directory and platform
        let (cache_dir, _) = Self::get_cache_dir()?;
        let platforms_dir = cache_dir.join(TLDR_PAGES_DIR).join("pages");
        let platform_dir = self.get_platform_dir();

        // Closure that allows the WalkDir instance to traverse platform
        // specific and common page directories, but not others.
        let should_walk = |entry: &DirEntry| -> bool {
            let file_type = entry.file_type();
            let file_name = match entry.file_name().to_str() {
                Some(name) => name,
                None => return false,
            };
            if file_type.is_dir() {
                if file_name == "common" {
                    return true;
                }
                if let Some(platform) = platform_dir {
                    return file_name == platform;
                }
            } else if file_type.is_file() {
                return true;
            }
            false
        };

        // Recursively walk through common and (if applicable) platform specific directory
        let mut pages = WalkDir::new(platforms_dir)
            .min_depth(1) // Skip root directory
            .into_iter()
            .filter_entry(|e| should_walk(e)) // Filter out pages for other architectures
            .filter_map(Result::ok) // Convert results to options, filter out errors
            .filter_map(|e| {
                let path = e.path();
                let extension = &path.extension().and_then(OsStr::to_str).unwrap_or("");
                if e.file_type().is_file() && extension == &"md" {
                    path.file_stem()
                        .and_then(|stem| stem.to_str().map(|s| s.into()))
                } else {
                    None
                }
            })
            .collect::<Vec<String>>();
        pages.sort();
        pages.dedup();
        Ok(pages)
    }

    /// Delete the cache directory.
    pub fn clear() -> Result<(), TealdeerError> {
        let (path, _) = Self::get_cache_dir()?;
        if path.exists() && path.is_dir() {
            // Delete old tldr-pages cache location as well if present
            // TODO: To be removed in the future
            for pages_dir_name in [TLDR_PAGES_DIR, TLDR_OLD_PAGES_DIR] {
                let pages_dir = path.join(pages_dir_name);

                if pages_dir.exists() {
                    fs::remove_dir_all(&pages_dir).map_err(|e| {
                        CacheError(format!(
                            "Could not remove cache directory ({}): {}",
                            pages_dir.display(),
                            e
                        ))
                    })?;
                }
            }
        } else if path.exists() {
            return Err(CacheError(format!(
                "Cache path ({}) is not a directory.",
                path.display()
            )));
        } else {
            return Err(CacheError(format!(
                "Cache path ({}) does not exist.",
                path.display()
            )));
        };
        Ok(())
    }
}

/// Unit Tests for cache module
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_page_lookup_result_iter_with_patch() {
        let lookup = PageLookupResult::with_page(PathBuf::from("test.page"))
            .with_optional_patch(Some(PathBuf::from("test.patch")));
        let mut iter = lookup.paths();
        assert_eq!(iter.next(), Some(Path::new("test.page")));
        assert_eq!(iter.next(), Some(Path::new("test.patch")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_page_lookup_result_iter_no_patch() {
        let lookup = PageLookupResult::with_page(PathBuf::from("test.page"));
        let mut iter = lookup.paths();
        assert_eq!(iter.next(), Some(Path::new("test.page")));
        assert_eq!(iter.next(), None);
    }
}
<|endoftext|>


<commit_before>//! Internal-use macros.

/// Reads n bytes into a buffer, or returns a `ReadError`.
macro_rules! try_read_n {
    ($reader:expr, $data:expr, $n:expr) => {{
        let took = try!($reader.take($n).read(&mut $data));
        if took != $n {
            return Err(Error::ReadError(format!("Tried to take {} bytes, only took {}", $n, took)));
        }
    }};
}

<commit_msg>Export the macro<commit_after>//! Internal-use macros.

/// Reads n bytes into a buffer, or returns a `ReadError`.
#[macro_export]
macro_rules! try_read_n {
    ($reader:expr, $data:expr, $n:expr) => {{
        let took = try!($reader.take($n).read(&mut $data));
        if took != $n {
            return Err(Error::ReadError(format!("Tried to take {} bytes, only took {}", $n, took)));
        }
    }};
}

<|endoftext|>


<commit_before>//! Helper macros for implementing NBT format encoding and decoding via `serde`.

macro_rules! return_expr_for_serialized_types_method {
    ($expr:expr, $func:ident($($arg:ty),*)) => {
        #[inline]
        fn $func(self, $(_: $arg,)*)
                 -> ::std::result::Result<Self::Ok, Self::Error>
        {
            $expr
        }
    };
    ($expr:expr, $func:ident($($arg:ty),*), result: $result:path) => {
        #[inline]
        fn $func(self, $(_: $arg,)*)
                 -> ::std::result::Result<$result, Self::Error>
        {
            $expr
        }
    };
    ($expr:expr, $func:ident($($arg:ty),*), where: $where:path) => {
        #[inline]
        fn $func<__T: ?Sized>(self, $(_: $arg,)*)
                              -> ::std::result::Result<Self::Ok, Self::Error>
            where __T: $where
        {
            $expr
        }
    };
}

macro_rules! return_expr_for_serialized_types_helper {
    ($expr:expr, bool) => {
        return_expr_for_serialized_types_method! {$expr, serialize_bool(bool)}
    };
    ($expr:expr, i8) => {
        return_expr_for_serialized_types_method! {$expr, serialize_i8(i8)}
    };
    ($expr:expr, i16) => {
        return_expr_for_serialized_types_method! {$expr, serialize_i16(i16)}
    };
    ($expr:expr, i32) => {
        return_expr_for_serialized_types_method! {$expr, serialize_i32(i32)}
    };
    ($expr:expr, i64) => {
        return_expr_for_serialized_types_method! {$expr, serialize_i64(i64)}
    };
    ($expr:expr, u8) => {
        return_expr_for_serialized_types_method! {$expr, serialize_u8(u8)}
    };
    ($expr:expr, u16) => {
        return_expr_for_serialized_types_method! {$expr, serialize_u16(u16)}
    };
    ($expr:expr, u32) => {
        return_expr_for_serialized_types_method! {$expr, serialize_u32(u32)}
    };
    ($expr:expr, u64) => {
        return_expr_for_serialized_types_method! {$expr, serialize_u64(u64)}
    };
    ($expr:expr, f32) => {
        return_expr_for_serialized_types_method! {$expr, serialize_f32(f32)}
    };
    ($expr:expr, f64) => {
        return_expr_for_serialized_types_method! {$expr, serialize_f64(f64)}
    };
    ($expr:expr, char) => {
        return_expr_for_serialized_types_method! {$expr, serialize_char(char)}
    };
    ($expr:expr, str) => {
        return_expr_for_serialized_types_method! {$expr, serialize_str(&str)}
    };
    ($expr:expr, bytes) => {
        return_expr_for_serialized_types_method! {$expr, serialize_bytes(&[u8])}
    };
    ($expr:expr, none) => {
        return_expr_for_serialized_types_method! {$expr, serialize_none()}
    };
    ($expr:expr, unit) => {
        return_expr_for_serialized_types_method! {$expr, serialize_unit()}
    };
    ($expr:expr, unit_struct) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_unit_struct(&'static str)
        }
    };
    ($expr:expr, unit_variant) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_unit_variant(&'static str, u32, &'static str)
        }
    };
    ($expr:expr, some) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_some(&__T),
            where: ::serde::ser::Serialize
        }
    };
    ($expr:expr, newtype_struct) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_newtype_struct(&'static str, &__T),
            where: ::serde::ser::Serialize
        }
    };
    ($expr:expr, newtype_variant) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_newtype_variant(&'static str, u32,
                                              &'static str, &__T),
            where: ::serde::ser::Serialize
        }
    };
    ($expr:expr, seq) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_seq(Option<usize>),
            result: Self::SerializeSeq
        }
    };
    ($expr:expr, tuple) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_tuple(usize),
            result: Self::SerializeTuple
        }
    };
    ($expr:expr, tuple_struct) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_tuple_struct(&'static str, usize),
            result: Self::SerializeTupleStruct
        }
    };
    ($expr:expr, tuple_variant) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_tuple_variant(&'static str, u32, &'static str,
                                            usize),
            result: Self::SerializeTupleVariant
        }
    };
    ($expr:expr, map) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_map(Option<usize>),
            result: Self::SerializeMap
        }
    };
    ($expr:expr, struct) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_struct(&'static str, usize),
            result: Self::SerializeStruct
        }
    };
    ($expr:expr, struct_variant) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_struct_variant(&'static str, u32, &'static str,
                                             usize),
            result: Self::SerializeStructVariant
        }
    };
}

/// Helper macro for implementing the `serde::se::Serializer` trait.
///
/// Implement the serializer methods for each `$type` with body `$expr`.
///
/// This macro is very similar to `serde::forward_to_deserialize`, but instead
/// of "forwarding" it allows arbitrary expressions in the body, so long as they
/// are all the same.
///
/// This macro can be used to stub out `Serializer` implementations.
macro_rules! return_expr_for_serialized_types {
    ($expr:expr; $($type:tt)*) => {
        $(return_expr_for_serialized_types_helper!{$expr, $type})*
    };
}

macro_rules! unrepresentable {
    ($($type:tt)*) => {
        $(return_expr_for_serialized_types_helper!{Err(Error::UnrepresentableType(stringify!($type))), $type})*
    };
}

/// Serde `serialize_with` implementation for array serialization.
///
/// This macro provides the function body for `i8_array`, `i32_array` and `i64_array`
/// in [`self::ser`], providing NBT `ByteArray`, `IntArray` and `LongArray`
/// serialization with serde.
macro_rules! array_serializer {
    ($func_name:literal, $arr: ident, $serializer: ident) => {{
        use serde::ser::SerializeTupleStruct;
        use std::borrow::Borrow;

        let error = concat!(
            $func_name,
            " serializer may only be used with known-length collections"
        );
        let magic = concat!("__hematite_nbt_", $func_name, "__");

        let mut iter = $arr.into_iter();
        let (length, max_length) = iter.size_hint();

        if max_length.is_none() || length != max_length.unwrap() {
            return Err(SerError::custom(error));
        }

        let mut seq = $serializer.serialize_tuple_struct(magic, length)?;
        for _i in 0..length {
            seq.serialize_field(iter.next().ok_or(SerError::custom(error))?.borrow())?;
        }

        if iter.next().is_some() {
            Err(SerError::custom(error))
        } else {
            seq.end()
        }
    }};
}
<commit_msg>Remove superfluous allocation in array serialization (#61)<commit_after>//! Helper macros for implementing NBT format encoding and decoding via `serde`.

macro_rules! return_expr_for_serialized_types_method {
    ($expr:expr, $func:ident($($arg:ty),*)) => {
        #[inline]
        fn $func(self, $(_: $arg,)*)
                 -> ::std::result::Result<Self::Ok, Self::Error>
        {
            $expr
        }
    };
    ($expr:expr, $func:ident($($arg:ty),*), result: $result:path) => {
        #[inline]
        fn $func(self, $(_: $arg,)*)
                 -> ::std::result::Result<$result, Self::Error>
        {
            $expr
        }
    };
    ($expr:expr, $func:ident($($arg:ty),*), where: $where:path) => {
        #[inline]
        fn $func<__T: ?Sized>(self, $(_: $arg,)*)
                              -> ::std::result::Result<Self::Ok, Self::Error>
            where __T: $where
        {
            $expr
        }
    };
}

macro_rules! return_expr_for_serialized_types_helper {
    ($expr:expr, bool) => {
        return_expr_for_serialized_types_method! {$expr, serialize_bool(bool)}
    };
    ($expr:expr, i8) => {
        return_expr_for_serialized_types_method! {$expr, serialize_i8(i8)}
    };
    ($expr:expr, i16) => {
        return_expr_for_serialized_types_method! {$expr, serialize_i16(i16)}
    };
    ($expr:expr, i32) => {
        return_expr_for_serialized_types_method! {$expr, serialize_i32(i32)}
    };
    ($expr:expr, i64) => {
        return_expr_for_serialized_types_method! {$expr, serialize_i64(i64)}
    };
    ($expr:expr, u8) => {
        return_expr_for_serialized_types_method! {$expr, serialize_u8(u8)}
    };
    ($expr:expr, u16) => {
        return_expr_for_serialized_types_method! {$expr, serialize_u16(u16)}
    };
    ($expr:expr, u32) => {
        return_expr_for_serialized_types_method! {$expr, serialize_u32(u32)}
    };
    ($expr:expr, u64) => {
        return_expr_for_serialized_types_method! {$expr, serialize_u64(u64)}
    };
    ($expr:expr, f32) => {
        return_expr_for_serialized_types_method! {$expr, serialize_f32(f32)}
    };
    ($expr:expr, f64) => {
        return_expr_for_serialized_types_method! {$expr, serialize_f64(f64)}
    };
    ($expr:expr, char) => {
        return_expr_for_serialized_types_method! {$expr, serialize_char(char)}
    };
    ($expr:expr, str) => {
        return_expr_for_serialized_types_method! {$expr, serialize_str(&str)}
    };
    ($expr:expr, bytes) => {
        return_expr_for_serialized_types_method! {$expr, serialize_bytes(&[u8])}
    };
    ($expr:expr, none) => {
        return_expr_for_serialized_types_method! {$expr, serialize_none()}
    };
    ($expr:expr, unit) => {
        return_expr_for_serialized_types_method! {$expr, serialize_unit()}
    };
    ($expr:expr, unit_struct) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_unit_struct(&'static str)
        }
    };
    ($expr:expr, unit_variant) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_unit_variant(&'static str, u32, &'static str)
        }
    };
    ($expr:expr, some) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_some(&__T),
            where: ::serde::ser::Serialize
        }
    };
    ($expr:expr, newtype_struct) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_newtype_struct(&'static str, &__T),
            where: ::serde::ser::Serialize
        }
    };
    ($expr:expr, newtype_variant) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_newtype_variant(&'static str, u32,
                                              &'static str, &__T),
            where: ::serde::ser::Serialize
        }
    };
    ($expr:expr, seq) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_seq(Option<usize>),
            result: Self::SerializeSeq
        }
    };
    ($expr:expr, tuple) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_tuple(usize),
            result: Self::SerializeTuple
        }
    };
    ($expr:expr, tuple_struct) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_tuple_struct(&'static str, usize),
            result: Self::SerializeTupleStruct
        }
    };
    ($expr:expr, tuple_variant) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_tuple_variant(&'static str, u32, &'static str,
                                            usize),
            result: Self::SerializeTupleVariant
        }
    };
    ($expr:expr, map) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_map(Option<usize>),
            result: Self::SerializeMap
        }
    };
    ($expr:expr, struct) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_struct(&'static str, usize),
            result: Self::SerializeStruct
        }
    };
    ($expr:expr, struct_variant) => {
        return_expr_for_serialized_types_method! {
            $expr, serialize_struct_variant(&'static str, u32, &'static str,
                                             usize),
            result: Self::SerializeStructVariant
        }
    };
}

/// Helper macro for implementing the `serde::se::Serializer` trait.
///
/// Implement the serializer methods for each `$type` with body `$expr`.
///
/// This macro is very similar to `serde::forward_to_deserialize`, but instead
/// of "forwarding" it allows arbitrary expressions in the body, so long as they
/// are all the same.
///
/// This macro can be used to stub out `Serializer` implementations.
macro_rules! return_expr_for_serialized_types {
    ($expr:expr; $($type:tt)*) => {
        $(return_expr_for_serialized_types_helper!{$expr, $type})*
    };
}

macro_rules! unrepresentable {
    ($($type:tt)*) => {
        $(return_expr_for_serialized_types_helper!{Err(Error::UnrepresentableType(stringify!($type))), $type})*
    };
}

/// Serde `serialize_with` implementation for array serialization.
///
/// This macro provides the function body for `i8_array`, `i32_array` and `i64_array`
/// in [`self::ser`], providing NBT `ByteArray`, `IntArray` and `LongArray`
/// serialization with serde.
macro_rules! array_serializer {
    ($func_name:literal, $arr: ident, $serializer: ident) => {{
        use serde::ser::SerializeTupleStruct;
        use std::borrow::Borrow;

        let error = concat!(
            $func_name,
            " serializer may only be used with known-length collections"
        );
        let magic = concat!("__hematite_nbt_", $func_name, "__");

        let mut iter = $arr.into_iter();
        let (length, max_length) = iter.size_hint();

        if max_length.is_none() || length != max_length.unwrap() {
            return Err(SerError::custom(error));
        }

        let mut seq = $serializer.serialize_tuple_struct(magic, length)?;
        for _i in 0..length {
            seq.serialize_field(iter.next().ok_or_else(|| SerError::custom(error))?.borrow())?;
        }

        if iter.next().is_some() {
            Err(SerError::custom(error))
        } else {
            seq.end()
        }
    }};
}
<|endoftext|>