<commit_before>use command::Command;
use vm::VM;

pub struct Program {
    instructions: Vec<Command>,
    instruction_pointer: Option<usize>,
    current_depth: u64,
    pub status: ProgramStatus,
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum ProgramStatus {
    Normal,
    Seeking(u64),
}

impl Program {
    pub fn new() -> Program {
        Program {
            instructions: Vec::new(),
            instruction_pointer: None,
            current_depth: 0,
            status: ProgramStatus::Normal,
        }
    }

    pub fn append(&mut self, instructions: &[Command]) {
        self.instructions.extend(instructions.iter().cloned());
        if self.instruction_pointer.is_none() {
            self.instruction_pointer = Some(0);
        }
    }

    pub fn execute(&mut self, vm: &mut VM) {
        match (self.instruction_pointer, self.status) {
            (Some(mut index), ProgramStatus::Seeking(_)) => {
                self.handle_jump_forward(vm, &mut index);
                self.instruction_pointer = Some(index);
                if index < self.instructions.len() { self.execute(vm); }
            }

            (Some(mut index), ProgramStatus::Normal) => {
                while index < self.instructions.len() {
                    match self.instructions[index] {
                        Command::JumpForward => self.handle_jump_forward(vm, &mut index),
                        Command::JumpBackward => self.handle_jump_backward(vm, &mut index),
                        command => {
                            vm.apply(command);
                            index += 1;
                        }
                    }
                }

                self.instruction_pointer = Some(index);
            }

            _ => {}
        }
    }

    fn handle_jump_backward(&mut self, vm: &VM, index_ref: &mut usize) {
        match vm.cells[vm.data_pointer] {
            0 => {
                self.current_depth -= 1;
                *index_ref += 1;
            }

            _ => {
                let goal_depth = self.current_depth;
                for index in (0..*index_ref).rev() {
                    match self.instructions[index] {
                        Command::JumpBackward => self.current_depth += 1,
                        Command::JumpForward => {
                            if self.current_depth == goal_depth {
                                *index_ref = index + 1;
                                return;
                            }
                            self.current_depth -= 1;
                        }
                        _ => {}
                    }
                }
                panic!("No starting brace found!");
            }
        }
    }

    fn seek_forward(&self, starting_index: usize) -> (usize, ProgramStatus) {
        match self.status {
            ProgramStatus::Seeking(goal_depth) => {
                for index in starting_index..self.instructions.len() {
                    match self.instructions[index] {
                        Command::JumpForward => self.current_depth += 1,
                        Command::JumpBackward => self.current_depth -= 1,
                        _ => {}
                    }
                    if self.current_depth == goal_depth {
                        return (index + 1, ProgramStatus::Normal);
                    }
                }

                return (self.instructions.len(), ProgramStatus::Seeking(goal_depth));
            }

            ProgramStatus::Normal => return (starting_index, self.status)
        }
    }

    fn handle_jump_forward(&mut self, vm: &VM, index: &mut usize) {
        match (vm.cells[vm.data_pointer], self.status) {
            (0, ProgramStatus::Normal) => {
                self.status = ProgramStatus::Seeking(self.current_depth);
                self.seek_forward(index);
            }

            (_, ProgramStatus::Normal) => {
                *index += 1;
                self.current_depth += 1;
                self.status = ProgramStatus::Normal;
            }

            (_, ProgramStatus::Seeking(_)) => self.seek_forward(index),

        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    pub fn test_seek_forward_normal() {
        let mut program = Program {
            instruction_pointer: None,
            current_depth: 0,
            status: ProgramStatus::Seeking(0),
            instructions: vec![
                Command::JumpForward,
                Command::JumpForward,
                Command::JumpBackward,
                Command::JumpBackward,
            ],
        };
        let mut index = 0;
        program.seek_forward(&mut index);
        assert_eq!(index, 4);
        assert_eq!(program.status, ProgramStatus::Normal);
    }

    #[test]
    pub fn test_seek_forward_seeking() {
        let mut program = Program {
            instruction_pointer: None,
            current_depth: 0,
            status: ProgramStatus::Seeking(0),
            instructions: vec![
                Command::JumpForward,
                Command::JumpForward,
                Command::JumpBackward,
            ],
        };
        let mut index = 0;
        program.seek_forward(&mut index);
        assert_eq!(index, 3);
        assert_eq!(program.status, ProgramStatus::Seeking(0));
    }
}
<commit_msg>Return current depth from seek_forward as well<commit_after>use command::Command;
use vm::VM;

pub struct Program {
    instructions: Vec<Command>,
    instruction_pointer: Option<usize>,
    current_depth: u64,
    pub status: ProgramStatus,
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum ProgramStatus {
    Normal,
    Seeking(u64),
}

impl Program {
    pub fn new() -> Program {
        Program {
            instructions: Vec::new(),
            instruction_pointer: None,
            current_depth: 0,
            status: ProgramStatus::Normal,
        }
    }

    pub fn append(&mut self, instructions: &[Command]) {
        self.instructions.extend(instructions.iter().cloned());
        if self.instruction_pointer.is_none() {
            self.instruction_pointer = Some(0);
        }
    }

    pub fn execute(&mut self, vm: &mut VM) {
        match (self.instruction_pointer, self.status) {
            (Some(mut index), ProgramStatus::Seeking(_)) => {
                self.handle_jump_forward(vm, &mut index);
                self.instruction_pointer = Some(index);
                if index < self.instructions.len() { self.execute(vm); }
            }

            (Some(mut index), ProgramStatus::Normal) => {
                while index < self.instructions.len() {
                    match self.instructions[index] {
                        Command::JumpForward => self.handle_jump_forward(vm, &mut index),
                        Command::JumpBackward => self.handle_jump_backward(vm, &mut index),
                        command => {
                            vm.apply(command);
                            index += 1;
                        }
                    }
                }

                self.instruction_pointer = Some(index);
            }

            _ => {}
        }
    }

    fn handle_jump_backward(&mut self, vm: &VM, index_ref: &mut usize) {
        match vm.cells[vm.data_pointer] {
            0 => {
                self.current_depth -= 1;
                *index_ref += 1;
            }

            _ => {
                let goal_depth = self.current_depth;
                for index in (0..*index_ref).rev() {
                    match self.instructions[index] {
                        Command::JumpBackward => self.current_depth += 1,
                        Command::JumpForward => {
                            if self.current_depth == goal_depth {
                                *index_ref = index + 1;
                                return;
                            }
                            self.current_depth -= 1;
                        }
                        _ => {}
                    }
                }
                panic!("No starting brace found!");
            }
        }
    }

    fn seek_forward(&self, starting_index: usize) -> (usize, u64, ProgramStatus) {
        match self.status {
            ProgramStatus::Seeking(goal_depth) => {
                for index in starting_index..self.instructions.len() {
                    match self.instructions[index] {
                        Command::JumpForward => self.current_depth += 1,
                        Command::JumpBackward => self.current_depth -= 1,
                        _ => {}
                    }
                    if self.current_depth == goal_depth {
                        return (index + 1, self.current_depth, ProgramStatus::Normal);
                    }
                }

                return (self.instructions.len()
                       , self.current_depth
                       , ProgramStatus::Seeking(goal_depth));
            }

            ProgramStatus::Normal => return (starting_index, self.current_depth, self.status)
        }
    }

    fn handle_jump_forward(&mut self, vm: &VM, index: &mut usize) {
        match (vm.cells[vm.data_pointer], self.status) {
            (0, ProgramStatus::Normal) => {
                self.status = ProgramStatus::Seeking(self.current_depth);
                self.seek_forward(index);
            }

            (_, ProgramStatus::Normal) => {
                *index += 1;
                self.current_depth += 1;
                self.status = ProgramStatus::Normal;
            }

            (_, ProgramStatus::Seeking(_)) => self.seek_forward(index),

        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    pub fn test_seek_forward_normal() {
        let mut program = Program {
            instruction_pointer: None,
            current_depth: 0,
            status: ProgramStatus::Seeking(0),
            instructions: vec![
                Command::JumpForward,
                Command::JumpForward,
                Command::JumpBackward,
                Command::JumpBackward,
            ],
        };
        let mut index = 0;
        program.seek_forward(&mut index);
        assert_eq!(index, 4);
        assert_eq!(program.status, ProgramStatus::Normal);
    }

    #[test]
    pub fn test_seek_forward_seeking() {
        let mut program = Program {
            instruction_pointer: None,
            current_depth: 0,
            status: ProgramStatus::Seeking(0),
            instructions: vec![
                Command::JumpForward,
                Command::JumpForward,
                Command::JumpBackward,
            ],
        };
        let mut index = 0;
        program.seek_forward(&mut index);
        assert_eq!(index, 3);
        assert_eq!(program.status, ProgramStatus::Seeking(0));
    }
}
<|endoftext|>


<commit_before>//! Endpoints for sending events.
<commit_msg>Add sending endpoints.<commit_after>//! Endpoints for sending events.

/// PUT /_matrix/client/r0/rooms/{roomId}/state/{eventType}
///
/// [Matrix spec link](http://matrix.org/docs/spec/client_server/r0.2.0.html#put-matrix-client-r0-rooms-roomid-state-eventtype)
pub mod send_state {
    use ruma_identifiers::{RoomId, EventId};
    use ruma_events::EventType;

    /// Details about this API endpoint.
    pub struct Endpoint;

    /// This API endpoint's path parameters.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct PathParams {
        pub room_id: RoomId,
        pub event_type: EventType
    }

    /// This API endpoint's response.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct Response {
        pub event_id: EventId,
    }

    impl ::Endpoint for Endpoint {
        type BodyParams = ::serde_json::Value;
        type PathParams = PathParams;
        type QueryParams = ();
        type Response = ();

        fn method() -> ::Method {
            ::Method::Put
        }

        fn request_path(params: Self::PathParams) -> String {
            format!(
                "/_matrix/client/r0/rooms/{}/state/{}",
                params.room_id,
                params.event_type
            )
        }

        fn router_path() -> String {
            "/_matrix/client/r0/rooms/:room_id/state/:event_type".to_string()
        }
    }
}

/// PUT /_matrix/client/r0/rooms/{roomId}/state/{eventType}/{stateKey}
///
/// [Matrix spec link](http://matrix.org/docs/spec/client_server/r0.2.0.html#put-matrix-client-r0-rooms-roomid-state-eventtype-statekey)
pub mod send_state_key {
    use ruma_identifiers::{RoomId, EventId};
    use ruma_events::EventType;

    /// Details about this API endpoint.
    pub struct Endpoint;

    /// This API endpoint's path parameters.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct PathParams {
        pub room_id: RoomId,
        pub event_type: EventType,
        pub state_key: String
    }

    /// This API endpoint's response.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct Response {
        pub event_id: EventId,
    }

    impl ::Endpoint for Endpoint {
        type BodyParams = ::serde_json::Value;
        type PathParams = PathParams;
        type QueryParams = ();
        type Response = ();

        fn method() -> ::Method {
            ::Method::Put
        }

        fn request_path(params: Self::PathParams) -> String {
            format!(
                "/_matrix/client/r0/rooms/{}/state/{}/{}",
                params.room_id,
                params.event_type,
                params.state_key
            )
        }

        fn router_path() -> String {
            "/_matrix/client/r0/rooms/:room_id/state/:event_type/:state_key".to_string()
        }
    }
}

/// PUT /_matrix/client/r0/rooms/{roomId}/send/{eventType}/{txnId}
///
/// [Matrix spec link](http://matrix.org/docs/spec/client_server/r0.2.0.html#put-matrix-client-r0-rooms-roomid-send-eventtype-txnid)
pub mod send_event {
    use ruma_identifiers::{RoomId, EventId};
    use ruma_events::EventType;

    /// Details about this API endpoint.
    pub struct Endpoint;

    /// This API endpoint's path parameters.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct PathParams {
        pub room_id: RoomId,
        pub event_type: EventType,
        pub txn_id: String
    }

    /// This API endpoint's response.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    pub struct Response {
        pub event_id: EventId,
    }

    impl ::Endpoint for Endpoint {
        type BodyParams = ::serde_json::Value;
        type PathParams = PathParams;
        type QueryParams = ();
        type Response = ();

        fn method() -> ::Method {
            ::Method::Put
        }

        fn request_path(params: Self::PathParams) -> String {
            format!(
                "/_matrix/client/r0/rooms/{}/send/{}/{}",
                params.room_id,
                params.event_type,
                params.txn_id
            )
        }

        fn router_path() -> String {
            "/_matrix/client/r0/rooms/:room_id/send/:event_type/:txn_id".to_string()
        }
    }
}
<|endoftext|>


<commit_before>intrinsics! {
    // Implementation from gcc
    // https://raw.githubusercontent.com/gcc-mirror/gcc/master/libgcc/config/epiphany/mulsi3.c
    pub extern "C" fn __mulsi3(mut a: u32, mut b: u32) -> u32 {
        let mut r: usize = 0;

        while a > 0 {
            if a & 1 > 0 {
                r += b;
            }
            a >>= 1;
            b <<= 1;
        }

        r
    }
}
<commit_msg>Fix compilation on riscv32<commit_after>intrinsics! {
    // Implementation from gcc
    // https://raw.githubusercontent.com/gcc-mirror/gcc/master/libgcc/config/epiphany/mulsi3.c
    pub extern "C" fn __mulsi3(a: u32, b: u32) -> u32 {
        let (mut a, mut b) = (a, b);
        let mut r: usize = 0;

        while a > 0 {
            if a & 1 > 0 {
                r += b;
            }
            a >>= 1;
            b <<= 1;
        }

        r
    }
}
<|endoftext|>


<commit_before>pub mod proto;
mod reader;
mod writer;
mod errors;
mod client;
mod processor;
mod transfer;

use std::{io, str, result, thread};
use std::io::Write;
use std::net::{TcpListener, TcpStream, Ipv4Addr, SocketAddrV4};

use amy;
use serde_json;

pub use self::proto::resource;
pub use self::errors::{Result, ResultExt, ErrorKind, Error};
use self::proto::message::{self, SMessage};
use self::proto::ws;
use self::client::{Incoming, IncomingStatus, Client};
use self::processor::{Processor, TransferKind};
use self::transfer::{Transfers, TransferResult};
use bencode;
use handle;
use torrent;
use disk;
use util::UHashMap;
use CONFIG;

const POLL_INT_MS: usize = 1000;
const CLEANUP_INT_MS: usize = 2000;

lazy_static! {
    pub static ref EMPTY_HTTP_RESP: Vec<u8> = {
        let lines = vec![
            format!("HTTP/1.1 {} {}", 204, "NO CONTENT"),
            format!("Connection: {}", "Close"),
            format!("Access-Control-Allow-Origin: {}", "*"),
            format!("Access-Control-Allow-Methods: {}", "OPTIONS, POST, GET"),
            format!(
                "Access-Control-Allow-Headers: {}, {}, {}, {}, {}, {}, {}, {}",
                "Access-Control-Allow-Headers",
                "Origin",
                "Accept",
                "X-Requested-With",
                "Content-Type",
                "Access-Control-Request-Method",
                "Access-Control-Request-Headers",
                "Authorization"
            ),
            format!("\r\n"),
        ];
        lines.join("\r\n").into_bytes()
    };
}

#[derive(Debug)]
pub enum CtlMessage {
    Extant(Vec<resource::Resource>),
    Update(Vec<resource::SResourceUpdate<'static>>),
    Removed(Vec<String>),
    ClientRemoved {
        id: String,
        client: usize,
        serial: u64,
    },
    Uploaded {
        id: String,
        client: usize,
        serial: u64,
    },
    Shutdown,
}

#[derive(Debug)]
pub enum Message {
    UpdateTorrent(resource::CResourceUpdate),
    UpdateServer {
        id: String,
        throttle_up: Option<Option<i64>>,
        throttle_down: Option<Option<i64>>,
    },
    UpdateFile {
        id: String,
        torrent_id: String,
        priority: u8,
    },
    RemoveTorrent {
        id: String,
        client: usize,
        serial: u64,
        artifacts: bool,
    },
    Pause(String),
    Resume(String),
    Validate(Vec<String>),
    RemovePeer {
        id: String,
        torrent_id: String,
        client: usize,
        serial: u64,
    },
    UpdateTracker { id: String, torrent_id: String },
    RemoveTracker {
        id: String,
        torrent_id: String,
        client: usize,
        serial: u64,
    },
    Torrent {
        info: torrent::Info,
        client: usize,
        serial: u64,
        path: Option<String>,
        start: bool,
    },
}

#[allow(dead_code)]
pub struct RPC {
    poll: amy::Poller,
    reg: amy::Registrar,
    ch: handle::Handle<CtlMessage, Message>,
    listener: TcpListener,
    lid: usize,
    cleanup: usize,
    processor: Processor,
    transfers: Transfers,
    clients: UHashMap<Client>,
    incoming: UHashMap<Incoming>,
    disk: amy::Sender<disk::Request>,
}

impl RPC {
    pub fn start(
        creg: &mut amy::Registrar,
        db: amy::Sender<disk::Request>,
    ) -> io::Result<(handle::Handle<Message, CtlMessage>, thread::JoinHandle<()>)> {
        let poll = amy::Poller::new()?;
        let mut reg = poll.get_registrar()?;
        let cleanup = reg.set_interval(CLEANUP_INT_MS)?;
        let (ch, dh) = handle::Handle::new(creg, &mut reg)?;

        let ip = if CONFIG.rpc.local {
            Ipv4Addr::new(127, 0, 0, 1)
        } else {
            Ipv4Addr::new(0, 0, 0, 0)
        };
        let port = CONFIG.rpc.port;
        let listener = TcpListener::bind(SocketAddrV4::new(ip, port))?;
        listener.set_nonblocking(true)?;
        let lid = reg.register(&listener, amy::Event::Both)?;

        let disk = db.try_clone()?;
        let th = dh.run("rpc", move |ch| {
            RPC {
                ch,
                disk,
                poll,
                reg,
                listener,
                lid,
                cleanup,
                clients: UHashMap::default(),
                incoming: UHashMap::default(),
                processor: Processor::new(db),
                transfers: Transfers::new(),
            }.run()
        })?;
        Ok((ch, th))
    }

    pub fn run(&mut self) {
        debug!("Running RPC!");
        'outer: while let Ok(res) = self.poll.wait(POLL_INT_MS) {
            for not in res {
                match not.id {
                    id if id == self.lid => self.handle_accept(),
                    id if id == self.ch.rx.get_id() => {
                        if self.handle_ctl() {
                            return;
                        }
                    }
                    id if self.incoming.contains_key(&id) => self.handle_incoming(id),
                    id if id == self.cleanup => self.cleanup(),
                    id if self.transfers.contains(id) => self.handle_transfer(id),
                    _ => self.handle_conn(not),
                }
            }
        }
    }

    fn handle_ctl(&mut self) -> bool {
        while let Ok(m) = self.ch.recv() {
            match m {
                CtlMessage::Shutdown => return true,
                m => {
                    let msgs: Vec<_> = {
                        self.processor
                            .handle_ctl(m)
                            .into_iter()
                            .map(|(c, m)| (c, serde_json::to_string(&m).unwrap()))
                            .collect()
                    };
                    for (c, m) in msgs {
                        let res = match self.clients.get_mut(&c) {
                            Some(client) => client.send(ws::Frame::Text(m)),
                            None => {
                                debug!("Processor referenced a nonexistent client!");
                                Ok(())
                            }
                        };
                        if res.is_err() {
                            let client = self.clients.remove(&c).unwrap();
                            self.remove_client(c, client);
                        }
                    }
                }
            }
        }
        false
    }

    fn handle_transfer(&mut self, id: usize) {
        match self.transfers.ready(id) {
            TransferResult::Incomplete => {}
            TransferResult::Torrent {
                conn,
                data,
                path,
                client,
                serial,
                start,
            } => {
                debug!("Got torrent via HTTP transfer!");
                self.reg.deregister(&conn).unwrap();
                // TODO: Send this to the client in an error msg
                match bencode::decode_buf(&data) {
                    Ok(b) => {
                        match torrent::info::Info::from_bencode(b) {
                            Ok(i) => {
                                if self.ch
                                    .send(Message::Torrent {
                                        info: i,
                                        path,
                                        start,
                                        client,
                                        serial,
                                    })
                                    .is_err()
                                {
                                    error!("Failed to pass message to ctrl!");
                                }
                            }
                            Err(e) => {
                                error!("Failed to parse torrent data: {}!", e);
                                self.clients.get_mut(&client).map(|c| {
                                    c.send(ws::Frame::Text(
                                        serde_json::to_string(
                                            &SMessage::TransferFailed(message::Error {
                                                serial: Some(serial),
                                                reason: format!(
                                                    "Invalid torrent file uploaded, {}.",
                                                    e
                                                ),
                                            }),
                                        ).unwrap(),
                                    ))
                                });
                            }
                        }
                    }
                    Err(e) => {
                        error!("Failed to decode BE data: {}!", e);
                        self.clients.get_mut(&client).map(|c| {
                            c.send(ws::Frame::Text(
                                serde_json::to_string(
                                    &SMessage::TransferFailed(message::Error {
                                        serial: Some(serial),
                                        reason: format!(
                                            "Invalid torrent file uploaded, bad bencoded data: {}.",
                                            e
                                        ),
                                    }),
                                ).unwrap(),
                            ))
                        });
                    }
                }
            }
            TransferResult::Error {
                conn,
                err,
                client: id,
            } => {
                self.reg.deregister(&conn).unwrap();
                let res = self.clients
                    .get_mut(&id)
                    .map(|c| {
                        c.send(ws::Frame::Text(
                            serde_json::to_string(&SMessage::TransferFailed(err))
                                .unwrap(),
                        ))
                    })
                    .unwrap_or(Ok(()));
                if res.is_err() {
                    let client = self.clients.remove(&id).unwrap();
                    self.remove_client(id, client);
                }
            }
        }
    }

    fn handle_accept(&mut self) {
        loop {
            match self.listener.accept() {
                Ok((conn, ip)) => {
                    debug!("Accepted new connection from {:?}!", ip);
                    let id = self.reg.register(&conn, amy::Event::Both).unwrap();
                    self.incoming.insert(id, Incoming::new(conn));
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    break;
                }
                Err(e) => {
                    error!("Failed to accept conn: {}", e);
                }
            }
        }
    }

    fn handle_incoming(&mut self, id: usize) {
        if let Some(mut i) = self.incoming.remove(&id) {
            match i.readable() {
                Ok(IncomingStatus::Upgrade) => {
                    debug!("Succesfully upgraded conn");
                    self.clients.insert(id, i.into());
                }
                Ok(IncomingStatus::Incomplete) => {
                    self.incoming.insert(id, i);
                }
                Ok(IncomingStatus::Transfer { data, token }) => {
                    debug!("File transfer requested, validating");
                    match self.processor.get_transfer(token) {
                        Some((client,
                              serial,
                              TransferKind::UploadTorrent { path, size, start })) => {
                            debug!("Torrent transfer initiated");
                            self.transfers.add_torrent(
                                id,
                                client,
                                serial,
                                i.into(),
                                data,
                                path,
                                size,
                                start,
                            );
                            // Since a succesful result means the buffer hasn't been flushed,
                            // immediatly attempt to handle the transfer as if it was ready
                            self.handle_transfer(id);
                        }
                        Some(_) => {
                            error!("Unimplemented transfer type ignored");
                        }
                        None => {
                            error!("Transfer used invalid token");
                            // TODO: Handle downloads and other uploads
                        }
                    }
                }
                Ok(IncomingStatus::DL(id)) => {
                    debug!("Attempting DL of {}", id);
                    let mut conn: TcpStream = i.into();
                    self.reg.deregister(&conn).is_ok();
                    if let Some(path) = self.processor.get_dl(&id) {
                        debug!("Initiating DL");
                        conn.set_nonblocking(false).is_ok();
                        self.disk.send(disk::Request::download(conn, path)).ok();
                    } else {
                        debug!("ID {} invalid, stopping DL", id);
                        conn.write(&EMPTY_HTTP_RESP).ok();
                    }
                }
                Err(e) => {
                    debug!("Incoming ws upgrade failed: {}", e);
                    self.reg.deregister::<TcpStream>(&i.into()).unwrap();
                }
            }
        }
    }

    fn handle_conn(&mut self, not: amy::Notification) {
        if let Some(mut c) = self.clients.remove(&not.id) {
            if not.event.readable() {
                loop {
                    match c.read() {
                        Ok(None) => break,
                        Ok(Some(ws::Frame::Text(data))) => {
                            if self.process_frame(not.id, &mut c, &data).is_err() {
                                debug!("Client error, disconnecting");
                                self.remove_client(not.id, c);
                                return;
                            }
                        }
                        Ok(Some(_)) | Err(_) => {
                            debug!("Client error, disconnecting");
                            self.remove_client(not.id, c);
                            return;
                        }
                    }
                }
            }
            if not.event.writable() && c.write().is_err() {
                self.remove_client(not.id, c);
                return;
            }
            self.clients.insert(not.id, c);
        }
    }

    fn process_frame(&mut self, id: usize, c: &mut Client, data: &str) -> result::Result<(), ()> {
        match serde_json::from_str(data) {
            Ok(m) => {
                let (msgs, rm) = self.processor.handle_client(id, m);
                if let Some(m) = rm {
                    self.ch.send(m).unwrap();
                }
                for msg in msgs {
                    if c.send(ws::Frame::Text(serde_json::to_string(&msg).unwrap()))
                        .is_err()
                    {
                        return Err(());
                    }
                }
            }
            Err(e) => {
                if e.is_syntax() || e.is_eof() {
                    let msg = SMessage::InvalidSchema(message::Error {
                        serial: None,
                        reason: format!("JSON decode error: {}", e),
                    });
                    if c.send(ws::Frame::Text(serde_json::to_string(&msg).unwrap()))
                        .is_err()
                    {}
                    return Err(());
                }
                if e.is_data() {
                    #[derive(Deserialize)]
                    struct Serial {
                        serial: u64,
                    }

                    let serial = match serde_json::from_str::<Serial>(data) {
                        Ok(s) => Some(s.serial),
                        Err(_) => None,
                    };

                    let msg = SMessage::InvalidSchema(message::Error {
                        serial,
                        reason: format!("Invalid message format: {}", e),
                    });
                    if c.send(ws::Frame::Text(serde_json::to_string(&msg).unwrap()))
                        .is_err()
                    {
                        return Err(());
                    }
                }
            }
        }
        Ok(())
    }

    fn cleanup(&mut self) {
        self.processor.remove_expired_tokens();
        let reg = &self.reg;
        let processor = &mut self.processor;
        self.clients.retain(|id, client| {
            let res = client.timed_out();
            if res {
                info!("client {} timed out", id);
                reg.deregister(&client.conn).unwrap();
                processor.remove_client(*id);
            }
            !res
        });
        self.incoming.retain(|_, inc| {
            let res = inc.timed_out();
            if res {
                reg.deregister(&inc.conn).unwrap();
            }
            !res
        });
        for (conn, id, err) in self.transfers.cleanup() {
            reg.deregister(&conn).unwrap();
            self.clients.get_mut(&id).map(|c| {
                c.send(ws::Frame::Text(
                    serde_json::to_string(&SMessage::TransferFailed(err))
                        .unwrap(),
                ))
            });
        }
    }

    fn remove_client(&mut self, id: usize, client: Client) {
        self.processor.remove_client(id);
        self.reg.deregister::<TcpStream>(&client.into()).unwrap();
    }
}
<commit_msg>Don't claim client error on disconnect<commit_after>pub mod proto;
mod reader;
mod writer;
mod errors;
mod client;
mod processor;
mod transfer;

use std::{io, str, result, thread};
use std::io::Write;
use std::net::{TcpListener, TcpStream, Ipv4Addr, SocketAddrV4};

use amy;
use serde_json;

pub use self::proto::resource;
pub use self::errors::{Result, ResultExt, ErrorKind, Error};
use self::proto::message::{self, SMessage};
use self::proto::ws;
use self::client::{Incoming, IncomingStatus, Client};
use self::processor::{Processor, TransferKind};
use self::transfer::{Transfers, TransferResult};
use bencode;
use handle;
use torrent;
use disk;
use util::UHashMap;
use CONFIG;

const POLL_INT_MS: usize = 1000;
const CLEANUP_INT_MS: usize = 2000;

lazy_static! {
    pub static ref EMPTY_HTTP_RESP: Vec<u8> = {
        let lines = vec![
            format!("HTTP/1.1 {} {}", 204, "NO CONTENT"),
            format!("Connection: {}", "Close"),
            format!("Access-Control-Allow-Origin: {}", "*"),
            format!("Access-Control-Allow-Methods: {}", "OPTIONS, POST, GET"),
            format!(
                "Access-Control-Allow-Headers: {}, {}, {}, {}, {}, {}, {}, {}",
                "Access-Control-Allow-Headers",
                "Origin",
                "Accept",
                "X-Requested-With",
                "Content-Type",
                "Access-Control-Request-Method",
                "Access-Control-Request-Headers",
                "Authorization"
            ),
            format!("\r\n"),
        ];
        lines.join("\r\n").into_bytes()
    };
}

#[derive(Debug)]
pub enum CtlMessage {
    Extant(Vec<resource::Resource>),
    Update(Vec<resource::SResourceUpdate<'static>>),
    Removed(Vec<String>),
    ClientRemoved {
        id: String,
        client: usize,
        serial: u64,
    },
    Uploaded {
        id: String,
        client: usize,
        serial: u64,
    },
    Shutdown,
}

#[derive(Debug)]
pub enum Message {
    UpdateTorrent(resource::CResourceUpdate),
    UpdateServer {
        id: String,
        throttle_up: Option<Option<i64>>,
        throttle_down: Option<Option<i64>>,
    },
    UpdateFile {
        id: String,
        torrent_id: String,
        priority: u8,
    },
    RemoveTorrent {
        id: String,
        client: usize,
        serial: u64,
        artifacts: bool,
    },
    Pause(String),
    Resume(String),
    Validate(Vec<String>),
    RemovePeer {
        id: String,
        torrent_id: String,
        client: usize,
        serial: u64,
    },
    UpdateTracker { id: String, torrent_id: String },
    RemoveTracker {
        id: String,
        torrent_id: String,
        client: usize,
        serial: u64,
    },
    Torrent {
        info: torrent::Info,
        client: usize,
        serial: u64,
        path: Option<String>,
        start: bool,
    },
}

#[allow(dead_code)]
pub struct RPC {
    poll: amy::Poller,
    reg: amy::Registrar,
    ch: handle::Handle<CtlMessage, Message>,
    listener: TcpListener,
    lid: usize,
    cleanup: usize,
    processor: Processor,
    transfers: Transfers,
    clients: UHashMap<Client>,
    incoming: UHashMap<Incoming>,
    disk: amy::Sender<disk::Request>,
}

impl RPC {
    pub fn start(
        creg: &mut amy::Registrar,
        db: amy::Sender<disk::Request>,
    ) -> io::Result<(handle::Handle<Message, CtlMessage>, thread::JoinHandle<()>)> {
        let poll = amy::Poller::new()?;
        let mut reg = poll.get_registrar()?;
        let cleanup = reg.set_interval(CLEANUP_INT_MS)?;
        let (ch, dh) = handle::Handle::new(creg, &mut reg)?;

        let ip = if CONFIG.rpc.local {
            Ipv4Addr::new(127, 0, 0, 1)
        } else {
            Ipv4Addr::new(0, 0, 0, 0)
        };
        let port = CONFIG.rpc.port;
        let listener = TcpListener::bind(SocketAddrV4::new(ip, port))?;
        listener.set_nonblocking(true)?;
        let lid = reg.register(&listener, amy::Event::Both)?;

        let disk = db.try_clone()?;
        let th = dh.run("rpc", move |ch| {
            RPC {
                ch,
                disk,
                poll,
                reg,
                listener,
                lid,
                cleanup,
                clients: UHashMap::default(),
                incoming: UHashMap::default(),
                processor: Processor::new(db),
                transfers: Transfers::new(),
            }.run()
        })?;
        Ok((ch, th))
    }

    pub fn run(&mut self) {
        debug!("Running RPC!");
        'outer: while let Ok(res) = self.poll.wait(POLL_INT_MS) {
            for not in res {
                match not.id {
                    id if id == self.lid => self.handle_accept(),
                    id if id == self.ch.rx.get_id() => {
                        if self.handle_ctl() {
                            return;
                        }
                    }
                    id if self.incoming.contains_key(&id) => self.handle_incoming(id),
                    id if id == self.cleanup => self.cleanup(),
                    id if self.transfers.contains(id) => self.handle_transfer(id),
                    _ => self.handle_conn(not),
                }
            }
        }
    }

    fn handle_ctl(&mut self) -> bool {
        while let Ok(m) = self.ch.recv() {
            match m {
                CtlMessage::Shutdown => return true,
                m => {
                    let msgs: Vec<_> = {
                        self.processor
                            .handle_ctl(m)
                            .into_iter()
                            .map(|(c, m)| (c, serde_json::to_string(&m).unwrap()))
                            .collect()
                    };
                    for (c, m) in msgs {
                        let res = match self.clients.get_mut(&c) {
                            Some(client) => client.send(ws::Frame::Text(m)),
                            None => {
                                debug!("Processor referenced a nonexistent client!");
                                Ok(())
                            }
                        };
                        if res.is_err() {
                            let client = self.clients.remove(&c).unwrap();
                            self.remove_client(c, client);
                        }
                    }
                }
            }
        }
        false
    }

    fn handle_transfer(&mut self, id: usize) {
        match self.transfers.ready(id) {
            TransferResult::Incomplete => {}
            TransferResult::Torrent {
                conn,
                data,
                path,
                client,
                serial,
                start,
            } => {
                debug!("Got torrent via HTTP transfer!");
                self.reg.deregister(&conn).unwrap();
                // TODO: Send this to the client in an error msg
                match bencode::decode_buf(&data) {
                    Ok(b) => {
                        match torrent::info::Info::from_bencode(b) {
                            Ok(i) => {
                                if self.ch
                                    .send(Message::Torrent {
                                        info: i,
                                        path,
                                        start,
                                        client,
                                        serial,
                                    })
                                    .is_err()
                                {
                                    error!("Failed to pass message to ctrl!");
                                }
                            }
                            Err(e) => {
                                error!("Failed to parse torrent data: {}!", e);
                                self.clients.get_mut(&client).map(|c| {
                                    c.send(ws::Frame::Text(
                                        serde_json::to_string(
                                            &SMessage::TransferFailed(message::Error {
                                                serial: Some(serial),
                                                reason: format!(
                                                    "Invalid torrent file uploaded, {}.",
                                                    e
                                                ),
                                            }),
                                        ).unwrap(),
                                    ))
                                });
                            }
                        }
                    }
                    Err(e) => {
                        error!("Failed to decode BE data: {}!", e);
                        self.clients.get_mut(&client).map(|c| {
                            c.send(ws::Frame::Text(
                                serde_json::to_string(
                                    &SMessage::TransferFailed(message::Error {
                                        serial: Some(serial),
                                        reason: format!(
                                            "Invalid torrent file uploaded, bad bencoded data: {}.",
                                            e
                                        ),
                                    }),
                                ).unwrap(),
                            ))
                        });
                    }
                }
            }
            TransferResult::Error {
                conn,
                err,
                client: id,
            } => {
                self.reg.deregister(&conn).unwrap();
                let res = self.clients
                    .get_mut(&id)
                    .map(|c| {
                        c.send(ws::Frame::Text(
                            serde_json::to_string(&SMessage::TransferFailed(err))
                                .unwrap(),
                        ))
                    })
                    .unwrap_or(Ok(()));
                if res.is_err() {
                    let client = self.clients.remove(&id).unwrap();
                    self.remove_client(id, client);
                }
            }
        }
    }

    fn handle_accept(&mut self) {
        loop {
            match self.listener.accept() {
                Ok((conn, ip)) => {
                    debug!("Accepted new connection from {:?}!", ip);
                    let id = self.reg.register(&conn, amy::Event::Both).unwrap();
                    self.incoming.insert(id, Incoming::new(conn));
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    break;
                }
                Err(e) => {
                    error!("Failed to accept conn: {}", e);
                }
            }
        }
    }

    fn handle_incoming(&mut self, id: usize) {
        if let Some(mut i) = self.incoming.remove(&id) {
            match i.readable() {
                Ok(IncomingStatus::Upgrade) => {
                    debug!("Succesfully upgraded conn");
                    self.clients.insert(id, i.into());
                }
                Ok(IncomingStatus::Incomplete) => {
                    self.incoming.insert(id, i);
                }
                Ok(IncomingStatus::Transfer { data, token }) => {
                    debug!("File transfer requested, validating");
                    match self.processor.get_transfer(token) {
                        Some((client,
                              serial,
                              TransferKind::UploadTorrent { path, size, start })) => {
                            debug!("Torrent transfer initiated");
                            self.transfers.add_torrent(
                                id,
                                client,
                                serial,
                                i.into(),
                                data,
                                path,
                                size,
                                start,
                            );
                            // Since a succesful result means the buffer hasn't been flushed,
                            // immediatly attempt to handle the transfer as if it was ready
                            self.handle_transfer(id);
                        }
                        Some(_) => {
                            error!("Unimplemented transfer type ignored");
                        }
                        None => {
                            error!("Transfer used invalid token");
                            // TODO: Handle downloads and other uploads
                        }
                    }
                }
                Ok(IncomingStatus::DL(id)) => {
                    debug!("Attempting DL of {}", id);
                    let mut conn: TcpStream = i.into();
                    self.reg.deregister(&conn).is_ok();
                    if let Some(path) = self.processor.get_dl(&id) {
                        debug!("Initiating DL");
                        conn.set_nonblocking(false).is_ok();
                        self.disk.send(disk::Request::download(conn, path)).ok();
                    } else {
                        debug!("ID {} invalid, stopping DL", id);
                        conn.write(&EMPTY_HTTP_RESP).ok();
                    }
                }
                Err(e) => {
                    debug!("Incoming ws upgrade failed: {}", e);
                    self.reg.deregister::<TcpStream>(&i.into()).unwrap();
                }
            }
        }
    }

    fn handle_conn(&mut self, not: amy::Notification) {
        if let Some(mut c) = self.clients.remove(&not.id) {
            if not.event.readable() {
                loop {
                    match c.read() {
                        Ok(None) => break,
                        Ok(Some(ws::Frame::Text(data))) => {
                            if self.process_frame(not.id, &mut c, &data).is_err() {
                                debug!("Client error, disconnecting");
                                self.remove_client(not.id, c);
                                return;
                            }
                        }
                        Err(Error(ErrorKind::Complete, _)) => {
                            info!("Client disconnected");
                            self.remove_client(not.id, c);
                            return;
                        }
                        Ok(_) | Err(_) => {
                            debug!("Client error, disconnecting");
                            self.remove_client(not.id, c);
                            return;
                        }
                    }
                }
            }
            if not.event.writable() && c.write().is_err() {
                self.remove_client(not.id, c);
                return;
            }
            self.clients.insert(not.id, c);
        }
    }

    fn process_frame(&mut self, id: usize, c: &mut Client, data: &str) -> result::Result<(), ()> {
        match serde_json::from_str(data) {
            Ok(m) => {
                let (msgs, rm) = self.processor.handle_client(id, m);
                if let Some(m) = rm {
                    self.ch.send(m).unwrap();
                }
                for msg in msgs {
                    if c.send(ws::Frame::Text(serde_json::to_string(&msg).unwrap()))
                        .is_err()
                    {
                        return Err(());
                    }
                }
            }
            Err(e) => {
                if e.is_syntax() || e.is_eof() {
                    let msg = SMessage::InvalidSchema(message::Error {
                        serial: None,
                        reason: format!("JSON decode error: {}", e),
                    });
                    if c.send(ws::Frame::Text(serde_json::to_string(&msg).unwrap()))
                        .is_err()
                    {}
                    return Err(());
                }
                if e.is_data() {
                    #[derive(Deserialize)]
                    struct Serial {
                        serial: u64,
                    }

                    let serial = match serde_json::from_str::<Serial>(data) {
                        Ok(s) => Some(s.serial),
                        Err(_) => None,
                    };

                    let msg = SMessage::InvalidSchema(message::Error {
                        serial,
                        reason: format!("Invalid message format: {}", e),
                    });
                    if c.send(ws::Frame::Text(serde_json::to_string(&msg).unwrap()))
                        .is_err()
                    {
                        return Err(());
                    }
                }
            }
        }
        Ok(())
    }

    fn cleanup(&mut self) {
        self.processor.remove_expired_tokens();
        let reg = &self.reg;
        let processor = &mut self.processor;
        self.clients.retain(|id, client| {
            let res = client.timed_out();
            if res {
                info!("client {} timed out", id);
                reg.deregister(&client.conn).unwrap();
                processor.remove_client(*id);
            }
            !res
        });
        self.incoming.retain(|_, inc| {
            let res = inc.timed_out();
            if res {
                reg.deregister(&inc.conn).unwrap();
            }
            !res
        });
        for (conn, id, err) in self.transfers.cleanup() {
            reg.deregister(&conn).unwrap();
            self.clients.get_mut(&id).map(|c| {
                c.send(ws::Frame::Text(
                    serde_json::to_string(&SMessage::TransferFailed(err))
                        .unwrap(),
                ))
            });
        }
    }

    fn remove_client(&mut self, id: usize, client: Client) {
        self.processor.remove_client(id);
        self.reg.deregister::<TcpStream>(&client.into()).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::path::{Path, PathBuf};
use std::fs::{OpenOptions, File};
use std::io::{self, Write};
use super::message::*;
use super::reader::*;
use super::Offset;

/// Number of bytes contained in the base name of the file.
pub static SEGMENT_FILE_NAME_LEN: usize = 20;
/// File extension for the segment file.
pub static SEGMENT_FILE_NAME_EXTENSION: &'static str = "log";

/// Magic that appears in the header of the segment for version 1.
///
/// There are a couple reasons for the magic. The primary reason is
/// to allow versioning, when the time comes. The second is to remove
/// the possibility of a 0 offset within the index. This helps to identity
/// the start of new index entries.
pub static VERSION_1_MAGIC: [u8; 2] = [0xff, 0xff];

/// A segment is a portion of the commit log. Segments are append-only logs written
/// until the maximum size is reached.
pub struct Segment {
    /// File descriptor
    file: File,

    /// Base offset of the log
    base_offset: u64,

    /// current file position for the write
    write_pos: usize,

    /// Maximum number of bytes permitted to be appended to the log
    max_bytes: usize,
}

#[derive(Debug)]
pub enum SegmentAppendError {
    LogFull,
    IoError(io::Error),
}

impl From<io::Error> for SegmentAppendError {
    #[inline]
    fn from(e: io::Error) -> SegmentAppendError {
        SegmentAppendError::IoError(e)
    }
}

impl Segment {
    pub fn new<P>(log_dir: P, base_offset: u64, max_bytes: usize) -> io::Result<Segment>
        where P: AsRef<Path>
    {
        let log_path = {
            // the log is of the form BASE_OFFSET.log
            let mut path_buf = PathBuf::new();
            path_buf.push(&log_dir);
            path_buf.push(format!("{:020}", base_offset));
            path_buf.set_extension(SEGMENT_FILE_NAME_EXTENSION);
            path_buf
        };

        let mut f = OpenOptions::new().write(true)
            .read(true)
            .create_new(true)
            .append(true)
            .open(&log_path)?;

        // add the magic
        f.write_all(&VERSION_1_MAGIC)?;

        Ok(Segment {
            file: f,
            base_offset: base_offset,
            write_pos: 2,
            max_bytes: max_bytes,
        })
    }

    pub fn open<P>(seg_path: P, max_bytes: usize) -> io::Result<Segment>
        where P: AsRef<Path>
    {
        let seg_file = OpenOptions::new().read(true)
            .write(true)
            .append(true)
            .open(&seg_path)?;

        let filename = seg_path.as_ref().file_name().unwrap().to_str().unwrap();
        let base_offset = match u64::from_str_radix(&filename[0..SEGMENT_FILE_NAME_LEN], 10) {
            Ok(v) => v,
            Err(_) => {
                return Err(io::Error::new(io::ErrorKind::InvalidData,
                                          "Segment file name does not parse as u64"))
            }
        };

        let meta = seg_file.metadata()?;
        // TODO: check magic

        info!("Opened segment {}", filename);

        Ok(Segment {
            file: seg_file,
            write_pos: meta.len() as usize,
            base_offset: base_offset,
            max_bytes: max_bytes,
        })
    }

    pub fn size(&self) -> usize {
        self.write_pos
    }

    #[inline]
    pub fn starting_offset(&self) -> u64 {
        self.base_offset
    }

    pub fn append<T: MessageSetMut>(&mut self,
                                    payload: &mut T,
                                    starting_offset: Offset)
                                    -> Result<Vec<LogEntryMetadata>, SegmentAppendError> {
        // ensure we have the capacity
        let payload_len = payload.bytes().len();
        if payload_len + self.write_pos > self.max_bytes {
            return Err(SegmentAppendError::LogFull);
        }

        let meta = super::message::set_offsets(payload, starting_offset, self.write_pos);

        self.file.write_all(payload.bytes())?;
        self.write_pos += payload_len;
        Ok(meta)
    }

    pub fn flush_sync(&mut self) -> io::Result<()> {
        self.file.flush()
    }

    pub fn read_slice<T: LogSliceReader>(&self,
                                         file_pos: u32,
                                         bytes: u32)
                                         -> Result<T::Result, MessageError> {
        T::read_from(&self.file, file_pos, bytes as usize)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::testutil::*;
    use test::Bencher;
    use std::path::PathBuf;

    #[test]
    pub fn log_append() {
        let path = TestDir::new();
        let mut f = Segment::new(path, 0, 1024).unwrap();

        {
            let mut buf = MessageBuf::default();
            buf.push("12345");
            let meta = f.append(&mut buf, 5).unwrap();

            assert_eq!(1, meta.len());
            let p0 = meta.iter().next().unwrap();
            assert_eq!(p0.offset, 5);
            assert_eq!(p0.file_pos, 2);
        }

        {
            let mut buf = MessageBuf::default();
            buf.push("66666");
            buf.push("77777");
            let meta = f.append(&mut buf, 6).unwrap();
            assert_eq!(2, meta.len());

            let mut it = meta.iter();
            let p0 = it.next().unwrap();
            assert_eq!(p0.offset, 6);
            assert_eq!(p0.file_pos, 27);

            let p1 = it.next().unwrap();
            assert_eq!(p1.offset, 7);
            assert_eq!(p1.file_pos, 52);
        }

        f.flush_sync().unwrap();
    }

    #[test]
    pub fn log_open() {
        let log_dir = TestDir::new();

        {
            let mut f = Segment::new(&log_dir, 0, 1024).unwrap();
            let mut buf = MessageBuf::default();
            buf.push("12345");
            buf.push("66666");
            f.append(&mut buf, 0).unwrap();
            f.flush_sync().unwrap();
        }

        // open it
        {
            let mut path_buf = PathBuf::new();
            path_buf.push(&log_dir);
            path_buf.push(format!("{:020}", 0));
            path_buf.set_extension(SEGMENT_FILE_NAME_EXTENSION);

            let res = Segment::open(&path_buf, 1024);
            assert!(res.is_ok(), "Err {:?}", res.err());

            let f = res.unwrap();
            assert_eq!(0, f.starting_offset());
        }
    }

    #[test]
    pub fn log_read() {
        let log_dir = TestDir::new();
        let mut f = Segment::new(&log_dir, 0, 1024).unwrap();

        {
            let mut buf = MessageBuf::default();
            buf.push("0123456789");
            buf.push("aaaaaaaaaa");
            buf.push("abc");
            f.append(&mut buf, 0).unwrap();
        }

        let msgs = f.read_slice::<MessageBufReader>(2, 83).unwrap();
        assert_eq!(3, msgs.len());

        for (i, m) in msgs.iter().enumerate() {
            assert_eq!(i as u64, m.offset());
        }
    }

    #[test]
    pub fn log_read_with_size_limit() {
        let log_dir = TestDir::new();
        let mut f = Segment::new(&log_dir, 0, 1024).unwrap();

        let meta = {
            let mut buf = MessageBuf::default();
            buf.push("0123456789");
            buf.push("aaaaaaaaaa");
            buf.push("abc");
            f.append(&mut buf, 0).unwrap()
        };

        // byte max contains message 0
        let msgs = f.read_slice::<MessageBufReader>(2, meta[1].file_pos - 2)
            .unwrap();

        assert_eq!(1, msgs.len());
    }

    #[test]
    pub fn log_read_from_write() {
        let log_dir = TestDir::new();
        let mut f = Segment::new(&log_dir, 0, 1024).unwrap();

        {
            let mut buf = MessageBuf::default();
            buf.push("0123456789");
            buf.push("aaaaaaaaaa");
            buf.push("abc");
            f.append(&mut buf, 0).unwrap();
        }

        let msgs = f.read_slice::<MessageBufReader>(2, 83).unwrap();
        assert_eq!(3, msgs.len());

        {
            let mut buf = MessageBuf::default();
            buf.push("foo");
            f.append(&mut buf, 3).unwrap();
        }

        let msgs = f.read_slice::<MessageBufReader>(2, 106).unwrap();
        assert_eq!(4, msgs.len());

        for (i, m) in msgs.iter().enumerate() {
            assert_eq!(i as u64, m.offset());
        }
    }

    #[bench]
    fn bench_segment_append(b: &mut Bencher) {
        let path = TestDir::new();

        let mut seg = Segment::new(path, 100u64, 100 * 1024 * 1024).unwrap();
        let payload = b"01234567891011121314151617181920";

        b.iter(|| {
            let mut buf = MessageBuf::default();
            buf.push(payload);
            seg.append(&mut buf, 0).unwrap();
        });
    }
}
<commit_msg>Checking magic when opening segment<commit_after>use std::path::{Path, PathBuf};
use std::fs::{OpenOptions, File};
use std::io::{self, Write};
use super::message::*;
use super::reader::*;
use super::Offset;
use std::os::unix::fs::FileExt;

/// Number of bytes contained in the base name of the file.
pub static SEGMENT_FILE_NAME_LEN: usize = 20;
/// File extension for the segment file.
pub static SEGMENT_FILE_NAME_EXTENSION: &'static str = "log";

/// Magic that appears in the header of the segment for version 1.
///
/// There are a couple reasons for the magic. The primary reason is
/// to allow versioning, when the time comes. The second is to remove
/// the possibility of a 0 offset within the index. This helps to identity
/// the start of new index entries.
pub static VERSION_1_MAGIC: [u8; 2] = [0xff, 0xff];

/// A segment is a portion of the commit log. Segments are append-only logs written
/// until the maximum size is reached.
pub struct Segment {
    /// File descriptor
    file: File,

    /// Base offset of the log
    base_offset: u64,

    /// current file position for the write
    write_pos: usize,

    /// Maximum number of bytes permitted to be appended to the log
    max_bytes: usize,
}

#[derive(Debug)]
pub enum SegmentAppendError {
    LogFull,
    IoError(io::Error),
}

impl From<io::Error> for SegmentAppendError {
    #[inline]
    fn from(e: io::Error) -> SegmentAppendError {
        SegmentAppendError::IoError(e)
    }
}

impl Segment {
    pub fn new<P>(log_dir: P, base_offset: u64, max_bytes: usize) -> io::Result<Segment>
        where P: AsRef<Path>
    {
        let log_path = {
            // the log is of the form BASE_OFFSET.log
            let mut path_buf = PathBuf::new();
            path_buf.push(&log_dir);
            path_buf.push(format!("{:020}", base_offset));
            path_buf.set_extension(SEGMENT_FILE_NAME_EXTENSION);
            path_buf
        };

        let mut f = OpenOptions::new().write(true)
            .read(true)
            .create_new(true)
            .append(true)
            .open(&log_path)?;

        // add the magic
        f.write_all(&VERSION_1_MAGIC)?;

        Ok(Segment {
            file: f,
            base_offset: base_offset,
            write_pos: 2,
            max_bytes: max_bytes,
        })
    }

    pub fn open<P>(seg_path: P, max_bytes: usize) -> io::Result<Segment>
        where P: AsRef<Path>
    {
        let seg_file = OpenOptions::new().read(true)
            .write(true)
            .append(true)
            .open(&seg_path)?;

        let filename = seg_path.as_ref().file_name().unwrap().to_str().unwrap();
        let base_offset = match u64::from_str_radix(&filename[0..SEGMENT_FILE_NAME_LEN], 10) {
            Ok(v) => v,
            Err(_) => {
                return Err(io::Error::new(io::ErrorKind::InvalidData,
                                          "Segment file name does not parse as u64"))
            }
        };

        let meta = seg_file.metadata()?;

        // check the magic
        {
            let mut bytes = [0u8; 2];
            let size = seg_file.read_at(&mut bytes, 0)?;
            if size < 2 || &bytes != &VERSION_1_MAGIC {
                return Err(io::Error::new(io::ErrorKind::InvalidData,
                                          format!("Segment file {} does not contain Version 1 magic",
                                          filename)));
            }
        }

        info!("Opened segment {}", filename);

        Ok(Segment {
            file: seg_file,
            write_pos: meta.len() as usize,
            base_offset: base_offset,
            max_bytes: max_bytes,
        })
    }

    pub fn size(&self) -> usize {
        self.write_pos
    }

    #[inline]
    pub fn starting_offset(&self) -> u64 {
        self.base_offset
    }

    pub fn append<T: MessageSetMut>(&mut self,
                                    payload: &mut T,
                                    starting_offset: Offset)
                                    -> Result<Vec<LogEntryMetadata>, SegmentAppendError> {
        // ensure we have the capacity
        let payload_len = payload.bytes().len();
        if payload_len + self.write_pos > self.max_bytes {
            return Err(SegmentAppendError::LogFull);
        }

        let meta = super::message::set_offsets(payload, starting_offset, self.write_pos);

        self.file.write_all(payload.bytes())?;
        self.write_pos += payload_len;
        Ok(meta)
    }

    pub fn flush_sync(&mut self) -> io::Result<()> {
        self.file.flush()
    }

    pub fn read_slice<T: LogSliceReader>(&self,
                                         file_pos: u32,
                                         bytes: u32)
                                         -> Result<T::Result, MessageError> {
        T::read_from(&self.file, file_pos, bytes as usize)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::testutil::*;
    use test::Bencher;
    use std::path::PathBuf;

    #[test]
    pub fn log_append() {
        let path = TestDir::new();
        let mut f = Segment::new(path, 0, 1024).unwrap();

        {
            let mut buf = MessageBuf::default();
            buf.push("12345");
            let meta = f.append(&mut buf, 5).unwrap();

            assert_eq!(1, meta.len());
            let p0 = meta.iter().next().unwrap();
            assert_eq!(p0.offset, 5);
            assert_eq!(p0.file_pos, 2);
        }

        {
            let mut buf = MessageBuf::default();
            buf.push("66666");
            buf.push("77777");
            let meta = f.append(&mut buf, 6).unwrap();
            assert_eq!(2, meta.len());

            let mut it = meta.iter();
            let p0 = it.next().unwrap();
            assert_eq!(p0.offset, 6);
            assert_eq!(p0.file_pos, 27);

            let p1 = it.next().unwrap();
            assert_eq!(p1.offset, 7);
            assert_eq!(p1.file_pos, 52);
        }

        f.flush_sync().unwrap();
    }

    #[test]
    pub fn log_open() {
        let log_dir = TestDir::new();

        {
            let mut f = Segment::new(&log_dir, 0, 1024).unwrap();
            let mut buf = MessageBuf::default();
            buf.push("12345");
            buf.push("66666");
            f.append(&mut buf, 0).unwrap();
            f.flush_sync().unwrap();
        }

        // open it
        {
            let mut path_buf = PathBuf::new();
            path_buf.push(&log_dir);
            path_buf.push(format!("{:020}", 0));
            path_buf.set_extension(SEGMENT_FILE_NAME_EXTENSION);

            let res = Segment::open(&path_buf, 1024);
            assert!(res.is_ok(), "Err {:?}", res.err());

            let f = res.unwrap();
            assert_eq!(0, f.starting_offset());
        }
    }

    #[test]
    pub fn log_read() {
        let log_dir = TestDir::new();
        let mut f = Segment::new(&log_dir, 0, 1024).unwrap();

        {
            let mut buf = MessageBuf::default();
            buf.push("0123456789");
            buf.push("aaaaaaaaaa");
            buf.push("abc");
            f.append(&mut buf, 0).unwrap();
        }

        let msgs = f.read_slice::<MessageBufReader>(2, 83).unwrap();
        assert_eq!(3, msgs.len());

        for (i, m) in msgs.iter().enumerate() {
            assert_eq!(i as u64, m.offset());
        }
    }

    #[test]
    pub fn log_read_with_size_limit() {
        let log_dir = TestDir::new();
        let mut f = Segment::new(&log_dir, 0, 1024).unwrap();

        let meta = {
            let mut buf = MessageBuf::default();
            buf.push("0123456789");
            buf.push("aaaaaaaaaa");
            buf.push("abc");
            f.append(&mut buf, 0).unwrap()
        };

        // byte max contains message 0
        let msgs = f.read_slice::<MessageBufReader>(2, meta[1].file_pos - 2)
            .unwrap();

        assert_eq!(1, msgs.len());
    }

    #[test]
    pub fn log_read_from_write() {
        let log_dir = TestDir::new();
        let mut f = Segment::new(&log_dir, 0, 1024).unwrap();

        {
            let mut buf = MessageBuf::default();
            buf.push("0123456789");
            buf.push("aaaaaaaaaa");
            buf.push("abc");
            f.append(&mut buf, 0).unwrap();
        }

        let msgs = f.read_slice::<MessageBufReader>(2, 83).unwrap();
        assert_eq!(3, msgs.len());

        {
            let mut buf = MessageBuf::default();
            buf.push("foo");
            f.append(&mut buf, 3).unwrap();
        }

        let msgs = f.read_slice::<MessageBufReader>(2, 106).unwrap();
        assert_eq!(4, msgs.len());

        for (i, m) in msgs.iter().enumerate() {
            assert_eq!(i as u64, m.offset());
        }
    }

    #[bench]
    fn bench_segment_append(b: &mut Bencher) {
        let path = TestDir::new();

        let mut seg = Segment::new(path, 100u64, 100 * 1024 * 1024).unwrap();
        let payload = b"01234567891011121314151617181920";

        b.iter(|| {
            let mut buf = MessageBuf::default();
            buf.push(payload);
            seg.append(&mut buf, 0).unwrap();
        });
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::sync::{Arc, Mutex};
use std::sync::atomic::AtomicUsize;

use futures::{future, Future, BoxFuture, Stream};
use tokio_proto::streaming::{Message, Body};
use tokio_service::Service;

use bytes::{BufMut, BytesMut};

use admin;
use codec::PacketHeader;
use packet::PacketMagic;
use queues::{HandleJobStorage, SharedJobStorage};
use worker::{SharedWorkers, Worker, Wake};
use constants::*;

pub type GearmanBody = Body<BytesMut, io::Error>;
pub type GearmanMessage = Message<PacketHeader, GearmanBody>;

fn new_res(ptype: u32, data: BytesMut) -> GearmanMessage {
    Message::WithBody(PacketHeader {
                          magic: PacketMagic::RES,
                          ptype: ptype,
                          psize: data.len() as u32,
                      },
                      Body::from(data))
}

pub struct GearmanService {
    pub conn_id: usize,
    pub queues: SharedJobStorage,
    pub workers: SharedWorkers,
    pub worker: Arc<Mutex<Worker>>,
    pub job_count: Arc<Mutex<AtomicUsize>>,
}

impl GearmanService {
    /// Things that don't require a body should use this
    fn response_from_header(&self,
                            header: &PacketHeader)
                            -> Message<PacketHeader, Body<BytesMut, io::Error>> {
        match header.ptype {
            ADMIN_VERSION => {
                let resp_str = "OK some-rustygear-version\n";
                let mut resp_body = BytesMut::with_capacity(resp_str.len());
                resp_body.put(&resp_str[..]);
                let resp_body = Body::from(resp_body);
                Message::WithBody(PacketHeader {
                                      magic: PacketMagic::TEXT,
                                      ptype: header.ptype,
                                      psize: resp_str.len() as u32,
                                  },
                                  resp_body)
            }
            ADMIN_STATUS => admin::admin_command_status(self.queues.clone(), self.workers.clone()),
            _ => {
                panic!("response_from_header called with invalid ptype: {}",
                       header.ptype)
            }
        }
    }

    pub fn new(conn_id: usize,
               queues: SharedJobStorage,
               workers: SharedWorkers,
               job_count: Arc<Mutex<AtomicUsize>>)
               -> GearmanService {
        GearmanService {
            conn_id: conn_id,
            queues: queues.clone(),
            worker: Arc::new(Mutex::new((Worker::new()))),
            workers: workers.clone(),
            job_count: job_count.clone(),
        }
    }

    fn handle_can_do(&self, body: GearmanBody) -> BoxFuture<GearmanMessage, io::Error> {
        let worker = self.worker.clone();
        let workers = self.workers.clone();
        let conn_id = self.conn_id;
        body.concat2()
            .and_then(move |fname| {
                let fname = fname.freeze();
                debug!("CAN_DO fname = {:?}", fname);
                let mut worker = worker.lock().unwrap();
                worker.can_do(fname);
                workers.clone().wakeup(&mut worker, conn_id);
                future::finished(Message::WithBody(PacketHeader::noop(),
                                                   Body::from(BytesMut::new())))
            })
            .boxed()
    }

    fn handle_grab_job_all(&self) -> BoxFuture<GearmanMessage, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        if queues.get_job(worker) {
            match worker.job() {
                Some(ref j) => {
                    let mut data = BytesMut::with_capacity(4 + j.handle.len() + j.fname.len() +
                                                           j.unique.len() +
                                                           j.data.len());
                    data.extend(&j.handle);
                    data.put_slice(b"\0");
                    data.extend(&j.fname);
                    data.put_slice(b"\0");
                    data.extend(&j.unique);
                    data.put_slice(b"\0");
                    // reducer not implemented
                    data.put_slice(b"\0");
                    data.extend(&j.data);
                    return future::finished(new_res(JOB_ASSIGN_ALL, data)).boxed();
                }
                None => {}
            }
        };
        future::finished(new_res(NO_JOB, BytesMut::new())).boxed()
    }
}

impl Service for GearmanService {
    type Request = GearmanMessage;
    type Response = GearmanMessage;
    type Error = io::Error;
    type Future = BoxFuture<Self::Response, Self::Error>;

    fn call(&self, req: Self::Request) -> Self::Future {
        debug!("Got a req {:?}", req);
        match req {
            Message::WithoutBody(header) => {
                match header.ptype {
                    ADMIN_VERSION | ADMIN_STATUS => {
                        future::ok(self.response_from_header(&header)).boxed()
                    }

                    _ => {
                        future::err(io::Error::new(io::ErrorKind::Other,
                                                   format!("Bodyless packet type = {}",
                                                           header.ptype)))
                            .boxed()
                    }
                }
            }
            Message::WithBody(header, body) => {
                match header.ptype {
                    /*
                    SUBMIT_JOB => self.handle_submit_job(PRIORITY_NORMAL, false),
                    SUBMIT_JOB_HIGH => self.handle_submit_job(PRIORITY_HIGH, false),
                    SUBMIT_JOB_LOW => self.handle_submit_job(PRIORITY_LOW, false),
                    SUBMIT_JOB_BG => self.handle_submit_job(PRIORITY_NORMAL, true),
                    SUBMIT_JOB_HIGH_BG => self.handle_submit_job(PRIORITY_HIGH, true),
                    SUBMIT_JOB_LOW_BG => self.handle_submit_job(PRIORITY_LOW, true),
                    PRE_SLEEP => self.handle_pre_sleep(),
                    */
                    CAN_DO => self.handle_can_do(body),/*
                    CANT_DO => self.handle_cant_do(),
                    GRAB_JOB => self.handle_grab_job(),
                    GRAB_JOB_UNIQ => self.handle_grab_job_uniq(),*/
                    GRAB_JOB_ALL => self.handle_grab_job_all(),/*
                    WORK_COMPLETE => self.handle_work_complete(),
                    WORK_STATUS | WORK_DATA | WORK_WARNING => self.handle_work_update(),
                    ECHO_REQ => self.handle_echo_req(&req),*/
                    _ => {
                        error!("Unimplemented: {:?} processing packet", header);
                        future::err(io::Error::new(io::ErrorKind::Other,
                                                   format!("Invalid packet type {}", header.ptype)))
                            .boxed()
                    }
                }
            }
        }
    }
}
<commit_msg>handle PRE_SLEEP<commit_after>use std::io;
use std::sync::{Arc, Mutex};
use std::sync::atomic::AtomicUsize;

use futures::{future, Future, BoxFuture, Stream};
use tokio_proto::streaming::{Message, Body};
use tokio_service::Service;

use bytes::{BufMut, BytesMut};

use admin;
use codec::PacketHeader;
use packet::PacketMagic;
use queues::{HandleJobStorage, SharedJobStorage};
use worker::{SharedWorkers, Worker, Wake};
use constants::*;

pub type GearmanBody = Body<BytesMut, io::Error>;
pub type GearmanMessage = Message<PacketHeader, GearmanBody>;

fn new_res(ptype: u32, data: BytesMut) -> GearmanMessage {
    Message::WithBody(PacketHeader {
                          magic: PacketMagic::RES,
                          ptype: ptype,
                          psize: data.len() as u32,
                      },
                      Body::from(data))
}

pub struct GearmanService {
    pub conn_id: usize,
    pub queues: SharedJobStorage,
    pub workers: SharedWorkers,
    pub worker: Arc<Mutex<Worker>>,
    pub job_count: Arc<Mutex<AtomicUsize>>,
}

impl GearmanService {
    /// Things that don't require a body should use this
    fn response_from_header(&self,
                            header: &PacketHeader)
                            -> Message<PacketHeader, Body<BytesMut, io::Error>> {
        match header.ptype {
            ADMIN_VERSION => {
                let resp_str = "OK some-rustygear-version\n";
                let mut resp_body = BytesMut::with_capacity(resp_str.len());
                resp_body.put(&resp_str[..]);
                let resp_body = Body::from(resp_body);
                Message::WithBody(PacketHeader {
                                      magic: PacketMagic::TEXT,
                                      ptype: header.ptype,
                                      psize: resp_str.len() as u32,
                                  },
                                  resp_body)
            }
            ADMIN_STATUS => admin::admin_command_status(self.queues.clone(), self.workers.clone()),
            _ => {
                panic!("response_from_header called with invalid ptype: {}",
                       header.ptype)
            }
        }
    }

    pub fn new(conn_id: usize,
               queues: SharedJobStorage,
               workers: SharedWorkers,
               job_count: Arc<Mutex<AtomicUsize>>)
               -> GearmanService {
        GearmanService {
            conn_id: conn_id,
            queues: queues.clone(),
            worker: Arc::new(Mutex::new((Worker::new()))),
            workers: workers.clone(),
            job_count: job_count.clone(),
        }
    }

    fn handle_can_do(&self, body: GearmanBody) -> BoxFuture<GearmanMessage, io::Error> {
        let worker = self.worker.clone();
        let workers = self.workers.clone();
        let conn_id = self.conn_id;
        body.concat2()
            .and_then(move |fname| {
                let fname = fname.freeze();
                debug!("CAN_DO fname = {:?}", fname);
                let mut worker = worker.lock().unwrap();
                worker.can_do(fname);
                workers.clone().wakeup(&mut worker, conn_id);
                future::finished(Message::WithBody(PacketHeader::noop(),
                                                   Body::from(BytesMut::new())))
            })
            .boxed()
    }

    fn handle_grab_job_all(&self) -> BoxFuture<GearmanMessage, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        if queues.get_job(worker) {
            match worker.job() {
                Some(ref j) => {
                    let mut data = BytesMut::with_capacity(4 + j.handle.len() + j.fname.len() +
                                                           j.unique.len() +
                                                           j.data.len());
                    data.extend(&j.handle);
                    data.put_slice(b"\0");
                    data.extend(&j.fname);
                    data.put_slice(b"\0");
                    data.extend(&j.unique);
                    data.put_slice(b"\0");
                    // reducer not implemented
                    data.put_slice(b"\0");
                    data.extend(&j.data);
                    return future::finished(new_res(JOB_ASSIGN_ALL, data)).boxed();
                }
                None => {}
            }
        };
        future::finished(new_res(NO_JOB, BytesMut::new())).boxed()
    }

    fn handle_pre_sleep(&self) -> BoxFuture<GearmanMessage, io::Error> {
        let worker = self.worker.clone();
        let ref mut w = worker.lock().unwrap();
        self.workers.clone().sleep(w, self.conn_id);
        future::finished(Message::WithBody(PacketHeader::noop(), Body::from(BytesMut::new()))).boxed()
    }
}

impl Service for GearmanService {
    type Request = GearmanMessage;
    type Response = GearmanMessage;
    type Error = io::Error;
    type Future = BoxFuture<Self::Response, Self::Error>;

    fn call(&self, req: Self::Request) -> Self::Future {
        debug!("Got a req {:?}", req);
        match req {
            Message::WithoutBody(header) => {
                match header.ptype {
                    ADMIN_VERSION | ADMIN_STATUS => {
                        future::ok(self.response_from_header(&header)).boxed()
                    }

                    _ => {
                        future::err(io::Error::new(io::ErrorKind::Other,
                                                   format!("Bodyless packet type = {}",
                                                           header.ptype)))
                            .boxed()
                    }
                }
            }
            Message::WithBody(header, body) => {
                match header.ptype {
                    /*
                    SUBMIT_JOB => self.handle_submit_job(PRIORITY_NORMAL, false),
                    SUBMIT_JOB_HIGH => self.handle_submit_job(PRIORITY_HIGH, false),
                    SUBMIT_JOB_LOW => self.handle_submit_job(PRIORITY_LOW, false),
                    SUBMIT_JOB_BG => self.handle_submit_job(PRIORITY_NORMAL, true),
                    SUBMIT_JOB_HIGH_BG => self.handle_submit_job(PRIORITY_HIGH, true),
                    SUBMIT_JOB_LOW_BG => self.handle_submit_job(PRIORITY_LOW, true),*/
                    PRE_SLEEP => self.handle_pre_sleep(),
                    CAN_DO => self.handle_can_do(body),/*
                    CANT_DO => self.handle_cant_do(),
                    GRAB_JOB => self.handle_grab_job(),
                    GRAB_JOB_UNIQ => self.handle_grab_job_uniq(),*/
                    GRAB_JOB_ALL => self.handle_grab_job_all(),/*
                    WORK_COMPLETE => self.handle_work_complete(),
                    WORK_STATUS | WORK_DATA | WORK_WARNING => self.handle_work_update(),
                    ECHO_REQ => self.handle_echo_req(&req),*/
                    _ => {
                        error!("Unimplemented: {:?} processing packet", header);
                        future::err(io::Error::new(io::ErrorKind::Other,
                                                   format!("Invalid packet type {}", header.ptype)))
                            .boxed()
                    }
                }
            }
        }
    }
}
<|endoftext|>


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

//! A simple test app for exercising and debugging the regex engine.

extern crate fancy_regex;

use fancy_regex::*;
use fancy_regex::analyze::analyze;
use fancy_regex::compile::compile;
use fancy_regex::vm::{Insn, Prog};
use std::env;
use std::str::FromStr;

fn main() {
    let mut args = env::args().skip(1);
    if let Some(cmd) = args.next() {
        if cmd == "parse" {
            if let Some(re) = args.next() {
                let e = Expr::parse(&re);
                println!("{:#?}", e);
            }
        } else if cmd == "analyze" {
            if let Some(re) = args.next() {
                let (e, backrefs) = Expr::parse(&re).unwrap();
                let a = analyze(&e, &backrefs);
                println!("{:#?}", a);
            }
        } else if cmd == "compile" {
            if let Some(re) = args.next() {
                let r = Regex::new(&re).unwrap();
                r.debug_print();
            }
        } else if cmd == "run" {
            let re = args.next().expect("expected regexp argument");
            let r = Regex::new(&re).unwrap();
            let text = args.next().expect("expected text argument");
            let mut pos = 0;
            if let Some(pos_str) = args.next() {
                pos = usize::from_str(&pos_str).unwrap();
            }
            if let Some(caps) = r.captures_from_pos(&text, pos).unwrap() {
                print!("captures:");
                for i in 0..caps.len() {
                    print!(" {}:", i);
                    if let Some((lo, hi)) = caps.pos(i) {
                        print!("[{}..{}] \"{}\"", lo, hi, caps.at(i).unwrap());
                    } else {
                        print!("_");
                    }
                }
                println!("");
                for cap in caps.iter() {
                    println!("iterate {:?}", cap);
                }
            }
        } else if cmd == "trace" {
            if let Some(re) = args.next() {
                let prog = prog(&re);
                if let Some(s) = args.next() {
                    vm::run(&prog, &s, 0, vm::OPTION_TRACE).unwrap();
                }
            }
        } else if cmd == "trace-inner" {
            if let Some(re) = args.next() {
                let (e, backrefs) = Expr::parse(&re).unwrap();
                let a = analyze(&e, &backrefs).unwrap();
                let p = compile(&a).unwrap();
                if let Some(s) = args.next() {
                    vm::run(&p, &s, 0, vm::OPTION_TRACE).unwrap();
                }
            }
        } else if cmd == "graph" {
            let re = args.next().expect("expected regexp argument");
            graph(&re);
        } else {
            println!("commands: parse|analyze|compile|graph <expr>, run|trace|trace-inner <expr> <input>");
        }
    }
}

fn graph(re: &str) {
    let prog = prog(re);
    println!("digraph G {{");
    for (i, insn) in prog.body.iter().enumerate() {
        let label = format!("{:?}", insn).replace("\"", "\\\"");
        println!(r#"{:3} [label="{}: {}"];"#, i, i, label);
        match *insn {
            Insn::Split(a, b) => {
                println!("{:3} -> {};", i, a);
                println!("{:3} -> {};", i, b);
            }
            Insn::Jmp(target) => {
                println!("{:3} -> {};", i, target);
            }
            Insn::End => {}
            _ => {
                println!("{:3} -> {};", i, i+1);
            }
        }
    }
    println!("}}");
}

fn prog(re: &str) -> Prog {
    let r = Regex::new(re).unwrap();
    match r {
        Regex::Impl { prog, .. } => prog,
        _ => {
            panic!("Expected regex to be fancy (not completely delegatable)");
        }
    }
}
<commit_msg>Some minor fixes for the toy example<commit_after>// Copyright 2016 Google Inc. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! A simple test app for exercising and debugging the regex engine.

extern crate fancy_regex;

use fancy_regex::*;
use fancy_regex::analyze::analyze;
use fancy_regex::compile::compile;
use fancy_regex::vm::{Insn, Prog};
use std::env;
use std::str::FromStr;

fn main() {
    let mut args = env::args().skip(1);
    if let Some(cmd) = args.next() {
        if cmd == "parse" {
            if let Some(re) = args.next() {
                let e = Expr::parse(&re);
                println!("{:#?}", e);
            }
        } else if cmd == "analyze" {
            if let Some(re) = args.next() {
                let (e, backrefs) = Expr::parse(&re).unwrap();
                let a = analyze(&e, &backrefs);
                println!("{:#?}", a);
            }
        } else if cmd == "compile" {
            if let Some(re) = args.next() {
                let r = Regex::new(&re).unwrap();
                r.debug_print();
            }
        } else if cmd == "run" {
            let re = args.next().expect("expected regexp argument");
            let r = Regex::new(&re).unwrap();
            let text = args.next().expect("expected text argument");
            let mut pos = 0;
            if let Some(pos_str) = args.next() {
                pos = usize::from_str(&pos_str).unwrap();
            }
            if let Some(caps) = r.captures_from_pos(&text, pos).unwrap() {
                print!("captures:");
                for i in 0..caps.len() {
                    print!(" {}:", i);
                    if let Some((lo, hi)) = caps.pos(i) {
                        print!("[{}..{}] \"{}\"", lo, hi, caps.at(i).unwrap());
                    } else {
                        print!("_");
                    }
                }
                println!("");
                for cap in caps.iter() {
                    println!("iterate {:?}", cap);
                }
            } else {
                println!("no match");
            }
        } else if cmd == "trace" {
            if let Some(re) = args.next() {
                let prog = prog(&re);
                if let Some(s) = args.next() {
                    vm::run(&prog, &s, 0, vm::OPTION_TRACE).unwrap();
                }
            }
        } else if cmd == "trace-inner" {
            if let Some(re) = args.next() {
                let (e, backrefs) = Expr::parse(&re).unwrap();
                let a = analyze(&e, &backrefs).unwrap();
                let p = compile(&a).unwrap();
                if let Some(s) = args.next() {
                    vm::run(&p, &s, 0, vm::OPTION_TRACE).unwrap();
                }
            }
        } else if cmd == "graph" {
            let re = args.next().expect("expected regexp argument");
            graph(&re);
        } else {
            println!("commands: parse|analyze|compile|graph <expr>, run|trace|trace-inner <expr> <input>");
        }
    }
}

fn graph(re: &str) {
    let prog = prog(re);
    println!("digraph G {{");
    for (i, insn) in prog.body.iter().enumerate() {
        let label = format!("{:?}", insn)
            .replace(r#"\"#, r#"\\"#)
            .replace(r#"""#, r#"\""#);
        println!(r#"{:3} [label="{}: {}"];"#, i, i, label);
        match *insn {
            Insn::Split(a, b) => {
                println!("{:3} -> {};", i, a);
                println!("{:3} -> {};", i, b);
            }
            Insn::Jmp(target) => {
                println!("{:3} -> {};", i, target);
            }
            Insn::End => {}
            _ => {
                println!("{:3} -> {};", i, i+1);
            }
        }
    }
    println!("}}");
}

fn prog(re: &str) -> Prog {
    let r = Regex::new(re).unwrap();
    match r {
        Regex::Impl { prog, .. } => prog,
        _ => {
            panic!("Expected regex to be fancy (not completely delegatable)");
        }
    }
}
<|endoftext|>