<commit_before>//! This crate helps to create a Fractalide component
//!
//! It provides the macro `component` which takes the high level view of the component, and creates the code for the scheduler.
//!
//! It also declare the Trait Agent and the shared methods needed by every component and the scheduler.

extern crate capnp;

// TODO : Add method to remove components
use ports::{IPSender, IPReceiver};
use scheduler::Signal;
use result::Result;

/// Provide the generic functions of agents
///
/// These three functions are used by the scheduler
pub trait Agent {
    /// Return true if there is at least one input port
    fn is_input_ports(&self) -> bool;
    /// Connect output port
    fn connect(&mut self, port: &str, sender: IPSender) -> Result<()>;
    /// Connect array output port
    fn connect_array(&mut self, port: &str, element: String, sender: IPSender) -> Result<()>;
    /// Add input element
    fn add_inarr_element(&mut self, port: &str, element: String, recv: IPReceiver) -> Result<()>;
    /// Run the method of the component, his personal logic
    fn run(&mut self) -> Result<Signal>;
}

/// The agent macro.
///
/// It helps to define a agent, by defining the input and output ports, if there is an option or an acc port, ...
///
/// Example :
///
/// ```rust,ignore
/// agent! {
///    inputs(input: any),
///    outputs(output: any),
///    option(generic_text),
///    fn run(&mut self) -> Result<Signal> {
///        // Receive an IP
///        let ip = try!(self.input.input.recv());
///
///        // Received an IP from the option port (a generic_text)
///        let opt = self.recv_opt();
///
///        // Get the capn'p reader
///        let reader: generic_text::Reader = try!(opt.read_schema());
///        // Print the option
///        println!("{}", try!(reader.get_text()));
///
///        // Send the received IP outside, but don't care about the success (drop on fail)
///        let _ = self.output.output.send(ip);
///
///        Ok(End)
///    }
/// }
/// ```
#[macro_export]
macro_rules! agent {
    (
        $( input($( $input_name:ident: $input_contract:ident ),*), )*
        $( inarr($( $input_a_name:ident: $input_a_contract:ident ),*), )*
        $( output($( $output_name:ident: $output_contract:ident ),*), )*
        $( outarr($( $output_a_name:ident: $output_a_contract:ident ),*), )*
        $( portal( $portal_type:ty => $portal_value:expr ), )*
        $( option($option:ident), )*
        $( accumulator($acc:ident ), )*
        fn run(&mut $arg:ident) -> Result<Signal> $fun:block
    )
        =>
    {
        use rustfbp::agent::Agent;

        use rustfbp::result;
        use rustfbp::result::Result;
        use rustfbp::scheduler::{CompMsg, Signal};
        use std::error::Error;

        use std::sync::mpsc::{Sender};
        use std::sync::mpsc::channel;

        use rustfbp::ports::{IP, IPSender, IPReceiver, OutputSend};

        #[allow(unused_imports)]
        use std::collections::HashMap;

        use capnp::serialize;
        use capnp::message;

        use std::io::{Read, Write};

        use rustfbp::scheduler::Signal::*;

        mod edge_capnp {
                include!("edge_capnp.rs");
        }

        use edge_capnp::*;

        impl ThisAgent {
            $(
            fn dummy() {
                if stringify!($option) != "" {}
            }
            pub fn recv_option(&mut self) -> IP {
                self.try_recv_option();
                if self.option_ip.is_none() {
                    self.option_ip = self.input.option.recv().ok();
                }
                match self.option_ip {
                    Some(ref ip) => ip.clone(),
                    None => unreachable!(),
                }
            }

            pub fn try_recv_option(&mut self) -> Option<IP> {
                loop {
                    match self.input.option.try_recv() {
                        Err(_) => { break; },
                        Ok(ip) => { self.option_ip = Some(ip); }
                    };
                }
                self.option_ip.as_ref().map(|ip|{ ip.clone() })
            }
        }

        // simple and array

            fn get_ports(&mut self) -> &mut Ports {
                &mut self.ports
            }
        impl Agent for ThisAgent {

            fn is_input_ports(&self) -> bool {
                $(
                    if true || stringify!($input_field_name) == "" { return true; }
                )*
                $(
                    if true || stringify!($input_array_name) == "" { return true; }
                )*
                false
            }

            fn run(&mut $arg) -> Result<()> $fun

        }

        /* Global agent */

        #[allow(dead_code)]
        #[allow(non_camel_case_types)]
        pub struct ThisAgent {
            name: String,
            pub ports: Ports,
            pub option_ip: Option<IP>,
            $(
            pub portal: $portal_type ,
            )*
        }

        #[allow(dead_code)]
        pub fn new(name: String, sched: Sender<CompMsg>) -> Result<(Box<Agent + Send>, HashMap<String, IPSender>)> {
            let (ports, senders) = try!(Ports::new(name.clone(), sched,
                                   vec!["option".into(), "acc".into(), $( stringify!($input_field_name).to_string() ),*],
                                   vec![$( stringify!($input_array_name).to_string() ),*],
                                   vec!["acc".into(), $( stringify!($output_field_name).to_string() ),*],
                                   vec![$( stringify!($output_array_name).to_string() ),*],));

            // Put it together
            let comp = $name{
                name: name,
                ports: ports,
                option_ip: None,
                $(
                    portal: $portal_value,
                )*
            };
            Ok((Box::new(comp) as Box<Agent + Send>, senders))
        }

        #[no_mangle]
        pub extern fn create_agent(name: String, sched: Sender<CompMsg>) -> Result<(Box<Agent + Send>, HashMap<String, IPSender>)> {
            new(name, sched)
        }

        #[no_mangle]
        pub extern fn get_schema_input(port: &str) -> Result<String> {
            match port {
                $(
                    stringify!($input_field_name)=> Ok(stringify!($input_edge_name).into()),
                )*
                _ => { Err(result::Error::PortNotFound("unknown".into(), port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_input_array(port: &str) -> Result<String> {
            match port {
                $(
                    stringify!($input_array_name) => Ok(stringify!($input_edge_array).into()),
                )*
                _ => { Err(result::Error::PortNotFound("unknown".into(), port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_output(port: &str) -> Result<String> {
            match port {
                $(
                    stringify!($output_field_name)=> Ok(stringify!($output_edge_name).into()),
                )*
                _ => { Err(result::Error::PortNotFound("unknown".into(), port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_output_array(port: &str) -> Result<String> {
            match port {
                $(
                    stringify!($output_array_name) => Ok(stringify!($output_edge_array).into()),
                )*
                _ => { Err(result::Error::PortNotFound("unknown".into(), port.into())) }
            }
        }
    }
}
<commit_msg>Agent: new stucts for ports<commit_after>//! This crate helps to create a Fractalide component
//!
//! It provides the macro `component` which takes the high level view of the component, and creates the code for the scheduler.
//!
//! It also declare the Trait Agent and the shared methods needed by every component and the scheduler.

extern crate capnp;

// TODO : Add method to remove components
use ports::{IPSender, IPReceiver};
use scheduler::Signal;
use result::Result;

/// Provide the generic functions of agents
///
/// These three functions are used by the scheduler
pub trait Agent {
    /// Return true if there is at least one input port
    fn is_input_ports(&self) -> bool;
    /// Connect output port
    fn connect(&mut self, port: &str, sender: IPSender) -> Result<()>;
    /// Connect array output port
    fn connect_array(&mut self, port: &str, element: String, sender: IPSender) -> Result<()>;
    /// Add input element
    fn add_inarr_element(&mut self, port: &str, element: String, recv: IPReceiver) -> Result<()>;
    /// Run the method of the component, his personal logic
    fn run(&mut self) -> Result<Signal>;
}

/// The agent macro.
///
/// It helps to define a agent, by defining the input and output ports, if there is an option or an acc port, ...
///
/// Example :
///
/// ```rust,ignore
/// agent! {
///    inputs(input: any),
///    outputs(output: any),
///    option(generic_text),
///    fn run(&mut self) -> Result<Signal> {
///        // Receive an IP
///        let ip = try!(self.input.input.recv());
///
///        // Received an IP from the option port (a generic_text)
///        let opt = self.recv_opt();
///
///        // Get the capn'p reader
///        let reader: generic_text::Reader = try!(opt.read_schema());
///        // Print the option
///        println!("{}", try!(reader.get_text()));
///
///        // Send the received IP outside, but don't care about the success (drop on fail)
///        let _ = self.output.output.send(ip);
///
///        Ok(End)
///    }
/// }
/// ```
#[macro_export]
macro_rules! agent {
    (
        $( input($( $input_name:ident: $input_contract:ident ),*), )*
        $( inarr($( $input_a_name:ident: $input_a_contract:ident ),*), )*
        $( output($( $output_name:ident: $output_contract:ident ),*), )*
        $( outarr($( $output_a_name:ident: $output_a_contract:ident ),*), )*
        $( portal( $portal_type:ty => $portal_value:expr ), )*
        $( option($option:ident), )*
        $( accumulator($acc:ident ), )*
        fn run(&mut $arg:ident) -> Result<Signal> $fun:block
    )
        =>
    {
        use rustfbp::agent::Agent;

        use rustfbp::result;
        use rustfbp::result::Result;
        use rustfbp::scheduler::{CompMsg, Signal};
        use std::error::Error;

        use std::sync::mpsc::{Sender};
        use std::sync::mpsc::channel;

        use rustfbp::ports::{IP, IPSender, IPReceiver, OutputSend};

        #[allow(unused_imports)]
        use std::collections::HashMap;

        use capnp::serialize;
        use capnp::message;

        use std::io::{Read, Write};

        use rustfbp::scheduler::Signal::*;

        mod edge_capnp {
                include!("edge_capnp.rs");
        }

        use edge_capnp::*;

        impl ThisAgent {
            $(
            fn dummy() {
                if stringify!($option) != "" {}
            }
            pub fn recv_option(&mut self) -> IP {
                self.try_recv_option();
                if self.option_ip.is_none() {
                    self.option_ip = self.input.option.recv().ok();
                }
                match self.option_ip {
                    Some(ref ip) => ip.clone(),
                    None => unreachable!(),
                }
            }

            pub fn try_recv_option(&mut self) -> Option<IP> {
                loop {
                    match self.input.option.try_recv() {
                        Err(_) => { break; },
                        Ok(ip) => { self.option_ip = Some(ip); }
                    };
                }
                self.option_ip.as_ref().map(|ip|{ ip.clone() })
            }
        }

        // simple and array

            fn get_ports(&mut self) -> &mut Ports {
                &mut self.ports
            }
        impl Agent for ThisAgent {

            fn is_input_ports(&self) -> bool {
                $(
                    if true || stringify!($input_field_name) == "" { return true; }
                )*
                $(
                    if true || stringify!($input_array_name) == "" { return true; }
                )*
                false
            }

            fn run(&mut $arg) -> Result<()> $fun
        pub struct Input {
            option: IPReceiver,
            acc: IPReceiver,
            $($(
                $input_name: IPReceiver,
            )*)*
        }

        pub struct Inarr {
            $($(
                $input_a_name: HashMap<String, IPReceiver>,
            )*)*
        }

        pub struct Output {
            acc: Option<IPSender>,
            $($(
                $output_name: Option<IPSender>,
            )*)*
        }

        pub struct Outarr {
            $($(
                $output_a_name: HashMap<String, IPSender>
            )*)*
        }

        /* Global component */

        #[allow(dead_code)]
        #[allow(non_camel_case_types)]
        pub struct ThisAgent {
            name: String,
            pub input: Input,
            pub inarr: Inarr,
            pub output: Output,
            pub outarr: Outarr,
            pub option_ip: Option<IP>,
            sched: Sender<CompMsg>,
            $(
            pub portal: $portal_type ,
            )*
        }

        #[allow(dead_code)]
        pub fn new(name: String, sched: Sender<CompMsg>) -> Result<(Box<Agent + Send>, HashMap<String, IPSender>)> {
            let (ports, senders) = try!(Ports::new(name.clone(), sched,
                                   vec!["option".into(), "acc".into(), $( stringify!($input_field_name).to_string() ),*],
                                   vec![$( stringify!($input_array_name).to_string() ),*],
                                   vec!["acc".into(), $( stringify!($output_field_name).to_string() ),*],
                                   vec![$( stringify!($output_array_name).to_string() ),*],));

            // Put it together
            let comp = $name{
                name: name,
                ports: ports,
                option_ip: None,
                $(
                    portal: $portal_value,
                )*
            };
            Ok((Box::new(comp) as Box<Agent + Send>, senders))
        }

        #[no_mangle]
        pub extern fn create_agent(name: String, sched: Sender<CompMsg>) -> Result<(Box<Agent + Send>, HashMap<String, IPSender>)> {
            new(name, sched)
        }

        #[no_mangle]
        pub extern fn get_schema_input(port: &str) -> Result<String> {
            match port {
                $(
                    stringify!($input_field_name)=> Ok(stringify!($input_edge_name).into()),
                )*
                _ => { Err(result::Error::PortNotFound("unknown".into(), port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_input_array(port: &str) -> Result<String> {
            match port {
                $(
                    stringify!($input_array_name) => Ok(stringify!($input_edge_array).into()),
                )*
                _ => { Err(result::Error::PortNotFound("unknown".into(), port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_output(port: &str) -> Result<String> {
            match port {
                $(
                    stringify!($output_field_name)=> Ok(stringify!($output_edge_name).into()),
                )*
                _ => { Err(result::Error::PortNotFound("unknown".into(), port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_output_array(port: &str) -> Result<String> {
            match port {
                $(
                    stringify!($output_array_name) => Ok(stringify!($output_edge_array).into()),
                )*
                _ => { Err(result::Error::PortNotFound("unknown".into(), port.into())) }
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::str::Chars;
use std::iter::Peekable;
use std::iter::Enumerate;

use ::syntax::lexer::{Token, Lexer};
use ::syntax::next_is;

pub struct PythonSyntax;

impl Lexer for PythonSyntax {
    fn handle_char(&self, ch: char, mut iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize, idx: usize) -> Option<Token> {
        match ch {
            '#' => {
                let st = idx;
                let mut end = idx;
                let mut s = String::from("#");
                while let Some(&(e, c)) = iter.peek() {
                    if c == '\n' { break }
                    end = e;
                    s.push(iter.next().unwrap().1)
                }
                return Some(Token::SingleLineComment(s));
            }

            '@' => {
                let st = idx;
                let mut end = idx;
                let mut s = String::from("@");
                while let Some(&(e, c)) = iter.peek() {
                    if c == '\n' { break }
                    end = e;
                    s.push(iter.next().unwrap().1)
                }
                return Some(Token::Attribute(s));
            }

            ch if ch == '"' || ch == '\'' => {
                let st = idx;
                let mut end = idx;
                let mut s= format!("{}", ch);
                while let Some(&(e, c)) = iter.peek() {
                    end = e;
                    s.push(iter.next().unwrap().1);
                    if c == ch {
                        break;
                    }
                }
                return Some(Token::String(s));
            }

            _ => None,
        }

    }

    fn handle_ident(&self, ch: char, iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize) -> Token {
        let mut ident = String::new();
        ident.push(ch);
        let start = iter.peek().unwrap().0 - 1;

        while self.is_ident(iter.peek()) {
            ident.push(iter.next().unwrap().1)
        }

        Token::Ident(ident)
    }
}

<commit_msg>Highlight python function calls / defs<commit_after>use std::str::Chars;
use std::iter::Peekable;
use std::iter::Enumerate;

use ::syntax::lexer::{Token, Lexer};
use ::syntax::next_is;

pub struct PythonSyntax;

impl Lexer for PythonSyntax {
    fn handle_char(&self, ch: char, mut iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize, idx: usize) -> Option<Token> {
        match ch {
            '#' => {
                let st = idx;
                let mut end = idx;
                let mut s = String::from("#");
                while let Some(&(e, c)) = iter.peek() {
                    if c == '\n' { break }
                    end = e;
                    s.push(iter.next().unwrap().1)
                }
                return Some(Token::SingleLineComment(s));
            }

            '@' => {
                let st = idx;
                let mut end = idx;
                let mut s = String::from("@");
                while let Some(&(e, c)) = iter.peek() {
                    if c == '\n' { break }
                    end = e;
                    s.push(iter.next().unwrap().1)
                }
                return Some(Token::Attribute(s));
            }

            ch if ch == '"' || ch == '\'' => {
                let st = idx;
                let mut end = idx;
                let mut s= format!("{}", ch);
                while let Some(&(e, c)) = iter.peek() {
                    end = e;
                    s.push(iter.next().unwrap().1);
                    if c == ch {
                        break;
                    }
                }
                return Some(Token::String(s));
            }

            _ => None,
        }

    }

    fn handle_ident(&self, ch: char, mut iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize) -> Token {
        let mut ident = String::new();
        ident.push(ch);
        let start = iter.peek().unwrap().0 - 1;

        while self.is_ident(iter.peek()) {
            ident.push(iter.next().unwrap().1)
        }

        let token;
        if next_is(&mut iter, '(') {
            // function calls or definitions
            token = Token::FunctionCallDef(ident);
        } else {
            // regular idents
            token = Token::Ident(ident);
        }

        token
    }
}

<|endoftext|>


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

//! This file defines transaction store APIs that are related to committed signed transactions.

use crate::{
    change_set::ChangeSet,
    errors::DiemDbError,
    schema::{
        transaction::TransactionSchema, transaction_by_account::TransactionByAccountSchema,
        transaction_by_hash::TransactionByHashSchema, write_set::WriteSetSchema,
    },
};
use anyhow::{ensure, format_err, Result};
use diem_crypto::{hash::CryptoHash, HashValue};
use diem_types::{
    account_address::AccountAddress,
    block_metadata::BlockMetadata,
    transaction::{Transaction, Version},
    write_set::WriteSet,
};
use schemadb::{ReadOptions, SchemaIterator, DB};
use std::sync::Arc;

#[derive(Debug)]
pub(crate) struct TransactionStore {
    db: Arc<DB>,
}

impl TransactionStore {
    pub fn new(db: Arc<DB>) -> Self {
        Self { db }
    }

    /// Gets the version of a transaction by the sender `address` and `sequence_number`.
    pub fn get_account_transaction_version(
        &self,
        address: AccountAddress,
        sequence_number: u64,
        ledger_version: Version,
    ) -> Result<Option<Version>> {
        if let Some(version) = self
            .db
            .get::<TransactionByAccountSchema>(&(address, sequence_number))?
        {
            if version <= ledger_version {
                return Ok(Some(version));
            }
        }

        Ok(None)
    }

    /// Gets the version of a transaction by its hash.
    pub fn get_transaction_version_by_hash(
        &self,
        hash: &HashValue,
        ledger_version: Version,
    ) -> Result<Option<Version>> {
        Ok(match self.db.get::<TransactionByHashSchema>(hash)? {
            Some(version) if version <= ledger_version => Some(version),
            _ => None,
        })
    }

    /// Gets an iterator that yields `(sequence_number, version)` for each
    /// transaction sent by an account, starting at `start_seq_num`, and returning
    /// at most `num_versions` results with `version <= ledger_version`.
    ///
    /// Guarantees that the returned sequence numbers are sequential, i.e.,
    /// `seq_num_{i} + 1 = seq_num_{i+1}`.
    pub fn get_account_transaction_version_iter(
        &self,
        address: AccountAddress,
        start_seq_num: u64,
        num_versions: u64,
        ledger_version: Version,
    ) -> Result<AccountTransactionVersionIter> {
        let mut iter = self
            .db
            .iter::<TransactionByAccountSchema>(ReadOptions::default())?;
        iter.seek(&(address, start_seq_num))?;
        Ok(AccountTransactionVersionIter {
            inner: iter,
            address,
            expected_next_seq_num: start_seq_num,
            end_seq_num: start_seq_num
                .checked_add(num_versions)
                .ok_or_else(|| format_err!("too many transactions requested"))?,
            prev_version: None,
            ledger_version,
        })
    }

    /// Get signed transaction given `version`
    pub fn get_transaction(&self, version: Version) -> Result<Transaction> {
        self.db
            .get::<TransactionSchema>(&version)?
            .ok_or_else(|| DiemDbError::NotFound(format!("Txn {}", version)).into())
    }

    /// Gets an iterator that yields `num_transactions` transactions starting from `start_version`.
    pub fn get_transaction_iter(
        &self,
        start_version: Version,
        num_transactions: usize,
    ) -> Result<TransactionIter> {
        let mut iter = self.db.iter::<TransactionSchema>(ReadOptions::default())?;
        iter.seek(&start_version)?;
        Ok(TransactionIter {
            inner: iter,
            expected_next_version: start_version,
            end_version: start_version
                .checked_add(num_transactions as u64)
                .ok_or_else(|| format_err!("too many transactions requested"))?,
        })
    }

    /// Get the first version that txn starts existent.
    pub fn get_first_txn_version(&self) -> Result<Option<Version>> {
        let mut iter = self.db.iter::<TransactionSchema>(Default::default())?;
        iter.seek_to_first();
        iter.next().map(|res| res.map(|(v, _)| v)).transpose()
    }

    /// Returns the block metadata carried on the block metadata transaction at or preceding
    /// `version`, together with the version of the block metadata transaction.
    /// Returns None if there's no such transaction at or preceding `version` (it's likely the genesis
    /// version 0).
    pub fn get_block_metadata(&self, version: Version) -> Result<Option<(Version, BlockMetadata)>> {
        // Maximum TPS from benchmark is around 1000.
        const MAX_VERSIONS_TO_SEARCH: usize = 1000 * 3;

        // Linear search via `DB::rev_iter()` here, NOT expecting performance hit, due to the fact
        // that the iterator caches data block and that there are limited number of transactions in
        // each block.
        let mut iter = self.db.rev_iter::<TransactionSchema>(Default::default())?;
        iter.seek(&version)?;
        for res in iter.take(MAX_VERSIONS_TO_SEARCH) {
            let (v, txn) = res?;
            if let Transaction::BlockMetadata(block_meta) = txn {
                return Ok(Some((v, block_meta)));
            } else if v == 0 {
                return Ok(None);
            }
        }

        Err(DiemDbError::NotFound(format!("BlockMetadata preceding version {}", version)).into())
    }

    /// Save signed transaction at `version`
    pub fn put_transaction(
        &self,
        version: Version,
        transaction: &Transaction,
        cs: &mut ChangeSet,
    ) -> Result<()> {
        if let Transaction::UserTransaction(txn) = transaction {
            cs.batch.put::<TransactionByAccountSchema>(
                &(txn.sender(), txn.sequence_number()),
                &version,
            )?;
        }
        cs.batch
            .put::<TransactionByHashSchema>(&transaction.hash(), &version)?;
        cs.batch.put::<TransactionSchema>(&version, transaction)?;

        Ok(())
    }

    /// Get executed transaction vm output given `version`
    pub fn get_write_set(&self, version: Version) -> Result<WriteSet> {
        self.db
            .get::<WriteSetSchema>(&version)?
            .ok_or_else(|| DiemDbError::NotFound(format!("WriteSet at version {}", version)).into())
    }

    /// Get the first version that write set starts existent.
    pub fn get_first_write_set_version(&self) -> Result<Option<Version>> {
        let mut iter = self.db.iter::<WriteSetSchema>(Default::default())?;
        iter.seek_to_first();
        iter.next().map(|res| res.map(|(v, _)| v)).transpose()
    }

    /// Save executed transaction vm output given `version`
    pub fn put_write_set(
        &self,
        version: Version,
        write_set: &WriteSet,
        cs: &mut ChangeSet,
    ) -> Result<()> {
        cs.batch.put::<WriteSetSchema>(&version, write_set)
    }
}

pub struct TransactionIter<'a> {
    inner: SchemaIterator<'a, TransactionSchema>,
    expected_next_version: Version,
    end_version: Version,
}

impl<'a> TransactionIter<'a> {
    fn next_impl(&mut self) -> Result<Option<Transaction>> {
        if self.expected_next_version >= self.end_version {
            return Ok(None);
        }

        let ret = match self.inner.next().transpose()? {
            Some((version, transaction)) => {
                ensure!(
                    version == self.expected_next_version,
                    "Transaction versions are not consecutive.",
                );
                self.expected_next_version += 1;
                Some(transaction)
            }
            None => None,
        };

        Ok(ret)
    }
}

impl<'a> Iterator for TransactionIter<'a> {
    type Item = Result<Transaction>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_impl().transpose()
    }
}

// TODO(philiphayes): this will need to change to support CRSNs
// (Conflict-Resistant Sequence Numbers)[https://github.com/diem/dip/blob/main/dips/dip-168.md].
//
// It depends on the implementation details, but we'll probably index by _requested_
// transaction sequence number rather than committed account sequence number.
// This would mean the property: `seq_num_{i+1} == seq_num_{i} + 1` would no longer
// be guaranteed and the check should be removed.
//
// This index would also no longer iterate over an account's transactions in
// committed order, meaning the outer method would need to overread by
// `CRSN_WINDOW_SIZE`, sort by version, and take only `limit` entries to get
// at most `limit` transactions in committed order. Alternatively, add another
// index for scanning an accounts transactions in committed order, e.g.,
// `(AccountAddress, Version) -> SeqNum`.

pub struct AccountTransactionVersionIter<'a> {
    inner: SchemaIterator<'a, TransactionByAccountSchema>,
    address: AccountAddress,
    expected_next_seq_num: u64,
    end_seq_num: u64,
    prev_version: Option<Version>,
    ledger_version: Version,
}

impl<'a> AccountTransactionVersionIter<'a> {
    fn next_impl(&mut self) -> Result<Option<(u64, Version)>> {
        if self.expected_next_seq_num >= self.end_seq_num {
            return Ok(None);
        }

        Ok(match self.inner.next().transpose()? {
            Some(((address, seq_num), version)) => {
                // No more transactions sent by this account.
                if address != self.address {
                    return Ok(None);
                }

                // Ensure seq_num_{i+1} == seq_num_{i} + 1
                ensure!(
                    seq_num == self.expected_next_seq_num,
                    "DB corruption: account transactions sequence numbers are not contiguous: \
                     actual: {}, expected: {}",
                    seq_num,
                    self.expected_next_seq_num,
                );

                // Ensure version_{i+1} > version_{i}
                if let Some(prev_version) = self.prev_version {
                    ensure!(
                        prev_version < version,
                        "DB corruption: account transaction versions are not strictly increasing: \
                         previous version: {}, current version: {}",
                        prev_version,
                        version,
                    );
                }

                // No more transactions (in this view of the ledger).
                if version > self.ledger_version {
                    return Ok(None);
                }

                self.expected_next_seq_num += 1;
                self.prev_version = Some(version);
                Some((seq_num, version))
            }
            None => None,
        })
    }
}

impl<'a> Iterator for AccountTransactionVersionIter<'a> {
    type Item = Result<(u64, Version)>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_impl().transpose()
    }
}

#[cfg(test)]
mod test;
<commit_msg>[diemdb] increase the MAX_VERSIONS_TO_SEARCH<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! This file defines transaction store APIs that are related to committed signed transactions.

use crate::{
    change_set::ChangeSet,
    errors::DiemDbError,
    schema::{
        transaction::TransactionSchema, transaction_by_account::TransactionByAccountSchema,
        transaction_by_hash::TransactionByHashSchema, write_set::WriteSetSchema,
    },
};
use anyhow::{ensure, format_err, Result};
use diem_crypto::{hash::CryptoHash, HashValue};
use diem_types::{
    account_address::AccountAddress,
    block_metadata::BlockMetadata,
    transaction::{Transaction, Version},
    write_set::WriteSet,
};
use schemadb::{ReadOptions, SchemaIterator, DB};
use std::sync::Arc;

#[derive(Debug)]
pub(crate) struct TransactionStore {
    db: Arc<DB>,
}

impl TransactionStore {
    pub fn new(db: Arc<DB>) -> Self {
        Self { db }
    }

    /// Gets the version of a transaction by the sender `address` and `sequence_number`.
    pub fn get_account_transaction_version(
        &self,
        address: AccountAddress,
        sequence_number: u64,
        ledger_version: Version,
    ) -> Result<Option<Version>> {
        if let Some(version) = self
            .db
            .get::<TransactionByAccountSchema>(&(address, sequence_number))?
        {
            if version <= ledger_version {
                return Ok(Some(version));
            }
        }

        Ok(None)
    }

    /// Gets the version of a transaction by its hash.
    pub fn get_transaction_version_by_hash(
        &self,
        hash: &HashValue,
        ledger_version: Version,
    ) -> Result<Option<Version>> {
        Ok(match self.db.get::<TransactionByHashSchema>(hash)? {
            Some(version) if version <= ledger_version => Some(version),
            _ => None,
        })
    }

    /// Gets an iterator that yields `(sequence_number, version)` for each
    /// transaction sent by an account, starting at `start_seq_num`, and returning
    /// at most `num_versions` results with `version <= ledger_version`.
    ///
    /// Guarantees that the returned sequence numbers are sequential, i.e.,
    /// `seq_num_{i} + 1 = seq_num_{i+1}`.
    pub fn get_account_transaction_version_iter(
        &self,
        address: AccountAddress,
        start_seq_num: u64,
        num_versions: u64,
        ledger_version: Version,
    ) -> Result<AccountTransactionVersionIter> {
        let mut iter = self
            .db
            .iter::<TransactionByAccountSchema>(ReadOptions::default())?;
        iter.seek(&(address, start_seq_num))?;
        Ok(AccountTransactionVersionIter {
            inner: iter,
            address,
            expected_next_seq_num: start_seq_num,
            end_seq_num: start_seq_num
                .checked_add(num_versions)
                .ok_or_else(|| format_err!("too many transactions requested"))?,
            prev_version: None,
            ledger_version,
        })
    }

    /// Get signed transaction given `version`
    pub fn get_transaction(&self, version: Version) -> Result<Transaction> {
        self.db
            .get::<TransactionSchema>(&version)?
            .ok_or_else(|| DiemDbError::NotFound(format!("Txn {}", version)).into())
    }

    /// Gets an iterator that yields `num_transactions` transactions starting from `start_version`.
    pub fn get_transaction_iter(
        &self,
        start_version: Version,
        num_transactions: usize,
    ) -> Result<TransactionIter> {
        let mut iter = self.db.iter::<TransactionSchema>(ReadOptions::default())?;
        iter.seek(&start_version)?;
        Ok(TransactionIter {
            inner: iter,
            expected_next_version: start_version,
            end_version: start_version
                .checked_add(num_transactions as u64)
                .ok_or_else(|| format_err!("too many transactions requested"))?,
        })
    }

    /// Get the first version that txn starts existent.
    pub fn get_first_txn_version(&self) -> Result<Option<Version>> {
        let mut iter = self.db.iter::<TransactionSchema>(Default::default())?;
        iter.seek_to_first();
        iter.next().map(|res| res.map(|(v, _)| v)).transpose()
    }

    /// Returns the block metadata carried on the block metadata transaction at or preceding
    /// `version`, together with the version of the block metadata transaction.
    /// Returns None if there's no such transaction at or preceding `version` (it's likely the genesis
    /// version 0).
    pub fn get_block_metadata(&self, version: Version) -> Result<Option<(Version, BlockMetadata)>> {
        // Must be larger than a block size, otherwise a NotFound error will be raised wrongly.
        const MAX_VERSIONS_TO_SEARCH: usize = 1000 * 100;

        // Linear search via `DB::rev_iter()` here, NOT expecting performance hit, due to the fact
        // that the iterator caches data block and that there are limited number of transactions in
        // each block.
        let mut iter = self.db.rev_iter::<TransactionSchema>(Default::default())?;
        iter.seek(&version)?;
        for res in iter.take(MAX_VERSIONS_TO_SEARCH) {
            let (v, txn) = res?;
            if let Transaction::BlockMetadata(block_meta) = txn {
                return Ok(Some((v, block_meta)));
            } else if v == 0 {
                return Ok(None);
            }
        }

        Err(DiemDbError::NotFound(format!("BlockMetadata preceding version {}", version)).into())
    }

    /// Save signed transaction at `version`
    pub fn put_transaction(
        &self,
        version: Version,
        transaction: &Transaction,
        cs: &mut ChangeSet,
    ) -> Result<()> {
        if let Transaction::UserTransaction(txn) = transaction {
            cs.batch.put::<TransactionByAccountSchema>(
                &(txn.sender(), txn.sequence_number()),
                &version,
            )?;
        }
        cs.batch
            .put::<TransactionByHashSchema>(&transaction.hash(), &version)?;
        cs.batch.put::<TransactionSchema>(&version, transaction)?;

        Ok(())
    }

    /// Get executed transaction vm output given `version`
    pub fn get_write_set(&self, version: Version) -> Result<WriteSet> {
        self.db
            .get::<WriteSetSchema>(&version)?
            .ok_or_else(|| DiemDbError::NotFound(format!("WriteSet at version {}", version)).into())
    }

    /// Get the first version that write set starts existent.
    pub fn get_first_write_set_version(&self) -> Result<Option<Version>> {
        let mut iter = self.db.iter::<WriteSetSchema>(Default::default())?;
        iter.seek_to_first();
        iter.next().map(|res| res.map(|(v, _)| v)).transpose()
    }

    /// Save executed transaction vm output given `version`
    pub fn put_write_set(
        &self,
        version: Version,
        write_set: &WriteSet,
        cs: &mut ChangeSet,
    ) -> Result<()> {
        cs.batch.put::<WriteSetSchema>(&version, write_set)
    }
}

pub struct TransactionIter<'a> {
    inner: SchemaIterator<'a, TransactionSchema>,
    expected_next_version: Version,
    end_version: Version,
}

impl<'a> TransactionIter<'a> {
    fn next_impl(&mut self) -> Result<Option<Transaction>> {
        if self.expected_next_version >= self.end_version {
            return Ok(None);
        }

        let ret = match self.inner.next().transpose()? {
            Some((version, transaction)) => {
                ensure!(
                    version == self.expected_next_version,
                    "Transaction versions are not consecutive.",
                );
                self.expected_next_version += 1;
                Some(transaction)
            }
            None => None,
        };

        Ok(ret)
    }
}

impl<'a> Iterator for TransactionIter<'a> {
    type Item = Result<Transaction>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_impl().transpose()
    }
}

// TODO(philiphayes): this will need to change to support CRSNs
// (Conflict-Resistant Sequence Numbers)[https://github.com/diem/dip/blob/main/dips/dip-168.md].
//
// It depends on the implementation details, but we'll probably index by _requested_
// transaction sequence number rather than committed account sequence number.
// This would mean the property: `seq_num_{i+1} == seq_num_{i} + 1` would no longer
// be guaranteed and the check should be removed.
//
// This index would also no longer iterate over an account's transactions in
// committed order, meaning the outer method would need to overread by
// `CRSN_WINDOW_SIZE`, sort by version, and take only `limit` entries to get
// at most `limit` transactions in committed order. Alternatively, add another
// index for scanning an accounts transactions in committed order, e.g.,
// `(AccountAddress, Version) -> SeqNum`.

pub struct AccountTransactionVersionIter<'a> {
    inner: SchemaIterator<'a, TransactionByAccountSchema>,
    address: AccountAddress,
    expected_next_seq_num: u64,
    end_seq_num: u64,
    prev_version: Option<Version>,
    ledger_version: Version,
}

impl<'a> AccountTransactionVersionIter<'a> {
    fn next_impl(&mut self) -> Result<Option<(u64, Version)>> {
        if self.expected_next_seq_num >= self.end_seq_num {
            return Ok(None);
        }

        Ok(match self.inner.next().transpose()? {
            Some(((address, seq_num), version)) => {
                // No more transactions sent by this account.
                if address != self.address {
                    return Ok(None);
                }

                // Ensure seq_num_{i+1} == seq_num_{i} + 1
                ensure!(
                    seq_num == self.expected_next_seq_num,
                    "DB corruption: account transactions sequence numbers are not contiguous: \
                     actual: {}, expected: {}",
                    seq_num,
                    self.expected_next_seq_num,
                );

                // Ensure version_{i+1} > version_{i}
                if let Some(prev_version) = self.prev_version {
                    ensure!(
                        prev_version < version,
                        "DB corruption: account transaction versions are not strictly increasing: \
                         previous version: {}, current version: {}",
                        prev_version,
                        version,
                    );
                }

                // No more transactions (in this view of the ledger).
                if version > self.ledger_version {
                    return Ok(None);
                }

                self.expected_next_seq_num += 1;
                self.prev_version = Some(version);
                Some((seq_num, version))
            }
            None => None,
        })
    }
}

impl<'a> Iterator for AccountTransactionVersionIter<'a> {
    type Item = Result<(u64, Version)>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_impl().transpose()
    }
}

#[cfg(test)]
mod test;
<|endoftext|>


<commit_before>// Copyright 2012 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.

/*!
 * A simple map based on a vector for small integer keys. Space requirements
 * are O(highest integer key).
 */

use core::container::{Container, Mutable, Map, Set};
use core::option::{Some, None};
use core::prelude::*;

pub struct SmallIntMap<T> {
    priv v: ~[Option<T>],
}

impl<V> SmallIntMap<V>: Container {
    /// Return the number of elements in the map
    pure fn len(&self) -> uint {
        let mut sz = 0;
        for self.v.each |item| {
            if item.is_some() {
                sz += 1;
            }
        }
        sz
    }

    /// Return true if the map contains no elements
    pure fn is_empty(&self) -> bool { self.len() == 0 }
}

impl<V> SmallIntMap<V>: Mutable {
    /// Clear the map, removing all key-value pairs.
    fn clear(&mut self) { self.v.clear() }
}

impl<V> SmallIntMap<V>: Map<uint, V> {
    /// Return true if the map contains a value for the specified key
    pure fn contains_key(&self, key: &uint) -> bool {
        self.find(key).is_some()
    }

    /// Visit all keys
    pure fn each_key(&self, blk: fn(key: &uint) -> bool) {
        self.each(|k, _| blk(k))
    }

    /// Visit all values
    pure fn each_value(&self, blk: fn(value: &V) -> bool) {
        self.each(|_, v| blk(v))
    }

    /// Return the value corresponding to the key in the map
    pure fn find(&self, key: &uint) -> Option<&self/V> {
        if *key < self.v.len() {
            match self.v[*key] {
              Some(ref value) => Some(value),
              None => None
            }
        } else {
            None
        }
    }

    /// Insert a key-value pair into the map. An existing value for a
    /// key is replaced by the new value. Return true if the key did
    /// not already exist in the map.
    fn insert(&mut self, key: uint, value: V) -> bool {
        let exists = self.contains_key(&key);
        let len = self.v.len();
        if len <= key {
            vec::grow_fn(&mut self.v, key - len + 1, |_| None);
        }
        self.v[key] = Some(value);
        !exists
    }

    /// Remove a key-value pair from the map. Return true if the key
    /// was present in the map, otherwise false.
    fn remove(&mut self, key: &uint) -> bool {
        if *key >= self.v.len() {
            return false;
        }
        let removed = self.v[*key].is_some();
        self.v[*key] = None;
        removed
    }
}

pub impl<V> SmallIntMap<V> {
    /// Create an empty SmallIntMap
    static pure fn new() -> SmallIntMap<V> { SmallIntMap{v: ~[]} }

    /// Visit all key-value pairs
    pure fn each(&self, it: fn(key: &uint, value: &V) -> bool) {
        for uint::range(0, self.v.len()) |i| {
            match self.v[i] {
              Some(ref elt) => if !it(&i, elt) { break },
              None => ()
            }
        }
    }

    pure fn get(&self, key: &uint) -> &self/V {
        self.find(key).expect("key not present")
    }
}

pub impl<V: Copy> SmallIntMap<V> {
    fn update_with_key(&mut self, key: uint, val: V,
                       ff: fn(uint, V, V) -> V) -> bool {
        match self.find(&key) {
          None => self.insert(key, val),
          Some(orig) => self.insert(key, ff(key, copy *orig, val)),
        }
    }

    fn update(&mut self, key: uint, newval: V, ff: fn(V, V) -> V) -> bool {
        self.update_with_key(key, newval, |_k, v, v1| ff(v,v1))
    }
}

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

    #[test]
    fn test_len() {
        let mut map = SmallIntMap::new();
        assert map.len() == 0;
        assert map.is_empty();
        assert map.insert(5, 20);
        assert map.len() == 1;
        assert !map.is_empty();
        assert map.insert(11, 12);
        assert map.len() == 2;
        assert !map.is_empty();
        assert map.insert(14, 22);
        assert map.len() == 3;
        assert !map.is_empty();
    }

    #[test]
    fn test_clear() {
        let mut map = SmallIntMap::new();
        assert map.insert(5, 20);
        assert map.insert(11, 12);
        assert map.insert(14, 22);
        map.clear();
        assert map.is_empty();
        assert map.find(&5).is_none();
        assert map.find(&11).is_none();
        assert map.find(&14).is_none();
    }

    #[test]
    fn test_insert_with_key() {
        let mut map = SmallIntMap::new();

        // given a new key, initialize it with this new count, given
        // given an existing key, add more to its count
        fn addMoreToCount(_k: uint, v0: uint, v1: uint) -> uint {
            v0 + v1
        }

        fn addMoreToCount_simple(v0: uint, v1: uint) -> uint {
            v0 + v1
        }

        // count integers
        map.update(3, 1, addMoreToCount_simple);
        map.update_with_key(9, 1, addMoreToCount);
        map.update(3, 7, addMoreToCount_simple);
        map.update_with_key(5, 3, addMoreToCount);
        map.update_with_key(3, 2, addMoreToCount);

        // check the total counts
        assert map.find(&3).get() == &10;
        assert map.find(&5).get() == &3;
        assert map.find(&9).get() == &1;

        // sadly, no sevens were counted
        assert map.find(&7).is_none();
    }
}
<commit_msg>implement BaseIter for SmallIntMap<commit_after>// Copyright 2012 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.

/*!
 * A simple map based on a vector for small integer keys. Space requirements
 * are O(highest integer key).
 */

use core::container::{Container, Mutable, Map, Set};
use core::iter::BaseIter;
use core::option::{Some, None};
use core::prelude::*;

pub struct SmallIntMap<T> {
    priv v: ~[Option<T>],
}

impl<V> SmallIntMap<V>: BaseIter<(uint, &V)> {
    /// Visit all key-value pairs in order
    pure fn each(&self, it: fn(&(uint, &self/V)) -> bool) {
        for uint::range(0, self.v.len()) |i| {
            match self.v[i] {
              Some(ref elt) => if !it(&(i, elt)) { break },
              None => ()
            }
        }
    }

    pure fn size_hint(&self) -> Option<uint> { Some(self.len()) }
}

impl<V> SmallIntMap<V>: Container {
    /// Return the number of elements in the map
    pure fn len(&self) -> uint {
        let mut sz = 0;
        for self.v.each |item| {
            if item.is_some() {
                sz += 1;
            }
        }
        sz
    }

    /// Return true if the map contains no elements
    pure fn is_empty(&self) -> bool { self.len() == 0 }
}

impl<V> SmallIntMap<V>: Mutable {
    /// Clear the map, removing all key-value pairs.
    fn clear(&mut self) { self.v.clear() }
}

impl<V> SmallIntMap<V>: Map<uint, V> {
    /// Return true if the map contains a value for the specified key
    pure fn contains_key(&self, key: &uint) -> bool {
        self.find(key).is_some()
    }

    /// Visit all keys in order
    pure fn each_key(&self, blk: fn(key: &uint) -> bool) {
        self.each(|&(k, _)| blk(&k))
    }

    /// Visit all values in order
    pure fn each_value(&self, blk: fn(value: &V) -> bool) {
        self.each(|&(_, v)| blk(v))
    }

    /// Return the value corresponding to the key in the map
    pure fn find(&self, key: &uint) -> Option<&self/V> {
        if *key < self.v.len() {
            match self.v[*key] {
              Some(ref value) => Some(value),
              None => None
            }
        } else {
            None
        }
    }

    /// Insert a key-value pair into the map. An existing value for a
    /// key is replaced by the new value. Return true if the key did
    /// not already exist in the map.
    fn insert(&mut self, key: uint, value: V) -> bool {
        let exists = self.contains_key(&key);
        let len = self.v.len();
        if len <= key {
            vec::grow_fn(&mut self.v, key - len + 1, |_| None);
        }
        self.v[key] = Some(value);
        !exists
    }

    /// Remove a key-value pair from the map. Return true if the key
    /// was present in the map, otherwise false.
    fn remove(&mut self, key: &uint) -> bool {
        if *key >= self.v.len() {
            return false;
        }
        let removed = self.v[*key].is_some();
        self.v[*key] = None;
        removed
    }
}

pub impl<V> SmallIntMap<V> {
    /// Create an empty SmallIntMap
    static pure fn new() -> SmallIntMap<V> { SmallIntMap{v: ~[]} }

    pure fn get(&self, key: &uint) -> &self/V {
        self.find(key).expect("key not present")
    }
}

pub impl<V: Copy> SmallIntMap<V> {
    fn update_with_key(&mut self, key: uint, val: V,
                       ff: fn(uint, V, V) -> V) -> bool {
        match self.find(&key) {
          None => self.insert(key, val),
          Some(orig) => self.insert(key, ff(key, copy *orig, val)),
        }
    }

    fn update(&mut self, key: uint, newval: V, ff: fn(V, V) -> V) -> bool {
        self.update_with_key(key, newval, |_k, v, v1| ff(v,v1))
    }
}

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

    #[test]
    fn test_len() {
        let mut map = SmallIntMap::new();
        assert map.len() == 0;
        assert map.is_empty();
        assert map.insert(5, 20);
        assert map.len() == 1;
        assert !map.is_empty();
        assert map.insert(11, 12);
        assert map.len() == 2;
        assert !map.is_empty();
        assert map.insert(14, 22);
        assert map.len() == 3;
        assert !map.is_empty();
    }

    #[test]
    fn test_clear() {
        let mut map = SmallIntMap::new();
        assert map.insert(5, 20);
        assert map.insert(11, 12);
        assert map.insert(14, 22);
        map.clear();
        assert map.is_empty();
        assert map.find(&5).is_none();
        assert map.find(&11).is_none();
        assert map.find(&14).is_none();
    }

    #[test]
    fn test_insert_with_key() {
        let mut map = SmallIntMap::new();

        // given a new key, initialize it with this new count, given
        // given an existing key, add more to its count
        fn addMoreToCount(_k: uint, v0: uint, v1: uint) -> uint {
            v0 + v1
        }

        fn addMoreToCount_simple(v0: uint, v1: uint) -> uint {
            v0 + v1
        }

        // count integers
        map.update(3, 1, addMoreToCount_simple);
        map.update_with_key(9, 1, addMoreToCount);
        map.update(3, 7, addMoreToCount_simple);
        map.update_with_key(5, 3, addMoreToCount);
        map.update_with_key(3, 2, addMoreToCount);

        // check the total counts
        assert map.find(&3).get() == &10;
        assert map.find(&5).get() == &3;
        assert map.find(&9).get() == &1;

        // sadly, no sevens were counted
        assert map.find(&7).is_none();
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{self, BufReader, BufRead};
use std::path::Path;
use std::collections::HashMap;

#[derive(Debug, PartialEq)]
enum ConfigType {
    String,
    Boolean,
    Vector,
}

#[derive(Debug)]
struct ConfigVariable {
    var_type: ConfigType,
    value: String,
}

impl ConfigVariable {
    fn extract_value_as_string(&self, typ: ConfigType) -> Result<String, String> {
        if self.var_type == typ {
            Ok(self.value.to_owned())
        } else {
            Err(format!("Incorrect variable type. Expected: {:?}. Actual: {:?}",
                        typ,
                        self.var_type))
        }
    }
}

#[derive(Debug)]
struct ConfigParams {
    params: HashMap<String, ConfigVariable>,
}

impl ConfigParams {
    fn new() -> ConfigParams {
        ConfigParams { params: HashMap::new() }
    }

    fn parse<P: AsRef<Path>>(path: P) -> ConfigParams {
        let conf_file = File::open(path).unwrap();
        let content = BufReader::new(conf_file);

        let is_not_comment = |x: &Result<String, io::Error>| {
            match x {
                &Err(_) => false,
                &Ok(ref line) => {
                    !(line.starts_with("#") || line.starts_with(";") || line.is_empty())
                }
            }

        };
        let mut params = ConfigParams::new();
        for line in content.lines().filter(is_not_comment) {
            if line.is_ok() {
                params.update_config(line.unwrap());
            }

        }

        params
    }
    // Will parse the line and update the internal structure
    fn update_config(&mut self, line: String) {
        // First space splits the key from value, then if a list is separated by a comma
        let mut parts = line.splitn(2, " ").map(|x| x.to_owned());
        let key = parts.next().unwrap().to_lowercase();
        match parts.next() {
            None => {
                self.params.insert(key,
                                   ConfigVariable {
                                       var_type: ConfigType::Boolean,
                                       value: "true".to_owned(),
                                   });
            }
            Some(value) => {
                if value.contains(",") {
                    self.params.insert(key,
                                       ConfigVariable {
                                           var_type: ConfigType::Vector,
                                           value: value,
                                       });
                } else {
                    self.params.insert(key,
                                       ConfigVariable {
                                           var_type: ConfigType::String,
                                           value: value,
                                       });
                }
            }
        }

    }

    fn string_param(&self, key: &str) -> Result<String, String> {
        key.to_lowercase();
        match self.params.get(key) {
            Some(variable) => variable.extract_value_as_string(ConfigType::String),
            None => Err(format!("No such configuration paramter {:?}", key)),
        }

    }

    fn bool_param(&self, key: &str) -> Result<bool, String> {
        key.to_lowercase();
        match self.params.get(key) {
            Some(variable) => {
                variable.extract_value_as_string(ConfigType::Boolean)
                    .map(|x| x.parse::<bool>().unwrap())
            }
            None => Ok(false),
        }
    }

    fn vec_param(&self, key: &str) -> Result<Vec<String>, String> {
        key.to_lowercase();
        match self.params.get(key) {
            Some(variable) => {
                // Get param as string
                // split string at comma
                // trim each item and convert to String
                // collect to vector
                variable.extract_value_as_string(ConfigType::Vector)
                    .map(|string| string.split(",").map(|item| item.trim().to_owned()).collect())
            }
            None => Err(format!("No such configuration paramter {:?}", key)),
        }
    }
}

fn main() {
    const CONF: &'static str = "test.conf";
    let params = ConfigParams::parse(CONF);

    println!("{:?}", params.string_param("fullname"));
    println!("{:?}", params.string_param("favouritefruit"));
    println!("{:?}", params.bool_param("needspeeling"));
    println!("{:?}", params.bool_param("seedsremoved"));
    println!("{:?}", params.vec_param("otherfamily"));
}

#[cfg(test)]
mod tests {

    #[test]
    fn main_test() {
        const CONF: &'static str = "test.conf";
        let params = super::ConfigParams::parse(CONF);
        assert_eq!(params.string_param("fullname").unwrap(), "Foo Barber");
        assert_eq!(params.string_param("favouritefruit").unwrap(), "banana");
        assert!(params.bool_param("needspeeling").unwrap());
        assert!(!params.bool_param("seedsremoved").unwrap());
        assert_eq!(params.vec_param("otherfamily").unwrap(),
                   vec!["Rhu Barber", "Harry Barber"]);
    }
}
<commit_msg>Changed to use generics.<commit_after>use std::fs::File;
use std::io::{self, BufReader, BufRead};
use std::path::Path;
use std::collections::HashMap;

#[derive(Debug, PartialEq)]
enum ConfigType {
    String,
    Boolean,
    Vector,
}

#[derive(Debug)]
struct ConfigVariable {
    var_type: ConfigType,
    value: String,
}

impl ConfigVariable {
    fn extract_value_as_string(&self, typ: ConfigType) -> Result<String, String> {
        if self.var_type == typ {
            Ok(self.value.to_owned())
        } else {
            Err(format!("Incorrect variable type. Expected: {:?}. Actual: {:?}",
                        typ,
                        self.var_type))
        }
    }
}

trait FromConfig {
    fn from_config(param: &ConfigParams, key: &str) -> Result<Self, String>
        where Self: std::marker::Sized;
}

impl FromConfig for String {
    fn from_config(param: &ConfigParams, key: &str) -> Result<String, String> {
        match param.params.get(key) {
            Some(variable) => variable.extract_value_as_string(ConfigType::String),
            None => Err(format!("Unknown configuration parameter: {:?}", key)),
        }
    }
}

impl FromConfig for bool {
    fn from_config(param: &ConfigParams, key: &str) -> Result<bool, String> {
        match param.params.get(key) {
            Some(variable) => {
                variable.extract_value_as_string(ConfigType::Boolean)
                    .map(|x| x.parse::<bool>().unwrap())
            }
            None => Ok(false),
        }

    }
}

impl FromConfig for Vec<String> {
    fn from_config(param: &ConfigParams, key: &str) -> Result<Vec<String>, String> {
        match param.params.get(key) {
            Some(variable) => {
                variable.extract_value_as_string(ConfigType::Vector)
                    .map(|string| string.split(",").map(|item| item.trim().to_owned()).collect())
            }
            None => Err(format!("Unknown configuration parameter: {:?}", key)),
        }

    }
}

#[derive(Debug)]
struct ConfigParams {
    params: HashMap<String, ConfigVariable>,
}

impl ConfigParams {
    fn new() -> ConfigParams {
        ConfigParams { params: HashMap::new() }
    }

    fn parse<P: AsRef<Path>>(path: P) -> ConfigParams {
        let conf_file = File::open(path).unwrap();
        let content = BufReader::new(conf_file);

        let is_not_comment = |x: &Result<String, io::Error>| {
            match x {
                &Err(_) => false,
                &Ok(ref line) => {
                    !(line.starts_with("#") || line.starts_with(";") || line.is_empty())
                }
            }
        };
        let mut params = ConfigParams::new();
        for line in content.lines().filter(is_not_comment) {
            if line.is_ok() {
                params.update_config(line.unwrap());
            }
        }

        params
    }
    // Will parse the line and update the internal structure
    fn update_config(&mut self, line: String) {
        let mut parts = line.splitn(2, " ").map(|x| x.to_owned());
        let key = parts.next().unwrap().to_lowercase();
        match parts.next() {
            None => {
                self.params.insert(key,
                                   ConfigVariable {
                                       var_type: ConfigType::Boolean,
                                       value: "true".to_owned(),
                                   });
            }
            Some(value) => {
                if value.contains(",") {
                    self.params.insert(key,
                                       ConfigVariable {
                                           var_type: ConfigType::Vector,
                                           value: value,
                                       });
                } else {
                    self.params.insert(key,
                                       ConfigVariable {
                                           var_type: ConfigType::String,
                                           value: value,
                                       });
                }
            }
        }

    }

    fn param<T: FromConfig>(&self, key: &str) -> Result<T, String> {
        key.to_lowercase();
        FromConfig::from_config(&self, key)
    }
}

fn main() {
    const CONF: &'static str = "test.conf";
    let params = ConfigParams::parse(CONF);

    println!("{:?}", params.param::<String>("fullname"));
    println!("{:?}", params.param::<String>("favouritefruit"));
    println!("{:?}", params.param::<bool>("needspeeling"));
    println!("{:?}", params.param::<bool>("seedsremoved"));
    println!("{:?}", params.param::<Vec<String>>("otherfamily"));
}

#[cfg(test)]
mod tests {

    #[test]
    fn main_test() {
        const CONF: &'static str = "test.conf";
        let params = super::ConfigParams::parse(CONF);
        assert_eq!(params.param::<String>("fullname").unwrap(), "Foo Barber");
        assert_eq!(params.param::<String>("favouritefruit").unwrap(), "banana");
        assert!(params.param::<bool>("needspeeling").unwrap());
        assert!(!params.param::<bool>("seedsremoved").unwrap());
        assert_eq!(params.param::<Vec<String>>("otherfamily").unwrap(),
                   vec!["Rhu Barber", "Harry Barber"]);
    }
}
<|endoftext|>


<commit_before>pub const POMBASE_ANN_EXT_TERM_CV_NAME: &'static str = "PomBase annotation extension terms";
pub const ANNOTATION_EXT_REL_PREFIX: &'static str = "annotation_extension_relation-";

pub enum FeatureRelAnnotationType {
    Interaction,
    Ortholog,
    Paralog,
}
pub struct FeatureRelConfig {
    pub rel_type_name: &'static str,
    pub annotation_type: FeatureRelAnnotationType,
}
pub const FEATURE_REL_CONFIGS: [FeatureRelConfig; 4] =
    [
        FeatureRelConfig {
            rel_type_name: "interacts_physically",
            annotation_type: FeatureRelAnnotationType::Interaction,
        },
        FeatureRelConfig {
            rel_type_name: "interacts_genetically",
            annotation_type: FeatureRelAnnotationType::Interaction,
        },
        FeatureRelConfig {
            rel_type_name: "orthologous_to",
            annotation_type: FeatureRelAnnotationType::Ortholog,
        },
        FeatureRelConfig {
            rel_type_name: "paralogous_to",
            annotation_type: FeatureRelAnnotationType::Paralog,
        },
    ];

// relations to use when copy annotation to parents (ie. adding the
// annotation of child terms to parents)
pub const DESCENDANT_REL_NAMES: [&'static str; 2] = ["is_a", "part_of", "regulates"];
<commit_msg>Fix compilation error<commit_after>pub const POMBASE_ANN_EXT_TERM_CV_NAME: &'static str = "PomBase annotation extension terms";
pub const ANNOTATION_EXT_REL_PREFIX: &'static str = "annotation_extension_relation-";

pub enum FeatureRelAnnotationType {
    Interaction,
    Ortholog,
    Paralog,
}
pub struct FeatureRelConfig {
    pub rel_type_name: &'static str,
    pub annotation_type: FeatureRelAnnotationType,
}
pub const FEATURE_REL_CONFIGS: [FeatureRelConfig; 4] =
    [
        FeatureRelConfig {
            rel_type_name: "interacts_physically",
            annotation_type: FeatureRelAnnotationType::Interaction,
        },
        FeatureRelConfig {
            rel_type_name: "interacts_genetically",
            annotation_type: FeatureRelAnnotationType::Interaction,
        },
        FeatureRelConfig {
            rel_type_name: "orthologous_to",
            annotation_type: FeatureRelAnnotationType::Ortholog,
        },
        FeatureRelConfig {
            rel_type_name: "paralogous_to",
            annotation_type: FeatureRelAnnotationType::Paralog,
        },
    ];

// relations to use when copy annotation to parents (ie. adding the
// annotation of child terms to parents)
pub const DESCENDANT_REL_NAMES: [&'static str; 3] = ["is_a", "part_of", "regulates"];
<|endoftext|>


<commit_before>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || std::str::from_utf8(iter.next().unwrap_or(&[48])).unwrap_or("*");
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = acl.check_name_in_proxy_list(query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match acl.check_name_in_proxy_list(name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::from(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::from(*ip)),
                    RData::PTR(_) => panic!("PTR records should not reach here"),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

async fn acl_lookup<Local, Remote>(
    acl: &Option<AccessControl>,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match should_forward_by_query(acl, query) {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (query): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (query): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    let decider = async {
        let local_response = local_response_fut.await;
        if should_forward_by_response(acl, &local_response, query) {
            None
        } else {
            Some(local_response)
        }
    };
    tokio::pin!(remote_response_fut);
    tokio::pin!(decider);
    let mut use_remote = false;
    let mut remote_response = None;
    loop {
        tokio::select! {
            response = &mut remote_response_fut => {
                if use_remote {
                    debug!("pick remote response (response): {:?}", response);
                    return (response, true);
                } else {
                    remote_response = Some(response);
                }
            }
            decision = &mut decider => {
                if let Some(local_response) = decision {
                    debug!("pick local response (response): {:?}", local_response);
                    return (local_response, false);
                } else if let Some(remote_response) = remote_response {
                    debug!("pick remote response (response): {:?}", remote_response);
                    return (remote_response, true);
                } else {
                    use_remote = true;
                }
            }
        }
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    #[cfg(target_os = "android")]
        let local_upstream = Arc::new(upstream::UnixSocketUpstream {
        path: config.local_dns_path.clone().expect("local query DNS path"),
    });
    #[cfg(not(target_os = "android"))]
        let local_upstream = Arc::new(upstream::UdpUpstream {
        server: config.local_dns_addr.clone().expect("local query DNS address"),
    });
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let local_upstream = Arc::clone(&local_upstream);
        let remote_upstream = Arc::clone(&remote_upstream);
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if !request.recursion_desired() {
                message.set_recursion_desired(false);
                message.set_response_code(ResponseCode::NotImp);
            } else if request.query_count() > 0 {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.acl(), local_upstream, remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<commit_msg>Refine code style<commit_after>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || std::str::from_utf8(iter.next().unwrap_or(&[48])).unwrap_or("*");
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = acl.check_name_in_proxy_list(query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match acl.check_name_in_proxy_list(name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                    RData::PTR(_) => unreachable!(),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

async fn acl_lookup<Local, Remote>(
    acl: &Option<AccessControl>,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match should_forward_by_query(acl, query) {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (query): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (query): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    let decider = async {
        let local_response = local_response_fut.await;
        if should_forward_by_response(acl, &local_response, query) {
            None
        } else {
            Some(local_response)
        }
    };
    tokio::pin!(remote_response_fut, decider);
    let mut use_remote = false;
    let mut remote_response = None;
    loop {
        tokio::select! {
            response = &mut remote_response_fut => {
                if use_remote {
                    debug!("pick remote response (response): {:?}", response);
                    return (response, true);
                } else {
                    remote_response = Some(response);
                }
            }
            decision = &mut decider => {
                if let Some(local_response) = decision {
                    debug!("pick local response (response): {:?}", local_response);
                    return (local_response, false);
                } else if let Some(remote_response) = remote_response {
                    debug!("pick remote response (response): {:?}", remote_response);
                    return (remote_response, true);
                } else {
                    use_remote = true;
                }
            }
        }
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    #[cfg(target_os = "android")]
        let local_upstream = Arc::new(upstream::UnixSocketUpstream {
        path: config.local_dns_path.clone().expect("local query DNS path"),
    });
    #[cfg(not(target_os = "android"))]
        let local_upstream = Arc::new(upstream::UdpUpstream {
        server: config.local_dns_addr.clone().expect("local query DNS address"),
    });
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let local_upstream = Arc::clone(&local_upstream);
        let remote_upstream = Arc::clone(&remote_upstream);
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if !request.recursion_desired() {
                message.set_recursion_desired(false);
                message.set_response_code(ResponseCode::NotImp);
            } else if request.query_count() > 0 {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.acl(), local_upstream, remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<|endoftext|>