<commit_before>use error_chain::ChainedError;
use ffmpeg::{Rational, format};
use ffmpeg::frame::Video as VideoFrame;
use fine_grained::Stopwatch;
use std::cell::RefCell;
use std::ptr;
use std::sync::{Mutex, ONCE_INIT, Once, RwLock};
use std::sync::mpsc::{Receiver, Sender, channel};
use std::thread;

use encode::{Encoder, EncoderParameters};
use errors::*;

lazy_static! {
    static ref CAPTURING: RwLock<bool> = RwLock::new(false);

    static ref ENCODER: Mutex<Option<Encoder>> = Mutex::new(None);

    /// Receives buffers to write pixels to.
    static ref BUF_RECEIVER: Mutex<Option<Receiver<Buffer>>> = Mutex::new(None);

    /// Sends events and frames to encode to the capture thread.
    static ref SEND_TO_CAPTURE_THREAD: Mutex<Option<Sender<CaptureThreadEvent>>> = Mutex::new(None);
}

thread_local! {
    static STOPWATCH: RefCell<Option<Stopwatch>> = RefCell::new(None);
}

enum CaptureThreadEvent {
    CaptureStart(EncoderParameters),
    CaptureStop,
    Frame((Buffer, f64)),
}

pub struct Buffer {
    data: Vec<u8>,
    width: u32,
    height: u32,
}

impl Buffer {
    fn new() -> Self {
        Self {
            data: Vec::new(),
            width: 0,
            height: 0,
        }
    }

    pub fn set_resolution(&mut self, width: u32, height: u32) {
        if self.width != width || self.height != height {
            println!("Changing resolution from {}×{} to {}×{}.",
                     self.width,
                     self.height,
                     width,
                     height);

            self.data.resize((width * height * 3) as usize, 0);
            self.width = width;
            self.height = height;
        }
    }

    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        self.data.as_mut_slice()
    }

    pub fn copy_to_frame(&self, frame: &mut VideoFrame) {
        // Make sure frame is of correct size.
        if self.width != frame.width() || self.height != frame.height() {
            *frame = VideoFrame::new(format::Pixel::RGB24, self.width, self.height);
        }

        // Copy the pixel data into the frame.
        let linesize = unsafe { ((*frame.as_ptr()).linesize[0]) as u32 };
        let mut data = frame.data_mut(0);

        for y in 0..self.height {
            unsafe {
                ptr::copy_nonoverlapping(self.data.as_ptr().offset((y * self.width * 3) as isize),
                                         data.as_mut_ptr()
                                             .offset(((self.height - y - 1) * linesize) as
                                                     isize),
                                         (self.width * 3) as usize);
            }
        }
    }
}

fn capture_thread(buf_sender: &Sender<Buffer>, event_receiver: &Receiver<CaptureThreadEvent>) {
    // Send the buffer to the game thread right away.
    buf_sender.send(Buffer::new()).unwrap();

    // This is our frame which will only be reallocated on resolution changes.
    let mut frame = VideoFrame::empty();

    // This is set to true on encoding error or cap_stop and reset on cap_start.
    // When this is true, ignore any received frames.
    let mut drop_frames = true;

    // Encoding parameters, set on CaptureStart.
    let mut parameters = None;

    // Event loop for the capture thread.
    loop {
        match event_receiver.recv().unwrap() {
            CaptureThreadEvent::CaptureStart(params) => {
                drop_frames = false;
                parameters = Some(params);
            }

            CaptureThreadEvent::CaptureStop => {
                *ENCODER.lock().unwrap() = None;
                drop_frames = true;
            }

            CaptureThreadEvent::Frame((buffer, frametime)) => {
                if drop_frames {
                    continue;
                }

                if let Err(ref e) = encode(buffer,
                                           frametime,
                                           &buf_sender,
                                           parameters.as_ref().unwrap(),
                                           &mut frame) {
                    *CAPTURING.write().unwrap() = false;
                    drop_frames = true;
                    println!("Encoding error: {}", e.display());
                }
            }
        }
    }
}

fn start_encoder(parameters: &EncoderParameters,
                 (width, height): (u32, u32))
                 -> Result<Encoder> {
    Encoder::start(&parameters, (width, height))
}

fn encode(buf: Buffer,
          frametime: f64,
          buf_sender: &Sender<Buffer>,
          parameters: &EncoderParameters,
          frame: &mut VideoFrame)
          -> Result<()> {
    buf.copy_to_frame(frame);

    // We're done with buf, now it can receive the next pack of pixels.
    buf_sender.send(buf).unwrap();

    // Let's encode the frame we just received.
    let mut encoder = ENCODER.lock().unwrap();

    // If the encoder wasn't initialized or if the frame size changed, initialize it.
    if encoder.as_ref()
              .map_or(true,
                      |enc| enc.width() != frame.width() || enc.height() != frame.height()) {
        *encoder = Some(start_encoder(&parameters,
                                      (frame.width(), frame.height()))
                            .chain_err(|| "could not start the video encoder")?);
    }

    // Encode the frame.
    encoder.as_mut()
           .unwrap()
           .take(frame, frametime)
           .chain_err(|| "could not encode the frame")?;

    Ok(())
}

pub fn initialize() {
    static INIT: Once = ONCE_INIT;
    INIT.call_once(|| {
        let (tx, rx) = channel::<Buffer>();
        let (tx2, rx2) = channel::<CaptureThreadEvent>();

        *BUF_RECEIVER.lock().unwrap() = Some(rx);
        *SEND_TO_CAPTURE_THREAD.lock().unwrap() = Some(tx2);

        thread::spawn(move || capture_thread(&tx, &rx2));
    });
}

pub fn get_buffer((width, height): (u32, u32)) -> Buffer {
    let mut buf = BUF_RECEIVER.lock()
                              .unwrap()
                              .as_ref()
                              .unwrap()
                              .recv()
                              .unwrap();

    buf.set_resolution(width, height);

    buf
}

pub fn capture(buf: Buffer, frametime: f64) {
    SEND_TO_CAPTURE_THREAD.lock()
                          .unwrap()
                          .as_ref()
                          .unwrap()
                          .send(CaptureThreadEvent::Frame((buf, frametime)))
                          .unwrap();
}

pub fn is_capturing() -> bool {
    *CAPTURING.read().unwrap()
}

pub fn capture_block_start() {
    STOPWATCH.with(|sw| if let Some(sw) = sw.borrow_mut().as_mut() {
                       sw.start();
                   });
}

pub fn capture_block_end() {
    STOPWATCH.with(|sw| if let Some(sw) = sw.borrow_mut().as_mut() {
                       sw.lap();
                       sw.stop();
                   });
}

/// Parses the given string and returns a time base.
///
/// The string can be in one of the two formats:
/// - `<i32 FPS>` - treated as an integer FPS value,
/// - `<i32 a> <i32 b>` - treated as a fractional `a/b` FPS value.
fn parse_fps(string: &str) -> Option<Rational> {
    if let Ok(fps) = string.parse() {
        return Some((1, fps).into());
    }

    let mut split = string.splitn(2, ' ');
    if let Some(den) = split.next().and_then(|s| s.parse().ok()) {
        if let Some(num) = split.next().and_then(|s| s.parse().ok()) {
            return Some((num, den).into());
        }
    }

    None
}

command!(cap_start, |mut engine| {
    let mut parameters = EncoderParameters {
        bitrate: 0,
        crf: String::new(),
        filename: String::new(),
        preset: String::new(),
        time_base: Rational::new(1, 1),
    };

    match cap_filename.get(&engine)
                      .to_string(&mut engine)
                      .chain_err(|| "invalid cap_filename") {
        Ok(filename) => parameters.filename = filename,
        Err(ref e) => {
            engine.con_print(&format!("{}", e.display()));
            return;
        }
    };

    if let Some(time_base) = cap_fps.get(&engine)
                                    .to_string(&mut engine)
                                    .ok()
                                    .and_then(|s| parse_fps(&s)) {
        parameters.time_base = time_base;
    } else {
        engine.con_print("Invalid cap_fps.\n");
        return;
    }

    match cap_crf.get(&engine)
                 .parse(&mut engine)
                 .chain_err(|| "invalid cap_crf") {
        Ok(crf) => parameters.crf = crf,
        Err(ref e) => {
            engine.con_print(&format!("{}", e.display()));
            return;
        }
    };

    match cap_bitrate.get(&engine)
                     .parse(&mut engine)
                     .chain_err(|| "invalid cap_bitrate") {
        Ok(bitrate) => parameters.bitrate = bitrate,
        Err(ref e) => {
            engine.con_print(&format!("{}", e.display()));
            return;
        }
    };

    match cap_preset.get(&engine)
                    .parse(&mut engine)
                    .chain_err(|| "invalid cap_preset") {
        Ok(preset) => parameters.preset = preset,
        Err(ref e) => {
            engine.con_print(&format!("{}", e.display()));
            return;
        }
    };

    *CAPTURING.write().unwrap() = true;

    SEND_TO_CAPTURE_THREAD.lock()
                          .unwrap()
                          .as_ref()
                          .unwrap()
                          .send(CaptureThreadEvent::CaptureStart(parameters))
                          .unwrap();

    STOPWATCH.with(|sw| *sw.borrow_mut() = Some(Stopwatch::new()));
});

command!(cap_stop, |engine| {
    *CAPTURING.write().unwrap() = false;

    SEND_TO_CAPTURE_THREAD.lock().unwrap()
        .as_ref().unwrap()
        .send(CaptureThreadEvent::CaptureStop).unwrap();

    STOPWATCH.with(|sw| if let Some(sw) = sw.borrow_mut().take() {
        let frames = sw.number_of_laps();

        if frames > 0 {
            engine.con_print(&format!("Captured {} frames in {} seconds (~{} msec of overhead per frame)\n",
                                      frames,
                                      sw.total_time() as f64 / 1_000_000_000f64,
                                      (sw.total_time() / frames as u64) as f64 / 1_000_000f64));
        }
    });
});

cvar!(cap_bitrate, "0");
cvar!(cap_crf, "15");
cvar!(cap_filename, "capture.mp4");
cvar!(cap_fps, "60");
cvar!(cap_preset, "veryfast");
<commit_msg>Rename cap_preset to cap_x264_preset<commit_after>use error_chain::ChainedError;
use ffmpeg::{Rational, format};
use ffmpeg::frame::Video as VideoFrame;
use fine_grained::Stopwatch;
use std::cell::RefCell;
use std::ptr;
use std::sync::{Mutex, ONCE_INIT, Once, RwLock};
use std::sync::mpsc::{Receiver, Sender, channel};
use std::thread;

use encode::{Encoder, EncoderParameters};
use errors::*;

lazy_static! {
    static ref CAPTURING: RwLock<bool> = RwLock::new(false);

    static ref ENCODER: Mutex<Option<Encoder>> = Mutex::new(None);

    /// Receives buffers to write pixels to.
    static ref BUF_RECEIVER: Mutex<Option<Receiver<Buffer>>> = Mutex::new(None);

    /// Sends events and frames to encode to the capture thread.
    static ref SEND_TO_CAPTURE_THREAD: Mutex<Option<Sender<CaptureThreadEvent>>> = Mutex::new(None);
}

thread_local! {
    static STOPWATCH: RefCell<Option<Stopwatch>> = RefCell::new(None);
}

enum CaptureThreadEvent {
    CaptureStart(EncoderParameters),
    CaptureStop,
    Frame((Buffer, f64)),
}

pub struct Buffer {
    data: Vec<u8>,
    width: u32,
    height: u32,
}

impl Buffer {
    fn new() -> Self {
        Self {
            data: Vec::new(),
            width: 0,
            height: 0,
        }
    }

    pub fn set_resolution(&mut self, width: u32, height: u32) {
        if self.width != width || self.height != height {
            println!("Changing resolution from {}×{} to {}×{}.",
                     self.width,
                     self.height,
                     width,
                     height);

            self.data.resize((width * height * 3) as usize, 0);
            self.width = width;
            self.height = height;
        }
    }

    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        self.data.as_mut_slice()
    }

    pub fn copy_to_frame(&self, frame: &mut VideoFrame) {
        // Make sure frame is of correct size.
        if self.width != frame.width() || self.height != frame.height() {
            *frame = VideoFrame::new(format::Pixel::RGB24, self.width, self.height);
        }

        // Copy the pixel data into the frame.
        let linesize = unsafe { ((*frame.as_ptr()).linesize[0]) as u32 };
        let mut data = frame.data_mut(0);

        for y in 0..self.height {
            unsafe {
                ptr::copy_nonoverlapping(self.data.as_ptr().offset((y * self.width * 3) as isize),
                                         data.as_mut_ptr()
                                             .offset(((self.height - y - 1) * linesize) as
                                                     isize),
                                         (self.width * 3) as usize);
            }
        }
    }
}

fn capture_thread(buf_sender: &Sender<Buffer>, event_receiver: &Receiver<CaptureThreadEvent>) {
    // Send the buffer to the game thread right away.
    buf_sender.send(Buffer::new()).unwrap();

    // This is our frame which will only be reallocated on resolution changes.
    let mut frame = VideoFrame::empty();

    // This is set to true on encoding error or cap_stop and reset on cap_start.
    // When this is true, ignore any received frames.
    let mut drop_frames = true;

    // Encoding parameters, set on CaptureStart.
    let mut parameters = None;

    // Event loop for the capture thread.
    loop {
        match event_receiver.recv().unwrap() {
            CaptureThreadEvent::CaptureStart(params) => {
                drop_frames = false;
                parameters = Some(params);
            }

            CaptureThreadEvent::CaptureStop => {
                *ENCODER.lock().unwrap() = None;
                drop_frames = true;
            }

            CaptureThreadEvent::Frame((buffer, frametime)) => {
                if drop_frames {
                    continue;
                }

                if let Err(ref e) = encode(buffer,
                                           frametime,
                                           &buf_sender,
                                           parameters.as_ref().unwrap(),
                                           &mut frame) {
                    *CAPTURING.write().unwrap() = false;
                    drop_frames = true;
                    println!("Encoding error: {}", e.display());
                }
            }
        }
    }
}

fn start_encoder(parameters: &EncoderParameters,
                 (width, height): (u32, u32))
                 -> Result<Encoder> {
    Encoder::start(&parameters, (width, height))
}

fn encode(buf: Buffer,
          frametime: f64,
          buf_sender: &Sender<Buffer>,
          parameters: &EncoderParameters,
          frame: &mut VideoFrame)
          -> Result<()> {
    buf.copy_to_frame(frame);

    // We're done with buf, now it can receive the next pack of pixels.
    buf_sender.send(buf).unwrap();

    // Let's encode the frame we just received.
    let mut encoder = ENCODER.lock().unwrap();

    // If the encoder wasn't initialized or if the frame size changed, initialize it.
    if encoder.as_ref()
              .map_or(true,
                      |enc| enc.width() != frame.width() || enc.height() != frame.height()) {
        *encoder = Some(start_encoder(&parameters,
                                      (frame.width(), frame.height()))
                            .chain_err(|| "could not start the video encoder")?);
    }

    // Encode the frame.
    encoder.as_mut()
           .unwrap()
           .take(frame, frametime)
           .chain_err(|| "could not encode the frame")?;

    Ok(())
}

pub fn initialize() {
    static INIT: Once = ONCE_INIT;
    INIT.call_once(|| {
        let (tx, rx) = channel::<Buffer>();
        let (tx2, rx2) = channel::<CaptureThreadEvent>();

        *BUF_RECEIVER.lock().unwrap() = Some(rx);
        *SEND_TO_CAPTURE_THREAD.lock().unwrap() = Some(tx2);

        thread::spawn(move || capture_thread(&tx, &rx2));
    });
}

pub fn get_buffer((width, height): (u32, u32)) -> Buffer {
    let mut buf = BUF_RECEIVER.lock()
                              .unwrap()
                              .as_ref()
                              .unwrap()
                              .recv()
                              .unwrap();

    buf.set_resolution(width, height);

    buf
}

pub fn capture(buf: Buffer, frametime: f64) {
    SEND_TO_CAPTURE_THREAD.lock()
                          .unwrap()
                          .as_ref()
                          .unwrap()
                          .send(CaptureThreadEvent::Frame((buf, frametime)))
                          .unwrap();
}

pub fn is_capturing() -> bool {
    *CAPTURING.read().unwrap()
}

pub fn capture_block_start() {
    STOPWATCH.with(|sw| if let Some(sw) = sw.borrow_mut().as_mut() {
                       sw.start();
                   });
}

pub fn capture_block_end() {
    STOPWATCH.with(|sw| if let Some(sw) = sw.borrow_mut().as_mut() {
                       sw.lap();
                       sw.stop();
                   });
}

/// Parses the given string and returns a time base.
///
/// The string can be in one of the two formats:
/// - `<i32 FPS>` - treated as an integer FPS value,
/// - `<i32 a> <i32 b>` - treated as a fractional `a/b` FPS value.
fn parse_fps(string: &str) -> Option<Rational> {
    if let Ok(fps) = string.parse() {
        return Some((1, fps).into());
    }

    let mut split = string.splitn(2, ' ');
    if let Some(den) = split.next().and_then(|s| s.parse().ok()) {
        if let Some(num) = split.next().and_then(|s| s.parse().ok()) {
            return Some((num, den).into());
        }
    }

    None
}

command!(cap_start, |mut engine| {
    let mut parameters = EncoderParameters {
        bitrate: 0,
        crf: String::new(),
        filename: String::new(),
        preset: String::new(),
        time_base: Rational::new(1, 1),
    };

    match cap_filename.get(&engine)
                      .to_string(&mut engine)
                      .chain_err(|| "invalid cap_filename") {
        Ok(filename) => parameters.filename = filename,
        Err(ref e) => {
            engine.con_print(&format!("{}", e.display()));
            return;
        }
    };

    if let Some(time_base) = cap_fps.get(&engine)
                                    .to_string(&mut engine)
                                    .ok()
                                    .and_then(|s| parse_fps(&s)) {
        parameters.time_base = time_base;
    } else {
        engine.con_print("Invalid cap_fps.\n");
        return;
    }

    match cap_crf.get(&engine)
                 .parse(&mut engine)
                 .chain_err(|| "invalid cap_crf") {
        Ok(crf) => parameters.crf = crf,
        Err(ref e) => {
            engine.con_print(&format!("{}", e.display()));
            return;
        }
    };

    match cap_bitrate.get(&engine)
                     .parse(&mut engine)
                     .chain_err(|| "invalid cap_bitrate") {
        Ok(bitrate) => parameters.bitrate = bitrate,
        Err(ref e) => {
            engine.con_print(&format!("{}", e.display()));
            return;
        }
    };

    match cap_x264_preset.get(&engine)
                         .parse(&mut engine)
                         .chain_err(|| "invalid cap_x264_preset") {
        Ok(preset) => parameters.preset = preset,
        Err(ref e) => {
            engine.con_print(&format!("{}", e.display()));
            return;
        }
    };

    *CAPTURING.write().unwrap() = true;

    SEND_TO_CAPTURE_THREAD.lock()
                          .unwrap()
                          .as_ref()
                          .unwrap()
                          .send(CaptureThreadEvent::CaptureStart(parameters))
                          .unwrap();

    STOPWATCH.with(|sw| *sw.borrow_mut() = Some(Stopwatch::new()));
});

command!(cap_stop, |engine| {
    *CAPTURING.write().unwrap() = false;

    SEND_TO_CAPTURE_THREAD.lock().unwrap()
        .as_ref().unwrap()
        .send(CaptureThreadEvent::CaptureStop).unwrap();

    STOPWATCH.with(|sw| if let Some(sw) = sw.borrow_mut().take() {
        let frames = sw.number_of_laps();

        if frames > 0 {
            engine.con_print(&format!("Captured {} frames in {} seconds (~{} msec of overhead per frame)\n",
                                      frames,
                                      sw.total_time() as f64 / 1_000_000_000f64,
                                      (sw.total_time() / frames as u64) as f64 / 1_000_000f64));
        }
    });
});

cvar!(cap_bitrate, "0");
cvar!(cap_crf, "15");
cvar!(cap_filename, "capture.mp4");
cvar!(cap_fps, "60");
cvar!(cap_x264_preset, "veryfast");
<|endoftext|>


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

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

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

/// The log level enum
///
/// It is used with the logger struct, and is used to determine the severity of the message.
pub enum LogLevel {

    /// Used for logging debug messages (stripped out of release builds)
    Debug   (String),
    /// Used for normal logging messages
    Message (String),
    /// Used for warnings (non critical errors)
    Warning (String),
    /// Used for critical errors
    Error   (String)
}

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

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

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

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

impl Logger {

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

        match File::create (&log_file_path) {

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

            Err (_) => {

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

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

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

    /// Starts file logging
    // TODO: Write begin log string
    pub fn begin_log (&mut self) {

        self.is_logging = true;
    }

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

    /// Logs a message
    ///
    /// # Examples
    /// ```
    /// let logger = Logger::new ("log.txt".to_string ());
    ///
    /// logger.log (LogLevel::Debug   ("This is a log entry".to_string ())) // Debug
    /// logger.log (LogLevel::Message ("This is a log entry".to_string ())) // Message
    /// logger.log (LogLevel::Warning ("This is a log entry".to_string ())) // Warning
    /// logger.log (LogLevel::Error   ("This is a log entry".to_string ())) // Error
    // TODO : Add log level specific messages (ie: DEBUG: "Insert message here")
    pub fn log (&mut self, log_level : LogLevel) {

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

            // Get the buffer from the logger struct
            let buffer = self.log_file.as_mut ().unwrap ();

            match log_level {

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

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

                        buffer.write (m.as_bytes ()).unwrap ();
                    }
                },

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

                    buffer.write (m.as_bytes ()).unwrap ();
                },

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

                    buffer.write (m.as_bytes ()).unwrap ();
                },

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

                    buffer.write (m.as_bytes ()).unwrap ();
                }
            }
        }
    }

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

    /// Stops file logging
    // TODO: Write end log string
    pub fn end_log (&mut self) {

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

extern crate time;

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

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

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

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

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

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

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

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

impl Logger {

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

        match File::create (&log_file_path) {

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

            Err (_) => {

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

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

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

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

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

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

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

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

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

            match log_level {

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

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

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

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

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

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

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

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

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

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

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

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


<commit_before>use buffer;
use buffer::Buffer;
use std::old_io::IoError;

pub struct Workspace {
    path: Path,
    buffers: Vec<Buffer>,
    current_buffer_index: Option<usize>,
}

impl Workspace {
    pub fn add_buffer(&mut self, buf: Buffer) {
        self.buffers.push(buf);
        self.current_buffer_index = Some(self.buffers.len()-1);
    }

    pub fn current_buffer(&mut self) -> Option<&mut Buffer> {
        match self.current_buffer_index {
            Some(index) => Some(&mut self.buffers[index]),
            None => None,
        }
    }

    pub fn close_current_buffer(&mut self) {
        match self.current_buffer_index {
            Some(index) => {
                self.buffers.remove(index);

                if self.buffers.is_empty() {
                    self.current_buffer_index = None;
                } else {
                    self.current_buffer_index = Some(self.buffers.len()-1);
                }
            }
            None => return,
        };
    }
}

pub fn new(path: Path) -> Workspace {
    Workspace{ path: path, buffers: Vec::new(), current_buffer_index: None }
}

#[cfg(test)]
mod tests {
    use super::new;
    use buffer;

    #[test]
    fn add_buffer_adds_and_selects_the_passed_buffer() {
        let mut workspace = new(Path::new("tests/sample"));
        let buf = buffer::from_file(Path::new("tests/sample/file")).unwrap();
        workspace.add_buffer(buf);

        assert_eq!(workspace.buffers.len(), 1);
        assert_eq!(workspace.current_buffer().unwrap().data(), "it works!\n");
    }

    #[test]
    fn current_buffer_returns_none_when_there_are_no_buffers() {
        let mut workspace = new(Path::new("tests/sample"));
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn current_buffer_returns_one_when_there_are_buffers() {
        let mut workspace = new(Path::new("tests/sample"));
        let buf = buffer::from_file(Path::new("tests/sample/file")).unwrap();
        workspace.add_buffer(buf);
        assert!(workspace.current_buffer().is_some());
    }

    #[test]
    fn close_current_buffer_does_nothing_when_none_are_open() {
        let mut workspace = new(Path::new("tests/sample"));
        workspace.close_current_buffer();
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn close_current_buffer_cleans_up_when_only_one_buffer_is_open() {
        let mut workspace = new(Path::new("tests/sample"));
        workspace.add_buffer(buffer::new());
        workspace.close_current_buffer();
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn close_current_buffer_when_two_are_open_selects_the_other() {
        let mut workspace = new(Path::new("tests/sample"));

        // Create two buffers and add them to the workspace.
        let mut first_buffer = buffer::new();
        let mut second_buffer = buffer::new();
        first_buffer.insert("first buffer");
        second_buffer.insert("second buffer");
        workspace.add_buffer(first_buffer);
        workspace.add_buffer(second_buffer);

        // Ensure that the second buffer is currently selected.
        assert_eq!(workspace.current_buffer().unwrap().data(), "second buffer");

        workspace.close_current_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "first buffer");
    }
}
<commit_msg>Ensure that current_buffer_index is cleared when last buffer is closed.<commit_after>use buffer;
use buffer::Buffer;
use std::old_io::IoError;

pub struct Workspace {
    path: Path,
    buffers: Vec<Buffer>,
    current_buffer_index: Option<usize>,
}

impl Workspace {
    pub fn add_buffer(&mut self, buf: Buffer) {
        self.buffers.push(buf);
        self.current_buffer_index = Some(self.buffers.len()-1);
    }

    pub fn current_buffer(&mut self) -> Option<&mut Buffer> {
        match self.current_buffer_index {
            Some(index) => Some(&mut self.buffers[index]),
            None => None,
        }
    }

    pub fn close_current_buffer(&mut self) {
        match self.current_buffer_index {
            Some(index) => {
                self.buffers.remove(index);

                if self.buffers.is_empty() {
                    self.current_buffer_index = None;
                } else {
                    self.current_buffer_index = Some(self.buffers.len()-1);
                }
            }
            None => return,
        };
    }
}

pub fn new(path: Path) -> Workspace {
    Workspace{ path: path, buffers: Vec::new(), current_buffer_index: None }
}

#[cfg(test)]
mod tests {
    use super::new;
    use buffer;

    #[test]
    fn add_buffer_adds_and_selects_the_passed_buffer() {
        let mut workspace = new(Path::new("tests/sample"));
        let buf = buffer::from_file(Path::new("tests/sample/file")).unwrap();
        workspace.add_buffer(buf);

        assert_eq!(workspace.buffers.len(), 1);
        assert_eq!(workspace.current_buffer().unwrap().data(), "it works!\n");
    }

    #[test]
    fn current_buffer_returns_none_when_there_are_no_buffers() {
        let mut workspace = new(Path::new("tests/sample"));
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn current_buffer_returns_one_when_there_are_buffers() {
        let mut workspace = new(Path::new("tests/sample"));
        let buf = buffer::from_file(Path::new("tests/sample/file")).unwrap();
        workspace.add_buffer(buf);
        assert!(workspace.current_buffer().is_some());
    }

    #[test]
    fn close_current_buffer_does_nothing_when_none_are_open() {
        let mut workspace = new(Path::new("tests/sample"));
        workspace.close_current_buffer();
        assert!(workspace.current_buffer().is_none());
    }

    #[test]
    fn close_current_buffer_cleans_up_when_only_one_buffer_is_open() {
        let mut workspace = new(Path::new("tests/sample"));
        workspace.add_buffer(buffer::new());
        workspace.close_current_buffer();
        assert!(workspace.current_buffer().is_none());
        assert!(workspace.current_buffer_index.is_none());
    }

    #[test]
    fn close_current_buffer_when_two_are_open_selects_the_other() {
        let mut workspace = new(Path::new("tests/sample"));

        // Create two buffers and add them to the workspace.
        let mut first_buffer = buffer::new();
        let mut second_buffer = buffer::new();
        first_buffer.insert("first buffer");
        second_buffer.insert("second buffer");
        workspace.add_buffer(first_buffer);
        workspace.add_buffer(second_buffer);

        // Ensure that the second buffer is currently selected.
        assert_eq!(workspace.current_buffer().unwrap().data(), "second buffer");

        workspace.close_current_buffer();
        assert_eq!(workspace.current_buffer().unwrap().data(), "first buffer");
    }
}
<|endoftext|>


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

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

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

#[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>>>,
}

impl<A: Arch + 'static> 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> {
        WorkspaceBuilder {
            loader: Some(loader),
            .. 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;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::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) = match self.loader {
            Some(ldr) => {
                let module = ldr.load(&self.buf)?;
                (ldr, module)
            }
            None => loader::load::<A>(&self.buf)?,
        };

        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()
        };

        Ok(Workspace::<A>{
            filename: self.filename,
            buf: self.buf,

            loader: ldr,
            module: module,

            decoder: decoder,

            analysis: analysis,
        })
    }
}

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,
        }
    }

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

    /// 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::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     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);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    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(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(WorkspaceError::InvalidAddress.into())
                    .and_then(|offset| {
                        if offset + length > section.buf.len() {
                            Err(WorkspaceError::BufferOverrun.into())
                        } else {
                            Ok(&section.buf[offset..offset+length])
                        }
                    })
            })
    }

    /// Read a byte from the given RVA.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     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);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    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::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     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);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    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::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\x00\x11\x22\x33")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.read_rva(0x0).unwrap(), 0x33221100);
    ///     assert_eq!(ws.read_rva(0x1).is_err(), true);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    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::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     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);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    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(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(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 `probe()` accessor<commit_after>use num::{ToPrimitive, FromPrimitive};

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

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

#[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>>>,
}

impl<A: Arch + 'static> 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> {
        WorkspaceBuilder {
            loader: Some(loader),
            .. 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;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::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) = match self.loader {
            Some(ldr) => {
                let module = ldr.load(&self.buf)?;
                (ldr, module)
            }
            None => loader::load::<A>(&self.buf)?,
        };

        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()
        };

        Ok(Workspace::<A>{
            filename: self.filename,
            buf: self.buf,

            loader: ldr,
            module: module,

            decoder: decoder,

            analysis: analysis,
        })
    }
}

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,
        }
    }

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

    /// 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::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     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);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    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(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(WorkspaceError::InvalidAddress.into())
                    .and_then(|offset| {
                        if offset + length > section.buf.len() {
                            Err(WorkspaceError::BufferOverrun.into())
                        } else {
                            Ok(&section.buf[offset..offset+length])
                        }
                    })
            })
    }

    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::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     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);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    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::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     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);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    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::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\x00\x11\x22\x33")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.read_rva(0x0).unwrap(), 0x33221100);
    ///     assert_eq!(ws.read_rva(0x1).is_err(), true);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    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::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     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);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    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(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(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`
}
<|endoftext|>


<commit_before>//! This modules contains an implementation of [r2d2](https://github.com/sfackler/r2d2)
//! functionality of connection pools. To get more details about creating r2d2 pools
//! please refer to original documentation.
use std::iter::Iterator;
use query::QueryBuilder;
use client::{CDRS, Session};
use error::{Error as CError, Result as CResult};
use authenticators::Authenticator;
use compression::Compression;
use r2d2;
use transport::CDRSTransport;
use rand;
use std::sync::atomic::{AtomicIsize, Ordering};

/// Load balancing strategy
#[derive(PartialEq)]
pub enum LoadBalancingStrategy {
    /// Round Robin balancing strategy
    RoundRobin,
    /// Random balancing strategy
    Random,
}

impl LoadBalancingStrategy {
    /// Returns next value for selected load balancing strategy
    pub fn next<'a, N>(&'a self, nodes: &'a Vec<N>, i: usize) -> Option<&N> {
        println!("node# {:?}", i);
        match *self {
            LoadBalancingStrategy::Random => nodes.get(self.rnd_idx((0, Some(nodes.len())))),
            LoadBalancingStrategy::RoundRobin => {
                let mut cycle = nodes.iter().cycle().skip(i);
                cycle.next()
            }
        }
    }

    /// Returns random number from a range
    fn rnd_idx(&self, bounds: (usize, Option<usize>)) -> usize {
        let min = bounds.0;
        let max = bounds.1.unwrap_or(u8::max_value() as usize);
        let rnd = rand::random::<usize>();
        min + rnd * (max - min) / (u8::max_value() as usize)
    }
}

/// Load balancer
///
/// #Example
///
/// ```no_run
/// use cdrs::cluster::{LoadBalancingStrategy, LoadBalancer};
/// use cdrs::transport::TransportTcp;
/// let transports = vec![TransportTcp::new("127.0.0.1:9042"), TransportTcp::new("127.0.0.1:9042")];
/// let load_balancer = LoadBalancer::new(transports, LoadBalancingStrategy::RoundRobin);
/// let node = load_balancer.next().unwrap();
/// ```
pub struct LoadBalancer<T> {
    strategy: LoadBalancingStrategy,
    nodes: Vec<T>,
    i: AtomicIsize,
}

impl<T> LoadBalancer<T> {
    /// Factory function which creates new `LoadBalancer` with provided strategy.
    pub fn new(nodes: Vec<T>, strategy: LoadBalancingStrategy) -> LoadBalancer<T> {
        LoadBalancer {
            nodes: nodes,
            strategy: strategy,
            i: AtomicIsize::new(0),
        }
    }

    /// Returns next node basing on provided strategy.
    pub fn next(&self) -> Option<&T> {
        let next = self.strategy
            .next(&self.nodes, self.i.load(Ordering::Relaxed) as usize);
        if self.strategy == LoadBalancingStrategy::RoundRobin {
            self.i.fetch_add(1, Ordering::Relaxed);
        }
        next
    }
}

/// [r2d2](https://github.com/sfackler/r2d2) `ManageConnection`.
pub struct ClusterConnectionManager<T, X> {
    load_balancer: LoadBalancer<X>,
    authenticator: T,
    compression: Compression,
}

impl<T, X> ClusterConnectionManager<T, X>
    where T: Authenticator + Send + Sync + 'static
{
    /// Creates a new instance of `ConnectionManager`.
    /// It requires transport, authenticator and compression as inputs.
    pub fn new(load_balancer: LoadBalancer<X>,
               authenticator: T,
               compression: Compression)
               -> ClusterConnectionManager<T, X> {
        ClusterConnectionManager {
            load_balancer: load_balancer,
            authenticator: authenticator,
            compression: compression,
        }
    }
}

impl<T: Authenticator + Send + Sync + 'static, X: CDRSTransport + Send + Sync + 'static>
r2d2::ManageConnection for ClusterConnectionManager<T, X> {
    type Connection = Session<T,X>;
    type Error = CError;

    fn connect(&self) -> Result<Self::Connection, Self::Error> {
        let transport_res: CResult<X> = self.load_balancer.next()
            .ok_or_else(|| "Cannot get next node".into())
            .and_then(|x| x.try_clone().map_err(|e| e.into()));
        let transport = try!(transport_res);
        let compression = self.compression;
        let cdrs = CDRS::new(transport, self.authenticator.clone());

        cdrs.start(compression)
    }

    fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), Self::Error> {
        let query = QueryBuilder::new("SELECT * FROM system.peers;").finalize();

        connection.query(query, false, false).map(|_| ())
    }

    fn has_broken(&self, _connection: &mut Self::Connection) -> bool {
        false
    }
}

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

    #[test]
    fn round_robin() {
        let nodes = vec!["a", "b", "c"];
        let nodes_c = nodes.clone();
        let load_balancer = LoadBalancer::new(nodes, LoadBalancingStrategy::RoundRobin);
        for i in 0..10 {
            assert_eq!(&nodes_c[i % 3], load_balancer.next().unwrap());
        }
    }
}
<commit_msg>Fix round robin next value evaluation<commit_after>//! This modules contains an implementation of [r2d2](https://github.com/sfackler/r2d2)
//! functionality of connection pools. To get more details about creating r2d2 pools
//! please refer to original documentation.
use std::iter::Iterator;
use query::QueryBuilder;
use client::{CDRS, Session};
use error::{Error as CError, Result as CResult};
use authenticators::Authenticator;
use compression::Compression;
use r2d2;
use transport::CDRSTransport;
use rand;
use std::sync::atomic::{AtomicUsize, Ordering};

/// Load balancing strategy
#[derive(PartialEq)]
pub enum LoadBalancingStrategy {
    /// Round Robin balancing strategy
    RoundRobin,
    /// Random balancing strategy
    Random,
}

impl LoadBalancingStrategy {
    /// Returns next value for selected load balancing strategy
    pub fn next<'a, N>(&'a self, nodes: &'a Vec<N>, i: usize) -> Option<&N> {
        println!("node# {:?}", i);
        match *self {
            LoadBalancingStrategy::Random => nodes.get(self.rnd_idx((0, Some(nodes.len())))),
            LoadBalancingStrategy::RoundRobin => {
                let mut cycle = nodes.iter().cycle().skip(i);
                cycle.next()
            }
        }
    }

    /// Returns random number from a range
    fn rnd_idx(&self, bounds: (usize, Option<usize>)) -> usize {
        let min = bounds.0;
        let max = bounds.1.unwrap_or(u8::max_value() as usize);
        let rnd = rand::random::<usize>();
        min + rnd * (max - min) / (u8::max_value() as usize)
    }
}

/// Load balancer
///
/// #Example
///
/// ```no_run
/// use cdrs::cluster::{LoadBalancingStrategy, LoadBalancer};
/// use cdrs::transport::TransportTcp;
/// let transports = vec![TransportTcp::new("127.0.0.1:9042"), TransportTcp::new("127.0.0.1:9042")];
/// let load_balancer = LoadBalancer::new(transports, LoadBalancingStrategy::RoundRobin);
/// let node = load_balancer.next().unwrap();
/// ```
pub struct LoadBalancer<T> {
    strategy: LoadBalancingStrategy,
    nodes: Vec<T>,
    i: AtomicUsize,
}

impl<T> LoadBalancer<T> {
    /// Factory function which creates new `LoadBalancer` with provided strategy.
    pub fn new(nodes: Vec<T>, strategy: LoadBalancingStrategy) -> LoadBalancer<T> {
        LoadBalancer {
            nodes: nodes,
            strategy: strategy,
            i: AtomicUsize::new(0),
        }
    }

    /// Returns next node basing on provided strategy.
    pub fn next(&self) -> Option<&T> {
        let next = self.strategy
            .next(&self.nodes, self.i.load(Ordering::Relaxed) as usize);
        if self.strategy == LoadBalancingStrategy::RoundRobin {
            self.i.fetch_add(1, Ordering::Relaxed);
            // prevent overflow
            let i = self.i.load(Ordering::Relaxed);
            match i.checked_rem(self.nodes.len() as usize) {
                Some(rem) => self.i.store(rem, Ordering::Relaxed),
                None => return None,
            }
        }
        next
    }
}

/// [r2d2](https://github.com/sfackler/r2d2) `ManageConnection`.
pub struct ClusterConnectionManager<T, X> {
    load_balancer: LoadBalancer<X>,
    authenticator: T,
    compression: Compression,
}

impl<T, X> ClusterConnectionManager<T, X>
    where T: Authenticator + Send + Sync + 'static
{
    /// Creates a new instance of `ConnectionManager`.
    /// It requires transport, authenticator and compression as inputs.
    pub fn new(load_balancer: LoadBalancer<X>,
               authenticator: T,
               compression: Compression)
               -> ClusterConnectionManager<T, X> {
        ClusterConnectionManager {
            load_balancer: load_balancer,
            authenticator: authenticator,
            compression: compression,
        }
    }
}

impl<T: Authenticator + Send + Sync + 'static, X: CDRSTransport + Send + Sync + 'static>
r2d2::ManageConnection for ClusterConnectionManager<T, X> {
    type Connection = Session<T,X>;
    type Error = CError;

    fn connect(&self) -> Result<Self::Connection, Self::Error> {
        let transport_res: CResult<X> = self.load_balancer.next()
            .ok_or_else(|| "Cannot get next node".into())
            .and_then(|x| x.try_clone().map_err(|e| e.into()));
        let transport = try!(transport_res);
        let compression = self.compression;
        let cdrs = CDRS::new(transport, self.authenticator.clone());

        cdrs.start(compression)
    }

    fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), Self::Error> {
        let query = QueryBuilder::new("SELECT * FROM system.peers;").finalize();

        connection.query(query, false, false).map(|_| ())
    }

    fn has_broken(&self, _connection: &mut Self::Connection) -> bool {
        false
    }
}

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

    #[test]
    fn round_robin() {
        let nodes = vec!["a", "b", "c"];
        let nodes_c = nodes.clone();
        let load_balancer = LoadBalancer::new(nodes, LoadBalancingStrategy::RoundRobin);
        for i in 0..10 {
            assert_eq!(&nodes_c[i % 3], load_balancer.next().unwrap());
        }
    }
}
<|endoftext|>


<commit_before>use crate::syntax::{
    attrs, error, Api, Atom, Doc, ExternFn, ExternType, Lang, Receiver, Ref, Signature, Struct,
    Ty1, Type, Var,
};
use proc_macro2::Ident;
use quote::quote;
use syn::{
    Abi, Error, Fields, FnArg, ForeignItem, ForeignItemFn, ForeignItemType, GenericArgument, Item,
    ItemForeignMod, ItemStruct, Pat, PathArguments, Result, ReturnType, Type as RustType,
};

pub fn parse_items(items: Vec<Item>) -> Result<Vec<Api>> {
    let mut apis = Vec::new();
    for item in items {
        match item {
            Item::Struct(item) => {
                let strct = parse_struct(item)?;
                apis.push(strct);
            }
            Item::ForeignMod(foreign_mod) => {
                let functions = parse_foreign_mod(foreign_mod)?;
                apis.extend(functions);
            }
            Item::Use(item) => return Err(Error::new_spanned(item, error::USE_NOT_ALLOWED)),
            _ => return Err(Error::new_spanned(item, "unsupported item")),
        }
    }
    Ok(apis)
}

fn parse_struct(item: ItemStruct) -> Result<Api> {
    let generics = &item.generics;
    if !generics.params.is_empty() || generics.where_clause.is_some() {
        let struct_token = item.struct_token;
        let ident = &item.ident;
        let where_clause = &generics.where_clause;
        let span = quote!(#struct_token #ident #generics #where_clause);
        return Err(Error::new_spanned(
            span,
            "struct with generic parameters is not supported yet",
        ));
    }

    let mut doc = Doc::new();
    let mut derives = Vec::new();
    attrs::parse(&item.attrs, &mut doc, Some(&mut derives))?;
    check_reserved_name(&item.ident)?;
    match item.fields {
        Fields::Named(fields) => Ok(Api::Struct(Struct {
            doc,
            derives,
            struct_token: item.struct_token,
            ident: item.ident,
            fields: fields
                .named
                .into_iter()
                .map(|field| {
                    Ok(Var {
                        ident: field.ident.unwrap(),
                        ty: parse_type(&field.ty)?,
                    })
                })
                .collect::<Result<_>>()?,
        })),
        Fields::Unit => Err(Error::new_spanned(item, "unit structs are not supported")),
        Fields::Unnamed(_) => Err(Error::new_spanned(item, "tuple structs are not supported")),
    }
}

fn parse_foreign_mod(foreign_mod: ItemForeignMod) -> Result<Vec<Api>> {
    let lang = parse_lang(foreign_mod.abi)?;
    let api_type = match lang {
        Lang::Cxx => Api::CxxType,
        Lang::Rust => Api::RustType,
    };
    let api_function = match lang {
        Lang::Cxx => Api::CxxFunction,
        Lang::Rust => Api::RustFunction,
    };

    let mut items = Vec::new();
    for foreign in &foreign_mod.items {
        match foreign {
            ForeignItem::Type(foreign) => {
                check_reserved_name(&foreign.ident)?;
                let ety = parse_extern_type(foreign)?;
                items.push(api_type(ety));
            }
            ForeignItem::Fn(foreign) => {
                let efn = parse_extern_fn(foreign, lang)?;
                items.push(api_function(efn));
            }
            ForeignItem::Macro(foreign) if foreign.mac.path.is_ident("include") => {
                let include = foreign.mac.parse_body()?;
                items.push(Api::Include(include));
            }
            _ => return Err(Error::new_spanned(foreign, "unsupported foreign item")),
        }
    }
    Ok(items)
}

fn parse_lang(abi: Abi) -> Result<Lang> {
    let name = match &abi.name {
        Some(name) => name,
        None => {
            return Err(Error::new_spanned(
                abi,
                "ABI name is required, extern \"C\" or extern \"Rust\"",
            ));
        }
    };
    match name.value().as_str() {
        "C" => Ok(Lang::Cxx),
        "Rust" => Ok(Lang::Rust),
        _ => Err(Error::new_spanned(abi, "unrecognized ABI")),
    }
}

fn parse_extern_type(foreign_type: &ForeignItemType) -> Result<ExternType> {
    let doc = attrs::parse_doc(&foreign_type.attrs)?;
    let type_token = foreign_type.type_token;
    let ident = foreign_type.ident.clone();
    Ok(ExternType {
        doc,
        type_token,
        ident,
    })
}

fn parse_extern_fn(foreign_fn: &ForeignItemFn, lang: Lang) -> Result<ExternFn> {
    let generics = &foreign_fn.sig.generics;
    if !generics.params.is_empty() || generics.where_clause.is_some() {
        return Err(Error::new_spanned(
            foreign_fn,
            "extern function with generic parameters is not supported yet",
        ));
    }
    if let Some(variadic) = &foreign_fn.sig.variadic {
        return Err(Error::new_spanned(
            variadic,
            "variadic function is not supported yet",
        ));
    }

    let mut receiver = None;
    let mut args = Vec::new();
    for arg in &foreign_fn.sig.inputs {
        match arg {
            FnArg::Receiver(receiver) => {
                return Err(Error::new_spanned(receiver, "unsupported signature"))
            }
            FnArg::Typed(arg) => {
                let ident = match arg.pat.as_ref() {
                    Pat::Ident(pat) => pat.ident.clone(),
                    _ => return Err(Error::new_spanned(arg, "unsupported signature")),
                };
                let ty = parse_type(&arg.ty)?;
                if ident != "self" {
                    args.push(Var { ident, ty });
                    continue;
                }
                if let Type::Ref(reference) = ty {
                    if let Type::Ident(ident) = reference.inner {
                        receiver = Some(Receiver {
                            mutability: reference.mutability,
                            ident,
                        });
                        continue;
                    }
                }
                return Err(Error::new_spanned(arg, "unsupported method receiver"));
            }
        }
    }

    let mut throws = false;
    let ret = match &foreign_fn.sig.output {
        ReturnType::Default => None,
        ReturnType::Type(_, ret) => {
            let mut ret = ret.as_ref();
            if let RustType::Path(ty) = ret {
                let path = &ty.path;
                if ty.qself.is_none() && path.leading_colon.is_none() && path.segments.len() == 1 {
                    let segment = &path.segments[0];
                    let ident = segment.ident.clone();
                    if let PathArguments::AngleBracketed(generic) = &segment.arguments {
                        if ident == "Result" && generic.args.len() == 1 {
                            if let GenericArgument::Type(arg) = &generic.args[0] {
                                ret = arg;
                                throws = true;
                            }
                        }
                    }
                }
            }
            match parse_type(ret)? {
                Type::Void(_) => None,
                ty => Some(ty),
            }
        }
    };

    let doc = attrs::parse_doc(&foreign_fn.attrs)?;
    let fn_token = foreign_fn.sig.fn_token;
    let ident = foreign_fn.sig.ident.clone();
    let semi_token = foreign_fn.semi_token;
    Ok(ExternFn {
        lang,
        doc,
        ident,
        sig: Signature {
            fn_token,
            receiver,
            args,
            ret,
            throws,
        },
        semi_token,
    })
}

fn parse_type(ty: &RustType) -> Result<Type> {
    match ty {
        RustType::Reference(ty) => {
            let inner = parse_type(&ty.elem)?;
            let which = match &inner {
                Type::Ident(ident) if ident == "str" => {
                    if ty.mutability.is_some() {
                        return Err(Error::new_spanned(ty, "unsupported type"));
                    } else {
                        Type::Str
                    }
                }
                _ => Type::Ref,
            };
            return Ok(which(Box::new(Ref {
                ampersand: ty.and_token,
                mutability: ty.mutability,
                inner,
            })));
        }
        RustType::Path(ty) => {
            let path = &ty.path;
            if ty.qself.is_none() && path.leading_colon.is_none() && path.segments.len() == 1 {
                let segment = &path.segments[0];
                let ident = segment.ident.clone();
                match &segment.arguments {
                    PathArguments::None => return Ok(Type::Ident(ident)),
                    PathArguments::AngleBracketed(generic) => {
                        if ident == "UniquePtr" && generic.args.len() == 1 {
                            if let GenericArgument::Type(arg) = &generic.args[0] {
                                let inner = parse_type(arg)?;
                                return Ok(Type::UniquePtr(Box::new(Ty1 {
                                    name: ident,
                                    langle: generic.lt_token,
                                    inner,
                                    rangle: generic.gt_token,
                                })));
                            }
                        } else if ident == "Box" && generic.args.len() == 1 {
                            if let GenericArgument::Type(arg) = &generic.args[0] {
                                let inner = parse_type(arg)?;
                                return Ok(Type::RustBox(Box::new(Ty1 {
                                    name: ident,
                                    langle: generic.lt_token,
                                    inner,
                                    rangle: generic.gt_token,
                                })));
                            }
                        }
                    }
                    PathArguments::Parenthesized(_) => {}
                }
            }
        }
        RustType::Tuple(ty) if ty.elems.is_empty() => {
            return Ok(Type::Void(ty.paren_token.span));
        }
        _ => {}
    }
    Err(Error::new_spanned(ty, "unsupported type"))
}

fn check_reserved_name(ident: &Ident) -> Result<()> {
    if ident == "Box" || ident == "UniquePtr" || Atom::from(ident).is_some() {
        Err(Error::new(ident.span(), "reserved name"))
    } else {
        Ok(())
    }
}
<commit_msg>Allow extern abi to be specified as "C++"<commit_after>use crate::syntax::{
    attrs, error, Api, Atom, Doc, ExternFn, ExternType, Lang, Receiver, Ref, Signature, Struct,
    Ty1, Type, Var,
};
use proc_macro2::Ident;
use quote::quote;
use syn::{
    Abi, Error, Fields, FnArg, ForeignItem, ForeignItemFn, ForeignItemType, GenericArgument, Item,
    ItemForeignMod, ItemStruct, Pat, PathArguments, Result, ReturnType, Type as RustType,
};

pub fn parse_items(items: Vec<Item>) -> Result<Vec<Api>> {
    let mut apis = Vec::new();
    for item in items {
        match item {
            Item::Struct(item) => {
                let strct = parse_struct(item)?;
                apis.push(strct);
            }
            Item::ForeignMod(foreign_mod) => {
                let functions = parse_foreign_mod(foreign_mod)?;
                apis.extend(functions);
            }
            Item::Use(item) => return Err(Error::new_spanned(item, error::USE_NOT_ALLOWED)),
            _ => return Err(Error::new_spanned(item, "unsupported item")),
        }
    }
    Ok(apis)
}

fn parse_struct(item: ItemStruct) -> Result<Api> {
    let generics = &item.generics;
    if !generics.params.is_empty() || generics.where_clause.is_some() {
        let struct_token = item.struct_token;
        let ident = &item.ident;
        let where_clause = &generics.where_clause;
        let span = quote!(#struct_token #ident #generics #where_clause);
        return Err(Error::new_spanned(
            span,
            "struct with generic parameters is not supported yet",
        ));
    }

    let mut doc = Doc::new();
    let mut derives = Vec::new();
    attrs::parse(&item.attrs, &mut doc, Some(&mut derives))?;
    check_reserved_name(&item.ident)?;
    match item.fields {
        Fields::Named(fields) => Ok(Api::Struct(Struct {
            doc,
            derives,
            struct_token: item.struct_token,
            ident: item.ident,
            fields: fields
                .named
                .into_iter()
                .map(|field| {
                    Ok(Var {
                        ident: field.ident.unwrap(),
                        ty: parse_type(&field.ty)?,
                    })
                })
                .collect::<Result<_>>()?,
        })),
        Fields::Unit => Err(Error::new_spanned(item, "unit structs are not supported")),
        Fields::Unnamed(_) => Err(Error::new_spanned(item, "tuple structs are not supported")),
    }
}

fn parse_foreign_mod(foreign_mod: ItemForeignMod) -> Result<Vec<Api>> {
    let lang = parse_lang(foreign_mod.abi)?;
    let api_type = match lang {
        Lang::Cxx => Api::CxxType,
        Lang::Rust => Api::RustType,
    };
    let api_function = match lang {
        Lang::Cxx => Api::CxxFunction,
        Lang::Rust => Api::RustFunction,
    };

    let mut items = Vec::new();
    for foreign in &foreign_mod.items {
        match foreign {
            ForeignItem::Type(foreign) => {
                check_reserved_name(&foreign.ident)?;
                let ety = parse_extern_type(foreign)?;
                items.push(api_type(ety));
            }
            ForeignItem::Fn(foreign) => {
                let efn = parse_extern_fn(foreign, lang)?;
                items.push(api_function(efn));
            }
            ForeignItem::Macro(foreign) if foreign.mac.path.is_ident("include") => {
                let include = foreign.mac.parse_body()?;
                items.push(Api::Include(include));
            }
            _ => return Err(Error::new_spanned(foreign, "unsupported foreign item")),
        }
    }
    Ok(items)
}

fn parse_lang(abi: Abi) -> Result<Lang> {
    let name = match &abi.name {
        Some(name) => name,
        None => {
            return Err(Error::new_spanned(
                abi,
                "ABI name is required, extern \"C\" or extern \"Rust\"",
            ));
        }
    };
    match name.value().as_str() {
        "C" | "C++" => Ok(Lang::Cxx),
        "Rust" => Ok(Lang::Rust),
        _ => Err(Error::new_spanned(abi, "unrecognized ABI")),
    }
}

fn parse_extern_type(foreign_type: &ForeignItemType) -> Result<ExternType> {
    let doc = attrs::parse_doc(&foreign_type.attrs)?;
    let type_token = foreign_type.type_token;
    let ident = foreign_type.ident.clone();
    Ok(ExternType {
        doc,
        type_token,
        ident,
    })
}

fn parse_extern_fn(foreign_fn: &ForeignItemFn, lang: Lang) -> Result<ExternFn> {
    let generics = &foreign_fn.sig.generics;
    if !generics.params.is_empty() || generics.where_clause.is_some() {
        return Err(Error::new_spanned(
            foreign_fn,
            "extern function with generic parameters is not supported yet",
        ));
    }
    if let Some(variadic) = &foreign_fn.sig.variadic {
        return Err(Error::new_spanned(
            variadic,
            "variadic function is not supported yet",
        ));
    }

    let mut receiver = None;
    let mut args = Vec::new();
    for arg in &foreign_fn.sig.inputs {
        match arg {
            FnArg::Receiver(receiver) => {
                return Err(Error::new_spanned(receiver, "unsupported signature"))
            }
            FnArg::Typed(arg) => {
                let ident = match arg.pat.as_ref() {
                    Pat::Ident(pat) => pat.ident.clone(),
                    _ => return Err(Error::new_spanned(arg, "unsupported signature")),
                };
                let ty = parse_type(&arg.ty)?;
                if ident != "self" {
                    args.push(Var { ident, ty });
                    continue;
                }
                if let Type::Ref(reference) = ty {
                    if let Type::Ident(ident) = reference.inner {
                        receiver = Some(Receiver {
                            mutability: reference.mutability,
                            ident,
                        });
                        continue;
                    }
                }
                return Err(Error::new_spanned(arg, "unsupported method receiver"));
            }
        }
    }

    let mut throws = false;
    let ret = match &foreign_fn.sig.output {
        ReturnType::Default => None,
        ReturnType::Type(_, ret) => {
            let mut ret = ret.as_ref();
            if let RustType::Path(ty) = ret {
                let path = &ty.path;
                if ty.qself.is_none() && path.leading_colon.is_none() && path.segments.len() == 1 {
                    let segment = &path.segments[0];
                    let ident = segment.ident.clone();
                    if let PathArguments::AngleBracketed(generic) = &segment.arguments {
                        if ident == "Result" && generic.args.len() == 1 {
                            if let GenericArgument::Type(arg) = &generic.args[0] {
                                ret = arg;
                                throws = true;
                            }
                        }
                    }
                }
            }
            match parse_type(ret)? {
                Type::Void(_) => None,
                ty => Some(ty),
            }
        }
    };

    let doc = attrs::parse_doc(&foreign_fn.attrs)?;
    let fn_token = foreign_fn.sig.fn_token;
    let ident = foreign_fn.sig.ident.clone();
    let semi_token = foreign_fn.semi_token;
    Ok(ExternFn {
        lang,
        doc,
        ident,
        sig: Signature {
            fn_token,
            receiver,
            args,
            ret,
            throws,
        },
        semi_token,
    })
}

fn parse_type(ty: &RustType) -> Result<Type> {
    match ty {
        RustType::Reference(ty) => {
            let inner = parse_type(&ty.elem)?;
            let which = match &inner {
                Type::Ident(ident) if ident == "str" => {
                    if ty.mutability.is_some() {
                        return Err(Error::new_spanned(ty, "unsupported type"));
                    } else {
                        Type::Str
                    }
                }
                _ => Type::Ref,
            };
            return Ok(which(Box::new(Ref {
                ampersand: ty.and_token,
                mutability: ty.mutability,
                inner,
            })));
        }
        RustType::Path(ty) => {
            let path = &ty.path;
            if ty.qself.is_none() && path.leading_colon.is_none() && path.segments.len() == 1 {
                let segment = &path.segments[0];
                let ident = segment.ident.clone();
                match &segment.arguments {
                    PathArguments::None => return Ok(Type::Ident(ident)),
                    PathArguments::AngleBracketed(generic) => {
                        if ident == "UniquePtr" && generic.args.len() == 1 {
                            if let GenericArgument::Type(arg) = &generic.args[0] {
                                let inner = parse_type(arg)?;
                                return Ok(Type::UniquePtr(Box::new(Ty1 {
                                    name: ident,
                                    langle: generic.lt_token,
                                    inner,
                                    rangle: generic.gt_token,
                                })));
                            }
                        } else if ident == "Box" && generic.args.len() == 1 {
                            if let GenericArgument::Type(arg) = &generic.args[0] {
                                let inner = parse_type(arg)?;
                                return Ok(Type::RustBox(Box::new(Ty1 {
                                    name: ident,
                                    langle: generic.lt_token,
                                    inner,
                                    rangle: generic.gt_token,
                                })));
                            }
                        }
                    }
                    PathArguments::Parenthesized(_) => {}
                }
            }
        }
        RustType::Tuple(ty) if ty.elems.is_empty() => {
            return Ok(Type::Void(ty.paren_token.span));
        }
        _ => {}
    }
    Err(Error::new_spanned(ty, "unsupported type"))
}

fn check_reserved_name(ident: &Ident) -> Result<()> {
    if ident == "Box" || ident == "UniquePtr" || Atom::from(ident).is_some() {
        Err(Error::new(ident.span(), "reserved name"))
    } else {
        Ok(())
    }
}
<|endoftext|>