<commit_before>//! Coumpound File Binary format MS-CFB

use std::borrow::Cow;
use std::cmp::min;
use std::io::Read;

use encoding_rs::{Encoding, UTF_16LE};

use errors::*;
use utils::*;

const ENDOFCHAIN: u32 = 0xFFFFFFFE;
const FREESECT: u32 = 0xFFFFFFFF;
const RESERVED_SECTORS: u32 = 0xFFFFFFFA;

/// A struct for managing Compound File Binary format
#[derive(Debug, Clone)]
pub struct Cfb {
    directories: Vec<Directory>,
    sectors: Sectors,
    fats: Vec<u32>,
    mini_sectors: Sectors,
    mini_fats: Vec<u32>,
}

impl Cfb {
    /// Create a new `Cfb`
    ///
    /// Starts reading project metadata (header, directories, sectors and minisectors).
    pub fn new<R: Read>(mut f: &mut R, len: usize) -> Result<Cfb> {

        // load header
        let (h, mut difat) = Header::from_reader(&mut f)?;
        let mut sectors = Sectors::new(h.sector_size, Vec::with_capacity(len));

        // load fat and dif sectors
        debug!("load difat");
        let mut sector_id = h.difat_start;
        while sector_id < RESERVED_SECTORS {
            difat.extend_from_slice(to_u32(sectors.get(sector_id, f)?));
            sector_id = difat.pop().unwrap(); //TODO: check if in infinite loop
        }

        // load the FATs
        debug!("load fat");
        let mut fats = Vec::with_capacity(h.fat_len);
        for id in difat.into_iter().filter(|id| *id != FREESECT) {
            fats.extend_from_slice(to_u32(sectors.get(id, f)?));
        }

        // get the list of directory sectors
        debug!("load directories");
        let dirs = sectors
            .get_chain(h.dir_start, &fats, f, h.dir_len * h.sector_size)?;
        let dirs = dirs.chunks(128)
            .map(|c| Directory::from_slice(c, h.sector_size))
            .collect::<Result<Vec<_>>>()?;

        if dirs.is_empty() || (h.version != 3 && dirs[0].start == ENDOFCHAIN) {
            bail!("Unexpected empty root directory");
        }
        debug!("{:?}", dirs);

        // load the mini streams
        debug!("load minis");
        let ministream = sectors.get_chain(dirs[0].start, &fats, f, dirs[0].len)?;
        let minifat = sectors
            .get_chain(h.mini_fat_start, &fats, f, h.mini_fat_len * h.sector_size)?;
        let minifat = to_u32(&minifat).to_vec();
        Ok(Cfb {
               directories: dirs,
               sectors: sectors,
               fats: fats,
               mini_sectors: Sectors::new(64, ministream),
               mini_fats: minifat,
           })
    }

    /// Checks if directory exists
    pub fn has_directory(&self, name: &str) -> bool {
        self.directories.iter().any(|d| &*d.name == name)
    }

    /// Gets a stream by name out of directories
    pub fn get_stream<R: Read>(&mut self, name: &str, r: &mut R) -> Result<Vec<u8>> {
        match self.directories.iter().find(|d| &*d.name == name) {
            None => Err(format!("Cannot find {} stream", name).into()),
            Some(d) => {
                if d.len < 4096 {
                    // TODO: Study the possibility to return a `VecArray` (stack allocated)
                    self.mini_sectors
                        .get_chain(d.start, &self.mini_fats, r, d.len)
                } else {
                    self.sectors.get_chain(d.start, &self.fats, r, d.len)
                }
            }
        }
    }
}

/// A hidden struct which defines cfb files structure
#[derive(Debug)]
struct Header {
    version: u16,
    sector_size: usize,
    dir_len: usize,
    dir_start: u32,
    fat_len: usize,
    mini_fat_len: usize,
    mini_fat_start: u32,
    difat_start: u32,
}

impl Header {
    fn from_reader<R: Read>(f: &mut R) -> Result<(Header, Vec<u32>)> {
        let mut buf = [0u8; 512];
        f.read_exact(&mut buf)?;

        // check ole signature
        if read_slice::<u64>(buf.as_ref()) != 0xE11AB1A1E011CFD0 {
            bail!("invalid OLE signature (not an office document?)");
        }

        let version = read_u16(&buf[26..28]);

        let sector_size = match read_u16(&buf[30..32]) {
            0x0009 => 512,
            0x000C => {
                // sector size is 4096 bytes, but header is 512 bytes,
                // so the remaining sector bytes have to be read
                let mut buf_end = [0u8; 3584];
                f.read_exact(&mut buf_end)?;
                4096
            }
            s => bail!("Invalid sector shift, expecting 0x09 or 0x0C, got {:x}", s),
        };

        if read_u16(&buf[32..34]) != 0x0006 {
            bail!("Invalid minisector shift");
        }

        let dir_len = read_usize(&buf[40..44]);
        let fat_len = read_usize(&buf[44..48]);
        let dir_start = read_u32(&buf[48..52]);
        let mini_fat_start = read_u32(&buf[60..64]);
        let mini_fat_len = read_usize(&buf[64..68]);
        let difat_start = read_u32(&buf[68..72]);
        let difat_len = read_usize(&buf[62..76]);

        let mut difat = Vec::with_capacity(difat_len);
        difat.extend_from_slice(to_u32(&buf[76..512]));

        Ok((Header {
                version: version,
                sector_size: sector_size,
                dir_len: dir_len,
                fat_len: fat_len,
                dir_start: dir_start,
                mini_fat_len: mini_fat_len,
                mini_fat_start: mini_fat_start,
                difat_start: difat_start,
            },
            difat))
    }
}

/// A struct corresponding to the elementary block of memory
///
/// `data` will persist in memory to ensure the file is read once
#[derive(Debug, Clone)]
struct Sectors {
    data: Vec<u8>,
    size: usize,
}

impl Sectors {
    fn new(size: usize, data: Vec<u8>) -> Sectors {
        Sectors {
            data: data,
            size: size,
        }
    }

    fn get<R: Read>(&mut self, id: u32, r: &mut R) -> Result<&[u8]> {
        let start = id as usize * self.size;
        let end = start + self.size;
        if end > self.data.len() {
            let len = self.data.len();
            unsafe {
                self.data.set_len(end);
            }
            r.read_exact(&mut self.data[len..end])?;
        }
        Ok(&self.data[start..end])
    }

    fn get_chain<R: Read>(&mut self,
                          mut sector_id: u32,
                          fats: &[u32],
                          r: &mut R,
                          len: usize)
                          -> Result<Vec<u8>> {
        let mut chain = if len > 0 {
            Vec::with_capacity(len)
        } else {
            Vec::new()
        };
        while sector_id != ENDOFCHAIN {
            chain.extend_from_slice(self.get(sector_id, r)?);
            sector_id = fats[sector_id as usize];
        }
        if len > 0 {
            chain.truncate(len);
        }
        Ok(chain)
    }
}

/// A struct representing sector organizations, behaves similarly to a tree
#[derive(Debug, Clone)]
struct Directory {
    name: String,
    start: u32,
    len: usize,
}

impl Directory {
    fn from_slice(buf: &[u8], sector_size: usize) -> Result<Directory> {
        let mut name = UTF_16LE.decode(&buf[..64]).0.into_owned();
        if let Some(l) = name.as_bytes().iter().position(|b| *b == 0) {
            name.truncate(l);
        }
        let start = read_u32(&buf[116..120]);
        let len = if sector_size == 512 {
            read_slice::<u32>(&buf[120..124]) as usize
        } else {
            read_slice::<u64>(&buf[120..128]) as usize
        };

        Ok(Directory {
               start: start,
               len: len,
               name: name,
           })
    }
}

/// Decompresses stream
pub fn decompress_stream(s: &[u8]) -> Result<Vec<u8>> {

    const POWER_2: [usize; 16] = [1, 1 << 1, 1 << 2, 1 << 3, 1 << 4, 1 << 5, 1 << 6, 1 << 7,
                                  1 << 8, 1 << 9, 1 << 10, 1 << 11, 1 << 12, 1 << 13, 1 << 14,
                                  1 << 15];

    debug!("decompress stream");
    let mut res = Vec::new();

    if s[0] != 0x01 {
        bail!("invalid signature byte");
    }

    let mut i = 1;
    while i < s.len() {

        let chunk_header = read_u16(&s[i..]);
        i += 2;

        // each 'chunk' is 4096 wide, let's reserve that space
        let start = res.len();
        res.reserve(4096);

        let chunk_size = chunk_header & 0x0FFF;
        let chunk_signature = (chunk_header & 0x7000) >> 12;
        let chunk_flag = (chunk_header & 0x8000) >> 15;

        assert_eq!(chunk_signature, 0b011, "i={}, len={}", i, s.len());

        if chunk_flag == 0 {
            // uncompressed
            res.extend_from_slice(&s[i..i + 4096]);
            i += 4096;
        } else {

            let mut chunk_len = 0;
            let mut buf = [0u8; 4096];
            'chunk: loop {

                let bit_flags = s[i];
                i += 1;
                chunk_len += 1;

                for bit_index in 0..8 {

                    if chunk_len > chunk_size {
                        break 'chunk;
                    }

                    if (bit_flags & (1 << bit_index)) == 0 {
                        // literal token
                        res.push(s[i]);
                        i += 1;
                        chunk_len += 1;
                    } else {
                        // copy token
                        let token = read_u16(&s[i..]);
                        i += 2;
                        chunk_len += 2;

                        let decomp_len = res.len() - start;
                        let bit_count = (4..16).find(|i| POWER_2[*i] >= decomp_len).unwrap();
                        let len_mask = 0xFFFF >> bit_count;
                        let mut len = (token & len_mask) as usize + 3;
                        let offset = ((token & !len_mask) >> (16 - bit_count)) as usize + 1;

                        while len > offset {
                            buf[..offset].copy_from_slice(&res[res.len() - offset..]);
                            res.extend_from_slice(&buf[..offset]);
                            len -= offset;
                        }
                        buf[..len]
                            .copy_from_slice(&res[res.len() - offset..res.len() - offset + len]);
                        res.extend_from_slice(&buf[..len]);
                    }
                }
            }
        }
    }
    Ok(res)
}

#[derive(Clone)]
pub struct XlsEncoding {
    encoding: &'static Encoding,
    pub high_byte: Option<bool>, // None if single byte encoding
}

impl XlsEncoding {
    pub fn from_codepage(codepage: u16) -> Result<XlsEncoding> {
        let e = match encoding_from_windows_code_page(codepage as usize) {
            Some(e) => e,
            None => bail!("Cannot find {} codepage", codepage),
        };
        let high_byte = match codepage {
            20127 | 65000 | 65001 | 20866 | 21866 | 10000 | 10007 | 874 | 1250...1258 |
            28591...28605 => None, // SingleByte encodings
            _ => Some(false),
        };

        Ok(XlsEncoding {
               encoding: e,
               high_byte: high_byte,
           })
    }

    pub fn decode_to(&self, stream: &[u8], len: usize, s: &mut String) -> Result<(usize, usize)> {
        let (l, ub, bytes) = match self.high_byte {
            None => {
                let l = min(stream.len(), len);
                (l, l, Cow::Borrowed(&stream[..l]))
            }
            Some(false) => {
                let l = min(stream.len(), len);

                // add 0x00 high bytes to unicodes
                let mut bytes = vec![0; l * 2];
                for (i, sce) in stream.iter().take(l).enumerate() {
                    bytes[2 * i] = *sce;
                }
                (l, l, Cow::Owned(bytes))
            }
            Some(true) => {
                let l = min(stream.len() / 2, len);
                (l, 2 * l, Cow::Borrowed(&stream[..2 * l]))
            }
        };

        s.push_str(&self.encoding.decode(&bytes).0);
        Ok((l, ub))
    }

    pub fn decode_all(&self, stream: &[u8]) -> Result<String> {
        let mut s = String::with_capacity(stream.len());
        let _ = self.decode_to(stream, stream.len(), &mut s)?;
        Ok(s)
    }
}
<commit_msg>Add an additional bound check (#94)<commit_after>//! Coumpound File Binary format MS-CFB

use std::borrow::Cow;
use std::cmp::min;
use std::io::Read;

use encoding_rs::{Encoding, UTF_16LE};

use errors::*;
use utils::*;

const ENDOFCHAIN: u32 = 0xFFFFFFFE;
const FREESECT: u32 = 0xFFFFFFFF;
const RESERVED_SECTORS: u32 = 0xFFFFFFFA;

/// A struct for managing Compound File Binary format
#[derive(Debug, Clone)]
pub struct Cfb {
    directories: Vec<Directory>,
    sectors: Sectors,
    fats: Vec<u32>,
    mini_sectors: Sectors,
    mini_fats: Vec<u32>,
}

impl Cfb {
    /// Create a new `Cfb`
    ///
    /// Starts reading project metadata (header, directories, sectors and minisectors).
    pub fn new<R: Read>(mut f: &mut R, len: usize) -> Result<Cfb> {

        // load header
        let (h, mut difat) = Header::from_reader(&mut f)?;
        let mut sectors = Sectors::new(h.sector_size, Vec::with_capacity(len));

        // load fat and dif sectors
        debug!("load difat");
        let mut sector_id = h.difat_start;
        while sector_id < RESERVED_SECTORS {
            difat.extend_from_slice(to_u32(sectors.get(sector_id, f)?));
            sector_id = difat.pop().unwrap(); //TODO: check if in infinite loop
        }

        // load the FATs
        debug!("load fat");
        let mut fats = Vec::with_capacity(h.fat_len);
        for id in difat.into_iter().filter(|id| *id != FREESECT) {
            fats.extend_from_slice(to_u32(sectors.get(id, f)?));
        }

        // get the list of directory sectors
        debug!("load directories");
        let dirs = sectors
            .get_chain(h.dir_start, &fats, f, h.dir_len * h.sector_size)?;
        let dirs = dirs.chunks(128)
            .map(|c| Directory::from_slice(c, h.sector_size))
            .collect::<Result<Vec<_>>>()?;

        if dirs.is_empty() || (h.version != 3 && dirs[0].start == ENDOFCHAIN) {
            bail!("Unexpected empty root directory");
        }
        debug!("{:?}", dirs);

        // load the mini streams
        debug!("load minis");
        let ministream = sectors.get_chain(dirs[0].start, &fats, f, dirs[0].len)?;
        let minifat = sectors
            .get_chain(h.mini_fat_start, &fats, f, h.mini_fat_len * h.sector_size)?;
        let minifat = to_u32(&minifat).to_vec();
        Ok(Cfb {
               directories: dirs,
               sectors: sectors,
               fats: fats,
               mini_sectors: Sectors::new(64, ministream),
               mini_fats: minifat,
           })
    }

    /// Checks if directory exists
    pub fn has_directory(&self, name: &str) -> bool {
        self.directories.iter().any(|d| &*d.name == name)
    }

    /// Gets a stream by name out of directories
    pub fn get_stream<R: Read>(&mut self, name: &str, r: &mut R) -> Result<Vec<u8>> {
        match self.directories.iter().find(|d| &*d.name == name) {
            None => Err(format!("Cannot find {} stream", name).into()),
            Some(d) => {
                if d.len < 4096 {
                    // TODO: Study the possibility to return a `VecArray` (stack allocated)
                    self.mini_sectors
                        .get_chain(d.start, &self.mini_fats, r, d.len)
                } else {
                    self.sectors.get_chain(d.start, &self.fats, r, d.len)
                }
            }
        }
    }
}

/// A hidden struct which defines cfb files structure
#[derive(Debug)]
struct Header {
    version: u16,
    sector_size: usize,
    dir_len: usize,
    dir_start: u32,
    fat_len: usize,
    mini_fat_len: usize,
    mini_fat_start: u32,
    difat_start: u32,
}

impl Header {
    fn from_reader<R: Read>(f: &mut R) -> Result<(Header, Vec<u32>)> {
        let mut buf = [0u8; 512];
        f.read_exact(&mut buf)?;

        // check ole signature
        if read_slice::<u64>(buf.as_ref()) != 0xE11AB1A1E011CFD0 {
            bail!("invalid OLE signature (not an office document?)");
        }

        let version = read_u16(&buf[26..28]);

        let sector_size = match read_u16(&buf[30..32]) {
            0x0009 => 512,
            0x000C => {
                // sector size is 4096 bytes, but header is 512 bytes,
                // so the remaining sector bytes have to be read
                let mut buf_end = [0u8; 3584];
                f.read_exact(&mut buf_end)?;
                4096
            }
            s => bail!("Invalid sector shift, expecting 0x09 or 0x0C, got {:x}", s),
        };

        if read_u16(&buf[32..34]) != 0x0006 {
            bail!("Invalid minisector shift");
        }

        let dir_len = read_usize(&buf[40..44]);
        let fat_len = read_usize(&buf[44..48]);
        let dir_start = read_u32(&buf[48..52]);
        let mini_fat_start = read_u32(&buf[60..64]);
        let mini_fat_len = read_usize(&buf[64..68]);
        let difat_start = read_u32(&buf[68..72]);
        let difat_len = read_usize(&buf[62..76]);

        let mut difat = Vec::with_capacity(difat_len);
        difat.extend_from_slice(to_u32(&buf[76..512]));

        Ok((Header {
                version: version,
                sector_size: sector_size,
                dir_len: dir_len,
                fat_len: fat_len,
                dir_start: dir_start,
                mini_fat_len: mini_fat_len,
                mini_fat_start: mini_fat_start,
                difat_start: difat_start,
            },
            difat))
    }
}

/// A struct corresponding to the elementary block of memory
///
/// `data` will persist in memory to ensure the file is read once
#[derive(Debug, Clone)]
struct Sectors {
    data: Vec<u8>,
    size: usize,
}

impl Sectors {
    fn new(size: usize, data: Vec<u8>) -> Sectors {
        Sectors {
            data: data,
            size: size,
        }
    }

    fn get<R: Read>(&mut self, id: u32, r: &mut R) -> Result<&[u8]> {
        let start = id as usize * self.size;
        let end = start + self.size;
        if end > self.data.len() {
            let len = self.data.len();
            unsafe {
                self.data.set_len(end);
            }
            r.read_exact(&mut self.data[len..end])?;
        }
        Ok(&self.data[start..end])
    }

    fn get_chain<R: Read>(&mut self,
                          mut sector_id: u32,
                          fats: &[u32],
                          r: &mut R,
                          len: usize)
                          -> Result<Vec<u8>> {
        let mut chain = if len > 0 {
            Vec::with_capacity(len)
        } else {
            Vec::new()
        };
        while sector_id != ENDOFCHAIN {
            chain.extend_from_slice(self.get(sector_id, r)?);
            sector_id = fats[sector_id as usize];
        }
        if len > 0 {
            chain.truncate(len);
        }
        Ok(chain)
    }
}

/// A struct representing sector organizations, behaves similarly to a tree
#[derive(Debug, Clone)]
struct Directory {
    name: String,
    start: u32,
    len: usize,
}

impl Directory {
    fn from_slice(buf: &[u8], sector_size: usize) -> Result<Directory> {
        let mut name = UTF_16LE.decode(&buf[..64]).0.into_owned();
        if let Some(l) = name.as_bytes().iter().position(|b| *b == 0) {
            name.truncate(l);
        }
        let start = read_u32(&buf[116..120]);
        let len = if sector_size == 512 {
            read_slice::<u32>(&buf[120..124]) as usize
        } else {
            read_slice::<u64>(&buf[120..128]) as usize
        };

        Ok(Directory {
               start: start,
               len: len,
               name: name,
           })
    }
}

/// Decompresses stream
pub fn decompress_stream(s: &[u8]) -> Result<Vec<u8>> {

    const POWER_2: [usize; 16] = [1, 1 << 1, 1 << 2, 1 << 3, 1 << 4, 1 << 5, 1 << 6, 1 << 7,
                                  1 << 8, 1 << 9, 1 << 10, 1 << 11, 1 << 12, 1 << 13, 1 << 14,
                                  1 << 15];

    debug!("decompress stream");
    let mut res = Vec::new();

    if s[0] != 0x01 {
        bail!("invalid signature byte");
    }

    let mut i = 1;
    while i < s.len() {

        let chunk_header = read_u16(&s[i..]);
        i += 2;

        // each 'chunk' is 4096 wide, let's reserve that space
        let start = res.len();
        res.reserve(4096);

        let chunk_size = chunk_header & 0x0FFF;
        let chunk_signature = (chunk_header & 0x7000) >> 12;
        let chunk_flag = (chunk_header & 0x8000) >> 15;

        assert_eq!(chunk_signature, 0b011, "i={}, len={}", i, s.len());

        if chunk_flag == 0 {
            // uncompressed
            res.extend_from_slice(&s[i..i + 4096]);
            i += 4096;
        } else {

            let mut chunk_len = 0;
            let mut buf = [0u8; 4096];
            'chunk: loop {

                if i >= s.len() {
                    break;
                }

                let bit_flags = s[i];
                i += 1;
                chunk_len += 1;

                for bit_index in 0..8 {

                    if chunk_len > chunk_size {
                        break 'chunk;
                    }

                    if (bit_flags & (1 << bit_index)) == 0 {
                        // literal token
                        res.push(s[i]);
                        i += 1;
                        chunk_len += 1;
                    } else {
                        // copy token
                        let token = read_u16(&s[i..]);
                        i += 2;
                        chunk_len += 2;

                        let decomp_len = res.len() - start;
                        let bit_count = (4..16).find(|i| POWER_2[*i] >= decomp_len).unwrap();
                        let len_mask = 0xFFFF >> bit_count;
                        let mut len = (token & len_mask) as usize + 3;
                        let offset = ((token & !len_mask) >> (16 - bit_count)) as usize + 1;

                        while len > offset {
                            buf[..offset].copy_from_slice(&res[res.len() - offset..]);
                            res.extend_from_slice(&buf[..offset]);
                            len -= offset;
                        }
                        buf[..len]
                            .copy_from_slice(&res[res.len() - offset..res.len() - offset + len]);
                        res.extend_from_slice(&buf[..len]);
                    }
                }
            }
        }
    }
    Ok(res)
}

#[derive(Clone)]
pub struct XlsEncoding {
    encoding: &'static Encoding,
    pub high_byte: Option<bool>, // None if single byte encoding
}

impl XlsEncoding {
    pub fn from_codepage(codepage: u16) -> Result<XlsEncoding> {
        let e = match encoding_from_windows_code_page(codepage as usize) {
            Some(e) => e,
            None => bail!("Cannot find {} codepage", codepage),
        };
        let high_byte = match codepage {
            20127 | 65000 | 65001 | 20866 | 21866 | 10000 | 10007 | 874 | 1250...1258 |
            28591...28605 => None, // SingleByte encodings
            _ => Some(false),
        };

        Ok(XlsEncoding {
               encoding: e,
               high_byte: high_byte,
           })
    }

    pub fn decode_to(&self, stream: &[u8], len: usize, s: &mut String) -> Result<(usize, usize)> {
        let (l, ub, bytes) = match self.high_byte {
            None => {
                let l = min(stream.len(), len);
                (l, l, Cow::Borrowed(&stream[..l]))
            }
            Some(false) => {
                let l = min(stream.len(), len);

                // add 0x00 high bytes to unicodes
                let mut bytes = vec![0; l * 2];
                for (i, sce) in stream.iter().take(l).enumerate() {
                    bytes[2 * i] = *sce;
                }
                (l, l, Cow::Owned(bytes))
            }
            Some(true) => {
                let l = min(stream.len() / 2, len);
                (l, 2 * l, Cow::Borrowed(&stream[..2 * l]))
            }
        };

        s.push_str(&self.encoding.decode(&bytes).0);
        Ok((l, ub))
    }

    pub fn decode_all(&self, stream: &[u8]) -> Result<String> {
        let mut s = String::with_capacity(stream.len());
        let _ = self.decode_to(stream, stream.len(), &mut s)?;
        Ok(s)
    }
}
<|endoftext|>


<commit_before>#![allow(unsafe_code)]
/// We use this because we never use the weak count on the std
/// `Arc`, but we use a LOT of `Arc`'s, so the extra 8 bytes
/// turn into a huge overhead.
use std::{
    alloc::{alloc, dealloc, Layout},
    convert::TryFrom,
    fmt::{self, Debug},
    mem,
    ops::Deref,
    ptr,
    sync::atomic::{AtomicUsize, Ordering},
};

struct ArcInner<T: ?Sized> {
    rc: AtomicUsize,
    inner: T,
}

pub struct Arc<T: ?Sized> {
    ptr: *mut ArcInner<T>,
}

unsafe impl<T: Send + Sync + ?Sized> Send for Arc<T> {}
unsafe impl<T: Send + Sync + ?Sized> Sync for Arc<T> {}

impl<T: Debug> Debug for Arc<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        f.debug_map().entry(&"inner", unsafe { &(*self.ptr).inner }).finish()
    }
}

impl<T> Arc<T> {
    pub fn new(inner: T) -> Arc<T> {
        let bx = Box::new(ArcInner { inner, rc: AtomicUsize::new(1) });
        let ptr = Box::into_raw(bx);
        Arc { ptr }
    }

    pub fn copy_from_slice(s: &[T]) -> Arc<[T]> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());
        let data_width = mem::size_of::<T>().checked_mul(s.len()).unwrap();

        let size = rc_width.checked_add(data_width).unwrap();

        let layout = Layout::from_size_align(size, align).unwrap();

        unsafe {
            let ptr = alloc(layout);

            assert!(!ptr.is_null(), "failed to allocate Arc");
            #[allow(clippy::cast_ptr_alignment)]
            ptr::write(ptr as _, AtomicUsize::new(1));

            let data_ptr = ptr.add(rc_width);
            ptr::copy_nonoverlapping(s.as_ptr(), data_ptr as _, s.len());

            let fat_ptr: *const ArcInner<[T]> = Arc::fatten(ptr, s.len());

            Arc { ptr: fat_ptr as *mut _ }
        }
    }

    /// <https://users.rust-lang.org/t/construct-fat-pointer-to-struct/29198/9>
    #[allow(trivial_casts)]
    fn fatten(data: *const u8, len: usize) -> *const ArcInner<[T]> {
        // Requirements of slice::from_raw_parts.
        assert!(!data.is_null());
        assert!(isize::try_from(len).is_ok());

        let slice =
            unsafe { core::slice::from_raw_parts(data as *const (), len) };
        slice as *const [()] as *const _
    }

    pub fn into_raw(arc: Arc<T>) -> *const T {
        let ptr = unsafe { &(*arc.ptr).inner };
        #[allow(clippy::mem_forget)]
        mem::forget(arc);
        ptr
    }

    pub unsafe fn from_raw(ptr: *const T) -> Arc<T> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());

        let sub_ptr = (ptr as *const u8).sub(rc_width) as *mut ArcInner<T>;

        Arc { ptr: sub_ptr }
    }
}

impl<T: ?Sized> Arc<T> {
    pub fn strong_count(arc: &Arc<T>) -> usize {
        unsafe { (*arc.ptr).rc.load(Ordering::Acquire) }
    }

    pub fn get_mut(arc: &mut Arc<T>) -> Option<&mut T> {
        if Arc::strong_count(arc) == 1 {
            Some(unsafe { &mut arc.ptr.as_mut().unwrap().inner })
        } else {
            None
        }
    }
}

impl<T: ?Sized + Clone> Arc<T> {
    pub fn make_mut(arc: &mut Arc<T>) -> &mut T {
        if Arc::strong_count(arc) != 1 {
            *arc = Arc::new((**arc).clone());
            assert_eq!(Arc::strong_count(arc), 1);
        }
        Arc::get_mut(arc).unwrap()
    }
}

impl<T: Default> Default for Arc<T> {
    fn default() -> Arc<T> {
        Arc::new(T::default())
    }
}

impl<T: ?Sized> Clone for Arc<T> {
    fn clone(&self) -> Arc<T> {
        unsafe {
            (*self.ptr).rc.fetch_add(1, Ordering::Release);
        }
        Arc { ptr: self.ptr }
    }
}

impl<T: ?Sized> Drop for Arc<T> {
    fn drop(&mut self) {
        unsafe {
            let rc = (*self.ptr).rc.fetch_sub(1, Ordering::Release) - 1;
            if rc == 0 {
                Box::from_raw(self.ptr);
            }
        }
    }
}

impl<T: Copy> From<&[T]> for Arc<[T]> {
    #[inline]
    fn from(s: &[T]) -> Arc<[T]> {
        Arc::copy_from_slice(s)
    }
}

#[allow(clippy::fallible_impl_from)]
impl<T> From<Box<T>> for Arc<T> {
    #[inline]
    fn from(b: Box<T>) -> Arc<T> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());

        let size = rc_width.checked_add(mem::size_of::<T>()).unwrap();

        unsafe {
            let dst_layout = Layout::from_size_align(size, align).unwrap();
            let dst = alloc(dst_layout);
            assert!(!dst.is_null(), "failed to allocate Arc");

            let data_ptr = dst.add(rc_width);
            #[allow(clippy::cast_ptr_alignment)]
            ptr::write(dst as _, AtomicUsize::new(1));
            let src = Box::into_raw(b);
            ptr::copy(src, data_ptr as *mut T, 1);

            // free the old box memory without running Drop
            let src_layout = Layout::new::<T>();
            dealloc(src as *mut u8, src_layout);

            Arc { ptr: dst as *mut _ }
        }
    }
}

impl<T> From<Vec<T>> for Arc<[T]> {
    #[inline]
    fn from(mut v: Vec<T>) -> Arc<[T]> {
        unsafe {
            let arc = Arc::copy_from_slice(&v);

            // Allow the Vec to free its memory, but not destroy its contents
            v.set_len(0);

            arc
        }
    }
}

impl<T: ?Sized> Deref for Arc<T> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { &(*self.ptr).inner }
    }
}

impl<T: ?Sized> std::borrow::Borrow<T> for Arc<T> {
    fn borrow(&self) -> &T {
        &**self
    }
}

impl<T: ?Sized> AsRef<T> for Arc<T> {
    fn as_ref(&self) -> &T {
        &**self
    }
}
<commit_msg>make ArcInner repr(C) due to raw memory writing<commit_after>#![allow(unsafe_code)]
/// We use this because we never use the weak count on the std
/// `Arc`, but we use a LOT of `Arc`'s, so the extra 8 bytes
/// turn into a huge overhead.
use std::{
    alloc::{alloc, dealloc, Layout},
    convert::TryFrom,
    fmt::{self, Debug},
    mem,
    ops::Deref,
    ptr,
    sync::atomic::{AtomicUsize, Ordering},
};

// we make this repr(C) because we do a raw
// write to the beginning where we expect
// the rc to be.
#[repr(C)]
struct ArcInner<T: ?Sized> {
    rc: AtomicUsize,
    inner: T,
}

pub struct Arc<T: ?Sized> {
    ptr: *mut ArcInner<T>,
}

unsafe impl<T: Send + Sync + ?Sized> Send for Arc<T> {}
unsafe impl<T: Send + Sync + ?Sized> Sync for Arc<T> {}

impl<T: Debug> Debug for Arc<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        f.debug_map().entry(&"inner", unsafe { &(*self.ptr).inner }).finish()
    }
}

impl<T> Arc<T> {
    pub fn new(inner: T) -> Arc<T> {
        let bx = Box::new(ArcInner { inner, rc: AtomicUsize::new(1) });
        let ptr = Box::into_raw(bx);
        Arc { ptr }
    }

    pub fn copy_from_slice(s: &[T]) -> Arc<[T]> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());
        let data_width = mem::size_of::<T>().checked_mul(s.len()).unwrap();

        let size = rc_width.checked_add(data_width).unwrap();

        let layout = Layout::from_size_align(size, align).unwrap();

        unsafe {
            let ptr = alloc(layout);

            assert!(!ptr.is_null(), "failed to allocate Arc");
            #[allow(clippy::cast_ptr_alignment)]
            ptr::write(ptr as _, AtomicUsize::new(1));

            let data_ptr = ptr.add(rc_width);
            ptr::copy_nonoverlapping(s.as_ptr(), data_ptr as _, s.len());

            let fat_ptr: *const ArcInner<[T]> = Arc::fatten(ptr, s.len());

            Arc { ptr: fat_ptr as *mut _ }
        }
    }

    /// <https://users.rust-lang.org/t/construct-fat-pointer-to-struct/29198/9>
    #[allow(trivial_casts)]
    fn fatten(data: *const u8, len: usize) -> *const ArcInner<[T]> {
        // Requirements of slice::from_raw_parts.
        assert!(!data.is_null());
        assert!(isize::try_from(len).is_ok());

        let slice =
            unsafe { core::slice::from_raw_parts(data as *const (), len) };
        slice as *const [()] as *const _
    }

    pub fn into_raw(arc: Arc<T>) -> *const T {
        let ptr = unsafe { &(*arc.ptr).inner };
        #[allow(clippy::mem_forget)]
        mem::forget(arc);
        ptr
    }

    pub unsafe fn from_raw(ptr: *const T) -> Arc<T> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());

        let sub_ptr = (ptr as *const u8).sub(rc_width) as *mut ArcInner<T>;

        Arc { ptr: sub_ptr }
    }
}

impl<T: ?Sized> Arc<T> {
    pub fn strong_count(arc: &Arc<T>) -> usize {
        unsafe { (*arc.ptr).rc.load(Ordering::Acquire) }
    }

    pub fn get_mut(arc: &mut Arc<T>) -> Option<&mut T> {
        if Arc::strong_count(arc) == 1 {
            Some(unsafe { &mut arc.ptr.as_mut().unwrap().inner })
        } else {
            None
        }
    }
}

impl<T: ?Sized + Clone> Arc<T> {
    pub fn make_mut(arc: &mut Arc<T>) -> &mut T {
        if Arc::strong_count(arc) != 1 {
            *arc = Arc::new((**arc).clone());
            assert_eq!(Arc::strong_count(arc), 1);
        }
        Arc::get_mut(arc).unwrap()
    }
}

impl<T: Default> Default for Arc<T> {
    fn default() -> Arc<T> {
        Arc::new(T::default())
    }
}

impl<T: ?Sized> Clone for Arc<T> {
    fn clone(&self) -> Arc<T> {
        unsafe {
            (*self.ptr).rc.fetch_add(1, Ordering::Release);
        }
        Arc { ptr: self.ptr }
    }
}

impl<T: ?Sized> Drop for Arc<T> {
    fn drop(&mut self) {
        unsafe {
            let rc = (*self.ptr).rc.fetch_sub(1, Ordering::Release) - 1;
            if rc == 0 {
                Box::from_raw(self.ptr);
            }
        }
    }
}

impl<T: Copy> From<&[T]> for Arc<[T]> {
    #[inline]
    fn from(s: &[T]) -> Arc<[T]> {
        Arc::copy_from_slice(s)
    }
}

#[allow(clippy::fallible_impl_from)]
impl<T> From<Box<T>> for Arc<T> {
    #[inline]
    fn from(b: Box<T>) -> Arc<T> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());

        let size = rc_width.checked_add(mem::size_of::<T>()).unwrap();

        unsafe {
            let dst_layout = Layout::from_size_align(size, align).unwrap();
            let dst = alloc(dst_layout);
            assert!(!dst.is_null(), "failed to allocate Arc");

            let data_ptr = dst.add(rc_width);
            #[allow(clippy::cast_ptr_alignment)]
            ptr::write(dst as _, AtomicUsize::new(1));
            let src = Box::into_raw(b);
            ptr::copy(src, data_ptr as *mut T, 1);

            // free the old box memory without running Drop
            let src_layout = Layout::new::<T>();
            dealloc(src as *mut u8, src_layout);

            Arc { ptr: dst as *mut _ }
        }
    }
}

impl<T> From<Vec<T>> for Arc<[T]> {
    #[inline]
    fn from(mut v: Vec<T>) -> Arc<[T]> {
        unsafe {
            let arc = Arc::copy_from_slice(&v);

            // Allow the Vec to free its memory, but not destroy its contents
            v.set_len(0);

            arc
        }
    }
}

impl<T: ?Sized> Deref for Arc<T> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { &(*self.ptr).inner }
    }
}

impl<T: ?Sized> std::borrow::Borrow<T> for Arc<T> {
    fn borrow(&self) -> &T {
        &**self
    }
}

impl<T: ?Sized> AsRef<T> for Arc<T> {
    fn as_ref(&self) -> &T {
        &**self
    }
}
<|endoftext|>


<commit_before>//! Compare two iterator-based trees.

// This clippy seems to be broken, as it has some false triggers in this code.
#![allow(clippy::if_same_then_else)]

use crate::{node::SureNode, Error, Result};
use log::error;
use std::{collections::HashSet, path::Path};

/// This is the mutable state that is threaded through the recursive
/// traversal of the two trees.
struct State<IA, IB> {
    left: SureNode,
    right: SureNode,
    left_iter: IA,
    right_iter: IB,

    // Track warning messages about added and deleted attributes.
    adds: HashSet<String>,
    missings: HashSet<String>,

    // Attributes to be ignored
    ignore: HashSet<String>,
}

pub fn compare_trees<P: AsRef<Path>, IA, IB>(
    mut left: IA,
    mut right: IB,
    dir: P,
    ignore: &[&str],
) -> Result<()>
where
    IA: Iterator<Item = Result<SureNode>>,
    IB: Iterator<Item = Result<SureNode>>,
{
    let mut ignore: HashSet<String> = ignore.iter().map(|x| (*x).to_owned()).collect();
    // The ctime and ino will be different if a backup is restored, and we'd still like to get
    // meaningful results.  Add these to the list of ignored attributes.
    ignore.insert("ctime".to_owned());
    ignore.insert("ino".to_owned());

    let ln = match left.next() {
        None => return Err(Error::EmptyLeftIterator),
        Some(Err(e)) => return Err(e),
        Some(Ok(node)) => node,
    };
    let rn = match right.next() {
        None => return Err(Error::EmptyRightIterator),
        Some(Err(e)) => return Err(e),
        Some(Ok(node)) => node,
    };
    let mut state = State {
        left: ln,
        right: rn,
        left_iter: left,
        right_iter: right,
        adds: HashSet::new(),
        missings: HashSet::new(),
        ignore,
    };

    state.walk_root(dir.as_ref())
}

impl<IA, IB> State<IA, IB>
where
    IA: Iterator<Item = Result<SureNode>>,
    IB: Iterator<Item = Result<SureNode>>,
{
    /// Advance the left iterator.  If it sees the end, it will drop in a
    /// "Leave" node, which shouldn't be visited as long as the tree is
    /// well-formed.
    fn next_left(&mut self) -> Result<()> {
        let next = match self.left_iter.next() {
            None => SureNode::Leave,
            Some(Ok(node)) => node,
            Some(Err(e)) => return Err(e),
        };

        self.left = next;
        Ok(())
    }

    /// Advance the right iterator.  If it sees the end, it will drop in a
    /// "Leave" node, which shouldn't be visited as long as the tree is
    /// well-formed.
    fn next_right(&mut self) -> Result<()> {
        let next = match self.right_iter.next() {
            None => SureNode::Leave,
            Some(Ok(node)) => node,
            Some(Err(e)) => return Err(e),
        };

        self.right = next;
        Ok(())
    }

    fn walk_root(&mut self, dir: &Path) -> Result<()> {
        if !self.left.is_enter() {
            Err(Error::UnexpectedLeftNode)
        } else if !self.right.is_enter() {
            Err(Error::UnexpectedRightNode)
        } else if self.left.name() != "__root__" {
            Err(Error::IncorrectName)
        } else if self.right.name() != "__root__" {
            Err(Error::IncorrectName)
        } else {
            self.compare_enter(dir)?;
            self.next_left()?;
            self.next_right()?;
            self.walk_samedir(dir)
        }
    }

    /// We are within a directory (of the given name) where both trees have
    /// the same directory.  This will recursively compare any children,
    /// and once both have reached the separator, move to `walk_samefiles`.
    fn walk_samedir(&mut self, dir: &Path) -> Result<()> {
        loop {
            match (self.left.is_sep(), self.right.is_sep()) {
                (true, true) => {
                    self.next_left()?;
                    self.next_right()?;
                    return self.walk_samefiles(dir);
                }
                (false, true) => {
                    // The old trees has subdirectories not in this
                    // directory.
                    self.show_delete(dir);
                    self.next_left()?;
                    self.walk_leftdir()?;
                }
                (true, false) => {
                    // The new tree has a newly added directory.
                    self.show_add(dir);
                    self.next_right()?;
                    self.walk_rightdir()?;
                }
                _ if self.left.name() < self.right.name() => {
                    // Old subdirectory.
                    self.show_delete(dir);
                    self.next_left()?;
                    self.walk_leftdir()?;
                }
                _ if self.left.name() > self.right.name() => {
                    // The new tree has a newly added directory.
                    self.show_add(dir);
                    self.next_right()?;
                    self.walk_rightdir()?;
                }
                _ => {
                    // Same named directory.
                    let dirname = dir.join(self.left.name());
                    self.compare_enter(&dirname)?;
                    self.next_left()?;
                    self.next_right()?;
                    self.walk_samedir(&dirname)?;
                }
            }
        }
    }

    /// We are within the files section of the same directory in the two
    /// trees.  Walk through the nodes, reading the Leave node in both, and
    /// returning.
    fn walk_samefiles(&mut self, dir: &Path) -> Result<()> {
        loop {
            match (self.left.is_leave(), self.right.is_leave()) {
                (true, true) => {
                    self.next_left()?;
                    self.next_right()?;
                    return Ok(());
                }
                (false, true) => {
                    self.show_delete(dir);
                    self.next_left()?;
                }
                (true, false) => {
                    self.show_add(dir);
                    self.next_right()?;
                }
                _ if self.left.name() < self.right.name() => {
                    self.show_delete(dir);
                    self.next_left()?;
                }
                _ if self.left.name() > self.right.name() => {
                    self.show_add(dir);
                    self.next_right()?;
                }
                _ => {
                    // Same file.
                    let nodename = dir.join(self.left.name());
                    self.compare_file(&nodename)?;
                    self.next_left()?;
                    self.next_right()?;
                }
            }
        }
    }

    /// Old directory on the left tree.  Walk through nodes recursively to
    /// discard entire tree.
    fn walk_leftdir(&mut self) -> Result<()> {
        loop {
            if self.left.is_enter() {
                self.next_left()?;
                self.walk_leftdir()?;
            } else if self.left.is_leave() {
                self.next_left()?;
                return Ok(());
            } else {
                self.next_left()?;
            }
        }
    }

    /// New directory on the right tree.  Walk through nodes recursively to
    /// discard entire tree.
    fn walk_rightdir(&mut self) -> Result<()> {
        loop {
            if self.right.is_enter() {
                self.next_right()?;
                self.walk_rightdir()?;
            } else if self.right.is_leave() {
                self.next_right()?;
                return Ok(());
            } else {
                self.next_right()?;
            }
        }
    }

    /// Print a message about something added (the name will be the thing
    /// on the right.
    fn show_add(&self, dir: &Path) {
        println!(
            "+ {:22} {:?}",
            self.right.kind(),
            dir.join(self.right.name())
        );
    }

    /// Print a message about something removed (the name will be the thing
    /// on the left.
    fn show_delete(&self, dir: &Path) {
        println!("- {:22} {:?}", self.left.kind(), dir.join(self.left.name()));
    }

    /// Compare the two "Enter" nodes we are visiting.
    fn compare_enter(&mut self, dir: &Path) -> Result<()> {
        self.compare_atts('d', dir)
    }

    /// Compare two file nodes.
    fn compare_file(&mut self, dir: &Path) -> Result<()> {
        self.compare_atts('f', dir)
    }

    /// Attribute comparison.
    fn compare_atts(&mut self, _kind: char, dir: &Path) -> Result<()> {
        let mut old = self.left.atts().unwrap().clone();
        let mut new = self.right.atts().unwrap().clone();
        let mut diffs = vec![];

        for att in self.ignore.iter() {
            old.remove(att);
            new.remove(att);
        }

        for (k, v) in &new {
            match old.get(k) {
                None => {
                    // This attribute is in the new tree, but not the old
                    // one, warn, but only once.
                    if !self.adds.contains(k) {
                        error!("Added attribute: {}", k);
                        self.adds.insert(k.clone());
                    }
                }
                Some(ov) => {
                    if v != ov {
                        diffs.push(k.clone());
                    }
                }
            }
            old.remove(k);
        }

        for k in old.keys() {
            if !self.missings.contains(k) {
                error!("Missing attribute: {}", k);
                self.missings.insert(k.clone());
            }
        }

        if diffs.len() > 0 {
            let mut buf = String::new();
            diffs.sort();
            for d in &diffs {
                if !buf.is_empty() {
                    buf.extend(",".chars());
                }
                buf.extend(d.chars());
            }
            println!("  [{:<20}] {:?}", buf, dir);
        }

        Ok(())
    }
}
<commit_msg>clippy: Use is_empty() instead of len()<commit_after>//! Compare two iterator-based trees.

// This clippy seems to be broken, as it has some false triggers in this code.
#![allow(clippy::if_same_then_else)]

use crate::{node::SureNode, Error, Result};
use log::error;
use std::{collections::HashSet, path::Path};

/// This is the mutable state that is threaded through the recursive
/// traversal of the two trees.
struct State<IA, IB> {
    left: SureNode,
    right: SureNode,
    left_iter: IA,
    right_iter: IB,

    // Track warning messages about added and deleted attributes.
    adds: HashSet<String>,
    missings: HashSet<String>,

    // Attributes to be ignored
    ignore: HashSet<String>,
}

pub fn compare_trees<P: AsRef<Path>, IA, IB>(
    mut left: IA,
    mut right: IB,
    dir: P,
    ignore: &[&str],
) -> Result<()>
where
    IA: Iterator<Item = Result<SureNode>>,
    IB: Iterator<Item = Result<SureNode>>,
{
    let mut ignore: HashSet<String> = ignore.iter().map(|x| (*x).to_owned()).collect();
    // The ctime and ino will be different if a backup is restored, and we'd still like to get
    // meaningful results.  Add these to the list of ignored attributes.
    ignore.insert("ctime".to_owned());
    ignore.insert("ino".to_owned());

    let ln = match left.next() {
        None => return Err(Error::EmptyLeftIterator),
        Some(Err(e)) => return Err(e),
        Some(Ok(node)) => node,
    };
    let rn = match right.next() {
        None => return Err(Error::EmptyRightIterator),
        Some(Err(e)) => return Err(e),
        Some(Ok(node)) => node,
    };
    let mut state = State {
        left: ln,
        right: rn,
        left_iter: left,
        right_iter: right,
        adds: HashSet::new(),
        missings: HashSet::new(),
        ignore,
    };

    state.walk_root(dir.as_ref())
}

impl<IA, IB> State<IA, IB>
where
    IA: Iterator<Item = Result<SureNode>>,
    IB: Iterator<Item = Result<SureNode>>,
{
    /// Advance the left iterator.  If it sees the end, it will drop in a
    /// "Leave" node, which shouldn't be visited as long as the tree is
    /// well-formed.
    fn next_left(&mut self) -> Result<()> {
        let next = match self.left_iter.next() {
            None => SureNode::Leave,
            Some(Ok(node)) => node,
            Some(Err(e)) => return Err(e),
        };

        self.left = next;
        Ok(())
    }

    /// Advance the right iterator.  If it sees the end, it will drop in a
    /// "Leave" node, which shouldn't be visited as long as the tree is
    /// well-formed.
    fn next_right(&mut self) -> Result<()> {
        let next = match self.right_iter.next() {
            None => SureNode::Leave,
            Some(Ok(node)) => node,
            Some(Err(e)) => return Err(e),
        };

        self.right = next;
        Ok(())
    }

    fn walk_root(&mut self, dir: &Path) -> Result<()> {
        if !self.left.is_enter() {
            Err(Error::UnexpectedLeftNode)
        } else if !self.right.is_enter() {
            Err(Error::UnexpectedRightNode)
        } else if self.left.name() != "__root__" {
            Err(Error::IncorrectName)
        } else if self.right.name() != "__root__" {
            Err(Error::IncorrectName)
        } else {
            self.compare_enter(dir)?;
            self.next_left()?;
            self.next_right()?;
            self.walk_samedir(dir)
        }
    }

    /// We are within a directory (of the given name) where both trees have
    /// the same directory.  This will recursively compare any children,
    /// and once both have reached the separator, move to `walk_samefiles`.
    fn walk_samedir(&mut self, dir: &Path) -> Result<()> {
        loop {
            match (self.left.is_sep(), self.right.is_sep()) {
                (true, true) => {
                    self.next_left()?;
                    self.next_right()?;
                    return self.walk_samefiles(dir);
                }
                (false, true) => {
                    // The old trees has subdirectories not in this
                    // directory.
                    self.show_delete(dir);
                    self.next_left()?;
                    self.walk_leftdir()?;
                }
                (true, false) => {
                    // The new tree has a newly added directory.
                    self.show_add(dir);
                    self.next_right()?;
                    self.walk_rightdir()?;
                }
                _ if self.left.name() < self.right.name() => {
                    // Old subdirectory.
                    self.show_delete(dir);
                    self.next_left()?;
                    self.walk_leftdir()?;
                }
                _ if self.left.name() > self.right.name() => {
                    // The new tree has a newly added directory.
                    self.show_add(dir);
                    self.next_right()?;
                    self.walk_rightdir()?;
                }
                _ => {
                    // Same named directory.
                    let dirname = dir.join(self.left.name());
                    self.compare_enter(&dirname)?;
                    self.next_left()?;
                    self.next_right()?;
                    self.walk_samedir(&dirname)?;
                }
            }
        }
    }

    /// We are within the files section of the same directory in the two
    /// trees.  Walk through the nodes, reading the Leave node in both, and
    /// returning.
    fn walk_samefiles(&mut self, dir: &Path) -> Result<()> {
        loop {
            match (self.left.is_leave(), self.right.is_leave()) {
                (true, true) => {
                    self.next_left()?;
                    self.next_right()?;
                    return Ok(());
                }
                (false, true) => {
                    self.show_delete(dir);
                    self.next_left()?;
                }
                (true, false) => {
                    self.show_add(dir);
                    self.next_right()?;
                }
                _ if self.left.name() < self.right.name() => {
                    self.show_delete(dir);
                    self.next_left()?;
                }
                _ if self.left.name() > self.right.name() => {
                    self.show_add(dir);
                    self.next_right()?;
                }
                _ => {
                    // Same file.
                    let nodename = dir.join(self.left.name());
                    self.compare_file(&nodename)?;
                    self.next_left()?;
                    self.next_right()?;
                }
            }
        }
    }

    /// Old directory on the left tree.  Walk through nodes recursively to
    /// discard entire tree.
    fn walk_leftdir(&mut self) -> Result<()> {
        loop {
            if self.left.is_enter() {
                self.next_left()?;
                self.walk_leftdir()?;
            } else if self.left.is_leave() {
                self.next_left()?;
                return Ok(());
            } else {
                self.next_left()?;
            }
        }
    }

    /// New directory on the right tree.  Walk through nodes recursively to
    /// discard entire tree.
    fn walk_rightdir(&mut self) -> Result<()> {
        loop {
            if self.right.is_enter() {
                self.next_right()?;
                self.walk_rightdir()?;
            } else if self.right.is_leave() {
                self.next_right()?;
                return Ok(());
            } else {
                self.next_right()?;
            }
        }
    }

    /// Print a message about something added (the name will be the thing
    /// on the right.
    fn show_add(&self, dir: &Path) {
        println!(
            "+ {:22} {:?}",
            self.right.kind(),
            dir.join(self.right.name())
        );
    }

    /// Print a message about something removed (the name will be the thing
    /// on the left.
    fn show_delete(&self, dir: &Path) {
        println!("- {:22} {:?}", self.left.kind(), dir.join(self.left.name()));
    }

    /// Compare the two "Enter" nodes we are visiting.
    fn compare_enter(&mut self, dir: &Path) -> Result<()> {
        self.compare_atts('d', dir)
    }

    /// Compare two file nodes.
    fn compare_file(&mut self, dir: &Path) -> Result<()> {
        self.compare_atts('f', dir)
    }

    /// Attribute comparison.
    fn compare_atts(&mut self, _kind: char, dir: &Path) -> Result<()> {
        let mut old = self.left.atts().unwrap().clone();
        let mut new = self.right.atts().unwrap().clone();
        let mut diffs = vec![];

        for att in self.ignore.iter() {
            old.remove(att);
            new.remove(att);
        }

        for (k, v) in &new {
            match old.get(k) {
                None => {
                    // This attribute is in the new tree, but not the old
                    // one, warn, but only once.
                    if !self.adds.contains(k) {
                        error!("Added attribute: {}", k);
                        self.adds.insert(k.clone());
                    }
                }
                Some(ov) => {
                    if v != ov {
                        diffs.push(k.clone());
                    }
                }
            }
            old.remove(k);
        }

        for k in old.keys() {
            if !self.missings.contains(k) {
                error!("Missing attribute: {}", k);
                self.missings.insert(k.clone());
            }
        }

        if !diffs.is_empty() {
            let mut buf = String::new();
            diffs.sort();
            for d in &diffs {
                if !buf.is_empty() {
                    buf.extend(",".chars());
                }
                buf.extend(d.chars());
            }
            println!("  [{:<20}] {:?}", buf, dir);
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// error-pattern: requires `copy` lang_item

#![feature(lang_items, start, no_std)]
#![no_std]

#[lang = "sized"]
trait Sized { }

#[start]
fn main(_: isize, _: *const *const u8) -> isize {
    0
}
<commit_msg>fix test fallout<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// error-pattern: requires `copy` lang_item

#![feature(lang_items, start, no_std)]
#![no_std]

#[lang = "sized"]
trait Sized { }

struct S;

#[start]
fn main(_: isize, _: *const *const u8) -> isize {
    let _ = S;
    0
}
<|endoftext|>


<commit_before>use super::ArrayLength;
use core::ops::Add;
use typenum::U1;

/// Helper trait for `arr!` macro
pub trait AddLength<T, N: ArrayLength<T>>: ArrayLength<T> {
    type Output: ArrayLength<T>;
}

impl<T, N1, N2> AddLength<T, N2> for N1
    where N1: ArrayLength<T> + Add<N2>,
          N2: ArrayLength<T>,
          <N1 as Add<N2>>::Output: ArrayLength<T>
{
    type Output = <N1 as Add<N2>>::Output;
}

/// Helper type for `arr!` macro
pub type Inc<T, U> = <U as AddLength<T, U1>>::Output;

/// Proper implementation of `arr!`
#[macro_export]
macro_rules! arr_impl {
    ($T:ty; $N:ty, [$($x:expr),*], []) => ({
	unsafe { $crate::transmute::<_, $crate::GenericArray<$T, $N>>([$($x),*]) }
    });
    ($T:ty; $N:ty, [], [$x1:expr]) => (
        arr_impl!($T; $crate::arr::Inc<$T, $N>, [$x1 as $T], [])
    );
    ($T:ty; $N:ty, [], [$x1:expr, $($x:expr),+]) => (
        arr_impl!($T; $crate::arr::Inc<$T, $N>, [$x1 as $T], [$($x),*])
    );
    ($T:ty; $N:ty, [$($y:expr),+], [$x1:expr]) => (
        arr_impl!($T; $crate::arr::Inc<$T, $N>, [$($y),*, $x1 as $T], [])
    );
    ($T:ty; $N:ty, [$($y:expr),+], [$x1:expr, $($x:expr),+]) => (
        arr_impl!($T; $crate::arr::Inc<$T, $N>, [$($y),*, $x1 as $T], [$($x),*])
    );
}

/// Macro allowing for easy generation of Generic Arrays.
/// Example: `let test = arr![u32; 1, 2, 3];`
#[macro_export]
macro_rules! arr {
    ($T:ty;) => ({
        unsafe { $crate::transmute::<[$T; 0], $crate::GenericArray<$T, $crate::typenum::U0>>([]) }
    });
    ($T:ty; $($x:expr),*) => (
        arr_impl!($T; $crate::typenum::U0, [], [$($x),*])
    );
    ($($x:expr,)+) => (arr![$($x),*]);
    () => ("""Macro requires a type, e.g. `let array = arr![u32; 1, 2, 3];`")
}
<commit_msg>Additional docs<commit_after>//! Implementation for `arr!` macro.

use super::ArrayLength;
use core::ops::Add;
use typenum::U1;

/// Helper trait for `arr!` macro
pub trait AddLength<T, N: ArrayLength<T>>: ArrayLength<T> {
    /// Resulting length
    type Output: ArrayLength<T>;
}

impl<T, N1, N2> AddLength<T, N2> for N1
where
    N1: ArrayLength<T> + Add<N2>,
    N2: ArrayLength<T>,
    <N1 as Add<N2>>::Output: ArrayLength<T>,
{
    type Output = <N1 as Add<N2>>::Output;
}

/// Helper type for `arr!` macro
pub type Inc<T, U> = <U as AddLength<T, U1>>::Output;

/// Proper implementation of `arr!`
#[macro_export]
macro_rules! arr_impl {
    ($T:ty; $N:ty, [$($x:expr),*], []) => ({
	unsafe { $crate::transmute::<_, $crate::GenericArray<$T, $N>>([$($x),*]) }
    });
    ($T:ty; $N:ty, [], [$x1:expr]) => (
        arr_impl!($T; $crate::arr::Inc<$T, $N>, [$x1 as $T], [])
    );
    ($T:ty; $N:ty, [], [$x1:expr, $($x:expr),+]) => (
        arr_impl!($T; $crate::arr::Inc<$T, $N>, [$x1 as $T], [$($x),*])
    );
    ($T:ty; $N:ty, [$($y:expr),+], [$x1:expr]) => (
        arr_impl!($T; $crate::arr::Inc<$T, $N>, [$($y),*, $x1 as $T], [])
    );
    ($T:ty; $N:ty, [$($y:expr),+], [$x1:expr, $($x:expr),+]) => (
        arr_impl!($T; $crate::arr::Inc<$T, $N>, [$($y),*, $x1 as $T], [$($x),*])
    );
}

/// Macro allowing for easy generation of Generic Arrays.
/// Example: `let test = arr![u32; 1, 2, 3];`
#[macro_export]
macro_rules! arr {
    ($T:ty;) => ({
        unsafe { $crate::transmute::<[$T; 0], $crate::GenericArray<$T, $crate::typenum::U0>>([]) }
    });
    ($T:ty; $($x:expr),*) => (
        arr_impl!($T; $crate::typenum::U0, [], [$($x),*])
    );
    ($($x:expr,)+) => (arr![$($x),*]);
    () => ("""Macro requires a type, e.g. `let array = arr![u32; 1, 2, 3];`")
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;

use cabot;

use std::collections::HashMap;
use std::iter::FromIterator;
use std::net::{AddrParseError, SocketAddr};
use std::pin::Pin;

use async_std;
use async_std::fs::OpenOptions;
use async_std::io::{self, Write};
use async_std::prelude::*;
use async_std::task::Context;
use async_std::task::Poll;
use clap::{App, Arg};
use log::Level::Info;

use cabot::constants;
use cabot::http;
use cabot::request::RequestBuilder;
use cabot::results::CabotResult;

pub async fn run() -> CabotResult<()> {
    let user_agent: String = constants::user_agent();
    let matches = App::new(env!("CARGO_PKG_NAME"))
        .version(constants::VERSION)
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .arg(
            Arg::with_name("URL")
                .index(1)
                .required(true)
                .help("URL to request"),
        )
        .arg(
            Arg::with_name("REQUEST")
                .short("X")
                .long("request")
                .default_value("GET")
                .help("Specify request command to use"),
        )
        .arg(
            Arg::with_name("HEADER")
                .short("H")
                .long("header")
                .takes_value(true)
                .multiple(true)
                .help("Pass custom header to server"),
        )
        .arg(
            Arg::with_name("FILE")
                .short("o")
                .long("output")
                .takes_value(true)
                .help("Write to FILE instead of stdout"),
        )
        .arg(
            Arg::with_name("VERBOSE")
                .short("v")
                .long("verbose")
                .help("Make the operation more talkative"),
        )
        .arg(
            Arg::with_name("BODY")
                .short("d")
                .long("data")
                .takes_value(true)
                .help("Post Data (Using utf-8 encoding)"),
        )
        .arg(
            Arg::with_name("IPv4")
                .short("4")
                .long("ipv4")
                .help("Resolve host names to IPv4 addresses"),
        )
        .arg(
            Arg::with_name("IPv6")
                .short("6")
                .long("ipv6")
                .help("Resolve host names to IPv6 addresses"),
        )
        .arg(
            Arg::with_name("UA")
                .short("A")
                .long("user-agent")
                .default_value(user_agent.as_str())
                .help("The user-agent HTTP header to use"),
        )
        .arg(
            Arg::with_name("DNS_LOOKUP_TIMEOUT")
                .long("dns-timeout")
                .takes_value(true)
                .default_value("5")
                .help("timeout for the dns lookup resolution in seconds"),
        )
        .arg(
            Arg::with_name("CONNECT_TIMEOUT")
                .long("connect-timeout")
                .takes_value(true)
                .default_value("15")
                .help("timeout for the tcp connection"),
        )
        .arg(
            Arg::with_name("READ_TIMEOUT")
                .long("read-timeout")
                .takes_value(true)
                .default_value("10")
                .help("timeout for the tcp read in seconds"),
        )
        .arg(
            Arg::with_name("REQUEST_TIMEOUT")
                .long("timeout")
                .takes_value(true)
                .default_value("30")
                .help("timeout for the whole http request in seconds (0 means no timeout)"),
        )
        .arg(
            Arg::with_name("RESOLVE")
                .long("resolve")
                .takes_value(true)
                .multiple(true)
                .help("<host:port:address> Resolve the host+port to this address"),
        )
        .get_matches();

    let url = matches.value_of("URL").unwrap();
    let http_method = matches.value_of("REQUEST").unwrap();
    let verbose = matches.is_present("VERBOSE");
    let body = matches.value_of("BODY");
    let ua = matches.value_of("UA").unwrap();

    let mut ipv4 = matches.is_present("IPv4");
    let mut ipv6 = matches.is_present("IPv6");

    if !ipv4 && !ipv6 {
        ipv4 = true;
        ipv6 = true;
    }

    let headers: Vec<&str> = match matches.values_of("HEADER") {
        Some(headers) => headers.collect(),
        None => Vec::new(),
    };

    let resolved: HashMap<String, SocketAddr> = match matches.values_of("RESOLVE") {
        Some(headers) => HashMap::from_iter(
            headers
                .map(|resolv| resolv.splitn(3, ':'))
                .map(|resolv| {
                    let count = resolv.clone().count();
                    if count != 3 {
                        let resolv: Vec<&str> = resolv.collect();
                        let _ =
                            eprintln!("Invalid format in resolve argument: {}", resolv.join(":"));
                        std::process::exit(1);
                    }
                    resolv
                })
                .map(|mut resolv| {
                    (
                        resolv.next().unwrap(),
                        resolv.next().unwrap(),
                        resolv.next().unwrap(),
                    )
                })
                .map(|(host, port, addr)| {
                    let parsed_port = port.parse::<usize>();
                    if parsed_port.is_err() {
                        let _ = eprintln!(
                            "Invalid port in resolve argument: {}:{}:{}",
                            host, port, addr
                        );
                        std::process::exit(1);
                    }
                    let sockaddr: Result<SocketAddr, AddrParseError> =
                        format!("{}:{}", addr, port).parse();
                    if sockaddr.is_err() {
                        let _ = eprintln!(
                            "Invalid address in resolve argument: {}:{}:{}",
                            host, port, addr
                        );
                        std::process::exit(1);
                    }
                    let sockaddr = sockaddr.unwrap();
                    (format!("{}:{}", host, port), sockaddr)
                }),
        ),
        None => HashMap::new(),
    };

    let mut builder = RequestBuilder::new(url)
        .set_http_method(http_method)
        .set_user_agent(ua)
        .add_headers(&headers.as_slice());

    if body.is_some() {
        builder = builder.set_body_as_str(body.unwrap());
    }

    let request = builder.build()?;

    let dns_timeout = u64::from_str_radix(matches.value_of("DNS_LOOKUP_TIMEOUT").unwrap(), 10)
        .expect("DNS_LOOKUP_TIMEOUT must be an integer")
        * 1_000;
    let connect_timeout = u64::from_str_radix(matches.value_of("CONNECT_TIMEOUT").unwrap(), 10)
        .expect("CONNECT_TIMEOUT must be an integer")
        * 1_000;
    let read_timeout = u64::from_str_radix(matches.value_of("READ_TIMEOUT").unwrap(), 10)
        .expect("READ_TIMEOUT must be an integer")
        * 1_000;
    let request_timeout = u64::from_str_radix(matches.value_of("REQUEST_TIMEOUT").unwrap(), 10)
        .expect("REQUEST_TIMEOUT must be an integer")
        * 1_000;

    if let Some(path) = matches.value_of("FILE") {
        let mut f = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(path)
            .await
            .unwrap();
        http::http_query(
            &request,
            &mut CabotBinWrite::new(&mut f, verbose),
            &resolved,
            verbose,
            ipv4,
            ipv6,
            dns_timeout,
            connect_timeout,
            read_timeout,
            request_timeout,
        )
        .await?
    } else {
        http::http_query(
            &request,
            &mut CabotBinWrite::new(&mut io::stdout(), verbose),
            &resolved,
            verbose,
            ipv4,
            ipv6,
            dns_timeout,
            connect_timeout,
            read_timeout,
            request_timeout,
        )
        .await?
    };
    Ok(())
}

#[async_std::main]
async fn main() {
    #[cfg(feature = "pretty_log")]
    pretty_env_logger::init();
    debug!("Starting cabot");

    run()
        .await
        .map(|ok| {
            debug!("Command cabot ended succesfully");
            ok
        })
        .map_err(|err| {
            eprintln!("{}", err);
            std::process::exit(1);
        })
        .unwrap();
}

// Internal Of the Binary

struct CabotBinWrite<'a> {
    out: &'a mut (dyn Write + Unpin),
    header_read: bool,
    verbose: bool,
}

impl<'a> CabotBinWrite<'a> {
    pub fn new(out: &'a mut (dyn Write + Unpin), verbose: bool) -> Self {
        CabotBinWrite {
            out,
            verbose,
            header_read: false,
        }
    }
    fn display_headers(&self, buf: &[u8]) {
        let headers = String::from_utf8_lossy(buf);
        let split: Vec<&str> = constants::SPLIT_HEADER_RE.split(&headers).collect();
        if log_enabled!(Info) {
            for part in split {
                info!("< {}", part);
            }
        } else if self.verbose {
            for part in split {
                eprintln!("< {}", part);
            }
        }
    }
}

impl<'a> Write for CabotBinWrite<'a> {
    // may receive headers
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
        let self_ = Pin::get_mut(self);
        if !self_.header_read {
            // the first time write is called, all headers are sent
            // in the buffer. there is no need to parse it again.
            if log_enabled!(Info) || self_.verbose {
                self_.display_headers(&buf);
            }
            self_.header_read = true;
            Poll::Ready(Ok(0))
        } else {
            Pin::new(&mut self_.out).poll_write(cx, buf)
        }
    }

    /// this function is called when the request is done
    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<io::Result<()>> {
        let self_ = Pin::get_mut(self);
        self_.out.flush();
        Poll::Ready(Ok(()))
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
        let self_ = Pin::get_mut(self);
        Pin::new(&mut self_.out).poll_close(cx)
    }
}
<commit_msg>Fix the -o parameter that does not write all bytes<commit_after>#[macro_use]
extern crate log;

use cabot;

use std::collections::HashMap;
use std::iter::FromIterator;
use std::net::{AddrParseError, SocketAddr};
use std::pin::Pin;

use async_std;
use async_std::fs::OpenOptions;
use async_std::io::{self, Write};
use async_std::prelude::*;
use async_std::task::Context;
use async_std::task::Poll;
use clap::{App, Arg};
use log::Level::Info;

use cabot::constants;
use cabot::http;
use cabot::request::RequestBuilder;
use cabot::results::CabotResult;

pub async fn run() -> CabotResult<()> {
    let user_agent: String = constants::user_agent();
    let matches = App::new(env!("CARGO_PKG_NAME"))
        .version(constants::VERSION)
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .arg(
            Arg::with_name("URL")
                .index(1)
                .required(true)
                .help("URL to request"),
        )
        .arg(
            Arg::with_name("REQUEST")
                .short("X")
                .long("request")
                .default_value("GET")
                .help("Specify request command to use"),
        )
        .arg(
            Arg::with_name("HEADER")
                .short("H")
                .long("header")
                .takes_value(true)
                .multiple(true)
                .help("Pass custom header to server"),
        )
        .arg(
            Arg::with_name("FILE")
                .short("o")
                .long("output")
                .takes_value(true)
                .help("Write to FILE instead of stdout"),
        )
        .arg(
            Arg::with_name("VERBOSE")
                .short("v")
                .long("verbose")
                .help("Make the operation more talkative"),
        )
        .arg(
            Arg::with_name("BODY")
                .short("d")
                .long("data")
                .takes_value(true)
                .help("Post Data (Using utf-8 encoding)"),
        )
        .arg(
            Arg::with_name("IPv4")
                .short("4")
                .long("ipv4")
                .help("Resolve host names to IPv4 addresses"),
        )
        .arg(
            Arg::with_name("IPv6")
                .short("6")
                .long("ipv6")
                .help("Resolve host names to IPv6 addresses"),
        )
        .arg(
            Arg::with_name("UA")
                .short("A")
                .long("user-agent")
                .default_value(user_agent.as_str())
                .help("The user-agent HTTP header to use"),
        )
        .arg(
            Arg::with_name("DNS_LOOKUP_TIMEOUT")
                .long("dns-timeout")
                .takes_value(true)
                .default_value("5")
                .help("timeout for the dns lookup resolution in seconds"),
        )
        .arg(
            Arg::with_name("CONNECT_TIMEOUT")
                .long("connect-timeout")
                .takes_value(true)
                .default_value("15")
                .help("timeout for the tcp connection"),
        )
        .arg(
            Arg::with_name("READ_TIMEOUT")
                .long("read-timeout")
                .takes_value(true)
                .default_value("10")
                .help("timeout for the tcp read in seconds"),
        )
        .arg(
            Arg::with_name("REQUEST_TIMEOUT")
                .long("timeout")
                .takes_value(true)
                .default_value("30")
                .help("timeout for the whole http request in seconds (0 means no timeout)"),
        )
        .arg(
            Arg::with_name("RESOLVE")
                .long("resolve")
                .takes_value(true)
                .multiple(true)
                .help("<host:port:address> Resolve the host+port to this address"),
        )
        .get_matches();

    let url = matches.value_of("URL").unwrap();
    let http_method = matches.value_of("REQUEST").unwrap();
    let verbose = matches.is_present("VERBOSE");
    let body = matches.value_of("BODY");
    let ua = matches.value_of("UA").unwrap();

    let mut ipv4 = matches.is_present("IPv4");
    let mut ipv6 = matches.is_present("IPv6");

    if !ipv4 && !ipv6 {
        ipv4 = true;
        ipv6 = true;
    }

    let headers: Vec<&str> = match matches.values_of("HEADER") {
        Some(headers) => headers.collect(),
        None => Vec::new(),
    };

    let resolved: HashMap<String, SocketAddr> = match matches.values_of("RESOLVE") {
        Some(headers) => HashMap::from_iter(
            headers
                .map(|resolv| resolv.splitn(3, ':'))
                .map(|resolv| {
                    let count = resolv.clone().count();
                    if count != 3 {
                        let resolv: Vec<&str> = resolv.collect();
                        let _ =
                            eprintln!("Invalid format in resolve argument: {}", resolv.join(":"));
                        std::process::exit(1);
                    }
                    resolv
                })
                .map(|mut resolv| {
                    (
                        resolv.next().unwrap(),
                        resolv.next().unwrap(),
                        resolv.next().unwrap(),
                    )
                })
                .map(|(host, port, addr)| {
                    let parsed_port = port.parse::<usize>();
                    if parsed_port.is_err() {
                        let _ = eprintln!(
                            "Invalid port in resolve argument: {}:{}:{}",
                            host, port, addr
                        );
                        std::process::exit(1);
                    }
                    let sockaddr: Result<SocketAddr, AddrParseError> =
                        format!("{}:{}", addr, port).parse();
                    if sockaddr.is_err() {
                        let _ = eprintln!(
                            "Invalid address in resolve argument: {}:{}:{}",
                            host, port, addr
                        );
                        std::process::exit(1);
                    }
                    let sockaddr = sockaddr.unwrap();
                    (format!("{}:{}", host, port), sockaddr)
                }),
        ),
        None => HashMap::new(),
    };

    let mut builder = RequestBuilder::new(url)
        .set_http_method(http_method)
        .set_user_agent(ua)
        .add_headers(&headers.as_slice());

    if body.is_some() {
        builder = builder.set_body_as_str(body.unwrap());
    }

    let request = builder.build()?;

    let dns_timeout = u64::from_str_radix(matches.value_of("DNS_LOOKUP_TIMEOUT").unwrap(), 10)
        .expect("DNS_LOOKUP_TIMEOUT must be an integer")
        * 1_000;
    let connect_timeout = u64::from_str_radix(matches.value_of("CONNECT_TIMEOUT").unwrap(), 10)
        .expect("CONNECT_TIMEOUT must be an integer")
        * 1_000;
    let read_timeout = u64::from_str_radix(matches.value_of("READ_TIMEOUT").unwrap(), 10)
        .expect("READ_TIMEOUT must be an integer")
        * 1_000;
    let request_timeout = u64::from_str_radix(matches.value_of("REQUEST_TIMEOUT").unwrap(), 10)
        .expect("REQUEST_TIMEOUT must be an integer")
        * 1_000;

    if let Some(path) = matches.value_of("FILE") {
        let mut f = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(path)
            .await
            .unwrap();
        http::http_query(
            &request,
            &mut CabotBinWrite::new(&mut f, verbose),
            &resolved,
            verbose,
            ipv4,
            ipv6,
            dns_timeout,
            connect_timeout,
            read_timeout,
            request_timeout,
        )
        .await?
    } else {
        http::http_query(
            &request,
            &mut CabotBinWrite::new(&mut io::stdout(), verbose),
            &resolved,
            verbose,
            ipv4,
            ipv6,
            dns_timeout,
            connect_timeout,
            read_timeout,
            request_timeout,
        )
        .await?
    };
    Ok(())
}

#[async_std::main]
async fn main() {
    #[cfg(feature = "pretty_log")]
    pretty_env_logger::init();
    debug!("Starting cabot");

    run()
        .await
        .map(|ok| {
            debug!("Command cabot ended succesfully");
            ok
        })
        .map_err(|err| {
            eprintln!("{}", err);
            std::process::exit(1);
        })
        .unwrap();
}

// Internal Of the Binary

struct CabotBinWrite<'a> {
    out: &'a mut (dyn Write + Unpin),
    header_read: bool,
    verbose: bool,
}

impl<'a> CabotBinWrite<'a> {
    pub fn new(out: &'a mut (dyn Write + Unpin), verbose: bool) -> Self {
        CabotBinWrite {
            out,
            verbose,
            header_read: false,
        }
    }
    fn display_headers(&self, buf: &[u8]) {
        let headers = String::from_utf8_lossy(buf);
        let split: Vec<&str> = constants::SPLIT_HEADER_RE.split(&headers).collect();
        if log_enabled!(Info) {
            for part in split {
                info!("< {}", part);
            }
        } else if self.verbose {
            for part in split {
                eprintln!("< {}", part);
            }
        }
    }
}

impl<'a> Write for CabotBinWrite<'a> {
    // may receive headers
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
        let self_ = Pin::get_mut(self);
        if !self_.header_read {
            // the first time write is called, all headers are sent
            // in the buffer. there is no need to parse it again.
            if log_enabled!(Info) || self_.verbose {
                self_.display_headers(&buf);
            }
            self_.header_read = true;
            Poll::Ready(Ok(0))
        } else {
            let towrite = buf.len();
            let mut written = 0;
            loop {
                let res = Pin::new(&mut self_.out).poll_write(cx, &buf[written..towrite]);
                match res {
                    Poll::Ready(Ok(l)) => written += l,
                    _ => {}
                }
                if written >= towrite {
                    break;
                }
            }
            Poll::Ready(Ok(written))
        }
    }

    /// this function is called when the request is done
    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<io::Result<()>> {
        let self_ = Pin::get_mut(self);
        self_.out.flush();
        Poll::Ready(Ok(()))
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
        let self_ = Pin::get_mut(self);
        Pin::new(&mut self_.out).poll_close(cx)
    }
}
<|endoftext|>


<commit_before>use core::clone::Clone;
use core::ops::Deref;
use diesel;
use diesel::expression::sql;
use diesel::prelude::*;
use diesel::BelongingToDsl;
use diesel::types::*;
use std::path::Path;
#[cfg(test)]
use db;
use db::ConnectionSource;
use db::{playlists, playlist_songs, songs, users};
use index::{self, Song};
use vfs::VFSSource;
use errors::*;

#[derive(Insertable)]
#[table_name="playlists"]
struct NewPlaylist {
	name: String,
	owner: i32,
}

#[derive(Identifiable, Queryable)]
pub struct User {
	id: i32,
}

#[derive(Identifiable, Queryable, Associations)]
#[belongs_to(User, foreign_key="owner")]
pub struct Playlist {
	id: i32,
	owner: i32,
	name: String,
}

#[derive(Identifiable, Queryable, Associations)]
#[belongs_to(Playlist, foreign_key="playlist")]
pub struct PlaylistSong {
	id: i32,
	playlist: i32,
	path: String,
	ordering: i32,
}

#[derive(Insertable)]
#[table_name="playlist_songs"]
pub struct NewPlaylistSong {
	playlist: i32,
	path: String,
	ordering: i32,
}

pub fn list_playlists<T>(owner: &str, db: &T) -> Result<Vec<String>>
	where T: ConnectionSource + VFSSource
{
	let connection = db.get_connection();

	let user: User;
	{
		use self::users::dsl::*;
		user = users
			.filter(name.eq(owner))
			.select((id,))
			.first(connection.deref())?;
	}

	{
		use self::playlists::dsl::*;
		let found_playlists: Vec<String> = Playlist::belonging_to(&user)
			.select(name)
			.load(connection.deref())?;
		Ok(found_playlists)
	}
}

pub fn save_playlist<T>(name: &str, owner: &str, content: &Vec<String>, db: &T) -> Result<()>
	where T: ConnectionSource + VFSSource
{
	let user: User;
	let new_playlist: NewPlaylist;
	let playlist: Playlist;
	let vfs = db.get_vfs()?;

	{
		let connection = db.get_connection();

		// Find owner
		{
			use self::users::dsl::*;
			user = users
				.filter(name.eq(owner))
				.select((id,))
				.get_result(connection.deref())?;
		}

		// Create playlist
		new_playlist = NewPlaylist {
			name: name.into(),
			owner: user.id,
		};

		diesel::insert(&new_playlist)
			.into(playlists::table)
			.execute(connection.deref())?;

		{
			use self::playlists::dsl::*;
			playlist = playlists
				.filter(name.eq(name).and(owner.eq(user.id)))
				.get_result(connection.deref())?;
		}
	}

	let mut new_songs: Vec<NewPlaylistSong> = Vec::new();
	new_songs.reserve(content.len());

	for (i, path) in content.iter().enumerate() {
		let virtual_path = Path::new(&path);
		if let Some(real_path) = vfs.virtual_to_real(virtual_path)
		       .ok()
		       .and_then(|p| p.to_str().map(|s| s.to_owned())) {
			new_songs.push(NewPlaylistSong {
			                   playlist: playlist.id,
			                   path: real_path.into(),
			                   ordering: i as i32,
			               });
		}
	}

	{
		let connection = db.get_connection();
		connection
			.deref()
			.transaction::<_, diesel::result::Error, _>(|| {
				// Delete old content (if any)
				let old_songs = PlaylistSong::belonging_to(&playlist);
				diesel::delete(old_songs).execute(connection.deref())?;

				// Insert content
				diesel::insert(&new_songs)
					.into(playlist_songs::table)
					.execute(connection.deref())?;
				Ok(())
			})?;
	}

	Ok(())
}

pub fn read_playlist<T>(playlist_name: &str, owner: &str, db: &T) -> Result<Vec<Song>>
	where T: ConnectionSource + VFSSource
{
	let vfs = db.get_vfs()?;
	let songs: Vec<Song>;

	{
		let connection = db.get_connection();
		let user: User;
		let playlist: Playlist;

		// Find owner
		{
			use self::users::dsl::*;
			user = users
				.filter(name.eq(owner))
				.select((id,))
				.get_result(connection.deref())?;
		}

		// Find playlist
		{
			use self::playlists::dsl::*;
			playlist = playlists
				.filter(name.eq(playlist_name).and(owner.eq(user.id)))
				.get_result(connection.deref())?;
		}

		// Select songs. Not using Diesel because we need to LEFT JOIN using a custom column
		let query = sql::<songs::SqlType>(r#"
			SELECT s.id, s.path, s.parent, s.track_number, s.disc_number, s.title, s.artist, s.album_artist, s.year, s.album, s.artwork
			FROM playlist_songs ps
			LEFT JOIN songs s ON ps.path = s.path
			WHERE ps.playlist = ?
			ORDER BY ps.ordering
		"#);
		let query = query.clone().bind::<Integer, _>(playlist.id);
		songs = query.get_results(connection.deref())?;
	}

	// Map real path to virtual paths
	let virtual_songs = songs
		.into_iter()
		.filter_map(|s| index::virtualize_song(&vfs, s))
		.collect();

	Ok(virtual_songs)
}

pub fn delete_playlist<T>(playlist_name: &str, owner: &str, db: &T) -> Result<()>
	where T: ConnectionSource + VFSSource
{
	let connection = db.get_connection();

	let user: User;
	{
		use self::users::dsl::*;
		user = users
			.filter(name.eq(owner))
			.select((id,))
			.first(connection.deref())?;
	}

	{
		use self::playlists::dsl::*;
		let q = Playlist::belonging_to(&user).filter(name.eq(playlist_name));
		diesel::delete(q).execute(connection.deref())?;
	}

	Ok(())
}

#[test]
fn test_create_playlist() {
	let db = db::_get_test_db("create_playlist.sqlite");

	let found_playlists = list_playlists("test_user", &db).unwrap();
	assert!(found_playlists.is_empty());

	save_playlist("chill_and_grill", "test_user", &Vec::new(), &db).unwrap();
	let found_playlists = list_playlists("test_user", &db).unwrap();
	assert_eq!(found_playlists.len(), 1);
	assert_eq!(found_playlists[0], "chill_and_grill");

	let found_playlists = list_playlists("someone_else", &db);
	assert!(found_playlists.is_err());
}

#[test]
fn test_delete_playlist() {
	let db = db::_get_test_db("delete_playlist.sqlite");
	let playlist_content = Vec::new();

	save_playlist("chill_and_grill", "test_user", &playlist_content, &db).unwrap();
	save_playlist("mellow_bungalow", "test_user", &playlist_content, &db).unwrap();
	let found_playlists = list_playlists("test_user", &db).unwrap();
	assert_eq!(found_playlists.len(), 2);

	delete_playlist("chill_and_grill", "test_user", &db).unwrap();
	let found_playlists = list_playlists("test_user", &db).unwrap();
	assert_eq!(found_playlists.len(), 1);
	assert_eq!(found_playlists[0], "mellow_bungalow");

	let delete_result = delete_playlist("mellow_bungalow", "someone_else", &db);
	assert!(delete_result.is_err());
}

#[test]
fn test_fill_playlist() {
	use index;

	let db = db::_get_test_db("fill_playlist.sqlite");
	index::update(&db).unwrap();

	let mut playlist_content: Vec<String> = index::flatten(&db, Path::new("root"))
		.unwrap()
		.into_iter()
		.map(|s| s.path)
		.collect();
	assert_eq!(playlist_content.len(), 12);

	let first_song = playlist_content[0].clone();
	playlist_content.push(first_song);
	assert_eq!(playlist_content.len(), 13);

	save_playlist("all_the_music", "test_user", &playlist_content, &db).unwrap();

	let songs = read_playlist("all_the_music", "test_user", &db).unwrap();
	assert_eq!(songs.len(), 13);
	assert_eq!(songs[0].title, Some("Above The Water".to_owned()));
	assert_eq!(songs[12].title, Some("Above The Water".to_owned()));

	use std::path::PathBuf;
	let mut first_song_path = PathBuf::new();
	first_song_path.push("root");
	first_song_path.push("Khemmis");
	first_song_path.push("Hunted");
	first_song_path.push("01 - Above The Water.mp3");
	assert_eq!(songs[0].path, first_song_path.to_str().unwrap());

	// Save again to verify that we don't dupe the content
	save_playlist("all_the_music", "test_user", &playlist_content, &db).unwrap();
	let songs = read_playlist("all_the_music", "test_user", &db).unwrap();
	assert_eq!(songs.len(), 13);
}
<commit_msg>Fixed a bug where playlists were overwriting each other<commit_after>use core::clone::Clone;
use core::ops::Deref;
use diesel;
use diesel::expression::sql;
use diesel::prelude::*;
use diesel::BelongingToDsl;
use diesel::types::*;
use std::path::Path;
#[cfg(test)]
use db;
use db::ConnectionSource;
use db::{playlists, playlist_songs, songs, users};
use index::{self, Song};
use vfs::VFSSource;
use errors::*;

#[derive(Insertable)]
#[table_name="playlists"]
struct NewPlaylist {
	name: String,
	owner: i32,
}

#[derive(Identifiable, Queryable)]
pub struct User {
	id: i32,
}

#[derive(Identifiable, Queryable, Associations)]
#[belongs_to(User, foreign_key="owner")]
pub struct Playlist {
	id: i32,
	owner: i32,
	name: String,
}

#[derive(Identifiable, Queryable, Associations)]
#[belongs_to(Playlist, foreign_key="playlist")]
pub struct PlaylistSong {
	id: i32,
	playlist: i32,
	path: String,
	ordering: i32,
}

#[derive(Insertable)]
#[table_name="playlist_songs"]
pub struct NewPlaylistSong {
	playlist: i32,
	path: String,
	ordering: i32,
}

pub fn list_playlists<T>(owner: &str, db: &T) -> Result<Vec<String>>
	where T: ConnectionSource + VFSSource
{
	let connection = db.get_connection();

	let user: User;
	{
		use self::users::dsl::*;
		user = users
			.filter(name.eq(owner))
			.select((id,))
			.first(connection.deref())?;
	}

	{
		use self::playlists::dsl::*;
		let found_playlists: Vec<String> = Playlist::belonging_to(&user)
			.select(name)
			.load(connection.deref())?;
		Ok(found_playlists)
	}
}

pub fn save_playlist<T>(playlist_name: &str, owner: &str, content: &Vec<String>, db: &T) -> Result<()>
	where T: ConnectionSource + VFSSource
{
	let user: User;
	let new_playlist: NewPlaylist;
	let playlist: Playlist;
	let vfs = db.get_vfs()?;

	{
		let connection = db.get_connection();

		// Find owner
		{
			use self::users::dsl::*;
			user = users
				.filter(name.eq(owner))
				.select((id,))
				.get_result(connection.deref())?;
		}

		// Create playlist
		new_playlist = NewPlaylist {
			name: playlist_name.into(),
			owner: user.id,
		};

		diesel::insert(&new_playlist)
			.into(playlists::table)
			.execute(connection.deref())?;

		{
			use self::playlists::dsl::*;
			playlist = playlists
				.filter(name.eq(playlist_name).and(owner.eq(user.id)))
				.get_result(connection.deref())?;
		}
	}

	let mut new_songs: Vec<NewPlaylistSong> = Vec::new();
	new_songs.reserve(content.len());

	for (i, path) in content.iter().enumerate() {
		let virtual_path = Path::new(&path);
		if let Some(real_path) = vfs.virtual_to_real(virtual_path)
		       .ok()
		       .and_then(|p| p.to_str().map(|s| s.to_owned())) {
			new_songs.push(NewPlaylistSong {
			                   playlist: playlist.id,
			                   path: real_path.into(),
			                   ordering: i as i32,
			               });
		}
	}

	{
		let connection = db.get_connection();
		connection
			.deref()
			.transaction::<_, diesel::result::Error, _>(|| {
				// Delete old content (if any)
				let old_songs = PlaylistSong::belonging_to(&playlist);
				diesel::delete(old_songs).execute(connection.deref())?;

				// Insert content
				diesel::insert(&new_songs)
					.into(playlist_songs::table)
					.execute(connection.deref())?;
				Ok(())
			})?;
	}

	Ok(())
}

pub fn read_playlist<T>(playlist_name: &str, owner: &str, db: &T) -> Result<Vec<Song>>
	where T: ConnectionSource + VFSSource
{
	let vfs = db.get_vfs()?;
	let songs: Vec<Song>;

	{
		let connection = db.get_connection();
		let user: User;
		let playlist: Playlist;

		// Find owner
		{
			use self::users::dsl::*;
			user = users
				.filter(name.eq(owner))
				.select((id,))
				.get_result(connection.deref())?;
		}

		// Find playlist
		{
			use self::playlists::dsl::*;
			playlist = playlists
				.filter(name.eq(playlist_name).and(owner.eq(user.id)))
				.get_result(connection.deref())?;
		}

		// Select songs. Not using Diesel because we need to LEFT JOIN using a custom column
		let query = sql::<songs::SqlType>(r#"
			SELECT s.id, s.path, s.parent, s.track_number, s.disc_number, s.title, s.artist, s.album_artist, s.year, s.album, s.artwork
			FROM playlist_songs ps
			LEFT JOIN songs s ON ps.path = s.path
			WHERE ps.playlist = ?
			ORDER BY ps.ordering
		"#);
		let query = query.clone().bind::<Integer, _>(playlist.id);
		songs = query.get_results(connection.deref())?;
	}

	// Map real path to virtual paths
	let virtual_songs = songs
		.into_iter()
		.filter_map(|s| index::virtualize_song(&vfs, s))
		.collect();

	Ok(virtual_songs)
}

pub fn delete_playlist<T>(playlist_name: &str, owner: &str, db: &T) -> Result<()>
	where T: ConnectionSource + VFSSource
{
	let connection = db.get_connection();

	let user: User;
	{
		use self::users::dsl::*;
		user = users
			.filter(name.eq(owner))
			.select((id,))
			.first(connection.deref())?;
	}

	{
		use self::playlists::dsl::*;
		let q = Playlist::belonging_to(&user).filter(name.eq(playlist_name));
		diesel::delete(q).execute(connection.deref())?;
	}

	Ok(())
}

#[test]
fn test_create_playlist() {
	let db = db::_get_test_db("create_playlist.sqlite");

	let found_playlists = list_playlists("test_user", &db).unwrap();
	assert!(found_playlists.is_empty());

	save_playlist("chill_and_grill", "test_user", &Vec::new(), &db).unwrap();
	let found_playlists = list_playlists("test_user", &db).unwrap();
	assert_eq!(found_playlists.len(), 1);
	assert_eq!(found_playlists[0], "chill_and_grill");

	let found_playlists = list_playlists("someone_else", &db);
	assert!(found_playlists.is_err());
}

#[test]
fn test_delete_playlist() {
	let db = db::_get_test_db("delete_playlist.sqlite");
	let playlist_content = Vec::new();

	save_playlist("chill_and_grill", "test_user", &playlist_content, &db).unwrap();
	save_playlist("mellow_bungalow", "test_user", &playlist_content, &db).unwrap();
	let found_playlists = list_playlists("test_user", &db).unwrap();
	assert_eq!(found_playlists.len(), 2);

	delete_playlist("chill_and_grill", "test_user", &db).unwrap();
	let found_playlists = list_playlists("test_user", &db).unwrap();
	assert_eq!(found_playlists.len(), 1);
	assert_eq!(found_playlists[0], "mellow_bungalow");

	let delete_result = delete_playlist("mellow_bungalow", "someone_else", &db);
	assert!(delete_result.is_err());
}

#[test]
fn test_fill_playlist() {
	use index;

	let db = db::_get_test_db("fill_playlist.sqlite");
	index::update(&db).unwrap();

	let mut playlist_content: Vec<String> = index::flatten(&db, Path::new("root"))
		.unwrap()
		.into_iter()
		.map(|s| s.path)
		.collect();
	assert_eq!(playlist_content.len(), 12);

	let first_song = playlist_content[0].clone();
	playlist_content.push(first_song);
	assert_eq!(playlist_content.len(), 13);

	save_playlist("all_the_music", "test_user", &playlist_content, &db).unwrap();

	let songs = read_playlist("all_the_music", "test_user", &db).unwrap();
	assert_eq!(songs.len(), 13);
	assert_eq!(songs[0].title, Some("Above The Water".to_owned()));
	assert_eq!(songs[12].title, Some("Above The Water".to_owned()));

	use std::path::PathBuf;
	let mut first_song_path = PathBuf::new();
	first_song_path.push("root");
	first_song_path.push("Khemmis");
	first_song_path.push("Hunted");
	first_song_path.push("01 - Above The Water.mp3");
	assert_eq!(songs[0].path, first_song_path.to_str().unwrap());

	// Save again to verify that we don't dupe the content
	save_playlist("all_the_music", "test_user", &playlist_content, &db).unwrap();
	let songs = read_playlist("all_the_music", "test_user", &db).unwrap();
	assert_eq!(songs.len(), 13);
}
<|endoftext|>