<commit_before>use std::net::ToSocketAddrs;
use std::time::Duration;
use std::env;
use router::{Router, HttpRouter, Matcher};
use middleware::{MiddlewareStack, Middleware, ErrorHandler};
use server::Server;
use hyper::method::Method;
use hyper::status::StatusCode;

//pre defined middleware
use default_error_handler::DefaultErrorHandler;

/// Configuration options for the server.
///
/// This struct provides a builder-style API for constructing the desired options.
///
/// # Examples
/// ```{rust}
/// use nickel::{Nickel, Options};
/// let mut server = Nickel::new();
///
/// // Don't print to stdout when starting the server
/// // and force using 8 threads.
/// server.options = Options::default()
///                      .output_on_listen(false)
///                      .thread_count(Some(8));
/// ```
pub struct Options {
    output_on_listen: bool,
    thread_count: Option<usize>,
}

impl Options {
    /// Whether the server should print the local address it is listening on when starting.
    ///
    /// Defaults to `true`.
    pub fn output_on_listen(mut self, output: bool) -> Self {
        self.output_on_listen = output;
        self
    }

    /// The number of threads the server should use or `None` to use the
    /// `hyper`'s default of `1.25 * core_count`.
    ///
    /// Defaults to `None`.
    pub fn thread_count(mut self, thread_count: Option<usize>) -> Self {
        self.thread_count = thread_count;
        self
    }
}

impl Default for Options {
    fn default() -> Self {
        Options {
            output_on_listen: true,
            thread_count: None,
        }
    }
}

/// Nickel is the application object. It's the surface that
/// holds all public APIs.
pub struct Nickel<D: Sync + Send + 'static = ()> {
    middleware_stack: MiddlewareStack<D>,
    data: D,
    keep_alive_timeout: Option<Duration>,

    /// Configuration options for the server.
    pub options: Options,
}

impl<D: Sync + Send + 'static> HttpRouter<D> for Nickel<D> {
    fn add_route<M: Into<Matcher>, H: Middleware<D>>(&mut self, method: Method, matcher: M, handler: H) -> &mut Self {
        let mut router = Router::new();
        router.add_route(method, matcher, handler);
        self.utilize(router);
        self
    }
}

impl Nickel<()> {
    /// Creates an instance of Nickel with default error handling.
    pub fn new() -> Nickel<()> {
        Nickel::with_data(())
    }
}

impl<D: Sync + Send + 'static> Nickel<D> {
    /// Creates an instance of Nickel with default error handling and custom data.
    pub fn with_data(data: D) -> Nickel<D> {
        let mut middleware_stack = MiddlewareStack::new();

        // Hook up the default error handler by default. Users are
        // free to cancel it out from their custom error handler if
        // they don't like the default behaviour.
        middleware_stack.add_error_handler(DefaultErrorHandler);

        Nickel {
            middleware_stack: middleware_stack,
            options: Options::default(),
            data: data,
            // Default value from nginx
            keep_alive_timeout: Some(Duration::from_secs(75))
        }
    }

    /// Registers a middleware handler which will be invoked among other middleware
    /// handlers before each request. Middleware can be stacked and is invoked in the
    /// same order it was registered.
    ///
    /// A middleware handler is nearly identical to a regular route handler with the only
    /// difference that it expects a result of either Action or NickelError.
    /// That is to indicate whether other middleware handlers (if any) further
    /// down the stack should continue or if the middleware invocation should
    /// be stopped after the current handler.
    ///
    /// # Examples
    /// ```{rust}
    /// # #[macro_use] extern crate nickel;
    /// # fn main() {
    /// use nickel::Nickel;
    /// let mut server = Nickel::new();
    /// server.utilize(middleware! { |req|
    ///     println!("logging request: {:?}", req.origin.uri);
    /// });
    /// # }
    /// ```
    pub fn utilize<T: Middleware<D>>(&mut self, handler: T){
        self.middleware_stack.add_middleware(handler);
    }

    /// Registers an error handler which will be invoked among other error handler
    /// as soon as any regular handler returned an error
    ///
    /// A error handler is nearly identical to a regular middleware handler with the only
    /// difference that it takes an additional error parameter or type `NickelError.
    ///
    /// # Examples
    ///
    /// ```{rust}
    /// # extern crate nickel;
    /// # fn main() {
    /// use std::io::Write;
    /// use nickel::{Nickel, Request, Continue, Halt};
    /// use nickel::{NickelError, Action};
    /// use nickel::status::StatusCode::NotFound;
    ///
    /// fn error_handler<D>(err: &mut NickelError<D>, _req: &mut Request<D>) -> Action {
    ///    if let Some(ref mut res) = err.stream {
    ///        if res.status() == NotFound {
    ///            let _ = res.write_all(b"<h1>Call the police!</h1>");
    ///            return Halt(())
    ///        }
    ///    }
    ///
    ///     Continue(())
    /// }
    ///
    /// let mut server = Nickel::new();
    ///
    /// let ehandler: fn(&mut NickelError<()>, &mut Request<()>) -> Action = error_handler;
    ///
    /// server.handle_error(ehandler)
    /// # }
    /// ```
    pub fn handle_error<T: ErrorHandler<D>>(&mut self, handler: T){
        self.middleware_stack.add_error_handler(handler);
    }

    /// Create a new middleware to serve as a router.
    ///
    ///
    /// # Examples
    /// ```{rust}
    /// #[macro_use] extern crate nickel;
    /// use nickel::{Nickel, HttpRouter};
    ///
    /// fn main() {
    ///     let mut server = Nickel::new();
    ///     let mut router = Nickel::router();
    ///
    ///     router.get("/foo", middleware! {
    ///         "Hi from /foo"
    ///     });
    ///
    ///     server.utilize(router);
    /// }
    /// ```
    pub fn router() -> Router<D> {
        Router::new()
    }

    /// Bind and listen for connections on the given host and port.
    ///
    /// # Panics
    ///
    /// Panics if `addr` is an invalid `SocketAddr`.
    ///
    /// # Examples
    /// ```{rust,no_run}
    /// use nickel::Nickel;
    ///
    /// let server = Nickel::new();
    /// server.listen("127.0.0.1:6767");
    /// ```
    pub fn listen<T: ToSocketAddrs>(mut self, addr: T) {
        self.middleware_stack.add_middleware(middleware! {
            (StatusCode::NotFound, "File Not Found")
        });

        let server = Server::new(self.middleware_stack, self.data);

        let is_test_harness = env::vars().any(|(ref k, _)| k == "NICKEL_TEST_HARNESS");

        let listener = if is_test_harness {
            // If we're under a test harness, we'll pass zero to get assigned a random
            // port. See http://doc.rust-lang.org/std/net/struct.TcpListener.html#method.bind
            server.serve("localhost:0",
                         self.keep_alive_timeout,
                         self.options.thread_count)
                  .unwrap()
        } else {
            server.serve(addr,
                         self.keep_alive_timeout,
                         self.options.thread_count)
                  .unwrap()
        };

        if self.options.output_on_listen {
            println!("Listening on http://{}", listener.socket);
            println!("Ctrl-C to shutdown server");
        }
    }

    /// Set the timeout for the keep-alive loop
    ///
    /// # Performance
    ///
    /// Setting this to `None` can have significant performance impact, but if
    /// you need to use a version of rustc < 1.4, then it may be a good choice.
    ///
    /// Alternatively, setting this too high, can lead to thread exhaustion,
    /// see [this thread](https://github.com/hyperium/hyper/issues/368) for more.
    ///
    /// # Default
    ///
    /// The default value is 75 seconds.
    pub fn keep_alive_timeout(&mut self, timeout: Option<Duration>){
        self.keep_alive_timeout = timeout;
    }
}

#[cfg(feature = "ssl")]
mod ssl {
    use std::net::ToSocketAddrs;
    use std::env;
    use server::Server;
    use hyper::status::StatusCode;
    use hyper::net::Ssl;

    use super::Nickel;

    /// Bind and listen for connections on the given host and port.
    /// Only accepts SSL connections
    ///
    /// # Panics
    ///
    /// Panics if `addr` is an invalid `SocketAddr`.
    ///
    /// # Examples
    /// ```{rust,no_run}
    /// # extern crate nickel;
    /// extern crate hyper;
    ///
    /// use nickel::Nickel;
    /// use hyper::net::Openssl;
    ///
    /// # fn main() {
    /// let server = Nickel::new();
    /// let ssl = Openssl::with_cert_and_key("foo.crt", "key.pem").unwrap();
    /// server.listen_https("127.0.0.1:6767", ssl);
    /// # }
    /// ```
    impl <D: Sync + Send + 'static> Nickel<D> {
        pub fn listen_https<T,S>(mut self, addr: T, ssl: S)
        where T: ToSocketAddrs,
              S: Ssl + Send + Clone + 'static {
            self.middleware_stack.add_middleware(middleware! {
                (StatusCode::NotFound, "File Not Found")
            });

            let server = Server::new(self.middleware_stack, self.data);

            let is_test_harness = env::vars().any(|(ref k, _)| k == "NICKEL_TEST_HARNESS");

            let listener = if is_test_harness {
                // If we're under a test harness, we'll pass zero to get assigned a random
                // port. See http://doc.rust-lang.org/std/net/struct.TcpListener.html#method.bind
                server.serve_https("localhost:0",
                                   self.keep_alive_timeout,
                                   self.options.thread_count,
                                   ssl)
                      .unwrap()
            } else {
                server.serve_https(addr,
                                   self.keep_alive_timeout,
                                   self.options.thread_count,
                                   ssl)
                      .unwrap()
            };

            println!("Listening on https://{}", listener.socket);
            println!("Ctrl-C to shutdown server");
        }
    }
}

#[test]
#[should_panic]
fn invalid_listen_addr() {
    Nickel::new().listen("127.0.0.1.6667");
}
<commit_msg>fix(options): obey `output_on_listen` for HTTPS<commit_after>use std::net::ToSocketAddrs;
use std::time::Duration;
use std::env;
use router::{Router, HttpRouter, Matcher};
use middleware::{MiddlewareStack, Middleware, ErrorHandler};
use server::Server;
use hyper::method::Method;
use hyper::status::StatusCode;

//pre defined middleware
use default_error_handler::DefaultErrorHandler;

/// Configuration options for the server.
///
/// This struct provides a builder-style API for constructing the desired options.
///
/// # Examples
/// ```{rust}
/// use nickel::{Nickel, Options};
/// let mut server = Nickel::new();
///
/// // Don't print to stdout when starting the server
/// // and force using 8 threads.
/// server.options = Options::default()
///                      .output_on_listen(false)
///                      .thread_count(Some(8));
/// ```
pub struct Options {
    output_on_listen: bool,
    thread_count: Option<usize>,
}

impl Options {
    /// Whether the server should print the local address it is listening on when starting.
    ///
    /// Defaults to `true`.
    pub fn output_on_listen(mut self, output: bool) -> Self {
        self.output_on_listen = output;
        self
    }

    /// The number of threads the server should use or `None` to use the
    /// `hyper`'s default of `1.25 * core_count`.
    ///
    /// Defaults to `None`.
    pub fn thread_count(mut self, thread_count: Option<usize>) -> Self {
        self.thread_count = thread_count;
        self
    }
}

impl Default for Options {
    fn default() -> Self {
        Options {
            output_on_listen: true,
            thread_count: None,
        }
    }
}

/// Nickel is the application object. It's the surface that
/// holds all public APIs.
pub struct Nickel<D: Sync + Send + 'static = ()> {
    middleware_stack: MiddlewareStack<D>,
    data: D,
    keep_alive_timeout: Option<Duration>,

    /// Configuration options for the server.
    pub options: Options,
}

impl<D: Sync + Send + 'static> HttpRouter<D> for Nickel<D> {
    fn add_route<M: Into<Matcher>, H: Middleware<D>>(&mut self, method: Method, matcher: M, handler: H) -> &mut Self {
        let mut router = Router::new();
        router.add_route(method, matcher, handler);
        self.utilize(router);
        self
    }
}

impl Nickel<()> {
    /// Creates an instance of Nickel with default error handling.
    pub fn new() -> Nickel<()> {
        Nickel::with_data(())
    }
}

impl<D: Sync + Send + 'static> Nickel<D> {
    /// Creates an instance of Nickel with default error handling and custom data.
    pub fn with_data(data: D) -> Nickel<D> {
        let mut middleware_stack = MiddlewareStack::new();

        // Hook up the default error handler by default. Users are
        // free to cancel it out from their custom error handler if
        // they don't like the default behaviour.
        middleware_stack.add_error_handler(DefaultErrorHandler);

        Nickel {
            middleware_stack: middleware_stack,
            options: Options::default(),
            data: data,
            // Default value from nginx
            keep_alive_timeout: Some(Duration::from_secs(75))
        }
    }

    /// Registers a middleware handler which will be invoked among other middleware
    /// handlers before each request. Middleware can be stacked and is invoked in the
    /// same order it was registered.
    ///
    /// A middleware handler is nearly identical to a regular route handler with the only
    /// difference that it expects a result of either Action or NickelError.
    /// That is to indicate whether other middleware handlers (if any) further
    /// down the stack should continue or if the middleware invocation should
    /// be stopped after the current handler.
    ///
    /// # Examples
    /// ```{rust}
    /// # #[macro_use] extern crate nickel;
    /// # fn main() {
    /// use nickel::Nickel;
    /// let mut server = Nickel::new();
    /// server.utilize(middleware! { |req|
    ///     println!("logging request: {:?}", req.origin.uri);
    /// });
    /// # }
    /// ```
    pub fn utilize<T: Middleware<D>>(&mut self, handler: T){
        self.middleware_stack.add_middleware(handler);
    }

    /// Registers an error handler which will be invoked among other error handler
    /// as soon as any regular handler returned an error
    ///
    /// A error handler is nearly identical to a regular middleware handler with the only
    /// difference that it takes an additional error parameter or type `NickelError.
    ///
    /// # Examples
    ///
    /// ```{rust}
    /// # extern crate nickel;
    /// # fn main() {
    /// use std::io::Write;
    /// use nickel::{Nickel, Request, Continue, Halt};
    /// use nickel::{NickelError, Action};
    /// use nickel::status::StatusCode::NotFound;
    ///
    /// fn error_handler<D>(err: &mut NickelError<D>, _req: &mut Request<D>) -> Action {
    ///    if let Some(ref mut res) = err.stream {
    ///        if res.status() == NotFound {
    ///            let _ = res.write_all(b"<h1>Call the police!</h1>");
    ///            return Halt(())
    ///        }
    ///    }
    ///
    ///     Continue(())
    /// }
    ///
    /// let mut server = Nickel::new();
    ///
    /// let ehandler: fn(&mut NickelError<()>, &mut Request<()>) -> Action = error_handler;
    ///
    /// server.handle_error(ehandler)
    /// # }
    /// ```
    pub fn handle_error<T: ErrorHandler<D>>(&mut self, handler: T){
        self.middleware_stack.add_error_handler(handler);
    }

    /// Create a new middleware to serve as a router.
    ///
    ///
    /// # Examples
    /// ```{rust}
    /// #[macro_use] extern crate nickel;
    /// use nickel::{Nickel, HttpRouter};
    ///
    /// fn main() {
    ///     let mut server = Nickel::new();
    ///     let mut router = Nickel::router();
    ///
    ///     router.get("/foo", middleware! {
    ///         "Hi from /foo"
    ///     });
    ///
    ///     server.utilize(router);
    /// }
    /// ```
    pub fn router() -> Router<D> {
        Router::new()
    }

    /// Bind and listen for connections on the given host and port.
    ///
    /// # Panics
    ///
    /// Panics if `addr` is an invalid `SocketAddr`.
    ///
    /// # Examples
    /// ```{rust,no_run}
    /// use nickel::Nickel;
    ///
    /// let server = Nickel::new();
    /// server.listen("127.0.0.1:6767");
    /// ```
    pub fn listen<T: ToSocketAddrs>(mut self, addr: T) {
        self.middleware_stack.add_middleware(middleware! {
            (StatusCode::NotFound, "File Not Found")
        });

        let server = Server::new(self.middleware_stack, self.data);

        let is_test_harness = env::vars().any(|(ref k, _)| k == "NICKEL_TEST_HARNESS");

        let listener = if is_test_harness {
            // If we're under a test harness, we'll pass zero to get assigned a random
            // port. See http://doc.rust-lang.org/std/net/struct.TcpListener.html#method.bind
            server.serve("localhost:0",
                         self.keep_alive_timeout,
                         self.options.thread_count)
                  .unwrap()
        } else {
            server.serve(addr,
                         self.keep_alive_timeout,
                         self.options.thread_count)
                  .unwrap()
        };

        if self.options.output_on_listen {
            println!("Listening on http://{}", listener.socket);
            println!("Ctrl-C to shutdown server");
        }
    }

    /// Set the timeout for the keep-alive loop
    ///
    /// # Performance
    ///
    /// Setting this to `None` can have significant performance impact, but if
    /// you need to use a version of rustc < 1.4, then it may be a good choice.
    ///
    /// Alternatively, setting this too high, can lead to thread exhaustion,
    /// see [this thread](https://github.com/hyperium/hyper/issues/368) for more.
    ///
    /// # Default
    ///
    /// The default value is 75 seconds.
    pub fn keep_alive_timeout(&mut self, timeout: Option<Duration>){
        self.keep_alive_timeout = timeout;
    }
}

#[cfg(feature = "ssl")]
mod ssl {
    use std::net::ToSocketAddrs;
    use std::env;
    use server::Server;
    use hyper::status::StatusCode;
    use hyper::net::Ssl;

    use super::Nickel;

    /// Bind and listen for connections on the given host and port.
    /// Only accepts SSL connections
    ///
    /// # Panics
    ///
    /// Panics if `addr` is an invalid `SocketAddr`.
    ///
    /// # Examples
    /// ```{rust,no_run}
    /// # extern crate nickel;
    /// extern crate hyper;
    ///
    /// use nickel::Nickel;
    /// use hyper::net::Openssl;
    ///
    /// # fn main() {
    /// let server = Nickel::new();
    /// let ssl = Openssl::with_cert_and_key("foo.crt", "key.pem").unwrap();
    /// server.listen_https("127.0.0.1:6767", ssl);
    /// # }
    /// ```
    impl <D: Sync + Send + 'static> Nickel<D> {
        pub fn listen_https<T,S>(mut self, addr: T, ssl: S)
        where T: ToSocketAddrs,
              S: Ssl + Send + Clone + 'static {
            self.middleware_stack.add_middleware(middleware! {
                (StatusCode::NotFound, "File Not Found")
            });

            let server = Server::new(self.middleware_stack, self.data);

            let is_test_harness = env::vars().any(|(ref k, _)| k == "NICKEL_TEST_HARNESS");

            let listener = if is_test_harness {
                // If we're under a test harness, we'll pass zero to get assigned a random
                // port. See http://doc.rust-lang.org/std/net/struct.TcpListener.html#method.bind
                server.serve_https("localhost:0",
                                   self.keep_alive_timeout,
                                   self.options.thread_count,
                                   ssl)
                      .unwrap()
            } else {
                server.serve_https(addr,
                                   self.keep_alive_timeout,
                                   self.options.thread_count,
                                   ssl)
                      .unwrap()
            };

            if self.options.output_on_listen {
                println!("Listening on https://{}", listener.socket);
                println!("Ctrl-C to shutdown server");
            }
        }
    }
}

#[test]
#[should_panic]
fn invalid_listen_addr() {
    Nickel::new().listen("127.0.0.1.6667");
}
<|endoftext|>


<commit_before>use std::hash::{Hash};
use std::slice::{
    Items,
};
use std::fmt;

#[deriving(Copy, Clone, Show, PartialEq, PartialOrd)]
pub struct NodeIndex(uint);
#[deriving(Copy, Clone, Show, PartialEq)]
pub struct EdgeIndex(uint);

const InvalidEdge: EdgeIndex = EdgeIndex(::std::uint::MAX);
const InvalidNode: NodeIndex = NodeIndex(::std::uint::MAX);

/// Index into the EdgeIndex arrays
enum InOut {
    Out = 0,
    In = 1
}

#[deriving(Show)]
pub struct Node<N> {
    pub data: N,
    next: [EdgeIndex, ..2],
}

#[deriving(Show, Copy)]
pub struct Edge<E> {
    pub data: E,
    next: [EdgeIndex, ..2],
    a: NodeIndex,
    b: NodeIndex,
}

/// **OGraph\<N, E\>** is a graph.
//#[deriving(Show)]
pub struct OGraph<N> {
    nodes: Vec<Node<N>>,
    edges: Vec<Edge<()>>,
}

impl<N: fmt::Show> fmt::Show for OGraph<N>
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for n in self.nodes.iter() {
            try!(writeln!(f, "{}", n));
        }
        for n in self.edges.iter() {
            try!(writeln!(f, "{}", n));
        }
        Ok(())
    }
}

pub enum Pair<'a, T: 'a> {
    Both(&'a mut T, &'a mut T),
    One(&'a mut T),
    None,
}

pub fn index_twice<T>(slc: &mut [T], a: uint, b: uint) -> Pair<T>
{
    if a == b {
        slc.get_mut(a).map_or(Pair::None, Pair::One)
    } else {
        if a >= slc.len() || b >= slc.len() {
            Pair::None
        } else {
            // safe because a, b are in bounds and distinct
            unsafe {
                let ar = &mut *(slc.unsafe_mut(a) as *mut _);
                let br = &mut *(slc.unsafe_mut(b) as *mut _);
                Pair::Both(ar, br)
            }
        }
    }
}

impl<N> OGraph<N>
where N: fmt::Show
{
    pub fn new() -> OGraph<N>
    {
        OGraph{nodes: Vec::new(), edges: Vec::new()}
    }

    pub fn add_node(&mut self, data: N) -> NodeIndex
    {
        let node = Node{data: data, next: [InvalidEdge, InvalidEdge]};
        let node_idx = NodeIndex(self.nodes.len());
        self.nodes.push(node);
        node_idx
    }

    pub fn add_edge(&mut self, a: NodeIndex, b: NodeIndex) -> EdgeIndex
    {
        let edge_idx = EdgeIndex(self.edges.len());
        match index_twice(self.nodes[mut], a.0, b.0) {
            Pair::None => panic!("NodeIndices out of bounds"),
            Pair::One(an) => {
                let edge = Edge {
                    data: (),
                    a: a,
                    b: b,
                    next: an.next,
                };
                an.next[0] = edge_idx;
                an.next[1] = edge_idx;
                self.edges.push(edge);
            }
            Pair::Both(an, bn) => {
                // a and b are different indices
                let edge = Edge {
                    data: (),
                    a: a,
                    b: b,
                    next: [an.next[0], bn.next[1]],
                };
                an.next[0] = edge_idx;
                bn.next[1] = edge_idx;
                self.edges.push(edge);
            }
        }
        edge_idx
    }

    pub fn remove_node(&mut self, a: NodeIndex) -> Option<N>
    {
        let remove_node = match self.nodes.remove(a.0) {
            None => return None,
            Some(n) => n,
        };

        // Adjust all node indices affected
        // Mark edges to be removed with InvalidNode links
        for edge in self.edges.iter_mut() {
            if edge.a == a {
                edge.a = InvalidNode;
            } else if edge.a > a {
                edge.a = NodeIndex(edge.a.0 - 1);
            }
            if edge.b == a {
                edge.b = InvalidNode;
            } else if edge.b > a {
                edge.b = NodeIndex(edge.b.0 - 1);
            }
        }

        // Rewrite edge chains to skip edges to be removed
        for node in self.nodes.iter_mut() {
            let mut fst = node.next[0];
            loop {
                println!("Examining {} for node {}", fst, node);
                match self.edges.get_mut(fst.0) {
                    None => break,
                    Some(edge) => {
                        if edge.a == InvalidNode || edge.b == InvalidNode {
                            println!("Edge to SKIP: {}", edge);
                        }
                        fst = edge.next[0];
                    }
                }
            }

            // "in" chain
            let mut fst = node.next[1];
            loop {
                println!("Examining {} for node {}", fst, node);
                match self.edges.get_mut(fst.0) {
                    None => break,
                    Some(edge) => {
                        if edge.a == InvalidNode || edge.b == InvalidNode {
                            println!("Edge to SKIP: {}", edge);
                        }
                        fst = edge.next[1];
                    }
                }
            }

        }

        Some(remove_node.data)
    }

    pub fn edge_mut(&mut self, e: EdgeIndex) -> &mut Edge<()>
    {
        &mut self.edges[e.0]
    }
    pub fn remove_edge(&mut self, e: EdgeIndex)
    {
        // every edge is part of two lists,
        // outgoing and incoming edges.
        // Remove it from both
        let edge = match self.edges.get(e.0) {
            None => return,
            Some(x) => *x,
        };
        // update start node
        if self.nodes[edge.a.0].next[0] == e {
            self.nodes[edge.a.0].next[0] = edge.next[0];
        } else {
            // walk through edge list
            let mut cur = self.nodes[edge.a.0].next[0];
            while cur != InvalidEdge {
                let curedge = &mut self.edges[cur.0];
                if curedge.next[0] == e {
                    println!("Have to replace link in {}", curedge);
                }
                if curedge.next[0] == e {
                    curedge.next[0] = edge.next[0];
                    break
                } else {
                    cur = curedge.next[0];
                }
            }
        }
        /*
        if self.nodes[edge.a.0].next[1] == e {
            self.nodes[edge.a.0].next[1] = edge.next[1];
        }
        */
    }
}
<commit_msg>OGraph: edge removal<commit_after>use std::hash::{Hash};
use std::slice::{
    Items,
};
use std::fmt;

#[deriving(Copy, Clone, Show, PartialEq, PartialOrd)]
pub struct NodeIndex(uint);
#[deriving(Copy, Clone, Show, PartialEq, PartialOrd)]
pub struct EdgeIndex(uint);

const InvalidEdge: EdgeIndex = EdgeIndex(::std::uint::MAX);
const InvalidNode: NodeIndex = NodeIndex(::std::uint::MAX);

/// Index into the EdgeIndex arrays
enum Dir {
    Out = 0,
    In = 1
}

#[deriving(Show)]
pub struct Node<N> {
    pub data: N,
    next: [EdgeIndex, ..2],
}

#[deriving(Show, Copy)]
pub struct Edge<E> {
    pub data: E,
    next: [EdgeIndex, ..2],
    a: NodeIndex,
    b: NodeIndex,
}

/// **OGraph\<N, E\>** is a graph.
//#[deriving(Show)]
pub struct OGraph<N> {
    nodes: Vec<Node<N>>,
    edges: Vec<Edge<()>>,
}

impl<N: fmt::Show> fmt::Show for OGraph<N>
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for n in self.nodes.iter() {
            try!(writeln!(f, "{}", n));
        }
        for n in self.edges.iter() {
            try!(writeln!(f, "{}", n));
        }
        Ok(())
    }
}

pub enum Pair<'a, T: 'a> {
    Both(&'a mut T, &'a mut T),
    One(&'a mut T),
    None,
}

pub fn index_twice<T>(slc: &mut [T], a: uint, b: uint) -> Pair<T>
{
    if a == b {
        slc.get_mut(a).map_or(Pair::None, Pair::One)
    } else {
        if a >= slc.len() || b >= slc.len() {
            Pair::None
        } else {
            // safe because a, b are in bounds and distinct
            unsafe {
                let ar = &mut *(slc.unsafe_mut(a) as *mut _);
                let br = &mut *(slc.unsafe_mut(b) as *mut _);
                Pair::Both(ar, br)
            }
        }
    }
}

impl<N> OGraph<N>
where N: fmt::Show
{
    pub fn new() -> OGraph<N>
    {
        OGraph{nodes: Vec::new(), edges: Vec::new()}
    }

    pub fn add_node(&mut self, data: N) -> NodeIndex
    {
        let node = Node{data: data, next: [InvalidEdge, InvalidEdge]};
        let node_idx = NodeIndex(self.nodes.len());
        self.nodes.push(node);
        node_idx
    }

    pub fn add_edge(&mut self, a: NodeIndex, b: NodeIndex) -> EdgeIndex
    {
        let edge_idx = EdgeIndex(self.edges.len());
        match index_twice(self.nodes[mut], a.0, b.0) {
            Pair::None => panic!("NodeIndices out of bounds"),
            Pair::One(an) => {
                let edge = Edge {
                    data: (),
                    a: a,
                    b: b,
                    next: an.next,
                };
                an.next[0] = edge_idx;
                an.next[1] = edge_idx;
                self.edges.push(edge);
            }
            Pair::Both(an, bn) => {
                // a and b are different indices
                let edge = Edge {
                    data: (),
                    a: a,
                    b: b,
                    next: [an.next[0], bn.next[1]],
                };
                an.next[0] = edge_idx;
                bn.next[1] = edge_idx;
                self.edges.push(edge);
            }
        }
        edge_idx
    }

    pub fn remove_node(&mut self, a: NodeIndex) -> Option<N>
    {
        let remove_node = match self.nodes.remove(a.0) {
            None => return None,
            Some(n) => n,
        };

        // Adjust all node indices affected
        // Mark edges to be removed with InvalidNode links
        for edge in self.edges.iter_mut() {
            if edge.a == a {
                edge.a = InvalidNode;
            } else if edge.a > a {
                edge.a = NodeIndex(edge.a.0 - 1);
            }
            if edge.b == a {
                edge.b = InvalidNode;
            } else if edge.b > a {
                edge.b = NodeIndex(edge.b.0 - 1);
            }
        }

        // Rewrite edge chains to skip edges to be removed
        for node in self.nodes.iter_mut() {
            let mut fst = node.next[0];
            loop {
                println!("Examining {} for node {}", fst, node);
                match self.edges.get_mut(fst.0) {
                    None => break,
                    Some(edge) => {
                        if edge.a == InvalidNode || edge.b == InvalidNode {
                            println!("Edge to SKIP: {}", edge);
                        }
                        fst = edge.next[0];
                    }
                }
            }

            // "in" chain
            let mut fst = node.next[1];
            loop {
                println!("Examining {} for node {}", fst, node);
                match self.edges.get_mut(fst.0) {
                    None => break,
                    Some(edge) => {
                        if edge.a == InvalidNode || edge.b == InvalidNode {
                            println!("Edge to SKIP: {}", edge);
                        }
                        fst = edge.next[1];
                    }
                }
            }

        }

        Some(remove_node.data)
    }

    pub fn edge_mut(&mut self, e: EdgeIndex) -> &mut Edge<()>
    {
        &mut self.edges[e.0]
    }

    pub fn remove_edge(&mut self, e: EdgeIndex)
    {
        fn update_edge_list(replace: EdgeIndex, fst: EdgeIndex, edges: &mut [Edge<()>], d: Dir) {
            debug_assert!(fst != replace);
            let k = d as uint;
            let edge_next = edges[replace.0].next[k];
            // walk through edge list
            let mut cur = fst;
            while cur != InvalidEdge {
                let curedge = &mut edges[cur.0];
                if curedge.next[k] == replace {
                    println!("Have to replace link in {}", curedge);
                    curedge.next[k] = edge_next;
                    break
                } else {
                    cur = curedge.next[k];
                }
            }
        }
        // every edge is part of two lists,
        // outgoing and incoming edges.
        // Remove it from both
        let (edge_a, edge_b, edge_next) = match self.edges.get(e.0) {
            None => return,
            Some(x) => (x.a, x.b, x.next),
        };
        {
            // List out from A
            let node = &mut self.nodes[edge_a.0];
            let fst = node.next[0];
            if fst == e {
                node.next[0] = edge_next[0];
            } else {
                update_edge_list(e, fst, self.edges[mut], Dir::Out);
            }
        }
        {
            // List in to B
            let node = &mut self.nodes[edge_b.0];
            let fst = node.next[1];
            if fst == e {
                node.next[1] = edge_next[1];
            } else {
                update_edge_list(e, fst, self.edges[mut], Dir::In);
            }
        }

        // Edge lists are fine, so remove the edge and adjust all edge indices
        let edge_data = self.edges.remove(e.0).unwrap().data;
        for node in self.nodes.iter_mut() {
            for &k in [0u, 1].iter() {
                if node.next[k] != InvalidEdge && node.next[k] > e {
                    node.next[k].0 -= 1;
                }
            }
        }
        for edge in self.edges.iter_mut() {
            for &k in [0u, 1].iter() {
                if edge.next[k] != InvalidEdge && edge.next[k] > e {
                    edge.next[k].0 -= 1;
                }
            }
        }
        //Some(edge_data)
    }
}
<|endoftext|>


<commit_before>pub use self::ProtType::*;

use libc;

use std::ptr;

static mut PAGE_SIZE: u32 = 0;

#[cfg(target_family = "unix")]
pub fn init_page_size() {
    unsafe {
        let val = libc::sysconf(libc::_SC_PAGESIZE);

        if val <= 0 {
            panic!("could not determine page size.");
        }

        PAGE_SIZE = val as u32;
    }
}

#[cfg(target_family = "unix")]
pub fn page_size() -> u32 {
    unsafe { PAGE_SIZE }
}

#[cfg(target_family = "unix")]
pub fn page_size_bits() -> u32 {
    // FIXME

    12
}

#[cfg(target_family = "windows")]
pub fn page_size() -> u32 {
    use kernel32::GetSystemInfo;
    use winapi::sysinfoapi::SYSTEM_INFO;
    use std::mem;

    unsafe {
        let mut system_info: SYSTEM_INFO = mem::uninitialized();
        GetSystemInfo(&mut system_info);

        system_info.dwPageSize
    }
}

#[derive(Copy, Clone, PartialEq, Eq)]
pub enum ProtType {
    Executable,
    Writable,
}

#[cfg(target_family = "unix")]
pub fn mmap(size: usize, exec: ProtType) -> *const u8 {
    let prot_exec = if exec == Executable {
        libc::PROT_EXEC
    } else {
        0
    };

    let ptr = unsafe {
        libc::mmap(ptr::null_mut(),
                   size,
                   libc::PROT_READ | libc::PROT_WRITE | prot_exec,
                   libc::MAP_PRIVATE | libc::MAP_ANON,
                   -1,
                   0) as *mut libc::c_void
    };

    if ptr == libc::MAP_FAILED {
        panic!("mmap failed");
    }

    ptr as *const u8
}

#[cfg(target_family = "windows")]
pub fn mmap(size: usize, exec: ProtType) -> *const u8 {
    use kernel32::VirtualAlloc;
    use winapi::winnt::{MEM_COMMIT, MEM_RESERVE, PAGE_EXECUTE_READWRITE, PAGE_READWRITE};

    let prot = if exec == Executable {
        PAGE_EXECUTE_READWRITE
    } else {
        PAGE_READWRITE
    };

    let ptr = unsafe { VirtualAlloc(ptr::null_mut(), size as u64, MEM_COMMIT | MEM_RESERVE, prot) };

    if ptr.is_null() {
        panic!("VirtualAlloc failed");
    }

    ptr as *const u8
}

#[cfg(target_family = "unix")]
pub fn munmap(ptr: *const u8, size: usize) {
    let res = unsafe { libc::munmap(ptr as *mut libc::c_void, size) };

    if res != 0 {
        panic!("munmap failed");
    }
}

#[cfg(target_family = "windows")]
pub fn munmap(ptr: *const u8, size: usize) {
    use kernel32::VirtualFree;
    use winapi::winnt::MEM_RELEASE;
    use winapi;

    let res = unsafe { VirtualFree(ptr as *mut winapi::c_void, 0, MEM_RELEASE) };

    if res == 0 {
        panic!("VirtualFree failed");
    }
}
<commit_msg>determine page size in bits<commit_after>pub use self::ProtType::*;

use libc;

use std::ptr;

static mut PAGE_SIZE: u32 = 0;
static mut PAGE_SIZE_BITS: u32 = 0;

#[cfg(target_family = "unix")]
pub fn init_page_size() {
    unsafe {
        let val = libc::sysconf(libc::_SC_PAGESIZE);

        if val <= 0 {
            panic!("could not determine page size.");
        }

        PAGE_SIZE = val as u32;

        assert!((PAGE_SIZE & (PAGE_SIZE - 1)) == 0);

        PAGE_SIZE_BITS = log2(PAGE_SIZE);
    }
}

/// determine log_2 of given value
fn log2(mut val: u32) -> u32 {
    let mut log = 0;
    if (val & 0xFFFF0000) != 0 { val >>= 16; log += 16; }
    if val >= 256 { val >>= 8; log += 8; }
    if val >= 16 { val >>= 4; log += 4; }
    if val >= 4 { val >>= 2; log += 2; }

    log + (val >> 1)
}

#[test]
fn test_log2() {
    assert_eq!(0, log2(0));
    assert_eq!(0, log2(1));
    assert_eq!(1, log2(2));
    assert_eq!(2, log2(4));
    assert_eq!(3, log2(8));
    assert_eq!(4, log2(16));
    assert_eq!(5, log2(32));
    assert_eq!(6, log2(64));
    assert_eq!(7, log2(128));
    assert_eq!(8, log2(256));
    assert_eq!(9, log2(512));
    assert_eq!(10, log2(1024));
    assert_eq!(11, log2(2048));
    assert_eq!(12, log2(4096));
    assert_eq!(16, log2(65536));
    assert_eq!(31, log2(1 << 31));
}

#[cfg(target_family = "unix")]
pub fn page_size() -> u32 {
    unsafe { PAGE_SIZE }
}

#[cfg(target_family = "unix")]
pub fn page_size_bits() -> u32 {
    unsafe { PAGE_SIZE_BITS }
}

#[cfg(target_family = "windows")]
pub fn page_size() -> u32 {
    use kernel32::GetSystemInfo;
    use winapi::sysinfoapi::SYSTEM_INFO;
    use std::mem;

    unsafe {
        let mut system_info: SYSTEM_INFO = mem::uninitialized();
        GetSystemInfo(&mut system_info);

        system_info.dwPageSize
    }
}

#[derive(Copy, Clone, PartialEq, Eq)]
pub enum ProtType {
    Executable,
    Writable,
}

#[cfg(target_family = "unix")]
pub fn mmap(size: usize, exec: ProtType) -> *const u8 {
    let prot_exec = if exec == Executable {
        libc::PROT_EXEC
    } else {
        0
    };

    let ptr = unsafe {
        libc::mmap(ptr::null_mut(),
                   size,
                   libc::PROT_READ | libc::PROT_WRITE | prot_exec,
                   libc::MAP_PRIVATE | libc::MAP_ANON,
                   -1,
                   0) as *mut libc::c_void
    };

    if ptr == libc::MAP_FAILED {
        panic!("mmap failed");
    }

    ptr as *const u8
}

#[cfg(target_family = "windows")]
pub fn mmap(size: usize, exec: ProtType) -> *const u8 {
    use kernel32::VirtualAlloc;
    use winapi::winnt::{MEM_COMMIT, MEM_RESERVE, PAGE_EXECUTE_READWRITE, PAGE_READWRITE};

    let prot = if exec == Executable {
        PAGE_EXECUTE_READWRITE
    } else {
        PAGE_READWRITE
    };

    let ptr = unsafe { VirtualAlloc(ptr::null_mut(), size as u64, MEM_COMMIT | MEM_RESERVE, prot) };

    if ptr.is_null() {
        panic!("VirtualAlloc failed");
    }

    ptr as *const u8
}

#[cfg(target_family = "unix")]
pub fn munmap(ptr: *const u8, size: usize) {
    let res = unsafe { libc::munmap(ptr as *mut libc::c_void, size) };

    if res != 0 {
        panic!("munmap failed");
    }
}

#[cfg(target_family = "windows")]
pub fn munmap(ptr: *const u8, size: usize) {
    use kernel32::VirtualFree;
    use winapi::winnt::MEM_RELEASE;
    use winapi;

    let res = unsafe { VirtualFree(ptr as *mut winapi::c_void, 0, MEM_RELEASE) };

    if res == 0 {
        panic!("VirtualFree failed");
    }
}
<|endoftext|>


<commit_before>#![cfg(target_os = "windows")]

pub mod ffi {
   #![allow(non_snake_case)]

   use std::mem;

   pub use libc::{
      c_char, c_int, c_uint, c_void
   };
   pub use libc::types::os::arch::extra::{
      WORD, DWORD, BYTE, BOOL, HANDLE
   };
   pub use libc::consts::os::extra::{
      TRUE
   };
   pub use os::win::ffi::{
      HDC
   };

   pub type HGLRC = HANDLE;

   pub const PFD_TYPE_RGBA:          BYTE = 0;

   pub const PFD_MAIN_PLANE:         BYTE = 0;

   pub const PFD_DOUBLEBUFFER:      DWORD = 0x00000001;
   pub const PFD_DRAW_TO_WINDOW:    DWORD = 0x00000004;
   pub const PFD_SUPPORT_OPENGL:    DWORD = 0x00000020;

   #[repr(C)]
   #[derive(Copy)]
   pub struct PIXELFORMATDESCRIPTOR {
      pub nSize: WORD,
      pub nVersion: WORD,
      pub dwFlags: DWORD,
      pub iPixelType: BYTE,
      pub cColorBits: BYTE,
      pub cRedBits: BYTE,
      pub cRedShift: BYTE,
      pub cGreenBits: BYTE,
      pub cGreenShift: BYTE,
      pub cBlueBits: BYTE,
      pub cBlueShift: BYTE,
      pub cAlphaBits: BYTE,
      pub cAlphaShift: BYTE,
      pub cAccumBits: BYTE,
      pub cAccumRedBits: BYTE,
      pub cAccumGreenBits: BYTE,
      pub cAccumBlueBits: BYTE,
      pub cAccumAlphaBits: BYTE,
      pub cDepthBits: BYTE,
      pub cStencilBits: BYTE,
      pub cAuxBuffers: BYTE,
      pub iLayerType: BYTE,
      pub bReserved: BYTE,
      pub dwLayerMask: DWORD,
      pub dwVisibleMask: DWORD,
      pub dwDamageMask: DWORD,
   }
   impl Clone for PIXELFORMATDESCRIPTOR {
      fn clone(&self) -> Self { *self }
   }
   impl Default for PIXELFORMATDESCRIPTOR {
      fn default() -> Self { unsafe { mem::zeroed() } }
   }

   #[link(name="gdi32")]
   extern "C" {
      pub fn ChoosePixelFormat(hdc: HDC, ppfd: *const PIXELFORMATDESCRIPTOR) -> c_int;

      pub fn DescribePixelFormat(hdc: HDC, iPixelFormat: c_int, nBytes: c_uint, ppfd: *mut PIXELFORMATDESCRIPTOR) -> c_int;

      pub fn SetPixelFormat(hdc: HDC, iPixelFormat: c_int, ppfd: *const PIXELFORMATDESCRIPTOR) -> BOOL;

      pub fn SwapBuffers(rc: HDC) -> BOOL;
   }

   #[link(name="opengl32")]
   extern "C" {
      pub fn wglCreateContext(hdc: HDC) -> HGLRC;

      pub fn wglMakeCurrent(hdc: HDC, rc: HGLRC) -> BOOL;

      pub fn wglDeleteContext(rc: HGLRC) -> BOOL;

      pub fn wglGetCurrentContext() -> HGLRC;

      pub fn wglGetProcAddress(name: *const c_char) -> *const c_void;
   }
}

use std::io;
use std::mem;
use std::ptr;
use std::ffi::CString;

use error::{RuntimeError, ErrorKind};
use super::utils::fn_ptr::{FnPtrLoader, FnPtr};

pub fn init_pixel_format(
   hdc: ffi::HDC,
) -> Result<(), RuntimeError> {
   let mut pfd = ffi::PIXELFORMATDESCRIPTOR {
      nSize: mem::size_of::<ffi::PIXELFORMATDESCRIPTOR>() as ffi::WORD,
      nVersion: 1,
      dwFlags: ffi::PFD_DRAW_TO_WINDOW | ffi::PFD_SUPPORT_OPENGL | ffi::PFD_DOUBLEBUFFER,
      iPixelType: ffi::PFD_TYPE_RGBA,
      cColorBits: 24,
      cRedBits: 0, cRedShift: 0, cGreenBits: 0, cGreenShift: 0, cBlueBits: 0, cBlueShift: 0,
      cAlphaBits: 0, cAlphaShift: 0, cAccumBits: 0,
      cAccumRedBits: 0, cAccumGreenBits: 0, cAccumBlueBits: 0, cAccumAlphaBits: 0,
      cDepthBits: 0,
      cStencilBits: 0,
      cAuxBuffers: 0,
      iLayerType: ffi::PFD_MAIN_PLANE,
      bReserved: 0,
      dwLayerMask: 0, dwVisibleMask: 0, dwDamageMask: 0
   };

   let pixel_format = unsafe { ffi::ChoosePixelFormat(hdc, &pfd) };

   if pixel_format == 0 {
      return Err(RuntimeError::new(
         ErrorKind::WGL,
         "Choosing pixel format failed".to_string()
      ));
   }

   println!("PIXEL FORMAT: {:?}", pixel_format);

   unsafe { ffi::DescribePixelFormat(hdc, pixel_format, mem::size_of::<ffi::PIXELFORMATDESCRIPTOR>() as ffi::c_uint, &mut pfd) };

   let result = unsafe { ffi::SetPixelFormat(hdc, pixel_format, &pfd) };

   if result != ffi::TRUE {
      return Err(RuntimeError::new(
         ErrorKind::WGL,
         "Setting pixel format failed".to_string()
      ));
   }

   Ok(())
}

pub struct Context {
   pub rc: ffi::HGLRC
}

impl Context {
   pub fn create(hdc: ffi::HDC) -> Result<Self, RuntimeError> {
      let rc = unsafe { ffi::wglCreateContext(hdc) };

      if rc == ptr::null_mut() {
         return Err(RuntimeError::new(
            ErrorKind::WGL,
            format!("Create WGL context failed: {}", io::Error::last_os_error())
         ));
      }

      let result = unsafe { ffi::wglMakeCurrent(hdc, rc) };

      if result != ffi::TRUE {
         return Err(RuntimeError::new(
            ErrorKind::WGL,
            "wglMakeCurrent failed".to_string()
         ));
      }

      Ok(Context {
         rc: rc,
      })
   }

   pub fn current() -> Result<Self, RuntimeError> {
      let rc = unsafe { ffi::wglGetCurrentContext() };

      if rc == ptr::null_mut() {
         return Err(RuntimeError::new(
            ErrorKind::WGL,
            "Getting current WGL context failed".to_string()
         ));
      }

      Ok(Context {
         rc: rc,
      })
   }
}

impl Drop for Context {
   fn drop (&mut self) {
      unsafe {
         ffi::wglDeleteContext(self.rc);
      }
   }
}

pub fn swap_buffers(hdc: ffi::HDC) {
   unsafe { ffi::SwapBuffers(hdc) };
}

pub struct Loader;

impl Loader {
   pub fn new() -> Self {
      Loader
   }
}

impl FnPtrLoader for Loader {
   fn get_proc_addr(&self, name: &str) -> FnPtr {
      let cname = CString::new(name).unwrap().as_ptr();

      let addr = unsafe {
         ffi::wglGetProcAddress(cname)
      };
      println!("{}: {:?}", name, addr);

      addr
   }
}
<commit_msg>Using super in WGL import<commit_after>#![cfg(target_os = "windows")]

pub mod ffi {
   #![allow(non_snake_case)]

   use std::mem;

   pub use libc::{
      c_char, c_int, c_uint, c_void
   };
   pub use libc::types::os::arch::extra::{
      WORD, DWORD, BYTE, BOOL, HANDLE
   };
   pub use libc::consts::os::extra::{
      TRUE
   };
   pub use super::super::win::ffi::{
      HDC
   };

   pub type HGLRC = HANDLE;

   pub const PFD_TYPE_RGBA:          BYTE = 0;

   pub const PFD_MAIN_PLANE:         BYTE = 0;

   pub const PFD_DOUBLEBUFFER:      DWORD = 0x00000001;
   pub const PFD_DRAW_TO_WINDOW:    DWORD = 0x00000004;
   pub const PFD_SUPPORT_OPENGL:    DWORD = 0x00000020;

   #[repr(C)]
   #[derive(Copy)]
   pub struct PIXELFORMATDESCRIPTOR {
      pub nSize: WORD,
      pub nVersion: WORD,
      pub dwFlags: DWORD,
      pub iPixelType: BYTE,
      pub cColorBits: BYTE,
      pub cRedBits: BYTE,
      pub cRedShift: BYTE,
      pub cGreenBits: BYTE,
      pub cGreenShift: BYTE,
      pub cBlueBits: BYTE,
      pub cBlueShift: BYTE,
      pub cAlphaBits: BYTE,
      pub cAlphaShift: BYTE,
      pub cAccumBits: BYTE,
      pub cAccumRedBits: BYTE,
      pub cAccumGreenBits: BYTE,
      pub cAccumBlueBits: BYTE,
      pub cAccumAlphaBits: BYTE,
      pub cDepthBits: BYTE,
      pub cStencilBits: BYTE,
      pub cAuxBuffers: BYTE,
      pub iLayerType: BYTE,
      pub bReserved: BYTE,
      pub dwLayerMask: DWORD,
      pub dwVisibleMask: DWORD,
      pub dwDamageMask: DWORD,
   }
   impl Clone for PIXELFORMATDESCRIPTOR {
      fn clone(&self) -> Self { *self }
   }
   impl Default for PIXELFORMATDESCRIPTOR {
      fn default() -> Self { unsafe { mem::zeroed() } }
   }

   #[link(name="gdi32")]
   extern "C" {
      pub fn ChoosePixelFormat(hdc: HDC, ppfd: *const PIXELFORMATDESCRIPTOR) -> c_int;

      pub fn DescribePixelFormat(hdc: HDC, iPixelFormat: c_int, nBytes: c_uint, ppfd: *mut PIXELFORMATDESCRIPTOR) -> c_int;

      pub fn SetPixelFormat(hdc: HDC, iPixelFormat: c_int, ppfd: *const PIXELFORMATDESCRIPTOR) -> BOOL;

      pub fn SwapBuffers(rc: HDC) -> BOOL;
   }

   #[link(name="opengl32")]
   extern "C" {
      pub fn wglCreateContext(hdc: HDC) -> HGLRC;

      pub fn wglMakeCurrent(hdc: HDC, rc: HGLRC) -> BOOL;

      pub fn wglDeleteContext(rc: HGLRC) -> BOOL;

      pub fn wglGetCurrentContext() -> HGLRC;

      pub fn wglGetProcAddress(name: *const c_char) -> *const c_void;
   }
}

use std::io;
use std::mem;
use std::ptr;
use std::ffi::CString;

use error::{RuntimeError, ErrorKind};
use super::utils::fn_ptr::{FnPtrLoader, FnPtr};

pub fn init_pixel_format(
   hdc: ffi::HDC,
) -> Result<(), RuntimeError> {
   let mut pfd = ffi::PIXELFORMATDESCRIPTOR {
      nSize: mem::size_of::<ffi::PIXELFORMATDESCRIPTOR>() as ffi::WORD,
      nVersion: 1,
      dwFlags: ffi::PFD_DRAW_TO_WINDOW | ffi::PFD_SUPPORT_OPENGL | ffi::PFD_DOUBLEBUFFER,
      iPixelType: ffi::PFD_TYPE_RGBA,
      cColorBits: 24,
      cRedBits: 0, cRedShift: 0, cGreenBits: 0, cGreenShift: 0, cBlueBits: 0, cBlueShift: 0,
      cAlphaBits: 0, cAlphaShift: 0, cAccumBits: 0,
      cAccumRedBits: 0, cAccumGreenBits: 0, cAccumBlueBits: 0, cAccumAlphaBits: 0,
      cDepthBits: 0,
      cStencilBits: 0,
      cAuxBuffers: 0,
      iLayerType: ffi::PFD_MAIN_PLANE,
      bReserved: 0,
      dwLayerMask: 0, dwVisibleMask: 0, dwDamageMask: 0
   };

   let pixel_format = unsafe { ffi::ChoosePixelFormat(hdc, &pfd) };

   if pixel_format == 0 {
      return Err(RuntimeError::new(
         ErrorKind::WGL,
         "Choosing pixel format failed".to_string()
      ));
   }

   println!("PIXEL FORMAT: {:?}", pixel_format);

   unsafe { ffi::DescribePixelFormat(hdc, pixel_format, mem::size_of::<ffi::PIXELFORMATDESCRIPTOR>() as ffi::c_uint, &mut pfd) };

   let result = unsafe { ffi::SetPixelFormat(hdc, pixel_format, &pfd) };

   if result != ffi::TRUE {
      return Err(RuntimeError::new(
         ErrorKind::WGL,
         "Setting pixel format failed".to_string()
      ));
   }

   Ok(())
}

pub struct Context {
   pub rc: ffi::HGLRC
}

impl Context {
   pub fn create(hdc: ffi::HDC) -> Result<Self, RuntimeError> {
      let rc = unsafe { ffi::wglCreateContext(hdc) };

      if rc == ptr::null_mut() {
         return Err(RuntimeError::new(
            ErrorKind::WGL,
            format!("Create WGL context failed: {}", io::Error::last_os_error())
         ));
      }

      let result = unsafe { ffi::wglMakeCurrent(hdc, rc) };

      if result != ffi::TRUE {
         return Err(RuntimeError::new(
            ErrorKind::WGL,
            "wglMakeCurrent failed".to_string()
         ));
      }

      Ok(Context {
         rc: rc,
      })
   }

   pub fn current() -> Result<Self, RuntimeError> {
      let rc = unsafe { ffi::wglGetCurrentContext() };

      if rc == ptr::null_mut() {
         return Err(RuntimeError::new(
            ErrorKind::WGL,
            "Getting current WGL context failed".to_string()
         ));
      }

      Ok(Context {
         rc: rc,
      })
   }
}

impl Drop for Context {
   fn drop (&mut self) {
      unsafe {
         ffi::wglDeleteContext(self.rc);
      }
   }
}

pub fn swap_buffers(hdc: ffi::HDC) {
   unsafe { ffi::SwapBuffers(hdc) };
}

pub struct Loader;

impl Loader {
   pub fn new() -> Self {
      Loader
   }
}

impl FnPtrLoader for Loader {
   fn get_proc_addr(&self, name: &str) -> FnPtr {
      let cname = CString::new(name).unwrap().as_ptr();

      let addr = unsafe {
         ffi::wglGetProcAddress(cname)
      };
      println!("{}: {:?}", name, addr);

      addr
   }
}
<|endoftext|>


<commit_before>use std::io::{BufRead, Cursor, Read};
use std::mem::size_of;
use byteorder::{BigEndian, ReadBytesExt};

use bits::test_bit;

pub fn bytes_of_mac(addr: u64) -> [u8; 6] {
    let mut arr = [0; 6];
    for i in 0..6 {
        arr[i] = ((addr >> (8 * i)) & 0xff) as u8;
    }
    arr
}

pub fn mac_of_bytes(addr: [u8; 6]) -> u64 {
    fn byte(u: &[u8; 6], i: usize) -> u64 {
        u[i] as u64
    };
    (byte(&addr, 0) << 8 * 5) | (byte(&addr, 1) << 8 * 4) | (byte(&addr, 2) << 8 * 3) |
    (byte(&addr, 3) << 8 * 2) | (byte(&addr, 4) << 8 * 1) | (byte(&addr, 5))
}

/// TCP Header flags.
pub struct TcpFlags {
    /// ECN-nonce concealment protection.
    pub ns: bool,
    /// Congestion window reduced.
    pub cwr: bool,
    /// ECN-Echo.
    pub ece: bool,
    /// Indicates the Urgent pointer field is significant.
    pub urg: bool,
    /// Indicates that the Acknowledgment field is significant.
    pub ack: bool,
    /// Asks to push the buffered data to the receiving application.
    pub psh: bool,
    /// Reset the connection.
    pub rst: bool,
    /// Synchronize sequence numbers.
    pub syn: bool,
    /// No more data from sender.
    pub fin: bool,
}

impl TcpFlags {
    fn of_int(d: u16) -> TcpFlags {
        TcpFlags {
            ns: test_bit(0, d as u64),
            cwr: test_bit(1, d as u64),
            ece: test_bit(2, d as u64),
            urg: test_bit(3, d as u64),
            ack: test_bit(4, d as u64),
            psh: test_bit(5, d as u64),
            rst: test_bit(6, d as u64),
            syn: test_bit(7, d as u64),
            fin: test_bit(8, d as u64),
        }
    }
}

/// TCP frame of a packet.
pub struct Tcp {
    pub src: u16,
    pub dst: u16,
    pub seq: u32,
    pub ack: u32,
    pub offset: u8,
    pub flags: TcpFlags,
    pub window: u16,
    pub chksum: u16,
    pub urgent: u16,
    pub payload: Vec<u8>,
}

#[repr(packed)]
struct TcpNet(u16, u16, u32, u32, u16, u16, u16, u16);

impl Tcp {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Tcp> {
        if bytes.get_ref().len() < size_of::<TcpNet>() {
            return None;
        }
        let src = bytes.read_u16::<BigEndian>().unwrap();
        let dst = bytes.read_u16::<BigEndian>().unwrap();
        let seq = bytes.read_u32::<BigEndian>().unwrap();
        let ack = bytes.read_u32::<BigEndian>().unwrap();
        let offset = bytes.read_u16::<BigEndian>().unwrap();
        let flags = TcpFlags::of_int(offset);
        let offset = (offset >> 12) as u8 & 0x0f;
        let window = bytes.read_u16::<BigEndian>().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let urgent = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Tcp {
            src: src,
            dst: dst,
            seq: seq,
            ack: ack,
            offset: offset,
            flags: flags,
            window: window,
            chksum: chksum,
            urgent: urgent,
            payload: payload,
        })
    }
}

/// UDP frame of a packet.
pub struct Udp {
    pub src: u16,
    pub dst: u16,
    pub chksum: u16,
    pub payload: Vec<u8>,
}

impl Udp {
    fn size_of() -> usize {
        8
    }

    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Udp> {
        if bytes.get_ref().len() < Self::size_of() {
            return None;
        }
        let src = bytes.read_u16::<BigEndian>().unwrap();
        let dst = bytes.read_u16::<BigEndian>().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Udp {
            src: src,
            dst: dst,
            chksum: chksum,
            payload: payload,
        })
    }
}

/// ICMP frame of a packet.
pub struct Icmp {
    pub typ: u8,
    pub code: u8,
    pub chksum: u16,
    pub payload: Vec<u8>,
}

impl Icmp {
    fn size_of() -> usize {
        4
    }

    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Icmp> {
        if bytes.get_ref().len() < Self::size_of() {
            return None;
        }
        let typ = bytes.read_u8().unwrap();
        let code = bytes.read_u8().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Icmp {
            typ: typ,
            code: code,
            chksum: chksum,
            payload: payload,
        })
    }
}

/// Represents packets at the transport protocol level, which are encapsulated
/// within the IPv4 payload. At present, we only support TCP, UDP, and ICMP
/// explicitly; otherwise, the raw bytes and IPv4 protocol number are provided.
pub enum Tp {
    Tcp(Tcp),
    Udp(Udp),
    Icmp(Icmp),
    Unparsable(u8, Vec<u8>),
}

/// The type of IPv4 flags.
pub struct Flags {
    pub dont_fragment: bool,
    pub more_fragments: bool,
}

impl Flags {
    fn of_int(flags: u32) -> Flags {
        Flags {
            dont_fragment: test_bit(1, flags as u64),
            more_fragments: test_bit(2, flags as u64),
        }
    }
}

/// IPv4 frame of a packet.
pub struct Ip {
    pub tos: u8,
    pub ident: u16,
    pub flags: Flags,
    pub frag: u16,
    pub ttl: u8,
    pub chksum: u16,
    pub src: u32,
    pub dst: u32,
    pub options: Vec<u8>,
    pub tp: Tp,
}

#[repr(u8)]
enum IpProto {
    IpICMP = 0x01,
    IpTCP = 0x06,
    IpUDP = 0x11,
}

#[repr(packed)]
struct IpNet(u8, u8, u16, u16, u16, u8, u8, u16, u32, u32);

impl Ip {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Ip> {
        if bytes.get_ref().len() < size_of::<IpNet>() {
            return None;
        }
        let vhl = bytes.read_u8().unwrap();
        if (vhl >> 4) != 4 {
            return None;
        }
        let ihl = vhl & 0x0f;
        let tos = bytes.read_u8().unwrap();
        bytes.consume(2);
        let ident = bytes.read_u16::<BigEndian>().unwrap();
        let frag = bytes.read_u16::<BigEndian>().unwrap();
        let flags = Flags::of_int((frag as u32) >> 13);
        let ttl = bytes.read_u8().unwrap();
        let proto = bytes.read_u8().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let src = bytes.read_u32::<BigEndian>().unwrap();
        let dst = bytes.read_u32::<BigEndian>().unwrap();
        let options_len = (ihl * 4) as usize - size_of::<IpNet>();
        let mut options = vec![0; options_len];
        bytes.read_exact(&mut options).unwrap();
        let tp = match proto {
            t if t == (IpProto::IpICMP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let icmp = Icmp::parse(bytes);
                if icmp.is_some() {
                    Tp::Icmp(icmp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            t if t == (IpProto::IpTCP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let tcp = Tcp::parse(bytes);
                if tcp.is_some() {
                    Tp::Tcp(tcp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            t if t == (IpProto::IpUDP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let udp = Udp::parse(bytes);
                if udp.is_some() {
                    Tp::Udp(udp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            _ => Tp::Unparsable(proto, bytes.get_ref().clone()),
        };
        Some(Ip {
            tos: tos,
            ident: ident,
            flags: flags,
            frag: frag,
            ttl: ttl,
            chksum: chksum,
            src: src,
            dst: dst,
            options: options,
            tp: tp,
        })
    }
}

/// Address resolution protocol (ARP) packet payload.
pub enum Arp {
    Query(u64, u32, u32),
    Reply(u64, u32, u64, u32),
}

#[repr(packed)]
struct ArpNet(u16, u16, u8, u8, u16, [u8; 6], u32, [u8; 6], u32);

impl Arp {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Arp> {
        if bytes.get_ref().len() < size_of::<ArpNet>() {
            return None;
        }
        bytes.consume(6);
        let oper = bytes.read_u16::<BigEndian>().unwrap();
        let mut sha: [u8; 6] = [0; 6];
        for i in 0..6 {
            sha[i] = bytes.read_u8().unwrap();
        }
        let spa = bytes.read_u32::<BigEndian>().unwrap();
        let mut tha: [u8; 6] = [0; 6];
        for i in 0..6 {
            tha[i] = bytes.read_u8().unwrap();
        }
        let tpa = bytes.read_u32::<BigEndian>().unwrap();
        match oper {
            0x0001 => Some(Arp::Query(mac_of_bytes(sha), spa, tpa)),
            0x0002 => Some(Arp::Reply(mac_of_bytes(sha), spa, mac_of_bytes(tha), tpa)),
            _ => None,
        }
    }
}

/// Represents a packet at the network protocol level.
pub enum Nw {
    Ip(Ip),
    Arp(Arp),
    Unparsable(u16, Vec<u8>),
}

/// Represents a packet at the ethernet protocol level.
pub struct Packet {
    pub dl_src: u64,
    pub dl_dst: u64,
    pub dl_vlan: Option<u16>,
    pub dl_vlan_dei: bool,
    pub dl_vlan_pcp: u8,
    pub nw: Nw,
}

#[repr(u16)]
enum EthTyp {
    EthTypIP = 0x0800,
    EthTypARP = 0x0806,
    EthTypVLAN = 0x8100,
}

impl Packet {
    pub fn parse(buf: &[u8]) -> Packet {
        let mut bytes = Cursor::new(buf.to_vec());
        let mut dst: [u8; 6] = [0; 6];
        let mut src: [u8; 6] = [0; 6];
        for i in 0..6 {
            dst[i] = bytes.read_u8().unwrap();
        }
        for i in 0..6 {
            src[i] = bytes.read_u8().unwrap();
        }
        let typ = bytes.read_u16::<BigEndian>().unwrap();
        let (tag, dei, pcp, typ) = match typ {
            t if t == (EthTyp::EthTypVLAN as u16) => {
                let tag_and_pcp = bytes.read_u16::<BigEndian>().unwrap();
                let tag = tag_and_pcp & 0xfff;
                let dei = (tag_and_pcp & 0x1000) > 0;
                let pcp = tag_and_pcp >> 13;
                let typ = bytes.read_u16::<BigEndian>().unwrap();
                (Some(tag), dei, pcp as u8, typ)
            }
            _ => (None, false, 0x0, typ),
        };
        let nw_header = match typ {
            t if t == (EthTyp::EthTypIP as u16) => {
                let ip = Ip::parse(&mut bytes);
                if ip.is_some() {
                    Nw::Ip(ip.unwrap())
                } else {
                    Nw::Unparsable(typ, bytes.fill_buf().unwrap().to_vec())
                }
            }
            t if t == (EthTyp::EthTypARP as u16) => {
                let arp = Arp::parse(&mut bytes);
                if arp.is_some() {
                    Nw::Arp(arp.unwrap())
                } else {
                    Nw::Unparsable(typ, bytes.fill_buf().unwrap().to_vec())
                }
            }
            _ => Nw::Unparsable(typ, bytes.fill_buf().unwrap().to_vec()),
        };
        Packet {
            dl_src: mac_of_bytes(src),
            dl_dst: mac_of_bytes(dst),
            dl_vlan: tag,
            dl_vlan_dei: dei,
            dl_vlan_pcp: pcp,
            nw: nw_header,
        }
    }
}
<commit_msg>Derive debug for some types<commit_after>use std::io::{BufRead, Cursor, Read};
use std::mem::size_of;
use byteorder::{BigEndian, ReadBytesExt};

use bits::test_bit;

pub fn bytes_of_mac(addr: u64) -> [u8; 6] {
    let mut arr = [0; 6];
    for i in 0..6 {
        arr[i] = ((addr >> (8 * i)) & 0xff) as u8;
    }
    arr
}

pub fn mac_of_bytes(addr: [u8; 6]) -> u64 {
    fn byte(u: &[u8; 6], i: usize) -> u64 {
        u[i] as u64
    };
    (byte(&addr, 0) << 8 * 5) | (byte(&addr, 1) << 8 * 4) | (byte(&addr, 2) << 8 * 3) |
    (byte(&addr, 3) << 8 * 2) | (byte(&addr, 4) << 8 * 1) | (byte(&addr, 5))
}

/// TCP Header flags.
#[derive(Debug)]
pub struct TcpFlags {
    /// ECN-nonce concealment protection.
    pub ns: bool,
    /// Congestion window reduced.
    pub cwr: bool,
    /// ECN-Echo.
    pub ece: bool,
    /// Indicates the Urgent pointer field is significant.
    pub urg: bool,
    /// Indicates that the Acknowledgment field is significant.
    pub ack: bool,
    /// Asks to push the buffered data to the receiving application.
    pub psh: bool,
    /// Reset the connection.
    pub rst: bool,
    /// Synchronize sequence numbers.
    pub syn: bool,
    /// No more data from sender.
    pub fin: bool,
}

impl TcpFlags {
    fn of_int(d: u16) -> TcpFlags {
        TcpFlags {
            ns: test_bit(0, d as u64),
            cwr: test_bit(1, d as u64),
            ece: test_bit(2, d as u64),
            urg: test_bit(3, d as u64),
            ack: test_bit(4, d as u64),
            psh: test_bit(5, d as u64),
            rst: test_bit(6, d as u64),
            syn: test_bit(7, d as u64),
            fin: test_bit(8, d as u64),
        }
    }
}

/// TCP frame of a packet.
#[derive(Debug)]
pub struct Tcp {
    pub src: u16,
    pub dst: u16,
    pub seq: u32,
    pub ack: u32,
    pub offset: u8,
    pub flags: TcpFlags,
    pub window: u16,
    pub chksum: u16,
    pub urgent: u16,
    pub payload: Vec<u8>,
}

#[repr(packed)]
struct TcpNet(u16, u16, u32, u32, u16, u16, u16, u16);

impl Tcp {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Tcp> {
        if bytes.get_ref().len() < size_of::<TcpNet>() {
            return None;
        }
        let src = bytes.read_u16::<BigEndian>().unwrap();
        let dst = bytes.read_u16::<BigEndian>().unwrap();
        let seq = bytes.read_u32::<BigEndian>().unwrap();
        let ack = bytes.read_u32::<BigEndian>().unwrap();
        let offset = bytes.read_u16::<BigEndian>().unwrap();
        let flags = TcpFlags::of_int(offset);
        let offset = (offset >> 12) as u8 & 0x0f;
        let window = bytes.read_u16::<BigEndian>().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let urgent = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Tcp {
            src: src,
            dst: dst,
            seq: seq,
            ack: ack,
            offset: offset,
            flags: flags,
            window: window,
            chksum: chksum,
            urgent: urgent,
            payload: payload,
        })
    }
}

/// UDP frame of a packet.
#[derive(Debug)]
pub struct Udp {
    pub src: u16,
    pub dst: u16,
    pub chksum: u16,
    pub payload: Vec<u8>,
}

impl Udp {
    fn size_of() -> usize {
        8
    }

    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Udp> {
        if bytes.get_ref().len() < Self::size_of() {
            return None;
        }
        let src = bytes.read_u16::<BigEndian>().unwrap();
        let dst = bytes.read_u16::<BigEndian>().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Udp {
            src: src,
            dst: dst,
            chksum: chksum,
            payload: payload,
        })
    }
}

/// ICMP frame of a packet.
#[derive(Debug)]
pub struct Icmp {
    pub typ: u8,
    pub code: u8,
    pub chksum: u16,
    pub payload: Vec<u8>,
}

impl Icmp {
    fn size_of() -> usize {
        4
    }

    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Icmp> {
        if bytes.get_ref().len() < Self::size_of() {
            return None;
        }
        let typ = bytes.read_u8().unwrap();
        let code = bytes.read_u8().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Icmp {
            typ: typ,
            code: code,
            chksum: chksum,
            payload: payload,
        })
    }
}

/// Represents packets at the transport protocol level, which are encapsulated
/// within the IPv4 payload. At present, we only support TCP, UDP, and ICMP
/// explicitly; otherwise, the raw bytes and IPv4 protocol number are provided.
#[derive(Debug)]
pub enum Tp {
    Tcp(Tcp),
    Udp(Udp),
    Icmp(Icmp),
    Unparsable(u8, Vec<u8>),
}

/// The type of IPv4 flags.
#[derive(Debug)]
pub struct Flags {
    pub dont_fragment: bool,
    pub more_fragments: bool,
}

impl Flags {
    fn of_int(flags: u32) -> Flags {
        Flags {
            dont_fragment: test_bit(1, flags as u64),
            more_fragments: test_bit(2, flags as u64),
        }
    }
}

/// IPv4 frame of a packet.
#[derive(Debug)]
pub struct Ip {
    pub tos: u8,
    pub ident: u16,
    pub flags: Flags,
    pub frag: u16,
    pub ttl: u8,
    pub chksum: u16,
    pub src: u32,
    pub dst: u32,
    pub options: Vec<u8>,
    pub tp: Tp,
}

#[repr(u8)]
enum IpProto {
    IpICMP = 0x01,
    IpTCP = 0x06,
    IpUDP = 0x11,
}

#[repr(packed)]
struct IpNet(u8, u8, u16, u16, u16, u8, u8, u16, u32, u32);

impl Ip {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Ip> {
        if bytes.get_ref().len() < size_of::<IpNet>() {
            return None;
        }
        let vhl = bytes.read_u8().unwrap();
        if (vhl >> 4) != 4 {
            return None;
        }
        let ihl = vhl & 0x0f;
        let tos = bytes.read_u8().unwrap();
        bytes.consume(2);
        let ident = bytes.read_u16::<BigEndian>().unwrap();
        let frag = bytes.read_u16::<BigEndian>().unwrap();
        let flags = Flags::of_int((frag as u32) >> 13);
        let ttl = bytes.read_u8().unwrap();
        let proto = bytes.read_u8().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let src = bytes.read_u32::<BigEndian>().unwrap();
        let dst = bytes.read_u32::<BigEndian>().unwrap();
        let options_len = (ihl * 4) as usize - size_of::<IpNet>();
        let mut options = vec![0; options_len];
        bytes.read_exact(&mut options).unwrap();
        let tp = match proto {
            t if t == (IpProto::IpICMP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let icmp = Icmp::parse(bytes);
                if icmp.is_some() {
                    Tp::Icmp(icmp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            t if t == (IpProto::IpTCP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let tcp = Tcp::parse(bytes);
                if tcp.is_some() {
                    Tp::Tcp(tcp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            t if t == (IpProto::IpUDP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let udp = Udp::parse(bytes);
                if udp.is_some() {
                    Tp::Udp(udp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            _ => Tp::Unparsable(proto, bytes.get_ref().clone()),
        };
        Some(Ip {
            tos: tos,
            ident: ident,
            flags: flags,
            frag: frag,
            ttl: ttl,
            chksum: chksum,
            src: src,
            dst: dst,
            options: options,
            tp: tp,
        })
    }
}

/// Address resolution protocol (ARP) packet payload.
#[derive(Debug)]
pub enum Arp {
    Query(u64, u32, u32),
    Reply(u64, u32, u64, u32),
}

#[repr(packed)]
struct ArpNet(u16, u16, u8, u8, u16, [u8; 6], u32, [u8; 6], u32);

impl Arp {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Arp> {
        if bytes.get_ref().len() < size_of::<ArpNet>() {
            return None;
        }
        bytes.consume(6);
        let oper = bytes.read_u16::<BigEndian>().unwrap();
        let mut sha: [u8; 6] = [0; 6];
        for i in 0..6 {
            sha[i] = bytes.read_u8().unwrap();
        }
        let spa = bytes.read_u32::<BigEndian>().unwrap();
        let mut tha: [u8; 6] = [0; 6];
        for i in 0..6 {
            tha[i] = bytes.read_u8().unwrap();
        }
        let tpa = bytes.read_u32::<BigEndian>().unwrap();
        match oper {
            0x0001 => Some(Arp::Query(mac_of_bytes(sha), spa, tpa)),
            0x0002 => Some(Arp::Reply(mac_of_bytes(sha), spa, mac_of_bytes(tha), tpa)),
            _ => None,
        }
    }
}

/// Represents a packet at the network protocol level.
#[derive(Debug)]
pub enum Nw {
    Ip(Ip),
    Arp(Arp),
    Unparsable(u16, Vec<u8>),
}

/// Represents a packet at the ethernet protocol level.
#[derive(Debug)]
pub struct Packet {
    pub dl_src: u64,
    pub dl_dst: u64,
    pub dl_vlan: Option<u16>,
    pub dl_vlan_dei: bool,
    pub dl_vlan_pcp: u8,
    pub nw: Nw,
}

#[repr(u16)]
enum EthTyp {
    EthTypIP = 0x0800,
    EthTypARP = 0x0806,
    EthTypVLAN = 0x8100,
}

impl Packet {
    pub fn parse(buf: &[u8]) -> Packet {
        let mut bytes = Cursor::new(buf.to_vec());
        let mut dst: [u8; 6] = [0; 6];
        let mut src: [u8; 6] = [0; 6];
        for i in 0..6 {
            dst[i] = bytes.read_u8().unwrap();
        }
        for i in 0..6 {
            src[i] = bytes.read_u8().unwrap();
        }
        let typ = bytes.read_u16::<BigEndian>().unwrap();
        let (tag, dei, pcp, typ) = match typ {
            t if t == (EthTyp::EthTypVLAN as u16) => {
                let tag_and_pcp = bytes.read_u16::<BigEndian>().unwrap();
                let tag = tag_and_pcp & 0xfff;
                let dei = (tag_and_pcp & 0x1000) > 0;
                let pcp = tag_and_pcp >> 13;
                let typ = bytes.read_u16::<BigEndian>().unwrap();
                (Some(tag), dei, pcp as u8, typ)
            }
            _ => (None, false, 0x0, typ),
        };
        let nw_header = match typ {
            t if t == (EthTyp::EthTypIP as u16) => {
                let ip = Ip::parse(&mut bytes);
                if ip.is_some() {
                    Nw::Ip(ip.unwrap())
                } else {
                    Nw::Unparsable(typ, bytes.fill_buf().unwrap().to_vec())
                }
            }
            t if t == (EthTyp::EthTypARP as u16) => {
                let arp = Arp::parse(&mut bytes);
                if arp.is_some() {
                    Nw::Arp(arp.unwrap())
                } else {
                    Nw::Unparsable(typ, bytes.fill_buf().unwrap().to_vec())
                }
            }
            _ => Nw::Unparsable(typ, bytes.fill_buf().unwrap().to_vec()),
        };
        Packet {
            dl_src: mac_of_bytes(src),
            dl_dst: mac_of_bytes(dst),
            dl_vlan: tag,
            dl_vlan_dei: dei,
            dl_vlan_pcp: pcp,
            nw: nw_header,
        }
    }
}
<|endoftext|>


<commit_before>use std::str::from_utf8;
use std::result::Result;
use nom::*;

/*
 * Core structs
 */

#[derive(Debug,PartialEq,Eq)]
pub struct TarEntry<'a> {
    pub header:   PosixHeader<'a>,
    pub contents: &'a [u8]
}

#[derive(Debug,PartialEq,Eq)]
pub struct PosixHeader<'a> {
    pub name:     &'a str,
    pub mode:     &'a str,
    pub uid:      u64,
    pub gid:      u64,
    pub size:     u64,
    pub mtime:    u64,
    pub chksum:   &'a str,
    pub typeflag: TypeFlag,
    pub linkname: &'a str,
    pub ustar:    ExtraHeader<'a>
}

/* TODO: support vendor specific */
#[derive(Debug,PartialEq,Eq)]
pub enum TypeFlag {
    NormalFile,
    HardLink,
    SymbolicLink,
    CharacterSpecial,
    BlockSpecial,
    Directory,
    FIFO,
    ContiguousFile,
    PaxInterexchangeFormat,
    PaxExtendedAttributes,
    VendorSpecific
}

#[derive(Debug,PartialEq,Eq)]
pub enum ExtraHeader<'a> {
    UStar(UStarHeader<'a>),
    Padding
}

#[derive(Debug,PartialEq,Eq)]
pub struct UStarHeader<'a> {
    pub magic:    &'a str,
    pub version:  &'a str,
    pub uname:    &'a str,
    pub gname:    &'a str,
    pub devmajor: u64,
    pub devminor: u64,
    pub extra:    UStarExtraHeader<'a>
}

#[derive(Debug,PartialEq,Eq)]
pub enum UStarExtraHeader<'a> {
    PosixUStar(PosixUStarHeader<'a>),
    Pax(PaxHeader<'a>)
}

#[derive(Debug,PartialEq,Eq)]
pub struct PosixUStarHeader<'a> {
    pub prefix: &'a str
}

#[derive(Debug,PartialEq,Eq)]
pub struct PaxHeader<'a> {
    pub atime:      u64,
    pub ctime:      u64,
    pub offset:     u64,
    pub longnames:  &'a str,
    pub sparses:    Vec<Sparse>,
    pub isextended: bool,
    pub realsize:   u64,
}

#[derive(Debug,PartialEq,Eq)]
pub struct Sparse {
    pub offset:   u64,
    pub numbytes: u64
}

#[derive(Debug,PartialEq,Eq)]
pub struct Padding;

/*
 * Useful macros
 */

macro_rules! take_str_eat_garbage (
    ( $i:expr, $size:expr ) => (
        chain!($i,
            s: map_res!(take_until!("\0"), from_utf8)                  ~
            length: expr_opt!({($size as usize).checked_sub(s.len())}) ~
            take!(length),
            ||{
                s
            }
        )
    );
);

named!(parse_str4<&[u8], &str>, take_str_eat_garbage!(4));
named!(parse_str8<&[u8], &str>, take_str_eat_garbage!(8));
named!(parse_str32<&[u8], &str>, take_str_eat_garbage!(32));
named!(parse_str100<&[u8], &str>, take_str_eat_garbage!(100));
named!(parse_str155<&[u8], &str>, take_str_eat_garbage!(155));

macro_rules! take_until_expr_with_limit_consume(
  ($i:expr, $submac:ident!( $($args:tt)* ), $stop: expr, $limit: expr) => (
    {
      let mut begin = 0;
      let mut remaining = $i.len();
      let mut res = Vec::new();
      let mut cnt = 0;
      let mut err = false;
      let mut append = true;
      loop {
        match $submac!(&$i[begin..], $($args)*) {
          IResult::Done(i,o) => {
            if append {
              if $stop(&o) {
                append = false;
              } else {
                res.push(o);
              }
            }
            begin += remaining - i.len();
            remaining = i.len();
            cnt = cnt + 1;
            if cnt == $limit {
              break
            }
          },
          IResult::Error(_)  => {
            err = true;
            break;
          },
          IResult::Incomplete(_) => {
            break;
          }
        }
      }
      if err {
        IResult::Error(Err::Position(ErrorKind::TakeUntil,$i))
      } else if cnt == $limit {
        IResult::Done(&$i[begin..], res)
      } else {
        IResult::Incomplete(Needed::Unknown)
      }
    }
  );
  ($i:expr, $f:expr, $stop: expr, $limit: expr) => (
    take_until_expr_with_limit_consume!($i, call!($f), $stop, $limit);
  );
);

/*
 * Octal string parsing
 */

pub fn octal_to_u64(s: &str) -> Result<u64, &'static str> {
    let mut u = 0;

    for c in s.chars() {
        if c < '0' || c > '7' {
            return Err("invalid octal string received");
        }
        u *= 8;
        u += (c as u64) - ('0' as u64);
    }

    Ok(u)
}

fn parse_octal(i: &[u8], n: usize) -> IResult<&[u8], u64> {
    map_res!(i, take_str_eat_garbage!(n), octal_to_u64)
}

named!(parse_octal8<&[u8], u64>, apply!(parse_octal, 8));
named!(parse_octal12<&[u8], u64>, apply!(parse_octal, 12));

/*
 * TypeFlag parsing
 */

fn char_to_type_flag(c: char) -> TypeFlag {
    match c {
        '0' | '\0' => TypeFlag::NormalFile,
        '1' => TypeFlag::HardLink,
        '2' => TypeFlag::SymbolicLink,
        '3' => TypeFlag::CharacterSpecial,
        '4' => TypeFlag::BlockSpecial,
        '5' => TypeFlag::Directory,
        '6' => TypeFlag::FIFO,
        '7' => TypeFlag::ContiguousFile,
        'g' => TypeFlag::PaxInterexchangeFormat,
        'x' => TypeFlag::PaxExtendedAttributes,
        'A' ... 'Z' => TypeFlag::VendorSpecific,
        _ => TypeFlag::NormalFile
    }
}

fn bytes_to_type_flag(i: &[u8]) -> TypeFlag {
    char_to_type_flag(i[0] as char)
}

named!(parse_type_flag<&[u8], TypeFlag>, map!(take!(1), bytes_to_type_flag));

/*
 * Sparse parsing
 */

named!(parse_one_sparse<&[u8], Sparse>, map!(pair!(parse_octal12, parse_octal12), |(offset, numbytes)| Sparse { offset: offset, numbytes: numbytes }));

fn parse_sparses_with_limit(i: &[u8], limit: usize) -> IResult<&[u8], Vec<Sparse>> {
    take_until_expr_with_limit_consume!(i, parse_one_sparse, |s: &Sparse| s.offset == 0 && s.numbytes == 0, limit)
}

fn add_to_vec<'a, 'b>(sparses: &'a mut Vec<Sparse>, extra: &'b mut Vec<Sparse>) -> &'a mut Vec<Sparse> {
    while sparses.len() != 0 {
      extra.pop().map(|s| sparses.push(s));
    }
    sparses
}

fn parse_extra_sparses<'a, 'b>(i: &'a [u8], isextended: bool, sparses: &'b mut Vec<Sparse>) -> IResult<&'a [u8], &'b mut Vec<Sparse>> {
    if isextended {
        chain!(i,
            mut sps: apply!(parse_sparses_with_limit, 21) ~
            extended:      parse_bool                     ~
            take!(7) /* padding to 512 */                 ~
            extra_sparses: apply!(parse_extra_sparses, extended, add_to_vec(sparses, &mut sps)),
            ||{
                extra_sparses
            }
        )
    } else {
        IResult::Done(i, sparses)
    }
}

fn parse_pax_extra_sparses<'a, 'b>(i: &'a [u8], h: &'b mut PaxHeader) -> IResult<&'a [u8], &'b mut Vec<Sparse>> {
    parse_extra_sparses(i, h.isextended, &mut h.sparses)
}

/*
 * Boolean parsing
 */

fn to_bool(i: &[u8]) -> bool {
    i[0] != 0
}

named!(parse_bool<&[u8], bool>, map!(take!(1), to_bool));

/*
 * UStar PAX extended parsing
 */

fn parse_ustar00_extra_pax(i: &[u8]) -> IResult<&[u8], PaxHeader> {
    chain!(i,
        atime:       parse_octal12                       ~
        ctime:       parse_octal12                       ~
        offset:      parse_octal12                       ~
        longnames:   parse_str4                          ~
        take!(1)                                         ~
        sparses:     apply!(parse_sparses_with_limit, 4) ~
        isextended:  parse_bool                          ~
        realsize:    parse_octal12                       ~
        take!(17), /* padding to 512 */
        ||{
            PaxHeader {
                atime:         atime,
                ctime:         ctime,
                offset:        offset,
                longnames:     longnames,
                sparses:       sparses,
                isextended:    isextended,
                realsize:      realsize,
            }
        }
    )
}

/*
 * UStar Posix parsing
 */

named!(parse_ustar00_extra_posix<&[u8], UStarExtraHeader>, map!(pair!(parse_str155, take!(12)), |(prefix, _)| UStarExtraHeader::PosixUStar(PosixUStarHeader { prefix: prefix })));

fn parse_ustar00_extra<'a, 'b>(i: &'a [u8], flag: &'b TypeFlag) -> IResult<&'a [u8], UStarExtraHeader<'a>> {
    match *flag {
        TypeFlag::PaxInterexchangeFormat => {
            chain!(i,
                mut header: parse_ustar00_extra_pax ~
                apply!(parse_pax_extra_sparses, &mut header),
                ||{
                    UStarExtraHeader::Pax(header)
                }
            )
        },
        _ => parse_ustar00_extra_posix(i)
    }
}

fn parse_ustar00<'a, 'b>(i: &'a [u8], flag: &'b TypeFlag) -> IResult<&'a [u8], ExtraHeader<'a>> {
    chain!(i,
        tag!("00")             ~
        uname:    parse_str32  ~
        gname:    parse_str32  ~
        devmajor: parse_octal8 ~
        devminor: parse_octal8 ~
        extra:    apply!(parse_ustar00_extra, flag),
        ||{
            ExtraHeader::UStar(UStarHeader {
                magic:    "ustar\0",
                version:  "00",
                uname:    uname,
                gname:    gname,
                devmajor: devmajor,
                devminor: devminor,
                extra:    extra
            })
        }
    )
}

fn parse_ustar<'a, 'b>(i: &'a [u8], flag: &'b TypeFlag) -> IResult<&'a [u8], ExtraHeader<'a>> {
    chain!(i,
        tag!("ustar\0") ~
        ustar: apply!(parse_ustar00, flag),
        ||{
            ustar
        }
    )
}

/*
 * Posix tar archive header parsing
 */

fn parse_posix(i: &[u8]) -> IResult<&[u8], ExtraHeader> {
    chain!(i,
        take!(255), /* padding to 512 */
        ||{
            ExtraHeader::Padding
        }
    )
}

fn parse_header(i: &[u8]) -> IResult<&[u8], PosixHeader> {
    chain!(i,
        name:     parse_str100    ~
        mode:     parse_str8      ~
        uid:      parse_octal8    ~
        gid:      parse_octal8    ~
        size:     parse_octal12   ~
        mtime:    parse_octal12   ~
        chksum:   parse_str8      ~
        typeflag: parse_type_flag ~
        linkname: parse_str100    ~
        ustar:    alt!(apply!(parse_ustar, &typeflag) | parse_posix),
        ||{
            PosixHeader {
                name:     name,
                mode:     mode,
                uid:      uid,
                gid:      gid,
                size:     size,
                mtime:    mtime,
                chksum:   chksum,
                typeflag: typeflag,
                linkname: linkname,
                ustar:    ustar
            }
        }
    )
}

/*
 * Contents parsing
 */

fn parse_contents(i: &[u8], size: u64) -> IResult<&[u8], &[u8]> {
    let trailing = size % 512;
    let padding = match trailing {
        0 => 0,
        t => 512 - t
    };
    chain!(i,
        contents: take!(size as usize) ~
        take!(padding as usize),
        ||{
            contents
        }
    )
}

/*
 * Tar entry header + contents parsing
 */

fn parse_entry(i: &[u8]) -> IResult<&[u8], TarEntry> {
    chain!(i,
        header:   parse_header ~
        contents: apply!(parse_contents, header.size),
        ||{
            TarEntry {
                header: header,
                contents: contents
            }
        }
    )
}

/*
 * Tar archive parsing
 */

fn filter_entries(entries: Vec<TarEntry>) -> Vec<TarEntry> {
    /* Filter out empty entries */
    entries.into_iter().filter(|e| e.header.name != "").collect::<Vec<TarEntry>>()
}

pub fn parse_tar(i: &[u8]) -> IResult<&[u8], Vec<TarEntry>> {
    chain!(i,
        entries: map!(many0!(parse_entry), filter_entries) ~
        eof,
        ||{
            entries
        }
    )
}

/*
 * Tests
 */

#[cfg(test)]
mod tests {
    use super::*;
    use std::str::from_utf8;
    use nom::IResult;

    #[test]
    fn octal_to_u64_ok_test() {
        assert_eq!(octal_to_u64("756"), Ok(494));
        assert_eq!(octal_to_u64(""), Ok(0));
    }

    #[test]
    fn octal_to_u64_error_test() {
        assert_eq!(octal_to_u64("1238"), Err("invalid octal string received"));
        assert_eq!(octal_to_u64("a"), Err("invalid octal string received"));
        assert_eq!(octal_to_u64("A"), Err("invalid octal string received"));
    }

    #[test]
    fn take_str_eat_garbage_test() {
        let s = b"foobar\0\0\0\0baz";
        let baz = b"baz";
        assert_eq!(take_str_eat_garbage!(&s[..], 10), IResult::Done(&baz[..], "foobar"));
    }
}
<commit_msg>simplify parse_posix<commit_after>use std::str::from_utf8;
use std::result::Result;
use nom::*;

/*
 * Core structs
 */

#[derive(Debug,PartialEq,Eq)]
pub struct TarEntry<'a> {
    pub header:   PosixHeader<'a>,
    pub contents: &'a [u8]
}

#[derive(Debug,PartialEq,Eq)]
pub struct PosixHeader<'a> {
    pub name:     &'a str,
    pub mode:     &'a str,
    pub uid:      u64,
    pub gid:      u64,
    pub size:     u64,
    pub mtime:    u64,
    pub chksum:   &'a str,
    pub typeflag: TypeFlag,
    pub linkname: &'a str,
    pub ustar:    ExtraHeader<'a>
}

/* TODO: support vendor specific */
#[derive(Debug,PartialEq,Eq)]
pub enum TypeFlag {
    NormalFile,
    HardLink,
    SymbolicLink,
    CharacterSpecial,
    BlockSpecial,
    Directory,
    FIFO,
    ContiguousFile,
    PaxInterexchangeFormat,
    PaxExtendedAttributes,
    VendorSpecific
}

#[derive(Debug,PartialEq,Eq)]
pub enum ExtraHeader<'a> {
    UStar(UStarHeader<'a>),
    Padding
}

#[derive(Debug,PartialEq,Eq)]
pub struct UStarHeader<'a> {
    pub magic:    &'a str,
    pub version:  &'a str,
    pub uname:    &'a str,
    pub gname:    &'a str,
    pub devmajor: u64,
    pub devminor: u64,
    pub extra:    UStarExtraHeader<'a>
}

#[derive(Debug,PartialEq,Eq)]
pub enum UStarExtraHeader<'a> {
    PosixUStar(PosixUStarHeader<'a>),
    Pax(PaxHeader<'a>)
}

#[derive(Debug,PartialEq,Eq)]
pub struct PosixUStarHeader<'a> {
    pub prefix: &'a str
}

#[derive(Debug,PartialEq,Eq)]
pub struct PaxHeader<'a> {
    pub atime:      u64,
    pub ctime:      u64,
    pub offset:     u64,
    pub longnames:  &'a str,
    pub sparses:    Vec<Sparse>,
    pub isextended: bool,
    pub realsize:   u64,
}

#[derive(Debug,PartialEq,Eq)]
pub struct Sparse {
    pub offset:   u64,
    pub numbytes: u64
}

#[derive(Debug,PartialEq,Eq)]
pub struct Padding;

/*
 * Useful macros
 */

macro_rules! take_str_eat_garbage (
    ( $i:expr, $size:expr ) => (
        chain!($i,
            s: map_res!(take_until!("\0"), from_utf8)                  ~
            length: expr_opt!({($size as usize).checked_sub(s.len())}) ~
            take!(length),
            ||{
                s
            }
        )
    );
);

named!(parse_str4<&[u8], &str>, take_str_eat_garbage!(4));
named!(parse_str8<&[u8], &str>, take_str_eat_garbage!(8));
named!(parse_str32<&[u8], &str>, take_str_eat_garbage!(32));
named!(parse_str100<&[u8], &str>, take_str_eat_garbage!(100));
named!(parse_str155<&[u8], &str>, take_str_eat_garbage!(155));

macro_rules! take_until_expr_with_limit_consume(
  ($i:expr, $submac:ident!( $($args:tt)* ), $stop: expr, $limit: expr) => (
    {
      let mut begin = 0;
      let mut remaining = $i.len();
      let mut res = Vec::new();
      let mut cnt = 0;
      let mut err = false;
      let mut append = true;
      loop {
        match $submac!(&$i[begin..], $($args)*) {
          IResult::Done(i,o) => {
            if append {
              if $stop(&o) {
                append = false;
              } else {
                res.push(o);
              }
            }
            begin += remaining - i.len();
            remaining = i.len();
            cnt = cnt + 1;
            if cnt == $limit {
              break
            }
          },
          IResult::Error(_)  => {
            err = true;
            break;
          },
          IResult::Incomplete(_) => {
            break;
          }
        }
      }
      if err {
        IResult::Error(Err::Position(ErrorKind::TakeUntil,$i))
      } else if cnt == $limit {
        IResult::Done(&$i[begin..], res)
      } else {
        IResult::Incomplete(Needed::Unknown)
      }
    }
  );
  ($i:expr, $f:expr, $stop: expr, $limit: expr) => (
    take_until_expr_with_limit_consume!($i, call!($f), $stop, $limit);
  );
);

/*
 * Octal string parsing
 */

pub fn octal_to_u64(s: &str) -> Result<u64, &'static str> {
    let mut u = 0;

    for c in s.chars() {
        if c < '0' || c > '7' {
            return Err("invalid octal string received");
        }
        u *= 8;
        u += (c as u64) - ('0' as u64);
    }

    Ok(u)
}

fn parse_octal(i: &[u8], n: usize) -> IResult<&[u8], u64> {
    map_res!(i, take_str_eat_garbage!(n), octal_to_u64)
}

named!(parse_octal8<&[u8], u64>, apply!(parse_octal, 8));
named!(parse_octal12<&[u8], u64>, apply!(parse_octal, 12));

/*
 * TypeFlag parsing
 */

fn char_to_type_flag(c: char) -> TypeFlag {
    match c {
        '0' | '\0' => TypeFlag::NormalFile,
        '1' => TypeFlag::HardLink,
        '2' => TypeFlag::SymbolicLink,
        '3' => TypeFlag::CharacterSpecial,
        '4' => TypeFlag::BlockSpecial,
        '5' => TypeFlag::Directory,
        '6' => TypeFlag::FIFO,
        '7' => TypeFlag::ContiguousFile,
        'g' => TypeFlag::PaxInterexchangeFormat,
        'x' => TypeFlag::PaxExtendedAttributes,
        'A' ... 'Z' => TypeFlag::VendorSpecific,
        _ => TypeFlag::NormalFile
    }
}

fn bytes_to_type_flag(i: &[u8]) -> TypeFlag {
    char_to_type_flag(i[0] as char)
}

named!(parse_type_flag<&[u8], TypeFlag>, map!(take!(1), bytes_to_type_flag));

/*
 * Sparse parsing
 */

named!(parse_one_sparse<&[u8], Sparse>, map!(pair!(parse_octal12, parse_octal12), |(offset, numbytes)| Sparse { offset: offset, numbytes: numbytes }));

fn parse_sparses_with_limit(i: &[u8], limit: usize) -> IResult<&[u8], Vec<Sparse>> {
    take_until_expr_with_limit_consume!(i, parse_one_sparse, |s: &Sparse| s.offset == 0 && s.numbytes == 0, limit)
}

fn add_to_vec<'a, 'b>(sparses: &'a mut Vec<Sparse>, extra: &'b mut Vec<Sparse>) -> &'a mut Vec<Sparse> {
    while sparses.len() != 0 {
      extra.pop().map(|s| sparses.push(s));
    }
    sparses
}

fn parse_extra_sparses<'a, 'b>(i: &'a [u8], isextended: bool, sparses: &'b mut Vec<Sparse>) -> IResult<&'a [u8], &'b mut Vec<Sparse>> {
    if isextended {
        chain!(i,
            mut sps: apply!(parse_sparses_with_limit, 21) ~
            extended:      parse_bool                     ~
            take!(7) /* padding to 512 */                 ~
            extra_sparses: apply!(parse_extra_sparses, extended, add_to_vec(sparses, &mut sps)),
            ||{
                extra_sparses
            }
        )
    } else {
        IResult::Done(i, sparses)
    }
}

fn parse_pax_extra_sparses<'a, 'b>(i: &'a [u8], h: &'b mut PaxHeader) -> IResult<&'a [u8], &'b mut Vec<Sparse>> {
    parse_extra_sparses(i, h.isextended, &mut h.sparses)
}

/*
 * Boolean parsing
 */

fn to_bool(i: &[u8]) -> bool {
    i[0] != 0
}

named!(parse_bool<&[u8], bool>, map!(take!(1), to_bool));

/*
 * UStar PAX extended parsing
 */

fn parse_ustar00_extra_pax(i: &[u8]) -> IResult<&[u8], PaxHeader> {
    chain!(i,
        atime:       parse_octal12                       ~
        ctime:       parse_octal12                       ~
        offset:      parse_octal12                       ~
        longnames:   parse_str4                          ~
        take!(1)                                         ~
        sparses:     apply!(parse_sparses_with_limit, 4) ~
        isextended:  parse_bool                          ~
        realsize:    parse_octal12                       ~
        take!(17), /* padding to 512 */
        ||{
            PaxHeader {
                atime:         atime,
                ctime:         ctime,
                offset:        offset,
                longnames:     longnames,
                sparses:       sparses,
                isextended:    isextended,
                realsize:      realsize,
            }
        }
    )
}

/*
 * UStar Posix parsing
 */

named!(parse_ustar00_extra_posix<&[u8], UStarExtraHeader>, map!(pair!(parse_str155, take!(12)), |(prefix, _)| UStarExtraHeader::PosixUStar(PosixUStarHeader { prefix: prefix })));

fn parse_ustar00_extra<'a, 'b>(i: &'a [u8], flag: &'b TypeFlag) -> IResult<&'a [u8], UStarExtraHeader<'a>> {
    match *flag {
        TypeFlag::PaxInterexchangeFormat => {
            chain!(i,
                mut header: parse_ustar00_extra_pax ~
                apply!(parse_pax_extra_sparses, &mut header),
                ||{
                    UStarExtraHeader::Pax(header)
                }
            )
        },
        _ => parse_ustar00_extra_posix(i)
    }
}

fn parse_ustar00<'a, 'b>(i: &'a [u8], flag: &'b TypeFlag) -> IResult<&'a [u8], ExtraHeader<'a>> {
    chain!(i,
        tag!("00")             ~
        uname:    parse_str32  ~
        gname:    parse_str32  ~
        devmajor: parse_octal8 ~
        devminor: parse_octal8 ~
        extra:    apply!(parse_ustar00_extra, flag),
        ||{
            ExtraHeader::UStar(UStarHeader {
                magic:    "ustar\0",
                version:  "00",
                uname:    uname,
                gname:    gname,
                devmajor: devmajor,
                devminor: devminor,
                extra:    extra
            })
        }
    )
}

fn parse_ustar<'a, 'b>(i: &'a [u8], flag: &'b TypeFlag) -> IResult<&'a [u8], ExtraHeader<'a>> {
    chain!(i,
        tag!("ustar\0") ~
        ustar: apply!(parse_ustar00, flag),
        ||{
            ustar
        }
    )
}

/*
 * Posix tar archive header parsing
 */

fn parse_posix(i: &[u8]) -> IResult<&[u8], ExtraHeader> {
    map!(i, take!(255), |_| ExtraHeader::Padding) /* padding to 512 */
}

fn parse_header(i: &[u8]) -> IResult<&[u8], PosixHeader> {
    chain!(i,
        name:     parse_str100    ~
        mode:     parse_str8      ~
        uid:      parse_octal8    ~
        gid:      parse_octal8    ~
        size:     parse_octal12   ~
        mtime:    parse_octal12   ~
        chksum:   parse_str8      ~
        typeflag: parse_type_flag ~
        linkname: parse_str100    ~
        ustar:    alt!(apply!(parse_ustar, &typeflag) | parse_posix),
        ||{
            PosixHeader {
                name:     name,
                mode:     mode,
                uid:      uid,
                gid:      gid,
                size:     size,
                mtime:    mtime,
                chksum:   chksum,
                typeflag: typeflag,
                linkname: linkname,
                ustar:    ustar
            }
        }
    )
}

/*
 * Contents parsing
 */

fn parse_contents(i: &[u8], size: u64) -> IResult<&[u8], &[u8]> {
    let trailing = size % 512;
    let padding = match trailing {
        0 => 0,
        t => 512 - t
    };
    chain!(i,
        contents: take!(size as usize) ~
        take!(padding as usize),
        ||{
            contents
        }
    )
}

/*
 * Tar entry header + contents parsing
 */

fn parse_entry(i: &[u8]) -> IResult<&[u8], TarEntry> {
    chain!(i,
        header:   parse_header ~
        contents: apply!(parse_contents, header.size),
        ||{
            TarEntry {
                header: header,
                contents: contents
            }
        }
    )
}

/*
 * Tar archive parsing
 */

fn filter_entries(entries: Vec<TarEntry>) -> Vec<TarEntry> {
    /* Filter out empty entries */
    entries.into_iter().filter(|e| e.header.name != "").collect::<Vec<TarEntry>>()
}

pub fn parse_tar(i: &[u8]) -> IResult<&[u8], Vec<TarEntry>> {
    chain!(i,
        entries: map!(many0!(parse_entry), filter_entries) ~
        eof,
        ||{
            entries
        }
    )
}

/*
 * Tests
 */

#[cfg(test)]
mod tests {
    use super::*;
    use std::str::from_utf8;
    use nom::IResult;

    #[test]
    fn octal_to_u64_ok_test() {
        assert_eq!(octal_to_u64("756"), Ok(494));
        assert_eq!(octal_to_u64(""), Ok(0));
    }

    #[test]
    fn octal_to_u64_error_test() {
        assert_eq!(octal_to_u64("1238"), Err("invalid octal string received"));
        assert_eq!(octal_to_u64("a"), Err("invalid octal string received"));
        assert_eq!(octal_to_u64("A"), Err("invalid octal string received"));
    }

    #[test]
    fn take_str_eat_garbage_test() {
        let s = b"foobar\0\0\0\0baz";
        let baz = b"baz";
        assert_eq!(take_str_eat_garbage!(&s[..], 10), IResult::Done(&baz[..], "foobar"));
    }
}
<|endoftext|>