<commit_before>use crate::core::shell::Verbosity;
use crate::core::{GitReference, Workspace};
use crate::ops;
use crate::sources::path::PathSource;
use crate::util::Sha256;
use crate::util::{paths, CargoResult, CargoResultExt, Config};
use anyhow::bail;
use serde::Serialize;
use std::collections::HashSet;
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::fs;
use std::path::{Path, PathBuf};

pub struct VendorOptions<'a> {
    pub no_delete: bool,
    pub versioned_dirs: bool,
    pub destination: &'a Path,
    pub extra: Vec<PathBuf>,
}

pub fn vendor(ws: &Workspace<'_>, opts: &VendorOptions<'_>) -> CargoResult<()> {
    let config = ws.config();
    let mut extra_workspaces = Vec::new();
    for extra in opts.extra.iter() {
        let extra = config.cwd().join(extra);
        let ws = Workspace::new(&extra, config)?;
        extra_workspaces.push(ws);
    }
    let workspaces = extra_workspaces.iter().chain(Some(ws)).collect::<Vec<_>>();
    let vendor_config =
        sync(config, &workspaces, opts).chain_err(|| "failed to sync".to_string())?;

    if config.shell().verbosity() != Verbosity::Quiet {
        crate::drop_eprint!(
            config,
            "To use vendored sources, add this to your .cargo/config for this project:\n\n"
        );
        crate::drop_print!(config, "{}", &toml::to_string(&vendor_config).unwrap());
    }

    Ok(())
}

#[derive(Serialize)]
struct VendorConfig {
    source: BTreeMap<String, VendorSource>,
}

#[derive(Serialize)]
#[serde(rename_all = "lowercase", untagged)]
enum VendorSource {
    Directory {
        directory: PathBuf,
    },
    Registry {
        registry: Option<String>,
        #[serde(rename = "replace-with")]
        replace_with: String,
    },
    Git {
        git: String,
        branch: Option<String>,
        tag: Option<String>,
        rev: Option<String>,
        #[serde(rename = "replace-with")]
        replace_with: String,
    },
}

fn sync(
    config: &Config,
    workspaces: &[&Workspace<'_>],
    opts: &VendorOptions<'_>,
) -> CargoResult<VendorConfig> {
    let canonical_destination = opts.destination.canonicalize();
    let canonical_destination = canonical_destination
        .as_ref()
        .map(|p| &**p)
        .unwrap_or(opts.destination);

    paths::create_dir_all(&canonical_destination)?;
    let mut to_remove = HashSet::new();
    if !opts.no_delete {
        for entry in canonical_destination.read_dir()? {
            let entry = entry?;
            if !entry
                .file_name()
                .to_str()
                .map_or(false, |s| s.starts_with('.'))
            {
                to_remove.insert(entry.path());
            }
        }
    }

    // First up attempt to work around rust-lang/cargo#5956. Apparently build
    // artifacts sprout up in Cargo's global cache for whatever reason, although
    // it's unsure what tool is causing these issues at this time. For now we
    // apply a heavy-hammer approach which is to delete Cargo's unpacked version
    // of each crate to start off with. After we do this we'll re-resolve and
    // redownload again, which should trigger Cargo to re-extract all the
    // crates.
    //
    // Note that errors are largely ignored here as this is a best-effort
    // attempt. If anything fails here we basically just move on to the next
    // crate to work with.
    for ws in workspaces {
        let (packages, resolve) =
            ops::resolve_ws(ws).chain_err(|| "failed to load pkg lockfile")?;

        packages
            .get_many(resolve.iter())
            .chain_err(|| "failed to download packages")?;

        for pkg in resolve.iter() {
            // Don't delete actual source code!
            if pkg.source_id().is_path() {
                if let Ok(path) = pkg.source_id().url().to_file_path() {
                    if let Ok(path) = path.canonicalize() {
                        to_remove.remove(&path);
                    }
                }
                continue;
            }
            if pkg.source_id().is_git() {
                continue;
            }
            if let Ok(pkg) = packages.get_one(pkg) {
                drop(fs::remove_dir_all(pkg.manifest_path().parent().unwrap()));
            }
        }
    }

    let mut checksums = HashMap::new();
    let mut ids = BTreeMap::new();

    // Next up let's actually download all crates and start storing internal
    // tables about them.
    for ws in workspaces {
        let (packages, resolve) =
            ops::resolve_ws(ws).chain_err(|| "failed to load pkg lockfile")?;

        packages
            .get_many(resolve.iter())
            .chain_err(|| "failed to download packages")?;

        for pkg in resolve.iter() {
            // No need to vendor path crates since they're already in the
            // repository
            if pkg.source_id().is_path() {
                continue;
            }
            ids.insert(
                pkg,
                packages
                    .get_one(pkg)
                    .chain_err(|| "failed to fetch package")?
                    .clone(),
            );

            checksums.insert(pkg, resolve.checksums().get(&pkg).cloned());
        }
    }

    let mut versions = HashMap::new();
    for id in ids.keys() {
        let map = versions.entry(id.name()).or_insert_with(BTreeMap::default);
        if let Some(prev) = map.get(&id.version()) {
            bail!(
                "found duplicate version of package `{} v{}` \
                 vendored from two sources:\n\
                 \n\
                 \tsource 1: {}\n\
                 \tsource 2: {}",
                id.name(),
                id.version(),
                prev,
                id.source_id()
            );
        }
        map.insert(id.version(), id.source_id());
    }

    let mut sources = BTreeSet::new();
    for (id, pkg) in ids.iter() {
        // Next up, copy it to the vendor directory
        let src = pkg
            .manifest_path()
            .parent()
            .expect("manifest_path should point to a file");
        let max_version = *versions[&id.name()].iter().rev().next().unwrap().0;
        let dir_has_version_suffix = opts.versioned_dirs || id.version() != max_version;
        let dst_name = if dir_has_version_suffix {
            // Eg vendor/futures-0.1.13
            format!("{}-{}", id.name(), id.version())
        } else {
            // Eg vendor/futures
            id.name().to_string()
        };

        sources.insert(id.source_id());
        let dst = canonical_destination.join(&dst_name);
        to_remove.remove(&dst);
        let cksum = dst.join(".cargo-checksum.json");
        if dir_has_version_suffix && cksum.exists() {
            // Always re-copy directory without version suffix in case the version changed
            continue;
        }

        config.shell().status(
            "Vendoring",
            &format!("{} ({}) to {}", id, src.to_string_lossy(), dst.display()),
        )?;

        let _ = fs::remove_dir_all(&dst);
        let pathsource = PathSource::new(src, id.source_id(), config);
        let paths = pathsource.list_files(pkg)?;
        let mut map = BTreeMap::new();
        cp_sources(src, &paths, &dst, &mut map)
            .chain_err(|| format!("failed to copy over vendored sources for: {}", id))?;

        // Finally, emit the metadata about this package
        let json = serde_json::json!({
            "package": checksums.get(id),
            "files": map,
        });

        paths::write(&cksum, json.to_string())?;
    }

    for path in to_remove {
        if path.is_dir() {
            paths::remove_dir_all(&path)?;
        } else {
            paths::remove_file(&path)?;
        }
    }

    // add our vendored source
    let mut config = BTreeMap::new();

    let merged_source_name = "vendored-sources";
    config.insert(
        merged_source_name.to_string(),
        VendorSource::Directory {
            directory: opts.destination.to_path_buf(),
        },
    );

    // replace original sources with vendor
    for source_id in sources {
        let name = if source_id.is_default_registry() {
            "crates-io".to_string()
        } else {
            source_id.url().to_string()
        };

        let source = if source_id.is_default_registry() {
            VendorSource::Registry {
                registry: None,
                replace_with: merged_source_name.to_string(),
            }
        } else if source_id.is_remote_registry() {
            let registry = source_id.url().to_string();
            VendorSource::Registry {
                registry: Some(registry),
                replace_with: merged_source_name.to_string(),
            }
        } else if source_id.is_git() {
            let mut branch = None;
            let mut tag = None;
            let mut rev = None;
            if let Some(reference) = source_id.git_reference() {
                match *reference {
                    GitReference::Branch(ref b) => branch = Some(b.clone()),
                    GitReference::Tag(ref t) => tag = Some(t.clone()),
                    GitReference::Rev(ref r) => rev = Some(r.clone()),
                    GitReference::DefaultBranch => {}
                }
            }
            VendorSource::Git {
                git: source_id.url().to_string(),
                branch,
                tag,
                rev,
                replace_with: merged_source_name.to_string(),
            }
        } else {
            panic!("Invalid source ID: {}", source_id)
        };
        config.insert(name, source);
    }

    Ok(VendorConfig { source: config })
}

fn cp_sources(
    src: &Path,
    paths: &[PathBuf],
    dst: &Path,
    cksums: &mut BTreeMap<String, String>,
) -> CargoResult<()> {
    for p in paths {
        let relative = p.strip_prefix(&src).unwrap();

        match relative.to_str() {
            // Skip git config files as they're not relevant to builds most of
            // the time and if we respect them (e.g.  in git) then it'll
            // probably mess with the checksums when a vendor dir is checked
            // into someone else's source control
            Some(".gitattributes") | Some(".gitignore") | Some(".git") => continue,

            // Temporary Cargo files
            Some(".cargo-ok") => continue,

            // Skip patch-style orig/rej files. Published crates on crates.io
            // have `Cargo.toml.orig` which we don't want to use here and
            // otherwise these are rarely used as part of the build process.
            Some(filename) => {
                if filename.ends_with(".orig") || filename.ends_with(".rej") {
                    continue;
                }
            }
            _ => {}
        };

        // Join pathname components individually to make sure that the joined
        // path uses the correct directory separators everywhere, since
        // `relative` may use Unix-style and `dst` may require Windows-style
        // backslashes.
        let dst = relative
            .iter()
            .fold(dst.to_owned(), |acc, component| acc.join(&component));

        paths::create_dir_all(dst.parent().unwrap())?;

        paths::copy(&p, &dst)?;
        let cksum = Sha256::new().update_path(dst)?.finish_hex();
        cksums.insert(relative.to_str().unwrap().replace("\\", "/"), cksum);
    }
    Ok(())
}
<commit_msg>Auto merge of #8822 - skrap:master, r=ehuss<commit_after>use crate::core::shell::Verbosity;
use crate::core::{GitReference, Workspace};
use crate::ops;
use crate::sources::path::PathSource;
use crate::util::Sha256;
use crate::util::{paths, CargoResult, CargoResultExt, Config};
use anyhow::bail;
use serde::Serialize;
use std::collections::HashSet;
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::fs;
use std::path::{Path, PathBuf};

pub struct VendorOptions<'a> {
    pub no_delete: bool,
    pub versioned_dirs: bool,
    pub destination: &'a Path,
    pub extra: Vec<PathBuf>,
}

pub fn vendor(ws: &Workspace<'_>, opts: &VendorOptions<'_>) -> CargoResult<()> {
    let config = ws.config();
    let mut extra_workspaces = Vec::new();
    for extra in opts.extra.iter() {
        let extra = config.cwd().join(extra);
        let ws = Workspace::new(&extra, config)?;
        extra_workspaces.push(ws);
    }
    let workspaces = extra_workspaces.iter().chain(Some(ws)).collect::<Vec<_>>();
    let vendor_config =
        sync(config, &workspaces, opts).chain_err(|| "failed to sync".to_string())?;

    if config.shell().verbosity() != Verbosity::Quiet {
        crate::drop_eprint!(
            config,
            "To use vendored sources, add this to your .cargo/config.toml for this project:\n\n"
        );
        crate::drop_print!(config, "{}", &toml::to_string(&vendor_config).unwrap());
    }

    Ok(())
}

#[derive(Serialize)]
struct VendorConfig {
    source: BTreeMap<String, VendorSource>,
}

#[derive(Serialize)]
#[serde(rename_all = "lowercase", untagged)]
enum VendorSource {
    Directory {
        directory: PathBuf,
    },
    Registry {
        registry: Option<String>,
        #[serde(rename = "replace-with")]
        replace_with: String,
    },
    Git {
        git: String,
        branch: Option<String>,
        tag: Option<String>,
        rev: Option<String>,
        #[serde(rename = "replace-with")]
        replace_with: String,
    },
}

fn sync(
    config: &Config,
    workspaces: &[&Workspace<'_>],
    opts: &VendorOptions<'_>,
) -> CargoResult<VendorConfig> {
    let canonical_destination = opts.destination.canonicalize();
    let canonical_destination = canonical_destination
        .as_ref()
        .map(|p| &**p)
        .unwrap_or(opts.destination);

    paths::create_dir_all(&canonical_destination)?;
    let mut to_remove = HashSet::new();
    if !opts.no_delete {
        for entry in canonical_destination.read_dir()? {
            let entry = entry?;
            if !entry
                .file_name()
                .to_str()
                .map_or(false, |s| s.starts_with('.'))
            {
                to_remove.insert(entry.path());
            }
        }
    }

    // First up attempt to work around rust-lang/cargo#5956. Apparently build
    // artifacts sprout up in Cargo's global cache for whatever reason, although
    // it's unsure what tool is causing these issues at this time. For now we
    // apply a heavy-hammer approach which is to delete Cargo's unpacked version
    // of each crate to start off with. After we do this we'll re-resolve and
    // redownload again, which should trigger Cargo to re-extract all the
    // crates.
    //
    // Note that errors are largely ignored here as this is a best-effort
    // attempt. If anything fails here we basically just move on to the next
    // crate to work with.
    for ws in workspaces {
        let (packages, resolve) =
            ops::resolve_ws(ws).chain_err(|| "failed to load pkg lockfile")?;

        packages
            .get_many(resolve.iter())
            .chain_err(|| "failed to download packages")?;

        for pkg in resolve.iter() {
            // Don't delete actual source code!
            if pkg.source_id().is_path() {
                if let Ok(path) = pkg.source_id().url().to_file_path() {
                    if let Ok(path) = path.canonicalize() {
                        to_remove.remove(&path);
                    }
                }
                continue;
            }
            if pkg.source_id().is_git() {
                continue;
            }
            if let Ok(pkg) = packages.get_one(pkg) {
                drop(fs::remove_dir_all(pkg.manifest_path().parent().unwrap()));
            }
        }
    }

    let mut checksums = HashMap::new();
    let mut ids = BTreeMap::new();

    // Next up let's actually download all crates and start storing internal
    // tables about them.
    for ws in workspaces {
        let (packages, resolve) =
            ops::resolve_ws(ws).chain_err(|| "failed to load pkg lockfile")?;

        packages
            .get_many(resolve.iter())
            .chain_err(|| "failed to download packages")?;

        for pkg in resolve.iter() {
            // No need to vendor path crates since they're already in the
            // repository
            if pkg.source_id().is_path() {
                continue;
            }
            ids.insert(
                pkg,
                packages
                    .get_one(pkg)
                    .chain_err(|| "failed to fetch package")?
                    .clone(),
            );

            checksums.insert(pkg, resolve.checksums().get(&pkg).cloned());
        }
    }

    let mut versions = HashMap::new();
    for id in ids.keys() {
        let map = versions.entry(id.name()).or_insert_with(BTreeMap::default);
        if let Some(prev) = map.get(&id.version()) {
            bail!(
                "found duplicate version of package `{} v{}` \
                 vendored from two sources:\n\
                 \n\
                 \tsource 1: {}\n\
                 \tsource 2: {}",
                id.name(),
                id.version(),
                prev,
                id.source_id()
            );
        }
        map.insert(id.version(), id.source_id());
    }

    let mut sources = BTreeSet::new();
    for (id, pkg) in ids.iter() {
        // Next up, copy it to the vendor directory
        let src = pkg
            .manifest_path()
            .parent()
            .expect("manifest_path should point to a file");
        let max_version = *versions[&id.name()].iter().rev().next().unwrap().0;
        let dir_has_version_suffix = opts.versioned_dirs || id.version() != max_version;
        let dst_name = if dir_has_version_suffix {
            // Eg vendor/futures-0.1.13
            format!("{}-{}", id.name(), id.version())
        } else {
            // Eg vendor/futures
            id.name().to_string()
        };

        sources.insert(id.source_id());
        let dst = canonical_destination.join(&dst_name);
        to_remove.remove(&dst);
        let cksum = dst.join(".cargo-checksum.json");
        if dir_has_version_suffix && cksum.exists() {
            // Always re-copy directory without version suffix in case the version changed
            continue;
        }

        config.shell().status(
            "Vendoring",
            &format!("{} ({}) to {}", id, src.to_string_lossy(), dst.display()),
        )?;

        let _ = fs::remove_dir_all(&dst);
        let pathsource = PathSource::new(src, id.source_id(), config);
        let paths = pathsource.list_files(pkg)?;
        let mut map = BTreeMap::new();
        cp_sources(src, &paths, &dst, &mut map)
            .chain_err(|| format!("failed to copy over vendored sources for: {}", id))?;

        // Finally, emit the metadata about this package
        let json = serde_json::json!({
            "package": checksums.get(id),
            "files": map,
        });

        paths::write(&cksum, json.to_string())?;
    }

    for path in to_remove {
        if path.is_dir() {
            paths::remove_dir_all(&path)?;
        } else {
            paths::remove_file(&path)?;
        }
    }

    // add our vendored source
    let mut config = BTreeMap::new();

    let merged_source_name = "vendored-sources";
    config.insert(
        merged_source_name.to_string(),
        VendorSource::Directory {
            directory: opts.destination.to_path_buf(),
        },
    );

    // replace original sources with vendor
    for source_id in sources {
        let name = if source_id.is_default_registry() {
            "crates-io".to_string()
        } else {
            source_id.url().to_string()
        };

        let source = if source_id.is_default_registry() {
            VendorSource::Registry {
                registry: None,
                replace_with: merged_source_name.to_string(),
            }
        } else if source_id.is_remote_registry() {
            let registry = source_id.url().to_string();
            VendorSource::Registry {
                registry: Some(registry),
                replace_with: merged_source_name.to_string(),
            }
        } else if source_id.is_git() {
            let mut branch = None;
            let mut tag = None;
            let mut rev = None;
            if let Some(reference) = source_id.git_reference() {
                match *reference {
                    GitReference::Branch(ref b) => branch = Some(b.clone()),
                    GitReference::Tag(ref t) => tag = Some(t.clone()),
                    GitReference::Rev(ref r) => rev = Some(r.clone()),
                    GitReference::DefaultBranch => {}
                }
            }
            VendorSource::Git {
                git: source_id.url().to_string(),
                branch,
                tag,
                rev,
                replace_with: merged_source_name.to_string(),
            }
        } else {
            panic!("Invalid source ID: {}", source_id)
        };
        config.insert(name, source);
    }

    Ok(VendorConfig { source: config })
}

fn cp_sources(
    src: &Path,
    paths: &[PathBuf],
    dst: &Path,
    cksums: &mut BTreeMap<String, String>,
) -> CargoResult<()> {
    for p in paths {
        let relative = p.strip_prefix(&src).unwrap();

        match relative.to_str() {
            // Skip git config files as they're not relevant to builds most of
            // the time and if we respect them (e.g.  in git) then it'll
            // probably mess with the checksums when a vendor dir is checked
            // into someone else's source control
            Some(".gitattributes") | Some(".gitignore") | Some(".git") => continue,

            // Temporary Cargo files
            Some(".cargo-ok") => continue,

            // Skip patch-style orig/rej files. Published crates on crates.io
            // have `Cargo.toml.orig` which we don't want to use here and
            // otherwise these are rarely used as part of the build process.
            Some(filename) => {
                if filename.ends_with(".orig") || filename.ends_with(".rej") {
                    continue;
                }
            }
            _ => {}
        };

        // Join pathname components individually to make sure that the joined
        // path uses the correct directory separators everywhere, since
        // `relative` may use Unix-style and `dst` may require Windows-style
        // backslashes.
        let dst = relative
            .iter()
            .fold(dst.to_owned(), |acc, component| acc.join(&component));

        paths::create_dir_all(dst.parent().unwrap())?;

        paths::copy(&p, &dst)?;
        let cksum = Sha256::new().update_path(dst)?.finish_hex();
        cksums.insert(relative.to_str().unwrap().replace("\\", "/"), cksum);
    }
    Ok(())
}
<|endoftext|>


<commit_before>use super::{directory_stack::DirectoryStack, variables::Variables};
use glob::glob;
use liner::{Completer, FilenameCompleter};

/// Performs escaping to an inner `FilenameCompleter` to enable a handful of special cases
/// needed by the shell, such as expanding '~' to a home directory, or adding a backslash
/// when a special character is contained within an expanded filename.
pub(crate) struct IonFileCompleter {
    /// The completer that this completer is  handling.
    inner: FilenameCompleter,
    /// A pointer to the directory stack in the shell.
    dir_stack: *const DirectoryStack,
    /// A pointer to the variables map in the shell.
    vars: *const Variables,
}

impl IonFileCompleter {
    pub(crate) fn new(
        path: Option<&str>,
        dir_stack: *const DirectoryStack,
        vars: *const Variables,
    ) -> IonFileCompleter {
        IonFileCompleter {
            inner: FilenameCompleter::new(path),
            dir_stack,
            vars,
        }
    }
}

impl Completer for IonFileCompleter {
    /// When the tab key is pressed, **Liner** will use this method to perform completions of
    /// filenames. As our `IonFileCompleter` is a wrapper around **Liner**'s
    /// `FilenameCompleter`,
    /// the purpose of our custom `Completer` is to expand possible `~` characters in the
    /// `start`
    /// value that we receive from the prompt, grab completions from the inner
    /// `FilenameCompleter`,
    /// and then escape the resulting filenames, as well as remove the expanded form of the `~`
    /// character and re-add the `~` character in it's place.
    fn completions(&self, start: &str) -> Vec<String> {
        // Only if the first character is a tilde character will we perform expansions
        if start.starts_with('~') {
            // Dereferencing the raw pointers here should be entirely safe, theoretically,
            // because no changes will occur to either of the underlying references in the
            // duration between creation of the completers and execution of their
            // completions.
            if let Some(expanded) = unsafe { (*self.vars).tilde_expansion(start, &*self.dir_stack) }
            {
                // Now we obtain completions for the `expanded` form of the `start` value.
                let completions = filename_completion(&expanded, |x| self.inner.completions(x));
                let mut iterator = completions.iter();

                // And then we will need to take those completions and remove the expanded form
                // of the tilde pattern and replace it with that pattern yet again.
                let mut completions = Vec::new();

                // We can do that by obtaining the index position where the tilde character
                // ends. We don't search with `~` because we also want to
                // handle other tilde variants.
                let t_index = start.find('/').unwrap_or(1);
                // `tilde` is the tilde pattern, and `search` is the pattern that follows.
                let (tilde, search) = start.split_at(t_index as usize);

                if search.len() < 2 {
                    // If the length of the search pattern is less than 2, the search pattern is
                    // empty, and thus the completions actually contain files and directories in
                    // the home directory.

                    // The tilde pattern will actually be our `start` command in itself,
                    // and the completed form will be all of the characters beyond the length of
                    // the expanded form of the tilde pattern.
                    for completion in iterator {
                        completions.push([start, &completion[expanded.len()..]].concat());
                    }
                } else {
                    // To save processing time, we should get obtain the index position where our
                    // search pattern begins, and re-use that index to slice the completions so
                    // that we may re-add the tilde character with the completion that follows.
                    if let Some(completion) = iterator.next() {
                        if let Some(e_index) = expanded.rfind(search) {
                            completions.push(escape(&[tilde, &completion[e_index..]].concat()));
                            for completion in iterator {
                                let expanded = &completion[e_index..];
                                completions.push(escape(&[tilde, expanded].concat()));
                            }
                        }
                    }
                }

                return completions;
            }
        }

        filename_completion(&start, |x| self.inner.completions(x))
    }
}

fn filename_completion<LC>(start: &str, liner_complete: LC) -> Vec<String>
where
    LC: Fn(&str) -> Vec<String>,
{
    let unescaped_start = unescape(start);

    let start_split: Vec<&str> = unescaped_start.split("/").collect();

    // When 'start' is an absolute path, "/..." gets split to ["", "..."]
    // So we ignore the first element and add "/" to the start of the string
    let start_for_glob = match unescaped_start.starts_with("/") {
        true => ["/", &start_split[1..].join("*/"), "*"].concat(),
        false => [&start_split.join("*/"), "*"].concat(),
    };

    let mut inner_glob: Vec<String> = match glob(&start_for_glob) {
        Ok(completions) => completions
            .filter_map(Result::ok)
            .map(|x| x.to_string_lossy().into_owned())
            .collect(),
        _ => vec![],
    };
    if inner_glob.len() == 0 {
        inner_glob.push(escape(&start.to_string()));
    }

    let mut completions = vec![];

    // Use Liner::Completer as well, to preserve the previous behaviour
    // around single-directory completions
    for path in inner_glob {
        let liner_completions: Vec<String> = liner_complete(&path)
            .iter()
            .map(|x| escape(x.as_str()))
            .collect();
        for c in liner_completions {
            completions.push(c)
        }
    }
    completions
}

/// Escapes filenames from the completer so that special characters will be properly escaped.
///
/// NOTE: Perhaps we should submit a PR to Liner to add a &'static [u8] field to
/// `FilenameCompleter` so that we don't have to perform the escaping ourselves?
fn escape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    for character in input.bytes() {
        match character {
            b'(' | b')' | b'[' | b']' | b'&' | b'$' | b'@' | b'{' | b'}' | b'<' | b'>' | b';'
            | b'"' | b'\'' | b'#' | b'^' | b'*' => output.push(b'\\'),
            _ => (),
        }
        output.push(character);
    }
    unsafe { String::from_utf8_unchecked(output) }
}

/// Unescapes filenames to be passed into the completer
fn unescape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    let mut bytes = input.bytes();
    while let Some(b) = bytes.next() {
        match b {
            b'\\' => if let Some(next) = bytes.next() {
                output.push(next);
            } else {
                output.push(b'\\')
            },
            _ => output.push(b),
        }
    }
    unsafe { String::from_utf8_unchecked(output) }
}

/// A completer that combines suggestions from multiple completers.
#[derive(Clone, Eq, PartialEq)]
pub(crate) struct MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    a: Vec<A>,
    b: B,
}

impl<A, B> MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    pub(crate) fn new(a: Vec<A>, b: B) -> MultiCompleter<A, B> { MultiCompleter { a, b } }
}

impl<A, B> Completer for MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    fn completions(&self, start: &str) -> Vec<String> {
        let mut completions = self.b.completions(start);
        for x in &self.a {
            completions.extend_from_slice(&x.completions(start));
        }
        completions
    }
}

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

    #[test]
    fn filename_completion() {
        let current_dir = env::current_dir().expect("Unable to get current directory");

        let completer = IonFileCompleter::new(
            current_dir.to_str(),
            &DirectoryStack::new(),
            &Variables::default(),
        );
        assert_eq!(completer.completions("testing"), vec!["testing/"]);
        assert_eq!(
            completer.completions("testing/file"),
            vec!["testing/file_with_text"]
        );

        assert_eq!(completer.completions("~"), vec!["~/"]);

        assert_eq!(
            completer.completions("tes/fil"),
            vec!["testing/file_with_text"]
        );
    }
}
<commit_msg>First step to refactoring completer.rs (#742)<commit_after>use super::{directory_stack::DirectoryStack, variables::Variables};
use glob::glob;
use liner::{Completer, FilenameCompleter};

/// Performs escaping to an inner `FilenameCompleter` to enable a handful of special cases
/// needed by the shell, such as expanding '~' to a home directory, or adding a backslash
/// when a special character is contained within an expanded filename.
pub(crate) struct IonFileCompleter {
    /// The completer that this completer is  handling.
    inner: FilenameCompleter,
    /// A pointer to the directory stack in the shell.
    dir_stack: *const DirectoryStack,
    /// A pointer to the variables map in the shell.
    vars: *const Variables,
}

impl IonFileCompleter {
    pub(crate) fn new(
        path: Option<&str>,
        dir_stack: *const DirectoryStack,
        vars: *const Variables,
    ) -> IonFileCompleter {
        IonFileCompleter {
            inner: FilenameCompleter::new(path),
            dir_stack,
            vars,
        }
    }
}

impl Completer for IonFileCompleter {
    /// When the tab key is pressed, **Liner** will use this method to perform completions of
    /// filenames. As our `IonFileCompleter` is a wrapper around **Liner**'s
    /// `FilenameCompleter`,
    /// the purpose of our custom `Completer` is to expand possible `~` characters in the
    /// `start`
    /// value that we receive from the prompt, grab completions from the inner
    /// `FilenameCompleter`,
    /// and then escape the resulting filenames, as well as remove the expanded form of the `~`
    /// character and re-add the `~` character in it's place.
    fn completions(&self, start: &str) -> Vec<String> {
        // Only if the first character is a tilde character will we perform expansions
        if start.starts_with('~') {
            // Dereferencing the raw pointers here should be entirely safe, theoretically,
            // because no changes will occur to either of the underlying references in the
            // duration between creation of the completers and execution of their
            // completions.
            if let Some(expanded) = unsafe { (*self.vars).tilde_expansion(start, &*self.dir_stack) }
            {
                // Now we obtain completions for the `expanded` form of the `start` value.
                let completions = filename_completion(&expanded, |x| self.inner.completions(x));
                let mut iterator = completions.iter();

                // And then we will need to take those completions and remove the expanded form
                // of the tilde pattern and replace it with that pattern yet again.
                let mut completions = Vec::new();

                // We can do that by obtaining the index position where the tilde character
                // ends. We don't search with `~` because we also want to
                // handle other tilde variants.
                let t_index = start.find('/').unwrap_or(1);
                // `tilde` is the tilde pattern, and `search` is the pattern that follows.
                let (tilde, search) = start.split_at(t_index as usize);

                if search.len() < 2 {
                    // If the length of the search pattern is less than 2, the search pattern is
                    // empty, and thus the completions actually contain files and directories in
                    // the home directory.

                    // The tilde pattern will actually be our `start` command in itself,
                    // and the completed form will be all of the characters beyond the length of
                    // the expanded form of the tilde pattern.
                    for completion in iterator {
                        completions.push([start, &completion[expanded.len()..]].concat());
                    }
                } else {
                    // To save processing time, we should get obtain the index position where our
                    // search pattern begins, and re-use that index to slice the completions so
                    // that we may re-add the tilde character with the completion that follows.
                    if let Some(completion) = iterator.next() {
                        if let Some(e_index) = expanded.rfind(search) {
                            completions.push(escape(&[tilde, &completion[e_index..]].concat()));
                            for completion in iterator {
                                let expanded = &completion[e_index..];
                                completions.push(escape(&[tilde, expanded].concat()));
                            }
                        }
                    }
                }

                return completions;
            }
        }

        filename_completion(&start, |x| self.inner.completions(x))
    }
}

fn filename_completion<LC>(start: &str, liner_complete: LC) -> Vec<String>
where
    LC: Fn(&str) -> Vec<String>,
{
    let unescaped_start = unescape(start);

    let split_start = unescaped_start.split("/");
    let mut string = String::with_capacity(5);

    // When 'start' is an absolute path, "/..." gets split to ["", "..."]
    // So we skip the first element and add "/" to the start of the string
    let mut start_for_glob = if unescaped_start.starts_with("/") {
        string.push('/');
        split_start.skip(1).fold(string, |mut state, element| {
            state.push_str(element);
            state.push_str("*/");
            state
        })
    } else {
        split_start.fold(string, |mut state, element| {
            state.push_str(element);
            state.push_str("*/");
            state
        })
    };
    start_for_glob.pop(); // pop out the last '/' character

    let mut inner_glob: Vec<String> = match glob(&start_for_glob) {
        Ok(completions) => completions
            .filter_map(Result::ok)
            .map(|x| x.to_string_lossy().into_owned())
            .collect(),
        _ => vec![],
    };
    if inner_glob.len() == 0 {
        inner_glob.push(escape(&start.to_string()));
    }

    let mut completions = vec![];

    // Use Liner::Completer as well, to preserve the previous behaviour
    // around single-directory completions
    for path in inner_glob {
        let liner_completions: Vec<String> = liner_complete(&path)
            .iter()
            .map(|x| escape(x.as_str()))
            .collect();
        for c in liner_completions {
            completions.push(c)
        }
    }
    completions
}

/// Escapes filenames from the completer so that special characters will be properly escaped.
///
/// NOTE: Perhaps we should submit a PR to Liner to add a &'static [u8] field to
/// `FilenameCompleter` so that we don't have to perform the escaping ourselves?
fn escape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    for character in input.bytes() {
        match character {
            b'(' | b')' | b'[' | b']' | b'&' | b'$' | b'@' | b'{' | b'}' | b'<' | b'>' | b';'
            | b'"' | b'\'' | b'#' | b'^' | b'*' => output.push(b'\\'),
            _ => (),
        }
        output.push(character);
    }
    unsafe { String::from_utf8_unchecked(output) }
}

/// Unescapes filenames to be passed into the completer
fn unescape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    let mut bytes = input.bytes();
    while let Some(b) = bytes.next() {
        match b {
            b'\\' => if let Some(next) = bytes.next() {
                output.push(next);
            } else {
                output.push(b'\\')
            },
            _ => output.push(b),
        }
    }
    unsafe { String::from_utf8_unchecked(output) }
}

/// A completer that combines suggestions from multiple completers.
#[derive(Clone, Eq, PartialEq)]
pub(crate) struct MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    a: Vec<A>,
    b: B,
}

impl<A, B> MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    pub(crate) fn new(a: Vec<A>, b: B) -> MultiCompleter<A, B> { MultiCompleter { a, b } }
}

impl<A, B> Completer for MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    fn completions(&self, start: &str) -> Vec<String> {
        let mut completions = self.b.completions(start);
        for x in &self.a {
            completions.extend_from_slice(&x.completions(start));
        }
        completions
    }
}

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

    #[test]
    fn filename_completion() {
        let current_dir = env::current_dir().expect("Unable to get current directory");

        let completer = IonFileCompleter::new(
            current_dir.to_str(),
            &DirectoryStack::new(),
            &Variables::default(),
        );
        assert_eq!(completer.completions("testing"), vec!["testing/"]);
        assert_eq!(
            completer.completions("testing/file"),
            vec!["testing/file_with_text"]
        );

        assert_eq!(completer.completions("~"), vec!["~/"]);

        assert_eq!(
            completer.completions("tes/fil"),
            vec!["testing/file_with_text"]
        );
    }
}
<|endoftext|>


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

#![no_std]
#![allow(unused_attributes)]
#![unstable(feature = "alloc_system",
            reason = "this library is unlikely to be stabilized in its current \
                      form or name",
            issue = "32838")]
#![feature(allocator_api)]
#![feature(core_intrinsics)]
#![feature(staged_api)]
#![feature(rustc_attrs)]
#![cfg_attr(any(unix, target_os = "cloudabi", target_os = "redox"), feature(libc))]
#![rustc_alloc_kind = "lib"]

// The minimum alignment guaranteed by the architecture. This value is used to
// add fast paths for low alignment values.
#[cfg(all(any(target_arch = "x86",
              target_arch = "arm",
              target_arch = "mips",
              target_arch = "powerpc",
              target_arch = "powerpc64",
              target_arch = "asmjs",
              target_arch = "wasm32")))]
#[allow(dead_code)]
const MIN_ALIGN: usize = 8;
#[cfg(all(any(target_arch = "x86_64",
              target_arch = "aarch64",
              target_arch = "mips64",
              target_arch = "s390x",
              target_arch = "sparc64")))]
#[allow(dead_code)]
const MIN_ALIGN: usize = 16;

use core::alloc::{Alloc, GlobalAlloc, AllocErr, Layout};
use core::ptr::NonNull;

/// The default memory allocator provided by the operating system.
///
/// This is based on `malloc` on Unix platforms and `HeapAlloc` on Windows,
/// plus related functions.
///
/// This type can be used in a `static` item
/// with the `#[global_allocator]` attribute
/// to force the global allocator to be the system’s one.
/// (The default is jemalloc for executables, on some platforms.)
///
/// ```rust
/// use std::alloc::System;
///
/// #[global_allocator]
/// static A: System = System;
///
/// fn main() {
///     let a = Box::new(4); // Allocates from the system allocator.
///     println!("{}", a);
/// }
/// ```
///
/// It can also be used directly to allocate memory
/// independently of the standard library’s global allocator.
#[stable(feature = "alloc_system_type", since = "1.28.0")]
pub struct System;

#[unstable(feature = "allocator_api", issue = "32838")]
unsafe impl Alloc for System {
    #[inline]
    unsafe fn alloc(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::alloc(self, layout)).ok_or(AllocErr)
    }

    #[inline]
    unsafe fn alloc_zeroed(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::alloc_zeroed(self, layout)).ok_or(AllocErr)
    }

    #[inline]
    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        GlobalAlloc::dealloc(self, ptr.as_ptr(), layout)
    }

    #[inline]
    unsafe fn realloc(&mut self,
                      ptr: NonNull<u8>,
                      layout: Layout,
                      new_size: usize) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::realloc(self, ptr.as_ptr(), layout, new_size)).ok_or(AllocErr)
    }
}

#[cfg(any(windows, unix, target_os = "cloudabi", target_os = "redox"))]
mod realloc_fallback {
    use core::alloc::{GlobalAlloc, Layout};
    use core::cmp;
    use core::ptr;

    impl super::System {
        pub(crate) unsafe fn realloc_fallback(&self, ptr: *mut u8, old_layout: Layout,
                                              new_size: usize) -> *mut u8 {
            // Docs for GlobalAlloc::realloc require this to be valid:
            let new_layout = Layout::from_size_align_unchecked(new_size, old_layout.align());

            let new_ptr = GlobalAlloc::alloc(self, new_layout);
            if !new_ptr.is_null() {
                let size = cmp::min(old_layout.size(), new_size);
                ptr::copy_nonoverlapping(ptr, new_ptr, size);
                GlobalAlloc::dealloc(self, ptr, old_layout);
            }
            new_ptr
        }
    }
}

#[cfg(any(unix, target_os = "cloudabi", target_os = "redox"))]
mod platform {
    extern crate libc;

    use core::ptr;

    use MIN_ALIGN;
    use System;
    use core::alloc::{GlobalAlloc, Layout};

    #[stable(feature = "alloc_system_type", since = "1.28.0")]
    unsafe impl GlobalAlloc for System {
        #[inline]
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            if layout.align() <= MIN_ALIGN && layout.align() <= layout.size() {
                libc::malloc(layout.size()) as *mut u8
            } else {
                #[cfg(target_os = "macos")]
                {
                    if layout.align() > (1 << 31) {
                        return ptr::null_mut()
                    }
                }
                aligned_malloc(&layout)
            }
        }

        #[inline]
        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
            if layout.align() <= MIN_ALIGN && layout.align() <= layout.size() {
                libc::calloc(layout.size(), 1) as *mut u8
            } else {
                let ptr = self.alloc(layout.clone());
                if !ptr.is_null() {
                    ptr::write_bytes(ptr, 0, layout.size());
                }
                ptr
            }
        }

        #[inline]
        unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
            libc::free(ptr as *mut libc::c_void)
        }

        #[inline]
        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
            if layout.align() <= MIN_ALIGN && layout.align() <= new_size {
                libc::realloc(ptr as *mut libc::c_void, new_size) as *mut u8
            } else {
                self.realloc_fallback(ptr, layout, new_size)
            }
        }
    }

    #[cfg(any(target_os = "android",
              target_os = "hermit",
              target_os = "redox",
              target_os = "solaris"))]
    #[inline]
    unsafe fn aligned_malloc(layout: &Layout) -> *mut u8 {
        // On android we currently target API level 9 which unfortunately
        // doesn't have the `posix_memalign` API used below. Instead we use
        // `memalign`, but this unfortunately has the property on some systems
        // where the memory returned cannot be deallocated by `free`!
        //
        // Upon closer inspection, however, this appears to work just fine with
        // Android, so for this platform we should be fine to call `memalign`
        // (which is present in API level 9). Some helpful references could
        // possibly be chromium using memalign [1], attempts at documenting that
        // memalign + free is ok [2] [3], or the current source of chromium
        // which still uses memalign on android [4].
        //
        // [1]: https://codereview.chromium.org/10796020/
        // [2]: https://code.google.com/p/android/issues/detail?id=35391
        // [3]: https://bugs.chromium.org/p/chromium/issues/detail?id=138579
        // [4]: https://chromium.googlesource.com/chromium/src/base/+/master/
        //                                       /memory/aligned_memory.cc
        libc::memalign(layout.align(), layout.size()) as *mut u8
    }

    #[cfg(not(any(target_os = "android",
                  target_os = "hermit",
                  target_os = "redox",
                  target_os = "solaris")))]
    #[inline]
    unsafe fn aligned_malloc(layout: &Layout) -> *mut u8 {
        let mut out = ptr::null_mut();
        let ret = libc::posix_memalign(&mut out, layout.align(), layout.size());
        if ret != 0 {
            ptr::null_mut()
        } else {
            out as *mut u8
        }
    }
}

#[cfg(windows)]
#[allow(bad_style)]
mod platform {
    use MIN_ALIGN;
    use System;
    use core::alloc::{GlobalAlloc, Layout};

    type LPVOID = *mut u8;
    type HANDLE = LPVOID;
    type SIZE_T = usize;
    type DWORD = u32;
    type BOOL = i32;

    extern "system" {
        fn GetProcessHeap() -> HANDLE;
        fn HeapAlloc(hHeap: HANDLE, dwFlags: DWORD, dwBytes: SIZE_T) -> LPVOID;
        fn HeapReAlloc(hHeap: HANDLE, dwFlags: DWORD, lpMem: LPVOID, dwBytes: SIZE_T) -> LPVOID;
        fn HeapFree(hHeap: HANDLE, dwFlags: DWORD, lpMem: LPVOID) -> BOOL;
        fn GetLastError() -> DWORD;
    }

    #[repr(C)]
    struct Header(*mut u8);

    const HEAP_ZERO_MEMORY: DWORD = 0x00000008;

    unsafe fn get_header<'a>(ptr: *mut u8) -> &'a mut Header {
        &mut *(ptr as *mut Header).offset(-1)
    }

    unsafe fn align_ptr(ptr: *mut u8, align: usize) -> *mut u8 {
        let aligned = ptr.offset((align - (ptr as usize & (align - 1))) as isize);
        *get_header(aligned) = Header(ptr);
        aligned
    }

    #[inline]
    unsafe fn allocate_with_flags(layout: Layout, flags: DWORD) -> *mut u8 {
        let ptr = if layout.align() <= MIN_ALIGN {
            HeapAlloc(GetProcessHeap(), flags, layout.size())
        } else {
            let size = layout.size() + layout.align();
            let ptr = HeapAlloc(GetProcessHeap(), flags, size);
            if ptr.is_null() {
                ptr
            } else {
                align_ptr(ptr, layout.align())
            }
        };
        ptr as *mut u8
    }

    #[stable(feature = "alloc_system_type", since = "1.28.0")]
    unsafe impl GlobalAlloc for System {
        #[inline]
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            allocate_with_flags(layout, 0)
        }

        #[inline]
        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
            allocate_with_flags(layout, HEAP_ZERO_MEMORY)
        }

        #[inline]
        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
            if layout.align() <= MIN_ALIGN {
                let err = HeapFree(GetProcessHeap(), 0, ptr as LPVOID);
                debug_assert!(err != 0, "Failed to free heap memory: {}",
                              GetLastError());
            } else {
                let header = get_header(ptr);
                let err = HeapFree(GetProcessHeap(), 0, header.0 as LPVOID);
                debug_assert!(err != 0, "Failed to free heap memory: {}",
                              GetLastError());
            }
        }

        #[inline]
        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
            if layout.align() <= MIN_ALIGN {
                HeapReAlloc(GetProcessHeap(), 0, ptr as LPVOID, new_size) as *mut u8
            } else {
                self.realloc_fallback(ptr, layout, new_size)
            }
        }
    }
}

// This is an implementation of a global allocator on the wasm32 platform when
// emscripten is not in use. In that situation there's no actual runtime for us
// to lean on for allocation, so instead we provide our own!
//
// The wasm32 instruction set has two instructions for getting the current
// amount of memory and growing the amount of memory. These instructions are the
// foundation on which we're able to build an allocator, so we do so! Note that
// the instructions are also pretty "global" and this is the "global" allocator
// after all!
//
// The current allocator here is the `dlmalloc` crate which we've got included
// in the rust-lang/rust repository as a submodule. The crate is a port of
// dlmalloc.c from C to Rust and is basically just so we can have "pure Rust"
// for now which is currently technically required (can't link with C yet).
//
// The crate itself provides a global allocator which on wasm has no
// synchronization as there are no threads!
#[cfg(all(target_arch = "wasm32", not(target_os = "emscripten")))]
mod platform {
    extern crate dlmalloc;

    use core::alloc::{GlobalAlloc, Layout};
    use System;

    // No need for synchronization here as wasm is currently single-threaded
    static mut DLMALLOC: dlmalloc::Dlmalloc = dlmalloc::DLMALLOC_INIT;

    #[stable(feature = "alloc_system_type", since = "1.28.0")]
    unsafe impl GlobalAlloc for System {
        #[inline]
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            DLMALLOC.malloc(layout.size(), layout.align())
        }

        #[inline]
        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
            DLMALLOC.calloc(layout.size(), layout.align())
        }

        #[inline]
        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
            DLMALLOC.free(ptr, layout.size(), layout.align())
        }

        #[inline]
        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
            DLMALLOC.realloc(ptr, layout.size(), layout.align(), new_size)
        }
    }
}
<commit_msg>liballoc_system: enable feature(nll) for bootstrap<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![no_std]
#![allow(unused_attributes)]
#![unstable(feature = "alloc_system",
            reason = "this library is unlikely to be stabilized in its current \
                      form or name",
            issue = "32838")]

#![feature(allocator_api)]
#![feature(core_intrinsics)]
#![cfg_attr(not(stage0), feature(nll))]
#![feature(staged_api)]
#![feature(rustc_attrs)]
#![cfg_attr(any(unix, target_os = "cloudabi", target_os = "redox"), feature(libc))]
#![rustc_alloc_kind = "lib"]

// The minimum alignment guaranteed by the architecture. This value is used to
// add fast paths for low alignment values.
#[cfg(all(any(target_arch = "x86",
              target_arch = "arm",
              target_arch = "mips",
              target_arch = "powerpc",
              target_arch = "powerpc64",
              target_arch = "asmjs",
              target_arch = "wasm32")))]
#[allow(dead_code)]
const MIN_ALIGN: usize = 8;
#[cfg(all(any(target_arch = "x86_64",
              target_arch = "aarch64",
              target_arch = "mips64",
              target_arch = "s390x",
              target_arch = "sparc64")))]
#[allow(dead_code)]
const MIN_ALIGN: usize = 16;

use core::alloc::{Alloc, GlobalAlloc, AllocErr, Layout};
use core::ptr::NonNull;

/// The default memory allocator provided by the operating system.
///
/// This is based on `malloc` on Unix platforms and `HeapAlloc` on Windows,
/// plus related functions.
///
/// This type can be used in a `static` item
/// with the `#[global_allocator]` attribute
/// to force the global allocator to be the system’s one.
/// (The default is jemalloc for executables, on some platforms.)
///
/// ```rust
/// use std::alloc::System;
///
/// #[global_allocator]
/// static A: System = System;
///
/// fn main() {
///     let a = Box::new(4); // Allocates from the system allocator.
///     println!("{}", a);
/// }
/// ```
///
/// It can also be used directly to allocate memory
/// independently of the standard library’s global allocator.
#[stable(feature = "alloc_system_type", since = "1.28.0")]
pub struct System;

#[unstable(feature = "allocator_api", issue = "32838")]
unsafe impl Alloc for System {
    #[inline]
    unsafe fn alloc(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::alloc(self, layout)).ok_or(AllocErr)
    }

    #[inline]
    unsafe fn alloc_zeroed(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::alloc_zeroed(self, layout)).ok_or(AllocErr)
    }

    #[inline]
    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        GlobalAlloc::dealloc(self, ptr.as_ptr(), layout)
    }

    #[inline]
    unsafe fn realloc(&mut self,
                      ptr: NonNull<u8>,
                      layout: Layout,
                      new_size: usize) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::realloc(self, ptr.as_ptr(), layout, new_size)).ok_or(AllocErr)
    }
}

#[cfg(any(windows, unix, target_os = "cloudabi", target_os = "redox"))]
mod realloc_fallback {
    use core::alloc::{GlobalAlloc, Layout};
    use core::cmp;
    use core::ptr;

    impl super::System {
        pub(crate) unsafe fn realloc_fallback(&self, ptr: *mut u8, old_layout: Layout,
                                              new_size: usize) -> *mut u8 {
            // Docs for GlobalAlloc::realloc require this to be valid:
            let new_layout = Layout::from_size_align_unchecked(new_size, old_layout.align());

            let new_ptr = GlobalAlloc::alloc(self, new_layout);
            if !new_ptr.is_null() {
                let size = cmp::min(old_layout.size(), new_size);
                ptr::copy_nonoverlapping(ptr, new_ptr, size);
                GlobalAlloc::dealloc(self, ptr, old_layout);
            }
            new_ptr
        }
    }
}

#[cfg(any(unix, target_os = "cloudabi", target_os = "redox"))]
mod platform {
    extern crate libc;

    use core::ptr;

    use MIN_ALIGN;
    use System;
    use core::alloc::{GlobalAlloc, Layout};

    #[stable(feature = "alloc_system_type", since = "1.28.0")]
    unsafe impl GlobalAlloc for System {
        #[inline]
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            if layout.align() <= MIN_ALIGN && layout.align() <= layout.size() {
                libc::malloc(layout.size()) as *mut u8
            } else {
                #[cfg(target_os = "macos")]
                {
                    if layout.align() > (1 << 31) {
                        return ptr::null_mut()
                    }
                }
                aligned_malloc(&layout)
            }
        }

        #[inline]
        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
            if layout.align() <= MIN_ALIGN && layout.align() <= layout.size() {
                libc::calloc(layout.size(), 1) as *mut u8
            } else {
                let ptr = self.alloc(layout.clone());
                if !ptr.is_null() {
                    ptr::write_bytes(ptr, 0, layout.size());
                }
                ptr
            }
        }

        #[inline]
        unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
            libc::free(ptr as *mut libc::c_void)
        }

        #[inline]
        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
            if layout.align() <= MIN_ALIGN && layout.align() <= new_size {
                libc::realloc(ptr as *mut libc::c_void, new_size) as *mut u8
            } else {
                self.realloc_fallback(ptr, layout, new_size)
            }
        }
    }

    #[cfg(any(target_os = "android",
              target_os = "hermit",
              target_os = "redox",
              target_os = "solaris"))]
    #[inline]
    unsafe fn aligned_malloc(layout: &Layout) -> *mut u8 {
        // On android we currently target API level 9 which unfortunately
        // doesn't have the `posix_memalign` API used below. Instead we use
        // `memalign`, but this unfortunately has the property on some systems
        // where the memory returned cannot be deallocated by `free`!
        //
        // Upon closer inspection, however, this appears to work just fine with
        // Android, so for this platform we should be fine to call `memalign`
        // (which is present in API level 9). Some helpful references could
        // possibly be chromium using memalign [1], attempts at documenting that
        // memalign + free is ok [2] [3], or the current source of chromium
        // which still uses memalign on android [4].
        //
        // [1]: https://codereview.chromium.org/10796020/
        // [2]: https://code.google.com/p/android/issues/detail?id=35391
        // [3]: https://bugs.chromium.org/p/chromium/issues/detail?id=138579
        // [4]: https://chromium.googlesource.com/chromium/src/base/+/master/
        //                                       /memory/aligned_memory.cc
        libc::memalign(layout.align(), layout.size()) as *mut u8
    }

    #[cfg(not(any(target_os = "android",
                  target_os = "hermit",
                  target_os = "redox",
                  target_os = "solaris")))]
    #[inline]
    unsafe fn aligned_malloc(layout: &Layout) -> *mut u8 {
        let mut out = ptr::null_mut();
        let ret = libc::posix_memalign(&mut out, layout.align(), layout.size());
        if ret != 0 {
            ptr::null_mut()
        } else {
            out as *mut u8
        }
    }
}

#[cfg(windows)]
#[allow(bad_style)]
mod platform {
    use MIN_ALIGN;
    use System;
    use core::alloc::{GlobalAlloc, Layout};

    type LPVOID = *mut u8;
    type HANDLE = LPVOID;
    type SIZE_T = usize;
    type DWORD = u32;
    type BOOL = i32;

    extern "system" {
        fn GetProcessHeap() -> HANDLE;
        fn HeapAlloc(hHeap: HANDLE, dwFlags: DWORD, dwBytes: SIZE_T) -> LPVOID;
        fn HeapReAlloc(hHeap: HANDLE, dwFlags: DWORD, lpMem: LPVOID, dwBytes: SIZE_T) -> LPVOID;
        fn HeapFree(hHeap: HANDLE, dwFlags: DWORD, lpMem: LPVOID) -> BOOL;
        fn GetLastError() -> DWORD;
    }

    #[repr(C)]
    struct Header(*mut u8);

    const HEAP_ZERO_MEMORY: DWORD = 0x00000008;

    unsafe fn get_header<'a>(ptr: *mut u8) -> &'a mut Header {
        &mut *(ptr as *mut Header).offset(-1)
    }

    unsafe fn align_ptr(ptr: *mut u8, align: usize) -> *mut u8 {
        let aligned = ptr.offset((align - (ptr as usize & (align - 1))) as isize);
        *get_header(aligned) = Header(ptr);
        aligned
    }

    #[inline]
    unsafe fn allocate_with_flags(layout: Layout, flags: DWORD) -> *mut u8 {
        let ptr = if layout.align() <= MIN_ALIGN {
            HeapAlloc(GetProcessHeap(), flags, layout.size())
        } else {
            let size = layout.size() + layout.align();
            let ptr = HeapAlloc(GetProcessHeap(), flags, size);
            if ptr.is_null() {
                ptr
            } else {
                align_ptr(ptr, layout.align())
            }
        };
        ptr as *mut u8
    }

    #[stable(feature = "alloc_system_type", since = "1.28.0")]
    unsafe impl GlobalAlloc for System {
        #[inline]
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            allocate_with_flags(layout, 0)
        }

        #[inline]
        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
            allocate_with_flags(layout, HEAP_ZERO_MEMORY)
        }

        #[inline]
        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
            if layout.align() <= MIN_ALIGN {
                let err = HeapFree(GetProcessHeap(), 0, ptr as LPVOID);
                debug_assert!(err != 0, "Failed to free heap memory: {}",
                              GetLastError());
            } else {
                let header = get_header(ptr);
                let err = HeapFree(GetProcessHeap(), 0, header.0 as LPVOID);
                debug_assert!(err != 0, "Failed to free heap memory: {}",
                              GetLastError());
            }
        }

        #[inline]
        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
            if layout.align() <= MIN_ALIGN {
                HeapReAlloc(GetProcessHeap(), 0, ptr as LPVOID, new_size) as *mut u8
            } else {
                self.realloc_fallback(ptr, layout, new_size)
            }
        }
    }
}

// This is an implementation of a global allocator on the wasm32 platform when
// emscripten is not in use. In that situation there's no actual runtime for us
// to lean on for allocation, so instead we provide our own!
//
// The wasm32 instruction set has two instructions for getting the current
// amount of memory and growing the amount of memory. These instructions are the
// foundation on which we're able to build an allocator, so we do so! Note that
// the instructions are also pretty "global" and this is the "global" allocator
// after all!
//
// The current allocator here is the `dlmalloc` crate which we've got included
// in the rust-lang/rust repository as a submodule. The crate is a port of
// dlmalloc.c from C to Rust and is basically just so we can have "pure Rust"
// for now which is currently technically required (can't link with C yet).
//
// The crate itself provides a global allocator which on wasm has no
// synchronization as there are no threads!
#[cfg(all(target_arch = "wasm32", not(target_os = "emscripten")))]
mod platform {
    extern crate dlmalloc;

    use core::alloc::{GlobalAlloc, Layout};
    use System;

    // No need for synchronization here as wasm is currently single-threaded
    static mut DLMALLOC: dlmalloc::Dlmalloc = dlmalloc::DLMALLOC_INIT;

    #[stable(feature = "alloc_system_type", since = "1.28.0")]
    unsafe impl GlobalAlloc for System {
        #[inline]
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            DLMALLOC.malloc(layout.size(), layout.align())
        }

        #[inline]
        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
            DLMALLOC.calloc(layout.size(), layout.align())
        }

        #[inline]
        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
            DLMALLOC.free(ptr, layout.size(), layout.align())
        }

        #[inline]
        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
            DLMALLOC.realloc(ptr, layout.size(), layout.align(), new_size)
        }
    }
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::net::Ipv4Addr;
use std::num::NonZeroU64;
use std::sync::Arc;

use actix_cors::Cors;
use actix_web::http::header::{HeaderName, CONTENT_TYPE};
use actix_web::{get, web, App, HttpRequest, HttpResponse, HttpServer, Responder};
use jsonwebtoken::errors::ErrorKind;
use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
use lazy_static::lazy_static;
use serde::{Deserialize, Serialize};
use serenity::model::id::{GuildId, UserId};
use serenity::model::Permissions;
use tracing::error;
use typemap_rev::TypeMap;
use uuid::Uuid;

use crate::database::settings::{PRIVACY_ADMIN_ONLY, PRIVACY_EVERYONE, PRIVACY_STAFF_ONLY};
use crate::discord::util::is_staff;
use crate::server::endpoint::ban::{get_ban, get_bans, update_ban};
use crate::server::endpoint::captcha::{get_captcha_page, submit_captcha};
use crate::server::endpoint::hardban::{get_hardban, get_hardbans, update_hardban};
use crate::server::endpoint::kick::{get_kick, get_kicks, update_kick};
use crate::server::endpoint::login::login;
use crate::server::endpoint::mute::{get_mute, get_mutes, update_mute};
use crate::server::endpoint::self_user::get_self;
use crate::server::endpoint::settings::{get_setting, reset_setting, update_setting};
use crate::server::endpoint::softban::{get_softban, get_softbans, update_softban};
use crate::server::endpoint::warn::{get_warn, get_warns, update_warn};
use crate::service::guild::{CachedMember, GuildService};
use crate::service::invalid_uuid::InvalidUUIDService;
use crate::service::setting::SettingService;
use crate::util::now;
use crate::Config;

mod endpoint;
mod model;

lazy_static! {
    static ref VALIDATION: Validation = Validation::new(Algorithm::HS512);
    static ref ENCODING_HEADER: Header = Header::new(Algorithm::HS512);
}

#[derive(Deserialize, Serialize)]
struct JwtClaims {
    #[serde(rename = "userId")]
    user_id: String,
    uuid: Uuid,
    exp: u64,
}

fn verify_token(secret: &str, token: &str) -> Option<JwtClaims> {
    match decode::<JwtClaims>(
        token,
        &DecodingKey::from_secret(secret.as_bytes()),
        &VALIDATION,
    ) {
        Ok(data) => Some(data.claims),
        Err(error) => match error.kind() {
            ErrorKind::InvalidSignature | ErrorKind::ExpiredSignature => None,
            _ => {
                error!("failed to decode jwt token {:?}", error.kind());
                None
            }
        },
    }
}

fn generate_token(secret: &str, user_id: UserId) -> Option<String> {
    let user_id = user_id.to_string();
    let expires_at = now() + 7 * 24 * 60 * 60;
    let uuid = Uuid::new_v4();
    let claims = JwtClaims {
        user_id,
        uuid,
        exp: expires_at,
    };

    match encode(
        &ENCODING_HEADER,
        &claims,
        &EncodingKey::from_secret(secret.as_bytes()),
    ) {
        Ok(token) => Some(token),
        Err(err) => {
            error!("failed to create jwt token {:?}", err);
            None
        }
    }
}

async fn is_authenticated(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
) -> Option<UserId> {
    let token = req
        .headers()
        .get(HeaderName::from_static("token")) // optimize into constant?
        .map(|value| value.to_str())
        .and_then(|result| result.ok())?;

    let claims = verify_token(&config.server_secret, token)?;

    let invalid_uuid_service = services.get::<InvalidUUIDService>()?;

    if invalid_uuid_service.is_uuid_invalid(claims.uuid).await {
        None
    } else {
        match claims.user_id.parse::<NonZeroU64>() {
            Ok(id) => Some(UserId(id)),
            Err(_) => {
                // secret leaked or we have a problem with token generation
                error!("received a token with valid signature with invalid user id");
                None
            }
        }
    }
}

pub async fn check_authentication(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
) -> Result<UserId, HttpResponse> {
    is_authenticated(config, services, req)
        .await
        .ok_or_else(|| HttpResponse::Unauthorized().json("Invalid token!"))
}

pub enum PrivateEndpointKind {
    ModLog,
    Settings,
}

async fn is_authorized_to_view_private_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    permissions: Permissions,
    endpoint_kind: PrivateEndpointKind,
) -> bool {
    let setting_service = if let Some(service) = services.get::<SettingService>() {
        service
    } else {
        return false;
    };

    let setting = setting_service.get_setting(guild_id).await;
    let privacy_setting = match endpoint_kind {
        PrivateEndpointKind::Settings => setting.privacy_settings,
        PrivateEndpointKind::ModLog => setting.privacy_mod_log,
    };

    if privacy_setting == PRIVACY_EVERYONE {
        true
    } else if privacy_setting == PRIVACY_STAFF_ONLY {
        is_staff(permissions)
    } else if privacy_setting == PRIVACY_ADMIN_ONLY {
        permissions.administrator()
    } else {
        false
    }
}

pub async fn check_authorization_to_view_private_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    permissions: Permissions,
    endpoint_kind: PrivateEndpointKind,
) -> Option<HttpResponse> {
    if is_authorized_to_view_private_endpoint(services, guild_id, permissions, endpoint_kind).await
    {
        None
    } else {
        Some(
            HttpResponse::Forbidden()
                .json("Server settings prevent you from viewing private information!"),
        )
    }
}

pub async fn is_guild_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    user_id: UserId,
) -> Option<(Arc<CachedMember>, Permissions)> {
    let guild_service = services.get::<GuildService>()?;

    let member = guild_service.get_member(guild_id, user_id).await.ok()?;

    let permissions = if let Ok(permissions) = guild_service
        .get_permissions(user_id, &member.roles, guild_id)
        .await
    {
        permissions
    } else {
        return None;
    };

    Some((member, permissions))
}

pub async fn check_guild_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    user_id: UserId,
) -> Result<(Arc<CachedMember>, Permissions), HttpResponse> {
    is_guild_endpoint(services, guild_id, user_id)
        .await
        .ok_or_else(|| {
            HttpResponse::Forbidden().json("This server either doesn't exist or you aren't in it!")
        })
}

pub async fn apply_private_endpoint_fetch_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
    endpoint_kind: PrivateEndpointKind,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if let Some(response) =
        check_authorization_to_view_private_endpoint(services, guild_id, permissions, endpoint_kind)
            .await
    {
        return Err(response);
    }

    Ok(())
}

pub async fn apply_mod_log_update_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
    required_permission: Permissions,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if permissions.contains(Permissions::ADMINISTRATOR) || permissions.contains(required_permission)
    {
        Ok(())
    } else {
        Err(HttpResponse::Forbidden()
            .json("You don't have permissions to change moderator log history!"))
    }
}

pub async fn apply_setting_update_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if permissions.contains(Permissions::ADMINISTRATOR) {
        Ok(())
    } else {
        Err(HttpResponse::Forbidden().json("You don't have permissions to change server settings!"))
    }
}

#[get("/")]
async fn index() -> impl Responder {
    "Welcome to Safety Jim API."
}

pub async fn run_server(config: Arc<Config>, services: Arc<TypeMap>) -> Result<(), Box<dyn Error>> {
    let port = config.server_port;
    HttpServer::new(move || {
        App::new()
            .app_data(web::Data::from(config.clone()))
            .app_data(web::Data::from(services.clone()))
            .service(index)
            .service(get_self)
            .service(get_bans)
            .service(get_ban)
            .service(update_ban)
            .service(get_captcha_page)
            .service(submit_captcha)
            .service(get_hardbans)
            .service(get_hardban)
            .service(update_hardban)
            .service(get_kicks)
            .service(get_kick)
            .service(update_kick)
            .service(login)
            .service(get_mutes)
            .service(get_mute)
            .service(update_mute)
            .service(get_setting)
            .service(update_setting)
            .service(reset_setting)
            .service(get_softbans)
            .service(get_softban)
            .service(update_softban)
            .service(get_warns)
            .service(get_warn)
            .service(update_warn)
            .wrap(
                Cors::default()
                    .allowed_origin(&config.cors_origin)
                    // needed for captcha endpoint, actix-cors doesn't actually check if request is for CORS
                    .allowed_origin(&config.self_url)
                    .allowed_methods(vec!["GET", "POST", "DELETE"])
                    .allowed_headers(vec![CONTENT_TYPE, HeaderName::from_static("token")])
                    .max_age(None),
            )
    })
    .bind((Ipv4Addr::LOCALHOST, port))?
    .run()
    .await?;

    Ok(())
}
<commit_msg>Listen on all ips<commit_after>use std::error::Error;
use std::net::Ipv6Addr;
use std::num::NonZeroU64;
use std::sync::Arc;

use actix_cors::Cors;
use actix_web::http::header::{HeaderName, CONTENT_TYPE};
use actix_web::{get, web, App, HttpRequest, HttpResponse, HttpServer, Responder};
use jsonwebtoken::errors::ErrorKind;
use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
use lazy_static::lazy_static;
use serde::{Deserialize, Serialize};
use serenity::model::id::{GuildId, UserId};
use serenity::model::Permissions;
use tracing::error;
use typemap_rev::TypeMap;
use uuid::Uuid;

use crate::database::settings::{PRIVACY_ADMIN_ONLY, PRIVACY_EVERYONE, PRIVACY_STAFF_ONLY};
use crate::discord::util::is_staff;
use crate::server::endpoint::ban::{get_ban, get_bans, update_ban};
use crate::server::endpoint::captcha::{get_captcha_page, submit_captcha};
use crate::server::endpoint::hardban::{get_hardban, get_hardbans, update_hardban};
use crate::server::endpoint::kick::{get_kick, get_kicks, update_kick};
use crate::server::endpoint::login::login;
use crate::server::endpoint::mute::{get_mute, get_mutes, update_mute};
use crate::server::endpoint::self_user::get_self;
use crate::server::endpoint::settings::{get_setting, reset_setting, update_setting};
use crate::server::endpoint::softban::{get_softban, get_softbans, update_softban};
use crate::server::endpoint::warn::{get_warn, get_warns, update_warn};
use crate::service::guild::{CachedMember, GuildService};
use crate::service::invalid_uuid::InvalidUUIDService;
use crate::service::setting::SettingService;
use crate::util::now;
use crate::Config;

mod endpoint;
mod model;

lazy_static! {
    static ref VALIDATION: Validation = Validation::new(Algorithm::HS512);
    static ref ENCODING_HEADER: Header = Header::new(Algorithm::HS512);
}

#[derive(Deserialize, Serialize)]
struct JwtClaims {
    #[serde(rename = "userId")]
    user_id: String,
    uuid: Uuid,
    exp: u64,
}

fn verify_token(secret: &str, token: &str) -> Option<JwtClaims> {
    match decode::<JwtClaims>(
        token,
        &DecodingKey::from_secret(secret.as_bytes()),
        &VALIDATION,
    ) {
        Ok(data) => Some(data.claims),
        Err(error) => match error.kind() {
            ErrorKind::InvalidSignature | ErrorKind::ExpiredSignature => None,
            _ => {
                error!("failed to decode jwt token {:?}", error.kind());
                None
            }
        },
    }
}

fn generate_token(secret: &str, user_id: UserId) -> Option<String> {
    let user_id = user_id.to_string();
    let expires_at = now() + 7 * 24 * 60 * 60;
    let uuid = Uuid::new_v4();
    let claims = JwtClaims {
        user_id,
        uuid,
        exp: expires_at,
    };

    match encode(
        &ENCODING_HEADER,
        &claims,
        &EncodingKey::from_secret(secret.as_bytes()),
    ) {
        Ok(token) => Some(token),
        Err(err) => {
            error!("failed to create jwt token {:?}", err);
            None
        }
    }
}

async fn is_authenticated(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
) -> Option<UserId> {
    let token = req
        .headers()
        .get(HeaderName::from_static("token")) // optimize into constant?
        .map(|value| value.to_str())
        .and_then(|result| result.ok())?;

    let claims = verify_token(&config.server_secret, token)?;

    let invalid_uuid_service = services.get::<InvalidUUIDService>()?;

    if invalid_uuid_service.is_uuid_invalid(claims.uuid).await {
        None
    } else {
        match claims.user_id.parse::<NonZeroU64>() {
            Ok(id) => Some(UserId(id)),
            Err(_) => {
                // secret leaked or we have a problem with token generation
                error!("received a token with valid signature with invalid user id");
                None
            }
        }
    }
}

pub async fn check_authentication(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
) -> Result<UserId, HttpResponse> {
    is_authenticated(config, services, req)
        .await
        .ok_or_else(|| HttpResponse::Unauthorized().json("Invalid token!"))
}

pub enum PrivateEndpointKind {
    ModLog,
    Settings,
}

async fn is_authorized_to_view_private_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    permissions: Permissions,
    endpoint_kind: PrivateEndpointKind,
) -> bool {
    let setting_service = if let Some(service) = services.get::<SettingService>() {
        service
    } else {
        return false;
    };

    let setting = setting_service.get_setting(guild_id).await;
    let privacy_setting = match endpoint_kind {
        PrivateEndpointKind::Settings => setting.privacy_settings,
        PrivateEndpointKind::ModLog => setting.privacy_mod_log,
    };

    if privacy_setting == PRIVACY_EVERYONE {
        true
    } else if privacy_setting == PRIVACY_STAFF_ONLY {
        is_staff(permissions)
    } else if privacy_setting == PRIVACY_ADMIN_ONLY {
        permissions.administrator()
    } else {
        false
    }
}

pub async fn check_authorization_to_view_private_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    permissions: Permissions,
    endpoint_kind: PrivateEndpointKind,
) -> Option<HttpResponse> {
    if is_authorized_to_view_private_endpoint(services, guild_id, permissions, endpoint_kind).await
    {
        None
    } else {
        Some(
            HttpResponse::Forbidden()
                .json("Server settings prevent you from viewing private information!"),
        )
    }
}

pub async fn is_guild_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    user_id: UserId,
) -> Option<(Arc<CachedMember>, Permissions)> {
    let guild_service = services.get::<GuildService>()?;

    let member = guild_service.get_member(guild_id, user_id).await.ok()?;

    let permissions = if let Ok(permissions) = guild_service
        .get_permissions(user_id, &member.roles, guild_id)
        .await
    {
        permissions
    } else {
        return None;
    };

    Some((member, permissions))
}

pub async fn check_guild_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    user_id: UserId,
) -> Result<(Arc<CachedMember>, Permissions), HttpResponse> {
    is_guild_endpoint(services, guild_id, user_id)
        .await
        .ok_or_else(|| {
            HttpResponse::Forbidden().json("This server either doesn't exist or you aren't in it!")
        })
}

pub async fn apply_private_endpoint_fetch_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
    endpoint_kind: PrivateEndpointKind,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if let Some(response) =
        check_authorization_to_view_private_endpoint(services, guild_id, permissions, endpoint_kind)
            .await
    {
        return Err(response);
    }

    Ok(())
}

pub async fn apply_mod_log_update_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
    required_permission: Permissions,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if permissions.contains(Permissions::ADMINISTRATOR) || permissions.contains(required_permission)
    {
        Ok(())
    } else {
        Err(HttpResponse::Forbidden()
            .json("You don't have permissions to change moderator log history!"))
    }
}

pub async fn apply_setting_update_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if permissions.contains(Permissions::ADMINISTRATOR) {
        Ok(())
    } else {
        Err(HttpResponse::Forbidden().json("You don't have permissions to change server settings!"))
    }
}

#[get("/")]
async fn index() -> impl Responder {
    "Welcome to Safety Jim API."
}

pub async fn run_server(config: Arc<Config>, services: Arc<TypeMap>) -> Result<(), Box<dyn Error>> {
    let port = config.server_port;
    HttpServer::new(move || {
        App::new()
            .app_data(web::Data::from(config.clone()))
            .app_data(web::Data::from(services.clone()))
            .service(index)
            .service(get_self)
            .service(get_bans)
            .service(get_ban)
            .service(update_ban)
            .service(get_captcha_page)
            .service(submit_captcha)
            .service(get_hardbans)
            .service(get_hardban)
            .service(update_hardban)
            .service(get_kicks)
            .service(get_kick)
            .service(update_kick)
            .service(login)
            .service(get_mutes)
            .service(get_mute)
            .service(update_mute)
            .service(get_setting)
            .service(update_setting)
            .service(reset_setting)
            .service(get_softbans)
            .service(get_softban)
            .service(update_softban)
            .service(get_warns)
            .service(get_warn)
            .service(update_warn)
            .wrap(
                Cors::default()
                    .allowed_origin(&config.cors_origin)
                    // needed for captcha endpoint, actix-cors doesn't actually check if request is for CORS
                    .allowed_origin(&config.self_url)
                    .allowed_methods(vec!["GET", "POST", "DELETE"])
                    .allowed_headers(vec![CONTENT_TYPE, HeaderName::from_static("token")])
                    .max_age(None),
            )
    })
    .bind((Ipv6Addr::UNSPECIFIED, port))?
    .run()
    .await?;

    Ok(())
}
<|endoftext|>