<commit_before>
use hyper::StatusCode;
use hyper::header::ByteRangeSpec;
use hyper::header::ByteRangeSpec::FromTo;
use hyper::header::Range::Bytes;
use hyper::header::{ContentLength, ContentRange, ContentRangeSpec};

use reqwest;
use reqwest::Client;
use reqwest::Error;
use reqwest::header;

use reader::Reader;
use buffer::Buffer;

use std::io::SeekFrom;
use std::cmp;
use std::time::Instant;

fn get_head(filename: &String) -> Result<reqwest::Response, Error> {
  if filename.contains("s3.amazonaws.com") {
    let range = vec![FromTo(0, 0)];

    let mut headers = header::Headers::new();
    headers.set(Bytes(range));
    let client = reqwest::Client::builder()
      .default_headers(headers)
      .build()
      .unwrap();

    client.get(filename).send()
  } else {
    let client = Client::new();
    client.head(filename).send()
  }
}

#[derive(Debug)]
struct ResponseData {
  body_data: Vec<u8>,
  file_size: u64
}

fn get_data(filename: &String, range: Vec<ByteRangeSpec>) -> Result<ResponseData, String> {
  let mut headers = header::Headers::new();
  headers.set(Bytes(range));
  let client = reqwest::Client::builder()
    .default_headers(headers)
    .build()
    .unwrap();

  let mut response = client.get(filename).send().unwrap();

  let status = response.status();

  if !(status == StatusCode::Ok || status == StatusCode::PartialContent) {
    println!("ERROR {:?}", response);
    return Err("bad response status".to_string());
  }

  let mut body: Vec<u8> = vec![];
  response.copy_to(&mut body).unwrap();

  let file_size =
    match get_content_range(&response) {
      Ok(length) => length.unwrap(),
      Err(_msg) => return Err("bad response header".to_string()),
    };

  Ok(ResponseData{
    body_data: body,
    file_size: file_size
  })
}

fn get_content_length(response: &reqwest::Response) -> Option<u64> {
  match response.headers().get::<ContentLength>() {
    Some(length) => Some(**length as u64),
    None => None
  }
}

fn get_content_range(response: &reqwest::Response) -> Result<Option<u64>, String> {
  match response.headers().get::<ContentRange>() {
    Some(content_range) => {
      match content_range.clone() {
        ContentRange(ContentRangeSpec::Bytes{
            range: _range,
            instance_length: length
          }) => {
          Ok(length)
        },
        ContentRange(ContentRangeSpec::Unregistered{unit: _unit, resp: _resp}) => {
          Err("Unregistered, actually unsupported".to_string())
        }
      }
    },
    None =>
      Err("Missing content_range".to_string()),
  }
}

#[derive(Debug)]
pub struct HttpReader {
  pub filename: String,
  pub file_size: Option<u64>,
  pub position: u64,
  pub buffer: Buffer
}

pub fn exists(filename: &String) -> bool {
  match get_head(filename) {
    Ok(resp) => {
      resp.status().is_success()
    },
    Err(_msg) => {
      false
    }
  }
}

fn get_data_range(position: u64, size: usize, max_end_position: Option<u64>) -> Vec<ByteRangeSpec> {
  let start = position;
  let end =
    match (position, size) {
      (0, 0) => 0,
      (_, _) => {
        match max_end_position {
          Some(max) => {
            let max_size = max - position;
            position + cmp::min((size - 1) as u64, max_size)
          },
          None => position + (size - 1) as u64,
        }
      },
    };

  vec![FromTo(start, end)]
}

fn load_data(reader: &mut HttpReader, size: usize) -> Result<Option<Vec<u8>>, String> {
  let start = Instant::now();
  info!("make HTTP request with request {:?} bytes", size);

  let position =
    match reader.buffer.size {
      Some(_) => reader.buffer.position,
      None => reader.position,
    };

  match reader.file_size {
    Some(size) => {
      if position >= size {
        return Ok(None)
      }
    },
    None => (),
  };

  let range = get_data_range(position, size, reader.buffer.max_end_position);
  let response = get_data(&reader.filename, range).map_err(|e| e.to_string()).unwrap();

  let elapsed = start.elapsed();
  if elapsed.as_secs() > 0 {
    println!("WARNING: Request duration {} seconds", elapsed.as_secs());
  }

  let new_position = position + response.body_data.len() as u64;
  match reader.buffer.size {
    Some(_) => {
      reader.buffer.position = new_position;
    },
    None => {
      reader.position = new_position;
    }
  };
  Ok(Some(response.body_data))
}

impl Reader for HttpReader {
  fn open(filename: &String) -> HttpReader {
    match get_head(filename) {
      Err(msg) => panic!(msg.to_string()),
      Ok(response) => {
        let content_length = get_content_length(&response);

        HttpReader {
          filename: filename.to_string(),
          file_size: content_length,
          position: 0,
          buffer: Buffer {
            size: None,
            position: 0,
            max_end_position: None,
            buffer: vec![]
          }
        }
      }
    }
  }

  fn get_position(&mut self) -> Result<u64, String> {
    Ok(self.position)
  }

  fn get_cache_size(&self) -> Option<usize> {
    self.buffer.size
  }

  fn set_cache_size(&mut self, cache_size: Option<usize>) {
    self.buffer.size = cache_size;
  }

  fn get_max_end_position(&self) -> Option<u64> {
    self.buffer.max_end_position
  }

  fn set_max_end_position(&mut self, max_end_position: Option<u64>) {
    self.buffer.max_end_position = max_end_position;
  }

  fn get_size(&mut self) -> Result<u64, String> {
    match self.file_size {
        Some(length) => Ok(length),
        None => Err("No length detected".to_string()),
    }
  }

  fn read(&mut self, size: usize) -> Result<Vec<u8>, String> {
    if self.buffer.get_cached_size() >= size {
      self.position += size as u64;
      Ok(self.buffer.get_data(size))
    } else {
      match self.buffer.size {
        Some(buffer_size) => {
          match load_data(self, buffer_size) {
            Err(msg) => Err(msg),
            Ok(some_data) => {
              match some_data {
                Some(data) => {
                  self.buffer.append_data(&data.to_vec());
                  self.position += size as u64;
                  Ok(self.buffer.get_data(size))
                },
                None => Ok(Vec::new()),
              }
            }
          }
        },
        None => {
          match load_data(self, size) {
            Err(msg) => Err(msg),
            Ok(some_data) => {
              match some_data {
                Some(data) => {
                  // println!("{:?} vs {:?}", data.len(), size);
                  match data.len() >= size {
                    true => {
                      Ok(data)
                    },
                    false => Ok(Vec::new()),
                  }
                },
                None => Ok(Vec::new()),
              }
            }
          }
        }
      }
    }
  }

  fn seek(&mut self, seek: SeekFrom) -> Result<u64, String> {
    match seek {
      SeekFrom::Current(offset) => {
        self.position = self.position + offset as u64;
        if self.buffer.size.is_some() {
          if offset > 0 && self.buffer.get_cached_size() > offset as usize {
            let _skiped_data = self.buffer.get_data(offset as usize);
          } else {
            self.buffer.reset();
          }
        }
      },
      SeekFrom::Start(offset) => {
        self.buffer.reset();
        self.position = offset;
        if self.buffer.size.is_some() {
          self.buffer.position = self.position;
        }
      },
      SeekFrom::End(offset) => {
        self.buffer.reset();
        match self.file_size {
          Some(size) => {
            self.position = size - offset as u64;
            if self.buffer.size.is_some() {
              self.buffer.position = self.position;
            }
          },
          None => return Err("Missing file size".to_string())
        }
      }
    }
    Ok(self.position)
  }
}
<commit_msg>fix get content size with AWS<commit_after>
use hyper::StatusCode;
use hyper::header::ByteRangeSpec;
use hyper::header::ByteRangeSpec::FromTo;
use hyper::header::Range::Bytes;
use hyper::header::{ContentLength, ContentRange, ContentRangeSpec};

use reqwest;
use reqwest::Client;
use reqwest::Error;
use reqwest::header;

use reader::Reader;
use buffer::Buffer;

use std::io::SeekFrom;
use std::cmp;
use std::time::Instant;

fn get_head(filename: &String) -> Result<reqwest::Response, Error> {
  if filename.contains("s3.amazonaws.com") {
    let range = vec![FromTo(0, 0)];

    let mut headers = header::Headers::new();
    headers.set(Bytes(range));
    let client = reqwest::Client::builder()
      .default_headers(headers)
      .build()
      .unwrap();

    client.get(filename).send()
  } else {
    let client = Client::new();
    client.head(filename).send()
  }
}

#[derive(Debug)]
struct ResponseData {
  body_data: Vec<u8>,
  file_size: u64
}

fn get_data(filename: &String, range: Vec<ByteRangeSpec>) -> Result<ResponseData, String> {
  let mut headers = header::Headers::new();
  headers.set(Bytes(range));
  let client = reqwest::Client::builder()
    .default_headers(headers)
    .build()
    .unwrap();

  let mut response = client.get(filename).send().unwrap();

  let status = response.status();

  if !(status == StatusCode::Ok || status == StatusCode::PartialContent) {
    println!("ERROR {:?}", response);
    return Err("bad response status".to_string());
  }

  let mut body: Vec<u8> = vec![];
  response.copy_to(&mut body).unwrap();

  let file_size =
    match get_content_range(&response) {
      Ok(length) => length.unwrap(),
      Err(_msg) => return Err("bad response header".to_string()),
    };

  Ok(ResponseData{
    body_data: body,
    file_size: file_size
  })
}

fn get_content_length(response: &reqwest::Response) -> Option<u64> {
  match response.headers().get::<ContentRange>() {
    Some(&ContentRange(ContentRangeSpec::Bytes{range: _range, instance_length})) => {
      instance_length
    },
    _ => {
      match response.headers().get::<ContentLength>() {
        Some(length) => Some(**length as u64),
        None => None
      }
    },
  }
}

fn get_content_range(response: &reqwest::Response) -> Result<Option<u64>, String> {
  match response.headers().get::<ContentRange>() {
    Some(content_range) => {
      match content_range.clone() {
        ContentRange(ContentRangeSpec::Bytes{
            range: _range,
            instance_length: length
          }) => {
          Ok(length)
        },
        ContentRange(ContentRangeSpec::Unregistered{unit: _unit, resp: _resp}) => {
          Err("Unregistered, actually unsupported".to_string())
        }
      }
    },
    None =>
      Err("Missing content_range".to_string()),
  }
}

#[derive(Debug)]
pub struct HttpReader {
  pub filename: String,
  pub file_size: Option<u64>,
  pub position: u64,
  pub buffer: Buffer
}

pub fn exists(filename: &String) -> bool {
  match get_head(filename) {
    Ok(resp) => {
      resp.status().is_success()
    },
    Err(_msg) => {
      false
    }
  }
}

fn get_data_range(position: u64, size: usize, max_end_position: Option<u64>) -> Vec<ByteRangeSpec> {
  let start = position;
  let end =
    match (position, size) {
      (0, 0) => 0,
      (_, _) => {
        match max_end_position {
          Some(max) => {
            let max_size = max - position;
            position + cmp::min((size - 1) as u64, max_size)
          },
          None => position + (size - 1) as u64,
        }
      },
    };

  vec![FromTo(start, end)]
}

fn load_data(reader: &mut HttpReader, size: usize) -> Result<Option<Vec<u8>>, String> {
  let start = Instant::now();
  info!("make HTTP request with request {:?} bytes", size);

  let position =
    match reader.buffer.size {
      Some(_) => reader.buffer.position,
      None => reader.position,
    };

  match reader.file_size {
    Some(total_file_size) => {
      if position >= total_file_size {
        info!("request range out of range: {} > {}", position, total_file_size);
        return Ok(None)
      }
    },
    None => (),
  };

  let range = get_data_range(position, size, reader.buffer.max_end_position);
  let response = get_data(&reader.filename, range).map_err(|e| e.to_string()).unwrap();

  let elapsed = start.elapsed();
  if elapsed.as_secs() > 0 {
    println!("WARNING: Request duration {} seconds", elapsed.as_secs());
  }

  let new_position = position + response.body_data.len() as u64;
  match reader.buffer.size {
    Some(_) => {
      reader.buffer.position = new_position;
    },
    None => {
      reader.position = new_position;
    }
  };
  Ok(Some(response.body_data))
}

impl Reader for HttpReader {
  fn open(filename: &String) -> HttpReader {
    match get_head(filename) {
      Err(msg) => panic!(msg.to_string()),
      Ok(response) => {
        let content_length = get_content_length(&response);

        HttpReader {
          filename: filename.to_string(),
          file_size: content_length,
          position: 0,
          buffer: Buffer {
            size: None,
            position: 0,
            max_end_position: None,
            buffer: vec![]
          }
        }
      }
    }
  }

  fn get_position(&mut self) -> Result<u64, String> {
    Ok(self.position)
  }

  fn get_cache_size(&self) -> Option<usize> {
    self.buffer.size
  }

  fn set_cache_size(&mut self, cache_size: Option<usize>) {
    self.buffer.size = cache_size;
  }

  fn get_max_end_position(&self) -> Option<u64> {
    self.buffer.max_end_position
  }

  fn set_max_end_position(&mut self, max_end_position: Option<u64>) {
    self.buffer.max_end_position = max_end_position;
  }

  fn get_size(&mut self) -> Result<u64, String> {
    match self.file_size {
        Some(length) => Ok(length),
        None => Err("No length detected".to_string()),
    }
  }

  fn read(&mut self, size: usize) -> Result<Vec<u8>, String> {
    if self.buffer.get_cached_size() >= size {
      self.position += size as u64;
      Ok(self.buffer.get_data(size))
    } else {
      match self.buffer.size {
        Some(buffer_size) => {
          match load_data(self, buffer_size) {
            Err(msg) => Err(msg),
            Ok(some_data) => {
              match some_data {
                Some(data) => {
                  self.buffer.append_data(&data.to_vec());
                  self.position += size as u64;
                  Ok(self.buffer.get_data(size))
                },
                None => Ok(Vec::new()),
              }
            }
          }
        },
        None => {
          match load_data(self, size) {
            Err(msg) => Err(msg),
            Ok(some_data) => {
              match some_data {
                Some(data) => {
                  // println!("{:?} vs {:?}", data.len(), size);
                  match data.len() >= size {
                    true => {
                      Ok(data)
                    },
                    false => Ok(Vec::new()),
                  }
                },
                None => Ok(Vec::new()),
              }
            }
          }
        }
      }
    }
  }

  fn seek(&mut self, seek: SeekFrom) -> Result<u64, String> {
    match seek {
      SeekFrom::Current(offset) => {
        self.position = self.position + offset as u64;
        if self.buffer.size.is_some() {
          if offset > 0 && self.buffer.get_cached_size() > offset as usize {
            let _skiped_data = self.buffer.get_data(offset as usize);
          } else {
            self.buffer.reset();
          }
        }
      },
      SeekFrom::Start(offset) => {
        self.buffer.reset();
        self.position = offset;
        if self.buffer.size.is_some() {
          self.buffer.position = self.position;
        }
      },
      SeekFrom::End(offset) => {
        self.buffer.reset();
        match self.file_size {
          Some(size) => {
            self.position = size - offset as u64;
            if self.buffer.size.is_some() {
              self.buffer.position = self.position;
            }
          },
          None => return Err("Missing file size".to_string())
        }
      }
    }
    Ok(self.position)
  }
}
<|endoftext|>


<commit_before>extern crate libc;
use std::ptr;
use std::str;
use std::ffi::CStr;
use std::fmt::{self, Display, Debug, Formatter, Error};
use capstone_sys::{cs_free, cs_insn, cs_detail};

/// Representation of the array of instructions returned by disasm
#[derive(Debug)]
pub struct Instructions {
    ptr: *mut cs_insn,
    len: isize,
}

impl Instructions {
    pub unsafe fn from_raw_parts(ptr: *mut cs_insn, len: isize) -> Instructions {
        Instructions {
            ptr: ptr,
            len: len,
        }
    }

    pub fn len(&self) -> isize {
        self.len
    }

     pub fn iter(&self) -> InstructionIterator {
         InstructionIterator {
             insns: &self,
             cur: 0,
         }
     }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }
}

impl Drop for Instructions {
     fn drop(&mut self) {
         unsafe {
             cs_free(self.ptr, self.len as libc::size_t);
         }
     }
}

/// An iterator over the instructions returned by disasm
///
/// This is currently the only supported interface for reading them.
pub struct InstructionIterator<'a> {
    insns: &'a Instructions,
    cur: isize,
}

impl<'a> Iterator for InstructionIterator<'a> {
    type Item = Insn;

    fn next(&mut self) -> Option<Insn> {
        if self.cur == self.insns.len {
            None
        } else {
            let obj = unsafe { self.insns.ptr.offset(self.cur) };
            self.cur += 1;
            Some(unsafe { Insn::from(ptr::read(obj)) })
        }
    }
}

/// A wrapper for the raw capstone-sys instruction
pub struct Insn(cs_insn);

/// A wrapper for the raw capstone-sys detail struct, which contains register information in addition
/// to architecture specific information (not implemented in capstone-sys yet)
pub struct Detail(cs_detail);

use std::convert::From;

impl From<cs_insn> for Insn {
    fn from(insn: cs_insn) -> Self {
        Insn (insn)
    }
}

impl From<cs_detail> for Detail{
    fn from(detail: cs_detail) -> Self {
        Detail (detail)
    }
}

impl Insn {
    /// The mnemonic for the instruction
    pub fn mnemonic(&self) -> Option<&str> {
        let cstr = unsafe { CStr::from_ptr(self.0.mnemonic.as_ptr()) };
        str::from_utf8(cstr.to_bytes()).ok()
    }

    /// The operand string associated with the instruction
    pub fn op_str(&self) -> Option<&str> {
        let cstr = unsafe { CStr::from_ptr(self.0.op_str.as_ptr()) };
        str::from_utf8(cstr.to_bytes()).ok()
    }

    fn size (&self) -> usize {
        self.0.size as usize
    }

    pub fn address (&self) -> u64 {
        self.0.address as u64
    }

    pub fn bytes (&self) -> &[u8] {
        &self.0.bytes[..self.size()]
    }

    /// Returns the detail object, if there is one.
    ///
    /// Be careful this is still in early stages and largely untested with various `cs_option` and
    /// architecture matrices
    pub fn detail (&self) -> Option<Detail> {
        if self.0.detail.is_null() {
            None
        } else {
            Some(Detail::from(unsafe { ptr::read(self.0.detail) }))
        }
    }
}

impl Debug for Insn {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        fmt.debug_struct("Insn")
            .field("address", &self.address())
            .field("size", &self.size())
            .field("mnemonic", &self.mnemonic())
            .field("op_str", &self.op_str())
            .finish()
    }
}

impl Display for Insn {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "{:#x}: ", self.address())?;
        if let Some(mnemonic) = self.mnemonic() {
            write!(fmt, "{} ", mnemonic)?;
            if let Some(op_str) = self.op_str() {
                write!(fmt, "{}", op_str)?;
            }
        }
        Ok(())
    }
}

impl Detail {
    /// list of implicit registers read by this insn
    pub fn regs_read(&self) -> &[libc::uint8_t; 12] {
        &self.0.regs_read
    }
    /// number of implicit registers read by this insn
    pub fn regs_read_count(&self) -> libc::uint8_t {
        self.0.regs_read_count
    }
    /// list of implicit registers modified by this insn
    pub fn regs_write(&self) -> &[libc::uint8_t; 20] {
        &self.0.regs_write
    }
    /// number of implicit registers modified by this insn
    pub fn regs_write_count(&self) -> libc::uint8_t {
        self.0.regs_write_count
    }
    /// list of group this instruction belong to
    pub fn groups(&self) -> &[libc::uint8_t; 8] {
        &self.0.groups
    }
    /// number of groups this insn belongs to
    pub fn groups_count(&self) -> libc::uint8_t {
        self.0.groups_count
    }

}

impl Debug for Detail {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        fmt.debug_struct("Detail")
            .field("regs_read", &self.regs_read())
            .field("regs_read_count", &self.regs_read_count())
            .field("regs_write", &self.regs_write())
            .field("regs_write_count", &self.regs_write_count())
            .field("groups", &self.groups())
            .field("groups_count", &self.groups_count())
            .finish()
    }
}

impl Display for Instructions {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        for sym in self.iter() {
            write!(fmt, "{:x}:\t", sym.address())?;
            for byte in sym.bytes() {
                write!(fmt, " {:02x}", byte)?;
            }
            let remainder = 16*3 - (sym.bytes().len())*3;
            for _ in 0..remainder {
                write!(fmt, " ")?;
            }
            if let Some(mnemonic) = sym.mnemonic() {
                write!(fmt, " {}", mnemonic)?;
                if let Some(op_str) = sym.op_str() {
                    write!(fmt, " {}", op_str)?;
                }
            }
            write!(fmt, "\n")?;
        }
        Ok(())
    }
}
<commit_msg>insn: add len public function and bytes to debug repr<commit_after>extern crate libc;
use std::ptr;
use std::str;
use std::ffi::CStr;
use std::fmt::{self, Display, Debug, Formatter, Error};
use capstone_sys::{cs_free, cs_insn, cs_detail};

/// Representation of the array of instructions returned by disasm
#[derive(Debug)]
pub struct Instructions {
    ptr: *mut cs_insn,
    len: isize,
}

impl Instructions {
    pub unsafe fn from_raw_parts(ptr: *mut cs_insn, len: isize) -> Instructions {
        Instructions {
            ptr: ptr,
            len: len,
        }
    }

    pub fn len(&self) -> isize {
        self.len
    }

     pub fn iter(&self) -> InstructionIterator {
         InstructionIterator {
             insns: &self,
             cur: 0,
         }
     }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }
}

impl Drop for Instructions {
     fn drop(&mut self) {
         unsafe {
             cs_free(self.ptr, self.len as libc::size_t);
         }
     }
}

/// An iterator over the instructions returned by disasm
///
/// This is currently the only supported interface for reading them.
pub struct InstructionIterator<'a> {
    insns: &'a Instructions,
    cur: isize,
}

impl<'a> Iterator for InstructionIterator<'a> {
    type Item = Insn;

    fn next(&mut self) -> Option<Insn> {
        if self.cur == self.insns.len {
            None
        } else {
            let obj = unsafe { self.insns.ptr.offset(self.cur) };
            self.cur += 1;
            Some(unsafe { Insn::from(ptr::read(obj)) })
        }
    }
}

/// A wrapper for the raw capstone-sys instruction
pub struct Insn(cs_insn);

/// A wrapper for the raw capstone-sys detail struct, which contains register information in addition
/// to architecture specific information (not implemented in capstone-sys yet)
pub struct Detail(cs_detail);

use std::convert::From;

impl From<cs_insn> for Insn {
    fn from(insn: cs_insn) -> Self {
        Insn (insn)
    }
}

impl From<cs_detail> for Detail{
    fn from(detail: cs_detail) -> Self {
        Detail (detail)
    }
}

impl Insn {
    /// The mnemonic for the instruction
    pub fn mnemonic(&self) -> Option<&str> {
        let cstr = unsafe { CStr::from_ptr(self.0.mnemonic.as_ptr()) };
        str::from_utf8(cstr.to_bytes()).ok()
    }

    /// The operand string associated with the instruction
    pub fn op_str(&self) -> Option<&str> {
        let cstr = unsafe { CStr::from_ptr(self.0.op_str.as_ptr()) };
        str::from_utf8(cstr.to_bytes()).ok()
    }

    /// Size of instruction (in bytes)
    pub fn len (&self) -> usize {
        self.0.size as usize
    }

    pub fn address (&self) -> u64 {
        self.0.address as u64
    }

    /// Byte-level representation of the instruction
    pub fn bytes (&self) -> &[u8] {
        &self.0.bytes[..self.len()]
    }

    /// Returns the detail object, if there is one.
    ///
    /// Be careful this is still in early stages and largely untested with various `cs_option` and
    /// architecture matrices
    pub fn detail (&self) -> Option<Detail> {
        if self.0.detail.is_null() {
            None
        } else {
            Some(Detail::from(unsafe { ptr::read(self.0.detail) }))
        }
    }
}

impl Debug for Insn {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        fmt.debug_struct("Insn")
            .field("address", &self.address())
            .field("len", &self.len())
            .field("bytes", &self.bytes())
            .field("mnemonic", &self.mnemonic())
            .field("op_str", &self.op_str())
            .finish()
    }
}

impl Display for Insn {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "{:#x}: ", self.address())?;
        if let Some(mnemonic) = self.mnemonic() {
            write!(fmt, "{} ", mnemonic)?;
            if let Some(op_str) = self.op_str() {
                write!(fmt, "{}", op_str)?;
            }
        }
        Ok(())
    }
}

impl Detail {
    /// list of implicit registers read by this insn
    pub fn regs_read(&self) -> &[libc::uint8_t; 12] {
        &self.0.regs_read
    }
    /// number of implicit registers read by this insn
    pub fn regs_read_count(&self) -> libc::uint8_t {
        self.0.regs_read_count
    }
    /// list of implicit registers modified by this insn
    pub fn regs_write(&self) -> &[libc::uint8_t; 20] {
        &self.0.regs_write
    }
    /// number of implicit registers modified by this insn
    pub fn regs_write_count(&self) -> libc::uint8_t {
        self.0.regs_write_count
    }
    /// list of group this instruction belong to
    pub fn groups(&self) -> &[libc::uint8_t; 8] {
        &self.0.groups
    }
    /// number of groups this insn belongs to
    pub fn groups_count(&self) -> libc::uint8_t {
        self.0.groups_count
    }

}

impl Debug for Detail {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        fmt.debug_struct("Detail")
            .field("regs_read", &self.regs_read())
            .field("regs_read_count", &self.regs_read_count())
            .field("regs_write", &self.regs_write())
            .field("regs_write_count", &self.regs_write_count())
            .field("groups", &self.groups())
            .field("groups_count", &self.groups_count())
            .finish()
    }
}

impl Display for Instructions {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        for sym in self.iter() {
            write!(fmt, "{:x}:\t", sym.address())?;
            for byte in sym.bytes() {
                write!(fmt, " {:02x}", byte)?;
            }
            let remainder = 16*3 - (sym.bytes().len())*3;
            for _ in 0..remainder {
                write!(fmt, " ")?;
            }
            if let Some(mnemonic) = sym.mnemonic() {
                write!(fmt, " {}", mnemonic)?;
                if let Some(op_str) = sym.op_str() {
                    write!(fmt, " {}", op_str)?;
                }
            }
            write!(fmt, "\n")?;
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>extern crate libc;
use std::ptr;
use std::str;
use std::ffi::CStr;
use std::fmt::{self, Display, Debug, Formatter, Error};
use capstone_sys::{cs_free, cs_insn, cs_detail};

/// Representation of the array of instructions returned by disasm
pub struct Instructions {
    ptr: *mut cs_insn,
    len: isize,
}

impl Instructions {
    pub unsafe fn from_raw_parts(ptr: *mut cs_insn, len: isize) -> Instructions {
        Instructions {
            ptr: ptr,
            len: len,
        }
    }

    pub fn len(&self) -> isize {
        self.len
    }

     pub fn iter(&self) -> InstructionIterator {
         InstructionIterator {
             insns: &self,
             cur: 0,
         }
     }
}

impl Drop for Instructions {
     fn drop(&mut self) {
         unsafe {
             cs_free(self.ptr, self.len as libc::size_t);
         }
     }
}

/// An iterator over the instructions returned by disasm
///
/// This is currently the only supported interface for reading them.
pub struct InstructionIterator<'a> {
    insns: &'a Instructions,
    cur: isize,
}

impl<'a> Iterator for InstructionIterator<'a> {
    type Item = Insn;

    fn next(&mut self) -> Option<Insn> {
        if self.cur == self.insns.len {
            None
        } else {
            let obj = unsafe { self.insns.ptr.offset(self.cur) };
            self.cur += 1;
            Some(unsafe { Insn::from(ptr::read(obj)) })
        }
    }
}

/// A wrapper for the raw capstone-sys instruction
pub struct Insn(cs_insn);

/// A wrapper for the raw capstone-sys detail struct, which contains register information in addition
/// to architecture specific information (not implemented in capstone-sys yet)
pub struct Detail(cs_detail);

use std::convert::From;

impl From<cs_insn> for Insn {
    fn from(insn: cs_insn) -> Self {
        Insn (insn)
    }
}

impl From<cs_detail> for Detail{
    fn from(detail: cs_detail) -> Self {
        Detail (detail)
    }
}

impl Insn {
    /// The mnemonic for the instruction
    pub fn mnemonic(&self) -> Option<&str> {
        let cstr = unsafe { CStr::from_ptr(self.0.mnemonic.as_ptr()) };
        str::from_utf8(cstr.to_bytes()).ok()
    }

    /// The operand string associated with the instruction
    pub fn op_str(&self) -> Option<&str> {
        let cstr = unsafe { CStr::from_ptr(self.0.op_str.as_ptr()) };
        str::from_utf8(cstr.to_bytes()).ok()
    }

    pub fn size (&self) -> usize {
        self.0.size as usize
    }

    pub fn address (&self) -> u64 {
        self.0.address as u64
    }

    /// Returns the detail object, if there is one.
    ///
    /// Be careful this is still in early stages and largely untested with various `cs_option` and
    /// architecture matrices
    pub fn detail (&self) -> Option<Detail> {
        if self.0.detail.is_null() {
            None
        } else {
            Some(Detail::from(unsafe { ptr::read(self.0.detail) }))
        }
    }
}

impl Debug for Insn {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        fmt.debug_struct("Insn")
            .field("address", &self.address())
            .field("size", &self.size())
            .field("mnemonic", &self.mnemonic())
            .field("op_str", &self.op_str())
            .finish()
    }
}

impl Display for Insn {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "{:#x}: ", self.address())?;
        if let Some(mnemonic) = self.mnemonic() {
            write!(fmt, "{} ", mnemonic)?;
            if let Some(op_str) = self.op_str() {
                write!(fmt, "{}", op_str)?;
            }
        }
        Ok(())
    }
}

impl Detail {
    /// list of implicit registers read by this insn
    pub fn regs_read(&self) -> &[libc::uint8_t; 12] {
        &self.0.regs_read
    }
    /// number of implicit registers read by this insn
    pub fn regs_read_count(&self) -> libc::uint8_t {
        self.0.regs_read_count
    }
    /// list of implicit registers modified by this insn
    pub fn regs_write(&self) -> &[libc::uint8_t; 20] {
        &self.0.regs_write
    }
    /// number of implicit registers modified by this insn
    pub fn regs_write_count(&self) -> libc::uint8_t {
        self.0.regs_write_count
    }
    /// list of group this instruction belong to
    pub fn groups(&self) -> &[libc::uint8_t; 8] {
        &self.0.groups
    }
    /// number of groups this insn belongs to
    pub fn groups_count(&self) -> libc::uint8_t {
        self.0.groups_count
    }

}

impl Debug for Detail {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        fmt.debug_struct("Detail")
            .field("regs_read", &self.regs_read())
            .field("regs_read_count", &self.regs_read_count())
            .field("regs_write", &self.regs_write())
            .field("regs_write_count", &self.regs_write_count())
            .field("groups", &self.groups())
            .field("groups_count", &self.groups_count())
            .finish()
    }
}
<commit_msg>instructions: add display for instructions. add is_empty. add bytes delimited by size for instruction<commit_after>extern crate libc;
use std::ptr;
use std::str;
use std::ffi::CStr;
use std::fmt::{self, Display, Debug, Formatter, Error};
use capstone_sys::{cs_free, cs_insn, cs_detail};

/// Representation of the array of instructions returned by disasm
#[derive(Debug)]
pub struct Instructions {
    ptr: *mut cs_insn,
    len: isize,
}

impl Instructions {
    pub unsafe fn from_raw_parts(ptr: *mut cs_insn, len: isize) -> Instructions {
        Instructions {
            ptr: ptr,
            len: len,
        }
    }

    pub fn len(&self) -> isize {
        self.len
    }

     pub fn iter(&self) -> InstructionIterator {
         InstructionIterator {
             insns: &self,
             cur: 0,
         }
     }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }
}

impl Drop for Instructions {
     fn drop(&mut self) {
         unsafe {
             cs_free(self.ptr, self.len as libc::size_t);
         }
     }
}

/// An iterator over the instructions returned by disasm
///
/// This is currently the only supported interface for reading them.
pub struct InstructionIterator<'a> {
    insns: &'a Instructions,
    cur: isize,
}

impl<'a> Iterator for InstructionIterator<'a> {
    type Item = Insn;

    fn next(&mut self) -> Option<Insn> {
        if self.cur == self.insns.len {
            None
        } else {
            let obj = unsafe { self.insns.ptr.offset(self.cur) };
            self.cur += 1;
            Some(unsafe { Insn::from(ptr::read(obj)) })
        }
    }
}

/// A wrapper for the raw capstone-sys instruction
pub struct Insn(cs_insn);

/// A wrapper for the raw capstone-sys detail struct, which contains register information in addition
/// to architecture specific information (not implemented in capstone-sys yet)
pub struct Detail(cs_detail);

use std::convert::From;

impl From<cs_insn> for Insn {
    fn from(insn: cs_insn) -> Self {
        Insn (insn)
    }
}

impl From<cs_detail> for Detail{
    fn from(detail: cs_detail) -> Self {
        Detail (detail)
    }
}

impl Insn {
    /// The mnemonic for the instruction
    pub fn mnemonic(&self) -> Option<&str> {
        let cstr = unsafe { CStr::from_ptr(self.0.mnemonic.as_ptr()) };
        str::from_utf8(cstr.to_bytes()).ok()
    }

    /// The operand string associated with the instruction
    pub fn op_str(&self) -> Option<&str> {
        let cstr = unsafe { CStr::from_ptr(self.0.op_str.as_ptr()) };
        str::from_utf8(cstr.to_bytes()).ok()
    }

    fn size (&self) -> usize {
        self.0.size as usize
    }

    pub fn address (&self) -> u64 {
        self.0.address as u64
    }

    pub fn bytes (&self) -> &[u8] {
        &self.0.bytes[..self.size()]
    }

    /// Returns the detail object, if there is one.
    ///
    /// Be careful this is still in early stages and largely untested with various `cs_option` and
    /// architecture matrices
    pub fn detail (&self) -> Option<Detail> {
        if self.0.detail.is_null() {
            None
        } else {
            Some(Detail::from(unsafe { ptr::read(self.0.detail) }))
        }
    }
}

impl Debug for Insn {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        fmt.debug_struct("Insn")
            .field("address", &self.address())
            .field("size", &self.size())
            .field("mnemonic", &self.mnemonic())
            .field("op_str", &self.op_str())
            .finish()
    }
}

impl Display for Insn {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "{:#x}: ", self.address())?;
        if let Some(mnemonic) = self.mnemonic() {
            write!(fmt, "{} ", mnemonic)?;
            if let Some(op_str) = self.op_str() {
                write!(fmt, "{}", op_str)?;
            }
        }
        Ok(())
    }
}

impl Detail {
    /// list of implicit registers read by this insn
    pub fn regs_read(&self) -> &[libc::uint8_t; 12] {
        &self.0.regs_read
    }
    /// number of implicit registers read by this insn
    pub fn regs_read_count(&self) -> libc::uint8_t {
        self.0.regs_read_count
    }
    /// list of implicit registers modified by this insn
    pub fn regs_write(&self) -> &[libc::uint8_t; 20] {
        &self.0.regs_write
    }
    /// number of implicit registers modified by this insn
    pub fn regs_write_count(&self) -> libc::uint8_t {
        self.0.regs_write_count
    }
    /// list of group this instruction belong to
    pub fn groups(&self) -> &[libc::uint8_t; 8] {
        &self.0.groups
    }
    /// number of groups this insn belongs to
    pub fn groups_count(&self) -> libc::uint8_t {
        self.0.groups_count
    }

}

impl Debug for Detail {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        fmt.debug_struct("Detail")
            .field("regs_read", &self.regs_read())
            .field("regs_read_count", &self.regs_read_count())
            .field("regs_write", &self.regs_write())
            .field("regs_write_count", &self.regs_write_count())
            .field("groups", &self.groups())
            .field("groups_count", &self.groups_count())
            .finish()
    }
}

impl Display for Instructions {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        for sym in self.iter() {
            write!(fmt, "{:x}:\t", sym.address())?;
            for byte in sym.bytes() {
                write!(fmt, " {:02x}", byte)?;
            }
            let remainder = 16*3 - (sym.bytes().len())*3;
            for _ in 0..remainder {
                write!(fmt, " ")?;
            }
            if let Some(mnemonic) = sym.mnemonic() {
                write!(fmt, " {}", mnemonic)?;
                if let Some(op_str) = sym.op_str() {
                    write!(fmt, " {}", op_str)?;
                }
            }
            write!(fmt, "\n")?;
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate redis_async;
#[macro_use]
extern crate serde_derive;

use std::io;

use actix::prelude::Addr;
use actix_files::{Files, NamedFile};
use actix_redis::{Command, RedisActor, RespValue};
use actix_web::{middleware, web, App, HttpResponse, HttpServer, Result};
use dotenv::dotenv;
use futures::future;
use futures::future::Future;
use serde::Serialize;
use uuid::Uuid;

#[derive(Serialize, Deserialize)]
pub struct SecretBody {
    content: String,
}

#[derive(Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct CreateSecretPostBody {
    content: String,
    expiration_seconds: usize,
}

#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub struct CreateSecretResponseBody {
    id: Uuid,
}

#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub struct ApiError<T: Serialize> {
    error: T,
}

impl<T: Serialize> ApiError<T> {
    pub fn from(err: T) -> ApiError<T> {
        ApiError { error: err }
    }
}

#[derive(Debug, Clone)]
struct ServerConfig {
    redis_host: String,
    bind_host: String,

    app_config: AppConfig,
}

#[derive(Debug, Clone)]
struct AppConfig {
    redis_key_prefix: String,
    max_secret_length: usize,
    max_secret_expiration: usize,
}

struct AppState {
    redis: Addr<RedisActor>,
    config: AppConfig,
}

impl AppState {
    fn format_key(self: &AppState, id: &str) -> String {
        format!("{}.{}", self.config.redis_key_prefix, id)
    }

    fn create_secret(
        self: &AppState,
        id: Uuid,
        body: &str,
        expiration_seconds: usize,
    ) -> impl Future<Item = bool, Error = actix_web::Error> {
        let redis_key = self.format_key(&id.to_string());

        self.redis
            .send(Command(resp_array![
                "SETEX",
                redis_key,
                expiration_seconds.to_string(),
                body
            ]))
            .map_err(actix_web::Error::from)
            .and_then(|res: Result<RespValue, actix_redis::Error>| match res {
                Ok(RespValue::SimpleString(_)) => Ok(true),

                err => {
                    println!("unexpected response to creating secret: {:?}", err);
                    Err(actix_web::error::ErrorInternalServerError(
                        "failed to create secret",
                    ))
                }
            })
    }

    fn fetch_secret(
        self: &AppState,
        id: &str,
    ) -> impl Future<Item = Option<String>, Error = actix_web::Error> {
        let redis_key = self.format_key(id);

        self.redis
            .send(Command(resp_array!["GET", redis_key]))
            .map_err(actix_web::Error::from)
            .map(|res| match res {
                Ok(RespValue::BulkString(bytes)) => {
                    let str = String::from_utf8(bytes).expect("non-utf8 string in redis");
                    Some(str)
                }

                _ => None,
            })
    }

    fn consume_secret(
        self: &AppState,
        id: &str,
    ) -> impl Future<Item = Option<String>, Error = actix_web::Error> {
        let fetch_secret = self.fetch_secret(id);
        let delete_secret = self
            .redis
            .send(Command(resp_array!["DEL", self.format_key(id)]))
            .map_err(actix_web::Error::from);

        fetch_secret.and_then(|res| delete_secret.map(|_| res))
    }
}

fn api_create_secret(
    body: web::Json<CreateSecretPostBody>,
    state: web::Data<AppState>,
) -> impl Future<Item = HttpResponse, Error = actix_web::Error> {
    let body = body.into_inner();

    let id = Uuid::new_v4();
    let expiration = body.expiration_seconds.to_string();

    let validate = if body.content.len() > state.config.max_secret_length {
        Err("content too large")
    } else if body.expiration_seconds > state.config.max_secret_expiration {
        Err("expiration too long")
    } else {
        Ok(())
    };

    future::result(validate)
        .map_err(|err| HttpResponse::BadRequest().json(ApiError::from(err)))
        .map_err(actix_web::Error::from)
        .and_then(move |_| {
            println!("Writing: {:?}: {:?}", id.to_string(), expiration);

            // TODO: Possibly re-encrypt secret here? libsodium or something.

            state
                .create_secret(id, body.content.as_ref(), body.expiration_seconds)
                .map(move |_| HttpResponse::Ok().json(CreateSecretResponseBody { id }))
        })
}

fn api_get_secret(
    path: web::Path<String>,
    state: web::Data<AppState>,
) -> impl Future<Item = HttpResponse, Error = actix_web::Error> {
    let secret_id = path.into_inner();
    println!("Fetching: {:?}", &secret_id);
    state
        .consume_secret(&secret_id)
        .map(|res: Option<String>| match res {
            Some(secret) => HttpResponse::Ok().json(SecretBody { content: secret }),
            None => HttpResponse::NotFound().json(ApiError::from("not found")),
        })
}

fn view_secret_page(
    path: web::Path<String>,
    state: web::Data<AppState>,
) -> impl Future<Item = NamedFile, Error = actix_web::Error> {
    let secret_id = path.into_inner();

    state.fetch_secret(&secret_id).and_then(|res| match res {
        Some(_) => Ok(NamedFile::open("./static/view.html")?),

        None => {
            let file = NamedFile::open("./static/404.html")?
                // TODO: (requires HTTP crate).set_status_code()
                ;

            Ok(file)
        }
    })
}

fn config_from_env() -> ServerConfig {
    ServerConfig {
        redis_host: std::env::var("REDIS_HOST").unwrap_or("localhost:6379".to_string()),
        bind_host: std::env::var("BIND_HOST").unwrap_or("127.0.0.1:8080".to_string()),
        app_config: AppConfig {
            max_secret_length: std::env::var("MAX_SECRET_LENGTH")
                .unwrap_or("4096".to_string())
                .parse()
                .expect("unvalid integer value"),
            max_secret_expiration: std::env::var("MAX_SECRET_EXPIRATION")
                .unwrap_or("4096".to_string())
                .parse()
                .expect("unvalid integer value"),
            redis_key_prefix: std::env::var("REDIS_KEY_PREFIX")
                .unwrap_or("tofu:".to_string()),
        },
    }
}

fn main() -> io::Result<()> {
    env_logger::init();
    dotenv().ok();

    let config = config_from_env();
    println!("Booting with config: {:?}", config);

    let redis_addr = RedisActor::start(config.redis_host);
    let app_config = config.app_config;

    HttpServer::new(move || {
        let state = AppState {
            redis: redis_addr.clone(),
            config: app_config.clone(),
        };

        App::new()
            .data(state)
            .wrap(middleware::Logger::default())
            .service(
                web::resource("/api/secret")
                    .route(web::post().to_async(api_create_secret)),
            )
            .service(
                web::resource("/api/secret/{secret_id}")
                    .route(web::get().to_async(api_get_secret)),
            )
            .service(
                web::resource("/view/{secret_id}")
                    .route(web::get().to_async(view_secret_page)),
            )
            .service(Files::new("/", "./static/").index_file("index.html"))
    })
    .bind(config.bind_host)?
    .run()
}
<commit_msg>Need to create actix system before launching server<commit_after>#[macro_use]
extern crate redis_async;
#[macro_use]
extern crate serde_derive;

use std::io;

use actix::prelude::Addr;
use actix_files::{Files, NamedFile};
use actix_redis::{Command, RedisActor, RespValue};
use actix_web::{middleware, web, App, HttpResponse, HttpServer, Result};
use dotenv::dotenv;
use futures::future;
use futures::future::Future;
use serde::Serialize;
use uuid::Uuid;

#[derive(Serialize, Deserialize)]
pub struct SecretBody {
    content: String,
}

#[derive(Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct CreateSecretPostBody {
    content: String,
    expiration_seconds: usize,
}

#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub struct CreateSecretResponseBody {
    id: Uuid,
}

#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub struct ApiError<T: Serialize> {
    error: T,
}

impl<T: Serialize> ApiError<T> {
    pub fn from(err: T) -> ApiError<T> {
        ApiError { error: err }
    }
}

#[derive(Debug, Clone)]
struct ServerConfig {
    redis_host: String,
    bind_host: String,

    app_config: AppConfig,
}

#[derive(Debug, Clone)]
struct AppConfig {
    redis_key_prefix: String,
    max_secret_length: usize,
    max_secret_expiration: usize,
}

struct AppState {
    redis: Addr<RedisActor>,
    config: AppConfig,
}

impl AppState {
    fn format_key(self: &AppState, id: &str) -> String {
        format!("{}.{}", self.config.redis_key_prefix, id)
    }

    fn create_secret(
        self: &AppState,
        id: Uuid,
        body: &str,
        expiration_seconds: usize,
    ) -> impl Future<Item = bool, Error = actix_web::Error> {
        let redis_key = self.format_key(&id.to_string());

        self.redis
            .send(Command(resp_array![
                "SETEX",
                redis_key,
                expiration_seconds.to_string(),
                body
            ]))
            .map_err(actix_web::Error::from)
            .and_then(|res: Result<RespValue, actix_redis::Error>| match res {
                Ok(RespValue::SimpleString(_)) => Ok(true),

                err => {
                    println!("unexpected response to creating secret: {:?}", err);
                    Err(actix_web::error::ErrorInternalServerError(
                        "failed to create secret",
                    ))
                }
            })
    }

    fn fetch_secret(
        self: &AppState,
        id: &str,
    ) -> impl Future<Item = Option<String>, Error = actix_web::Error> {
        let redis_key = self.format_key(id);

        self.redis
            .send(Command(resp_array!["GET", redis_key]))
            .map_err(actix_web::Error::from)
            .map(|res| match res {
                Ok(RespValue::BulkString(bytes)) => {
                    let str = String::from_utf8(bytes).expect("non-utf8 string in redis");
                    Some(str)
                }

                _ => None,
            })
    }

    fn consume_secret(
        self: &AppState,
        id: &str,
    ) -> impl Future<Item = Option<String>, Error = actix_web::Error> {
        let fetch_secret = self.fetch_secret(id);
        let delete_secret = self
            .redis
            .send(Command(resp_array!["DEL", self.format_key(id)]))
            .map_err(actix_web::Error::from);

        fetch_secret.and_then(|res| delete_secret.map(|_| res))
    }
}

fn api_create_secret(
    body: web::Json<CreateSecretPostBody>,
    state: web::Data<AppState>,
) -> impl Future<Item = HttpResponse, Error = actix_web::Error> {
    let body = body.into_inner();

    let id = Uuid::new_v4();
    let expiration = body.expiration_seconds.to_string();

    let validate = if body.content.len() > state.config.max_secret_length {
        Err("content too large")
    } else if body.expiration_seconds > state.config.max_secret_expiration {
        Err("expiration too long")
    } else {
        Ok(())
    };

    future::result(validate)
        .map_err(|err| HttpResponse::BadRequest().json(ApiError::from(err)))
        .map_err(actix_web::Error::from)
        .and_then(move |_| {
            println!("Writing: {:?}: {:?}", id.to_string(), expiration);

            // TODO: Possibly re-encrypt secret here? libsodium or something.

            state
                .create_secret(id, body.content.as_ref(), body.expiration_seconds)
                .map(move |_| HttpResponse::Ok().json(CreateSecretResponseBody { id }))
        })
}

fn api_get_secret(
    path: web::Path<String>,
    state: web::Data<AppState>,
) -> impl Future<Item = HttpResponse, Error = actix_web::Error> {
    let secret_id = path.into_inner();
    println!("Fetching: {:?}", &secret_id);
    state
        .consume_secret(&secret_id)
        .map(|res: Option<String>| match res {
            Some(secret) => HttpResponse::Ok().json(SecretBody { content: secret }),
            None => HttpResponse::NotFound().json(ApiError::from("not found")),
        })
}

fn view_secret_page(
    path: web::Path<String>,
    state: web::Data<AppState>,
) -> impl Future<Item = NamedFile, Error = actix_web::Error> {
    let secret_id = path.into_inner();

    state.fetch_secret(&secret_id).and_then(|res| match res {
        Some(_) => Ok(NamedFile::open("./static/view.html")?),

        None => {
            let file = NamedFile::open("./static/404.html")?
                // TODO: (requires HTTP crate).set_status_code()
                ;

            Ok(file)
        }
    })
}

fn config_from_env() -> ServerConfig {
    ServerConfig {
        redis_host: std::env::var("REDIS_HOST").unwrap_or("localhost:6379".to_string()),
        bind_host: std::env::var("BIND_HOST").unwrap_or("127.0.0.1:8080".to_string()),
        app_config: AppConfig {
            max_secret_length: std::env::var("MAX_SECRET_LENGTH")
                .unwrap_or("4096".to_string())
                .parse()
                .expect("unvalid integer value"),
            max_secret_expiration: std::env::var("MAX_SECRET_EXPIRATION")
                .unwrap_or("4096".to_string())
                .parse()
                .expect("unvalid integer value"),
            redis_key_prefix: std::env::var("REDIS_KEY_PREFIX")
                .unwrap_or("tofu:".to_string()),
        },
    }
}

fn main() -> io::Result<()> {
    env_logger::init();
    dotenv().ok();

    let config = config_from_env();
    println!("Booting with config: {:?}", config);

    let sys = actix::System::new("tofu");

    let redis_addr = RedisActor::start(config.redis_host);
    let app_config = config.app_config;

    HttpServer::new(move || {
        let state = AppState {
            redis: redis_addr.clone(),
            config: app_config.clone(),
        };

        App::new()
            .data(state)
            .wrap(middleware::Logger::default())
            .service(
                web::resource("/api/secret")
                    .route(web::post().to_async(api_create_secret)),
            )
            .service(
                web::resource("/api/secret/{secret_id}")
                    .route(web::get().to_async(api_get_secret)),
            )
            .service(
                web::resource("/view/{secret_id}")
                    .route(web::get().to_async(view_secret_page)),
            )
            .service(Files::new("/", "./static/").index_file("index.html"))
    })
    .bind(config.bind_host)?
    .start();

    sys.run()
}
<|endoftext|>


<commit_before>// src/common/authentication.rs

/// Authenticate the user (JWT Token)

// Import
// External
use chrono::{DateTime, Local};
use crypto::sha2::Sha256;
use hyper::header::{self, Authorization, Bearer};
use jwt::{Header, Registered, Token};
use nickel::{MiddlewareResult, Response, Request};
use nickel::status::StatusCode::{Forbidden};
// Common Utilities
use ::common::mm_result::{MMResult, MMError, MMErrorKind};
use ::common::data_access::ServerData;
use ::common::config::Config;

/// Authentication Middleware
///
/// # Arguments
/// request - &mut Request<ServerData>
/// response - Response<'mw, ServerData>
///
/// # Returns
/// `MiddlewareResult<'mw, ServerData>`
pub fn authenticator<'mw>(request: &mut Request<ServerData>, response: Response<'mw, ServerData> ) -> MiddlewareResult<'mw, ServerData> {
    let server_data: &ServerData = request.server_data();

    // Check if we are getting an OPTIONS request
    if request.origin.method.to_string() == "OPTIONS".to_string() {
        // The middleware should not be used for OPTIONS, so continue
        response.next_middleware()
    } else {
        match request.origin.uri.to_string().as_ref(){
            // Public routes don't require authentication
            "/account/login" | "/account" => {
                response.next_middleware()
            },
            _ => {
                let token = match get_jwt_from_header(request){
                    Ok(jwt_token) => jwt_token,
                    Err(_) => {
                        return response.error(Forbidden, "Access denied. Unable to parse token.")
                    }
                };

                if let Some(ref auth_secret) = server_data.config.auth.auth_secret{
                    let secret = auth_secret.as_bytes();

                    match validate_token(token, &secret, server_data.config.clone()){
                        Ok(_) => {
                            return response.next_middleware();
                        },
                        Err(e) => {
                            return response.error(Forbidden, e.get_message());
                        }
                    }
                }else{
                    error!("Authentication failure. Unable to verify JWT Token. No auth_secret key.");
                }
                response.error(Forbidden, "Access denied. Invalid token format.")
            }
        }
    }
}// end authenticator

/// Retrieve the sub claim value from a JWT Token
/// Verify the signature before returning the value
///
/// # Arguments
/// token - Token<Header, Registered>
/// secret - &[u8]
///
/// # Returns
/// `MMResult<()>` Ok if token is verified, Err otherwise
pub fn get_session_id(request: &Request<ServerData>) -> MMResult<String>{
    let server_data: &ServerData = request.server_data();

    let token = match get_jwt_from_header(request){
        Ok(jwt_token) => jwt_token,
        Err(e) => {
            return Err(MMError::new(e.get_message(), MMErrorKind::Other));
        }
    };

    if let Some(ref auth_secret) = server_data.config.auth.auth_secret{
        let secret = auth_secret.as_bytes();

        return get_token_sub(token, &secret);
    }
    error!("Authentication failure. Unable to verify JWT Token. No auth_secret key.");

    Err(MMError::new("Token claim cannot be retrieved.", MMErrorKind::Other))
}// end get_session_id

/// Validate a JWT Token
/// Verify the signature and validate claims
///
/// # Arguments
/// token - Token<Header, Registered>
/// secret - &[u8]
/// config - Config
///
/// # Returns
/// `MMResult<()>` Ok if token is verified, Err otherwise
fn validate_token(token: Token<Header, Registered>, secret: &[u8], config: Config) -> MMResult<()>{
    // Verify the token
    if token.verify(secret, Sha256::new()) {
        // Verify The claims
        let current_time: DateTime<Local> = Local::now();

        // Verify Issuer claim
        if let Some(ref claim_iss) = config.auth.claim_iss{
            match token.claims.iss{
                Some(iss) => {
                    if iss.ne(claim_iss){
                        //return response.error(Forbidden, "Access denied. Invalid token issuer.");
                        return Err(MMError::new("Access denied. Invalid token issuer.", MMErrorKind::Other));
                    }
                },
                None => {
                    //return response.error(Forbidden, "Access denied. Token invalid. Claim iss is required).");
                    return Err(MMError::new("Access denied. Token invalid. Claim iss is required).", MMErrorKind::Other));
                }
            }
        }

        // Verify Issued-At claim
        if let Some(ref iat_ack_s) = config.auth.iat_ack{
            match token.claims.iat{
                Some(iat) => {
                    match DateTime::parse_from_rfc3339(iat_ack_s.as_str()){
                        Ok(iat_ack) => {
                            if iat as i64 <= iat_ack.timestamp(){
                                // Token was issued before iat_ack
                                //return response.error(Forbidden, "Access denied. Expired token (iat).");
                                return Err(MMError::new("Access denied. Expired token (iat).", MMErrorKind::Other));
                            }
                        },
                        Err(e) => {
                            error!("Authentication: iat_inception format invalid. {}", e);
                        }
                    }
                },
                None => {
                    //return response.error(Forbidden, "Access denied. Token invalid. Claim iat is required).");
                    return Err(MMError::new("Access denied. Token invalid. Claim iat is required).", MMErrorKind::Other));
                }
            }
        }

        // Verify Expiration claim
        match token.claims.exp{
            Some(exp) => {
                if exp as i64 <= current_time.timestamp(){
                    //return response.error(Forbidden, "Access denied. Expired token (exp).");
                    return Err(MMError::new("Access denied. Expired token (exp).", MMErrorKind::Other));
                }
            },
            None => {
                //return response.error(Forbidden, "Access denied. Token invalid. Claim exp is required).");
                return Err(MMError::new("Access denied. Token invalid. Claim exp is required).", MMErrorKind::Other));
            }
        }

        // Token signature and claims verified
        return Ok(());
    } else {
        //return response.error(Forbidden, "Access denied. Invalid token.");
        return Err(MMError::new("Access denied. Invalid token.", MMErrorKind::Other));
    }
}// end validate_token

/// Retrieve the JWT from request headers
///
/// # Arguments
/// request - &Request<ServerData>
///
/// # Returns
/// `MMResult<Token::<Header, Registered>>` The JWT token
fn get_jwt_from_header(request: &Request<ServerData>) -> MMResult<Token<Header, Registered>>{
    // Get the full Authorization header from the incoming request headers
    let auth_header = match request.origin.headers.get::<Authorization<Bearer>>() {
        Some(header) => header,
        None => {
            return Err(MMError::new("Token cannot be retrieved from header.", MMErrorKind::Other));
        }
    };

    // Format the header to only take the value
    let jwt = header::HeaderFormatter(auth_header).to_string();

    // Pass "Bearer" and retrieve just the token string
    let jwt_slice = &jwt[7..];

    // parse the JWT token
    match Token::<Header, Registered>::parse(jwt_slice){
        Ok(token) => Ok(token),
        Err(e) => {
            warn!("{:?}", e);
            Err(MMError::new("Unable to parse JWT Token.", MMErrorKind::Other))
        }
    }
}// end get_jwt_from_header

/// Retrieve the sub claim value from a JWT Token
/// Verify the signature before returning the value
///
/// # Arguments
/// token - Token<Header, Registered>
/// secret - &[u8]
///
/// # Returns
/// `MMResult<()>` Ok if token is verified, Err otherwise
fn get_token_sub(token: Token<Header, Registered>, secret: &[u8]) -> MMResult<String>{
    // Verify the token
    if token.verify(secret, Sha256::new()) {

        //return the sub claim
        if let Some(sub) = token.claims.sub{
            return Ok(sub);
        }
    }
    Err(MMError::new("Token claim cannot be retrieved.", MMErrorKind::Other))
}// end get_token_sub
<commit_msg>Make account api endpoint private<commit_after>// src/common/authentication.rs

/// Authenticate the user (JWT Token)

// Import
// External
use chrono::{DateTime, Local};
use crypto::sha2::Sha256;
use hyper::header::{self, Authorization, Bearer};
use jwt::{Header, Registered, Token};
use nickel::{MiddlewareResult, Response, Request};
use nickel::status::StatusCode::{Forbidden};
// Common Utilities
use ::common::mm_result::{MMResult, MMError, MMErrorKind};
use ::common::data_access::ServerData;
use ::common::config::Config;

/// Authentication Middleware
///
/// # Arguments
/// request - &mut Request<ServerData>
/// response - Response<'mw, ServerData>
///
/// # Returns
/// `MiddlewareResult<'mw, ServerData>`
pub fn authenticator<'mw>(request: &mut Request<ServerData>, response: Response<'mw, ServerData> ) -> MiddlewareResult<'mw, ServerData> {
    let server_data: &ServerData = request.server_data();

    // Check if we are getting an OPTIONS request
    if request.origin.method.to_string() == "OPTIONS".to_string() {
        // The middleware should not be used for OPTIONS, so continue
        response.next_middleware()
    } else {
        match request.origin.uri.to_string().as_ref(){
            // Public routes don't require authentication
            "/account/login" => {
                response.next_middleware()
            },
            _ => {
                let token = match get_jwt_from_header(request){
                    Ok(jwt_token) => jwt_token,
                    Err(_) => {
                        return response.error(Forbidden, "Access denied. Unable to parse token.")
                    }
                };

                if let Some(ref auth_secret) = server_data.config.auth.auth_secret{
                    let secret = auth_secret.as_bytes();

                    match validate_token(token, &secret, server_data.config.clone()){
                        Ok(_) => {
                            return response.next_middleware();
                        },
                        Err(e) => {
                            return response.error(Forbidden, e.get_message());
                        }
                    }
                }else{
                    error!("Authentication failure. Unable to verify JWT Token. No auth_secret key.");
                }
                response.error(Forbidden, "Access denied. Invalid token format.")
            }
        }
    }
}// end authenticator

/// Retrieve the sub claim value from a JWT Token
/// Verify the signature before returning the value
///
/// # Arguments
/// token - Token<Header, Registered>
/// secret - &[u8]
///
/// # Returns
/// `MMResult<()>` Ok if token is verified, Err otherwise
pub fn get_session_id(request: &Request<ServerData>) -> MMResult<String>{
    let server_data: &ServerData = request.server_data();

    let token = match get_jwt_from_header(request){
        Ok(jwt_token) => jwt_token,
        Err(e) => {
            return Err(MMError::new(e.get_message(), MMErrorKind::Other));
        }
    };

    if let Some(ref auth_secret) = server_data.config.auth.auth_secret{
        let secret = auth_secret.as_bytes();

        return get_token_sub(token, &secret);
    }
    error!("Authentication failure. Unable to verify JWT Token. No auth_secret key.");

    Err(MMError::new("Token claim cannot be retrieved.", MMErrorKind::Other))
}// end get_session_id

/// Validate a JWT Token
/// Verify the signature and validate claims
///
/// # Arguments
/// token - Token<Header, Registered>
/// secret - &[u8]
/// config - Config
///
/// # Returns
/// `MMResult<()>` Ok if token is verified, Err otherwise
fn validate_token(token: Token<Header, Registered>, secret: &[u8], config: Config) -> MMResult<()>{
    // Verify the token
    if token.verify(secret, Sha256::new()) {
        // Verify The claims
        let current_time: DateTime<Local> = Local::now();

        // Verify Issuer claim
        if let Some(ref claim_iss) = config.auth.claim_iss{
            match token.claims.iss{
                Some(iss) => {
                    if iss.ne(claim_iss){
                        //return response.error(Forbidden, "Access denied. Invalid token issuer.");
                        return Err(MMError::new("Access denied. Invalid token issuer.", MMErrorKind::Other));
                    }
                },
                None => {
                    //return response.error(Forbidden, "Access denied. Token invalid. Claim iss is required).");
                    return Err(MMError::new("Access denied. Token invalid. Claim iss is required).", MMErrorKind::Other));
                }
            }
        }

        // Verify Issued-At claim
        if let Some(ref iat_ack_s) = config.auth.iat_ack{
            match token.claims.iat{
                Some(iat) => {
                    match DateTime::parse_from_rfc3339(iat_ack_s.as_str()){
                        Ok(iat_ack) => {
                            if iat as i64 <= iat_ack.timestamp(){
                                // Token was issued before iat_ack
                                //return response.error(Forbidden, "Access denied. Expired token (iat).");
                                return Err(MMError::new("Access denied. Expired token (iat).", MMErrorKind::Other));
                            }
                        },
                        Err(e) => {
                            error!("Authentication: iat_inception format invalid. {}", e);
                        }
                    }
                },
                None => {
                    //return response.error(Forbidden, "Access denied. Token invalid. Claim iat is required).");
                    return Err(MMError::new("Access denied. Token invalid. Claim iat is required).", MMErrorKind::Other));
                }
            }
        }

        // Verify Expiration claim
        match token.claims.exp{
            Some(exp) => {
                if exp as i64 <= current_time.timestamp(){
                    //return response.error(Forbidden, "Access denied. Expired token (exp).");
                    return Err(MMError::new("Access denied. Expired token (exp).", MMErrorKind::Other));
                }
            },
            None => {
                //return response.error(Forbidden, "Access denied. Token invalid. Claim exp is required).");
                return Err(MMError::new("Access denied. Token invalid. Claim exp is required).", MMErrorKind::Other));
            }
        }

        // Token signature and claims verified
        return Ok(());
    } else {
        //return response.error(Forbidden, "Access denied. Invalid token.");
        return Err(MMError::new("Access denied. Invalid token.", MMErrorKind::Other));
    }
}// end validate_token

/// Retrieve the JWT from request headers
///
/// # Arguments
/// request - &Request<ServerData>
///
/// # Returns
/// `MMResult<Token::<Header, Registered>>` The JWT token
fn get_jwt_from_header(request: &Request<ServerData>) -> MMResult<Token<Header, Registered>>{
    // Get the full Authorization header from the incoming request headers
    let auth_header = match request.origin.headers.get::<Authorization<Bearer>>() {
        Some(header) => header,
        None => {
            return Err(MMError::new("Token cannot be retrieved from header.", MMErrorKind::Other));
        }
    };

    // Format the header to only take the value
    let jwt = header::HeaderFormatter(auth_header).to_string();

    // Pass "Bearer" and retrieve just the token string
    let jwt_slice = &jwt[7..];

    // parse the JWT token
    match Token::<Header, Registered>::parse(jwt_slice){
        Ok(token) => Ok(token),
        Err(e) => {
            warn!("{:?}", e);
            Err(MMError::new("Unable to parse JWT Token.", MMErrorKind::Other))
        }
    }
}// end get_jwt_from_header

/// Retrieve the sub claim value from a JWT Token
/// Verify the signature before returning the value
///
/// # Arguments
/// token - Token<Header, Registered>
/// secret - &[u8]
///
/// # Returns
/// `MMResult<()>` Ok if token is verified, Err otherwise
fn get_token_sub(token: Token<Header, Registered>, secret: &[u8]) -> MMResult<String>{
    // Verify the token
    if token.verify(secret, Sha256::new()) {

        //return the sub claim
        if let Some(sub) = token.claims.sub{
            return Ok(sub);
        }
    }
    Err(MMError::new("Token claim cannot be retrieved.", MMErrorKind::Other))
}// end get_token_sub
<|endoftext|>


<commit_before>use crate::size_hint;
use crate::Itertools;

use alloc::vec::Vec;
use std::iter::FusedIterator;
use std::mem::replace;
use std::fmt;

/// Head element and Tail iterator pair
///
/// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on
/// first items (which are guaranteed to exist).
///
/// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the heap used in
/// `KMerge` into a min-heap.
#[derive(Debug)]
struct HeadTail<I>
    where I: Iterator
{
    head: I::Item,
    tail: I,
}

impl<I> HeadTail<I>
    where I: Iterator
{
    /// Constructs a `HeadTail` from an `Iterator`. Returns `None` if the `Iterator` is empty.
    fn new(mut it: I) -> Option<HeadTail<I>> {
        let head = it.next();
        head.map(|h| {
            HeadTail {
                head: h,
                tail: it,
            }
        })
    }

    /// Get the next element and update `head`, returning the old head in `Some`.
    ///
    /// Returns `None` when the tail is exhausted (only `head` then remains).
    fn next(&mut self) -> Option<I::Item> {
        if let Some(next) = self.tail.next() {
            Some(replace(&mut self.head, next))
        } else {
            None
        }
    }

    /// Hints at the size of the sequence, same as the `Iterator` method.
    fn size_hint(&self) -> (usize, Option<usize>) {
        size_hint::add_scalar(self.tail.size_hint(), 1)
    }
}

impl<I> Clone for HeadTail<I>
    where I: Iterator + Clone,
          I::Item: Clone
{
    clone_fields!(head, tail);
}

/// Make `data` a heap (min-heap w.r.t the sorting).
fn heapify<T, S>(data: &mut [T], mut less_than: S)
    where S: FnMut(&T, &T) -> bool
{
    for i in (0..data.len() / 2).rev() {
        sift_down(data, i, &mut less_than);
    }
}

/// Sift down element at `index` (`heap` is a min-heap wrt the ordering)
fn sift_down<T, S>(heap: &mut [T], index: usize, mut less_than: S)
    where S: FnMut(&T, &T) -> bool
{
    debug_assert!(index <= heap.len());
    let mut pos = index;
    let mut child = 2 * pos + 1;
    // Require the right child to be present
    // This allows to find the index of the smallest child without a branch
    // that wouldn't be predicted if present
    while child + 1 < heap.len() {
        // pick the smaller of the two children
        // use aritmethic to avoid an unpredictable branch
        child += less_than(&heap[child+1], &heap[child]) as usize;

        // sift down is done if we are already in order
        if !less_than(&heap[child], &heap[pos]) {
            return;
        }
        heap.swap(pos, child);
        pos = child;
        child = 2 * pos + 1;
    }
    // Check if the last (left) child was an only child
    // if it is then it has to be compared with the parent
    if child + 1 == heap.len() && less_than(&heap[child], &heap[pos]) {
        heap.swap(pos, child);
    }
}

/// An iterator adaptor that merges an abitrary number of base iterators in ascending order.
/// If all base iterators are sorted (ascending), the result is sorted.
///
/// Iterator element type is `I::Item`.
///
/// See [`.kmerge()`](crate::Itertools::kmerge) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub type KMerge<I> = KMergeBy<I, KMergeByLt>;

pub trait KMergePredicate<T> {
    fn kmerge_pred(&mut self, a: &T, b: &T) -> bool;
}

#[derive(Clone, Debug)]
pub struct KMergeByLt;

impl<T: PartialOrd> KMergePredicate<T> for KMergeByLt {
    fn kmerge_pred(&mut self, a: &T, b: &T) -> bool {
        a < b
    }
}

impl<T, F: FnMut(&T, &T)->bool> KMergePredicate<T> for F {
    fn kmerge_pred(&mut self, a: &T, b: &T) -> bool {
        self(a, b)
    }
}

/// Create an iterator that merges elements of the contained iterators using
/// the ordering function.
///
/// [`IntoIterator`] enabled version of [`Itertools::kmerge`].
///
/// ```
/// use itertools::kmerge;
///
/// for elt in kmerge(vec![vec![0, 2, 4], vec![1, 3, 5], vec![6, 7]]) {
///     /* loop body */
/// }
/// ```
pub fn kmerge<I>(iterable: I) -> KMerge<<I::Item as IntoIterator>::IntoIter>
    where I: IntoIterator,
          I::Item: IntoIterator,
          <<I as IntoIterator>::Item as IntoIterator>::Item: PartialOrd
{
    kmerge_by(iterable, KMergeByLt)
}

/// An iterator adaptor that merges an abitrary number of base iterators
/// according to an ordering function.
///
/// Iterator element type is `I::Item`.
///
/// See [`.kmerge_by()`](crate::Itertools::kmerge_by) for more
/// information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct KMergeBy<I, F>
    where I: Iterator,
{
    heap: Vec<HeadTail<I>>,
    less_than: F,
}

impl<I, F> fmt::Debug for KMergeBy<I, F>
    where I: Iterator + fmt::Debug,
          I::Item: fmt::Debug,
{
    debug_fmt_fields!(KMergeBy, heap);
}

/// Create an iterator that merges elements of the contained iterators.
///
/// Equivalent to `iterable.into_iter().kmerge_by(less_than)`.
pub fn kmerge_by<I, F>(iterable: I, mut less_than: F)
    -> KMergeBy<<I::Item as IntoIterator>::IntoIter, F>
    where I: IntoIterator,
          I::Item: IntoIterator,
          F: KMergePredicate<<<I as IntoIterator>::Item as IntoIterator>::Item>,
{
    let iter = iterable.into_iter();
    let (lower, _) = iter.size_hint();
    let mut heap: Vec<_> = Vec::with_capacity(lower);
    heap.extend(iter.filter_map(|it| HeadTail::new(it.into_iter())));
    heapify(&mut heap, |a, b| less_than.kmerge_pred(&a.head, &b.head));
    KMergeBy { heap, less_than }
}

impl<I, F> Clone for KMergeBy<I, F>
    where I: Iterator + Clone,
          I::Item: Clone,
          F: Clone,
{
    clone_fields!(heap, less_than);
}

impl<I, F> Iterator for KMergeBy<I, F>
    where I: Iterator,
          F: KMergePredicate<I::Item>
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        if self.heap.is_empty() {
            return None;
        }
        let result = if let Some(next) = self.heap[0].next() {
            next
        } else {
            self.heap.swap_remove(0).head
        };
        let less_than = &mut self.less_than;
        sift_down(&mut self.heap, 0, |a, b| less_than.kmerge_pred(&a.head, &b.head));
        Some(result)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.heap.iter()
                 .map(|i| i.size_hint())
                 .fold1(size_hint::add)
                 .unwrap_or((0, Some(0)))
    }
}

impl<I, F> FusedIterator for KMergeBy<I, F>
    where I: Iterator,
          F: KMergePredicate<I::Item>
{}
<commit_msg>Canonicalize some free functions (3.2) kmerge_by: documentation<commit_after>use crate::size_hint;
use crate::Itertools;

use alloc::vec::Vec;
use std::iter::FusedIterator;
use std::mem::replace;
use std::fmt;

/// Head element and Tail iterator pair
///
/// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on
/// first items (which are guaranteed to exist).
///
/// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the heap used in
/// `KMerge` into a min-heap.
#[derive(Debug)]
struct HeadTail<I>
    where I: Iterator
{
    head: I::Item,
    tail: I,
}

impl<I> HeadTail<I>
    where I: Iterator
{
    /// Constructs a `HeadTail` from an `Iterator`. Returns `None` if the `Iterator` is empty.
    fn new(mut it: I) -> Option<HeadTail<I>> {
        let head = it.next();
        head.map(|h| {
            HeadTail {
                head: h,
                tail: it,
            }
        })
    }

    /// Get the next element and update `head`, returning the old head in `Some`.
    ///
    /// Returns `None` when the tail is exhausted (only `head` then remains).
    fn next(&mut self) -> Option<I::Item> {
        if let Some(next) = self.tail.next() {
            Some(replace(&mut self.head, next))
        } else {
            None
        }
    }

    /// Hints at the size of the sequence, same as the `Iterator` method.
    fn size_hint(&self) -> (usize, Option<usize>) {
        size_hint::add_scalar(self.tail.size_hint(), 1)
    }
}

impl<I> Clone for HeadTail<I>
    where I: Iterator + Clone,
          I::Item: Clone
{
    clone_fields!(head, tail);
}

/// Make `data` a heap (min-heap w.r.t the sorting).
fn heapify<T, S>(data: &mut [T], mut less_than: S)
    where S: FnMut(&T, &T) -> bool
{
    for i in (0..data.len() / 2).rev() {
        sift_down(data, i, &mut less_than);
    }
}

/// Sift down element at `index` (`heap` is a min-heap wrt the ordering)
fn sift_down<T, S>(heap: &mut [T], index: usize, mut less_than: S)
    where S: FnMut(&T, &T) -> bool
{
    debug_assert!(index <= heap.len());
    let mut pos = index;
    let mut child = 2 * pos + 1;
    // Require the right child to be present
    // This allows to find the index of the smallest child without a branch
    // that wouldn't be predicted if present
    while child + 1 < heap.len() {
        // pick the smaller of the two children
        // use aritmethic to avoid an unpredictable branch
        child += less_than(&heap[child+1], &heap[child]) as usize;

        // sift down is done if we are already in order
        if !less_than(&heap[child], &heap[pos]) {
            return;
        }
        heap.swap(pos, child);
        pos = child;
        child = 2 * pos + 1;
    }
    // Check if the last (left) child was an only child
    // if it is then it has to be compared with the parent
    if child + 1 == heap.len() && less_than(&heap[child], &heap[pos]) {
        heap.swap(pos, child);
    }
}

/// An iterator adaptor that merges an abitrary number of base iterators in ascending order.
/// If all base iterators are sorted (ascending), the result is sorted.
///
/// Iterator element type is `I::Item`.
///
/// See [`.kmerge()`](crate::Itertools::kmerge) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub type KMerge<I> = KMergeBy<I, KMergeByLt>;

pub trait KMergePredicate<T> {
    fn kmerge_pred(&mut self, a: &T, b: &T) -> bool;
}

#[derive(Clone, Debug)]
pub struct KMergeByLt;

impl<T: PartialOrd> KMergePredicate<T> for KMergeByLt {
    fn kmerge_pred(&mut self, a: &T, b: &T) -> bool {
        a < b
    }
}

impl<T, F: FnMut(&T, &T)->bool> KMergePredicate<T> for F {
    fn kmerge_pred(&mut self, a: &T, b: &T) -> bool {
        self(a, b)
    }
}

/// Create an iterator that merges elements of the contained iterators using
/// the ordering function.
///
/// [`IntoIterator`] enabled version of [`Itertools::kmerge`].
///
/// ```
/// use itertools::kmerge;
///
/// for elt in kmerge(vec![vec![0, 2, 4], vec![1, 3, 5], vec![6, 7]]) {
///     /* loop body */
/// }
/// ```
pub fn kmerge<I>(iterable: I) -> KMerge<<I::Item as IntoIterator>::IntoIter>
    where I: IntoIterator,
          I::Item: IntoIterator,
          <<I as IntoIterator>::Item as IntoIterator>::Item: PartialOrd
{
    kmerge_by(iterable, KMergeByLt)
}

/// An iterator adaptor that merges an abitrary number of base iterators
/// according to an ordering function.
///
/// Iterator element type is `I::Item`.
///
/// See [`.kmerge_by()`](crate::Itertools::kmerge_by) for more
/// information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct KMergeBy<I, F>
    where I: Iterator,
{
    heap: Vec<HeadTail<I>>,
    less_than: F,
}

impl<I, F> fmt::Debug for KMergeBy<I, F>
    where I: Iterator + fmt::Debug,
          I::Item: fmt::Debug,
{
    debug_fmt_fields!(KMergeBy, heap);
}

/// Create an iterator that merges elements of the contained iterators.
///
/// [`IntoIterator`] enabled version of [`Itertools::kmerge_by`].
pub fn kmerge_by<I, F>(iterable: I, mut less_than: F)
    -> KMergeBy<<I::Item as IntoIterator>::IntoIter, F>
    where I: IntoIterator,
          I::Item: IntoIterator,
          F: KMergePredicate<<<I as IntoIterator>::Item as IntoIterator>::Item>,
{
    let iter = iterable.into_iter();
    let (lower, _) = iter.size_hint();
    let mut heap: Vec<_> = Vec::with_capacity(lower);
    heap.extend(iter.filter_map(|it| HeadTail::new(it.into_iter())));
    heapify(&mut heap, |a, b| less_than.kmerge_pred(&a.head, &b.head));
    KMergeBy { heap, less_than }
}

impl<I, F> Clone for KMergeBy<I, F>
    where I: Iterator + Clone,
          I::Item: Clone,
          F: Clone,
{
    clone_fields!(heap, less_than);
}

impl<I, F> Iterator for KMergeBy<I, F>
    where I: Iterator,
          F: KMergePredicate<I::Item>
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        if self.heap.is_empty() {
            return None;
        }
        let result = if let Some(next) = self.heap[0].next() {
            next
        } else {
            self.heap.swap_remove(0).head
        };
        let less_than = &mut self.less_than;
        sift_down(&mut self.heap, 0, |a, b| less_than.kmerge_pred(&a.head, &b.head));
        Some(result)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.heap.iter()
                 .map(|i| i.size_hint())
                 .fold1(size_hint::add)
                 .unwrap_or((0, Some(0)))
    }
}

impl<I, F> FusedIterator for KMergeBy<I, F>
    where I: Iterator,
          F: KMergePredicate<I::Item>
{}
<|endoftext|>