<commit_before>use crate::{
    config, db,
    models::Version,
    schema::{crates, readme_renderings, versions},
    uploaders::Uploader,
};
use std::{io::Read, path::Path, sync::Arc, thread};

use cargo_registry_markdown::text_to_html;
use chrono::{TimeZone, Utc};
use clap::Clap;
use diesel::{dsl::any, prelude::*};
use flate2::read::GzDecoder;
use reqwest::{blocking::Client, header};
use tar::{self, Archive};

const CACHE_CONTROL_README: &str = "public,max-age=604800";
const USER_AGENT: &str = "crates-admin";

#[derive(Clap, Debug)]
#[clap(
    name = "render-readmes",
    about = "Iterates over every crate versions ever uploaded and (re-)renders their \
        readme using the readme renderer from the cargo_registry crate.",
    after_help = "Warning: this can take a lot of time."
)]
pub struct Opts {
    /// How many versions should be queried and processed at a time.
    #[clap(long, default_value = "25")]
    page_size: usize,

    /// Only rerender readmes that are older than this date.
    #[clap(long)]
    older_than: Option<String>,

    /// Only rerender readmes for the specified crate.
    #[clap(long = "crate")]
    crate_name: Option<String>,
}

pub fn run(opts: Opts) {
    let base_config = Arc::new(config::Base::from_environment());
    let conn = db::connect_now().unwrap();

    let start_time = Utc::now();

    let older_than = if let Some(ref time) = opts.older_than {
        Utc.datetime_from_str(time, "%Y-%m-%d %H:%M:%S")
            .expect("Could not parse --older-than argument as a time")
    } else {
        start_time
    };
    let older_than = older_than.naive_utc();

    println!("Start time:                   {}", start_time);
    println!("Rendering readmes older than: {}", older_than);

    let mut query = versions::table
        .inner_join(crates::table)
        .left_outer_join(readme_renderings::table)
        .filter(
            readme_renderings::rendered_at
                .lt(older_than)
                .or(readme_renderings::version_id.is_null()),
        )
        .select(versions::id)
        .into_boxed();

    if let Some(crate_name) = opts.crate_name {
        println!("Rendering readmes for {}", crate_name);
        query = query.filter(crates::name.eq(crate_name));
    }

    let version_ids: Vec<i32> = query.load(&conn).expect("error loading version ids");

    let total_versions = version_ids.len();
    println!("Rendering {} versions", total_versions);

    let page_size = opts.page_size;

    let total_pages = total_versions / page_size;
    let total_pages = if total_versions % page_size == 0 {
        total_pages
    } else {
        total_pages + 1
    };

    let client = Client::new();

    for (page_num, version_ids_chunk) in version_ids.chunks(page_size).enumerate() {
        println!(
            "= Page {} of {} ==================================",
            page_num + 1,
            total_pages
        );

        let versions: Vec<(Version, String)> = versions::table
            .inner_join(crates::table)
            .filter(versions::id.eq(any(version_ids_chunk)))
            .select((versions::all_columns, crates::name))
            .load(&conn)
            .expect("error loading versions");

        let mut tasks = Vec::with_capacity(page_size as usize);
        for (version, krate_name) in versions {
            Version::record_readme_rendering(version.id, &conn).unwrap_or_else(|_| {
                panic!(
                    "[{}-{}] Couldn't record rendering time",
                    krate_name, version.num
                )
            });
            let client = client.clone();
            let base_config = base_config.clone();
            let handle = thread::spawn(move || {
                println!("[{}-{}] Rendering README...", krate_name, version.num);
                let readme = get_readme(base_config.uploader(), &client, &version, &krate_name);
                if readme.is_none() {
                    return;
                }
                let readme = readme.unwrap();
                let content_length = readme.len() as u64;
                let content = std::io::Cursor::new(readme);
                let readme_path = format!("readmes/{0}/{0}-{1}.html", krate_name, version.num);
                let mut extra_headers = header::HeaderMap::new();
                extra_headers.insert(
                    header::CACHE_CONTROL,
                    header::HeaderValue::from_static(CACHE_CONTROL_README),
                );
                base_config
                    .uploader()
                    .upload(
                        &client,
                        &readme_path,
                        content,
                        content_length,
                        "text/html",
                        extra_headers,
                    )
                    .unwrap_or_else(|_| {
                        panic!(
                            "[{}-{}] Couldn't upload file to S3",
                            krate_name, version.num
                        )
                    });
            });
            tasks.push(handle);
        }
        for handle in tasks {
            if let Err(err) = handle.join() {
                println!("Thread panicked: {:?}", err);
            }
        }
    }
}

/// Renders the readme of an uploaded crate version.
fn get_readme(
    uploader: &Uploader,
    client: &Client,
    version: &Version,
    krate_name: &str,
) -> Option<String> {
    let pkg_name = format!("{}-{}", krate_name, version.num);

    let location = uploader.crate_location(krate_name, &version.num.to_string());

    let location = match uploader {
        Uploader::S3 { .. } => location,
        Uploader::Local => format!("http://localhost:8888/{}", location),
    };

    let mut extra_headers = header::HeaderMap::new();
    extra_headers.insert(
        header::USER_AGENT,
        header::HeaderValue::from_static(USER_AGENT),
    );
    let response = match client.get(&location).headers(extra_headers).send() {
        Ok(r) => r,
        Err(err) => {
            println!("[{}] Unable to fetch crate: {}", pkg_name, err);
            return None;
        }
    };

    if !response.status().is_success() {
        println!(
            "[{}] Failed to get a 200 response: {}",
            pkg_name,
            response.text().unwrap()
        );
        return None;
    }

    let reader = GzDecoder::new(response);
    let archive = Archive::new(reader);
    render_pkg_readme(archive, &pkg_name)
}

fn render_pkg_readme<R: Read>(mut archive: Archive<R>, pkg_name: &str) -> Option<String> {
    let mut entries = archive
        .entries()
        .unwrap_or_else(|_| panic!("[{}] Invalid tar archive entries", pkg_name));

    let manifest: Manifest = {
        let path = format!("{}/Cargo.toml", pkg_name);
        let contents = find_file_by_path(&mut entries, Path::new(&path), pkg_name);
        toml::from_str(&contents)
            .unwrap_or_else(|_| panic!("[{}] Syntax error in manifest file", pkg_name))
    };

    let rendered = {
        let readme_path = manifest.package.readme.as_ref()?;
        let path = format!("{}/{}", pkg_name, readme_path);
        let contents = find_file_by_path(&mut entries, Path::new(&path), pkg_name);
        text_to_html(
            &contents,
            readme_path,
            manifest.package.repository.as_deref(),
        )
    };
    return Some(rendered);

    #[derive(Debug, Deserialize)]
    struct Package {
        readme: Option<String>,
        repository: Option<String>,
    }

    #[derive(Debug, Deserialize)]
    struct Manifest {
        package: Package,
    }
}

/// Search an entry by its path in a Tar archive.
fn find_file_by_path<R: Read>(
    entries: &mut tar::Entries<'_, R>,
    path: &Path,
    pkg_name: &str,
) -> String {
    let mut file = entries
        .find(|entry| match *entry {
            Err(_) => false,
            Ok(ref file) => {
                let filepath = match file.path() {
                    Ok(p) => p,
                    Err(_) => return false,
                };
                filepath == path
            }
        })
        .unwrap_or_else(|| panic!("[{}] couldn't open file: {}", pkg_name, path.display()))
        .unwrap_or_else(|_| panic!("[{}] file is not present: {}", pkg_name, path.display()));

    let mut contents = String::new();
    file.read_to_string(&mut contents)
        .unwrap_or_else(|_| panic!("[{}] Couldn't read file contents", pkg_name));
    contents
}

#[cfg(test)]
mod tests {
    use std::io::Write;
    use tar;

    use super::render_pkg_readme;

    fn add_file<W: Write>(pkg: &mut tar::Builder<W>, path: &str, content: &[u8]) {
        let mut header = tar::Header::new_gnu();
        header.set_size(content.len() as u64);
        header.set_cksum();
        pkg.append_data(&mut header, path, content).unwrap();
    }

    #[test]
    fn test_render_pkg_readme() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "README.md"
"#,
        );
        add_file(&mut pkg, "foo-0.0.1/README.md", b"readme");
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("readme"))
    }

    #[test]
    fn test_render_pkg_readme_w_link() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "README.md"
repository = "https://github.com/foo/foo"
"#,
        );
        add_file(
            &mut pkg,
            "foo-0.0.1/README.md",
            b"readme [link](./Other.md)",
        );
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("\"https://github.com/foo/foo/blob/HEAD/./Other.md\""))
    }

    #[test]
    fn test_render_pkg_readme_not_at_root() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "docs/README.md"
repository = "https://github.com/foo/foo"
"#,
        );
        add_file(
            &mut pkg,
            "foo-0.0.1/docs/README.md",
            b"docs/readme [link](./Other.md)",
        );
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("docs/readme"));
        assert!(result.contains("\"https://github.com/foo/foo/blob/HEAD/docs/./Other.md\""))
    }
}
<commit_msg>admin::render_readmes: Extract file-not-found case of find_file_by_path<commit_after>use crate::{
    config, db,
    models::Version,
    schema::{crates, readme_renderings, versions},
    uploaders::Uploader,
};
use std::{io::Read, path::Path, sync::Arc, thread};

use cargo_registry_markdown::text_to_html;
use chrono::{TimeZone, Utc};
use clap::Clap;
use diesel::{dsl::any, prelude::*};
use flate2::read::GzDecoder;
use reqwest::{blocking::Client, header};
use tar::{self, Archive};

const CACHE_CONTROL_README: &str = "public,max-age=604800";
const USER_AGENT: &str = "crates-admin";

#[derive(Clap, Debug)]
#[clap(
    name = "render-readmes",
    about = "Iterates over every crate versions ever uploaded and (re-)renders their \
        readme using the readme renderer from the cargo_registry crate.",
    after_help = "Warning: this can take a lot of time."
)]
pub struct Opts {
    /// How many versions should be queried and processed at a time.
    #[clap(long, default_value = "25")]
    page_size: usize,

    /// Only rerender readmes that are older than this date.
    #[clap(long)]
    older_than: Option<String>,

    /// Only rerender readmes for the specified crate.
    #[clap(long = "crate")]
    crate_name: Option<String>,
}

pub fn run(opts: Opts) {
    let base_config = Arc::new(config::Base::from_environment());
    let conn = db::connect_now().unwrap();

    let start_time = Utc::now();

    let older_than = if let Some(ref time) = opts.older_than {
        Utc.datetime_from_str(time, "%Y-%m-%d %H:%M:%S")
            .expect("Could not parse --older-than argument as a time")
    } else {
        start_time
    };
    let older_than = older_than.naive_utc();

    println!("Start time:                   {}", start_time);
    println!("Rendering readmes older than: {}", older_than);

    let mut query = versions::table
        .inner_join(crates::table)
        .left_outer_join(readme_renderings::table)
        .filter(
            readme_renderings::rendered_at
                .lt(older_than)
                .or(readme_renderings::version_id.is_null()),
        )
        .select(versions::id)
        .into_boxed();

    if let Some(crate_name) = opts.crate_name {
        println!("Rendering readmes for {}", crate_name);
        query = query.filter(crates::name.eq(crate_name));
    }

    let version_ids: Vec<i32> = query.load(&conn).expect("error loading version ids");

    let total_versions = version_ids.len();
    println!("Rendering {} versions", total_versions);

    let page_size = opts.page_size;

    let total_pages = total_versions / page_size;
    let total_pages = if total_versions % page_size == 0 {
        total_pages
    } else {
        total_pages + 1
    };

    let client = Client::new();

    for (page_num, version_ids_chunk) in version_ids.chunks(page_size).enumerate() {
        println!(
            "= Page {} of {} ==================================",
            page_num + 1,
            total_pages
        );

        let versions: Vec<(Version, String)> = versions::table
            .inner_join(crates::table)
            .filter(versions::id.eq(any(version_ids_chunk)))
            .select((versions::all_columns, crates::name))
            .load(&conn)
            .expect("error loading versions");

        let mut tasks = Vec::with_capacity(page_size as usize);
        for (version, krate_name) in versions {
            Version::record_readme_rendering(version.id, &conn).unwrap_or_else(|_| {
                panic!(
                    "[{}-{}] Couldn't record rendering time",
                    krate_name, version.num
                )
            });
            let client = client.clone();
            let base_config = base_config.clone();
            let handle = thread::spawn(move || {
                println!("[{}-{}] Rendering README...", krate_name, version.num);
                let readme = get_readme(base_config.uploader(), &client, &version, &krate_name);
                if readme.is_none() {
                    return;
                }
                let readme = readme.unwrap();
                let content_length = readme.len() as u64;
                let content = std::io::Cursor::new(readme);
                let readme_path = format!("readmes/{0}/{0}-{1}.html", krate_name, version.num);
                let mut extra_headers = header::HeaderMap::new();
                extra_headers.insert(
                    header::CACHE_CONTROL,
                    header::HeaderValue::from_static(CACHE_CONTROL_README),
                );
                base_config
                    .uploader()
                    .upload(
                        &client,
                        &readme_path,
                        content,
                        content_length,
                        "text/html",
                        extra_headers,
                    )
                    .unwrap_or_else(|_| {
                        panic!(
                            "[{}-{}] Couldn't upload file to S3",
                            krate_name, version.num
                        )
                    });
            });
            tasks.push(handle);
        }
        for handle in tasks {
            if let Err(err) = handle.join() {
                println!("Thread panicked: {:?}", err);
            }
        }
    }
}

/// Renders the readme of an uploaded crate version.
fn get_readme(
    uploader: &Uploader,
    client: &Client,
    version: &Version,
    krate_name: &str,
) -> Option<String> {
    let pkg_name = format!("{}-{}", krate_name, version.num);

    let location = uploader.crate_location(krate_name, &version.num.to_string());

    let location = match uploader {
        Uploader::S3 { .. } => location,
        Uploader::Local => format!("http://localhost:8888/{}", location),
    };

    let mut extra_headers = header::HeaderMap::new();
    extra_headers.insert(
        header::USER_AGENT,
        header::HeaderValue::from_static(USER_AGENT),
    );
    let response = match client.get(&location).headers(extra_headers).send() {
        Ok(r) => r,
        Err(err) => {
            println!("[{}] Unable to fetch crate: {}", pkg_name, err);
            return None;
        }
    };

    if !response.status().is_success() {
        println!(
            "[{}] Failed to get a 200 response: {}",
            pkg_name,
            response.text().unwrap()
        );
        return None;
    }

    let reader = GzDecoder::new(response);
    let archive = Archive::new(reader);
    render_pkg_readme(archive, &pkg_name)
}

fn render_pkg_readme<R: Read>(mut archive: Archive<R>, pkg_name: &str) -> Option<String> {
    let mut entries = archive
        .entries()
        .unwrap_or_else(|_| panic!("[{}] Invalid tar archive entries", pkg_name));

    let manifest: Manifest = {
        let path = format!("{}/Cargo.toml", pkg_name);
        let contents = find_file_by_path(&mut entries, Path::new(&path), pkg_name)
            .unwrap_or_else(|| panic!("[{}] couldn't open file: Cargo.toml", pkg_name));
        toml::from_str(&contents)
            .unwrap_or_else(|_| panic!("[{}] Syntax error in manifest file", pkg_name))
    };

    let rendered = {
        let readme_path = manifest.package.readme.as_ref()?;
        let path = format!("{}/{}", pkg_name, readme_path);
        let contents = find_file_by_path(&mut entries, Path::new(&path), pkg_name)
            .unwrap_or_else(|| panic!("[{}] couldn't open file: {}", pkg_name, readme_path));
        text_to_html(
            &contents,
            readme_path,
            manifest.package.repository.as_deref(),
        )
    };
    return Some(rendered);

    #[derive(Debug, Deserialize)]
    struct Package {
        readme: Option<String>,
        repository: Option<String>,
    }

    #[derive(Debug, Deserialize)]
    struct Manifest {
        package: Package,
    }
}

/// Search an entry by its path in a Tar archive.
fn find_file_by_path<R: Read>(
    entries: &mut tar::Entries<'_, R>,
    path: &Path,
    pkg_name: &str,
) -> Option<String> {
    let mut file = entries
        .find(|entry| match *entry {
            Err(_) => false,
            Ok(ref file) => {
                let filepath = match file.path() {
                    Ok(p) => p,
                    Err(_) => return false,
                };
                filepath == path
            }
        })?
        .unwrap_or_else(|_| panic!("[{}] file is not present: {}", pkg_name, path.display()));

    let mut contents = String::new();
    file.read_to_string(&mut contents)
        .unwrap_or_else(|_| panic!("[{}] Couldn't read file contents", pkg_name));
    Some(contents)
}

#[cfg(test)]
mod tests {
    use std::io::Write;
    use tar;

    use super::render_pkg_readme;

    fn add_file<W: Write>(pkg: &mut tar::Builder<W>, path: &str, content: &[u8]) {
        let mut header = tar::Header::new_gnu();
        header.set_size(content.len() as u64);
        header.set_cksum();
        pkg.append_data(&mut header, path, content).unwrap();
    }

    #[test]
    fn test_render_pkg_readme() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "README.md"
"#,
        );
        add_file(&mut pkg, "foo-0.0.1/README.md", b"readme");
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("readme"))
    }

    #[test]
    fn test_render_pkg_readme_w_link() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "README.md"
repository = "https://github.com/foo/foo"
"#,
        );
        add_file(
            &mut pkg,
            "foo-0.0.1/README.md",
            b"readme [link](./Other.md)",
        );
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("\"https://github.com/foo/foo/blob/HEAD/./Other.md\""))
    }

    #[test]
    fn test_render_pkg_readme_not_at_root() {
        let mut pkg = tar::Builder::new(vec![]);
        add_file(
            &mut pkg,
            "foo-0.0.1/Cargo.toml",
            br#"
[package]
readme = "docs/README.md"
repository = "https://github.com/foo/foo"
"#,
        );
        add_file(
            &mut pkg,
            "foo-0.0.1/docs/README.md",
            b"docs/readme [link](./Other.md)",
        );
        let serialized_archive = pkg.into_inner().unwrap();
        let result =
            render_pkg_readme(tar::Archive::new(&*serialized_archive), "foo-0.0.1").unwrap();
        assert!(result.contains("docs/readme"));
        assert!(result.contains("\"https://github.com/foo/foo/blob/HEAD/docs/./Other.md\""))
    }
}
<|endoftext|>


<commit_before>//! WebAssembly function translation state.
//!
//! The `TranslationState` struct defined in this module is used to keep track of the WebAssembly
//! value and control stacks during the translation of a single function.

use cretonne::ir::{self, Ebb, Inst, Value};
use environ::{FuncEnvironment, GlobalValue};
use std::collections::HashMap;
use translation_utils::{GlobalIndex, MemoryIndex, SignatureIndex, FunctionIndex};

/// A control stack frame can be an `if`, a `block` or a `loop`, each one having the following
/// fields:
///
/// - `destination`: reference to the `Ebb` that will hold the code after the control block;
/// - `num_return_values`: number of values returned by the control block;
/// - `original_stack_size`: size of the value stack at the beginning of the control block.
///
/// Moreover, the `if` frame has the `branch_inst` field that points to the `brz` instruction
/// separating the `true` and `false` branch. The `loop` frame has a `header` field that references
/// the `Ebb` that contains the beginning of the body of the loop.
#[derive(Debug)]
pub enum ControlStackFrame {
    If {
        destination: Ebb,
        branch_inst: Inst,
        num_return_values: usize,
        original_stack_size: usize,
        reachable: bool,
    },
    Block {
        destination: Ebb,
        num_return_values: usize,
        original_stack_size: usize,
        reachable: bool,
    },
    Loop {
        destination: Ebb,
        header: Ebb,
        num_return_values: usize,
        original_stack_size: usize,
        reachable: bool,
    },
}

/// Helper methods for the control stack objects.
impl ControlStackFrame {
    pub fn num_return_values(&self) -> usize {
        match *self {
            ControlStackFrame::If { num_return_values, .. } |
            ControlStackFrame::Block { num_return_values, .. } |
            ControlStackFrame::Loop { num_return_values, .. } => num_return_values,
        }
    }
    pub fn following_code(&self) -> Ebb {
        match *self {
            ControlStackFrame::If { destination, .. } |
            ControlStackFrame::Block { destination, .. } |
            ControlStackFrame::Loop { destination, .. } => destination,
        }
    }
    pub fn br_destination(&self) -> Ebb {
        match *self {
            ControlStackFrame::If { destination, .. } |
            ControlStackFrame::Block { destination, .. } => destination,
            ControlStackFrame::Loop { header, .. } => header,
        }
    }
    pub fn original_stack_size(&self) -> usize {
        match *self {
            ControlStackFrame::If { original_stack_size, .. } |
            ControlStackFrame::Block { original_stack_size, .. } |
            ControlStackFrame::Loop { original_stack_size, .. } => original_stack_size,
        }
    }
    pub fn is_loop(&self) -> bool {
        match *self {
            ControlStackFrame::If { .. } |
            ControlStackFrame::Block { .. } => false,
            ControlStackFrame::Loop { .. } => true,
        }
    }

    pub fn is_reachable(&self) -> bool {
        match *self {
            ControlStackFrame::If { reachable, .. } |
            ControlStackFrame::Block { reachable, .. } |
            ControlStackFrame::Loop { reachable, .. } => reachable,
        }
    }

    pub fn set_reachable(&mut self) {
        match *self {
            ControlStackFrame::If { ref mut reachable, .. } |
            ControlStackFrame::Block { ref mut reachable, .. } |
            ControlStackFrame::Loop { ref mut reachable, .. } => *reachable = true,
        }
    }
}

/// Contains information passed along during the translation and that records:
///
/// - The current value and control stacks.
/// - The depth of the two unreachable control blocks stacks, that are manipulated when translating
///   unreachable code;
pub struct TranslationState {
    pub stack: Vec<Value>,
    pub control_stack: Vec<ControlStackFrame>,
    pub phantom_unreachable_stack_depth: usize,
    pub real_unreachable_stack_depth: usize,

    // Map of global variables that have already been created by `FuncEnvironment::make_global`.
    globals: HashMap<GlobalIndex, GlobalValue>,

    // Map of heaps that have been created by `FuncEnvironment::make_heap`.
    heaps: HashMap<MemoryIndex, ir::Heap>,

    // Map of indirect call signatures that have been created by
    // `FuncEnvironment::make_indirect_sig()`.
    // Stores both the signature reference and the number of WebAssembly arguments
    signatures: HashMap<SignatureIndex, (ir::SigRef, usize)>,

    // Imported and local functions that have been created by
    // `FuncEnvironment::make_direct_func()`.
    // Stores both the function reference and the number of WebAssembly arguments
    functions: HashMap<FunctionIndex, (ir::FuncRef, usize)>,
}

impl TranslationState {
    pub fn new() -> TranslationState {
        TranslationState {
            stack: Vec::new(),
            control_stack: Vec::new(),
            phantom_unreachable_stack_depth: 0,
            real_unreachable_stack_depth: 0,
            globals: HashMap::new(),
            heaps: HashMap::new(),
            signatures: HashMap::new(),
            functions: HashMap::new(),
        }
    }

    fn clear(&mut self) {
        self.stack.clear();
        self.control_stack.clear();
        self.phantom_unreachable_stack_depth = 0;
        self.real_unreachable_stack_depth = 0;
        self.globals.clear();
        self.heaps.clear();
        self.signatures.clear();
        self.functions.clear();
    }

    /// Initialize the state for compiling a function with the given signature.
    ///
    /// This resets the state to containing only a single block representing the whole function.
    /// The exit block is the last block in the function which will contain the return instruction.
    pub fn initialize(&mut self, sig: &ir::Signature, exit_block: Ebb) {
        self.clear();
        self.push_block(
            exit_block,
            sig.returns
                .iter()
                .filter(|arg| arg.purpose == ir::ArgumentPurpose::Normal)
                .count(),
        );
    }

    /// Push a value.
    pub fn push1(&mut self, val: Value) {
        self.stack.push(val);
    }

    /// Push multiple values.
    pub fn pushn(&mut self, vals: &[Value]) {
        self.stack.extend_from_slice(vals);
    }

    /// Pop one value.
    pub fn pop1(&mut self) -> Value {
        self.stack.pop().unwrap()
    }

    /// Peek at the top of the stack without popping it.
    pub fn peek1(&self) -> Value {
        *self.stack.last().unwrap()
    }

    /// Pop two values. Return them in the order they were pushed.
    pub fn pop2(&mut self) -> (Value, Value) {
        let v2 = self.stack.pop().unwrap();
        let v1 = self.stack.pop().unwrap();
        (v1, v2)
    }

    /// Pop three values. Return them in the order they were pushed.
    pub fn pop3(&mut self) -> (Value, Value, Value) {
        let v3 = self.stack.pop().unwrap();
        let v2 = self.stack.pop().unwrap();
        let v1 = self.stack.pop().unwrap();
        (v1, v2, v3)
    }

    /// Pop the top `n` values on the stack.
    ///
    /// The popped values are not returned. Use `peekn` to look at them before popping.
    pub fn popn(&mut self, n: usize) {
        let new_len = self.stack.len() - n;
        self.stack.truncate(new_len);
    }

    /// Peek at the top `n` values on the stack in the order they were pushed.
    pub fn peekn(&self, n: usize) -> &[Value] {
        &self.stack[self.stack.len() - n..]
    }

    // Push a block on the control stack.
    pub fn push_block(&mut self, following_code: Ebb, num_result_types: usize) {
        self.control_stack.push(ControlStackFrame::Block {
            destination: following_code,
            original_stack_size: self.stack.len(),
            num_return_values: num_result_types,
            reachable: false,
        });
    }

    // Push a loop on the control stack.
    pub fn push_loop(&mut self, header: Ebb, following_code: Ebb, num_result_types: usize) {
        self.control_stack.push(ControlStackFrame::Loop {
            header,
            destination: following_code,
            original_stack_size: self.stack.len(),
            num_return_values: num_result_types,
            reachable: false,
        });
    }

    // Push an if on the control stack.
    pub fn push_if(&mut self, branch_inst: Inst, following_code: Ebb, num_result_types: usize) {
        self.control_stack.push(ControlStackFrame::If {
            branch_inst,
            destination: following_code,
            original_stack_size: self.stack.len(),
            num_return_values: num_result_types,
            reachable: false,
        });
    }

    /// Test if the translation state is currently in unreachable code.
    pub fn in_unreachable_code(&self) -> bool {
        if self.real_unreachable_stack_depth > 0 {
            true
        } else {
            debug_assert_eq!(self.phantom_unreachable_stack_depth, 0, "in reachable code");
            false
        }
    }
}

/// Methods for handling entity references.
impl TranslationState {
    /// Get the `GlobalVar` reference that should be used to access the global variable `index`.
    /// Create the reference if necessary.
    /// Also return the WebAssembly type of the global.
    pub fn get_global<FE: FuncEnvironment + ?Sized>(
        &mut self,
        func: &mut ir::Function,
        index: u32,
        environ: &mut FE,
    ) -> GlobalValue {
        let index = index as GlobalIndex;
        *self.globals.entry(index).or_insert_with(
            || environ.make_global(func, index),
        )
    }

    /// Get the `Heap` reference that should be used to access linear memory `index`.
    /// Create the reference if necessary.
    pub fn get_heap<FE: FuncEnvironment + ?Sized>(
        &mut self,
        func: &mut ir::Function,
        index: u32,
        environ: &mut FE,
    ) -> ir::Heap {
        let index = index as MemoryIndex;
        *self.heaps.entry(index).or_insert_with(
            || environ.make_heap(func, index),
        )
    }

    /// Get the `SigRef` reference that should be used to make an indirect call with signature
    /// `index`. Also return the number of WebAssembly arguments in the signature.
    ///
    /// Create the signature if necessary.
    pub fn get_indirect_sig<FE: FuncEnvironment + ?Sized>(
        &mut self,
        func: &mut ir::Function,
        index: u32,
        environ: &mut FE,
    ) -> (ir::SigRef, usize) {
        let index = index as SignatureIndex;
        *self.signatures.entry(index).or_insert_with(|| {
            let sig = environ.make_indirect_sig(func, index);
            (sig, normal_args(&func.dfg.signatures[sig]))
        })
    }

    /// Get the `FuncRef` reference that should be used to make a direct call to function
    /// `index`. Also return the number of WebAssembly arguments in the signature.
    ///
    /// Create the function reference if necessary.
    pub fn get_direct_func<FE: FuncEnvironment + ?Sized>(
        &mut self,
        func: &mut ir::Function,
        index: u32,
        environ: &mut FE,
    ) -> (ir::FuncRef, usize) {
        let index = index as FunctionIndex;
        *self.functions.entry(index).or_insert_with(|| {
            let fref = environ.make_direct_func(func, index);
            let sig = func.dfg.ext_funcs[fref].signature;
            (fref, normal_args(&func.dfg.signatures[sig]))
        })
    }
}

/// Count the number of normal parameters in a signature.
/// Exclude special-purpose parameters that represent runtime stuff and not WebAssembly arguments.
fn normal_args(sig: &ir::Signature) -> usize {
    sig.params
        .iter()
        .filter(|arg| arg.purpose == ir::ArgumentPurpose::Normal)
        .count()
}
<commit_msg>TranslationState can assert that its stacks are empty between functions.<commit_after>//! WebAssembly function translation state.
//!
//! The `TranslationState` struct defined in this module is used to keep track of the WebAssembly
//! value and control stacks during the translation of a single function.

use cretonne::ir::{self, Ebb, Inst, Value};
use environ::{FuncEnvironment, GlobalValue};
use std::collections::HashMap;
use translation_utils::{GlobalIndex, MemoryIndex, SignatureIndex, FunctionIndex};

/// A control stack frame can be an `if`, a `block` or a `loop`, each one having the following
/// fields:
///
/// - `destination`: reference to the `Ebb` that will hold the code after the control block;
/// - `num_return_values`: number of values returned by the control block;
/// - `original_stack_size`: size of the value stack at the beginning of the control block.
///
/// Moreover, the `if` frame has the `branch_inst` field that points to the `brz` instruction
/// separating the `true` and `false` branch. The `loop` frame has a `header` field that references
/// the `Ebb` that contains the beginning of the body of the loop.
#[derive(Debug)]
pub enum ControlStackFrame {
    If {
        destination: Ebb,
        branch_inst: Inst,
        num_return_values: usize,
        original_stack_size: usize,
        reachable: bool,
    },
    Block {
        destination: Ebb,
        num_return_values: usize,
        original_stack_size: usize,
        reachable: bool,
    },
    Loop {
        destination: Ebb,
        header: Ebb,
        num_return_values: usize,
        original_stack_size: usize,
        reachable: bool,
    },
}

/// Helper methods for the control stack objects.
impl ControlStackFrame {
    pub fn num_return_values(&self) -> usize {
        match *self {
            ControlStackFrame::If { num_return_values, .. } |
            ControlStackFrame::Block { num_return_values, .. } |
            ControlStackFrame::Loop { num_return_values, .. } => num_return_values,
        }
    }
    pub fn following_code(&self) -> Ebb {
        match *self {
            ControlStackFrame::If { destination, .. } |
            ControlStackFrame::Block { destination, .. } |
            ControlStackFrame::Loop { destination, .. } => destination,
        }
    }
    pub fn br_destination(&self) -> Ebb {
        match *self {
            ControlStackFrame::If { destination, .. } |
            ControlStackFrame::Block { destination, .. } => destination,
            ControlStackFrame::Loop { header, .. } => header,
        }
    }
    pub fn original_stack_size(&self) -> usize {
        match *self {
            ControlStackFrame::If { original_stack_size, .. } |
            ControlStackFrame::Block { original_stack_size, .. } |
            ControlStackFrame::Loop { original_stack_size, .. } => original_stack_size,
        }
    }
    pub fn is_loop(&self) -> bool {
        match *self {
            ControlStackFrame::If { .. } |
            ControlStackFrame::Block { .. } => false,
            ControlStackFrame::Loop { .. } => true,
        }
    }

    pub fn is_reachable(&self) -> bool {
        match *self {
            ControlStackFrame::If { reachable, .. } |
            ControlStackFrame::Block { reachable, .. } |
            ControlStackFrame::Loop { reachable, .. } => reachable,
        }
    }

    pub fn set_reachable(&mut self) {
        match *self {
            ControlStackFrame::If { ref mut reachable, .. } |
            ControlStackFrame::Block { ref mut reachable, .. } |
            ControlStackFrame::Loop { ref mut reachable, .. } => *reachable = true,
        }
    }
}

/// Contains information passed along during the translation and that records:
///
/// - The current value and control stacks.
/// - The depth of the two unreachable control blocks stacks, that are manipulated when translating
///   unreachable code;
pub struct TranslationState {
    pub stack: Vec<Value>,
    pub control_stack: Vec<ControlStackFrame>,
    pub phantom_unreachable_stack_depth: usize,
    pub real_unreachable_stack_depth: usize,

    // Map of global variables that have already been created by `FuncEnvironment::make_global`.
    globals: HashMap<GlobalIndex, GlobalValue>,

    // Map of heaps that have been created by `FuncEnvironment::make_heap`.
    heaps: HashMap<MemoryIndex, ir::Heap>,

    // Map of indirect call signatures that have been created by
    // `FuncEnvironment::make_indirect_sig()`.
    // Stores both the signature reference and the number of WebAssembly arguments
    signatures: HashMap<SignatureIndex, (ir::SigRef, usize)>,

    // Imported and local functions that have been created by
    // `FuncEnvironment::make_direct_func()`.
    // Stores both the function reference and the number of WebAssembly arguments
    functions: HashMap<FunctionIndex, (ir::FuncRef, usize)>,
}

impl TranslationState {
    pub fn new() -> TranslationState {
        TranslationState {
            stack: Vec::new(),
            control_stack: Vec::new(),
            phantom_unreachable_stack_depth: 0,
            real_unreachable_stack_depth: 0,
            globals: HashMap::new(),
            heaps: HashMap::new(),
            signatures: HashMap::new(),
            functions: HashMap::new(),
        }
    }

    fn clear(&mut self) {
        debug_assert!(self.stack.is_empty());
        debug_assert!(self.control_stack.is_empty());
        debug_assert_eq!(self.phantom_unreachable_stack_depth, 0);
        debug_assert_eq!(self.real_unreachable_stack_depth, 0);
        self.globals.clear();
        self.heaps.clear();
        self.signatures.clear();
        self.functions.clear();
    }

    /// Initialize the state for compiling a function with the given signature.
    ///
    /// This resets the state to containing only a single block representing the whole function.
    /// The exit block is the last block in the function which will contain the return instruction.
    pub fn initialize(&mut self, sig: &ir::Signature, exit_block: Ebb) {
        self.clear();
        self.push_block(
            exit_block,
            sig.returns
                .iter()
                .filter(|arg| arg.purpose == ir::ArgumentPurpose::Normal)
                .count(),
        );
    }

    /// Push a value.
    pub fn push1(&mut self, val: Value) {
        self.stack.push(val);
    }

    /// Push multiple values.
    pub fn pushn(&mut self, vals: &[Value]) {
        self.stack.extend_from_slice(vals);
    }

    /// Pop one value.
    pub fn pop1(&mut self) -> Value {
        self.stack.pop().unwrap()
    }

    /// Peek at the top of the stack without popping it.
    pub fn peek1(&self) -> Value {
        *self.stack.last().unwrap()
    }

    /// Pop two values. Return them in the order they were pushed.
    pub fn pop2(&mut self) -> (Value, Value) {
        let v2 = self.stack.pop().unwrap();
        let v1 = self.stack.pop().unwrap();
        (v1, v2)
    }

    /// Pop three values. Return them in the order they were pushed.
    pub fn pop3(&mut self) -> (Value, Value, Value) {
        let v3 = self.stack.pop().unwrap();
        let v2 = self.stack.pop().unwrap();
        let v1 = self.stack.pop().unwrap();
        (v1, v2, v3)
    }

    /// Pop the top `n` values on the stack.
    ///
    /// The popped values are not returned. Use `peekn` to look at them before popping.
    pub fn popn(&mut self, n: usize) {
        let new_len = self.stack.len() - n;
        self.stack.truncate(new_len);
    }

    /// Peek at the top `n` values on the stack in the order they were pushed.
    pub fn peekn(&self, n: usize) -> &[Value] {
        &self.stack[self.stack.len() - n..]
    }

    // Push a block on the control stack.
    pub fn push_block(&mut self, following_code: Ebb, num_result_types: usize) {
        self.control_stack.push(ControlStackFrame::Block {
            destination: following_code,
            original_stack_size: self.stack.len(),
            num_return_values: num_result_types,
            reachable: false,
        });
    }

    // Push a loop on the control stack.
    pub fn push_loop(&mut self, header: Ebb, following_code: Ebb, num_result_types: usize) {
        self.control_stack.push(ControlStackFrame::Loop {
            header,
            destination: following_code,
            original_stack_size: self.stack.len(),
            num_return_values: num_result_types,
            reachable: false,
        });
    }

    // Push an if on the control stack.
    pub fn push_if(&mut self, branch_inst: Inst, following_code: Ebb, num_result_types: usize) {
        self.control_stack.push(ControlStackFrame::If {
            branch_inst,
            destination: following_code,
            original_stack_size: self.stack.len(),
            num_return_values: num_result_types,
            reachable: false,
        });
    }

    /// Test if the translation state is currently in unreachable code.
    pub fn in_unreachable_code(&self) -> bool {
        if self.real_unreachable_stack_depth > 0 {
            true
        } else {
            debug_assert_eq!(self.phantom_unreachable_stack_depth, 0, "in reachable code");
            false
        }
    }
}

/// Methods for handling entity references.
impl TranslationState {
    /// Get the `GlobalVar` reference that should be used to access the global variable `index`.
    /// Create the reference if necessary.
    /// Also return the WebAssembly type of the global.
    pub fn get_global<FE: FuncEnvironment + ?Sized>(
        &mut self,
        func: &mut ir::Function,
        index: u32,
        environ: &mut FE,
    ) -> GlobalValue {
        let index = index as GlobalIndex;
        *self.globals.entry(index).or_insert_with(
            || environ.make_global(func, index),
        )
    }

    /// Get the `Heap` reference that should be used to access linear memory `index`.
    /// Create the reference if necessary.
    pub fn get_heap<FE: FuncEnvironment + ?Sized>(
        &mut self,
        func: &mut ir::Function,
        index: u32,
        environ: &mut FE,
    ) -> ir::Heap {
        let index = index as MemoryIndex;
        *self.heaps.entry(index).or_insert_with(
            || environ.make_heap(func, index),
        )
    }

    /// Get the `SigRef` reference that should be used to make an indirect call with signature
    /// `index`. Also return the number of WebAssembly arguments in the signature.
    ///
    /// Create the signature if necessary.
    pub fn get_indirect_sig<FE: FuncEnvironment + ?Sized>(
        &mut self,
        func: &mut ir::Function,
        index: u32,
        environ: &mut FE,
    ) -> (ir::SigRef, usize) {
        let index = index as SignatureIndex;
        *self.signatures.entry(index).or_insert_with(|| {
            let sig = environ.make_indirect_sig(func, index);
            (sig, normal_args(&func.dfg.signatures[sig]))
        })
    }

    /// Get the `FuncRef` reference that should be used to make a direct call to function
    /// `index`. Also return the number of WebAssembly arguments in the signature.
    ///
    /// Create the function reference if necessary.
    pub fn get_direct_func<FE: FuncEnvironment + ?Sized>(
        &mut self,
        func: &mut ir::Function,
        index: u32,
        environ: &mut FE,
    ) -> (ir::FuncRef, usize) {
        let index = index as FunctionIndex;
        *self.functions.entry(index).or_insert_with(|| {
            let fref = environ.make_direct_func(func, index);
            let sig = func.dfg.ext_funcs[fref].signature;
            (fref, normal_args(&func.dfg.signatures[sig]))
        })
    }
}

/// Count the number of normal parameters in a signature.
/// Exclude special-purpose parameters that represent runtime stuff and not WebAssembly arguments.
fn normal_args(sig: &ir::Signature) -> usize {
    sig.params
        .iter()
        .filter(|arg| arg.purpose == ir::ArgumentPurpose::Normal)
        .count()
}
<|endoftext|>


<commit_before>use urlencoding::{encode as urlencode};
use urlparse::{Query, parse_qs, urlparse, urlunparse};

fn partial_urlencode(value: &str) -> String {
  // The urlencode from the urlencoding library goes further than I like,
  // and also URL encodes ASCII digits. Reverse that stuff.
  urlencode(value)
    .replace("%30", "0")
    .replace("%31", "1")
    .replace("%32", "2")
    .replace("%33", "3")
    .replace("%34", "4")
    .replace("%35", "5")
    .replace("%36", "6")
    .replace("%37", "7")
    .replace("%38", "8")
    .replace("%39", "9")
}

/// Re-encode a query string for Rust
fn encode_querystring(query: Query) -> Option<String> {
  let mut query_components: Vec<String> = vec![];
  for (key, value) in query {
    for v in value.iter() {
      query_components.push(format!("{}={}", key, partial_urlencode(v)));
    }
  }
  if query_components.len() > 0 {
    Some(query_components.join("&"))
  } else {
    None
  }
}

/// Strip tracking junk and URL suffixes.
pub fn tidy_url(url: &str) -> String {
    let mut parsed_url = urlparse(url);

    // Always get the desktop version of Twitter URLs
    if parsed_url.netloc == "mobile.twitter.com" {
        parsed_url.netloc = String::from("twitter.com");
    }

    // Remove any tracking junk from Amazon URLs so they're not a
    // ridiculous length
    if parsed_url.netloc == "www.amazon.co.uk" {
        parsed_url.query = None;
        let mut new_path = String::from("");
        for component in parsed_url.path.split("ref=") {
            new_path = String::from(component);
            break;
        }
        parsed_url.path = new_path;
    }

    // Strip tracking junk from Medium, Mashable and Buzzfeed
    if (parsed_url.netloc == "medium.com") ||
       (parsed_url.netloc == "www.buzzfeed.com") ||
       (parsed_url.netloc == "mashable.com") {
            parsed_url.fragment = None;
    }

    // Remove '#notes' from Tumblr URLs
    if parsed_url.netloc.ends_with("tumblr.com") {
        let new_fragment: Option<String> = match parsed_url.fragment {
            Some(fragment) => {
                if fragment == "notes" {
                    None
                } else {
                    Some(fragment)
                }
            }
            None => None,
        };
        parsed_url.fragment = new_fragment;
    }

  // Remove &feature=youtu.be from YouTube URLs
  if parsed_url.netloc.ends_with("youtube.com") {
    parsed_url.query = match parsed_url.query {
      Some(qs) => {
        let mut query = parse_qs(&qs);
        query.remove("feature");
        encode_querystring(query)
      },
      None => None,
    };
  }

  // Remove any UTM tracking parameters from URLs
  parsed_url.query = match parsed_url.query {
    Some(qs) => {
      let mut query = parse_qs(&qs);
      let utm_keys: Vec<_> = query
        .keys()
        .filter(|key| key.starts_with("utm_"))
        .map(|k| k.clone())
        .collect();
      for key in utm_keys {
        query.remove(&key);
      }
      encode_querystring(query)
    },
    None => None
  };

  urlunparse(parsed_url)
}

macro_rules! tidy_url_tests {
    ($($name:ident: $value:expr,)*) => {
    $(
        #[test]
        fn $name() {
            let (input, expected) = $value;
            assert_eq!(expected, tidy_url(input));
        }
    )*
    }
}

tidy_url_tests! {
    url_tw_0:  ("https://mobile.twitter.com/Breaking911/status/822589852596191235",
                "https://twitter.com/Breaking911/status/822589852596191235"),
    url_tw_1:  ("https://twitter.com/elise3aum/status/822903824268533762",
                "https://twitter.com/elise3aum/status/822903824268533762"),

    url_az_0:  ("https://www.amazon.co.uk/dp/B01DFKBL68/ref=gw_aucc_comb_AB_clean_2?pf_rd_r=N1MF2PWADRHS55427ETG&pf_rd_p=c11a2c11-c670-46ff-87a1-c1eef4fad652",
                "https://www.amazon.co.uk/dp/B01DFKBL68/"),

    url_md_0:  ("https://medium.com/@anildash/forget-why-its-time-to-get-to-work-c49ac5f0da20#.sjyskxdsz",
                "https://medium.com/@anildash/forget-why-its-time-to-get-to-work-c49ac5f0da20"),

    url_tm_0:  ("http://azurelunatic.tumblr.com/post/155525051123/things-about-hufflepuffs-539#notes",
                "http://azurelunatic.tumblr.com/post/155525051123/things-about-hufflepuffs-539"),

    url_8:   ("https://www.buzzfeed.com/katienotopoulos/the-end-of-apple-man#.biqmkzz8kn",
              "https://www.buzzfeed.com/katienotopoulos/the-end-of-apple-man"),
    url_9:   ("http://mashable.com/2016/03/21/apple-liam-recycling-robot/#b9y4lv3m4qqX",
              "http://mashable.com/2016/03/21/apple-liam-recycling-robot/"),

  youtube_plain_video: (
    "https://www.youtube.com/watch?v=zB4I68XVPzQ",
    "https://www.youtube.com/watch?v=zB4I68XVPzQ",
  ),

  youtube_with_timestamp: (
    "https://www.youtube.com/watch?v=D68cUzqcTrg&feature=youtu.be#t=1m",
    "https://www.youtube.com/watch?v=D68cUzqcTrg#t=1m",
  ),

  youtube_with_trailing_feature: (
    "https://www.youtube.com/watch?v=PbJqNa0_Oz0&feature=youtu.be",
    "https://www.youtube.com/watch?v=PbJqNa0_Oz0",
  ),

  youtube_with_leading_feature: (
    "https://www.youtube.com/watch?feature=youtu.be&v=oPo4n9tBPsk",
    "https://www.youtube.com/watch?v=oPo4n9tBPsk",
  ),

  youtube_with_only_feature: (
    "https://www.youtube.com/watch?feature=youtu.be",
    "https://www.youtube.com/watch",
  ),

  youtube_channel: (
    "https://www.youtube.com/user/TheQIElves",
    "https://www.youtube.com/user/TheQIElves"
  ),

  single_utm_tracker: (
    "https://example.com?utm_medium=social",
    "https://example.com",
  ),

  multiple_utm_tracker: (
    "https://example.com?utm_medium=social&utm_source=twitter",
    "https://example.com",
  ),

  multiple_utm_tracker_with_others: (
    "https://example.com?utm_medium=social&foo=bar&utm_source=twitter",
    "https://example.com?foo=bar",
  ),

  url_with_spaces: (
    "https://example.com?foo=bar%20baz",
    "https://example.com?foo=bar%20baz"
  ),

  url_with_numerals: (
    "https://example.com?foo=bar0baz",
    "https://example.com?foo=bar0baz"
  ),

  python_docs_bare: (
    "https://docs.python.org/3.5/library/subprocess.html",
    "https://docs.python.org/3.5/library/subprocess.html"
  ),

  python_docs_with_anchor: (
    "https://docs.python.org/3.5/library/subprocess.html#subprocess.run",
    "https://docs.python.org/3.5/library/subprocess.html#subprocess.run"
  ),

  python_docs_with_highlight: (
    "https://docs.python.org/3.5/library/subprocess.html?highlight=subprocess",
    "https://docs.python.org/3.5/library/subprocess.html"
  ),

  python_docs_with_highlight_anchor: (
    "https://docs.python.org/3.5/library/subprocess.html?highlight=subprocess#subprocess.run",
    "https://docs.python.org/3.5/library/subprocess.html#subprocess.run"
  ),

  python_docs_with_highlight_module_anchor: (
    "https://docs.python.org/3.5/library/subprocess.html?highlight=subprocess#module-subprocess",
    "https://docs.python.org/3.5/library/subprocess.html"
  ),
}
<commit_msg>Remove ?highlight= and certain fragments from docs.python.org URLs<commit_after>use urlencoding::{encode as urlencode};
use urlparse::{Query, parse_qs, urlparse, urlunparse};

fn partial_urlencode(value: &str) -> String {
  // The urlencode from the urlencoding library goes further than I like,
  // and also URL encodes ASCII digits. Reverse that stuff.
  urlencode(value)
    .replace("%30", "0")
    .replace("%31", "1")
    .replace("%32", "2")
    .replace("%33", "3")
    .replace("%34", "4")
    .replace("%35", "5")
    .replace("%36", "6")
    .replace("%37", "7")
    .replace("%38", "8")
    .replace("%39", "9")
}

/// Re-encode a query string for Rust
fn encode_querystring(query: Query) -> Option<String> {
  let mut query_components: Vec<String> = vec![];
  for (key, value) in query {
    for v in value.iter() {
      query_components.push(format!("{}={}", key, partial_urlencode(v)));
    }
  }
  if query_components.len() > 0 {
    Some(query_components.join("&"))
  } else {
    None
  }
}

/// Strip tracking junk and URL suffixes.
pub fn tidy_url(url: &str) -> String {
    let mut parsed_url = urlparse(url);

    // Always get the desktop version of Twitter URLs
    if parsed_url.netloc == "mobile.twitter.com" {
        parsed_url.netloc = String::from("twitter.com");
    }

    // Remove any tracking junk from Amazon URLs so they're not a
    // ridiculous length
    if parsed_url.netloc == "www.amazon.co.uk" {
        parsed_url.query = None;
        let mut new_path = String::from("");
        for component in parsed_url.path.split("ref=") {
            new_path = String::from(component);
            break;
        }
        parsed_url.path = new_path;
    }

    // Strip tracking junk from Medium, Mashable and Buzzfeed
    if (parsed_url.netloc == "medium.com") ||
       (parsed_url.netloc == "www.buzzfeed.com") ||
       (parsed_url.netloc == "mashable.com") {
            parsed_url.fragment = None;
    }

    // Remove '#notes' from Tumblr URLs
    if parsed_url.netloc.ends_with("tumblr.com") {
        let new_fragment: Option<String> = match parsed_url.fragment {
            Some(fragment) => {
                if fragment == "notes" {
                    None
                } else {
                    Some(fragment)
                }
            }
            None => None,
        };
        parsed_url.fragment = new_fragment;
    }

  // Remove &feature=youtu.be from YouTube URLs
  if parsed_url.netloc.ends_with("youtube.com") {
    parsed_url.query = match parsed_url.query {
      Some(qs) => {
        let mut query = parse_qs(&qs);
        query.remove("feature");
        encode_querystring(query)
      },
      None => None,
    };
  }

  // Remove any UTM tracking parameters from URLs
  parsed_url.query = match parsed_url.query {
    Some(qs) => {
      let mut query = parse_qs(&qs);
      let utm_keys: Vec<_> = query
        .keys()
        .filter(|key| key.starts_with("utm_"))
        .map(|k| k.clone())
        .collect();
      for key in utm_keys {
        query.remove(&key);
      }
      encode_querystring(query)
    },
    None => None
  };

  // Tidy up the query and anchor links in modules on docs.python.org
  if parsed_url.netloc == "docs.python.org" {

    // If this is a module page, scrap any module- fragment.
    parsed_url.fragment = match parsed_url.fragment {
      Some(fragment) => if fragment.starts_with("module-") { None } else { Some(fragment) },
      None => None,
    };

    // Scrap the highlight
    parsed_url.query = match parsed_url.query {
      Some(qs) => {
        let mut query = parse_qs(&qs);
        query.remove("highlight");
        encode_querystring(query)
      },
      None => None,
    };
  }

  urlunparse(parsed_url)
}

macro_rules! tidy_url_tests {
    ($($name:ident: $value:expr,)*) => {
    $(
        #[test]
        fn $name() {
            let (input, expected) = $value;
            assert_eq!(expected, tidy_url(input));
        }
    )*
    }
}

tidy_url_tests! {
    url_tw_0:  ("https://mobile.twitter.com/Breaking911/status/822589852596191235",
                "https://twitter.com/Breaking911/status/822589852596191235"),
    url_tw_1:  ("https://twitter.com/elise3aum/status/822903824268533762",
                "https://twitter.com/elise3aum/status/822903824268533762"),

    url_az_0:  ("https://www.amazon.co.uk/dp/B01DFKBL68/ref=gw_aucc_comb_AB_clean_2?pf_rd_r=N1MF2PWADRHS55427ETG&pf_rd_p=c11a2c11-c670-46ff-87a1-c1eef4fad652",
                "https://www.amazon.co.uk/dp/B01DFKBL68/"),

    url_md_0:  ("https://medium.com/@anildash/forget-why-its-time-to-get-to-work-c49ac5f0da20#.sjyskxdsz",
                "https://medium.com/@anildash/forget-why-its-time-to-get-to-work-c49ac5f0da20"),

    url_tm_0:  ("http://azurelunatic.tumblr.com/post/155525051123/things-about-hufflepuffs-539#notes",
                "http://azurelunatic.tumblr.com/post/155525051123/things-about-hufflepuffs-539"),

    url_8:   ("https://www.buzzfeed.com/katienotopoulos/the-end-of-apple-man#.biqmkzz8kn",
              "https://www.buzzfeed.com/katienotopoulos/the-end-of-apple-man"),
    url_9:   ("http://mashable.com/2016/03/21/apple-liam-recycling-robot/#b9y4lv3m4qqX",
              "http://mashable.com/2016/03/21/apple-liam-recycling-robot/"),

  youtube_plain_video: (
    "https://www.youtube.com/watch?v=zB4I68XVPzQ",
    "https://www.youtube.com/watch?v=zB4I68XVPzQ",
  ),

  youtube_with_timestamp: (
    "https://www.youtube.com/watch?v=D68cUzqcTrg&feature=youtu.be#t=1m",
    "https://www.youtube.com/watch?v=D68cUzqcTrg#t=1m",
  ),

  youtube_with_trailing_feature: (
    "https://www.youtube.com/watch?v=PbJqNa0_Oz0&feature=youtu.be",
    "https://www.youtube.com/watch?v=PbJqNa0_Oz0",
  ),

  youtube_with_leading_feature: (
    "https://www.youtube.com/watch?feature=youtu.be&v=oPo4n9tBPsk",
    "https://www.youtube.com/watch?v=oPo4n9tBPsk",
  ),

  youtube_with_only_feature: (
    "https://www.youtube.com/watch?feature=youtu.be",
    "https://www.youtube.com/watch",
  ),

  youtube_channel: (
    "https://www.youtube.com/user/TheQIElves",
    "https://www.youtube.com/user/TheQIElves"
  ),

  single_utm_tracker: (
    "https://example.com?utm_medium=social",
    "https://example.com",
  ),

  multiple_utm_tracker: (
    "https://example.com?utm_medium=social&utm_source=twitter",
    "https://example.com",
  ),

  multiple_utm_tracker_with_others: (
    "https://example.com?utm_medium=social&foo=bar&utm_source=twitter",
    "https://example.com?foo=bar",
  ),

  url_with_spaces: (
    "https://example.com?foo=bar%20baz",
    "https://example.com?foo=bar%20baz"
  ),

  url_with_numerals: (
    "https://example.com?foo=bar0baz",
    "https://example.com?foo=bar0baz"
  ),

  python_docs_bare: (
    "https://docs.python.org/3.5/library/subprocess.html",
    "https://docs.python.org/3.5/library/subprocess.html"
  ),

  python_docs_with_anchor: (
    "https://docs.python.org/3.5/library/subprocess.html#subprocess.run",
    "https://docs.python.org/3.5/library/subprocess.html#subprocess.run"
  ),

  python_docs_with_highlight: (
    "https://docs.python.org/3.5/library/subprocess.html?highlight=subprocess",
    "https://docs.python.org/3.5/library/subprocess.html"
  ),

  python_docs_with_highlight_anchor: (
    "https://docs.python.org/3.5/library/subprocess.html?highlight=subprocess#subprocess.run",
    "https://docs.python.org/3.5/library/subprocess.html#subprocess.run"
  ),

  python_docs_with_highlight_module_anchor: (
    "https://docs.python.org/3.5/library/subprocess.html?highlight=subprocess#module-subprocess",
    "https://docs.python.org/3.5/library/subprocess.html"
  ),
}
<|endoftext|>


<commit_before>//Copyright 2016 secret-service-rs Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

// Contains helpers for:
//   exec_prompt
//   interfaces
//   formatting secrets
//
//   Consider: What else should be in here? Should
//   formatting secrets be in crypto? Should interfaces
//   have their own module?

use error::SsError;
use proxy::prompt::PromptProxy;
use proxy::service::ServiceProxy;
use session::Session;
use ss::SS_DBUS_NAME;
use ss_crypto::encrypt;
use proxy::SecretStruct;
use proxy::item::SecretStructInput;

use rand::{Rng, rngs::OsRng};
use std::sync::mpsc::channel;
use zvariant::ObjectPath;

// Helper enum for locking
pub(crate) enum LockAction {
    Lock,
    Unlock,
}

pub(crate) fn lock_or_unlock(
    conn: zbus::Connection,
    service_interface: &ServiceProxy,
    object_path: &ObjectPath,
    lock_action: LockAction
    ) -> ::Result<()>
{
    let objects = vec![object_path];

    let lock_action_res = match lock_action {
        LockAction::Lock => service_interface.lock(objects)?,
        LockAction::Unlock => service_interface.unlock(objects)?,
    };

    if lock_action_res.object_paths.is_empty() {
        exec_prompt(conn.clone(), &lock_action_res.prompt)?;
    }
    Ok(())
}

pub(crate) fn format_secret(
    session: &Session,
    secret: &[u8],
    content_type: &str
    ) -> ::Result<SecretStructInput>
{
    let content_type = content_type.to_owned();

    if session.is_encrypted() {
        let mut rng = OsRng {};
        let mut aes_iv = [0;16];
        rng.fill(&mut aes_iv);

        let encrypted_secret = encrypt(secret, &session.get_aes_key()[..], &aes_iv)?;

        // Construct secret struct
        let parameters = aes_iv.to_vec();
        let value = encrypted_secret;

        Ok(SecretStructInput {
            inner: SecretStruct {
                session: session.object_path.clone(),
                parameters,
                value,
                content_type,
            }
        })
    } else {
        // just Plain for now
        let parameters = Vec::new();
        let value = secret.to_vec();

        Ok(SecretStructInput {
            inner: SecretStruct {
                session: session.object_path.clone(),
                parameters,
                value,
                content_type,
            }
        })
    }
}

pub(crate) fn exec_prompt(conn: zbus::Connection, prompt: &ObjectPath) -> ::Result<zvariant::OwnedValue> {
    let prompt_interface = PromptProxy::new_for(
        &conn,
        SS_DBUS_NAME,
        prompt,
        )?;

    let (tx, rx) = channel();

    // create a handler for `completed` signal
    prompt_interface
        .connect_completed(move |dismissed, result| {
            let res = if dismissed {
                Err(SsError::Prompt)
            } else {
                Ok(result.into())
            };

            // FIXME remove unwrap
            tx.send(res).expect("remove this unwrap");

            Ok(())
        })?;

    // FIXME figure out window_id
    let window_id = "";
    prompt_interface.prompt(window_id)?;

    // waits for next signal and calls the handler.
    // If message handled by above handler, `next_signal` returns `Ok(None)`, ending loop.
    while prompt_interface.next_signal()?.is_some() {};

    // FIXME remove unwrap
    rx.recv().expect("remove this unwrap")
}
<commit_msg>allow channel in exec_prompt to panic<commit_after>//Copyright 2016 secret-service-rs Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

// Contains helpers for:
//   exec_prompt
//   interfaces
//   formatting secrets
//
//   Consider: What else should be in here? Should
//   formatting secrets be in crypto? Should interfaces
//   have their own module?

use error::SsError;
use proxy::prompt::PromptProxy;
use proxy::service::ServiceProxy;
use session::Session;
use ss::SS_DBUS_NAME;
use ss_crypto::encrypt;
use proxy::SecretStruct;
use proxy::item::SecretStructInput;

use rand::{Rng, rngs::OsRng};
use std::sync::mpsc::channel;
use zvariant::ObjectPath;

// Helper enum for locking
pub(crate) enum LockAction {
    Lock,
    Unlock,
}

pub(crate) fn lock_or_unlock(
    conn: zbus::Connection,
    service_interface: &ServiceProxy,
    object_path: &ObjectPath,
    lock_action: LockAction
    ) -> ::Result<()>
{
    let objects = vec![object_path];

    let lock_action_res = match lock_action {
        LockAction::Lock => service_interface.lock(objects)?,
        LockAction::Unlock => service_interface.unlock(objects)?,
    };

    if lock_action_res.object_paths.is_empty() {
        exec_prompt(conn.clone(), &lock_action_res.prompt)?;
    }
    Ok(())
}

pub(crate) fn format_secret(
    session: &Session,
    secret: &[u8],
    content_type: &str
    ) -> ::Result<SecretStructInput>
{
    let content_type = content_type.to_owned();

    if session.is_encrypted() {
        let mut rng = OsRng {};
        let mut aes_iv = [0;16];
        rng.fill(&mut aes_iv);

        let encrypted_secret = encrypt(secret, &session.get_aes_key()[..], &aes_iv)?;

        // Construct secret struct
        let parameters = aes_iv.to_vec();
        let value = encrypted_secret;

        Ok(SecretStructInput {
            inner: SecretStruct {
                session: session.object_path.clone(),
                parameters,
                value,
                content_type,
            }
        })
    } else {
        // just Plain for now
        let parameters = Vec::new();
        let value = secret.to_vec();

        Ok(SecretStructInput {
            inner: SecretStruct {
                session: session.object_path.clone(),
                parameters,
                value,
                content_type,
            }
        })
    }
}

pub(crate) fn exec_prompt(conn: zbus::Connection, prompt: &ObjectPath) -> ::Result<zvariant::OwnedValue> {
    let prompt_interface = PromptProxy::new_for(
        &conn,
        SS_DBUS_NAME,
        prompt,
        )?;

    let (tx, rx) = channel();

    // create a handler for `completed` signal
    prompt_interface
        .connect_completed(move |dismissed, result| {
            let res = if dismissed {
                Err(SsError::Prompt)
            } else {
                Ok(result.into())
            };

            tx.send(res).expect("rx should not be dropped, channel scoped to this fn");

            Ok(())
        })?;

    // TODO figure out window_id
    let window_id = "";
    prompt_interface.prompt(window_id)?;

    // waits for next signal and calls the handler.
    // If message handled by above handler, `next_signal` returns `Ok(None)`, ending loop.
    while prompt_interface.next_signal()?.is_some() {};

    // See comment on allowing channels to panic: https://www.reddit.com/r/rust/comments/awh751/proposal_new_channels_for_rusts_standard_library/ehmk3lz/
    rx.recv().expect("tx shold not be dropped, channel scoped to this fn")
}
<|endoftext|>


<commit_before>use log::{debug, error};
use std::fs;
use std::io::prelude::*;

use failure::{Error, Fail};

#[derive(Debug, Fail)]
pub enum UtilError {
    #[fail(display = "insufficient file access")]
    FileAccess,
    #[fail(display = "invalid file format")]
    FileFormat,
}

/// Static cast the given 64-bit unsigned integer to a 64-bit signed integer.
/// This is probably only useful when some other code provides you a u64
///  that is meant to be an i64 (aka. uncommon).
///
/// In C: `*(int64_t *)&i`
///
/// # Examples
///
/// ```
/// use lancelot::util::*;
/// assert_eq!(0, u64_i64(0));
/// assert_eq!(1, u64_i64(0x1));
/// assert_eq!(-1, u64_i64(0xFFFF_FFFF_FFFF_FFFF));
/// ```
pub fn u64_i64(i: u64) -> i64 {
    if i & 1 << 63 > 0 {
        // TODO: there's probably some elegant rust-way to do this.
        // in the meantime, manually compute twos-complement.
        let bits = i & 0x7FFF_FFFF_FFFF_FFFF;
        let bits = !bits;
        let bits = bits + 1;
        let bits = bits & 0x7FFF_FFFF_FFFF_FFFF;
        -(bits as i64)
    } else {
        i as i64
    }
}

/// Round the given value up to the next multiple of the given base.
///
/// # Panics
///
///   - Base `b` must be at least `2`.
///
/// # Examples
///
/// ```
/// use lancelot::util::*;
/// assert_eq!(align(0, 2), 0);
/// assert_eq!(align(1, 2), 2);
/// assert_eq!(align(2, 2), 2);
/// assert_eq!(align(3, 2), 4);
/// assert_eq!(align(4, 2), 4);
/// ```
pub fn align(i: usize, b: usize) -> usize {
    if b < 2 {
        panic!("base `b` must be at least: 2");
    }
    let rem = i % b;
    if rem == 0 {
        i
    } else {
        i + (b - rem)
    }
}

pub fn hexdump_ascii(b: u8) -> char {
    if b.is_ascii_graphic() || b == b' ' {
        b as char
    } else {
        '.'
    }
}

pub fn hexdump(buf: &[u8], offset: usize) -> String {
    // 01234567:  00 01 02 03 04 05 06 07  ...............
    // <prefix>   <hex col>                <ascii col>

    let padding = "  ";

    let padding_size = 2;
    let hex_col_size = 3;
    let ascii_col_size = 1;
    let prefix_size = 8 + 1;
    let newline_size = 1;
    let line_size = prefix_size
        + padding_size
        + 16 * hex_col_size
        + padding_size
        + 16 * ascii_col_size
        + newline_size;
    let line_count = align(buf.len(), 0x10) / 0x10;

    let mut ret = String::with_capacity(line_count * line_size);

    let mut line = String::with_capacity(line_size);
    let mut remaining_count = buf.len();
    for line_index in 0..line_count {
        let line_elem_count = 0x10.min(remaining_count);
        let padding_elem_count = 0x10 - line_elem_count;

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        // ^^^^^^^^^
        line.push_str(format!("{:08x}:", offset + 0x10 * line_index).as_str());

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //          ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //            ^^^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push_str(format!("{:02x} ", elem).as_str());
        }
        for _ in 0..padding_elem_count {
            line.push_str("   ");
        }

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                   ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                     ^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push(hexdump_ascii(*elem))
        }
        for _ in 0..padding_elem_count {
            line.push(' ');
        }
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                                    ^
        line.push('\n');

        ret.push_str(line.as_str());
        line.truncate(0x0);
        remaining_count -= line_elem_count;
    }

    ret
}

pub fn read_file(filename: &str) -> Result<Vec<u8>, Error> {
    debug!("read_file: {:?}", filename);

    let mut buf = Vec::new();
    {
        debug!("reading file: {}", filename);
        let mut f = match fs::File::open(filename) {
            Ok(f) => f,
            Err(_) => {
                error!("failed to open file: {}", filename);
                return Err(UtilError::FileAccess.into());
            }
        };
        let bytes_read = match f.read_to_end(&mut buf) {
            Ok(c) => c,
            Err(_) => {
                error!("failed to read entire file: {}", filename);
                return Err(UtilError::FileAccess.into());
            }
        };
        debug!("read {} bytes", bytes_read);
        if bytes_read < 0x10 {
            error!("file too small: {}", filename);
            return Err(UtilError::FileFormat.into());
        }
    }

    Ok(buf)
}
<commit_msg>util: fix rendering of hexdump<commit_after>use log::{debug, error};
use std::fs;
use std::io::prelude::*;

use failure::{Error, Fail};

#[derive(Debug, Fail)]
pub enum UtilError {
    #[fail(display = "insufficient file access")]
    FileAccess,
    #[fail(display = "invalid file format")]
    FileFormat,
}

/// Static cast the given 64-bit unsigned integer to a 64-bit signed integer.
/// This is probably only useful when some other code provides you a u64
///  that is meant to be an i64 (aka. uncommon).
///
/// In C: `*(int64_t *)&i`
///
/// # Examples
///
/// ```
/// use lancelot::util::*;
/// assert_eq!(0, u64_i64(0));
/// assert_eq!(1, u64_i64(0x1));
/// assert_eq!(-1, u64_i64(0xFFFF_FFFF_FFFF_FFFF));
/// ```
pub fn u64_i64(i: u64) -> i64 {
    if i & 1 << 63 > 0 {
        // TODO: there's probably some elegant rust-way to do this.
        // in the meantime, manually compute twos-complement.
        let bits = i & 0x7FFF_FFFF_FFFF_FFFF;
        let bits = !bits;
        let bits = bits + 1;
        let bits = bits & 0x7FFF_FFFF_FFFF_FFFF;
        -(bits as i64)
    } else {
        i as i64
    }
}

/// Round the given value up to the next multiple of the given base.
///
/// # Panics
///
///   - Base `b` must be at least `2`.
///
/// # Examples
///
/// ```
/// use lancelot::util::*;
/// assert_eq!(align(0, 2), 0);
/// assert_eq!(align(1, 2), 2);
/// assert_eq!(align(2, 2), 2);
/// assert_eq!(align(3, 2), 4);
/// assert_eq!(align(4, 2), 4);
/// ```
pub fn align(i: usize, b: usize) -> usize {
    if b < 2 {
        panic!("base `b` must be at least: 2");
    }
    let rem = i % b;
    if rem == 0 {
        i
    } else {
        i + (b - rem)
    }
}

pub fn hexdump_ascii(b: u8) -> char {
    if b.is_ascii_graphic() || b == b' ' {
        b as char
    } else {
        '.'
    }
}

pub fn hexdump(buf: &[u8], offset: usize) -> String {
    // 01234567:  00 01 02 03 04 05 06 07  ...............
    // <prefix>   <hex col>                <ascii col>

    let padding = "  ";

    let padding_size = 2;
    let hex_col_size = 3;
    let ascii_col_size = 1;
    let prefix_size = 8 + 1;
    let newline_size = 1;
    let line_size = prefix_size
        + padding_size
        + 16 * hex_col_size
        + padding_size
        + 16 * ascii_col_size
        + newline_size;
    let line_count = align(buf.len(), 0x10) / 0x10;

    let mut ret = String::with_capacity(line_count * line_size);

    let mut line = String::with_capacity(line_size);
    let mut remaining_count = buf.len();
    for line_index in 0..line_count {
        let line_elem_count = 0x10.min(remaining_count);
        let padding_elem_count = 0x10 - line_elem_count;

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        // ^^^^^^^^^
        line.push_str(format!("{:08x}:", offset + 0x10 * line_index).as_str());

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //          ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //            ^^^
        for elem in &buf[(line_index * 0x10)..(line_index * 0x10) + line_elem_count] {
            line.push_str(format!("{:02x} ", elem).as_str());
        }
        for _ in 0..padding_elem_count {
            line.push_str("   ");
        }

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                   ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                     ^
        for elem in &buf[(line_index * 0x10)..(line_index * 0x10) + line_elem_count] {
            line.push(hexdump_ascii(*elem))
        }
        for _ in 0..padding_elem_count {
            line.push(' ');
        }
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                                    ^
        line.push('\n');

        ret.push_str(line.as_str());
        line.truncate(0x0);
        remaining_count -= line_elem_count;
    }

    ret
}

pub fn read_file(filename: &str) -> Result<Vec<u8>, Error> {
    debug!("read_file: {:?}", filename);

    let mut buf = Vec::new();
    {
        debug!("reading file: {}", filename);
        let mut f = match fs::File::open(filename) {
            Ok(f) => f,
            Err(_) => {
                error!("failed to open file: {}", filename);
                return Err(UtilError::FileAccess.into());
            }
        };
        let bytes_read = match f.read_to_end(&mut buf) {
            Ok(c) => c,
            Err(_) => {
                error!("failed to read entire file: {}", filename);
                return Err(UtilError::FileAccess.into());
            }
        };
        debug!("read {} bytes", bytes_read);
        if bytes_read < 0x10 {
            error!("file too small: {}", filename);
            return Err(UtilError::FileFormat.into());
        }
    }

    Ok(buf)
}
<|endoftext|>


<commit_before>use std::borrow::Cow;

/// Escape a filter literal.
///
/// Literal values appearing in an LDAP filter can contain any character,
/// but some characters (parentheses, asterisk, backslash, NUL) must be
/// escaped in the filter's string representation. This function does the
/// escaping.
///
/// The argument, `lit`, can be owned or borrowed. The function doesn't
/// allocate the return value unless there's need to escape the input.
pub fn ldap_escape<'a, S: Into<Cow<'a, str>>>(lit: S) -> Cow<'a, str> {
    #[inline]
    fn needs_escape(c: u8) -> bool {
        c == b'\\' || c == b'*' || c == b'(' || c == b')' || c == 0
    }

    #[inline]
    fn xdigit(c: u8) -> u8 {
        c + if c < 10 { b'0' } else { b'a' - 10 }
    }

    let lit = lit.into();
    let mut output = None;
    for (i, &c) in lit.as_bytes().iter().enumerate() {
        if needs_escape(c) {
            if output.is_none() {
                output = Some(Vec::with_capacity(lit.len() + 12)); // guess: up to 4 escaped chars
                output.as_mut().unwrap().extend(lit[..i].as_bytes());
            }
            let output = output.as_mut().unwrap();
            output.push(b'\\');
            output.push(xdigit(c >> 4));
            output.push(xdigit(c & 0xF));
        } else if let Some(ref mut output) = output {
            output.push(c);
        }
    }
    if let Some(output) = output {
        // unchecked conversion is safe here: we receive a valid
        // UTF-8 value, by definition, and only replace single ASCII
        // bytes with ASCII byte sequences
        Cow::Owned(unsafe { String::from_utf8_unchecked(output) })
    } else {
        lit
    }
}

/// Escape an attribute value in a relative distinguished name (RDN).
///
/// When a literal string is used to represent an attribute value in an RDN,
/// some of its characters might need to be escaped according to the rules
/// of [RFC 4514](https://tools.ietf.org/html/rfc4514).
///
/// The function is named `dn_escape()` instead of `rdn_escape()` because of
/// a long-standing association of its intended use with the handling of DNs.
///
/// The argument, `val`, can be owned or borrowed. The function doesn't
/// allocate the return value unless there's need to escape the input.
pub fn dn_escape<'a, S: Into<Cow<'a, str>>>(val: S) -> Cow<'a, str> {
    #[inline]
    fn always_escape(c: u8) -> bool {
        c == b'"'
            || c == b'+'
            || c == b','
            || c == b';'
            || c == b'<'
            || c == b'='
            || c == b'>'
            || c == b'\\'
            || c == 0
    }

    #[inline]
    fn escape_leading(c: u8) -> bool {
        c == b' ' || c == b'#'
    }

    #[inline]
    fn escape_trailing(c: u8) -> bool {
        c == b' '
    }

    #[inline]
    fn xdigit(c: u8) -> u8 {
        c + if c < 10 { b'0' } else { b'a' - 10 }
    }

    let val = val.into();
    let mut output = None;
    for (i, &c) in val.as_bytes().iter().enumerate() {
        if always_escape(c)
            || i == 0 && escape_leading(c)
            || i + 1 == val.len() && escape_trailing(c)
        {
            if output.is_none() {
                output = Some(Vec::with_capacity(val.len() + 12)); // guess: up to 4 escaped chars
                output.as_mut().unwrap().extend(val[..i].as_bytes());
            }
            let output = output.as_mut().unwrap();
            output.push(b'\\');
            output.push(xdigit(c >> 4));
            output.push(xdigit(c & 0xF));
        } else if let Some(ref mut output) = output {
            output.push(c);
        }
    }
    if let Some(output) = output {
        // see the rationale for the same construct in ldap_escape()
        Cow::Owned(unsafe { String::from_utf8_unchecked(output) })
    } else {
        val
    }
}

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

    #[test]
    fn dn_esc_leading_space() {
        assert_eq!(dn_escape(" foo"), "\\20foo");
    }

    #[test]
    fn dn_esc_trailing_space() {
        assert_eq!(dn_escape("foo "), "foo\\20");
    }

    #[test]
    fn dn_esc_inner_space() {
        assert_eq!(dn_escape("f o o"), "f o o");
    }

    #[test]
    fn dn_esc_single_space() {
        assert_eq!(dn_escape(" "), "\\20");
    }

    #[test]
    fn dn_esc_two_spaces() {
        assert_eq!(dn_escape("  "), "\\20\\20");
    }

    #[test]
    fn dn_esc_three_spaces() {
        assert_eq!(dn_escape("   "), "\\20 \\20");
    }

    #[test]
    fn dn_esc_leading_hash() {
        assert_eq!(dn_escape("#rust"), "\\23rust");
    }
}
<commit_msg>Eliminate unsafe string conversions<commit_after>use std::borrow::Cow;

/// Escape a filter literal.
///
/// Literal values appearing in an LDAP filter can contain any character,
/// but some characters (parentheses, asterisk, backslash, NUL) must be
/// escaped in the filter's string representation. This function does the
/// escaping.
///
/// The argument, `lit`, can be owned or borrowed. The function doesn't
/// allocate the return value unless there's need to escape the input.
pub fn ldap_escape<'a, S: Into<Cow<'a, str>>>(lit: S) -> Cow<'a, str> {
    #[inline]
    fn needs_escape(c: u8) -> bool {
        c == b'\\' || c == b'*' || c == b'(' || c == b')' || c == 0
    }

    #[inline]
    fn xdigit(c: u8) -> u8 {
        c + if c < 10 { b'0' } else { b'a' - 10 }
    }

    let lit = lit.into();
    let mut output = None;
    for (i, &c) in lit.as_bytes().iter().enumerate() {
        if needs_escape(c) {
            if output.is_none() {
                output = Some(Vec::with_capacity(lit.len() + 12)); // guess: up to 4 escaped chars
                output.as_mut().unwrap().extend(lit[..i].as_bytes());
            }
            let output = output.as_mut().unwrap();
            output.push(b'\\');
            output.push(xdigit(c >> 4));
            output.push(xdigit(c & 0xF));
        } else if let Some(ref mut output) = output {
            output.push(c);
        }
    }
    if let Some(output) = output {
        Cow::Owned(String::from_utf8(output).expect("ldap escaped"))
    } else {
        lit
    }
}

/// Escape an attribute value in a relative distinguished name (RDN).
///
/// When a literal string is used to represent an attribute value in an RDN,
/// some of its characters might need to be escaped according to the rules
/// of [RFC 4514](https://tools.ietf.org/html/rfc4514).
///
/// The function is named `dn_escape()` instead of `rdn_escape()` because of
/// a long-standing association of its intended use with the handling of DNs.
///
/// The argument, `val`, can be owned or borrowed. The function doesn't
/// allocate the return value unless there's need to escape the input.
pub fn dn_escape<'a, S: Into<Cow<'a, str>>>(val: S) -> Cow<'a, str> {
    #[inline]
    fn always_escape(c: u8) -> bool {
        c == b'"'
            || c == b'+'
            || c == b','
            || c == b';'
            || c == b'<'
            || c == b'='
            || c == b'>'
            || c == b'\\'
            || c == 0
    }

    #[inline]
    fn escape_leading(c: u8) -> bool {
        c == b' ' || c == b'#'
    }

    #[inline]
    fn escape_trailing(c: u8) -> bool {
        c == b' '
    }

    #[inline]
    fn xdigit(c: u8) -> u8 {
        c + if c < 10 { b'0' } else { b'a' - 10 }
    }

    let val = val.into();
    let mut output = None;
    for (i, &c) in val.as_bytes().iter().enumerate() {
        if always_escape(c)
            || i == 0 && escape_leading(c)
            || i + 1 == val.len() && escape_trailing(c)
        {
            if output.is_none() {
                output = Some(Vec::with_capacity(val.len() + 12)); // guess: up to 4 escaped chars
                output.as_mut().unwrap().extend(val[..i].as_bytes());
            }
            let output = output.as_mut().unwrap();
            output.push(b'\\');
            output.push(xdigit(c >> 4));
            output.push(xdigit(c & 0xF));
        } else if let Some(ref mut output) = output {
            output.push(c);
        }
    }
    if let Some(output) = output {
        Cow::Owned(String::from_utf8(output).expect("dn escaped"))
    } else {
        val
    }
}

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

    #[test]
    fn dn_esc_leading_space() {
        assert_eq!(dn_escape(" foo"), "\\20foo");
    }

    #[test]
    fn dn_esc_trailing_space() {
        assert_eq!(dn_escape("foo "), "foo\\20");
    }

    #[test]
    fn dn_esc_inner_space() {
        assert_eq!(dn_escape("f o o"), "f o o");
    }

    #[test]
    fn dn_esc_single_space() {
        assert_eq!(dn_escape(" "), "\\20");
    }

    #[test]
    fn dn_esc_two_spaces() {
        assert_eq!(dn_escape("  "), "\\20\\20");
    }

    #[test]
    fn dn_esc_three_spaces() {
        assert_eq!(dn_escape("   "), "\\20 \\20");
    }

    #[test]
    fn dn_esc_leading_hash() {
        assert_eq!(dn_escape("#rust"), "\\23rust");
    }
}
<|endoftext|>