<commit_before>mod events;

use winit::{WindowBuilder, get_available_monitors, get_primary_monitor, Event, ElementState};
use vulkano_win::{Window, VkSurfaceBuild, required_extensions};
use vulkano::instance::{Instance, PhysicalDevice};
use vulkano::device::{Queue, Device, DeviceExtensions};
use vulkano::swapchain::{Swapchain, SurfaceTransform, PresentMode};
use vulkano::image::SwapchainImage;
use vulkano::image::attachment::AttachmentImage;
use vulkano::framebuffer::Framebuffer;
use vulkano::command_buffer::{PrimaryCommandBufferBuilder, Submission, submit};
use vulkano::format;

use config::Config;
use config::WindowConfig;

use std::sync::Arc;
use std::time::Duration;
use std::collections::HashMap;

mod render_pass {

    use vulkano::format;
    use vulkano::format::Format;

    single_pass_renderpass!{
    attachments: {
        color: {
            load: Clear,
            store: Store,
            format: Format,
        },
        depth: {
            load: Clear,
            store: DontCare,
            format: format::D16Unorm,
            }
        },
        pass: {
            color: [color],
            depth_stencil: {depth}
        }
    }
}

pub struct Engine {
    window: Window,
    instance: Arc<Instance>,
    physical_device: usize,
    device: Arc<Device>,
    swapchain: Arc<Swapchain>,
    images: Vec<Arc<SwapchainImage>>,
    depth_buffer: Arc<AttachmentImage<format::D16Unorm>>,
    render_pass: Arc<render_pass::CustomRenderPass>,
    framebuffers: Vec<Arc<Framebuffer<render_pass::CustomRenderPass>>>,
    submissions: Vec<Arc<Submission>>,
    queue: Arc<Queue>,
    config: Config,
    inputs: HashMap<String, f32>,
}

impl Engine {
    pub fn new(config: Config) -> Engine {

        let window_builder = create_window(&config.window);

        let instance = {
            let extensions = required_extensions();
            Instance::new(None, &extensions, None).expect("Failed to create Vulkan instance.")
        };

        let ins = &instance.clone();

        let physical = PhysicalDevice::enumerate(&ins)
            .next()
            .expect("No vulkan device is available.");

        let physical_device = physical.index();

        let window = window_builder.build_vk_surface(&instance).unwrap();

        let queue = physical
            .queue_families()
            .find(|q| q.supports_graphics() && window.surface().is_supported(q).unwrap_or(false))
            .expect("Couldn't find a graphical queue family.");

        // Logical Device, Queues
        let (device, mut queues) = {
            let device_ext = DeviceExtensions {
                khr_swapchain: true,
                ..DeviceExtensions::none()
            };

            Device::new(&physical,
                        physical.supported_features(),
                        &device_ext,
                        [(queue, 0.5)].iter().cloned())
                    .expect("failed to create device")
        };

        // Device Queue
        let queue = queues.next().unwrap();

        // Swapchain, Swapchain Images
        let (swapchain, images) =
            create_swapchain(&window, &physical, &device, &queue, None, &config);

        let depth_buffer =
            AttachmentImage::transient(&device, images[0].dimensions(), format::D16Unorm).unwrap();

        // Render Pass
        let render_pass =
            render_pass::CustomRenderPass::new(&device,
                                               &render_pass::Formats {
                                                    // Use the format of the images and one sample.
                                                    color: (images[0].format(), 1),
                                                    depth: (format::D16Unorm, 1),
                                                })
                    .unwrap();

        let framebuffers = images
            .iter()
            .map(|image| {
                let attachments = render_pass::AList {
                    color: &image,
                    depth: &depth_buffer,
                };

                Framebuffer::new(&render_pass,
                                 [image.dimensions()[0], image.dimensions()[1], 1],
                                 attachments)
                        .unwrap()
            })
            .collect::<Vec<_>>();

        // Queue Submissions
        let submissions = Vec::new();

        // Input System
        let mut inputs: HashMap<String, f32> = config
            .input
            .keys()
            .map(|k| (k.clone(), 0.0f32))
            .collect();

        Engine {
            window: window,
            instance,
            physical_device,
            device,
            swapchain,
            images,
            depth_buffer,
            framebuffers,
            render_pass,
            submissions,
            queue,
            config,
            inputs,
        }
    }

    /// Handles input/output events from the window and any input middleware.
    pub fn io(&mut self) -> bool {

        for ev in self.window.window().poll_events() {

            // Axis Map
            for (string_key, axis) in self.config.input.iter() {
                for axis_value in axis {

                    let out = events::string_to_wevent(&axis_value.key, &ev);

                    // Check axis_value.meta for additional checks.

                    // Write to axis map
                    match out {
                        Some(x) => *self.inputs.get_mut(string_key).unwrap() = x,
                        None => (),
                    };

                }
            }

            // Core Events
            match &ev {
                &Event::Resized(w, h) => {
                    self.config.window.resolution = [w, h];
                    let (swapchain, images) =
                        create_swapchain(&self.window,
                                         &PhysicalDevice::from_index(&self.instance,
                                                                     self.physical_device)
                                                  .unwrap(),
                                         &self.device,
                                         &self.queue,
                                         Some(&self.swapchain),
                                         &self.config);
                    self.swapchain = swapchain;
                    self.images = images;
                    self.depth_buffer = AttachmentImage::transient(&self.device,
                                                                   self.images[0].dimensions(),
                                                                   format::D16Unorm)
                            .unwrap();
                    self.framebuffers = self.images
                        .iter()
                        .map(|image| {
                            let attachments = render_pass::AList {
                                color: &image,
                                depth: &self.depth_buffer,
                            };

                            Framebuffer::new(&self.render_pass,
                                             [image.dimensions()[0], image.dimensions()[1], 1],
                                             attachments)
                                    .unwrap()
                        })
                        .collect::<Vec<_>>();
                }
                &Event::Closed => return false,
                _ => (),
            };
        }

        true
    }

    /// Recursively updates application tree.
    pub fn update(&self) {

        //@TODO recursively map each update function from each node.
        //node.children.map(update)

    }

    /// Updates the display surface with a new image.
    pub fn render(&mut self) {
        let command_buffers = self.framebuffers
            .iter()
            .map(|framebuffer| {
                PrimaryCommandBufferBuilder::new(&self.device, self.queue.family())
                    .draw_inline(&self.render_pass,
                                 &framebuffer,
                                 render_pass::ClearValues {
                                     color: [0.2, 0.4, 0.8, 1.0],
                                     depth: 1.0,
                                 })
                    .draw_end()
                    .build()
            })
            .collect::<Vec<_>>();
        let image_num = self.swapchain
            .acquire_next_image(Duration::new(1, 0))
            .unwrap();

        // @TODO build command buffers with threads and submit the changes in main thread (here)
        self.submissions
            .push(submit(&command_buffers[image_num], &self.queue).unwrap());

        self.swapchain.present(&self.queue, image_num).unwrap();
    }
}

/// Sets up and creates a swapchain
fn create_swapchain(window: &Window,
                    physical_device: &PhysicalDevice,
                    device: &Arc<Device>,
                    queue: &Arc<Queue>,
                    old: Option<&Arc<Swapchain>>,
                    config: &Config)
                    -> (Arc<Swapchain>, Vec<Arc<SwapchainImage>>) {
    {
        let caps = window
            .surface()
            .get_capabilities(&physical_device)
            .expect("failed to get surface capabilities");

        let dimensions = if config.window.resolution[0] <= 240 ||
                            config.window.resolution[1] <= 240 {

            let min = caps.min_image_extent;

            let extent = caps.current_extent.unwrap_or([800, 600]);

            if extent[0] < min[0] || extent[1] < min[1] {
                min
            }
            else {
                extent
            }
        } else {
            config.window.resolution
        };

        let present = if config.graphics.vsync &&
                         caps.present_modes.supports(PresentMode::Mailbox) {
            PresentMode::Mailbox
        } else {
            caps.present_modes.iter().next().unwrap()
        };

        let alpha = caps.supported_composite_alpha.iter().next().unwrap();

        let format = caps.supported_formats[0].0;

        Swapchain::new(&device,
                       &window.surface(),
                       caps.min_image_count,
                       format,
                       dimensions,
                       1,
                       &caps.supported_usage_flags,
                       queue,
                       SurfaceTransform::Identity,
                       alpha,
                       present,
                       true,
                       old)
                .expect("failed to create swapchain")
    }
}

/// Creates a window builder with a given window configuration.
fn create_window(config: &WindowConfig) -> WindowBuilder {

    let mut w = get_primary_monitor();
    let dimensions = w.get_dimensions();

    // Set resolution
    let resolution = if config.resolution[0] <= 240 || config.resolution[1] <= 240 {
        [dimensions.0, dimensions.1]
    } else {
        config.resolution.clone()
    };

    // Set fullscreen
    if config.fullscreen {
        for (i, monitor) in get_available_monitors().enumerate() {
            if i == config.display as usize {
                w = monitor;
            }
        }
    }

    // Create manager
    let window_manager = WindowBuilder::new()
        .with_title("CodeVR")
        .with_dimensions(resolution[0], resolution[1]);

    if config.fullscreen {
        return window_manager.with_fullscreen(w);
    }

    window_manager
}
<commit_msg>Started fix<commit_after>mod events;

use winit::{WindowBuilder, get_available_monitors, get_primary_monitor, Event, ElementState};
use vulkano_win::{Window, VkSurfaceBuild, required_extensions};
use vulkano::instance::{Instance, PhysicalDevice};
use vulkano::device::{Queue, Device, DeviceExtensions};
use vulkano::swapchain::{Swapchain, SurfaceTransform, PresentMode};
use vulkano::image::SwapchainImage;
use vulkano::image::attachment::AttachmentImage;
use vulkano::framebuffer::Framebuffer;
use vulkano::command_buffer::{PrimaryCommandBufferBuilder, Submission, submit};
use vulkano::format;

use config::Config;
use config::WindowConfig;

use std::sync::Arc;
use std::time::Duration;
use std::collections::HashMap;

mod render_pass {

    use vulkano::format;
    use vulkano::format::Format;

    single_pass_renderpass!{
    attachments: {
        color: {
            load: Clear,
            store: Store,
            format: Format,
        },
        depth: {
            load: Clear,
            store: DontCare,
            format: format::D16Unorm,
            }
        },
        pass: {
            color: [color],
            depth_stencil: {depth}
        }
    }
}

pub struct Engine {
    window: Window,
    instance: Arc<Instance>,
    physical_device: usize,
    device: Arc<Device>,
    swapchain: Arc<Swapchain>,
    images: Vec<Arc<SwapchainImage>>,
    depth_buffer: Arc<AttachmentImage<format::D16Unorm>>,
    render_pass: Arc<render_pass::CustomRenderPass>,
    framebuffers: Vec<Arc<Framebuffer<render_pass::CustomRenderPass>>>,
    submissions: Vec<Arc<Submission>>,
    queue: Arc<Queue>,
    config: Config,
    inputs: HashMap<String, (f32, u8)>,
}

impl Engine {
    pub fn new(config: Config) -> Engine {

        let window_builder = create_window(&config.window);

        let instance = {
            let extensions = required_extensions();
            Instance::new(None, &extensions, None).expect("Failed to create Vulkan instance.")
        };

        let ins = &instance.clone();

        let physical = PhysicalDevice::enumerate(&ins)
            .next()
            .expect("No vulkan device is available.");

        let physical_device = physical.index();

        let window = window_builder.build_vk_surface(&instance).unwrap();

        let queue = physical
            .queue_families()
            .find(|q| q.supports_graphics() && window.surface().is_supported(q).unwrap_or(false))
            .expect("Couldn't find a graphical queue family.");

        // Logical Device, Queues
        let (device, mut queues) = {
            let device_ext = DeviceExtensions {
                khr_swapchain: true,
                ..DeviceExtensions::none()
            };

            Device::new(&physical,
                        physical.supported_features(),
                        &device_ext,
                        [(queue, 0.5)].iter().cloned())
                    .expect("failed to create device")
        };

        // Device Queue
        let queue = queues.next().unwrap();

        // Swapchain, Swapchain Images
        let (swapchain, images) =
            create_swapchain(&window, &physical, &device, &queue, None, &config);

        let depth_buffer =
            AttachmentImage::transient(&device, images[0].dimensions(), format::D16Unorm).unwrap();

        // Render Pass
        let render_pass =
            render_pass::CustomRenderPass::new(&device,
                                               &render_pass::Formats {
                                                    // Use the format of the images and one sample.
                                                    color: (images[0].format(), 1),
                                                    depth: (format::D16Unorm, 1),
                                                })
                    .unwrap();

        let framebuffers = images
            .iter()
            .map(|image| {
                let attachments = render_pass::AList {
                    color: &image,
                    depth: &depth_buffer,
                };

                Framebuffer::new(&render_pass,
                                 [image.dimensions()[0], image.dimensions()[1], 1],
                                 attachments)
                        .unwrap()
            })
            .collect::<Vec<_>>();

        // Queue Submissions
        let submissions = Vec::new();

        // Input System
        let mut inputs: HashMap<String, (f32, u8)> = config
            .input
            .keys()
            .map(|k| (k.clone(), (0.0, 0)))
            .collect();

        Engine {
            window: window,
            instance,
            physical_device,
            device,
            swapchain,
            images,
            depth_buffer,
            framebuffers,
            render_pass,
            submissions,
            queue,
            config,
            inputs,
        }
    }

    /// Handles input/output events from the window and any input middleware.
    pub fn io(&mut self) -> bool {

        println!("{:?}", self.inputs.clone());

        for ev in self.window.window().poll_events() {

            // Axis Map
            for (string_key, axis) in self.config.input.iter() {
                for axis_value in axis {

                    let out = events::string_to_wevent(&axis_value.key, &ev);

                    // Check axis_value.meta for additional checks.

                    // Write to axis map
                    match out {
                        Some(x) => {
                            let mut hash_ref = self.inputs.get_mut(string_key).unwrap();
                            *hash_ref = (x, 0);

                            // Now we need some logic to make sure we're not masking inputs...
                            },
                        None => (),
                    };

                }
            }

            // Core Events
            match &ev {
                &Event::Resized(w, h) => {
                    self.config.window.resolution = [w, h];
                    let (swapchain, images) =
                        create_swapchain(&self.window,
                                         &PhysicalDevice::from_index(&self.instance,
                                                                     self.physical_device)
                                                  .unwrap(),
                                         &self.device,
                                         &self.queue,
                                         Some(&self.swapchain),
                                         &self.config);
                    self.swapchain = swapchain;
                    self.images = images;
                    self.depth_buffer = AttachmentImage::transient(&self.device,
                                                                   self.images[0].dimensions(),
                                                                   format::D16Unorm)
                            .unwrap();
                    self.framebuffers = self.images
                        .iter()
                        .map(|image| {
                            let attachments = render_pass::AList {
                                color: &image,
                                depth: &self.depth_buffer,
                            };

                            Framebuffer::new(&self.render_pass,
                                             [image.dimensions()[0], image.dimensions()[1], 1],
                                             attachments)
                                    .unwrap()
                        })
                        .collect::<Vec<_>>();
                }
                &Event::Closed => return false,
                _ => (),
            };
        }

        true
    }

    /// Recursively updates application tree.
    pub fn update(&self) {

        //@TODO recursively map each update function from each node.
        //node.children.map(update)

    }

    /// Updates the display surface with a new image.
    pub fn render(&mut self) {
        let command_buffers = self.framebuffers
            .iter()
            .map(|framebuffer| {
                PrimaryCommandBufferBuilder::new(&self.device, self.queue.family())
                    .draw_inline(&self.render_pass,
                                 &framebuffer,
                                 render_pass::ClearValues {
                                     color: [0.2, 0.4, 0.8, 1.0],
                                     depth: 1.0,
                                 })
                    .draw_end()
                    .build()
            })
            .collect::<Vec<_>>();
        let image_num = self.swapchain
            .acquire_next_image(Duration::new(1, 0))
            .unwrap();

        // @TODO build command buffers with threads and submit the changes in main thread (here)
        self.submissions
            .push(submit(&command_buffers[image_num], &self.queue).unwrap());

        self.swapchain.present(&self.queue, image_num).unwrap();
    }
}

/// Sets up and creates a swapchain
fn create_swapchain(window: &Window,
                    physical_device: &PhysicalDevice,
                    device: &Arc<Device>,
                    queue: &Arc<Queue>,
                    old: Option<&Arc<Swapchain>>,
                    config: &Config)
                    -> (Arc<Swapchain>, Vec<Arc<SwapchainImage>>) {
    {
        let caps = window
            .surface()
            .get_capabilities(&physical_device)
            .expect("failed to get surface capabilities");

        let dimensions = if config.window.resolution[0] <= 240 ||
                            config.window.resolution[1] <= 240 {

            let min = caps.min_image_extent;

            let extent = caps.current_extent.unwrap_or([800, 600]);

            if extent[0] < min[0] || extent[1] < min[1] {
                min
            }
            else {
                extent
            }
        } else {
            config.window.resolution
        };

        let present = if config.graphics.vsync &&
                         caps.present_modes.supports(PresentMode::Mailbox) {
            PresentMode::Mailbox
        } else {
            caps.present_modes.iter().next().unwrap()
        };

        let alpha = caps.supported_composite_alpha.iter().next().unwrap();

        let format = caps.supported_formats[0].0;

        Swapchain::new(&device,
                       &window.surface(),
                       caps.min_image_count,
                       format,
                       dimensions,
                       1,
                       &caps.supported_usage_flags,
                       queue,
                       SurfaceTransform::Identity,
                       alpha,
                       present,
                       true,
                       old)
                .expect("failed to create swapchain")
    }
}

/// Creates a window builder with a given window configuration.
fn create_window(config: &WindowConfig) -> WindowBuilder {

    let mut w = get_primary_monitor();
    let dimensions = w.get_dimensions();

    // Set resolution
    let resolution = if config.resolution[0] <= 240 || config.resolution[1] <= 240 {
        [dimensions.0, dimensions.1]
    } else {
        config.resolution.clone()
    };

    // Set fullscreen
    if config.fullscreen {
        for (i, monitor) in get_available_monitors().enumerate() {
            if i == config.display as usize {
                w = monitor;
            }
        }
    }

    // Create manager
    let window_manager = WindowBuilder::new()
        .with_title("CodeVR")
        .with_dimensions(resolution[0], resolution[1]);

    if config.fullscreen {
        return window_manager.with_fullscreen(w);
    }

    window_manager
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;

extern crate abstract_ns;
extern crate futures;
extern crate getopts;
extern crate ns_dns_tokio;
extern crate rand;
extern crate tokio_core;
extern crate tokio_io;

// the device under test
extern crate net_interceptor;

use abstract_ns::Resolver;
use futures::{Future, Map, Stream, Async, Poll};
use futures::future;
use ns_dns_tokio::DnsResolver;
use tokio_core::net::{TcpStream, TcpListener};
use tokio_core::reactor::{Core, Handle};
use tokio_io::{AsyncRead, io};

use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::fmt::{Debug, Display};
use std::time::{Duration};

use std::io::Error;

use net_interceptor::logging::unittestlogger;

fn add_ten<F>(future: F) -> Map<F, fn(i32) -> i32>
    where F: Future<Item=i32>,
{
    error!("shutdown");
    fn add(a: i32) -> i32 { a + 10 }
    future.map(add)
}

pub struct AbortCondition {
    condition: bool,
    reactor_handle: futures::task::Task,
}

#[derive(Copy, Clone)]
pub struct Releaser {
    block: *const AbortCondition
}

impl AbortCondition {

    fn set(&mut self) {
        debug!("abort condition met");
        self.condition = true;
        self.reactor_handle.unpark();
    }

}

impl Future for AbortCondition {

    type Item = ();
    type Error = Error;

    fn poll(&mut self) -> Poll<(), Error> {
        if  self.condition {
            Ok(Async::Ready(()))
        } else {
            self.reactor_handle = futures::task::park();
            Ok(Async::NotReady)
        }
    }

}

/*impl Copy for Abort {
    fn clone(&self) -> Abort {
        *self
    }
}
*/
// Here we'll express the handling of `client` and return it as a future
// to be spawned onto the event loop.
/*fn process(client: TcpStream) -> Box<Future<Item = (i64), Error = error>> {
    warn!("received");
    Box::new(futures::future::ok(0));
}
*/
#[test]
fn test_infra_connect() {
    unittestlogger();

    // this is the main event loop, powered by tokio core
    let mut core = Core::new().unwrap();
    let handle = core.handle();

    let bind_socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 12000);
    let listener = TcpListener::bind(&bind_socket_addr, &handle)
        .expect(&format!("Unable to bind to {}", &bind_socket_addr));
    warn!("Listening on {}", listener.local_addr().unwrap());

    let addr = listener.local_addr().unwrap();
    assert!(addr.port() == 12000, "should accept {}");

    let mut abort = AbortCondition { condition: false, reactor_handle: futures::task::park() };

    let connections = listener.incoming();
    let server = connections.for_each(move |(socket, _peer_addr)| {
        warn!("connection received");
        /*let (writer, reader) = socket.framed(LineCodec).split();
        let service = s.new_service()?;

        let responses = reader.and_then(move |req| service.call(req));
        let server = writer.send_all(responses)
            .then(|_| Ok(()));
        handle.spawn(server);
        */
        abort.set();
        Ok(())
    });

    fn abortCondition() -> futures::Poll<(), std::io::Error> {
        warn!("poll");
        if true {
            Ok(Async::Ready(()))
        }
        else {
            futures::task::park().unpark();
            Ok(Async::NotReady)
        }
    }

    //let abort = futures::future::poll_fn(abortCondition);

    //let timeout = tokio_core::reactor::Timeout::new_at(std::time::Duration::seconds(1));
    // run server, stop on panic
//    core.run(server.select2(timeout)).unwrap();
//    core.run(abort).unwrap();
    // core.spawn(server);
    core.run(server.select(abort)).unwrap();
}
<commit_msg>broken!<commit_after>#[macro_use]
extern crate log;

extern crate abstract_ns;
extern crate futures;
extern crate getopts;
extern crate ns_dns_tokio;
extern crate rand;
extern crate tokio_core;
extern crate tokio_io;

// the device under test
extern crate net_interceptor;

use abstract_ns::Resolver;
use futures::{Future, Map, Stream, Async, Poll};
use futures::future;
use futures::sync::oneshot;
use futures::sync::oneshot::{Sender};
use ns_dns_tokio::DnsResolver;
use tokio_core::net::{TcpStream, TcpListener};
use tokio_core::reactor::{Core, Handle};
use tokio_io::{AsyncRead, io};

use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::fmt::{Debug, Display};
use std::time::{Duration};

use std::io::Error;

use net_interceptor::logging::unittestlogger;

fn add_ten<F>(future: F) -> Map<F, fn(i32) -> i32>
    where F: Future<Item=i32>,
{
    error!("shutdown");
    fn add(a: i32) -> i32 { a + 10 }
    future.map(add)
}

pub struct AbortCondition {
    condition: bool,
    reactor_handle: futures::task::Task,
}

#[derive(Copy, Clone)]
pub struct Releaser {
    block: *const AbortCondition
}

impl AbortCondition {

    fn set(&mut self) {
        debug!("abort condition met");
        self.condition = true;
        self.reactor_handle.unpark();
    }

}

impl Future for AbortCondition {

    type Item = ();
    type Error = Error;

    fn poll(&mut self) -> Poll<(), Error> {
        if  self.condition {
            Ok(Async::Ready(()))
        } else {
            self.reactor_handle = futures::task::park();
            Ok(Async::NotReady)
        }
    }

}
/*
#[derive(Clone)]
pub struct Stopper {
    signal: futures::sync::oneshot::Sender<i64>,
}

impl Stopper {

    pub fn stop(&mut self) {
        self.signal = 1;
    }

}
*/

/*impl Copy for Abort {
    fn clone(&self) -> Abort {
        *self
    }
}
*/
// Here we'll express the handling of `client` and return it as a future
// to be spawned onto the event loop.
/*fn process(client: TcpStream) -> Box<Future<Item = (i64), Error = error>> {
    warn!("received");
    Box::new(futures::future::ok(0));
}
*/
#[test]
fn test_infra_connect() {
    unittestlogger();

    // this is the main event loop, powered by tokio core
    let mut core = Core::new().unwrap();
    let handle = core.handle();

    let bind_socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 12000);
    let listener = TcpListener::bind(&bind_socket_addr, &handle)
        .expect(&format!("Unable to bind to {}", &bind_socket_addr));
    warn!("Listening on {}", listener.local_addr().unwrap());

    let addr = listener.local_addr().unwrap();
    assert!(addr.port() == 12000, "should accept {}");

    let mut abort = AbortCondition { condition: false, reactor_handle: futures::task::park() };

    let connections = listener.incoming();
    let (tx, stop) = oneshot::channel();

//    let mut stopper = Stopper { signal: tx };
    let server = connections.for_each(move |(socket, _peer_addr)| {
        warn!("connection received");
        /*let (writer, reader) = socket.framed(LineCodec).split();
        let service = s.new_service()?;

        let responses = reader.and_then(move |req| service.call(req));
        let server = writer.send_all(responses)
            .then(|_| Ok(()));
        handle.spawn(server);
        */
//        tx.send(1);
        tx.send(1).unwrap();
        abort.set();
        Ok(())
    });

    core.handle().spawn(server.map_err(|_|()));

    fn abortCondition() -> futures::Poll<(), std::io::Error> {
        warn!("poll");
        if true {
            Ok(Async::Ready(()))
        }
        else {
            futures::task::park().unpark();
            Ok(Async::NotReady)
        }
    }

    //let abort = futures::future::poll_fn(abortCondition);

    //let timeout = tokio_core::reactor::Timeout::new_at(std::time::Duration::seconds(1));
    // run server, stop on panic
//    core.run(server.select2(timeout)).unwrap();
//    core.run(abort).unwrap();
    // core.spawn(server);
//    core.run(server.select(abort)).unwrap();
    core.run(stop).unwrap();
}
<|endoftext|>


<commit_before>use axum::{
    body::{boxed, Body, BoxBody},
    response::{Html, IntoResponse},
    routing::{get, get_service},
    Router,
};
use axum_extra::middleware::{self, Next};
use headers::{ContentLength, HeaderMapExt as _};
use http::{header, HeaderValue, Request, Response, StatusCode};
use std::fmt::Write as _;
use structopt::StructOpt;
use tower::ServiceBuilder;
use tower_http::{services::ServeDir, trace::TraceLayer};
use webbundle::{Bundle, Version};

#[derive(StructOpt, Debug)]
struct Cli {
    // TODO: Support https.
    // #[structopt(long = "https")]
    // https: bool,
    #[structopt(short = "p", long = "port", default_value = "8000")]
    port: u16,
    #[structopt(long = "bind-all")]
    /// Bind all interfaces (default: only localhost - "127.0.0.1"),
    bind_all: bool,
}

#[tokio::main]
async fn main() {
    // Set the RUST_LOG, if it hasn't been explicitly defined
    if std::env::var_os("RUST_LOG").is_none() {
        std::env::set_var("RUST_LOG", "my_http_server=debug,tower_http=debug")
    }
    tracing_subscriber::fmt::init();
    let args = Cli::from_args();

    let app = Router::new()
        .nest("/wbn", get(webbundle_serve))
        .fallback(
            get_service(ServeDir::new("."))
                .handle_error(|error: std::io::Error| async move {
                    (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        format!("Unhandled internal error: {}", error),
                    )
                })
                .layer(middleware::from_fn(serve_dir_extra)),
        )
        .layer(ServiceBuilder::new().layer(TraceLayer::new_for_http()));

    let addr = std::net::SocketAddr::from((
        if args.bind_all {
            [0, 0, 0, 0]
        } else {
            [127, 0, 0, 1]
        },
        args.port,
    ));
    println!("Listening on http://{}/", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn webbundle_serve(req: Request) -> Result<Response<BoxBody>, (StatusCode, String)> {
    webbundle_serve_internal(req).await.map_err(|err| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Unhandled internal error {}", err),
        )
    })
}

async fn webbundle_serve_internal(req: Request) -> anyhow::Result<Response<BoxBody>> {
    let path = req.uri().path();
    let mut full_path = std::path::PathBuf::from(".");
    for seg in path.trim_start_matches('/').split('/') {
        anyhow::ensure!(
            !seg.starts_with("..") && !seg.contains('\\'),
            "Invalid request"
        );
        full_path.push(seg);
    }
    anyhow::ensure!(is_dir(&full_path).await, "Not found");

    let bundle = Bundle::builder()
        .version(Version::VersionB2)
        .exchanges_from_dir(full_path)
        .await?
        .build()?;

    let bytes = bundle.encode()?;
    let content_length = ContentLength(bytes.len() as u64);
    let mut response = Response::new(boxed(Body::from(bytes)));
    response.headers_mut().typed_insert(content_length);
    set_response_webbundle_headers(&mut response);
    Ok(response)
}

fn set_response_webbundle_headers(response: &mut Response<BoxBody>) {
    response.headers_mut().insert(
        header::CONTENT_TYPE,
        HeaderValue::from_static("application/webbundle"),
    );
    response.headers_mut().insert(
        header::X_CONTENT_TYPE_OPTIONS,
        HeaderValue::from_static("nosniff"),
    );
}

async fn is_dir(full_path: &std::path::Path) -> bool {
    tokio::fs::metadata(full_path)
        .await
        .map(|m| m.is_dir())
        .unwrap_or(false)
}

async fn serve_dir_extra(
    req: Request,
    next: Next,
) -> Result<Response<BoxBody>, (StatusCode, String)> {
    serve_dir_extra_internal(req, next).await.map_err(|err| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Unhandled internal error {}", err),
        )
    })
}

async fn serve_dir_extra_internal(
    req: Request,
    next: Next,
) -> anyhow::Result<Response<BoxBody>> {
    // Directory listing.
    // Ref: https://docs.rs/tower-http/0.1.0/src/tower_http/services/fs/serve_dir.rs.html
    let path = req.uri().path();
    let mut full_path = std::path::PathBuf::from(".");
    for seg in path.trim_start_matches('/').split('/') {
        anyhow::ensure!(!seg.starts_with("..") && !seg.contains('\\'));
        full_path.push(seg);
    }
    if is_dir(&full_path).await {
        let html = directory_list_files(full_path, path).await?;
        return Ok(Html(html).into_response());
    }

    if req.uri().path().ends_with(".wbn") {
        let mut res = next.run(req).await;
        set_response_webbundle_headers(&mut res);
        return Ok(res);
    }

    // default.
    Ok(next.run(req).await)
}

async fn directory_list_files(
    path: impl AsRef<std::path::Path>,
    display_name: &str,
) -> anyhow::Result<String> {
    let path = path.as_ref();

    let mut contents = String::new();
    // ReadDir is Stream
    let mut read_dir = tokio::fs::read_dir(path).await?;
    let mut files = Vec::new();
    while let Some(file) = read_dir.next_entry().await? {
        files.push(file.path());
    }
    files.sort();
    for p in files {
        let link_name = format!(
            "{}{}",
            p.file_name().unwrap().to_str().unwrap(),
            if is_dir(&p).await { "/" } else { "" }
        );
        write!(
            contents,
            "
{link}
",
            link = link_name
        )?;
    }

    let inline_style = r#"
body {
  box-sizing: border-box;
  min-width: 200px;
  max-width: 980px;
  margin: 0 auto;
  padding: 45px;
}
"#;

    Ok(format!(
        r#"

<meta charset="utf-8"/>
<title>{title}</title>
<link rel=stylesheet href="https://cdn.jsdelivr.net/npm/github-markdown-css">

webbundle-server: Directory listing for {display_name}

..

{contents}

"#,
        title = display_name,
        inline_style = inline_style,
        display_name = display_name,
        contents = contents
    ))
}
<commit_msg>Return 404 if a directory to serve WebBundle is not found<commit_after>use axum::{
    body::{boxed, Body, BoxBody},
    response::{Html, IntoResponse},
    routing::{get, get_service},
    Router,
};
use axum_extra::middleware::{self, Next};
use headers::{ContentLength, HeaderMapExt as _};
use http::{header, HeaderValue, Request, Response, StatusCode};
use std::fmt::Write as _;
use structopt::StructOpt;
use tower::ServiceBuilder;
use tower_http::{services::ServeDir, trace::TraceLayer};
use webbundle::{Bundle, Version};

#[derive(StructOpt, Debug)]
struct Cli {
    // TODO: Support https.
    // #[structopt(long = "https")]
    // https: bool,
    #[structopt(short = "p", long = "port", default_value = "8000")]
    port: u16,
    #[structopt(long = "bind-all")]
    /// Bind all interfaces (default: only localhost - "127.0.0.1"),
    bind_all: bool,
}

#[tokio::main]
async fn main() {
    // Set the RUST_LOG, if it hasn't been explicitly defined
    if std::env::var_os("RUST_LOG").is_none() {
        std::env::set_var("RUST_LOG", "my_http_server=debug,tower_http=debug")
    }
    tracing_subscriber::fmt::init();
    let args = Cli::from_args();

    let app = Router::new()
        .nest("/wbn", get(webbundle_serve))
        .fallback(
            get_service(ServeDir::new("."))
                .handle_error(|error: std::io::Error| async move {
                    (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        format!("Unhandled internal error: {}", error),
                    )
                })
                .layer(middleware::from_fn(serve_dir_extra)),
        )
        .layer(ServiceBuilder::new().layer(TraceLayer::new_for_http()));

    let addr = std::net::SocketAddr::from((
        if args.bind_all {
            [0, 0, 0, 0]
        } else {
            [127, 0, 0, 1]
        },
        args.port,
    ));
    println!("Listening on http://{}/", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn webbundle_serve(req: Request) -> Result<Response<BoxBody>, (StatusCode, String)> {
    match webbundle_serve_internal(req).await {
        Ok(WebBundleServeResponse::Body(body)) => Ok(body),
        Ok(WebBundleServeResponse::NotFound) => Err((StatusCode::NOT_FOUND, "".to_string())),
        Err(err) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Unhandled internal error {}", err),
        )),
    }
}

enum WebBundleServeResponse {
    Body(Response<BoxBody>),
    NotFound,
}

async fn webbundle_serve_internal(req: Request) -> anyhow::Result<WebBundleServeResponse> {
    let path = req.uri().path();
    let mut full_path = std::path::PathBuf::from(".");
    for seg in path.trim_start_matches('/').split('/') {
        anyhow::ensure!(
            !seg.starts_with("..") && !seg.contains('\\'),
            "Invalid request"
        );
        full_path.push(seg);
    }
    if !is_dir(&full_path).await {
        return Ok(WebBundleServeResponse::NotFound);
    }

    let bundle = Bundle::builder()
        .version(Version::VersionB2)
        .exchanges_from_dir(full_path)
        .await?
        .build()?;

    let bytes = bundle.encode()?;
    let content_length = ContentLength(bytes.len() as u64);
    let mut response = Response::new(boxed(Body::from(bytes)));
    response.headers_mut().typed_insert(content_length);
    set_response_webbundle_headers(&mut response);
    Ok(WebBundleServeResponse::Body(response))
}

fn set_response_webbundle_headers(response: &mut Response<BoxBody>) {
    response.headers_mut().insert(
        header::CONTENT_TYPE,
        HeaderValue::from_static("application/webbundle"),
    );
    response.headers_mut().insert(
        header::X_CONTENT_TYPE_OPTIONS,
        HeaderValue::from_static("nosniff"),
    );
}

async fn is_dir(full_path: &std::path::Path) -> bool {
    tokio::fs::metadata(full_path)
        .await
        .map(|m| m.is_dir())
        .unwrap_or(false)
}

async fn serve_dir_extra(
    req: Request,
    next: Next,
) -> Result<Response<BoxBody>, (StatusCode, String)> {
    serve_dir_extra_internal(req, next).await.map_err(|err| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Unhandled internal error {}", err),
        )
    })
}

async fn serve_dir_extra_internal(
    req: Request,
    next: Next,
) -> anyhow::Result<Response<BoxBody>> {
    // Directory listing.
    // Ref: https://docs.rs/tower-http/0.1.0/src/tower_http/services/fs/serve_dir.rs.html
    let path = req.uri().path();
    let mut full_path = std::path::PathBuf::from(".");
    for seg in path.trim_start_matches('/').split('/') {
        anyhow::ensure!(!seg.starts_with("..") && !seg.contains('\\'));
        full_path.push(seg);
    }
    if is_dir(&full_path).await {
        let html = directory_list_files(full_path, path).await?;
        return Ok(Html(html).into_response());
    }

    if req.uri().path().ends_with(".wbn") {
        let mut res = next.run(req).await;
        set_response_webbundle_headers(&mut res);
        return Ok(res);
    }

    // default.
    Ok(next.run(req).await)
}

async fn directory_list_files(
    path: impl AsRef<std::path::Path>,
    display_name: &str,
) -> anyhow::Result<String> {
    let path = path.as_ref();

    let mut contents = String::new();
    // ReadDir is Stream
    let mut read_dir = tokio::fs::read_dir(path).await?;
    let mut files = Vec::new();
    while let Some(file) = read_dir.next_entry().await? {
        files.push(file.path());
    }
    files.sort();
    for p in files {
        let link_name = format!(
            "{}{}",
            p.file_name().unwrap().to_str().unwrap(),
            if is_dir(&p).await { "/" } else { "" }
        );
        write!(
            contents,
            "
{link}
",
            link = link_name
        )?;
    }

    let inline_style = r#"
body {
  box-sizing: border-box;
  min-width: 200px;
  max-width: 980px;
  margin: 0 auto;
  padding: 45px;
}
"#;

    Ok(format!(
        r#"

<meta charset="utf-8"/>
<title>{title}</title>
<link rel=stylesheet href="https://cdn.jsdelivr.net/npm/github-markdown-css">

webbundle-server: Directory listing for {display_name}

..

{contents}

"#,
        title = display_name,
        inline_style = inline_style,
        display_name = display_name,
        contents = contents
    ))
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#![allow(unsafe_code)]

use app_units::Au;
use bindings::{RawGeckoDocument, RawGeckoElement, RawGeckoNode};
use bindings::{RawServoStyleSet, RawServoStyleSheet, ServoComputedValues, ServoNodeData};
use bindings::{nsIAtom};
use data::PerDocumentStyleData;
use env_logger;
use euclid::Size2D;
use gecko_style_structs::SheetParsingMode;
use properties::GeckoComputedValues;
use selector_impl::{GeckoSelectorImpl, PseudoElement, SharedStyleContext, Stylesheet};
use std::marker::PhantomData;
use std::mem::{forget, transmute};
use std::ptr;
use std::slice;
use std::str::from_utf8_unchecked;
use std::sync::{Arc, Mutex};
use style::context::{ReflowGoal};
use style::dom::{TDocument, TElement, TNode};
use style::error_reporting::StdoutErrorReporter;
use style::parallel;
use style::properties::ComputedValues;
use style::selector_impl::{SelectorImplExt, PseudoElementCascadeType};
use style::stylesheets::Origin;
use traversal::RecalcStyleOnly;
use url::Url;
use util::arc_ptr_eq;
use wrapper::{GeckoDocument, GeckoElement, GeckoNode, NonOpaqueStyleData};

// TODO: This is ugly and should go away once we get an atom back-end.
pub fn pseudo_element_from_atom(pseudo: *mut nsIAtom,
                                in_ua_stylesheet: bool) -> Result<PseudoElement, String> {
    use bindings::Gecko_GetAtomAsUTF16;
    use selectors::parser::{ParserContext, SelectorImpl};

    let pseudo_string = unsafe {
        let mut length = 0;
        let mut buff = Gecko_GetAtomAsUTF16(pseudo, &mut length);

        // Handle the annoying preceding colon in front of everything in nsCSSAnonBoxList.h.
        debug_assert!(length >= 2 && *buff == ':' as u16 && *buff.offset(1) != ':' as u16);
        buff = buff.offset(1);
        length -= 1;

        String::from_utf16(slice::from_raw_parts(buff, length as usize)).unwrap()
    };

    let mut context = ParserContext::new();
    context.in_user_agent_stylesheet = in_ua_stylesheet;
    GeckoSelectorImpl::parse_pseudo_element(&context, &pseudo_string).map_err(|_| pseudo_string)
}

/*
 * For Gecko->Servo function calls, we need to redeclare the same signature that was declared in
 * the C header in Gecko. In order to catch accidental mismatches, we run rust-bindgen against
 * those signatures as well, giving us a second declaration of all the Servo_* functions in this
 * crate. If there's a mismatch, LLVM will assert and abort, which is a rather awful thing to
 * depend on but good enough for our purposes.
 */

#[no_mangle]
pub extern "C" fn Servo_Initialize() -> () {
    // Enable standard Rust logging.
    //
    // See https://doc.rust-lang.org/log/env_logger/index.html for instructions.
    env_logger::init().unwrap();
}

#[no_mangle]
pub extern "C" fn Servo_RestyleDocument(doc: *mut RawGeckoDocument, raw_data: *mut RawServoStyleSet) -> () {
    let document = unsafe { GeckoDocument::from_raw(doc) };
    let node = match document.root_node() {
        Some(x) => x,
        None => return,
    };
    let data = unsafe { &mut *(raw_data as *mut PerDocumentStyleData) };

    // Force the creation of our lazily-constructed initial computed values on
    // the main thread, since it's not safe to call elsewhere.
    //
    // FIXME(bholley): this should move into Servo_Initialize as soon as we get
    // rid of the HackilyFindSomeDeviceContext stuff that happens during
    // initial_values computation, since that stuff needs to be called further
    // along in startup than the sensible place to call Servo_Initialize.
    GeckoComputedValues::initial_values();

    let _needs_dirtying = Arc::get_mut(&mut data.stylist).unwrap()
                              .update(&data.stylesheets, data.stylesheets_changed);
    data.stylesheets_changed = false;

    let shared_style_context = SharedStyleContext {
        viewport_size: Size2D::new(Au(0), Au(0)),
        screen_size_changed: false,
        generation: 0,
        goal: ReflowGoal::ForScriptQuery,
        stylist: data.stylist.clone(),
        new_animations_sender: Mutex::new(data.new_animations_sender.clone()),
        running_animations: data.running_animations.clone(),
        expired_animations: data.expired_animations.clone(),
        error_reporter: Box::new(StdoutErrorReporter),
    };

    if node.is_dirty() || node.has_dirty_descendants() {
        parallel::traverse_dom::<GeckoNode, RecalcStyleOnly>(node, &shared_style_context, &mut data.work_queue);
    }
}

#[no_mangle]
pub extern "C" fn Servo_DropNodeData(data: *mut ServoNodeData) -> () {
    unsafe {
        let _ = Box::<NonOpaqueStyleData>::from_raw(data as *mut NonOpaqueStyleData);
    }
}

#[no_mangle]
pub extern "C" fn Servo_StylesheetFromUTF8Bytes(bytes: *const u8,
                                                length: u32,
                                                mode: SheetParsingMode) -> *mut RawServoStyleSheet {

    let input = unsafe { from_utf8_unchecked(slice::from_raw_parts(bytes, length as usize)) };

    let origin = match mode {
        SheetParsingMode::eAuthorSheetFeatures => Origin::Author,
        SheetParsingMode::eUserSheetFeatures => Origin::User,
        SheetParsingMode::eAgentSheetFeatures => Origin::UserAgent,
    };

    // FIXME(heycam): Pass in the real base URL.
    let url = Url::parse("about:none").unwrap();
    let sheet = Arc::new(Stylesheet::from_str(input, url, origin, Box::new(StdoutErrorReporter)));
    unsafe {
        transmute(sheet)
    }
}

pub struct ArcHelpers<GeckoType, ServoType> {
    phantom1: PhantomData<GeckoType>,
    phantom2: PhantomData<ServoType>,
}

impl<GeckoType, ServoType> ArcHelpers<GeckoType, ServoType> {
    pub fn with<F, Output>(raw: *mut GeckoType, cb: F) -> Output
                           where F: FnOnce(&Arc<ServoType>) -> Output {
        debug_assert!(!raw.is_null());

        let owned = unsafe { Self::into(raw) };
        let result = cb(&owned);
        forget(owned);
        result
    }

    pub fn maybe_with<F, Output>(maybe_raw: *mut GeckoType, cb: F) -> Output
                                 where F: FnOnce(Option<&Arc<ServoType>>) -> Output {
        let owned = if maybe_raw.is_null() {
            None
        } else {
            Some(unsafe { Self::into(maybe_raw) })
        };

        let result = cb(owned.as_ref());
        forget(owned);

        result
    }

    pub unsafe fn into(ptr: *mut GeckoType) -> Arc<ServoType> {
        transmute(ptr)
    }

    pub fn from(owned: Arc<ServoType>) -> *mut GeckoType {
        unsafe { transmute(owned) }
    }

    pub unsafe fn addref(ptr: *mut GeckoType) {
        Self::with(ptr, |arc| forget(arc.clone()));
    }

    pub unsafe fn release(ptr: *mut GeckoType) {
        let _ = Self::into(ptr);
    }
}

#[no_mangle]
pub extern "C" fn Servo_AppendStyleSheet(raw_sheet: *mut RawServoStyleSheet,
                                         raw_data: *mut RawServoStyleSet) {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);
    Helpers::with(raw_sheet, |sheet| {
        data.stylesheets.retain(|x| !arc_ptr_eq(x, sheet));
        data.stylesheets.push(sheet.clone());
        data.stylesheets_changed = true;
    });
}

#[no_mangle]
pub extern "C" fn Servo_PrependStyleSheet(raw_sheet: *mut RawServoStyleSheet,
                                          raw_data: *mut RawServoStyleSet) {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);
    Helpers::with(raw_sheet, |sheet| {
        data.stylesheets.retain(|x| !arc_ptr_eq(x, sheet));
        data.stylesheets.insert(0, sheet.clone());
        data.stylesheets_changed = true;
    })
}

#[no_mangle]
pub extern "C" fn Servo_InsertStyleSheetBefore(raw_sheet: *mut RawServoStyleSheet,
                                               raw_reference: *mut RawServoStyleSheet,
                                               raw_data: *mut RawServoStyleSet) {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);
    Helpers::with(raw_sheet, |sheet| {
        Helpers::with(raw_reference, |reference| {
            data.stylesheets.retain(|x| !arc_ptr_eq(x, sheet));
            let index = data.stylesheets.iter().position(|x| arc_ptr_eq(x, reference)).unwrap();
            data.stylesheets.insert(index, sheet.clone());
            data.stylesheets_changed = true;
        })
    })
}

#[no_mangle]
pub extern "C" fn Servo_RemoveStyleSheet(raw_sheet: *mut RawServoStyleSheet,
                                         raw_data: *mut RawServoStyleSet) {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);
    Helpers::with(raw_sheet, |sheet| {
        data.stylesheets.retain(|x| !arc_ptr_eq(x, sheet));
        data.stylesheets_changed = true;
    });
}

#[no_mangle]
pub extern "C" fn Servo_StyleSheetHasRules(raw_sheet: *mut RawServoStyleSheet) -> bool {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    Helpers::with(raw_sheet, |sheet| !sheet.rules.is_empty())
}

#[no_mangle]
pub extern "C" fn Servo_AddRefStyleSheet(sheet: *mut RawServoStyleSheet) -> () {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    unsafe { Helpers::addref(sheet) };
}

#[no_mangle]
pub extern "C" fn Servo_ReleaseStyleSheet(sheet: *mut RawServoStyleSheet) -> () {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    unsafe { Helpers::release(sheet) };
}

#[no_mangle]
pub extern "C" fn Servo_GetComputedValues(node: *mut RawGeckoNode)
     -> *mut ServoComputedValues {
    let node = unsafe { GeckoNode::from_raw(node) };
    let arc_cv = node.borrow_data().map(|data| data.style.clone());
    arc_cv.map_or(ptr::null_mut(), |arc| unsafe { transmute(arc) })
}

#[no_mangle]
pub extern "C" fn Servo_GetComputedValuesForAnonymousBox(parent_style_or_null: *mut ServoComputedValues,
                                                         pseudo_tag: *mut nsIAtom,
                                                         raw_data: *mut RawServoStyleSet)
     -> *mut ServoComputedValues {
    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);

    let pseudo = match pseudo_element_from_atom(pseudo_tag, /* ua_stylesheet = */ true) {
        Ok(pseudo) => pseudo,
        Err(pseudo) => {
            warn!("stylo: Unable to parse anonymous-box pseudo-element: {}", pseudo);
            return ptr::null_mut();
        }
    };

    type Helpers = ArcHelpers<ServoComputedValues, GeckoComputedValues>;

    Helpers::maybe_with(parent_style_or_null, |maybe_parent| {
        let new_computed = data.stylist.precomputed_values_for_pseudo(&pseudo, maybe_parent);
        new_computed.map_or(ptr::null_mut(), |c| Helpers::from(c))
    })
}

#[no_mangle]
pub extern "C" fn Servo_GetComputedValuesForPseudoElement(parent_style: *mut ServoComputedValues,
                                                          match_element: *mut RawGeckoElement,
                                                          pseudo_tag: *mut nsIAtom,
                                                          raw_data: *mut RawServoStyleSet,
                                                          is_probe: bool)
     -> *mut ServoComputedValues {
    debug_assert!(!match_element.is_null());

    let parent_or_null = || {
        if is_probe {
            ptr::null_mut()
        } else {
            Servo_AddRefComputedValues(parent_style);
            parent_style
        }
    };

    let pseudo = match pseudo_element_from_atom(pseudo_tag, /* ua_stylesheet = */ true) {
        Ok(pseudo) => pseudo,
        Err(pseudo) => {
            warn!("stylo: Unable to parse anonymous-box pseudo-element: {}", pseudo);
            return parent_or_null();
        }
    };

    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);

    let element = unsafe { GeckoElement::from_raw(match_element) };

    type Helpers = ArcHelpers<ServoComputedValues, GeckoComputedValues>;

    match GeckoSelectorImpl::pseudo_element_cascade_type(&pseudo) {
        PseudoElementCascadeType::Eager => {
            let node = element.as_node();
            let maybe_computed = node.borrow_data()
                                     .and_then(|data| {
                                         data.per_pseudo.get(&pseudo).map(|c| c.clone())
                                     });
            maybe_computed.map_or_else(parent_or_null, Helpers::from)
        }
        PseudoElementCascadeType::Lazy => {
            Helpers::with(parent_style, |parent| {
                data.stylist
                    .lazily_compute_pseudo_element_style(&element, &pseudo, parent)
                    .map_or_else(parent_or_null, Helpers::from)
            })
        }
        PseudoElementCascadeType::Precomputed => {
            unreachable!("Anonymous pseudo found in \
                         Servo_GetComputedValuesForPseudoElement");
        }
    }
}

#[no_mangle]
pub extern "C" fn Servo_InheritComputedValues(parent_style: *mut ServoComputedValues)
     -> *mut ServoComputedValues {
    type Helpers = ArcHelpers<ServoComputedValues, GeckoComputedValues>;
    Helpers::with(parent_style, |parent| {
        let style = GeckoComputedValues::inherit_from(parent);
        Helpers::from(style)
    })
}

#[no_mangle]
pub extern "C" fn Servo_AddRefComputedValues(ptr: *mut ServoComputedValues) -> () {
    type Helpers = ArcHelpers<ServoComputedValues, GeckoComputedValues>;
    unsafe { Helpers::addref(ptr) };
}

#[no_mangle]
pub extern "C" fn Servo_ReleaseComputedValues(ptr: *mut ServoComputedValues) -> () {
    type Helpers = ArcHelpers<ServoComputedValues, GeckoComputedValues>;
    unsafe { Helpers::release(ptr) };
}

#[no_mangle]
pub extern "C" fn Servo_InitStyleSet() -> *mut RawServoStyleSet {
    let data = Box::new(PerDocumentStyleData::new());
    Box::into_raw(data) as *mut RawServoStyleSet
}

#[no_mangle]
pub extern "C" fn Servo_DropStyleSet(data: *mut RawServoStyleSet) -> () {
    unsafe {
        let _ = Box::<PerDocumentStyleData>::from_raw(data as *mut PerDocumentStyleData);
    }
}
<commit_msg>Hackily return default computed values for unstyled nodes.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#![allow(unsafe_code)]

use app_units::Au;
use bindings::{RawGeckoDocument, RawGeckoElement, RawGeckoNode};
use bindings::{RawServoStyleSet, RawServoStyleSheet, ServoComputedValues, ServoNodeData};
use bindings::{nsIAtom};
use data::PerDocumentStyleData;
use env_logger;
use euclid::Size2D;
use gecko_style_structs::SheetParsingMode;
use properties::GeckoComputedValues;
use selector_impl::{GeckoSelectorImpl, PseudoElement, SharedStyleContext, Stylesheet};
use std::marker::PhantomData;
use std::mem::{forget, transmute};
use std::ptr;
use std::slice;
use std::str::from_utf8_unchecked;
use std::sync::{Arc, Mutex};
use style::context::{ReflowGoal};
use style::dom::{TDocument, TElement, TNode};
use style::error_reporting::StdoutErrorReporter;
use style::parallel;
use style::properties::ComputedValues;
use style::selector_impl::{SelectorImplExt, PseudoElementCascadeType};
use style::stylesheets::Origin;
use traversal::RecalcStyleOnly;
use url::Url;
use util::arc_ptr_eq;
use wrapper::{GeckoDocument, GeckoElement, GeckoNode, NonOpaqueStyleData};

// TODO: This is ugly and should go away once we get an atom back-end.
pub fn pseudo_element_from_atom(pseudo: *mut nsIAtom,
                                in_ua_stylesheet: bool) -> Result<PseudoElement, String> {
    use bindings::Gecko_GetAtomAsUTF16;
    use selectors::parser::{ParserContext, SelectorImpl};

    let pseudo_string = unsafe {
        let mut length = 0;
        let mut buff = Gecko_GetAtomAsUTF16(pseudo, &mut length);

        // Handle the annoying preceding colon in front of everything in nsCSSAnonBoxList.h.
        debug_assert!(length >= 2 && *buff == ':' as u16 && *buff.offset(1) != ':' as u16);
        buff = buff.offset(1);
        length -= 1;

        String::from_utf16(slice::from_raw_parts(buff, length as usize)).unwrap()
    };

    let mut context = ParserContext::new();
    context.in_user_agent_stylesheet = in_ua_stylesheet;
    GeckoSelectorImpl::parse_pseudo_element(&context, &pseudo_string).map_err(|_| pseudo_string)
}

/*
 * For Gecko->Servo function calls, we need to redeclare the same signature that was declared in
 * the C header in Gecko. In order to catch accidental mismatches, we run rust-bindgen against
 * those signatures as well, giving us a second declaration of all the Servo_* functions in this
 * crate. If there's a mismatch, LLVM will assert and abort, which is a rather awful thing to
 * depend on but good enough for our purposes.
 */

#[no_mangle]
pub extern "C" fn Servo_Initialize() -> () {
    // Enable standard Rust logging.
    //
    // See https://doc.rust-lang.org/log/env_logger/index.html for instructions.
    env_logger::init().unwrap();
}

#[no_mangle]
pub extern "C" fn Servo_RestyleDocument(doc: *mut RawGeckoDocument, raw_data: *mut RawServoStyleSet) -> () {
    let document = unsafe { GeckoDocument::from_raw(doc) };
    let node = match document.root_node() {
        Some(x) => x,
        None => return,
    };
    let data = unsafe { &mut *(raw_data as *mut PerDocumentStyleData) };

    // Force the creation of our lazily-constructed initial computed values on
    // the main thread, since it's not safe to call elsewhere.
    //
    // FIXME(bholley): this should move into Servo_Initialize as soon as we get
    // rid of the HackilyFindSomeDeviceContext stuff that happens during
    // initial_values computation, since that stuff needs to be called further
    // along in startup than the sensible place to call Servo_Initialize.
    GeckoComputedValues::initial_values();

    let _needs_dirtying = Arc::get_mut(&mut data.stylist).unwrap()
                              .update(&data.stylesheets, data.stylesheets_changed);
    data.stylesheets_changed = false;

    let shared_style_context = SharedStyleContext {
        viewport_size: Size2D::new(Au(0), Au(0)),
        screen_size_changed: false,
        generation: 0,
        goal: ReflowGoal::ForScriptQuery,
        stylist: data.stylist.clone(),
        new_animations_sender: Mutex::new(data.new_animations_sender.clone()),
        running_animations: data.running_animations.clone(),
        expired_animations: data.expired_animations.clone(),
        error_reporter: Box::new(StdoutErrorReporter),
    };

    if node.is_dirty() || node.has_dirty_descendants() {
        parallel::traverse_dom::<GeckoNode, RecalcStyleOnly>(node, &shared_style_context, &mut data.work_queue);
    }
}

#[no_mangle]
pub extern "C" fn Servo_DropNodeData(data: *mut ServoNodeData) -> () {
    unsafe {
        let _ = Box::<NonOpaqueStyleData>::from_raw(data as *mut NonOpaqueStyleData);
    }
}

#[no_mangle]
pub extern "C" fn Servo_StylesheetFromUTF8Bytes(bytes: *const u8,
                                                length: u32,
                                                mode: SheetParsingMode) -> *mut RawServoStyleSheet {

    let input = unsafe { from_utf8_unchecked(slice::from_raw_parts(bytes, length as usize)) };

    let origin = match mode {
        SheetParsingMode::eAuthorSheetFeatures => Origin::Author,
        SheetParsingMode::eUserSheetFeatures => Origin::User,
        SheetParsingMode::eAgentSheetFeatures => Origin::UserAgent,
    };

    // FIXME(heycam): Pass in the real base URL.
    let url = Url::parse("about:none").unwrap();
    let sheet = Arc::new(Stylesheet::from_str(input, url, origin, Box::new(StdoutErrorReporter)));
    unsafe {
        transmute(sheet)
    }
}

pub struct ArcHelpers<GeckoType, ServoType> {
    phantom1: PhantomData<GeckoType>,
    phantom2: PhantomData<ServoType>,
}

impl<GeckoType, ServoType> ArcHelpers<GeckoType, ServoType> {
    pub fn with<F, Output>(raw: *mut GeckoType, cb: F) -> Output
                           where F: FnOnce(&Arc<ServoType>) -> Output {
        debug_assert!(!raw.is_null());

        let owned = unsafe { Self::into(raw) };
        let result = cb(&owned);
        forget(owned);
        result
    }

    pub fn maybe_with<F, Output>(maybe_raw: *mut GeckoType, cb: F) -> Output
                                 where F: FnOnce(Option<&Arc<ServoType>>) -> Output {
        let owned = if maybe_raw.is_null() {
            None
        } else {
            Some(unsafe { Self::into(maybe_raw) })
        };

        let result = cb(owned.as_ref());
        forget(owned);

        result
    }

    pub unsafe fn into(ptr: *mut GeckoType) -> Arc<ServoType> {
        transmute(ptr)
    }

    pub fn from(owned: Arc<ServoType>) -> *mut GeckoType {
        unsafe { transmute(owned) }
    }

    pub unsafe fn addref(ptr: *mut GeckoType) {
        Self::with(ptr, |arc| forget(arc.clone()));
    }

    pub unsafe fn release(ptr: *mut GeckoType) {
        let _ = Self::into(ptr);
    }
}

#[no_mangle]
pub extern "C" fn Servo_AppendStyleSheet(raw_sheet: *mut RawServoStyleSheet,
                                         raw_data: *mut RawServoStyleSet) {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);
    Helpers::with(raw_sheet, |sheet| {
        data.stylesheets.retain(|x| !arc_ptr_eq(x, sheet));
        data.stylesheets.push(sheet.clone());
        data.stylesheets_changed = true;
    });
}

#[no_mangle]
pub extern "C" fn Servo_PrependStyleSheet(raw_sheet: *mut RawServoStyleSheet,
                                          raw_data: *mut RawServoStyleSet) {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);
    Helpers::with(raw_sheet, |sheet| {
        data.stylesheets.retain(|x| !arc_ptr_eq(x, sheet));
        data.stylesheets.insert(0, sheet.clone());
        data.stylesheets_changed = true;
    })
}

#[no_mangle]
pub extern "C" fn Servo_InsertStyleSheetBefore(raw_sheet: *mut RawServoStyleSheet,
                                               raw_reference: *mut RawServoStyleSheet,
                                               raw_data: *mut RawServoStyleSet) {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);
    Helpers::with(raw_sheet, |sheet| {
        Helpers::with(raw_reference, |reference| {
            data.stylesheets.retain(|x| !arc_ptr_eq(x, sheet));
            let index = data.stylesheets.iter().position(|x| arc_ptr_eq(x, reference)).unwrap();
            data.stylesheets.insert(index, sheet.clone());
            data.stylesheets_changed = true;
        })
    })
}

#[no_mangle]
pub extern "C" fn Servo_RemoveStyleSheet(raw_sheet: *mut RawServoStyleSheet,
                                         raw_data: *mut RawServoStyleSet) {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);
    Helpers::with(raw_sheet, |sheet| {
        data.stylesheets.retain(|x| !arc_ptr_eq(x, sheet));
        data.stylesheets_changed = true;
    });
}

#[no_mangle]
pub extern "C" fn Servo_StyleSheetHasRules(raw_sheet: *mut RawServoStyleSheet) -> bool {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    Helpers::with(raw_sheet, |sheet| !sheet.rules.is_empty())
}

#[no_mangle]
pub extern "C" fn Servo_AddRefStyleSheet(sheet: *mut RawServoStyleSheet) -> () {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    unsafe { Helpers::addref(sheet) };
}

#[no_mangle]
pub extern "C" fn Servo_ReleaseStyleSheet(sheet: *mut RawServoStyleSheet) -> () {
    type Helpers = ArcHelpers<RawServoStyleSheet, Stylesheet>;
    unsafe { Helpers::release(sheet) };
}

#[no_mangle]
pub extern "C" fn Servo_GetComputedValues(node: *mut RawGeckoNode)
     -> *mut ServoComputedValues {
    use selectors::Element;
    let node = unsafe { GeckoNode::from_raw(node) };
    let arc_cv = match node.borrow_data().map_or(None, |data| data.style.clone()) {
        Some(style) => style,
        None => {
            // FIXME(bholley): This case subverts the intended semantics of this
            // function, and exists only to make stylo builds more robust corner-
            // cases where Gecko wants the style for a node that Servo never
            // traversed. We should remove this as soon as possible.
            error!("stylo: encountered unstyled node, substituting default values.");
            Arc::new(GeckoComputedValues::initial_values().clone())
        },
    };
    unsafe { transmute(arc_cv) }
}

#[no_mangle]
pub extern "C" fn Servo_GetComputedValuesForAnonymousBox(parent_style_or_null: *mut ServoComputedValues,
                                                         pseudo_tag: *mut nsIAtom,
                                                         raw_data: *mut RawServoStyleSet)
     -> *mut ServoComputedValues {
    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);

    let pseudo = match pseudo_element_from_atom(pseudo_tag, /* ua_stylesheet = */ true) {
        Ok(pseudo) => pseudo,
        Err(pseudo) => {
            warn!("stylo: Unable to parse anonymous-box pseudo-element: {}", pseudo);
            return ptr::null_mut();
        }
    };

    type Helpers = ArcHelpers<ServoComputedValues, GeckoComputedValues>;

    Helpers::maybe_with(parent_style_or_null, |maybe_parent| {
        let new_computed = data.stylist.precomputed_values_for_pseudo(&pseudo, maybe_parent);
        new_computed.map_or(ptr::null_mut(), |c| Helpers::from(c))
    })
}

#[no_mangle]
pub extern "C" fn Servo_GetComputedValuesForPseudoElement(parent_style: *mut ServoComputedValues,
                                                          match_element: *mut RawGeckoElement,
                                                          pseudo_tag: *mut nsIAtom,
                                                          raw_data: *mut RawServoStyleSet,
                                                          is_probe: bool)
     -> *mut ServoComputedValues {
    debug_assert!(!match_element.is_null());

    let parent_or_null = || {
        if is_probe {
            ptr::null_mut()
        } else {
            Servo_AddRefComputedValues(parent_style);
            parent_style
        }
    };

    let pseudo = match pseudo_element_from_atom(pseudo_tag, /* ua_stylesheet = */ true) {
        Ok(pseudo) => pseudo,
        Err(pseudo) => {
            warn!("stylo: Unable to parse anonymous-box pseudo-element: {}", pseudo);
            return parent_or_null();
        }
    };

    let data = PerDocumentStyleData::borrow_mut_from_raw(raw_data);

    let element = unsafe { GeckoElement::from_raw(match_element) };

    type Helpers = ArcHelpers<ServoComputedValues, GeckoComputedValues>;

    match GeckoSelectorImpl::pseudo_element_cascade_type(&pseudo) {
        PseudoElementCascadeType::Eager => {
            let node = element.as_node();
            let maybe_computed = node.borrow_data()
                                     .and_then(|data| {
                                         data.per_pseudo.get(&pseudo).map(|c| c.clone())
                                     });
            maybe_computed.map_or_else(parent_or_null, Helpers::from)
        }
        PseudoElementCascadeType::Lazy => {
            Helpers::with(parent_style, |parent| {
                data.stylist
                    .lazily_compute_pseudo_element_style(&element, &pseudo, parent)
                    .map_or_else(parent_or_null, Helpers::from)
            })
        }
        PseudoElementCascadeType::Precomputed => {
            unreachable!("Anonymous pseudo found in \
                         Servo_GetComputedValuesForPseudoElement");
        }
    }
}

#[no_mangle]
pub extern "C" fn Servo_InheritComputedValues(parent_style: *mut ServoComputedValues)
     -> *mut ServoComputedValues {
    type Helpers = ArcHelpers<ServoComputedValues, GeckoComputedValues>;
    Helpers::with(parent_style, |parent| {
        let style = GeckoComputedValues::inherit_from(parent);
        Helpers::from(style)
    })
}

#[no_mangle]
pub extern "C" fn Servo_AddRefComputedValues(ptr: *mut ServoComputedValues) -> () {
    type Helpers = ArcHelpers<ServoComputedValues, GeckoComputedValues>;
    unsafe { Helpers::addref(ptr) };
}

#[no_mangle]
pub extern "C" fn Servo_ReleaseComputedValues(ptr: *mut ServoComputedValues) -> () {
    type Helpers = ArcHelpers<ServoComputedValues, GeckoComputedValues>;
    unsafe { Helpers::release(ptr) };
}

#[no_mangle]
pub extern "C" fn Servo_InitStyleSet() -> *mut RawServoStyleSet {
    let data = Box::new(PerDocumentStyleData::new());
    Box::into_raw(data) as *mut RawServoStyleSet
}

#[no_mangle]
pub extern "C" fn Servo_DropStyleSet(data: *mut RawServoStyleSet) -> () {
    unsafe {
        let _ = Box::<PerDocumentStyleData>::from_raw(data as *mut PerDocumentStyleData);
    }
}
<|endoftext|>


<commit_before>use futures::stream;
use futures::stream::Stream;

use bytes::Bytes;

use data_or_headers::DataOrHeaders;
use data_or_trailers::DataOrTrailers;
use error;
use futures::Poll;
use misc::any_to_string;
use solicit::end_stream::EndStream;
use solicit::header::Headers;
use solicit_async::HttpFutureStreamSend;
use std::panic::AssertUnwindSafe;

/// Stream frame content with END_STREAM flag
pub struct DataOrHeadersWithFlag {
    pub content: DataOrHeaders,
    /// END_STREAM
    pub last: bool,
}

impl DataOrHeadersWithFlag {
    pub fn last_headers(headers: Headers) -> Self {
        DataOrHeadersWithFlag {
            content: DataOrHeaders::Headers(headers),
            last: true,
        }
    }

    pub fn intermediate_headers(headers: Headers) -> Self {
        DataOrHeadersWithFlag {
            content: DataOrHeaders::Headers(headers),
            last: false,
        }
    }

    pub fn intermediate_data(data: Bytes) -> Self {
        DataOrHeadersWithFlag {
            content: DataOrHeaders::Data(data),
            last: false,
        }
    }

    pub fn last_data(data: Bytes) -> Self {
        DataOrHeadersWithFlag {
            content: DataOrHeaders::Data(data),
            last: true,
        }
    }

    pub fn into_after_headers(self) -> DataOrTrailers {
        let DataOrHeadersWithFlag { content, last } = self;
        match (content, last) {
            (DataOrHeaders::Data(data), last) => {
                let end_stream = if last { EndStream::Yes } else { EndStream::No };
                DataOrTrailers::Data(data, end_stream)
            }
            (DataOrHeaders::Headers(headers), _) => DataOrTrailers::Trailers(headers),
        }
    }
}

/// Stream of DATA of HEADER frames
pub struct DataOrHeadersWithFlagStream(pub HttpFutureStreamSend<DataOrHeadersWithFlag>);

impl DataOrHeadersWithFlagStream {
    // constructors

    pub fn new<S>(s: S) -> DataOrHeadersWithFlagStream
    where
        S: Stream<Item = DataOrHeadersWithFlag, Error = error::Error> + Send + 'static,
    {
        DataOrHeadersWithFlagStream(Box::new(s))
    }

    pub fn empty() -> DataOrHeadersWithFlagStream {
        DataOrHeadersWithFlagStream::new(stream::empty())
    }

    pub fn bytes<S>(bytes: S) -> DataOrHeadersWithFlagStream
    where
        S: Stream<Item = Bytes, Error = error::Error> + Send + 'static,
    {
        DataOrHeadersWithFlagStream::new(bytes.map(DataOrHeadersWithFlag::intermediate_data))
    }

    pub fn once(part: DataOrHeaders) -> DataOrHeadersWithFlagStream {
        DataOrHeadersWithFlagStream::new(stream::once(Ok(DataOrHeadersWithFlag {
            content: part,
            last: true,
        })))
    }

    pub fn once_bytes<B>(bytes: B) -> DataOrHeadersWithFlagStream
    where
        B: Into<Bytes>,
    {
        DataOrHeadersWithFlagStream::once(DataOrHeaders::Data(bytes.into()))
    }

    // getters

    /// Create a stream without "last" flag
    pub fn drop_last_flag(self) -> HttpFutureStreamSend<DataOrHeaders> {
        Box::new(self.map(|DataOrHeadersWithFlag { content, .. }| content))
    }

    /// Take only `DATA` frames from the stream
    pub fn filter_data(self) -> HttpFutureStreamSend<Bytes> {
        Box::new(
            self.filter_map(|DataOrHeadersWithFlag { content, .. }| match content {
                DataOrHeaders::Data(data) => Some(data),
                _ => None,
            }),
        )
    }

    /// Wrap a stream with `catch_unwind` combinator.
    /// Transform panic into `error::Error`
    pub fn catch_unwind(self) -> DataOrHeadersWithFlagStream {
        DataOrHeadersWithFlagStream::new(AssertUnwindSafe(self.0).catch_unwind().then(|r| {
            match r {
                Ok(r) => r,
                Err(e) => {
                    let e = any_to_string(e);
                    // TODO: send plain text error if headers weren't sent yet
                    warn!("handler panicked: {}", e);
                    Err(error::Error::HandlerPanicked(e))
                }
            }
        }))
    }
}

impl Stream for DataOrHeadersWithFlagStream {
    type Item = DataOrHeadersWithFlag;
    type Error = error::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.0.poll()
    }
}
<commit_msg>Derive Debug for DataOrHeadersWithFlag<commit_after>use futures::stream;
use futures::stream::Stream;

use bytes::Bytes;

use data_or_headers::DataOrHeaders;
use data_or_trailers::DataOrTrailers;
use error;
use futures::Poll;
use misc::any_to_string;
use solicit::end_stream::EndStream;
use solicit::header::Headers;
use solicit_async::HttpFutureStreamSend;
use std::panic::AssertUnwindSafe;

/// Stream frame content with END_STREAM flag
#[derive(Debug)]
pub struct DataOrHeadersWithFlag {
    pub content: DataOrHeaders,
    /// END_STREAM
    pub last: bool,
}

impl DataOrHeadersWithFlag {
    pub fn last_headers(headers: Headers) -> Self {
        DataOrHeadersWithFlag {
            content: DataOrHeaders::Headers(headers),
            last: true,
        }
    }

    pub fn intermediate_headers(headers: Headers) -> Self {
        DataOrHeadersWithFlag {
            content: DataOrHeaders::Headers(headers),
            last: false,
        }
    }

    pub fn intermediate_data(data: Bytes) -> Self {
        DataOrHeadersWithFlag {
            content: DataOrHeaders::Data(data),
            last: false,
        }
    }

    pub fn last_data(data: Bytes) -> Self {
        DataOrHeadersWithFlag {
            content: DataOrHeaders::Data(data),
            last: true,
        }
    }

    pub fn into_after_headers(self) -> DataOrTrailers {
        let DataOrHeadersWithFlag { content, last } = self;
        match (content, last) {
            (DataOrHeaders::Data(data), last) => {
                let end_stream = if last { EndStream::Yes } else { EndStream::No };
                DataOrTrailers::Data(data, end_stream)
            }
            (DataOrHeaders::Headers(headers), _) => DataOrTrailers::Trailers(headers),
        }
    }
}

/// Stream of DATA of HEADER frames
pub struct DataOrHeadersWithFlagStream(pub HttpFutureStreamSend<DataOrHeadersWithFlag>);

impl DataOrHeadersWithFlagStream {
    // constructors

    pub fn new<S>(s: S) -> DataOrHeadersWithFlagStream
    where
        S: Stream<Item = DataOrHeadersWithFlag, Error = error::Error> + Send + 'static,
    {
        DataOrHeadersWithFlagStream(Box::new(s))
    }

    pub fn empty() -> DataOrHeadersWithFlagStream {
        DataOrHeadersWithFlagStream::new(stream::empty())
    }

    pub fn bytes<S>(bytes: S) -> DataOrHeadersWithFlagStream
    where
        S: Stream<Item = Bytes, Error = error::Error> + Send + 'static,
    {
        DataOrHeadersWithFlagStream::new(bytes.map(DataOrHeadersWithFlag::intermediate_data))
    }

    pub fn once(part: DataOrHeaders) -> DataOrHeadersWithFlagStream {
        DataOrHeadersWithFlagStream::new(stream::once(Ok(DataOrHeadersWithFlag {
            content: part,
            last: true,
        })))
    }

    pub fn once_bytes<B>(bytes: B) -> DataOrHeadersWithFlagStream
    where
        B: Into<Bytes>,
    {
        DataOrHeadersWithFlagStream::once(DataOrHeaders::Data(bytes.into()))
    }

    // getters

    /// Create a stream without "last" flag
    pub fn drop_last_flag(self) -> HttpFutureStreamSend<DataOrHeaders> {
        Box::new(self.map(|DataOrHeadersWithFlag { content, .. }| content))
    }

    /// Take only `DATA` frames from the stream
    pub fn filter_data(self) -> HttpFutureStreamSend<Bytes> {
        Box::new(
            self.filter_map(|DataOrHeadersWithFlag { content, .. }| match content {
                DataOrHeaders::Data(data) => Some(data),
                _ => None,
            }),
        )
    }

    /// Wrap a stream with `catch_unwind` combinator.
    /// Transform panic into `error::Error`
    pub fn catch_unwind(self) -> DataOrHeadersWithFlagStream {
        DataOrHeadersWithFlagStream::new(AssertUnwindSafe(self.0).catch_unwind().then(|r| {
            match r {
                Ok(r) => r,
                Err(e) => {
                    let e = any_to_string(e);
                    // TODO: send plain text error if headers weren't sent yet
                    warn!("handler panicked: {}", e);
                    Err(error::Error::HandlerPanicked(e))
                }
            }
        }))
    }
}

impl Stream for DataOrHeadersWithFlagStream {
    type Item = DataOrHeadersWithFlag;
    type Error = error::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.0.poll()
    }
}
<|endoftext|>


<commit_before>use getopts::Options;
use std::env;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::join;
use tokio::net::{TcpListener, TcpStream};

type BoxedError = Box<dyn std::error::Error + Sync + Send + 'static>;
static DEBUG: AtomicBool = AtomicBool::new(false);

async fn forward(
    bind_ip: &str,
    local_port: i32,
    remote_host: &str,
    remote_port: i32,
) -> Result<(), BoxedError> {
    // Listen on the specified IP and port
    let bind_addr = if bind_ip.contains(':') {
        format!("[{}]:{}", bind_ip, local_port)
    } else {
        format!("{}:{}", bind_ip, local_port)
    };
    let bind_sock = bind_addr
        .parse::<std::net::SocketAddr>()
        .expect("Failed to parse bind address");
    let listener = TcpListener::bind(&bind_sock).await?;
    println!("Listening on {}", listener.local_addr().unwrap());

    // We have either been provided an IP address or a host name.
    // Instead of trying to check its format, just trying creating a SocketAddr from it.
    let parse_result = format!("{}:{}", remote_host, remote_port).parse::<std::net::SocketAddr>();
    let remote_addr = match parse_result {
        Ok(s) => s,
        Err(_) => {
            // It's a hostname; we're going to need to resolve it.
            // Create an async dns resolver

            use trust_dns_resolver::config::*;
            use trust_dns_resolver::TokioAsyncResolver;

            let resolver =
                TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default())
                    .await
                    .expect("Failed to create DNS resolver");

            let resolutions = resolver
                .lookup_ip(remote_host)
                .await
                .expect("Failed to resolve server IP address!");
            let remote_addr = resolutions
                .iter()
                .nth(1)
                .expect("Failed to resolve server IP address!");
            println!("Resolved {} to {}", remote_host, remote_addr);
            format!("{}:{}", remote_addr, remote_port).parse()?
        }
    };

    loop {
        let (mut client, client_addr) = listener.accept().await?;

        tokio::spawn(async move {
            println!("New connection from {}", client_addr);

            // Establish connection to upstream for each incoming client connection
            let mut remote = TcpStream::connect(&remote_addr).await?;
            let (mut client_recv, mut client_send) = client.split();
            let (mut remote_recv, mut remote_send) = remote.split();

            let (remote_bytes_copied, client_bytes_copied) = join!(
                tokio::io::copy(&mut remote_recv, &mut client_send),
                tokio::io::copy(&mut client_recv, &mut remote_send),
            );

            match remote_bytes_copied {
                Ok(count) => {
                    if DEBUG.load(Ordering::Relaxed) {
                        eprintln!(
                            "Transferred {} bytes from remote client {} to upstream server",
                            count, client_addr
                        );
                    }
                }
                Err(err) => {
                    eprintln!(
                        "Error writing from remote client {} to upstream server!",
                        client_addr
                    );
                    eprintln!("{:?}", err);
                }
            };

            match client_bytes_copied {
                Ok(count) => {
                    if DEBUG.load(Ordering::Relaxed) {
                        eprintln!(
                            "Transferred {} bytes from upstream server to remote client {}",
                            count, client_addr
                        );
                    }
                }
                Err(err) => {
                    eprintln!(
                        "Error writing bytes from upstream server to remote client {}",
                        client_addr
                    );
                    eprintln!("{:?}", err);
                }
            };

            let r: Result<(), BoxedError> = Ok(());
            r
        });
    }
}

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

pub fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    opts.optopt(
        "o",
        "",
        "set remote host with port for the forwarding type: REMOTE_HOST:NUMBER_PORT",
        "REMOTE_HOST_WITH_PORT",
    );
    opts.optopt(
        "i",
        "",
        "set local host with  port for the forwarding type: LOCAL_HOST:NUMBER_PORT",
        "LOCAL_HOST_WITH_PORT",
    );
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            panic!(f.to_string())
        }
    };
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }
    if !matches.opt_present("i") || !matches.opt_present("o") {
        print_usage(&program, opts);
        return;
    }
    let str_output = matches.opt_str("o").expect("Error getting output");
    let output: Vec<&str> = str_output.split(":").collect();
    let str_input = matches.opt_str("i").expect("Error getting input");
    let input: Vec<&str> = str_input.split(":").collect();

    if input.len() != 2 && output.len() != 2 {
        print_usage(&program, opts);
        return;
    }
    if input[1].parse::<i32>().is_err() || output[1].parse::<i32>().is_err() {
        print_usage(&program, opts);
        return;
    }
    let input_port = input[1].parse::<i32>().unwrap();
    let output_port = output[1].parse::<i32>().unwrap();
    let input_host = input[0];
    let output_host = output[0];
    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async {
        forward(input_host, input_port, output_host, output_port)
            .await
            .expect("Thread could not be executed");
    });
}
<commit_msg>Modifications in tcp proxy<commit_after>use futures::future::join_all;
use getopts::Options;
use std::env;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::join;
use tokio::net::{TcpListener, TcpStream};
use tokio::runtime::Runtime;

type BoxedError = Box<dyn std::error::Error + Sync + Send + 'static>;
static DEBUG: AtomicBool = AtomicBool::new(false);

async fn forward(
    bind_ip: &str,
    local_port: i32,
    remote_host: &str,
    remote_port: i32,
) -> Result<(), BoxedError> {
    // Listen on the specified IP and port
    println!("Calling forward port");
    let bind_addr = if bind_ip.contains(':') {
        format!("[{}]:{}", bind_ip, local_port)
    } else {
        format!("{}:{}", bind_ip, local_port)
    };
    println!("Binding address {}", bind_addr);
    let bind_sock = bind_addr
        .parse::<std::net::SocketAddr>()
        .expect("Failed to parse bind address");
    println!("Trying to binding...");
    let listener = TcpListener::bind(&bind_sock).await?;
    println!("Listening on {}", listener.local_addr().unwrap());

    // We have either been provided an IP address or a host name.
    // Instead of trying to check its format, just trying creating a SocketAddr from it.
    let parse_result = format!("{}:{}", remote_host, remote_port).parse::<std::net::SocketAddr>();
    let remote_addr = match parse_result {
        Ok(s) => s,
        Err(_) => {
            // It's a hostname; we're going to need to resolve it.
            // Create an async dns resolver

            use trust_dns_resolver::config::*;
            use trust_dns_resolver::TokioAsyncResolver;

            let resolver =
                TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default())
                    .await
                    .expect("Failed to create DNS resolver");

            let resolutions = resolver
                .lookup_ip(remote_host)
                .await
                .expect("Failed to resolve server IP address!");
            let remote_addr = resolutions
                .iter()
                .nth(1)
                .expect("Failed to resolve server IP address!");
            println!("Resolved {} to {}", remote_host, remote_addr);
            format!("{}:{}", remote_addr, remote_port).parse()?
        }
    };

    loop {
        let (mut client, client_addr) = listener.accept().await?;

        tokio::spawn(async move {
            println!("New connection from {}", client_addr);

            // Establish connection to upstream for each incoming client connection
            let mut remote = TcpStream::connect(&remote_addr).await?;
            let (mut client_recv, mut client_send) = client.split();
            let (mut remote_recv, mut remote_send) = remote.split();

            let (remote_bytes_copied, client_bytes_copied) = join!(
                tokio::io::copy(&mut remote_recv, &mut client_send),
                tokio::io::copy(&mut client_recv, &mut remote_send),
            );

            match remote_bytes_copied {
                Ok(count) => {
                    if DEBUG.load(Ordering::Relaxed) {
                        eprintln!(
                            "Transferred {} bytes from remote client {} to upstream server",
                            count, client_addr
                        );
                    }
                }
                Err(err) => {
                    eprintln!(
                        "Error writing from remote client {} to upstream server!",
                        client_addr
                    );
                    eprintln!("{:?}", err);
                }
            };

            match client_bytes_copied {
                Ok(count) => {
                    if DEBUG.load(Ordering::Relaxed) {
                        eprintln!(
                            "Transferred {} bytes from upstream server to remote client {}",
                            count, client_addr
                        );
                    }
                }
                Err(err) => {
                    eprintln!(
                        "Error writing bytes from upstream server to remote client {}",
                        client_addr
                    );
                    eprintln!("{:?}", err);
                }
            };

            let r: Result<(), BoxedError> = Ok(());
            r
        });
    }
}

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    opts.optopt(
        "o",
        "",
        "set remote host with port for the forwarding type: REMOTE_HOST:NUMBER_PORT",
        "REMOTE_HOST_WITH_PORT",
    );
    opts.optopt(
        "i",
        "",
        "set local host with  port for the forwarding type: LOCAL_HOST:NUMBER_PORT",
        "LOCAL_HOST_WITH_PORT",
    );
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            panic!(f.to_string())
        }
    };
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }
    if !matches.opt_present("i") || !matches.opt_present("o") {
        print_usage(&program, opts);
        return;
    }
    let str_output = matches.opt_str("o").expect("Error getting output");
    let output: Vec<&str> = str_output.split(":").collect();
    let str_input = matches.opt_str("i").expect("Error getting input");
    let input: Vec<&str> = str_input.split(":").collect();

    if input.len() != 2 && output.len() != 2 {
        print_usage(&program, opts);
        return;
    }
    if input[1].parse::<i32>().is_err() || output[1].parse::<i32>().is_err() {
        print_usage(&program, opts);
        return;
    }
    let input_port = input[1].parse::<i32>().unwrap();
    let output_port = output[1].parse::<i32>().unwrap();
    let input_host = input[0];
    let output_host = output[0];
    //let mut futures: Vec<dyn Future<Output = Result<(), BoxedError>>> = Vec::new();
    //futures.push(forward(input_host, input_port, output_host, output_port));
    let futures = vec![forward(input_host, input_port, output_host, output_port)];
    let rt = Runtime::new().unwrap();
    rt.block_on(async { join_all(futures).await });
}
<|endoftext|>


<commit_before>use std::default;

use enum_iterator::IntoEnumIterator;
use int_enum::IntEnum;

#[repr(usize)]
#[derive(Clone, Copy, Debug, IntoEnumIterator, IntEnum, PartialEq)]
pub enum ResistorColor {
    Black = 0,
    Brown = 1,
    Red = 2,
    Orange = 3,
    Yellow = 4,
    Green = 5,
    Blue = 6,
    Violet = 7,
    Grey = 8,
    White = 9,
}

pub fn color_to_value(_color: ResistorColor) -> usize {
    return _color.int_value()
}

pub fn value_to_color_string(value: usize) -> String {
    match value {
        0 => "Black".to_string(),
        1 => "Brown".to_string(),
        2 => "Red".to_string(),
        3 => "Orange".to_string(),
        4 => "Yellow".to_string(),
        5 => "Green".to_string(),
        6 => "Blue".to_string(),
        7 => "Violet".to_string(),
        8 => "Grey".to_string(),
        9 => "White".to_string(),
        _ => "value out of range".to_string()
    }
}

pub fn colors() -> Vec<ResistorColor> {
    let mut result = Vec::new();
    for color in ResistorColor::into_enum_iter() {
        result.push(color)
    }
    return result;
}
<commit_msg>Remove unused import<commit_after>use enum_iterator::IntoEnumIterator;
use int_enum::IntEnum;

#[repr(usize)]
#[derive(Clone, Copy, Debug, IntoEnumIterator, IntEnum, PartialEq)]
pub enum ResistorColor {
    Black = 0,
    Brown = 1,
    Red = 2,
    Orange = 3,
    Yellow = 4,
    Green = 5,
    Blue = 6,
    Violet = 7,
    Grey = 8,
    White = 9,
}

pub fn color_to_value(_color: ResistorColor) -> usize {
    return _color.int_value()
}

pub fn value_to_color_string(value: usize) -> String {
    match value {
        0 => "Black".to_string(),
        1 => "Brown".to_string(),
        2 => "Red".to_string(),
        3 => "Orange".to_string(),
        4 => "Yellow".to_string(),
        5 => "Green".to_string(),
        6 => "Blue".to_string(),
        7 => "Violet".to_string(),
        8 => "Grey".to_string(),
        9 => "White".to_string(),
        _ => "value out of range".to_string()
    }
}

pub fn colors() -> Vec<ResistorColor> {
    let mut result = Vec::new();
    for color in ResistorColor::into_enum_iter() {
        result.push(color)
    }
    return result;
}
<|endoftext|>