<commit_before>/*
Copyright (c) 2016-2017, Robert Ou <rqou@robertou.com> and contributors
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

use std::hash::{Hash, Hasher};
use std::marker::PhantomData;

#[derive(Debug)]
pub struct ObjPoolIndex<T> {
    i: usize,
    type_marker: PhantomData<T>
}

impl<T> Copy for ObjPoolIndex<T> { }

impl<T> Clone for ObjPoolIndex<T> {
    fn clone(&self) -> ObjPoolIndex<T> {
        *self
    }
}

impl<T> PartialEq for ObjPoolIndex<T> {
    fn eq(&self, other: &ObjPoolIndex<T>) -> bool {
        self.i == other.i
    }
}

impl<T> Eq for ObjPoolIndex<T> { }

impl<T> Hash for ObjPoolIndex<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.i.hash(state);
    }
}

impl<T> ObjPoolIndex<T> {
    pub fn get_raw_i(&self) -> usize {
        self.i
    }
}

#[derive(Debug)]
pub struct ObjPool<T> {
    storage: Vec<T>
}

impl<T> ObjPool<T> {
    pub fn new() -> ObjPool<T> {
        ObjPool {storage: Vec::new()}
    }

    pub fn insert(&mut self, item: T) -> ObjPoolIndex<T> {
        let i = self.storage.len();

        self.storage.push(item);

        ObjPoolIndex::<T> {i: i, type_marker: PhantomData}
    }

    pub fn get(&self, i: ObjPoolIndex<T>) -> &T {
        &self.storage[i.i]
    }

    pub fn get_mut(&mut self, i: ObjPoolIndex<T>) -> &mut T {
        &mut self.storage[i.i]
    }

    pub fn iter_idx(&self) -> ObjPoolIdxIterator<T> {
        ObjPoolIdxIterator {
            pool: self,
            current_idx: 0,
        }
    }

    pub fn iter(&self) -> ObjPoolItemIterator<T> {
        ObjPoolItemIterator {
            pool: self,
            current_idx: 0,
        }
    }

    pub fn iter_mut(&mut self) -> ObjPoolMutItemIterator<T> {
        ObjPoolMutItemIterator {
            pool: self,
            current_idx: 0,
        }
    }

    pub fn len(&self) -> usize {
        self.storage.len()
    }
}

impl<T: Default> ObjPool<T> {
    pub fn alloc(&mut self) -> ObjPoolIndex<T> {
        let i = self.storage.len();
        let o = T::default();

        self.storage.push(o);

        ObjPoolIndex::<T> {i: i, type_marker: PhantomData}
    }
}

pub struct ObjPoolIdxIterator<'a, T: 'a> {
    pool: &'a ObjPool<T>,
    current_idx: usize,
}

impl<'a, T> Iterator for ObjPoolIdxIterator<'a, T> {
    type Item = ObjPoolIndex<T>;

    fn next(&mut self) -> Option<ObjPoolIndex<T>> {
        if self.current_idx == self.pool.storage.len() {
            None
        } else {
            let ret = ObjPoolIndex::<T> {i: self.current_idx, type_marker: PhantomData};
            self.current_idx += 1;
            Some(ret)
        }
    }
}

pub struct ObjPoolItemIterator<'a, T: 'a> {
    pool: &'a ObjPool<T>,
    current_idx: usize,
}

impl<'a, T> Iterator for ObjPoolItemIterator<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<&'a T> {
        if self.current_idx == self.pool.storage.len() {
            None
        } else {
            let ret = self.pool.get(ObjPoolIndex::<T> {i: self.current_idx, type_marker: PhantomData});
            self.current_idx += 1;
            Some(ret)
        }
    }
}

pub struct ObjPoolMutItemIterator<'a, T: 'a> {
    pool: &'a mut ObjPool<T>,
    current_idx: usize,
}

impl<'a, T> Iterator for ObjPoolMutItemIterator<'a, T> {
    type Item = &'a mut T;

    fn next(&mut self) -> Option<&'a mut T> {
        if self.current_idx == self.pool.storage.len() {
            None
        } else {
            // This is necessary because Rust can't otherwise check that we didn't end up with two overlapping elements
            // available at once.
            let ret = unsafe {
                &mut *self.pool.storage.as_mut_ptr().offset(self.current_idx as isize)
            };
            self.current_idx += 1;
            Some(ret)
        }
    }
}

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

    #[derive(Default)]
    struct ObjPoolTestObject {
        foo: u32
    }

    #[test]
    fn objpool_basic_works() {
        let mut pool = ObjPool::<ObjPoolTestObject>::new();
        let x = pool.alloc();
        let y = pool.alloc();
        {
            let o = pool.get_mut(x);
            o.foo = 123;
        }
        {
            let o = pool.get_mut(y);
            o.foo = 456;
        }
        let ox = pool.get(x);
        let oy = pool.get(y);
        assert_eq!(ox.foo, 123);
        assert_eq!(oy.foo, 456);
    }
}
<commit_msg>xc2par: Replace objectpool iterators with Vec's iterators<commit_after>/*
Copyright (c) 2016-2017, Robert Ou <rqou@robertou.com> and contributors
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use std::slice::{Iter, IterMut};

#[derive(Debug)]
pub struct ObjPoolIndex<T> {
    i: usize,
    type_marker: PhantomData<T>
}

impl<T> Copy for ObjPoolIndex<T> { }

impl<T> Clone for ObjPoolIndex<T> {
    fn clone(&self) -> ObjPoolIndex<T> {
        *self
    }
}

impl<T> PartialEq for ObjPoolIndex<T> {
    fn eq(&self, other: &ObjPoolIndex<T>) -> bool {
        self.i == other.i
    }
}

impl<T> Eq for ObjPoolIndex<T> { }

impl<T> Hash for ObjPoolIndex<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.i.hash(state);
    }
}

impl<T> ObjPoolIndex<T> {
    pub fn get_raw_i(&self) -> usize {
        self.i
    }
}

#[derive(Debug)]
pub struct ObjPool<T> {
    storage: Vec<T>
}

impl<T> ObjPool<T> {
    pub fn new() -> ObjPool<T> {
        ObjPool {storage: Vec::new()}
    }

    pub fn insert(&mut self, item: T) -> ObjPoolIndex<T> {
        let i = self.storage.len();

        self.storage.push(item);

        ObjPoolIndex::<T> {i: i, type_marker: PhantomData}
    }

    pub fn get(&self, i: ObjPoolIndex<T>) -> &T {
        &self.storage[i.i]
    }

    pub fn get_mut(&mut self, i: ObjPoolIndex<T>) -> &mut T {
        &mut self.storage[i.i]
    }

    pub fn iter_idx(&self) -> ObjPoolIdxIterator<T> {
        ObjPoolIdxIterator {
            pool: self,
            current_idx: 0,
        }
    }

    pub fn iter(&self) -> Iter<T> {
        self.storage.iter()
    }

    pub fn iter_mut(&mut self) -> IterMut<T> {
        self.storage.iter_mut()
    }

    pub fn len(&self) -> usize {
        self.storage.len()
    }
}

impl<T: Default> ObjPool<T> {
    pub fn alloc(&mut self) -> ObjPoolIndex<T> {
        let i = self.storage.len();
        let o = T::default();

        self.storage.push(o);

        ObjPoolIndex::<T> {i: i, type_marker: PhantomData}
    }
}

pub struct ObjPoolIdxIterator<'a, T: 'a> {
    pool: &'a ObjPool<T>,
    current_idx: usize,
}

impl<'a, T> Iterator for ObjPoolIdxIterator<'a, T> {
    type Item = ObjPoolIndex<T>;

    fn next(&mut self) -> Option<ObjPoolIndex<T>> {
        if self.current_idx == self.pool.storage.len() {
            None
        } else {
            let ret = ObjPoolIndex::<T> {i: self.current_idx, type_marker: PhantomData};
            self.current_idx += 1;
            Some(ret)
        }
    }
}

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

    #[derive(Default)]
    struct ObjPoolTestObject {
        foo: u32
    }

    #[test]
    fn objpool_basic_works() {
        let mut pool = ObjPool::<ObjPoolTestObject>::new();
        let x = pool.alloc();
        let y = pool.alloc();
        {
            let o = pool.get_mut(x);
            o.foo = 123;
        }
        {
            let o = pool.get_mut(y);
            o.foo = 456;
        }
        let ox = pool.get(x);
        let oy = pool.get(y);
        assert_eq!(ox.foo, 123);
        assert_eq!(oy.foo, 456);
    }
}
<|endoftext|>


<commit_before>// This benchmark suite contains some benchmarks along a set of dimensions:
//   Hasher: std default (SipHash) and crate default (FxHash).
//   Int key distribution: low bit heavy, top bit heavy, and random.
//   Task: basic functionality: insert, insert_erase, lookup, lookup_fail, iter
#![feature(test)]

extern crate test;

use test::{black_box, Bencher};

use hashbrown::hash_map::DefaultHashBuilder;
use hashbrown::HashMap;
use std::collections::hash_map::RandomState;

const SIZE: usize = 1000;

// The default hashmap when using this crate directly.
type FxHashMap<K, V> = HashMap<K, V, DefaultHashBuilder>;
// This uses the hashmap from this crate with the default hasher of the stdlib.
type StdHashMap<K, V> = HashMap<K, V, RandomState>;

// A random key iterator.
#[derive(Clone, Copy)]
struct RandomKeys {
    state: usize,
}

impl RandomKeys {
    fn new() -> Self {
        RandomKeys { state: 0 }
    }
}

impl Iterator for RandomKeys {
    type Item = usize;
    fn next(&mut self) -> Option<usize> {
        // Add 1 then multiply by some 32 bit prime.
        self.state = self.state.wrapping_add(1).wrapping_mul(3787392781);
        Some(self.state)
    }
}

macro_rules! bench_suite {
    ($bench_macro:ident, $bench_fx_serial:ident, $bench_std_serial:ident,
     $bench_fx_highbits:ident, $bench_std_highbits:ident,
     $bench_fx_random:ident, $bench_std_random:ident) => {
        $bench_macro!($bench_fx_serial, FxHashMap, 0..);
        $bench_macro!($bench_std_serial, StdHashMap, 0..);
        $bench_macro!($bench_fx_highbits, FxHashMap, (0..).map(usize::swap_bytes));
        $bench_macro!(
            $bench_std_highbits,
            StdHashMap,
            (0..).map(usize::swap_bytes)
        );
        $bench_macro!($bench_fx_random, FxHashMap, RandomKeys::new());
        $bench_macro!($bench_std_random, StdHashMap, RandomKeys::new());
    };
}

macro_rules! bench_insert {
    ($name:ident, $maptype:ident, $keydist:expr) => {
        #[bench]
        fn $name(b: &mut Bencher) {
            b.iter(|| {
                let mut m = $maptype::default();
                for i in ($keydist).take(SIZE) {
                    m.insert(i, i);
                }
                black_box(m);
            })
        }
    };
}

bench_suite!(
    bench_insert,
    insert_fx_serial,
    insert_std_serial,
    insert_fx_highbits,
    insert_std_highbits,
    insert_fx_random,
    insert_std_random
);

macro_rules! bench_insert_erase {
    ($name:ident, $maptype:ident, $keydist:expr) => {
        #[bench]
        fn $name(b: &mut Bencher) {
            let mut m = $maptype::default();
            let mut add_iter = $keydist;
            for i in (&mut add_iter).take(SIZE) {
                m.insert(i, i);
            }
            let mut remove_iter = $keydist;
            b.iter(|| {
                // While keeping the size constant,
                // replace the first keydist with the second.
                for (add, remove) in (&mut add_iter).zip(&mut remove_iter).take(SIZE) {
                    m.insert(add, add);
                    black_box(m.remove(&remove));
                }
            })
        }
    };
}

bench_suite!(
    bench_insert_erase,
    insert_erase_fx_serial,
    insert_erase_std_serial,
    insert_erase_fx_highbits,
    insert_erase_std_highbits,
    insert_erase_fx_random,
    insert_erase_std_random
);

macro_rules! bench_lookup {
    ($name:ident, $maptype:ident, $keydist:expr) => {
        #[bench]
        fn $name(b: &mut Bencher) {
            let mut m = $maptype::default();
            for i in $keydist.take(SIZE) {
                m.insert(i, i);
            }

            b.iter(|| {
                for i in $keydist.take(SIZE) {
                    black_box(m.get(&i));
                }
            })
        }
    };
}

bench_suite!(
    bench_lookup,
    lookup_fx_serial,
    lookup_std_serial,
    lookup_fx_highbits,
    lookup_std_highbits,
    lookup_fx_random,
    lookup_std_random
);

macro_rules! bench_lookup_fail {
    ($name:ident, $maptype:ident, $keydist:expr) => {
        #[bench]
        fn $name(b: &mut Bencher) {
            let mut m = $maptype::default();
            let mut iter = $keydist;
            for i in (&mut iter).take(SIZE) {
                m.insert(i, i);
            }

            b.iter(|| {
                for i in (&mut iter).take(SIZE) {
                    black_box(m.get(&i));
                }
            })
        }
    };
}

bench_suite!(
    bench_lookup_fail,
    lookup_fail_fx_serial,
    lookup_fail_std_serial,
    lookup_fail_fx_highbits,
    lookup_fail_std_highbits,
    lookup_fail_fx_random,
    lookup_fail_std_random
);

macro_rules! bench_iter {
    ($name:ident, $maptype:ident, $keydist:expr) => {
        #[bench]
        fn $name(b: &mut Bencher) {
            let mut m = $maptype::default();
            for i in ($keydist).take(SIZE) {
                m.insert(i, i);
            }

            b.iter(|| {
                for i in &m {
                    black_box(i);
                }
            })
        }
    };
}

bench_suite!(
    bench_iter,
    iter_fx_serial,
    iter_std_serial,
    iter_fx_highbits,
    iter_std_highbits,
    iter_fx_random,
    iter_std_random
);
<commit_msg>Update benchmarks<commit_after>// This benchmark suite contains some benchmarks along a set of dimensions:
//   Hasher: std default (SipHash) and crate default (AHash).
//   Int key distribution: low bit heavy, top bit heavy, and random.
//   Task: basic functionality: insert, insert_erase, lookup, lookup_fail, iter
#![feature(test)]

extern crate test;

use test::{black_box, Bencher};

use hashbrown::hash_map::DefaultHashBuilder;
use hashbrown::HashMap;
use std::collections::hash_map::RandomState;

const SIZE: usize = 1000;

// The default hashmap when using this crate directly.
type AHashMap<K, V> = HashMap<K, V, DefaultHashBuilder>;
// This uses the hashmap from this crate with the default hasher of the stdlib.
type StdHashMap<K, V> = HashMap<K, V, RandomState>;

// A random key iterator.
#[derive(Clone, Copy)]
struct RandomKeys {
    state: usize,
}

impl RandomKeys {
    fn new() -> Self {
        RandomKeys { state: 0 }
    }
}

impl Iterator for RandomKeys {
    type Item = usize;
    fn next(&mut self) -> Option<usize> {
        // Add 1 then multiply by some 32 bit prime.
        self.state = self.state.wrapping_add(1).wrapping_mul(3787392781);
        Some(self.state)
    }
}

macro_rules! bench_suite {
    ($bench_macro:ident, $bench_ahash_serial:ident, $bench_std_serial:ident,
     $bench_ahash_highbits:ident, $bench_std_highbits:ident,
     $bench_ahash_random:ident, $bench_std_random:ident) => {
        $bench_macro!($bench_ahash_serial, AHashMap, 0..);
        $bench_macro!($bench_std_serial, StdHashMap, 0..);
        $bench_macro!(
            $bench_ahash_highbits,
            AHashMap,
            (0..).map(usize::swap_bytes)
        );
        $bench_macro!(
            $bench_std_highbits,
            StdHashMap,
            (0..).map(usize::swap_bytes)
        );
        $bench_macro!($bench_ahash_random, AHashMap, RandomKeys::new());
        $bench_macro!($bench_std_random, StdHashMap, RandomKeys::new());
    };
}

macro_rules! bench_insert {
    ($name:ident, $maptype:ident, $keydist:expr) => {
        #[bench]
        fn $name(b: &mut Bencher) {
            let mut m = $maptype::with_capacity_and_hasher(SIZE, Default::default());
            b.iter(|| {
                m.clear();
                for i in ($keydist).take(SIZE) {
                    m.insert(i, i);
                }
                black_box(&mut m);
            })
        }
    };
}

bench_suite!(
    bench_insert,
    insert_ahash_serial,
    insert_std_serial,
    insert_ahash_highbits,
    insert_std_highbits,
    insert_ahash_random,
    insert_std_random
);

macro_rules! bench_insert_erase {
    ($name:ident, $maptype:ident, $keydist:expr) => {
        #[bench]
        fn $name(b: &mut Bencher) {
            let mut base = $maptype::default();
            for i in ($keydist).take(SIZE) {
                base.insert(i, i);
            }
            let skip = $keydist.skip(SIZE);
            b.iter(|| {
                let mut m = base.clone();
                let mut add_iter = skip.clone();
                let mut remove_iter = $keydist;
                // While keeping the size constant,
                // replace the first keydist with the second.
                for (add, remove) in (&mut add_iter).zip(&mut remove_iter).take(SIZE) {
                    m.insert(add, add);
                    black_box(m.remove(&remove));
                }
                black_box(m);
            })
        }
    };
}

bench_suite!(
    bench_insert_erase,
    insert_erase_ahash_serial,
    insert_erase_std_serial,
    insert_erase_ahash_highbits,
    insert_erase_std_highbits,
    insert_erase_ahash_random,
    insert_erase_std_random
);

macro_rules! bench_lookup {
    ($name:ident, $maptype:ident, $keydist:expr) => {
        #[bench]
        fn $name(b: &mut Bencher) {
            let mut m = $maptype::default();
            for i in $keydist.take(SIZE) {
                m.insert(i, i);
            }

            b.iter(|| {
                for i in $keydist.take(SIZE) {
                    black_box(m.get(&i));
                }
            })
        }
    };
}

bench_suite!(
    bench_lookup,
    lookup_ahash_serial,
    lookup_std_serial,
    lookup_ahash_highbits,
    lookup_std_highbits,
    lookup_ahash_random,
    lookup_std_random
);

macro_rules! bench_lookup_fail {
    ($name:ident, $maptype:ident, $keydist:expr) => {
        #[bench]
        fn $name(b: &mut Bencher) {
            let mut m = $maptype::default();
            let mut iter = $keydist;
            for i in (&mut iter).take(SIZE) {
                m.insert(i, i);
            }

            b.iter(|| {
                for i in (&mut iter).take(SIZE) {
                    black_box(m.get(&i));
                }
            })
        }
    };
}

bench_suite!(
    bench_lookup_fail,
    lookup_fail_ahash_serial,
    lookup_fail_std_serial,
    lookup_fail_ahash_highbits,
    lookup_fail_std_highbits,
    lookup_fail_ahash_random,
    lookup_fail_std_random
);

macro_rules! bench_iter {
    ($name:ident, $maptype:ident, $keydist:expr) => {
        #[bench]
        fn $name(b: &mut Bencher) {
            let mut m = $maptype::default();
            for i in ($keydist).take(SIZE) {
                m.insert(i, i);
            }

            b.iter(|| {
                for i in &m {
                    black_box(i);
                }
            })
        }
    };
}

bench_suite!(
    bench_iter,
    iter_ahash_serial,
    iter_std_serial,
    iter_ahash_highbits,
    iter_std_highbits,
    iter_ahash_random,
    iter_std_random
);
<|endoftext|>


<commit_before>use std::slice::Iter;
use std::vec::Vec;

use analysis::imports::Imports;
use analysis::ref_mode::RefMode;
use env::Env;
use library::*;
use super::conversion_type::ConversionType;
use super::rust_type::parameter_rust_type;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Mode {
    None,
    Normal,
    Optional,
}

impl Default for Mode {
    fn default() -> Mode {
        Mode::None
    }
}

#[derive(Debug, Default)]
pub struct Info {
    pub mode: Mode,
    pub params: Vec<Parameter>,
}

impl Info {
    pub fn is_empty(&self) -> bool {
        self.mode == Mode::None
    }

    pub fn iter(&self) -> Iter<Parameter> {
        self.params.iter()
    }

    pub fn len(&self) -> usize {
        self.params.len()
    }
}

pub fn analyze(env: &Env, type_: &Function) -> (Info, bool) {
    let mut info: Info = Default::default();
    let mut unsupported_outs = false;

    if type_.throws {
        //TODO: throwable functions
        return (info, true);
    } else if type_.ret.typ == TypeId::tid_none() {
        info.mode = Mode::Normal;
    } else if type_.ret.typ == TypeId::tid_bool() {
        info.mode = Mode::Optional;
    } else {
        return (info, false);
    }

    for par in &type_.parameters {
        if par.direction != ParameterDirection::Out { continue; }
        if can_as_return(env, par) {
            info.params.push(par.clone());
        } else {
            unsupported_outs = true;
        }
    }

    if info.params.is_empty() { info.mode = Mode::None }

    (info, unsupported_outs)
}

pub fn analyze_imports(env: &Env, func: &Function, imports: &mut Imports) {
    for par in &func.parameters {
        if par.direction == ParameterDirection::Out && !par.caller_allocates {
            match *env.library.type_(par.typ) {
                Type::Fundamental(..) |
                    Type::Bitfield(..) |
                    Type::Enumeration(..) => imports.add("std::mem".into(), func.version),
                _ => imports.add("std::ptr".into(), func.version),
            }
        }
    }
}

fn can_as_return(env: &Env, par: &Parameter) -> bool {
    use super::conversion_type::ConversionType::*;
    match ConversionType::of(&env.library, par.typ) {
        Direct => true,
        Scalar => true,
        Pointer => parameter_rust_type(env, par.typ, ParameterDirection::Out, Nullable(false), RefMode::None).is_ok(),
        Unknown => false,
    }
}
<commit_msg>Fix argument name<commit_after>use std::slice::Iter;
use std::vec::Vec;

use analysis::imports::Imports;
use analysis::ref_mode::RefMode;
use env::Env;
use library::*;
use super::conversion_type::ConversionType;
use super::rust_type::parameter_rust_type;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Mode {
    None,
    Normal,
    Optional,
}

impl Default for Mode {
    fn default() -> Mode {
        Mode::None
    }
}

#[derive(Debug, Default)]
pub struct Info {
    pub mode: Mode,
    pub params: Vec<Parameter>,
}

impl Info {
    pub fn is_empty(&self) -> bool {
        self.mode == Mode::None
    }

    pub fn iter(&self) -> Iter<Parameter> {
        self.params.iter()
    }

    pub fn len(&self) -> usize {
        self.params.len()
    }
}

pub fn analyze(env: &Env, func: &Function) -> (Info, bool) {
    let mut info: Info = Default::default();
    let mut unsupported_outs = false;

    if func.throws {
        //TODO: throwable functions
        return (info, true);
    } else if func.ret.typ == TypeId::tid_none() {
        info.mode = Mode::Normal;
    } else if func.ret.typ == TypeId::tid_bool() {
        info.mode = Mode::Optional;
    } else {
        return (info, false);
    }

    for par in &func.parameters {
        if par.direction != ParameterDirection::Out { continue; }
        if can_as_return(env, par) {
            info.params.push(par.clone());
        } else {
            unsupported_outs = true;
        }
    }

    if info.params.is_empty() { info.mode = Mode::None }

    (info, unsupported_outs)
}

pub fn analyze_imports(env: &Env, func: &Function, imports: &mut Imports) {
    for par in &func.parameters {
        if par.direction == ParameterDirection::Out && !par.caller_allocates {
            match *env.library.type_(par.typ) {
                Type::Fundamental(..) |
                    Type::Bitfield(..) |
                    Type::Enumeration(..) => imports.add("std::mem".into(), func.version),
                _ => imports.add("std::ptr".into(), func.version),
            }
        }
    }
}

fn can_as_return(env: &Env, par: &Parameter) -> bool {
    use super::conversion_type::ConversionType::*;
    match ConversionType::of(&env.library, par.typ) {
        Direct => true,
        Scalar => true,
        Pointer => parameter_rust_type(env, par.typ, ParameterDirection::Out, Nullable(false), RefMode::None).is_ok(),
        Unknown => false,
    }
}
<|endoftext|>


<commit_before>use syntax::{codemap};
use syntax::ext::base;
use syntax::ext::build::AstBuilder;
use syntax::util::small_vector;
use std::ascii::AsciiExt;

use super::super::helpers;

impl super::super::Generator<()> for super::MigrationState {
    fn generate<'a>(self, sp: codemap::Span, cx: &mut base::ExtCtxt, _: ()) -> Box<base::MacResult + 'a> {

        let pathes = ::std::fs::read_dir(&self.path).unwrap();
        let mut migrations = vec![];

        let path_checker = regex!(r"^_(\d{12})");
        let upcaser = regex!(r"_([a-z])");

        for path in pathes {
            let path = path.unwrap().path();

            let filestem = path.into_os_string().into_string().unwrap();
            let captures = path_checker.captures(&filestem[..]);

            if captures.is_none() { continue };

            let captures = captures.expect("Mailformed migration name");
            let tm = captures.at(1).expect("Timestamp must exists");
            let version: u64 = tm.parse().ok().expect("Timestamp must be valid u64");
            let name = filestem.replace(captures.at(0).unwrap(), "");

            let name = upcaser.replace_all(&name, |caps: &::regex::Captures| {
                caps.at(1).unwrap().to_ascii_uppercase()
            });

            migrations.push(format!("({}, {}, {})", filestem.to_string(), version, name.to_string()));
        }

        let macro_body = migrations.connect(", ");

        let mut impls = vec![];
        impls.push(helpers::generate_macro_invocation(cx, "migrations", macro_body, sp));
        base::MacEager::items(small_vector::SmallVector::many(impls))

    }
}
<commit_msg>fix(migration): Use file stem in migration generator<commit_after>use syntax::{codemap};
use syntax::ext::base;
use syntax::ext::build::AstBuilder;
use syntax::util::small_vector;
use std::ascii::AsciiExt;

use super::super::helpers;

impl super::super::Generator<()> for super::MigrationState {
    fn generate<'a>(self, sp: codemap::Span, cx: &mut base::ExtCtxt, _: ()) -> Box<base::MacResult + 'a> {

        let pathes = ::std::fs::read_dir(&self.path).unwrap();
        let mut migrations = vec![];

        let path_checker = regex!(r"^_(\d{12})");
        let upcaser = regex!(r"_([a-z])");

        for path in pathes {
            let path = path.unwrap().path();

            let filestem = match path.file_stem() {
                Some(f) => f.to_str().unwrap(),
                None => { continue }
            };
            let captures = path_checker.captures(&filestem[..]);

            if captures.is_none() { continue };

            let captures = captures.expect("Mailformed migration name");
            let tm = captures.at(1).expect("Timestamp must exists");
            let version: u64 = tm.parse().ok().expect("Timestamp must be valid u64");
            let name = filestem.replace(captures.at(0).unwrap(), "");

            let name = upcaser.replace_all(&name, |caps: &::regex::Captures| {
                caps.at(1).unwrap().to_ascii_uppercase()
            });

            migrations.push(format!("({}, {}, {})", filestem.to_string(), version, name.to_string()));
        }

        let macro_body = migrations.connect(", ");

        let mut impls = vec![];
        impls.push(helpers::generate_macro_invocation(cx, "migrations", macro_body, sp));
        base::MacEager::items(small_vector::SmallVector::many(impls))

    }
}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::ops::Not;

use cranelift_codegen::cursor::{Cursor, FuncCursor};
use cranelift_codegen::ir::{Opcode, InstructionData, ValueDef};
use cranelift_codegen::ir::immediates::Offset32;

use crate::prelude::*;

/// Workaround for `StackSlot` not implementing `Ord`.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
struct OrdStackSlot(StackSlot);

impl PartialOrd for OrdStackSlot {
    fn partial_cmp(&self, rhs: &Self) -> Option<std::cmp::Ordering> {
        self.0.as_u32().partial_cmp(&rhs.0.as_u32())
    }
}

impl Ord for OrdStackSlot {
    fn cmp(&self, rhs: &Self) -> std::cmp::Ordering {
        self.0.as_u32().cmp(&rhs.0.as_u32())
    }
}

#[derive(Debug, Default)]
struct StackSlotUsage {
    stack_addr: HashSet<Inst>,
    stack_load: HashSet<Inst>,
    stack_store: HashSet<Inst>,
}

pub(super) fn optimize_function(
    func: &mut Function,
    clif_comments: &mut crate::pretty_clif::CommentWriter,
    name: String, // FIXME remove
) {
    combine_stack_addr_with_load_store(func);

    // Record all stack_addr, stack_load and stack_store instructions.
    let mut stack_slot_usage_map = BTreeMap::<OrdStackSlot, StackSlotUsage>::new();

    let mut cursor = FuncCursor::new(func);
    while let Some(_ebb) = cursor.next_ebb() {
        while let Some(inst) = cursor.next_inst() {
            match cursor.func.dfg[inst] {
                InstructionData::StackLoad {
                    opcode: Opcode::StackAddr,
                    stack_slot,
                    offset: _,
                } => {
                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_addr.insert(inst);
                }
                InstructionData::StackLoad {
                    opcode: Opcode::StackLoad,
                    stack_slot,
                    offset: _,
                } => {
                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_load.insert(inst);
                }
                InstructionData::StackStore {
                    opcode: Opcode::StackStore,
                    arg: _,
                    stack_slot,
                    offset: _,
                } => {
                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_store.insert(inst);
                }
                _ => {}
            }
        }
    }

    println!("{}:\nstack slot usage: {:?}", name, stack_slot_usage_map);

    remove_unused_stack_addr_and_stack_load(func, &mut stack_slot_usage_map);

    println!("stack slot usage (after): {:?}", stack_slot_usage_map);

    for (stack_slot, users) in stack_slot_usage_map.iter_mut() {
        if users.stack_addr.is_empty().not() || (users.stack_load.is_empty().not() && users.stack_store.is_empty().not()) {
            continue;
        }

        if users.stack_load.is_empty().not() {
            println!("[{}] [BUG?] Reading uninitialized memory", name);
        } else {
            // Stored value never read; just remove reads.
            for user in users.stack_store.drain() {
                println!("[{}] Remove dead stack store {} of {}", name, user, stack_slot.0);
                func.dfg.replace(user).nop();
            }
        }
    }

    println!();
}

fn combine_stack_addr_with_load_store(func: &mut Function) {
    // Turn load and store into stack_load and stack_store when possible.
    let mut cursor = FuncCursor::new(func);
    while let Some(_ebb) = cursor.next_ebb() {
        while let Some(inst) = cursor.next_inst() {
            match cursor.func.dfg[inst] {
                InstructionData::Load { opcode: Opcode::Load, arg: addr, flags: _, offset } => {
                    if cursor.func.dfg.ctrl_typevar(inst) == types::I128 || cursor.func.dfg.ctrl_typevar(inst).is_vector() {
                        continue; // WORKAROUD: stack_load.i128 not yet implemented
                    }
                    if let Some((stack_slot, stack_addr_offset)) = try_get_stack_slot_and_offset_for_addr(cursor.func, addr) {
                        if let Some(combined_offset) = offset.try_add_i64(stack_addr_offset.into()) {
                            let ty = cursor.func.dfg.ctrl_typevar(inst);
                            cursor.func.dfg.replace(inst).stack_load(ty, stack_slot, combined_offset);
                        }
                    }
                }
                InstructionData::Store { opcode: Opcode::Store, args: [value, addr], flags: _, offset } => {
                    if cursor.func.dfg.ctrl_typevar(inst) == types::I128 || cursor.func.dfg.ctrl_typevar(inst).is_vector() {
                        continue; // WORKAROUND: stack_store.i128 not yet implemented
                    }
                    if let Some((stack_slot, stack_addr_offset)) = try_get_stack_slot_and_offset_for_addr(cursor.func, addr) {
                        if let Some(combined_offset) = offset.try_add_i64(stack_addr_offset.into()) {
                            cursor.func.dfg.replace(inst).stack_store(value, stack_slot, combined_offset);
                        }
                    }
                }
                _ => {}
            }
        }
    }
}

fn remove_unused_stack_addr_and_stack_load(func: &mut Function, stack_slot_usage_map: &mut BTreeMap<OrdStackSlot, StackSlotUsage>) {
    // FIXME incrementally rebuild on each call?
    let mut stack_addr_load_insts_users = HashMap::<Inst, HashSet<Inst>>::new();

    let mut cursor = FuncCursor::new(func);
    while let Some(_ebb) = cursor.next_ebb() {
        while let Some(inst) = cursor.next_inst() {
            for &arg in cursor.func.dfg.inst_args(inst) {
                if let ValueDef::Result(arg_origin, 0) = cursor.func.dfg.value_def(arg) {
                    match cursor.func.dfg[arg_origin].opcode() {
                        Opcode::StackAddr | Opcode::StackLoad => {
                            stack_addr_load_insts_users.entry(arg_origin).or_insert_with(HashSet::new).insert(inst);
                        }
                        _ => {}
                    }
                }
            }
        }
    }

    println!("stack_addr/stack_load users: {:?}", stack_addr_load_insts_users);

    for inst in stack_addr_load_insts_users.keys() {
        let mut is_recorded_stack_addr_or_stack_load = false;
        for stack_slot_users in stack_slot_usage_map.values() {
            is_recorded_stack_addr_or_stack_load |= stack_slot_users.stack_addr.contains(inst) || stack_slot_users.stack_load.contains(inst);
        }
        assert!(is_recorded_stack_addr_or_stack_load);
    }

    // Replace all unused stack_addr and stack_load instructions with nop.
    for stack_slot_users in stack_slot_usage_map.values_mut() {
        // FIXME remove clone
        for &inst in stack_slot_users.stack_addr.clone().iter() {
            if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
                println!("Removing unused stack_addr {}", inst);
                func.dfg.detach_results(inst);
                func.dfg.replace(inst).nop();
                stack_slot_users.stack_addr.remove(&inst);
            }
        }

        for &inst in stack_slot_users.stack_load.clone().iter() {
            if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
                println!("Removing unused stack_addr {}", inst);
                func.dfg.detach_results(inst);
                func.dfg.replace(inst).nop();
                stack_slot_users.stack_load.remove(&inst);
            }
        }
    }
}

fn try_get_stack_slot_and_offset_for_addr(func: &Function, addr: Value) -> Option<(StackSlot, Offset32)> {
    if let ValueDef::Result(addr_inst, 0) = func.dfg.value_def(addr) {
        if let InstructionData::StackLoad {
            opcode: Opcode::StackAddr,
            stack_slot,
            offset,
        } = func.dfg[addr_inst] {
            return Some((stack_slot, offset));
        }
    }
    None
}
<commit_msg>Refactor stack_store removal<commit_after>use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::ops::Not;

use cranelift_codegen::cursor::{Cursor, FuncCursor};
use cranelift_codegen::ir::{Opcode, InstructionData, ValueDef};
use cranelift_codegen::ir::immediates::Offset32;

use crate::prelude::*;

/// Workaround for `StackSlot` not implementing `Ord`.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
struct OrdStackSlot(StackSlot);

impl PartialOrd for OrdStackSlot {
    fn partial_cmp(&self, rhs: &Self) -> Option<std::cmp::Ordering> {
        self.0.as_u32().partial_cmp(&rhs.0.as_u32())
    }
}

impl Ord for OrdStackSlot {
    fn cmp(&self, rhs: &Self) -> std::cmp::Ordering {
        self.0.as_u32().cmp(&rhs.0.as_u32())
    }
}

#[derive(Debug, Default)]
struct StackSlotUsage {
    stack_addr: HashSet<Inst>,
    stack_load: HashSet<Inst>,
    stack_store: HashSet<Inst>,
}

pub(super) fn optimize_function(
    func: &mut Function,
    clif_comments: &mut crate::pretty_clif::CommentWriter,
    name: String, // FIXME remove
) {
    combine_stack_addr_with_load_store(func);

    // Record all stack_addr, stack_load and stack_store instructions.
    let mut stack_slot_usage_map = BTreeMap::<OrdStackSlot, StackSlotUsage>::new();

    let mut cursor = FuncCursor::new(func);
    while let Some(_ebb) = cursor.next_ebb() {
        while let Some(inst) = cursor.next_inst() {
            match cursor.func.dfg[inst] {
                InstructionData::StackLoad {
                    opcode: Opcode::StackAddr,
                    stack_slot,
                    offset: _,
                } => {
                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_addr.insert(inst);
                }
                InstructionData::StackLoad {
                    opcode: Opcode::StackLoad,
                    stack_slot,
                    offset: _,
                } => {
                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_load.insert(inst);
                }
                InstructionData::StackStore {
                    opcode: Opcode::StackStore,
                    arg: _,
                    stack_slot,
                    offset: _,
                } => {
                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_store.insert(inst);
                }
                _ => {}
            }
        }
    }

    println!("{}:\nstack slot usage: {:?}", name, stack_slot_usage_map);

    remove_unused_stack_addr_and_stack_load(func, &mut stack_slot_usage_map);

    println!("stack slot usage (after): {:?}", stack_slot_usage_map);

    for (stack_slot, users) in stack_slot_usage_map.iter_mut() {
        if users.stack_addr.is_empty().not() {
            // Stack addr leaked; there may be unknown loads and stores.
            // FIXME use stacked borrows to optimize
            continue;
        }

        let is_loaded = users.stack_load.is_empty().not();
        let is_stored = users.stack_store.is_empty().not();
        match (is_loaded, is_stored) {
            (true, true) => {} // FIXME perform store to load optimization
            (true, false) => println!("[{}] [BUG?] Reading uninitialized memory", name),
            (false, _) => {
                // Never loaded; can safely remove all stores and the stack slot.
                for user in users.stack_store.drain() {
                    println!("[{}] Remove dead stack store {} of {}", name, user, stack_slot.0);
                    func.dfg.replace(user).nop();
                }

                // FIXME make stack_slot zero sized.
            }
        }
    }

    println!();
}

fn combine_stack_addr_with_load_store(func: &mut Function) {
    // Turn load and store into stack_load and stack_store when possible.
    let mut cursor = FuncCursor::new(func);
    while let Some(_ebb) = cursor.next_ebb() {
        while let Some(inst) = cursor.next_inst() {
            match cursor.func.dfg[inst] {
                InstructionData::Load { opcode: Opcode::Load, arg: addr, flags: _, offset } => {
                    if cursor.func.dfg.ctrl_typevar(inst) == types::I128 || cursor.func.dfg.ctrl_typevar(inst).is_vector() {
                        continue; // WORKAROUD: stack_load.i128 not yet implemented
                    }
                    if let Some((stack_slot, stack_addr_offset)) = try_get_stack_slot_and_offset_for_addr(cursor.func, addr) {
                        if let Some(combined_offset) = offset.try_add_i64(stack_addr_offset.into()) {
                            let ty = cursor.func.dfg.ctrl_typevar(inst);
                            cursor.func.dfg.replace(inst).stack_load(ty, stack_slot, combined_offset);
                        }
                    }
                }
                InstructionData::Store { opcode: Opcode::Store, args: [value, addr], flags: _, offset } => {
                    if cursor.func.dfg.ctrl_typevar(inst) == types::I128 || cursor.func.dfg.ctrl_typevar(inst).is_vector() {
                        continue; // WORKAROUND: stack_store.i128 not yet implemented
                    }
                    if let Some((stack_slot, stack_addr_offset)) = try_get_stack_slot_and_offset_for_addr(cursor.func, addr) {
                        if let Some(combined_offset) = offset.try_add_i64(stack_addr_offset.into()) {
                            cursor.func.dfg.replace(inst).stack_store(value, stack_slot, combined_offset);
                        }
                    }
                }
                _ => {}
            }
        }
    }
}

fn remove_unused_stack_addr_and_stack_load(func: &mut Function, stack_slot_usage_map: &mut BTreeMap<OrdStackSlot, StackSlotUsage>) {
    // FIXME incrementally rebuild on each call?
    let mut stack_addr_load_insts_users = HashMap::<Inst, HashSet<Inst>>::new();

    let mut cursor = FuncCursor::new(func);
    while let Some(_ebb) = cursor.next_ebb() {
        while let Some(inst) = cursor.next_inst() {
            for &arg in cursor.func.dfg.inst_args(inst) {
                if let ValueDef::Result(arg_origin, 0) = cursor.func.dfg.value_def(arg) {
                    match cursor.func.dfg[arg_origin].opcode() {
                        Opcode::StackAddr | Opcode::StackLoad => {
                            stack_addr_load_insts_users.entry(arg_origin).or_insert_with(HashSet::new).insert(inst);
                        }
                        _ => {}
                    }
                }
            }
        }
    }

    println!("stack_addr/stack_load users: {:?}", stack_addr_load_insts_users);

    for inst in stack_addr_load_insts_users.keys() {
        let mut is_recorded_stack_addr_or_stack_load = false;
        for stack_slot_users in stack_slot_usage_map.values() {
            is_recorded_stack_addr_or_stack_load |= stack_slot_users.stack_addr.contains(inst) || stack_slot_users.stack_load.contains(inst);
        }
        assert!(is_recorded_stack_addr_or_stack_load);
    }

    // Replace all unused stack_addr and stack_load instructions with nop.
    for stack_slot_users in stack_slot_usage_map.values_mut() {
        // FIXME remove clone
        for &inst in stack_slot_users.stack_addr.clone().iter() {
            if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
                println!("Removing unused stack_addr {}", inst);
                func.dfg.detach_results(inst);
                func.dfg.replace(inst).nop();
                stack_slot_users.stack_addr.remove(&inst);
            }
        }

        for &inst in stack_slot_users.stack_load.clone().iter() {
            if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
                println!("Removing unused stack_addr {}", inst);
                func.dfg.detach_results(inst);
                func.dfg.replace(inst).nop();
                stack_slot_users.stack_load.remove(&inst);
            }
        }
    }
}

fn try_get_stack_slot_and_offset_for_addr(func: &Function, addr: Value) -> Option<(StackSlot, Offset32)> {
    if let ValueDef::Result(addr_inst, 0) = func.dfg.value_def(addr) {
        if let InstructionData::StackLoad {
            opcode: Opcode::StackAddr,
            stack_slot,
            offset,
        } = func.dfg[addr_inst] {
            return Some((stack_slot, offset));
        }
    }
    None
}
<|endoftext|>


<commit_before>
<commit_msg>Delete a.rs<commit_after><|endoftext|>


<commit_before>use std::io::Read;
use azure::core;
use azure::core::errors::{self, AzureError};
use azure::storage::client::Client;
use hyper::client::response::Response;
use hyper::status::StatusCode;
use rustc_serialize::{Decodable, Encodable, json};

const TABLE_SUFFIX: &'static str = "table.core.windows.net";
const TABLE_TABLES: &'static str = "TABLES";

pub struct TableClient {
    client: Client,
}

impl TableClient {
    pub fn new(client: Client) -> Self {
        TableClient { client: client }
    }

    pub fn list(&self) -> Result<Vec<String>, core::errors::AzureError> {
        Ok(self.query_range_entity(TABLE_TABLES, None)?
               .into_iter()
               .map(|x: TableEntry| x.TableName)
               .collect())
    }

    pub fn create_if_not_exists<T: Into<String>>(&self,
                                                 table_name: T)
                                                 -> Result<(), core::errors::AzureError> {
        let ref body = json::encode(&TableEntry { TableName: table_name.into() }).unwrap();
        let mut response = try!(self.do_request(TABLE_TABLES, core::HTTPMethod::Post, Some(body)));
        // TODO: Here treats conflict as existed, but could be reserved name, such as 'Tables',
        // should check table existence directly
        if !(StatusCode::Created == response.status || StatusCode::Conflict == response.status) {
            try!(errors::check_status(&mut response, StatusCode::Created));
        }

        Ok(())
    }

    pub fn insert_entity<T: Encodable>(&self,
                                       table_name: &str,
                                       partition_key: &str,
                                       row_key: &str,
                                       entity: &T)
                                       -> Result<(), core::errors::AzureError> {
        let ref body = json::encode(entity).unwrap();
        if !body.starts_with("{") {
            return Err(AzureError::InputParametersError("body not valid.".to_owned()));
        };

        let ref post_body = format!(r#"{{"PartitionKey":"{}","RowKey":"{}",{}"#,
                                    partition_key,
                                    row_key,
                                    &body[1..]);

        let mut resp = try!(self.do_request(table_name, core::HTTPMethod::Post, Some(post_body)));
        try!(errors::check_status(&mut resp, StatusCode::Created));
        Ok(())
    }

    pub fn update_entity<T: Encodable>(&self,
                                       table_name: &str,
                                       partition_key: &str,
                                       row_key: &str,
                                       entity: &T)
                                       -> Result<(), core::errors::AzureError> {
        let ref body = json::encode(entity).unwrap();
        if !body.starts_with("{") {
            return Err(AzureError::InputParametersError("body not valid.".to_owned()));
        };

        let ref send_body = format!(r#"{{"PartitionKey":"{}","RowKey":"{}",{}"#,
                                    partition_key,
                                    row_key,
                                    &body[1..]);

        let ref path = format!("{}(PartitionKey='{}',RowKey='{}')",
                               table_name,
                               partition_key,
                               row_key);

        let mut resp = try!(self.do_request(path, core::HTTPMethod::Put, Some(send_body)));
        try!(errors::check_status(&mut resp, StatusCode::NoContent));
        Ok(())
    }

    pub fn get_entity<T: Decodable>(&self,
                                    table_name: &str,
                                    partition_key: &str,
                                    row_key: &str)
                                    -> Result<Option<T>, core::errors::AzureError> {
        let ref path = format!("{}(PartitionKey='{}',RowKey='{}')",
                               table_name,
                               partition_key,
                               row_key);
        let mut response = try!(self.do_request(path, core::HTTPMethod::Get, None));
        if StatusCode::NotFound == response.status {
            return Ok(None);
        }
        try!(errors::check_status(&mut response, StatusCode::Ok));
        let ref body = try!(get_response_body(&mut response));
        Ok(json::decode(body).unwrap())
    }

    pub fn query_range_entity<T: Decodable>(&self,
                                            path: &str,
                                            query: Option<&str>)
                                            -> Result<Vec<T>, core::errors::AzureError> {
        let ref path = format!("{}?{}",
                               path,
                               match query {
                                   Some(clause) => clause,
                                   None => "",
                               });
        let mut response = try!(self.do_request(path, core::HTTPMethod::Get, None));
        try!(errors::check_status(&mut response, StatusCode::Ok));
        let ref body = try!(get_response_body(&mut response));
        let ec: EntryCollection<T> = json::decode(body).unwrap();
        Ok(ec.value)
    }

    pub fn do_request(&self,
                      segment: &str,
                      method: core::HTTPMethod,
                      request_str: Option<&str>)
                      -> Result<Response, core::errors::AzureError> {
        let client = &self.client;
        let uri = format!("{}://{}.{}/{}",
                          client.auth_scheme(),
                          client.account(),
                          TABLE_SUFFIX,
                          segment);
        trace!("{:?} {}", method, uri);
        if let Some(ref body) = request_str {
            trace!("Request: {}", body);
        }

        let resp = try!(client.perform_table_request(&uri, method, request_str));
        trace!("Response status: {:?}", resp.status);
        Ok(resp)
        // if StatusCode::NotFound == resp.status {
        //     return Ok(None);
        // }

        // let cs = errors::check_status(&mut resp, expected_status_code);
        // if cs.is_err() {
        //     trace!("Err Response:{:?}", cs);
        //     try!(cs);
        // }

        // let mut resp_s = String::new();
        // try!(resp.read_to_string(&mut resp_s));

        // trace!("Response: {}", resp_s);
        // Ok(Some(resp_s))
    }
}

fn get_response_body(resp: &mut Response) -> Result<String, core::errors::AzureError> {
    let mut body = String::new();
    try!(resp.read_to_string(&mut body));
    Ok(body)
}

#[allow(non_snake_case)]
#[derive(RustcEncodable, RustcDecodable)]
struct TableEntry {
    TableName: String,
}

#[derive(RustcDecodable)]
struct EntryCollection<T> {
    value: Vec<T>,
}
<commit_msg>Add log.<commit_after>use std::io::Read;
use azure::core;
use azure::core::errors::{self, AzureError};
use azure::storage::client::Client;
use hyper::client::response::Response;
use hyper::status::StatusCode;
use rustc_serialize::{Decodable, Encodable, json};

const TABLE_SUFFIX: &'static str = "table.core.windows.net";
const TABLE_TABLES: &'static str = "TABLES";

pub struct TableClient {
    client: Client,
}

impl TableClient {
    pub fn new(client: Client) -> Self {
        TableClient { client: client }
    }

    pub fn list(&self) -> Result<Vec<String>, core::errors::AzureError> {
        Ok(self.query_range_entity(TABLE_TABLES, None)?
               .into_iter()
               .map(|x: TableEntry| x.TableName)
               .collect())
    }

    pub fn create_if_not_exists<T: Into<String>>(&self,
                                                 table_name: T)
                                                 -> Result<(), core::errors::AzureError> {
        let ref body = json::encode(&TableEntry { TableName: table_name.into() }).unwrap();
        let mut response = try!(self.do_request(TABLE_TABLES, core::HTTPMethod::Post, Some(body)));
        // TODO: Here treats conflict as existed, but could be reserved name, such as 'Tables',
        // should check table existence directly
        if !(StatusCode::Created == response.status || StatusCode::Conflict == response.status) {
            try!(errors::check_status(&mut response, StatusCode::Created));
        }

        Ok(())
    }

    pub fn insert_entity<T: Encodable>(&self,
                                       table_name: &str,
                                       partition_key: &str,
                                       row_key: &str,
                                       entity: &T)
                                       -> Result<(), core::errors::AzureError> {
        let ref body = json::encode(entity).unwrap();
        if !body.starts_with("{") {
            return Err(AzureError::InputParametersError("body not valid.".to_owned()));
        };

        let ref post_body = format!(r#"{{"PartitionKey":"{}","RowKey":"{}",{}"#,
                                    partition_key,
                                    row_key,
                                    &body[1..]);

        let mut resp = try!(self.do_request(table_name, core::HTTPMethod::Post, Some(post_body)));
        try!(errors::check_status(&mut resp, StatusCode::Created));
        Ok(())
    }

    pub fn update_entity<T: Encodable>(&self,
                                       table_name: &str,
                                       partition_key: &str,
                                       row_key: &str,
                                       entity: &T)
                                       -> Result<(), core::errors::AzureError> {
        let ref body = json::encode(entity).unwrap();
        if !body.starts_with("{") {
            return Err(AzureError::InputParametersError("body not valid.".to_owned()));
        };

        let ref send_body = format!(r#"{{"PartitionKey":"{}","RowKey":"{}",{}"#,
                                    partition_key,
                                    row_key,
                                    &body[1..]);

        let ref path = format!("{}(PartitionKey='{}',RowKey='{}')",
                               table_name,
                               partition_key,
                               row_key);

        let mut resp = try!(self.do_request(path, core::HTTPMethod::Put, Some(send_body)));
        try!(errors::check_status(&mut resp, StatusCode::NoContent));
        Ok(())
    }

    pub fn get_entity<T: Decodable>(&self,
                                    table_name: &str,
                                    partition_key: &str,
                                    row_key: &str)
                                    -> Result<Option<T>, core::errors::AzureError> {
        let ref path = format!("{}(PartitionKey='{}',RowKey='{}')",
                               table_name,
                               partition_key,
                               row_key);
        let mut response = try!(self.do_request(path, core::HTTPMethod::Get, None));
        if StatusCode::NotFound == response.status {
            return Ok(None);
        }
        try!(errors::check_status(&mut response, StatusCode::Ok));
        let ref body = try!(get_response_body(&mut response));
        Ok(json::decode(body).unwrap())
    }

    pub fn query_range_entity<T: Decodable>(&self,
                                            path: &str,
                                            query: Option<&str>)
                                            -> Result<Vec<T>, core::errors::AzureError> {
        let ref path = format!("{}?{}",
                               path,
                               match query {
                                   Some(clause) => clause,
                                   None => "",
                               });
        let mut response = try!(self.do_request(path, core::HTTPMethod::Get, None));
        try!(errors::check_status(&mut response, StatusCode::Ok));
        let ref body = try!(get_response_body(&mut response));
        let ec: EntryCollection<T> = json::decode(body).unwrap();
        Ok(ec.value)
    }

    pub fn do_request(&self,
                      segment: &str,
                      method: core::HTTPMethod,
                      request_str: Option<&str>)
                      -> Result<Response, core::errors::AzureError> {
        let client = &self.client;
        let uri = format!("{}://{}.{}/{}",
                          client.auth_scheme(),
                          client.account(),
                          TABLE_SUFFIX,
                          segment);
        trace!("{:?} {}", method, uri);
        if let Some(ref body) = request_str {
            trace!("Request: {}", body);
        }

        let resp = try!(client.perform_table_request(&uri, method, request_str));
        trace!("Response status: {:?}", resp.status);
        Ok(resp)
        // if StatusCode::NotFound == resp.status {
        //     return Ok(None);
        // }

        // let cs = errors::check_status(&mut resp, expected_status_code);
        // if cs.is_err() {
        //     trace!("Err Response:{:?}", cs);
        //     try!(cs);
        // }

        // let mut resp_s = String::new();
        // try!(resp.read_to_string(&mut resp_s));

        // trace!("Response: {}", resp_s);
        // Ok(Some(resp_s))
    }
}

fn get_response_body(resp: &mut Response) -> Result<String, core::errors::AzureError> {
    let mut body = String::new();
    try!(resp.read_to_string(&mut body));
    trace!("Response Body:{}", body);
    Ok(body)
}

#[allow(non_snake_case)]
#[derive(RustcEncodable, RustcDecodable)]
struct TableEntry {
    TableName: String,
}

#[derive(RustcDecodable)]
struct EntryCollection<T> {
    value: Vec<T>,
}
<|endoftext|>


<commit_before>use std::io::Read;
use azure::core;
use azure::core::errors::{self, AzureError};
use azure::storage::client::Client;
use hyper::status::StatusCode;
use rustc_serialize::{Decodable, Encodable, json};

const TABLE_SUFFIX: &'static str = "table.core.windows.net";

pub struct TableClient {
    client: Client,
}

impl TableClient {
    pub fn new(client: Client) -> Self {
        TableClient { client: client }
    }

    pub fn list(&self) -> Result<Vec<String>, core::errors::AzureError> {
        Ok(self.query_range_entity("Tables", None)?
               .into_iter()
               .map(|x: TableEntry| x.TableName)
               .collect())
    }

    pub fn insert_entity<T: Encodable>(&self,
                                       table_name: &str,
                                       partition_key: &str,
                                       row_key: &str,
                                       entity: &T)
                                       -> Result<(), core::errors::AzureError> {
        let ref body = json::encode(entity).unwrap();
        if !body.starts_with("{") {
            return Err(AzureError::InputParametersError("body not valid.".to_owned()));
        };

        let ref post_body = format!(r#"{{"PartitionKey":"{}","RowKey":"{}",{}"#,
                                    partition_key,
                                    row_key,
                                    &body[1..]);

        try!(self.do_request(table_name,
                             core::HTTPMethod::Post,
                             Some(post_body),
                             StatusCode::Created));
        Ok(())
    }

    pub fn get_entity<T: Decodable>(&self,
                                    table_name: &str,
                                    partition_key: &str,
                                    row_key: &str)
                                    -> Result<Option<T>, core::errors::AzureError> {
        let ref path = format!("{}(PartitionKey='{}',RowKey='{}')",
                               table_name,
                               partition_key,
                               row_key);
        Ok(self.do_request(path,
                           core::HTTPMethod::Get,
                           None,
                           StatusCode::Ok)?
               .map(|x| json::decode(x.as_str()).unwrap()))
    }

    pub fn query_range_entity<T: Decodable>(&self,
                                            path: &str,
                                            query: Option<&str>)
                                            -> Result<Vec<T>, core::errors::AzureError> {
        let ref path = format!("{}?{}",
                               path,
                               match query {
                                   Some(clause) => clause,
                                   None => "",
                               });
        let result =
            self.do_request(path, core::HTTPMethod::Get, None, StatusCode::Ok).map(|x| x.unwrap());
        let ec: EntryCollection<T> = json::decode(result?.as_str()).unwrap();
        Ok(ec.value)
    }

    pub fn do_request(&self,
                      segment: &str,
                      method: core::HTTPMethod,
                      request_str: Option<&str>,
                      expected_status_code: StatusCode)
                      -> Result<Option<String>, core::errors::AzureError> {
        let client = &self.client;
        if let Some(ref body) = request_str {
            trace!("Request: {:?} {}", method, body);
        }

        let uri = format!("{}://{}.{}/{}",
                          client.auth_scheme(),
                          client.account(),
                          TABLE_SUFFIX,
                          segment);
        trace!("uri:{}", uri);
        let mut resp = try!(client.perform_table_request(&uri, method, request_str));
        if StatusCode::NotFound == resp.status {
            return Ok(None);
        }

        let cs = errors::check_status(&mut resp, expected_status_code);
        trace!("cs check: {:?}", cs);
        if cs.is_err() {
            trace!("Err Response:{:?}", cs);
            try!(cs);
        }

        let mut resp_s = String::new();
        try!(resp.read_to_string(&mut resp_s));

        trace!("Response: {}", resp_s);
        Ok(Some(resp_s))
    }
}

#[allow(non_snake_case)]
#[derive(RustcDecodable)]
struct TableEntry {
    TableName: String,
}

#[derive(RustcDecodable)]
struct EntryCollection<T> {
    value: Vec<T>,
}
<commit_msg>Update log.<commit_after>use std::io::Read;
use azure::core;
use azure::core::errors::{self, AzureError};
use azure::storage::client::Client;
use hyper::status::StatusCode;
use rustc_serialize::{Decodable, Encodable, json};

const TABLE_SUFFIX: &'static str = "table.core.windows.net";

pub struct TableClient {
    client: Client,
}

impl TableClient {
    pub fn new(client: Client) -> Self {
        TableClient { client: client }
    }

    pub fn list(&self) -> Result<Vec<String>, core::errors::AzureError> {
        Ok(self.query_range_entity("Tables", None)?
               .into_iter()
               .map(|x: TableEntry| x.TableName)
               .collect())
    }

    pub fn insert_entity<T: Encodable>(&self,
                                       table_name: &str,
                                       partition_key: &str,
                                       row_key: &str,
                                       entity: &T)
                                       -> Result<(), core::errors::AzureError> {
        let ref body = json::encode(entity).unwrap();
        if !body.starts_with("{") {
            return Err(AzureError::InputParametersError("body not valid.".to_owned()));
        };

        let ref post_body = format!(r#"{{"PartitionKey":"{}","RowKey":"{}",{}"#,
                                    partition_key,
                                    row_key,
                                    &body[1..]);

        try!(self.do_request(table_name,
                             core::HTTPMethod::Post,
                             Some(post_body),
                             StatusCode::Created));
        Ok(())
    }

    pub fn get_entity<T: Decodable>(&self,
                                    table_name: &str,
                                    partition_key: &str,
                                    row_key: &str)
                                    -> Result<Option<T>, core::errors::AzureError> {
        let ref path = format!("{}(PartitionKey='{}',RowKey='{}')",
                               table_name,
                               partition_key,
                               row_key);
        Ok(self.do_request(path,
                           core::HTTPMethod::Get,
                           None,
                           StatusCode::Ok)?
               .map(|x| json::decode(x.as_str()).unwrap()))
    }

    pub fn query_range_entity<T: Decodable>(&self,
                                            path: &str,
                                            query: Option<&str>)
                                            -> Result<Vec<T>, core::errors::AzureError> {
        let ref path = format!("{}?{}",
                               path,
                               match query {
                                   Some(clause) => clause,
                                   None => "",
                               });
        let result =
            self.do_request(path, core::HTTPMethod::Get, None, StatusCode::Ok).map(|x| x.unwrap());
        let ec: EntryCollection<T> = json::decode(result?.as_str()).unwrap();
        Ok(ec.value)
    }

    pub fn do_request(&self,
                      segment: &str,
                      method: core::HTTPMethod,
                      request_str: Option<&str>,
                      expected_status_code: StatusCode)
                      -> Result<Option<String>, core::errors::AzureError> {
        let client = &self.client;
        if let Some(ref body) = request_str {
            trace!("Request: {}", body);
        }

        let uri = format!("{}://{}.{}/{}",
                          client.auth_scheme(),
                          client.account(),
                          TABLE_SUFFIX,
                          segment);
        trace!("{:?} {}", method, uri);
        let mut resp = try!(client.perform_table_request(&uri, method, request_str));
        if StatusCode::NotFound == resp.status {
            return Ok(None);
        }

        let cs = errors::check_status(&mut resp, expected_status_code);
        trace!("cs check: {:?}", cs);
        if cs.is_err() {
            trace!("Err Response:{:?}", cs);
            try!(cs);
        }

        let mut resp_s = String::new();
        try!(resp.read_to_string(&mut resp_s));

        trace!("Response: {}", resp_s);
        Ok(Some(resp_s))
    }
}

#[allow(non_snake_case)]
#[derive(RustcDecodable)]
struct TableEntry {
    TableName: String,
}

#[derive(RustcDecodable)]
struct EntryCollection<T> {
    value: Vec<T>,
}
<|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,
    },
};

pub mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || match iter.next() {
        Some(label) => std::str::from_utf8(label).unwrap_or("*"),
        None => "0",    // zero fill the missing labels
    };
    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(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(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 = check_name_in_proxy_list(acl, 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 check_name_in_proxy_list(acl, 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<Remote: upstream::Upstream>(
    acl: &Option<AccessControl>,
    local: &upstream::LocalUpstream,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool) {
    // 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 remote_response.is_none() => {
                if use_remote {
                    debug!("pick remote response (response): {:?}", response);
                    return (response, true);
                } else {
                    remote_response = Some(response);
                }
            }
            decision = &mut decider, if !use_remote => {
                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;
                }
            }
            else => unreachable!(),
        }
    }
}

/// 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();
    // 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 remote_upstream = remote_upstream.clone();
        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(), context.local_dns(), 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");
                    }
                }
            }
        });
    }
}
<commit_msg>Move DnsRelay into a struct<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, OpCode, 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,
    },
};

pub mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || match iter.next() {
        Some(label) => std::str::from_utf8(label).unwrap_or("*"),
        None => "0",    // zero fill the missing labels
    };
    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(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(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 = check_name_in_proxy_list(acl, 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 check_name_in_proxy_list(acl, 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!()
    }
}

struct DnsRelay<Remote: upstream::Upstream> {
    context: SharedContext,
    remote_upstream: Remote,
}

impl<Remote: upstream::Upstream> DnsRelay<Remote> {
    async fn acl_lookup(&self, query: &Query) -> (io::Result<Message>, bool) {
        let acl = self.context.acl();
        let local = self.context.local_dns();
        let remote = &self.remote_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 remote_response.is_none() => {
                    if use_remote {
                        debug!("pick remote response (response): {:?}", response);
                        return (response, true);
                    } else {
                        remote_response = Some(response);
                    }
                }
                decision = &mut decider, if !use_remote => {
                    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;
                    }
                }
                else => unreachable!(),
            }
        }
    }

    async fn resolve(&self, request: Message) -> Message {
        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.op_code() != OpCode::Query || request.message_type() != MessageType::Query {
            message.set_response_code(ResponseCode::NotImp);
        } else if request.query_count() > 0 {
            let (r, forward) = self.acl_lookup(&request.queries()[0]).await;
            if let Ok(result) = r {
                for rec in result.answers() {
                    debug!("dns answer: {:?}", rec);
                    match rec.rdata() {
                        RData::A(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward),
                        RData::AAAA(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward),
                        _ => (),
                    }
                }
                message = result;
                message.set_id(request.id());
            } else {
                message.set_response_code(ResponseCode::ServFail);
            }
        }
        message
    }
}

/// 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();
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let relay = Arc::new(DnsRelay {
        context: context.clone(),
        remote_upstream: 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 relay = relay.clone();
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let message = relay.resolve(request).await;
            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|>