<commit_before>//! TODO Documentation

use libc::{c_float, c_int, c_void};

use Output;
use render::Texture;
use wlroots_sys::{wl_shm_format, wlr_backend, wlr_backend_get_egl, wlr_render_ellipse_with_matrix,
                  wlr_render_quad_with_matrix, wlr_render_texture, wlr_render_texture_with_matrix,
                  wlr_renderer, wlr_renderer_begin, wlr_renderer_clear, wlr_renderer_destroy,
                  wlr_renderer_end, wlr_texture_from_pixels, wlr_gles2_renderer_create};

/// A generic interface for rendering to the screen.
///
/// Note that it will technically be possible to have multiple renderers
/// at the same time.
#[derive(Debug)]
pub struct GenericRenderer {
    renderer: *mut wlr_renderer
}

/// The state machine type that allows you to manipulate a screen and
/// its buffer.
///
/// When this structure is dropped it automatically calls wlr_renderer_end
/// and swaps the buffers.
#[derive(Debug)]
pub struct Renderer<'output> {
    renderer: *mut wlr_renderer,
    pub output: &'output mut Output
}

impl GenericRenderer {
    /// Make a gles2 renderer.
    pub(crate) unsafe fn gles2_renderer(backend: *mut wlr_backend) -> Self {
        let egl = wlr_backend_get_egl(backend);
        if egl.is_null() {
            panic!("EGL not available for this backend");
        }
        let renderer = wlr_gles2_renderer_create(egl);
        if renderer.is_null() {
            panic!("Could not construct GLES2 renderer");
        }
        GenericRenderer { renderer }
    }

    /// Make the `Renderer` state machine type.
    ///
    /// This automatically makes the given output the current output.
    pub fn render<'output>(&mut self, output: &'output mut Output) -> Renderer<'output> {
        unsafe {
            output.make_current();
            let (width, height) = output.dimensions();
            wlr_renderer_begin(self.renderer, width, height);
            Renderer { renderer: self.renderer,
                       output }
        }
    }

    /// Create a texture using this renderer.
    pub fn create_texture_from_pixels(&mut self,
                                      format: wl_shm_format,
                                      stride: u32,
                                      width: u32,
                                      height: u32,
                                      data: &[u8])
                                      -> Option<Texture> {
        unsafe {
            create_texture_from_pixels(self.renderer,
                                       format,
                                       stride,
                                       width,
                                       height,
                                       data.as_ptr() as _)
        }
    }

    pub(crate) unsafe fn as_ptr(&self) -> *mut wlr_renderer {
        self.renderer
    }
}

impl Drop for GenericRenderer {
    fn drop(&mut self) {
        unsafe { wlr_renderer_destroy(self.renderer) }
    }
}

impl<'output> Renderer<'output> {
    pub fn clear(&mut self, float: [f32; 4]) {
        unsafe { wlr_renderer_clear(self.renderer, float.as_ptr()) }
    }

    /// Renders the requseted texture.
    pub fn render_texture(&mut self,
                          texture: &Texture,
                          projection: [f32; 9],
                          x: c_int,
                          y: c_int,
                          alpha: c_float)
                          -> bool {
        unsafe {
            wlr_render_texture(self.renderer,
                               texture.as_ptr(),
                               projection.as_ptr(),
                               x,
                               y,
                               alpha)
        }
    }

    /// Renders the requested texture using the provided matrix. A typical texture
    /// rendering goes like so:
    ///
    /// TODO FIXME Show how the typical rendering goes in Rust.
    ///
    /// ```c
    /// struct wlr_renderer *renderer;
    /// struct wlr_texture *texture;
    /// float projection[16];
    /// float matrix[16];
    /// wlr_texture_get_matrix(texture, &matrix, &projection, 123, 321);
    /// wlr_render_texture_with_matrix(renderer, texture, &matrix);
    /// ```
    ///
    /// This will render the texture at <123, 321>.
    pub fn render_texture_with_matrix(&mut self, texture: &Texture, matrix: [f32; 9]) -> bool {
        // TODO FIXME Add alpha as param
        unsafe {
            wlr_render_texture_with_matrix(self.renderer, texture.as_ptr(), matrix.as_ptr(), 1.0)
        }
    }

    /// Renders a solid quad in the specified color.
    pub fn render_colored_quad(&mut self, color: [f32; 4], matrix: [f32; 9]) {
        unsafe { wlr_render_quad_with_matrix(self.renderer, color.as_ptr(), matrix.as_ptr()) }
    }

    /// Renders a solid ellipse in the specified color.
    pub fn render_colored_ellipse(&mut self, color: [f32; 4], matrix: [f32; 9]) {
        unsafe { wlr_render_ellipse_with_matrix(self.renderer, color.as_ptr(), matrix.as_ptr()) }
    }
}

impl<'output> Drop for Renderer<'output> {
    fn drop(&mut self) {
        unsafe {
            wlr_renderer_end(self.renderer);
            // TODO What about damage tracking?
            self.output.swap_buffers(None, None);
        }
    }
}

unsafe fn create_texture_from_pixels(renderer: *mut wlr_renderer,
                                     format: wl_shm_format,
                                     stride: u32,
                                     width: u32,
                                     height: u32,
                                     // TODO Slice of u8? It's a void*, hmm
                                     data: *const c_void)
                                     -> Option<Texture> {
    let texture = wlr_texture_from_pixels(renderer, format, stride, width, height, data);
    if texture.is_null() {
        None
    } else {
        Some(Texture::from_ptr(texture))
    }
}
<commit_msg>Can specify damage in renderer<commit_after>//! TODO Documentation

use std::time::Duration;

use libc::{c_float, c_int, c_void};

use {Output, PixmanRegion};
use render::Texture;
use wlroots_sys::{timespec, wl_shm_format, wlr_backend, wlr_backend_get_egl,
                  wlr_render_ellipse_with_matrix, wlr_render_quad_with_matrix, wlr_render_texture,
                  wlr_render_texture_with_matrix, wlr_renderer, wlr_renderer_begin,
                  wlr_renderer_clear, wlr_renderer_destroy, wlr_renderer_end,
                  wlr_texture_from_pixels, wlr_gles2_renderer_create};

/// A generic interface for rendering to the screen.
///
/// Note that it will technically be possible to have multiple renderers
/// at the same time.
#[derive(Debug)]
pub struct GenericRenderer {
    renderer: *mut wlr_renderer
}

/// The state machine type that allows you to manipulate a screen and
/// its buffer.
///
/// When this structure is dropped it automatically calls wlr_renderer_end
/// and swaps the buffers.
#[derive(Debug)]
pub struct Renderer<'output> {
    renderer: *mut wlr_renderer,
    pub damage: Option<(PixmanRegion, Duration)>,
    pub output: &'output mut Output
}

impl GenericRenderer {
    /// Make a gles2 renderer.
    pub(crate) unsafe fn gles2_renderer(backend: *mut wlr_backend) -> Self {
        let egl = wlr_backend_get_egl(backend);
        if egl.is_null() {
            panic!("EGL not available for this backend");
        }
        let renderer = wlr_gles2_renderer_create(egl);
        if renderer.is_null() {
            panic!("Could not construct GLES2 renderer");
        }
        GenericRenderer { renderer }
    }

    /// Make the `Renderer` state machine type.
    ///
    /// This automatically makes the given output the current output.
    pub fn render<'output>(&mut self,
                           output: &'output mut Output,
                           damage: Option<(PixmanRegion, Duration)>)
                           -> Renderer<'output> {
        unsafe {
            output.make_current();
            let (width, height) = output.dimensions();
            wlr_renderer_begin(self.renderer, width, height);
            Renderer { renderer: self.renderer,
                       damage,
                       output }
        }
    }

    /// Create a texture using this renderer.
    pub fn create_texture_from_pixels(&mut self,
                                      format: wl_shm_format,
                                      stride: u32,
                                      width: u32,
                                      height: u32,
                                      data: &[u8])
                                      -> Option<Texture> {
        unsafe {
            create_texture_from_pixels(self.renderer,
                                       format,
                                       stride,
                                       width,
                                       height,
                                       data.as_ptr() as _)
        }
    }

    pub(crate) unsafe fn as_ptr(&self) -> *mut wlr_renderer {
        self.renderer
    }
}

impl Drop for GenericRenderer {
    fn drop(&mut self) {
        unsafe { wlr_renderer_destroy(self.renderer) }
    }
}

impl<'output> Renderer<'output> {
    pub fn clear(&mut self, float: [f32; 4]) {
        unsafe { wlr_renderer_clear(self.renderer, float.as_ptr()) }
    }

    /// Renders the requseted texture.
    pub fn render_texture(&mut self,
                          texture: &Texture,
                          projection: [f32; 9],
                          x: c_int,
                          y: c_int,
                          alpha: c_float)
                          -> bool {
        unsafe {
            wlr_render_texture(self.renderer,
                               texture.as_ptr(),
                               projection.as_ptr(),
                               x,
                               y,
                               alpha)
        }
    }

    /// Renders the requested texture using the provided matrix. A typical texture
    /// rendering goes like so:
    ///
    /// TODO FIXME Show how the typical rendering goes in Rust.
    ///
    /// ```c
    /// struct wlr_renderer *renderer;
    /// struct wlr_texture *texture;
    /// float projection[16];
    /// float matrix[16];
    /// wlr_texture_get_matrix(texture, &matrix, &projection, 123, 321);
    /// wlr_render_texture_with_matrix(renderer, texture, &matrix);
    /// ```
    ///
    /// This will render the texture at <123, 321>.
    pub fn render_texture_with_matrix(&mut self, texture: &Texture, matrix: [f32; 9]) -> bool {
        // TODO FIXME Add alpha as param
        unsafe {
            wlr_render_texture_with_matrix(self.renderer, texture.as_ptr(), matrix.as_ptr(), 1.0)
        }
    }

    /// Renders a solid quad in the specified color.
    pub fn render_colored_quad(&mut self, color: [f32; 4], matrix: [f32; 9]) {
        unsafe { wlr_render_quad_with_matrix(self.renderer, color.as_ptr(), matrix.as_ptr()) }
    }

    /// Renders a solid ellipse in the specified color.
    pub fn render_colored_ellipse(&mut self, color: [f32; 4], matrix: [f32; 9]) {
        unsafe { wlr_render_ellipse_with_matrix(self.renderer, color.as_ptr(), matrix.as_ptr()) }
    }
}

impl<'output> Drop for Renderer<'output> {
    fn drop(&mut self) {
        unsafe {
            wlr_renderer_end(self.renderer);
            if let Some((mut damage, when)) = self.damage.take() {
                self.output.swap_buffers(Some(when), Some(&mut damage));
            } else {
                self.output.swap_buffers(None, None);
            }
        }
    }
}

unsafe fn create_texture_from_pixels(renderer: *mut wlr_renderer,
                                     format: wl_shm_format,
                                     stride: u32,
                                     width: u32,
                                     height: u32,
                                     // TODO Slice of u8? It's a void*, hmm
                                     data: *const c_void)
                                     -> Option<Texture> {
    let texture = wlr_texture_from_pixels(renderer, format, stride, width, height, data);
    if texture.is_null() {
        None
    } else {
        Some(Texture::from_ptr(texture))
    }
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate bufstream;
extern crate cargo_registry;
extern crate conduit;
extern crate conduit_middleware;
extern crate conduit_test;
extern crate curl;
extern crate dotenv;
extern crate git2;
extern crate postgres;
extern crate rustc_serialize;
extern crate semver;
extern crate time;
extern crate url;

use std::collections::HashMap;
use std::process::Command;
use std::env;
use std::sync::{Once, ONCE_INIT, Arc};
use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};
use rustc_serialize::json::{self, Json};

use conduit::{Request, Method};
use conduit_test::MockRequest;
use cargo_registry::app::App;
use cargo_registry::db::{self, RequestTransaction};
use cargo_registry::dependency::Kind;
use cargo_registry::{User, Crate, Version, Keyword, Dependency};
use cargo_registry::upload as u;

macro_rules! t {
    ($e:expr) => (
        match $e {
            Ok(e) => e,
            Err(m) => panic!("{} failed with: {}", stringify!($e), m),
        }
    )
}

macro_rules! t_resp { ($e:expr) => (t!($e)) }

macro_rules! ok_resp {
    ($e:expr) => ({
        let resp = t_resp!($e);
        if !::ok_resp(&resp) { panic!("bad response: {:?}", resp.status); }
        resp
    })
}

macro_rules! bad_resp {
    ($e:expr) => ({
        let mut resp = t_resp!($e);
        match ::bad_resp(&mut resp) {
            None => panic!("ok response: {:?}", resp.status),
            Some(b) => b,
        }
    })
}

#[derive(RustcDecodable, Debug)]
struct Error { detail: String }
#[derive(RustcDecodable)]
struct Bad { errors: Vec<Error> }

mod keyword;
mod krate;
mod user;
mod record;
mod git;
mod version;
mod team;

fn app() -> (record::Bomb, Arc<App>, conduit_middleware::MiddlewareBuilder) {
    dotenv::dotenv().ok();
    static INIT: Once = ONCE_INIT;
    git::init();

    let (proxy, bomb) = record::proxy();
    let config = cargo_registry::Config {
        s3_bucket: env::var("S3_BUCKET").unwrap_or(String::new()),
        s3_access_key: env::var("S3_ACCESS_KEY").unwrap_or(String::new()),
        s3_secret_key: env::var("S3_SECRET_KEY").unwrap_or(String::new()),
        s3_region: env::var("S3_REGION").ok(),
        s3_proxy: Some(proxy),
        session_key: "test".to_string(),
        git_repo_checkout: git::checkout(),
        gh_client_id: env::var("GH_CLIENT_ID").unwrap_or(String::new()),
        gh_client_secret: env::var("GH_CLIENT_SECRET").unwrap_or(String::new()),
        db_url: env("TEST_DATABASE_URL"),
        env: cargo_registry::Env::Test,
        max_upload_size: 1000,
        mirror: false,
    };
    INIT.call_once(|| db_setup(&config.db_url));
    let app = App::new(&config);
    let app = Arc::new(app);
    let middleware = cargo_registry::middleware(app.clone());
    return (bomb, app, middleware);

    fn db_setup(db: &str) {
        let migrate = t!(env::current_exe()).parent().unwrap().join("migrate");
        assert!(t!(Command::new(&migrate).env("DATABASE_URL", db)
                           .status()).success());
    }
}

fn env(s: &str) -> String {
    match env::var(s).ok() {
        Some(s) => s,
        None => panic!("must have `{}` defined", s),
    }
}

fn req(app: Arc<App>, method: conduit::Method, path: &str) -> MockRequest {
    let mut req = MockRequest::new(method, path);
    req.mut_extensions().insert(db::Transaction::new(app));
    return req;
}

fn ok_resp(r: &conduit::Response) -> bool {
    r.status.0 == 200
}

fn bad_resp(r: &mut conduit::Response) -> Option<Bad> {
    let bad = json::<Bad>(r);
    if bad.errors.len() == 0 { return None }
    Some(bad)
}

fn json<T: rustc_serialize::Decodable>(r: &mut conduit::Response) -> T {
    let mut data = Vec::new();
    r.body.read_to_end(&mut data).unwrap();
    let s = std::str::from_utf8(&data).unwrap();
    let j = match Json::from_str(s) {
        Ok(t) => t,
        Err(e) => panic!("failed to decode: {:?}\n{}", e, s),
    };
    let j = fixup(j);
    let s = j.to_string();
    return match json::decode(&s) {
        Ok(t) => t,
        Err(e) => panic!("failed to decode: {:?}\n{}", e, s),
    };

    fn fixup(json: Json) -> Json {
        match json {
            Json::Object(object) => {
                Json::Object(object.into_iter().map(|(k, v)| {
                    let k = if k == "crate" {
                        "krate".to_string()
                    } else {
                        k
                    };
                    (k, fixup(v))
                }).collect())
            }
            Json::Array(list) => {
                Json::Array(list.into_iter().map(fixup).collect())
            }
            j => j,
        }
    }
}

static NEXT_ID: AtomicUsize = ATOMIC_USIZE_INIT;

fn user(login: &str) -> User {
    User {
        id: NEXT_ID.fetch_add(1, Ordering::SeqCst) as i32,
        gh_id: NEXT_ID.fetch_add(1, Ordering::SeqCst) as i32,
        gh_login: login.to_string(),
        email: None,
        name: None,
        avatar: None,
        gh_access_token: User::new_api_token(), // just randomize it
        api_token: User::new_api_token(),
    }
}

fn krate(name: &str) -> Crate {
    cargo_registry::krate::Crate {
        id: NEXT_ID.fetch_add(1, Ordering::SeqCst) as i32,
        name: name.to_string(),
        updated_at: time::now().to_timespec(),
        created_at: time::now().to_timespec(),
        downloads: 10,
        max_version: semver::Version::parse("0.0.0").unwrap(),
        documentation: None,
        homepage: None,
        description: None,
        readme: None,
        license: None,
        repository: None,
        max_upload_size: None,
    }
}

fn mock_user(req: &mut Request, u: User) -> User {
    let u = User::find_or_insert(req.tx().unwrap(),
                                 u.gh_id,
                                 &u.gh_login,
                                 u.email.as_ref().map(|s| &s[..]),
                                 u.name.as_ref().map(|s| &s[..]),
                                 u.avatar.as_ref().map(|s| &s[..]),
                                 &u.gh_access_token,
                                 &u.api_token).unwrap();
    req.mut_extensions().insert(u.clone());
    return u;
}

fn mock_crate(req: &mut Request, krate: Crate) -> (Crate, Version) {
    mock_crate_vers(req, krate, &semver::Version::parse("1.0.0").unwrap())
}

fn mock_crate_vers(req: &mut Request, krate: Crate, v: &semver::Version)
                   -> (Crate, Version) {
    let user = req.extensions().find::<User>().unwrap();
    let mut krate = Crate::find_or_insert(req.tx().unwrap(), &krate.name,
                                          user.id, &krate.description,
                                          &krate.homepage,
                                          &krate.documentation,
                                          &krate.readme,
                                          &krate.repository,
                                          &krate.license,
                                          &None,
                                          krate.max_upload_size).unwrap();
    let v = krate.add_version(req.tx().unwrap(), v, &HashMap::new(), &[]);
    (krate, v.unwrap())
}

fn mock_dep(req: &mut Request, version: &Version, krate: &Crate,
            target: Option<&str>) -> Dependency {
    Dependency::insert(req.tx().unwrap(),
                       version.id,
                       krate.id,
                       &semver::VersionReq::parse(">= 0").unwrap(),
                       Kind::Normal,
                       false, true, &[],
                       &target.map(|s| s.to_string())).unwrap()
}

fn mock_keyword(req: &mut Request, name: &str) -> Keyword {
    Keyword::find_or_insert(req.tx().unwrap(), name).unwrap()
}

fn logout(req: &mut Request) {
    req.mut_extensions().pop::<User>();
}

fn new_req(app: Arc<App>, krate: &str, version: &str) -> MockRequest {
    new_req_full(app, ::krate(krate), version, Vec::new())
}

fn new_req_full(app: Arc<App>, krate: Crate, version: &str,
                deps: Vec<u::CrateDependency>) -> MockRequest {
    let mut req = ::req(app, Method::Put, "/api/v1/crates/new");
    req.with_body(&new_req_body(krate, version, deps, Vec::new()));
    return req;
}

fn new_req_with_keywords(app: Arc<App>, krate: Crate, version: &str,
                         kws: Vec<String>) -> MockRequest {
    let mut req = ::req(app, Method::Put, "/api/v1/crates/new");
    req.with_body(&new_req_body(krate, version, Vec::new(), kws));
    return req;
}

fn new_req_body(krate: Crate, version: &str, deps: Vec<u::CrateDependency>,
                kws: Vec<String>) -> Vec<u8> {
    let kws = kws.into_iter().map(u::Keyword).collect();
    new_crate_to_body(&u::NewCrate {
        name: u::CrateName(krate.name),
        vers: u::CrateVersion(semver::Version::parse(version).unwrap()),
        features: HashMap::new(),
        deps: deps,
        authors: vec!["foo".to_string()],
        description: Some("description".to_string()),
        homepage: krate.homepage,
        documentation: krate.documentation,
        readme: krate.readme,
        keywords: Some(u::KeywordList(kws)),
        license: Some("MIT".to_string()),
        license_file: None,
        repository: krate.repository,
    }, &[])
}

fn new_crate_to_body(new_crate: &u::NewCrate, krate: &[u8]) -> Vec<u8> {
    let json = json::encode(&new_crate).unwrap();
    let mut body = Vec::new();
    body.extend([
        (json.len() >>  0) as u8,
        (json.len() >>  8) as u8,
        (json.len() >> 16) as u8,
        (json.len() >> 24) as u8,
    ].iter().cloned());
    body.extend(json.as_bytes().iter().cloned());
    body.extend(&[
        (krate.len() >>  0) as u8,
        (krate.len() >>  8) as u8,
        (krate.len() >> 16) as u8,
        (krate.len() >> 24) as u8,
    ]);
    body.extend(krate);
    body
}
<commit_msg>Fix tests on nightly<commit_after>#![deny(warnings)]

extern crate bufstream;
extern crate cargo_registry;
extern crate conduit;
extern crate conduit_middleware;
extern crate conduit_test;
extern crate curl;
extern crate dotenv;
extern crate git2;
extern crate postgres;
extern crate rustc_serialize;
extern crate semver;
extern crate time;
extern crate url;

use std::collections::HashMap;
use std::process::Command;
use std::env;
use std::sync::{Once, ONCE_INIT, Arc};
use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};
use rustc_serialize::json::{self, Json};

use conduit::{Request, Method};
use conduit_test::MockRequest;
use cargo_registry::app::App;
use cargo_registry::db::{self, RequestTransaction};
use cargo_registry::dependency::Kind;
use cargo_registry::{User, Crate, Version, Keyword, Dependency};
use cargo_registry::upload as u;

macro_rules! t {
    ($e:expr) => (
        match $e {
            Ok(e) => e,
            Err(m) => panic!("{} failed with: {}", stringify!($e), m),
        }
    )
}

macro_rules! t_resp { ($e:expr) => (t!($e)) }

macro_rules! ok_resp {
    ($e:expr) => ({
        let resp = t_resp!($e);
        if !::ok_resp(&resp) { panic!("bad response: {:?}", resp.status); }
        resp
    })
}

macro_rules! bad_resp {
    ($e:expr) => ({
        let mut resp = t_resp!($e);
        match ::bad_resp(&mut resp) {
            None => panic!("ok response: {:?}", resp.status),
            Some(b) => b,
        }
    })
}

#[derive(RustcDecodable, Debug)]
struct Error { detail: String }
#[derive(RustcDecodable)]
struct Bad { errors: Vec<Error> }

mod keyword;
mod krate;
mod user;
mod record;
mod git;
mod version;
mod team;

fn app() -> (record::Bomb, Arc<App>, conduit_middleware::MiddlewareBuilder) {
    dotenv::dotenv().ok();
    static INIT: Once = ONCE_INIT;
    git::init();

    let (proxy, bomb) = record::proxy();
    let config = cargo_registry::Config {
        s3_bucket: env::var("S3_BUCKET").unwrap_or(String::new()),
        s3_access_key: env::var("S3_ACCESS_KEY").unwrap_or(String::new()),
        s3_secret_key: env::var("S3_SECRET_KEY").unwrap_or(String::new()),
        s3_region: env::var("S3_REGION").ok(),
        s3_proxy: Some(proxy),
        session_key: "test".to_string(),
        git_repo_checkout: git::checkout(),
        gh_client_id: env::var("GH_CLIENT_ID").unwrap_or(String::new()),
        gh_client_secret: env::var("GH_CLIENT_SECRET").unwrap_or(String::new()),
        db_url: env("TEST_DATABASE_URL"),
        env: cargo_registry::Env::Test,
        max_upload_size: 1000,
        mirror: false,
    };
    INIT.call_once(|| db_setup(&config.db_url));
    let app = App::new(&config);
    let app = Arc::new(app);
    let middleware = cargo_registry::middleware(app.clone());
    return (bomb, app, middleware);

    fn db_setup(db: &str) {
        let mut me = t!(env::current_exe());
        me.pop();
        if me.ends_with("deps") {
            me.pop();
        }
        me.push("migrate");
        assert!(t!(Command::new(&me).env("DATABASE_URL", db)
                           .status()).success());
    }
}

fn env(s: &str) -> String {
    match env::var(s).ok() {
        Some(s) => s,
        None => panic!("must have `{}` defined", s),
    }
}

fn req(app: Arc<App>, method: conduit::Method, path: &str) -> MockRequest {
    let mut req = MockRequest::new(method, path);
    req.mut_extensions().insert(db::Transaction::new(app));
    return req;
}

fn ok_resp(r: &conduit::Response) -> bool {
    r.status.0 == 200
}

fn bad_resp(r: &mut conduit::Response) -> Option<Bad> {
    let bad = json::<Bad>(r);
    if bad.errors.len() == 0 { return None }
    Some(bad)
}

fn json<T: rustc_serialize::Decodable>(r: &mut conduit::Response) -> T {
    let mut data = Vec::new();
    r.body.read_to_end(&mut data).unwrap();
    let s = std::str::from_utf8(&data).unwrap();
    let j = match Json::from_str(s) {
        Ok(t) => t,
        Err(e) => panic!("failed to decode: {:?}\n{}", e, s),
    };
    let j = fixup(j);
    let s = j.to_string();
    return match json::decode(&s) {
        Ok(t) => t,
        Err(e) => panic!("failed to decode: {:?}\n{}", e, s),
    };

    fn fixup(json: Json) -> Json {
        match json {
            Json::Object(object) => {
                Json::Object(object.into_iter().map(|(k, v)| {
                    let k = if k == "crate" {
                        "krate".to_string()
                    } else {
                        k
                    };
                    (k, fixup(v))
                }).collect())
            }
            Json::Array(list) => {
                Json::Array(list.into_iter().map(fixup).collect())
            }
            j => j,
        }
    }
}

static NEXT_ID: AtomicUsize = ATOMIC_USIZE_INIT;

fn user(login: &str) -> User {
    User {
        id: NEXT_ID.fetch_add(1, Ordering::SeqCst) as i32,
        gh_id: NEXT_ID.fetch_add(1, Ordering::SeqCst) as i32,
        gh_login: login.to_string(),
        email: None,
        name: None,
        avatar: None,
        gh_access_token: User::new_api_token(), // just randomize it
        api_token: User::new_api_token(),
    }
}

fn krate(name: &str) -> Crate {
    cargo_registry::krate::Crate {
        id: NEXT_ID.fetch_add(1, Ordering::SeqCst) as i32,
        name: name.to_string(),
        updated_at: time::now().to_timespec(),
        created_at: time::now().to_timespec(),
        downloads: 10,
        max_version: semver::Version::parse("0.0.0").unwrap(),
        documentation: None,
        homepage: None,
        description: None,
        readme: None,
        license: None,
        repository: None,
        max_upload_size: None,
    }
}

fn mock_user(req: &mut Request, u: User) -> User {
    let u = User::find_or_insert(req.tx().unwrap(),
                                 u.gh_id,
                                 &u.gh_login,
                                 u.email.as_ref().map(|s| &s[..]),
                                 u.name.as_ref().map(|s| &s[..]),
                                 u.avatar.as_ref().map(|s| &s[..]),
                                 &u.gh_access_token,
                                 &u.api_token).unwrap();
    req.mut_extensions().insert(u.clone());
    return u;
}

fn mock_crate(req: &mut Request, krate: Crate) -> (Crate, Version) {
    mock_crate_vers(req, krate, &semver::Version::parse("1.0.0").unwrap())
}

fn mock_crate_vers(req: &mut Request, krate: Crate, v: &semver::Version)
                   -> (Crate, Version) {
    let user = req.extensions().find::<User>().unwrap();
    let mut krate = Crate::find_or_insert(req.tx().unwrap(), &krate.name,
                                          user.id, &krate.description,
                                          &krate.homepage,
                                          &krate.documentation,
                                          &krate.readme,
                                          &krate.repository,
                                          &krate.license,
                                          &None,
                                          krate.max_upload_size).unwrap();
    let v = krate.add_version(req.tx().unwrap(), v, &HashMap::new(), &[]);
    (krate, v.unwrap())
}

fn mock_dep(req: &mut Request, version: &Version, krate: &Crate,
            target: Option<&str>) -> Dependency {
    Dependency::insert(req.tx().unwrap(),
                       version.id,
                       krate.id,
                       &semver::VersionReq::parse(">= 0").unwrap(),
                       Kind::Normal,
                       false, true, &[],
                       &target.map(|s| s.to_string())).unwrap()
}

fn mock_keyword(req: &mut Request, name: &str) -> Keyword {
    Keyword::find_or_insert(req.tx().unwrap(), name).unwrap()
}

fn logout(req: &mut Request) {
    req.mut_extensions().pop::<User>();
}

fn new_req(app: Arc<App>, krate: &str, version: &str) -> MockRequest {
    new_req_full(app, ::krate(krate), version, Vec::new())
}

fn new_req_full(app: Arc<App>, krate: Crate, version: &str,
                deps: Vec<u::CrateDependency>) -> MockRequest {
    let mut req = ::req(app, Method::Put, "/api/v1/crates/new");
    req.with_body(&new_req_body(krate, version, deps, Vec::new()));
    return req;
}

fn new_req_with_keywords(app: Arc<App>, krate: Crate, version: &str,
                         kws: Vec<String>) -> MockRequest {
    let mut req = ::req(app, Method::Put, "/api/v1/crates/new");
    req.with_body(&new_req_body(krate, version, Vec::new(), kws));
    return req;
}

fn new_req_body(krate: Crate, version: &str, deps: Vec<u::CrateDependency>,
                kws: Vec<String>) -> Vec<u8> {
    let kws = kws.into_iter().map(u::Keyword).collect();
    new_crate_to_body(&u::NewCrate {
        name: u::CrateName(krate.name),
        vers: u::CrateVersion(semver::Version::parse(version).unwrap()),
        features: HashMap::new(),
        deps: deps,
        authors: vec!["foo".to_string()],
        description: Some("description".to_string()),
        homepage: krate.homepage,
        documentation: krate.documentation,
        readme: krate.readme,
        keywords: Some(u::KeywordList(kws)),
        license: Some("MIT".to_string()),
        license_file: None,
        repository: krate.repository,
    }, &[])
}

fn new_crate_to_body(new_crate: &u::NewCrate, krate: &[u8]) -> Vec<u8> {
    let json = json::encode(&new_crate).unwrap();
    let mut body = Vec::new();
    body.extend([
        (json.len() >>  0) as u8,
        (json.len() >>  8) as u8,
        (json.len() >> 16) as u8,
        (json.len() >> 24) as u8,
    ].iter().cloned());
    body.extend(json.as_bytes().iter().cloned());
    body.extend(&[
        (krate.len() >>  0) as u8,
        (krate.len() >>  8) as u8,
        (krate.len() >> 16) as u8,
        (krate.len() >> 24) as u8,
    ]);
    body.extend(krate);
    body
}
<|endoftext|>


<commit_before>use event::Event;
use sass::comment::SassComment;
use sass::rule::SassRule;
use sass::selector::SassSelector;
use sass::variable::SassVariable;
use top_level_event::TopLevelEvent;

use std::borrow::Cow::Borrowed;

fn is_ascii_whitespace(c: u8) -> bool {
   is_newline(c) || is_ascii_whitespace_no_nl(c)
}

fn is_ascii_whitespace_no_nl(c: u8) -> bool {
    c == b'\t' || c == 0x0b || c == 0x0c || c == b' '
}

fn is_newline(c: u8) -> bool {
    c == b'\n' || c == b'\r'
}

fn isnt_newline(c: u8) -> bool {
    !is_newline(c)
}

// unusual among "scan" functions in that it scans from the _back_ of the string
// TODO: should also scan unicode whitespace?
fn scan_trailing_whitespace(data: &str) -> usize {
    match data.as_bytes().iter().rev().position(|&c| !is_ascii_whitespace_no_nl(c)) {
        Some(i) => i,
        None => data.len()
    }
}

#[derive(Debug)]
pub struct Tokenizer<'a> {
    sass: &'a str,
    bytes: &'a [u8],
    offset: usize,
    state: State,
    sass_rule_stack: Vec<SassRule<'a>>,
    current_sass_rule_selectors_done: bool,
}

#[derive(PartialEq, Debug, Clone)]
enum State {
    OutsideRules,
    InVariable,
    InComment,
    InRule,
    InSelectors,
    InProperties,
    EndRule,
    Eof,
}

impl<'a> Tokenizer<'a> {
    pub fn new(sass: &'a str) -> Tokenizer<'a> {
        Tokenizer {
            sass: &sass,
            bytes: &sass.as_bytes(),
            offset: 0,
            state: State::OutsideRules,
            sass_rule_stack: Vec::new(),
            current_sass_rule_selectors_done: false,
        }
    }

    fn start_something(&mut self) -> Option<TopLevelEvent<'a>> {
        let mut current_sass_rule = SassRule::new();
        self.current_sass_rule_selectors_done = false;

        self.pick_something();
        while self.state != State::OutsideRules {
            if self.state == State::Eof { return None }

            if self.state == State::InSelectors {
                let sel = self.next_selector();
                if sel.is_some() {
                    current_sass_rule.selectors.push(sel.unwrap());
                }
            } else if self.state == State::InProperties {
                let prop = self.next_property();
                if prop.is_some() {
                    current_sass_rule.children.push(prop.unwrap());
                }
            } else if self.state == State::EndRule {
                // consume '}', should probably be checking that's what we actually have
                self.offset += 1;

                match self.sass_rule_stack.pop() {
                    Some(mut rule) => {
                        rule.children.push(Event::ChildRule(current_sass_rule));
                        current_sass_rule = rule;
                        self.pick_something();
                    },
                    None => self.state = State::OutsideRules,
                }
            } else if self.state == State::InRule {
                self.sass_rule_stack.push(current_sass_rule);
                current_sass_rule = SassRule::new();
                self.current_sass_rule_selectors_done = false;
                self.pick_something();
            } else if self.state == State::InVariable {
                let var = self.next_variable();
                if var.is_some() {
                    return var
                } else {
                    // is this really what we should be doing here? reachable?
                    self.state = State::Eof;
                }
            } else if self.state == State::InComment {
                let comment = self.next_comment();
                if comment.is_some() {
                    if self.sass_rule_stack.len() == 0 &&
                       current_sass_rule.selectors.len() == 0 &&
                       current_sass_rule.children.len() == 0 {
                           return Some(TopLevelEvent::Comment(SassComment { comment: comment.unwrap() }))
                    } else {
                        current_sass_rule.children.push(comment.unwrap());
                        self.pick_something();
                    }
                } else {
                    // is this really what we should be doing here? reachable?
                    self.state = State::Eof;
                }
            } else {
                println!("i dont know what to do for {:?}", self.state);
                println!("current sass rule = {:?}", current_sass_rule);
                self.state = State::Eof;
            }
        }

        Some(TopLevelEvent::Rule(current_sass_rule))
    }

    fn pick_something(&mut self) {
        self.skip_leading_whitespace();

        if self.offset == self.sass.len() {
            self.state = State::Eof;
            return
        }

        let c = self.bytes[self.offset];

        if c == b'}' {
            self.state = State::EndRule;
            return
        }

        if c == b'$' {
            self.state = State::InVariable;
            return
        }

        if c == b'/' && (self.offset + 1) < self.sass.len() {
            let d = self.bytes[self.offset + 1];
            if d == b'*' {
                self.state = State::InComment;
                return
            }
        }

        self.state = State::InSelectors;
    }

    fn scan_while<F>(&mut self, data: &str, f: F) -> usize
            where F: Fn(u8) -> bool {
        match data.as_bytes().iter().position(|&c| !f(c)) {
            Some(i) => i,
            None => data.len()
        }
    }

    fn skip_leading_whitespace(&mut self) {
        let mut i = self.offset;
        let limit = self.sass.len();

        while i < limit {
            let c = self.bytes[i];
            if is_ascii_whitespace(c) {
                i += self.scan_while(&self.sass[i..self.sass.len()], is_ascii_whitespace);
            } else if c == b'/' && i + 1 < limit && self.bytes[i + 1] == b'/' {
                i += self.scan_while(&self.sass[i..self.sass.len()], isnt_newline);
            } else {
                self.offset = i;
                return
            }
        }
        self.offset = limit;
    }

    fn next_comment(&mut self) -> Option<Event<'a>> {
        let comment_body_beginning = self.offset;
        let mut i = comment_body_beginning + 2;
        let limit = self.sass.len();

        while i < limit {
            match self.bytes[i..limit].iter().position(|&c| c == b'*' ) {
                Some(pos) => { i += pos; },
                None => { break; },
            }

            if self.bytes[i+1] == b'/' {
                self.offset = i + 2;
                return Some(Event::Comment(Borrowed(&self.sass[comment_body_beginning..i + 2])))
            } else {
                i += 1;
            }
        }
        None
    }

    fn next_variable(&mut self) -> Option<TopLevelEvent<'a>> {
        // TODO: can parts of this be deduplicated with properties?
        let name_beginning = self.offset;
        let mut i = name_beginning;
        let limit = self.sass.len();

        while i < limit {
            match self.bytes[i..limit].iter().position(|&c| c == b':' ) {
                Some(pos) => { i += pos; },
                None => { break; },
            }

            let name_end = i;

            i += 1;
            self.offset = i;
            self.skip_leading_whitespace();

            let value_beginning = self.offset;
            i = value_beginning;

            while i < limit {
                match self.bytes[i..limit].iter().position(|&c| c == b';') {
                    Some(pos) => { i += pos; },
                    None => { i = limit; break; },
                }

                let value_end = i;
                self.offset = i + 1;

                self.skip_leading_whitespace();

                return Some(TopLevelEvent::Variable(SassVariable{
                    name: Borrowed(&self.sass[name_beginning..name_end]),
                    value: Borrowed(&self.sass[value_beginning..value_end]),
                }))
            }
        }
        self.offset = self.sass.len();
        None
    }

    fn next_property(&mut self) -> Option<Event<'a>> {
        self.skip_leading_whitespace();

        if self.offset == self.sass.len() {
            self.state = State::Eof;
            return None
        }

        let name_beginning = self.offset;
        let mut i = name_beginning;
        let limit = self.sass.len();

        let c = self.bytes[i];
        if c == b'}' {
            self.state = State::EndRule;
            return None
        }

        let d = self.bytes[i + 1];
        if c == b'/' && d == b'*' {
            self.state = State::InComment;
            return None
        }

        while i < limit {
            match self.bytes[i..limit].iter().position(|&c| c == b':' || c == b'{') {
                Some(pos) => { i += pos; },
                None => { break; },
            }

            // Inefficient since we already skipped the whitespace and we'll have to
            // do it again but oh well
            let c = self.bytes[i];
            if c == b'{' {
                self.state = State::InRule;
                return None
            }

            let name_end = i;

            i += 1;
            self.offset = i;
            self.skip_leading_whitespace();

            let value_beginning = self.offset;
            i = value_beginning;

            while i < limit {
                match self.bytes[i..limit].iter().position(|&c| c == b';') {
                    Some(pos) => { i += pos; },
                    None => { i = limit; break; },
                }

                let value_end = i;
                self.offset = i + 1;

                self.skip_leading_whitespace();

                if self.bytes[name_beginning] == b'$' {
                    return Some(Event::Variable(SassVariable {
                        name: Borrowed(&self.sass[name_beginning..name_end]),
                        value: Borrowed(&self.sass[value_beginning..value_end])
                    }))
                } else {
                    return Some(Event::Property(
                        Borrowed(&self.sass[name_beginning..name_end]),
                        Borrowed(&self.sass[value_beginning..value_end])
                    ))
                }
            }
        }
        self.offset = self.sass.len();
        None
    }

    fn next_selector(&mut self) -> Option<SassSelector<'a>> {
        let beginning = self.offset;
        let mut i = beginning;
        let limit = self.sass.len();

        while i < limit {
            match self.bytes[i..limit].iter().position(|&c| c == b',' || c == b'{' || c == b':') {
                Some(pos) => { i += pos; },
                None => { i = limit; break; },
            }

            let c = self.bytes[i];

            if c == b':' {
                self.state = State::InProperties;
                return None
            }

            if c == b',' || c == b'{' {
                let n = scan_trailing_whitespace(&self.sass[beginning..i]);
                let end = i - n;
                if end > beginning {
                    if c == b'{' {
                        if self.current_sass_rule_selectors_done {
                            self.state = State::InRule;
                            return None
                        } else {
                            self.current_sass_rule_selectors_done = true;
                            self.state = State::InProperties;
                        }
                    }
                    self.offset = i + 1;
                    return Some(SassSelector { name: Borrowed(&self.sass[beginning..end]) })
                }
            }

            self.offset = i;
            if i > beginning {
                return Some(SassSelector { name: Borrowed(&self.sass[beginning..i]) })
            }
            i += 1;
        }

        if i > beginning {
            self.offset = i;
            Some(SassSelector { name: Borrowed(&self.sass[beginning..i]) })
        } else {
            self.state = State::Eof;
            None
        }
    }
}

impl<'a> Iterator for Tokenizer<'a> {
    type Item = TopLevelEvent<'a>;

    fn next(&mut self) -> Option<TopLevelEvent<'a>> {
        if self.offset < self.sass.len() {
            return self.start_something()
        }
        None
    }
}
<commit_msg>Handle spaces and extra commas in selectors<commit_after>use event::Event;
use sass::comment::SassComment;
use sass::rule::SassRule;
use sass::selector::SassSelector;
use sass::variable::SassVariable;
use top_level_event::TopLevelEvent;

use std::borrow::Cow::Borrowed;

fn is_ascii_whitespace(c: u8) -> bool {
   is_newline(c) || is_ascii_whitespace_no_nl(c)
}

fn is_ascii_whitespace_no_nl(c: u8) -> bool {
    c == b'\t' || c == 0x0b || c == 0x0c || c == b' '
}

fn is_newline(c: u8) -> bool {
    c == b'\n' || c == b'\r'
}

fn isnt_newline(c: u8) -> bool {
    !is_newline(c)
}

// unusual among "scan" functions in that it scans from the _back_ of the string
// TODO: should also scan unicode whitespace?
fn scan_trailing_whitespace(data: &str) -> usize {
    match data.as_bytes().iter().rev().position(|&c| !is_ascii_whitespace_no_nl(c)) {
        Some(i) => i,
        None => data.len()
    }
}

#[derive(Debug)]
pub struct Tokenizer<'a> {
    sass: &'a str,
    bytes: &'a [u8],
    offset: usize,
    state: State,
    sass_rule_stack: Vec<SassRule<'a>>,
    current_sass_rule_selectors_done: bool,
}

#[derive(PartialEq, Debug, Clone)]
enum State {
    OutsideRules,
    InVariable,
    InComment,
    InRule,
    InSelectors,
    InProperties,
    EndRule,
    Eof,
}

impl<'a> Tokenizer<'a> {
    pub fn new(sass: &'a str) -> Tokenizer<'a> {
        Tokenizer {
            sass: &sass,
            bytes: &sass.as_bytes(),
            offset: 0,
            state: State::OutsideRules,
            sass_rule_stack: Vec::new(),
            current_sass_rule_selectors_done: false,
        }
    }

    fn start_something(&mut self) -> Option<TopLevelEvent<'a>> {
        let mut current_sass_rule = SassRule::new();
        self.current_sass_rule_selectors_done = false;

        self.pick_something();
        while self.state != State::OutsideRules {
            if self.state == State::Eof { return None }

            if self.state == State::InSelectors {
                let sel = self.next_selector();
                if sel.is_some() {
                    current_sass_rule.selectors.push(sel.unwrap());
                }
            } else if self.state == State::InProperties {
                let prop = self.next_property();
                if prop.is_some() {
                    current_sass_rule.children.push(prop.unwrap());
                }
            } else if self.state == State::EndRule {
                // consume '}', should probably be checking that's what we actually have
                self.offset += 1;

                match self.sass_rule_stack.pop() {
                    Some(mut rule) => {
                        rule.children.push(Event::ChildRule(current_sass_rule));
                        current_sass_rule = rule;
                        self.pick_something();
                    },
                    None => self.state = State::OutsideRules,
                }
            } else if self.state == State::InRule {
                self.sass_rule_stack.push(current_sass_rule);
                current_sass_rule = SassRule::new();
                self.current_sass_rule_selectors_done = false;
                self.pick_something();
            } else if self.state == State::InVariable {
                let var = self.next_variable();
                if var.is_some() {
                    return var
                } else {
                    // is this really what we should be doing here? reachable?
                    self.state = State::Eof;
                }
            } else if self.state == State::InComment {
                let comment = self.next_comment();
                if comment.is_some() {
                    if self.sass_rule_stack.len() == 0 &&
                       current_sass_rule.selectors.len() == 0 &&
                       current_sass_rule.children.len() == 0 {
                           return Some(TopLevelEvent::Comment(SassComment { comment: comment.unwrap() }))
                    } else {
                        current_sass_rule.children.push(comment.unwrap());
                        self.pick_something();
                    }
                } else {
                    // is this really what we should be doing here? reachable?
                    self.state = State::Eof;
                }
            } else {
                println!("i dont know what to do for {:?}", self.state);
                println!("current sass rule = {:?}", current_sass_rule);
                self.state = State::Eof;
            }
        }

        Some(TopLevelEvent::Rule(current_sass_rule))
    }

    fn pick_something(&mut self) {
        self.skip_leading_whitespace();

        if self.offset == self.sass.len() {
            self.state = State::Eof;
            return
        }

        let c = self.bytes[self.offset];

        if c == b'}' {
            self.state = State::EndRule;
            return
        }

        if c == b'$' {
            self.state = State::InVariable;
            return
        }

        if c == b'/' && (self.offset + 1) < self.sass.len() {
            let d = self.bytes[self.offset + 1];
            if d == b'*' {
                self.state = State::InComment;
                return
            }
        }

        self.state = State::InSelectors;
    }

    fn scan_while<F>(&mut self, data: &str, f: F) -> usize
            where F: Fn(u8) -> bool {
        match data.as_bytes().iter().position(|&c| !f(c)) {
            Some(i) => i,
            None => data.len()
        }
    }

    fn skip_leading_whitespace(&mut self) {
        let mut i = self.offset;
        let limit = self.sass.len();

        while i < limit {
            let c = self.bytes[i];
            if is_ascii_whitespace(c) {
                i += self.scan_while(&self.sass[i..self.sass.len()], is_ascii_whitespace);
            } else if c == b'/' && i + 1 < limit && self.bytes[i + 1] == b'/' {
                i += self.scan_while(&self.sass[i..self.sass.len()], isnt_newline);
            } else {
                self.offset = i;
                return
            }
        }
        self.offset = limit;
    }

    fn next_comment(&mut self) -> Option<Event<'a>> {
        let comment_body_beginning = self.offset;
        let mut i = comment_body_beginning + 2;
        let limit = self.sass.len();

        while i < limit {
            match self.bytes[i..limit].iter().position(|&c| c == b'*' ) {
                Some(pos) => { i += pos; },
                None => { break; },
            }

            if self.bytes[i+1] == b'/' {
                self.offset = i + 2;
                return Some(Event::Comment(Borrowed(&self.sass[comment_body_beginning..i + 2])))
            } else {
                i += 1;
            }
        }
        None
    }

    fn next_variable(&mut self) -> Option<TopLevelEvent<'a>> {
        // TODO: can parts of this be deduplicated with properties?
        let name_beginning = self.offset;
        let mut i = name_beginning;
        let limit = self.sass.len();

        while i < limit {
            match self.bytes[i..limit].iter().position(|&c| c == b':' ) {
                Some(pos) => { i += pos; },
                None => { break; },
            }

            let name_end = i;

            i += 1;
            self.offset = i;
            self.skip_leading_whitespace();

            let value_beginning = self.offset;
            i = value_beginning;

            while i < limit {
                match self.bytes[i..limit].iter().position(|&c| c == b';') {
                    Some(pos) => { i += pos; },
                    None => { i = limit; break; },
                }

                let value_end = i;
                self.offset = i + 1;

                self.skip_leading_whitespace();

                return Some(TopLevelEvent::Variable(SassVariable{
                    name: Borrowed(&self.sass[name_beginning..name_end]),
                    value: Borrowed(&self.sass[value_beginning..value_end]),
                }))
            }
        }
        self.offset = self.sass.len();
        None
    }

    fn next_property(&mut self) -> Option<Event<'a>> {
        self.skip_leading_whitespace();

        if self.offset == self.sass.len() {
            self.state = State::Eof;
            return None
        }

        let name_beginning = self.offset;
        let mut i = name_beginning;
        let limit = self.sass.len();

        let c = self.bytes[i];
        if c == b'}' {
            self.state = State::EndRule;
            return None
        }

        let d = self.bytes[i + 1];
        if c == b'/' && d == b'*' {
            self.state = State::InComment;
            return None
        }

        while i < limit {
            match self.bytes[i..limit].iter().position(|&c| c == b':' || c == b'{') {
                Some(pos) => { i += pos; },
                None => { break; },
            }

            // Inefficient since we already skipped the whitespace and we'll have to
            // do it again but oh well
            let c = self.bytes[i];
            if c == b'{' {
                self.state = State::InRule;
                return None
            }

            let name_end = i;

            i += 1;
            self.offset = i;
            self.skip_leading_whitespace();

            let value_beginning = self.offset;
            i = value_beginning;

            while i < limit {
                match self.bytes[i..limit].iter().position(|&c| c == b';') {
                    Some(pos) => { i += pos; },
                    None => { i = limit; break; },
                }

                let value_end = i;
                self.offset = i + 1;

                self.skip_leading_whitespace();

                if self.bytes[name_beginning] == b'$' {
                    return Some(Event::Variable(SassVariable {
                        name: Borrowed(&self.sass[name_beginning..name_end]),
                        value: Borrowed(&self.sass[value_beginning..value_end])
                    }))
                } else {
                    return Some(Event::Property(
                        Borrowed(&self.sass[name_beginning..name_end]),
                        Borrowed(&self.sass[value_beginning..value_end])
                    ))
                }
            }
        }
        self.offset = self.sass.len();
        None
    }

    fn next_selector(&mut self) -> Option<SassSelector<'a>> {
        self.skip_leading_whitespace();

        let beginning = self.offset;
        let mut i = beginning;
        let limit = self.sass.len();

        while i < limit {
            match self.bytes[i..limit].iter().position(|&c| c == b',' || c == b'{' || c == b':') {
                Some(pos) => { i += pos; },
                None => { i = limit; break; },
            }

            let c = self.bytes[i];

            if c == b':' {
                self.state = State::InProperties;
                return None
            }

            if c == b',' || c == b'{' {
                let n = scan_trailing_whitespace(&self.sass[beginning..i]);
                let end = i - n;
                if end > beginning {
                    if c == b'{' {
                        if self.current_sass_rule_selectors_done {
                            self.state = State::InRule;
                            return None
                        } else {
                            self.current_sass_rule_selectors_done = true;
                            self.state = State::InProperties;
                        }
                    }
                    self.offset = i + 1;
                    if c == b',' && beginning + 1 == end {
                        self.offset += 1;
                        return self.next_selector()
                    }
                    return Some(SassSelector { name: Borrowed(&self.sass[beginning..end]) })
                } else {
                    // only whitespace between commas
                    self.offset += 1;
                    return self.next_selector()
                }
            }

            self.offset = i;
            if i > beginning {
                return Some(SassSelector { name: Borrowed(&self.sass[beginning..i]) })
            }
            i += 1;
        }

        if i > beginning {
            self.offset = i;
            Some(SassSelector { name: Borrowed(&self.sass[beginning..i]) })
        } else {
            self.state = State::Eof;
            None
        }
    }
}

impl<'a> Iterator for Tokenizer<'a> {
    type Item = TopLevelEvent<'a>;

    fn next(&mut self) -> Option<TopLevelEvent<'a>> {
        if self.offset < self.sass.len() {
            return self.start_something()
        }
        None
    }
}
<|endoftext|>


<commit_before>use num::{FromPrimitive, ToPrimitive};

use byteorder::{ByteOrder, LittleEndian};
use failure::{Error, Fail};
use zydis::gen::*;
use zydis::Decoder;
use log::{info};

use super::analysis::Analysis;
use super::arch::Arch;
use super::loader;
use super::loader::{LoadedModule, Loader};
use super::util;

#[derive(Debug, Fail)]
pub enum WorkspaceError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
    #[fail(display = "The given address is not mapped")]
    InvalidAddress,
    #[fail(display = "Buffer overrun")]
    BufferOverrun,
    #[fail(display = "The instruction at the given address is invalid")]
    InvalidInstruction,
}

pub struct WorkspaceBuilder<A: Arch> {
    filename: String,
    buf: Vec<u8>,

    loader: Option<Box<dyn Loader<A>>>,

    should_analyze: bool,
}

impl<A: Arch + 'static + std::fmt::Debug> WorkspaceBuilder<A> {
    /// Override the default loader picker with the given loader.
    pub fn with_loader(
        self: WorkspaceBuilder<A>,
        loader: Box<dyn Loader<A>>,
    ) -> WorkspaceBuilder<A> {
        info!("using explicitly chosen loader: {}", loader.get_name());
        WorkspaceBuilder {
            loader: Some(loader),
            ..self
        }
    }

    pub fn disable_analysis(self: WorkspaceBuilder<A>) -> WorkspaceBuilder<A> {
        info!("disabling analysis");
        WorkspaceBuilder {
            should_analyze: false,
            ..self
        }
    }

    /// Construct a workspace with the given builder configuration.
    ///
    /// This invokes the loaders, analyzers, and another other logic,
    ///  resulting in an initialized Workspace instance.
    ///
    /// Example (with default loader):
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.loader.get_name(),       "Windows/32/Raw");
    ///     assert_eq!(ws.module.base_address,     0x0);
    ///     assert_eq!(ws.module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    ///
    /// Example (with specific loader):
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader::Platform;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::loaders::sc::ShellcodeLoader;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(ShellcodeLoader::<Arch32>::new(Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.loader.get_name(),       "Windows/32/Raw");
    ///     assert_eq!(ws.module.base_address,     0x0);
    ///     assert_eq!(ws.module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn load(self: WorkspaceBuilder<A>) -> Result<Workspace<A>, Error> {
        // if the user provided a loader, use that.
        // otherwise, use the default detected loader.
        let (ldr, module, analyzers) = match self.loader {
            // TODO: let users specify analyzers via builder
            Some(ldr) => {
                let (module, analyzers) = ldr.load(&self.buf)?;
                (ldr, module, analyzers)
            }
            None => loader::load::<A>(&self.buf)?,
        };

        info!("loaded {} sections:", module.sections.len());
        module.sections.iter().for_each(|sec| {
            info!("  - {:8} {:>8x}", sec.name, sec.addr);
        });

        let analysis = Analysis::new(&module);

        let decoder = if A::get_bits() == 32 {
            Decoder::new(ZYDIS_MACHINE_MODE_LEGACY_32, ZYDIS_ADDRESS_WIDTH_32).unwrap()
        } else {
            Decoder::new(ZYDIS_MACHINE_MODE_LONG_64, ZYDIS_ADDRESS_WIDTH_64).unwrap()
        };

        let mut ws = Workspace::<A> {
            filename: self.filename,
            buf: self.buf,

            loader: ldr,
            module,

            decoder,

            analysis,
        };

        if self.should_analyze {
            for analyzer in analyzers.iter() {
                info!("analyzing with {}", analyzer.get_name());
                analyzer.analyze(&mut ws)?;
            }
        }

        Ok(ws)
    }
}

pub struct Workspace<A: Arch> {
    // name or source of the file
    pub filename: String,
    // raw bytes of the file
    pub buf: Vec<u8>,

    pub loader: Box<dyn Loader<A>>,
    pub module: LoadedModule<A>,

    pub decoder: Decoder,

    // pub only so that we can split the impl
    pub analysis: Analysis<A>,
}

impl<A: Arch + 'static> Workspace<A> {
    /// Create a workspace and load the given bytes.
    ///
    /// See example on `WorkspaceBuilder::load()`
    pub fn from_bytes(filename: &str, buf: &[u8]) -> WorkspaceBuilder<A> {
        WorkspaceBuilder {
            filename: filename.to_string(),
            buf: buf.to_vec(),
            loader: None,
            should_analyze: true,
        }
    }

    pub fn from_file(filename: &str) -> Result<WorkspaceBuilder<A>, Error> {
        Ok(WorkspaceBuilder {
            filename: filename.to_string(),
            buf: util::read_file(filename)?,
            loader: None,
            should_analyze: true,
        })
    }

    /// Read bytes from the given RVA.
    ///
    /// Errors:
    ///
    ///   - InvalidAddress - if the address is not mapped.
    ///   - BufferOverrun - if the requested region runs beyond the matching section.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\xEB\xFE");
    /// assert_eq!(ws.read_bytes(0x0, 0x1).unwrap().to_vec(), b"\xEB");
    /// assert_eq!(ws.read_bytes(0x1, 0x1).unwrap().to_vec(), b"\xFE");
    /// assert_eq!(ws.read_bytes(0x0, 0x2).unwrap().to_vec(), b"\xEB\xFE");
    /// assert_eq!(ws.read_bytes(0x0, 0x3).is_err(), true);
    /// assert_eq!(ws.read_bytes(0x2, 0x1).is_err(), true);
    /// ```
    pub fn read_bytes(&self, rva: A::RVA, length: usize) -> Result<&[u8], Error> {
        self.module
            .sections
            .iter()
            .filter(|section| section.contains(rva))
            .nth(0)
            .ok_or_else(|| WorkspaceError::InvalidAddress.into())
            .and_then(|section| -> Result<&[u8], Error> {
                // rva is guaranteed to be within this section,
                // so we can do an unchecked subtract here.
                let offset = rva - section.addr;
                A::RVA::to_usize(&offset)
                    .ok_or_else(|| WorkspaceError::InvalidAddress.into())
                    .and_then(|offset| {
                        if offset + length > section.buf.len() {
                            Err(WorkspaceError::BufferOverrun.into())
                        } else {
                            Ok(&section.buf[offset..offset + length])
                        }
                    })
            })
    }

    /// Is the given range mapped?
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\xEB\xFE");
    /// assert!( ws.probe(0x0, 1));
    /// assert!( ws.probe(0x0, 2));
    /// assert!(!ws.probe(0x0, 3));
    /// assert!( ws.probe(0x1, 1));
    /// assert!(!ws.probe(0x2, 1));
    /// ```
    pub fn probe(&self, rva: A::RVA, length: usize) -> bool {
        self.read_bytes(rva, length).is_ok()
    }

    /// Read a byte from the given RVA.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\xEB\xFE");
    /// assert_eq!(ws.read_u8(0x0).unwrap(), 0xEB);
    /// assert_eq!(ws.read_u8(0x1).unwrap(), 0xFE);
    /// assert_eq!(ws.read_u8(0x2).is_err(), true);
    /// ```
    pub fn read_u8(&self, rva: A::RVA) -> Result<u8, Error> {
        self.read_bytes(rva, 1).and_then(|buf| Ok(buf[0]))
    }

    /// Read a word from the given RVA.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\xEB\xFE");
    /// assert_eq!(ws.read_u16(0x0).unwrap(), 0xFEEB);
    /// assert_eq!(ws.read_u16(0x1).is_err(), true);
    /// assert_eq!(ws.read_u16(0x2).is_err(), true);
    /// ```
    pub fn read_u16(&self, rva: A::RVA) -> Result<u16, Error> {
        self.read_bytes(rva, 2)
            .and_then(|buf| Ok(LittleEndian::read_u16(buf)))
    }

    /// Read a dword from the given RVA.
    pub fn read_u32(&self, rva: A::RVA) -> Result<u32, Error> {
        self.read_bytes(rva, 4)
            .and_then(|buf| Ok(LittleEndian::read_u32(buf)))
    }

    /// Read a qword from the given RVA.
    pub fn read_u64(&self, rva: A::RVA) -> Result<u64, Error> {
        self.read_bytes(rva, 8)
            .and_then(|buf| Ok(LittleEndian::read_u64(buf)))
    }

    /// Read an RVA from the given RVA.
    /// Note that the size of the read is dependent on the architecture.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\x00\x11\x22\x33");
    /// assert_eq!(ws.read_rva(0x0).unwrap(), 0x33221100);
    /// assert_eq!(ws.read_rva(0x1).is_err(), true);
    /// ```
    pub fn read_rva(&self, rva: A::RVA) -> Result<A::RVA, Error> {
        match A::get_bits() {
            // these conversions will never fail
            32 => Ok(A::RVA::from_u32(self.read_u32(rva)?).unwrap()),
            64 => Ok(A::RVA::from_u64(self.read_u64(rva)?).unwrap()),
            _ => panic!("unexpected architecture"),
        }
    }

    /// Read a VA from the given RVA.
    /// Note that the size of the read is dependent on the architecture.
    pub fn read_va(&self, rva: A::RVA) -> Result<A::VA, Error> {
        match A::get_bits() {
            // these conversions will never fail
            32 => Ok(A::VA::from_u32(self.read_u32(rva)?).unwrap()),
            64 => Ok(A::VA::from_u64(self.read_u64(rva)?).unwrap()),
            _ => panic!("unexpected architecture"),
        }
    }

    /// Decode an instruction at the given RVA.
    ///
    /// Errors:
    ///
    ///   - InvalidAddress - if the address is not mapped.
    ///   - BufferOverrun - if the requested region runs beyond the matching section.
    ///   - InvalidInstruction - if an instruction cannot be decoded.
    ///
    /// Example:
    ///
    /// ```
    /// use zydis::gen::*;
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\xEB\xFE");
    /// assert_eq!(ws.read_insn(0x0).is_ok(), true);
    /// assert_eq!(ws.read_insn(0x0).unwrap().length, 2);
    /// assert_eq!(ws.read_insn(0x0).unwrap().mnemonic as i32, ZYDIS_MNEMONIC_JMP);
    /// ```
    pub fn read_insn(&self, rva: A::RVA) -> Result<ZydisDecodedInstruction, Error> {
        // this is `read_bytes` except that it reads at most 0x10 bytes.
        // if less are available, then less are returned.
        let buf = self
            .module
            .sections
            .iter()
            .filter(|section| section.contains(rva))
            .nth(0)
            .ok_or_else(|| WorkspaceError::InvalidAddress.into())
            .and_then(|section| -> Result<&[u8], Error> {
                // rva is guaranteed to be within this section,
                // so we can do an unchecked subtract here.
                let offset = rva - section.addr;
                A::RVA::to_usize(&offset)
                    .ok_or_else(|| WorkspaceError::InvalidAddress.into())
                    .and_then(|offset| {
                        if offset + 0x10 > section.buf.len() {
                            Ok(&section.buf[offset..])
                        } else {
                            Ok(&section.buf[offset..offset + 0x10])
                        }
                    })
            })?;

        // RVA will always be either u32 or u64 (never bigger),
        // so we can always fit this into a u64.
        let pc = A::RVA::to_u64(&rva).unwrap();

        match self.decoder.decode(&buf, pc) {
            Ok(Some(insn)) => Ok(insn),
            Ok(None) => Err(WorkspaceError::InvalidInstruction.into()),
            Err(_) => Err(WorkspaceError::InvalidInstruction.into()),
        }
    }

    // API:
    //   get_insn
    //   get_xrefs_to
    //   get_xrefs_from
    //   get_functions

    // elsewhere:
    //   call graph
    //   control flow graph

    // see: `analysis::impl Workspace`
}
<commit_msg>workspace: add `.rva()` helper<commit_after>use num::{FromPrimitive, ToPrimitive};

use byteorder::{ByteOrder, LittleEndian};
use failure::{Error, Fail};
use zydis::gen::*;
use zydis::Decoder;
use log::{info};

use super::analysis::Analysis;
use super::arch;
use super::arch::Arch;
use super::loader;
use super::loader::{LoadedModule, Loader};
use super::util;

#[derive(Debug, Fail)]
pub enum WorkspaceError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
    #[fail(display = "The given address is not mapped")]
    InvalidAddress,
    #[fail(display = "Buffer overrun")]
    BufferOverrun,
    #[fail(display = "The instruction at the given address is invalid")]
    InvalidInstruction,
}

pub struct WorkspaceBuilder<A: Arch> {
    filename: String,
    buf: Vec<u8>,

    loader: Option<Box<dyn Loader<A>>>,

    should_analyze: bool,
}

impl<A: Arch + 'static + std::fmt::Debug> WorkspaceBuilder<A> {
    /// Override the default loader picker with the given loader.
    pub fn with_loader(
        self: WorkspaceBuilder<A>,
        loader: Box<dyn Loader<A>>,
    ) -> WorkspaceBuilder<A> {
        info!("using explicitly chosen loader: {}", loader.get_name());
        WorkspaceBuilder {
            loader: Some(loader),
            ..self
        }
    }

    pub fn disable_analysis(self: WorkspaceBuilder<A>) -> WorkspaceBuilder<A> {
        info!("disabling analysis");
        WorkspaceBuilder {
            should_analyze: false,
            ..self
        }
    }

    /// Construct a workspace with the given builder configuration.
    ///
    /// This invokes the loaders, analyzers, and another other logic,
    ///  resulting in an initialized Workspace instance.
    ///
    /// Example (with default loader):
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.loader.get_name(),       "Windows/32/Raw");
    ///     assert_eq!(ws.module.base_address,     0x0);
    ///     assert_eq!(ws.module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    ///
    /// Example (with specific loader):
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader::Platform;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::loaders::sc::ShellcodeLoader;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(ShellcodeLoader::<Arch32>::new(Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.loader.get_name(),       "Windows/32/Raw");
    ///     assert_eq!(ws.module.base_address,     0x0);
    ///     assert_eq!(ws.module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn load(self: WorkspaceBuilder<A>) -> Result<Workspace<A>, Error> {
        // if the user provided a loader, use that.
        // otherwise, use the default detected loader.
        let (ldr, module, analyzers) = match self.loader {
            // TODO: let users specify analyzers via builder
            Some(ldr) => {
                let (module, analyzers) = ldr.load(&self.buf)?;
                (ldr, module, analyzers)
            }
            None => loader::load::<A>(&self.buf)?,
        };

        info!("loaded {} sections:", module.sections.len());
        module.sections.iter().for_each(|sec| {
            info!("  - {:8} {:>8x}", sec.name, sec.addr);
        });

        let analysis = Analysis::new(&module);

        let decoder = if A::get_bits() == 32 {
            Decoder::new(ZYDIS_MACHINE_MODE_LEGACY_32, ZYDIS_ADDRESS_WIDTH_32).unwrap()
        } else {
            Decoder::new(ZYDIS_MACHINE_MODE_LONG_64, ZYDIS_ADDRESS_WIDTH_64).unwrap()
        };

        let mut ws = Workspace::<A> {
            filename: self.filename,
            buf: self.buf,

            loader: ldr,
            module,

            decoder,

            analysis,
        };

        if self.should_analyze {
            for analyzer in analyzers.iter() {
                info!("analyzing with {}", analyzer.get_name());
                analyzer.analyze(&mut ws)?;
            }
        }

        Ok(ws)
    }
}

pub struct Workspace<A: Arch> {
    // name or source of the file
    pub filename: String,
    // raw bytes of the file
    pub buf: Vec<u8>,

    pub loader: Box<dyn Loader<A>>,
    pub module: LoadedModule<A>,

    pub decoder: Decoder,

    // pub only so that we can split the impl
    pub analysis: Analysis<A>,
}

impl<A: Arch + 'static> Workspace<A> {
    /// Create a workspace and load the given bytes.
    ///
    /// See example on `WorkspaceBuilder::load()`
    pub fn from_bytes(filename: &str, buf: &[u8]) -> WorkspaceBuilder<A> {
        WorkspaceBuilder {
            filename: filename.to_string(),
            buf: buf.to_vec(),
            loader: None,
            should_analyze: true,
        }
    }

    pub fn from_file(filename: &str) -> Result<WorkspaceBuilder<A>, Error> {
        Ok(WorkspaceBuilder {
            filename: filename.to_string(),
            buf: util::read_file(filename)?,
            loader: None,
            should_analyze: true,
        })
    }

    /// Read bytes from the given RVA.
    ///
    /// Errors:
    ///
    ///   - InvalidAddress - if the address is not mapped.
    ///   - BufferOverrun - if the requested region runs beyond the matching section.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\xEB\xFE");
    /// assert_eq!(ws.read_bytes(0x0, 0x1).unwrap().to_vec(), b"\xEB");
    /// assert_eq!(ws.read_bytes(0x1, 0x1).unwrap().to_vec(), b"\xFE");
    /// assert_eq!(ws.read_bytes(0x0, 0x2).unwrap().to_vec(), b"\xEB\xFE");
    /// assert_eq!(ws.read_bytes(0x0, 0x3).is_err(), true);
    /// assert_eq!(ws.read_bytes(0x2, 0x1).is_err(), true);
    /// ```
    pub fn read_bytes(&self, rva: A::RVA, length: usize) -> Result<&[u8], Error> {
        self.module
            .sections
            .iter()
            .filter(|section| section.contains(rva))
            .nth(0)
            .ok_or_else(|| WorkspaceError::InvalidAddress.into())
            .and_then(|section| -> Result<&[u8], Error> {
                // rva is guaranteed to be within this section,
                // so we can do an unchecked subtract here.
                let offset = rva - section.addr;
                A::RVA::to_usize(&offset)
                    .ok_or_else(|| WorkspaceError::InvalidAddress.into())
                    .and_then(|offset| {
                        if offset + length > section.buf.len() {
                            Err(WorkspaceError::BufferOverrun.into())
                        } else {
                            Ok(&section.buf[offset..offset + length])
                        }
                    })
            })
    }

    /// Is the given range mapped?
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\xEB\xFE");
    /// assert!( ws.probe(0x0, 1));
    /// assert!( ws.probe(0x0, 2));
    /// assert!(!ws.probe(0x0, 3));
    /// assert!( ws.probe(0x1, 1));
    /// assert!(!ws.probe(0x2, 1));
    /// ```
    pub fn probe(&self, rva: A::RVA, length: usize) -> bool {
        self.read_bytes(rva, length).is_ok()
    }

    /// Read a byte from the given RVA.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\xEB\xFE");
    /// assert_eq!(ws.read_u8(0x0).unwrap(), 0xEB);
    /// assert_eq!(ws.read_u8(0x1).unwrap(), 0xFE);
    /// assert_eq!(ws.read_u8(0x2).is_err(), true);
    /// ```
    pub fn read_u8(&self, rva: A::RVA) -> Result<u8, Error> {
        self.read_bytes(rva, 1).and_then(|buf| Ok(buf[0]))
    }

    /// Read a word from the given RVA.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\xEB\xFE");
    /// assert_eq!(ws.read_u16(0x0).unwrap(), 0xFEEB);
    /// assert_eq!(ws.read_u16(0x1).is_err(), true);
    /// assert_eq!(ws.read_u16(0x2).is_err(), true);
    /// ```
    pub fn read_u16(&self, rva: A::RVA) -> Result<u16, Error> {
        self.read_bytes(rva, 2)
            .and_then(|buf| Ok(LittleEndian::read_u16(buf)))
    }

    /// Read a dword from the given RVA.
    pub fn read_u32(&self, rva: A::RVA) -> Result<u32, Error> {
        self.read_bytes(rva, 4)
            .and_then(|buf| Ok(LittleEndian::read_u32(buf)))
    }

    /// Read a qword from the given RVA.
    pub fn read_u64(&self, rva: A::RVA) -> Result<u64, Error> {
        self.read_bytes(rva, 8)
            .and_then(|buf| Ok(LittleEndian::read_u64(buf)))
    }

    /// Read an RVA from the given RVA.
    /// Note that the size of the read is dependent on the architecture.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\x00\x11\x22\x33");
    /// assert_eq!(ws.read_rva(0x0).unwrap(), 0x33221100);
    /// assert_eq!(ws.read_rva(0x1).is_err(), true);
    /// ```
    pub fn read_rva(&self, rva: A::RVA) -> Result<A::RVA, Error> {
        match A::get_bits() {
            // these conversions will never fail
            32 => Ok(A::RVA::from_u32(self.read_u32(rva)?).unwrap()),
            64 => Ok(A::RVA::from_u64(self.read_u64(rva)?).unwrap()),
            _ => panic!("unexpected architecture"),
        }
    }

    /// Read a VA from the given RVA.
    /// Note that the size of the read is dependent on the architecture.
    pub fn read_va(&self, rva: A::RVA) -> Result<A::VA, Error> {
        match A::get_bits() {
            // these conversions will never fail
            32 => Ok(A::VA::from_u32(self.read_u32(rva)?).unwrap()),
            64 => Ok(A::VA::from_u64(self.read_u64(rva)?).unwrap()),
            _ => panic!("unexpected architecture"),
        }
    }

    /// Decode an instruction at the given RVA.
    ///
    /// Errors:
    ///
    ///   - InvalidAddress - if the address is not mapped.
    ///   - BufferOverrun - if the requested region runs beyond the matching section.
    ///   - InvalidInstruction - if an instruction cannot be decoded.
    ///
    /// Example:
    ///
    /// ```
    /// use zydis::gen::*;
    /// use lancelot::test;
    ///
    /// let ws = test::get_shellcode32_workspace(b"\xEB\xFE");
    /// assert_eq!(ws.read_insn(0x0).is_ok(), true);
    /// assert_eq!(ws.read_insn(0x0).unwrap().length, 2);
    /// assert_eq!(ws.read_insn(0x0).unwrap().mnemonic as i32, ZYDIS_MNEMONIC_JMP);
    /// ```
    pub fn read_insn(&self, rva: A::RVA) -> Result<ZydisDecodedInstruction, Error> {
        // this is `read_bytes` except that it reads at most 0x10 bytes.
        // if less are available, then less are returned.
        let buf = self
            .module
            .sections
            .iter()
            .filter(|section| section.contains(rva))
            .nth(0)
            .ok_or_else(|| WorkspaceError::InvalidAddress.into())
            .and_then(|section| -> Result<&[u8], Error> {
                // rva is guaranteed to be within this section,
                // so we can do an unchecked subtract here.
                let offset = rva - section.addr;
                A::RVA::to_usize(&offset)
                    .ok_or_else(|| WorkspaceError::InvalidAddress.into())
                    .and_then(|offset| {
                        if offset + 0x10 > section.buf.len() {
                            Ok(&section.buf[offset..])
                        } else {
                            Ok(&section.buf[offset..offset + 0x10])
                        }
                    })
            })?;

        // RVA will always be either u32 or u64 (never bigger),
        // so we can always fit this into a u64.
        let pc = A::RVA::to_u64(&rva).unwrap();

        match self.decoder.decode(&buf, pc) {
            Ok(Some(insn)) => Ok(insn),
            Ok(None) => Err(WorkspaceError::InvalidInstruction.into()),
            Err(_) => Err(WorkspaceError::InvalidInstruction.into()),
        }
    }

    pub fn rva(&self, va: A::VA) -> Option<A::RVA> {
        arch::va_compute_rva::<A>(self.module.base_address, va)
    }

    // API:
    //   get_insn
    //   get_xrefs_to
    //   get_xrefs_from
    //   get_functions

    // elsewhere:
    //   call graph
    //   control flow graph

    // see: `analysis::impl Workspace`
}
<|endoftext|>