<commit_before>extern crate piston_window;
extern crate gfx_device_gl;
extern crate gfx_graphics;
extern crate find_folder;

use piston_window::Transformed;

use piston_window as pw;

mod object;
use object::Object;

struct Game {
    player: Object,
    up: bool,
    down: bool,
    left: bool,
    right: bool,
}

impl Game {
    fn new() -> Game {
        Game { player: Object::new(), up: false, down: false, left: false, right: false }
    }

    fn on_update(&mut self, args: &pw::UpdateArgs) {
        if self.up   { self.player.mov(0.0, -150.0 * args.dt); }
        if self.down { self.player.mov(0.0, 150.0 * args.dt); }
        if self.left { self.player.mov(-150.0 * args.dt, 0.0); }
        if self.right { self.player.mov(150.0 * args.dt, 0.0); }
    }

    fn on_draw(&mut self, args: &pw::RenderArgs, window: &pw::PistonWindow) {
        window.draw_2d(|context, graphics| {
            pw::clear([0.0, 0.0, 0.0, 1.0], graphics);

            let center = context.transform.trans((args.width / 2) as f64, (args.height / 2) as f64);
            self.player.render(graphics, &center);
        });
    }

    fn on_input(&mut self, input: &pw::Input) {
        use piston_window::{Button, Key};

        match input {
            &pw::Input::Press(button) => {
                match button {
                    Button::Keyboard(Key::Up) => { self.up = true; }
                    Button::Keyboard(Key::Down) => { self.down = true; }
                    Button::Keyboard(Key::Left) => { self.left = true; }
                    Button::Keyboard(Key::Right) => { self.right = true; }
                    _ => {}
                }
            }
            &pw::Input::Release(button) => {
                match button {
                    Button::Keyboard(Key::Up) => { self.up = false; }
                    Button::Keyboard(Key::Down) => { self.down = false; }
                    Button::Keyboard(Key::Left) => { self.left = false; }
                    Button::Keyboard(Key::Right) => { self.right = false; }
                    _ => {}
                }
            }
            _ => {}
        }
    }
}

fn main() {
    use piston_window::Event;

    let window: pw::PistonWindow = pw::WindowSettings::new("rrthozopsi", [600, 600])
        .exit_on_esc(true)
        .build()
        .unwrap();

    let mut game = Game::new();
    {
        let assets = find_folder::Search::ParentsThenKids(3, 3).for_folder("assets").unwrap();
        let sprite = assets.join("skeleton.png");
        let sprite = pw::Texture::from_path(&mut *window.factory.borrow_mut(), &sprite, pw::Flip::None, &pw::TextureSettings::new()).unwrap();

        game.player.set_sprite(sprite);
    }

    for w in window {
        match w.event {
            Some(Event::Update(args)) => { game.on_update(&args); }
            Some(Event::Render(args)) => { game.on_draw(&args, &w); }
            Some(Event::Input(input)) => { game.on_input(&input); }
            _ => {}
        }
    }
}
<commit_msg>Use pw::Input<commit_after>extern crate piston_window;
extern crate gfx_device_gl;
extern crate gfx_graphics;
extern crate find_folder;

use piston_window::Transformed;

use piston_window as pw;

mod object;
use object::Object;

struct Game {
    player: Object,
    up: bool,
    down: bool,
    left: bool,
    right: bool,
}

impl Game {
    fn new() -> Game {
        Game { player: Object::new(), up: false, down: false, left: false, right: false }
    }

    fn on_update(&mut self, args: &pw::UpdateArgs) {
        if self.up   { self.player.mov(0.0, -150.0 * args.dt); }
        if self.down { self.player.mov(0.0, 150.0 * args.dt); }
        if self.left { self.player.mov(-150.0 * args.dt, 0.0); }
        if self.right { self.player.mov(150.0 * args.dt, 0.0); }
    }

    fn on_draw(&mut self, args: &pw::RenderArgs, window: &pw::PistonWindow) {
        window.draw_2d(|context, graphics| {
            pw::clear([0.0, 0.0, 0.0, 1.0], graphics);

            let center = context.transform.trans((args.width / 2) as f64, (args.height / 2) as f64);
            self.player.render(graphics, &center);
        });
    }

    fn on_input(&mut self, input: &pw::Input) {
        use piston_window::{Input, Button, Key};

        match input {
            &Input::Press(button) => {
                match button {
                    Button::Keyboard(Key::Up) => { self.up = true; }
                    Button::Keyboard(Key::Down) => { self.down = true; }
                    Button::Keyboard(Key::Left) => { self.left = true; }
                    Button::Keyboard(Key::Right) => { self.right = true; }
                    _ => {}
                }
            }
            &Input::Release(button) => {
                match button {
                    Button::Keyboard(Key::Up) => { self.up = false; }
                    Button::Keyboard(Key::Down) => { self.down = false; }
                    Button::Keyboard(Key::Left) => { self.left = false; }
                    Button::Keyboard(Key::Right) => { self.right = false; }
                    _ => {}
                }
            }
            _ => {}
        }
    }
}

fn main() {
    use piston_window::Event;

    let window: pw::PistonWindow = pw::WindowSettings::new("rrthozopsi", [600, 600])
        .exit_on_esc(true)
        .build()
        .unwrap();

    let mut game = Game::new();
    {
        let assets = find_folder::Search::ParentsThenKids(3, 3).for_folder("assets").unwrap();
        let sprite = assets.join("skeleton.png");
        let sprite = pw::Texture::from_path(&mut *window.factory.borrow_mut(), &sprite, pw::Flip::None, &pw::TextureSettings::new()).unwrap();

        game.player.set_sprite(sprite);
    }

    for w in window {
        match w.event {
            Some(Event::Update(args)) => { game.on_update(&args); }
            Some(Event::Render(args)) => { game.on_draw(&args, &w); }
            Some(Event::Input(input)) => { game.on_input(&input); }
            _ => {}
        }
    }
}
<|endoftext|>


<commit_before>use super::job_queue::*;
use super::queue_state::*;
use super::core::*;
use super::active_queue::*;
use super::wake_queue::*;

use futures::prelude::*;
use futures::channel::oneshot;
use futures::task;

use std::mem;
use std::sync::*;
use std::pin::{Pin};

///
/// The possible states of a future result
///
enum FutureResultState<T> {
    None,
    Some(T),
    ReturnedViaFuture
}

///
/// The possible states of a waker that will wake up a queue that's being drained as part of polling a future
///
enum DrainWakerState {
    /// The drain waker has never been woken before
    NotWoken,

    /// This has been woken but the waker is not set
    Woken,

    /// This has been woken and the waker is set
    WillWakeWithWaker(task::Waker)
}

///
/// Waker that will wake up a queue being drained. We set the waker with a delay so it's possible to switch to
/// a waker that will wake a future on a background queue if necessary.
///
struct DrainWaker {
    state: Mutex<DrainWakerState>
}

impl DrainWaker {
    ///
    /// Creates a new drain waker
    ///
    fn new() -> DrainWaker {
        DrainWaker {
            state: Mutex::new(DrainWakerState::NotWoken)
        }
    }

    ///
    /// Sets the waker to be called when this drain waker is woken
    ///
    fn wake_with(&self, new_waker: task::Waker) {
        use self::DrainWakerState::*;

        // Update the state and determine if we need to invoke the waker immediately (if it's been woken before the waker was set)
        let to_wake = {
            // Fetch the current state
            let mut new_state   = self.state.lock().expect("Drain waker state");
            let mut state       = Woken;
            mem::swap(&mut *new_state, &mut state);

            // Update the state based on this action
            match state {
                Woken                           => { *new_state = Woken; Some(new_waker) },
                NotWoken                        => { *new_state = WillWakeWithWaker(new_waker); None },
                WillWakeWithWaker(_old_waker)   => { *new_state = WillWakeWithWaker(new_waker); None }
            }
        };

        // Wake up the waker, if we need to call it immediately
        to_wake.map(|to_wake| to_wake.wake());
    }
}

impl task::ArcWake for DrainWaker {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        use self::DrainWakerState::*;

        // If the current state contains a waker, we'll call it once we've unlocked the mutex
        let to_wake = {
            // Fetch the current state
            let mut new_state   = arc_self.state.lock().expect("Drain waker state");
            let mut state       = Woken;
            mem::swap(&mut *new_state, &mut state);

            // Update the state based on this action
            match state {
                NotWoken                    => { *new_state = Woken; None },
                Woken                       => { *new_state = Woken; None },
                WillWakeWithWaker(waker)    => { *new_state = Woken; Some(waker) }
            }
        };

        // Wake up the waker
        to_wake.map(|to_wake| to_wake.wake());
    }
}

///
/// Signalling structure used to return the result of a scheduler future
///
struct SchedulerFutureResult<T> {
    /// The result of the future, or None if it has not been generated yet
    result: FutureResultState<Result<T, oneshot::Canceled>>,

    /// The waker to be called when the future is available
    waker: Option<task::Waker>
}

///
/// Wrapper that cancels the future if it is dropped before it is signalled
///
pub (super) struct SchedulerFutureSignaller<T>(Arc<Mutex<SchedulerFutureResult<T>>>);

///
/// Future representing a task pending on a scheduler
///
/// If polled when no threads are available, this future will run synchronously on the current thread
/// (stealing its execution time rather than blocking)
///
pub struct SchedulerFuture<T> {
    /// The unique ID of this future
    id: FutureId,

    /// The queue which will eventually evaluate the result of this future
    queue: Arc<JobQueue>,

    /// The scheduler core that this future belongs to
    scheduler: Arc<SchedulerCore>,

    /// A container for the result of this scheduler future
    result: Arc<Mutex<SchedulerFutureResult<T>>>
}

impl<T> FutureResultState<T> {
    ///
    /// Returns true if no result has ever been generated for this future
    ///
    fn is_none(&self) -> bool {
        match self {
            FutureResultState::None                 => true,
            FutureResultState::Some(_)              => false,
            FutureResultState::ReturnedViaFuture    => false
        }
    }

    ///
    /// If a result has been generated by this future and has not previously been taken, returns it. Panics if the result has already
    /// been returned and is no longer available.
    ///
    fn take(&mut self) -> Option<T> {
        // Move the value out of this object
        let mut new_value = FutureResultState::ReturnedViaFuture;
        mem::swap(self, &mut new_value);

        // Return it if it contains a value
        match new_value {
            FutureResultState::None                 => { *self = FutureResultState::None; None },
            FutureResultState::ReturnedViaFuture    => { *self = FutureResultState::ReturnedViaFuture; panic!("Future result has already been returned") },
            FutureResultState::Some(value)          => { Some(value) }
        }
    }
}

impl<T> Drop for SchedulerFutureSignaller<T> {
    fn drop(&mut self) {
        let waker = {
            let mut future_result = self.0.lock().expect("Scheduler future result");

            // If no result has been generated, then mark the future as canceled
            if future_result.result.is_none() {
                // Mark the future as canceled
                future_result.result = FutureResultState::Some(Err(oneshot::Canceled));

                // Wake up anything that was polling it
                let waker = future_result.waker.take();
                waker
            } else {
                // Result is already set, so don't wake anything up
                None
            }
        };

        // If we need to wake the future, then do so here (note that we're outside of the lock when we do this)
        waker.map(|waker| waker.wake());
    }
}

impl<T> SchedulerFutureSignaller<T> {
    ///
    /// Signals that the result of the calculation is available
    ///
    pub (super) fn signal(self, result: T) {
        let waker = {
            let mut future_result = self.0.lock().expect("Scheduler future result");

            // Set the result
            future_result.result = FutureResultState::Some(Ok(result));

            // Retrieve the waker
            future_result.waker.take()
        };

        // If we retrieved a waker from the result, wake it up
        waker.map(|waker| waker.wake());
    }
}

///
/// Possible actions we can take based on the state of the signal for the future
///
enum SchedulerAction<T> {
    /// Queue is running elsewhere: we should wait for it to reach the point where the future's result is available
    WaitForCompletion,

    /// Future has already completed: we should return the value to sender
    ReturnValue(T),

    /// We've claimed the 'running' state of the queue and should drain it
    DrainQueue,

    /// The queue has panicked
    Panic
}

impl<T> SchedulerFuture<T> {
    ///
    /// Creates a new scheduler future and the result needed to signal it
    ///
    pub (super) fn new(queue: &Arc<JobQueue>, core: Arc<SchedulerCore>) -> (SchedulerFuture<T>, SchedulerFutureSignaller<T>) {
        // Create an unfinished result
        let result = SchedulerFutureResult {
            result: FutureResultState::None,
            waker:  None
        };
        let result = Arc::new(Mutex::new(result));

        // Insert into a future
        let future = SchedulerFuture {
            id:         FutureId::new(),
            queue:      Arc::clone(queue),
            scheduler:  core,
            result:     Arc::clone(&result)
        };

        (future, SchedulerFutureSignaller(result))
    }

    ///
    /// We moved the queue into the running state and need to drain it until we've got a result
    ///
    fn drain_queue(&mut self, context: &mut task::Context) -> task::Poll<Result<T, oneshot::Canceled>> {
        debug_assert!(self.queue.core.lock().expect("JobQueue core lock").state.is_running());

        // Set the queue as active
        let _active     = ActiveQueue { queue: &*self.queue };
        let mut result;

        // While there is no result, run a job from the queue
        loop {
            // See if the result has arrived yet
            result = self.result.lock().expect("Scheduler future result").result.take();
            if !result.is_none() { break; }

            // Run the next job in the queue
            if let Some(mut job) = self.queue.dequeue() {
                // Queue is running
                debug_assert!(self.queue.core.lock().expect("Job queue core").state.is_running());

                // Create a context to poll in (we may need to reschedule in the background)
                let waker               = Arc::new(DrainWaker::new());
                let waker_ref           = task::waker_ref(&waker);
                let mut drain_context   = task::Context::from_waker(&waker_ref);

                // Poll the queue
                let poll_result = job.run(&mut drain_context);

                match poll_result {
                    task::Poll::Ready(())   => {
                        // Keep running jobs and checking the results if ready
                    },

                    task::Poll::Pending     => {
                        // Requeue the job
                        self.queue.requeue(job);

                        // If the result was supplied, break out of the loop and reschedule the queue
                        result = self.result.lock().expect("Scheduler future result").result.take();
                        if result.is_some() {
                            // Wake the queue in the background if needed (the result has arrived)
                            self.queue.core.lock().expect("JobQueue core lock").state = QueueState::WaitingForWake;

                            let queue_waker = WakeQueue(Arc::clone(&self.queue), Arc::clone(&self.scheduler));
                            let queue_waker = Arc::new(queue_waker);
                            let queue_waker = task::waker(queue_waker);

                            waker.wake_with(queue_waker);

                            // The future is ready (job will be rescheduled in the background)
                            return task::Poll::Ready(result.unwrap());
                        } else {
                            // Wait for the next poll
                            self.queue.core.lock().expect("JobQueue core lock").state = QueueState::WaitingForPoll(self.id);

                            // Use the context waker
                            waker.wake_with(context.waker().clone());

                            // Result is pending
                            return task::Poll::Pending;
                        }
                    }
                }
            } else {
                // Queue is empty and our result hasn't arrived yet?!

                // Assume the future will resolve eventually: move the queue in to the background
                self.result.lock().expect("Scheduler future result").waker = Some(context.waker().clone());

                // Reschedule the queue
                self.queue.core.lock().expect("JobQueue core lock").state = QueueState::Idle;
                self.scheduler.reschedule_queue(&self.queue, Arc::clone(&self.scheduler));

                return task::Poll::Pending;
            }
        }

        // Reschedule the queue if there are any events left pending
        // Note: the queue is already pending when we start running events from it here.
        // This means it'll get dequeued by a thread eventually: maybe while it's running
        // here. As we've set the queue state to running while we're busy, the thread won't
        // start the queue while it's already running.
        self.queue.core.lock().expect("JobQueue core lock").state = QueueState::Idle;
        self.scheduler.reschedule_queue(&self.queue, Arc::clone(&self.scheduler));

        // Result must be available by this point
        task::Poll::Ready(result.unwrap())
    }
}

impl<T> Future for SchedulerFuture<T> {
    type Output = Result<T, oneshot::Canceled>;

    ///
    /// Polls this future
    ///
    fn poll(mut self: Pin<&mut Self>, context: &mut task::Context) -> task::Poll<Self::Output> {
        // Lock the result and determine which action to take
        let next_action = {
            let mut future_result = self.result.lock().expect("Scheduler future result");

            if let Some(result) = future_result.result.take() {
                // The result is available: we should return it immediately
                SchedulerAction::ReturnValue(result)
            } else {
                // If the queue is idle when this is called, we need to schedule this task on this thread rather than one owned by the background process
                let run_action = {
                    let mut core = self.queue.core.lock().expect("JobQueue core lock");

                    match core.state {
                        QueueState::Running                     => SchedulerAction::WaitForCompletion,
                        QueueState::WaitingForWake              => SchedulerAction::WaitForCompletion,
                        QueueState::WaitingForUnpark            => SchedulerAction::WaitForCompletion,
                        QueueState::AwokenWhileRunning          => SchedulerAction::WaitForCompletion,
                        QueueState::Panicked                    => SchedulerAction::Panic,
                        QueueState::Pending                     => { core.state = QueueState::Running; SchedulerAction::DrainQueue },
                        QueueState::Idle                        => { core.state = QueueState::Running; SchedulerAction::DrainQueue }

                        QueueState::WaitingForPoll(owner_id)    => {
                            if owner_id == self.id || true {
                                // TODO: the '|| true' here should reproduce an old bug I'm writing tests for, remove it :-)
                                // Continue polling on this future
                                core.state = QueueState::Running; SchedulerAction::DrainQueue
                            } else {
                                // Wait for the owning future to complete
                                SchedulerAction::WaitForCompletion
                            }
                        },
                    }
                };

                if let SchedulerAction::WaitForCompletion = run_action {
                    // Wake us up when the future is available
                    future_result.waker = Some(context.waker().clone());
                }

                run_action
            }
        };

        match next_action {
            SchedulerAction::WaitForCompletion  => task::Poll::Pending,
            SchedulerAction::ReturnValue(value) => task::Poll::Ready(value),
            SchedulerAction::DrainQueue         => self.drain_queue(context),
            SchedulerAction::Panic              => panic!("Cannot schedule jobs on a panicked queue"),
        }
    }
}

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

    #[test]
    fn returns_immediately_if_signaled() {
        use futures::executor;

        let scheduler           = scheduler();
        let queue               = queue();
        let (future, signaller) = SchedulerFuture::new(&queue, Arc::clone(&scheduler.core));

        signaller.signal(42);

        assert!(executor::block_on(future) == Ok(42));
    }

    #[test]
    fn cancels_when_dropped() {
        use futures::executor;

        let scheduler           = scheduler();
        let queue               = queue();
        let future              = { SchedulerFuture::<i32>::new(&queue, Arc::clone(&scheduler.core)).0 };

        assert!(executor::block_on(future) == Err(oneshot::Canceled));
    }

    #[test]
    fn signals_from_another_thread() {
        use futures::executor;
        use std::thread;
        use std::time::Duration;

        let scheduler           = scheduler();
        let queue               = queue();
        let (future, signaller) = SchedulerFuture::new(&queue, Arc::clone(&scheduler.core));

        thread::spawn(move || {
            thread::sleep(Duration::from_millis(100));

            signaller.signal(42);
        });

        assert!(executor::block_on(future) == Ok(42));
    }

    #[test]
    fn forces_queue_drain() {
        use futures::executor;
        use std::thread;
        use std::time::Duration;

        let scheduler   = scheduler();
        let queue       = queue();

        scheduler.set_max_threads(0);
        scheduler.despawn_threads_if_overloaded();

        let (future, signaller) = SchedulerFuture::new(&queue, Arc::clone(&scheduler.core));

        scheduler.desync(&queue, move || {
            thread::sleep(Duration::from_millis(100));

            signaller.signal(42);
        });

        assert!(executor::block_on(future) == Ok(42));
    }
}
<commit_msg>Re-enable the fix<commit_after>use super::job_queue::*;
use super::queue_state::*;
use super::core::*;
use super::active_queue::*;
use super::wake_queue::*;

use futures::prelude::*;
use futures::channel::oneshot;
use futures::task;

use std::mem;
use std::sync::*;
use std::pin::{Pin};

///
/// The possible states of a future result
///
enum FutureResultState<T> {
    None,
    Some(T),
    ReturnedViaFuture
}

///
/// The possible states of a waker that will wake up a queue that's being drained as part of polling a future
///
enum DrainWakerState {
    /// The drain waker has never been woken before
    NotWoken,

    /// This has been woken but the waker is not set
    Woken,

    /// This has been woken and the waker is set
    WillWakeWithWaker(task::Waker)
}

///
/// Waker that will wake up a queue being drained. We set the waker with a delay so it's possible to switch to
/// a waker that will wake a future on a background queue if necessary.
///
struct DrainWaker {
    state: Mutex<DrainWakerState>
}

impl DrainWaker {
    ///
    /// Creates a new drain waker
    ///
    fn new() -> DrainWaker {
        DrainWaker {
            state: Mutex::new(DrainWakerState::NotWoken)
        }
    }

    ///
    /// Sets the waker to be called when this drain waker is woken
    ///
    fn wake_with(&self, new_waker: task::Waker) {
        use self::DrainWakerState::*;

        // Update the state and determine if we need to invoke the waker immediately (if it's been woken before the waker was set)
        let to_wake = {
            // Fetch the current state
            let mut new_state   = self.state.lock().expect("Drain waker state");
            let mut state       = Woken;
            mem::swap(&mut *new_state, &mut state);

            // Update the state based on this action
            match state {
                Woken                           => { *new_state = Woken; Some(new_waker) },
                NotWoken                        => { *new_state = WillWakeWithWaker(new_waker); None },
                WillWakeWithWaker(_old_waker)   => { *new_state = WillWakeWithWaker(new_waker); None }
            }
        };

        // Wake up the waker, if we need to call it immediately
        to_wake.map(|to_wake| to_wake.wake());
    }
}

impl task::ArcWake for DrainWaker {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        use self::DrainWakerState::*;

        // If the current state contains a waker, we'll call it once we've unlocked the mutex
        let to_wake = {
            // Fetch the current state
            let mut new_state   = arc_self.state.lock().expect("Drain waker state");
            let mut state       = Woken;
            mem::swap(&mut *new_state, &mut state);

            // Update the state based on this action
            match state {
                NotWoken                    => { *new_state = Woken; None },
                Woken                       => { *new_state = Woken; None },
                WillWakeWithWaker(waker)    => { *new_state = Woken; Some(waker) }
            }
        };

        // Wake up the waker
        to_wake.map(|to_wake| to_wake.wake());
    }
}

///
/// Signalling structure used to return the result of a scheduler future
///
struct SchedulerFutureResult<T> {
    /// The result of the future, or None if it has not been generated yet
    result: FutureResultState<Result<T, oneshot::Canceled>>,

    /// The waker to be called when the future is available
    waker: Option<task::Waker>
}

///
/// Wrapper that cancels the future if it is dropped before it is signalled
///
pub (super) struct SchedulerFutureSignaller<T>(Arc<Mutex<SchedulerFutureResult<T>>>);

///
/// Future representing a task pending on a scheduler
///
/// If polled when no threads are available, this future will run synchronously on the current thread
/// (stealing its execution time rather than blocking)
///
pub struct SchedulerFuture<T> {
    /// The unique ID of this future
    id: FutureId,

    /// The queue which will eventually evaluate the result of this future
    queue: Arc<JobQueue>,

    /// The scheduler core that this future belongs to
    scheduler: Arc<SchedulerCore>,

    /// A container for the result of this scheduler future
    result: Arc<Mutex<SchedulerFutureResult<T>>>
}

impl<T> FutureResultState<T> {
    ///
    /// Returns true if no result has ever been generated for this future
    ///
    fn is_none(&self) -> bool {
        match self {
            FutureResultState::None                 => true,
            FutureResultState::Some(_)              => false,
            FutureResultState::ReturnedViaFuture    => false
        }
    }

    ///
    /// If a result has been generated by this future and has not previously been taken, returns it. Panics if the result has already
    /// been returned and is no longer available.
    ///
    fn take(&mut self) -> Option<T> {
        // Move the value out of this object
        let mut new_value = FutureResultState::ReturnedViaFuture;
        mem::swap(self, &mut new_value);

        // Return it if it contains a value
        match new_value {
            FutureResultState::None                 => { *self = FutureResultState::None; None },
            FutureResultState::ReturnedViaFuture    => { *self = FutureResultState::ReturnedViaFuture; panic!("Future result has already been returned") },
            FutureResultState::Some(value)          => { Some(value) }
        }
    }
}

impl<T> Drop for SchedulerFutureSignaller<T> {
    fn drop(&mut self) {
        let waker = {
            let mut future_result = self.0.lock().expect("Scheduler future result");

            // If no result has been generated, then mark the future as canceled
            if future_result.result.is_none() {
                // Mark the future as canceled
                future_result.result = FutureResultState::Some(Err(oneshot::Canceled));

                // Wake up anything that was polling it
                let waker = future_result.waker.take();
                waker
            } else {
                // Result is already set, so don't wake anything up
                None
            }
        };

        // If we need to wake the future, then do so here (note that we're outside of the lock when we do this)
        waker.map(|waker| waker.wake());
    }
}

impl<T> SchedulerFutureSignaller<T> {
    ///
    /// Signals that the result of the calculation is available
    ///
    pub (super) fn signal(self, result: T) {
        let waker = {
            let mut future_result = self.0.lock().expect("Scheduler future result");

            // Set the result
            future_result.result = FutureResultState::Some(Ok(result));

            // Retrieve the waker
            future_result.waker.take()
        };

        // If we retrieved a waker from the result, wake it up
        waker.map(|waker| waker.wake());
    }
}

///
/// Possible actions we can take based on the state of the signal for the future
///
enum SchedulerAction<T> {
    /// Queue is running elsewhere: we should wait for it to reach the point where the future's result is available
    WaitForCompletion,

    /// Future has already completed: we should return the value to sender
    ReturnValue(T),

    /// We've claimed the 'running' state of the queue and should drain it
    DrainQueue,

    /// The queue has panicked
    Panic
}

impl<T> SchedulerFuture<T> {
    ///
    /// Creates a new scheduler future and the result needed to signal it
    ///
    pub (super) fn new(queue: &Arc<JobQueue>, core: Arc<SchedulerCore>) -> (SchedulerFuture<T>, SchedulerFutureSignaller<T>) {
        // Create an unfinished result
        let result = SchedulerFutureResult {
            result: FutureResultState::None,
            waker:  None
        };
        let result = Arc::new(Mutex::new(result));

        // Insert into a future
        let future = SchedulerFuture {
            id:         FutureId::new(),
            queue:      Arc::clone(queue),
            scheduler:  core,
            result:     Arc::clone(&result)
        };

        (future, SchedulerFutureSignaller(result))
    }

    ///
    /// We moved the queue into the running state and need to drain it until we've got a result
    ///
    fn drain_queue(&mut self, context: &mut task::Context) -> task::Poll<Result<T, oneshot::Canceled>> {
        debug_assert!(self.queue.core.lock().expect("JobQueue core lock").state.is_running());

        // Set the queue as active
        let _active     = ActiveQueue { queue: &*self.queue };
        let mut result;

        // While there is no result, run a job from the queue
        loop {
            // See if the result has arrived yet
            result = self.result.lock().expect("Scheduler future result").result.take();
            if !result.is_none() { break; }

            // Run the next job in the queue
            if let Some(mut job) = self.queue.dequeue() {
                // Queue is running
                debug_assert!(self.queue.core.lock().expect("Job queue core").state.is_running());

                // Create a context to poll in (we may need to reschedule in the background)
                let waker               = Arc::new(DrainWaker::new());
                let waker_ref           = task::waker_ref(&waker);
                let mut drain_context   = task::Context::from_waker(&waker_ref);

                // Poll the queue
                let poll_result = job.run(&mut drain_context);

                match poll_result {
                    task::Poll::Ready(())   => {
                        // Keep running jobs and checking the results if ready
                    },

                    task::Poll::Pending     => {
                        // Requeue the job
                        self.queue.requeue(job);

                        // If the result was supplied, break out of the loop and reschedule the queue
                        result = self.result.lock().expect("Scheduler future result").result.take();
                        if result.is_some() {
                            // Wake the queue in the background if needed (the result has arrived)
                            self.queue.core.lock().expect("JobQueue core lock").state = QueueState::WaitingForWake;

                            let queue_waker = WakeQueue(Arc::clone(&self.queue), Arc::clone(&self.scheduler));
                            let queue_waker = Arc::new(queue_waker);
                            let queue_waker = task::waker(queue_waker);

                            waker.wake_with(queue_waker);

                            // The future is ready (job will be rescheduled in the background)
                            return task::Poll::Ready(result.unwrap());
                        } else {
                            // Wait for the next poll
                            self.queue.core.lock().expect("JobQueue core lock").state = QueueState::WaitingForPoll(self.id);

                            // Use the context waker
                            waker.wake_with(context.waker().clone());

                            // Result is pending
                            return task::Poll::Pending;
                        }
                    }
                }
            } else {
                // Queue is empty and our result hasn't arrived yet?!

                // Assume the future will resolve eventually: move the queue in to the background
                self.result.lock().expect("Scheduler future result").waker = Some(context.waker().clone());

                // Reschedule the queue
                self.queue.core.lock().expect("JobQueue core lock").state = QueueState::Idle;
                self.scheduler.reschedule_queue(&self.queue, Arc::clone(&self.scheduler));

                return task::Poll::Pending;
            }
        }

        // Reschedule the queue if there are any events left pending
        // Note: the queue is already pending when we start running events from it here.
        // This means it'll get dequeued by a thread eventually: maybe while it's running
        // here. As we've set the queue state to running while we're busy, the thread won't
        // start the queue while it's already running.
        self.queue.core.lock().expect("JobQueue core lock").state = QueueState::Idle;
        self.scheduler.reschedule_queue(&self.queue, Arc::clone(&self.scheduler));

        // Result must be available by this point
        task::Poll::Ready(result.unwrap())
    }
}

impl<T> Future for SchedulerFuture<T> {
    type Output = Result<T, oneshot::Canceled>;

    ///
    /// Polls this future
    ///
    fn poll(mut self: Pin<&mut Self>, context: &mut task::Context) -> task::Poll<Self::Output> {
        // Lock the result and determine which action to take
        let next_action = {
            let mut future_result = self.result.lock().expect("Scheduler future result");

            if let Some(result) = future_result.result.take() {
                // The result is available: we should return it immediately
                SchedulerAction::ReturnValue(result)
            } else {
                // If the queue is idle when this is called, we need to schedule this task on this thread rather than one owned by the background process
                let run_action = {
                    let mut core = self.queue.core.lock().expect("JobQueue core lock");

                    match core.state {
                        QueueState::Running                     => SchedulerAction::WaitForCompletion,
                        QueueState::WaitingForWake              => SchedulerAction::WaitForCompletion,
                        QueueState::WaitingForUnpark            => SchedulerAction::WaitForCompletion,
                        QueueState::AwokenWhileRunning          => SchedulerAction::WaitForCompletion,
                        QueueState::Panicked                    => SchedulerAction::Panic,
                        QueueState::Pending                     => { core.state = QueueState::Running; SchedulerAction::DrainQueue },
                        QueueState::Idle                        => { core.state = QueueState::Running; SchedulerAction::DrainQueue }

                        QueueState::WaitingForPoll(owner_id)    => {
                            if owner_id == self.id {
                                // Continue polling on this future
                                core.state = QueueState::Running; SchedulerAction::DrainQueue
                            } else {
                                // Wait for the owning future to complete
                                SchedulerAction::WaitForCompletion
                            }
                        },
                    }
                };

                if let SchedulerAction::WaitForCompletion = run_action {
                    // Wake us up when the future is available
                    future_result.waker = Some(context.waker().clone());
                }

                run_action
            }
        };

        match next_action {
            SchedulerAction::WaitForCompletion  => task::Poll::Pending,
            SchedulerAction::ReturnValue(value) => task::Poll::Ready(value),
            SchedulerAction::DrainQueue         => self.drain_queue(context),
            SchedulerAction::Panic              => panic!("Cannot schedule jobs on a panicked queue"),
        }
    }
}

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

    #[test]
    fn returns_immediately_if_signaled() {
        use futures::executor;

        let scheduler           = scheduler();
        let queue               = queue();
        let (future, signaller) = SchedulerFuture::new(&queue, Arc::clone(&scheduler.core));

        signaller.signal(42);

        assert!(executor::block_on(future) == Ok(42));
    }

    #[test]
    fn cancels_when_dropped() {
        use futures::executor;

        let scheduler           = scheduler();
        let queue               = queue();
        let future              = { SchedulerFuture::<i32>::new(&queue, Arc::clone(&scheduler.core)).0 };

        assert!(executor::block_on(future) == Err(oneshot::Canceled));
    }

    #[test]
    fn signals_from_another_thread() {
        use futures::executor;
        use std::thread;
        use std::time::Duration;

        let scheduler           = scheduler();
        let queue               = queue();
        let (future, signaller) = SchedulerFuture::new(&queue, Arc::clone(&scheduler.core));

        thread::spawn(move || {
            thread::sleep(Duration::from_millis(100));

            signaller.signal(42);
        });

        assert!(executor::block_on(future) == Ok(42));
    }

    #[test]
    fn forces_queue_drain() {
        use futures::executor;
        use std::thread;
        use std::time::Duration;

        let scheduler   = scheduler();
        let queue       = queue();

        scheduler.set_max_threads(0);
        scheduler.despawn_threads_if_overloaded();

        let (future, signaller) = SchedulerFuture::new(&queue, Arc::clone(&scheduler.core));

        scheduler.desync(&queue, move || {
            thread::sleep(Duration::from_millis(100));

            signaller.signal(42);
        });

        assert!(executor::block_on(future) == Ok(42));
    }
}
<|endoftext|>


<commit_before>use std::fs::{File, OpenOptions};
use std::io::{Read, Write, Seek, SeekFrom};
use std::cell::RefCell;

trait Mbc {
    fn read(&self, address: u16) -> u8;
    fn write(&mut self, address: u16, v: u8);
}

struct Mbc0 {
    data: Vec<u8>,
    ram: [u8; 8192],
}

impl Mbc0 {
    pub fn new(data: Vec<u8>) -> Mbc0 {
        Mbc0 {
            data: data,
            ram: [0; 8192],
        }
    }
}

impl Mbc for Mbc0 {
    fn read(&self, address: u16) -> u8 {
        match address {
            0x0000 ... 0x7FFF => self.data[address as usize],
            0xA000 ... 0xBFFF => self.ram[(address as usize) - 0xA000],
            _ => unimplemented!(),
        }
    }

    fn write(&mut self, _: u16, _: u8) {
        // println!("Writing {:02X} to {:04X}", v, address);
        // match address {
        //    0xA000 ... 0xBFFF => self.ram[(address as usize) - 0xA000] = v,
        //    _ => unimplemented!(),
        //}
    }
}

#[allow(dead_code)]
#[derive(PartialEq, Eq)]
enum MemoryMode {
    C4_32,
    C16_8,
}

#[allow(dead_code)]
struct Mbc3 {
    selected_bank: usize,
    data: Vec<u8>,
    offset: usize,
    ram: RefCell<File>,
    memory_mode: MemoryMode,

    ram_offset: usize,
    ram_enabled: bool,
}

impl Mbc3 {
    pub fn new(data: Vec<u8>) -> Mbc3 {
        let mut disk = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open("save.dat")
            .unwrap();

        let mut ram = vec![];
        disk.read_to_end(&mut ram).unwrap();

        if ram.len() < 32768 {
            println!("Warning: invalid file size, blanking ram.");
            disk.seek(SeekFrom::Start(0)).unwrap();
            disk.write_all(&[0; 32768]).unwrap();
        }

        Mbc3 {
            selected_bank: 1,
            data: data,
            offset: 0,
            ram: RefCell::new(disk),
            memory_mode: MemoryMode::C4_32,

            ram_offset: 0,
            ram_enabled: false,
        }
    }

    fn write_ram(&mut self, address: u16, v: u8) {
        if !self.ram_enabled {
            panic!("MBC3 RAM has not been enabled.");
        }

        let address = (address - 0xA000) as usize + self.ram_offset;
        self.ram.borrow_mut().seek(SeekFrom::Start(address as u64)).unwrap();
        self.ram.borrow_mut().write_all(&[v]).unwrap();
    }

    fn read_ram(&self, address: u16) -> u8 {
        if !self.ram_enabled {
            panic!("MBC3 RAM has not been enabled.");
        }

        let address = (address - 0xA000) as usize + self.ram_offset;
        self.ram.borrow_mut().seek(SeekFrom::Start(address as u64)).unwrap();
        let mut buffer = [0; 1];
        self.ram.borrow_mut().read_exact(&mut buffer).unwrap();

        buffer[0]
    }
}

impl Mbc for Mbc3 {
    fn read(&self, address: u16) -> u8 {
        match address {
            0x0000 ... 0x3FFF => self.data[address as usize],
            0x4000 ... 0x7FFF => self.data[address as usize + self.offset],
            0xA000 ... 0xBFFF => self.read_ram(address),
            _ => unimplemented!(),
        }
    }

    fn write(&mut self, address: u16, v: u8) {
        match address {
            0x0000 ... 0x1FFF => {
                match v {
                    0x00 => self.ram_enabled = false,
                    0x0A => self.ram_enabled = true,
                    _    => panic!("Unrecognized value for RAM enablement."),
                }
            },
            0x2000 ... 0x3FFF => {
                // Any write to this area will enable the bank of memory contained in v
                // if 0 is selected we will select 1 because 0 is already available
                // statically in 0 0000 - 3FFF
                if v > 0 {
                    self.offset = ((v - 1) as usize) * 0x4000;
                } else {
                    self.offset = 0x0000;
                }
                // println!("Setting offset to {:06X}", self.offset);
            },
            0x4000 ... 0x5FFF => {
                match v {
                    0x00 => self.ram_offset = 0,
                    0x01 => self.ram_offset = 8192,
                    0x02 => self.ram_offset = 16384,
                    0x03 => self.ram_offset = 24576,
                    _ => unimplemented!(),
                }
            },
            0x6000 ... 0x7FFF => {
                // println!("Attempt to write {:02X} to {:04X}", v, address);
                /* match v & 0x1 {
                    0x0 => self.memory_mode = MemoryMode::C16_8, // Memory mode 16_8,
                    0x1 => self.memory_mode = MemoryMode::C4_32,
                    _   => panic!(),
                }*/
            },
            0xA000 ... 0xBFFF => { self.write_ram(address, v); },
            _ => unimplemented!(),
        }
    }
}

pub struct MemoryController {
    controller: Box<Mbc>,
}

impl MemoryController {
    pub fn from_bytes(bytes: Vec<u8>) -> MemoryController {
        let controller = match bytes[0x147] {
            0x00 => Box::new(Mbc0::new(bytes)) as Box<Mbc>,
            0x0F | 0x10 | 0x11 | 0x12 | 0x13 | 0x01 => Box::new(Mbc3::new(bytes)) as Box<Mbc>,
            _ => {
                println!("Unrecognized type {:02X}", bytes[0x147]);
                panic!();
            }
        };

        MemoryController {
            controller: controller,
        }
    }

    pub fn read(&self, address: u16) -> u8 {
        self.controller.read(address)
    }

    pub fn write(&mut self, address: u16, v: u8) {
        self.controller.write(address, v);
    }
}
<commit_msg>Only consider lowest two bits when writing to 0x4000-0x5FFF to switch bank.<commit_after>use std::fs::{File, OpenOptions};
use std::io::{Read, Write, Seek, SeekFrom};
use std::cell::RefCell;

trait Mbc {
    fn read(&self, address: u16) -> u8;
    fn write(&mut self, address: u16, v: u8);
}

struct Mbc0 {
    data: Vec<u8>,
    ram: [u8; 8192],
}

impl Mbc0 {
    pub fn new(data: Vec<u8>) -> Mbc0 {
        Mbc0 {
            data: data,
            ram: [0; 8192],
        }
    }
}

impl Mbc for Mbc0 {
    fn read(&self, address: u16) -> u8 {
        match address {
            0x0000 ... 0x7FFF => self.data[address as usize],
            0xA000 ... 0xBFFF => self.ram[(address as usize) - 0xA000],
            _ => unimplemented!(),
        }
    }

    fn write(&mut self, _: u16, _: u8) {
        // println!("Writing {:02X} to {:04X}", v, address);
        // match address {
        //    0xA000 ... 0xBFFF => self.ram[(address as usize) - 0xA000] = v,
        //    _ => unimplemented!(),
        //}
    }
}

#[allow(dead_code)]
#[derive(PartialEq, Eq)]
enum MemoryMode {
    C4_32,
    C16_8,
}

#[allow(dead_code)]
struct Mbc3 {
    selected_bank: usize,
    data: Vec<u8>,
    offset: usize,
    ram: RefCell<File>,
    memory_mode: MemoryMode,

    ram_offset: usize,
    ram_enabled: bool,
}

impl Mbc3 {
    pub fn new(data: Vec<u8>) -> Mbc3 {
        let mut disk = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open("save.dat")
            .unwrap();

        let mut ram = vec![];
        disk.read_to_end(&mut ram).unwrap();

        if ram.len() < 32768 {
            println!("Warning: invalid file size, blanking ram.");
            disk.seek(SeekFrom::Start(0)).unwrap();
            disk.write_all(&[0; 32768]).unwrap();
        }

        Mbc3 {
            selected_bank: 1,
            data: data,
            offset: 0,
            ram: RefCell::new(disk),
            memory_mode: MemoryMode::C4_32,

            ram_offset: 0,
            ram_enabled: false,
        }
    }

    fn write_ram(&mut self, address: u16, v: u8) {
        if !self.ram_enabled {
            panic!("MBC3 RAM has not been enabled.");
        }

        let address = (address - 0xA000) as usize + self.ram_offset;
        self.ram.borrow_mut().seek(SeekFrom::Start(address as u64)).unwrap();
        self.ram.borrow_mut().write_all(&[v]).unwrap();
    }

    fn read_ram(&self, address: u16) -> u8 {
        if !self.ram_enabled {
            panic!("MBC3 RAM has not been enabled.");
        }

        let address = (address - 0xA000) as usize + self.ram_offset;
        self.ram.borrow_mut().seek(SeekFrom::Start(address as u64)).unwrap();
        let mut buffer = [0; 1];
        self.ram.borrow_mut().read_exact(&mut buffer).unwrap();

        buffer[0]
    }
}

impl Mbc for Mbc3 {
    fn read(&self, address: u16) -> u8 {
        match address {
            0x0000 ... 0x3FFF => self.data[address as usize],
            0x4000 ... 0x7FFF => self.data[address as usize + self.offset],
            0xA000 ... 0xBFFF => self.read_ram(address),
            _ => unimplemented!(),
        }
    }

    fn write(&mut self, address: u16, v: u8) {
        match address {
            0x0000 ... 0x1FFF => {
                match v {
                    0x00 => self.ram_enabled = false,
                    0x0A => self.ram_enabled = true,
                    _    => panic!("Unrecognized value for RAM enablement."),
                }
            },
            0x2000 ... 0x3FFF => {
                // Any write to this area will enable the bank of memory contained in v
                // if 0 is selected we will select 1 because 0 is already available
                // statically in 0 0000 - 3FFF
                if v > 0 {
                    self.offset = ((v - 1) as usize) * 0x4000;
                } else {
                    self.offset = 0x0000;
                }
                // println!("Setting offset to {:06X}", self.offset);
            },
            0x4000 ... 0x5FFF => {
                // Writing to this area will cause the controller to switch
                // banks of in-cartridge RAM.
                // Only the lowest two bits are relevant for this register.
                match v & 0b11 {
                    0x00 => self.ram_offset = 0,
                    0x01 => self.ram_offset = 8192,
                    0x02 => self.ram_offset = 16384,
                    0x03 => self.ram_offset = 24576,
                    _ => unimplemented!(),
                }
            },
            0x6000 ... 0x7FFF => {
                // println!("Attempt to write {:02X} to {:04X}", v, address);
                /* match v & 0x1 {
                    0x0 => self.memory_mode = MemoryMode::C16_8, // Memory mode 16_8,
                    0x1 => self.memory_mode = MemoryMode::C4_32,
                    _   => panic!(),
                }*/
            },
            0xA000 ... 0xBFFF => { self.write_ram(address, v); },
            _ => unimplemented!(),
        }
    }
}

pub struct MemoryController {
    controller: Box<Mbc>,
}

impl MemoryController {
    pub fn from_bytes(bytes: Vec<u8>) -> MemoryController {
        let controller = match bytes[0x147] {
            0x00 => Box::new(Mbc0::new(bytes)) as Box<Mbc>,
            0x0F | 0x10 | 0x11 | 0x12 | 0x13 | 0x01 => Box::new(Mbc3::new(bytes)) as Box<Mbc>,
            _ => {
                println!("Unrecognized type {:02X}", bytes[0x147]);
                panic!();
            }
        };

        MemoryController {
            controller: controller,
        }
    }

    pub fn read(&self, address: u16) -> u8 {
        self.controller.read(address)
    }

    pub fn write(&mut self, address: u16, v: u8) {
        self.controller.write(address, v);
    }
}
<|endoftext|>


<commit_before>//! TCP stream for communicating with shadowsocks' proxy server

use std::{
    io::{self, ErrorKind},
    pin::Pin,
    task::{self, Poll},
};

use bytes::{BufMut, BytesMut};
use cfg_if::cfg_if;
use futures::ready;
use log::trace;
use once_cell::sync::Lazy;
use pin_project::pin_project;
use tokio::{
    io::{AsyncRead, AsyncWrite, ReadBuf},
    time,
};

#[cfg(feature = "aead-cipher-2022")]
use crate::relay::get_aead_2022_padding_size;
use crate::{
    config::ServerConfig,
    context::SharedContext,
    crypto::CipherKind,
    net::{ConnectOpts, TcpStream as OutboundTcpStream},
    relay::{
        socks5::Address,
        tcprelay::crypto_io::{CryptoRead, CryptoStream, CryptoWrite, StreamType},
    },
};

enum ProxyClientStreamWriteState {
    Connect(Address),
    Connecting(BytesMut),
    Connected,
}

enum ProxyClientStreamReadState {
    #[cfg(feature = "aead-cipher-2022")]
    CheckRequestNonce,
    Established,
}

/// A stream for sending / receiving data stream from remote server via shadowsocks' proxy server
#[pin_project]
pub struct ProxyClientStream<S> {
    #[pin]
    stream: CryptoStream<S>,
    writer_state: ProxyClientStreamWriteState,
    reader_state: ProxyClientStreamReadState,
    context: SharedContext,
}

static DEFAULT_CONNECT_OPTS: Lazy<ConnectOpts> = Lazy::new(Default::default);

impl ProxyClientStream<OutboundTcpStream> {
    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`
    pub async fn connect<A>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
    ) -> io::Result<ProxyClientStream<OutboundTcpStream>>
    where
        A: Into<Address>,
    {
        ProxyClientStream::connect_with_opts(context, svr_cfg, addr, &DEFAULT_CONNECT_OPTS).await
    }

    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`
    pub async fn connect_with_opts<A>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        opts: &ConnectOpts,
    ) -> io::Result<ProxyClientStream<OutboundTcpStream>>
    where
        A: Into<Address>,
    {
        ProxyClientStream::connect_with_opts_map(context, svr_cfg, addr, opts, |s| s).await
    }
}

impl<S> ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`, maps `TcpStream` to customized stream with `map_fn`
    pub async fn connect_map<A, F>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        map_fn: F,
    ) -> io::Result<ProxyClientStream<S>>
    where
        A: Into<Address>,
        F: FnOnce(OutboundTcpStream) -> S,
    {
        ProxyClientStream::connect_with_opts_map(context, svr_cfg, addr, &DEFAULT_CONNECT_OPTS, map_fn).await
    }

    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`, maps `TcpStream` to customized stream with `map_fn`
    pub async fn connect_with_opts_map<A, F>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        opts: &ConnectOpts,
        map_fn: F,
    ) -> io::Result<ProxyClientStream<S>>
    where
        A: Into<Address>,
        F: FnOnce(OutboundTcpStream) -> S,
    {
        let stream = match svr_cfg.timeout() {
            Some(d) => {
                match time::timeout(
                    d,
                    OutboundTcpStream::connect_server_with_opts(&context, svr_cfg.external_addr(), opts),
                )
                .await
                {
                    Ok(Ok(s)) => s,
                    Ok(Err(e)) => return Err(e),
                    Err(..) => {
                        return Err(io::Error::new(
                            ErrorKind::TimedOut,
                            format!("connect {} timeout", svr_cfg.addr()),
                        ))
                    }
                }
            }
            None => OutboundTcpStream::connect_server_with_opts(&context, svr_cfg.external_addr(), opts).await?,
        };

        trace!(
            "connected tcp remote {} (outbound: {}) with {:?}",
            svr_cfg.addr(),
            svr_cfg.external_addr(),
            opts
        );

        Ok(ProxyClientStream::from_stream(context, map_fn(stream), svr_cfg, addr))
    }

    /// Create a `ProxyClientStream` with a connected `stream` to a shadowsocks' server
    ///
    /// NOTE: `stream` must be connected to the server with the same configuration as `svr_cfg`, otherwise strange errors would occurs
    pub fn from_stream<A>(context: SharedContext, stream: S, svr_cfg: &ServerConfig, addr: A) -> ProxyClientStream<S>
    where
        A: Into<Address>,
    {
        let addr = addr.into();
        let stream = CryptoStream::from_stream(&context, stream, StreamType::Client, svr_cfg.method(), svr_cfg.key());

        #[cfg(not(feature = "aead-cipher-2022"))]
        let reader_state = ProxyClientStreamReadState::Established;

        #[cfg(feature = "aead-cipher-2022")]
        let reader_state = if svr_cfg.method().is_aead_2022() {
            // AEAD 2022 has a respond header
            ProxyClientStreamReadState::CheckRequestNonce
        } else {
            ProxyClientStreamReadState::Established
        };

        ProxyClientStream {
            stream,
            writer_state: ProxyClientStreamWriteState::Connect(addr),
            reader_state,
            context,
        }
    }

    /// Get reference to the underlying stream
    pub fn get_ref(&self) -> &S {
        self.stream.get_ref()
    }

    /// Get mutable reference to the underlying stream
    pub fn get_mut(&mut self) -> &mut S {
        self.stream.get_mut()
    }

    /// Consumes the `ProxyClientStream` and return the underlying stream
    pub fn into_inner(self) -> S {
        self.stream.into_inner()
    }
}

impl<S> AsyncRead for ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    #[inline]
    fn poll_read(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        #[allow(unused_mut)]
        let mut this = self.project();

        #[allow(clippy::never_loop)]
        loop {
            match this.reader_state {
                ProxyClientStreamReadState::Established => {
                    return this.stream.poll_read_decrypted(cx, this.context, buf);
                }
                #[cfg(feature = "aead-cipher-2022")]
                ProxyClientStreamReadState::CheckRequestNonce => {
                    ready!(this.stream.as_mut().poll_read_decrypted(cx, this.context, buf))?;

                    // REQUEST_NONCE should be in the respond packet (header) of AEAD-2022.
                    //
                    // If received_request_nonce() is None, then:
                    // 1. method.salt_len() == 0, no checking required.
                    // 2. TCP stream read() returns EOF before receiving the header, no checking required.
                    //
                    // poll_read_decrypted will wait until the first non-zero size data chunk.
                    let method = this.stream.method();
                    if method.salt_len() > 0 {
                        if let Some(request_nonce) = this.stream.received_request_nonce() {
                            if this.stream.sent_nonce() != request_nonce {
                                return Err(io::Error::new(
                                    ErrorKind::Other,
                                    "received TCP response header with unmatched salt",
                                ))
                                .into();
                            }
                        }
                    }

                    *(this.reader_state) = ProxyClientStreamReadState::Established;
                    return Ok(()).into();
                }
            }
        }
    }
}

#[inline]
fn make_first_packet_buffer(method: CipherKind, addr: &Address, buf: &[u8]) -> BytesMut {
    // Target Address should be sent with the first packet together,
    // which would prevent from being detected.

    let addr_length = addr.serialized_len();
    let mut buffer = BytesMut::new();

    cfg_if! {
        if #[cfg(feature = "aead-cipher-2022")] {
            let padding_size = get_aead_2022_padding_size(buf);
            let header_length = if method.is_aead_2022() {
                addr_length + 2 + padding_size + buf.len()
            } else {
                addr_length + buf.len()
            };
        } else {
            let _ = method;
            let header_length = addr_length + buf.len();
        }
    }

    buffer.reserve(header_length);

    // STREAM / AEAD / AEAD2022 protocol, append the Address before payload
    addr.write_to_buf(&mut buffer);

    #[cfg(feature = "aead-cipher-2022")]
    if method.is_aead_2022() {
        buffer.put_u16(padding_size as u16);

        if padding_size > 0 {
            unsafe {
                buffer.advance_mut(padding_size);
            }
        }
    }

    buffer.put_slice(buf);

    buffer
}

impl<S> AsyncWrite for ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_write(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
        let this = self.project();

        loop {
            match this.writer_state {
                ProxyClientStreamWriteState::Connect(ref addr) => {
                    let buffer = make_first_packet_buffer(this.stream.method(), addr, buf);

                    // Save the concatenated buffer before it is written successfully.
                    // APIs require buffer to be kept alive before Poll::Ready
                    //
                    // Proactor APIs like IOCP on Windows, pointers of buffers have to be kept alive
                    // before IO completion.
                    *(this.writer_state) = ProxyClientStreamWriteState::Connecting(buffer);
                }
                ProxyClientStreamWriteState::Connecting(ref buffer) => {
                    let n = ready!(this.stream.poll_write_encrypted(cx, buffer))?;

                    // In general, poll_write_encrypted should perform like write_all.
                    debug_assert!(n == buffer.len());

                    *(this.writer_state) = ProxyClientStreamWriteState::Connected;

                    // NOTE:
                    // poll_write will return Ok(0) if buf.len() == 0
                    // But for the first call, this function will eventually send the handshake packet (IV/Salt + ADDR) to the remote address.
                    //
                    // https://github.com/shadowsocks/shadowsocks-rust/issues/232
                    //
                    // For protocols that requires *Server Hello* message, like FTP, clients won't send anything to the server until server sends handshake messages.
                    // This could be achieved by calling poll_write with an empty input buffer.
                    return Ok(buf.len()).into();
                }
                ProxyClientStreamWriteState::Connected => {
                    return this.stream.poll_write_encrypted(cx, buf);
                }
            }
        }
    }

    #[inline]
    fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.project().stream.poll_flush(cx)
    }

    #[inline]
    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.project().stream.poll_shutdown(cx)
    }
}
<commit_msg>check REQUEST_SALT only when received at least 1 data chunk<commit_after>//! TCP stream for communicating with shadowsocks' proxy server

use std::{
    io::{self, ErrorKind},
    pin::Pin,
    task::{self, Poll},
};

use bytes::{BufMut, BytesMut};
use cfg_if::cfg_if;
use futures::ready;
use log::trace;
use once_cell::sync::Lazy;
use pin_project::pin_project;
use tokio::{
    io::{AsyncRead, AsyncWrite, ReadBuf},
    time,
};

#[cfg(feature = "aead-cipher-2022")]
use crate::relay::get_aead_2022_padding_size;
use crate::{
    config::ServerConfig,
    context::SharedContext,
    crypto::CipherKind,
    net::{ConnectOpts, TcpStream as OutboundTcpStream},
    relay::{
        socks5::Address,
        tcprelay::crypto_io::{CryptoRead, CryptoStream, CryptoWrite, StreamType},
    },
};

enum ProxyClientStreamWriteState {
    Connect(Address),
    Connecting(BytesMut),
    Connected,
}

enum ProxyClientStreamReadState {
    #[cfg(feature = "aead-cipher-2022")]
    CheckRequestNonce,
    Established,
}

/// A stream for sending / receiving data stream from remote server via shadowsocks' proxy server
#[pin_project]
pub struct ProxyClientStream<S> {
    #[pin]
    stream: CryptoStream<S>,
    writer_state: ProxyClientStreamWriteState,
    reader_state: ProxyClientStreamReadState,
    context: SharedContext,
}

static DEFAULT_CONNECT_OPTS: Lazy<ConnectOpts> = Lazy::new(Default::default);

impl ProxyClientStream<OutboundTcpStream> {
    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`
    pub async fn connect<A>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
    ) -> io::Result<ProxyClientStream<OutboundTcpStream>>
    where
        A: Into<Address>,
    {
        ProxyClientStream::connect_with_opts(context, svr_cfg, addr, &DEFAULT_CONNECT_OPTS).await
    }

    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`
    pub async fn connect_with_opts<A>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        opts: &ConnectOpts,
    ) -> io::Result<ProxyClientStream<OutboundTcpStream>>
    where
        A: Into<Address>,
    {
        ProxyClientStream::connect_with_opts_map(context, svr_cfg, addr, opts, |s| s).await
    }
}

impl<S> ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`, maps `TcpStream` to customized stream with `map_fn`
    pub async fn connect_map<A, F>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        map_fn: F,
    ) -> io::Result<ProxyClientStream<S>>
    where
        A: Into<Address>,
        F: FnOnce(OutboundTcpStream) -> S,
    {
        ProxyClientStream::connect_with_opts_map(context, svr_cfg, addr, &DEFAULT_CONNECT_OPTS, map_fn).await
    }

    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`, maps `TcpStream` to customized stream with `map_fn`
    pub async fn connect_with_opts_map<A, F>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        opts: &ConnectOpts,
        map_fn: F,
    ) -> io::Result<ProxyClientStream<S>>
    where
        A: Into<Address>,
        F: FnOnce(OutboundTcpStream) -> S,
    {
        let stream = match svr_cfg.timeout() {
            Some(d) => {
                match time::timeout(
                    d,
                    OutboundTcpStream::connect_server_with_opts(&context, svr_cfg.external_addr(), opts),
                )
                .await
                {
                    Ok(Ok(s)) => s,
                    Ok(Err(e)) => return Err(e),
                    Err(..) => {
                        return Err(io::Error::new(
                            ErrorKind::TimedOut,
                            format!("connect {} timeout", svr_cfg.addr()),
                        ))
                    }
                }
            }
            None => OutboundTcpStream::connect_server_with_opts(&context, svr_cfg.external_addr(), opts).await?,
        };

        trace!(
            "connected tcp remote {} (outbound: {}) with {:?}",
            svr_cfg.addr(),
            svr_cfg.external_addr(),
            opts
        );

        Ok(ProxyClientStream::from_stream(context, map_fn(stream), svr_cfg, addr))
    }

    /// Create a `ProxyClientStream` with a connected `stream` to a shadowsocks' server
    ///
    /// NOTE: `stream` must be connected to the server with the same configuration as `svr_cfg`, otherwise strange errors would occurs
    pub fn from_stream<A>(context: SharedContext, stream: S, svr_cfg: &ServerConfig, addr: A) -> ProxyClientStream<S>
    where
        A: Into<Address>,
    {
        let addr = addr.into();
        let stream = CryptoStream::from_stream(&context, stream, StreamType::Client, svr_cfg.method(), svr_cfg.key());

        #[cfg(not(feature = "aead-cipher-2022"))]
        let reader_state = ProxyClientStreamReadState::Established;

        #[cfg(feature = "aead-cipher-2022")]
        let reader_state = if svr_cfg.method().is_aead_2022() {
            // AEAD 2022 has a respond header
            ProxyClientStreamReadState::CheckRequestNonce
        } else {
            ProxyClientStreamReadState::Established
        };

        ProxyClientStream {
            stream,
            writer_state: ProxyClientStreamWriteState::Connect(addr),
            reader_state,
            context,
        }
    }

    /// Get reference to the underlying stream
    pub fn get_ref(&self) -> &S {
        self.stream.get_ref()
    }

    /// Get mutable reference to the underlying stream
    pub fn get_mut(&mut self) -> &mut S {
        self.stream.get_mut()
    }

    /// Consumes the `ProxyClientStream` and return the underlying stream
    pub fn into_inner(self) -> S {
        self.stream.into_inner()
    }
}

impl<S> AsyncRead for ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    #[inline]
    fn poll_read(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        #[allow(unused_mut)]
        let mut this = self.project();

        #[allow(clippy::never_loop)]
        loop {
            match this.reader_state {
                ProxyClientStreamReadState::Established => {
                    return this.stream.poll_read_decrypted(cx, this.context, buf);
                }
                #[cfg(feature = "aead-cipher-2022")]
                ProxyClientStreamReadState::CheckRequestNonce => {
                    ready!(this.stream.as_mut().poll_read_decrypted(cx, this.context, buf))?;

                    // REQUEST_NONCE should be in the respond packet (header) of AEAD-2022.
                    //
                    // If received_request_nonce() is None, then:
                    // 1. method.salt_len() == 0, no checking required.
                    // 2. TCP stream read() returns EOF before receiving the header, no checking required.
                    //
                    // poll_read_decrypted will wait until the first non-zero size data chunk.
                    let (data_chunk_count, _) = this.stream.current_data_chunk_remaining();
                    if data_chunk_count > 0 {
                        // data_chunk_count > 0, so the reader received at least 1 data chunk.

                        let sent_nonce = this.stream.sent_nonce();
                        let sent_nonce = if sent_nonce.is_empty() { None } else { Some(sent_nonce) };
                        if sent_nonce != this.stream.received_request_nonce() {
                            return Err(io::Error::new(
                                ErrorKind::Other,
                                "received TCP response header with unmatched salt",
                            ))
                            .into();
                        }

                        *(this.reader_state) = ProxyClientStreamReadState::Established;
                    }

                    return Ok(()).into();
                }
            }
        }
    }
}

#[inline]
fn make_first_packet_buffer(method: CipherKind, addr: &Address, buf: &[u8]) -> BytesMut {
    // Target Address should be sent with the first packet together,
    // which would prevent from being detected.

    let addr_length = addr.serialized_len();
    let mut buffer = BytesMut::new();

    cfg_if! {
        if #[cfg(feature = "aead-cipher-2022")] {
            let padding_size = get_aead_2022_padding_size(buf);
            let header_length = if method.is_aead_2022() {
                addr_length + 2 + padding_size + buf.len()
            } else {
                addr_length + buf.len()
            };
        } else {
            let _ = method;
            let header_length = addr_length + buf.len();
        }
    }

    buffer.reserve(header_length);

    // STREAM / AEAD / AEAD2022 protocol, append the Address before payload
    addr.write_to_buf(&mut buffer);

    #[cfg(feature = "aead-cipher-2022")]
    if method.is_aead_2022() {
        buffer.put_u16(padding_size as u16);

        if padding_size > 0 {
            unsafe {
                buffer.advance_mut(padding_size);
            }
        }
    }

    buffer.put_slice(buf);

    buffer
}

impl<S> AsyncWrite for ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_write(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
        let this = self.project();

        loop {
            match this.writer_state {
                ProxyClientStreamWriteState::Connect(ref addr) => {
                    let buffer = make_first_packet_buffer(this.stream.method(), addr, buf);

                    // Save the concatenated buffer before it is written successfully.
                    // APIs require buffer to be kept alive before Poll::Ready
                    //
                    // Proactor APIs like IOCP on Windows, pointers of buffers have to be kept alive
                    // before IO completion.
                    *(this.writer_state) = ProxyClientStreamWriteState::Connecting(buffer);
                }
                ProxyClientStreamWriteState::Connecting(ref buffer) => {
                    let n = ready!(this.stream.poll_write_encrypted(cx, buffer))?;

                    // In general, poll_write_encrypted should perform like write_all.
                    debug_assert!(n == buffer.len());

                    *(this.writer_state) = ProxyClientStreamWriteState::Connected;

                    // NOTE:
                    // poll_write will return Ok(0) if buf.len() == 0
                    // But for the first call, this function will eventually send the handshake packet (IV/Salt + ADDR) to the remote address.
                    //
                    // https://github.com/shadowsocks/shadowsocks-rust/issues/232
                    //
                    // For protocols that requires *Server Hello* message, like FTP, clients won't send anything to the server until server sends handshake messages.
                    // This could be achieved by calling poll_write with an empty input buffer.
                    return Ok(buf.len()).into();
                }
                ProxyClientStreamWriteState::Connected => {
                    return this.stream.poll_write_encrypted(cx, buf);
                }
            }
        }
    }

    #[inline]
    fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.project().stream.poll_flush(cx)
    }

    #[inline]
    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.project().stream.poll_shutdown(cx)
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2017 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#![feature(fn_traits, unboxed_closures)]

extern crate chrono;
extern crate gtk;
#[macro_use]
extern crate relm;
#[macro_use]
extern crate relm_derive;

use chrono::Local;
use gtk::{
    ContainerExt,
    Inhibit,
    Label,
    LabelExt,
    WidgetExt,
    Window,
    WindowType,
};
use relm::{Relm, Update, Widget, interval};

use self::Msg::*;

#[derive(SimpleMsg)]
enum Msg {
    Quit,
    Tick,
}

struct Win {
    label: Label,
    window: Window,
}

impl Update for Win {
    type Model = ();
    type ModelParam = ();
    type Msg = Msg;

    fn model(_: &Relm<Self>, _: ()) -> () {
        ()
    }

    fn subscriptions(&mut self, relm: &Relm<Self>) {
        interval(relm.stream(), 1000, || Tick);
    }

    fn update(&mut self, event: Msg) {
        match event {
            Tick => {
                let time = Local::now();
                self.label.set_text(&format!("{}", time.format("%H:%M:%S")));
            },
            Quit => gtk::main_quit(),
        }
    }
}

impl Widget for Win {
    type Root = Window;

    fn root(&self) -> Self::Root {
        self.window.clone()
    }

    fn view(relm: &Relm<Self>, _model: Self::Model) -> Self {
        let label = Label::new(None);

        let window = Window::new(WindowType::Toplevel);

        window.add(&label);

        window.show_all();

        connect!(relm, window, connect_delete_event(_, _), return (Some(Quit), Inhibit(false)));

        let mut win = Win {
            label: label,
            window: window,
        };
        win.update(Tick);
        win
    }
}

fn main() {
    Win::run(()).unwrap();
}
<commit_msg>Remove useless example<commit_after><|endoftext|>


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

extern crate time;

use std::fs::File;
use std::io::BufWriter;
use std::io::prelude::Write;

/*================================================================================================*/
/*------ENUMS-------------------------------------------------------------------------------------*/
/*================================================================================================*/

/// The log level enum
///
/// It is used with the logger struct, and is used to determine the severity of the message.
#[derive (Copy, Clone)]
pub enum LogLevel {

    /// Used for logging debug messages (stripped out of release builds)
    Debug (&'static str),
    /// Used for normal logging messages
    Info (&'static str),
    /// Used for warnings (non critical errors)
    Warning (&'static str),
    /// Used for critical errors
    Error (&'static str)
}

/*================================================================================================*/
/*------STRUCTS-----------------------------------------------------------------------------------*/
/*================================================================================================*/

/// The logger struct.
///
/// It allows for basic logging capabilities.
/// It can currently log strings, but others will be added in the future.
pub struct Logger {

    // Private
    log_file    : Option <BufWriter <File>>,
    is_enabled  : bool,
    is_logging  : bool
}

/*================================================================================================*/
/*------PUBLIC STATIC FUNCTIONS-------------------------------------------------------------------*/
/*================================================================================================*/

impl Logger {

    /// Creates a new logger instance
    ///
    /// Logging is disabled on faliure.
    ///
    /// # Examples
    /// ```
    /// let logger = Logger::new ("log.txt");
    pub fn new (log_file_path : &str) -> Logger {

        match File::create (&log_file_path) {

            Ok (file) => Logger {log_file   : Some (BufWriter::new (file)),
                                 is_enabled : true,
                                 is_logging : false},

            Err (_) => {

                println! ("WARNING: Log file could not be created\n
                           Logging has been disabled");

                Logger {log_file   : None,
                        is_enabled : false,
                        is_logging : false}
            }
        }
    }

/*================================================================================================*/
/*------PUBLIC FUNCTIONS--------------------------------------------------------------------------*/
/*================================================================================================*/

    /// Starts file logging
    pub fn begin_log (&mut self) {

        self.is_logging = true;
        self.log_file.as_mut ().unwrap ().write ("---BEGIN LOG---\n\n".as_bytes ()).unwrap ();
    }

/*================================================================================================*/

    /// Logs a message
    ///
    /// # Examples
    /// ```
    /// let mut logger = Logger::new ("log.txt");
    ///
    /// logger.log (LogLevel::Debug   ("This is a log entry")) // Debug
    /// logger.log (LogLevel::Message ("This is a log entry")) // Message
    /// logger.log (LogLevel::Warning ("This is a log entry")) // Warning
    /// logger.log (LogLevel::Error   ("This is a log entry")) // Error
    pub fn log (&mut self, log_level : LogLevel) {

        // Check if logger is enabled and logging
        if self.is_enabled && self.is_logging {

            // Get the buffer from the logger struct, and get current time
            let buffer       = self.log_file.as_mut ().unwrap ();
            let current_time = time::now ();

            match log_level {

                // Debug
                LogLevel::Debug (m) => {

                    // Only call if run in debug
                    if cfg! (debug_assertions) {

                        let log_string = format! ("DEBUG   ({}) {}\n", current_time.asctime (), m);
                        buffer.write (log_string.as_bytes ()).unwrap ();
                    }
                },

                // Message
                LogLevel::Info (m) => {

                    let log_string = format! ("INFO    ({}) {}\n", current_time.asctime (), m);
                    buffer.write (log_string.as_bytes ()).unwrap ();
                },

                // Warning
                LogLevel::Warning (m) => {

                    let log_string = format! ("WARNING ({}) {}\n", current_time.asctime (), m);
                    buffer.write (log_string.as_bytes ()).unwrap ();
                },

                // Error
                LogLevel::Error (m) => {

                    let log_string = format! ("ERROR   ({}) {}\n", current_time.asctime (), m);
                    buffer.write (log_string.as_bytes ()).unwrap ();
                }
            }
        }
    }

/*================================================================================================*/

    /// Stops file logging
    pub fn end_log (&mut self) {

        self.is_logging = false;
        self.log_file.as_mut ().unwrap ().write ("\n---END LOG---\n".as_bytes ()).unwrap ();
    }
}
<commit_msg>Added extra checks to prevent panicking if logging fails to start<commit_after>/*================================================================================================*/
// Copyright 2016 Kyle Finlay
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/*================================================================================================*/

extern crate time;

use std::fs::File;
use std::io::BufWriter;
use std::io::prelude::Write;

/*================================================================================================*/
/*------ENUMS-------------------------------------------------------------------------------------*/
/*================================================================================================*/

/// The log level enum
///
/// It is used with the logger struct, and is used to determine the severity of the message.
#[derive (Copy, Clone)]
pub enum LogLevel {

    /// Used for logging debug messages (stripped out of release builds)
    Debug (&'static str),
    /// Used for normal logging messages
    Info (&'static str),
    /// Used for warnings (non critical errors)
    Warning (&'static str),
    /// Used for critical errors
    Error (&'static str)
}

/*================================================================================================*/
/*------STRUCTS-----------------------------------------------------------------------------------*/
/*================================================================================================*/

/// The logger struct.
///
/// It allows for basic logging capabilities.
/// It can currently log strings, but others will be added in the future.
pub struct Logger {

    // Private
    log_file    : Option <BufWriter <File>>,
    is_enabled  : bool,
    is_logging  : bool
}

/*================================================================================================*/
/*------PUBLIC STATIC FUNCTIONS-------------------------------------------------------------------*/
/*================================================================================================*/

impl Logger {

    /// Creates a new logger instance
    ///
    /// Logging is disabled on faliure.
    ///
    /// # Examples
    /// ```
    /// let logger = Logger::new ("log.txt");
    pub fn new (log_file_path : &str) -> Logger {

        match File::create (&log_file_path) {

            Ok (file) => Logger {log_file   : Some (BufWriter::new (file)),
                                 is_enabled : true,
                                 is_logging : false},

            Err (_) => {

                println! ("{}\n{}",
                          "WARNING: Log file could not be created",
                          "         Logging has been disabled\n");

                Logger {log_file   : None,
                        is_enabled : false,
                        is_logging : false}
            }
        }
    }

/*================================================================================================*/
/*------PUBLIC FUNCTIONS--------------------------------------------------------------------------*/
/*================================================================================================*/

    /// Starts file logging
    pub fn begin_log (&mut self) {

        // Check if enabled
        if self.is_enabled {

            self.is_logging = true;
            self.log_file.as_mut ().unwrap ().write ("---BEGIN LOG---\n\n".as_bytes ()).unwrap ();
        }
    }

/*================================================================================================*/

    /// Logs a message
    ///
    /// # Examples
    /// ```
    /// let mut logger = Logger::new ("log.txt");
    ///
    /// logger.log (LogLevel::Debug   ("This is a log entry")) // Debug
    /// logger.log (LogLevel::Message ("This is a log entry")) // Message
    /// logger.log (LogLevel::Warning ("This is a log entry")) // Warning
    /// logger.log (LogLevel::Error   ("This is a log entry")) // Error
    pub fn log (&mut self, log_level : LogLevel) {

        // Check if logger is enabled and logging
        if self.is_enabled && self.is_logging {

            // Get the buffer from the logger struct, and get current time
            let buffer       = self.log_file.as_mut ().unwrap ();
            let current_time = time::now ();

            match log_level {

                // Debug
                LogLevel::Debug (m) => {

                    // Only call if run in debug
                    if cfg! (debug_assertions) {

                        let log_string = format! ("DEBUG   ({}) {}\n", current_time.asctime (), m);
                        buffer.write (log_string.as_bytes ()).unwrap ();
                    }
                },

                // Message
                LogLevel::Info (m) => {

                    let log_string = format! ("INFO    ({}) {}\n", current_time.asctime (), m);
                    buffer.write (log_string.as_bytes ()).unwrap ();
                },

                // Warning
                LogLevel::Warning (m) => {

                    let log_string = format! ("WARNING ({}) {}\n", current_time.asctime (), m);
                    buffer.write (log_string.as_bytes ()).unwrap ();
                },

                // Error
                LogLevel::Error (m) => {

                    let log_string = format! ("ERROR   ({}) {}\n", current_time.asctime (), m);
                    buffer.write (log_string.as_bytes ()).unwrap ();
                }
            }
        }
    }

/*================================================================================================*/

    /// Stops file logging
    pub fn end_log (&mut self) {

        // Check if enabled
        if self.is_enabled {

            self.is_logging = false;
            self.log_file.as_mut ().unwrap ().write ("\n---END LOG---\n".as_bytes ()).unwrap ();
        }
    }
}
<|endoftext|>


<commit_before>//! `Source` data description.
use crate::data::RoomName;

with_base_fields_and_update_struct! {
    /// A source object, which creeps can gain energy by mining from.
    #[derive(Clone, Debug, PartialEq)]
    #[serde(rename_all = "camelCase")]
    pub struct Source {
        /// The source's current energy - available to be mined be creeps.
        pub energy: u32,
        /// The source's maximum energy - what `energy` resets to on regeneration.
        pub energy_capacity: u32,
        /// The amount of energy either harvested on this source specifically or for the room since
        /// the last invasion (not sure which it is).
        #[serde(default)]
        pub invader_harvested: u32,
        /// The game time at which the source will next regenerate. `None` when waiting on a creep to
        /// first harvest it.
        pub next_regeneration_time: Option<u32>,
        /// The number of ticks between when a source is first harvested after regeneration and when it next
        /// regenerates.
        pub ticks_to_regeneration: u32,
    }

    /// The update structure for a source object.
    #[derive(Clone, Debug)]
    #[serde(rename_all = "camelCase")]
    pub struct SourceUpdate {
        - energy: u32,
        - energy_capacity: u32,
        - invader_harvested: u32,
        - next_regeneration_time: Option<u32>,
        - ticks_to_regeneration: u32,
    }
}

#[cfg(test)]
mod test {
    use serde::Deserialize;
    use serde_json;

    use crate::data::RoomName;

    use super::Source;

    #[test]
    fn parse_source_and_update() {
        let json = json!({
            "_id": "57ef9dba86f108ae6e60e2fc",
            "energy": 260,
            "energyCapacity": 3000,
            "invaderHarvested": 29240,
            "nextRegenerationTime": 19894026,
            "room": "E4S61",
            "ticksToRegeneration": 300,
            "type": "source",
            "x": 26,
            "y": 9,
        });

        let mut obj = Source::deserialize(&json).unwrap();

        assert_eq!(
            obj,
            Source {
                id: "57ef9dba86f108ae6e60e2fc".to_owned(),
                room: RoomName::new("E4S61").unwrap(),
                x: 26,
                y: 9,
                energy: 260,
                energy_capacity: 3000,
                invader_harvested: 29240,
                next_regeneration_time: Some(19894026),
                ticks_to_regeneration: 300,
            }
        );

        obj.update(
            serde_json::from_value(json!({
                "x": 40,
                "y": 50,
                "energy": 0,
            }))
            .unwrap(),
        );

        assert_eq!(
            obj,
            Source {
                id: "57ef9dba86f108ae6e60e2fc".to_owned(),
                room: RoomName::new("E4S61").unwrap(),
                x: 40,
                y: 50,
                energy: 0,
                energy_capacity: 3000,
                invader_harvested: 29240,
                next_regeneration_time: Some(19894026),
                ticks_to_regeneration: 300,
            }
        );
    }

    #[test]
    fn handle_no_invader_harvested() {
        let json = json!({
            "_id": "5bbcad499099fc012e6370ab",
            "energy": 4000,
            "energyCapacity": 4000,
            "nextRegenerationTime": null,
            "room": "E6S26",
            "ticksToRegeneration": 300,
            "type": "source",
            "x": 40,
            "y": 34,
        });

        let obj = Source::deserialize(&json).unwrap();

        assert_eq!(
            obj,
            Source {
                id: "5bbcad499099fc012e6370ab".to_owned(),
                room: RoomName::new("E6S26").unwrap(),
                x: 40,
                y: 34,
                energy: 4000,
                energy_capacity: 4000,
                invader_harvested: 0,
                next_regeneration_time: None,
                ticks_to_regeneration: 300,
            }
        );
    }
}
<commit_msg>Update source.energy and capacity type to match others<commit_after>//! `Source` data description.
use crate::data::RoomName;

with_base_fields_and_update_struct! {
    /// A source object, which creeps can gain energy by mining from.
    #[derive(Clone, Debug, PartialEq)]
    #[serde(rename_all = "camelCase")]
    pub struct Source {
        /// The source's current energy - available to be mined be creeps.
        pub energy: i32,
        /// The source's maximum energy - what `energy` resets to on regeneration.
        pub energy_capacity: i32,
        /// The amount of energy either harvested on this source specifically or for the room since
        /// the last invasion (not sure which it is).
        #[serde(default)]
        pub invader_harvested: u32,
        /// The game time at which the source will next regenerate. `None` when waiting on a creep to
        /// first harvest it.
        pub next_regeneration_time: Option<u32>,
        /// The number of ticks between when a source is first harvested after regeneration and when it next
        /// regenerates.
        pub ticks_to_regeneration: u32,
    }

    /// The update structure for a source object.
    #[derive(Clone, Debug)]
    #[serde(rename_all = "camelCase")]
    pub struct SourceUpdate {
        - energy: i32,
        - energy_capacity: i32,
        - invader_harvested: u32,
        - next_regeneration_time: Option<u32>,
        - ticks_to_regeneration: u32,
    }
}

#[cfg(test)]
mod test {
    use serde::Deserialize;
    use serde_json;

    use crate::data::RoomName;

    use super::Source;

    #[test]
    fn parse_source_and_update() {
        let json = json!({
            "_id": "57ef9dba86f108ae6e60e2fc",
            "energy": 260,
            "energyCapacity": 3000,
            "invaderHarvested": 29240,
            "nextRegenerationTime": 19894026,
            "room": "E4S61",
            "ticksToRegeneration": 300,
            "type": "source",
            "x": 26,
            "y": 9,
        });

        let mut obj = Source::deserialize(&json).unwrap();

        assert_eq!(
            obj,
            Source {
                id: "57ef9dba86f108ae6e60e2fc".to_owned(),
                room: RoomName::new("E4S61").unwrap(),
                x: 26,
                y: 9,
                energy: 260,
                energy_capacity: 3000,
                invader_harvested: 29240,
                next_regeneration_time: Some(19894026),
                ticks_to_regeneration: 300,
            }
        );

        obj.update(
            serde_json::from_value(json!({
                "x": 40,
                "y": 50,
                "energy": 0,
            }))
            .unwrap(),
        );

        assert_eq!(
            obj,
            Source {
                id: "57ef9dba86f108ae6e60e2fc".to_owned(),
                room: RoomName::new("E4S61").unwrap(),
                x: 40,
                y: 50,
                energy: 0,
                energy_capacity: 3000,
                invader_harvested: 29240,
                next_regeneration_time: Some(19894026),
                ticks_to_regeneration: 300,
            }
        );
    }

    #[test]
    fn handle_no_invader_harvested() {
        let json = json!({
            "_id": "5bbcad499099fc012e6370ab",
            "energy": 4000,
            "energyCapacity": 4000,
            "nextRegenerationTime": null,
            "room": "E6S26",
            "ticksToRegeneration": 300,
            "type": "source",
            "x": 40,
            "y": 34,
        });

        let obj = Source::deserialize(&json).unwrap();

        assert_eq!(
            obj,
            Source {
                id: "5bbcad499099fc012e6370ab".to_owned(),
                room: RoomName::new("E6S26").unwrap(),
                x: 40,
                y: 34,
                energy: 4000,
                energy_capacity: 4000,
                invader_harvested: 0,
                next_regeneration_time: None,
                ticks_to_regeneration: 300,
            }
        );
    }
}
<|endoftext|>