<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use backup_service::start_backup_service;
use consensus::{consensus_provider::start_consensus, gen_consensus_reconfig_subscription};
use debug_interface::node_debug_service::NodeDebugService;
use executor::{db_bootstrapper::bootstrap_db_if_empty, Executor};
use executor_types::ChunkExecutor;
use futures::{channel::mpsc::channel, executor::block_on};
use libra_config::{
    config::{NetworkConfig, NodeConfig, RoleType},
    utils::get_genesis_txn,
};
use libra_json_rpc::bootstrap_from_config as bootstrap_rpc;
use libra_logger::prelude::*;
use libra_mempool::gen_mempool_reconfig_subscription;
use libra_metrics::metric_server;
use libra_secure_storage::config;
use libra_vm::LibraVM;
use libradb::LibraDB;
use network_simple_onchain_discovery::{
    gen_simple_discovery_reconfig_subscription, ConfigurationChangeListener,
};
use state_synchronizer::StateSynchronizer;
use std::{boxed::Box, net::ToSocketAddrs, sync::Arc, thread, time::Instant};
use storage_interface::DbReaderWriter;
use storage_service::start_storage_service_with_db;
use tokio::runtime::Runtime;

use debug_interface::libra_trace;

const AC_SMP_CHANNEL_BUFFER_SIZE: usize = 1_024;
const INTRA_NODE_CHANNEL_BUFFER_SIZE: usize = 1;

pub struct LibraHandle {
    _rpc: Runtime,
    _mempool: Runtime,
    _state_synchronizer: StateSynchronizer,
    _network_runtimes: Vec<Runtime>,
    _consensus_runtime: Option<Runtime>,
    _debug: NodeDebugService,
    _backup: Runtime,
}

fn setup_chunk_executor(db: DbReaderWriter) -> Box<dyn ChunkExecutor> {
    Box::new(Executor::<LibraVM>::new(db))
}

fn setup_debug_interface(config: &NodeConfig) -> NodeDebugService {
    let addr = format!(
        "{}:{}",
        config.debug_interface.address, config.debug_interface.admission_control_node_debug_port,
    )
    .to_socket_addrs()
    .unwrap()
    .next()
    .unwrap();

    libra_trace::set_libra_trace(&config.debug_interface.libra_trace.sampling)
        .expect("Failed to set libra trace sampling rate.");

    NodeDebugService::new(addr)
}

pub fn setup_environment(node_config: &mut NodeConfig) -> LibraHandle {
    crash_handler::setup_panic_handler();

    // Some of our code uses the rayon global thread pool. Name the rayon threads so it doesn't
    // cause confusion, otherwise the threads would have their parent's name.
    rayon::ThreadPoolBuilder::new()
        .thread_name(|index| format!("rayon-global-{}", index))
        .build_global()
        .expect("Building rayon global thread pool should work.");

    let mut instant = Instant::now();
    let (libra_db, db_rw) = DbReaderWriter::wrap(
        LibraDB::open(
            &node_config.storage.dir(),
            false, /* readonly */
            node_config.storage.prune_window,
        )
        .expect("DB should open."),
    );
    let _simple_storage_service =
        start_storage_service_with_db(&node_config, Arc::clone(&libra_db));
    let backup_service = start_backup_service(
        node_config.storage.backup_service_port,
        Arc::clone(&libra_db),
    );

    bootstrap_db_if_empty::<LibraVM>(&db_rw, get_genesis_txn(&node_config).unwrap())
        .expect("Db-bootstrapper should not fail.");

    debug!(
        "Storage service started in {} ms",
        instant.elapsed().as_millis()
    );

    instant = Instant::now();
    let chunk_executor = setup_chunk_executor(db_rw.clone());
    debug!(
        "ChunkExecutor setup in {} ms",
        instant.elapsed().as_millis()
    );
    let mut network_runtimes = vec![];
    let mut state_sync_network_handles = vec![];
    let mut mempool_network_handles = vec![];
    let mut consensus_network_handles = None;
    let mut reconfig_subscriptions = vec![];

    let (mempool_reconfig_subscription, mempool_reconfig_events) =
        gen_mempool_reconfig_subscription();
    reconfig_subscriptions.push(mempool_reconfig_subscription);
    // consensus has to subscribe to ALL on-chain configs
    let (consensus_reconfig_subscription, consensus_reconfig_events) =
        gen_consensus_reconfig_subscription();
    reconfig_subscriptions.push(consensus_reconfig_subscription);

    let waypoint = config::waypoint(&node_config.base.waypoint);

    // Gather all network configs into a single vector.
    // TODO:  consider explicitly encoding the role in the NetworkConfig
    let mut network_configs: Vec<(RoleType, &mut NetworkConfig)> = node_config
        .full_node_networks
        .iter_mut()
        .map(|network_config| (RoleType::FullNode, network_config))
        .collect();
    if let Some(network_config) = node_config.validator_network.as_mut() {
        network_configs.push((RoleType::Validator, network_config));
    }

    // Instantiate every network and collect the requisite endpoints for state_sync, mempool, and consensus.
    for (role, network_config) in network_configs {
        // Perform common instantiation steps
        let (runtime, mut network_builder) = network_builder::builder::setup_network(
            &node_config.base.chain_id,
            role,
            network_config,
            Arc::clone(&db_rw.reader),
            waypoint,
        );
        let peer_id = network_builder.peer_id();

        // Create the endpoints to connect the Network to StateSynchronizer.
        let (state_sync_sender, state_sync_events) = network_builder
            .add_protocol_handler(state_synchronizer::network::network_endpoint_config());
        state_sync_network_handles.push((peer_id, state_sync_sender, state_sync_events));

        // Create the endpoints t connect the Network to MemPool.
        let (mempool_sender, mempool_events) =
            network_builder.add_protocol_handler(libra_mempool::network::network_endpoint_config(
                // TODO:  Make this configuration option more clear.
                node_config.mempool.max_broadcasts_per_peer,
            ));
        mempool_network_handles.push((peer_id, mempool_sender, mempool_events));

        match role {
            // Perform steps relevant specifically to Validator networks.
            RoleType::Validator => {
                // A valid config is allowed to have at most one ValidatorNetwork
                // TODO:  `expect_none` would be perfect here, once it is stable.
                if consensus_network_handles.is_some() {
                    panic!("There can be at most one validator network!");
                }

                // Set up to listen for network configuration changes from StateSync.
                // TODO:  move this inside network_builder.
                if let Some(conn_mgr_reqs_tx) = network_builder.conn_mgr_reqs_tx() {
                    let (simple_discovery_reconfig_subscription, simple_discovery_reconfig_rx) =
                        gen_simple_discovery_reconfig_subscription();
                    reconfig_subscriptions.push(simple_discovery_reconfig_subscription);
                    let network_config_listener =
                        ConfigurationChangeListener::new(conn_mgr_reqs_tx, RoleType::Validator);
                    runtime
                        .handle()
                        .spawn(network_config_listener.start(simple_discovery_reconfig_rx));
                };

                consensus_network_handles =
                    Some(network_builder.add_protocol_handler(
                        consensus::network_interface::network_endpoint_config(),
                    ));
            }
            // Currently no FullNode network specific steps.
            RoleType::FullNode => (),
        }

        // Start the network and cache the runtime so it does not go out of scope.
        // TODO:  move all 'start' commands to a second phase at the end of setup_environment.  Target is to have one pass to wire the pieces together and a second pass to start processing in an appropriate order.
        let _listen_addr = network_builder.build();
        network_runtimes.push(runtime);
        debug!("Network started for peer_id: {}", peer_id);
    }

    // TODO set up on-chain discovery network based on UpstreamConfig.fallback_network
    // and pass network handles to mempool/state sync

    // for state sync to send requests to mempool
    let (state_sync_to_mempool_sender, state_sync_requests) =
        channel(INTRA_NODE_CHANNEL_BUFFER_SIZE);
    let state_synchronizer = StateSynchronizer::bootstrap(
        state_sync_network_handles,
        state_sync_to_mempool_sender,
        Arc::clone(&db_rw.reader),
        chunk_executor,
        &node_config,
        waypoint,
        reconfig_subscriptions,
    );
    let (mp_client_sender, mp_client_events) = channel(AC_SMP_CHANNEL_BUFFER_SIZE);

    let rpc_runtime = bootstrap_rpc(&node_config, libra_db.clone(), mp_client_sender);

    let mut consensus_runtime = None;
    let (consensus_to_mempool_sender, consensus_requests) = channel(INTRA_NODE_CHANNEL_BUFFER_SIZE);

    instant = Instant::now();
    let mempool = libra_mempool::bootstrap(
        node_config,
        Arc::clone(&db_rw.reader),
        mempool_network_handles,
        mp_client_events,
        consensus_requests,
        state_sync_requests,
        mempool_reconfig_events,
    );
    debug!("Mempool started in {} ms", instant.elapsed().as_millis());

    // StateSync should be instantiated and started before Consensus to avoid a cyclic dependency:
    // network provider -> consensus -> state synchronizer -> network provider.  This has resulted
    // in a deadlock as observed in GitHub issue #749.
    if let Some((consensus_network_sender, consensus_network_events)) = consensus_network_handles {
        // Make sure that state synchronizer is caught up at least to its waypoint
        // (in case it's present). There is no sense to start consensus prior to that.
        // TODO: Note that we need the networking layer to be able to discover & connect to the
        // peers with potentially outdated network identity public keys.
        debug!("Wait until state synchronizer is initialized");
        block_on(state_synchronizer.wait_until_initialized())
            .expect("State synchronizer initialization failure");
        debug!("State synchronizer initialization complete.");

        // Initialize and start consensus.
        instant = Instant::now();
        consensus_runtime = Some(start_consensus(
            node_config,
            consensus_network_sender,
            consensus_network_events,
            state_synchronizer.create_client(),
            consensus_to_mempool_sender,
            libra_db,
            consensus_reconfig_events,
        ));
        debug!("Consensus started in {} ms", instant.elapsed().as_millis());
    }

    let debug_if = setup_debug_interface(&node_config);

    let metrics_port = node_config.debug_interface.metrics_server_port;
    let metric_host = node_config.debug_interface.address.clone();
    thread::spawn(move || metric_server::start_server(metric_host, metrics_port, false));
    let public_metrics_port = node_config.debug_interface.public_metrics_server_port;
    let public_metric_host = node_config.debug_interface.address.clone();
    thread::spawn(move || {
        metric_server::start_server(public_metric_host, public_metrics_port, true)
    });

    LibraHandle {
        _network_runtimes: network_runtimes,
        _rpc: rpc_runtime,
        _mempool: mempool,
        _state_synchronizer: state_synchronizer,
        _consensus_runtime: consensus_runtime,
        _debug: debug_if,
        _backup: backup_service,
    }
}
<commit_msg>[libra-node] Remove duplicate setup_panic_handler<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use backup_service::start_backup_service;
use consensus::{consensus_provider::start_consensus, gen_consensus_reconfig_subscription};
use debug_interface::node_debug_service::NodeDebugService;
use executor::{db_bootstrapper::bootstrap_db_if_empty, Executor};
use executor_types::ChunkExecutor;
use futures::{channel::mpsc::channel, executor::block_on};
use libra_config::{
    config::{NetworkConfig, NodeConfig, RoleType},
    utils::get_genesis_txn,
};
use libra_json_rpc::bootstrap_from_config as bootstrap_rpc;
use libra_logger::prelude::*;
use libra_mempool::gen_mempool_reconfig_subscription;
use libra_metrics::metric_server;
use libra_secure_storage::config;
use libra_vm::LibraVM;
use libradb::LibraDB;
use network_simple_onchain_discovery::{
    gen_simple_discovery_reconfig_subscription, ConfigurationChangeListener,
};
use state_synchronizer::StateSynchronizer;
use std::{boxed::Box, net::ToSocketAddrs, sync::Arc, thread, time::Instant};
use storage_interface::DbReaderWriter;
use storage_service::start_storage_service_with_db;
use tokio::runtime::Runtime;

use debug_interface::libra_trace;

const AC_SMP_CHANNEL_BUFFER_SIZE: usize = 1_024;
const INTRA_NODE_CHANNEL_BUFFER_SIZE: usize = 1;

pub struct LibraHandle {
    _rpc: Runtime,
    _mempool: Runtime,
    _state_synchronizer: StateSynchronizer,
    _network_runtimes: Vec<Runtime>,
    _consensus_runtime: Option<Runtime>,
    _debug: NodeDebugService,
    _backup: Runtime,
}

fn setup_chunk_executor(db: DbReaderWriter) -> Box<dyn ChunkExecutor> {
    Box::new(Executor::<LibraVM>::new(db))
}

fn setup_debug_interface(config: &NodeConfig) -> NodeDebugService {
    let addr = format!(
        "{}:{}",
        config.debug_interface.address, config.debug_interface.admission_control_node_debug_port,
    )
    .to_socket_addrs()
    .unwrap()
    .next()
    .unwrap();

    libra_trace::set_libra_trace(&config.debug_interface.libra_trace.sampling)
        .expect("Failed to set libra trace sampling rate.");

    NodeDebugService::new(addr)
}

pub fn setup_environment(node_config: &mut NodeConfig) -> LibraHandle {
    // Some of our code uses the rayon global thread pool. Name the rayon threads so it doesn't
    // cause confusion, otherwise the threads would have their parent's name.
    rayon::ThreadPoolBuilder::new()
        .thread_name(|index| format!("rayon-global-{}", index))
        .build_global()
        .expect("Building rayon global thread pool should work.");

    let mut instant = Instant::now();
    let (libra_db, db_rw) = DbReaderWriter::wrap(
        LibraDB::open(
            &node_config.storage.dir(),
            false, /* readonly */
            node_config.storage.prune_window,
        )
        .expect("DB should open."),
    );
    let _simple_storage_service =
        start_storage_service_with_db(&node_config, Arc::clone(&libra_db));
    let backup_service = start_backup_service(
        node_config.storage.backup_service_port,
        Arc::clone(&libra_db),
    );

    bootstrap_db_if_empty::<LibraVM>(&db_rw, get_genesis_txn(&node_config).unwrap())
        .expect("Db-bootstrapper should not fail.");

    debug!(
        "Storage service started in {} ms",
        instant.elapsed().as_millis()
    );

    instant = Instant::now();
    let chunk_executor = setup_chunk_executor(db_rw.clone());
    debug!(
        "ChunkExecutor setup in {} ms",
        instant.elapsed().as_millis()
    );
    let mut network_runtimes = vec![];
    let mut state_sync_network_handles = vec![];
    let mut mempool_network_handles = vec![];
    let mut consensus_network_handles = None;
    let mut reconfig_subscriptions = vec![];

    let (mempool_reconfig_subscription, mempool_reconfig_events) =
        gen_mempool_reconfig_subscription();
    reconfig_subscriptions.push(mempool_reconfig_subscription);
    // consensus has to subscribe to ALL on-chain configs
    let (consensus_reconfig_subscription, consensus_reconfig_events) =
        gen_consensus_reconfig_subscription();
    reconfig_subscriptions.push(consensus_reconfig_subscription);

    let waypoint = config::waypoint(&node_config.base.waypoint);

    // Gather all network configs into a single vector.
    // TODO:  consider explicitly encoding the role in the NetworkConfig
    let mut network_configs: Vec<(RoleType, &mut NetworkConfig)> = node_config
        .full_node_networks
        .iter_mut()
        .map(|network_config| (RoleType::FullNode, network_config))
        .collect();
    if let Some(network_config) = node_config.validator_network.as_mut() {
        network_configs.push((RoleType::Validator, network_config));
    }

    // Instantiate every network and collect the requisite endpoints for state_sync, mempool, and consensus.
    for (role, network_config) in network_configs {
        // Perform common instantiation steps
        let (runtime, mut network_builder) = network_builder::builder::setup_network(
            &node_config.base.chain_id,
            role,
            network_config,
            Arc::clone(&db_rw.reader),
            waypoint,
        );
        let peer_id = network_builder.peer_id();

        // Create the endpoints to connect the Network to StateSynchronizer.
        let (state_sync_sender, state_sync_events) = network_builder
            .add_protocol_handler(state_synchronizer::network::network_endpoint_config());
        state_sync_network_handles.push((peer_id, state_sync_sender, state_sync_events));

        // Create the endpoints t connect the Network to MemPool.
        let (mempool_sender, mempool_events) =
            network_builder.add_protocol_handler(libra_mempool::network::network_endpoint_config(
                // TODO:  Make this configuration option more clear.
                node_config.mempool.max_broadcasts_per_peer,
            ));
        mempool_network_handles.push((peer_id, mempool_sender, mempool_events));

        match role {
            // Perform steps relevant specifically to Validator networks.
            RoleType::Validator => {
                // A valid config is allowed to have at most one ValidatorNetwork
                // TODO:  `expect_none` would be perfect here, once it is stable.
                if consensus_network_handles.is_some() {
                    panic!("There can be at most one validator network!");
                }

                // Set up to listen for network configuration changes from StateSync.
                // TODO:  move this inside network_builder.
                if let Some(conn_mgr_reqs_tx) = network_builder.conn_mgr_reqs_tx() {
                    let (simple_discovery_reconfig_subscription, simple_discovery_reconfig_rx) =
                        gen_simple_discovery_reconfig_subscription();
                    reconfig_subscriptions.push(simple_discovery_reconfig_subscription);
                    let network_config_listener =
                        ConfigurationChangeListener::new(conn_mgr_reqs_tx, RoleType::Validator);
                    runtime
                        .handle()
                        .spawn(network_config_listener.start(simple_discovery_reconfig_rx));
                };

                consensus_network_handles =
                    Some(network_builder.add_protocol_handler(
                        consensus::network_interface::network_endpoint_config(),
                    ));
            }
            // Currently no FullNode network specific steps.
            RoleType::FullNode => (),
        }

        // Start the network and cache the runtime so it does not go out of scope.
        // TODO:  move all 'start' commands to a second phase at the end of setup_environment.  Target is to have one pass to wire the pieces together and a second pass to start processing in an appropriate order.
        let _listen_addr = network_builder.build();
        network_runtimes.push(runtime);
        debug!("Network started for peer_id: {}", peer_id);
    }

    // TODO set up on-chain discovery network based on UpstreamConfig.fallback_network
    // and pass network handles to mempool/state sync

    // for state sync to send requests to mempool
    let (state_sync_to_mempool_sender, state_sync_requests) =
        channel(INTRA_NODE_CHANNEL_BUFFER_SIZE);
    let state_synchronizer = StateSynchronizer::bootstrap(
        state_sync_network_handles,
        state_sync_to_mempool_sender,
        Arc::clone(&db_rw.reader),
        chunk_executor,
        &node_config,
        waypoint,
        reconfig_subscriptions,
    );
    let (mp_client_sender, mp_client_events) = channel(AC_SMP_CHANNEL_BUFFER_SIZE);

    let rpc_runtime = bootstrap_rpc(&node_config, libra_db.clone(), mp_client_sender);

    let mut consensus_runtime = None;
    let (consensus_to_mempool_sender, consensus_requests) = channel(INTRA_NODE_CHANNEL_BUFFER_SIZE);

    instant = Instant::now();
    let mempool = libra_mempool::bootstrap(
        node_config,
        Arc::clone(&db_rw.reader),
        mempool_network_handles,
        mp_client_events,
        consensus_requests,
        state_sync_requests,
        mempool_reconfig_events,
    );
    debug!("Mempool started in {} ms", instant.elapsed().as_millis());

    // StateSync should be instantiated and started before Consensus to avoid a cyclic dependency:
    // network provider -> consensus -> state synchronizer -> network provider.  This has resulted
    // in a deadlock as observed in GitHub issue #749.
    if let Some((consensus_network_sender, consensus_network_events)) = consensus_network_handles {
        // Make sure that state synchronizer is caught up at least to its waypoint
        // (in case it's present). There is no sense to start consensus prior to that.
        // TODO: Note that we need the networking layer to be able to discover & connect to the
        // peers with potentially outdated network identity public keys.
        debug!("Wait until state synchronizer is initialized");
        block_on(state_synchronizer.wait_until_initialized())
            .expect("State synchronizer initialization failure");
        debug!("State synchronizer initialization complete.");

        // Initialize and start consensus.
        instant = Instant::now();
        consensus_runtime = Some(start_consensus(
            node_config,
            consensus_network_sender,
            consensus_network_events,
            state_synchronizer.create_client(),
            consensus_to_mempool_sender,
            libra_db,
            consensus_reconfig_events,
        ));
        debug!("Consensus started in {} ms", instant.elapsed().as_millis());
    }

    let debug_if = setup_debug_interface(&node_config);

    let metrics_port = node_config.debug_interface.metrics_server_port;
    let metric_host = node_config.debug_interface.address.clone();
    thread::spawn(move || metric_server::start_server(metric_host, metrics_port, false));
    let public_metrics_port = node_config.debug_interface.public_metrics_server_port;
    let public_metric_host = node_config.debug_interface.address.clone();
    thread::spawn(move || {
        metric_server::start_server(public_metric_host, public_metrics_port, true)
    });

    LibraHandle {
        _network_runtimes: network_runtimes,
        _rpc: rpc_runtime,
        _mempool: mempool,
        _state_synchronizer: state_synchronizer,
        _consensus_runtime: consensus_runtime,
        _debug: debug_if,
        _backup: backup_service,
    }
}
<|endoftext|>


<commit_before>use display_delegate::display;
use lifetime_remover::remove_lifetimes_from_path;
use proc_macro2::{Group, Span, TokenTree};
use quote::{ToTokens};
use std::fmt::{Error, Formatter};
use syn::{self, ArgCaptured, Block, Expr, ExprVerbatim, FnArg, FnDecl, GenericParam, Ident, Pat, PatIdent, PathSegment,
          Stmt};
use syn::punctuated::Punctuated;
use syn::token::{Colon2, Semi};

const MOCKTOPUS_CRATE_NAME:     &str = "__mocktopus_crate__";
const STD_CRATE_NAME:           &str = "__mocktopus_std__";
const ARGS_TO_CONTINUE_NAME:    &str = "__mocktopus_args_to_continue__";
const UNWIND_DATA_NAME:         &str = "__mocktopus_unwind_data__";

macro_rules! error_msg {
    ($msg:expr) => { concat!("Mocktopus internal error: ", $msg) }
}

pub enum FnHeaderBuilder<'a> {
    StaticFn,
    StructImpl,
    TraitDefault,
    TraitImpl(&'a Punctuated<PathSegment, Colon2>),
}

impl<'a> FnHeaderBuilder<'a> {
    pub fn build(&self, fn_ident: &Ident, fn_decl: &FnDecl, fn_block_span: Span) -> Stmt {
        let fn_args = &fn_decl.inputs;
        let header_str = format!(
r#"{{
    extern crate mocktopus as {mocktopus};
    extern crate std as {std_crate};
    match {std_crate}::panic::catch_unwind({std_crate}::panic::AssertUnwindSafe (
            || {mocktopus}::mocking::Mockable::call_mock(&{full_fn_name}, {extract_args}))) {{
        Ok({mocktopus}::mocking::MockResult::Continue({args_to_continue})) => {restore_args},
        Ok({mocktopus}::mocking::MockResult::Return(result)) => {{
            {forget_args}
            let returned = unsafe {{ {std_crate}::mem::transmute_copy(&result) }};
            {std_crate}::mem::forget(result);
            return returned;
        }},
        Err({unwind}) => {{
            {forget_args}
            {std_crate}::panic::resume_unwind({unwind});
        }},
    }}
}}"#,
        mocktopus           = MOCKTOPUS_CRATE_NAME,
        std_crate           = STD_CRATE_NAME,
        full_fn_name        = display(|f| write_full_fn_name(f, self, fn_ident, fn_decl)),
        extract_args        = display(|f| write_extract_args(f, fn_args)),
        args_to_continue    = ARGS_TO_CONTINUE_NAME,
        restore_args        = display(|f| write_restore_args(f, fn_args)),
        forget_args         = display(|f| write_forget_args(f, fn_args)),
        unwind              = UNWIND_DATA_NAME);
        let header_block = syn::parse_str::<Block>(&header_str)
            .expect(error_msg!("generated header unparsable"));
        create_call_site_spanned_stmt(header_block, fn_block_span)
    }
}

fn create_call_site_spanned_stmt(block: Block, span: Span) -> Stmt {
    let token_stream = block.into_tokens()
        .into_iter()
        .map(|tt| make_token_tree_span_call_site(tt, span))
        .collect();
    Stmt::Semi(
        Expr::Verbatim(
            ExprVerbatim {
                tts: token_stream,
            }
        ),
        Semi([span])
    )
}

fn make_token_tree_span_call_site(mut token_tree: TokenTree, span: Span) -> TokenTree {
    token_tree.set_span(span);
    if let TokenTree::Group(ref mut group) = token_tree {
        let tokens = group.stream().into_iter()
            .map(|tt| make_token_tree_span_call_site(tt, span))
            .collect();
        *group = Group::new(group.delimiter(), tokens);
    }
    token_tree
}

fn write_full_fn_name(f: &mut Formatter, builder: &FnHeaderBuilder, fn_ident: &Ident, fn_decl: &FnDecl)
        -> Result<(), Error> {
    match *builder {
        FnHeaderBuilder::StaticFn               => (),
        FnHeaderBuilder::StructImpl |
        FnHeaderBuilder::TraitDefault           => write!(f, "Self::")?,
        FnHeaderBuilder::TraitImpl(ref path)    => write!(f, "<Self as {}>::", display(|f| write_trait_path(f, path)))?,
    }
    write!(f, "{}::<{}>", fn_ident.as_ref(), display(|f| write_fn_generics(f, fn_decl)))
}

fn write_trait_path<T: ToTokens + Clone>(f: &mut Formatter, trait_path: &Punctuated<PathSegment, T>) -> Result<(), Error> {
    let mut trait_path_without_lifetimes = trait_path.clone();
    remove_lifetimes_from_path(&mut trait_path_without_lifetimes);
    write!(f, "{}", trait_path_without_lifetimes.into_tokens())
}

fn write_fn_generics(f: &mut Formatter, fn_decl: &FnDecl) -> Result<(), Error> {
    fn_decl.generics.params.iter()
        .filter_map(get_generic_param_name)
        .map(|param| write!(f, "{},", param))
        .collect()
}

fn get_generic_param_name<'a>(param: &'a GenericParam) -> Option<&'a str> {
    match *param {
        GenericParam::Type(ref type_param) => Some(type_param.ident.as_ref()),
        _ => None,
    }
}

fn write_extract_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    if fn_args.is_empty() {
        return write!(f, "()");
    }
    write!(f, "unsafe {{ (")?;
    for fn_arg_name in iter_fn_arg_names(fn_args) {
        write!(f, "{}::mem::replace(&mut *(&{} as *const _ as *mut _), {}::mem::uninitialized()), ",
               STD_CRATE_NAME, fn_arg_name, STD_CRATE_NAME)?;
    }
    write!(f, ") }}")
}

fn write_restore_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    if fn_args.is_empty() {
        return writeln!(f, "()");
    }
    writeln!(f, "unsafe {{")?;
    for (fn_arg_index, fn_arg_name) in iter_fn_arg_names(fn_args).enumerate() {
        writeln!(f, "{}::mem::replace({}::mocking_utils::as_mut(&{}), {}.{});",
                 STD_CRATE_NAME, MOCKTOPUS_CRATE_NAME, fn_arg_name, ARGS_TO_CONTINUE_NAME, fn_arg_index)?;
    }
    writeln!(f, "}}")
}

fn write_forget_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    for fn_arg_name in iter_fn_arg_names(fn_args) {
        writeln!(f, "{}::mem::forget({});", STD_CRATE_NAME, fn_arg_name)?;
    }
    Ok(())
}

fn iter_fn_arg_names<'a, T>(input_args: &'a Punctuated<FnArg, T>) -> impl Iterator<Item = &'a str> {
    input_args.iter()
        .map(|fn_arg| match *fn_arg {
            FnArg::SelfRef(_) | FnArg::SelfValue(_) => "self",
            FnArg::Captured(
                ArgCaptured {
                    pat: Pat::Ident(
                        PatIdent {
                            ref ident,
                            subpat: None,
                            ..
                        }
                    ),
                    ..
                }
            ) => ident.as_ref(),
            _ => panic!("{}: '{}'", error_msg!("invalid fn arg type"), fn_arg.clone().into_tokens()),
        })
}
<commit_msg>Fix uninitialised memory access when fn arg has drop and mock returns continue<commit_after>use display_delegate::display;
use lifetime_remover::remove_lifetimes_from_path;
use proc_macro2::{Group, Span, TokenTree};
use quote::{ToTokens};
use std::fmt::{Error, Formatter};
use syn::{self, ArgCaptured, Block, Expr, ExprVerbatim, FnArg, FnDecl, GenericParam, Ident, Pat, PatIdent, PathSegment,
          Stmt};
use syn::punctuated::Punctuated;
use syn::token::{Colon2, Semi};

const MOCKTOPUS_CRATE_NAME:     &str = "__mocktopus_crate__";
const STD_CRATE_NAME:           &str = "__mocktopus_std__";
const ARGS_TO_CONTINUE_NAME:    &str = "__mocktopus_args_to_continue__";
const UNWIND_DATA_NAME:         &str = "__mocktopus_unwind_data__";

macro_rules! error_msg {
    ($msg:expr) => { concat!("Mocktopus internal error: ", $msg) }
}

pub enum FnHeaderBuilder<'a> {
    StaticFn,
    StructImpl,
    TraitDefault,
    TraitImpl(&'a Punctuated<PathSegment, Colon2>),
}

impl<'a> FnHeaderBuilder<'a> {
    pub fn build(&self, fn_ident: &Ident, fn_decl: &FnDecl, fn_block_span: Span) -> Stmt {
        let fn_args = &fn_decl.inputs;
        let header_str = format!(
r#"{{
    extern crate mocktopus as {mocktopus};
    extern crate std as {std_crate};
    match {std_crate}::panic::catch_unwind({std_crate}::panic::AssertUnwindSafe (
            || {mocktopus}::mocking::Mockable::call_mock(&{full_fn_name}, {extract_args}))) {{
        Ok({mocktopus}::mocking::MockResult::Continue({args_to_continue})) => {restore_args},
        Ok({mocktopus}::mocking::MockResult::Return(result)) => {{
            {forget_args}
            let returned = unsafe {{ {std_crate}::mem::transmute_copy(&result) }};
            {std_crate}::mem::forget(result);
            return returned;
        }},
        Err({unwind}) => {{
            {forget_args}
            {std_crate}::panic::resume_unwind({unwind});
        }},
    }}
}}"#,
        mocktopus           = MOCKTOPUS_CRATE_NAME,
        std_crate           = STD_CRATE_NAME,
        full_fn_name        = display(|f| write_full_fn_name(f, self, fn_ident, fn_decl)),
        extract_args        = display(|f| write_extract_args(f, fn_args)),
        args_to_continue    = ARGS_TO_CONTINUE_NAME,
        restore_args        = display(|f| write_restore_args(f, fn_args)),
        forget_args         = display(|f| write_forget_args(f, fn_args)),
        unwind              = UNWIND_DATA_NAME);
        let header_block = syn::parse_str::<Block>(&header_str)
            .expect(error_msg!("generated header unparsable"));
        create_call_site_spanned_stmt(header_block, fn_block_span)
    }
}

fn create_call_site_spanned_stmt(block: Block, span: Span) -> Stmt {
    let token_stream = block.into_tokens()
        .into_iter()
        .map(|tt| make_token_tree_span_call_site(tt, span))
        .collect();
    Stmt::Semi(
        Expr::Verbatim(
            ExprVerbatim {
                tts: token_stream,
            }
        ),
        Semi([span])
    )
}

fn make_token_tree_span_call_site(mut token_tree: TokenTree, span: Span) -> TokenTree {
    token_tree.set_span(span);
    if let TokenTree::Group(ref mut group) = token_tree {
        let tokens = group.stream().into_iter()
            .map(|tt| make_token_tree_span_call_site(tt, span))
            .collect();
        *group = Group::new(group.delimiter(), tokens);
    }
    token_tree
}

fn write_full_fn_name(f: &mut Formatter, builder: &FnHeaderBuilder, fn_ident: &Ident, fn_decl: &FnDecl)
        -> Result<(), Error> {
    match *builder {
        FnHeaderBuilder::StaticFn               => (),
        FnHeaderBuilder::StructImpl |
        FnHeaderBuilder::TraitDefault           => write!(f, "Self::")?,
        FnHeaderBuilder::TraitImpl(ref path)    => write!(f, "<Self as {}>::", display(|f| write_trait_path(f, path)))?,
    }
    write!(f, "{}::<{}>", fn_ident.as_ref(), display(|f| write_fn_generics(f, fn_decl)))
}

fn write_trait_path<T: ToTokens + Clone>(f: &mut Formatter, trait_path: &Punctuated<PathSegment, T>) -> Result<(), Error> {
    let mut trait_path_without_lifetimes = trait_path.clone();
    remove_lifetimes_from_path(&mut trait_path_without_lifetimes);
    write!(f, "{}", trait_path_without_lifetimes.into_tokens())
}

fn write_fn_generics(f: &mut Formatter, fn_decl: &FnDecl) -> Result<(), Error> {
    fn_decl.generics.params.iter()
        .filter_map(get_generic_param_name)
        .map(|param| write!(f, "{},", param))
        .collect()
}

fn get_generic_param_name<'a>(param: &'a GenericParam) -> Option<&'a str> {
    match *param {
        GenericParam::Type(ref type_param) => Some(type_param.ident.as_ref()),
        _ => None,
    }
}

fn write_extract_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    if fn_args.is_empty() {
        return write!(f, "()");
    }
    write!(f, "unsafe {{ (")?;
    for fn_arg_name in iter_fn_arg_names(fn_args) {
        write!(f, "{}::mem::replace(&mut *(&{} as *const _ as *mut _), {}::mem::uninitialized()), ",
               STD_CRATE_NAME, fn_arg_name, STD_CRATE_NAME)?;
    }
    write!(f, ") }}")
}

fn write_restore_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    if fn_args.is_empty() {
        return writeln!(f, "()");
    }
    writeln!(f, "unsafe {{")?;
    for (fn_arg_index, fn_arg_name) in iter_fn_arg_names(fn_args).enumerate() {
        writeln!(f, "{}::mem::forget({}::mem::replace({}::mocking_utils::as_mut(&{}), {}.{}));",
            STD_CRATE_NAME, STD_CRATE_NAME, MOCKTOPUS_CRATE_NAME, fn_arg_name, ARGS_TO_CONTINUE_NAME, fn_arg_index)?;
    }
    writeln!(f, "}}")
}

fn write_forget_args<T>(f: &mut Formatter, fn_args: &Punctuated<FnArg, T>) -> Result<(), Error> {
    for fn_arg_name in iter_fn_arg_names(fn_args) {
        writeln!(f, "{}::mem::forget({});", STD_CRATE_NAME, fn_arg_name)?;
    }
    Ok(())
}

fn iter_fn_arg_names<'a, T>(input_args: &'a Punctuated<FnArg, T>) -> impl Iterator<Item = &'a str> {
    input_args.iter()
        .map(|fn_arg| match *fn_arg {
            FnArg::SelfRef(_) | FnArg::SelfValue(_) => "self",
            FnArg::Captured(
                ArgCaptured {
                    pat: Pat::Ident(
                        PatIdent {
                            ref ident,
                            subpat: None,
                            ..
                        }
                    ),
                    ..
                }
            ) => ident.as_ref(),
            _ => panic!("{}: '{}'", error_msg!("invalid fn arg type"), fn_arg.clone().into_tokens()),
        })
}
<|endoftext|>


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::io::stdout;

use lister::Lister;
use error::Result;
use error::ResultExt;

use libimagstore::store::FileLockEntry;

use prettytable::Table;
use prettytable::cell::Cell;
use prettytable::row::Row;

pub struct TableLister<F: Fn(&FileLockEntry) -> Vec<String>> {
    line_generator: F,
    header: Option<Vec<String>>,

    with_idx: bool,
}

impl<F: Fn(&FileLockEntry) -> Vec<String>> TableLister<F> {

    pub fn new(gen: F) -> TableLister<F> {
        TableLister {
            line_generator: gen,
            header: None,
            with_idx: true,
        }
    }

    pub fn with_header(mut self, hdr: Vec<String>) -> TableLister<F> {
        self.header = Some(hdr);
        self
    }

    pub fn with_idx(mut self, b: bool) -> TableLister<F> {
        self.with_idx = b;
        self
    }

}

impl<F: Fn(&FileLockEntry) -> Vec<String>> Lister for TableLister<F> {

    fn list<'b, I: Iterator<Item = FileLockEntry<'b>>>(&self, entries: I) -> Result<()> {
        use error::ListErrorKind as LEK;
        use error::ListError as LE;

        let mut table = Table::new();
        let mut header_len : Option<usize> = None;
        match self.header {
            Some(ref s) => {
                debug!("We have a header... preparing");
                let mut cells : Vec<Cell> = s.iter().map(|s| Cell::new(s)).collect();
                if self.with_idx {
                    cells.insert(0, Cell::new("#"));
                }
                table.set_titles(Row::new(cells));
                header_len = Some(s.len());
            },
            None => {
                debug!("No header for table found... continuing without");
            },
        }

        entries.fold(Ok(table), |table, entry| {
            table.and_then(|mut table| {
                let mut v = (self.line_generator)(&entry);
                {
                    let v_len = v.len();
                    if header_len.is_none() {
                        header_len = Some(v_len);
                    }
                    if header_len.map(|l| v_len > l).unwrap_or(false) {
                        return Err(LE::from_kind(LEK::FormatError));
                    }
                    while header_len.map(|l| v.len() != l).unwrap_or(false) {
                        v.push(String::from(""));
                    }
                }

                table.add_row(v.iter().map(|s| Cell::new(s)).collect());
                Ok(table)
            })
        })
        .and_then(|tbl| {
            let mut io = stdout();
            tbl.print(&mut io).chain_err(|| LEK::IOError)
        })
    }

}
<commit_msg>TableLister: If configured with index, also print index row<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::io::stdout;

use lister::Lister;
use error::Result;
use error::ResultExt;

use libimagstore::store::FileLockEntry;

use prettytable::Table;
use prettytable::cell::Cell;
use prettytable::row::Row;

pub struct TableLister<F: Fn(&FileLockEntry) -> Vec<String>> {
    line_generator: F,
    header: Option<Vec<String>>,

    with_idx: bool,
}

impl<F: Fn(&FileLockEntry) -> Vec<String>> TableLister<F> {

    pub fn new(gen: F) -> TableLister<F> {
        TableLister {
            line_generator: gen,
            header: None,
            with_idx: true,
        }
    }

    pub fn with_header(mut self, hdr: Vec<String>) -> TableLister<F> {
        self.header = Some(hdr);
        self
    }

    pub fn with_idx(mut self, b: bool) -> TableLister<F> {
        self.with_idx = b;
        self
    }

}

impl<F: Fn(&FileLockEntry) -> Vec<String>> Lister for TableLister<F> {

    fn list<'b, I: Iterator<Item = FileLockEntry<'b>>>(&self, entries: I) -> Result<()> {
        use error::ListErrorKind as LEK;
        use error::ListError as LE;

        let mut table = Table::new();
        let mut header_len : Option<usize> = None;
        match self.header {
            Some(ref s) => {
                debug!("We have a header... preparing");
                let mut cells : Vec<Cell> = s.iter().map(|s| Cell::new(s)).collect();
                if self.with_idx {
                    cells.insert(0, Cell::new("#"));
                }
                table.set_titles(Row::new(cells));
                header_len = Some(s.len());
            },
            None => {
                debug!("No header for table found... continuing without");
            },
        }

        entries.enumerate().fold(Ok(table), |table, (i, entry)| {
            table.and_then(|mut table| {
                let mut v = (self.line_generator)(&entry);
                {
                    let v_len = v.len();
                    if header_len.is_none() {
                        header_len = Some(v_len);
                    }
                    if header_len.map(|l| v_len > l).unwrap_or(false) {
                        return Err(LE::from_kind(LEK::FormatError));
                    }
                    while header_len.map(|l| v.len() != l).unwrap_or(false) {
                        v.push(String::from(""));
                    }
                }

                if self.with_idx {
                    v.insert(0, format!("{}", i));
                }

                table.add_row(v.iter().map(|s| Cell::new(s)).collect());
                Ok(table)
            })
        })
        .and_then(|tbl| {
            let mut io = stdout();
            tbl.print(&mut io).chain_err(|| LEK::IOError)
        })
    }

}
<|endoftext|>


<commit_before>// Copyright 2021 lowRISC 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
//
//     https://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.
//
// SPDX-License-Identifier: Apache-2.0

use crate::manticore_support::Identity;
use crate::manticore_support::NoRsa;
use crate::manticore_support::Reset;
use crate::spi_host;
use crate::spi_host_h1;
use crate::spi_device;

use core::cmp::min;
use core::convert::TryFrom;
use core::fmt::Write;

use libtock::console::Console;
use libtock::result::TockError;

use manticore::io::Cursor as ManticoreCursor;
use manticore::server::pa_rot::PaRot;

use spiutils::io::Cursor as SpiutilsCursor;
use spiutils::io::Write as SpiutilsWrite;
use spiutils::protocol::flash;
use spiutils::protocol::flash::Address;
use spiutils::protocol::flash::AddressMode;
use spiutils::protocol::flash::OpCode;
use spiutils::protocol::payload;
use spiutils::protocol::wire::FromWire;
use spiutils::protocol::wire::FromWireError;
use spiutils::protocol::wire::ToWire;
use spiutils::protocol::wire::ToWireError;

// Size of the SPI flash chip.
// Hard-coded to 64 MiB for now.
// TODO(osenft): Modify this to be read from the actual SPI flash chip at runtime.
pub const SPI_FLASH_SIZE: u32 = 0x4000000;

// The location of the mailbox.
// TODO(osenft): Make this configurable, possibly by reading it from the SPI flash chip.
pub const SPI_MAILBOX_ADDRESS: u32 = 0xf00000;

// The size of the mailbox.
const SPI_MAILBOX_SIZE: u32 = spi_device::MAX_READ_BUFFER_SIZE as u32;

#[derive(Copy, Clone, Debug)]
pub enum SpiProcessorError {
    FromWire(FromWireError),
    ToWire(ToWireError),
    Tock,
    Manticore(manticore::server::Error),
    UnsupportedContentType(payload::ContentType),
    UnsupportedOpCode(OpCode),
    InvalidAddress(Option<u32>),
    Format(core::fmt::Error),
}

impl From<FromWireError> for SpiProcessorError {
    fn from(err: FromWireError) -> Self {
        SpiProcessorError::FromWire(err)
    }
}

impl From<ToWireError> for SpiProcessorError {
    fn from(err: ToWireError) -> Self {
        SpiProcessorError::ToWire(err)
    }
}

impl From<TockError> for SpiProcessorError {
    fn from(_err: TockError) -> Self {
        SpiProcessorError::Tock
    }
}

impl From<manticore::server::Error> for SpiProcessorError {
    fn from(err: manticore::server::Error) -> Self {
        SpiProcessorError::Manticore(err)
    }
}

impl From<core::fmt::Error> for SpiProcessorError {
    fn from(err: core::fmt::Error) -> Self {
        SpiProcessorError::Format(err)
    }
}

//////////////////////////////////////////////////////////////////////////////

pub struct SpiProcessor<'a> {
    // The Manticore protocol server.
    pub server: PaRot<'a, Identity, Reset, NoRsa>,

    // Whether to print incoming flash headers.
    pub print_flash_headers: bool,
}

const SPI_TX_BUF_SIZE : usize = 512;

pub type SpiProcessorResult<T> = Result<T, SpiProcessorError>;

impl<'a> SpiProcessor<'a> {

    fn send_data(&mut self, tx_header: &payload::Header, tx_buf: &mut[u8]) -> SpiProcessorResult<()> {
        {
            // Scope for tx_cursor (which doesn't implement Drop).
            // We need tx_cursor to go out of scope so that we can use tx_buf further down.
            let tx_cursor = SpiutilsCursor::new(tx_buf);
            tx_header.to_wire(tx_cursor)?;
        }
        spi_device::get().end_transaction_with_data(tx_buf, true, true)?;

        Ok(())
    }

    fn process_manticore(&mut self, mut data: &[u8]) -> SpiProcessorResult<()> {
        let mut console = Console::new();
        writeln!(console, "Device: Manticore!")?;

        let mut tx_buf : [u8; SPI_TX_BUF_SIZE] = [0xff; SPI_TX_BUF_SIZE];
        let payload_len : u16;
        {
            let mut tx_cursor = ManticoreCursor::new(&mut tx_buf[payload::HEADER_LEN..]);
            self.server.process_request(&mut data, &mut tx_cursor)?;
            payload_len = u16::try_from(tx_cursor.consumed_len())
                .map_err(|_| SpiProcessorError::FromWire(FromWireError::OutOfRange))?;
        }
        let tx_header = payload::Header {
            content: payload::ContentType::Manticore,
            content_len: payload_len,
        };
        self.send_data(&tx_header, &mut tx_buf)?;
        writeln!(console, "Device: Data sent")?;
        Ok(())
    }

    fn process_spi_payload(&mut self, mut data: &[u8]) -> SpiProcessorResult<()> {
        let mut console = Console::new();
        let header = payload::Header::from_wire(&mut data)?;
        writeln!(console, "Device: payload header: {:?}", header)?;
        match header.content {
            payload::ContentType::Manticore => {
                self.process_manticore(&data[..header.content_len as usize])
            }
            _ => {
                Err(SpiProcessorError::UnsupportedContentType(header.content))
            }
        }
    }

    // Send data via the SPI host.
    // The transaction is split into smaller transactions that fit into the SPI host's buffer.
    // The write enable status bit is set before each transaction is executed.
    // The `pre_transaction_fn` is executed prior to each transaction.
    fn spi_host_send<AddrType, F>(&self, header: &flash::Header::<AddrType>, mut data: &[u8], pre_transaction_fn: &F) -> SpiProcessorResult<()>
    where AddrType: Address,
        F: Fn() -> SpiProcessorResult<()>
    {
        // We need to update the header so copy it.
        let mut header = *header;
        loop {
            pre_transaction_fn()?;

            let mut tx_buf = [0xff; spi_host::MAX_READ_BUFFER_LENGTH];
            let tx_len : usize;
            let data_len_to_send : usize;
            {
                let mut tx_cursor = SpiutilsCursor::new(&mut tx_buf);
                header.to_wire(&mut tx_cursor)?;
                if header.opcode.has_dummy_byte() {
                    // Skip one dummy byte (send 0x0)
                    tx_cursor.write_bytes(&[0x0; 1])
                        .map_err(|err| SpiProcessorError::ToWire(ToWireError::Io(err)))?;
                }

                data_len_to_send = min(spi_host::MAX_READ_BUFFER_LENGTH - tx_cursor.consumed_len(), data.len());
                tx_cursor.write_bytes(&data[..data_len_to_send])
                    .map_err(|err| SpiProcessorError::ToWire(ToWireError::Io(err)))?;

                tx_len = tx_cursor.consumed_len()
            }

            spi_host_h1::get().set_wait_busy_clear_in_transactions(header.opcode.wait_busy_clear())?;
            spi_host::get().read_write_bytes(&mut tx_buf, tx_len)?;
            spi_host::get().wait_read_write_done();

            // Move data and address forward
            data = &data[data_len_to_send..];
            if let Some(addr) = header.address {
                let delta : u32 = core::convert::TryFrom::<usize>::try_from(data_len_to_send)
                    .map_err(|_| SpiProcessorError::FromWire(FromWireError::OutOfRange))?;
                let next_addr = addr.into() + delta;
                header.address = Some(AddrType::try_from(next_addr)
                    .map_err(|_| SpiProcessorError::FromWire(FromWireError::OutOfRange))?);
            }

            if data.len() == 0 { break; }
        }
        Ok(())
    }

    // Send a "write enable" command via the SPI host.
    fn spi_host_write_enable(&self) -> SpiProcessorResult<()> {
        let header = flash::Header::<u32> {
            opcode: OpCode::WriteEnable,
            address: None,
        };

        // The command has no data.
        let data : [u8; 0] = [0; 0];
        self.spi_host_send(&header, &data, &|| Ok(()))
    }

    // Send a "write" type command (e.g. PageProgram, *Erase) via the SPI host.
    // This splits the data into smaller transactions as needed and executes
    // "enable write" for each transaction.
    fn spi_host_write<AddrType>(&self, header: &flash::Header::<AddrType>, data: &[u8]) -> SpiProcessorResult<()>
    where AddrType: Address {
        self.spi_host_send(header, data, &|| self.spi_host_write_enable())
    }

    fn clear_device_status(&self, clear_busy: bool, clear_write_enable: bool) -> SpiProcessorResult<()> {
        spi_device::get().end_transaction_with_status(clear_busy, clear_write_enable)?;
        Ok(())
    }

    // Check if the specified address is within the mailbox address space.
    fn is_mailbox_address(&self, addr: u32) -> bool {
        addr >= SPI_MAILBOX_ADDRESS && addr < SPI_MAILBOX_ADDRESS + SPI_MAILBOX_SIZE
    }

    fn process_spi_header<AddrType>(&mut self, header: &flash::Header::<AddrType>, rx_buf: &[u8]) -> SpiProcessorResult<()>
    where AddrType: Address {
        let mut data: &[u8] = rx_buf;
        if header.opcode.has_dummy_byte() {
            // Skip dummy byte
            data = &rx_buf[1..];
        }
        match header.opcode {
            OpCode::PageProgram => {
                match header.get_address() {
                    Some(addr) if self.is_mailbox_address(addr) => {
                        if spi_device::get().is_write_enable_set() {
                            self.process_spi_payload(data)?;
                        }
                        self.clear_device_status(true, true)
                    }
                    Some(addr) if !self.is_mailbox_address(addr) => {
                        if spi_device::get().is_write_enable_set() {
                            // Pass through to SPI host
                            self.spi_host_write(header, data)?;
                        }
                        self.clear_device_status(true, true)
                    }
                    _ => return Err(SpiProcessorError::InvalidAddress(header.get_address())),
                }
            }
            OpCode::SectorErase | OpCode::BlockErase32KB | OpCode::BlockErase64KB => {
                match header.get_address() {
                    Some(addr) if self.is_mailbox_address(addr) => {
                        // Nothing to do.
                        self.clear_device_status(true, true)
                    }
                    Some(addr) if !self.is_mailbox_address(addr) => {
                        if spi_device::get().is_write_enable_set() {
                            // Pass through to SPI host
                            self.spi_host_write(header, data)?;
                        }
                        self.clear_device_status(true, true)
                    }
                    _ => return Err(SpiProcessorError::InvalidAddress(header.get_address())),
                }
            }
            OpCode::ChipErase | OpCode::ChipErase2 => {
                if spi_device::get().is_write_enable_set() {
                    // Pass through to SPI host
                    self.spi_host_write(header, data)?;
                }
                self.clear_device_status(true, true)
            }
            _ => return Err(SpiProcessorError::UnsupportedOpCode(header.opcode)),
        }
    }

    pub fn process_spi_packet(&mut self, mut rx_buf: &[u8]) -> SpiProcessorResult<()> {
        let mut console = Console::new();
        match spi_device::get().get_address_mode() {
            AddressMode::ThreeByte => {
                let header = flash::Header::<ux::u24>::from_wire(&mut rx_buf)?;
                if self.print_flash_headers {
                    writeln!(console, "Device: flash header (3B): {:?}", header)?;
                }
                self.process_spi_header(&header, rx_buf)
            }
            AddressMode::FourByte => {
                let header = flash::Header::<u32>::from_wire(&mut rx_buf)?;
                if self.print_flash_headers {
                    writeln!(console, "Device: flash header (4B): {:?}", header)?;
                }
                self.process_spi_header(&header, rx_buf)
            }
        }
    }
}
<commit_msg>otpilot: Change Mailbox address to 0x80000. (#247)<commit_after>// Copyright 2021 lowRISC 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
//
//     https://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.
//
// SPDX-License-Identifier: Apache-2.0

use crate::manticore_support::Identity;
use crate::manticore_support::NoRsa;
use crate::manticore_support::Reset;
use crate::spi_host;
use crate::spi_host_h1;
use crate::spi_device;

use core::cmp::min;
use core::convert::TryFrom;
use core::fmt::Write;

use libtock::console::Console;
use libtock::result::TockError;

use manticore::io::Cursor as ManticoreCursor;
use manticore::server::pa_rot::PaRot;

use spiutils::io::Cursor as SpiutilsCursor;
use spiutils::io::Write as SpiutilsWrite;
use spiutils::protocol::flash;
use spiutils::protocol::flash::Address;
use spiutils::protocol::flash::AddressMode;
use spiutils::protocol::flash::OpCode;
use spiutils::protocol::payload;
use spiutils::protocol::wire::FromWire;
use spiutils::protocol::wire::FromWireError;
use spiutils::protocol::wire::ToWire;
use spiutils::protocol::wire::ToWireError;

// Size of the SPI flash chip.
// Hard-coded to 64 MiB for now.
// TODO(osenft): Modify this to be read from the actual SPI flash chip at runtime.
pub const SPI_FLASH_SIZE: u32 = 0x4000000;

// The location of the mailbox.
// TODO(osenft): Make this configurable, possibly by reading it from the SPI flash chip.
pub const SPI_MAILBOX_ADDRESS: u32 = 0x80000;

// The size of the mailbox.
const SPI_MAILBOX_SIZE: u32 = spi_device::MAX_READ_BUFFER_SIZE as u32;

#[derive(Copy, Clone, Debug)]
pub enum SpiProcessorError {
    FromWire(FromWireError),
    ToWire(ToWireError),
    Tock,
    Manticore(manticore::server::Error),
    UnsupportedContentType(payload::ContentType),
    UnsupportedOpCode(OpCode),
    InvalidAddress(Option<u32>),
    Format(core::fmt::Error),
}

impl From<FromWireError> for SpiProcessorError {
    fn from(err: FromWireError) -> Self {
        SpiProcessorError::FromWire(err)
    }
}

impl From<ToWireError> for SpiProcessorError {
    fn from(err: ToWireError) -> Self {
        SpiProcessorError::ToWire(err)
    }
}

impl From<TockError> for SpiProcessorError {
    fn from(_err: TockError) -> Self {
        SpiProcessorError::Tock
    }
}

impl From<manticore::server::Error> for SpiProcessorError {
    fn from(err: manticore::server::Error) -> Self {
        SpiProcessorError::Manticore(err)
    }
}

impl From<core::fmt::Error> for SpiProcessorError {
    fn from(err: core::fmt::Error) -> Self {
        SpiProcessorError::Format(err)
    }
}

//////////////////////////////////////////////////////////////////////////////

pub struct SpiProcessor<'a> {
    // The Manticore protocol server.
    pub server: PaRot<'a, Identity, Reset, NoRsa>,

    // Whether to print incoming flash headers.
    pub print_flash_headers: bool,
}

const SPI_TX_BUF_SIZE : usize = 512;

pub type SpiProcessorResult<T> = Result<T, SpiProcessorError>;

impl<'a> SpiProcessor<'a> {

    fn send_data(&mut self, tx_header: &payload::Header, tx_buf: &mut[u8]) -> SpiProcessorResult<()> {
        {
            // Scope for tx_cursor (which doesn't implement Drop).
            // We need tx_cursor to go out of scope so that we can use tx_buf further down.
            let tx_cursor = SpiutilsCursor::new(tx_buf);
            tx_header.to_wire(tx_cursor)?;
        }
        spi_device::get().end_transaction_with_data(tx_buf, true, true)?;

        Ok(())
    }

    fn process_manticore(&mut self, mut data: &[u8]) -> SpiProcessorResult<()> {
        let mut console = Console::new();
        writeln!(console, "Device: Manticore!")?;

        let mut tx_buf : [u8; SPI_TX_BUF_SIZE] = [0xff; SPI_TX_BUF_SIZE];
        let payload_len : u16;
        {
            let mut tx_cursor = ManticoreCursor::new(&mut tx_buf[payload::HEADER_LEN..]);
            self.server.process_request(&mut data, &mut tx_cursor)?;
            payload_len = u16::try_from(tx_cursor.consumed_len())
                .map_err(|_| SpiProcessorError::FromWire(FromWireError::OutOfRange))?;
        }
        let tx_header = payload::Header {
            content: payload::ContentType::Manticore,
            content_len: payload_len,
        };
        self.send_data(&tx_header, &mut tx_buf)?;
        writeln!(console, "Device: Data sent")?;
        Ok(())
    }

    fn process_spi_payload(&mut self, mut data: &[u8]) -> SpiProcessorResult<()> {
        let mut console = Console::new();
        let header = payload::Header::from_wire(&mut data)?;
        writeln!(console, "Device: payload header: {:?}", header)?;
        match header.content {
            payload::ContentType::Manticore => {
                self.process_manticore(&data[..header.content_len as usize])
            }
            _ => {
                Err(SpiProcessorError::UnsupportedContentType(header.content))
            }
        }
    }

    // Send data via the SPI host.
    // The transaction is split into smaller transactions that fit into the SPI host's buffer.
    // The write enable status bit is set before each transaction is executed.
    // The `pre_transaction_fn` is executed prior to each transaction.
    fn spi_host_send<AddrType, F>(&self, header: &flash::Header::<AddrType>, mut data: &[u8], pre_transaction_fn: &F) -> SpiProcessorResult<()>
    where AddrType: Address,
        F: Fn() -> SpiProcessorResult<()>
    {
        // We need to update the header so copy it.
        let mut header = *header;
        loop {
            pre_transaction_fn()?;

            let mut tx_buf = [0xff; spi_host::MAX_READ_BUFFER_LENGTH];
            let tx_len : usize;
            let data_len_to_send : usize;
            {
                let mut tx_cursor = SpiutilsCursor::new(&mut tx_buf);
                header.to_wire(&mut tx_cursor)?;
                if header.opcode.has_dummy_byte() {
                    // Skip one dummy byte (send 0x0)
                    tx_cursor.write_bytes(&[0x0; 1])
                        .map_err(|err| SpiProcessorError::ToWire(ToWireError::Io(err)))?;
                }

                data_len_to_send = min(spi_host::MAX_READ_BUFFER_LENGTH - tx_cursor.consumed_len(), data.len());
                tx_cursor.write_bytes(&data[..data_len_to_send])
                    .map_err(|err| SpiProcessorError::ToWire(ToWireError::Io(err)))?;

                tx_len = tx_cursor.consumed_len()
            }

            spi_host_h1::get().set_wait_busy_clear_in_transactions(header.opcode.wait_busy_clear())?;
            spi_host::get().read_write_bytes(&mut tx_buf, tx_len)?;
            spi_host::get().wait_read_write_done();

            // Move data and address forward
            data = &data[data_len_to_send..];
            if let Some(addr) = header.address {
                let delta : u32 = core::convert::TryFrom::<usize>::try_from(data_len_to_send)
                    .map_err(|_| SpiProcessorError::FromWire(FromWireError::OutOfRange))?;
                let next_addr = addr.into() + delta;
                header.address = Some(AddrType::try_from(next_addr)
                    .map_err(|_| SpiProcessorError::FromWire(FromWireError::OutOfRange))?);
            }

            if data.len() == 0 { break; }
        }
        Ok(())
    }

    // Send a "write enable" command via the SPI host.
    fn spi_host_write_enable(&self) -> SpiProcessorResult<()> {
        let header = flash::Header::<u32> {
            opcode: OpCode::WriteEnable,
            address: None,
        };

        // The command has no data.
        let data : [u8; 0] = [0; 0];
        self.spi_host_send(&header, &data, &|| Ok(()))
    }

    // Send a "write" type command (e.g. PageProgram, *Erase) via the SPI host.
    // This splits the data into smaller transactions as needed and executes
    // "enable write" for each transaction.
    fn spi_host_write<AddrType>(&self, header: &flash::Header::<AddrType>, data: &[u8]) -> SpiProcessorResult<()>
    where AddrType: Address {
        self.spi_host_send(header, data, &|| self.spi_host_write_enable())
    }

    fn clear_device_status(&self, clear_busy: bool, clear_write_enable: bool) -> SpiProcessorResult<()> {
        spi_device::get().end_transaction_with_status(clear_busy, clear_write_enable)?;
        Ok(())
    }

    // Check if the specified address is within the mailbox address space.
    fn is_mailbox_address(&self, addr: u32) -> bool {
        addr >= SPI_MAILBOX_ADDRESS && addr < SPI_MAILBOX_ADDRESS + SPI_MAILBOX_SIZE
    }

    fn process_spi_header<AddrType>(&mut self, header: &flash::Header::<AddrType>, rx_buf: &[u8]) -> SpiProcessorResult<()>
    where AddrType: Address {
        let mut data: &[u8] = rx_buf;
        if header.opcode.has_dummy_byte() {
            // Skip dummy byte
            data = &rx_buf[1..];
        }
        match header.opcode {
            OpCode::PageProgram => {
                match header.get_address() {
                    Some(addr) if self.is_mailbox_address(addr) => {
                        if spi_device::get().is_write_enable_set() {
                            self.process_spi_payload(data)?;
                        }
                        self.clear_device_status(true, true)
                    }
                    Some(addr) if !self.is_mailbox_address(addr) => {
                        if spi_device::get().is_write_enable_set() {
                            // Pass through to SPI host
                            self.spi_host_write(header, data)?;
                        }
                        self.clear_device_status(true, true)
                    }
                    _ => return Err(SpiProcessorError::InvalidAddress(header.get_address())),
                }
            }
            OpCode::SectorErase | OpCode::BlockErase32KB | OpCode::BlockErase64KB => {
                match header.get_address() {
                    Some(addr) if self.is_mailbox_address(addr) => {
                        // Nothing to do.
                        self.clear_device_status(true, true)
                    }
                    Some(addr) if !self.is_mailbox_address(addr) => {
                        if spi_device::get().is_write_enable_set() {
                            // Pass through to SPI host
                            self.spi_host_write(header, data)?;
                        }
                        self.clear_device_status(true, true)
                    }
                    _ => return Err(SpiProcessorError::InvalidAddress(header.get_address())),
                }
            }
            OpCode::ChipErase | OpCode::ChipErase2 => {
                if spi_device::get().is_write_enable_set() {
                    // Pass through to SPI host
                    self.spi_host_write(header, data)?;
                }
                self.clear_device_status(true, true)
            }
            _ => return Err(SpiProcessorError::UnsupportedOpCode(header.opcode)),
        }
    }

    pub fn process_spi_packet(&mut self, mut rx_buf: &[u8]) -> SpiProcessorResult<()> {
        let mut console = Console::new();
        match spi_device::get().get_address_mode() {
            AddressMode::ThreeByte => {
                let header = flash::Header::<ux::u24>::from_wire(&mut rx_buf)?;
                if self.print_flash_headers {
                    writeln!(console, "Device: flash header (3B): {:?}", header)?;
                }
                self.process_spi_header(&header, rx_buf)
            }
            AddressMode::FourByte => {
                let header = flash::Header::<u32>::from_wire(&mut rx_buf)?;
                if self.print_flash_headers {
                    writeln!(console, "Device: flash header (4B): {:?}", header)?;
                }
                self.process_spi_header(&header, rx_buf)
            }
        }
    }
}
<|endoftext|>


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

use ast::InputValue;
use value::Value;
use schema::model::RootNode;
use types::scalars::EmptyMutation;
use executor::{FieldResult, Context};

/*

Syntax to validate:

* Order of items: fields, description, interfaces
* Optional Generics/lifetimes
* Custom name vs. default name
* Optional commas between items
* Nullable/fallible context switching

 */

struct Interface;

struct CustomName;

#[allow(dead_code)]
struct WithLifetime<'a> {
    data: PhantomData<&'a i32>,
}

#[allow(dead_code)]
struct WithGenerics<T> {
    data: T,
}

struct DescriptionFirst;
struct FieldsFirst;
struct InterfacesFirst;

struct CommasWithTrailing;
struct CommasOnMeta;

struct Root;

graphql_object!(CustomName: () as "ACustomNamedType" |&self| {
    field simple() -> i32 { 0 }
});

graphql_object!(<'a> WithLifetime<'a>: () as "WithLifetime" |&self| {
    field simple() -> i32 { 0 }
});

graphql_object!(<T> WithGenerics<T>: () as "WithGenerics" |&self| {
    field simple() -> i32 { 0 }
});

graphql_interface!(Interface: () |&self| {
    field simple() -> i32 { 0 }

    instance_resolvers: |_| {
        DescriptionFirst => Some(DescriptionFirst {}),
    }
});

graphql_object!(DescriptionFirst: () |&self| {
    description: "A description"

    field simple() -> i32 { 0 }

    interfaces: [Interface]
});

graphql_object!(FieldsFirst: () |&self| {
    field simple() -> i32 { 0 }

    description: "A description"

    interfaces: [Interface]
});

graphql_object!(InterfacesFirst: ()|&self| {
    interfaces: [Interface]

    field simple() -> i32 { 0 }

    description: "A description"
});

graphql_object!(CommasWithTrailing: () |&self| {
    interfaces: [Interface],

    field simple() -> i32 { 0 },

    description: "A description",
});

graphql_object!(CommasOnMeta: () |&self| {
    interfaces: [Interface],
    description: "A description",

    field simple() -> i32 { 0 }
});

struct InnerContext;
const INNER_CONTEXT: InnerContext = InnerContext;
impl Context for InnerContext {}

struct InnerType;
graphql_object!(InnerType: InnerContext |&self| {
});

struct CtxSwitcher;
graphql_object!(CtxSwitcher: () |&self| {
    field ctx_switch_always() -> (&InnerContext, InnerType) { (&INNER_CONTEXT, InnerType) }
    field ctx_switch_opt() -> Option<(&InnerContext, InnerType)> { Some((&INNER_CONTEXT, InnerType)) }
    field ctx_switch_res() -> FieldResult<(&InnerContext, InnerType)> { Ok((&INNER_CONTEXT, InnerType)) }
    field ctx_switch_res_opt() -> FieldResult<Option<(&InnerContext, InnerType)>> { Ok(Some((&INNER_CONTEXT, InnerType))) }
});

graphql_object!(<'a> Root: () as "Root" |&self| {
    field custom_name() -> CustomName { CustomName {} }

    field with_lifetime() -> WithLifetime<'a> { WithLifetime { data: PhantomData } }
    field with_generics() -> WithGenerics<i32> { WithGenerics { data: 123 } }

    field description_first() -> DescriptionFirst { DescriptionFirst {} }
    field fields_first() -> FieldsFirst { FieldsFirst {} }
    field interfaces_first() -> InterfacesFirst { InterfacesFirst {} }

    field commas_with_trailing() -> CommasWithTrailing { CommasWithTrailing {} }
    field commas_on_meta() -> CommasOnMeta { CommasOnMeta {} }

    field ctx_switcher() -> CtxSwitcher { CtxSwitcher {} }
});

fn run_type_info_query<F>(type_name: &str, f: F)
where
    F: Fn(&OrderMap<String, Value>, &Vec<Value>) -> (),
{
    let doc = r#"
    query ($typeName: String!) {
        __type(name: $typeName) {
            name
            description
            fields(includeDeprecated: true) {
                name
                type {
                    kind
                    name
                    ofType {
                        kind
                        name
                    }
                }
            }
            interfaces {
                name
                kind
            }
        }
    }
    "#;
    let schema = RootNode::new(Root {}, EmptyMutation::<()>::new());
    let vars = vec![
        ("typeName".to_owned(), InputValue::string(type_name)),
    ].into_iter()
        .collect();

    let (result, errs) = ::execute(doc, None, &schema, &vars, &()).expect("Execution failed");

    assert_eq!(errs, []);

    println!("Result: {:?}", result);

    let type_info = result
        .as_object_value()
        .expect("Result is not an object")
        .get("__type")
        .expect("__type field missing")
        .as_object_value()
        .expect("__type field not an object value");

    let fields = type_info
        .get("fields")
        .expect("fields field missing")
        .as_list_value()
        .expect("fields field not a list value");

    f(type_info, fields);
}

#[test]
fn introspect_custom_name() {
    run_type_info_query("ACustomNamedType", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("ACustomNamedType")));
        assert_eq!(object.get("description"), Some(&Value::null()));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_with_lifetime() {
    run_type_info_query("WithLifetime", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("WithLifetime")));
        assert_eq!(object.get("description"), Some(&Value::null()));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_with_generics() {
    run_type_info_query("WithGenerics", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("WithGenerics")));
        assert_eq!(object.get("description"), Some(&Value::null()));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_description_first() {
    run_type_info_query("DescriptionFirst", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("DescriptionFirst")));
        assert_eq!(object.get("description"), Some(&Value::string("A description")));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![
            Value::object(vec![
                ("name", Value::string("Interface")),
                ("kind", Value::string("INTERFACE")),
            ].into_iter().collect()),
        ])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_fields_first() {
    run_type_info_query("FieldsFirst", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("FieldsFirst")));
        assert_eq!(object.get("description"), Some(&Value::string("A description")));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![
            Value::object(vec![
                ("name", Value::string("Interface")),
                ("kind", Value::string("INTERFACE")),
            ].into_iter().collect()),
        ])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_interfaces_first() {
    run_type_info_query("InterfacesFirst", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("InterfacesFirst")));
        assert_eq!(object.get("description"), Some(&Value::string("A description")));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![
            Value::object(vec![
                ("name", Value::string("Interface")),
                ("kind", Value::string("INTERFACE")),
            ].into_iter().collect()),
        ])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_commas_with_trailing() {
    run_type_info_query("CommasWithTrailing", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("CommasWithTrailing")));
        assert_eq!(object.get("description"), Some(&Value::string("A description")));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![
            Value::object(vec![
                ("name", Value::string("Interface")),
                ("kind", Value::string("INTERFACE")),
            ].into_iter().collect()),
        ])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_commas_on_meta() {
    run_type_info_query("CommasOnMeta", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("CommasOnMeta")));
        assert_eq!(object.get("description"), Some(&Value::string("A description")));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![
            Value::object(vec![
                ("name", Value::string("Interface")),
                ("kind", Value::string("INTERFACE")),
            ].into_iter().collect()),
        ])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_ctx_switch() {
    run_type_info_query("CtxSwitcher", |_, fields| {
        assert!(fields.contains(&graphql_value!({
            "name": "ctxSwitchAlways",
            "type": {
                "kind": "NON_NULL",
                "name": None,
                "ofType": {
                    "kind": "OBJECT",
                    "name": "InnerType",
                }
            }
        })));

        assert!(fields.contains(&graphql_value!({
            "name": "ctxSwitchOpt",
            "type": {
                "kind": "OBJECT",
                "name": "InnerType",
                "ofType": None
            }
        })));

        assert!(fields.contains(&graphql_value!({
            "name": "ctxSwitchRes",
            "type": {
                "kind": "NON_NULL",
                "name": None,
                "ofType": {
                    "kind": "OBJECT",
                    "name": "InnerType",
                }
            }
        })));

        assert!(fields.contains(&graphql_value!({
            "name": "ctxSwitchResOpt",
            "type": {
                "kind": "OBJECT",
                "name": "InnerType",
                "ofType": None
            }
        })));
    });
}
<commit_msg>Don't use global variable in test<commit_after>use ordermap::OrderMap;
use std::marker::PhantomData;

use ast::InputValue;
use value::Value;
use schema::model::RootNode;
use types::scalars::EmptyMutation;
use executor::{FieldResult, Context};

/*

Syntax to validate:

* Order of items: fields, description, interfaces
* Optional Generics/lifetimes
* Custom name vs. default name
* Optional commas between items
* Nullable/fallible context switching

 */

struct Interface;

struct CustomName;

#[allow(dead_code)]
struct WithLifetime<'a> {
    data: PhantomData<&'a i32>,
}

#[allow(dead_code)]
struct WithGenerics<T> {
    data: T,
}

struct DescriptionFirst;
struct FieldsFirst;
struct InterfacesFirst;

struct CommasWithTrailing;
struct CommasOnMeta;

struct Root;

graphql_object!(CustomName: () as "ACustomNamedType" |&self| {
    field simple() -> i32 { 0 }
});

graphql_object!(<'a> WithLifetime<'a>: () as "WithLifetime" |&self| {
    field simple() -> i32 { 0 }
});

graphql_object!(<T> WithGenerics<T>: () as "WithGenerics" |&self| {
    field simple() -> i32 { 0 }
});

graphql_interface!(Interface: () |&self| {
    field simple() -> i32 { 0 }

    instance_resolvers: |_| {
        DescriptionFirst => Some(DescriptionFirst {}),
    }
});

graphql_object!(DescriptionFirst: () |&self| {
    description: "A description"

    field simple() -> i32 { 0 }

    interfaces: [Interface]
});

graphql_object!(FieldsFirst: () |&self| {
    field simple() -> i32 { 0 }

    description: "A description"

    interfaces: [Interface]
});

graphql_object!(InterfacesFirst: ()|&self| {
    interfaces: [Interface]

    field simple() -> i32 { 0 }

    description: "A description"
});

graphql_object!(CommasWithTrailing: () |&self| {
    interfaces: [Interface],

    field simple() -> i32 { 0 },

    description: "A description",
});

graphql_object!(CommasOnMeta: () |&self| {
    interfaces: [Interface],
    description: "A description",

    field simple() -> i32 { 0 }
});

struct InnerContext;
impl Context for InnerContext {}

struct InnerType;
graphql_object!(InnerType: InnerContext |&self| {
});

struct CtxSwitcher;
graphql_object!(CtxSwitcher: InnerContext |&self| {
    field ctx_switch_always(&executor) -> (&InnerContext, InnerType) { (executor.context(), InnerType) }
    field ctx_switch_opt(&executor) -> Option<(&InnerContext, InnerType)> { Some((executor.context(), InnerType)) }
    field ctx_switch_res(&executor) -> FieldResult<(&InnerContext, InnerType)> { Ok((executor.context(), InnerType)) }
    field ctx_switch_res_opt(&executor) -> FieldResult<Option<(&InnerContext, InnerType)>> { Ok(Some((executor.context(), InnerType))) }
});

graphql_object!(<'a> Root: InnerContext as "Root" |&self| {
    field custom_name() -> CustomName { CustomName {} }

    field with_lifetime() -> WithLifetime<'a> { WithLifetime { data: PhantomData } }
    field with_generics() -> WithGenerics<i32> { WithGenerics { data: 123 } }

    field description_first() -> DescriptionFirst { DescriptionFirst {} }
    field fields_first() -> FieldsFirst { FieldsFirst {} }
    field interfaces_first() -> InterfacesFirst { InterfacesFirst {} }

    field commas_with_trailing() -> CommasWithTrailing { CommasWithTrailing {} }
    field commas_on_meta() -> CommasOnMeta { CommasOnMeta {} }

    field ctx_switcher() -> CtxSwitcher { CtxSwitcher {} }
});

fn run_type_info_query<F>(type_name: &str, f: F)
where
    F: Fn(&OrderMap<String, Value>, &Vec<Value>) -> (),
{
    let doc = r#"
    query ($typeName: String!) {
        __type(name: $typeName) {
            name
            description
            fields(includeDeprecated: true) {
                name
                type {
                    kind
                    name
                    ofType {
                        kind
                        name
                    }
                }
            }
            interfaces {
                name
                kind
            }
        }
    }
    "#;
    let schema = RootNode::new(Root {}, EmptyMutation::<InnerContext>::new());
    let vars = vec![
        ("typeName".to_owned(), InputValue::string(type_name)),
    ].into_iter()
        .collect();

    let (result, errs) = ::execute(doc, None, &schema, &vars, &InnerContext).expect("Execution failed");

    assert_eq!(errs, []);

    println!("Result: {:?}", result);

    let type_info = result
        .as_object_value()
        .expect("Result is not an object")
        .get("__type")
        .expect("__type field missing")
        .as_object_value()
        .expect("__type field not an object value");

    let fields = type_info
        .get("fields")
        .expect("fields field missing")
        .as_list_value()
        .expect("fields field not a list value");

    f(type_info, fields);
}

#[test]
fn introspect_custom_name() {
    run_type_info_query("ACustomNamedType", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("ACustomNamedType")));
        assert_eq!(object.get("description"), Some(&Value::null()));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_with_lifetime() {
    run_type_info_query("WithLifetime", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("WithLifetime")));
        assert_eq!(object.get("description"), Some(&Value::null()));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_with_generics() {
    run_type_info_query("WithGenerics", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("WithGenerics")));
        assert_eq!(object.get("description"), Some(&Value::null()));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_description_first() {
    run_type_info_query("DescriptionFirst", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("DescriptionFirst")));
        assert_eq!(object.get("description"), Some(&Value::string("A description")));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![
            Value::object(vec![
                ("name", Value::string("Interface")),
                ("kind", Value::string("INTERFACE")),
            ].into_iter().collect()),
        ])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_fields_first() {
    run_type_info_query("FieldsFirst", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("FieldsFirst")));
        assert_eq!(object.get("description"), Some(&Value::string("A description")));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![
            Value::object(vec![
                ("name", Value::string("Interface")),
                ("kind", Value::string("INTERFACE")),
            ].into_iter().collect()),
        ])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_interfaces_first() {
    run_type_info_query("InterfacesFirst", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("InterfacesFirst")));
        assert_eq!(object.get("description"), Some(&Value::string("A description")));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![
            Value::object(vec![
                ("name", Value::string("Interface")),
                ("kind", Value::string("INTERFACE")),
            ].into_iter().collect()),
        ])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_commas_with_trailing() {
    run_type_info_query("CommasWithTrailing", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("CommasWithTrailing")));
        assert_eq!(object.get("description"), Some(&Value::string("A description")));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![
            Value::object(vec![
                ("name", Value::string("Interface")),
                ("kind", Value::string("INTERFACE")),
            ].into_iter().collect()),
        ])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_commas_on_meta() {
    run_type_info_query("CommasOnMeta", |object, fields| {
        assert_eq!(object.get("name"), Some(&Value::string("CommasOnMeta")));
        assert_eq!(object.get("description"), Some(&Value::string("A description")));
        assert_eq!(object.get("interfaces"), Some(&Value::list(vec![
            Value::object(vec![
                ("name", Value::string("Interface")),
                ("kind", Value::string("INTERFACE")),
            ].into_iter().collect()),
        ])));

        assert!(fields.contains(&graphql_value!({
            "name": "simple",
            "type": { "kind": "NON_NULL", "name": None, "ofType": { "kind": "SCALAR", "name": "Int" } }
        })));
    });
}

#[test]
fn introspect_ctx_switch() {
    run_type_info_query("CtxSwitcher", |_, fields| {
        assert!(fields.contains(&graphql_value!({
            "name": "ctxSwitchAlways",
            "type": {
                "kind": "NON_NULL",
                "name": None,
                "ofType": {
                    "kind": "OBJECT",
                    "name": "InnerType",
                }
            }
        })));

        assert!(fields.contains(&graphql_value!({
            "name": "ctxSwitchOpt",
            "type": {
                "kind": "OBJECT",
                "name": "InnerType",
                "ofType": None
            }
        })));

        assert!(fields.contains(&graphql_value!({
            "name": "ctxSwitchRes",
            "type": {
                "kind": "NON_NULL",
                "name": None,
                "ofType": {
                    "kind": "OBJECT",
                    "name": "InnerType",
                }
            }
        })));

        assert!(fields.contains(&graphql_value!({
            "name": "ctxSwitchResOpt",
            "type": {
                "kind": "OBJECT",
                "name": "InnerType",
                "ofType": None
            }
        })));
    });
}
<|endoftext|>