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

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]

extern crate protobuf;
extern crate protobuf_codegen;

mod convert;

use std::fs;
use std::path::Path;
use std::path::PathBuf;
use std::process;

mod linked_hash_map;
mod model;
mod parse_dependencies;
mod parser;
mod path;
mod rel_path;

use linked_hash_map::LinkedHashMap;
use protobuf::descriptor::FileDescriptorProto;
use protobuf_codegen::amend_io_error;
pub use protobuf_codegen::Customize;
use protobuf_codegen::ProtoPath;
use protobuf_codegen::ProtoPathBuf;

use crate::parser::ParserErrorWithLocation;

#[cfg(test)]
mod test_against_protobuf_protos;

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

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

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

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

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

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

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

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

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

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

    /// Similar to `run`, but prints the message to stderr and exits the process on error.
    pub fn run_from_script(&self) {
        if let Err(e) = self.run() {
            eprintln!("pure rust codegen failed: {}", e);
            process::exit(1);
        }
    }
}

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

#[derive(Debug, thiserror::Error)]
#[error("error in `{file}`: {error}")]
struct WithFileError {
    file: String,
    #[source]
    error: anyhow::Error,
}

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("protobuf path `{0}` is not found in import path {1}")]
    FileNotFoundInImportPath(String, String),
    #[error("file `{0}` must reside in include path {1}")]
    FileMustResideInImportPath(String, String),
}

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

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

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

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

    fn get_all_deps_already_parsed(
        &self,
        parsed: &model::FileDescriptor,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        for import in &parsed.imports {
            self.get_file_and_all_deps_already_parsed(&import.path, result);
        }
    }

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

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

        self.add_file_content(protobuf_path, fs_path, &content)
    }

    fn add_file_content(
        &mut self,
        protobuf_path: &ProtoPath,
        fs_path: &Path,
        content: &str,
    ) -> anyhow::Result<()> {
        let parsed = model::FileDescriptor::parse(content).map_err(|e| WithFileError {
            file: format!("{}", fs_path.display()),
            error: e.into(),
        })?;

        for import in &parsed.imports {
            self.add_imported_file(&import.path)?;
        }

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

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

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

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

        Ok(())
    }

    fn add_imported_file(&mut self, protobuf_path: &ProtoPath) -> anyhow::Result<()> {
        for include_dir in self.includes {
            if let Some(_) = self.parsed_files.get(protobuf_path) {
                return Ok(());
            }

            let fs_path = include_dir.join(protobuf_path.to_path());
            if fs_path.exists() {
                return self.add_file(protobuf_path, &fs_path);
            }
        }

        let embedded = match protobuf_path.to_str() {
            "rustproto.proto" => Some(RUSTPROTO_PROTO),
            "google/protobuf/any.proto" => Some(ANY_PROTO),
            "google/protobuf/api.proto" => Some(API_PROTO),
            "google/protobuf/descriptor.proto" => Some(DESCRIPTOR_PROTO),
            "google/protobuf/duration.proto" => Some(DURATION_PROTO),
            "google/protobuf/empty.proto" => Some(EMPTY_PROTO),
            "google/protobuf/field_mask.proto" => Some(FIELD_MASK_PROTO),
            "google/protobuf/source_context.proto" => Some(SOURCE_CONTEXT_PROTO),
            "google/protobuf/struct.proto" => Some(STRUCT_PROTO),
            "google/protobuf/timestamp.proto" => Some(TIMESTAMP_PROTO),
            "google/protobuf/type.proto" => Some(TYPE_PROTO),
            "google/protobuf/wrappers.proto" => Some(WRAPPERS_PROTO),
            _ => None,
        };

        match embedded {
            Some(content) => self.add_file_content(protobuf_path, protobuf_path.to_path(), content),
            None => Err(Error::FileNotFoundInImportPath(
                protobuf_path.to_string(),
                format!("{:?}", self.includes),
            )
            .into()),
        }
    }

    fn strip_prefix<'b>(path: &'b Path, prefix: &Path) -> anyhow::Result<ProtoPathBuf> {
        // special handling of `.` to allow successful `strip_prefix("foo.proto", ".")
        if prefix == Path::new(".") && path.is_relative() {
            ProtoPathBuf::from_path(path)
        } else {
            ProtoPathBuf::from_path(path.strip_prefix(prefix)?)
        }
    }

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

        match relative_path {
            Some(relative_path) => {
                self.add_file(&relative_path, fs_path)?;
                Ok(relative_path)
            }
            None => Err(Error::FileMustResideInImportPath(
                fs_path.display().to_string(),
                format!("{:?}", self.includes),
            )
            .into()),
        }
    }
}

/// Result of parsing `.proto` files.
#[doc(hidden)]
pub struct ParsedAndTypechecked {
    /// One entry for each input `.proto` file.
    pub relative_paths: Vec<ProtoPathBuf>,
    /// All parsed `.proto` files including dependencies of input files.
    pub file_descriptors: Vec<protobuf::descriptor::FileDescriptorProto>,
}

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

    let mut relative_paths = Vec::new();

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

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

    Ok(ParsedAndTypechecked {
        relative_paths,
        file_descriptors,
    })
}

/// Parse imports from a `.proto` file.
///
/// The result is [`FileDescriptorProto`] object with only `*dependency` fields filled.
#[doc(hidden)]
pub fn parse_dependencies(content: &str) -> Result<FileDescriptorProto, ParserErrorWithLocation> {
    parse_dependencies::parse_dependencies(content)
}

// TODO: these include don't work when publishing to crates
const RUSTPROTO_PROTO: &str = include_str!("proto/rustproto.proto");
const ANY_PROTO: &str = include_str!("proto/google/protobuf/any.proto");
const API_PROTO: &str = include_str!("proto/google/protobuf/api.proto");
const DESCRIPTOR_PROTO: &str = include_str!("proto/google/protobuf/descriptor.proto");
const DURATION_PROTO: &str = include_str!("proto/google/protobuf/duration.proto");
const EMPTY_PROTO: &str = include_str!("proto/google/protobuf/empty.proto");
const FIELD_MASK_PROTO: &str = include_str!("proto/google/protobuf/field_mask.proto");
const SOURCE_CONTEXT_PROTO: &str = include_str!("proto/google/protobuf/source_context.proto");
const STRUCT_PROTO: &str = include_str!("proto/google/protobuf/struct.proto");
const TIMESTAMP_PROTO: &str = include_str!("proto/google/protobuf/timestamp.proto");
const TYPE_PROTO: &str = include_str!("proto/google/protobuf/type.proto");
const WRAPPERS_PROTO: &str = include_str!("proto/google/protobuf/wrappers.proto");
<commit_msg>More thiserror<commit_after>//! API to generate `.rs` files.
//!
//! This API does not require `protoc` command present in `$PATH`.
//!
//! ```
//! extern crate protoc_rust;
//!
//! fn main() {
//!     protobuf_codegen_pure::Codegen::new()
//!         .out_dir("src/protos")
//!         .inputs(&["protos/a.proto", "protos/b.proto"])
//!         .include("protos")
//!         .run()
//!         .expect("Codegen failed.");
//! }
//! ```
//!
//! It is advisable that `protobuf-codegen-pure` build-dependecy version be the same as
//! `protobuf` dependency.
//!
//! The alternative is to use `protoc-rust` crate.

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]

extern crate protobuf;
extern crate protobuf_codegen;

mod convert;

use std::fs;
use std::io;
use std::path::Path;
use std::path::PathBuf;
use std::process;

mod linked_hash_map;
mod model;
mod parse_dependencies;
mod parser;
mod path;
mod rel_path;

use linked_hash_map::LinkedHashMap;
use protobuf::descriptor::FileDescriptorProto;
pub use protobuf_codegen::Customize;
use protobuf_codegen::ProtoPath;
use protobuf_codegen::ProtoPathBuf;

use crate::parser::ParserErrorWithLocation;

#[cfg(test)]
mod test_against_protobuf_protos;

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

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

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

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

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

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

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

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

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

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

    /// Similar to `run`, but prints the message to stderr and exits the process on error.
    pub fn run_from_script(&self) {
        if let Err(e) = self.run() {
            eprintln!("pure rust codegen failed: {}", e);
            process::exit(1);
        }
    }
}

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

#[derive(Debug, thiserror::Error)]
#[error("error in `{file}`: {error}")]
struct WithFileError {
    file: String,
    #[source]
    error: anyhow::Error,
}

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("protobuf path `{0}` is not found in import path {1}")]
    FileNotFoundInImportPath(String, String),
    #[error("file `{0}` must reside in include path {1}")]
    FileMustResideInImportPath(String, String),
    #[error("could not read file `{0}`: {1}")]
    CouldNotReadFile(String, io::Error),
}

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

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

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

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

    fn get_all_deps_already_parsed(
        &self,
        parsed: &model::FileDescriptor,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        for import in &parsed.imports {
            self.get_file_and_all_deps_already_parsed(&import.path, result);
        }
    }

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

        let content = fs::read_to_string(fs_path)
            .map_err(|e| Error::CouldNotReadFile(fs_path.display().to_string(), e))?;

        self.add_file_content(protobuf_path, fs_path, &content)
    }

    fn add_file_content(
        &mut self,
        protobuf_path: &ProtoPath,
        fs_path: &Path,
        content: &str,
    ) -> anyhow::Result<()> {
        let parsed = model::FileDescriptor::parse(content).map_err(|e| WithFileError {
            file: format!("{}", fs_path.display()),
            error: e.into(),
        })?;

        for import in &parsed.imports {
            self.add_imported_file(&import.path)?;
        }

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

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

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

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

        Ok(())
    }

    fn add_imported_file(&mut self, protobuf_path: &ProtoPath) -> anyhow::Result<()> {
        for include_dir in self.includes {
            if let Some(_) = self.parsed_files.get(protobuf_path) {
                return Ok(());
            }

            let fs_path = include_dir.join(protobuf_path.to_path());
            if fs_path.exists() {
                return self.add_file(protobuf_path, &fs_path);
            }
        }

        let embedded = match protobuf_path.to_str() {
            "rustproto.proto" => Some(RUSTPROTO_PROTO),
            "google/protobuf/any.proto" => Some(ANY_PROTO),
            "google/protobuf/api.proto" => Some(API_PROTO),
            "google/protobuf/descriptor.proto" => Some(DESCRIPTOR_PROTO),
            "google/protobuf/duration.proto" => Some(DURATION_PROTO),
            "google/protobuf/empty.proto" => Some(EMPTY_PROTO),
            "google/protobuf/field_mask.proto" => Some(FIELD_MASK_PROTO),
            "google/protobuf/source_context.proto" => Some(SOURCE_CONTEXT_PROTO),
            "google/protobuf/struct.proto" => Some(STRUCT_PROTO),
            "google/protobuf/timestamp.proto" => Some(TIMESTAMP_PROTO),
            "google/protobuf/type.proto" => Some(TYPE_PROTO),
            "google/protobuf/wrappers.proto" => Some(WRAPPERS_PROTO),
            _ => None,
        };

        match embedded {
            Some(content) => self.add_file_content(protobuf_path, protobuf_path.to_path(), content),
            None => Err(Error::FileNotFoundInImportPath(
                protobuf_path.to_string(),
                format!("{:?}", self.includes),
            )
            .into()),
        }
    }

    fn strip_prefix<'b>(path: &'b Path, prefix: &Path) -> anyhow::Result<ProtoPathBuf> {
        // special handling of `.` to allow successful `strip_prefix("foo.proto", ".")
        if prefix == Path::new(".") && path.is_relative() {
            ProtoPathBuf::from_path(path)
        } else {
            ProtoPathBuf::from_path(path.strip_prefix(prefix)?)
        }
    }

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

        match relative_path {
            Some(relative_path) => {
                self.add_file(&relative_path, fs_path)?;
                Ok(relative_path)
            }
            None => Err(Error::FileMustResideInImportPath(
                fs_path.display().to_string(),
                format!("{:?}", self.includes),
            )
            .into()),
        }
    }
}

/// Result of parsing `.proto` files.
#[doc(hidden)]
pub struct ParsedAndTypechecked {
    /// One entry for each input `.proto` file.
    pub relative_paths: Vec<ProtoPathBuf>,
    /// All parsed `.proto` files including dependencies of input files.
    pub file_descriptors: Vec<protobuf::descriptor::FileDescriptorProto>,
}

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

    let mut relative_paths = Vec::new();

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

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

    Ok(ParsedAndTypechecked {
        relative_paths,
        file_descriptors,
    })
}

/// Parse imports from a `.proto` file.
///
/// The result is [`FileDescriptorProto`] object with only `*dependency` fields filled.
#[doc(hidden)]
pub fn parse_dependencies(content: &str) -> Result<FileDescriptorProto, ParserErrorWithLocation> {
    parse_dependencies::parse_dependencies(content)
}

// TODO: these include don't work when publishing to crates
const RUSTPROTO_PROTO: &str = include_str!("proto/rustproto.proto");
const ANY_PROTO: &str = include_str!("proto/google/protobuf/any.proto");
const API_PROTO: &str = include_str!("proto/google/protobuf/api.proto");
const DESCRIPTOR_PROTO: &str = include_str!("proto/google/protobuf/descriptor.proto");
const DURATION_PROTO: &str = include_str!("proto/google/protobuf/duration.proto");
const EMPTY_PROTO: &str = include_str!("proto/google/protobuf/empty.proto");
const FIELD_MASK_PROTO: &str = include_str!("proto/google/protobuf/field_mask.proto");
const SOURCE_CONTEXT_PROTO: &str = include_str!("proto/google/protobuf/source_context.proto");
const STRUCT_PROTO: &str = include_str!("proto/google/protobuf/struct.proto");
const TIMESTAMP_PROTO: &str = include_str!("proto/google/protobuf/timestamp.proto");
const TYPE_PROTO: &str = include_str!("proto/google/protobuf/type.proto");
const WRAPPERS_PROTO: &str = include_str!("proto/google/protobuf/wrappers.proto");
<|endoftext|>


<commit_before>// This file is part of the uutils coreutils package.
//
// (c) Vsevolod Velichko <torkvemada@sorokdva.net>
// (c) Jian Zeng <anonymousknight96 AT gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

// spell-checker:ignore (ToDO) NEWROOT Userspec pstatus
mod error;

use crate::error::ChrootError;
use clap::{crate_version, Arg, Command};
use std::ffi::CString;
use std::io::Error;
use std::path::Path;
use std::process;
use uucore::error::{set_exit_code, UResult};
use uucore::libc::{self, chroot, setgid, setgroups, setuid};
use uucore::{entries, format_usage};

static ABOUT: &str = "Run COMMAND with root directory set to NEWROOT.";
static USAGE: &str = "{} [OPTION]... NEWROOT [COMMAND [ARG]...]";

mod options {
    pub const NEWROOT: &str = "newroot";
    pub const USER: &str = "user";
    pub const GROUP: &str = "group";
    pub const GROUPS: &str = "groups";
    pub const USERSPEC: &str = "userspec";
    pub const COMMAND: &str = "command";
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let args = args.collect_lossy();

    let matches = uu_app().get_matches_from(args);

    let default_shell: &'static str = "/bin/sh";
    let default_option: &'static str = "-i";
    let user_shell = std::env::var("SHELL");

    let newroot: &Path = match matches.value_of(options::NEWROOT) {
        Some(v) => Path::new(v),
        None => return Err(ChrootError::MissingNewRoot.into()),
    };

    if !newroot.is_dir() {
        return Err(ChrootError::NoSuchDirectory(format!("{}", newroot.display())).into());
    }

    let commands = match matches.get_many::<String>(options::COMMAND) {
        Some(v) => v.map(|s| s.as_str()).collect(),
        None => vec![],
    };

    // TODO: refactor the args and command matching
    // See: https://github.com/uutils/coreutils/pull/2365#discussion_r647849967
    let command: Vec<&str> = match commands.len() {
        0 => {
            let shell: &str = match user_shell {
                Err(_) => default_shell,
                Ok(ref s) => s.as_ref(),
            };
            vec![shell, default_option]
        }
        _ => commands,
    };

    assert!(!command.is_empty());
    let chroot_command = command[0];
    let chroot_args = &command[1..];

    // NOTE: Tests can only trigger code beyond this point if they're invoked with root permissions
    set_context(newroot, &matches)?;

    let pstatus = match process::Command::new(chroot_command)
        .args(chroot_args)
        .status()
    {
        Ok(status) => status,
        Err(e) => return Err(ChrootError::CommandFailed(command[0].to_string(), e).into()),
    };

    let code = if pstatus.success() {
        0
    } else {
        pstatus.code().unwrap_or(-1)
    };
    set_exit_code(code);
    Ok(())
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .infer_long_args(true)
        .arg(
            Arg::new(options::NEWROOT)
                .value_hint(clap::ValueHint::DirPath)
                .hide(true)
                .required(true)
                .index(1),
        )
        .arg(
            Arg::new(options::USER)
                .short('u')
                .long(options::USER)
                .help("User (ID or name) to switch before running the program")
                .value_name("USER"),
        )
        .arg(
            Arg::new(options::GROUP)
                .short('g')
                .long(options::GROUP)
                .help("Group (ID or name) to switch to")
                .value_name("GROUP"),
        )
        .arg(
            Arg::new(options::GROUPS)
                .short('G')
                .long(options::GROUPS)
                .help("Comma-separated list of groups to switch to")
                .value_name("GROUP1,GROUP2..."),
        )
        .arg(
            Arg::new(options::USERSPEC)
                .long(options::USERSPEC)
                .help(
                    "Colon-separated user and group to switch to. \
                     Same as -u USER -g GROUP. \
                     Userspec has higher preference than -u and/or -g",
                )
                .value_name("USER:GROUP"),
        )
        .arg(
            Arg::new(options::COMMAND)
                .value_hint(clap::ValueHint::CommandName)
                .hide(true)
                .multiple_occurrences(true)
                .index(2),
        )
}

fn set_context(root: &Path, options: &clap::ArgMatches) -> UResult<()> {
    let userspec_str = options.value_of(options::USERSPEC);
    let user_str = options.value_of(options::USER).unwrap_or_default();
    let group_str = options.value_of(options::GROUP).unwrap_or_default();
    let groups_str = options.value_of(options::GROUPS).unwrap_or_default();
    let userspec = match userspec_str {
        Some(u) => {
            let s: Vec<&str> = u.split(':').collect();
            if s.len() != 2 || s.iter().any(|&spec| spec.is_empty()) {
                return Err(ChrootError::InvalidUserspec(u.to_string()).into());
            };
            s
        }
        None => Vec::new(),
    };

    let (user, group) = if userspec.is_empty() {
        (user_str, group_str)
    } else {
        (userspec[0], userspec[1])
    };

    enter_chroot(root)?;

    set_groups_from_str(groups_str)?;
    set_main_group(group)?;
    set_user(user)?;
    Ok(())
}

fn enter_chroot(root: &Path) -> UResult<()> {
    std::env::set_current_dir(root).unwrap();
    let err = unsafe {
        chroot(CString::new(".").unwrap().as_bytes_with_nul().as_ptr() as *const libc::c_char)
    };
    if err == 0 {
        Ok(())
    } else {
        Err(ChrootError::CannotEnter(format!("{}", root.display()), Error::last_os_error()).into())
    }
}

fn set_main_group(group: &str) -> UResult<()> {
    if !group.is_empty() {
        let group_id = match entries::grp2gid(group) {
            Ok(g) => g,
            _ => return Err(ChrootError::NoSuchGroup(group.to_string()).into()),
        };
        let err = unsafe { setgid(group_id) };
        if err != 0 {
            return Err(
                ChrootError::SetGidFailed(group_id.to_string(), Error::last_os_error()).into(),
            );
        }
    }
    Ok(())
}

#[cfg(any(target_vendor = "apple", target_os = "freebsd"))]
fn set_groups(groups: &[libc::gid_t]) -> libc::c_int {
    unsafe { setgroups(groups.len() as libc::c_int, groups.as_ptr()) }
}

#[cfg(any(target_os = "linux", target_os = "android"))]
fn set_groups(groups: &[libc::gid_t]) -> libc::c_int {
    unsafe { setgroups(groups.len() as libc::size_t, groups.as_ptr()) }
}

fn set_groups_from_str(groups: &str) -> UResult<()> {
    if !groups.is_empty() {
        let mut groups_vec = vec![];
        for group in groups.split(',') {
            let gid = match entries::grp2gid(group) {
                Ok(g) => g,
                Err(_) => return Err(ChrootError::NoSuchGroup(group.to_string()).into()),
            };
            groups_vec.push(gid);
        }
        let err = set_groups(&groups_vec);
        if err != 0 {
            return Err(ChrootError::SetGroupsFailed(Error::last_os_error()).into());
        }
    }
    Ok(())
}

fn set_user(user: &str) -> UResult<()> {
    if !user.is_empty() {
        let user_id = entries::usr2uid(user).unwrap();
        let err = unsafe { setuid(user_id as libc::uid_t) };
        if err != 0 {
            return Err(
                ChrootError::SetUserFailed(user.to_string(), Error::last_os_error()).into(),
            );
        }
    }
    Ok(())
}
<commit_msg>chroot: set trailing var arg<commit_after>// This file is part of the uutils coreutils package.
//
// (c) Vsevolod Velichko <torkvemada@sorokdva.net>
// (c) Jian Zeng <anonymousknight96 AT gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

// spell-checker:ignore (ToDO) NEWROOT Userspec pstatus
mod error;

use crate::error::ChrootError;
use clap::{crate_version, Arg, Command};
use std::ffi::CString;
use std::io::Error;
use std::path::Path;
use std::process;
use uucore::error::{set_exit_code, UResult};
use uucore::libc::{self, chroot, setgid, setgroups, setuid};
use uucore::{entries, format_usage};

static ABOUT: &str = "Run COMMAND with root directory set to NEWROOT.";
static USAGE: &str = "{} [OPTION]... NEWROOT [COMMAND [ARG]...]";

mod options {
    pub const NEWROOT: &str = "newroot";
    pub const USER: &str = "user";
    pub const GROUP: &str = "group";
    pub const GROUPS: &str = "groups";
    pub const USERSPEC: &str = "userspec";
    pub const COMMAND: &str = "command";
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let args = args.collect_lossy();

    let matches = uu_app().get_matches_from(args);

    let default_shell: &'static str = "/bin/sh";
    let default_option: &'static str = "-i";
    let user_shell = std::env::var("SHELL");

    let newroot: &Path = match matches.value_of(options::NEWROOT) {
        Some(v) => Path::new(v),
        None => return Err(ChrootError::MissingNewRoot.into()),
    };

    if !newroot.is_dir() {
        return Err(ChrootError::NoSuchDirectory(format!("{}", newroot.display())).into());
    }

    let commands = match matches.get_many::<String>(options::COMMAND) {
        Some(v) => v.map(|s| s.as_str()).collect(),
        None => vec![],
    };

    // TODO: refactor the args and command matching
    // See: https://github.com/uutils/coreutils/pull/2365#discussion_r647849967
    let command: Vec<&str> = match commands.len() {
        0 => {
            let shell: &str = match user_shell {
                Err(_) => default_shell,
                Ok(ref s) => s.as_ref(),
            };
            vec![shell, default_option]
        }
        _ => commands,
    };

    assert!(!command.is_empty());
    let chroot_command = command[0];
    let chroot_args = &command[1..];

    // NOTE: Tests can only trigger code beyond this point if they're invoked with root permissions
    set_context(newroot, &matches)?;

    let pstatus = match process::Command::new(chroot_command)
        .args(chroot_args)
        .status()
    {
        Ok(status) => status,
        Err(e) => return Err(ChrootError::CommandFailed(command[0].to_string(), e).into()),
    };

    let code = if pstatus.success() {
        0
    } else {
        pstatus.code().unwrap_or(-1)
    };
    set_exit_code(code);
    Ok(())
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .infer_long_args(true)
        .trailing_var_arg(true)
        .arg(
            Arg::new(options::NEWROOT)
                .value_hint(clap::ValueHint::DirPath)
                .hide(true)
                .required(true)
                .index(1),
        )
        .arg(
            Arg::new(options::USER)
                .short('u')
                .long(options::USER)
                .help("User (ID or name) to switch before running the program")
                .value_name("USER"),
        )
        .arg(
            Arg::new(options::GROUP)
                .short('g')
                .long(options::GROUP)
                .help("Group (ID or name) to switch to")
                .value_name("GROUP"),
        )
        .arg(
            Arg::new(options::GROUPS)
                .short('G')
                .long(options::GROUPS)
                .help("Comma-separated list of groups to switch to")
                .value_name("GROUP1,GROUP2..."),
        )
        .arg(
            Arg::new(options::USERSPEC)
                .long(options::USERSPEC)
                .help(
                    "Colon-separated user and group to switch to. \
                     Same as -u USER -g GROUP. \
                     Userspec has higher preference than -u and/or -g",
                )
                .value_name("USER:GROUP"),
        )
        .arg(
            Arg::new(options::COMMAND)
                .value_hint(clap::ValueHint::CommandName)
                .hide(true)
                .multiple_occurrences(true)
                .index(2),
        )
}

fn set_context(root: &Path, options: &clap::ArgMatches) -> UResult<()> {
    let userspec_str = options.value_of(options::USERSPEC);
    let user_str = options.value_of(options::USER).unwrap_or_default();
    let group_str = options.value_of(options::GROUP).unwrap_or_default();
    let groups_str = options.value_of(options::GROUPS).unwrap_or_default();
    let userspec = match userspec_str {
        Some(u) => {
            let s: Vec<&str> = u.split(':').collect();
            if s.len() != 2 || s.iter().any(|&spec| spec.is_empty()) {
                return Err(ChrootError::InvalidUserspec(u.to_string()).into());
            };
            s
        }
        None => Vec::new(),
    };

    let (user, group) = if userspec.is_empty() {
        (user_str, group_str)
    } else {
        (userspec[0], userspec[1])
    };

    enter_chroot(root)?;

    set_groups_from_str(groups_str)?;
    set_main_group(group)?;
    set_user(user)?;
    Ok(())
}

fn enter_chroot(root: &Path) -> UResult<()> {
    std::env::set_current_dir(root).unwrap();
    let err = unsafe {
        chroot(CString::new(".").unwrap().as_bytes_with_nul().as_ptr() as *const libc::c_char)
    };
    if err == 0 {
        Ok(())
    } else {
        Err(ChrootError::CannotEnter(format!("{}", root.display()), Error::last_os_error()).into())
    }
}

fn set_main_group(group: &str) -> UResult<()> {
    if !group.is_empty() {
        let group_id = match entries::grp2gid(group) {
            Ok(g) => g,
            _ => return Err(ChrootError::NoSuchGroup(group.to_string()).into()),
        };
        let err = unsafe { setgid(group_id) };
        if err != 0 {
            return Err(
                ChrootError::SetGidFailed(group_id.to_string(), Error::last_os_error()).into(),
            );
        }
    }
    Ok(())
}

#[cfg(any(target_vendor = "apple", target_os = "freebsd"))]
fn set_groups(groups: &[libc::gid_t]) -> libc::c_int {
    unsafe { setgroups(groups.len() as libc::c_int, groups.as_ptr()) }
}

#[cfg(any(target_os = "linux", target_os = "android"))]
fn set_groups(groups: &[libc::gid_t]) -> libc::c_int {
    unsafe { setgroups(groups.len() as libc::size_t, groups.as_ptr()) }
}

fn set_groups_from_str(groups: &str) -> UResult<()> {
    if !groups.is_empty() {
        let mut groups_vec = vec![];
        for group in groups.split(',') {
            let gid = match entries::grp2gid(group) {
                Ok(g) => g,
                Err(_) => return Err(ChrootError::NoSuchGroup(group.to_string()).into()),
            };
            groups_vec.push(gid);
        }
        let err = set_groups(&groups_vec);
        if err != 0 {
            return Err(ChrootError::SetGroupsFailed(Error::last_os_error()).into());
        }
    }
    Ok(())
}

fn set_user(user: &str) -> UResult<()> {
    if !user.is_empty() {
        let user_id = entries::usr2uid(user).unwrap();
        let err = unsafe { setuid(user_id as libc::uid_t) };
        if err != 0 {
            return Err(
                ChrootError::SetUserFailed(user.to_string(), Error::last_os_error()).into(),
            );
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod handlers;

use std::collections::HashMap;
use std::path::PathBuf;
use std::process::{Child, Command, Stdio};
use std::str::FromStr;
use std::thread;
use std::time::Duration;

use core;
use core::package::{PackageIdent, PackageInstall};
use core::os::process::Signal;
use core::os::signals::{self, SignalEvent};
use ipc_channel::ipc::{IpcOneShotServer, IpcReceiver, IpcSender};
use protobuf;
use protocol::{self, ERR_NO_RETRY_EXCODE, OK_NO_RETRY_EXCODE};

use self::handlers::Handler;
use {SUP_CMD, SUP_PACKAGE_IDENT};
use error::{Error, Result};
use service::Service;

const SUP_CMD_ENVVAR: &'static str = "HAB_SUP_BINARY";
static LOGKEY: &'static str = "SV";

type Receiver = IpcReceiver<Vec<u8>>;
type Sender = IpcSender<Vec<u8>>;

enum TickState {
    Continue,
    Exit(i32),
}

pub struct Server {
    services: ServiceTable,
    tx: Sender,
    rx: Receiver,
    supervisor: Child,
    args: Vec<String>,
}

impl Server {
    pub fn new(args: Vec<String>) -> Result<Self> {
        let ((rx, tx), supervisor) = Self::init(&args)?;
        Ok(Server {
            services: ServiceTable::default(),
            tx: tx,
            rx: rx,
            supervisor: supervisor,
            args: args,
        })
    }

    fn init(args: &[String]) -> Result<((Receiver, Sender), Child)> {
        let (server, pipe) = IpcOneShotServer::new().map_err(Error::OpenPipe)?;
        let supervisor = spawn_supervisor(&pipe, args)?;
        let channel = setup_connection(server)?;
        Ok((channel, supervisor))
    }

    #[allow(unused_must_use)]
    fn reload(&mut self) -> Result<()> {
        self.supervisor.kill();
        self.supervisor.wait();
        let ((rx, tx), supervisor) = Self::init(&self.args)?;
        self.tx = tx;
        self.rx = rx;
        self.supervisor = supervisor;
        Ok(())
    }

    fn forward_signal(&self, signal: Signal) {
        if let Err(err) = core::os::process::signal(self.supervisor.id(), signal) {
            error!(
                "Unable to signal Supervisor, {}, {}",
                self.supervisor.id(),
                err
            );
        }
    }

    fn handle_message(&mut self) -> Result<TickState> {
        match self.rx.try_recv() {
            Ok(bytes) => {
                dispatch(&self.tx, &bytes, &mut self.services);
                Ok(TickState::Continue)
            }
            Err(_) => {
                match self.supervisor.try_wait() {
                    Ok(None) => Ok(TickState::Continue),
                    Ok(Some(status)) => {
                        debug!("Supervisor exited: {}", status);
                        match status.code() {
                            Some(ERR_NO_RETRY_EXCODE) => {
                                return Ok(TickState::Exit(ERR_NO_RETRY_EXCODE))
                            }
                            Some(OK_NO_RETRY_EXCODE) => return Ok(TickState::Exit(0)),
                            _ => (),
                        }
                        Err(Error::SupShutdown)
                    }
                    Err(err) => {
                        warn!("Unable to wait for Supervisor, {}", err);
                        Err(Error::SupShutdown)
                    }
                }
            }
        }
    }

    fn reap_zombies(&mut self) {
        self.services.reap_zombies()
    }

    fn shutdown(&mut self) {
        debug!("Shutting down...");
        if send(&self.tx, &protocol::Shutdown::new()).is_err() {
            warn!("Forcefully stopping Supervisor: {}", self.supervisor.id());
            if let Err(err) = self.supervisor.kill() {
                warn!(
                    "Unable to kill Supervisor, {}, {}",
                    self.supervisor.id(),
                    err
                );
            }
        }
        self.supervisor.wait().ok();
        self.services.kill_all();
        outputln!("Hasta la vista, services.");
    }

    fn tick(&mut self) -> Result<TickState> {
        self.reap_zombies();
        match signals::check_for_signal() {
            Some(SignalEvent::Shutdown) => {
                self.shutdown();
                return Ok(TickState::Exit(0));
            }
            Some(SignalEvent::Passthrough(signal)) => self.forward_signal(signal),
            None => (),
        }
        self.handle_message()
    }
}

#[derive(Debug, Default)]
pub struct ServiceTable(HashMap<u32, Service>);

impl ServiceTable {
    pub fn get(&self, pid: u32) -> Option<&Service> {
        self.0.get(&pid)
    }

    pub fn get_mut(&mut self, pid: u32) -> Option<&mut Service> {
        self.0.get_mut(&pid)
    }

    pub fn insert(&mut self, service: Service) {
        self.0.insert(service.id(), service);
    }

    pub fn remove(&mut self, pid: u32) -> Option<Service> {
        self.0.remove(&pid)
    }

    fn kill_all(&mut self) {
        for service in self.0.values_mut() {
            outputln!(preamble service.name(), "Stopping...");
            let shutdown_method = service.kill();
            outputln!(preamble service.name(), "Shutdown OK: {}", shutdown_method);
        }
    }

    fn reap_zombies(&mut self) {
        let mut dead: Vec<u32> = vec![];
        for service in self.0.values_mut() {
            match service.try_wait() {
                Ok(None) => (),
                Ok(Some(code)) => {
                    debug!("Child exited, {}, {}", service.id(), code);
                    dead.push(service.id());
                }
                Err(err) => {
                    warn!("Error waiting for child, {}, {}", service.id(), err);
                    dead.push(service.id());
                }
            }
        }
        for pid in dead {
            self.0.remove(&pid);
        }
    }
}

////////////////////////
// Public Func
//

pub fn reply<T>(tx: &Sender, txn: &protocol::NetTxn, msg: &T) -> Result<()>
where
    T: protobuf::MessageStatic,
{
    let bytes = txn.build_reply(msg)
        .map_err(Error::Serialize)?
        .to_bytes()
        .map_err(Error::Serialize)?;
    tx.send(bytes).map_err(Error::Send)?;
    Ok(())
}

pub fn run(args: Vec<String>) -> Result<i32> {
    let mut server = Server::new(args)?;
    signals::init();
    loop {
        match server.tick() {
            Ok(TickState::Continue) => thread::sleep(Duration::from_millis(100)),
            Ok(TickState::Exit(code)) => return Ok(code),
            Err(_) => {
                while server.reload().is_err() {
                    thread::sleep(Duration::from_millis(1_000));
                }
            }
        }
    }
}

pub fn send<T>(tx: &Sender, msg: &T) -> Result<()>
where
    T: protobuf::MessageStatic,
{
    let bytes = protocol::NetTxn::build(msg)
        .map_err(Error::Serialize)?
        .to_bytes()
        .map_err(Error::Serialize)?;
    tx.send(bytes).map_err(Error::Send)?;
    Ok(())
}

////////////////////////
// Private Func
//

fn dispatch(tx: &Sender, bytes: &[u8], services: &mut ServiceTable) {
    let msg = match protocol::NetTxn::from_bytes(bytes) {
        Ok(msg) => msg,
        Err(err) => {
            error!("Unable to decode NetTxn from Supervisor, {}", err);
            return;
        }
    };
    let func = match msg.message_id() {
        "Restart" => handlers::RestartHandler::run,
        "Spawn" => handlers::SpawnHandler::run,
        "Terminate" => handlers::TerminateHandler::run,
        unknown => {
            warn!("Received unknown message from Supervisor, {}", unknown);
            return;
        }
    };
    func(tx, msg, services);
}

fn setup_connection(server: IpcOneShotServer<Vec<u8>>) -> Result<(Receiver, Sender)> {
    let (rx, raw) = server.accept().map_err(|_| Error::AcceptConn)?;
    let txn = protocol::NetTxn::from_bytes(&raw).map_err(
        Error::Deserialize,
    )?;
    let mut msg = txn.decode::<protocol::Register>().map_err(
        Error::Deserialize,
    )?;
    let tx = IpcSender::connect(msg.take_pipe()).map_err(Error::Connect)?;
    send(&tx, &protocol::NetOk::new())?;
    Ok((rx, tx))
}

fn spawn_supervisor(pipe: &str, args: &[String]) -> Result<Child> {
    let binary = supervisor_cmd()?;
    let mut command = Command::new(&binary);
    debug!("Starting Supervisor...");
    let child = command
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .env(protocol::LAUNCHER_PIPE_ENV, pipe)
        .args(args)
        .spawn()
        .map_err(Error::SupSpawn)?;
    Ok(child)
}

fn supervisor_cmd() -> Result<PathBuf> {
    if let Ok(command) = core::env::var(SUP_CMD_ENVVAR) {
        return Ok(PathBuf::from(command));
    }
    let ident = PackageIdent::from_str(SUP_PACKAGE_IDENT).unwrap();
    match PackageInstall::load_at_least(&ident, None) {
        Ok(install) => {
            match core::fs::find_command_in_pkg(SUP_CMD, &install, "/") {
                Ok(Some(cmd)) => Ok(cmd),
                _ => Err(Error::SupBinaryNotFound),
            }
        }
        Err(_) => Err(Error::SupPackageNotFound),
    }
}
<commit_msg>Stop all services on Launcher shutdown<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod handlers;

use std::collections::HashMap;
use std::path::PathBuf;
use std::process::{Child, Command, Stdio};
use std::str::FromStr;
use std::thread;
use std::time::Duration;

use core;
use core::package::{PackageIdent, PackageInstall};
use core::os::process::Signal;
use core::os::signals::{self, SignalEvent};
use ipc_channel::ipc::{IpcOneShotServer, IpcReceiver, IpcSender};
use protobuf;
use protocol::{self, ERR_NO_RETRY_EXCODE, OK_NO_RETRY_EXCODE};

use self::handlers::Handler;
use {SUP_CMD, SUP_PACKAGE_IDENT};
use error::{Error, Result};
use service::Service;

const SUP_CMD_ENVVAR: &'static str = "HAB_SUP_BINARY";
static LOGKEY: &'static str = "SV";

type Receiver = IpcReceiver<Vec<u8>>;
type Sender = IpcSender<Vec<u8>>;

enum TickState {
    Continue,
    Exit(i32),
}

pub struct Server {
    services: ServiceTable,
    tx: Sender,
    rx: Receiver,
    supervisor: Child,
    args: Vec<String>,
}

impl Server {
    pub fn new(args: Vec<String>) -> Result<Self> {
        let ((rx, tx), supervisor) = Self::init(&args)?;
        Ok(Server {
            services: ServiceTable::default(),
            tx: tx,
            rx: rx,
            supervisor: supervisor,
            args: args,
        })
    }

    fn init(args: &[String]) -> Result<((Receiver, Sender), Child)> {
        let (server, pipe) = IpcOneShotServer::new().map_err(Error::OpenPipe)?;
        let supervisor = spawn_supervisor(&pipe, args)?;
        let channel = setup_connection(server)?;
        Ok((channel, supervisor))
    }

    #[allow(unused_must_use)]
    fn reload(&mut self) -> Result<()> {
        self.supervisor.kill();
        self.supervisor.wait();
        let ((rx, tx), supervisor) = Self::init(&self.args)?;
        self.tx = tx;
        self.rx = rx;
        self.supervisor = supervisor;
        Ok(())
    }

    fn forward_signal(&self, signal: Signal) {
        if let Err(err) = core::os::process::signal(self.supervisor.id(), signal) {
            error!(
                "Unable to signal Supervisor, {}, {}",
                self.supervisor.id(),
                err
            );
        }
    }

    fn handle_message(&mut self) -> Result<TickState> {
        match self.rx.try_recv() {
            Ok(bytes) => {
                dispatch(&self.tx, &bytes, &mut self.services);
                Ok(TickState::Continue)
            }
            Err(_) => {
                match self.supervisor.try_wait() {
                    Ok(None) => Ok(TickState::Continue),
                    Ok(Some(status)) => {
                        debug!("Supervisor exited: {}", status);
                        match status.code() {
                            Some(ERR_NO_RETRY_EXCODE) => {
                                self.services.kill_all();
                                return Ok(TickState::Exit(ERR_NO_RETRY_EXCODE));
                            }
                            Some(OK_NO_RETRY_EXCODE) => {
                                self.services.kill_all();
                                return Ok(TickState::Exit(0));
                            }
                            _ => (),
                        }
                        Err(Error::SupShutdown)
                    }
                    Err(err) => {
                        warn!("Unable to wait for Supervisor, {}", err);
                        Err(Error::SupShutdown)
                    }
                }
            }
        }
    }

    fn reap_zombies(&mut self) {
        self.services.reap_zombies()
    }

    fn shutdown(&mut self) {
        debug!("Shutting down...");
        if send(&self.tx, &protocol::Shutdown::new()).is_err() {
            warn!("Forcefully stopping Supervisor: {}", self.supervisor.id());
            if let Err(err) = self.supervisor.kill() {
                warn!(
                    "Unable to kill Supervisor, {}, {}",
                    self.supervisor.id(),
                    err
                );
            }
        }
        self.supervisor.wait().ok();
        self.services.kill_all();
        outputln!("Hasta la vista, services.");
    }

    fn tick(&mut self) -> Result<TickState> {
        self.reap_zombies();
        match signals::check_for_signal() {
            Some(SignalEvent::Shutdown) => {
                self.shutdown();
                return Ok(TickState::Exit(0));
            }
            Some(SignalEvent::Passthrough(signal)) => self.forward_signal(signal),
            None => (),
        }
        self.handle_message()
    }
}

#[derive(Debug, Default)]
pub struct ServiceTable(HashMap<u32, Service>);

impl ServiceTable {
    pub fn get(&self, pid: u32) -> Option<&Service> {
        self.0.get(&pid)
    }

    pub fn get_mut(&mut self, pid: u32) -> Option<&mut Service> {
        self.0.get_mut(&pid)
    }

    pub fn insert(&mut self, service: Service) {
        self.0.insert(service.id(), service);
    }

    pub fn remove(&mut self, pid: u32) -> Option<Service> {
        self.0.remove(&pid)
    }

    fn kill_all(&mut self) {
        for service in self.0.values_mut() {
            outputln!(preamble service.name(), "Stopping...");
            let shutdown_method = service.kill();
            outputln!(preamble service.name(), "Shutdown OK: {}", shutdown_method);
        }
    }

    fn reap_zombies(&mut self) {
        let mut dead: Vec<u32> = vec![];
        for service in self.0.values_mut() {
            match service.try_wait() {
                Ok(None) => (),
                Ok(Some(code)) => {
                    debug!("Child exited, {}, {}", service.id(), code);
                    dead.push(service.id());
                }
                Err(err) => {
                    warn!("Error waiting for child, {}, {}", service.id(), err);
                    dead.push(service.id());
                }
            }
        }
        for pid in dead {
            self.0.remove(&pid);
        }
    }
}

////////////////////////
// Public Func
//

pub fn reply<T>(tx: &Sender, txn: &protocol::NetTxn, msg: &T) -> Result<()>
where
    T: protobuf::MessageStatic,
{
    let bytes = txn.build_reply(msg)
        .map_err(Error::Serialize)?
        .to_bytes()
        .map_err(Error::Serialize)?;
    tx.send(bytes).map_err(Error::Send)?;
    Ok(())
}

pub fn run(args: Vec<String>) -> Result<i32> {
    let mut server = Server::new(args)?;
    signals::init();
    loop {
        match server.tick() {
            Ok(TickState::Continue) => thread::sleep(Duration::from_millis(100)),
            Ok(TickState::Exit(code)) => return Ok(code),
            Err(_) => {
                while server.reload().is_err() {
                    thread::sleep(Duration::from_millis(1_000));
                }
            }
        }
    }
}

pub fn send<T>(tx: &Sender, msg: &T) -> Result<()>
where
    T: protobuf::MessageStatic,
{
    let bytes = protocol::NetTxn::build(msg)
        .map_err(Error::Serialize)?
        .to_bytes()
        .map_err(Error::Serialize)?;
    tx.send(bytes).map_err(Error::Send)?;
    Ok(())
}

////////////////////////
// Private Func
//

fn dispatch(tx: &Sender, bytes: &[u8], services: &mut ServiceTable) {
    let msg = match protocol::NetTxn::from_bytes(bytes) {
        Ok(msg) => msg,
        Err(err) => {
            error!("Unable to decode NetTxn from Supervisor, {}", err);
            return;
        }
    };
    let func = match msg.message_id() {
        "Restart" => handlers::RestartHandler::run,
        "Spawn" => handlers::SpawnHandler::run,
        "Terminate" => handlers::TerminateHandler::run,
        unknown => {
            warn!("Received unknown message from Supervisor, {}", unknown);
            return;
        }
    };
    func(tx, msg, services);
}

fn setup_connection(server: IpcOneShotServer<Vec<u8>>) -> Result<(Receiver, Sender)> {
    let (rx, raw) = server.accept().map_err(|_| Error::AcceptConn)?;
    let txn = protocol::NetTxn::from_bytes(&raw).map_err(
        Error::Deserialize,
    )?;
    let mut msg = txn.decode::<protocol::Register>().map_err(
        Error::Deserialize,
    )?;
    let tx = IpcSender::connect(msg.take_pipe()).map_err(Error::Connect)?;
    send(&tx, &protocol::NetOk::new())?;
    Ok((rx, tx))
}

fn spawn_supervisor(pipe: &str, args: &[String]) -> Result<Child> {
    let binary = supervisor_cmd()?;
    let mut command = Command::new(&binary);
    debug!("Starting Supervisor...");
    let child = command
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .env(protocol::LAUNCHER_PIPE_ENV, pipe)
        .args(args)
        .spawn()
        .map_err(Error::SupSpawn)?;
    Ok(child)
}

fn supervisor_cmd() -> Result<PathBuf> {
    if let Ok(command) = core::env::var(SUP_CMD_ENVVAR) {
        return Ok(PathBuf::from(command));
    }
    let ident = PackageIdent::from_str(SUP_PACKAGE_IDENT).unwrap();
    match PackageInstall::load_at_least(&ident, None) {
        Ok(install) => {
            match core::fs::find_command_in_pkg(SUP_CMD, &install, "/") {
                Ok(Some(cmd)) => Ok(cmd),
                _ => Err(Error::SupBinaryNotFound),
            }
        }
        Err(_) => Err(Error::SupPackageNotFound),
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// https://www.khronos.org/registry/webgl/specs/latest/1.0/webgl.idl
use canvas_traits::webgl::{WebGLCommand, WebGLError, WebGLMsgSender, WebGLParameter, WebGLProgramId, WebGLResult};
use canvas_traits::webgl::webgl_channel;
use dom::bindings::codegen::Bindings::WebGLProgramBinding;
use dom::bindings::codegen::Bindings::WebGLRenderingContextBinding::WebGLRenderingContextConstants as constants;
use dom::bindings::reflector::{DomObject, reflect_dom_object};
use dom::bindings::root::{DomRoot, MutNullableDom};
use dom::bindings::str::DOMString;
use dom::webglactiveinfo::WebGLActiveInfo;
use dom::webglobject::WebGLObject;
use dom::webglrenderingcontext::MAX_UNIFORM_AND_ATTRIBUTE_LEN;
use dom::webglshader::WebGLShader;
use dom::window::Window;
use dom_struct::dom_struct;
use std::cell::Cell;

#[dom_struct]
pub struct WebGLProgram {
    webgl_object: WebGLObject,
    id: WebGLProgramId,
    is_deleted: Cell<bool>,
    link_called: Cell<bool>,
    linked: Cell<bool>,
    fragment_shader: MutNullableDom<WebGLShader>,
    vertex_shader: MutNullableDom<WebGLShader>,
    #[ignore_malloc_size_of = "Defined in ipc-channel"]
    renderer: WebGLMsgSender,
}

/// ANGLE adds a `_u` prefix to variable names:
///
/// https://chromium.googlesource.com/angle/angle/+/855d964bd0d05f6b2cb303f625506cf53d37e94f
///
/// To avoid hard-coding this we would need to use the `sh::GetAttributes` and `sh::GetUniforms`
/// API to look up the `x.name` and `x.mappedName` members,
/// then build a data structure for bi-directional lookup (so either linear scan or two hashmaps).
/// Even then, this would probably only support plain variable names like "foo".
/// Strings passed to e.g. `GetUniformLocation` can be expressions like "foo[0].bar",
/// with the mapping for that "bar" name in yet another part of ANGLE’s API.
const ANGLE_NAME_PREFIX: &'static str = "_u";

fn to_name_in_compiled_shader(s: &str) -> String {
    map_dot_separated(s, |s, mapped| {
        mapped.push_str(ANGLE_NAME_PREFIX);
        mapped.push_str(s);
    })
}

fn from_name_in_compiled_shader(s: &str) -> String {
    map_dot_separated(s, |s, mapped| {
        mapped.push_str(if s.starts_with(ANGLE_NAME_PREFIX) {
            &s[ANGLE_NAME_PREFIX.len()..]
        } else {
            s
        })
    })
}

fn map_dot_separated<F: Fn(&str, &mut String)>(s: &str, f: F) -> String {
    let mut iter = s.split('.');
    let mut mapped = String::new();
    f(iter.next().unwrap(), &mut mapped);
    for s in iter {
        mapped.push('.');
        f(s, &mut mapped);
    }
    mapped
}

impl WebGLProgram {
    fn new_inherited(renderer: WebGLMsgSender,
                     id: WebGLProgramId)
                     -> WebGLProgram {
        WebGLProgram {
            webgl_object: WebGLObject::new_inherited(),
            id: id,
            is_deleted: Cell::new(false),
            link_called: Cell::new(false),
            linked: Cell::new(false),
            fragment_shader: Default::default(),
            vertex_shader: Default::default(),
            renderer: renderer,
        }
    }

    pub fn maybe_new(window: &Window, renderer: WebGLMsgSender)
                     -> Option<DomRoot<WebGLProgram>> {
        let (sender, receiver) = webgl_channel().unwrap();
        renderer.send(WebGLCommand::CreateProgram(sender)).unwrap();

        let result = receiver.recv().unwrap();
        result.map(|program_id| WebGLProgram::new(window, renderer, program_id))
    }

    pub fn new(window: &Window,
               renderer: WebGLMsgSender,
               id: WebGLProgramId)
               -> DomRoot<WebGLProgram> {
        reflect_dom_object(Box::new(WebGLProgram::new_inherited(renderer, id)),
                           window,
                           WebGLProgramBinding::Wrap)
    }
}

impl WebGLProgram {
    pub fn id(&self) -> WebGLProgramId {
        self.id
    }

    /// glDeleteProgram
    pub fn delete(&self) {
        if !self.is_deleted.get() {
            self.is_deleted.set(true);
            let _ = self.renderer.send(WebGLCommand::DeleteProgram(self.id));

            if let Some(shader) = self.fragment_shader.get() {
                shader.decrement_attached_counter();
            }

            if let Some(shader) = self.vertex_shader.get() {
                shader.decrement_attached_counter();
            }
        }
    }

    pub fn is_deleted(&self) -> bool {
        self.is_deleted.get()
    }

    pub fn is_linked(&self) -> bool {
        self.linked.get()
    }

    /// glLinkProgram
    pub fn link(&self) -> WebGLResult<()>  {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        self.linked.set(false);
        self.link_called.set(true);

        match self.fragment_shader.get() {
            Some(ref shader) if shader.successfully_compiled() => {},
            _ => return Ok(()), // callers use gl.LINK_STATUS to check link errors
        }

        match self.vertex_shader.get() {
            Some(ref shader) if shader.successfully_compiled() => {},
            _ => return Ok(()), // callers use gl.LINK_STATUS to check link errors
        }

        self.linked.set(true);
        self.renderer.send(WebGLCommand::LinkProgram(self.id)).unwrap();
        Ok(())
    }

    /// glUseProgram
    pub fn use_program(&self) -> WebGLResult<()> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        if !self.linked.get() {
            return Err(WebGLError::InvalidOperation);
        }

        self.renderer.send(WebGLCommand::UseProgram(self.id)).unwrap();
        Ok(())
    }

    /// glValidateProgram
    pub fn validate(&self) -> WebGLResult<()> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        self.renderer.send(WebGLCommand::ValidateProgram(self.id)).unwrap();
        Ok(())
    }

    /// glAttachShader
    pub fn attach_shader(&self, shader: &WebGLShader) -> WebGLResult<()> {
        if self.is_deleted() || shader.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        let shader_slot = match shader.gl_type() {
            constants::FRAGMENT_SHADER => &self.fragment_shader,
            constants::VERTEX_SHADER => &self.vertex_shader,
            _ => {
                error!("detachShader: Unexpected shader type");
                return Err(WebGLError::InvalidValue);
            }
        };

        // TODO(emilio): Differentiate between same shader already assigned and other previous
        // shader.
        if shader_slot.get().is_some() {
            return Err(WebGLError::InvalidOperation);
        }

        shader_slot.set(Some(shader));
        shader.increment_attached_counter();

        self.renderer.send(WebGLCommand::AttachShader(self.id, shader.id())).unwrap();

        Ok(())
    }

    /// glDetachShader
    pub fn detach_shader(&self, shader: &WebGLShader) -> WebGLResult<()> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        let shader_slot = match shader.gl_type() {
            constants::FRAGMENT_SHADER => &self.fragment_shader,
            constants::VERTEX_SHADER => &self.vertex_shader,
            _ => {
                error!("detachShader: Unexpected shader type");
                return Err(WebGLError::InvalidValue);
            }
        };

        match shader_slot.get() {
            Some(ref attached_shader) if attached_shader.id() != shader.id() =>
                return Err(WebGLError::InvalidOperation),
            None =>
                return Err(WebGLError::InvalidOperation),
            _ => {}
        }

        shader_slot.set(None);
        shader.decrement_attached_counter();

        self.renderer.send(WebGLCommand::DetachShader(self.id, shader.id())).unwrap();

        Ok(())
    }

    /// glBindAttribLocation
    pub fn bind_attrib_location(&self, index: u32, name: DOMString) -> WebGLResult<()> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        if name.len() > MAX_UNIFORM_AND_ATTRIBUTE_LEN {
            return Err(WebGLError::InvalidValue);
        }

        // Check if the name is reserved
        if name.starts_with("gl_") || name.starts_with("webgl") || name.starts_with("_webgl_") {
            return Err(WebGLError::InvalidOperation);
        }

        let name = to_name_in_compiled_shader(&name);

        self.renderer
            .send(WebGLCommand::BindAttribLocation(self.id, index, name))
            .unwrap();
        Ok(())
    }

    pub fn get_active_uniform(&self, index: u32) -> WebGLResult<DomRoot<WebGLActiveInfo>> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidValue);
        }
        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer
            .send(WebGLCommand::GetActiveUniform(self.id, index, sender))
            .unwrap();

        receiver.recv().unwrap().map(|(size, ty, name)| {
            let name = DOMString::from(from_name_in_compiled_shader(&name));
            WebGLActiveInfo::new(self.global().as_window(), size, ty, name)
        })
    }

    /// glGetActiveAttrib
    pub fn get_active_attrib(&self, index: u32) -> WebGLResult<DomRoot<WebGLActiveInfo>> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidValue);
        }
        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer
            .send(WebGLCommand::GetActiveAttrib(self.id, index, sender))
            .unwrap();

        receiver.recv().unwrap().map(|(size, ty, name)| {
            let name = DOMString::from(from_name_in_compiled_shader(&name));
            WebGLActiveInfo::new(self.global().as_window(), size, ty, name)
        })
    }

    /// glGetAttribLocation
    pub fn get_attrib_location(&self, name: DOMString) -> WebGLResult<Option<i32>> {
        if !self.is_linked() || self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        if name.len() > MAX_UNIFORM_AND_ATTRIBUTE_LEN {
            return Err(WebGLError::InvalidValue);
        }

        // Check if the name is reserved
        if name.starts_with("gl_") {
            return Err(WebGLError::InvalidOperation);
        }

        if name.starts_with("webgl") || name.starts_with("_webgl_") {
            return Ok(None);
        }

        let name = to_name_in_compiled_shader(&name);

        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer
            .send(WebGLCommand::GetAttribLocation(self.id, name, sender))
            .unwrap();
        Ok(receiver.recv().unwrap())
    }

    /// glGetUniformLocation
    pub fn get_uniform_location(&self, name: DOMString) -> WebGLResult<Option<i32>> {
        if !self.is_linked() || self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        if name.len() > MAX_UNIFORM_AND_ATTRIBUTE_LEN {
            return Err(WebGLError::InvalidValue);
        }

        // Check if the name is reserved
        if name.starts_with("webgl") || name.starts_with("_webgl_") {
            return Ok(None);
        }

        let name = to_name_in_compiled_shader(&name);

        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer
            .send(WebGLCommand::GetUniformLocation(self.id, name, sender))
            .unwrap();
        Ok(receiver.recv().unwrap())
    }

    /// glGetProgramInfoLog
    pub fn get_info_log(&self) -> WebGLResult<String> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        if self.link_called.get() {
            let shaders_compiled = match (self.fragment_shader.get(), self.vertex_shader.get()) {
                (Some(fs), Some(vs)) => fs.successfully_compiled() && vs.successfully_compiled(),
                _ => false
            };
            if !shaders_compiled {
                return Ok("One or more shaders failed to compile".to_string());
            }
        }
        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer.send(WebGLCommand::GetProgramInfoLog(self.id, sender)).unwrap();
        Ok(receiver.recv().unwrap())
    }

    /// glGetProgramParameter
    pub fn parameter(&self, param_id: u32) -> WebGLResult<WebGLParameter> {
        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer.send(WebGLCommand::GetProgramParameter(self.id, param_id, sender)).unwrap();
        receiver.recv().unwrap()
    }

    pub fn attached_shaders(&self) -> WebGLResult<Vec<DomRoot<WebGLShader>>> {
        if self.is_deleted.get() {
            return Err(WebGLError::InvalidValue);
        }
        Ok(match (self.vertex_shader.get(), self.fragment_shader.get()) {
            (Some(vertex_shader), Some(fragment_shader)) => {
                vec![vertex_shader, fragment_shader]
            }
            (Some(shader), None) | (None, Some(shader)) => vec![shader],
            (None, None) => vec![]
        })
    }
}

impl Drop for WebGLProgram {
    fn drop(&mut self) {
        self.delete();
    }
}
<commit_msg>Fix filtering of reserved WebGL names in gl.getAttribLocation<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// https://www.khronos.org/registry/webgl/specs/latest/1.0/webgl.idl
use canvas_traits::webgl::{WebGLCommand, WebGLError, WebGLMsgSender, WebGLParameter, WebGLProgramId, WebGLResult};
use canvas_traits::webgl::webgl_channel;
use dom::bindings::codegen::Bindings::WebGLProgramBinding;
use dom::bindings::codegen::Bindings::WebGLRenderingContextBinding::WebGLRenderingContextConstants as constants;
use dom::bindings::reflector::{DomObject, reflect_dom_object};
use dom::bindings::root::{DomRoot, MutNullableDom};
use dom::bindings::str::DOMString;
use dom::webglactiveinfo::WebGLActiveInfo;
use dom::webglobject::WebGLObject;
use dom::webglrenderingcontext::MAX_UNIFORM_AND_ATTRIBUTE_LEN;
use dom::webglshader::WebGLShader;
use dom::window::Window;
use dom_struct::dom_struct;
use std::cell::Cell;

#[dom_struct]
pub struct WebGLProgram {
    webgl_object: WebGLObject,
    id: WebGLProgramId,
    is_deleted: Cell<bool>,
    link_called: Cell<bool>,
    linked: Cell<bool>,
    fragment_shader: MutNullableDom<WebGLShader>,
    vertex_shader: MutNullableDom<WebGLShader>,
    #[ignore_malloc_size_of = "Defined in ipc-channel"]
    renderer: WebGLMsgSender,
}

/// ANGLE adds a `_u` prefix to variable names:
///
/// https://chromium.googlesource.com/angle/angle/+/855d964bd0d05f6b2cb303f625506cf53d37e94f
///
/// To avoid hard-coding this we would need to use the `sh::GetAttributes` and `sh::GetUniforms`
/// API to look up the `x.name` and `x.mappedName` members,
/// then build a data structure for bi-directional lookup (so either linear scan or two hashmaps).
/// Even then, this would probably only support plain variable names like "foo".
/// Strings passed to e.g. `GetUniformLocation` can be expressions like "foo[0].bar",
/// with the mapping for that "bar" name in yet another part of ANGLE’s API.
const ANGLE_NAME_PREFIX: &'static str = "_u";

fn to_name_in_compiled_shader(s: &str) -> String {
    map_dot_separated(s, |s, mapped| {
        mapped.push_str(ANGLE_NAME_PREFIX);
        mapped.push_str(s);
    })
}

fn from_name_in_compiled_shader(s: &str) -> String {
    map_dot_separated(s, |s, mapped| {
        mapped.push_str(if s.starts_with(ANGLE_NAME_PREFIX) {
            &s[ANGLE_NAME_PREFIX.len()..]
        } else {
            s
        })
    })
}

fn map_dot_separated<F: Fn(&str, &mut String)>(s: &str, f: F) -> String {
    let mut iter = s.split('.');
    let mut mapped = String::new();
    f(iter.next().unwrap(), &mut mapped);
    for s in iter {
        mapped.push('.');
        f(s, &mut mapped);
    }
    mapped
}

impl WebGLProgram {
    fn new_inherited(renderer: WebGLMsgSender,
                     id: WebGLProgramId)
                     -> WebGLProgram {
        WebGLProgram {
            webgl_object: WebGLObject::new_inherited(),
            id: id,
            is_deleted: Cell::new(false),
            link_called: Cell::new(false),
            linked: Cell::new(false),
            fragment_shader: Default::default(),
            vertex_shader: Default::default(),
            renderer: renderer,
        }
    }

    pub fn maybe_new(window: &Window, renderer: WebGLMsgSender)
                     -> Option<DomRoot<WebGLProgram>> {
        let (sender, receiver) = webgl_channel().unwrap();
        renderer.send(WebGLCommand::CreateProgram(sender)).unwrap();

        let result = receiver.recv().unwrap();
        result.map(|program_id| WebGLProgram::new(window, renderer, program_id))
    }

    pub fn new(window: &Window,
               renderer: WebGLMsgSender,
               id: WebGLProgramId)
               -> DomRoot<WebGLProgram> {
        reflect_dom_object(Box::new(WebGLProgram::new_inherited(renderer, id)),
                           window,
                           WebGLProgramBinding::Wrap)
    }
}

impl WebGLProgram {
    pub fn id(&self) -> WebGLProgramId {
        self.id
    }

    /// glDeleteProgram
    pub fn delete(&self) {
        if !self.is_deleted.get() {
            self.is_deleted.set(true);
            let _ = self.renderer.send(WebGLCommand::DeleteProgram(self.id));

            if let Some(shader) = self.fragment_shader.get() {
                shader.decrement_attached_counter();
            }

            if let Some(shader) = self.vertex_shader.get() {
                shader.decrement_attached_counter();
            }
        }
    }

    pub fn is_deleted(&self) -> bool {
        self.is_deleted.get()
    }

    pub fn is_linked(&self) -> bool {
        self.linked.get()
    }

    /// glLinkProgram
    pub fn link(&self) -> WebGLResult<()>  {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        self.linked.set(false);
        self.link_called.set(true);

        match self.fragment_shader.get() {
            Some(ref shader) if shader.successfully_compiled() => {},
            _ => return Ok(()), // callers use gl.LINK_STATUS to check link errors
        }

        match self.vertex_shader.get() {
            Some(ref shader) if shader.successfully_compiled() => {},
            _ => return Ok(()), // callers use gl.LINK_STATUS to check link errors
        }

        self.linked.set(true);
        self.renderer.send(WebGLCommand::LinkProgram(self.id)).unwrap();
        Ok(())
    }

    /// glUseProgram
    pub fn use_program(&self) -> WebGLResult<()> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        if !self.linked.get() {
            return Err(WebGLError::InvalidOperation);
        }

        self.renderer.send(WebGLCommand::UseProgram(self.id)).unwrap();
        Ok(())
    }

    /// glValidateProgram
    pub fn validate(&self) -> WebGLResult<()> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        self.renderer.send(WebGLCommand::ValidateProgram(self.id)).unwrap();
        Ok(())
    }

    /// glAttachShader
    pub fn attach_shader(&self, shader: &WebGLShader) -> WebGLResult<()> {
        if self.is_deleted() || shader.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        let shader_slot = match shader.gl_type() {
            constants::FRAGMENT_SHADER => &self.fragment_shader,
            constants::VERTEX_SHADER => &self.vertex_shader,
            _ => {
                error!("detachShader: Unexpected shader type");
                return Err(WebGLError::InvalidValue);
            }
        };

        // TODO(emilio): Differentiate between same shader already assigned and other previous
        // shader.
        if shader_slot.get().is_some() {
            return Err(WebGLError::InvalidOperation);
        }

        shader_slot.set(Some(shader));
        shader.increment_attached_counter();

        self.renderer.send(WebGLCommand::AttachShader(self.id, shader.id())).unwrap();

        Ok(())
    }

    /// glDetachShader
    pub fn detach_shader(&self, shader: &WebGLShader) -> WebGLResult<()> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        let shader_slot = match shader.gl_type() {
            constants::FRAGMENT_SHADER => &self.fragment_shader,
            constants::VERTEX_SHADER => &self.vertex_shader,
            _ => {
                error!("detachShader: Unexpected shader type");
                return Err(WebGLError::InvalidValue);
            }
        };

        match shader_slot.get() {
            Some(ref attached_shader) if attached_shader.id() != shader.id() =>
                return Err(WebGLError::InvalidOperation),
            None =>
                return Err(WebGLError::InvalidOperation),
            _ => {}
        }

        shader_slot.set(None);
        shader.decrement_attached_counter();

        self.renderer.send(WebGLCommand::DetachShader(self.id, shader.id())).unwrap();

        Ok(())
    }

    /// glBindAttribLocation
    pub fn bind_attrib_location(&self, index: u32, name: DOMString) -> WebGLResult<()> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        if name.len() > MAX_UNIFORM_AND_ATTRIBUTE_LEN {
            return Err(WebGLError::InvalidValue);
        }

        // Check if the name is reserved
        if name.starts_with("gl_") || name.starts_with("webgl") || name.starts_with("_webgl_") {
            return Err(WebGLError::InvalidOperation);
        }

        let name = to_name_in_compiled_shader(&name);

        self.renderer
            .send(WebGLCommand::BindAttribLocation(self.id, index, name))
            .unwrap();
        Ok(())
    }

    pub fn get_active_uniform(&self, index: u32) -> WebGLResult<DomRoot<WebGLActiveInfo>> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidValue);
        }
        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer
            .send(WebGLCommand::GetActiveUniform(self.id, index, sender))
            .unwrap();

        receiver.recv().unwrap().map(|(size, ty, name)| {
            let name = DOMString::from(from_name_in_compiled_shader(&name));
            WebGLActiveInfo::new(self.global().as_window(), size, ty, name)
        })
    }

    /// glGetActiveAttrib
    pub fn get_active_attrib(&self, index: u32) -> WebGLResult<DomRoot<WebGLActiveInfo>> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidValue);
        }
        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer
            .send(WebGLCommand::GetActiveAttrib(self.id, index, sender))
            .unwrap();

        receiver.recv().unwrap().map(|(size, ty, name)| {
            let name = DOMString::from(from_name_in_compiled_shader(&name));
            WebGLActiveInfo::new(self.global().as_window(), size, ty, name)
        })
    }

    /// glGetAttribLocation
    pub fn get_attrib_location(&self, name: DOMString) -> WebGLResult<Option<i32>> {
        if !self.is_linked() || self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        if name.len() > MAX_UNIFORM_AND_ATTRIBUTE_LEN {
            return Err(WebGLError::InvalidValue);
        }

        // Check if the name is reserved
        if name.starts_with("gl_") {
            return Err(WebGLError::InvalidOperation);
        }

        // https://www.khronos.org/registry/webgl/specs/latest/1.0/#GLSL_CONSTRUCTS
        if name.starts_with("webgl_") || name.starts_with("_webgl_") {
            return Ok(None);
        }

        let name = to_name_in_compiled_shader(&name);

        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer
            .send(WebGLCommand::GetAttribLocation(self.id, name, sender))
            .unwrap();
        Ok(receiver.recv().unwrap())
    }

    /// glGetUniformLocation
    pub fn get_uniform_location(&self, name: DOMString) -> WebGLResult<Option<i32>> {
        if !self.is_linked() || self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        if name.len() > MAX_UNIFORM_AND_ATTRIBUTE_LEN {
            return Err(WebGLError::InvalidValue);
        }

        // Check if the name is reserved
        if name.starts_with("webgl") || name.starts_with("_webgl_") {
            return Ok(None);
        }

        let name = to_name_in_compiled_shader(&name);

        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer
            .send(WebGLCommand::GetUniformLocation(self.id, name, sender))
            .unwrap();
        Ok(receiver.recv().unwrap())
    }

    /// glGetProgramInfoLog
    pub fn get_info_log(&self) -> WebGLResult<String> {
        if self.is_deleted() {
            return Err(WebGLError::InvalidOperation);
        }
        if self.link_called.get() {
            let shaders_compiled = match (self.fragment_shader.get(), self.vertex_shader.get()) {
                (Some(fs), Some(vs)) => fs.successfully_compiled() && vs.successfully_compiled(),
                _ => false
            };
            if !shaders_compiled {
                return Ok("One or more shaders failed to compile".to_string());
            }
        }
        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer.send(WebGLCommand::GetProgramInfoLog(self.id, sender)).unwrap();
        Ok(receiver.recv().unwrap())
    }

    /// glGetProgramParameter
    pub fn parameter(&self, param_id: u32) -> WebGLResult<WebGLParameter> {
        let (sender, receiver) = webgl_channel().unwrap();
        self.renderer.send(WebGLCommand::GetProgramParameter(self.id, param_id, sender)).unwrap();
        receiver.recv().unwrap()
    }

    pub fn attached_shaders(&self) -> WebGLResult<Vec<DomRoot<WebGLShader>>> {
        if self.is_deleted.get() {
            return Err(WebGLError::InvalidValue);
        }
        Ok(match (self.vertex_shader.get(), self.fragment_shader.get()) {
            (Some(vertex_shader), Some(fragment_shader)) => {
                vec![vertex_shader, fragment_shader]
            }
            (Some(shader), None) | (None, Some(shader)) => vec![shader],
            (None, None) => vec![]
        })
    }
}

impl Drop for WebGLProgram {
    fn drop(&mut self) {
        self.delete();
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 click2stream, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#[cfg(feature="discovery")]
use std::thread;

use std::thread::JoinHandle;

use context::ApplicationContext;

#[cfg(feature = "discovery")]
use scanner::discovery;

#[cfg(feature="discovery")]
use utils;

use utils::logger::{BoxedLogger, Logger};

#[cfg(feature="discovery")]
use utils::logger::Severity;

use futures::{BoxFuture, Future, Poll, Stream};

use futures::sync::mpsc;
use futures::sync::mpsc::{UnboundedReceiver, UnboundedSender};

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Command {
    ResetServiceTable,
    ScanNetwork,
}

#[derive(Debug, Copy, Clone)]
enum Event {
    Command(Command),
    ScanCompleted,
}

type CommandReceiver = UnboundedReceiver<Event>;

type CommandSender = UnboundedSender<Event>;

/// A channel for sending commands across threads. The channel is cloneable and every copy of it
/// will send commands to the same handler.
#[derive(Clone)]
pub struct CommandChannel {
    tx: CommandSender,
}

impl CommandChannel {
    /// Create a new command channel.
    fn new(tx: CommandSender) -> CommandChannel {
        CommandChannel {
            tx: tx,
        }
    }

    /// Send a given command.
    pub fn send(&self, cmd: Command) {
        self.tx.send(Event::Command(cmd))
            .expect("broken command channel");
    }
}

/// Command handler context.
struct CommandHandlerContext {
    app_context: ApplicationContext,
    logger:      BoxedLogger,
    scanner:     Option<JoinHandle<()>>,
}

impl CommandHandlerContext {
    /// Create a new command handler context.
    fn new(app_context: ApplicationContext) -> CommandHandlerContext {
        let logger = app_context.get_logger();

        CommandHandlerContext {
            app_context: app_context,
            logger:      logger,
            scanner:     None,
        }
    }

    /// Process a given command handler event.
    fn proces_event(&mut self, event: Event) {
        match event {
            Event::Command(cmd)  => self.process_command(cmd),
            Event::ScanCompleted => self.scan_completed(),
        }
    }

    /// Process a given command.
    fn process_command(&mut self, cmd: Command) {
        match cmd {
            Command::ResetServiceTable => self.reset_service_table(),
            Command::ScanNetwork       => self.scan_network(),
        }
    }

    /// Reset service table.
    fn reset_service_table(&mut self) {
        self.app_context.reset_service_table()
    }

    #[cfg(feature="discovery")]
    /// Trigger a network scan.
    fn scan_network(&mut self) {
        if self.scanner.is_some() {
            return;
        }

        let app_context = self.app_context.clone();

        let handle = thread::spawn(move || {
            network_scanner_thread(app_context);
        });

        self.scanner = Some(handle);
    }

    #[cfg(not(feature="discovery"))]
    /// Dummy network scan.
    fn scan_network(&mut self) {
    }

    /// Cleanup the scanner context.
    fn scan_completed(&mut self) {
        if let Some(handle) = self.scanner.take() {
            if let Err(_) = handle.join() {
                log_warn!(self.logger, "network scanner thread panicked");
            }
        }
    }
}

/// Command handler. It implements the future trait and it's designed to be used in combination
/// with tokio event loop.
pub struct CommandHandler {
    handler: BoxFuture<(), ()>,
}

impl CommandHandler {
    /// Create a new command handler.
    fn new(app_context: ApplicationContext, rx: CommandReceiver) -> CommandHandler {
        let mut context = CommandHandlerContext::new(app_context);

        let handler = rx.for_each(move |event| {
            context.proces_event(event);
            Ok(())
        });

        CommandHandler {
            handler: handler.boxed(),
        }
    }
}

impl Future for CommandHandler {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.handler.poll()
    }
}

/// Create a new channel-handler pair.
pub fn new(app_context: ApplicationContext) -> (CommandChannel, CommandHandler) {
    let (tx, rx) = mpsc::unbounded();

    let rx = CommandHandler::new(app_context, rx);
    let tx = CommandChannel::new(tx);

    (tx, rx)
}

#[cfg(feature = "discovery")]
/// Run device discovery and update the service table.
fn network_scanner_thread(mut app_context: ApplicationContext) {
    let mut logger = app_context.get_logger();

    let rtsp_paths_file = app_context.get_rtsp_paths_file();
    let mjpeg_paths_file = app_context.get_mjpeg_paths_file();

    app_context.set_scanning(true);

    log_info!(logger, "looking for local services...");

    let result = utils::result_or_log(&mut logger, Severity::WARN,
        "network scanner error",
        discovery::scan_network(
            &rtsp_paths_file,
            &mjpeg_paths_file));

    if let Some(result) = result {
        let services = result.services()
            .map(|svc| svc.clone())
            .collect::<Vec<_>>();

        let count = services.len();

        app_context.update_service_table(services);
        app_context.set_scan_result(result);

        log_info!(logger, "{} services found, current service table: {}",
            count, app_context.get_service_table());
    }

    app_context.set_scanning(false);
}
<commit_msg>Minor fix in CommandHandler<commit_after>// Copyright 2017 click2stream, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#[cfg(feature="discovery")]
use std::thread;

use std::thread::JoinHandle;

use context::ApplicationContext;

#[cfg(feature = "discovery")]
use scanner::discovery;

#[cfg(feature="discovery")]
use utils;

use utils::logger::{BoxedLogger, Logger};

#[cfg(feature="discovery")]
use utils::logger::Severity;

use futures::{BoxFuture, Future, Poll, Stream};

use futures::sync::mpsc;
use futures::sync::mpsc::{UnboundedReceiver, UnboundedSender};

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Command {
    ResetServiceTable,
    ScanNetwork,
}

#[derive(Debug, Copy, Clone)]
enum Event {
    Command(Command),
    ScanCompleted,
}

type CommandReceiver = UnboundedReceiver<Event>;

type CommandSender = UnboundedSender<Event>;

/// A channel for sending commands across threads. The channel is cloneable and every copy of it
/// will send commands to the same handler.
#[derive(Clone)]
pub struct CommandChannel {
    tx: CommandSender,
}

impl CommandChannel {
    /// Create a new command channel.
    fn new(tx: CommandSender) -> CommandChannel {
        CommandChannel {
            tx: tx,
        }
    }

    /// Send a given command.
    pub fn send(&self, cmd: Command) {
        self.tx.send(Event::Command(cmd))
            .expect("broken command channel");
    }
}

/// Command handler context.
struct CommandHandlerContext {
    app_context: ApplicationContext,
    logger:      BoxedLogger,
    cmd_sender:  CommandSender,
    scanner:     Option<JoinHandle<()>>,
}

impl CommandHandlerContext {
    /// Create a new command handler context.
    fn new(app_context: ApplicationContext, cmd_sender: CommandSender) -> CommandHandlerContext {
        let logger = app_context.get_logger();

        CommandHandlerContext {
            app_context: app_context,
            logger:      logger,
            cmd_sender:  cmd_sender,
            scanner:     None,
        }
    }

    /// Process a given command handler event.
    fn proces_event(&mut self, event: Event) {
        match event {
            Event::Command(cmd)  => self.process_command(cmd),
            Event::ScanCompleted => self.scan_completed(),
        }
    }

    /// Process a given command.
    fn process_command(&mut self, cmd: Command) {
        match cmd {
            Command::ResetServiceTable => self.reset_service_table(),
            Command::ScanNetwork       => self.scan_network(),
        }
    }

    /// Reset service table.
    fn reset_service_table(&mut self) {
        self.app_context.reset_service_table()
    }

    #[cfg(feature="discovery")]
    /// Trigger a network scan.
    fn scan_network(&mut self) {
        if self.scanner.is_some() {
            return;
        }

        let app_context = self.app_context.clone();
        let cmd_sender  = self.cmd_sender.clone();

        let handle = thread::spawn(move || {
            network_scanner_thread(app_context);

            cmd_sender.send(Event::ScanCompleted)
                .unwrap();
        });

        self.scanner = Some(handle);
    }

    #[cfg(not(feature="discovery"))]
    /// Dummy network scan.
    fn scan_network(&mut self) {
    }

    /// Cleanup the scanner context.
    fn scan_completed(&mut self) {
        if let Some(handle) = self.scanner.take() {
            if let Err(_) = handle.join() {
                log_warn!(self.logger, "network scanner thread panicked");
            }
        }
    }
}

/// Command handler. It implements the future trait and it's designed to be used in combination
/// with tokio event loop.
pub struct CommandHandler {
    handler: BoxFuture<(), ()>,
}

impl CommandHandler {
    /// Create a new command handler.
    fn new(
        app_context: ApplicationContext,
        rx: CommandReceiver,
        tx: CommandSender) -> CommandHandler {
        let mut context = CommandHandlerContext::new(app_context, tx);

        let handler = rx.for_each(move |event| {
            context.proces_event(event);
            Ok(())
        });

        CommandHandler {
            handler: handler.boxed(),
        }
    }
}

impl Future for CommandHandler {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.handler.poll()
    }
}

/// Create a new channel-handler pair.
pub fn new(app_context: ApplicationContext) -> (CommandChannel, CommandHandler) {
    let (tx, rx) = mpsc::unbounded();

    let rx = CommandHandler::new(app_context, rx, tx.clone());
    let tx = CommandChannel::new(tx);

    (tx, rx)
}

#[cfg(feature = "discovery")]
/// Run device discovery and update the service table.
fn network_scanner_thread(mut app_context: ApplicationContext) {
    let mut logger = app_context.get_logger();

    let rtsp_paths_file = app_context.get_rtsp_paths_file();
    let mjpeg_paths_file = app_context.get_mjpeg_paths_file();

    app_context.set_scanning(true);

    log_info!(logger, "looking for local services...");

    let result = utils::result_or_log(&mut logger, Severity::WARN,
        "network scanner error",
        discovery::scan_network(
            &rtsp_paths_file,
            &mjpeg_paths_file));

    if let Some(result) = result {
        let services = result.services()
            .map(|svc| svc.clone())
            .collect::<Vec<_>>();

        let count = services.len();

        app_context.update_service_table(services);
        app_context.set_scan_result(result);

        log_info!(logger, "{} services found, current service table: {}",
            count, app_context.get_service_table());
    }

    app_context.set_scanning(false);
}
<|endoftext|>