<commit_before>use super::Time;
use crate::airport_data::AirportType;
use serde_derive::{Deserialize, Serialize};
use smol_str::SmolStr;

#[derive(Debug, Deserialize, Serialize)]
pub struct Filters {
    pub speed: SpeedFilter,
    pub departure: Option<AirportFilters>,
    pub arrival: Option<AirportFilters>,
    #[serde(rename = "timeRange", default)]
    pub time_range: TimeRange,
}

// TODO: implement as discriminated union and implement Serialize / Deserialize manually
#[derive(Debug, Deserialize, Serialize)]
pub struct SpeedFilter {
    pub value: SmolStr,
    #[serde(rename = "type")]
    pub variant: SpeedType,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum SpeedType {
    Mach,
    Knots,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct AirportFilters {
    pub icao: SmolStr,
    #[serde(rename = "airportType", default)]
    pub airport_type: AirportType,
    #[serde(rename = "runwayLength")]
    pub runway_length: Option<RunwayLength>,
    #[serde(default)]
    pub countries: Vec<String>,
}

// TODO: implement as discriminated union and implement Serialize / Deserialize manually
#[derive(Debug, Deserialize, Serialize)]
pub struct RunwayLength {
    pub length: u32,
    pub selector: LengthSelector,
}

#[derive(Debug, Deserialize, Serialize)]
pub enum LengthSelector {
    #[serde(rename = "eq")]
    Equal,
    #[serde(rename = "gt")]
    GreaterThan,
    #[serde(rename = "lt")]
    LessThan,
}

#[derive(Debug, Default, Deserialize, Serialize)]
pub struct TimeRange {
    pub min: Option<Time>,
    pub max: Option<Time>,
}
<commit_msg>backend: Condense SpeedFilter type into discriminated union<commit_after>use super::Time;
use crate::airport_data::AirportType;
use serde_derive::Deserialize;
use smol_str::SmolStr;

#[derive(Debug, Deserialize)]
pub struct Filters {
    pub speed: SpeedFilter,
    pub departure: Option<AirportFilters>,
    pub arrival: Option<AirportFilters>,
    #[serde(rename = "timeRange", default)]
    pub time_range: TimeRange,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "lowercase")]
#[serde(tag = "type", content = "value")]
pub enum SpeedFilter {
    Mach(f32),
    Knots(u32),
}

#[derive(Debug, Deserialize)]
pub struct AirportFilters {
    pub icao: SmolStr,
    #[serde(rename = "airportType", default)]
    pub airport_type: AirportType,
    #[serde(rename = "runwayLength")]
    pub runway_length: Option<RunwayLength>,
    #[serde(default)]
    pub countries: Vec<String>,
}

// TODO: implement as discriminated union and implement Serialize / Deserialize manually
#[derive(Debug, Deserialize)]
pub struct RunwayLength {
    pub length: u32,
    pub selector: LengthSelector,
}

#[derive(Debug, Deserialize)]
pub enum LengthSelector {
    #[serde(rename = "eq")]
    Equal,
    #[serde(rename = "gt")]
    GreaterThan,
    #[serde(rename = "lt")]
    LessThan,
}

#[derive(Debug, Default, Deserialize)]
pub struct TimeRange {
    pub min: Option<Time>,
    pub max: Option<Time>,
}
<|endoftext|>


<commit_before>use std::ops::BitXor;

use super::active_state::ActiveState;
use super::crystallized_state::CrystallizedState;
use super::recent_proposer_record::RecentPropserRecord;
use super::block::Block;
use super::utils::types::Sha256Digest;
use super::utils::logging::Logger;
use super::config::Config;
use super::rlp;

use super::attestors::{
    process_attestations,
    get_attesters_and_proposer
};
use super::ffg::update_ffg_and_crosslink_progress;

pub fn compute_new_active_state(
    cry_state: &CrystallizedState,
    act_state: &ActiveState,
    parent_block: &Block,
    block: &Block,
    config: &Config,
    log: &Logger)
    -> ActiveState
{
    let (attestation_indicies, proposer) = get_attesters_and_proposer(
        &cry_state,
        &act_state,
        &block.skip_count,
        &config,
        &log);

    info!(log, "calculated attesters and proposers";
          "attesters_count" => attestation_indicies.len(),
          "proposer_index" => proposer);

    let parent_block_rlp = rlp::encode(parent_block);
    let attesters_option = process_attestations(
        &cry_state.active_validators,
        &attestation_indicies,
        &block.attestation_bitfield,
        &parent_block_rlp.to_vec(),
        &block.attestation_aggregate_sig);

    // TODO: bls verify signature here.

    let (partial_crosslinks, ffg_voter_bitfield, total_new_voters) =
        update_ffg_and_crosslink_progress(
            &cry_state,
            &act_state.partial_crosslinks,
            &act_state.ffg_voter_bitfield,
            &block.shard_aggregate_votes,
            &config);

    let attesters = match attesters_option {
        None => panic!("No attestors available"),
        Some(x) => x
    };

    let proposer = RecentPropserRecord {
        index: proposer,
        randao_commitment: Sha256Digest::zero(),
        balance_delta: (attesters.len() + total_new_voters) as i64
    };

    let height = act_state.height + 1;
    let randao = act_state.randao.bitxor(block.randao_reveal);
    let mut recent_attesters = act_state.recent_attesters.to_vec();
    recent_attesters.extend_from_slice(&attesters);
    let total_skip_count = act_state.total_skip_count + block.skip_count;
    let mut recent_proposers = act_state.recent_proposers.to_vec();
    recent_proposers.push(proposer);

    ActiveState {
        height,
        randao,
        ffg_voter_bitfield,
        recent_attesters,
        partial_crosslinks,
        total_skip_count,
        recent_proposers
    }
}
<commit_msg>Tidy up new_active_state<commit_after>use std::ops::BitXor;

use super::active_state::ActiveState;
use super::crystallized_state::CrystallizedState;
use super::recent_proposer_record::RecentPropserRecord;
use super::block::Block;
use super::utils::types::Sha256Digest;
use super::utils::logging::Logger;
use super::config::Config;
use super::super::pubkeystore::verify_block;
use super::rlp;

use super::attestors::{
    process_attestations,
    get_attesters_and_proposer
};
use super::ffg::update_ffg_and_crosslink_progress;

/// Computes the active state generated by applying a new block
/// to a parent block.
///
/// The `CrystallizedState` and `ActiveState` represented by
/// parent_block.state_hash must also be supplied.
///
/// # Panics
/// Will panic if the block verification or state transition fails.
/// Some situations may include; signature verification failure or
/// insufficient active validators.
pub fn compute_new_active_state(
    cry_state: &CrystallizedState,
    act_state: &ActiveState,
    parent_block: &Block,
    new_block: &Block,
    config: &Config,
    log: &Logger)
    -> ActiveState
{
    /*
     * For the present state pair, determine the proposer and attester
     * set for the newly presented block (`new_block`).
     */
    let (attestation_indicies, proposer) = get_attesters_and_proposer(
        &cry_state,
        &act_state,
        &new_block.skip_count,
        &config,
        &log);

    info!(log, "calculated attesters and proposers";
          "attesters_count" => attestation_indicies.len(),
          "proposer_index" => proposer);

    /*
     *  Ensure that the new_block bitfield accurately represents the
     *  public keys which have signed the previous block. I.e., for each
     *  `true` bit in the new_block bitfield, ensure the corresponding
     *  attester has signed the aggregate signature.
     *
     *  As it is possible for there to be no attesters, an option
     *  is returned. The option unwraps to a vec of validator indicies
     *  which attested to the block.
     */
    let parent_block_rlp = rlp::encode(parent_block);
    let attesters_option = process_attestations(
        &cry_state.active_validators,
        &attestation_indicies,
        &new_block.attestation_bitfield,
        &parent_block_rlp.to_vec(),
        &new_block.attestation_aggregate_sig);
    let attesters = attesters_option.expect("No attesters available.");

    /*
     * Ensure the new_block has been signed by the proposer dictated
     * by the state pair.
     */
    let verified = verify_block(&new_block, &proposer);
    assert!(verified, "Block sig verification failed.");

    info!(log, "verified block signature";
          "proposer_index" => proposer);

    /*
     * Collect all the aggregate votes and partial crosslinks
     * into a new vec of partial crosslinks. Each entry in this
     * new vec will be unique for the "shard_id" field -- each of these
     * entries will contain the binary AND of all attestation bitfields
     * for that relevant shard_id.
     *
     * During this process, the aggregate signatures are tested
     * against each attester in the respective bitfields (i.e.,
     * signatures are verified).
     *
     * Also returned is a count of all validators who attested to a
     * shard, as well as a bitfield representing the same thing.
     */
    let (partial_crosslinks, ffg_voter_bitfield, total_new_voters) =
        update_ffg_and_crosslink_progress(
            &cry_state,
            &act_state.partial_crosslinks,
            &act_state.ffg_voter_bitfield,
            &new_block.shard_aggregate_votes,
            &config);

    /*
     * Create a new RecentProposerRecord, with a reward determined
     * by:
     *
     *  - The number of attestations to last_block.
     *  - The number of crosslink attestations present in the
     *    previous active_state and new_block.
     */
    let proposer = RecentPropserRecord {
        index: proposer,
        randao_commitment: Sha256Digest::zero(),
        balance_delta: (attesters.len() + total_new_voters) as i64
    };

    /*
     * Determine fields for the new active state, then return it.
     */
    let height = act_state.height + 1;
    let randao = act_state.randao.bitxor(new_block.randao_reveal);
    let mut recent_attesters = act_state.recent_attesters.to_vec();
    recent_attesters.extend_from_slice(&attesters);
    let total_skip_count = act_state.total_skip_count + new_block.skip_count;
    let mut recent_proposers = act_state.recent_proposers.to_vec();
    recent_proposers.push(proposer);

    ActiveState {
        height,
        randao,
        ffg_voter_bitfield,
        recent_attesters,
        partial_crosslinks,
        total_skip_count,
        recent_proposers
    } }
<|endoftext|>


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

// Test struct inheritance on structs from another crate.
#![feature(struct_inherit)]

pub virtual struct S1 {
    pub f1: isize,
}

pub struct S2 : S1 {
    pub f2: isize,
}

pub fn test_s2(s2: S2) {
    assert!(s2.f1 == 115);
    assert!(s2.f2 == 113);
}

pub static glob_s: S2 = S2 { f1: 32, f2: -45 };
<commit_msg>Remove auxiliary file not used since 17da4c7<commit_after><|endoftext|>


<commit_before>use crate::compiler::Compiler;
use byteorder::{LittleEndian, WriteBytesExt};
use cranelift_module::{DataContext, Linkage};
use failure::Error;

/// Formats and stores WASM data segment initializers[0] from the program to
/// symbols in the data section of obj.
///
/// - The segment initialization data is stored to the "wasm_data_segment"
///     symbol
/// - The total number of bytes stored to "wasm_data_segment" symbol is
///     stored in the "wasm_data_segment_length" symbol
///
/// The program that consumes the resulting ELF object is responsible for
/// using it to initialize WASM linear memory regions.
///
/// [0] https://webassembly.github.io/spec/syntax/modules.html#data-segments
///
/// WARNING: At present, this code
/// - Does limited validation of data segments
/// - Does not coalesce data segments
/// - Uses an implicit data format for the serialized segment data, defined
///    only in the code below
pub fn compile_data_initializers(compiler: &mut Compiler) -> Result<(), Error> {
    let mut serialized: Vec<u8> = Vec::new();

    for initializer in compiler.prog.data_initializers()? {
        // Data segment has been validated in program::data.
        // memory_index is always 0 per spec, so we dont put it in data.
        let memory_index: u32 = 0;
        serialized.write_u32::<LittleEndian>(memory_index).unwrap();
        serialized
            .write_u32::<LittleEndian>(initializer.offset)
            .unwrap();
        serialized
            .write_u32::<LittleEndian>(initializer.data.len() as u32)
            .unwrap();
        serialized.extend_from_slice(initializer.data);

        // Pad to 8 bytes: this aligns the data for architectures with 4 or
        // 8 byte word sizes (i.e. everything we are likely to support at
        // least until we replaced this with an engineered serialization
        // format)

        let pad = vec![0u8; (8 - serialized.len() % 8) % 8];
        serialized.extend(pad);
    }

    let mut serialized_len: Vec<u8> = Vec::new();
    serialized_len
        .write_u32::<LittleEndian>(serialized.len() as u32)
        .unwrap();
    let mut seg_len_ctx = DataContext::new();
    seg_len_ctx.define(serialized_len.into_boxed_slice());

    let writeable = false;
    let seg_len_decl =
        compiler
            .module
            .declare_data("wasm_data_segments_len", Linkage::Export, writeable)?;
    compiler.module.define_data(seg_len_decl, &seg_len_ctx)?;

    let mut seg_ctx = DataContext::new();
    seg_ctx.define(serialized.into_boxed_slice());
    let seg_decl =
        compiler
            .module
            .declare_data("wasm_data_segments", Linkage::Export, writeable)?;
    compiler.module.define_data(seg_decl, &seg_ctx)?;

    Ok(())
}

use std::io::Cursor;

pub fn compile_sparse_page_data(compiler: &mut Compiler) -> Result<(), Error> {
    use crate::program::data::sparse::OwnedSparseData;
    let owned_data = OwnedSparseData::new(
        &compiler.prog.data_initializers()?,
        compiler.prog.heap_spec(),
    );
    let sparse_data = owned_data.sparse_data();

    let mut table_ctx = DataContext::new();
    let mut table_data: Cursor<Vec<u8>> =
        Cursor::new(Vec::with_capacity(sparse_data.pages().len() * 8 + 8));

    // The table is an array of 64-bit elements:
    //  [0] the number subsequent elements
    //  [1..] a pointer to a 4096-byte array of the contents of the page,
    //        or null if it is initialized as zero.

    table_data
        .write_u64::<LittleEndian>(sparse_data.pages().len() as u64)
        .unwrap();
    for (dix, d) in sparse_data.pages().iter().enumerate() {
        if let Some(vs) = d {
            // Define the 4096-byte array for the contents of the page
            let seg_decl = compiler.module.declare_data(
                &format!("guest_sparse_page_data_{}", dix),
                Linkage::Local,
                false,
            )?;
            let mut seg_ctx = DataContext::new();
            seg_ctx.define(vs.to_vec().into_boxed_slice());
            compiler.module.define_data(seg_decl, &seg_ctx)?;

            // Put a relocation to that array into the table:
            let seg_gv = compiler
                .module
                .declare_data_in_data(seg_decl, &mut table_ctx);
            table_ctx.write_data_addr(table_data.position() as u32, seg_gv, 0);
        }
        table_data.write_u64::<LittleEndian>(0)?;
    }

    table_ctx.define(table_data.into_inner().into_boxed_slice());
    let table_decl =
        compiler
            .module
            .declare_data("guest_sparse_page_data", Linkage::Export, false)?;
    compiler.module.define_data(table_decl, &table_ctx)?;

    Ok(())
}
<commit_msg>Fix url in comments<commit_after>use crate::compiler::Compiler;
use byteorder::{LittleEndian, WriteBytesExt};
use cranelift_module::{DataContext, Linkage};
use failure::Error;

/// Formats and stores WASM data segment initializers[0] from the program to
/// symbols in the data section of obj.
///
/// - The segment initialization data is stored to the "wasm_data_segment"
///     symbol
/// - The total number of bytes stored to "wasm_data_segment" symbol is
///     stored in the "wasm_data_segment_length" symbol
///
/// The program that consumes the resulting ELF object is responsible for
/// using it to initialize WASM linear memory regions.
///
/// [0] https://webassembly.github.io/spec/core/syntax/modules.html#data-segments
///
/// WARNING: At present, this code
/// - Does limited validation of data segments
/// - Does not coalesce data segments
/// - Uses an implicit data format for the serialized segment data, defined
///    only in the code below
pub fn compile_data_initializers(compiler: &mut Compiler) -> Result<(), Error> {
    let mut serialized: Vec<u8> = Vec::new();

    for initializer in compiler.prog.data_initializers()? {
        // Data segment has been validated in program::data.
        // memory_index is always 0 per spec, so we dont put it in data.
        let memory_index: u32 = 0;
        serialized.write_u32::<LittleEndian>(memory_index).unwrap();
        serialized
            .write_u32::<LittleEndian>(initializer.offset)
            .unwrap();
        serialized
            .write_u32::<LittleEndian>(initializer.data.len() as u32)
            .unwrap();
        serialized.extend_from_slice(initializer.data);

        // Pad to 8 bytes: this aligns the data for architectures with 4 or
        // 8 byte word sizes (i.e. everything we are likely to support at
        // least until we replaced this with an engineered serialization
        // format)

        let pad = vec![0u8; (8 - serialized.len() % 8) % 8];
        serialized.extend(pad);
    }

    let mut serialized_len: Vec<u8> = Vec::new();
    serialized_len
        .write_u32::<LittleEndian>(serialized.len() as u32)
        .unwrap();
    let mut seg_len_ctx = DataContext::new();
    seg_len_ctx.define(serialized_len.into_boxed_slice());

    let writeable = false;
    let seg_len_decl =
        compiler
            .module
            .declare_data("wasm_data_segments_len", Linkage::Export, writeable)?;
    compiler.module.define_data(seg_len_decl, &seg_len_ctx)?;

    let mut seg_ctx = DataContext::new();
    seg_ctx.define(serialized.into_boxed_slice());
    let seg_decl =
        compiler
            .module
            .declare_data("wasm_data_segments", Linkage::Export, writeable)?;
    compiler.module.define_data(seg_decl, &seg_ctx)?;

    Ok(())
}

use std::io::Cursor;

pub fn compile_sparse_page_data(compiler: &mut Compiler) -> Result<(), Error> {
    use crate::program::data::sparse::OwnedSparseData;
    let owned_data = OwnedSparseData::new(
        &compiler.prog.data_initializers()?,
        compiler.prog.heap_spec(),
    );
    let sparse_data = owned_data.sparse_data();

    let mut table_ctx = DataContext::new();
    let mut table_data: Cursor<Vec<u8>> =
        Cursor::new(Vec::with_capacity(sparse_data.pages().len() * 8 + 8));

    // The table is an array of 64-bit elements:
    //  [0] the number subsequent elements
    //  [1..] a pointer to a 4096-byte array of the contents of the page,
    //        or null if it is initialized as zero.

    table_data
        .write_u64::<LittleEndian>(sparse_data.pages().len() as u64)
        .unwrap();
    for (dix, d) in sparse_data.pages().iter().enumerate() {
        if let Some(vs) = d {
            // Define the 4096-byte array for the contents of the page
            let seg_decl = compiler.module.declare_data(
                &format!("guest_sparse_page_data_{}", dix),
                Linkage::Local,
                false,
            )?;
            let mut seg_ctx = DataContext::new();
            seg_ctx.define(vs.to_vec().into_boxed_slice());
            compiler.module.define_data(seg_decl, &seg_ctx)?;

            // Put a relocation to that array into the table:
            let seg_gv = compiler
                .module
                .declare_data_in_data(seg_decl, &mut table_ctx);
            table_ctx.write_data_addr(table_data.position() as u32, seg_gv, 0);
        }
        table_data.write_u64::<LittleEndian>(0)?;
    }

    table_ctx.define(table_data.into_inner().into_boxed_slice());
    let table_decl =
        compiler
            .module
            .declare_data("guest_sparse_page_data", Linkage::Export, false)?;
    compiler.module.define_data(table_decl, &table_ctx)?;

    Ok(())
}
<|endoftext|>


<commit_before>use log::{debug, error, info};
use std::cell;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::env;
use std::fmt::Write;
use std::os::raw::{c_char, c_int, c_void};
use std::slice;
use std::str;

use cpython::*;
use lldb::*;

use crate::debug_protocol::{DisplayHtmlEventBody, EventBody};
use crate::error::Error;

#[derive(Debug)]
pub enum PythonValue {
    SBValue(SBValue),
    Int(i64),
    Bool(bool),
    String(String),
    Object(String),
}

pub struct PythonInterface {
    interpreter: SBCommandInterpreter,
    pymod_lldb: PyModule,
    pymod_codelldb: PyModule,
    pymod_traceback: PyModule,
    pyty_sbexec_context: PyType,
    pyty_sbmodule: PyType,
    pyty_sbvalue: PyType,
    pyfn_evaluate_in_frame: PyObject,
    pyfn_modules_loaded: PyObject,
}

impl PythonInterface {
    pub fn new(
        interpreter: SBCommandInterpreter, send_event: Box<Fn(EventBody) + Send>,
    ) -> Result<Box<PythonInterface>, Error> {
        let mut init_script = env::current_exe()?;
        init_script.set_file_name("codelldb.py");

        // Import codelldb.py
        let mut command_result = SBCommandReturnObject::new();
        let command = format!("command script import '{}'", init_script.to_str()?);
        interpreter.handle_command(&command, &mut command_result, false);
        info!("{:?}", command_result);

        // Import rust.py
        let mut rust_formatters = init_script.clone();
        rust_formatters.set_file_name("rust.py");
        let command = format!("command script import '{}'", rust_formatters.to_str()?);
        interpreter.handle_command(&command, &mut command_result, false);
        info!("{:?}", command_result);

        // Cache some objects
        let gil = Python::acquire_gil();
        let py = gil.python();

        let pymod_lldb = py.import("lldb").unwrap();
        let pyty_sbvalue = PyType::downcast_from(py, pymod_lldb.get(py, "SBValue").unwrap()).unwrap();
        let pyty_sbmodule = PyType::downcast_from(py, pymod_lldb.get(py, "SBModule").unwrap()).unwrap();
        let pyty_sbexec_context = PyType::downcast_from(py, pymod_lldb.get(py, "SBExecutionContext").unwrap()).unwrap();

        let pymod_codelldb = py.import("codelldb").unwrap();
        let pyfn_evaluate_in_frame = pymod_codelldb.get(py, "evaluate_in_frame").unwrap();
        let pyfn_modules_loaded = pymod_codelldb.get(py, "modules_loaded").unwrap();

        let pymod_traceback = py.import("traceback").unwrap();

        let callback = RustClosure::new(py, move |py: Python, args, kwargs| {
            py_argparse!(py, None, args, kwargs, (
                    html: String,
                    title: Option<String> = None,
                    position: Option<i32> = None,
                    reveal: bool = false
            ) {
                send_event(EventBody::displayHtml(DisplayHtmlEventBody {
                    html, title, position, reveal
                }));
                Ok(py.None())
            })
        });
        pymod_codelldb.as_object().setattr(py, "display_html", callback.unwrap());

        let mut py_interface = Box::new(PythonInterface {
            interpreter,
            pymod_lldb,
            pymod_codelldb,
            pymod_traceback,
            pyty_sbexec_context,
            pyty_sbmodule,
            pyty_sbvalue,
            pyfn_evaluate_in_frame,
            pyfn_modules_loaded,
        });
        Ok(py_interface)
    }

    pub fn evaluate(&self, expr: &str, is_simple_expr: bool, context: &SBExecutionContext) -> Result<SBValue, String> {
        let gil = Python::acquire_gil();
        let py = gil.python();
        let pysb_exec_context =
            unsafe { into_swig_wrapper::<SBExecutionContext>(py, context.clone(), &self.pyty_sbexec_context) };
        let result = self.pyfn_evaluate_in_frame.call(py, (expr, is_simple_expr, pysb_exec_context), None);
        let target = context.target().unwrap();
        let result = self.to_sbvalue(py, &target, result);
        debug!("evaluate {} -> {:?}", expr, result);
        result
    }

    pub fn evaluate_as_bool(
        &self, expr: &str, is_simple_expr: bool, context: &SBExecutionContext,
    ) -> Result<bool, String> {
        let gil = Python::acquire_gil();
        let py = gil.python();
        let pysb_exec_context =
            unsafe { into_swig_wrapper::<SBExecutionContext>(py, context.clone(), &self.pyty_sbexec_context) };
        let result = self.pyfn_evaluate_in_frame.call(py, (expr, is_simple_expr, pysb_exec_context), None);
        let result = match result {
            Ok(value) => Ok(value.is_true(py).unwrap()),
            Err(pyerr) => Err(self.format_exception(py, &pyerr)),
        };
        debug!("evaluate_as_bool {} -> {:?}", expr, result);
        result
    }

    pub fn modules_loaded(&self, modules: &mut Iterator<Item = &SBModule>) {
        let gil = Python::acquire_gil();
        let py = gil.python();

        let list = PyList::new(py, &[]);
        for module in modules {
            let pysbmodule = unsafe { into_swig_wrapper::<SBModule>(py, module.clone(), &self.pyty_sbmodule) };
            list.insert_item(py, list.len(py), pysbmodule);
        }
        self.pyfn_modules_loaded.call(py, (list,), None).unwrap();
    }

    fn to_sbvalue(&self, py: Python, target: &SBTarget, result: PyResult<PyObject>) -> Result<SBValue, String> {
        match result {
            Ok(value) => {
                if self.pyty_sbvalue.is_instance(py, &value) {
                    let sbvalue = unsafe { from_swig_wrapper::<SBValue>(py, &value) };
                    Ok(sbvalue)
                } else if PyBool::type_object(py).is_instance(py, &value) {
                    let value = bool::extract(py, &value).unwrap();
                    Ok(sbvalue_from_bool(value, target))
                } else if PyInt::type_object(py).is_instance(py, &value) {
                    let value = i64::extract(py, &value).unwrap();
                    Ok(sbvalue_from_i64(value, target))
                } else if PyLong::type_object(py).is_instance(py, &value) {
                    let value = i64::extract(py, &value).unwrap();
                    Ok(sbvalue_from_i64(value, target))
                } else if PyString::type_object(py).is_instance(py, &value) {
                    let value = String::extract(py, &value).unwrap();
                    Ok(sbvalue_from_str(&value, target))
                } else {
                    let value = value.to_string();
                    Ok(sbvalue_from_str(&value, target))
                }
            }
            Err(pyerr) => Err(self.format_exception(py, &pyerr)),
        }
    }

    fn format_exception(&self, py: Python, err: &PyErr) -> String {
        let tb = self
            .pymod_traceback
            .call(py, "format_exception", (&err.ptype, &err.pvalue, &err.ptraceback), None)
            .unwrap();
        let lines = Vec::<String>::extract(py, &tb).unwrap();
        lines.concat()
    }
}

// Creates SWIG wrapper containing native SB object.
// `pytype` is Python type object of the wrapper.
// Obviously, `SBT` and `pytype` must match, hence `unsafe`.
unsafe fn into_swig_wrapper<SBT>(py: Python, obj: SBT, pytype: &PyType) -> PyObject {
    // SWIG does not provide an API for creating Python wrapper from a native object, so we have to employ a bit of trickery:
    // First, we call SB wrapper's constructor on the Python side, which creates an instance wrapping dummy native SB object,
    let pysb = pytype.call(py, NoArgs, None).unwrap();
    // then, we retrieve a pointer to the native object via wrapper's `this` attribute,
    let this = pysb.getattr(py, "this").unwrap().extract::<usize>(py).unwrap();
    // finally, we replace it with the actual SB object we wanted to wrap.
    std::ptr::replace(this as *mut SBT, obj);
    pysb
}

// Extracts native SB object from SWIG wrapper.
unsafe fn from_swig_wrapper<SBT>(py: Python, pyobj: &PyObject) -> SBT
where
    SBT: Clone,
{
    let this = pyobj.getattr(py, "this").unwrap().extract::<usize>(py).unwrap();
    let mut sb = this as *const SBT;
    (*sb).clone()
}

fn sbvalue_from_bool(value: bool, target: &SBTarget) -> SBValue {
    let ty = target.get_basic_type(BasicType::Bool);
    let slice = [value as u8, 0, 0, 0, 0, 0, 0, 0]; // 8 bytes ought to be enough to hold a bool on any platform
    let data = SBData::borrow_bytes(&slice, ByteOrder::Little, 8);
    target.create_value_from_data("result", &data, &ty)
}

fn sbvalue_from_i64(value: i64, target: &SBTarget) -> SBValue {
    let ty = target.get_basic_type(BasicType::LongLong);
    let bytes = value.to_le_bytes();
    let data = SBData::borrow_bytes(&bytes, ByteOrder::Little, 8);
    target.create_value_from_data("result", &data, &ty)
}

fn sbvalue_from_str(value: &str, target: &SBTarget) -> SBValue {
    let ty = target.get_basic_type(BasicType::Char).array_type(value.len() as u64);
    let bytes = value.as_bytes();
    let data = SBData::borrow_bytes(bytes, ByteOrder::Little, 8);
    target.create_value_from_data("result", &data, &ty)
}

// Python wrapper for Rust closures
py_class!(class RustClosure |py| {
    data closure: cell::RefCell<Box<FnMut(Python, &PyTuple, Option<&PyDict>) -> PyResult<PyObject> + Send>>;

    def __call__(&self, *args, **kwargs) -> PyResult<PyObject> {
        use std::ops::DerefMut;
        let mut mut_ref = self.closure(py).borrow_mut();
        mut_ref.deref_mut()(py, &args, kwargs)
    }
});

impl RustClosure {
    fn new<F>(py: Python, closure: F) -> PyResult<RustClosure>
    where
        F: FnMut(Python, &PyTuple, Option<&PyDict>) -> PyResult<PyObject> + Send + 'static,
    {
        let closure = cell::RefCell::new(Box::new(closure));
        RustClosure::create_instance(py, closure)
    }
}
<commit_msg>Don't panic on Python exception formatting errors.<commit_after>use log::{debug, error, info};
use std::cell;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::env;
use std::fmt::Write;
use std::os::raw::{c_char, c_int, c_void};
use std::slice;
use std::str;

use cpython::*;
use lldb::*;

use crate::debug_protocol::{DisplayHtmlEventBody, EventBody};
use crate::error::Error;

#[derive(Debug)]
pub enum PythonValue {
    SBValue(SBValue),
    Int(i64),
    Bool(bool),
    String(String),
    Object(String),
}

pub struct PythonInterface {
    interpreter: SBCommandInterpreter,
    pymod_lldb: PyModule,
    pymod_codelldb: PyModule,
    pymod_traceback: PyModule,
    pyty_sbexec_context: PyType,
    pyty_sbmodule: PyType,
    pyty_sbvalue: PyType,
    pyfn_evaluate_in_frame: PyObject,
    pyfn_modules_loaded: PyObject,
}

impl PythonInterface {
    pub fn new(
        interpreter: SBCommandInterpreter, send_event: Box<Fn(EventBody) + Send>,
    ) -> Result<Box<PythonInterface>, Error> {
        let mut init_script = env::current_exe()?;
        init_script.set_file_name("codelldb.py");

        // Import codelldb.py
        let mut command_result = SBCommandReturnObject::new();
        let command = format!("command script import '{}'", init_script.to_str()?);
        interpreter.handle_command(&command, &mut command_result, false);
        info!("{:?}", command_result);

        // Import rust.py
        let mut rust_formatters = init_script.clone();
        rust_formatters.set_file_name("rust.py");
        let command = format!("command script import '{}'", rust_formatters.to_str()?);
        interpreter.handle_command(&command, &mut command_result, false);
        info!("{:?}", command_result);

        // Cache some objects
        let gil = Python::acquire_gil();
        let py = gil.python();

        let pymod_lldb = py.import("lldb").unwrap();
        let pyty_sbvalue = PyType::downcast_from(py, pymod_lldb.get(py, "SBValue").unwrap()).unwrap();
        let pyty_sbmodule = PyType::downcast_from(py, pymod_lldb.get(py, "SBModule").unwrap()).unwrap();
        let pyty_sbexec_context = PyType::downcast_from(py, pymod_lldb.get(py, "SBExecutionContext").unwrap()).unwrap();

        let pymod_codelldb = py.import("codelldb").unwrap();
        let pyfn_evaluate_in_frame = pymod_codelldb.get(py, "evaluate_in_frame").unwrap();
        let pyfn_modules_loaded = pymod_codelldb.get(py, "modules_loaded").unwrap();

        let pymod_traceback = py.import("traceback").unwrap();

        let callback = RustClosure::new(py, move |py: Python, args, kwargs| {
            py_argparse!(py, None, args, kwargs, (
                    html: String,
                    title: Option<String> = None,
                    position: Option<i32> = None,
                    reveal: bool = false
            ) {
                send_event(EventBody::displayHtml(DisplayHtmlEventBody {
                    html, title, position, reveal
                }));
                Ok(py.None())
            })
        });
        pymod_codelldb.as_object().setattr(py, "display_html", callback.unwrap());

        let mut py_interface = Box::new(PythonInterface {
            interpreter,
            pymod_lldb,
            pymod_codelldb,
            pymod_traceback,
            pyty_sbexec_context,
            pyty_sbmodule,
            pyty_sbvalue,
            pyfn_evaluate_in_frame,
            pyfn_modules_loaded,
        });
        Ok(py_interface)
    }

    pub fn evaluate(&self, expr: &str, is_simple_expr: bool, context: &SBExecutionContext) -> Result<SBValue, String> {
        let gil = Python::acquire_gil();
        let py = gil.python();
        let pysb_exec_context =
            unsafe { into_swig_wrapper::<SBExecutionContext>(py, context.clone(), &self.pyty_sbexec_context) };
        let result = self.pyfn_evaluate_in_frame.call(py, (expr, is_simple_expr, pysb_exec_context), None);
        let target = context.target().unwrap();
        let result = self.to_sbvalue(py, &target, result);
        debug!("evaluate {} -> {:?}", expr, result);
        result
    }

    pub fn evaluate_as_bool(
        &self, expr: &str, is_simple_expr: bool, context: &SBExecutionContext,
    ) -> Result<bool, String> {
        let gil = Python::acquire_gil();
        let py = gil.python();
        let pysb_exec_context =
            unsafe { into_swig_wrapper::<SBExecutionContext>(py, context.clone(), &self.pyty_sbexec_context) };
        let result = self.pyfn_evaluate_in_frame.call(py, (expr, is_simple_expr, pysb_exec_context), None);
        let result = match result {
            Ok(value) => Ok(value.is_true(py).unwrap()),
            Err(pyerr) => Err(self.format_exception(py, pyerr)),
        };
        debug!("evaluate_as_bool {} -> {:?}", expr, result);
        result
    }

    pub fn modules_loaded(&self, modules: &mut Iterator<Item = &SBModule>) {
        let gil = Python::acquire_gil();
        let py = gil.python();

        let list = PyList::new(py, &[]);
        for module in modules {
            let pysbmodule = unsafe { into_swig_wrapper::<SBModule>(py, module.clone(), &self.pyty_sbmodule) };
            list.insert_item(py, list.len(py), pysbmodule);
        }
        self.pyfn_modules_loaded.call(py, (list,), None).unwrap();
    }

    fn to_sbvalue(&self, py: Python, target: &SBTarget, result: PyResult<PyObject>) -> Result<SBValue, String> {
        match result {
            Ok(value) => {
                if self.pyty_sbvalue.is_instance(py, &value) {
                    let sbvalue = unsafe { from_swig_wrapper::<SBValue>(py, &value) };
                    Ok(sbvalue)
                } else if PyBool::type_object(py).is_instance(py, &value) {
                    let value = bool::extract(py, &value).unwrap();
                    Ok(sbvalue_from_bool(value, target))
                } else if PyInt::type_object(py).is_instance(py, &value) {
                    let value = i64::extract(py, &value).unwrap();
                    Ok(sbvalue_from_i64(value, target))
                } else if PyLong::type_object(py).is_instance(py, &value) {
                    let value = i64::extract(py, &value).unwrap();
                    Ok(sbvalue_from_i64(value, target))
                } else if PyString::type_object(py).is_instance(py, &value) {
                    let value = String::extract(py, &value).unwrap();
                    Ok(sbvalue_from_str(&value, target))
                } else {
                    let value = value.to_string();
                    Ok(sbvalue_from_str(&value, target))
                }
            }
            Err(pyerr) => Err(self.format_exception(py, pyerr)),
        }
    }

    fn format_exception(&self, py: Python, mut err: PyErr) -> String {
        err.normalize(py);
        match self.pymod_traceback.call(py, "format_exception", (&err.ptype, &err.pvalue, &err.ptraceback), None) {
            Ok(tb) => {
                let lines = Vec::<String>::extract(py, &tb).unwrap();
                lines.concat()
            }
            Err(_) => format!("Could not format exception: {:?}", err),
        }
    }
}

// Creates a SWIG wrapper containing native SB object.
// `pytype` is the Python type object of the wrapper.
// Obviously, `SBT` and `pytype` must match, hence `unsafe`.
unsafe fn into_swig_wrapper<SBT>(py: Python, obj: SBT, pytype: &PyType) -> PyObject {
    // SWIG does not provide an API for creating Python wrapper from a native object, so we have to employ a bit of trickery:
    // First, we call SB wrapper's constructor on the Python side, which creates an instance wrapping dummy native SB object,
    let pysb = pytype.call(py, NoArgs, None).unwrap();
    // then, we retrieve a pointer to the native object via wrapper's `this` attribute,
    let this = pysb.getattr(py, "this").unwrap().extract::<usize>(py).unwrap();
    // finally, we replace it with the actual SB object we wanted to wrap.
    std::ptr::replace(this as *mut SBT, obj);
    pysb
}

// Extracts native SB object from a SWIG wrapper.
unsafe fn from_swig_wrapper<SBT>(py: Python, pyobj: &PyObject) -> SBT
where
    SBT: Clone,
{
    let this = pyobj.getattr(py, "this").unwrap().extract::<usize>(py).unwrap();
    let mut sb = this as *const SBT;
    (*sb).clone()
}

fn sbvalue_from_bool(value: bool, target: &SBTarget) -> SBValue {
    let ty = target.get_basic_type(BasicType::Bool);
    let slice = [value as u8, 0, 0, 0, 0, 0, 0, 0]; // 8 bytes ought to be enough to hold a bool on any platform
    let data = SBData::borrow_bytes(&slice, ByteOrder::Little, 8);
    target.create_value_from_data("result", &data, &ty)
}

fn sbvalue_from_i64(value: i64, target: &SBTarget) -> SBValue {
    let ty = target.get_basic_type(BasicType::LongLong);
    let bytes = value.to_le_bytes();
    let data = SBData::borrow_bytes(&bytes, ByteOrder::Little, 8);
    target.create_value_from_data("result", &data, &ty)
}

fn sbvalue_from_str(value: &str, target: &SBTarget) -> SBValue {
    let ty = target.get_basic_type(BasicType::Char).array_type(value.len() as u64);
    let bytes = value.as_bytes();
    let data = SBData::borrow_bytes(bytes, ByteOrder::Little, 8);
    target.create_value_from_data("result", &data, &ty)
}

// Python wrapper for Rust closures
py_class!(class RustClosure |py| {
    data closure: cell::RefCell<Box<FnMut(Python, &PyTuple, Option<&PyDict>) -> PyResult<PyObject> + Send>>;

    def __call__(&self, *args, **kwargs) -> PyResult<PyObject> {
        use std::ops::DerefMut;
        let mut mut_ref = self.closure(py).borrow_mut();
        mut_ref.deref_mut()(py, &args, kwargs)
    }
});

impl RustClosure {
    fn new<F>(py: Python, closure: F) -> PyResult<RustClosure>
    where
        F: FnMut(Python, &PyTuple, Option<&PyDict>) -> PyResult<PyObject> + Send + 'static,
    {
        let closure = cell::RefCell::new(Box::new(closure));
        RustClosure::create_instance(py, closure)
    }
}
<|endoftext|>


<commit_before>use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{error, info, warn};
use winapi::shared::minwindef::FALSE;
use winapi::shared::ntdef::NULL;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::processthreadsapi::{
    OpenProcess, OpenThread, ResumeThread, SuspendThread, TerminateProcess,
};
use winapi::um::tlhelp32::{
    CreateToolhelp32Snapshot, Process32First, Process32Next, Thread32First, Thread32Next,
    PROCESSENTRY32, TH32CS_SNAPPROCESS, TH32CS_SNAPTHREAD, THREADENTRY32,
};
use winapi::um::winnt::{PROCESS_TERMINATE, THREAD_SUSPEND_RESUME};

use crate::task_handler::ProcessAction;
use pueue_lib::network::message::Signal as InternalSignal;

pub fn compile_shell_command(command_string: &str) -> Command {
    // Chain two `powershell` commands, one that sets the output encoding to utf8 and then the user provided one.
    let mut command = Command::new("powershell");
    command.arg("-c").arg(format!(
        "[Console]::OutputEncoding = [Text.UTF8Encoding]::UTF8; {}",
        command_string
    ));

    command
}

pub fn send_internal_signal_to_child(
    child: &Child,
    signal: InternalSignal,
    send_to_children: bool,
) -> Result<bool> {
    bail!("Trying to send unix signal on a windows machine. This isn't supported.");
}

/// Send a signal to a windows process.
pub fn run_action_on_child(child: &Child, action: &ProcessAction, _children: bool) -> Result<bool> {
    let pids = get_cur_task_processes(child.id());
    if pids.is_empty() {
        bail!("Process has just gone away");
    }

    match action {
        ProcessAction::Pause => {
            for pid in pids {
                for thread in get_threads(pid) {
                    suspend_thread(thread);
                }
            }
        }
        ProcessAction::Resume => {
            for pid in pids {
                for thread in get_threads(pid) {
                    resume_thread(thread);
                }
            }
        }
    }

    Ok(true)
}

/// Kill a child process
pub fn kill_child(task_id: usize, child: &mut Child, _kill_children: bool) -> bool {
    match child.kill() {
        Err(_) => {
            info!("Task {} has already finished by itself", task_id);
            false
        }
        Ok(_) => {
            let pids = get_cur_task_processes(child.id());

            for pid in pids {
                terminate_process(pid);
            }
            true
        }
    }
}

/// Get current task pid, all child pid and all children's children
fn get_cur_task_processes(task_pid: u32) -> Vec<u32> {
    let mut all_pids = Vec::new();

    // Get all pids by BFS
    let mut parent_pids = vec![task_pid];
    while let Some(pid) = parent_pids.pop() {
        all_pids.push(pid);

        get_child_pids(pid, &mut parent_pids);
    }

    // Keep parent pid ahead of child. We need execute action for parent process first.
    all_pids.reverse();
    all_pids
}

/// Get child pids of a specific process.
fn get_child_pids(target_pid: u32, pid_list: &mut Vec<u32>) {
    unsafe {
        // Take a snapshot of all processes in the system.
        // While enumerating the set of processes, new processes can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, target_pid);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return;
        }

        // Walk the list of processes.
        let mut process_entry = PROCESSENTRY32 {
            dwSize: std::mem::size_of::<PROCESSENTRY32>() as u32,
            ..Default::default()
        };
        if Process32First(snapshot_handle, &mut process_entry) == FALSE {
            error!("Couldn't get first process.");
            CloseHandle(snapshot_handle);
            return;
        }

        loop {
            if process_entry.th32ParentProcessID == target_pid {
                pid_list.push(process_entry.th32ProcessID);
            }

            if Process32Next(snapshot_handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }
}

/// Get all thread id of a specific process
fn get_threads(target_pid: u32) -> Vec<u32> {
    let mut threads = Vec::new();

    unsafe {
        // Take a snapshot of all threads in the system.
        // While enumerating the set of threads, new threads can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return threads;
        }

        // Walk the list of threads.
        let mut thread_entry = THREADENTRY32 {
            dwSize: std::mem::size_of::<THREADENTRY32>() as u32,
            ..Default::default()
        };
        if Thread32First(snapshot_handle, &mut thread_entry) == FALSE {
            error!("Couldn't get first thread.");
            CloseHandle(snapshot_handle);
            return threads;
        }

        loop {
            if thread_entry.th32OwnerProcessID == target_pid {
                threads.push(thread_entry.th32ThreadID);
            }

            if Thread32Next(snapshot_handle, &mut thread_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }

    threads
}

/// Suspend a thread
/// Each thread has a suspend count (with a maximum value of `MAXIMUM_SUSPEND_COUNT`).
/// If the suspend count is greater than zero, the thread is suspended; otherwise, the thread is not suspended and is eligible for execution.
/// Calling `SuspendThread` causes the target thread's suspend count to be incremented.
/// Attempting to increment past the maximum suspend count causes an error without incrementing the count.
/// [SuspendThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-suspendthread)
fn suspend_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If SuspendThread fails, the return value is (DWORD) -1
            if u32::max_value() == SuspendThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to suspend thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Resume a thread
/// ResumeThread checks the suspend count of the subject thread.
/// If the suspend count is zero, the thread is not currently suspended. Otherwise, the subject thread's suspend count is decremented.
/// If the resulting value is zero, then the execution of the subject thread is resumed.
/// [ResumeThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-resumethread)
fn resume_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If ResumeThread fails, the return value is (DWORD) -1
            if u32::max_value() == ResumeThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to resume thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Terminate a process
/// [TerminateProcess](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-terminateprocess)
fn terminate_process(pid: u32) {
    unsafe {
        println!("{}", pid);
        // Get a handle for the target process
        let process_handle = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
        // If TerminateProcess fails, the return value is zero.
        if 0 == TerminateProcess(process_handle, 1) {
            let err_code = GetLastError();
            warn!(
                "Failed to terminate process {} with error code {}",
                pid, err_code
            );
        }

        CloseHandle(process_handle);
    }
}

#[cfg(test)]
mod test {
    use std::thread::sleep;
    use std::time::Duration;

    use pretty_assertions::assert_eq;

    use super::*;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: u32) -> bool {
        unsafe {
            let handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

            let mut process_entry = PROCESSENTRY32::default();
            process_entry.dwSize = std::mem::size_of::<PROCESSENTRY32>() as u32;

            loop {
                if process_entry.th32ProcessID == pid {
                    CloseHandle(handle);
                    return false;
                }

                if Process32Next(handle, &mut process_entry) == FALSE {
                    break;
                }
            }

            CloseHandle(handle);
        }

        true
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("sleep 0.1")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[test]
    /// Ensure a `powershell -c` command will be properly killed without detached processes.
    fn test_shell_command_is_killed() {
        let mut child = compile_shell_command("sleep 60; sleep 60; echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }

    #[test]
    /// Ensure that a `powershell -c` process with a child process that has children of it's own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() {
        let mut child = compile_shell_command("powershell -c 'sleep 60; sleep 60'; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Assert that the direct child (powershell -c) has been killed.
        sleep(Duration::from_millis(500));
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }

    #[test]
    /// Ensure a normal command without `powershell -c` will be killed.
    fn test_normal_command_is_killed() {
        let mut child = Command::new("ping")
            .arg("localhost")
            .arg("-t")
            .spawn()
            .expect("Failed to spawn ping");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all processes
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));
    }

    #[test]
    /// Ensure a normal command and all it's children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() {
        let mut child = Command::new("powershell")
            .arg("-c")
            .arg("sleep 60; sleep 60; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }
}
<commit_msg>Longer timeout for flaky windows process tests<commit_after>use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{error, info, warn};
use winapi::shared::minwindef::FALSE;
use winapi::shared::ntdef::NULL;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::processthreadsapi::{
    OpenProcess, OpenThread, ResumeThread, SuspendThread, TerminateProcess,
};
use winapi::um::tlhelp32::{
    CreateToolhelp32Snapshot, Process32First, Process32Next, Thread32First, Thread32Next,
    PROCESSENTRY32, TH32CS_SNAPPROCESS, TH32CS_SNAPTHREAD, THREADENTRY32,
};
use winapi::um::winnt::{PROCESS_TERMINATE, THREAD_SUSPEND_RESUME};

use crate::task_handler::ProcessAction;
use pueue_lib::network::message::Signal as InternalSignal;

pub fn compile_shell_command(command_string: &str) -> Command {
    // Chain two `powershell` commands, one that sets the output encoding to utf8 and then the user provided one.
    let mut command = Command::new("powershell");
    command.arg("-c").arg(format!(
        "[Console]::OutputEncoding = [Text.UTF8Encoding]::UTF8; {}",
        command_string
    ));

    command
}

pub fn send_internal_signal_to_child(
    child: &Child,
    signal: InternalSignal,
    send_to_children: bool,
) -> Result<bool> {
    bail!("Trying to send unix signal on a windows machine. This isn't supported.");
}

/// Send a signal to a windows process.
pub fn run_action_on_child(child: &Child, action: &ProcessAction, _children: bool) -> Result<bool> {
    let pids = get_cur_task_processes(child.id());
    if pids.is_empty() {
        bail!("Process has just gone away");
    }

    match action {
        ProcessAction::Pause => {
            for pid in pids {
                for thread in get_threads(pid) {
                    suspend_thread(thread);
                }
            }
        }
        ProcessAction::Resume => {
            for pid in pids {
                for thread in get_threads(pid) {
                    resume_thread(thread);
                }
            }
        }
    }

    Ok(true)
}

/// Kill a child process
pub fn kill_child(task_id: usize, child: &mut Child, _kill_children: bool) -> bool {
    match child.kill() {
        Err(_) => {
            info!("Task {} has already finished by itself", task_id);
            false
        }
        Ok(_) => {
            let pids = get_cur_task_processes(child.id());

            for pid in pids {
                terminate_process(pid);
            }
            true
        }
    }
}

/// Get current task pid, all child pid and all children's children
fn get_cur_task_processes(task_pid: u32) -> Vec<u32> {
    let mut all_pids = Vec::new();

    // Get all pids by BFS
    let mut parent_pids = vec![task_pid];
    while let Some(pid) = parent_pids.pop() {
        all_pids.push(pid);

        get_child_pids(pid, &mut parent_pids);
    }

    // Keep parent pid ahead of child. We need execute action for parent process first.
    all_pids.reverse();
    all_pids
}

/// Get child pids of a specific process.
fn get_child_pids(target_pid: u32, pid_list: &mut Vec<u32>) {
    unsafe {
        // Take a snapshot of all processes in the system.
        // While enumerating the set of processes, new processes can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, target_pid);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return;
        }

        // Walk the list of processes.
        let mut process_entry = PROCESSENTRY32 {
            dwSize: std::mem::size_of::<PROCESSENTRY32>() as u32,
            ..Default::default()
        };
        if Process32First(snapshot_handle, &mut process_entry) == FALSE {
            error!("Couldn't get first process.");
            CloseHandle(snapshot_handle);
            return;
        }

        loop {
            if process_entry.th32ParentProcessID == target_pid {
                pid_list.push(process_entry.th32ProcessID);
            }

            if Process32Next(snapshot_handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }
}

/// Get all thread id of a specific process
fn get_threads(target_pid: u32) -> Vec<u32> {
    let mut threads = Vec::new();

    unsafe {
        // Take a snapshot of all threads in the system.
        // While enumerating the set of threads, new threads can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return threads;
        }

        // Walk the list of threads.
        let mut thread_entry = THREADENTRY32 {
            dwSize: std::mem::size_of::<THREADENTRY32>() as u32,
            ..Default::default()
        };
        if Thread32First(snapshot_handle, &mut thread_entry) == FALSE {
            error!("Couldn't get first thread.");
            CloseHandle(snapshot_handle);
            return threads;
        }

        loop {
            if thread_entry.th32OwnerProcessID == target_pid {
                threads.push(thread_entry.th32ThreadID);
            }

            if Thread32Next(snapshot_handle, &mut thread_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }

    threads
}

/// Suspend a thread
/// Each thread has a suspend count (with a maximum value of `MAXIMUM_SUSPEND_COUNT`).
/// If the suspend count is greater than zero, the thread is suspended; otherwise, the thread is not suspended and is eligible for execution.
/// Calling `SuspendThread` causes the target thread's suspend count to be incremented.
/// Attempting to increment past the maximum suspend count causes an error without incrementing the count.
/// [SuspendThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-suspendthread)
fn suspend_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If SuspendThread fails, the return value is (DWORD) -1
            if u32::max_value() == SuspendThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to suspend thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Resume a thread
/// ResumeThread checks the suspend count of the subject thread.
/// If the suspend count is zero, the thread is not currently suspended. Otherwise, the subject thread's suspend count is decremented.
/// If the resulting value is zero, then the execution of the subject thread is resumed.
/// [ResumeThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-resumethread)
fn resume_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If ResumeThread fails, the return value is (DWORD) -1
            if u32::max_value() == ResumeThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to resume thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Terminate a process
/// [TerminateProcess](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-terminateprocess)
fn terminate_process(pid: u32) {
    unsafe {
        println!("{}", pid);
        // Get a handle for the target process
        let process_handle = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
        // If TerminateProcess fails, the return value is zero.
        if 0 == TerminateProcess(process_handle, 1) {
            let err_code = GetLastError();
            warn!(
                "Failed to terminate process {} with error code {}",
                pid, err_code
            );
        }

        CloseHandle(process_handle);
    }
}

#[cfg(test)]
mod test {
    use std::thread::sleep;
    use std::time::Duration;

    use pretty_assertions::assert_eq;

    use super::*;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: u32) -> bool {
        unsafe {
            let handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

            let mut process_entry = PROCESSENTRY32::default();
            process_entry.dwSize = std::mem::size_of::<PROCESSENTRY32>() as u32;

            loop {
                if process_entry.th32ProcessID == pid {
                    CloseHandle(handle);
                    return false;
                }

                if Process32Next(handle, &mut process_entry) == FALSE {
                    break;
                }
            }

            CloseHandle(handle);
        }

        true
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("sleep 0.1")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[test]
    /// Ensure a `powershell -c` command will be properly killed without detached processes.
    fn test_shell_command_is_killed() {
        let mut child = compile_shell_command("sleep 60; sleep 60; echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(1000));

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }

    #[test]
    /// Ensure that a `powershell -c` process with a child process that has children of it's own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() {
        let mut child = compile_shell_command("powershell -c 'sleep 60; sleep 60'; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(1000));

        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Assert that the direct child (powershell -c) has been killed.
        sleep(Duration::from_millis(500));
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }

    #[test]
    /// Ensure a normal command without `powershell -c` will be killed.
    fn test_normal_command_is_killed() {
        let mut child = Command::new("ping")
            .arg("localhost")
            .arg("-t")
            .spawn()
            .expect("Failed to spawn ping");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all processes
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));
    }

    #[test]
    /// Ensure a normal command and all it's children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() {
        let mut child = Command::new("powershell")
            .arg("-c")
            .arg("sleep 60; sleep 60; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }
}
<|endoftext|>


<commit_before>//! This module provides a microfacet model for transmission (a BTDF), see
//! [Walter et al. 07](https://www.cs.cornell.edu/~srm/publications/EGSR07-btdf.pdf)
//! for details.

use std::f32;
use enum_set::EnumSet;

use linalg::{self, Vector};
use film::Colorf;
use bxdf::{self, BxDF, BxDFType};
use bxdf::fresnel::{Dielectric, Fresnel};
use bxdf::microfacet::{MicrofacetDistribution};

/// Struct providing the microfacet BTDF, implemented as described in
/// [Walter et al. 07](https://www.cs.cornell.edu/~srm/publications/EGSR07-btdf.pdf)
#[derive(Copy, Clone)]
pub struct MicrofacetTransmission<'a> {
    reflectance: Colorf,
    fresnel: &'a Dielectric,
    /// Microfacet distribution describing the structure of the microfacets of
    /// the material
    microfacet: &'a MicrofacetDistribution,
}

impl<'a> MicrofacetTransmission<'a> {
    /// Create a new transmissive microfacet BRDF
    pub fn new(c: &Colorf, fresnel: &'a Dielectric, microfacet: &'a MicrofacetDistribution)
            -> MicrofacetTransmission<'a> {
        MicrofacetTransmission { reflectance: *c, fresnel: fresnel, microfacet: microfacet }
    }
    /// Convenience method for getting `eta_i` and `eta_t` in the right order for if
    /// we're entering or exiting this material based on the direction of the outgoing
    /// ray.
    fn eta_for_interaction(&self, w_o: &Vector) -> (f32, f32) {
        if bxdf::cos_theta(w_o) > 0.0 {
            (self.fresnel.eta_i, self.fresnel.eta_t)
        } else {
            (self.fresnel.eta_t, self.fresnel.eta_i)
        }
    }
    /// Compute the Jacobian for the change of variables (see [Walter et al 07] section 4.2),
    /// here we compute equation 17 in that section.
    fn jacobian(w_o: &Vector, w_i: &Vector, w_h: &Vector, eta: (f32, f32)) -> f32 {
        let wi_dot_h = linalg::dot(w_i, w_h);
        let wo_dot_h = linalg::dot(w_o, w_h);
        let denom = f32::powf(eta.0 * wi_dot_h + eta.1 * wo_dot_h, 2.0);
        if denom != 0.0 {
            f32::abs(f32::powf(eta.1, 2.0) * f32::abs(wo_dot_h) / denom)
        } else {
            0.0
        }
    }
    fn half_vector(w_o: &Vector, w_i: &Vector, eta: (f32, f32)) -> Option<Vector> {
        let w_h = -eta.1 * *w_i - eta.0 * *w_o;
        if w_h.x == 0.0 && w_h.y == 0.0 && w_h.z == 0.0 {
            None
        } else {
            Some(w_h.normalized())
        }
    }
}

impl<'a> BxDF for MicrofacetTransmission<'a> {
    fn bxdf_type(&self) -> EnumSet<BxDFType> {
        let mut e = EnumSet::new();
        e.insert(BxDFType::Glossy);
        e.insert(BxDFType::Transmission);
        e
    }
    fn eval(&self, w_o: &Vector, w_i: &Vector) -> Colorf {
        if bxdf::same_hemisphere(w_o, w_i) {
            return Colorf::black();
        }
        let cos_to = bxdf::cos_theta(w_o);
        let cos_ti = bxdf::cos_theta(w_i);
        if cos_to == 0.0 || cos_ti == 0.0 {
            return Colorf::black();
        }
        let eta = self.eta_for_interaction(w_o);
        if let Some(w_h) = MicrofacetTransmission::half_vector(w_o, w_i, eta) {
            let d = self.microfacet.normal_distribution(&w_h);
            let f = Colorf::broadcast(1.0) - self.fresnel.fresnel(linalg::dot(w_o, &w_h));
            let g = self.microfacet.shadowing_masking(w_i, w_o, &w_h);
            let wi_dot_h = linalg::dot(w_i, &w_h);
            let jacobian = MicrofacetTransmission::jacobian(w_o, w_i, &w_h, eta);
            self.reflectance * (f32::abs(wi_dot_h) / (f32::abs(w_i.z) * f32::abs(w_o.z)))
                * (f * g * d) * jacobian
        } else {
            Colorf::black()
        }
    }
    fn sample(&self, w_o: &Vector, samples: &(f32, f32)) -> (Colorf, Vector, f32) {
        let mut w_h = self.microfacet.sample(w_o, samples);
        if !bxdf::same_hemisphere(w_o, &w_h) {
            w_h = -w_h;
        }
        let eta = self.eta_for_interaction(w_o);
        if let Some(w_i) = linalg::refract(w_o, &w_h, eta.0 / eta.1) {
            if bxdf::same_hemisphere(w_o, &w_i) {
                (Colorf::black(), Vector::broadcast(0.0), 0.0)
            } else {
                let pdf = self.microfacet.pdf(&w_h)
                    * MicrofacetTransmission::jacobian(w_o, &w_i, &w_h, eta);
                (self.eval(w_o, &w_i), w_i, pdf)
            }
        } else {
            (Colorf::black(), Vector::broadcast(0.0), 0.0)
        }
    }
    fn pdf(&self, w_o: &Vector, w_i: &Vector) -> f32 {
        if bxdf::same_hemisphere(w_o, w_i) {
            0.0
        } else {
            let eta = self.eta_for_interaction(w_o);
            if let Some(w_h) = MicrofacetTransmission::half_vector(w_o, w_i, eta) {
                self.microfacet.pdf(&w_h) * MicrofacetTransmission::jacobian(w_o, w_i, &w_h, eta)
            } else {
                0.0
            }
        }
    }
}

<commit_msg>Fix typo in jacobian computation<commit_after>//! This module provides a microfacet model for transmission (a BTDF), see
//! [Walter et al. 07](https://www.cs.cornell.edu/~srm/publications/EGSR07-btdf.pdf)
//! for details.

use std::f32;
use enum_set::EnumSet;

use linalg::{self, Vector};
use film::Colorf;
use bxdf::{self, BxDF, BxDFType};
use bxdf::fresnel::{Dielectric, Fresnel};
use bxdf::microfacet::{MicrofacetDistribution};

/// Struct providing the microfacet BTDF, implemented as described in
/// [Walter et al. 07](https://www.cs.cornell.edu/~srm/publications/EGSR07-btdf.pdf)
#[derive(Copy, Clone)]
pub struct MicrofacetTransmission<'a> {
    reflectance: Colorf,
    fresnel: &'a Dielectric,
    /// Microfacet distribution describing the structure of the microfacets of
    /// the material
    microfacet: &'a MicrofacetDistribution,
}

impl<'a> MicrofacetTransmission<'a> {
    /// Create a new transmissive microfacet BRDF
    pub fn new(c: &Colorf, fresnel: &'a Dielectric, microfacet: &'a MicrofacetDistribution)
            -> MicrofacetTransmission<'a> {
        MicrofacetTransmission { reflectance: *c, fresnel: fresnel, microfacet: microfacet }
    }
    /// Convenience method for getting `eta_i` and `eta_t` in the right order for if
    /// we're entering or exiting this material based on the direction of the outgoing
    /// ray.
    fn eta_for_interaction(&self, w_o: &Vector) -> (f32, f32) {
        if bxdf::cos_theta(w_o) > 0.0 {
            (self.fresnel.eta_i, self.fresnel.eta_t)
        } else {
            (self.fresnel.eta_t, self.fresnel.eta_i)
        }
    }
    /// Compute the Jacobian for the change of variables (see [Walter et al 07] section 4.2),
    /// here we compute equation 17 in that section.
    fn jacobian(w_o: &Vector, w_i: &Vector, w_h: &Vector, eta: (f32, f32)) -> f32 {
        let wi_dot_h = linalg::dot(w_i, w_h);
        let wo_dot_h = linalg::dot(w_o, w_h);
        let denom = f32::powf(eta.1 * wi_dot_h + eta.0 * wo_dot_h, 2.0);
        if denom != 0.0 {
            f32::abs(f32::powf(eta.0, 2.0) * f32::abs(wo_dot_h) / denom)
        } else {
            0.0
        }
    }
    fn half_vector(w_o: &Vector, w_i: &Vector, eta: (f32, f32)) -> Vector {
        (-eta.1 * *w_i - eta.0 * *w_o).normalized()
    }
}

impl<'a> BxDF for MicrofacetTransmission<'a> {
    fn bxdf_type(&self) -> EnumSet<BxDFType> {
        let mut e = EnumSet::new();
        e.insert(BxDFType::Glossy);
        e.insert(BxDFType::Transmission);
        e
    }
    fn eval(&self, w_o: &Vector, w_i: &Vector) -> Colorf {
        if bxdf::same_hemisphere(w_o, w_i) {
            return Colorf::black();
        }
        let cos_to = bxdf::cos_theta(w_o);
        let cos_ti = bxdf::cos_theta(w_i);
        if cos_to == 0.0 || cos_ti == 0.0 {
            return Colorf::black();
        }
        let eta = self.eta_for_interaction(w_o);
        let w_h = MicrofacetTransmission::half_vector(w_o, w_i, eta);
        let d = self.microfacet.normal_distribution(&w_h);
        let f = Colorf::broadcast(1.0) - self.fresnel.fresnel(linalg::dot(w_o, &w_h));
        let g = self.microfacet.shadowing_masking(w_i, w_o, &w_h);
        let wi_dot_h = linalg::dot(w_i, &w_h);
        let jacobian = MicrofacetTransmission::jacobian(w_o, w_i, &w_h, eta);
        self.reflectance * (f32::abs(wi_dot_h) / (f32::abs(w_i.z) * f32::abs(w_o.z)))
            * (f * g * d) * jacobian
    }
    fn sample(&self, w_o: &Vector, samples: &(f32, f32)) -> (Colorf, Vector, f32) {
        let mut w_h = self.microfacet.sample(w_o, samples);
        if !bxdf::same_hemisphere(w_o, &w_h) {
            w_h = -w_h;
        }
        let eta = self.eta_for_interaction(w_o);
        if let Some(w_i) = linalg::refract(w_o, &w_h, eta.0 / eta.1) {
            if bxdf::same_hemisphere(w_o, &w_i) {
                (Colorf::black(), Vector::broadcast(0.0), 0.0)
            } else {
                (self.eval(w_o, &w_i), w_i, self.pdf(w_o, &w_i))
            }
        } else {
            (Colorf::black(), Vector::broadcast(0.0), 0.0)
        }
    }
    fn pdf(&self, w_o: &Vector, w_i: &Vector) -> f32 {
        if bxdf::same_hemisphere(w_o, w_i) {
            0.0
        } else {
            let eta = self.eta_for_interaction(w_o);
            let w_h = MicrofacetTransmission::half_vector(w_o, w_i, eta);
            self.microfacet.pdf(&w_h) * MicrofacetTransmission::jacobian(w_o, w_i, &w_h, eta)
        }
    }
}

<|endoftext|>


<commit_before>use std::marker::PhantomData;

#[derive(Copy, Clone)]
pub enum ColorType {
    RGBA(u8),
}
pub struct ImageBuffer<T, U = ()>(U, PhantomData<T>);
pub struct Rgba<T>(PhantomData<T>);

pub mod png {
    use super::ColorType;
    use super::super::png::Encoder;
    use std::io::Write;

    pub struct PNGEncoder<W>(W);

    impl<W: Write> PNGEncoder<W> {
        pub fn new(w: W) -> Self {
            PNGEncoder(w)
        }

        pub fn encode(&mut self, d: &[u8], w: u32, h: u32, _: ColorType) -> Result<(), ()> {
            Encoder::new(&mut self.0, w, h)
                .write_header()
                .map_err(|_| ())?
                .write_image_data(d)
                .map_err(|_| ())
        }
    }
}

impl<T, U> ImageBuffer<T, U> {
    pub fn from_raw(_: u32, _: u32, d: U) -> Option<ImageBuffer<Rgba<u8>, U>> {
        Some(ImageBuffer(d, PhantomData))
    }

    pub fn as_ref(&self) -> &[u8]
    where
        U: AsRef<[u8]>,
    {
        self.0.as_ref()
    }
}
<commit_msg>Fix Bug in PNG Encoding<commit_after>use std::marker::PhantomData;

#[derive(Copy, Clone)]
pub enum ColorType {
    RGBA(u8),
}
pub struct ImageBuffer<T, U = ()>(U, PhantomData<T>);
pub struct Rgba<T>(PhantomData<T>);

pub mod png {
    use super::ColorType;
    use super::super::png::{self as pnglib, HasParameters};
    use std::io::Write;

    pub struct PNGEncoder<W>(W);

    impl<W: Write> PNGEncoder<W> {
        pub fn new(w: W) -> Self {
            PNGEncoder(w)
        }

        pub fn encode(&mut self, d: &[u8], w: u32, h: u32, _: ColorType) -> Result<(), ()> {
            let mut encoder = pnglib::Encoder::new(&mut self.0, w, h);
            encoder
                .set(pnglib::ColorType::RGBA)
                .set(pnglib::BitDepth::Eight);
            encoder
                .write_header()
                .map_err(|_| ())?
                .write_image_data(d)
                .map_err(|_| ())
        }
    }
}

impl<T, U> ImageBuffer<T, U> {
    pub fn from_raw(_: u32, _: u32, d: U) -> Option<ImageBuffer<Rgba<u8>, U>> {
        Some(ImageBuffer(d, PhantomData))
    }

    pub fn as_ref(&self) -> &[u8]
    where
        U: AsRef<[u8]>,
    {
        self.0.as_ref()
    }
}
<|endoftext|>


<commit_before>use std::cmp;
use std::io::Read;
use std::io::Write;
use std::io::Result;

// A stream that is seeded with incoming data which can be consumed and
// records data written to it (like a socket).
//
// This allows us to unit test a transport without using sockets. :)
pub struct TestStream {
    pub incoming: Vec<u8>,
    pub outgoing: Vec<u8>,
}

impl TestStream {
    pub fn new(incoming: Vec<u8>) -> TestStream {
        TestStream {
            incoming: incoming,
            outgoing: vec![],
        }
    }
}

impl Read for TestStream {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        let read_len = cmp::min(buf.len(), self.incoming.len());

        for i in 0..read_len {
            let byte = self.incoming.remove(0);
            buf[i] = byte;
        }

        Ok(read_len)
    }
}

impl Write for TestStream {
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        for byte in buf.iter() {
            self.outgoing.push(*byte);
        }

        Ok(buf.len())
    }

    fn flush(&mut self) -> Result<()> {
        Ok(())
    }
}

#[test]
fn test_stream_read_whole() {
    let mut ts = TestStream::new(vec![1, 2, 3]);

    let mut buf = [0; 4];
    let read_cnt = ts.read(&mut buf).unwrap();
    assert_eq!(buf, [1, 2, 3, 0]);
    assert_eq!(read_cnt, 3);
    assert_eq!(ts.incoming, []);
}

#[test]
fn test_stream_read_incremental() {
    let mut ts = TestStream::new(vec![1, 2, 3]);

    // Read once
    let mut buf = [0; 2];
    let read_cnt = ts.read(&mut buf).unwrap();
    assert_eq!(read_cnt, 2);
    assert_eq!(buf, [1, 2]);
    assert_eq!(ts.incoming, [3]);

    // Read once more
    let mut buf = [0; 2];
    let read_cnt = ts.read(&mut buf).unwrap();
    assert_eq!(read_cnt, 1);
    assert_eq!(buf, [3, 0]);
    assert_eq!(ts.incoming, []);
}

#[test]
fn test_stream_write() {
    let mut ts = TestStream::new(vec![]);

    // Write once
    let buf = [1, 2];
    let write_cnt = ts.write(&buf).unwrap();
    assert_eq!(write_cnt, 2);
    assert_eq!(ts.outgoing, [1, 2]);

    // Write once more
    let buf = [3, 4, 5];
    let write_cnt = ts.write(&buf).unwrap();
    assert_eq!(write_cnt, 3);
    assert_eq!(ts.outgoing, [1, 2, 3, 4, 5]);
}
<commit_msg>make TestStream use a cursor when reading instead of popping from the front<commit_after>use std::cmp;
use std::io::Read;
use std::io::Write;
use std::io::Result;

// A stream that is seeded with incoming data which can be consumed and
// records data written to it (like a socket).
//
// This allows us to unit test a transport without using sockets. :)
pub struct TestStream {
    pub incoming: Vec<u8>,
    pub incoming_cursor: usize,

    pub outgoing: Vec<u8>,
}

impl TestStream {
    pub fn new(incoming: Vec<u8>) -> TestStream {
        TestStream {
            incoming: incoming,
            incoming_cursor: 0,

            // Should be a good fit for most test responses
            outgoing: Vec::with_capacity(200),
        }
    }
}

impl Read for TestStream {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        // We can only read either as much as we have in incoming, or as big as
        // the output buffer is.
        let read_len = cmp::min(buf.len(),
                                self.incoming.len() - self.incoming_cursor);

        for i in 0..read_len {
            buf[i] = self.incoming[self.incoming_cursor + i];
        }
        self.incoming_cursor += read_len;

        Ok(read_len)
    }
}

impl Write for TestStream {
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        self.outgoing.extend(buf.iter().cloned());

        Ok(buf.len())
    }

    fn flush(&mut self) -> Result<()> {
        Ok(())
    }
}

#[test]
fn test_stream_read_whole() {
    let mut ts = TestStream::new(vec![1, 2, 3]);

    let mut buf = [0; 4];
    let read_cnt = ts.read(&mut buf).unwrap();
    assert_eq!(buf, [1, 2, 3, 0]);
    assert_eq!(read_cnt, 3);
    assert_eq!(ts.incoming, [1, 2, 3]);
    assert_eq!(ts.incoming_cursor, 3);
}

#[test]
fn test_stream_read_incremental() {
    let mut ts = TestStream::new(vec![1, 2, 3]);

    // Read once
    let mut buf = [0; 2];
    let read_cnt = ts.read(&mut buf).unwrap();
    assert_eq!(read_cnt, 2);
    assert_eq!(buf, [1, 2]);
    assert_eq!(ts.incoming, [1, 2, 3]);
    assert_eq!(ts.incoming_cursor, 2);

    // Read once more
    let mut buf = [0; 2];
    let read_cnt = ts.read(&mut buf).unwrap();
    assert_eq!(read_cnt, 1);
    assert_eq!(buf, [3, 0]);
    assert_eq!(ts.incoming, [1, 2, 3]);
    assert_eq!(ts.incoming_cursor, 3);
}

#[test]
fn test_stream_write() {
    let mut ts = TestStream::new(vec![]);

    // Write once
    let buf = [1, 2];
    let write_cnt = ts.write(&buf).unwrap();
    assert_eq!(write_cnt, 2);
    assert_eq!(ts.outgoing, [1, 2]);

    // Write once more
    let buf = [3, 4, 5];
    let write_cnt = ts.write(&buf).unwrap();
    assert_eq!(write_cnt, 3);
    assert_eq!(ts.outgoing, [1, 2, 3, 4, 5]);
}
<|endoftext|>