<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The implementations of the `Standard` distribution for other built-in types.

use core::char;
use core::num::Wrapping;

use crate::distributions::{Distribution, Standard, Uniform};
use crate::Rng;

#[cfg(feature = "serde1")]
use serde::{Serialize, Deserialize};
#[cfg(feature = "min_const_gen")]
use std::mem::{self, MaybeUninit};

// ----- Sampling distributions -----

/// Sample a `u8`, uniformly distributed over ASCII letters and numbers:
/// a-z, A-Z and 0-9.
///
/// # Example
///
/// ```
/// use std::iter;
/// use rand::{Rng, thread_rng};
/// use rand::distributions::Alphanumeric;
///
/// let mut rng = thread_rng();
/// let chars: String = iter::repeat(())
///         .map(|()| rng.sample(Alphanumeric))
///         .map(char::from)
///         .take(7)
///         .collect();
/// println!("Random chars: {}", chars);
/// ```
///
/// # Passwords
///
/// Users sometimes ask whether it is safe to use a string of random characters
/// as a password. In principle, all RNGs in Rand implementing `CryptoRng` are
/// suitable as a source of randomness for generating passwords (if they are
/// properly seeded), but it is more conservative to only use randomness
/// directly from the operating system via the `getrandom` crate, or the
/// corresponding bindings of a crypto library.
///
/// When generating passwords or keys, it is important to consider the threat
/// model and in some cases the memorability of the password. This is out of
/// scope of the Rand project, and therefore we defer to the following
/// references:
///
/// - [Wikipedia article on Password Strength](https://en.wikipedia.org/wiki/Password_strength)
/// - [Diceware for generating memorable passwords](https://en.wikipedia.org/wiki/Diceware)
#[derive(Debug)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub struct Alphanumeric;

// ----- Implementations of distributions -----

impl Distribution<char> for Standard {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> char {
        // A valid `char` is either in the interval `[0, 0xD800)` or
        // `(0xDFFF, 0x11_0000)`. All `char`s must therefore be in
        // `[0, 0x11_0000)` but not in the "gap" `[0xD800, 0xDFFF]` which is
        // reserved for surrogates. This is the size of that gap.
        const GAP_SIZE: u32 = 0xDFFF - 0xD800 + 1;

        // Uniform::new(0, 0x11_0000 - GAP_SIZE) can also be used but it
        // seemed slower.
        let range = Uniform::new(GAP_SIZE, 0x11_0000);

        let mut n = range.sample(rng);
        if n <= 0xDFFF {
            n -= GAP_SIZE;
        }
        unsafe { char::from_u32_unchecked(n) }
    }
}

impl Distribution<u8> for Alphanumeric {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u8 {
        const RANGE: u32 = 26 + 26 + 10;
        const GEN_ASCII_STR_CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ\
                abcdefghijklmnopqrstuvwxyz\
                0123456789";
        // We can pick from 62 characters. This is so close to a power of 2, 64,
        // that we can do better than `Uniform`. Use a simple bitshift and
        // rejection sampling. We do not use a bitmask, because for small RNGs
        // the most significant bits are usually of higher quality.
        loop {
            let var = rng.next_u32() >> (32 - 6);
            if var < RANGE {
                return GEN_ASCII_STR_CHARSET[var as usize];
            }
        }
    }
}

impl Distribution<bool> for Standard {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> bool {
        // We can compare against an arbitrary bit of an u32 to get a bool.
        // Because the least significant bits of a lower quality RNG can have
        // simple patterns, we compare against the most significant bit. This is
        // easiest done using a sign test.
        (rng.next_u32() as i32) < 0
    }
}

macro_rules! tuple_impl {
    // use variables to indicate the arity of the tuple
    ($($tyvar:ident),* ) => {
        // the trailing commas are for the 1 tuple
        impl< $( $tyvar ),* >
            Distribution<( $( $tyvar ),* , )>
            for Standard
            where $( Standard: Distribution<$tyvar> ),*
        {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> ( $( $tyvar ),* , ) {
                (
                    // use the $tyvar's to get the appropriate number of
                    // repeats (they're not actually needed)
                    $(
                        _rng.gen::<$tyvar>()
                    ),*
                    ,
                )
            }
        }
    }
}

impl Distribution<()> for Standard {
    #[allow(clippy::unused_unit)]
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, _: &mut R) -> () {
        ()
    }
}
tuple_impl! {A}
tuple_impl! {A, B}
tuple_impl! {A, B, C}
tuple_impl! {A, B, C, D}
tuple_impl! {A, B, C, D, E}
tuple_impl! {A, B, C, D, E, F}
tuple_impl! {A, B, C, D, E, F, G}
tuple_impl! {A, B, C, D, E, F, G, H}
tuple_impl! {A, B, C, D, E, F, G, H, I}
tuple_impl! {A, B, C, D, E, F, G, H, I, J}
tuple_impl! {A, B, C, D, E, F, G, H, I, J, K}
tuple_impl! {A, B, C, D, E, F, G, H, I, J, K, L}

#[cfg(feature = "min_const_gen")]
impl<T, const N: usize> Distribution<[T; N]> for Standard
where Standard: Distribution<T>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; N] {
        let mut buff: [MaybeUninit<T>; N] = unsafe { MaybeUninit::uninit().assume_init() };

        for elem in &mut buff {
            *elem = MaybeUninit::new(_rng.gen());
        }

        unsafe { mem::transmute_copy::<_, _>(&buff) }
    }
}

#[cfg(not(feature = "min_const_gen"))]
macro_rules! array_impl {
    // recursive, given at least one type parameter:
    {$n:expr, $t:ident, $($ts:ident,)*} => {
        array_impl!{($n - 1), $($ts,)*}

        impl<T> Distribution<[T; $n]> for Standard where Standard: Distribution<T> {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] {
                [_rng.gen::<$t>(), $(_rng.gen::<$ts>()),*]
            }
        }
    };
    // empty case:
    {$n:expr,} => {
        impl<T> Distribution<[T; $n]> for Standard {
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] { [] }
        }
    };
}

#[cfg(not(feature = "min_const_gen"))]
array_impl! {32, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T,}

impl<T> Distribution<Option<T>> for Standard
where Standard: Distribution<T>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Option<T> {
        // UFCS is needed here: https://github.com/rust-lang/rust/issues/24066
        if rng.gen::<bool>() {
            Some(rng.gen())
        } else {
            None
        }
    }
}

impl<T> Distribution<Wrapping<T>> for Standard
where Standard: Distribution<T>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Wrapping<T> {
        Wrapping(rng.gen())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::RngCore;
    #[cfg(feature = "alloc")] use alloc::string::String;

    #[test]
    fn test_misc() {
        let rng: &mut dyn RngCore = &mut crate::test::rng(820);

        rng.sample::<char, _>(Standard);
        rng.sample::<bool, _>(Standard);
    }

    #[cfg(feature = "alloc")]
    #[test]
    fn test_chars() {
        use core::iter;
        let mut rng = crate::test::rng(805);

        // Test by generating a relatively large number of chars, so we also
        // take the rejection sampling path.
        let word: String = iter::repeat(())
            .map(|()| rng.gen::<char>())
            .take(1000)
            .collect();
        assert!(!word.is_empty());
    }

    #[test]
    fn test_alphanumeric() {
        let mut rng = crate::test::rng(806);

        // Test by generating a relatively large number of chars, so we also
        // take the rejection sampling path.
        let mut incorrect = false;
        for _ in 0..100 {
            let c: char = rng.sample(Alphanumeric).into();
            incorrect |= !(('0'..='9').contains(&c) ||
                           ('A'..='Z').contains(&c) ||
                           ('a'..='z').contains(&c) );
        }
        assert!(!incorrect);
    }

    #[test]
    fn value_stability() {
        fn test_samples<T: Copy + core::fmt::Debug + PartialEq, D: Distribution<T>>(
            distr: &D, zero: T, expected: &[T],
        ) {
            let mut rng = crate::test::rng(807);
            let mut buf = [zero; 5];
            for x in &mut buf {
                *x = rng.sample(&distr);
            }
            assert_eq!(&buf, expected);
        }

        test_samples(&Standard, 'a', &[
            '\u{8cdac}',
            '\u{a346a}',
            '\u{80120}',
            '\u{ed692}',
            '\u{35888}',
        ]);
        test_samples(&Alphanumeric, 0, &[104, 109, 101, 51, 77]);
        test_samples(&Standard, false, &[true, true, false, true, false]);
        test_samples(&Standard, None as Option<bool>, &[
            Some(true),
            None,
            Some(false),
            None,
            Some(false),
        ]);
        test_samples(&Standard, Wrapping(0i32), &[
            Wrapping(-2074640887),
            Wrapping(-1719949321),
            Wrapping(2018088303),
            Wrapping(-547181756),
            Wrapping(838957336),
        ]);

        // We test only sub-sets of tuple and array impls
        test_samples(&Standard, (), &[(), (), (), (), ()]);
        test_samples(&Standard, (false,), &[
            (true,),
            (true,),
            (false,),
            (true,),
            (false,),
        ]);
        test_samples(&Standard, (false, false), &[
            (true, true),
            (false, true),
            (false, false),
            (true, false),
            (false, false),
        ]);

        test_samples(&Standard, [0u8; 0], &[[], [], [], [], []]);
        test_samples(&Standard, [0u8; 3], &[
            [9, 247, 111],
            [68, 24, 13],
            [174, 19, 194],
            [172, 69, 213],
            [149, 207, 29],
        ]);
    }
}
<commit_msg>Implement `Clone` and `Debug` for `Alphanumeric`<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The implementations of the `Standard` distribution for other built-in types.

use core::char;
use core::num::Wrapping;

use crate::distributions::{Distribution, Standard, Uniform};
use crate::Rng;

#[cfg(feature = "serde1")]
use serde::{Serialize, Deserialize};
#[cfg(feature = "min_const_gen")]
use std::mem::{self, MaybeUninit};

// ----- Sampling distributions -----

/// Sample a `u8`, uniformly distributed over ASCII letters and numbers:
/// a-z, A-Z and 0-9.
///
/// # Example
///
/// ```
/// use std::iter;
/// use rand::{Rng, thread_rng};
/// use rand::distributions::Alphanumeric;
///
/// let mut rng = thread_rng();
/// let chars: String = iter::repeat(())
///         .map(|()| rng.sample(Alphanumeric))
///         .map(char::from)
///         .take(7)
///         .collect();
/// println!("Random chars: {}", chars);
/// ```
///
/// # Passwords
///
/// Users sometimes ask whether it is safe to use a string of random characters
/// as a password. In principle, all RNGs in Rand implementing `CryptoRng` are
/// suitable as a source of randomness for generating passwords (if they are
/// properly seeded), but it is more conservative to only use randomness
/// directly from the operating system via the `getrandom` crate, or the
/// corresponding bindings of a crypto library.
///
/// When generating passwords or keys, it is important to consider the threat
/// model and in some cases the memorability of the password. This is out of
/// scope of the Rand project, and therefore we defer to the following
/// references:
///
/// - [Wikipedia article on Password Strength](https://en.wikipedia.org/wiki/Password_strength)
/// - [Diceware for generating memorable passwords](https://en.wikipedia.org/wiki/Diceware)
#[derive(Debug, Clone, Copy)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub struct Alphanumeric;

// ----- Implementations of distributions -----

impl Distribution<char> for Standard {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> char {
        // A valid `char` is either in the interval `[0, 0xD800)` or
        // `(0xDFFF, 0x11_0000)`. All `char`s must therefore be in
        // `[0, 0x11_0000)` but not in the "gap" `[0xD800, 0xDFFF]` which is
        // reserved for surrogates. This is the size of that gap.
        const GAP_SIZE: u32 = 0xDFFF - 0xD800 + 1;

        // Uniform::new(0, 0x11_0000 - GAP_SIZE) can also be used but it
        // seemed slower.
        let range = Uniform::new(GAP_SIZE, 0x11_0000);

        let mut n = range.sample(rng);
        if n <= 0xDFFF {
            n -= GAP_SIZE;
        }
        unsafe { char::from_u32_unchecked(n) }
    }
}

impl Distribution<u8> for Alphanumeric {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u8 {
        const RANGE: u32 = 26 + 26 + 10;
        const GEN_ASCII_STR_CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ\
                abcdefghijklmnopqrstuvwxyz\
                0123456789";
        // We can pick from 62 characters. This is so close to a power of 2, 64,
        // that we can do better than `Uniform`. Use a simple bitshift and
        // rejection sampling. We do not use a bitmask, because for small RNGs
        // the most significant bits are usually of higher quality.
        loop {
            let var = rng.next_u32() >> (32 - 6);
            if var < RANGE {
                return GEN_ASCII_STR_CHARSET[var as usize];
            }
        }
    }
}

impl Distribution<bool> for Standard {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> bool {
        // We can compare against an arbitrary bit of an u32 to get a bool.
        // Because the least significant bits of a lower quality RNG can have
        // simple patterns, we compare against the most significant bit. This is
        // easiest done using a sign test.
        (rng.next_u32() as i32) < 0
    }
}

macro_rules! tuple_impl {
    // use variables to indicate the arity of the tuple
    ($($tyvar:ident),* ) => {
        // the trailing commas are for the 1 tuple
        impl< $( $tyvar ),* >
            Distribution<( $( $tyvar ),* , )>
            for Standard
            where $( Standard: Distribution<$tyvar> ),*
        {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> ( $( $tyvar ),* , ) {
                (
                    // use the $tyvar's to get the appropriate number of
                    // repeats (they're not actually needed)
                    $(
                        _rng.gen::<$tyvar>()
                    ),*
                    ,
                )
            }
        }
    }
}

impl Distribution<()> for Standard {
    #[allow(clippy::unused_unit)]
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, _: &mut R) -> () {
        ()
    }
}
tuple_impl! {A}
tuple_impl! {A, B}
tuple_impl! {A, B, C}
tuple_impl! {A, B, C, D}
tuple_impl! {A, B, C, D, E}
tuple_impl! {A, B, C, D, E, F}
tuple_impl! {A, B, C, D, E, F, G}
tuple_impl! {A, B, C, D, E, F, G, H}
tuple_impl! {A, B, C, D, E, F, G, H, I}
tuple_impl! {A, B, C, D, E, F, G, H, I, J}
tuple_impl! {A, B, C, D, E, F, G, H, I, J, K}
tuple_impl! {A, B, C, D, E, F, G, H, I, J, K, L}

#[cfg(feature = "min_const_gen")]
impl<T, const N: usize> Distribution<[T; N]> for Standard
where Standard: Distribution<T>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; N] {
        let mut buff: [MaybeUninit<T>; N] = unsafe { MaybeUninit::uninit().assume_init() };

        for elem in &mut buff {
            *elem = MaybeUninit::new(_rng.gen());
        }

        unsafe { mem::transmute_copy::<_, _>(&buff) }
    }
}

#[cfg(not(feature = "min_const_gen"))]
macro_rules! array_impl {
    // recursive, given at least one type parameter:
    {$n:expr, $t:ident, $($ts:ident,)*} => {
        array_impl!{($n - 1), $($ts,)*}

        impl<T> Distribution<[T; $n]> for Standard where Standard: Distribution<T> {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] {
                [_rng.gen::<$t>(), $(_rng.gen::<$ts>()),*]
            }
        }
    };
    // empty case:
    {$n:expr,} => {
        impl<T> Distribution<[T; $n]> for Standard {
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] { [] }
        }
    };
}

#[cfg(not(feature = "min_const_gen"))]
array_impl! {32, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T,}

impl<T> Distribution<Option<T>> for Standard
where Standard: Distribution<T>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Option<T> {
        // UFCS is needed here: https://github.com/rust-lang/rust/issues/24066
        if rng.gen::<bool>() {
            Some(rng.gen())
        } else {
            None
        }
    }
}

impl<T> Distribution<Wrapping<T>> for Standard
where Standard: Distribution<T>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Wrapping<T> {
        Wrapping(rng.gen())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::RngCore;
    #[cfg(feature = "alloc")] use alloc::string::String;

    #[test]
    fn test_misc() {
        let rng: &mut dyn RngCore = &mut crate::test::rng(820);

        rng.sample::<char, _>(Standard);
        rng.sample::<bool, _>(Standard);
    }

    #[cfg(feature = "alloc")]
    #[test]
    fn test_chars() {
        use core::iter;
        let mut rng = crate::test::rng(805);

        // Test by generating a relatively large number of chars, so we also
        // take the rejection sampling path.
        let word: String = iter::repeat(())
            .map(|()| rng.gen::<char>())
            .take(1000)
            .collect();
        assert!(!word.is_empty());
    }

    #[test]
    fn test_alphanumeric() {
        let mut rng = crate::test::rng(806);

        // Test by generating a relatively large number of chars, so we also
        // take the rejection sampling path.
        let mut incorrect = false;
        for _ in 0..100 {
            let c: char = rng.sample(Alphanumeric).into();
            incorrect |= !(('0'..='9').contains(&c) ||
                           ('A'..='Z').contains(&c) ||
                           ('a'..='z').contains(&c) );
        }
        assert!(!incorrect);
    }

    #[test]
    fn value_stability() {
        fn test_samples<T: Copy + core::fmt::Debug + PartialEq, D: Distribution<T>>(
            distr: &D, zero: T, expected: &[T],
        ) {
            let mut rng = crate::test::rng(807);
            let mut buf = [zero; 5];
            for x in &mut buf {
                *x = rng.sample(&distr);
            }
            assert_eq!(&buf, expected);
        }

        test_samples(&Standard, 'a', &[
            '\u{8cdac}',
            '\u{a346a}',
            '\u{80120}',
            '\u{ed692}',
            '\u{35888}',
        ]);
        test_samples(&Alphanumeric, 0, &[104, 109, 101, 51, 77]);
        test_samples(&Standard, false, &[true, true, false, true, false]);
        test_samples(&Standard, None as Option<bool>, &[
            Some(true),
            None,
            Some(false),
            None,
            Some(false),
        ]);
        test_samples(&Standard, Wrapping(0i32), &[
            Wrapping(-2074640887),
            Wrapping(-1719949321),
            Wrapping(2018088303),
            Wrapping(-547181756),
            Wrapping(838957336),
        ]);

        // We test only sub-sets of tuple and array impls
        test_samples(&Standard, (), &[(), (), (), (), ()]);
        test_samples(&Standard, (false,), &[
            (true,),
            (true,),
            (false,),
            (true,),
            (false,),
        ]);
        test_samples(&Standard, (false, false), &[
            (true, true),
            (false, true),
            (false, false),
            (true, false),
            (false, false),
        ]);

        test_samples(&Standard, [0u8; 0], &[[], [], [], [], []]);
        test_samples(&Standard, [0u8; 3], &[
            [9, 247, 111],
            [68, 24, 13],
            [174, 19, 194],
            [172, 69, 213],
            [149, 207, 29],
        ]);
    }
}
<|endoftext|>


<commit_before>//! Defines `Module` and `Function` that act as containers.

use std::fmt::Debug;
use std::hash::Hash;
use std::collections::{BTreeSet, HashMap};

use r2pipe::structs::{LOpInfo, LRegInfo};

use petgraph::graph::NodeIndex;

use frontend::source::Source;
use frontend::ssaconstructor::SSAConstruct;
use middle::ssa::ssastorage::SSAStorage;

#[derive(Clone, Debug, Default)]
pub struct RadecoModule {
    functions: HashMap<u64, RadecoFunction>,
    fname: HashMap<String, u64>,
}

#[derive(Clone, Debug)]
pub struct RadecoFunction {
    ssa: SSAStorage,
    name: String,
    call_ctx: Vec<CallContext>,
    callrefs: BTreeSet<u64>,
    callxrefs: BTreeSet<u64>,
}

impl Default for RadecoFunction {
    fn default() -> RadecoFunction {
        RadecoFunction {
            ssa: SSAStorage::new(),
            call_ctx: Vec::new(),
            callrefs: BTreeSet::new(),
            callxrefs: BTreeSet::new(),
            name: String::new(),
        }
    }
}

impl RadecoFunction {
    fn new() -> RadecoFunction {
        RadecoFunction::default()
    }

    pub fn construct(reg_profile: &LRegInfo, insts: Vec<LOpInfo>) -> RadecoFunction {
        let mut rfn = RadecoFunction::default();
        {
            let mut constructor = SSAConstruct::new(&mut rfn.ssa, &reg_profile);
            constructor.run(insts);
        }
        rfn
    }
}

#[derive(Clone, Debug)]
pub struct CallContext {
    /// Start offset of caller (uniquely identifies a function).
    caller: u64,
    /// Start offset of callee (uniquely identifies a function).
    callee: u64,
    /// Offset the call. Note that this will belong to the caller function.
    call_site: u64,
    /// Translate a node in callees context into a node in caller's context.
    ctx_translate: HashMap<NodeIndex, NodeIndex>,
}

impl CallContext {
    pub fn new(caller: u64, callee: u64, call_site: u64) -> CallContext {
        CallContext {
            caller: caller,
            callee: callee,
            call_site: call_site,
            ctx_translate: HashMap::new(),
        }
    }
}

// From trait to construct a module from `Source`.
// Note that this conversion is expensive as the source is used to construct
// the SSA for all the
// function that it holds and perform basic analysis.
impl<'a, T: 'a + Source> From<&'a mut T> for RadecoModule {
    fn from(source: &'a mut T) -> RadecoModule {
        let reg_info = source.register_profile();
        let mut rmod = RadecoModule::default();
        for f in source.functions() {
            // TODO: Thread this part.
            if f.name.as_ref().unwrap().contains("imp") {
                continue
            }

            let offset = f.offset.expect("Invalid offset");
            let instructions = source.instructions_at(offset);
            let mut rfn = RadecoFunction::construct(&reg_info, instructions);
            rfn.name = f.name.as_ref().unwrap().clone();
            if let Some(ref callrefs) = f.callrefs {
                rfn.callrefs = callrefs.iter()
                                       .filter(|x| {
                                           match x.call_type {
                                               Some(ref c) if c == "C" => true,
                                               _ => false,
                                           }
                                       })
                                       .map(|x| x.addr.expect("Invalid address"))
                                       .collect::<BTreeSet<_>>();
            }

            if let Some(ref callxrefs) = f.codexrefs {
                rfn.callxrefs = callxrefs.iter()
                                       .filter(|x| {
                                           match x.call_type {
                                               Some(ref c) if c == "C" => true,
                                               _ => false,
                                           }
                                       })
                                       .map(|x| x.addr.expect("Invalid address"))
                                       .collect::<BTreeSet<_>>();
            }

            rmod.functions.insert(offset, rfn);
            rmod.fname.insert(f.name.unwrap(), offset);
        }
        rmod
    }
}

impl RadecoModule {
    fn construct<S: Source>(mut source: S) -> RadecoModule {
        RadecoModule::from(&mut source)
    }
}

pub trait RFunction { }
pub trait RModule {
    type FnRef: Copy + Clone + Debug + Hash + Eq;

    fn callees_of(&self, &Self::FnRef) -> Vec<Self::FnRef>;
    fn callers_of(&self, &Self::FnRef) -> Vec<Self::FnRef>;
    fn functions(&self) -> Vec<Self::FnRef>;
    fn function_by_ref(&self, &Self::FnRef) -> Box<&RFunction>;
}

#[cfg(test)]
mod test {
    use super::*;
    use r2pipe::r2::R2;
    use frontend::source::Source;
    use middle::ir_writer::IRWriter;
    use std::io;
    use middle::dce;

    #[test]
    fn module_test() {
        let mut r2 = R2::open(Some("./crowell_example.o")).expect("Failed to open r2");
        r2.init();
        let mut rmod = RadecoModule::from(&mut r2);
        for (ref addr, ref mut rfn) in rmod.functions.iter_mut() {
            {
                dce::collect(&mut rfn.ssa);
            }
            let mut writer: IRWriter = Default::default();
            writer.emit_il(Some(rfn.name.clone()), &rfn.ssa, &mut io::stdout());
        }
    }
}
<commit_msg>Adds threading to module ssa construction<commit_after>//! Defines `Module` and `Function` that act as containers.

use std::fmt::Debug;
use std::hash::Hash;
use std::collections::{BTreeSet, HashMap};
use std::thread;
use std::sync::mpsc;

use r2pipe::structs::{LOpInfo, LRegInfo};

use petgraph::graph::NodeIndex;

use frontend::source::Source;
use frontend::ssaconstructor::SSAConstruct;
use middle::ssa::ssastorage::SSAStorage;

#[derive(Clone, Debug, Default)]
pub struct RadecoModule {
    functions: HashMap<u64, RadecoFunction>,
    fname: HashMap<String, u64>,
}

#[derive(Clone, Debug)]
pub struct RadecoFunction {
    ssa: SSAStorage,
    name: String,
    call_ctx: Vec<CallContext>,
    callrefs: BTreeSet<u64>,
    callxrefs: BTreeSet<u64>,
}

impl Default for RadecoFunction {
    fn default() -> RadecoFunction {
        RadecoFunction {
            ssa: SSAStorage::new(),
            call_ctx: Vec::new(),
            callrefs: BTreeSet::new(),
            callxrefs: BTreeSet::new(),
            name: String::new(),
        }
    }
}

impl RadecoFunction {
    fn new() -> RadecoFunction {
        RadecoFunction::default()
    }

    pub fn construct(reg_profile: &LRegInfo, insts: Vec<LOpInfo>) -> RadecoFunction {
        let mut rfn = RadecoFunction::default();
        {
            let mut constructor = SSAConstruct::new(&mut rfn.ssa, &reg_profile);
            constructor.run(insts);
        }
        rfn
    }
}

#[derive(Clone, Debug)]
pub struct CallContext {
    /// Start offset of caller (uniquely identifies a function).
    caller: u64,
    /// Start offset of callee (uniquely identifies a function).
    callee: u64,
    /// Offset the call. Note that this will belong to the caller function.
    call_site: u64,
    /// Translate a node in callees context into a node in caller's context.
    ctx_translate: HashMap<NodeIndex, NodeIndex>,
}

impl CallContext {
    pub fn new(caller: u64, callee: u64, call_site: u64) -> CallContext {
        CallContext {
            caller: caller,
            callee: callee,
            call_site: call_site,
            ctx_translate: HashMap::new(),
        }
    }
}

// TODO
fn ssa_single_fn() -> () {
}

// From trait to construct a module from `Source`.
// Note that this conversion is expensive as the source is used to construct
// the SSA for all the
// function that it holds and perform basic analysis.
impl<'a, T: 'a + Source> From<&'a mut T> for RadecoModule {
    fn from(source: &'a mut T) -> RadecoModule {
        let reg_info = source.register_profile();
        let mut rmod = RadecoModule::default();
        //let mut childs = Vec::new();
        let (tx, rx) = mpsc::channel();
        let mut i = 0;
        let mut handles = Vec::new();
        for f in source.functions() {
            if f.name.as_ref().unwrap().contains("sym.imp") {
                continue;
            }
            let offset = f.offset.expect("Invalid offset");
            let instructions = source.instructions_at(offset);
            let tx = tx.clone();
            let reg_info = reg_info.clone();

            let handle = thread::spawn(move || {
                let mut rfn = RadecoFunction::construct(&reg_info, instructions);
                rfn.name = f.name.as_ref().unwrap().clone();
                if let Some(ref callrefs) = f.callrefs {
                    rfn.callrefs = callrefs.iter()
                                           .filter(|x| {
                                               match x.call_type {
                                                   Some(ref c) if c == "C" => true,
                                                   _ => false,
                                               }
                                           })
                                           .map(|x| x.addr.expect("Invalid address"))
                                           .collect::<BTreeSet<_>>();
                }

                if let Some(ref callxrefs) = f.codexrefs {
                    rfn.callxrefs = callxrefs.iter()
                                             .filter(|x| {
                                                 match x.call_type {
                                                     Some(ref c) if c == "C" => true,
                                                     _ => false,
                                                 }
                                             })
                                             .map(|x| x.addr.expect("Invalid address"))
                                             .collect::<BTreeSet<_>>();
                }
                tx.send((offset, f.name.unwrap(), rfn)).unwrap();
            });
            handles.push(handle);
        }

        let mut success = 0;
        for h in handles {
            if let Ok(_) = h.join() {
                success += 1;
            }
        }

        for _ in 0..success {
            let (offset, name, rfn) = rx.recv().unwrap();
            rmod.functions.insert(offset, rfn);
            rmod.fname.insert(name, offset);
        }
        rmod
    }
}

impl RadecoModule {
    fn construct<S: Source>(mut source: S) -> RadecoModule {
        RadecoModule::from(&mut source)
    }
}

pub trait RFunction { }
pub trait RModule {
    type FnRef: Copy + Clone + Debug + Hash + Eq;

    fn callees_of(&self, &Self::FnRef) -> Vec<Self::FnRef>;
    fn callers_of(&self, &Self::FnRef) -> Vec<Self::FnRef>;
    fn functions(&self) -> Vec<Self::FnRef>;
    fn function_ref(&self, &Self::FnRef) -> &RFunction;
    fn function_ref_mut(&mut self, &Self::FnRef) -> &mut RFunction;
}

#[cfg(test)]
mod test {
    use super::*;
    use r2pipe::r2::R2;
    use frontend::source::Source;
    use middle::ir_writer::IRWriter;
    use std::io;
    use middle::dce;

    #[test]
    fn module_test() {
        let mut r2 = R2::open(Some("./crowell_example.o")).expect("Failed to open r2");
        r2.init();
        let mut rmod = RadecoModule::from(&mut r2);
        for (ref addr, ref mut rfn) in rmod.functions.iter_mut() {
            {
                dce::collect(&mut rfn.ssa);
            }
            let mut writer: IRWriter = Default::default();
            writer.emit_il(Some(rfn.name.clone()), &rfn.ssa, &mut io::stdout());
        }
    }
}
<|endoftext|>


<commit_before>use super::super::request::Request;
use super::super::response::Response;
use super::super::ingot::{Ingot, Continue, Unwind};
use super::super::anymap::AnyMap;
use super::super::alloy::Alloy;

use super::Furnace;

struct IronFurnace<Rq, Rs> {
    stack: Vec<Box<Ingot<Rq, Rs>: Send>>
}

impl<Rq: Request, Rs: Response> Clone for IronFurnace<Rq, Rs> {
    fn clone(&self) -> IronFurnace<Rq, Rs> { IronFurnace { stack: self.stack.clone() } }
}

impl<Rq: Request, Rs: Response> Furnace<Rq, Rs> for IronFurnace<Rq, Rs> {
    fn forge(&mut self, request: &mut Rq, response: &mut Rs, opt_alloy: Option<&mut Alloy>) {
        let mut alloy = &mut AnyMap::new();

        match opt_alloy {
            Some(a) => alloy = a,
            None => ()
        };

        let mut exit_stack = vec![];

        'enter: for ingot in self.stack.mut_iter() {
            match ingot.enter(request, response, alloy) {
                Unwind   => break 'enter,
                Continue => exit_stack.push(ingot)
            }
        }

        exit_stack.reverse();
        'exit: for ingot in exit_stack.mut_iter() {
            let _ = ingot.exit(request, response, alloy);
        }
    }
    fn smelt<I: Ingot<Rq, Rs>>(&mut self, ingot: I) {
        self.stack.push(box ingot);
    }
}

<commit_msg>(doc) Documented IronFurnace.<commit_after>use super::super::request::Request;
use super::super::response::Response;
use super::super::ingot::{Ingot, Continue, Unwind};
use super::super::anymap::AnyMap;
use super::super::alloy::Alloy;

use super::Furnace;

/// The default Furnace used by Iron.
/// IronFurnace just runs each request through all Ingots in its stack,
/// then, when it hits an Ingot which returns `Unwind`, it will
/// pass the request back up through all Ingots it has hit
/// so far.
struct IronFurnace<Rq, Rs> {
    /// The storage used by IronFurnace to hold all Ingots
    /// that have been smelted on to it.
    stack: Vec<Box<Ingot<Rq, Rs>: Send>>
}

impl<Rq: Request, Rs: Response> Clone for IronFurnace<Rq, Rs> {
    fn clone(&self) -> IronFurnace<Rq, Rs> { IronFurnace { stack: self.stack.clone() } }
}

/// IronFurnace is a Furnace
impl<Rq: Request, Rs: Response> Furnace<Rq, Rs> for IronFurnace<Rq, Rs> {
    fn forge(&mut self, request: &mut Rq, response: &mut Rs, opt_alloy: Option<&mut Alloy>) {
        /// Create a placeholder alloy.
        let mut alloy = &mut AnyMap::new();

        /// See if we were passed one.
        match opt_alloy {
            /// If so, forget about our placeholder.
            Some(a) => alloy = a,
            /// Else just use our new Alloy.
            None => ()
        };

        /// The exit_stack will hold all Ingots that are passed through
        /// in the enter loop. This is so we know to take exactly the same
        /// path through ingots in reverse order that we did on the way in.
        let mut exit_stack = vec![];

        'enter: for ingot in self.stack.mut_iter() {
            match ingot.enter(request, response, alloy) {
                Unwind   => break 'enter,
                /// Mark the ingot for traversal on exit.
                Continue => exit_stack.push(ingot)
            }
        }

        /// Reverse the stack so we go through in the reverse order.
        /// i.e. LIFO.
        exit_stack.reverse();
        /// Call each ingots exit handler.
        'exit: for ingot in exit_stack.mut_iter() {
            let _ = ingot.exit(request, response, alloy);
        }
    }
    /// Add an Ingot to the furnace.
    fn smelt<I: Ingot<Rq, Rs>>(&mut self, ingot: I) {
        self.stack.push(box ingot);
    }
}

<|endoftext|>


<commit_before>use std::rc::Rc;

use glib;

use gtk;
use gtk::prelude::*;
use gtk::{Window, MenuBar, MenuItem, Menu};

use app::{AppContextPointer, AppContext};
use app::config;
use gui::Gui;

mod menu_callbacks;

pub fn layout(gui: &Gui, app_context: &AppContextPointer) {

    let window = gui.get_window();

    // Build the menu bar
    let menu_bar = build_menu_bar(app_context, &window);

    // Layout main gui areas
    let frames = layout_frames(gui);

    // Layout everything else
    let v_box = gtk::Box::new(gtk::Orientation::Vertical, 0);
    v_box.pack_start(&menu_bar, false, false, 0);
    v_box.pack_start(&frames, true, true, 0);
    window.add(&v_box);

    configure_main_window(&window, &app_context.borrow());

}

fn build_menu_bar(ac: &AppContextPointer, win: &Window) -> MenuBar {

    //
    // The file menu.
    //

    // The open item
    let open_item = MenuItem::new_with_label("Open");
    let win1 = win.clone();
    let ac1 = Rc::clone(ac);
    open_item.connect_activate(move |mi| menu_callbacks::open_callback(mi, &ac1, &win1));

    // The quit item
    let quit_item = MenuItem::new_with_label("Quit");
    quit_item.connect_activate(|_| { gtk::main_quit(); });

    // Build the file menu
    let file_menu = Menu::new();
    file_menu.append(&open_item);
    file_menu.append(&quit_item);

    // Build the file menu item
    let file = MenuItem::new_with_label("File");
    file.set_submenu(&file_menu);

    //
    // End the file menu
    //

    //
    // Build the menu bar
    //
    let menu_bar = MenuBar::new();
    menu_bar.append(&file);
    menu_bar
}

fn layout_frames(gui: &Gui) -> gtk::Grid {

    const TOTAL_WIDTH: i32 = 4;
    const TOTAL_HEIGHT: i32 = 4;

    const SKEW_T_WIDTH_FRACTION: f32 = 0.75;
    const SKEW_T_WIDTH: i32 = (TOTAL_WIDTH as f32 * SKEW_T_WIDTH_FRACTION) as i32;
    const OTHER_WIDTH: i32 = TOTAL_WIDTH - SKEW_T_WIDTH;

    const HODO_FRACTION: f32 = 0.4;
    const HODO_HEIGHT: i32 = (TOTAL_HEIGHT as f32 * HODO_FRACTION) as i32;

    const INDEX_FRACTION: f32 = 0.4;
    const INDEX_HEIGHT: i32 = (TOTAL_HEIGHT as f32 * INDEX_FRACTION) as i32;

    const CONTROL_HEIGHT: i32 = TOTAL_HEIGHT - INDEX_HEIGHT - HODO_HEIGHT;

    let h_box = gtk::Box::new(gtk::Orientation::Horizontal, 0);
    h_box.pack_start(&gui.get_omega_area(), false, true, 0);
    h_box.pack_start(&gui.get_sounding_area(), true, true, 0);

    let grid = gtk::Grid::new();
    grid.attach(&add_border_frame(&h_box), 0, 0, SKEW_T_WIDTH, TOTAL_HEIGHT);
    grid.attach(
        &add_border_frame(&gui.get_hodograph_area()),
        SKEW_T_WIDTH,
        0,
        OTHER_WIDTH,
        HODO_HEIGHT,
    );
    grid.attach(
        &add_border_frame(&gui.get_index_area()),
        SKEW_T_WIDTH,
        HODO_HEIGHT,
        OTHER_WIDTH,
        INDEX_HEIGHT,
    );
    grid.attach(
        &add_border_frame(&gui.get_control_area()),
        SKEW_T_WIDTH,
        HODO_HEIGHT + INDEX_HEIGHT,
        OTHER_WIDTH,
        CONTROL_HEIGHT,
    );

    grid.set_row_homogeneous(true);
    grid.set_column_homogeneous(true);

    grid
}

fn configure_main_window(window: &Window, ac: &AppContext) {
    // Set up window title, size, etc
    window.set_title("Sonde");
    window.set_default_size(ac.config.window_width, ac.config.window_height);
    window.set_decorated(true);
    window.show_all();
    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(false)
    });
}

fn add_border_frame<P: glib::IsA<gtk::Widget>>(widget: &P) -> gtk::Frame {
    let f = gtk::Frame::new(None);
    f.add(widget);
    f.set_border_width(config::BORDER_WIDTH);
    f.set_hexpand(true);
    f.set_vexpand(true);

    f
}
<commit_msg>Better window layout.<commit_after>use std::rc::Rc;

use glib;

use gtk;
use gtk::prelude::*;
use gtk::{Window, MenuBar, MenuItem, Menu};

use app::{AppContextPointer, AppContext};
use app::config;
use gui::Gui;

mod menu_callbacks;

pub fn layout(gui: &Gui, app_context: &AppContextPointer) {

    let ac = app_context.borrow();
    let window = gui.get_window();

    // Build the menu bar
    let menu_bar = build_menu_bar(app_context, &window);

    // Layout main gui areas
    let frames = layout_frames(gui, &ac);

    // Layout everything else
    let v_box = gtk::Box::new(gtk::Orientation::Vertical, 0);
    v_box.pack_start(&menu_bar, false, false, 0);
    v_box.pack_start(&frames, true, true, 0);
    window.add(&v_box);

    configure_main_window(&window, &ac);
}

fn build_menu_bar(ac: &AppContextPointer, win: &Window) -> MenuBar {

    //
    // The file menu.
    //

    // The open item
    let open_item = MenuItem::new_with_label("Open");
    let win1 = win.clone();
    let ac1 = Rc::clone(ac);
    open_item.connect_activate(move |mi| menu_callbacks::open_callback(mi, &ac1, &win1));

    // The quit item
    let quit_item = MenuItem::new_with_label("Quit");
    quit_item.connect_activate(|_| { gtk::main_quit(); });

    // Build the file menu
    let file_menu = Menu::new();
    file_menu.append(&open_item);
    file_menu.append(&quit_item);

    // Build the file menu item
    let file = MenuItem::new_with_label("File");
    file.set_submenu(&file_menu);

    //
    // End the file menu
    //

    //
    // Build the menu bar
    //
    let menu_bar = MenuBar::new();
    menu_bar.append(&file);
    menu_bar
}

fn layout_frames(gui: &Gui, ac: &AppContext) -> gtk::Paned {

    const BOX_SPACING:i32 = 0;
    const BOX_PADDING:u32 = 0;

    let skew_t = gui.get_sounding_area();
    let rh_omega = gui.get_omega_area();
    let hodo = gui.get_hodograph_area();
    let index = gui.get_index_area();
    let controls = gui.get_control_area();

    let main_pane = gtk::Paned::new(gtk::Orientation::Horizontal);
    let h_box = gtk::Box::new(gtk::Orientation::Horizontal, BOX_SPACING);

    h_box.pack_start(&rh_omega, false, true, BOX_PADDING);
    h_box.pack_start(&skew_t, true, true, BOX_PADDING);

    let v_pane_inner = gtk::Paned::new(gtk::Orientation::Vertical);
    let v_pane_outer = gtk::Paned::new(gtk::Orientation::Vertical);

    v_pane_inner.add1(&add_border_frame(&hodo));
    v_pane_inner.add2(&add_border_frame(&index));
    v_pane_outer.add1(&v_pane_inner);
    v_pane_outer.add2(&add_border_frame(&controls));
    v_pane_inner.set_position(ac.config.window_height / 3);
    v_pane_outer.set_position(ac.config.window_height * 2 / 3);

    main_pane.add1(&add_border_frame(&h_box));
    main_pane.add2(&v_pane_outer);
    main_pane.set_position(ac.config.window_width * 11 / 16);

    main_pane
}

fn configure_main_window(window: &Window, ac: &AppContext) {
    // Set up window title, size, etc
    window.set_title("Sonde");
    window.set_default_size(ac.config.window_width, ac.config.window_height);
    window.set_decorated(true);
    window.show_all();
    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(false)
    });
}

fn add_border_frame<P: glib::IsA<gtk::Widget>>(widget: &P) -> gtk::Frame {
    let f = gtk::Frame::new(None);
    f.add(widget);
    f.set_border_width(config::BORDER_WIDTH);
    f.set_hexpand(true);
    f.set_vexpand(true);
    f.set_shadow_type(gtk::ShadowType::In);

    f
}
<|endoftext|>


<commit_before>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interface listing implementation for all non-Windows platforms

use std::mem;
use std::ffi::{CStr, CString};
use std::str::from_utf8_unchecked;
use std::net::IpAddr;

use libc;

use util::MacAddr;
use datalink::NetworkInterface;

use internal;

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    fn merge(old: &mut NetworkInterface, new: &NetworkInterface) {
        old.mac = match new.mac {
            None => old.mac,
            _ => new.mac,
        };
        match (&mut old.ips, &new.ips) {
            (&mut Some(ref mut old_ips), &Some(ref new_ips)) => {
                old_ips.extend_from_slice(&new_ips[..])
            }
            (&mut ref mut old_ips @ None, &Some(ref new_ips)) => *old_ips = Some(new_ips.clone()),
            _ => {}
        };
        old.flags = old.flags | new.flags;
    }

    let mut ifaces: Vec<NetworkInterface> = Vec::new();
    unsafe {
        let mut addrs: *mut libc::ifaddrs = mem::uninitialized();
        if libc::getifaddrs(&mut addrs) != 0 {
            return ifaces;
        }
        let mut addr = addrs;
        while !addr.is_null() {
            let c_str = (*addr).ifa_name as *const i8;
            let bytes = CStr::from_ptr(c_str).to_bytes();
            let name = from_utf8_unchecked(bytes).to_owned();
            let (mac, ip) = sockaddr_to_network_addr((*addr).ifa_addr as *const libc::sockaddr);
            let ni = NetworkInterface {
                name: name.clone(),
                index: 0,
                mac: mac,
                ips: ip.map(|ip| [ip].to_vec()),
                flags: (*addr).ifa_flags,
            };
            let mut found: bool = false;
            for iface in &mut ifaces {
                if name == iface.name {
                    merge(iface, &ni);
                    found = true;
                }
            }
            if !found {
                ifaces.push(ni);
            }

            addr = (*addr).ifa_next;
        }
        libc::freeifaddrs(addrs);

        for iface in &mut ifaces {
            let name = CString::new(iface.name.as_bytes());
            iface.index = libc::if_nametoindex(name.unwrap().as_ptr());
        }

        ifaces
    }
}

#[cfg(target_os = "linux")]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == libc::AF_PACKET {
            let sll: *const libc::sockaddr_ll = mem::transmute(sa);
            let mac = MacAddr((*sll).sll_addr[0],
                              (*sll).sll_addr[1],
                              (*sll).sll_addr[2],
                              (*sll).sll_addr[3],
                              (*sll).sll_addr[4],
                              (*sll).sll_addr[5]);

            (Some(mac), None)
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                                  mem::size_of::<libc::sockaddr_storage>());

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}

#[cfg(any(target_os = "freebsd", target_os = "macos"))]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use bindings::bpf;
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == bpf::AF_LINK {
            let sdl: *const bpf::sockaddr_dl = mem::transmute(sa);
            let nlen = (*sdl).sdl_nlen as usize;
            let mac = MacAddr((*sdl).sdl_data[nlen] as u8,
                              (*sdl).sdl_data[nlen + 1] as u8,
                              (*sdl).sdl_data[nlen + 2] as u8,
                              (*sdl).sdl_data[nlen + 3] as u8,
                              (*sdl).sdl_data[nlen + 4] as u8,
                              (*sdl).sdl_data[nlen + 5] as u8);

            (Some(mac), None)
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                                  mem::size_of::<libc::sockaddr_storage>());

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}
<commit_msg>Update unix_interfaces with CStr compatibility<commit_after>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interface listing implementation for all non-Windows platforms

use std::mem;
use std::ffi::{CStr, CString};
use std::os::raw::c_char;
use std::str::from_utf8_unchecked;
use std::net::IpAddr;
use libc;

use util::MacAddr;
use datalink::NetworkInterface;

use internal;

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    fn merge(old: &mut NetworkInterface, new: &NetworkInterface) {
        old.mac = match new.mac {
            None => old.mac,
            _ => new.mac,
        };
        match (&mut old.ips, &new.ips) {
            (&mut Some(ref mut old_ips), &Some(ref new_ips)) => {
                old_ips.extend_from_slice(&new_ips[..])
            }
            (&mut ref mut old_ips @ None, &Some(ref new_ips)) => *old_ips = Some(new_ips.clone()),
            _ => {}
        };
        old.flags = old.flags | new.flags;
    }

    let mut ifaces: Vec<NetworkInterface> = Vec::new();
    unsafe {
        let mut addrs: *mut libc::ifaddrs = mem::uninitialized();
        if libc::getifaddrs(&mut addrs) != 0 {
            return ifaces;
        }
        let mut addr = addrs;
        while !addr.is_null() {
            let c_str = (*addr).ifa_name as *const c_char;
            let bytes = CStr::from_ptr(c_str).to_bytes();
            let name = from_utf8_unchecked(bytes).to_owned();
            let (mac, ip) = sockaddr_to_network_addr((*addr).ifa_addr as *const libc::sockaddr);
            let ni = NetworkInterface {
                name: name.clone(),
                index: 0,
                mac: mac,
                ips: ip.map(|ip| [ip].to_vec()),
                flags: (*addr).ifa_flags,
            };
            let mut found: bool = false;
            for iface in &mut ifaces {
                if name == iface.name {
                    merge(iface, &ni);
                    found = true;
                }
            }
            if !found {
                ifaces.push(ni);
            }

            addr = (*addr).ifa_next;
        }
        libc::freeifaddrs(addrs);

        for iface in &mut ifaces {
            let name = CString::new(iface.name.as_bytes());
            iface.index = libc::if_nametoindex(name.unwrap().as_ptr());
        }

        ifaces
    }
}

#[cfg(target_os = "linux")]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == libc::AF_PACKET {
            let sll: *const libc::sockaddr_ll = mem::transmute(sa);
            let mac = MacAddr((*sll).sll_addr[0],
                              (*sll).sll_addr[1],
                              (*sll).sll_addr[2],
                              (*sll).sll_addr[3],
                              (*sll).sll_addr[4],
                              (*sll).sll_addr[5]);

            (Some(mac), None)
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                                  mem::size_of::<libc::sockaddr_storage>());

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}

#[cfg(any(target_os = "freebsd", target_os = "macos"))]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use bindings::bpf;
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == bpf::AF_LINK {
            let sdl: *const bpf::sockaddr_dl = mem::transmute(sa);
            let nlen = (*sdl).sdl_nlen as usize;
            let mac = MacAddr((*sdl).sdl_data[nlen] as u8,
                              (*sdl).sdl_data[nlen + 1] as u8,
                              (*sdl).sdl_data[nlen + 2] as u8,
                              (*sdl).sdl_data[nlen + 3] as u8,
                              (*sdl).sdl_data[nlen + 4] as u8,
                              (*sdl).sdl_data[nlen + 5] as u8);

            (Some(mac), None)
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                                  mem::size_of::<libc::sockaddr_storage>());

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}
<|endoftext|>


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

//! Configuration file management

use std::env;
use std::io::Read;
use std::path::PathBuf;
use std::fs::File;

use toml;
use servers::{ServerConfig, StratumServerConfig};
use util::LoggingConfig;
use types::{ConfigError, ConfigMembers, GlobalConfig};
use wallet::WalletConfig;

/// The default file name to use when trying to derive
/// the config file location

const CONFIG_FILE_NAME: &'static str = "grin.toml";
const GRIN_HOME: &'static str = ".grin";

/// Returns the defaults, as strewn throughout the code

impl Default for ConfigMembers {
	fn default() -> ConfigMembers {
		ConfigMembers {
			server: ServerConfig::default(),
			mining_server: Some(StratumServerConfig::default()),
			logging: Some(LoggingConfig::default()),
			wallet: WalletConfig::default(),
		}
	}
}

impl Default for GlobalConfig {
	fn default() -> GlobalConfig {
		GlobalConfig {
			config_file_path: None,
			using_config_file: false,
			members: Some(ConfigMembers::default()),
		}
	}
}

impl GlobalConfig {
	/// Need to decide on rules where to read the config file from,
	/// but will take a stab at logic for now

	fn derive_config_location(&mut self) -> Result<(), ConfigError> {
		// First, check working directory
		let mut config_path = env::current_dir().unwrap();
		config_path.push(CONFIG_FILE_NAME);
		if config_path.exists() {
			self.config_file_path = Some(config_path);
			return Ok(());
		}
		// Next, look in directory of executable
		let mut config_path = env::current_exe().unwrap();
		config_path.pop();
		config_path.push(CONFIG_FILE_NAME);
		if config_path.exists() {
			self.config_file_path = Some(config_path);
			return Ok(());
		}
		// Then look in {user_home}/.grin
		let config_path = env::home_dir();
		if let Some(mut p) = config_path {
			p.push(GRIN_HOME);
			p.push(CONFIG_FILE_NAME);
			if p.exists() {
				self.config_file_path = Some(p);
				return Ok(());
			}
		}

		// Give up
		Err(ConfigError::FileNotFoundError(String::from("")))
	}

	/// Takes the path to a config file, or if NONE, tries to determine a config
	/// file based on rules in derive_config_location
	pub fn new(file_path: Option<&str>) -> Result<GlobalConfig, ConfigError> {
		let mut return_value = GlobalConfig::default();
		if let Some(fp) = file_path {
			return_value.config_file_path = Some(PathBuf::from(&fp));
		} else {
			let _result = return_value.derive_config_location();
		}

		// No attempt at a config file, just return defaults
		if let None = return_value.config_file_path {
			return Ok(return_value);
		}

		// Config file path is given but not valid
		if !return_value.config_file_path.as_mut().unwrap().exists() {
			return Err(ConfigError::FileNotFoundError(String::from(
				return_value
					.config_file_path
					.as_mut()
					.unwrap()
					.to_str()
					.unwrap()
					.clone(),
			)));
		}

		// Try to parse the config file if it exists explode if it does exist but
		// something's wrong with it
		return_value.read_config()
	}

	/// Read config
	fn read_config(mut self) -> Result<GlobalConfig, ConfigError> {
		let mut file = File::open(self.config_file_path.as_mut().unwrap())?;
		let mut contents = String::new();
		file.read_to_string(&mut contents)?;
		let decoded: Result<ConfigMembers, toml::de::Error> = toml::from_str(&contents);
		match decoded {
			Ok(mut gc) => {
				// Put the struct back together, because the config
				// file was flattened a bit
				gc.server.stratum_mining_config = gc.mining_server.clone();
				self.using_config_file = true;
				self.members = Some(gc);
				return Ok(self);
			}
			Err(e) => {
				return Err(ConfigError::ParseError(
					String::from(
						self.config_file_path
							.as_mut()
							.unwrap()
							.to_str()
							.unwrap()
							.clone(),
					),
					String::from(format!("{}", e)),
				));
			}
		}
	}

	/// Serialize config
	pub fn ser_config(&mut self) -> Result<String, ConfigError> {
		let encoded: Result<String, toml::ser::Error> =
			toml::to_string(self.members.as_mut().unwrap());
		match encoded {
			Ok(enc) => return Ok(enc),
			Err(e) => {
				return Err(ConfigError::SerializationError(String::from(format!(
					"{}",
					e
				))));
			}
		}
	}

	/*pub fn wallet_enabled(&mut self) -> bool {
        return self.members.as_mut().unwrap().wallet.as_mut().unwrap().enable_wallet;
    }*/

	/// Enable mining
	pub fn stratum_enabled(&mut self) -> bool {
		return self.members
			.as_mut()
			.unwrap()
			.mining_server
			.as_mut()
			.unwrap()
			.enable_stratum_server
			.unwrap();
	}
}
<commit_msg>Minor config file handling cleanup<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Configuration file management

use std::env;
use std::io::Read;
use std::path::PathBuf;
use std::fs::File;

use toml;
use servers::{ServerConfig, StratumServerConfig};
use util::LoggingConfig;
use types::{ConfigError, ConfigMembers, GlobalConfig};
use wallet::WalletConfig;

/// The default file name to use when trying to derive
/// the config file location

const CONFIG_FILE_NAME: &'static str = "grin.toml";
const GRIN_HOME: &'static str = ".grin";

/// Returns the defaults, as strewn throughout the code

impl Default for ConfigMembers {
	fn default() -> ConfigMembers {
		ConfigMembers {
			server: ServerConfig::default(),
			mining_server: Some(StratumServerConfig::default()),
			logging: Some(LoggingConfig::default()),
			wallet: WalletConfig::default(),
		}
	}
}

impl Default for GlobalConfig {
	fn default() -> GlobalConfig {
		GlobalConfig {
			config_file_path: None,
			using_config_file: false,
			members: Some(ConfigMembers::default()),
		}
	}
}

impl GlobalConfig {
	/// Need to decide on rules where to read the config file from,
	/// but will take a stab at logic for now

	fn derive_config_location(&mut self) -> Result<(), ConfigError> {
		// First, check working directory
		let mut config_path = env::current_dir().unwrap();
		config_path.push(CONFIG_FILE_NAME);
		if config_path.exists() {
			self.config_file_path = Some(config_path);
			return Ok(());
		}
		// Next, look in directory of executable
		let mut config_path = env::current_exe().unwrap();
		config_path.pop();
		config_path.push(CONFIG_FILE_NAME);
		if config_path.exists() {
			self.config_file_path = Some(config_path);
			return Ok(());
		}
		// Then look in {user_home}/.grin
		let config_path = env::home_dir();
		if let Some(mut p) = config_path {
			p.push(GRIN_HOME);
			p.push(CONFIG_FILE_NAME);
			if p.exists() {
				self.config_file_path = Some(p);
				return Ok(());
			}
		}

		// Give up
		Err(ConfigError::FileNotFoundError(String::from("")))
	}

	/// Takes the path to a config file, or if NONE, tries to determine a config
	/// file based on rules in derive_config_location
	pub fn new(file_path: Option<&str>) -> Result<GlobalConfig, ConfigError> {
		let mut return_value = GlobalConfig::default();
		if let Some(fp) = file_path {
			return_value.config_file_path = Some(PathBuf::from(&fp));
		} else {
			let _result = return_value.derive_config_location();
		}

		// No attempt at a config file, just return defaults
		if let None = return_value.config_file_path {
			return Ok(return_value);
		}

		// Config file path is given but not valid
		let config_file = return_value.config_file_path.clone().unwrap();
		if !config_file.exists() {
			return Err(ConfigError::FileNotFoundError(String::from(
						config_file.to_str().unwrap()
			)));
		}

		// Try to parse the config file if it exists explode if it does exist but
		// something's wrong with it
		return_value.read_config()
	}

	/// Read config
	fn read_config(mut self) -> Result<GlobalConfig, ConfigError> {
		let mut file = File::open(self.config_file_path.as_mut().unwrap())?;
		let mut contents = String::new();
		file.read_to_string(&mut contents)?;
		let decoded: Result<ConfigMembers, toml::de::Error> = toml::from_str(&contents);
		match decoded {
			Ok(mut gc) => {
				// Put the struct back together, because the config
				// file was flattened a bit
				gc.server.stratum_mining_config = gc.mining_server.clone();
				self.using_config_file = true;
				self.members = Some(gc);
				return Ok(self);
			}
			Err(e) => {
				return Err(ConfigError::ParseError(
					String::from(
						self.config_file_path
							.as_mut()
							.unwrap()
							.to_str()
							.unwrap()
							.clone(),
					),
					String::from(format!("{}", e)),
				));
			}
		}
	}

	/// Serialize config
	pub fn ser_config(&mut self) -> Result<String, ConfigError> {
		let encoded: Result<String, toml::ser::Error> =
			toml::to_string(self.members.as_mut().unwrap());
		match encoded {
			Ok(enc) => return Ok(enc),
			Err(e) => {
				return Err(ConfigError::SerializationError(String::from(format!(
					"{}",
					e
				))));
			}
		}
	}

	/*pub fn wallet_enabled(&mut self) -> bool {
        return self.members.as_mut().unwrap().wallet.as_mut().unwrap().enable_wallet;
    }*/

	/// Enable mining
	pub fn stratum_enabled(&mut self) -> bool {
		return self.members
			.as_mut()
			.unwrap()
			.mining_server
			.as_mut()
			.unwrap()
			.enable_stratum_server
			.unwrap();
	}
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use crate::controllers::frontend_prelude::*;

use crate::controllers::helpers::*;

use crate::controllers::helpers::pagination::{Paginated, PaginationOptions};
use crate::models::{
    CrateOwner, Email, Follow, NewEmail, OwnerKind, User, Version, VersionOwnerAction,
};
use crate::schema::{crate_owners, crates, emails, follows, users, versions};
use crate::views::{EncodableMe, EncodablePrivateUser, EncodableVersion, OwnedCrate};

/// Handles the `GET /me` route.
pub fn me(req: &mut dyn RequestExt) -> EndpointResult {
    let user_id = req.authenticate()?.forbid_api_token_auth()?.user_id();
    let conn = req.db_write()?;

    let (user, verified, email, verification_sent): (User, Option<bool>, Option<String>, bool) =
        users::table
            .find(user_id)
            .left_join(emails::table)
            .select((
                users::all_columns,
                emails::verified.nullable(),
                emails::email.nullable(),
                emails::token_generated_at.nullable().is_not_null(),
            ))
            .first(&*conn)?;

    let owned_crates = CrateOwner::by_owner_kind(OwnerKind::User)
        .inner_join(crates::table)
        .filter(crate_owners::owner_id.eq(user_id))
        .select((crates::id, crates::name, crate_owners::email_notifications))
        .order(crates::name.asc())
        .load(&*conn)?
        .into_iter()
        .map(|(id, name, email_notifications)| OwnedCrate {
            id,
            name,
            email_notifications,
        })
        .collect();

    let verified = verified.unwrap_or(false);
    let verification_sent = verified || verification_sent;
    Ok(req.json(&EncodableMe {
        user: EncodablePrivateUser::from(user, email, verified, verification_sent),
        owned_crates,
    }))
}

/// Handles the `GET /me/updates` route.
pub fn updates(req: &mut dyn RequestExt) -> EndpointResult {
    use diesel::dsl::any;

    let authenticated_user = req.authenticate()?.forbid_api_token_auth()?;
    let user = authenticated_user.user();

    let followed_crates = Follow::belonging_to(&user).select(follows::crate_id);
    let query = versions::table
        .inner_join(crates::table)
        .left_outer_join(users::table)
        .filter(crates::id.eq(any(followed_crates)))
        .order(versions::created_at.desc())
        .select((
            versions::all_columns,
            crates::name,
            users::all_columns.nullable(),
        ))
        .pages_pagination(PaginationOptions::builder().gather(req)?);
    let conn = req.db_write()?;
    let data: Paginated<(Version, String, Option<User>)> = query.load(&*conn)?;
    let more = data.next_page_params().is_some();
    let versions = data.iter().map(|(v, _, _)| v).cloned().collect::<Vec<_>>();
    let data = data
        .into_iter()
        .zip(VersionOwnerAction::for_versions(&conn, &versions)?.into_iter())
        .map(|((v, cn, pb), voas)| (v, cn, pb, voas));

    let versions = data
        .into_iter()
        .map(|(version, crate_name, published_by, actions)| {
            EncodableVersion::from(version, &crate_name, published_by, actions)
        })
        .collect::<Vec<_>>();

    Ok(req.json(&json!({
        "versions": versions,
        "meta": { "more": more },
    })))
}

/// Handles the `PUT /users/:user_id` route.
pub fn update_user(req: &mut dyn RequestExt) -> EndpointResult {
    use self::emails::user_id;
    use diesel::insert_into;

    let authenticated_user = req.authenticate()?;

    let mut body = String::new();
    req.body().read_to_string(&mut body)?;

    let param_user_id = &req.params()["user_id"];
    let conn = req.db_write()?;
    let user = authenticated_user.user();

    // need to check if current user matches user to be updated
    if &user.id.to_string() != param_user_id {
        return Err(bad_request("current user does not match requested user"));
    }

    #[derive(Deserialize)]
    struct UserUpdate {
        user: User,
    }

    #[derive(Deserialize)]
    struct User {
        email: Option<String>,
    }

    let user_update: UserUpdate =
        serde_json::from_str(&body).map_err(|_| bad_request("invalid json request"))?;

    let user_email = match &user_update.user.email {
        Some(email) => email.trim(),
        None => return Err(bad_request("empty email rejected")),
    };

    if user_email.is_empty() {
        return Err(bad_request("empty email rejected"));
    }

    conn.transaction::<_, Box<dyn AppError>, _>(|| {
        let new_email = NewEmail {
            user_id: user.id,
            email: user_email,
        };

        let token: String = insert_into(emails::table)
            .values(&new_email)
            .on_conflict(user_id)
            .do_update()
            .set(&new_email)
            .returning(emails::token)
            .get_result(&*conn)
            .map_err(|_| server_error("Error in creating token"))?;

        // This swallows any errors that occur while attempting to send the email. Some users have
        // an invalid email set in their GitHub profile, and we should let them sign in even though
        // we're trying to silently use their invalid address during signup and can't send them an
        // email. They'll then have to provide a valid email address.
        let _ = req
            .app()
            .emails
            .send_user_confirm(user_email, &user.gh_login, &token);

        Ok(())
    })?;

    ok_true()
}

/// Handles the `PUT /confirm/:email_token` route
pub fn confirm_user_email(req: &mut dyn RequestExt) -> EndpointResult {
    use diesel::update;

    let conn = req.db_write()?;
    let req_token = &req.params()["email_token"];

    let updated_rows = update(emails::table.filter(emails::token.eq(req_token)))
        .set(emails::verified.eq(true))
        .execute(&*conn)?;

    if updated_rows == 0 {
        return Err(bad_request("Email belonging to token not found."));
    }

    ok_true()
}

/// Handles `PUT /user/:user_id/resend` route
pub fn regenerate_token_and_send(req: &mut dyn RequestExt) -> EndpointResult {
    use diesel::dsl::sql;
    use diesel::update;

    let param_user_id = req.params()["user_id"]
        .parse::<i32>()
        .map_err(|err| err.chain(bad_request("invalid user_id")))?;
    let authenticated_user = req.authenticate()?;
    let conn = req.db_write()?;
    let user = authenticated_user.user();

    // need to check if current user matches user to be updated
    if user.id != param_user_id {
        return Err(bad_request("current user does not match requested user"));
    }

    conn.transaction(|| {
        let email: Email = update(Email::belonging_to(&user))
            .set(emails::token.eq(sql("DEFAULT")))
            .get_result(&*conn)
            .map_err(|_| bad_request("Email could not be found"))?;

        req.app()
            .emails
            .send_user_confirm(&email.email, &user.gh_login, &email.token)
    })?;

    ok_true()
}

/// Handles `PUT /me/email_notifications` route
pub fn update_email_notifications(req: &mut dyn RequestExt) -> EndpointResult {
    use self::crate_owners::dsl::*;
    use diesel::pg::upsert::excluded;

    #[derive(Deserialize)]
    struct CrateEmailNotifications {
        id: i32,
        email_notifications: bool,
    }

    let mut body = String::new();
    req.body().read_to_string(&mut body)?;
    let updates: HashMap<i32, bool> = serde_json::from_str::<Vec<CrateEmailNotifications>>(&body)
        .map_err(|_| bad_request("invalid json request"))?
        .iter()
        .map(|c| (c.id, c.email_notifications))
        .collect();

    let user_id = req.authenticate()?.user_id();
    let conn = req.db_write()?;

    // Build inserts from existing crates belonging to the current user
    let to_insert = CrateOwner::by_owner_kind(OwnerKind::User)
        .filter(owner_id.eq(user_id))
        .select((crate_id, owner_id, owner_kind, email_notifications))
        .load(&*conn)?
        .into_iter()
        // Remove records whose `email_notifications` will not change from their current value
        .map(
            |(c_id, o_id, o_kind, e_notifications): (i32, i32, i32, bool)| {
                let current_e_notifications = *updates.get(&c_id).unwrap_or(&e_notifications);
                (
                    crate_id.eq(c_id),
                    owner_id.eq(o_id),
                    owner_kind.eq(o_kind),
                    email_notifications.eq(current_e_notifications),
                )
            },
        )
        .collect::<Vec<_>>();

    // Upsert crate owners; this should only actually exectute updates
    diesel::insert_into(crate_owners)
        .values(&to_insert)
        .on_conflict((crate_id, owner_id, owner_kind))
        .do_update()
        .set(email_notifications.eq(excluded(email_notifications)))
        .execute(&*conn)?;

    ok_true()
}
<commit_msg>Use `db_read_prefer_primary()` for `GET /me` endpoint<commit_after>use std::collections::HashMap;

use crate::controllers::frontend_prelude::*;

use crate::controllers::helpers::*;

use crate::controllers::helpers::pagination::{Paginated, PaginationOptions};
use crate::models::{
    CrateOwner, Email, Follow, NewEmail, OwnerKind, User, Version, VersionOwnerAction,
};
use crate::schema::{crate_owners, crates, emails, follows, users, versions};
use crate::views::{EncodableMe, EncodablePrivateUser, EncodableVersion, OwnedCrate};

/// Handles the `GET /me` route.
pub fn me(req: &mut dyn RequestExt) -> EndpointResult {
    let user_id = req.authenticate()?.forbid_api_token_auth()?.user_id();
    let conn = req.db_read_prefer_primary()?;

    let (user, verified, email, verification_sent): (User, Option<bool>, Option<String>, bool) =
        users::table
            .find(user_id)
            .left_join(emails::table)
            .select((
                users::all_columns,
                emails::verified.nullable(),
                emails::email.nullable(),
                emails::token_generated_at.nullable().is_not_null(),
            ))
            .first(&*conn)?;

    let owned_crates = CrateOwner::by_owner_kind(OwnerKind::User)
        .inner_join(crates::table)
        .filter(crate_owners::owner_id.eq(user_id))
        .select((crates::id, crates::name, crate_owners::email_notifications))
        .order(crates::name.asc())
        .load(&*conn)?
        .into_iter()
        .map(|(id, name, email_notifications)| OwnedCrate {
            id,
            name,
            email_notifications,
        })
        .collect();

    let verified = verified.unwrap_or(false);
    let verification_sent = verified || verification_sent;
    Ok(req.json(&EncodableMe {
        user: EncodablePrivateUser::from(user, email, verified, verification_sent),
        owned_crates,
    }))
}

/// Handles the `GET /me/updates` route.
pub fn updates(req: &mut dyn RequestExt) -> EndpointResult {
    use diesel::dsl::any;

    let authenticated_user = req.authenticate()?.forbid_api_token_auth()?;
    let user = authenticated_user.user();

    let followed_crates = Follow::belonging_to(&user).select(follows::crate_id);
    let query = versions::table
        .inner_join(crates::table)
        .left_outer_join(users::table)
        .filter(crates::id.eq(any(followed_crates)))
        .order(versions::created_at.desc())
        .select((
            versions::all_columns,
            crates::name,
            users::all_columns.nullable(),
        ))
        .pages_pagination(PaginationOptions::builder().gather(req)?);
    let conn = req.db_write()?;
    let data: Paginated<(Version, String, Option<User>)> = query.load(&*conn)?;
    let more = data.next_page_params().is_some();
    let versions = data.iter().map(|(v, _, _)| v).cloned().collect::<Vec<_>>();
    let data = data
        .into_iter()
        .zip(VersionOwnerAction::for_versions(&conn, &versions)?.into_iter())
        .map(|((v, cn, pb), voas)| (v, cn, pb, voas));

    let versions = data
        .into_iter()
        .map(|(version, crate_name, published_by, actions)| {
            EncodableVersion::from(version, &crate_name, published_by, actions)
        })
        .collect::<Vec<_>>();

    Ok(req.json(&json!({
        "versions": versions,
        "meta": { "more": more },
    })))
}

/// Handles the `PUT /users/:user_id` route.
pub fn update_user(req: &mut dyn RequestExt) -> EndpointResult {
    use self::emails::user_id;
    use diesel::insert_into;

    let authenticated_user = req.authenticate()?;

    let mut body = String::new();
    req.body().read_to_string(&mut body)?;

    let param_user_id = &req.params()["user_id"];
    let conn = req.db_write()?;
    let user = authenticated_user.user();

    // need to check if current user matches user to be updated
    if &user.id.to_string() != param_user_id {
        return Err(bad_request("current user does not match requested user"));
    }

    #[derive(Deserialize)]
    struct UserUpdate {
        user: User,
    }

    #[derive(Deserialize)]
    struct User {
        email: Option<String>,
    }

    let user_update: UserUpdate =
        serde_json::from_str(&body).map_err(|_| bad_request("invalid json request"))?;

    let user_email = match &user_update.user.email {
        Some(email) => email.trim(),
        None => return Err(bad_request("empty email rejected")),
    };

    if user_email.is_empty() {
        return Err(bad_request("empty email rejected"));
    }

    conn.transaction::<_, Box<dyn AppError>, _>(|| {
        let new_email = NewEmail {
            user_id: user.id,
            email: user_email,
        };

        let token: String = insert_into(emails::table)
            .values(&new_email)
            .on_conflict(user_id)
            .do_update()
            .set(&new_email)
            .returning(emails::token)
            .get_result(&*conn)
            .map_err(|_| server_error("Error in creating token"))?;

        // This swallows any errors that occur while attempting to send the email. Some users have
        // an invalid email set in their GitHub profile, and we should let them sign in even though
        // we're trying to silently use their invalid address during signup and can't send them an
        // email. They'll then have to provide a valid email address.
        let _ = req
            .app()
            .emails
            .send_user_confirm(user_email, &user.gh_login, &token);

        Ok(())
    })?;

    ok_true()
}

/// Handles the `PUT /confirm/:email_token` route
pub fn confirm_user_email(req: &mut dyn RequestExt) -> EndpointResult {
    use diesel::update;

    let conn = req.db_write()?;
    let req_token = &req.params()["email_token"];

    let updated_rows = update(emails::table.filter(emails::token.eq(req_token)))
        .set(emails::verified.eq(true))
        .execute(&*conn)?;

    if updated_rows == 0 {
        return Err(bad_request("Email belonging to token not found."));
    }

    ok_true()
}

/// Handles `PUT /user/:user_id/resend` route
pub fn regenerate_token_and_send(req: &mut dyn RequestExt) -> EndpointResult {
    use diesel::dsl::sql;
    use diesel::update;

    let param_user_id = req.params()["user_id"]
        .parse::<i32>()
        .map_err(|err| err.chain(bad_request("invalid user_id")))?;
    let authenticated_user = req.authenticate()?;
    let conn = req.db_write()?;
    let user = authenticated_user.user();

    // need to check if current user matches user to be updated
    if user.id != param_user_id {
        return Err(bad_request("current user does not match requested user"));
    }

    conn.transaction(|| {
        let email: Email = update(Email::belonging_to(&user))
            .set(emails::token.eq(sql("DEFAULT")))
            .get_result(&*conn)
            .map_err(|_| bad_request("Email could not be found"))?;

        req.app()
            .emails
            .send_user_confirm(&email.email, &user.gh_login, &email.token)
    })?;

    ok_true()
}

/// Handles `PUT /me/email_notifications` route
pub fn update_email_notifications(req: &mut dyn RequestExt) -> EndpointResult {
    use self::crate_owners::dsl::*;
    use diesel::pg::upsert::excluded;

    #[derive(Deserialize)]
    struct CrateEmailNotifications {
        id: i32,
        email_notifications: bool,
    }

    let mut body = String::new();
    req.body().read_to_string(&mut body)?;
    let updates: HashMap<i32, bool> = serde_json::from_str::<Vec<CrateEmailNotifications>>(&body)
        .map_err(|_| bad_request("invalid json request"))?
        .iter()
        .map(|c| (c.id, c.email_notifications))
        .collect();

    let user_id = req.authenticate()?.user_id();
    let conn = req.db_write()?;

    // Build inserts from existing crates belonging to the current user
    let to_insert = CrateOwner::by_owner_kind(OwnerKind::User)
        .filter(owner_id.eq(user_id))
        .select((crate_id, owner_id, owner_kind, email_notifications))
        .load(&*conn)?
        .into_iter()
        // Remove records whose `email_notifications` will not change from their current value
        .map(
            |(c_id, o_id, o_kind, e_notifications): (i32, i32, i32, bool)| {
                let current_e_notifications = *updates.get(&c_id).unwrap_or(&e_notifications);
                (
                    crate_id.eq(c_id),
                    owner_id.eq(o_id),
                    owner_kind.eq(o_kind),
                    email_notifications.eq(current_e_notifications),
                )
            },
        )
        .collect::<Vec<_>>();

    // Upsert crate owners; this should only actually exectute updates
    diesel::insert_into(crate_owners)
        .values(&to_insert)
        .on_conflict((crate_id, owner_id, owner_kind))
        .do_update()
        .set(email_notifications.eq(excluded(email_notifications)))
        .execute(&*conn)?;

    ok_true()
}
<|endoftext|>