<commit_before>// This file is part of libfringe, a low-level green threading library.
// Copyright (c) whitequark <whitequark@whitequark.org>
// See the LICENSE file included in this distribution.

//! Generators.
//!
//! Generators allow repeatedly suspending the execution of a function,
//! returning a value to the caller, and resuming the suspended function
//! afterwards.

use core::marker::PhantomData;
use core::iter::Iterator;
use core::{ptr, mem};

use stack;
use context;

#[derive(Debug, Clone, Copy)]
pub enum State {
  /// Generator can be resumed. This is the initial state.
  Suspended,
  /// Generator cannot be resumed. This is the state of the generator after
  /// the generator function has returned.
  Finished
}

/// Generator wraps a function and allows suspending its execution more than
/// once, return a value each time.
///
/// It implements the Iterator trait. The first time `next()` is called,
/// the function is called as `f(yielder)`; every time `next()` is called afterwards,
/// the function is resumed. In either case, it runs until it suspends its execution
/// through `yielder.generate(val)`), in which case `next()` returns `Some(val)`, or
/// returns, in which case `next()` returns `None`. `next()` will return `None`
/// every time after that.
///
/// After the generator function returns, it is safe to reclaim the generator
/// stack using `unwrap()`.
///
/// `state()` can be used to determine whether the generator function has returned;
/// the state is `State::Suspended` after creation and suspension, and `State::Finished`
/// once the generator function returns.
///
/// # Example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut gen = Generator::new(stack, move |yielder| {
///   for i in 1..4 {
///     yielder.generate(i);
///   }
/// });
/// println!("{:?}", gen.next()); // prints Some(1)
/// println!("{:?}", gen.next()); // prints Some(2)
/// println!("{:?}", gen.next()); // prints Some(3)
/// println!("{:?}", gen.next()); // prints None
/// ```
#[derive(Debug)]
pub struct Generator<Item: Send, Stack: stack::Stack> {
  state:   State,
  context: context::Context<Stack>,
  phantom: PhantomData<Item>
}

impl<Item, Stack> Generator<Item, Stack>
    where Item: Send, Stack: stack::Stack {
  /// Creates a new generator.
  pub fn new<F>(stack: Stack, f: F) -> Generator<Item, Stack>
      where Stack: stack::GuardedStack, F: FnOnce(&mut Yielder<Item, Stack>) + Send {
    unsafe { Generator::unsafe_new(stack, f) }
  }

  /// Same as `new`, but does not require `stack` to have a guard page.
  ///
  /// This function is unsafe because the generator function can easily violate
  /// memory safety by overflowing the stack. It is useful in environments where
  /// guarded stacks do not exist, e.g. in absence of an MMU.
  pub unsafe fn unsafe_new<F>(stack: Stack, f: F) -> Generator<Item, Stack>
      where F: FnOnce(&mut Yielder<Item, Stack>) + Send {
    unsafe extern "C" fn generator_wrapper<Item, Stack, F>(info: usize) -> !
        where Item: Send, Stack: stack::Stack, F: FnOnce(&mut Yielder<Item, Stack>) {
      // Retrieve our environment from the callee and return control to it.
      let (mut yielder, f) = ptr::read(info as *mut (Yielder<Item, Stack>, F));
      let new_context = context::Context::swap(yielder.context, yielder.context, 0);
      // See Yielder::return_.
      yielder.context = new_context as *mut context::Context<Stack>;
      // Run the body of the generator.
      f(&mut yielder);
      // Past this point, the generator has dropped everything it has held.
      loop { yielder.return_(None) }
    }

    let mut generator = Generator {
      state:   State::Suspended,
      context: context::Context::new(stack, generator_wrapper::<Item, Stack, F>),
      phantom: PhantomData
    };

    // Transfer environment to the callee.
    let mut data = (Yielder::new(&mut generator.context), f);
    context::Context::swap(&mut generator.context, &generator.context,
                           &mut data as *mut (Yielder<Item, Stack>, F) as usize);
    mem::forget(data);

    generator
  }

  /// Returns the state of the generator.
  pub fn state(&self) -> State { self.state }

  /// Extracts the stack from a generator when the generator function has returned.
  /// If the generator function has not returned
  /// (i.e. `self.state() == State::Suspended`), panics.
  pub fn unwrap(self) -> Stack {
    match self.state {
      State::Suspended => panic!("Argh! Bastard! Don't touch that!"),
      State::Finished  => unsafe { self.context.unwrap() }
    }
  }
}

/// Yielder is an interface provided to every generator through which it
/// returns a value.
#[derive(Debug)]
pub struct Yielder<Item: Send, Stack: stack::Stack> {
  context: *mut context::Context<Stack>,
  phantom: PhantomData<Item>
}

impl<Item, Stack> Yielder<Item, Stack>
    where Item: Send, Stack: stack::Stack {
  fn new(context: *mut context::Context<Stack>) -> Yielder<Item, Stack> {
    Yielder {
      context: context,
      phantom: PhantomData
    }
  }

  #[inline(always)]
  fn return_(&mut self, mut val: Option<Item>) {
    unsafe {
      let new_context = context::Context::swap(self.context, self.context,
                                               &mut val as *mut Option<Item> as usize);
      // The generator can be moved (and with it, the context).
      // This changes the address of the context.
      // Thus, we update it after each swap.
      self.context = new_context as *mut context::Context<Stack>;
      // However, between this point and the next time we enter return_
      // the generator cannot be moved, as a &mut Generator is necessary
      // to resume the generator function.
    }
  }

  /// Suspends the generator and returns `Some(item)` from the `next()`
  /// invocation that resumed the generator.
  #[inline(always)]
  pub fn generate(&mut self, item: Item) {
    self.return_(Some(item))
  }
}

impl<Item, Stack> Iterator for Generator<Item, Stack>
    where Item: Send, Stack: stack::Stack {
  type Item = Item;

  /// Resumes the generator and return the next value it yields.
  /// If the generator function has returned, returns `None`.
  #[inline]
  fn next(&mut self) -> Option<Self::Item> {
    match self.state {
      State::Suspended => {
        let new_context = &mut self.context as *mut context::Context<Stack> as usize;
        let val = unsafe {
          ptr::read(context::Context::swap(&mut self.context, &self.context, new_context)
                    as *mut Option<Item>)
        };
        if let None = val { self.state = State::Finished }
        val
      }
      State::Finished => None
    }
  }
}
<commit_msg>Allow inlining Generator::state<commit_after>// This file is part of libfringe, a low-level green threading library.
// Copyright (c) whitequark <whitequark@whitequark.org>
// See the LICENSE file included in this distribution.

//! Generators.
//!
//! Generators allow repeatedly suspending the execution of a function,
//! returning a value to the caller, and resuming the suspended function
//! afterwards.

use core::marker::PhantomData;
use core::iter::Iterator;
use core::{ptr, mem};

use stack;
use context;

#[derive(Debug, Clone, Copy)]
pub enum State {
  /// Generator can be resumed. This is the initial state.
  Suspended,
  /// Generator cannot be resumed. This is the state of the generator after
  /// the generator function has returned.
  Finished
}

/// Generator wraps a function and allows suspending its execution more than
/// once, return a value each time.
///
/// It implements the Iterator trait. The first time `next()` is called,
/// the function is called as `f(yielder)`; every time `next()` is called afterwards,
/// the function is resumed. In either case, it runs until it suspends its execution
/// through `yielder.generate(val)`), in which case `next()` returns `Some(val)`, or
/// returns, in which case `next()` returns `None`. `next()` will return `None`
/// every time after that.
///
/// After the generator function returns, it is safe to reclaim the generator
/// stack using `unwrap()`.
///
/// `state()` can be used to determine whether the generator function has returned;
/// the state is `State::Suspended` after creation and suspension, and `State::Finished`
/// once the generator function returns.
///
/// # Example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut gen = Generator::new(stack, move |yielder| {
///   for i in 1..4 {
///     yielder.generate(i);
///   }
/// });
/// println!("{:?}", gen.next()); // prints Some(1)
/// println!("{:?}", gen.next()); // prints Some(2)
/// println!("{:?}", gen.next()); // prints Some(3)
/// println!("{:?}", gen.next()); // prints None
/// ```
#[derive(Debug)]
pub struct Generator<Item: Send, Stack: stack::Stack> {
  state:   State,
  context: context::Context<Stack>,
  phantom: PhantomData<Item>
}

impl<Item, Stack> Generator<Item, Stack>
    where Item: Send, Stack: stack::Stack {
  /// Creates a new generator.
  pub fn new<F>(stack: Stack, f: F) -> Generator<Item, Stack>
      where Stack: stack::GuardedStack, F: FnOnce(&mut Yielder<Item, Stack>) + Send {
    unsafe { Generator::unsafe_new(stack, f) }
  }

  /// Same as `new`, but does not require `stack` to have a guard page.
  ///
  /// This function is unsafe because the generator function can easily violate
  /// memory safety by overflowing the stack. It is useful in environments where
  /// guarded stacks do not exist, e.g. in absence of an MMU.
  pub unsafe fn unsafe_new<F>(stack: Stack, f: F) -> Generator<Item, Stack>
      where F: FnOnce(&mut Yielder<Item, Stack>) + Send {
    unsafe extern "C" fn generator_wrapper<Item, Stack, F>(info: usize) -> !
        where Item: Send, Stack: stack::Stack, F: FnOnce(&mut Yielder<Item, Stack>) {
      // Retrieve our environment from the callee and return control to it.
      let (mut yielder, f) = ptr::read(info as *mut (Yielder<Item, Stack>, F));
      let new_context = context::Context::swap(yielder.context, yielder.context, 0);
      // See Yielder::return_.
      yielder.context = new_context as *mut context::Context<Stack>;
      // Run the body of the generator.
      f(&mut yielder);
      // Past this point, the generator has dropped everything it has held.
      loop { yielder.return_(None) }
    }

    let mut generator = Generator {
      state:   State::Suspended,
      context: context::Context::new(stack, generator_wrapper::<Item, Stack, F>),
      phantom: PhantomData
    };

    // Transfer environment to the callee.
    let mut data = (Yielder::new(&mut generator.context), f);
    context::Context::swap(&mut generator.context, &generator.context,
                           &mut data as *mut (Yielder<Item, Stack>, F) as usize);
    mem::forget(data);

    generator
  }

  /// Returns the state of the generator.
  #[inline]
  pub fn state(&self) -> State { self.state }

  /// Extracts the stack from a generator when the generator function has returned.
  /// If the generator function has not returned
  /// (i.e. `self.state() == State::Suspended`), panics.
  pub fn unwrap(self) -> Stack {
    match self.state {
      State::Suspended => panic!("Argh! Bastard! Don't touch that!"),
      State::Finished  => unsafe { self.context.unwrap() }
    }
  }
}

/// Yielder is an interface provided to every generator through which it
/// returns a value.
#[derive(Debug)]
pub struct Yielder<Item: Send, Stack: stack::Stack> {
  context: *mut context::Context<Stack>,
  phantom: PhantomData<Item>
}

impl<Item, Stack> Yielder<Item, Stack>
    where Item: Send, Stack: stack::Stack {
  fn new(context: *mut context::Context<Stack>) -> Yielder<Item, Stack> {
    Yielder {
      context: context,
      phantom: PhantomData
    }
  }

  #[inline(always)]
  fn return_(&mut self, mut val: Option<Item>) {
    unsafe {
      let new_context = context::Context::swap(self.context, self.context,
                                               &mut val as *mut Option<Item> as usize);
      // The generator can be moved (and with it, the context).
      // This changes the address of the context.
      // Thus, we update it after each swap.
      self.context = new_context as *mut context::Context<Stack>;
      // However, between this point and the next time we enter return_
      // the generator cannot be moved, as a &mut Generator is necessary
      // to resume the generator function.
    }
  }

  /// Suspends the generator and returns `Some(item)` from the `next()`
  /// invocation that resumed the generator.
  #[inline(always)]
  pub fn generate(&mut self, item: Item) {
    self.return_(Some(item))
  }
}

impl<Item, Stack> Iterator for Generator<Item, Stack>
    where Item: Send, Stack: stack::Stack {
  type Item = Item;

  /// Resumes the generator and return the next value it yields.
  /// If the generator function has returned, returns `None`.
  #[inline]
  fn next(&mut self) -> Option<Self::Item> {
    match self.state {
      State::Suspended => {
        let new_context = &mut self.context as *mut context::Context<Stack> as usize;
        let val = unsafe {
          ptr::read(context::Context::swap(&mut self.context, &self.context, new_context)
                    as *mut Option<Item>)
        };
        if let None = val { self.state = State::Finished }
        val
      }
      State::Finished => None
    }
  }
}
<|endoftext|>


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

pub const MAXMSGLEN : usize = 65535;

#[derive(Copy, Clone)]
pub enum Token {E, S, Dhee, Dhes, Dhse, Dhss, Empty}

pub trait HandshakePattern {
    fn name(out: &mut [u8]) -> usize;
    fn get(premsg_pattern_i: &mut [Token],
           premsg_pattern_r: &mut [Token],
           msg_patterns: &mut [[Token; 8]; 5]);
}

#[derive(Debug)]
pub enum NoiseError {DecryptError}

struct SymmetricState<C: Cipher, H: Hash> {
    cipherstate : CipherState<C>,
    has_k: bool,
    has_psk: bool,
    h : [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    ck: [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    wtf : PhantomData<H>, /* So rust thinks I'm using H, this is ugly */
}

pub struct HandshakeState<P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> {
    symmetricstate: SymmetricState<C, H>,
    s: Option<D>,
    e: Option<D>,
    rs: Option<[u8; DHLEN]>,
    re: Option<[u8; DHLEN]>,
    my_turn_to_send : bool,
    message_patterns : [[Token; 8]; 5],
    message_index: usize,
    initiator: bool,
    rng : R,
    wtf : PhantomData<P>, /* So rust thinks I'm using P, this is ugly */
}

impl <C: Cipher, H: Hash> SymmetricState<C, H> {

    fn new(handshake_name: &[u8]) -> SymmetricState<C, H> {
        let mut hname = [0u8; MAXHASHLEN];
        if handshake_name.len() <= H::hash_len() {
            copy_memory(handshake_name, &mut hname);
        } else {
            let mut hasher = H::new();
            hasher.input(handshake_name);
            hasher.result(&mut hname);
        }
        SymmetricState{
            cipherstate: CipherState::new(&[0u8; CIPHERKEYLEN], 0),
            has_k : false,
            has_psk: false,
            h: hname,
            ck : hname,
            wtf: PhantomData::<H>
        }
    }

    fn mix_key(&mut self, data: &[u8]) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], data, &mut hkdf_output.0, &mut hkdf_output.1);
        copy_memory(&hkdf_output.0, &mut self.ck);
        self.cipherstate = CipherState::new(&hkdf_output.1[..CIPHERKEYLEN], 0);
        self.has_k = true;
    }

    fn mix_hash(&mut self, data: &[u8]) {
        let mut hasher = H::new();
        hasher.input(&self.h[..H::hash_len()]);
        hasher.input(data);
        hasher.result(&mut self.h);
    }

    fn mix_preshared_key(&mut self, psk: &[u8]) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], psk, &mut hkdf_output.0, &mut hkdf_output.1);
        copy_memory(&hkdf_output.0, &mut self.ck);
        self.mix_hash(&hkdf_output.1[..H::hash_len()]);
        self.has_psk = true;
    }

    fn encrypt_and_hash(&mut self, plaintext: &[u8], out: &mut [u8]) -> usize {
        if self.has_k {
            self.cipherstate.encrypt_ad(&self.h[..H::hash_len()], plaintext, out);
            self.mix_hash(&out[..plaintext.len() + TAGLEN]);
            return plaintext.len() + TAGLEN;
        } else {
            copy_memory(plaintext, out);
            self.mix_hash(plaintext);
            return plaintext.len();
        }
    }

    fn decrypt_and_hash(&mut self, data: &[u8], out: &mut [u8]) -> bool {
        if self.has_k {
            if !self.cipherstate.decrypt_ad(&self.h[..H::hash_len()], data, out) {
                return false;
            }
        }
        else {
            copy_memory(data, out);
        }
        self.mix_hash(data);
        true
    }

    fn split(&self, initiator: bool) -> (CipherState<C>, CipherState<C>) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], &[0u8; 0], &mut hkdf_output.0, &mut hkdf_output.1);
        let c1 = CipherState::<C>::new(&hkdf_output.0[..CIPHERKEYLEN], 0);
        let c2 = CipherState::<C>::new(&hkdf_output.1[..CIPHERKEYLEN], 0);
        if initiator { (c1, c2) } else { (c2, c1) }
    }

}

impl <P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> HandshakeState<P, D, C, H, R> {

    pub fn new(rng: R,
               initiator: bool,
               prologue: &[u8],
               new_preshared_key: Option<&[u8]>,
               new_s : Option<D>,
               new_e : Option<D>,
               new_rs: Option<[u8; DHLEN]>,
               new_re: Option<[u8; DHLEN]>) -> HandshakeState<P, D, C, H, R> {
        let mut handshake_name = [0u8; 128];
        let mut name_len: usize;
        if let Some(_) = new_preshared_key {
            copy_memory("NoisePSK_".as_bytes(), &mut handshake_name);
            name_len = 9;
        } else {
            copy_memory("Noise_".as_bytes(), &mut handshake_name);
            name_len = 6;
        }
        name_len += P::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += D::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += C::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += H::name(&mut handshake_name[name_len..]);

        let mut symmetricstate = SymmetricState::new(&handshake_name[..name_len]);

        symmetricstate.mix_hash(prologue);

        if let Some(preshared_key) = new_preshared_key {
            symmetricstate.mix_preshared_key(preshared_key);
        }

        let mut premsg_pattern_i = [Token::Empty; 2];
        let mut premsg_pattern_r = [Token::Empty; 2];
        let mut message_patterns = [[Token::Empty; 8]; 5];
        P::get(&mut premsg_pattern_i, &mut premsg_pattern_r, &mut message_patterns);
        if initiator {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        } else {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        }

        HandshakeState{
            symmetricstate: symmetricstate,
            s: new_s, e: new_e, rs: new_rs, re: new_re,
            my_turn_to_send: initiator,
            message_patterns: message_patterns,
            message_index: 0,
            initiator: initiator,
            rng: rng,
            wtf: PhantomData::<P>}
    }

    pub fn write_message(&mut self,
                         payload: &[u8],
                         message: &mut [u8]) -> (usize, Option<(CipherState<C>, CipherState<C>)>) {
        assert!(self.my_turn_to_send);
        let tokens = self.message_patterns[self.message_index];
        let mut last = false;
        if let Token::Empty = self.message_patterns[self.message_index+1][0] {
            last = true;
        }
        self.message_index += 1;

        let mut byte_index = 0;
        for token in &tokens {
            match *token {
                Token::E => {
                    self.e = Some(D::generate(&mut self.rng));
                    let pubkey = self.e.as_ref().unwrap().pubkey();
                    copy_memory(pubkey, &mut message[byte_index..]);
                    byte_index += DHLEN;
                    self.symmetricstate.mix_hash(&pubkey);
                    if self.symmetricstate.has_psk {
                        self.symmetricstate.mix_key(&pubkey);
                    }
                },
                Token::S => {
                    byte_index += self.symmetricstate.encrypt_and_hash(
                                        &self.s.as_ref().unwrap().pubkey(),
                                        &mut message[byte_index..]);
                },
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        self.my_turn_to_send = false;
        byte_index += self.symmetricstate.encrypt_and_hash(payload, &mut message[byte_index..]);
        assert!(byte_index <= MAXMSGLEN);
        match last {
            true => (byte_index, Some(self.symmetricstate.split(self.initiator))),
            false => (byte_index, None)
        }
    }

    pub fn read_message(&mut self,
                        message: &[u8],
                        payload: &mut [u8]) ->
                            Result<(usize, Option<(CipherState<C>, CipherState<C>)>),
                                    NoiseError> {
        assert!(self.my_turn_to_send == false);
        assert!(message.len() <= MAXMSGLEN);

        let tokens = self.message_patterns[self.message_index];
        let mut last = false;
        if let Token::Empty = self.message_patterns[self.message_index+1][0] {
            last = true;
        }
        self.message_index += 1;

        let mut ptr = message;
        for token in &tokens {
            match *token {
                Token::E => {
                    let mut pubkey = [0u8; DHLEN];
                    copy_memory(&ptr[..DHLEN], &mut pubkey);
                    ptr = &ptr[DHLEN..];
                    self.re = Some(pubkey);
                    self.symmetricstate.mix_hash(&pubkey);
                    if self.symmetricstate.has_psk {
                        self.symmetricstate.mix_key(&pubkey);
                    }
                },
                Token::S => {
                    let data = match self.symmetricstate.has_k {
                        true =>  {
                            let temp = &ptr[..DHLEN + TAGLEN];
                            ptr = &ptr[DHLEN + TAGLEN..];
                            temp
                        }
                        false => {
                            let temp = &ptr[..DHLEN];
                            ptr = &ptr[DHLEN..];
                            temp
                        }
                    };
                    let mut pubkey = [0u8; DHLEN];
                    if !self.symmetricstate.decrypt_and_hash(data, &mut pubkey) {
                        return Err(NoiseError::DecryptError);
                    }
                    self.rs = Some(pubkey);
                },
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        let payload_len = match self.symmetricstate.has_k {
            true => ptr.len() - TAGLEN,
            false => ptr.len()
        };
        if !self.symmetricstate.decrypt_and_hash(ptr, payload) {
            return Err(NoiseError::DecryptError);
        }
        self.my_turn_to_send = true;
        match last {
            true => Ok( (payload_len, Some(self.symmetricstate.split(self.initiator)) ) ),
            false => Ok( (payload_len, None) )
        }
    }

}

<commit_msg>Remove has_k<commit_after>
use std::marker::PhantomData;
use crypto_stuff::*;

pub const MAXMSGLEN : usize = 65535;

#[derive(Copy, Clone)]
pub enum Token {E, S, Dhee, Dhes, Dhse, Dhss, Empty}

pub trait HandshakePattern {
    fn name(out: &mut [u8]) -> usize;
    fn get(premsg_pattern_i: &mut [Token],
           premsg_pattern_r: &mut [Token],
           msg_patterns: &mut [[Token; 8]; 5]);
}

#[derive(Debug)]
pub enum NoiseError {DecryptError}

struct SymmetricState<C: Cipher, H: Hash> {
    cipherstate : Option< CipherState<C> >,
    h : [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    ck: [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    has_psk: bool,
    wtf : PhantomData<H>, /* So rust thinks I'm using H, this is ugly */
}

pub struct HandshakeState<P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> {
    symmetricstate: SymmetricState<C, H>,
    s: Option<D>,
    e: Option<D>,
    rs: Option<[u8; DHLEN]>,
    re: Option<[u8; DHLEN]>,
    my_turn_to_send : bool,
    message_patterns : [[Token; 8]; 5],
    message_index: usize,
    initiator: bool,
    rng : R,
    wtf : PhantomData<P>, /* So rust thinks I'm using P, this is ugly */
}

impl <C: Cipher, H: Hash> SymmetricState<C, H> {

    fn new(handshake_name: &[u8]) -> SymmetricState<C, H> {
        let mut hname = [0u8; MAXHASHLEN];
        if handshake_name.len() <= H::hash_len() {
            copy_memory(handshake_name, &mut hname);
        } else {
            let mut hasher = H::new();
            hasher.input(handshake_name);
            hasher.result(&mut hname);
        }
        SymmetricState{
            cipherstate: None,
            h: hname,
            ck : hname,
            has_psk: false,
            wtf: PhantomData::<H>
        }
    }

    fn mix_key(&mut self, data: &[u8]) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], data, &mut hkdf_output.0, &mut hkdf_output.1);
        copy_memory(&hkdf_output.0, &mut self.ck);
        self.cipherstate = Some(CipherState::new(&hkdf_output.1[..CIPHERKEYLEN], 0));
    }

    fn mix_hash(&mut self, data: &[u8]) {
        let mut hasher = H::new();
        hasher.input(&self.h[..H::hash_len()]);
        hasher.input(data);
        hasher.result(&mut self.h);
    }

    fn mix_preshared_key(&mut self, psk: &[u8]) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], psk, &mut hkdf_output.0, &mut hkdf_output.1);
        copy_memory(&hkdf_output.0, &mut self.ck);
        self.mix_hash(&hkdf_output.1[..H::hash_len()]);
        self.has_psk = true;
    }

    fn encrypt_and_hash(&mut self, plaintext: &[u8], out: &mut [u8]) -> usize {
        let output_len:usize;
        if let Some(ref mut cipherstate) = self.cipherstate {
            cipherstate.encrypt_ad(&self.h[..H::hash_len()], plaintext, out);
            output_len = plaintext.len() + TAGLEN;
        }
        else {
            copy_memory(plaintext, out);
            output_len = plaintext.len();
        }
        self.mix_hash(&out[..output_len]);
        output_len
    }

    fn decrypt_and_hash(&mut self, data: &[u8], out: &mut [u8]) -> bool {
        if let Some(ref mut cipherstate) = self.cipherstate {
            if !cipherstate.decrypt_ad(&self.h[..H::hash_len()], data, out) {
                return false;
            }
        }
        else {
            copy_memory(data, out);
        }
        self.mix_hash(data);
        true
    }

    fn split(&self, initiator: bool) -> (CipherState<C>, CipherState<C>) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], &[0u8; 0], &mut hkdf_output.0, &mut hkdf_output.1);
        let c1 = CipherState::<C>::new(&hkdf_output.0[..CIPHERKEYLEN], 0);
        let c2 = CipherState::<C>::new(&hkdf_output.1[..CIPHERKEYLEN], 0);
        if initiator { (c1, c2) } else { (c2, c1) }
    }

}

impl <P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> HandshakeState<P, D, C, H, R> {

    pub fn new(rng: R,
               initiator: bool,
               prologue: &[u8],
               new_preshared_key: Option<&[u8]>,
               new_s : Option<D>,
               new_e : Option<D>,
               new_rs: Option<[u8; DHLEN]>,
               new_re: Option<[u8; DHLEN]>) -> HandshakeState<P, D, C, H, R> {
        let mut handshake_name = [0u8; 128];
        let mut name_len: usize;
        if let Some(_) = new_preshared_key {
            copy_memory("NoisePSK_".as_bytes(), &mut handshake_name);
            name_len = 9;
        } else {
            copy_memory("Noise_".as_bytes(), &mut handshake_name);
            name_len = 6;
        }
        name_len += P::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += D::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += C::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += H::name(&mut handshake_name[name_len..]);

        let mut symmetricstate = SymmetricState::new(&handshake_name[..name_len]);

        symmetricstate.mix_hash(prologue);

        if let Some(preshared_key) = new_preshared_key {
            symmetricstate.mix_preshared_key(preshared_key);
        }

        let mut premsg_pattern_i = [Token::Empty; 2];
        let mut premsg_pattern_r = [Token::Empty; 2];
        let mut message_patterns = [[Token::Empty; 8]; 5];
        P::get(&mut premsg_pattern_i, &mut premsg_pattern_r, &mut message_patterns);
        if initiator {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        } else {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        }

        HandshakeState{
            symmetricstate: symmetricstate,
            s: new_s, e: new_e, rs: new_rs, re: new_re,
            my_turn_to_send: initiator,
            message_patterns: message_patterns,
            message_index: 0,
            initiator: initiator,
            rng: rng,
            wtf: PhantomData::<P>}
    }

    pub fn write_message(&mut self,
                         payload: &[u8],
                         message: &mut [u8]) -> (usize, Option<(CipherState<C>, CipherState<C>)>) {
        assert!(self.my_turn_to_send);
        let tokens = self.message_patterns[self.message_index];
        let mut last = false;
        if let Token::Empty = self.message_patterns[self.message_index+1][0] {
            last = true;
        }
        self.message_index += 1;

        let mut byte_index = 0;
        for token in &tokens {
            match *token {
                Token::E => {
                    self.e = Some(D::generate(&mut self.rng));
                    let pubkey = self.e.as_ref().unwrap().pubkey();
                    copy_memory(pubkey, &mut message[byte_index..]);
                    byte_index += DHLEN;
                    self.symmetricstate.mix_hash(&pubkey);
                    if self.symmetricstate.has_psk {
                        self.symmetricstate.mix_key(&pubkey);
                    }
                },
                Token::S => {
                    byte_index += self.symmetricstate.encrypt_and_hash(
                                        &self.s.as_ref().unwrap().pubkey(),
                                        &mut message[byte_index..]);
                },
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        self.my_turn_to_send = false;
        byte_index += self.symmetricstate.encrypt_and_hash(payload, &mut message[byte_index..]);
        assert!(byte_index <= MAXMSGLEN);
        match last {
            true => (byte_index, Some(self.symmetricstate.split(self.initiator))),
            false => (byte_index, None)
        }
    }

    pub fn read_message(&mut self,
                        message: &[u8],
                        payload: &mut [u8]) ->
                            Result<(usize, Option<(CipherState<C>, CipherState<C>)>),
                                    NoiseError> {
        assert!(self.my_turn_to_send == false);
        assert!(message.len() <= MAXMSGLEN);

        let tokens = self.message_patterns[self.message_index];
        let mut last = false;
        if let Token::Empty = self.message_patterns[self.message_index+1][0] {
            last = true;
        }
        self.message_index += 1;

        let mut ptr = message;
        for token in &tokens {
            match *token {
                Token::E => {
                    let mut pubkey = [0u8; DHLEN];
                    copy_memory(&ptr[..DHLEN], &mut pubkey);
                    ptr = &ptr[DHLEN..];
                    self.re = Some(pubkey);
                    self.symmetricstate.mix_hash(&pubkey);
                    if self.symmetricstate.has_psk {
                        self.symmetricstate.mix_key(&pubkey);
                    }
                },
                Token::S => {
                    let data = match self.symmetricstate.cipherstate {
                        Some(_) =>  {
                            let temp = &ptr[..DHLEN + TAGLEN];
                            ptr = &ptr[DHLEN + TAGLEN..];
                            temp
                        }
                        None => {
                            let temp = &ptr[..DHLEN];
                            ptr = &ptr[DHLEN..];
                            temp
                        }
                    };
                    let mut pubkey = [0u8; DHLEN];
                    if !self.symmetricstate.decrypt_and_hash(data, &mut pubkey) {
                        return Err(NoiseError::DecryptError);
                    }
                    self.rs = Some(pubkey);
                },
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        let payload_len = match self.symmetricstate.cipherstate {
            Some(_) => ptr.len() - TAGLEN,
            None => ptr.len()
        };
        if !self.symmetricstate.decrypt_and_hash(ptr, payload) {
            return Err(NoiseError::DecryptError);
        }
        self.my_turn_to_send = true;
        match last {
            true => Ok( (payload_len, Some(self.symmetricstate.split(self.initiator)) ) ),
            false => Ok( (payload_len, None) )
        }
    }

}

<|endoftext|>


<commit_before>use std::cmp::Ordering;

/// An *objective* defines a total ordering relation and a distance
/// metric on `solutions` in order to be able to compare any two
/// solutions according to the following two questions:
///
/// - "which solution is better" (total order)
///
/// - "how much better" (distance metric)
///
/// Objectives can be seen as a projection of a (possibly) multi-variate
/// solution value to a scalar value. There can be any number of
/// different projections (objectives) for any given solution value. Of
/// course solution values need not be multi-variate.
///
/// We use the term "solution" here, ignoring the fact that in pratice
/// we often have to evaluate the "fitness" of a solution prior of being
/// able to define any useful ordering relation or distance metric. As
/// the fitness generally is a function of the solution, this is more or
/// less an implementation detail or that of an optimization. Nothing
/// prevents you from using the fitness value here as the solution
/// value.

pub trait Objective {
    /// The solution value type that we define the objective on.
    type Solution;

    /// The output type of the distance metric.
    type Distance: Sized;

    /// An objective defines a total order between two solution values.
    fn total_order(&self, a: &Self::Solution, b: &Self::Solution) -> Ordering;

    /// An objective defines a distance metric between two solution
    /// values. Distance values can be negative, i.e. the caller is
    /// responsible for obtaining absolute values.
    fn distance(&self, a: &Self::Solution, b: &Self::Solution) -> Self::Distance;
}

#[test]
fn test_objectives() {
    use test_helper_objective::{Objective1, Objective2, Objective3, Tuple};

    let a = &Tuple(1, 2);
    let b = &Tuple(2, 1);
    assert_eq!(Ordering::Less, Objective1.total_order(a, b));
    assert_eq!(Ordering::Greater, Objective2.total_order(a, b));
    assert_eq!(Ordering::Equal, Objective3.total_order(a, b));

    assert_eq!(-1.0, Objective1.distance(a, b));
    assert_eq!(1.0, Objective2.distance(a, b));
    assert_eq!(0.0, Objective3.distance(a, b));
}
<commit_msg>Reword documentation a bit<commit_after>use std::cmp::Ordering;

/// An *objective* defines a *total ordering relation* and a *distance
/// metric* on a set of `solutions`. Given any two solutions, an
/// objective answers the following two questions:
///
/// - "which solution is the better one" (total order)
///
/// - "how similar are the two solutions" (distance metric)
///
/// Objectives can be seen as a projection of a (possibly) multi-variate
/// solution value to a scalar value. There can be any number of
/// different projections (objectives) for any given solution value. Of
/// course solution values need not be multi-variate.
///
/// We use the term "solution" here, ignoring the fact that in pratice
/// we often have to evaluate the "fitness" of a solution prior of being
/// able to define any useful ordering relation or distance metric. As
/// the fitness generally is a function of the solution, this is more or
/// less an implementation detail or that of an optimization. Nothing
/// prevents you from using the fitness value here as the solution
/// value.

pub trait Objective {
    /// The solution value type that we define the objective on.
    type Solution;

    /// The output type of the distance metric.
    type Distance: Sized;

    /// An objective defines a total ordering between any two solution
    /// values.
    ///
    /// This answers the question, is solution `a` better, equal or
    /// worse than solution `b`, according to the objective.
    fn total_order(&self, a: &Self::Solution, b: &Self::Solution) -> Ordering;

    /// An objective defines a distance metric between any two solution
    /// values.
    ///
    /// The distance metric answer the question, how similar the
    /// solutions `a` and `b` are, according to the objective.  A zero
    /// value would mean, that both solutions are in fact the same,
    /// according to the objective. Larger magnitudes would mean "less
    /// similar".
    ///
    /// Note: Distance values can be negative, i.e. the caller is
    /// responsible for obtaining absolute values.
    fn distance(&self, a: &Self::Solution, b: &Self::Solution) -> Self::Distance;
}

#[test]
fn test_objectives() {
    use test_helper_objective::{Objective1, Objective2, Objective3, Tuple};

    let a = &Tuple(1, 2);
    let b = &Tuple(2, 1);
    assert_eq!(Ordering::Less, Objective1.total_order(a, b));
    assert_eq!(Ordering::Greater, Objective2.total_order(a, b));
    assert_eq!(Ordering::Equal, Objective3.total_order(a, b));

    assert_eq!(-1.0, Objective1.distance(a, b));
    assert_eq!(1.0, Objective2.distance(a, b));
    assert_eq!(0.0, Objective3.distance(a, b));
}
<|endoftext|>


<commit_before>use winapi::shared::minwindef::{WPARAM, LPARAM};
use winapi::um::winuser::{ES_AUTOVSCROLL, ES_AUTOHSCROLL, WS_VISIBLE, WS_DISABLED, WS_TABSTOP};
use crate::win32::window_helper as wh;
use crate::{Font, NwgError};
use super::{ControlBase, ControlHandle};
use std::ops::Range;

const NOT_BOUND: &'static str = "TextBox is not yet bound to a winapi object";
const BAD_HANDLE: &'static str = "INTERNAL ERROR: TextBox handle is not HWND!";

bitflags! {

    /**
        The text box flags

        * VSCROLL:  The text box has a vertical scrollbar
        * HSCROLL:  The text box has a horizontal scrollbar
        * VISIBLE:  The text box is immediatly visible after creation
        * DISABLED: The text box cannot be interacted with by the user. It also has a grayed out look.
        * TAB_STOP: The text box can be selected using tab navigation
    */
    pub struct TextBoxFlags: u32 {
        const VSCROLL = ES_AUTOVSCROLL;
        const HSCROLL = ES_AUTOHSCROLL;
        const VISIBLE = WS_VISIBLE;
        const DISABLED = WS_DISABLED;
        const TAB_STOP = WS_TABSTOP;
    }
}

/**
An edit control is a rectangular control window to permit the user to enter and edit text by typing on the keyboard
This control allow multi line input. For a single line of text, use `TextInput`.

Requires the `textbox` feature.

Note: Use `\r\n` to input a new line not just `\n`.

**Builder parameters:**
  * `parent`:   **Required.** The text box parent container.
  * `text`:     The text box text.
  * `size`:     The text box size.
  * `position`: The text box position.
  * `flags`:    A combination of the TextBoxFlags values.
  * `font`:     The font used for the text box text
  * `limit`:    The maximum number of character that can be inserted in the control
  * `readonly`: If the textbox should allow user input or not
  * `focus`:    The control receive focus after being created

**Control events:**
  * `OnTextInput`: When a TextBox value is changed
  * `MousePress(_)`: Generic mouse press events on the button
  * `OnMouseMove`: Generic mouse mouse event
  * `OnMouseWheel`: Generic mouse wheel event

```rust
use native_windows_gui as nwg;
fn build_box(tbox: &mut nwg::TextBox, window: &nwg::Window, font: &nwg::Font) {
    nwg::TextBox::builder()
        .text("Hello")
        .font(Some(font))
        .parent(window)
        .build(tbox);
}
```
*/
#[derive(Default, PartialEq, Eq)]
pub struct TextBox {
    pub handle: ControlHandle
}

impl TextBox {

    pub fn builder<'a>() -> TextBoxBuilder<'a> {
        TextBoxBuilder {
            text: "",
            size: (100, 25),
            position: (0, 0),
            flags: None,
            limit: 0,
            readonly: false,
            focus: false,
            font: None,
            parent: None
        }
    }

    /// Return the font of the control
    pub fn font(&self) -> Option<Font> {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        let font_handle = wh::get_window_font(handle);
        if font_handle.is_null() {
            None
        } else {
            Some(Font { handle: font_handle })
        }
    }

    /// Set the font of the control
    pub fn set_font(&self, font: Option<&Font>) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_font(handle, font.map(|f| f.handle), true); }
    }

    /// Return the number of maximum character allowed in this text input
    pub fn limit(&self) -> u32 {
        use winapi::um::winuser::EM_GETLIMITTEXT;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_GETLIMITTEXT as u32, 0, 0) as u32
    }

    /// Set the number of maximum character allowed in this text input
    pub fn set_limit(&self, limit: usize) {
        use winapi::um::winuser::EM_SETLIMITTEXT;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_SETLIMITTEXT as u32, limit, 0);
    }

    /// Check if the content of the text input was modified after it's creation
    pub fn modified(&self) -> bool {
        use winapi::um::winuser::EM_GETMODIFY;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_GETMODIFY as u32, 0, 0) != 0
    }

    /// Manually set modified flag of the text input
    pub fn set_modified(&self, e: bool) {
        use winapi::um::winuser::EM_SETMODIFY;
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_SETMODIFY as u32, e as usize, 0);
    }

    /// Undo the last action by the user in the control
    pub fn undo(&self) {
        use winapi::um::winuser::EM_UNDO;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_UNDO as u32, 0, 0);
    }

    /// Return the selected range of characters by the user in the text input
    pub fn selection(&self) -> Range<u32> {
        use winapi::um::winuser::EM_GETSEL;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        let (mut out1, mut out2) = (0u32, 0u32);
        let (ptr1, ptr2) = (&mut out1 as *mut u32, &mut out2 as *mut u32);
        wh::send_message(handle, EM_GETSEL as u32, ptr1 as WPARAM, ptr2 as LPARAM);

        Range { start: out1 as u32, end: out2 as u32 }
    }

    /// Return the selected range of characters by the user in the text input
    pub fn set_selection(&self, r: Range<u32>) {
        use winapi::um::winuser::EM_SETSEL;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        wh::send_message(handle, EM_SETSEL as u32, r.start as usize, r.end as isize);
    }

    /// Return the length of the user input in the control. This is better than test.len() as it
    /// does not allocate a string in memory
    pub fn len(&self) -> u32 {
        use winapi::um::winuser::EM_LINELENGTH;
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_LINELENGTH as u32, 0, 0) as u32
    }

    /// Return true if the TextInput value cannot be edited. Retrurn false otherwise.
    /// A user can still copy text from a readonly TextEdit (unlike disabled)
    pub fn readonly(&self) -> bool {
        use winapi::um::winuser::ES_READONLY;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        wh::get_style(handle) & ES_READONLY == ES_READONLY
    }

    /// Set the readonly flag of the text input
    /// A user can still copy text from a readonly TextEdit (unlike disabled)
    pub fn set_readonly(&self, r: bool) {
        use winapi::um::winuser::EM_SETREADONLY;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        wh::send_message(handle, EM_SETREADONLY as u32, r as WPARAM, 0);
    }

    /// Remove all text from the textbox
    pub fn clear(&self) {
        self.set_text("");
    }

    /// Return true if the control currently has the keyboard focus
    pub fn focus(&self) -> bool {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_focus(handle) }
    }

    /// Set the keyboard focus on the button
    pub fn set_focus(&self) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_focus(handle); }
    }

    /// Return true if the control user can interact with the control, return false otherwise
    pub fn enabled(&self) -> bool {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_window_enabled(handle) }
    }

    /// Enable or disable the control
    pub fn set_enabled(&self, v: bool) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_enabled(handle, v) }
    }

    /// Return true if the control is visible to the user. Will return true even if the
    /// control is outside of the parent client view (ex: at the position (10000, 10000))
    pub fn visible(&self) -> bool {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_window_visibility(handle) }
    }

    /// Show or hide the control to the user
    pub fn set_visible(&self, v: bool) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_visibility(handle, v) }
    }

    /// Return the size of the button in the parent window
    pub fn size(&self) -> (u32, u32) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_window_size(handle) }
    }

    /// Set the size of the button in the parent window
    pub fn set_size(&self, x: u32, y: u32) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_size(handle, x, y, false) }
    }

    /// Return the position of the button in the parent window
    pub fn position(&self) -> (i32, i32) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_window_position(handle) }
    }

    /// Set the position of the button in the parent window
    pub fn set_position(&self, x: i32, y: i32) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_position(handle, x, y) }
    }

    /// Return the text displayed in the TextInput
    pub fn text(&self) -> String {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_window_text(handle) }
    }

    /// Set the text displayed in the TextInput
    pub fn set_text<'a>(&self, v: &'a str) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_text(handle, v) }
    }

    /// Winapi class name used during control creation
    pub fn class_name(&self) -> &'static str {
        "EDIT"
    }

    /// Winapi base flags used during window creation
    pub fn flags(&self) -> u32 {
        WS_VISIBLE | ES_AUTOVSCROLL | ES_AUTOHSCROLL | WS_TABSTOP
    }

    /// Winapi flags required by the control
    pub fn forced_flags(&self) -> u32 {
        use winapi::um::winuser::{WS_BORDER, WS_CHILD, ES_MULTILINE};

        WS_BORDER | WS_CHILD | ES_MULTILINE
    }

}

impl Drop for TextBox {
    fn drop(&mut self) {
        self.handle.destroy();
    }
}
pub struct TextBoxBuilder<'a> {
    text: &'a str,
    size: (i32, i32),
    position: (i32, i32),
    flags: Option<TextBoxFlags>,
    limit: usize,
    readonly: bool,
    focus: bool,
    font: Option<&'a Font>,
    parent: Option<ControlHandle>
}

impl<'a> TextBoxBuilder<'a> {

    pub fn flags(mut self, flags: TextBoxFlags) -> TextBoxBuilder<'a> {
        self.flags = Some(flags);
        self
    }

    pub fn text(mut self, text: &'a str) -> TextBoxBuilder<'a> {
        self.text = text;
        self
    }

    pub fn size(mut self, size: (i32, i32)) -> TextBoxBuilder<'a> {
        self.size = size;
        self
    }

    pub fn position(mut self, pos: (i32, i32)) -> TextBoxBuilder<'a> {
        self.position = pos;
        self
    }

    pub fn limit(mut self, limit: usize) -> TextBoxBuilder<'a> {
        self.limit = limit;
        self
    }

    pub fn readonly(mut self, read: bool) -> TextBoxBuilder<'a> {
        self.readonly = read;
        self
    }

    pub fn focus(mut self, focus: bool) -> TextBoxBuilder<'a> {
        self.focus = focus;
        self
    }

    pub fn font(mut self, font: Option<&'a Font>) -> TextBoxBuilder<'a> {
        self.font = font;
        self
    }

    pub fn parent<C: Into<ControlHandle>>(mut self, p: C) -> TextBoxBuilder<'a> {
        self.parent = Some(p.into());
        self
    }

    pub fn build(self, out: &mut TextBox) -> Result<(), NwgError> {
        let flags = self.flags.map(|f| f.bits()).unwrap_or(out.flags());

        let parent = match self.parent {
            Some(p) => Ok(p),
            None => Err(NwgError::no_parent("TextBox"))
        }?;

        out.handle = ControlBase::build_hwnd()
            .class_name(out.class_name())
            .forced_flags(out.forced_flags())
            .flags(flags)
            .size(self.size)
            .position(self.position)
            .text(self.text)
            .parent(Some(parent))
            .build()?;

        if self.limit > 0 {
            out.set_limit(self.limit);
        }

        if self.readonly {
            out.set_readonly(self.readonly);
        }

        if self.focus {
            out.set_focus();
        }

        if self.font.is_some() {
            out.set_font(self.font);
        } else {
            out.set_font(Font::global_default().as_ref());
        }

        Ok(())
    }

}
<commit_msg>Fix Scroll Bar Flags<commit_after>use winapi::shared::minwindef::{WPARAM, LPARAM};
use winapi::um::winuser::{WS_VSCROLL, WS_HSCROLL, WS_VISIBLE, WS_DISABLED, WS_TABSTOP};
use crate::win32::window_helper as wh;
use crate::{Font, NwgError};
use super::{ControlBase, ControlHandle};
use std::ops::Range;

const NOT_BOUND: &'static str = "TextBox is not yet bound to a winapi object";
const BAD_HANDLE: &'static str = "INTERNAL ERROR: TextBox handle is not HWND!";

bitflags! {

    /**
        The text box flags

        * VSCROLL:  The text box has a vertical scrollbar
        * HSCROLL:  The text box has a horizontal scrollbar
        * VISIBLE:  The text box is immediatly visible after creation
        * DISABLED: The text box cannot be interacted with by the user. It also has a grayed out look.
        * TAB_STOP: The text box can be selected using tab navigation
    */
    pub struct TextBoxFlags: u32 {
        const VSCROLL = WS_VSCROLL;
        const HSCROLL = WS_HSCROLL;
        const VISIBLE = WS_VISIBLE;
        const DISABLED = WS_DISABLED;
        const TAB_STOP = WS_TABSTOP;
    }
}

/**
An edit control is a rectangular control window to permit the user to enter and edit text by typing on the keyboard
This control allow multi line input. For a single line of text, use `TextInput`.

Requires the `textbox` feature.

Note: Use `\r\n` to input a new line not just `\n`.

**Builder parameters:**
  * `parent`:   **Required.** The text box parent container.
  * `text`:     The text box text.
  * `size`:     The text box size.
  * `position`: The text box position.
  * `flags`:    A combination of the TextBoxFlags values.
  * `font`:     The font used for the text box text
  * `limit`:    The maximum number of character that can be inserted in the control
  * `readonly`: If the textbox should allow user input or not
  * `focus`:    The control receive focus after being created

**Control events:**
  * `OnTextInput`: When a TextBox value is changed
  * `MousePress(_)`: Generic mouse press events on the button
  * `OnMouseMove`: Generic mouse mouse event
  * `OnMouseWheel`: Generic mouse wheel event

```rust
use native_windows_gui as nwg;
fn build_box(tbox: &mut nwg::TextBox, window: &nwg::Window, font: &nwg::Font) {
    nwg::TextBox::builder()
        .text("Hello")
        .font(Some(font))
        .parent(window)
        .build(tbox);
}
```
*/
#[derive(Default, PartialEq, Eq)]
pub struct TextBox {
    pub handle: ControlHandle
}

impl TextBox {

    pub fn builder<'a>() -> TextBoxBuilder<'a> {
        TextBoxBuilder {
            text: "",
            size: (100, 25),
            position: (0, 0),
            flags: None,
            limit: 0,
            readonly: false,
            focus: false,
            font: None,
            parent: None
        }
    }

    /// Return the font of the control
    pub fn font(&self) -> Option<Font> {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        let font_handle = wh::get_window_font(handle);
        if font_handle.is_null() {
            None
        } else {
            Some(Font { handle: font_handle })
        }
    }

    /// Set the font of the control
    pub fn set_font(&self, font: Option<&Font>) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_font(handle, font.map(|f| f.handle), true); }
    }

    /// Return the number of maximum character allowed in this text input
    pub fn limit(&self) -> u32 {
        use winapi::um::winuser::EM_GETLIMITTEXT;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_GETLIMITTEXT as u32, 0, 0) as u32
    }

    /// Set the number of maximum character allowed in this text input
    pub fn set_limit(&self, limit: usize) {
        use winapi::um::winuser::EM_SETLIMITTEXT;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_SETLIMITTEXT as u32, limit, 0);
    }

    /// Check if the content of the text input was modified after it's creation
    pub fn modified(&self) -> bool {
        use winapi::um::winuser::EM_GETMODIFY;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_GETMODIFY as u32, 0, 0) != 0
    }

    /// Manually set modified flag of the text input
    pub fn set_modified(&self, e: bool) {
        use winapi::um::winuser::EM_SETMODIFY;
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_SETMODIFY as u32, e as usize, 0);
    }

    /// Undo the last action by the user in the control
    pub fn undo(&self) {
        use winapi::um::winuser::EM_UNDO;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_UNDO as u32, 0, 0);
    }

    /// Return the selected range of characters by the user in the text input
    pub fn selection(&self) -> Range<u32> {
        use winapi::um::winuser::EM_GETSEL;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        let (mut out1, mut out2) = (0u32, 0u32);
        let (ptr1, ptr2) = (&mut out1 as *mut u32, &mut out2 as *mut u32);
        wh::send_message(handle, EM_GETSEL as u32, ptr1 as WPARAM, ptr2 as LPARAM);

        Range { start: out1 as u32, end: out2 as u32 }
    }

    /// Return the selected range of characters by the user in the text input
    pub fn set_selection(&self, r: Range<u32>) {
        use winapi::um::winuser::EM_SETSEL;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        wh::send_message(handle, EM_SETSEL as u32, r.start as usize, r.end as isize);
    }

    /// Return the length of the user input in the control. This is better than test.len() as it
    /// does not allocate a string in memory
    pub fn len(&self) -> u32 {
        use winapi::um::winuser::EM_LINELENGTH;
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        wh::send_message(handle, EM_LINELENGTH as u32, 0, 0) as u32
    }

    /// Return true if the TextInput value cannot be edited. Retrurn false otherwise.
    /// A user can still copy text from a readonly TextEdit (unlike disabled)
    pub fn readonly(&self) -> bool {
        use winapi::um::winuser::ES_READONLY;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        wh::get_style(handle) & ES_READONLY == ES_READONLY
    }

    /// Set the readonly flag of the text input
    /// A user can still copy text from a readonly TextEdit (unlike disabled)
    pub fn set_readonly(&self, r: bool) {
        use winapi::um::winuser::EM_SETREADONLY;

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        wh::send_message(handle, EM_SETREADONLY as u32, r as WPARAM, 0);
    }

    /// Remove all text from the textbox
    pub fn clear(&self) {
        self.set_text("");
    }

    /// Return true if the control currently has the keyboard focus
    pub fn focus(&self) -> bool {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_focus(handle) }
    }

    /// Set the keyboard focus on the button
    pub fn set_focus(&self) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_focus(handle); }
    }

    /// Return true if the control user can interact with the control, return false otherwise
    pub fn enabled(&self) -> bool {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_window_enabled(handle) }
    }

    /// Enable or disable the control
    pub fn set_enabled(&self, v: bool) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_enabled(handle, v) }
    }

    /// Return true if the control is visible to the user. Will return true even if the
    /// control is outside of the parent client view (ex: at the position (10000, 10000))
    pub fn visible(&self) -> bool {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_window_visibility(handle) }
    }

    /// Show or hide the control to the user
    pub fn set_visible(&self, v: bool) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_visibility(handle, v) }
    }

    /// Return the size of the button in the parent window
    pub fn size(&self) -> (u32, u32) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_window_size(handle) }
    }

    /// Set the size of the button in the parent window
    pub fn set_size(&self, x: u32, y: u32) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_size(handle, x, y, false) }
    }

    /// Return the position of the button in the parent window
    pub fn position(&self) -> (i32, i32) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_window_position(handle) }
    }

    /// Set the position of the button in the parent window
    pub fn set_position(&self, x: i32, y: i32) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_position(handle, x, y) }
    }

    /// Return the text displayed in the TextInput
    pub fn text(&self) -> String {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::get_window_text(handle) }
    }

    /// Set the text displayed in the TextInput
    pub fn set_text<'a>(&self, v: &'a str) {
        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);
        unsafe { wh::set_window_text(handle, v) }
    }

    /// Winapi class name used during control creation
    pub fn class_name(&self) -> &'static str {
        "EDIT"
    }

    /// Winapi base flags used during window creation
    pub fn flags(&self) -> u32 {
        WS_VISIBLE | WS_VSCROLL | WS_HSCROLL | WS_TABSTOP
    }

    /// Winapi flags required by the control
    pub fn forced_flags(&self) -> u32 {
        use winapi::um::winuser::{WS_BORDER, WS_CHILD, ES_MULTILINE};

        WS_BORDER | WS_CHILD | ES_MULTILINE
    }

}

impl Drop for TextBox {
    fn drop(&mut self) {
        self.handle.destroy();
    }
}
pub struct TextBoxBuilder<'a> {
    text: &'a str,
    size: (i32, i32),
    position: (i32, i32),
    flags: Option<TextBoxFlags>,
    limit: usize,
    readonly: bool,
    focus: bool,
    font: Option<&'a Font>,
    parent: Option<ControlHandle>
}

impl<'a> TextBoxBuilder<'a> {

    pub fn flags(mut self, flags: TextBoxFlags) -> TextBoxBuilder<'a> {
        self.flags = Some(flags);
        self
    }

    pub fn text(mut self, text: &'a str) -> TextBoxBuilder<'a> {
        self.text = text;
        self
    }

    pub fn size(mut self, size: (i32, i32)) -> TextBoxBuilder<'a> {
        self.size = size;
        self
    }

    pub fn position(mut self, pos: (i32, i32)) -> TextBoxBuilder<'a> {
        self.position = pos;
        self
    }

    pub fn limit(mut self, limit: usize) -> TextBoxBuilder<'a> {
        self.limit = limit;
        self
    }

    pub fn readonly(mut self, read: bool) -> TextBoxBuilder<'a> {
        self.readonly = read;
        self
    }

    pub fn focus(mut self, focus: bool) -> TextBoxBuilder<'a> {
        self.focus = focus;
        self
    }

    pub fn font(mut self, font: Option<&'a Font>) -> TextBoxBuilder<'a> {
        self.font = font;
        self
    }

    pub fn parent<C: Into<ControlHandle>>(mut self, p: C) -> TextBoxBuilder<'a> {
        self.parent = Some(p.into());
        self
    }

    pub fn build(self, out: &mut TextBox) -> Result<(), NwgError> {
        let flags = self.flags.map(|f| f.bits()).unwrap_or(out.flags());

        let parent = match self.parent {
            Some(p) => Ok(p),
            None => Err(NwgError::no_parent("TextBox"))
        }?;

        out.handle = ControlBase::build_hwnd()
            .class_name(out.class_name())
            .forced_flags(out.forced_flags())
            .flags(flags)
            .size(self.size)
            .position(self.position)
            .text(self.text)
            .parent(Some(parent))
            .build()?;

        if self.limit > 0 {
            out.set_limit(self.limit);
        }

        if self.readonly {
            out.set_readonly(self.readonly);
        }

        if self.focus {
            out.set_focus();
        }

        if self.font.is_some() {
            out.set_font(self.font);
        } else {
            out.set_font(Font::global_default().as_ref());
        }

        Ok(())
    }

}
<|endoftext|>


<commit_before>use alloc::borrow::Cow;
use alloc::vec::Vec;
use core::fmt::{LowerHex, Debug};

use scroll::{self, Pread};
use scroll::ctx::TryFromCtx;
use error;

use pe::section_table;
use pe::utils;
use pe::data_directories;

pub const IMPORT_BY_ORDINAL_32: u32 = 0x8000_0000;
pub const IMPORT_BY_ORDINAL_64: u64 = 0x8000_0000_0000_0000;
pub const IMPORT_RVA_MASK_32: u32 = 0x8fff_ffff;
pub const IMPORT_RVA_MASK_64: u64 = 0x0000_0000_8fff_ffff;
pub const SIZEOF_IMPORT_ADDRESS_TABLE_ENTRY_32: usize = 4;
pub const SIZEOF_IMPORT_ADDRESS_TABLE_ENTRY_64: usize = 8;

pub trait Bitfield<'a>: PartialEq + Eq + LowerHex + Debug + TryFromCtx<'a, scroll::Endian, Error=scroll::Error, Size=usize> {
    fn is_ordinal(&self) -> bool;
    fn to_ordinal(&self) -> u16;
    fn to_rva(&self) -> usize;
    fn size_of() -> usize;
    fn is_zero(&self) -> bool;
}

impl<'a> Bitfield<'a> for u64 {
    fn is_ordinal(&self) -> bool { self & IMPORT_BY_ORDINAL_64 == IMPORT_BY_ORDINAL_64 }
    fn to_ordinal(&self) -> u16 { (0xffff & self) as u16 }
    fn to_rva(&self) -> usize { (self & IMPORT_RVA_MASK_64) as usize }
    fn size_of() -> usize { 8 }
    fn is_zero(&self) -> bool { *self == 0}
}

impl<'a> Bitfield<'a> for u32 {
    fn is_ordinal(&self) -> bool { self & IMPORT_BY_ORDINAL_32 == IMPORT_BY_ORDINAL_32 }
    fn to_ordinal(&self) -> u16 { (0xffff & self) as u16 }
    fn to_rva(&self) -> usize { (self & IMPORT_RVA_MASK_32) as usize }
    fn size_of() -> usize { 4 }
    fn is_zero(&self) -> bool { *self == 0}
}

#[derive(Debug, Clone)]
pub struct HintNameTableEntry<'a> {
    pub hint: u16,
    pub name: &'a str,
}

impl<'a> HintNameTableEntry<'a> {
    fn parse(bytes: &'a [u8], mut offset: usize) -> error::Result<Self> {
        let offset = &mut offset;
        let hint = bytes.gread_with(offset, scroll::LE)?;
        let name = bytes.pread::<&'a str>(*offset)?;
        Ok(HintNameTableEntry { hint: hint, name: name })
    }
}

#[derive(Debug, Clone)]
pub enum SyntheticImportLookupTableEntry<'a> {
    OrdinalNumber(u16),
    HintNameTableRVA ((usize, HintNameTableEntry<'a>)), // [u8; 31] bitfield :/
}

#[derive(Debug)]
pub struct ImportLookupTableEntry<'a> {
    pub synthetic: SyntheticImportLookupTableEntry<'a>,
}

pub type ImportLookupTable<'a> = Vec<ImportLookupTableEntry<'a>>;

impl<'a> ImportLookupTableEntry<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a [u8], mut offset: usize, sections: &[section_table::SectionTable])
                                                                      -> error::Result<ImportLookupTable<'a>> {
        let le = scroll::LE;
        let offset = &mut offset;
        let mut table = Vec::new();
        loop {
            let bitfield: T = bytes.gread_with(offset, le)?;
            if bitfield.is_zero() {
                debug!("imports done");
                break;
            } else {
                let synthetic = {
                    debug!("bitfield {:#x}", bitfield);
                    use self::SyntheticImportLookupTableEntry::*;
                    if bitfield.is_ordinal() {
                        let ordinal = bitfield.to_ordinal();
                        debug!("importing by ordinal {:#x}", ordinal);
                        OrdinalNumber(ordinal)
                    } else {
                        let rva = bitfield.to_rva();
                        let hentry = {
                            debug!("searching for RVA {:#x}", rva);
                            if let Some(offset) = utils::find_offset(rva as usize, sections) {
                                debug!("offset {:#x}", offset);
                                HintNameTableEntry::parse(bytes, offset)?
                            } else {
                                warn!("Entry {} has bad RVA: {:#x}", table.len(), rva);
                                continue
                            }
                        };
                        HintNameTableRVA ((rva, hentry))
                    }
                };
                let entry = ImportLookupTableEntry { synthetic: synthetic };
                table.push(entry);
            }
        }
        Ok(table)
    }
}

// get until entry is 0
pub type ImportAddressTable = Vec<usize>;

#[repr(C)]
#[derive(Debug)]
#[derive(Pread, Pwrite, SizeWith)]
pub struct ImportDirectoryEntry {
    pub import_lookup_table_rva: u32,
    pub time_date_stamp: u32,
    pub forwarder_chain: u32,
    pub name_rva: u32,
    pub import_address_table_rva: u32,
}

pub const SIZEOF_IMPORT_DIRECTORY_ENTRY: usize = 20;

impl ImportDirectoryEntry {
    pub fn is_null (&self) -> bool {
        (self.import_lookup_table_rva == 0) &&
            (self.time_date_stamp == 0) &&
            (self.forwarder_chain == 0) &&
            (self.name_rva == 0) &&
            (self.import_address_table_rva == 0)
    }
}

#[derive(Debug)]
pub struct SyntheticImportDirectoryEntry<'a> {
    pub import_directory_entry: ImportDirectoryEntry,
    /// Computed
    pub name: &'a str,
    /// The import lookup table is a vector of either ordinals, or RVAs + import names
    pub import_lookup_table: Option<ImportLookupTable<'a>>,
    /// Computed
    pub import_address_table: ImportAddressTable,
}

impl<'a> SyntheticImportDirectoryEntry<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a [u8], import_directory_entry: ImportDirectoryEntry, sections: &[section_table::SectionTable]) -> error::Result<SyntheticImportDirectoryEntry<'a>> {
        const LE: scroll::Endian = scroll::LE;
        let name_rva = import_directory_entry.name_rva;
        let name = utils::try_name(bytes, name_rva as usize, sections)?;
        let import_lookup_table = {
            let import_lookup_table_rva = import_directory_entry.import_lookup_table_rva;
            debug!("Synthesizing lookup table imports for {} lib, with import lookup table rva: {:#x}", name, import_lookup_table_rva);
            if let Some(import_lookup_table_offset) = utils::find_offset(import_lookup_table_rva as usize, sections) {
                let import_lookup_table = ImportLookupTableEntry::parse::<T>(bytes, import_lookup_table_offset, sections)?;
                debug!("Successfully synthesized import lookup table entry: {:#?}", import_lookup_table);
                Some(import_lookup_table)
            } else {
                None
            }
        };

        let import_address_table_offset = &mut utils::find_offset(import_directory_entry.import_address_table_rva as usize, sections).ok_or(error::Error::Malformed(format!("Cannot map import_address_table_rva {:#x} into offset for {}", import_directory_entry.import_address_table_rva, name)))?;
        let mut import_address_table = Vec::new();
        loop {
            let import_address = bytes.gread_with(import_address_table_offset, LE)?;
            if import_address == 0 { break } else { import_address_table.push(import_address); }
        }
        Ok(SyntheticImportDirectoryEntry {
            import_directory_entry: import_directory_entry,
            name: name,
            import_lookup_table: import_lookup_table,
            import_address_table: import_address_table
        })
    }
}

#[derive(Debug)]
/// Contains a list of synthesized import data for this binary, e.g., which symbols from which libraries it is importing from
pub struct ImportData<'a> {
    pub import_data: Vec<SyntheticImportDirectoryEntry<'a>>,
}

impl<'a> ImportData<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a[u8], dd: &data_directories::DataDirectory, sections: &[section_table::SectionTable]) -> error::Result<ImportData<'a>> {
        let import_directory_table_rva = dd.virtual_address as usize;
        debug!("import_directory_table_rva {:#x}", import_directory_table_rva);
        let offset = &mut utils::find_offset(import_directory_table_rva, sections).ok_or(error::Error::Malformed(format!("Cannot create ImportData; cannot map import_directory_table_rva {:#x} into offset", import_directory_table_rva)))?;;
        debug!("import data offset {:#x}", offset);
        let mut import_data = Vec::new();
        loop {
            let import_directory_entry: ImportDirectoryEntry = bytes.gread_with(offset, scroll::LE)?;
            debug!("{:#?}", import_directory_entry);
            if import_directory_entry.is_null() {
                break;
            } else {
                let entry = SyntheticImportDirectoryEntry::parse::<T>(bytes, import_directory_entry, sections)?;
                debug!("entry {:#?}", entry);
                import_data.push(entry);
            }
        }
        debug!("finished ImportData");
        Ok(ImportData { import_data: import_data})
    }
}

#[derive(Debug)]
/// A synthesized symbol import, the name is pre-indexed, and the binary offset is computed, as well as which dll it belongs to
pub struct Import<'a> {
    pub name: Cow<'a, str>,
    pub dll: &'a str,
    pub ordinal: u16,
    pub offset: usize,
    pub rva: usize,
    pub size: usize,
}

impl<'a> Import<'a> {
    pub fn parse<T: Bitfield<'a>>(_bytes: &'a [u8], import_data: &ImportData<'a>, _sections: &[section_table::SectionTable]) -> error::Result<Vec<Import<'a>>> {
        let mut imports = Vec::new();
        for data in &import_data.import_data {
            if let Some(ref import_lookup_table) = data.import_lookup_table {
                let dll = data.name;
                let import_base = data.import_directory_entry.import_address_table_rva as usize;
                debug!("Getting imports from {}", &dll);
                for (i, entry) in import_lookup_table.iter().enumerate() {
                    let offset = import_base + (i * T::size_of());
                    use self::SyntheticImportLookupTableEntry::*;
                    let (rva, name, ordinal) =
                        match &entry.synthetic {
                            &HintNameTableRVA ((rva, ref hint_entry)) => {
                                // if hint_entry.name = "" && hint_entry.hint = 0 {
                                //     println!("<PE.Import> warning hint/name table rva from {} without hint {:#x}", dll, rva);
                                // }
                                (rva, Cow::Borrowed(hint_entry.name), hint_entry.hint.clone())
                            },
                            &OrdinalNumber(ordinal) => {
                                let name = format!("ORDINAL {}", ordinal);
                                (0x0, Cow::Owned(name), ordinal)
                            },
                        };
                    let import =
                        Import {
                            name: name,
                            ordinal: ordinal, dll: dll,
                            size: T::size_of(), offset: offset, rva: rva as usize
                        };
                    imports.push(import);
                }
            }
        }
        Ok (imports)
    }
}
<commit_msg>Remove unused constants<commit_after>use alloc::borrow::Cow;
use alloc::vec::Vec;
use core::fmt::{LowerHex, Debug};

use scroll::{self, Pread};
use scroll::ctx::TryFromCtx;
use error;

use pe::section_table;
use pe::utils;
use pe::data_directories;

pub const IMPORT_BY_ORDINAL_32: u32 = 0x8000_0000;
pub const IMPORT_BY_ORDINAL_64: u64 = 0x8000_0000_0000_0000;
pub const IMPORT_RVA_MASK_32: u32 = 0x8fff_ffff;
pub const IMPORT_RVA_MASK_64: u64 = 0x0000_0000_8fff_ffff;

pub trait Bitfield<'a>: PartialEq + Eq + LowerHex + Debug + TryFromCtx<'a, scroll::Endian, Error=scroll::Error, Size=usize> {
    fn is_ordinal(&self) -> bool;
    fn to_ordinal(&self) -> u16;
    fn to_rva(&self) -> usize;
    fn size_of() -> usize;
    fn is_zero(&self) -> bool;
}

impl<'a> Bitfield<'a> for u64 {
    fn is_ordinal(&self) -> bool { self & IMPORT_BY_ORDINAL_64 == IMPORT_BY_ORDINAL_64 }
    fn to_ordinal(&self) -> u16 { (0xffff & self) as u16 }
    fn to_rva(&self) -> usize { (self & IMPORT_RVA_MASK_64) as usize }
    fn size_of() -> usize { 8 }
    fn is_zero(&self) -> bool { *self == 0}
}

impl<'a> Bitfield<'a> for u32 {
    fn is_ordinal(&self) -> bool { self & IMPORT_BY_ORDINAL_32 == IMPORT_BY_ORDINAL_32 }
    fn to_ordinal(&self) -> u16 { (0xffff & self) as u16 }
    fn to_rva(&self) -> usize { (self & IMPORT_RVA_MASK_32) as usize }
    fn size_of() -> usize { 4 }
    fn is_zero(&self) -> bool { *self == 0}
}

#[derive(Debug, Clone)]
pub struct HintNameTableEntry<'a> {
    pub hint: u16,
    pub name: &'a str,
}

impl<'a> HintNameTableEntry<'a> {
    fn parse(bytes: &'a [u8], mut offset: usize) -> error::Result<Self> {
        let offset = &mut offset;
        let hint = bytes.gread_with(offset, scroll::LE)?;
        let name = bytes.pread::<&'a str>(*offset)?;
        Ok(HintNameTableEntry { hint: hint, name: name })
    }
}

#[derive(Debug, Clone)]
pub enum SyntheticImportLookupTableEntry<'a> {
    OrdinalNumber(u16),
    HintNameTableRVA ((usize, HintNameTableEntry<'a>)), // [u8; 31] bitfield :/
}

#[derive(Debug)]
pub struct ImportLookupTableEntry<'a> {
    pub synthetic: SyntheticImportLookupTableEntry<'a>,
}

pub type ImportLookupTable<'a> = Vec<ImportLookupTableEntry<'a>>;

impl<'a> ImportLookupTableEntry<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a [u8], mut offset: usize, sections: &[section_table::SectionTable])
                                                                      -> error::Result<ImportLookupTable<'a>> {
        let le = scroll::LE;
        let offset = &mut offset;
        let mut table = Vec::new();
        loop {
            let bitfield: T = bytes.gread_with(offset, le)?;
            if bitfield.is_zero() {
                debug!("imports done");
                break;
            } else {
                let synthetic = {
                    debug!("bitfield {:#x}", bitfield);
                    use self::SyntheticImportLookupTableEntry::*;
                    if bitfield.is_ordinal() {
                        let ordinal = bitfield.to_ordinal();
                        debug!("importing by ordinal {:#x}", ordinal);
                        OrdinalNumber(ordinal)
                    } else {
                        let rva = bitfield.to_rva();
                        let hentry = {
                            debug!("searching for RVA {:#x}", rva);
                            if let Some(offset) = utils::find_offset(rva as usize, sections) {
                                debug!("offset {:#x}", offset);
                                HintNameTableEntry::parse(bytes, offset)?
                            } else {
                                warn!("Entry {} has bad RVA: {:#x}", table.len(), rva);
                                continue
                            }
                        };
                        HintNameTableRVA ((rva, hentry))
                    }
                };
                let entry = ImportLookupTableEntry { synthetic: synthetic };
                table.push(entry);
            }
        }
        Ok(table)
    }
}

// get until entry is 0
pub type ImportAddressTable = Vec<usize>;

#[repr(C)]
#[derive(Debug)]
#[derive(Pread, Pwrite, SizeWith)]
pub struct ImportDirectoryEntry {
    pub import_lookup_table_rva: u32,
    pub time_date_stamp: u32,
    pub forwarder_chain: u32,
    pub name_rva: u32,
    pub import_address_table_rva: u32,
}

pub const SIZEOF_IMPORT_DIRECTORY_ENTRY: usize = 20;

impl ImportDirectoryEntry {
    pub fn is_null (&self) -> bool {
        (self.import_lookup_table_rva == 0) &&
            (self.time_date_stamp == 0) &&
            (self.forwarder_chain == 0) &&
            (self.name_rva == 0) &&
            (self.import_address_table_rva == 0)
    }
}

#[derive(Debug)]
pub struct SyntheticImportDirectoryEntry<'a> {
    pub import_directory_entry: ImportDirectoryEntry,
    /// Computed
    pub name: &'a str,
    /// The import lookup table is a vector of either ordinals, or RVAs + import names
    pub import_lookup_table: Option<ImportLookupTable<'a>>,
    /// Computed
    pub import_address_table: ImportAddressTable,
}

impl<'a> SyntheticImportDirectoryEntry<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a [u8], import_directory_entry: ImportDirectoryEntry, sections: &[section_table::SectionTable]) -> error::Result<SyntheticImportDirectoryEntry<'a>> {
        const LE: scroll::Endian = scroll::LE;
        let name_rva = import_directory_entry.name_rva;
        let name = utils::try_name(bytes, name_rva as usize, sections)?;
        let import_lookup_table = {
            let import_lookup_table_rva = import_directory_entry.import_lookup_table_rva;
            debug!("Synthesizing lookup table imports for {} lib, with import lookup table rva: {:#x}", name, import_lookup_table_rva);
            if let Some(import_lookup_table_offset) = utils::find_offset(import_lookup_table_rva as usize, sections) {
                let import_lookup_table = ImportLookupTableEntry::parse::<T>(bytes, import_lookup_table_offset, sections)?;
                debug!("Successfully synthesized import lookup table entry: {:#?}", import_lookup_table);
                Some(import_lookup_table)
            } else {
                None
            }
        };

        let import_address_table_offset = &mut utils::find_offset(import_directory_entry.import_address_table_rva as usize, sections).ok_or(error::Error::Malformed(format!("Cannot map import_address_table_rva {:#x} into offset for {}", import_directory_entry.import_address_table_rva, name)))?;
        let mut import_address_table = Vec::new();
        loop {
            let import_address = bytes.gread_with(import_address_table_offset, LE)?;
            if import_address == 0 { break } else { import_address_table.push(import_address); }
        }
        Ok(SyntheticImportDirectoryEntry {
            import_directory_entry: import_directory_entry,
            name: name,
            import_lookup_table: import_lookup_table,
            import_address_table: import_address_table
        })
    }
}

#[derive(Debug)]
/// Contains a list of synthesized import data for this binary, e.g., which symbols from which libraries it is importing from
pub struct ImportData<'a> {
    pub import_data: Vec<SyntheticImportDirectoryEntry<'a>>,
}

impl<'a> ImportData<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a[u8], dd: &data_directories::DataDirectory, sections: &[section_table::SectionTable]) -> error::Result<ImportData<'a>> {
        let import_directory_table_rva = dd.virtual_address as usize;
        debug!("import_directory_table_rva {:#x}", import_directory_table_rva);
        let offset = &mut utils::find_offset(import_directory_table_rva, sections).ok_or(error::Error::Malformed(format!("Cannot create ImportData; cannot map import_directory_table_rva {:#x} into offset", import_directory_table_rva)))?;;
        debug!("import data offset {:#x}", offset);
        let mut import_data = Vec::new();
        loop {
            let import_directory_entry: ImportDirectoryEntry = bytes.gread_with(offset, scroll::LE)?;
            debug!("{:#?}", import_directory_entry);
            if import_directory_entry.is_null() {
                break;
            } else {
                let entry = SyntheticImportDirectoryEntry::parse::<T>(bytes, import_directory_entry, sections)?;
                debug!("entry {:#?}", entry);
                import_data.push(entry);
            }
        }
        debug!("finished ImportData");
        Ok(ImportData { import_data: import_data})
    }
}

#[derive(Debug)]
/// A synthesized symbol import, the name is pre-indexed, and the binary offset is computed, as well as which dll it belongs to
pub struct Import<'a> {
    pub name: Cow<'a, str>,
    pub dll: &'a str,
    pub ordinal: u16,
    pub offset: usize,
    pub rva: usize,
    pub size: usize,
}

impl<'a> Import<'a> {
    pub fn parse<T: Bitfield<'a>>(_bytes: &'a [u8], import_data: &ImportData<'a>, _sections: &[section_table::SectionTable]) -> error::Result<Vec<Import<'a>>> {
        let mut imports = Vec::new();
        for data in &import_data.import_data {
            if let Some(ref import_lookup_table) = data.import_lookup_table {
                let dll = data.name;
                let import_base = data.import_directory_entry.import_address_table_rva as usize;
                debug!("Getting imports from {}", &dll);
                for (i, entry) in import_lookup_table.iter().enumerate() {
                    let offset = import_base + (i * T::size_of());
                    use self::SyntheticImportLookupTableEntry::*;
                    let (rva, name, ordinal) =
                        match &entry.synthetic {
                            &HintNameTableRVA ((rva, ref hint_entry)) => {
                                // if hint_entry.name = "" && hint_entry.hint = 0 {
                                //     println!("<PE.Import> warning hint/name table rva from {} without hint {:#x}", dll, rva);
                                // }
                                (rva, Cow::Borrowed(hint_entry.name), hint_entry.hint.clone())
                            },
                            &OrdinalNumber(ordinal) => {
                                let name = format!("ORDINAL {}", ordinal);
                                (0x0, Cow::Owned(name), ordinal)
                            },
                        };
                    let import =
                        Import {
                            name: name,
                            ordinal: ordinal, dll: dll,
                            size: T::size_of(), offset: offset, rva: rva as usize
                        };
                    imports.push(import);
                }
            }
        }
        Ok (imports)
    }
}
<|endoftext|>


<commit_before><commit_msg>Fix lint casting to the same type is unnecessary<commit_after><|endoftext|>