<commit_before>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate rbpf;

use self::rbpf::{disassembler, ebpf};
use super::address::Address;
use super::instruction::Instruction;

#[allow(missing_docs)]
#[derive(Debug)]
pub struct BpfInstruction {
    idx: u64,
    insn: disassembler::HLInsn,
}

impl BpfInstruction {
    #[allow(missing_docs)]
    pub fn new(idx: u64, insn: disassembler::HLInsn) -> Self {
        BpfInstruction { idx, insn }
    }
}

impl Instruction for BpfInstruction {
    fn address(&self) -> Address {
        Address::new(self.idx)
    }

    fn comment(&self) -> Option<String> {
        None
    }

    fn mnemonic(&self) -> &str {
        &*self.insn.name
    }

    fn cycle_count(&self) -> Option<u32> {
        None
    }

    fn is_call(&self) -> bool {
        self.insn.opc == ebpf::CALL
    }

    fn is_local_conditional_jump(&self) -> bool {
        self.is_local_jump() && (self.insn.opc != ebpf::JA)
    }

    fn is_local_jump(&self) -> bool {
        (self.insn.opc & 0x07) == ebpf::BPF_JMP && (self.insn.opc != ebpf::CALL) &&
            (self.insn.opc != ebpf::TAIL_CALL) && (self.insn.opc != ebpf::EXIT)
    }

    fn is_return(&self) -> bool {
        self.insn.opc == ebpf::EXIT
    }

    fn target_address(&self) -> Option<Address> {
        if self.is_local_jump() {
            Some(Address::new(
                (self.idx as i64 + i64::from(self.insn.off)) as u64,
            ))
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{rbpf, BpfInstruction};
    use super::super::{Address, Function, Symbol};

    #[test]
    fn test() {
        let prog = &[
            0xb7, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // mov r0, 0
            0x79, 0x12, 0x40, 0x00, 0x00, 0x00, 0x00, 0x00, // load ptr from r1[0x40] to r2
            0x07, 0x02, 0x00, 0x00, 0x05, 0x00, 0x00, 0x00, // add r2, 5
            0x79, 0x11, 0x50, 0x00, 0x00, 0x00, 0x00, 0x00, // load ptr from r1[0x50] to r1
            0x2d, 0x12, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, // if r2 > r1 skip 3 instructions
            0x71, 0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // load r2 (= *(mem + 5)) into r0
            0x67, 0x00, 0x00, 0x00, 0x38, 0x00, 0x00, 0x00, // r0 >>= 56
            0xc7, 0x00, 0x00, 0x00, 0x38, 0x00, 0x00, 0x00, // r0 <<= 56 (arsh)
            0x95, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00  // exit
        ];

        let v = rbpf::disassembler::to_insn_vec(prog);
        let is = v.into_iter()
            .enumerate()
            .map(|(idx, insn)| BpfInstruction::new(idx as u64, insn))
            .collect::<Vec<_>>();

        let f = Function::new(Symbol::new(Address::new(100000), Some("test")), is);

        assert!(f.control_flow_graph.entry_block.is_some());
        assert_eq!(f.control_flow_graph.graph.node_count(), 3);
    }
}
<commit_msg>Add Function<BpfInstruction>::from_bpf().<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate rbpf;

use self::rbpf::{disassembler, ebpf};
use super::address::Address;
use super::function::Function;
use super::instruction::Instruction;
use super::symbol::Symbol;

#[allow(missing_docs)]
#[derive(Debug)]
pub struct BpfInstruction {
    idx: u64,
    insn: disassembler::HLInsn,
}

impl BpfInstruction {
    #[allow(missing_docs)]
    pub fn new(idx: u64, insn: disassembler::HLInsn) -> Self {
        BpfInstruction { idx, insn }
    }
}

impl Instruction for BpfInstruction {
    fn address(&self) -> Address {
        Address::new(self.idx)
    }

    fn comment(&self) -> Option<String> {
        None
    }

    fn mnemonic(&self) -> &str {
        &*self.insn.name
    }

    fn cycle_count(&self) -> Option<u32> {
        None
    }

    fn is_call(&self) -> bool {
        self.insn.opc == ebpf::CALL
    }

    fn is_local_conditional_jump(&self) -> bool {
        self.is_local_jump() && (self.insn.opc != ebpf::JA)
    }

    fn is_local_jump(&self) -> bool {
        (self.insn.opc & 0x07) == ebpf::BPF_JMP && (self.insn.opc != ebpf::CALL) &&
            (self.insn.opc != ebpf::TAIL_CALL) && (self.insn.opc != ebpf::EXIT)
    }

    fn is_return(&self) -> bool {
        self.insn.opc == ebpf::EXIT
    }

    fn target_address(&self) -> Option<Address> {
        if self.is_local_jump() {
            Some(Address::new(
                (self.idx as i64 + i64::from(self.insn.off)) as u64,
            ))
        } else {
            None
        }
    }
}

impl Function<BpfInstruction> {
    /// Create a function from eBPF bytecode.
    pub fn from_bpf(symbol: Symbol, data: &[u8]) -> Function<BpfInstruction> {
        let v = rbpf::disassembler::to_insn_vec(data);
        let is = v.into_iter()
            .enumerate()
            .map(|(idx, insn)| BpfInstruction::new(idx as u64, insn))
            .collect::<Vec<_>>();
        Function::new(symbol, is)
    }
}

#[cfg(test)]
mod tests {
    use super::super::{Address, Function, Symbol};

    #[test]
    fn test() {
        let prog = &[
            0xb7, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // mov r0, 0
            0x79, 0x12, 0x40, 0x00, 0x00, 0x00, 0x00, 0x00, // load ptr from r1[0x40] to r2
            0x07, 0x02, 0x00, 0x00, 0x05, 0x00, 0x00, 0x00, // add r2, 5
            0x79, 0x11, 0x50, 0x00, 0x00, 0x00, 0x00, 0x00, // load ptr from r1[0x50] to r1
            0x2d, 0x12, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, // if r2 > r1 skip 3 instructions
            0x71, 0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // load r2 (= *(mem + 5)) into r0
            0x67, 0x00, 0x00, 0x00, 0x38, 0x00, 0x00, 0x00, // r0 >>= 56
            0xc7, 0x00, 0x00, 0x00, 0x38, 0x00, 0x00, 0x00, // r0 <<= 56 (arsh)
            0x95, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00  // exit
        ];

        let f = Function::from_bpf(Symbol::new(Address::new(100000), Some("test")), prog);

        assert!(f.control_flow_graph.entry_block.is_some());
        assert_eq!(f.control_flow_graph.graph.node_count(), 3);
    }
}
<|endoftext|>


<commit_before>// src/cli.rs
// Argument parser and CLI (Command Line Interface)
// for BreakFast

use clap::{App, Arg, SubCommand};

pub fn build_cli() -> App<'static, 'static> {
    App::new("breakfastrs")
        .version("0.1")
        .about("BreakFast is a toolkit for detecting chromosomal rearrangements based on whole genome sequencing data.")
        .usage("\tbreakfast detect <bam_file> <genome> <out_prefix> [-a N] [-f N] [-q N] [-O OR] [--discard-duplicates=METHOD]\n\
        \tbreakfast detect specific [-A] <bam_file> <donors> <acceptors> <genome> <out_prefix>\n\
        \tbreakfast filter <sv_file> [-r P-S-A]... [--blacklist=PATH]\n\
        \tbreakfast annotate <sv_file> <bed_file>\n\
        \tbreakfast blacklist [--freq-above=FREQ] <sv_files>...\n\
        \tbreakfast visualize <sv_file>\n\
        \tbreakfast tabulate rearranged genes <sv_files>...\n\
        \tbreakfast tabulate fusions <sv_files>...\n\
        \tbreakfast statistics <sv_files>...\n\
        \tbreakfast filter by region <sv_file> <region>\n\
        \tbreakfast filter by distance <min_distance> <sv_file>\n\
        \tbreakfast align junction <reads>")

        .arg(Arg::with_name("anchor-len")
            .short("a")
            .long("anchor-len")
            .global(true)
            .takes_value(true)
            .value_name("N")
            .display_order(1)
            .help("Anchor length for split read analysis. When zero, split reads are not used [default: 0]"))

        .arg(Arg::with_name("max-frag-len")
            .short("f")
            .long("max-frag-len")
            .global(true)
            .takes_value(true)
            .value_name("N")
            .display_order(2)
            .help("Maximum fragment length [default: 5000]"))

        .arg(Arg::with_name("min-mapq")
            .short("q")
            .long("min-mapq")
            .global(true)
            .takes_value(true)
            .value_name("N")
            .display_order(3)
            .help("Minimum mapping quality to consider [default: 15]"))

        .arg(Arg::with_name("orientation")
            .short("O")
            .long("orientation")
            .global(true)
            .takes_value(true)
            .value_name("OR")
            .display_order(4)
            .help("Read pair orientation produced by sequencer. Either 'fr' (converging), 'rf' (diverging) or 'ff' [default: fr]"))

        .arg(Arg::with_name("all-reads")
            .short("A")
            .long("all-reads")
            .global(true)
            .display_order(5)
            .help("Use all reads for rearrangement detection, not just unaligned reads."))

        .arg(Arg::with_name("discard-duplicates")
            .long("discard-duplicates")
            .global(true)
            .takes_value(true)
            .value_name("Method")
            .display_order(6)
            .help("Method to use when discarding duplicate reads.'both-ends' considers a read pair (or unaligned read) to be a duplicate of another if the positions"))

        .arg(Arg::with_name("blacklist")
            .long("blacklist")
            .global(true)
            .takes_value(true)
            .value_name("list")
            .display_order(7)
            .help("Path to a file containing blacklisted regions."))

        .arg(Arg::with_name("min-reads")
            .short("r")
            .long("min-reads")
            .global(true)
            .takes_value(true)
            .value_name("P-S-A")
            .display_order(8)
            .help("Minimum number of spanning reads required to accept a breakpoint. Specified in the format P-S-A, where P=paired, S=split, A=either. For example, -r 1-2-0 would require at least one mate pair and two split reads of evidence [default: 0-0-0]."))

        .arg(Arg::with_name("freq-above")
            .long("freq-above")
            .global(true)
            .takes_value(true)
            .value_name("FREQ")
            .display_order(9)
            .help("Minimum frequency at which a variant must be present among the control samples to be considered a false positive [default: 0]."))

        .subcommand(SubCommand::with_name("detect")
            .usage("\tbreakfast detect <bam_file> <genome> <out_prefix> [-a N] [-f N] [-q N] [-O OR] [--discard-duplicates=METHOD]")
            .display_order(1)
            .arg(Arg::with_name("bam_file")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("genome")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("out_prefix")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("detect-specific")
            .usage("breakfast detect specific [-A] <bam_file> <donors> <acceptors> <genome> <out_prefix>")
            .display_order(2)
            .arg(Arg::with_name("bam_file")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("donors")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("acceptors")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("genome")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("out_prefix")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("filter")
            .usage("breakfast filter <sv_file> [-r P-S-A]... [--blacklist=PATH]")
            .display_order(3)
            .arg(Arg::with_name("sv_file")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("annotate")
            .usage("breakfast annotate <sv_file> <bed_file>")
            .display_order(4)
            .arg(Arg::with_name("sv_file")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("bed_file")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("blacklist")
            .usage("breakfast blacklist [--freq-above=FREQ] <sv_files>...")
            .display_order(5)
            .arg(Arg::with_name("sv_files")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("visualize")
            .usage("breakfast visualize <sv_file>")
            .display_order(6)
            .arg(Arg::with_name("sv_file")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("tabulate-rearranged-genes")
            .usage("breakfast tabulate rearranged genes <sv_files>...")
            .display_order(7)
            .arg(Arg::with_name("sv_files")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("tabulate-fusions")
            .usage("breakfast tabulate fusions <sv_files>...")
            .display_order(8)
            .arg(Arg::with_name("sv_files")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("statistics")
            .usage("breakfast statistics <sv_files>..")
            .display_order(9)
            .arg(Arg::with_name("sv_files")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("filter-by-region")
            .usage("breakfast filter by region <sv_file> <region>")
            .display_order(10)
            .arg(Arg::with_name("sv_file")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("region")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("filter-by-distance")
            .usage("breakfast filter by distance <min_distance> <sv_file>")
            .display_order(11)
            .arg(Arg::with_name("min_distance")
                .takes_value(true)
                .required(true))
            .arg(Arg::with_name("sv_file")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("align-junction")
            .usage("breakfast align junction <reads>")
            .display_order(12)
            .arg(Arg::with_name("reads")
                .takes_value(true)
                .required(true)))
}
<commit_msg>CLI Complete<commit_after>// src/cli.rs
// Argument parser and CLI (Command Line Interface)
// for BreakFast

use clap::{App, Arg, SubCommand};

pub fn build_cli() -> App<'static, 'static> {
    App::new("breakfastrs")
        .version("0.1")
        .about("BreakFast is a toolkit for detecting chromosomal rearrangements based on whole genome sequencing data.")
        .usage("\tbreakfast detect <bam_file> <genome> <out_prefix> [-a N] [-f N] [-q N] [-O OR] [--discard-duplicates=METHOD]\n\
        \tbreakfast detect specific [-A] <bam_file> <donors> <acceptors> <genome> <out_prefix>\n\
        \tbreakfast filter <sv_file> [-r P-S-A]... [--blacklist=PATH]\n\
        \tbreakfast annotate <sv_file> <bed_file>\n\
        \tbreakfast blacklist [--freq-above=FREQ] <sv_files>...\n\
        \tbreakfast visualize <sv_file>\n\
        \tbreakfast tabulate rearranged genes <sv_files>...\n\
        \tbreakfast tabulate fusions <sv_files>...\n\
        \tbreakfast statistics <sv_files>...\n\
        \tbreakfast filter by region <sv_file> <region>\n\
        \tbreakfast filter by distance <min_distance> <sv_file>\n\
        \tbreakfast align junction <reads>")

        .arg(Arg::with_name("anchor-len")
            .short("a")
            .long("anchor-len")
            .takes_value(true)
            .value_name("N")
            .display_order(1)
            .help("Anchor length for split read analysis. When zero, split reads are not used [default: 0]"))

        .arg(Arg::with_name("max-frag-len")
            .short("f")
            .long("max-frag-len")
            .takes_value(true)
            .value_name("N")
            .display_order(2)
            .help("Maximum fragment length [default: 5000]"))

        .arg(Arg::with_name("min-mapq")
            .short("q")
            .long("min-mapq")
            .takes_value(true)
            .value_name("N")
            .display_order(3)
            .help("Minimum mapping quality to consider [default: 15]"))

        .arg(Arg::with_name("orientation")
            .short("O")
            .long("orientation")
            .takes_value(true)
            .value_name("OR")
            .display_order(4)
            .help("Read pair orientation produced by sequencer. Either 'fr' (converging), 'rf' (diverging) or 'ff' [default: fr]"))

        .arg(Arg::with_name("all-reads")
            .short("A")
            .long("all-reads")
            .display_order(5)
            .help("Use all reads for rearrangement detection, not just unaligned reads."))

        .arg(Arg::with_name("discard-duplicates")
            .long("discard-duplicates")
            .takes_value(true)
            .value_name("Method")
            .display_order(6)
            .help("Method to use when discarding duplicate reads.'both-ends' considers a read pair (or unaligned read) to be a duplicate of another if the positions"))

        .arg(Arg::with_name("blacklist")
            .long("blacklist")
            .takes_value(true)
            .value_name("list")
            .display_order(7)
            .help("Path to a file containing blacklisted regions."))

        .arg(Arg::with_name("min-reads")
            .short("r")
            .long("min-reads")
            .takes_value(true)
            .value_name("P-S-A")
            .display_order(8)
            .help("Minimum number of spanning reads required to accept a breakpoint. Specified in the format P-S-A, where P=paired, S=split, A=either. For example, -r 1-2-0 would require at least one mate pair and two split reads of evidence [default: 0-0-0]."))

        .arg(Arg::with_name("freq-above")
            .long("freq-above")
            .takes_value(true)
            .value_name("FREQ")
            .display_order(9)
            .help("Minimum frequency at which a variant must be present among the control samples to be considered a false positive [default: 0]."))

        .subcommand(SubCommand::with_name("detect")
            .usage("\tbreakfast detect <bam_file> <genome> <out_prefix> [-a N] [-f N] [-q N] [-O OR] [--discard-duplicates=METHOD]")
            .display_order(1)
            .arg(Arg::with_name("bam_file")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("genome")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("out_prefix")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("anchor-len")
                .short("a")
                .long("anchor-len")
                .takes_value(true)
                .value_name("N")
                .display_order(1)
                .help("Anchor length for split read analysis. When zero, split reads are not used [default: 0]"))
            .arg(Arg::with_name("max-frag-len")
                .short("f")
                .long("max-frag-len")
                .takes_value(true)
                .value_name("N")
                .display_order(2)
                .help("Maximum fragment length [default: 5000]"))
            .arg(Arg::with_name("min-mapq")
                .short("q")
                .long("min-mapq")
                .takes_value(true)
                .value_name("N")
                .display_order(3)
                .help("Minimum mapping quality to consider [default: 15]"))
            .arg(Arg::with_name("orientation")
                .short("O")
                .long("orientation")
                .takes_value(true)
                .value_name("OR")
                .display_order(4)
                .help("Read pair orientation produced by sequencer. Either 'fr' (converging), 'rf' (diverging) or 'ff' [default: fr]"))
            .arg(Arg::with_name("discard-duplicates")
                .long("discard-duplicates")
                .takes_value(true)
                .value_name("Method")
                .display_order(6)
                .help("Method to use when discarding duplicate reads.'both-ends' considers a read pair (or unaligned read) to be a duplicate of another if the positions")))

        .subcommand(SubCommand::with_name("detect-specific")
            .usage("breakfast detect specific [-A] <bam_file> <donors> <acceptors> <genome> <out_prefix>")
            .display_order(2)
            .arg(Arg::with_name("all-reads")
                .short("A")
                .long("all-reads")
                .display_order(5)
                .help("Use all reads for rearrangement detection, not just unaligned reads."))
            .arg(Arg::with_name("bam_file")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("donors")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("acceptors")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("genome")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("out_prefix")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("filter")
            .usage("breakfast filter <sv_file> [-r P-S-A]... [--blacklist=PATH]")
            .display_order(3)
            .arg(Arg::with_name("sv_file")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("min-reads")
                .short("r")
                .long("min-reads")
                .takes_value(true)
                .value_name("P-S-A")
                .display_order(8)
                .help("Minimum number of spanning reads required to accept a breakpoint. Specified in the format P-S-A, where P=paired, S=split, A=either. For example, -r 1-2-0 would require at least one mate pair and two split reads of evidence [default: 0-0-0]."))
            .arg(Arg::with_name("blacklist")
                .long("blacklist")
                .takes_value(true)
                .value_name("list")
                .display_order(7)
                .help("Path to a file containing blacklisted regions.")))

        .subcommand(SubCommand::with_name("annotate")
            .usage("breakfast annotate <sv_file> <bed_file>")
            .display_order(4)
            .arg(Arg::with_name("sv_file")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("bed_file")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("blacklist")
            .usage("breakfast blacklist [--freq-above=FREQ] <sv_files>...")
            .display_order(5)
            .arg(Arg::with_name("sv_files")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("freq-above")
                .long("freq-above")
                .takes_value(true)
                .value_name("FREQ")
                .display_order(9)
                .help("Minimum frequency at which a variant must be present among the control samples to be considered a false positive [default: 0].")))

        .subcommand(SubCommand::with_name("visualize")
            .usage("breakfast visualize <sv_file>")
            .display_order(6)
            .arg(Arg::with_name("sv_file")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("tabulate-rearranged-genes")
            .usage("breakfast tabulate rearranged genes <sv_files>...")
            .display_order(7)
            .arg(Arg::with_name("sv_files")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("tabulate-fusions")
            .usage("breakfast tabulate fusions <sv_files>...")
            .display_order(8)
            .arg(Arg::with_name("sv_files")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("statistics")
            .usage("breakfast statistics <sv_files>..")
            .display_order(9)
            .arg(Arg::with_name("sv_files")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("filter-by-region")
            .usage("breakfast filter by region <sv_file> <region>")
            .display_order(10)
            .arg(Arg::with_name("sv_file")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("region")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("filter-by-distance")
            .usage("breakfast filter by distance <min_distance> <sv_file>")
            .display_order(11)
            .arg(Arg::with_name("min_distance")
                .takes_value(true)
                .required(true))
            .arg(Arg::with_name("sv_file")
                .required(true)
                .takes_value(true)))

        .subcommand(SubCommand::with_name("align-junction")
            .usage("breakfast align junction <reads>")
            .display_order(12)
            .arg(Arg::with_name("reads")
                .takes_value(true)
                .required(true)))
}
<|endoftext|>


<commit_before>/// Module encapsules all command constants.
macro_rules! commands {
    {$(
        $ident:ident
        $bytes:expr
        #[$doc:meta];
    )*} => {
        /// Enumeration of all supported IRC commands (mainly RFC1459)
        #[deriving(Show)]
        pub enum Command<'a> {
            $(#[$doc] $ident,)*
            /// Numeric reply codes, see `ResponseCode`
            REPLY(ResponseCode),
            /// Catch all unknown/unsupported commands
            UNKNOWN(&'a[u8]),
        }

        impl<'a> Command<'a> {
            /// Converts bytestring to Command
            pub fn from_bytes(bytes: &'a [u8]) -> Command<'a> {
                // TODO add REPLY(...)
                $(if bytes == $bytes { $ident } else)* {
                    UNKNOWN(bytes)
                }
            }

            /// Converts the command into bytes
            pub fn to_bytes(&'a self) -> Vec<u8> {
                match *self {
                    $($ident => Vec::from_slice($bytes),)*
                    REPLY(reply) => format!("{:03u}", reply as u16).into_bytes(),
                    UNKNOWN(bytes) => Vec::from_slice(bytes)
                }
            }

            /// Converts the command into a string
            pub fn to_string(&'a self) -> String {
                String::from_utf8_lossy(self.to_bytes().as_slice()).into_string()
            }
        }
    }
}

commands!{
    PRIVMSG     b"PRIVMSG"  #[doc = "`PRIVMSG` command"];
    MODE        b"MODE"     #[doc = "`MODE` command"];
    JOIN        b"JOIN"     #[doc = "`JOIN` command, see http://tools.ietf.org/html/rfc1459.html#section-4.2.1"];
    PING        b"PING"     #[doc = "`PING` command"];
    NAMES       b"NAMES"    #[doc = "`NAMES` command"];
    QUIT        b"QUIT"     #[doc = "`QUIT` command"];
    PONG        b"PONG"     #[doc = "`PONG` command"];
    NICK        b"NICK"     #[doc = "`NICK` command"];
    USER        b"USER"     #[doc = "`USER` command"];
}

#[allow(non_camel_case_types)]
#[deriving(Show)]
pub enum ResponseCode {
    RPL_WELCOME = 001,
    RPL_YOURHOST = 002,
    RPL_CREATED = 003,
    RPL_MYINFO = 004,
    RPL_BOUNCE = 005,
    RPL_USERHOST = 302,
    RPL_ISON = 303,
    RPL_AWAY = 301,
    RPL_UNAWAY = 305,
    RPL_NOWAWAY = 306,
    RPL_WHOISUSER = 311,
    RPL_WHOISSERVER = 312,
    RPL_WHOISOPERATOR = 313,
    RPL_WHOISIDLE = 317,
    RPL_ENDOFWHOIS = 318,
    RPL_WHOISCHANNELS = 319,
    RPL_WHOWASUSER = 314,
    RPL_ENDOFWHOWAS = 369,
    RPL_LISTSTART = 321,
    RPL_LIST = 322,
    RPL_LISTEND = 323,
    RPL_UNIQOPIS = 325,
    RPL_CHANNELMODEIS = 324,
    RPL_NOTOPIC = 331,
    RPL_TOPIC = 332,
    RPL_INVITING = 341,
    RPL_SUMMONING = 342,
    RPL_INVITELIST = 346,
    RPL_ENDOFINVITELIST = 347,
    RPL_EXCEPTLIST = 348,
    RPL_ENDOFEXCEPTLIST = 349,
    RPL_VERSION = 351,
    RPL_WHOREPLY = 352,
    RPL_ENDOFWHO = 315,
    RPL_NAMREPLY = 353,
    RPL_ENDOFNAMES = 366,
    RPL_LINKS = 364,
    RPL_ENDOFLINKS = 365,
    RPL_BANLIST = 367,
    RPL_ENDOFBANLIST = 368,
    RPL_INFO = 371,
    RPL_ENDOFINFO = 374,
    RPL_MOTDSTART = 375,
    RPL_MOTD = 372,
    RPL_ENDOFMOTD = 376,
    RPL_YOUREOPER = 381,
    RPL_REHASHING = 382,
    RPL_YOURESERVICE = 383,
    RPL_TIME = 391,
    RPL_USERSSTART = 392,
    RPL_USERS = 393,
    RPL_ENDOFUSERS = 394,
    RPL_NOUSERS = 395,
    RPL_TRACELINK = 200,
    RPL_TRACECONNECTING = 201,
    RPL_TRACEHANDSHAKE = 202,
    RPL_TRACEUNKNOWN = 203,
    RPL_TRACEOPERATOR = 204,
    RPL_TRACEUSER = 205,
    RPL_TRACESERVER = 206,
    RPL_TRACESERVICE = 207,
    RPL_TRACENEWTYPE = 208,
    RPL_TRACECLASS = 209,
    RPL_TRACERECONNECT = 210,
    RPL_TRACELOG = 261,
    RPL_TRACEEND = 262,
    RPL_STATSLINKINFO = 211,
    RPL_STATSCOMMANDS = 212,
    RPL_ENDOFSTATS = 219,
    RPL_STATSUPTIME = 242,
    RPL_STATSOLINE = 243,
    RPL_UMODEIS = 221,
    RPL_SERVLIST = 234,
    RPL_SERVLISTEND = 235,
    RPL_LUSERCLIENT = 251,
    RPL_LUSEROP = 252,
    RPL_LUSERUNKNOWN = 253,
    RPL_LUSERCHANNELS = 254,
    RPL_LUSERME = 255,
    RPL_ADMINME = 256,
    RPL_ADMINLOC1 = 257,
    RPL_ADMINLOC2 = 258,
    RPL_ADMINEMAIL = 259,
    RPL_TRYAGAIN = 263,
    ERR_NOSUCHNICK = 401,
    ERR_NOSUCHSERVER = 402,
    ERR_NOSUCHCHANNEL = 403,
    ERR_CANNOTSENDTOCHAN = 404,
    ERR_TOOMANYCHANNELS = 405,
    ERR_WASNOSUCHNICK = 406,
    ERR_TOOMANYTARGETS = 407,
    ERR_NOSUCHSERVICE = 408,
    ERR_NOORIGIN = 409,
    ERR_NORECIPIENT = 411,
    ERR_NOTEXTTOSEND = 412,
    ERR_NOTOPLEVEL = 413,
    ERR_WILDTOPLEVEL = 414,
    ERR_BADMASK = 415,
    ERR_UNKNOWNCOMMAND = 421,
    ERR_NOMOTD = 422,
    ERR_NOADMININFO = 423,
    ERR_FILEERROR = 424,
    ERR_NONICKNAMEGIVEN = 431,
    ERR_ERRONEUSNICKNAME = 432,
    ERR_NICKNAMEINUSE = 433,
    ERR_NICKCOLLISION = 436,
    ERR_UNAVAILRESOURCE = 437,
    ERR_USERNOTINCHANNEL = 441,
    ERR_NOTONCHANNEL = 442,
    ERR_USERONCHANNEL = 443,
    ERR_NOLOGIN = 444,
    ERR_SUMMONDISABLED = 445,
    ERR_USERSDISABLED = 446,
    ERR_NOTREGISTERED = 451,
    ERR_NEEDMOREPARAMS = 461,
    ERR_ALREADYREGISTRED = 462,
    ERR_NOPERMFORHOST = 463,
    ERR_PASSWDMISMATCH = 464,
    ERR_YOUREBANNEDCREEP = 465,
    ERR_YOUWILLBEBANNED = 466,
    ERR_KEYSET = 467,
    ERR_CHANNELISFULL = 471,
    ERR_UNKNOWNMODE = 472,
    ERR_INVITEONLYCHAN = 473,
    ERR_BANNEDFROMCHAN = 474,
    ERR_BADCHANNELKEY = 475,
    ERR_BADCHANMASK = 476,
    ERR_NOCHANMODES = 477,
    ERR_BANLISTFULL = 478,
    ERR_NOPRIVILEGES = 481,
    ERR_CHANOPRIVSNEEDED = 482,
    ERR_CANTKILLSERVER = 483,
    ERR_RESTRICTED = 484,
    ERR_UNIQOPPRIVSNEEDED = 485,
    ERR_NOOPERHOST = 491,
    ERR_UMODEUNKNOWNFLAG = 501,
    ERR_USERSDONTMATCH = 502,
}<commit_msg>new macro<commit_after>/// Module encapsules all command constants.
macro_rules! commands {
    {$(
        $ident:ident
        #[$doc:meta];
    )*} => {
        /// Enumeration of all supported IRC commands (mainly RFC1459)
        #[deriving(Show)]
        pub enum Command<'a> {
            $(#[$doc] $ident,)*
            /// Numeric reply codes, see `ResponseCode`
            REPLY(ResponseCode),
            /// Catch all unknown/unsupported commands
            UNKNOWN(&'a[u8]),
        }

        impl<'a> Command<'a> {
            /// Converts bytestring to Command
            pub fn from_bytes(bytes: &'a [u8]) -> Command<'a> {
                // TODO add REPLY(...)
                $(if bytes == stringify!($ident).as_bytes() { $ident } else)* {
                    UNKNOWN(bytes)
                }
            }

            /// Converts the command into bytes
            pub fn to_bytes(&'a self) -> Vec<u8> {
                match *self {
                    $($ident => Vec::from_slice(stringify!($ident).as_bytes()),)*
                    REPLY(reply) => format!("{:03u}", reply as u16).into_bytes(),
                    UNKNOWN(bytes) => Vec::from_slice(bytes)
                }
            }

            /// Converts the command into a string
            pub fn to_string(&'a self) -> String {
                String::from_utf8_lossy(self.to_bytes().as_slice()).into_string()
            }
        }
    }
}

commands!{
    PRIVMSG     #[doc = "`PRIVMSG` command"];
    MODE        #[doc = "`MODE` command"];
    JOIN        #[doc = "`JOIN` command, see http://tools.ietf.org/html/rfc1459.html#section-4.2.1"];
    PING        #[doc = "`PING` command"];
    NAMES       #[doc = "`NAMES` command"];
    QUIT        #[doc = "`QUIT` command"];
    PONG        #[doc = "`PONG` command"];
    NICK        #[doc = "`NICK` command"];
    USER        #[doc = "`USER` command"];
}

macro_rules! messages {
    {$(
        $keyword:ident as
        $name:ident
        { $($attr:ident: $ty:ty),+ }
        <- $parser:expr;
    )*} => {
    }
}

messages!{
    JOIN as Join { targets: Vec<String>, passwords: Vec<Vec<u8>> } <- {
        PRIVMSG
    };
}

#[allow(non_camel_case_types)]
#[deriving(Show)]
pub enum ResponseCode {
    RPL_WELCOME = 001,
    RPL_YOURHOST = 002,
    RPL_CREATED = 003,
    RPL_MYINFO = 004,
    RPL_BOUNCE = 005,
    RPL_USERHOST = 302,
    RPL_ISON = 303,
    RPL_AWAY = 301,
    RPL_UNAWAY = 305,
    RPL_NOWAWAY = 306,
    RPL_WHOISUSER = 311,
    RPL_WHOISSERVER = 312,
    RPL_WHOISOPERATOR = 313,
    RPL_WHOISIDLE = 317,
    RPL_ENDOFWHOIS = 318,
    RPL_WHOISCHANNELS = 319,
    RPL_WHOWASUSER = 314,
    RPL_ENDOFWHOWAS = 369,
    RPL_LISTSTART = 321,
    RPL_LIST = 322,
    RPL_LISTEND = 323,
    RPL_UNIQOPIS = 325,
    RPL_CHANNELMODEIS = 324,
    RPL_NOTOPIC = 331,
    RPL_TOPIC = 332,
    RPL_INVITING = 341,
    RPL_SUMMONING = 342,
    RPL_INVITELIST = 346,
    RPL_ENDOFINVITELIST = 347,
    RPL_EXCEPTLIST = 348,
    RPL_ENDOFEXCEPTLIST = 349,
    RPL_VERSION = 351,
    RPL_WHOREPLY = 352,
    RPL_ENDOFWHO = 315,
    RPL_NAMREPLY = 353,
    RPL_ENDOFNAMES = 366,
    RPL_LINKS = 364,
    RPL_ENDOFLINKS = 365,
    RPL_BANLIST = 367,
    RPL_ENDOFBANLIST = 368,
    RPL_INFO = 371,
    RPL_ENDOFINFO = 374,
    RPL_MOTDSTART = 375,
    RPL_MOTD = 372,
    RPL_ENDOFMOTD = 376,
    RPL_YOUREOPER = 381,
    RPL_REHASHING = 382,
    RPL_YOURESERVICE = 383,
    RPL_TIME = 391,
    RPL_USERSSTART = 392,
    RPL_USERS = 393,
    RPL_ENDOFUSERS = 394,
    RPL_NOUSERS = 395,
    RPL_TRACELINK = 200,
    RPL_TRACECONNECTING = 201,
    RPL_TRACEHANDSHAKE = 202,
    RPL_TRACEUNKNOWN = 203,
    RPL_TRACEOPERATOR = 204,
    RPL_TRACEUSER = 205,
    RPL_TRACESERVER = 206,
    RPL_TRACESERVICE = 207,
    RPL_TRACENEWTYPE = 208,
    RPL_TRACECLASS = 209,
    RPL_TRACERECONNECT = 210,
    RPL_TRACELOG = 261,
    RPL_TRACEEND = 262,
    RPL_STATSLINKINFO = 211,
    RPL_STATSCOMMANDS = 212,
    RPL_ENDOFSTATS = 219,
    RPL_STATSUPTIME = 242,
    RPL_STATSOLINE = 243,
    RPL_UMODEIS = 221,
    RPL_SERVLIST = 234,
    RPL_SERVLISTEND = 235,
    RPL_LUSERCLIENT = 251,
    RPL_LUSEROP = 252,
    RPL_LUSERUNKNOWN = 253,
    RPL_LUSERCHANNELS = 254,
    RPL_LUSERME = 255,
    RPL_ADMINME = 256,
    RPL_ADMINLOC1 = 257,
    RPL_ADMINLOC2 = 258,
    RPL_ADMINEMAIL = 259,
    RPL_TRYAGAIN = 263,
    ERR_NOSUCHNICK = 401,
    ERR_NOSUCHSERVER = 402,
    ERR_NOSUCHCHANNEL = 403,
    ERR_CANNOTSENDTOCHAN = 404,
    ERR_TOOMANYCHANNELS = 405,
    ERR_WASNOSUCHNICK = 406,
    ERR_TOOMANYTARGETS = 407,
    ERR_NOSUCHSERVICE = 408,
    ERR_NOORIGIN = 409,
    ERR_NORECIPIENT = 411,
    ERR_NOTEXTTOSEND = 412,
    ERR_NOTOPLEVEL = 413,
    ERR_WILDTOPLEVEL = 414,
    ERR_BADMASK = 415,
    ERR_UNKNOWNCOMMAND = 421,
    ERR_NOMOTD = 422,
    ERR_NOADMININFO = 423,
    ERR_FILEERROR = 424,
    ERR_NONICKNAMEGIVEN = 431,
    ERR_ERRONEUSNICKNAME = 432,
    ERR_NICKNAMEINUSE = 433,
    ERR_NICKCOLLISION = 436,
    ERR_UNAVAILRESOURCE = 437,
    ERR_USERNOTINCHANNEL = 441,
    ERR_NOTONCHANNEL = 442,
    ERR_USERONCHANNEL = 443,
    ERR_NOLOGIN = 444,
    ERR_SUMMONDISABLED = 445,
    ERR_USERSDISABLED = 446,
    ERR_NOTREGISTERED = 451,
    ERR_NEEDMOREPARAMS = 461,
    ERR_ALREADYREGISTRED = 462,
    ERR_NOPERMFORHOST = 463,
    ERR_PASSWDMISMATCH = 464,
    ERR_YOUREBANNEDCREEP = 465,
    ERR_YOUWILLBEBANNED = 466,
    ERR_KEYSET = 467,
    ERR_CHANNELISFULL = 471,
    ERR_UNKNOWNMODE = 472,
    ERR_INVITEONLYCHAN = 473,
    ERR_BANNEDFROMCHAN = 474,
    ERR_BADCHANNELKEY = 475,
    ERR_BADCHANMASK = 476,
    ERR_NOCHANMODES = 477,
    ERR_BANLISTFULL = 478,
    ERR_NOPRIVILEGES = 481,
    ERR_CHANOPRIVSNEEDED = 482,
    ERR_CANTKILLSERVER = 483,
    ERR_RESTRICTED = 484,
    ERR_UNIQOPPRIVSNEEDED = 485,
    ERR_NOOPERHOST = 491,
    ERR_UMODEUNKNOWNFLAG = 501,
    ERR_USERSDONTMATCH = 502,
}<|endoftext|>


<commit_before>use num::integer::Integer;
use memory::Memory;
use memory::low_nibble;
use memory::high_nibble;
use extensions::Incrementor;
use ops::{mod};
use std::num::One;

pub struct Cpu {
    reg: Registers,
    mem: Memory
}

impl Cpu {
    pub fn new(memory: Memory) -> Cpu {
        return Cpu{reg: Registers::new(), mem: memory}
    }

    /// Execute a cycle on the cpu
    pub fn tick(&mut self) {
        let instr = self.fetch_instruction();
        self.reg.pc.increment();

        match high_nibble(instr) {
            0x0 => self.zero(instr),
            _ => return
        }
    }

    /// Fetches the instruction pointed to by the program counter
    /// and increments the pc by 1
    fn fetch_instruction(&mut self) -> u8 {
        let instr = self.mem.read_byte(self.reg.pc.val);
        self.reg.pc.increment();
        return instr;
    }

    fn zero(&mut self, op_code: u8) {
        match low_nibble(op_code) {
            0x0 => ops::nop(),
            0x6 => ops::ld_next_byte_to_reg(self.mem, &mut self.reg.pc, &mut self.reg.b),
            _ => return
        }
    }
}

struct Registers {
    a: Register<u8>, b: Register<u8>, c: Register<u8>,
    d: Register<u8>, e: Register<u8>, f: Register<u8>,
    h: Register<u8>, l: Register<u8>, // 8-bit registers

    pc: Register<u16>, sp: Register<u16>, // 16-bit registers
    m: Register<u16>, t: Register<u16> // clock
}

impl Registers {
    fn new() -> Registers {
        return Registers{
            a: Register::new(0),
            b: Register::new(0),
            c: Register::new(0),
            d: Register::new(0),
            e: Register::new(0),
            f: Register::new(0),
            h: Register::new(0),
            l: Register::new(0),

            pc: Register::new(0),
            sp: Register::new(0),

            m: Register::new(0),
            t: Register::new(0)
        }
    }
}

pub struct Register<T: Copy + Unsigned> {
    val: T
}

impl<T: Copy + Unsigned> Register<T> {
    pub fn new(i: T) -> Register<T> {
        return Register { val: i };
    }

    pub fn read(&self) -> T {
        return self.val;
    }

    pub fn write(&mut self, i: T) {
        self.val = i;
    }

    pub fn increment(&mut self) {
        let i = self.val;
        self.write(i + One::one());
    }
}
<commit_msg>Add tests for Register<commit_after>use num::integer::Integer;
use memory::Memory;
use memory::low_nibble;
use memory::high_nibble;
use extensions::Incrementor;
use ops::{mod};
use std::num::One;

pub struct Cpu {
    reg: Registers,
    mem: Memory
}

impl Cpu {
    pub fn new(memory: Memory) -> Cpu {
        return Cpu{reg: Registers::new(), mem: memory}
    }

    /// Execute a cycle on the cpu
    pub fn tick(&mut self) {
        let instr = self.fetch_instruction();
        self.reg.pc.increment();

        match high_nibble(instr) {
            0x0 => self.zero(instr),
            _ => return
        }
    }

    /// Fetches the instruction pointed to by the program counter
    /// and increments the pc by 1
    fn fetch_instruction(&mut self) -> u8 {
        let instr = self.mem.read_byte(self.reg.pc.val);
        self.reg.pc.increment();
        return instr;
    }

    fn zero(&mut self, op_code: u8) {
        match low_nibble(op_code) {
            0x0 => ops::nop(),
            0x6 => ops::ld_next_byte_to_reg(self.mem, &mut self.reg.pc, &mut self.reg.b),
            _ => return
        }
    }
}

struct Registers {
    a: Register<u8>, b: Register<u8>, c: Register<u8>,
    d: Register<u8>, e: Register<u8>, f: Register<u8>,
    h: Register<u8>, l: Register<u8>, // 8-bit registers

    pc: Register<u16>, sp: Register<u16>, // 16-bit registers
    m: Register<u16>, t: Register<u16> // clock
}

impl Registers {
    fn new() -> Registers {
        return Registers{
            a: Register::new(0),
            b: Register::new(0),
            c: Register::new(0),
            d: Register::new(0),
            e: Register::new(0),
            f: Register::new(0),
            h: Register::new(0),
            l: Register::new(0),

            pc: Register::new(0),
            sp: Register::new(0xFFFE),

            m: Register::new(0),
            t: Register::new(0)
        }
    }
}

pub struct Register<T: Copy + Unsigned> {
    val: T
}

impl<T: Copy + Unsigned> Register<T> {
    pub fn new(i: T) -> Register<T> {
        return Register { val: i };
    }

    pub fn read(&self) -> T {
        return self.val;
    }

    pub fn write(&mut self, i: T) {
        self.val = i;
    }

    pub fn increment(&mut self) {
        let i = self.val;
        self.write(i + One::one());
    }
}

#[test]
fn test_Register_new() {
    let reg: Register<u8> = Register::new(10);
    assert!(reg.val == 10);
}

fn test_Register_read() {
    let reg: Register<u32> = Register::new(5);

    assert!(reg.read() == 5);
}

fn test_Register_write() {
    let mut reg: Register<u16> = Register::new(483);

    reg.write(5);
    assert!(reg.val == 5);
}
<|endoftext|>


<commit_before>extern crate bit_vec;

use self::bit_vec::BitVec;
use std::mem;
use std::fmt;

use nfa::{AUTO_START, AUTO_STUCK};
use automaton::{Automaton, Match};

pub type Input = u8;
pub type StateNumber = usize;
pub type PatternNumber = usize;

#[derive(Default)]
pub struct DFAState {
    transitions: Box<[StateNumber]>,
    pattern_ends: Box<[PatternNumber]>,
}

#[derive(Default)]
pub struct DFA {
    states: Box<[DFAState]>,
    finals: BitVec,
    dict: Vec<Vec<Input>>,
}

#[derive(Default)]
pub struct DDFA {
    states: Box<[DDFAState]>,
    dict: Vec<Vec<Input>>,
}

// Living dangerously: raw pointers baby
#[derive(Clone, PartialEq)]
struct DDFAState {
    transitions: Box<[*const DDFAState]>,
    is_final: bool,
}

impl DFAState {
    pub fn new(transitions: Box<[StateNumber]>, pattern_ends: Box<[PatternNumber]>) -> Self {
        DFAState {
            transitions: transitions,
            pattern_ends: pattern_ends,
        }
    }
}

impl DFA {
    pub fn new(states: Box<[DFAState]>, finals: BitVec, dict: Vec<Vec<Input>>) -> Self {
        DFA {
            states: states,
            finals: finals,
            dict: dict,
        }
    }

    pub fn into_ddfa(self) -> Result<DDFA, ()> {
        let len = self.states.len();
        let mut states = vec![DDFAState { transitions: Box::new([]), is_final: false }; len]
            .into_boxed_slice();

        let states_start: *mut DDFAState = (*states).as_mut_ptr();

        for (i, ref st) in self.states.iter().enumerate() {
            let mut v: Vec<*const DDFAState> = Vec::with_capacity(st.transitions.len());
            for &offset in st.transitions.iter() {
                if offset >= len {
                    return Err(());
                }
                unsafe {
                    v.push(states_start.offset(offset as isize) as *const DDFAState);
                }
            }
            states[i].transitions = v.into_boxed_slice();
            states[i].is_final = self.finals[i];
        }
        Ok(DDFA {
            states: states,
            dict: self.dict,
        })
    }

    pub fn apply(&self, input: &[u8]) -> bool {
        let mut cur_state = AUTO_START;
        for &byte in input {
            cur_state = self.states[cur_state].transitions[byte as usize];
            if cur_state == AUTO_STUCK {
                break;
            }
        }
        self.finals[cur_state]
    }
}

impl Automaton<Input> for DFA {
    type State = StateNumber;

    fn start_state() -> Self::State {
        AUTO_START
    }

    fn stuck_state() -> Self::State {
        AUTO_STUCK
    }

    #[inline]
    fn next_state(&self, &state: &Self::State, &input: &Input) -> Self::State {
        self.states[state].transitions[input as usize]
    }

    #[inline]
    fn has_match(&self, &state: &Self::State, patt_no_offset: usize) -> bool {
        patt_no_offset < self.states[state].pattern_ends.len()
    }

    #[inline]
    fn get_match(&self, &state: &Self::State, patt_no_offset: usize, text_offset: usize) -> Match {
        let patt_no = self.states[state].pattern_ends[patt_no_offset];
        Match {
            patt_no: patt_no,
            start: text_offset - self.dict[patt_no].len(),
            end: text_offset,
        }
    }
}

impl fmt::Debug for DFA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        macro_rules! w {
            ($($tt:tt)*) => { try!(write!(f, $($tt)*)) }
        }
        for (i, state) in (*self.states).into_iter().enumerate() {
            if i == AUTO_STUCK {
                w!("{} (stuck),\n", AUTO_STUCK);
                continue;
            }
            w!("{}", i);
            if i == AUTO_START {
                w!(" (start)");
            }
            if self.finals[i] {
                w!(" (final)");
            }
            w!(": {{");
            if !state.transitions.is_empty() {
                w!("\n");
            }
            let mut last_c = 0;
            let mut iter = (*state.transitions)
                .into_iter()
                .enumerate()
                .peekable();
            while let Some((c, tr)) = iter.next() {
                if let Some(&(c2, tr2)) = iter.peek() {
                    if tr == tr2 {
                        continue;
                    }
                    if c == last_c {
                        w!("  {:?}: {:?},\n", c as u8 as char, tr);
                    } else {
                        w!("  [{:?}-{:?}]: {:?},\n",
                           last_c as u8 as char,
                           (c as u8) as char,
                           tr);
                    }
                    last_c = c2;
                }
            }
            w!("}},\n");
        }
        Ok(())
    }
}

impl DDFA {
    pub fn apply(&self, input: &[u8]) -> bool {
        let mut cur_state: *const DDFAState = &self.states[AUTO_START];
        let stuck = &self.states[AUTO_STUCK];
        for &byte in input {
            cur_state = unsafe { (*cur_state).transitions[byte as usize] };
            if cur_state == stuck {
                break;
            }
        }
        unsafe { (*cur_state).is_final }
    }
}

impl fmt::Debug for DDFA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        macro_rules! w {
            ($($tt:tt)*) => { try!(write!(f, $($tt)*)) }
        }
        let start = &self.states[0] as *const DDFAState;
        for (i, state) in (*self.states).into_iter().enumerate() {
            if i == AUTO_STUCK {
                w!("{} (stuck),\n", AUTO_STUCK);
                continue;
            }
            w!("{}", i);
            if i == AUTO_START {
                w!(" (start)");
            }
            if state.is_final {
                w!(" (final)");
            }
            w!(": {{");
            if !state.transitions.is_empty() {
                w!("\n");
            }
            let mut last_c = 0;
            let mut iter = (*state.transitions)
                .into_iter()
                .enumerate()
                .peekable();
            while let Some((c, tr)) = iter.next() {
                if let Some(&(c2, tr2)) = iter.peek() {
                    if tr == tr2 {
                        continue;
                    }
                    let tr_no = (*tr as usize - start as usize) / mem::size_of::<DDFAState>();
                    if c == last_c {
                        w!("  {:?}: {:?},\n", c as u8 as char, tr_no);
                    } else {
                        w!("  [{:?}-{:?}]: {:?},\n",
                           last_c as u8 as char,
                           (c as u8) as char,
                           tr_no);
                    }
                    last_c = c2;
                }
            }
            w!("}},\n");
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {}
<commit_msg>Sharing Debug::fmt implementation between DFA and DDFA through a macro<commit_after>extern crate bit_vec;

use self::bit_vec::BitVec;
use std::mem;
use std::fmt;

use nfa::{AUTO_START, AUTO_STUCK};
use automaton::{Automaton, Match};

pub type Input = u8;
pub type StateNumber = usize;
pub type PatternNumber = usize;

#[derive(Default)]
pub struct DFAState {
    transitions: Box<[StateNumber]>,
    pattern_ends: Box<[PatternNumber]>,
}

#[derive(Default)]
pub struct DFA {
    states: Box<[DFAState]>,
    finals: BitVec,
    dict: Vec<Vec<Input>>,
}

#[derive(Default)]
pub struct DDFA {
    states: Box<[DDFAState]>,
    dict: Vec<Vec<Input>>,
}

// Living dangerously: raw pointers baby
#[derive(Clone, PartialEq)]
struct DDFAState {
    transitions: Box<[*const DDFAState]>,
    is_final: bool,
}

impl DFAState {
    pub fn new(transitions: Box<[StateNumber]>, pattern_ends: Box<[PatternNumber]>) -> Self {
        DFAState {
            transitions: transitions,
            pattern_ends: pattern_ends,
        }
    }
}

impl DFA {
    pub fn new(states: Box<[DFAState]>, finals: BitVec, dict: Vec<Vec<Input>>) -> Self {
        DFA {
            states: states,
            finals: finals,
            dict: dict,
        }
    }

    pub fn into_ddfa(self) -> Result<DDFA, ()> {
        let len = self.states.len();
        let mut states = vec![DDFAState { transitions: Box::new([]), is_final: false }; len]
            .into_boxed_slice();

        let states_start: *mut DDFAState = (*states).as_mut_ptr();

        for (i, ref st) in self.states.iter().enumerate() {
            let mut v: Vec<*const DDFAState> = Vec::with_capacity(st.transitions.len());
            for &offset in st.transitions.iter() {
                if offset >= len {
                    return Err(());
                }
                unsafe {
                    v.push(states_start.offset(offset as isize) as *const DDFAState);
                }
            }
            states[i].transitions = v.into_boxed_slice();
            states[i].is_final = self.finals[i];
        }
        Ok(DDFA {
            states: states,
            dict: self.dict,
        })
    }

    pub fn apply(&self, input: &[u8]) -> bool {
        let mut cur_state = AUTO_START;
        for &byte in input {
            cur_state = self.states[cur_state].transitions[byte as usize];
            if cur_state == AUTO_STUCK {
                break;
            }
        }
        self.finals[cur_state]
    }
}

impl Automaton<Input> for DFA {
    type State = StateNumber;

    fn start_state() -> Self::State {
        AUTO_START
    }

    fn stuck_state() -> Self::State {
        AUTO_STUCK
    }

    #[inline]
    fn next_state(&self, &state: &Self::State, &input: &Input) -> Self::State {
        self.states[state].transitions[input as usize]
    }

    #[inline]
    fn has_match(&self, &state: &Self::State, patt_no_offset: usize) -> bool {
        patt_no_offset < self.states[state].pattern_ends.len()
    }

    #[inline]
    fn get_match(&self, &state: &Self::State, patt_no_offset: usize, text_offset: usize) -> Match {
        let patt_no = self.states[state].pattern_ends[patt_no_offset];
        Match {
            patt_no: patt_no,
            start: text_offset - self.dict[patt_no].len(),
            end: text_offset,
        }
    }
}

impl DDFA {
    pub fn apply(&self, input: &[u8]) -> bool {
        let mut cur_state: *const DDFAState = &self.states[AUTO_START];
        let stuck = &self.states[AUTO_STUCK];
        for &byte in input {
            cur_state = unsafe { (*cur_state).transitions[byte as usize] };
            if cur_state == stuck {
                break;
            }
        }
        unsafe { (*cur_state).is_final }
    }
}

// The Debug::fmt implementation for DFA and DDFA are extremely similar. The only differences are in
//  computing the finality of a state and computing the index of a state in the states array.
// Therefore we share these with a macro:
macro_rules! debug_impl {
    ($struct_name:ident, $compute_finality:item, $compute_start:item, $compute_tr_no:item) => {
        impl fmt::Debug for $struct_name {
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                $compute_finality
                $compute_start
                $compute_tr_no
                let start = compute_start(&self);
                for (i, state) in (*self.states).into_iter().enumerate() {
                    if i == AUTO_STUCK {
                        try!(write!(f, "{} (stuck),\n", AUTO_STUCK));
                        continue;
                    }
                    try!(write!(f, "{}", i));
                    if i == AUTO_START {
                        try!(write!(f, " (start)"));
                    }
                    if compute_finality(&self, state, i) {
                        try!(write!(f, " (final)"));
                    }
                    try!(write!(f, ": {{"));
                    if !state.transitions.is_empty() {
                        try!(write!(f, "\n"));
                    }
                    let mut last_c = 0;
                    let mut iter = (*state.transitions)
                        .into_iter()
                        .enumerate()
                        .peekable();
                    while let Some((c, tr)) = iter.next() {
                        if let Some(&(c2, tr2)) = iter.peek() {
                            if tr == tr2 {
                                continue;
                            }
                            let tr_no = compute_tr_no(tr, start);
                            if c == last_c {
                                try!(write!(f, "  {:?}: {:?},\n", c as u8 as char, tr_no));
                            } else {
                                try!(write!(f, "  [{:?}-{:?}]: {:?},\n",
                                   last_c as u8 as char,
                                   (c as u8) as char,
                                   tr_no));
                            }
                            last_c = c2;
                        }
                    }
                    try!(write!(f, "}},\n"));
                }
                Ok(())
            }
        }
    }
}

debug_impl!(
    DFA,
    #[allow(unused_variables)]
    fn compute_finality(dfa: &DFA, state: &DFAState, i: usize) -> bool {
        dfa.finals[i]
    },
    #[allow(unused_variables)]
    fn compute_start(dfa: &DFA) -> () {
        ()
    },
    #[allow(unused_variables)]
    fn compute_tr_no(tr: &StateNumber, start: ()) -> &StateNumber {
        tr
    }
);

debug_impl!(
    DDFA,
    #[allow(unused_variables)]
    fn compute_finality(ddfa: &DDFA, state: &DDFAState, i: usize) -> bool {
        state.is_final
    },
    fn compute_start(ddfa: &DDFA) -> *const DDFAState {
        &ddfa.states[0] as *const DDFAState
    },
    fn compute_tr_no(tr: &*const DDFAState, start: *const DDFAState) -> usize {
        (*tr as usize - start as usize) / mem::size_of::<DDFAState>()
    }
);

#[cfg(test)]
mod tests {}
<|endoftext|>


<commit_before>use yup_oauth2;

use futures::prelude::*;
use yup_oauth2::GetToken;

use hyper::client::Client;
use hyper_tls::HttpsConnector;
use tokio;

use std::path;

fn main() {
    let creds =
        yup_oauth2::service_account_key_from_file(path::Path::new("serviceaccount.json")).unwrap();
    let https = HttpsConnector::new(1).expect("tls");
    let client = Client::builder()
        .keep_alive(false)
        .build::<_, hyper::Body>(https);

    let mut sa = yup_oauth2::ServiceAccountAccess::new(creds, client);

    let fut = sa
        .token(["https://www.googleapis.com/auth/pubsub"].iter())
        .and_then(|tok| {
            println!("token is: {:?}", tok);
            Ok(())
        });
    let fut2 = sa
        .token(["https://www.googleapis.com/auth/pubsub"].iter())
        .and_then(|tok| {
            println!("cached token is {:?} and should be identical", tok);
            Ok(())
        });
    let all = fut.join(fut2).then(|_| Ok(()));
    tokio::run(all)
}
<commit_msg>fix(test-svc-acct): Non-lazy cache lookup resulted in two requests.<commit_after>use yup_oauth2;

use futures::prelude::*;
use yup_oauth2::GetToken;

use hyper::client::Client;
use hyper_tls::HttpsConnector;
use tokio;

use std::path;

fn main() {
    let creds =
        yup_oauth2::service_account_key_from_file(path::Path::new("serviceaccount.json")).unwrap();
    let https = HttpsConnector::new(1).expect("tls");
    let client = Client::builder()
        .keep_alive(false)
        .build::<_, hyper::Body>(https);

    let mut sa = yup_oauth2::ServiceAccountAccess::new(creds, client);

    let fut = sa
        .token(["https://www.googleapis.com/auth/pubsub"].iter())
        .and_then(|tok| {
            println!("token is: {:?}", tok);
            Ok(())
        });
    let mut sa2 = sa.clone();
    let all = fut
        .then(move |_| {
            sa2.token(["https://www.googleapis.com/auth/pubsub"].iter())
                .and_then(|tok| {
                    println!("cached token is {:?} and should be identical", tok);
                    Ok(())
                })
        })
        .then(|_| Ok(()));
    tokio::run(all)
}
<|endoftext|>


<commit_before>use crate::chain::Chain;
use crate::error::ErrorImpl;
use core::fmt::{self, Debug};

impl ErrorImpl<()> {
    pub(crate) fn display(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.error())?;

        if f.alternate() {
            for cause in self.chain().skip(1) {
                write!(f, ": {}", cause)?;
            }
        }

        Ok(())
    }

    pub(crate) fn debug(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let error = self.error();

        if f.alternate() {
            return Debug::fmt(error, f);
        }

        write!(f, "{}", error)?;

        if let Some(cause) = error.source() {
            write!(f, "\n\nCaused by:")?;
            let multiple = cause.source().is_some();
            for (n, error) in Chain::new(cause).enumerate() {
                write!(f, "\n    ")?;
                if multiple {
                    write!(f, "{}: ", n)?;
                }
                write!(f, "{}", error)?;
            }
        }

        #[cfg(backtrace)]
        {
            use std::backtrace::BacktraceStatus;

            let backtrace = self.backtrace();
            if let BacktraceStatus::Captured = backtrace.status() {
                let mut backtrace = backtrace.to_string();
                if backtrace.starts_with("stack backtrace:") {
                    // Capitalize to match "Caused by:"
                    backtrace.replace_range(0..1, "S");
                }
                backtrace.truncate(backtrace.trim_end().len());
                write!(f, "\n\n{}", backtrace)?;
            }
        }

        Ok(())
    }
}
<commit_msg>Apply indent in causes on each new line (#54)<commit_after>use crate::chain::Chain;
use crate::error::ErrorImpl;
use core::fmt::{self, Debug, Write};

impl ErrorImpl<()> {
    pub(crate) fn display(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.error())?;

        if f.alternate() {
            for cause in self.chain().skip(1) {
                write!(f, ": {}", cause)?;
            }
        }

        Ok(())
    }

    pub(crate) fn debug(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let error = self.error();

        if f.alternate() {
            return Debug::fmt(error, f);
        }

        write!(f, "{}", error)?;

        if let Some(cause) = error.source() {
            write!(f, "\n\nCaused by:")?;
            let multiple = cause.source().is_some();
            for (n, error) in Chain::new(cause).enumerate() {
                writeln!(f)?;
                write!(
                    if multiple {
                        Indented::numbered(f, n)
                    } else {
                        Indented::new(f)
                    },
                    "{}",
                    error
                )?;
            }
        }

        #[cfg(backtrace)]
        {
            use std::backtrace::BacktraceStatus;

            let backtrace = self.backtrace();
            if let BacktraceStatus::Captured = backtrace.status() {
                let mut backtrace = backtrace.to_string();
                if backtrace.starts_with("stack backtrace:") {
                    // Capitalize to match "Caused by:"
                    backtrace.replace_range(0..1, "S");
                }
                backtrace.truncate(backtrace.trim_end().len());
                write!(f, "\n\n{}", backtrace)?;
            }
        }

        Ok(())
    }
}

struct Indented<'a, D> {
    inner: &'a mut D,
    ind: Option<usize>,
    started: bool,
}

impl<'a, D> Indented<'a, D> {
    fn numbered(inner: &'a mut D, ind: usize) -> Self {
        Self {
            inner,
            ind: Some(ind),
            started: false,
        }
    }

    fn new(inner: &'a mut D) -> Self {
        Self {
            inner,
            ind: None,
            started: false,
        }
    }
}

impl<T> fmt::Write for Indented<'_, T>
where
    T: fmt::Write,
{
    fn write_str(&mut self, s: &str) -> fmt::Result {
        for (ind, mut line) in s.split('\n').enumerate() {
            if !self.started {
                // trim first line to ensure it lines up with the number nicely
                line = line.trim();
                // Don't render the first line unless its actually got text on it
                if line.is_empty() {
                    continue;
                }

                self.started = true;
                match self.ind {
                    Some(ind) => self.inner.write_fmt(format_args!("{: >5}: ", ind))?,
                    None => self.inner.write_fmt(format_args!("    "))?,
                }
            } else if ind > 0 {
                self.inner.write_char('\n')?;
                if self.ind.is_some() {
                    self.inner.write_fmt(format_args!("       "))?;
                } else {
                    self.inner.write_fmt(format_args!("    "))?;
                }
            }

            self.inner.write_fmt(format_args!("{}", line))?;
        }

        Ok(())
    }
}

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

    #[test]
    fn one_digit() {
        let input = "verify\nthis";
        let expected = "    2: verify\n       this";
        let mut output = String::new();

        Indented {
            inner: &mut output,
            ind: Some(2),
            started: false,
        }
        .write_str(input)
        .unwrap();

        assert_eq!(expected, output);
    }

    #[test]
    fn two_digits() {
        let input = "verify\nthis";
        let expected = "   12: verify\n       this";
        let mut output = String::new();

        Indented {
            inner: &mut output,
            ind: Some(12),
            started: false,
        }
        .write_str(input)
        .unwrap();

        assert_eq!(expected, output);
    }

    #[test]
    fn no_digits() {
        let input = "verify\nthis";
        let expected = "    verify\n    this";
        let mut output = String::new();

        Indented {
            inner: &mut output,
            ind: None,
            started: false,
        }
        .write_str(input)
        .unwrap();

        assert_eq!(expected, output);
    }
}
<|endoftext|>


<commit_before>use std::{num, io, fmt, error, result, thread, str, net, borrow};
use std::sync::mpsc;
use std::io::{Read, Write};
use std::borrow::Borrow;

#[derive(Clone,Debug)]
pub enum Error {
    NotEnoughData { needed: usize },
    IncorrectChecksum { got: u8, expected: u8 },
    MissingStart,
    MissingEnd,
    ParseError,
    UnexpectedCharacter { chr: char },
}
pub type Result<T> = result::Result<T, Error>;

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}", error::Error::description(self))
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match self {
            &Error::NotEnoughData { needed: _ } => "not enough data",
            &Error::IncorrectChecksum { got: _, expected: _ } => "incorrect checksum",
            &Error::MissingStart => "missing start character",
            &Error::MissingEnd => "missing end character",
            &Error::ParseError => "parse error",
            &Error::UnexpectedCharacter { chr: _ } => "unexpected character",
        }
    }
}

impl From<num::ParseIntError> for Error {
    fn from(_src: num::ParseIntError) -> Error {
        Error::ParseError
    }
}

#[derive(Clone,Debug)]
pub struct Command<'a>(borrow::Cow<'a, str>);

impl<'a> Command<'a> {
    pub fn new(src: &'a str) -> Command<'a> {
        Command(borrow::Cow::Borrowed(src))
    }

    pub fn new_owned(src: String) -> Command<'static> {
        Command(borrow::Cow::Owned(src))
    }

    pub fn checksum_str(src: &str) -> u8 {
        let mut ret = 0u8;

        for i in src.bytes() {
            ret = ret.wrapping_add(i)
        }

        ret
    }

    pub fn contents(&self) -> &str {
        let contents: &str = self.0.borrow();
        contents
    }

    pub fn checksum(&self) -> u8 {
        Command::checksum_str(self.contents())
    }

    pub fn write<T: Write>(&self, out: &mut T) -> io::Result<()> {
        write!(out, "${}#{:02x}", self.contents(), self.checksum())
    }

    pub fn fmt<T: fmt::Write>(&self, out: &mut T) -> fmt::Result {
        write!(out, "${}#{:02x}", self.contents(), self.checksum())
    }

    pub fn to_owned(&self) -> Command<'static> {
        Command(borrow::Cow::Owned(self.contents().to_owned()))
    }
}

impl<'a> fmt::Display for Command<'a> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.fmt(fmt)
    }
}

#[derive(Clone,Debug)]
pub enum Packet<'a> {
    Ack,
    Nack,
    Command(Command<'a>),
}

impl<'a> Packet<'a> {
    pub fn parse(src: &'a str) -> Result<(&'a str, Packet<'a>)> {
        if src.len() < 1 {
            return Err(Error::NotEnoughData { needed: 1 })
        }

        match src.chars().next() {
            Some('$') => {},
            Some('+') => return Ok((&src[1..], Packet::Ack)),
            Some('-') => return Ok((&src[1..], Packet::Nack)),
            Some(x) => return Err(Error::UnexpectedCharacter { chr: x }),
            _ => return Err(Error::ParseError),
        }

        if src.chars().next() != Some('$') {
            return Err(Error::MissingStart)
        }

        if src.len() < 4 {
            return Err(Error::NotEnoughData { needed: 4 })
        }

        let end = match src.find('#') {
            None => { return Err(Error::MissingEnd) },
            Some(x) => x,
        };

        if src.len() < end + 3 {
            return Err(Error::NotEnoughData { needed: end + 3 })
        }

        let contents = &src[1..end];
        let chksum_calc = Command::checksum_str(contents);
        let chksum_existing = try!(u8::from_str_radix(&src[end+1..end+3], 16));

        if chksum_calc != chksum_existing {
            return Err(Error::IncorrectChecksum { got: chksum_existing, expected: chksum_calc })
        }

        Ok((&src[end+3..], Packet::Command(Command(borrow::Cow::Borrowed(contents)))))
    }

    pub fn write<T: Write>(&self, out: &mut T) -> io::Result<()> {
        match self {
            &Packet::Ack => write!(out, "+"),
            &Packet::Nack => write!(out, "-"),
            &Packet::Command(ref x) => x.write(out),
        }
    }

    pub fn fmt<T: fmt::Write>(&self, out: &mut T) -> fmt::Result {
        match self {
            &Packet::Ack => write!(out, "+"),
            &Packet::Nack => write!(out, "-"),
            &Packet::Command(ref x) => x.fmt(out),
        }
    }

    pub fn to_owned(&self) -> Packet<'static> {
        match self {
            &Packet::Ack => Packet::Ack,
            &Packet::Nack => Packet::Nack,
            &Packet::Command(ref x) => {
                Packet::Command(Command(borrow::Cow::Owned(x.contents().to_owned())))
            },
        }
    }
}

impl<'a> fmt::Display for Packet<'a> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.fmt(fmt)
    }
}

enum ServerCommand {
    ClientReady(net::TcpStream),
    ClientDisconnected,
    Packet(Packet<'static>),
}
pub struct Server {
    receiver: Option<mpsc::Receiver<ServerCommand>>,
    stream: Option<net::TcpStream>,
    to_send: Vec<u8>,
    no_ack: bool,

    last_response: Command<'static>,
}

impl Server {
    pub fn new() -> Server {
        Server {
            receiver: None,
            stream: None,
            to_send: Vec::new(),
            no_ack: false,

            last_response: Command::new(""),
        }
    }

    pub fn start_listening<T: net::ToSocketAddrs + Send + Clone>(&mut self, bind_to: T) {
        let (send, recv) = mpsc::channel();
        self.receiver = Some(recv);

        let listener = net::TcpListener::bind(bind_to).unwrap();

        thread::spawn(move || {
            for stream in listener.incoming() {
                let mut stream = stream.unwrap();
                send.send(ServerCommand::ClientReady(stream.try_clone().unwrap())).unwrap();

                let mut vec = Vec::new();
                let mut buf = [0u8;1024];
                while let Ok(amt) = stream.read(&mut buf) {
                    vec.write_all(&buf[0..amt]).unwrap();
                    let mut done = 0;

                    loop {
                        let string = str::from_utf8(&vec[done..]).unwrap();
                        let ret = Packet::parse(string);
                        match ret {
                            Ok((new_slice, cmd)) => {
                                send.send(ServerCommand::Packet(cmd.to_owned())).unwrap();
                                done = vec.len() - new_slice.len();
                            },
                            _=> break,
                        }
                    }

                    vec = Vec::from(&vec[done..vec.len()]);
                }

                send.send(ServerCommand::ClientDisconnected).unwrap();
            }
        });
    }

    pub fn send_command(&mut self, cmd: &Command) {
        self.last_response = cmd.to_owned();
        self.send_packet(&Packet::Command(cmd.clone()))
    }

    pub fn send_packet(&mut self, cmd: &Packet) {
        //println!("<<< {}", cmd);
        if let Err(_) = match self.stream.as_mut() {
            Some(x) => cmd.write(x),
            None => cmd.write(&mut self.to_send),
        } {
            cmd.write(&mut self.to_send).ok();
        }
    }

    pub fn process_commands<T: for<'a> FnMut(&'a Command<'a>)>(&mut self, mut callback: T) {
        while let Ok(command) = self.receiver.as_mut().map_or(Err(mpsc::TryRecvError::Empty), |x| x.try_recv()) {
            match command {
                ServerCommand::ClientReady(stream) => {
                    let mut stream = stream;
                    stream.write(&self.to_send[..]).ok();
                    self.stream = Some(stream);
                    self.to_send.clear()
                },
                ServerCommand::ClientDisconnected => {
                    self.stream = None;
                    self.no_ack = false;
                },
                ServerCommand::Packet(src) => {
                    //println!(">>> {}", &src);

                    match &src {
                        &Packet::Ack => {},
                        &Packet::Nack => {
                            let packet = &Packet::Command(self.last_response.clone());
                            self.send_packet(packet);
                        },

                        &Packet::Command(ref x) => {
                            if !self.no_ack {
                                self.send_packet(&Packet::Ack);
                            }

                            if regex!("qSupported.*").is_match(x.contents()) {
                                self.send_command(&Command::new("QStartNoAckMode+;PacketSize=1023"));
                            } else if x.contents() == "QStartNoAckMode" {
                                self.no_ack = true;
                                self.send_command(&Command::new("OK"));
                            } else {
                                callback(x);
                            }
                        },
                    }
                },

            }
        }
    }
}
<commit_msg>Don't need Clone/Send on our bind_to address any more.<commit_after>use std::{num, io, fmt, error, result, thread, str, net, borrow};
use std::sync::mpsc;
use std::io::{Read, Write};
use std::borrow::Borrow;

#[derive(Clone,Debug)]
pub enum Error {
    NotEnoughData { needed: usize },
    IncorrectChecksum { got: u8, expected: u8 },
    MissingStart,
    MissingEnd,
    ParseError,
    UnexpectedCharacter { chr: char },
}
pub type Result<T> = result::Result<T, Error>;

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}", error::Error::description(self))
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match self {
            &Error::NotEnoughData { needed: _ } => "not enough data",
            &Error::IncorrectChecksum { got: _, expected: _ } => "incorrect checksum",
            &Error::MissingStart => "missing start character",
            &Error::MissingEnd => "missing end character",
            &Error::ParseError => "parse error",
            &Error::UnexpectedCharacter { chr: _ } => "unexpected character",
        }
    }
}

impl From<num::ParseIntError> for Error {
    fn from(_src: num::ParseIntError) -> Error {
        Error::ParseError
    }
}

#[derive(Clone,Debug)]
pub struct Command<'a>(borrow::Cow<'a, str>);

impl<'a> Command<'a> {
    pub fn new(src: &'a str) -> Command<'a> {
        Command(borrow::Cow::Borrowed(src))
    }

    pub fn new_owned(src: String) -> Command<'static> {
        Command(borrow::Cow::Owned(src))
    }

    pub fn checksum_str(src: &str) -> u8 {
        let mut ret = 0u8;

        for i in src.bytes() {
            ret = ret.wrapping_add(i)
        }

        ret
    }

    pub fn contents(&self) -> &str {
        let contents: &str = self.0.borrow();
        contents
    }

    pub fn checksum(&self) -> u8 {
        Command::checksum_str(self.contents())
    }

    pub fn write<T: Write>(&self, out: &mut T) -> io::Result<()> {
        write!(out, "${}#{:02x}", self.contents(), self.checksum())
    }

    pub fn fmt<T: fmt::Write>(&self, out: &mut T) -> fmt::Result {
        write!(out, "${}#{:02x}", self.contents(), self.checksum())
    }

    pub fn to_owned(&self) -> Command<'static> {
        Command(borrow::Cow::Owned(self.contents().to_owned()))
    }
}

impl<'a> fmt::Display for Command<'a> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.fmt(fmt)
    }
}

#[derive(Clone,Debug)]
pub enum Packet<'a> {
    Ack,
    Nack,
    Command(Command<'a>),
}

impl<'a> Packet<'a> {
    pub fn parse(src: &'a str) -> Result<(&'a str, Packet<'a>)> {
        if src.len() < 1 {
            return Err(Error::NotEnoughData { needed: 1 })
        }

        match src.chars().next() {
            Some('$') => {},
            Some('+') => return Ok((&src[1..], Packet::Ack)),
            Some('-') => return Ok((&src[1..], Packet::Nack)),
            Some(x) => return Err(Error::UnexpectedCharacter { chr: x }),
            _ => return Err(Error::ParseError),
        }

        if src.chars().next() != Some('$') {
            return Err(Error::MissingStart)
        }

        if src.len() < 4 {
            return Err(Error::NotEnoughData { needed: 4 })
        }

        let end = match src.find('#') {
            None => { return Err(Error::MissingEnd) },
            Some(x) => x,
        };

        if src.len() < end + 3 {
            return Err(Error::NotEnoughData { needed: end + 3 })
        }

        let contents = &src[1..end];
        let chksum_calc = Command::checksum_str(contents);
        let chksum_existing = try!(u8::from_str_radix(&src[end+1..end+3], 16));

        if chksum_calc != chksum_existing {
            return Err(Error::IncorrectChecksum { got: chksum_existing, expected: chksum_calc })
        }

        Ok((&src[end+3..], Packet::Command(Command(borrow::Cow::Borrowed(contents)))))
    }

    pub fn write<T: Write>(&self, out: &mut T) -> io::Result<()> {
        match self {
            &Packet::Ack => write!(out, "+"),
            &Packet::Nack => write!(out, "-"),
            &Packet::Command(ref x) => x.write(out),
        }
    }

    pub fn fmt<T: fmt::Write>(&self, out: &mut T) -> fmt::Result {
        match self {
            &Packet::Ack => write!(out, "+"),
            &Packet::Nack => write!(out, "-"),
            &Packet::Command(ref x) => x.fmt(out),
        }
    }

    pub fn to_owned(&self) -> Packet<'static> {
        match self {
            &Packet::Ack => Packet::Ack,
            &Packet::Nack => Packet::Nack,
            &Packet::Command(ref x) => {
                Packet::Command(Command(borrow::Cow::Owned(x.contents().to_owned())))
            },
        }
    }
}

impl<'a> fmt::Display for Packet<'a> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.fmt(fmt)
    }
}

enum ServerCommand {
    ClientReady(net::TcpStream),
    ClientDisconnected,
    Packet(Packet<'static>),
}
pub struct Server {
    receiver: Option<mpsc::Receiver<ServerCommand>>,
    stream: Option<net::TcpStream>,
    to_send: Vec<u8>,
    no_ack: bool,

    last_response: Command<'static>,
}

impl Server {
    pub fn new() -> Server {
        Server {
            receiver: None,
            stream: None,
            to_send: Vec::new(),
            no_ack: false,

            last_response: Command::new(""),
        }
    }

    pub fn start_listening<T: net::ToSocketAddrs>(&mut self, bind_to: T) {
        let (send, recv) = mpsc::channel();
        self.receiver = Some(recv);

        let listener = net::TcpListener::bind(bind_to).unwrap();

        thread::spawn(move || {
            for stream in listener.incoming() {
                let mut stream = stream.unwrap();
                send.send(ServerCommand::ClientReady(stream.try_clone().unwrap())).unwrap();

                let mut vec = Vec::new();
                let mut buf = [0u8;1024];
                while let Ok(amt) = stream.read(&mut buf) {
                    vec.write_all(&buf[0..amt]).unwrap();
                    let mut done = 0;

                    loop {
                        let string = str::from_utf8(&vec[done..]).unwrap();
                        let ret = Packet::parse(string);
                        match ret {
                            Ok((new_slice, cmd)) => {
                                send.send(ServerCommand::Packet(cmd.to_owned())).unwrap();
                                done = vec.len() - new_slice.len();
                            },
                            _=> break,
                        }
                    }

                    vec = Vec::from(&vec[done..vec.len()]);
                }

                send.send(ServerCommand::ClientDisconnected).unwrap();
            }
        });
    }

    pub fn send_command(&mut self, cmd: &Command) {
        self.last_response = cmd.to_owned();
        self.send_packet(&Packet::Command(cmd.clone()))
    }

    pub fn send_packet(&mut self, cmd: &Packet) {
        //println!("<<< {}", cmd);
        if let Err(_) = match self.stream.as_mut() {
            Some(x) => cmd.write(x),
            None => cmd.write(&mut self.to_send),
        } {
            cmd.write(&mut self.to_send).ok();
        }
    }

    pub fn process_commands<T: for<'a> FnMut(&'a Command<'a>)>(&mut self, mut callback: T) {
        while let Ok(command) = self.receiver.as_mut().map_or(Err(mpsc::TryRecvError::Empty), |x| x.try_recv()) {
            match command {
                ServerCommand::ClientReady(stream) => {
                    let mut stream = stream;
                    stream.write(&self.to_send[..]).ok();
                    self.stream = Some(stream);
                    self.to_send.clear()
                },
                ServerCommand::ClientDisconnected => {
                    self.stream = None;
                    self.no_ack = false;
                },
                ServerCommand::Packet(src) => {
                    //println!(">>> {}", &src);

                    match &src {
                        &Packet::Ack => {},
                        &Packet::Nack => {
                            let packet = &Packet::Command(self.last_response.clone());
                            self.send_packet(packet);
                        },

                        &Packet::Command(ref x) => {
                            if !self.no_ack {
                                self.send_packet(&Packet::Ack);
                            }

                            if regex!("qSupported.*").is_match(x.contents()) {
                                self.send_command(&Command::new("QStartNoAckMode+;PacketSize=1023"));
                            } else if x.contents() == "QStartNoAckMode" {
                                self.no_ack = true;
                                self.send_command(&Command::new("OK"));
                            } else {
                                callback(x);
                            }
                        },
                    }
                },

            }
        }
    }
}
<|endoftext|>


<commit_before>//! SOCKS proxy support for Hyper clients
#![doc(html_root_url="https://sfackler.github.io/rust-hyper-socks/doc/v0.1.0")]
#![warn(missing_docs)]

extern crate socks;
extern crate hyper;

use hyper::net::{NetworkConnector, HttpStream, HttpsStream, Ssl};
use socks::Socks4Stream;
use std::io;
use std::net::{SocketAddr, ToSocketAddrs};
use std::vec;

#[derive(Debug)]
struct CachedAddrs(Vec<SocketAddr>);

impl ToSocketAddrs for CachedAddrs {
    type Iter = vec::IntoIter<SocketAddr>;

    fn to_socket_addrs(&self) -> io::Result<Self::Iter> {
        Ok(self.0.clone().into_iter())
    }
}

/// A connector that will produce proxied HttpStreams.
#[derive(Debug)]
pub struct Socks4HttpConnector {
    addrs:  CachedAddrs,
    userid: String,
}

impl Socks4HttpConnector {
    /// Creates a new `Socks4HttpConnector` which will connect to the specified
    /// proxy with the specified userid.
    pub fn new<T: ToSocketAddrs>(proxy: T, userid: &str) -> io::Result<Socks4HttpConnector> {
        Ok(Socks4HttpConnector {
            addrs: CachedAddrs(try!(proxy.to_socket_addrs()).collect()),
            userid: userid.to_owned(),
        })
    }
}

impl NetworkConnector for Socks4HttpConnector {
    type Stream = HttpStream;

    fn connect(&self, host: &str, port: u16, scheme: &str) -> hyper::Result<HttpStream> {
        if scheme != "http" {
            return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                      "invalid scheme for HTTP").into());
        }

        let socket = try!(Socks4Stream::connect(&self.addrs, (host, port), &self.userid));
        Ok(HttpStream(socket.into_inner()))
    }
}

/// A connector that will produce protected, proxied HTTP streams using SSL.
#[derive(Debug)]
pub struct Socks4HttpsConnector<S> {
    addrs: CachedAddrs,
    userid: String,
    ssl: S,
}

impl<S: Ssl> Socks4HttpsConnector<S> {
    /// Creates a new `Socks4HttpsConnector` which will connect to the specified
    /// proxy with the specified userid, and use the provided SSL implementation
    /// to encrypt the resulting stream.
    pub fn new<T: ToSocketAddrs>(proxy: T, userid: &str, ssl: S) -> io::Result<Self> {
        Ok(Socks4HttpsConnector {
            addrs: CachedAddrs(try!(proxy.to_socket_addrs()).collect()),
            userid: userid.to_owned(),
            ssl: ssl,
        })
    }
}

impl<S: Ssl> NetworkConnector for Socks4HttpsConnector<S> {
    type Stream = HttpsStream<S::Stream>;

    fn connect(&self, host: &str, port: u16, scheme: &str) -> hyper::Result<Self::Stream> {
        if scheme != "http" && scheme != "https" {
            return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                      "invalid scheme for HTTPS").into());
        }

        let socket = try!(Socks4Stream::connect(&self.addrs, (host, port), &self.userid));
        let stream = HttpStream(socket.into_inner());

        if scheme == "http" {
            Ok(HttpsStream::Http(stream))
        } else {
            Ok(HttpsStream::Https(try!(self.ssl.wrap_client(stream, host))))
        }
    }
}

#[cfg(test)]
mod test {
    use hyper;
    use hyper::net::Openssl;
    use std::io::Read;

    use super::*;

    #[test]
    fn google() {
        let connector = Socks4HttpConnector::new("127.0.0.1:8080", "").unwrap();
        let client = hyper::Client::with_connector(connector);
        let mut response = client.get("http://www.google.com").send().unwrap();

        assert!(response.status.is_success());
        let mut body = vec![];
        response.read_to_end(&mut body).unwrap();
    }

    #[test]
    fn google_ssl_http() {
        let connector = Socks4HttpsConnector::new("127.0.0.1:8080", "", Openssl::default()).unwrap();
        let client = hyper::Client::with_connector(connector);
        let mut response = client.get("http://www.google.com").send().unwrap();

        assert!(response.status.is_success());
        let mut body = vec![];
        response.read_to_end(&mut body).unwrap();
    }

    #[test]
    fn google_ssl_https() {
        let connector = Socks4HttpsConnector::new("127.0.0.1:8080", "", Openssl::default()).unwrap();
        let client = hyper::Client::with_connector(connector);
        let mut response = client.get("https://www.google.com").send().unwrap();

        assert!(response.status.is_success());
        let mut body = vec![];
        response.read_to_end(&mut body).unwrap();
    }
}
<commit_msg>Don't clone addrs on each connection<commit_after>//! SOCKS proxy support for Hyper clients
#![doc(html_root_url="https://sfackler.github.io/rust-hyper-socks/doc/v0.1.0")]
#![warn(missing_docs)]

extern crate socks;
extern crate hyper;

use hyper::net::{NetworkConnector, HttpStream, HttpsStream, Ssl};
use socks::Socks4Stream;
use std::io;
use std::net::{SocketAddr, ToSocketAddrs};
use std::slice;
use std::iter;

struct AddrSlice<'a>(&'a [SocketAddr]);

impl<'a> ToSocketAddrs for AddrSlice<'a> {
    type Iter = iter::Cloned<slice::Iter<'a, SocketAddr>>;

    fn to_socket_addrs(&self) -> io::Result<Self::Iter> {
        Ok(self.0.iter().cloned())
    }
}

/// A connector that will produce proxied HttpStreams.
#[derive(Debug)]
pub struct Socks4HttpConnector {
    addrs: Vec<SocketAddr>,
    userid: String,
}

impl Socks4HttpConnector {
    /// Creates a new `Socks4HttpConnector` which will connect to the specified
    /// proxy with the specified userid.
    pub fn new<T: ToSocketAddrs>(proxy: T, userid: &str) -> io::Result<Socks4HttpConnector> {
        Ok(Socks4HttpConnector {
            addrs: try!(proxy.to_socket_addrs()).collect(),
            userid: userid.to_owned(),
        })
    }
}

impl NetworkConnector for Socks4HttpConnector {
    type Stream = HttpStream;

    fn connect(&self, host: &str, port: u16, scheme: &str) -> hyper::Result<HttpStream> {
        if scheme != "http" {
            return Err(io::Error::new(io::ErrorKind::InvalidInput, "invalid scheme for HTTP")
                           .into());
        }

        let socket = try!(Socks4Stream::connect(AddrSlice(&self.addrs),
                                                (host, port),
                                                &self.userid));
        Ok(HttpStream(socket.into_inner()))
    }
}

/// A connector that will produce protected, proxied HTTP streams using SSL.
#[derive(Debug)]
pub struct Socks4HttpsConnector<S> {
    addrs: Vec<SocketAddr>,
    userid: String,
    ssl: S,
}

impl<S: Ssl> Socks4HttpsConnector<S> {
    /// Creates a new `Socks4HttpsConnector` which will connect to the specified
    /// proxy with the specified userid, and use the provided SSL implementation
    /// to encrypt the resulting stream.
    pub fn new<T: ToSocketAddrs>(proxy: T, userid: &str, ssl: S) -> io::Result<Self> {
        Ok(Socks4HttpsConnector {
            addrs: try!(proxy.to_socket_addrs()).collect(),
            userid: userid.to_owned(),
            ssl: ssl,
        })
    }
}

impl<S: Ssl> NetworkConnector for Socks4HttpsConnector<S> {
    type Stream = HttpsStream<S::Stream>;

    fn connect(&self, host: &str, port: u16, scheme: &str) -> hyper::Result<Self::Stream> {
        if scheme != "http" && scheme != "https" {
            return Err(io::Error::new(io::ErrorKind::InvalidInput, "invalid scheme for HTTPS")
                           .into());
        }

        let socket = try!(Socks4Stream::connect(AddrSlice(&self.addrs),
                                                (host, port),
                                                &self.userid));
        let stream = HttpStream(socket.into_inner());

        if scheme == "http" {
            Ok(HttpsStream::Http(stream))
        } else {
            Ok(HttpsStream::Https(try!(self.ssl.wrap_client(stream, host))))
        }
    }
}

#[cfg(test)]
mod test {
    use hyper;
    use hyper::net::Openssl;
    use std::io::Read;

    use super::*;

    #[test]
    fn google() {
        let connector = Socks4HttpConnector::new("127.0.0.1:8080", "").unwrap();
        let client = hyper::Client::with_connector(connector);
        let mut response = client.get("http://www.google.com").send().unwrap();

        assert!(response.status.is_success());
        let mut body = vec![];
        response.read_to_end(&mut body).unwrap();
    }

    #[test]
    fn google_ssl_http() {
        let connector = Socks4HttpsConnector::new("127.0.0.1:8080", "", Openssl::default())
                            .unwrap();
        let client = hyper::Client::with_connector(connector);
        let mut response = client.get("http://www.google.com").send().unwrap();

        assert!(response.status.is_success());
        let mut body = vec![];
        response.read_to_end(&mut body).unwrap();
    }

    #[test]
    fn google_ssl_https() {
        let connector = Socks4HttpsConnector::new("127.0.0.1:8080", "", Openssl::default())
                            .unwrap();
        let client = hyper::Client::with_connector(connector);
        let mut response = client.get("https://www.google.com").send().unwrap();

        assert!(response.status.is_success());
        let mut body = vec![];
        response.read_to_end(&mut body).unwrap();
    }
}
<|endoftext|>