<commit_before>extern crate random_choice;

use std::collections::{HashSet, HashMap};
use std::iter::FromIterator;
use std::cmp::max;
use self::random_choice::random_choice;

pub struct GSDMM {
    alpha: f64,
    beta: f64,
    K:usize,
    V:f64,
    D:usize,
    maxit:isize,
    clusters: Vec<usize>,
    pub docs: Vec<Vec<String>>,
    pub labels: Vec<usize>,
    pub cluster_counts: Vec<u32>,
    pub cluster_word_counts:Vec<u32>,
    pub cluster_word_distributions: Vec<HashMap<String,u32>>
}

impl GSDMM {
    pub fn new(alpha:f64, beta:f64, K: usize, maxit:isize, vocab:HashSet<String>, docs:Vec<Vec<String>>) -> GSDMM {
        let D = docs.len();

        // compute utilized vocabulary size.
        let mut utilized_vocab = HashSet::with_capacity(vocab.len());
        for doc in &docs {
            for word in doc {
                utilized_vocab.insert(word.clone());
            }
        }
        let V = utilized_vocab.len() as f64;
        println!("Fitting with alpha={}, beta={}, K={}, maxit={}, vocab size={}", alpha, beta, K, maxit, V as u32);

        let clusters = (0_usize..K).collect::<Vec<usize>>();
        let mut d_z: Vec<usize> = (0_usize..D).map(|_| 0_usize).collect::<Vec<usize>>(); // doc labels
        let mut m_z: Vec<u32> = GSDMM::zero_vector(K);  // cluster sizes
        let mut n_z: Vec<u32> = GSDMM::zero_vector(K);  // cluster word counts
        let mut n_z_w = Vec::<HashMap<String, u32>>::with_capacity(K);  // container for cluster word distributions
        for _ in 0_usize..K {
            let mut m = HashMap::<String, u32>::with_capacity(max(vocab.len() / 10, 100));
            &n_z_w.push(m);
        }

        // randomly initialize cluster assignment
        let p = (0..K).map(|_| 1_f64 / (K as f64)).collect::<Vec<f64>>();

        let choices = random_choice().random_choice_f64(&clusters, &p, D) ;
        for i in 0..D {
            let z = choices[i].clone();
            let ref doc = docs[i];
            d_z[i] = z;
            m_z[z] += 1;
            n_z[z] += doc.len() as u32;
            let ref mut clust_words: HashMap<String, u32> = n_z_w[z];
            for word in doc {
                if !clust_words.contains_key(word) {
                    clust_words.insert(word.clone(), 0_u32);
                }
                    * clust_words.get_mut(word).unwrap() += 1_u32;
            }
        }

        GSDMM {
            alpha: alpha,
            beta: beta,
            K: K,
            V: V,
            D: D,
            maxit:maxit,
            docs:docs,
            clusters: clusters.clone(),     // Don't totally get why we need the clone here!
            labels: d_z,
            cluster_counts: m_z,
            cluster_word_counts: n_z,
            cluster_word_distributions: n_z_w
        }
    }

    pub fn fit(&mut self) {
        let mut number_clusters = self.K;
        for it in 0..self.maxit {
            let mut total_transfers = 0;
            for i in 0..self.D {
                let ref doc = self.docs[i];
                let doc_size = doc.len() as u32;

                // remove the doc from its current cluster
                let z_old = self.labels[i];
                self.cluster_counts[z_old] -= 1;
                self.cluster_word_counts[z_old] -= doc_size;

                // modify the map: enclose it in a block so we can borrow views again.
                {
                    let ref mut old_clust_words: HashMap<String, u32> = self.cluster_word_distributions[z_old];
                    for word in doc {
                        *old_clust_words.get_mut(word).unwrap() -= 1_u32;

                        // compact dictionary once a key is exausted.
                        if old_clust_words[word] == 0_u32 {
                            old_clust_words.remove(word);
                        }
                    }
                }

                // update the probability vector
                let p = self.score(&doc);

                // choose the next cluster randomly according to the computed probability
                let z_new: usize = random_choice().random_choice_f64(&self.clusters, &p, 1)[0].clone();

                // transfer document to the new cluster
                if z_new != z_old {
                    total_transfers += 1;
                }
                self.labels[i] = z_new;
                self.cluster_counts[z_new] += 1_u32;
                self.cluster_word_counts[z_new] += doc_size;

                {
                    let ref mut new_clust_words: HashMap<String, u32> = self.cluster_word_distributions[z_new];
                    for word in doc {
                        if !new_clust_words.contains_key(word) {
                            new_clust_words.insert(word.clone(), 0_u32);
                        }
                            *new_clust_words.get_mut(word).unwrap() += 1_u32;
                    }
                }
            }
            let new_number_clusters = self.cluster_word_distributions.iter().map(|c| if c.len()>0 {1} else {0} ).sum();
            println!("Iteration {}: {} docs transferred with {} clusters populated.", it, total_transfers, new_number_clusters);

            // apply ad-hoc convergence test
            if total_transfers==0 && new_number_clusters==number_clusters {
                println!("Converged after {} iterations. Solution has {} clusters.", it, new_number_clusters);
                break
            }
            number_clusters = new_number_clusters;
        }
    }

    pub fn score(&self, doc:&Vec<String>) -> Vec<f64> {
        let mut p = (0..self.K).map(|_| 0_f64).collect::<Vec<f64>>();
        let lD1 = ((self.D - 1) as f64 + (self.K as f64) * self.alpha).ln();
        let doc_size = doc.len() as u32;
        for label in 0_usize..self.K {
            let lN1 = (self.cluster_counts[label] as f64 + self.alpha).ln();
            let mut lN2 = 0_f64;
            let mut lD2 = 0_f64;

            let ref cluster: HashMap<String, u32> = self.cluster_word_distributions[label];

            for word in doc {
                lN2 += (*cluster.get(word).unwrap_or(&0_u32) as f64 + self.beta).ln();
            }
            for j in 1_u32..(doc_size+1) {
                lD2 += ((self.cluster_word_counts[label] + j) as f64 - 1_f64 + self.V * self.beta).ln();
            }
            p[label] = (lN1 - lD1 + lN2 - lD2).exp();
        }
        let pnorm: f64 = p.iter().sum();
        for label in 0_usize..self.K {
            p[label] = p[label] / pnorm;
        }
        p
    }

    fn zero_vector(size:usize) -> Vec<u32>
    {
        let mut v = Vec::<u32>::with_capacity(size);
        for _ in 0_usize..size {
            v.push(0_u32)
        }
        v
    }

}

#[test]
fn simple_run() {
    let mut vocab = HashSet::<String>::new();
    vocab.insert("A".to_string());
    vocab.insert("B".to_string());
    vocab.insert("C".to_string());

    let mut docs = Vec::<Vec<String>>::new();
    docs.push(vec!("A".to_string()));
    docs.push(vec!("A".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));

    let mut model = GSDMM::new(0.1, 0.00001, 10, 30, vocab, docs);
    model.fit();

    // check the total number across all partitions is equal to the number of docs
    assert_eq!(18, model.cluster_counts.iter().sum::<u32>());

    // check that we get three clusters
    assert_eq!(3, model.cluster_counts.into_iter().filter(|x| x>&0_u32 ).collect::<Vec<u32>>().len());

    // check that the clusters are pure
    println!("{:?}", model.cluster_word_distributions);
    let mut check_map = HashMap::<usize,String>::new();
    for (i,label) in vec!("A","A","B","B","B","B","B","B","B","B","C","C","C","C","C","C","C","C").into_iter().enumerate() {
        if check_map.contains_key(&model.labels[i]) {
            assert_eq!(check_map[&model.labels[i]], label);
        } else {
            check_map.insert(model.labels[i], label.to_string());
        }

    }
}
<commit_msg>New documentation<commit_after>extern crate random_choice;

use std::collections::{HashSet, HashMap};
use std::iter::FromIterator;
use std::cmp::max;
use self::random_choice::random_choice;

pub struct GSDMM {
    alpha: f64,
    beta: f64,
    K:usize,
    V:f64,
    D:usize,
    maxit:isize,
    clusters: Vec<usize>,
    pub docs: Vec<Vec<String>>,
    pub labels: Vec<usize>,
    pub cluster_counts: Vec<u32>,
    pub cluster_word_counts:Vec<u32>,
    pub cluster_word_distributions: Vec<HashMap<String,u32>>
}

impl GSDMM {
    pub fn new(alpha:f64, beta:f64, K: usize, maxit:isize, vocab:HashSet<String>, docs:Vec<Vec<String>>) -> GSDMM {
        let D = docs.len();

        // compute utilized vocabulary size.
        let mut utilized_vocab = HashSet::with_capacity(vocab.len());
        for doc in &docs {
            for word in doc {
                utilized_vocab.insert(word.clone());
            }
        }
        let V = utilized_vocab.len() as f64;
        println!("Fitting with alpha={}, beta={}, K={}, maxit={}, vocab size={}", alpha, beta, K, maxit, V as u32);

        let clusters = (0_usize..K).collect::<Vec<usize>>();
        let mut d_z: Vec<usize> = (0_usize..D).map(|_| 0_usize).collect::<Vec<usize>>(); // doc labels
        let mut m_z: Vec<u32> = GSDMM::zero_vector(K);  // cluster sizes
        let mut n_z: Vec<u32> = GSDMM::zero_vector(K);  // cluster word counts
        let mut n_z_w = Vec::<HashMap<String, u32>>::with_capacity(K);  // container for cluster word distributions
        for _ in 0_usize..K {
            let mut m = HashMap::<String, u32>::with_capacity(max(vocab.len() / 10, 100));
            &n_z_w.push(m);
        }

        // randomly initialize cluster assignment
        let p = (0..K).map(|_| 1_f64 / (K as f64)).collect::<Vec<f64>>();

        let choices = random_choice().random_choice_f64(&clusters, &p, D) ;
        for i in 0..D {
            let z = choices[i].clone();
            let ref doc = docs[i];
            d_z[i] = z;
            m_z[z] += 1;
            n_z[z] += doc.len() as u32;
            let ref mut clust_words: HashMap<String, u32> = n_z_w[z];
            for word in doc {
                if !clust_words.contains_key(word) {
                    clust_words.insert(word.clone(), 0_u32);
                }
                    * clust_words.get_mut(word).unwrap() += 1_u32;
            }
        }

        GSDMM {
            alpha: alpha,
            beta: beta,
            K: K,
            V: V,
            D: D,
            maxit:maxit,
            docs:docs,
            clusters: clusters.clone(),     // Don't totally get why we need the clone here!
            labels: d_z,
            cluster_counts: m_z,
            cluster_word_counts: n_z,
            cluster_word_distributions: n_z_w
        }
    }

    pub fn fit(&mut self) {
        let mut number_clusters = self.K;
        for it in 0..self.maxit {
            let mut total_transfers = 0;
            for i in 0..self.D {
                let ref doc = self.docs[i];
                let doc_size = doc.len() as u32;

                // remove the doc from its current cluster
                let z_old = self.labels[i];
                self.cluster_counts[z_old] -= 1;
                self.cluster_word_counts[z_old] -= doc_size;

                // modify the map: enclose it in a block so we can borrow views again.
                {
                    let ref mut old_clust_words: HashMap<String, u32> = self.cluster_word_distributions[z_old];
                    for word in doc {
                        *old_clust_words.get_mut(word).unwrap() -= 1_u32;

                        // compact dictionary once a key is exausted.
                        if old_clust_words[word] == 0_u32 {
                            old_clust_words.remove(word);
                        }
                    }
                }

                // update the probability vector
                let p = self.score(&doc);

                // choose the next cluster randomly according to the computed probability
                let z_new: usize = random_choice().random_choice_f64(&self.clusters, &p, 1)[0].clone();

                // transfer document to the new cluster
                if z_new != z_old {
                    total_transfers += 1;
                }
                self.labels[i] = z_new;
                self.cluster_counts[z_new] += 1_u32;
                self.cluster_word_counts[z_new] += doc_size;

                {
                    let ref mut new_clust_words: HashMap<String, u32> = self.cluster_word_distributions[z_new];
                    for word in doc {
                        if !new_clust_words.contains_key(word) {
                            new_clust_words.insert(word.clone(), 0_u32);
                        }
                            *new_clust_words.get_mut(word).unwrap() += 1_u32;
                    }
                }
            }
            let new_number_clusters = self.cluster_word_distributions.iter().map(|c| if c.len()>0 {1} else {0} ).sum();
            println!("Iteration {}: {} docs transferred with {} clusters populated.", it, total_transfers, new_number_clusters);

            // apply ad-hoc convergence test
            if total_transfers==0 && new_number_clusters==number_clusters {
                println!("Converged after {} iterations. Solution has {} clusters.", it, new_number_clusters);
                break
            }
            number_clusters = new_number_clusters;
        }
    }

    pub fn score(&self, doc:&Vec<String>) -> Vec<f64> {
        /// Score an input document using the formula of Yin and Wang 2014 (equation 3)
        /// http://dbgroup.cs.tsinghua.edu.cn/wangjy/papers/KDD14-GSDMM.pdf
        ///
        /// # Arguments
        ///
        /// * `doc` - A vector of unique string tokens characterizing the document
        ///
        /// # Value
        ///
        /// Vec<f64> - A length K probability vector where each component represents the probability
        /// of the doc belonging to a particular cluster.
        ///

        // We break the formula into the following pieces
        // p = N1*N2/(D1*D2) = exp(lN1 - lD1 + lN2 - lD2)
        // lN1 = log(m_z[z] + alpha)
        // lN2 = log(D - 1 + K*alpha)
        // lN2 = log(product(n_z_w[w] + beta)) = sum(log(n_z_w[w] + beta))
        // lD2 = log(product(n_z[d] + V*beta + i -1)) = sum(log(n_z[d] + V*beta + i -1))
        let mut p = (0..self.K).map(|_| 0_f64).collect::<Vec<f64>>();
        let lD1 = ((self.D - 1) as f64 + (self.K as f64) * self.alpha).ln();
        let doc_size = doc.len() as u32;
        for label in 0_usize..self.K {
            let lN1 = (self.cluster_counts[label] as f64 + self.alpha).ln();
            let mut lN2 = 0_f64;
            let mut lD2 = 0_f64;

            let ref cluster: HashMap<String, u32> = self.cluster_word_distributions[label];

            for word in doc {
                lN2 += (*cluster.get(word).unwrap_or(&0_u32) as f64 + self.beta).ln();
            }
            for j in 1_u32..(doc_size+1) {
                lD2 += ((self.cluster_word_counts[label] + j) as f64 - 1_f64 + self.V * self.beta).ln();
            }
            p[label] = (lN1 - lD1 + lN2 - lD2).exp();
        }

        // normalize the probability
        let pnorm: f64 = p.iter().sum();
        for label in 0_usize..self.K {
            p[label] = p[label] / pnorm;
        }
        p
    }

    fn zero_vector(size:usize) -> Vec<u32>
    {
        let mut v = Vec::<u32>::with_capacity(size);
        for _ in 0_usize..size {
            v.push(0_u32)
        }
        v
    }

}

#[test]
fn simple_run() {
    let mut vocab = HashSet::<String>::new();
    vocab.insert("A".to_string());
    vocab.insert("B".to_string());
    vocab.insert("C".to_string());

    let mut docs = Vec::<Vec<String>>::new();
    docs.push(vec!("A".to_string()));
    docs.push(vec!("A".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("B".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));
    docs.push(vec!("C".to_string()));

    let mut model = GSDMM::new(0.1, 0.00001, 10, 30, vocab, docs);
    model.fit();

    // check the total number across all partitions is equal to the number of docs
    assert_eq!(18, model.cluster_counts.iter().sum::<u32>());

    // check that we get three clusters
    assert_eq!(3, model.cluster_counts.into_iter().filter(|x| x>&0_u32 ).collect::<Vec<u32>>().len());

    // check that the clusters are pure
    println!("{:?}", model.cluster_word_distributions);
    let mut check_map = HashMap::<usize,String>::new();
    for (i,label) in vec!("A","A","B","B","B","B","B","B","B","B","C","C","C","C","C","C","C","C").into_iter().enumerate() {
        if check_map.contains_key(&model.labels[i]) {
            assert_eq!(check_map[&model.labels[i]], label);
        } else {
            check_map.insert(model.labels[i], label.to_string());
        }

    }
}
<|endoftext|>


<commit_before>// Copyright 2012-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.

//! Crate docs

#![allow(bad_style, raw_pointer_derive, overflowing_literals, improper_ctypes)]
#![crate_type = "rlib"]
#![crate_name = "libc"]
#![cfg_attr(dox, feature(no_core, lang_items))]
#![cfg_attr(dox, no_core)]
#![doc(html_logo_url = "https://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "https://doc.rust-lang.org/favicon.ico")]

#![cfg_attr(all(target_os = "linux", target_arch = "x86_64"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/x86_64-unknown-linux-gnu"
))]
#![cfg_attr(all(target_os = "linux", target_arch = "x86"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/i686-unknown-linux-gnu"
))]
#![cfg_attr(all(target_os = "linux", target_arch = "arm"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/arm-unknown-linux-gnueabihf"
))]
#![cfg_attr(all(target_os = "linux", target_arch = "mips"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/mips-unknown-linux-gnu"
))]
#![cfg_attr(all(target_os = "linux", target_arch = "aarch64"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/aarch64-unknown-linux-gnu"
))]
#![cfg_attr(all(target_os = "linux", target_env = "musl"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/x86_64-unknown-linux-musl"
))]
#![cfg_attr(all(target_os = "macos", target_arch = "x86_64"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/x86_64-apple-darwin"
))]
#![cfg_attr(all(target_os = "macos", target_arch = "x86"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/i686-apple-darwin"
))]
#![cfg_attr(all(windows, target_arch = "x86_64", target_env = "gnu"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/x86_64-pc-windows-gnu"
))]
#![cfg_attr(all(windows, target_arch = "x86", target_env = "gnu"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/i686-pc-windows-gnu"
))]
#![cfg_attr(all(windows, target_arch = "x86_64", target_env = "msvc"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/x86_64-pc-windows-msvc"
))]
#![cfg_attr(all(windows, target_arch = "x86", target_env = "msvc"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/i686-pc-windows-msvc"
))]
#![cfg_attr(all(target_os = "android"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/arm-linux-androideabi"
))]

// Attributes needed when building as part of the standard library
#![cfg_attr(stdbuild, feature(no_std, core, core_slice_ext, staged_api))]
#![cfg_attr(stdbuild, no_std)]
#![cfg_attr(stdbuild, staged_api)]
#![cfg_attr(stdbuild, unstable(feature = "libc",
                               reason = "use `libc` from crates.io",
                               issue = "27783"))]

#[cfg(all(not(stdbuild), not(dox)))]
extern crate std as core;

#[macro_use] mod macros;
mod dox;

// Use repr(u8) as LLVM expects `void*` to be the same as `i8*` to help enable
// more optimization opportunities around it recognizing things like
// malloc/free.
#[repr(u8)]
pub enum c_void {
    // Two dummy variants so the #[repr] attribute can be used.
    #[doc(hidden)]
    __variant1,
    #[doc(hidden)]
    __variant2,
}

pub type int8_t = i8;
pub type int16_t = i16;
pub type int32_t = i32;
pub type int64_t = i64;
pub type uint8_t = u8;
pub type uint16_t = u16;
pub type uint32_t = u32;
pub type uint64_t = u64;

pub type c_schar = i8;
pub type c_uchar = u8;
pub type c_short = i16;
pub type c_ushort = u16;
pub type c_int = i32;
pub type c_uint = u32;
pub type c_float = f32;
pub type c_double = f64;
pub type c_longlong = i64;
pub type c_ulonglong = u64;
pub type intmax_t = i64;
pub type uintmax_t = u64;

pub type size_t = usize;
pub type ptrdiff_t = isize;
pub type intptr_t = isize;
pub type uintptr_t = usize;
pub type ssize_t = isize;

pub enum FILE {}
pub enum fpos_t {} // TODO: fill this out with a struct

extern {
    pub fn isalnum(c: c_int) -> c_int;
    pub fn isalpha(c: c_int) -> c_int;
    pub fn iscntrl(c: c_int) -> c_int;
    pub fn isdigit(c: c_int) -> c_int;
    pub fn isgraph(c: c_int) -> c_int;
    pub fn islower(c: c_int) -> c_int;
    pub fn isprint(c: c_int) -> c_int;
    pub fn ispunct(c: c_int) -> c_int;
    pub fn isspace(c: c_int) -> c_int;
    pub fn isupper(c: c_int) -> c_int;
    pub fn isxdigit(c: c_int) -> c_int;
    pub fn tolower(c: c_int) -> c_int;
    pub fn toupper(c: c_int) -> c_int;

    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "fopen$UNIX2003")]
    pub fn fopen(filename: *const c_char,
                 mode: *const c_char) -> *mut FILE;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "freopen$UNIX2003")]
    pub fn freopen(filename: *const c_char, mode: *const c_char,
                   file: *mut FILE) -> *mut FILE;
    pub fn fflush(file: *mut FILE) -> c_int;
    pub fn fclose(file: *mut FILE) -> c_int;
    pub fn remove(filename: *const c_char) -> c_int;
    pub fn rename(oldname: *const c_char, newname: *const c_char) -> c_int;
    pub fn tmpfile() -> *mut FILE;
    pub fn setvbuf(stream: *mut FILE,
                   buffer: *mut c_char,
                   mode: c_int,
                   size: size_t) -> c_int;
    pub fn setbuf(stream: *mut FILE, buf: *mut c_char);
    pub fn fgetc(stream: *mut FILE) -> c_int;
    pub fn fgets(buf: *mut c_char, n: c_int, stream: *mut FILE) -> *mut c_char;
    pub fn fputc(c: c_int, stream: *mut FILE) -> c_int;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "fputs$UNIX2003")]
    pub fn fputs(s: *const c_char, stream: *mut FILE)-> c_int;
    pub fn puts(s: *const c_char) -> c_int;
    pub fn ungetc(c: c_int, stream: *mut FILE) -> c_int;
    pub fn fread(ptr: *mut c_void,
                 size: size_t,
                 nobj: size_t,
                 stream: *mut FILE)
                 -> size_t;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "fwrite$UNIX2003")]
    pub fn fwrite(ptr: *const c_void,
                  size: size_t,
                  nobj: size_t,
                  stream: *mut FILE)
                  -> size_t;
    pub fn fseek(stream: *mut FILE, offset: c_long, whence: c_int) -> c_int;
    pub fn ftell(stream: *mut FILE) -> c_long;
    pub fn rewind(stream: *mut FILE);
    pub fn fgetpos(stream: *mut FILE, ptr: *mut fpos_t) -> c_int;
    pub fn fsetpos(stream: *mut FILE, ptr: *const fpos_t) -> c_int;
    pub fn feof(stream: *mut FILE) -> c_int;
    pub fn ferror(stream: *mut FILE) -> c_int;
    pub fn perror(s: *const c_char);
    pub fn atoi(s: *const c_char) -> c_int;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "strtod$UNIX2003")]
    pub fn strtod(s: *const c_char, endp: *mut *mut c_char) -> c_double;
    pub fn strtol(s: *const c_char,
                  endp: *mut *mut c_char, base: c_int) -> c_long;
    pub fn strtoul(s: *const c_char, endp: *mut *mut c_char,
                   base: c_int) -> c_ulong;
    pub fn calloc(nobj: size_t, size: size_t) -> *mut c_void;
    pub fn malloc(size: size_t) -> *mut c_void;
    pub fn realloc(p: *mut c_void, size: size_t) -> *mut c_void;
    pub fn free(p: *mut c_void);
    pub fn exit(status: c_int) -> !;
    pub fn _exit(status: c_int) -> !;
    pub fn atexit(cb: extern fn()) -> c_int;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "system$UNIX2003")]
    pub fn system(s: *const c_char) -> c_int;
    pub fn getenv(s: *const c_char) -> *mut c_char;

    pub fn strcpy(dst: *mut c_char, src: *const c_char) -> *mut c_char;
    pub fn strncpy(dst: *mut c_char, src: *const c_char, n: size_t)
                   -> *mut c_char;
    pub fn strcat(s: *mut c_char, ct: *const c_char) -> *mut c_char;
    pub fn strncat(s: *mut c_char, ct: *const c_char, n: size_t) -> *mut c_char;
    pub fn strcmp(cs: *const c_char, ct: *const c_char) -> c_int;
    pub fn strncmp(cs: *const c_char, ct: *const c_char, n: size_t) -> c_int;
    pub fn strcoll(cs: *const c_char, ct: *const c_char) -> c_int;
    pub fn strchr(cs: *const c_char, c: c_int) -> *mut c_char;
    pub fn strrchr(cs: *const c_char, c: c_int) -> *mut c_char;
    pub fn strspn(cs: *const c_char, ct: *const c_char) -> size_t;
    pub fn strcspn(cs: *const c_char, ct: *const c_char) -> size_t;
    pub fn strpbrk(cs: *const c_char, ct: *const c_char) -> *mut c_char;
    pub fn strstr(cs: *const c_char, ct: *const c_char) -> *mut c_char;
    pub fn strlen(cs: *const c_char) -> size_t;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "strerror$UNIX2003")]
    pub fn strerror(n: c_int) -> *mut c_char;
    pub fn strtok(s: *mut c_char, t: *const c_char) -> *mut c_char;
    pub fn strxfrm(s: *mut c_char, ct: *const c_char, n: size_t) -> size_t;
    pub fn wcslen(buf: *const wchar_t) -> size_t;

    pub fn memcmp(cx: *const c_void, ct: *const c_void, n: size_t) -> c_int;
    pub fn memchr(cx: *const c_void, c: c_int, n: size_t) -> *mut c_void;
}

// These are all inline functions on android, so they end up just being entirely
// missing on that platform.
#[cfg(not(target_os = "android"))]
extern {
    pub fn abs(i: c_int) -> c_int;
    pub fn atof(s: *const c_char) -> c_double;
    pub fn labs(i: c_long) -> c_long;
    pub fn rand() -> c_int;
    pub fn srand(seed: c_uint);
}

cfg_if! {
    if #[cfg(windows)] {
        mod windows;
        pub use windows::*;
    } else {
        mod unix;
        pub use unix::*;
    }
}
<commit_msg>Allow warnings during a libstd build<commit_after>// Copyright 2012-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.

//! Crate docs

#![allow(bad_style, raw_pointer_derive, overflowing_literals, improper_ctypes)]
#![crate_type = "rlib"]
#![crate_name = "libc"]
#![cfg_attr(dox, feature(no_core, lang_items))]
#![cfg_attr(dox, no_core)]
#![doc(html_logo_url = "https://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "https://doc.rust-lang.org/favicon.ico")]

#![cfg_attr(all(target_os = "linux", target_arch = "x86_64"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/x86_64-unknown-linux-gnu"
))]
#![cfg_attr(all(target_os = "linux", target_arch = "x86"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/i686-unknown-linux-gnu"
))]
#![cfg_attr(all(target_os = "linux", target_arch = "arm"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/arm-unknown-linux-gnueabihf"
))]
#![cfg_attr(all(target_os = "linux", target_arch = "mips"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/mips-unknown-linux-gnu"
))]
#![cfg_attr(all(target_os = "linux", target_arch = "aarch64"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/aarch64-unknown-linux-gnu"
))]
#![cfg_attr(all(target_os = "linux", target_env = "musl"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/x86_64-unknown-linux-musl"
))]
#![cfg_attr(all(target_os = "macos", target_arch = "x86_64"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/x86_64-apple-darwin"
))]
#![cfg_attr(all(target_os = "macos", target_arch = "x86"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/i686-apple-darwin"
))]
#![cfg_attr(all(windows, target_arch = "x86_64", target_env = "gnu"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/x86_64-pc-windows-gnu"
))]
#![cfg_attr(all(windows, target_arch = "x86", target_env = "gnu"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/i686-pc-windows-gnu"
))]
#![cfg_attr(all(windows, target_arch = "x86_64", target_env = "msvc"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/x86_64-pc-windows-msvc"
))]
#![cfg_attr(all(windows, target_arch = "x86", target_env = "msvc"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/i686-pc-windows-msvc"
))]
#![cfg_attr(all(target_os = "android"), doc(
    html_root_url = "https://doc.rust-lang.org/libc/arm-linux-androideabi"
))]

// Attributes needed when building as part of the standard library
#![cfg_attr(stdbuild, feature(no_std, core, core_slice_ext, staged_api))]
#![cfg_attr(stdbuild, no_std)]
#![cfg_attr(stdbuild, staged_api)]
#![cfg_attr(stdbuild, allow(warnings))]
#![cfg_attr(stdbuild, unstable(feature = "libc",
                               reason = "use `libc` from crates.io",
                               issue = "27783"))]

#[cfg(all(not(stdbuild), not(dox)))]
extern crate std as core;

#[macro_use] mod macros;
mod dox;

// Use repr(u8) as LLVM expects `void*` to be the same as `i8*` to help enable
// more optimization opportunities around it recognizing things like
// malloc/free.
#[repr(u8)]
pub enum c_void {
    // Two dummy variants so the #[repr] attribute can be used.
    #[doc(hidden)]
    __variant1,
    #[doc(hidden)]
    __variant2,
}

pub type int8_t = i8;
pub type int16_t = i16;
pub type int32_t = i32;
pub type int64_t = i64;
pub type uint8_t = u8;
pub type uint16_t = u16;
pub type uint32_t = u32;
pub type uint64_t = u64;

pub type c_schar = i8;
pub type c_uchar = u8;
pub type c_short = i16;
pub type c_ushort = u16;
pub type c_int = i32;
pub type c_uint = u32;
pub type c_float = f32;
pub type c_double = f64;
pub type c_longlong = i64;
pub type c_ulonglong = u64;
pub type intmax_t = i64;
pub type uintmax_t = u64;

pub type size_t = usize;
pub type ptrdiff_t = isize;
pub type intptr_t = isize;
pub type uintptr_t = usize;
pub type ssize_t = isize;

pub enum FILE {}
pub enum fpos_t {} // TODO: fill this out with a struct

extern {
    pub fn isalnum(c: c_int) -> c_int;
    pub fn isalpha(c: c_int) -> c_int;
    pub fn iscntrl(c: c_int) -> c_int;
    pub fn isdigit(c: c_int) -> c_int;
    pub fn isgraph(c: c_int) -> c_int;
    pub fn islower(c: c_int) -> c_int;
    pub fn isprint(c: c_int) -> c_int;
    pub fn ispunct(c: c_int) -> c_int;
    pub fn isspace(c: c_int) -> c_int;
    pub fn isupper(c: c_int) -> c_int;
    pub fn isxdigit(c: c_int) -> c_int;
    pub fn tolower(c: c_int) -> c_int;
    pub fn toupper(c: c_int) -> c_int;

    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "fopen$UNIX2003")]
    pub fn fopen(filename: *const c_char,
                 mode: *const c_char) -> *mut FILE;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "freopen$UNIX2003")]
    pub fn freopen(filename: *const c_char, mode: *const c_char,
                   file: *mut FILE) -> *mut FILE;
    pub fn fflush(file: *mut FILE) -> c_int;
    pub fn fclose(file: *mut FILE) -> c_int;
    pub fn remove(filename: *const c_char) -> c_int;
    pub fn rename(oldname: *const c_char, newname: *const c_char) -> c_int;
    pub fn tmpfile() -> *mut FILE;
    pub fn setvbuf(stream: *mut FILE,
                   buffer: *mut c_char,
                   mode: c_int,
                   size: size_t) -> c_int;
    pub fn setbuf(stream: *mut FILE, buf: *mut c_char);
    pub fn fgetc(stream: *mut FILE) -> c_int;
    pub fn fgets(buf: *mut c_char, n: c_int, stream: *mut FILE) -> *mut c_char;
    pub fn fputc(c: c_int, stream: *mut FILE) -> c_int;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "fputs$UNIX2003")]
    pub fn fputs(s: *const c_char, stream: *mut FILE)-> c_int;
    pub fn puts(s: *const c_char) -> c_int;
    pub fn ungetc(c: c_int, stream: *mut FILE) -> c_int;
    pub fn fread(ptr: *mut c_void,
                 size: size_t,
                 nobj: size_t,
                 stream: *mut FILE)
                 -> size_t;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "fwrite$UNIX2003")]
    pub fn fwrite(ptr: *const c_void,
                  size: size_t,
                  nobj: size_t,
                  stream: *mut FILE)
                  -> size_t;
    pub fn fseek(stream: *mut FILE, offset: c_long, whence: c_int) -> c_int;
    pub fn ftell(stream: *mut FILE) -> c_long;
    pub fn rewind(stream: *mut FILE);
    pub fn fgetpos(stream: *mut FILE, ptr: *mut fpos_t) -> c_int;
    pub fn fsetpos(stream: *mut FILE, ptr: *const fpos_t) -> c_int;
    pub fn feof(stream: *mut FILE) -> c_int;
    pub fn ferror(stream: *mut FILE) -> c_int;
    pub fn perror(s: *const c_char);
    pub fn atoi(s: *const c_char) -> c_int;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "strtod$UNIX2003")]
    pub fn strtod(s: *const c_char, endp: *mut *mut c_char) -> c_double;
    pub fn strtol(s: *const c_char,
                  endp: *mut *mut c_char, base: c_int) -> c_long;
    pub fn strtoul(s: *const c_char, endp: *mut *mut c_char,
                   base: c_int) -> c_ulong;
    pub fn calloc(nobj: size_t, size: size_t) -> *mut c_void;
    pub fn malloc(size: size_t) -> *mut c_void;
    pub fn realloc(p: *mut c_void, size: size_t) -> *mut c_void;
    pub fn free(p: *mut c_void);
    pub fn exit(status: c_int) -> !;
    pub fn _exit(status: c_int) -> !;
    pub fn atexit(cb: extern fn()) -> c_int;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "system$UNIX2003")]
    pub fn system(s: *const c_char) -> c_int;
    pub fn getenv(s: *const c_char) -> *mut c_char;

    pub fn strcpy(dst: *mut c_char, src: *const c_char) -> *mut c_char;
    pub fn strncpy(dst: *mut c_char, src: *const c_char, n: size_t)
                   -> *mut c_char;
    pub fn strcat(s: *mut c_char, ct: *const c_char) -> *mut c_char;
    pub fn strncat(s: *mut c_char, ct: *const c_char, n: size_t) -> *mut c_char;
    pub fn strcmp(cs: *const c_char, ct: *const c_char) -> c_int;
    pub fn strncmp(cs: *const c_char, ct: *const c_char, n: size_t) -> c_int;
    pub fn strcoll(cs: *const c_char, ct: *const c_char) -> c_int;
    pub fn strchr(cs: *const c_char, c: c_int) -> *mut c_char;
    pub fn strrchr(cs: *const c_char, c: c_int) -> *mut c_char;
    pub fn strspn(cs: *const c_char, ct: *const c_char) -> size_t;
    pub fn strcspn(cs: *const c_char, ct: *const c_char) -> size_t;
    pub fn strpbrk(cs: *const c_char, ct: *const c_char) -> *mut c_char;
    pub fn strstr(cs: *const c_char, ct: *const c_char) -> *mut c_char;
    pub fn strlen(cs: *const c_char) -> size_t;
    #[cfg_attr(all(target_os = "macos", target_arch = "x86"),
               link_name = "strerror$UNIX2003")]
    pub fn strerror(n: c_int) -> *mut c_char;
    pub fn strtok(s: *mut c_char, t: *const c_char) -> *mut c_char;
    pub fn strxfrm(s: *mut c_char, ct: *const c_char, n: size_t) -> size_t;
    pub fn wcslen(buf: *const wchar_t) -> size_t;

    pub fn memcmp(cx: *const c_void, ct: *const c_void, n: size_t) -> c_int;
    pub fn memchr(cx: *const c_void, c: c_int, n: size_t) -> *mut c_void;
}

// These are all inline functions on android, so they end up just being entirely
// missing on that platform.
#[cfg(not(target_os = "android"))]
extern {
    pub fn abs(i: c_int) -> c_int;
    pub fn atof(s: *const c_char) -> c_double;
    pub fn labs(i: c_long) -> c_long;
    pub fn rand() -> c_int;
    pub fn srand(seed: c_uint);
}

cfg_if! {
    if #[cfg(windows)] {
        mod windows;
        pub use windows::*;
    } else {
        mod unix;
        pub use unix::*;
    }
}
<|endoftext|>


<commit_before>//! API to generate `.rs` files.
//!
//! This API requires `protoc` command present in `$PATH`.
//!
//! ```
//! extern crate protoc_rust;
//!
//! fn main() {
//!     protobuf_codegen_pure::Codegen::new()
//!         .out_dir("src/protos")
//!         .inputs(&["protos/a.proto", "protos/b.proto"]),
//!         .include("protos")
//!         .run()
//!         .expect("Codegen failed.");
//! }
//! ```
//!
//! It is advisable that `protobuf-codegen-pure` build-dependecy version be the same as
//! `protobuf` dependency.
//!
//! The alternative is to use `protoc-rust` crate.

#![deny(missing_docs)]
#![deny(intra_doc_link_resolution_failure)]

extern crate protobuf;
extern crate protobuf_codegen;

mod convert;

use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::fs;
use std::io;
use std::path::Path;
use std::path::PathBuf;

mod model;
mod parser;

use protobuf_codegen::amend_io_error;
pub use protobuf_codegen::Customize;

#[cfg(test)]
mod test_against_protobuf_protos;

/// Invoke pure rust codegen. See [crate docs](crate) for example.
// TODO: merge with protoc-rust def
#[derive(Debug, Default)]
pub struct Codegen {
    /// --lang_out= param
    out_dir: PathBuf,
    /// -I args
    includes: Vec<PathBuf>,
    /// List of .proto files to compile
    inputs: Vec<PathBuf>,
    /// Customize code generation
    customize: Customize,
}

impl Codegen {
    /// Fresh new codegen object.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the output directory for codegen.
    pub fn out_dir(&mut self, out_dir: impl AsRef<Path>) -> &mut Self {
        self.out_dir = out_dir.as_ref().to_owned();
        self
    }

    /// Add an include directory.
    pub fn include(&mut self, include: impl AsRef<Path>) -> &mut Self {
        self.includes.push(include.as_ref().to_owned());
        self
    }

    /// Add include directories.
    pub fn includes(&mut self, includes: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for include in includes {
            self.include(include);
        }
        self
    }

    /// Add an input (`.proto` file).
    pub fn input(&mut self, input: impl AsRef<Path>) -> &mut Self {
        self.inputs.push(input.as_ref().to_owned());
        self
    }

    /// Add inputs (`.proto` files).
    pub fn inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for input in inputs {
            self.input(input);
        }
        self
    }

    /// Specify generated code [`Customize`] object.
    pub fn customize(&mut self, customize: Customize) -> &mut Self {
        self.customize = customize;
        self
    }

    /// Like `protoc --rust_out=...` but without requiring `protoc` or `protoc-gen-rust`
    /// commands in `$PATH`.
    pub fn run(&self) -> io::Result<()> {
        let p = parse_and_typecheck(&self.includes, &self.inputs)?;

        protobuf_codegen::gen_and_write(
            &p.file_descriptors,
            &format!("protobuf-codegen-pure={}", env!("CARGO_PKG_VERSION")),
            &p.relative_paths,
            &self.out_dir,
            &self.customize,
        )
    }
}

#[derive(Clone)]
struct FileDescriptorPair {
    parsed: model::FileDescriptor,
    descriptor: protobuf::descriptor::FileDescriptorProto,
}

#[derive(Debug)]
enum CodegenError {
    ParserErrorWithLocation(parser::ParserErrorWithLocation),
    ConvertError(convert::ConvertError),
}

impl From<parser::ParserErrorWithLocation> for CodegenError {
    fn from(e: parser::ParserErrorWithLocation) -> Self {
        CodegenError::ParserErrorWithLocation(e)
    }
}

impl From<convert::ConvertError> for CodegenError {
    fn from(e: convert::ConvertError) -> Self {
        CodegenError::ConvertError(e)
    }
}

#[derive(Debug)]
struct WithFileError {
    file: String,
    error: CodegenError,
}

impl fmt::Display for WithFileError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "WithFileError")
    }
}

impl Error for WithFileError {
    fn description(&self) -> &str {
        "WithFileError"
    }
}

struct Run<'a> {
    parsed_files: HashMap<PathBuf, FileDescriptorPair>,
    includes: &'a [PathBuf],
}

impl<'a> Run<'a> {
    fn get_file_and_all_deps_already_parsed(
        &self,
        protobuf_path: &Path,
        result: &mut HashMap<PathBuf, FileDescriptorPair>,
    ) {
        if let Some(_) = result.get(protobuf_path) {
            return;
        }

        let pair = self
            .parsed_files
            .get(protobuf_path)
            .expect("must be already parsed");
        result.insert(protobuf_path.to_owned(), pair.clone());

        self.get_all_deps_already_parsed(&pair.parsed, result);
    }

    fn get_all_deps_already_parsed(
        &self,
        parsed: &model::FileDescriptor,
        result: &mut HashMap<PathBuf, FileDescriptorPair>,
    ) {
        for import in &parsed.import_paths {
            self.get_file_and_all_deps_already_parsed(Path::new(import), result);
        }
    }

    fn add_file(&mut self, protobuf_path: &Path, fs_path: &Path) -> io::Result<()> {
        if let Some(_) = self.parsed_files.get(protobuf_path) {
            return Ok(());
        }

        let content = fs::read_to_string(fs_path)
            .map_err(|e| amend_io_error(e, format!("failed to read {:?}", fs_path)))?;

        let parsed = model::FileDescriptor::parse(content).map_err(|e| {
            io::Error::new(
                io::ErrorKind::Other,
                WithFileError {
                    file: format!("{}", fs_path.display()),
                    error: e.into(),
                },
            )
        })?;

        for import_path in &parsed.import_paths {
            self.add_imported_file(Path::new(import_path))?;
        }

        let mut this_file_deps = HashMap::new();
        self.get_all_deps_already_parsed(&parsed, &mut this_file_deps);

        let this_file_deps: Vec<_> = this_file_deps.into_iter().map(|(_, v)| v.parsed).collect();

        let descriptor = convert::file_descriptor(protobuf_path, &parsed, &this_file_deps)
            .map_err(|e| {
                io::Error::new(
                    io::ErrorKind::Other,
                    WithFileError {
                        file: format!("{}", fs_path.display()),
                        error: e.into(),
                    },
                )
            })?;

        self.parsed_files.insert(
            protobuf_path.to_owned(),
            FileDescriptorPair { parsed, descriptor },
        );

        Ok(())
    }

    fn add_imported_file(&mut self, protobuf_path: &Path) -> io::Result<()> {
        for include_dir in self.includes {
            let fs_path = include_dir.join(protobuf_path);
            if fs_path.exists() {
                return self.add_file(protobuf_path, &fs_path);
            }
        }

        Err(io::Error::new(
            io::ErrorKind::Other,
            format!(
                "protobuf path {:?} is not found in import path {:?}",
                protobuf_path, self.includes
            ),
        ))
    }

    fn add_fs_file(&mut self, fs_path: &Path) -> io::Result<PathBuf> {
        let relative_path = self
            .includes
            .iter()
            .filter_map(|include_dir| fs_path.strip_prefix(include_dir).ok())
            .next();

        match relative_path {
            Some(relative_path) => {
                assert!(relative_path.is_relative());
                self.add_file(relative_path, fs_path)?;
                Ok(relative_path.to_owned())
            }
            None => Err(io::Error::new(
                io::ErrorKind::Other,
                format!(
                    "file {:?} must reside in include path {:?}",
                    fs_path, self.includes
                ),
            )),
        }
    }
}

#[doc(hidden)]
pub struct ParsedAndTypechecked {
    pub relative_paths: Vec<PathBuf>,
    pub file_descriptors: Vec<protobuf::descriptor::FileDescriptorProto>,
}

#[doc(hidden)]
pub fn parse_and_typecheck(
    includes: &[PathBuf],
    input: &[PathBuf],
) -> io::Result<ParsedAndTypechecked> {
    let mut run = Run {
        parsed_files: HashMap::new(),
        includes: includes,
    };

    let mut relative_paths = Vec::new();

    for input in input {
        relative_paths.push(run.add_fs_file(input)?);
    }

    let file_descriptors: Vec<_> = run
        .parsed_files
        .into_iter()
        .map(|(_, v)| v.descriptor)
        .collect();

    Ok(ParsedAndTypechecked {
        relative_paths,
        file_descriptors,
    })
}
<commit_msg>Fix doc<commit_after>//! API to generate `.rs` files.
//!
//! This API does not require `protoc` command present in `$PATH`.
//!
//! ```
//! extern crate protoc_rust;
//!
//! fn main() {
//!     protobuf_codegen_pure::Codegen::new()
//!         .out_dir("src/protos")
//!         .inputs(&["protos/a.proto", "protos/b.proto"]),
//!         .include("protos")
//!         .run()
//!         .expect("Codegen failed.");
//! }
//! ```
//!
//! It is advisable that `protobuf-codegen-pure` build-dependecy version be the same as
//! `protobuf` dependency.
//!
//! The alternative is to use `protoc-rust` crate.

#![deny(missing_docs)]
#![deny(intra_doc_link_resolution_failure)]

extern crate protobuf;
extern crate protobuf_codegen;

mod convert;

use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::fs;
use std::io;
use std::path::Path;
use std::path::PathBuf;

mod model;
mod parser;

use protobuf_codegen::amend_io_error;
pub use protobuf_codegen::Customize;

#[cfg(test)]
mod test_against_protobuf_protos;

/// Invoke pure rust codegen. See [crate docs](crate) for example.
// TODO: merge with protoc-rust def
#[derive(Debug, Default)]
pub struct Codegen {
    /// --lang_out= param
    out_dir: PathBuf,
    /// -I args
    includes: Vec<PathBuf>,
    /// List of .proto files to compile
    inputs: Vec<PathBuf>,
    /// Customize code generation
    customize: Customize,
}

impl Codegen {
    /// Fresh new codegen object.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the output directory for codegen.
    pub fn out_dir(&mut self, out_dir: impl AsRef<Path>) -> &mut Self {
        self.out_dir = out_dir.as_ref().to_owned();
        self
    }

    /// Add an include directory.
    pub fn include(&mut self, include: impl AsRef<Path>) -> &mut Self {
        self.includes.push(include.as_ref().to_owned());
        self
    }

    /// Add include directories.
    pub fn includes(&mut self, includes: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for include in includes {
            self.include(include);
        }
        self
    }

    /// Add an input (`.proto` file).
    pub fn input(&mut self, input: impl AsRef<Path>) -> &mut Self {
        self.inputs.push(input.as_ref().to_owned());
        self
    }

    /// Add inputs (`.proto` files).
    pub fn inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for input in inputs {
            self.input(input);
        }
        self
    }

    /// Specify generated code [`Customize`] object.
    pub fn customize(&mut self, customize: Customize) -> &mut Self {
        self.customize = customize;
        self
    }

    /// Like `protoc --rust_out=...` but without requiring `protoc` or `protoc-gen-rust`
    /// commands in `$PATH`.
    pub fn run(&self) -> io::Result<()> {
        let p = parse_and_typecheck(&self.includes, &self.inputs)?;

        protobuf_codegen::gen_and_write(
            &p.file_descriptors,
            &format!("protobuf-codegen-pure={}", env!("CARGO_PKG_VERSION")),
            &p.relative_paths,
            &self.out_dir,
            &self.customize,
        )
    }
}

#[derive(Clone)]
struct FileDescriptorPair {
    parsed: model::FileDescriptor,
    descriptor: protobuf::descriptor::FileDescriptorProto,
}

#[derive(Debug)]
enum CodegenError {
    ParserErrorWithLocation(parser::ParserErrorWithLocation),
    ConvertError(convert::ConvertError),
}

impl From<parser::ParserErrorWithLocation> for CodegenError {
    fn from(e: parser::ParserErrorWithLocation) -> Self {
        CodegenError::ParserErrorWithLocation(e)
    }
}

impl From<convert::ConvertError> for CodegenError {
    fn from(e: convert::ConvertError) -> Self {
        CodegenError::ConvertError(e)
    }
}

#[derive(Debug)]
struct WithFileError {
    file: String,
    error: CodegenError,
}

impl fmt::Display for WithFileError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "WithFileError")
    }
}

impl Error for WithFileError {
    fn description(&self) -> &str {
        "WithFileError"
    }
}

struct Run<'a> {
    parsed_files: HashMap<PathBuf, FileDescriptorPair>,
    includes: &'a [PathBuf],
}

impl<'a> Run<'a> {
    fn get_file_and_all_deps_already_parsed(
        &self,
        protobuf_path: &Path,
        result: &mut HashMap<PathBuf, FileDescriptorPair>,
    ) {
        if let Some(_) = result.get(protobuf_path) {
            return;
        }

        let pair = self
            .parsed_files
            .get(protobuf_path)
            .expect("must be already parsed");
        result.insert(protobuf_path.to_owned(), pair.clone());

        self.get_all_deps_already_parsed(&pair.parsed, result);
    }

    fn get_all_deps_already_parsed(
        &self,
        parsed: &model::FileDescriptor,
        result: &mut HashMap<PathBuf, FileDescriptorPair>,
    ) {
        for import in &parsed.import_paths {
            self.get_file_and_all_deps_already_parsed(Path::new(import), result);
        }
    }

    fn add_file(&mut self, protobuf_path: &Path, fs_path: &Path) -> io::Result<()> {
        if let Some(_) = self.parsed_files.get(protobuf_path) {
            return Ok(());
        }

        let content = fs::read_to_string(fs_path)
            .map_err(|e| amend_io_error(e, format!("failed to read {:?}", fs_path)))?;

        let parsed = model::FileDescriptor::parse(content).map_err(|e| {
            io::Error::new(
                io::ErrorKind::Other,
                WithFileError {
                    file: format!("{}", fs_path.display()),
                    error: e.into(),
                },
            )
        })?;

        for import_path in &parsed.import_paths {
            self.add_imported_file(Path::new(import_path))?;
        }

        let mut this_file_deps = HashMap::new();
        self.get_all_deps_already_parsed(&parsed, &mut this_file_deps);

        let this_file_deps: Vec<_> = this_file_deps.into_iter().map(|(_, v)| v.parsed).collect();

        let descriptor = convert::file_descriptor(protobuf_path, &parsed, &this_file_deps)
            .map_err(|e| {
                io::Error::new(
                    io::ErrorKind::Other,
                    WithFileError {
                        file: format!("{}", fs_path.display()),
                        error: e.into(),
                    },
                )
            })?;

        self.parsed_files.insert(
            protobuf_path.to_owned(),
            FileDescriptorPair { parsed, descriptor },
        );

        Ok(())
    }

    fn add_imported_file(&mut self, protobuf_path: &Path) -> io::Result<()> {
        for include_dir in self.includes {
            let fs_path = include_dir.join(protobuf_path);
            if fs_path.exists() {
                return self.add_file(protobuf_path, &fs_path);
            }
        }

        Err(io::Error::new(
            io::ErrorKind::Other,
            format!(
                "protobuf path {:?} is not found in import path {:?}",
                protobuf_path, self.includes
            ),
        ))
    }

    fn add_fs_file(&mut self, fs_path: &Path) -> io::Result<PathBuf> {
        let relative_path = self
            .includes
            .iter()
            .filter_map(|include_dir| fs_path.strip_prefix(include_dir).ok())
            .next();

        match relative_path {
            Some(relative_path) => {
                assert!(relative_path.is_relative());
                self.add_file(relative_path, fs_path)?;
                Ok(relative_path.to_owned())
            }
            None => Err(io::Error::new(
                io::ErrorKind::Other,
                format!(
                    "file {:?} must reside in include path {:?}",
                    fs_path, self.includes
                ),
            )),
        }
    }
}

#[doc(hidden)]
pub struct ParsedAndTypechecked {
    pub relative_paths: Vec<PathBuf>,
    pub file_descriptors: Vec<protobuf::descriptor::FileDescriptorProto>,
}

#[doc(hidden)]
pub fn parse_and_typecheck(
    includes: &[PathBuf],
    input: &[PathBuf],
) -> io::Result<ParsedAndTypechecked> {
    let mut run = Run {
        parsed_files: HashMap::new(),
        includes: includes,
    };

    let mut relative_paths = Vec::new();

    for input in input {
        relative_paths.push(run.add_fs_file(input)?);
    }

    let file_descriptors: Vec<_> = run
        .parsed_files
        .into_iter()
        .map(|(_, v)| v.descriptor)
        .collect();

    Ok(ParsedAndTypechecked {
        relative_paths,
        file_descriptors,
    })
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;
extern crate getopts;
extern crate unix_daemonize;
extern crate byteorder;
extern crate udt;
extern crate time;
extern crate sodiumoxide;
extern crate rustc_serialize;
extern crate colored;

pub mod connection;

use unix_daemonize::{daemonize_redirect, ChdirMode};
use std::process::Command;
use std::net::{SocketAddr, IpAddr};
use std::str;
use std::env;
use std::thread;
use std::str::FromStr;
use std::fs::{OpenOptions, File};
use std::path::{Path, PathBuf};
use std::time::{Instant, Duration};
use std::io::{Cursor, Error, Seek, SeekFrom, stderr, Read, Write};
use std::sync::mpsc;
use log::{LogRecord, LogLevel, LogMetadata};
use colored::*;
use byteorder::{ReadBytesExt, WriteBytesExt, LittleEndian};
use sodiumoxide::crypto::secretbox;
use sodiumoxide::crypto::secretbox::xsalsa20poly1305::Key;
use rustc_serialize::hex::{FromHex, ToHex};

// TODO config
const INITIAL_ACCEPT_TIMEOUT_SECONDS: u64 = 60;

macro_rules! overprint {
    ($fmt: expr) => {
        print!(concat!("\x1b[2K\r", $fmt));
        std::io::stdout().flush().unwrap();
    };
    ($fmt:expr, $($arg:tt)*) => {
        print!(concat!("\x1b[2K\r", $fmt) , $($arg)*);
        std::io::stdout().flush().unwrap();
    };
}

#[macro_export]
macro_rules! die {
    ($fmt: expr) => {
        error!($fmt);
        panic!($fmt)
    };
    ($fmt:expr, $($arg:tt)*) => {
        error!($fmt, $($arg)*);
        panic!($fmt, $($arg)*)
    };
}

pub struct Server<'a> {
    pub ip: String,
    filename: &'a str,
    conn: connection::Server,
}

#[allow(dead_code)]
enum ShoopErrKind {
    Severed,
    Fatal,
}

struct ShoopErr {
    kind: ShoopErrKind,
    msg: Option<String>,
    finished: u64,
}

pub struct ShoopLogger;

impl log::Log for ShoopLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= LogLevel::Info
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            let prefix_symbol = match record.level() {
                LogLevel::Error => "E".red().bold(),
                LogLevel::Warn => "W".yellow().bold(),
                LogLevel::Info => "I".normal(),
                LogLevel::Debug => "D".dimmed(),
                LogLevel::Trace => "T".dimmed(),
            };

            println!("[{}] {}", prefix_symbol, record.args());
        }
    }
}

impl ShoopLogger {
    pub fn init() -> Result<(), log::SetLoggerError> {
        log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Info);
            Box::new(ShoopLogger)
        })
    }
}

impl ShoopErr {
    pub fn new(kind: ShoopErrKind, msg: &str, finished: u64) -> ShoopErr {
        ShoopErr {
            kind: kind,
            msg: Some(String::from(msg)),
            finished: finished,
        }
    }

    #[allow(dead_code)]
    pub fn from(err: Error, finished: u64) -> ShoopErr {
        ShoopErr {
            kind: ShoopErrKind::Severed,
            msg: Some(format!("{:?}", err)),
            finished: finished,
        }
    }
}

impl<'a> Server<'a> {
    pub fn new(port_range: connection::PortRange, filename: &str) -> Server {
        let sshconnstr = match env::var("SSH_CONNECTION") {
            Ok(s) => s.trim().to_owned(),
            Err(_) => {
                die!("SSH_CONNECTION env variable unset and required. Quitting.");
            }
        };
        let sshconn: Vec<&str> = sshconnstr.split(' ').collect();
        let ip = sshconn[2].to_owned();
        let key = secretbox::gen_key();
        let Key(keybytes) = key;
        let port = connection::Server::get_open_port(&port_range).unwrap();
        println!("shoop 0 {} {} {}", ip, port, keybytes.to_hex());

        let stdout = Some(Path::new(&env::var("HOME").unwrap()).join(".shoop.log"));
        let stderr = Some(Path::new(&env::var("HOME").unwrap()).join(".shoop.log"));
        daemonize_redirect(stdout, stderr, ChdirMode::ChdirRoot).unwrap();

        let conn = connection::Server::new(IpAddr::from_str(&ip).unwrap(), port, key);
        Server {
            ip: ip,
            conn: conn,
            filename: filename,
        }
    }

    pub fn start(&self) {
        self.conn.listen().unwrap();

        let mut connection_count: usize = 0;
        info!("listening...");
        loop {
            info!("waiting for connection...");
            let (tx, rx) = mpsc::channel();
            if connection_count == 0 {
                thread::spawn(move || {
                    thread::sleep(Duration::from_secs(INITIAL_ACCEPT_TIMEOUT_SECONDS));
                    if let Err(_) = rx.try_recv() {
                        error!("timed out waiting for initial connection. exiting.");
                        std::process::exit(1);
                    }
                });
            }
            let client = match self.conn.accept() {
                Ok(client) => client,
                Err(e) => {
                    die!("error on sock accept() {:?}", e);
                }
            };
            connection_count += 1;
            tx.send(()).unwrap();
            info!("accepted connection!");
            if let Ok(starthdr) = client.recv() {
                let version = starthdr[0];
                let mut rdr = Cursor::new(starthdr);
                rdr.set_position(1);
                let offset = rdr.read_u64::<LittleEndian>().unwrap();
                if version == 0x00 {
                    match self.send_file(&client, offset) {
                        Ok(_) => {
                            info!("done sending file");
                            let _ = client.close();
                            break;
                        }
                        Err(ShoopErr { kind: ShoopErrKind::Severed, msg, finished }) => {
                            info!("connection severed, msg: {:?}, finished: {}", msg, finished);
                            let _ = client.close();
                            continue;
                        }
                        Err(ShoopErr { kind: ShoopErrKind::Fatal, msg, finished }) => {
                            die!("connection fatal, msg: {:?}, finished: {}", msg, finished);
                        }
                    }
                } else {
                    die!("unrecognized version");
                }
            } else {
                die!("failed to receive version byte from client");
            }
        }
        info!("exiting listen loop.");
    }

    fn send_file(&self,
                 client: &connection::ServerConnection,
                 offset: u64)
                 -> Result<(), ShoopErr> {
        let mut f = File::open(self.filename).unwrap();
        f.seek(SeekFrom::Start(offset)).unwrap();
        let metadata = f.metadata().unwrap();

        let remaining = metadata.len() - offset;
        info!("total {} bytes", remaining);

        let mut wtr = vec![];
        wtr.write_u64::<LittleEndian>(remaining).unwrap();
        match client.send(&wtr[..]) {
            Ok(()) => info!("wrote filesize header."),
            Err(e) => {
                return Err(ShoopErr::new(ShoopErrKind::Severed, &format!("{:?}", e), remaining))
            }
        }
        let mut payload = vec![0; 1300];
        f.seek(SeekFrom::Start(offset)).unwrap();
        info!("sending file...");
        loop {
            match f.read(&mut payload) {
                Ok(0) => {
                    break;
                }
                Ok(read) => {
                    if let Err(e) = client.send(&payload[0..read]) {
                        return Err(ShoopErr::new(ShoopErrKind::Severed,
                                                 &format!("{:?}", e),
                                                 remaining));
                    }
                }
                Err(e) => {
                    client.close().expect("Error closing stream");
                    die!("failed to read from file: {:?}", e);
                }
            }
        }

        client.close().expect("Error closing stream.");
        Ok(())
    }
}

pub fn download(remote_ssh_host: &str,
                port_range: connection::PortRange,
                remote_path: &str,
                local_path: PathBuf) {
    let cmd = format!("shoop -s '{}' -p {}", remote_path, port_range);
    // println!("addr: {}, path: {}, cmd: {}", addr, path, cmd);

    overprint!(" - establishing SSH session...");
    assert!(command_exists("ssh"), "`ssh` is required!");
    let output = Command::new("ssh")
        .arg(remote_ssh_host.to_owned())
        .arg(cmd)
        .output()
        .unwrap_or_else(|e| {
            die!("failed to execute process: {}", e);
        });
    let infostring = String::from_utf8_lossy(&output.stdout).to_owned().trim().to_owned();
    let info: Vec<&str> = infostring.split(' ').collect();
    if info.len() != 5 {
        die!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
    }

    let (magic, version, ip, port, keyhex) = (info[0], info[1], info[2], info[3], info[4]);
    overprint!(" - opening UDT connection...");
    if magic != "shoop" || version != "0" {
        die!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
    }

    let mut keybytes = [0u8; 32];
    keybytes.copy_from_slice(&keyhex.from_hex().unwrap()[..]);
    let key = Key(keybytes);
    let addr: SocketAddr = SocketAddr::from_str(&format!("{}:{}", ip, port)[..])
        .unwrap();
    let conn = connection::Client::new(addr, key);

    let mut offset = 0u64;
    let mut filesize = None;
    let start_ts = Instant::now();
    loop {
        match conn.connect() {
            Ok(()) => {
                overprint!(" - connection opened, shakin' hands, makin' frands");
            }
            Err(e) => {
                die!("errrrrrrr connecting to {}:{} - {:?}", ip, port, e);
            }
        }
        let mut wtr = vec![];
        wtr.push(0);
        wtr.write_u64::<LittleEndian>(offset).unwrap();
        if let Err(_) = conn.send(&wtr[..]) {
            conn.close().unwrap();
            continue;
        }

        if let Ok(msg) = conn.recv() {
            if msg.is_empty() {
                die!("failed to get filesize from server, probable timeout.");
            }
            let mut rdr = Cursor::new(msg);
            filesize = filesize.or_else(|| Some(rdr.read_u64::<LittleEndian>().unwrap()));
            overprint!("downloading {} ({:.1}MB)\n",
                       local_path.to_string_lossy(),
                       (filesize.unwrap() as f64) / (1024f64 * 1024f64));
            match recv_file(&conn, filesize.unwrap(), &local_path, offset) {
                Ok(_) => {
                    break;
                }
                Err(ShoopErr { kind: ShoopErrKind::Severed, finished, .. }) => {
                    println!("{}", " * [[SEVERED]]".yellow().bold());
                    offset = finished;
                }
                Err(ShoopErr { kind: ShoopErrKind::Fatal, msg, .. }) => {
                    die!("{:?}", msg);
                }
            }
        }
        let _ = conn.close();
    }

    let elapsed = start_ts.elapsed().as_secs();
    let fmt_time = if elapsed < 60 {
        format!("{}s", elapsed)
    } else if elapsed < 60 * 60 {
        format!("{}m{}s", elapsed / 60, elapsed % 60)
    } else {
        format!("{}h{}m{}s", elapsed / (60 * 60), elapsed / 60, elapsed % 60)
    };
    println!("shooped it all up in {}", fmt_time.green().bold());
}

fn command_exists(command: &str) -> bool {
    match Command::new("which").arg(command).output() {
        Ok(output) => output.status.success(),
        Err(_) => false,
    }
}

fn recv_file(conn: &connection::Client,
             filesize: u64,
             filename: &PathBuf,
             offset: u64)
             -> Result<(), ShoopErr> {
    let mut f = OpenOptions::new().write(true).create(true).truncate(false).open(filename).unwrap();
    f.seek(SeekFrom::Start(offset)).unwrap();
    let start = Instant::now();
    let mut ts = Instant::now();
    let mut total = offset;
    let mut speed_ts = Instant::now();
    let mut speed_total = total;
    let mut speed = 0u64;
    loop {
        let buf = try!(conn.recv()
            .map_err(|e| ShoopErr::new(ShoopErrKind::Severed, &format!("{:?}", e), total)));
        if buf.len() < 1 {
            return Err(ShoopErr::new(ShoopErrKind::Severed, "empty msg", total));
        }

        f.write_all(&buf[..]).unwrap();
        total += buf.len() as u64;
        let speed_elapsed = speed_ts.elapsed();
        if speed_elapsed > Duration::new(1, 0) {
            speed = ((total - speed_total) as f64 /
                     ((speed_elapsed.as_secs() as f64) +
                      (speed_elapsed.subsec_nanos() as f64) /
                      1_000_000_000f64)) as u64;
            speed_ts = Instant::now();
            speed_total = total;
        }
        let speedfmt = if speed < 1024 {
            format!("{} b/s", speed)
        } else if speed < 1024 * 1024 {
            format!("{} kb/s", speed / 1024)
        } else {
            format!("{:.1} MB/s", ((speed / 1024) as f64) / 1024f64)
        };

        if ts.elapsed() > Duration::new(0, 100_000_000) {
            overprint!("   {:.1}M / {:.1}M ({:.1}%) [ {} ]",
                       (total as f64) / (1024f64 * 1024f64),
                       (filesize as f64) / (1024f64 * 1024f64),
                       (total as f64) / (filesize as f64) * 100f64,
                       speedfmt);
            ts = Instant::now();
        }
        if total >= filesize {
            overprint!("   {0:.1}M / {0:.1}M (100%) [ avg {1:.1} MB/s ]\n",
                       (filesize as f64) / (1024f64 * 1024f64),
                       ((total - offset) / start.elapsed().as_secs() / 1024) as f64 / 1024f64);
            break;
        }
    }
    let _ = conn.close();
    Ok(())
}
<commit_msg>(protocol break) remove version byte from header<commit_after>#[macro_use]
extern crate log;
extern crate getopts;
extern crate unix_daemonize;
extern crate byteorder;
extern crate udt;
extern crate time;
extern crate sodiumoxide;
extern crate rustc_serialize;
extern crate colored;

pub mod connection;

use unix_daemonize::{daemonize_redirect, ChdirMode};
use std::process::Command;
use std::net::{SocketAddr, IpAddr};
use std::str;
use std::env;
use std::thread;
use std::str::FromStr;
use std::fs::{OpenOptions, File};
use std::path::{Path, PathBuf};
use std::time::{Instant, Duration};
use std::io::{Cursor, Error, Seek, SeekFrom, stderr, Read, Write};
use std::sync::mpsc;
use log::{LogRecord, LogLevel, LogMetadata};
use colored::*;
use byteorder::{ReadBytesExt, WriteBytesExt, LittleEndian};
use sodiumoxide::crypto::secretbox;
use sodiumoxide::crypto::secretbox::xsalsa20poly1305::Key;
use rustc_serialize::hex::{FromHex, ToHex};

// TODO config
const INITIAL_ACCEPT_TIMEOUT_SECONDS: u64 = 60;

macro_rules! overprint {
    ($fmt: expr) => {
        print!(concat!("\x1b[2K\r", $fmt));
        std::io::stdout().flush().unwrap();
    };
    ($fmt:expr, $($arg:tt)*) => {
        print!(concat!("\x1b[2K\r", $fmt) , $($arg)*);
        std::io::stdout().flush().unwrap();
    };
}

#[macro_export]
macro_rules! die {
    ($fmt: expr) => {
        error!($fmt);
        panic!($fmt)
    };
    ($fmt:expr, $($arg:tt)*) => {
        error!($fmt, $($arg)*);
        panic!($fmt, $($arg)*)
    };
}

pub struct Server<'a> {
    pub ip: String,
    filename: &'a str,
    conn: connection::Server,
}

#[allow(dead_code)]
enum ShoopErrKind {
    Severed,
    Fatal,
}

struct ShoopErr {
    kind: ShoopErrKind,
    msg: Option<String>,
    finished: u64,
}

pub struct ShoopLogger;

impl log::Log for ShoopLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= LogLevel::Info
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            let prefix_symbol = match record.level() {
                LogLevel::Error => "E".red().bold(),
                LogLevel::Warn => "W".yellow().bold(),
                LogLevel::Info => "I".normal(),
                LogLevel::Debug => "D".dimmed(),
                LogLevel::Trace => "T".dimmed(),
            };

            println!("[{}] {}", prefix_symbol, record.args());
        }
    }
}

impl ShoopLogger {
    pub fn init() -> Result<(), log::SetLoggerError> {
        log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Info);
            Box::new(ShoopLogger)
        })
    }
}

impl ShoopErr {
    pub fn new(kind: ShoopErrKind, msg: &str, finished: u64) -> ShoopErr {
        ShoopErr {
            kind: kind,
            msg: Some(String::from(msg)),
            finished: finished,
        }
    }

    #[allow(dead_code)]
    pub fn from(err: Error, finished: u64) -> ShoopErr {
        ShoopErr {
            kind: ShoopErrKind::Severed,
            msg: Some(format!("{:?}", err)),
            finished: finished,
        }
    }
}

impl<'a> Server<'a> {
    pub fn new(port_range: connection::PortRange, filename: &str) -> Server {
        let sshconnstr = match env::var("SSH_CONNECTION") {
            Ok(s) => s.trim().to_owned(),
            Err(_) => {
                die!("SSH_CONNECTION env variable unset and required. Quitting.");
            }
        };
        let sshconn: Vec<&str> = sshconnstr.split(' ').collect();
        let ip = sshconn[2].to_owned();
        let key = secretbox::gen_key();
        let Key(keybytes) = key;
        let port = connection::Server::get_open_port(&port_range).unwrap();
        println!("shoop 0 {} {} {}", ip, port, keybytes.to_hex());

        let stdout = Some(Path::new(&env::var("HOME").unwrap()).join(".shoop.log"));
        let stderr = Some(Path::new(&env::var("HOME").unwrap()).join(".shoop.log"));
        daemonize_redirect(stdout, stderr, ChdirMode::ChdirRoot).unwrap();

        let conn = connection::Server::new(IpAddr::from_str(&ip).unwrap(), port, key);
        Server {
            ip: ip,
            conn: conn,
            filename: filename,
        }
    }

    pub fn start(&self) {
        self.conn.listen().unwrap();

        let mut connection_count: usize = 0;
        info!("listening...");
        loop {
            info!("waiting for connection...");
            let (tx, rx) = mpsc::channel();
            if connection_count == 0 {
                thread::spawn(move || {
                    thread::sleep(Duration::from_secs(INITIAL_ACCEPT_TIMEOUT_SECONDS));
                    if let Err(_) = rx.try_recv() {
                        error!("timed out waiting for initial connection. exiting.");
                        std::process::exit(1);
                    }
                });
            }
            let client = match self.conn.accept() {
                Ok(client) => client,
                Err(e) => {
                    die!("error on sock accept() {:?}", e);
                }
            };
            connection_count += 1;
            tx.send(()).unwrap();
            info!("accepted connection!");
            if let Ok(starthdr) = client.recv() {
                let mut rdr = Cursor::new(starthdr);
                let offset = rdr.read_u64::<LittleEndian>().unwrap();
                match self.send_file(&client, offset) {
                    Ok(_) => {
                        info!("done sending file");
                        let _ = client.close();
                        break;
                    }
                    Err(ShoopErr { kind: ShoopErrKind::Severed, msg, finished }) => {
                        info!("connection severed, msg: {:?}, finished: {}", msg, finished);
                        let _ = client.close();
                        continue;
                    }
                    Err(ShoopErr { kind: ShoopErrKind::Fatal, msg, finished }) => {
                        die!("connection fatal, msg: {:?}, finished: {}", msg, finished);
                    }
                }
            } else {
                die!("failed to receive version byte from client");
            }
        }
        info!("exiting listen loop.");
    }

    fn send_file(&self,
                 client: &connection::ServerConnection,
                 offset: u64)
                 -> Result<(), ShoopErr> {
        let mut f = File::open(self.filename).unwrap();
        f.seek(SeekFrom::Start(offset)).unwrap();
        let metadata = f.metadata().unwrap();

        let remaining = metadata.len() - offset;
        info!("total {} bytes", remaining);

        let mut wtr = vec![];
        wtr.write_u64::<LittleEndian>(remaining).unwrap();
        match client.send(&wtr[..]) {
            Ok(()) => info!("wrote filesize header."),
            Err(e) => {
                return Err(ShoopErr::new(ShoopErrKind::Severed, &format!("{:?}", e), remaining))
            }
        }
        let mut payload = vec![0; 1300];
        f.seek(SeekFrom::Start(offset)).unwrap();
        info!("sending file...");
        loop {
            match f.read(&mut payload) {
                Ok(0) => {
                    break;
                }
                Ok(read) => {
                    if let Err(e) = client.send(&payload[0..read]) {
                        return Err(ShoopErr::new(ShoopErrKind::Severed,
                                                 &format!("{:?}", e),
                                                 remaining));
                    }
                }
                Err(e) => {
                    client.close().expect("Error closing stream");
                    die!("failed to read from file: {:?}", e);
                }
            }
        }

        client.close().expect("Error closing stream.");
        Ok(())
    }
}

pub fn download(remote_ssh_host: &str,
                port_range: connection::PortRange,
                remote_path: &str,
                local_path: PathBuf) {
    let cmd = format!("shoop -s '{}' -p {}", remote_path, port_range);
    // println!("addr: {}, path: {}, cmd: {}", addr, path, cmd);

    overprint!(" - establishing SSH session...");
    assert!(command_exists("ssh"), "`ssh` is required!");
    let output = Command::new("ssh")
        .arg(remote_ssh_host.to_owned())
        .arg(cmd)
        .output()
        .unwrap_or_else(|e| {
            die!("failed to execute process: {}", e);
        });
    let infostring = String::from_utf8_lossy(&output.stdout).to_owned().trim().to_owned();
    let info: Vec<&str> = infostring.split(' ').collect();
    if info.len() != 5 {
        die!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
    }

    let (magic, version, ip, port, keyhex) = (info[0], info[1], info[2], info[3], info[4]);
    overprint!(" - opening UDT connection...");
    if magic != "shoop" || version != "0" {
        die!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
    }

    let mut keybytes = [0u8; 32];
    keybytes.copy_from_slice(&keyhex.from_hex().unwrap()[..]);
    let key = Key(keybytes);
    let addr: SocketAddr = SocketAddr::from_str(&format!("{}:{}", ip, port)[..])
        .unwrap();
    let conn = connection::Client::new(addr, key);

    let mut offset = 0u64;
    let mut filesize = None;
    let start_ts = Instant::now();
    loop {
        match conn.connect() {
            Ok(()) => {
                overprint!(" - connection opened, shakin' hands, makin' frands");
            }
            Err(e) => {
                die!("errrrrrrr connecting to {}:{} - {:?}", ip, port, e);
            }
        }
        let mut wtr = vec![];
        wtr.write_u64::<LittleEndian>(offset).unwrap();
        if let Err(_) = conn.send(&wtr[..]) {
            conn.close().unwrap();
            continue;
        }

        if let Ok(msg) = conn.recv() {
            if msg.is_empty() {
                die!("failed to get filesize from server, probable timeout.");
            }
            let mut rdr = Cursor::new(msg);
            filesize = filesize.or_else(|| Some(rdr.read_u64::<LittleEndian>().unwrap()));
            overprint!("downloading {} ({:.1}MB)\n",
                       local_path.to_string_lossy(),
                       (filesize.unwrap() as f64) / (1024f64 * 1024f64));
            match recv_file(&conn, filesize.unwrap(), &local_path, offset) {
                Ok(_) => {
                    break;
                }
                Err(ShoopErr { kind: ShoopErrKind::Severed, finished, .. }) => {
                    println!("{}", " * [[SEVERED]]".yellow().bold());
                    offset = finished;
                }
                Err(ShoopErr { kind: ShoopErrKind::Fatal, msg, .. }) => {
                    die!("{:?}", msg);
                }
            }
        }
        let _ = conn.close();
    }

    let elapsed = start_ts.elapsed().as_secs();
    let fmt_time = if elapsed < 60 {
        format!("{}s", elapsed)
    } else if elapsed < 60 * 60 {
        format!("{}m{}s", elapsed / 60, elapsed % 60)
    } else {
        format!("{}h{}m{}s", elapsed / (60 * 60), elapsed / 60, elapsed % 60)
    };
    println!("shooped it all up in {}", fmt_time.green().bold());
}

fn command_exists(command: &str) -> bool {
    match Command::new("which").arg(command).output() {
        Ok(output) => output.status.success(),
        Err(_) => false,
    }
}

fn recv_file(conn: &connection::Client,
             filesize: u64,
             filename: &PathBuf,
             offset: u64)
             -> Result<(), ShoopErr> {
    let mut f = OpenOptions::new().write(true).create(true).truncate(false).open(filename).unwrap();
    f.seek(SeekFrom::Start(offset)).unwrap();
    let start = Instant::now();
    let mut ts = Instant::now();
    let mut total = offset;
    let mut speed_ts = Instant::now();
    let mut speed_total = total;
    let mut speed = 0u64;
    loop {
        let buf = try!(conn.recv()
            .map_err(|e| ShoopErr::new(ShoopErrKind::Severed, &format!("{:?}", e), total)));
        if buf.len() < 1 {
            return Err(ShoopErr::new(ShoopErrKind::Severed, "empty msg", total));
        }

        f.write_all(&buf[..]).unwrap();
        total += buf.len() as u64;
        let speed_elapsed = speed_ts.elapsed();
        if speed_elapsed > Duration::new(1, 0) {
            speed = ((total - speed_total) as f64 /
                     ((speed_elapsed.as_secs() as f64) +
                      (speed_elapsed.subsec_nanos() as f64) /
                      1_000_000_000f64)) as u64;
            speed_ts = Instant::now();
            speed_total = total;
        }
        let speedfmt = if speed < 1024 {
            format!("{} b/s", speed)
        } else if speed < 1024 * 1024 {
            format!("{} kb/s", speed / 1024)
        } else {
            format!("{:.1} MB/s", ((speed / 1024) as f64) / 1024f64)
        };

        if ts.elapsed() > Duration::new(0, 100_000_000) {
            overprint!("   {:.1}M / {:.1}M ({:.1}%) [ {} ]",
                       (total as f64) / (1024f64 * 1024f64),
                       (filesize as f64) / (1024f64 * 1024f64),
                       (total as f64) / (filesize as f64) * 100f64,
                       speedfmt);
            ts = Instant::now();
        }
        if total >= filesize {
            overprint!("   {0:.1}M / {0:.1}M (100%) [ avg {1:.1} MB/s ]\n",
                       (filesize as f64) / (1024f64 * 1024f64),
                       ((total - offset) / start.elapsed().as_secs() / 1024) as f64 / 1024f64);
            break;
        }
    }
    let _ = conn.close();
    Ok(())
}
<|endoftext|>