<commit_before>use parser::{Cond, Expr, Node, Target};
use std::str;
use std::collections::HashSet;
use syn;

fn annotations(generics: &syn::Generics) -> String {
    if generics.lifetimes.len() < 1 {
        return String::new();
    }
    let mut res = String::new();
    res.push('<');
    for lt in &generics.lifetimes {
        res.push_str(lt.lifetime.ident.as_ref());
    }
    res.push('>');
    res
}

fn path_as_identifier(s: &str) -> String {
    let mut res = String::new();
    for c in s.chars() {
        if c.is_alphanumeric() {
            res.push(c);
        } else {
            res.push_str(&format!("{:x}", c as u32));
        }
    }
    res
}

struct Generator {
    buf: String,
    indent: u8,
    start: bool,
    locals: HashSet<String>,
}

impl Generator {

    fn new() -> Generator {
        Generator {
            buf: String::new(),
            indent: 0,
            start: true,
            locals: HashSet::new(),
        }
    }

    fn indent(&mut self) {
        self.indent += 1;
    }

    fn dedent(&mut self) {
        self.indent -= 1;
    }

    fn write(&mut self, s: &str) {
        if self.start {
            for _ in 0..(self.indent * 4) {
                self.buf.push(' ');
            }
            self.start = false;
        }
        self.buf.push_str(s);
    }

    fn writeln(&mut self, s: &str) {
        if s.is_empty() {
            return;
        }
        self.write(s);
        self.buf.push('\n');
        self.start = true;
    }

    fn visit_str_lit(&mut self, s: &str) {
        self.write(&format!("\"{}\"", s));
    }

    fn visit_var(&mut self, s: &str) {
        if self.locals.contains(s) {
            self.write(s);
        } else {
            self.write(&format!("self.{}", s));
        }
    }

    fn visit_filter(&mut self, name: &str, val: &Expr) {
        self.write(&format!("askama::filters::{}(&", name));
        self.visit_expr(val);
        self.write(")");
    }

    fn visit_compare(&mut self, op: &str, left: &Expr, right: &Expr) {
        self.visit_expr(left);
        self.write(&format!(" {} ", op));
        self.visit_expr(right);
    }

    fn visit_expr(&mut self, expr: &Expr) {
        match *expr {
            Expr::StrLit(s) => self.visit_str_lit(s),
            Expr::Var(s) => self.visit_var(s),
            Expr::Filter(name, ref val) => self.visit_filter(name, val),
            Expr::Compare(op, ref left, ref right) =>
                self.visit_compare(op, left, right),
        }
    }

    fn visit_target_single(&mut self, name: &str) -> Vec<String> {
        vec![name.to_string()]
    }

    fn visit_target(&mut self, target: &Target) -> Vec<String> {
        match *target {
            Target::Name(s) => { self.visit_target_single(s) },
        }
    }

    fn write_lit(&mut self, s: &str) {
        self.write(&format!("writer.write_str({:#?}).unwrap();", s));
    }

    fn write_expr(&mut self, s: &Expr) {
        self.write("writer.write_fmt(format_args!(\"{}\", ");
        self.visit_expr(s);
        self.writeln(")).unwrap();");
    }

    fn write_cond(&mut self, conds: &[Cond]) {
        for (i, &(_, ref cond, ref nodes)) in conds.iter().enumerate() {
            match *cond {
                Some(ref expr) => {
                    if i == 0 {
                        self.write("if ");
                    } else {
                        self.write("} else if ");
                    }
                    self.visit_expr(expr);
                },
                None => { self.writeln("} else"); },
            }
            self.writeln(" {");
            self.indent();
            self.handle(nodes);
            self.dedent();
        }
        self.writeln("}");
    }

    fn write_loop(&mut self, var: &Target, iter: &Expr, body: &[Node]) {

        self.write("for ");
        let targets = self.visit_target(var);
        for name in &targets {
            self.locals.insert(name.clone());
            self.write(name);
        }
        self.write(" in &");
        self.visit_expr(iter);
        self.writeln(" {");

        self.indent();
        self.handle(body);
        self.dedent();
        self.writeln("}");
        for name in &targets {
            self.locals.remove(name);
        }
    }

    fn write_block(&mut self, name: &str) {
        self.writeln(&format!("self.render_block_{}_into(writer);", name));
    }

    fn write_block_def(&mut self, name: &str, nodes: &[Node]) {
        self.writeln("#[allow(unused_variables)]");
        self.writeln(&format!(
            "fn render_block_{}_into(&self, writer: &mut std::fmt::Write) {{",
            name));
        self.indent();
        self.handle(nodes);
        self.dedent();
        self.writeln("}");
    }

    fn handle(&mut self, nodes: &[Node]) {
        for n in nodes {
            match *n {
                Node::Lit(lws, val, rws) => {
                    self.write_lit(lws);
                    self.write_lit(val);
                    self.write_lit(rws);
                },
                Node::Expr(_, ref val) => { self.write_expr(val); },
                Node::Cond(ref conds, _) => { self.write_cond(conds); },
                Node::Loop(_, ref var, ref iter, ref body, _) => {
                    self.write_loop(var, iter, body);
                },
                Node::Block(_, name, _) => { self.write_block(name) },
                Node::BlockDef(_, name, ref block_nodes, _) => {
                    self.write_block_def(name, block_nodes);
                }
                Node::Extends(_) => {
                    panic!("no extends or block definition allowed in content");
                },
            }
        }
    }

    fn path_based_name(&mut self, ast: &syn::DeriveInput, path: &str) {
        let encoded = path_as_identifier(path);
        let original = ast.ident.as_ref();
        let anno = annotations(&ast.generics);
        self.writeln("#[allow(dead_code, non_camel_case_types)]");
        let s = format!("type TemplateFrom{}{} = {}{};",
                        encoded, &anno, original, &anno);
        self.writeln(&s);
    }

    fn template_impl(&mut self, ast: &syn::DeriveInput, nodes: &[Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.handle(nodes);
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn trait_impl(&mut self, ast: &syn::DeriveInput, base: &str, blocks: &[Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} TraitFrom{} for {}{} {{",
                              anno, path_as_identifier(base),
                              ast.ident.as_ref(), anno));
        self.indent();
        self.handle(blocks);
        self.dedent();
        self.writeln("}");
    }

    fn trait_based_impl(&mut self, ast: &syn::DeriveInput) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.writeln("self.render_trait_into(writer);");
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn template_trait(&mut self, ast: &syn::DeriveInput, path: &str,
                      blocks: &[Node], nodes: &[Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("trait{} TraitFrom{}{} {{", anno,
                              path_as_identifier(path), anno));
        self.indent();

        self.handle(blocks);

        self.writeln("fn render_trait_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.handle(nodes);
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn result(self) -> String {
        self.buf
    }

}

pub fn generate(ast: &syn::DeriveInput, path: &str, mut nodes: Vec<Node>) -> String {
    let mut base: Option<Expr> = None;
    let mut blocks = Vec::new();
    let mut content = Vec::new();
    for n in nodes.drain(..) {
        match n {
            Node::Extends(path) => {
                match base {
                    Some(_) => panic!("multiple extend blocks found"),
                    None => { base = Some(path); },
                }
            },
            Node::BlockDef(ws1, name, _, ws2) => {
                blocks.push(n);
                content.push(Node::Block(ws1, name, ws2));
            },
            _ => { content.push(n); },
        }
    }

    let mut gen = Generator::new();
    gen.path_based_name(ast, path);
    if !blocks.is_empty() {
        if let Some(extends) = base {
            if let Expr::StrLit(base_path) = extends {
                gen.trait_impl(ast, base_path, &blocks);
            }
        } else {
            gen.template_trait(ast, path, &blocks, &content);
            gen.trait_impl(ast, path, &Vec::new());
        }
        gen.trait_based_impl(ast);
    } else {
        gen.template_impl(ast, &content);
    }
    gen.result()
}
<commit_msg>Move handling of Node::Lit parts into write_lit() method<commit_after>use parser::{Cond, Expr, Node, Target};
use std::str;
use std::collections::HashSet;
use syn;

fn annotations(generics: &syn::Generics) -> String {
    if generics.lifetimes.len() < 1 {
        return String::new();
    }
    let mut res = String::new();
    res.push('<');
    for lt in &generics.lifetimes {
        res.push_str(lt.lifetime.ident.as_ref());
    }
    res.push('>');
    res
}

fn path_as_identifier(s: &str) -> String {
    let mut res = String::new();
    for c in s.chars() {
        if c.is_alphanumeric() {
            res.push(c);
        } else {
            res.push_str(&format!("{:x}", c as u32));
        }
    }
    res
}

struct Generator {
    buf: String,
    indent: u8,
    start: bool,
    locals: HashSet<String>,
}

impl Generator {

    fn new() -> Generator {
        Generator {
            buf: String::new(),
            indent: 0,
            start: true,
            locals: HashSet::new(),
        }
    }

    fn indent(&mut self) {
        self.indent += 1;
    }

    fn dedent(&mut self) {
        self.indent -= 1;
    }

    fn write(&mut self, s: &str) {
        if self.start {
            for _ in 0..(self.indent * 4) {
                self.buf.push(' ');
            }
            self.start = false;
        }
        self.buf.push_str(s);
    }

    fn writeln(&mut self, s: &str) {
        if s.is_empty() {
            return;
        }
        self.write(s);
        self.buf.push('\n');
        self.start = true;
    }

    fn visit_str_lit(&mut self, s: &str) {
        self.write(&format!("\"{}\"", s));
    }

    fn visit_var(&mut self, s: &str) {
        if self.locals.contains(s) {
            self.write(s);
        } else {
            self.write(&format!("self.{}", s));
        }
    }

    fn visit_filter(&mut self, name: &str, val: &Expr) {
        self.write(&format!("askama::filters::{}(&", name));
        self.visit_expr(val);
        self.write(")");
    }

    fn visit_compare(&mut self, op: &str, left: &Expr, right: &Expr) {
        self.visit_expr(left);
        self.write(&format!(" {} ", op));
        self.visit_expr(right);
    }

    fn visit_expr(&mut self, expr: &Expr) {
        match *expr {
            Expr::StrLit(s) => self.visit_str_lit(s),
            Expr::Var(s) => self.visit_var(s),
            Expr::Filter(name, ref val) => self.visit_filter(name, val),
            Expr::Compare(op, ref left, ref right) =>
                self.visit_compare(op, left, right),
        }
    }

    fn visit_target_single(&mut self, name: &str) -> Vec<String> {
        vec![name.to_string()]
    }

    fn visit_target(&mut self, target: &Target) -> Vec<String> {
        match *target {
            Target::Name(s) => { self.visit_target_single(s) },
        }
    }

    fn write_lit(&mut self, lws: &str, val: &str, rws: &str) {
        self.write(&format!("writer.write_str({:#?}).unwrap();", lws));
        self.write(&format!("writer.write_str({:#?}).unwrap();", val));
        self.write(&format!("writer.write_str({:#?}).unwrap();", rws));
    }

    fn write_expr(&mut self, s: &Expr) {
        self.write("writer.write_fmt(format_args!(\"{}\", ");
        self.visit_expr(s);
        self.writeln(")).unwrap();");
    }

    fn write_cond(&mut self, conds: &[Cond]) {
        for (i, &(_, ref cond, ref nodes)) in conds.iter().enumerate() {
            match *cond {
                Some(ref expr) => {
                    if i == 0 {
                        self.write("if ");
                    } else {
                        self.write("} else if ");
                    }
                    self.visit_expr(expr);
                },
                None => { self.writeln("} else"); },
            }
            self.writeln(" {");
            self.indent();
            self.handle(nodes);
            self.dedent();
        }
        self.writeln("}");
    }

    fn write_loop(&mut self, var: &Target, iter: &Expr, body: &[Node]) {

        self.write("for ");
        let targets = self.visit_target(var);
        for name in &targets {
            self.locals.insert(name.clone());
            self.write(name);
        }
        self.write(" in &");
        self.visit_expr(iter);
        self.writeln(" {");

        self.indent();
        self.handle(body);
        self.dedent();
        self.writeln("}");
        for name in &targets {
            self.locals.remove(name);
        }
    }

    fn write_block(&mut self, name: &str) {
        self.writeln(&format!("self.render_block_{}_into(writer);", name));
    }

    fn write_block_def(&mut self, name: &str, nodes: &[Node]) {
        self.writeln("#[allow(unused_variables)]");
        self.writeln(&format!(
            "fn render_block_{}_into(&self, writer: &mut std::fmt::Write) {{",
            name));
        self.indent();
        self.handle(nodes);
        self.dedent();
        self.writeln("}");
    }

    fn handle(&mut self, nodes: &[Node]) {
        for n in nodes {
            match *n {
                Node::Lit(lws, val, rws) => { self.write_lit(lws, val, rws); }
                Node::Expr(_, ref val) => { self.write_expr(val); },
                Node::Cond(ref conds, _) => { self.write_cond(conds); },
                Node::Loop(_, ref var, ref iter, ref body, _) => {
                    self.write_loop(var, iter, body);
                },
                Node::Block(_, name, _) => { self.write_block(name) },
                Node::BlockDef(_, name, ref block_nodes, _) => {
                    self.write_block_def(name, block_nodes);
                }
                Node::Extends(_) => {
                    panic!("no extends or block definition allowed in content");
                },
            }
        }
    }

    fn path_based_name(&mut self, ast: &syn::DeriveInput, path: &str) {
        let encoded = path_as_identifier(path);
        let original = ast.ident.as_ref();
        let anno = annotations(&ast.generics);
        self.writeln("#[allow(dead_code, non_camel_case_types)]");
        let s = format!("type TemplateFrom{}{} = {}{};",
                        encoded, &anno, original, &anno);
        self.writeln(&s);
    }

    fn template_impl(&mut self, ast: &syn::DeriveInput, nodes: &[Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.handle(nodes);
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn trait_impl(&mut self, ast: &syn::DeriveInput, base: &str, blocks: &[Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} TraitFrom{} for {}{} {{",
                              anno, path_as_identifier(base),
                              ast.ident.as_ref(), anno));
        self.indent();
        self.handle(blocks);
        self.dedent();
        self.writeln("}");
    }

    fn trait_based_impl(&mut self, ast: &syn::DeriveInput) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.writeln("self.render_trait_into(writer);");
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn template_trait(&mut self, ast: &syn::DeriveInput, path: &str,
                      blocks: &[Node], nodes: &[Node]) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("trait{} TraitFrom{}{} {{", anno,
                              path_as_identifier(path), anno));
        self.indent();

        self.handle(blocks);

        self.writeln("fn render_trait_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.handle(nodes);
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn result(self) -> String {
        self.buf
    }

}

pub fn generate(ast: &syn::DeriveInput, path: &str, mut nodes: Vec<Node>) -> String {
    let mut base: Option<Expr> = None;
    let mut blocks = Vec::new();
    let mut content = Vec::new();
    for n in nodes.drain(..) {
        match n {
            Node::Extends(path) => {
                match base {
                    Some(_) => panic!("multiple extend blocks found"),
                    None => { base = Some(path); },
                }
            },
            Node::BlockDef(ws1, name, _, ws2) => {
                blocks.push(n);
                content.push(Node::Block(ws1, name, ws2));
            },
            _ => { content.push(n); },
        }
    }

    let mut gen = Generator::new();
    gen.path_based_name(ast, path);
    if !blocks.is_empty() {
        if let Some(extends) = base {
            if let Expr::StrLit(base_path) = extends {
                gen.trait_impl(ast, base_path, &blocks);
            }
        } else {
            gen.template_trait(ast, path, &blocks, &content);
            gen.trait_impl(ast, path, &Vec::new());
        }
        gen.trait_based_impl(ast);
    } else {
        gen.template_impl(ast, &content);
    }
    gen.result()
}
<|endoftext|>


<commit_before>use std::time::{Duration, Instant, SystemTime};

use crate::*;
use thread::Time;

/// Returns the time elapsed between the provided time and the unix epoch as a `Duration`.
pub fn system_time_to_duration<'tcx>(time: &SystemTime) -> InterpResult<'tcx, Duration> {
    time.duration_since(SystemTime::UNIX_EPOCH)
        .map_err(|_| err_unsup_format!("times before the Unix epoch are not supported").into())
}

impl<'mir, 'tcx: 'mir> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn clock_gettime(
        &mut self,
        clk_id_op: &OpTy<'tcx, Provenance>,
        tp_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx, Scalar<Provenance>> {
        // This clock support is deliberately minimal because a lot of clock types have fiddly
        // properties (is it possible for Miri to be suspended independently of the host?). If you
        // have a use for another clock type, please open an issue.

        let this = self.eval_context_mut();

        this.assert_target_os("linux", "clock_gettime");
        this.check_no_isolation("`clock_gettime`")?;

        let clk_id = this.read_scalar(clk_id_op)?.to_i32()?;

        // Linux has two main kinds of clocks. REALTIME clocks return the actual time since the
        // Unix epoch, including effects which may cause time to move backwards such as NTP.
        // Linux further distinguishes regular and "coarse" clocks, but the "coarse" version
        // is just specified to be "faster and less precise", so we implement both the same way.
        let absolute_clocks =
            [this.eval_libc_i32("CLOCK_REALTIME")?, this.eval_libc_i32("CLOCK_REALTIME_COARSE")?];
        // The second kind is MONOTONIC clocks for which 0 is an arbitrary time point, but they are
        // never allowed to go backwards. We don't need to do any additonal monotonicity
        // enforcement because std::time::Instant already guarantees that it is monotonic.
        let relative_clocks =
            [this.eval_libc_i32("CLOCK_MONOTONIC")?, this.eval_libc_i32("CLOCK_MONOTONIC_COARSE")?];

        let duration = if absolute_clocks.contains(&clk_id) {
            system_time_to_duration(&SystemTime::now())?
        } else if relative_clocks.contains(&clk_id) {
            Instant::now().duration_since(this.machine.time_anchor)
        } else {
            let einval = this.eval_libc("EINVAL")?;
            this.set_last_error(einval)?;
            return Ok(Scalar::from_i32(-1));
        };

        let tv_sec = duration.as_secs();
        let tv_nsec = duration.subsec_nanos();

        this.write_int_fields(&[tv_sec.into(), tv_nsec.into()], &this.deref_operand(tp_op)?)?;

        Ok(Scalar::from_i32(0))
    }

    fn gettimeofday(
        &mut self,
        tv_op: &OpTy<'tcx, Provenance>,
        tz_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.assert_target_os_is_unix("gettimeofday");
        this.check_no_isolation("`gettimeofday`")?;

        // Using tz is obsolete and should always be null
        let tz = this.read_pointer(tz_op)?;
        if !this.ptr_is_null(tz)? {
            let einval = this.eval_libc("EINVAL")?;
            this.set_last_error(einval)?;
            return Ok(-1);
        }

        let duration = system_time_to_duration(&SystemTime::now())?;
        let tv_sec = duration.as_secs();
        let tv_usec = duration.subsec_micros();

        this.write_int_fields(&[tv_sec.into(), tv_usec.into()], &this.deref_operand(tv_op)?)?;

        Ok(0)
    }

    #[allow(non_snake_case)]
    fn GetSystemTimeAsFileTime(
        &mut self,
        LPFILETIME_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx> {
        let this = self.eval_context_mut();

        this.assert_target_os("windows", "GetSystemTimeAsFileTime");
        this.check_no_isolation("`GetSystemTimeAsFileTime`")?;

        let NANOS_PER_SEC = this.eval_windows_u64("time", "NANOS_PER_SEC")?;
        let INTERVALS_PER_SEC = this.eval_windows_u64("time", "INTERVALS_PER_SEC")?;
        let INTERVALS_TO_UNIX_EPOCH = this.eval_windows_u64("time", "INTERVALS_TO_UNIX_EPOCH")?;
        let NANOS_PER_INTERVAL = NANOS_PER_SEC / INTERVALS_PER_SEC;
        let SECONDS_TO_UNIX_EPOCH = INTERVALS_TO_UNIX_EPOCH / INTERVALS_PER_SEC;

        let duration = system_time_to_duration(&SystemTime::now())?
            + Duration::from_secs(SECONDS_TO_UNIX_EPOCH);
        let duration_ticks = u64::try_from(duration.as_nanos() / u128::from(NANOS_PER_INTERVAL))
            .map_err(|_| err_unsup_format!("programs running more than 2^64 Windows ticks after the Windows epoch are not supported"))?;

        let dwLowDateTime = u32::try_from(duration_ticks & 0x00000000FFFFFFFF).unwrap();
        let dwHighDateTime = u32::try_from((duration_ticks & 0xFFFFFFFF00000000) >> 32).unwrap();
        this.write_int_fields(
            &[dwLowDateTime.into(), dwHighDateTime.into()],
            &this.deref_operand(LPFILETIME_op)?,
        )?;

        Ok(())
    }

    #[allow(non_snake_case)]
    fn QueryPerformanceCounter(
        &mut self,
        lpPerformanceCount_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.assert_target_os("windows", "QueryPerformanceCounter");
        this.check_no_isolation("`QueryPerformanceCounter`")?;

        // QueryPerformanceCounter uses a hardware counter as its basis.
        // Miri will emulate a counter with a resolution of 1 nanosecond.
        let duration = Instant::now().duration_since(this.machine.time_anchor);
        let qpc = i64::try_from(duration.as_nanos()).map_err(|_| {
            err_unsup_format!("programs running longer than 2^63 nanoseconds are not supported")
        })?;
        this.write_scalar(
            Scalar::from_i64(qpc),
            &this.deref_operand(lpPerformanceCount_op)?.into(),
        )?;
        Ok(-1) // return non-zero on success
    }

    #[allow(non_snake_case)]
    fn QueryPerformanceFrequency(
        &mut self,
        lpFrequency_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.assert_target_os("windows", "QueryPerformanceFrequency");
        this.check_no_isolation("`QueryPerformanceFrequency`")?;

        // Retrieves the frequency of the hardware performance counter.
        // The frequency of the performance counter is fixed at system boot and
        // is consistent across all processors.
        // Miri emulates a "hardware" performance counter with a resolution of 1ns,
        // and thus 10^9 counts per second.
        this.write_scalar(
            Scalar::from_i64(1_000_000_000),
            &this.deref_operand(lpFrequency_op)?.into(),
        )?;
        Ok(-1) // Return non-zero on success
    }

    fn mach_absolute_time(&self) -> InterpResult<'tcx, Scalar<Provenance>> {
        let this = self.eval_context_ref();

        this.assert_target_os("macos", "mach_absolute_time");
        this.check_no_isolation("`mach_absolute_time`")?;

        // This returns a u64, with time units determined dynamically by `mach_timebase_info`.
        // We return plain nanoseconds.
        let duration = Instant::now().duration_since(this.machine.time_anchor);
        let res = u64::try_from(duration.as_nanos()).map_err(|_| {
            err_unsup_format!("programs running longer than 2^64 nanoseconds are not supported")
        })?;
        Ok(Scalar::from_u64(res))
    }

    fn mach_timebase_info(
        &mut self,
        info_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx, Scalar<Provenance>> {
        let this = self.eval_context_mut();

        this.assert_target_os("macos", "mach_timebase_info");
        this.check_no_isolation("`mach_timebase_info`")?;

        let info = this.deref_operand(info_op)?;

        // Since our emulated ticks in `mach_absolute_time` *are* nanoseconds,
        // no scaling needs to happen.
        let (numer, denom) = (1, 1);
        this.write_int_fields(&[numer.into(), denom.into()], &info)?;

        Ok(Scalar::from_i32(0)) // KERN_SUCCESS
    }

    fn nanosleep(
        &mut self,
        req_op: &OpTy<'tcx, Provenance>,
        _rem: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx, i32> {
        // Signal handlers are not supported, so rem will never be written to.

        let this = self.eval_context_mut();

        this.check_no_isolation("`nanosleep`")?;

        let duration = match this.read_timespec(&this.deref_operand(req_op)?)? {
            Some(duration) => duration,
            None => {
                let einval = this.eval_libc("EINVAL")?;
                this.set_last_error(einval)?;
                return Ok(-1);
            }
        };
        // If adding the duration overflows, let's just sleep for an hour. Waking up early is always acceptable.
        let timeout_time = Instant::now()
            .checked_add(duration)
            .unwrap_or_else(|| Instant::now().checked_add(Duration::from_secs(3600)).unwrap());
        let timeout_time = Time::Monotonic(timeout_time);

        let active_thread = this.get_active_thread();
        this.block_thread(active_thread);

        this.register_timeout_callback(
            active_thread,
            timeout_time,
            Box::new(move |ecx| {
                ecx.unblock_thread(active_thread);
                Ok(())
            }),
        );

        Ok(0)
    }

    #[allow(non_snake_case)]
    fn Sleep(&mut self, timeout: &OpTy<'tcx, Provenance>) -> InterpResult<'tcx> {
        let this = self.eval_context_mut();

        this.check_no_isolation("`Sleep`")?;

        let timeout_ms = this.read_scalar(timeout)?.to_u32()?;

        let duration = Duration::from_millis(timeout_ms.into());
        let timeout_time = Time::Monotonic(Instant::now().checked_add(duration).unwrap());

        let active_thread = this.get_active_thread();
        this.block_thread(active_thread);

        this.register_timeout_callback(
            active_thread,
            timeout_time,
            Box::new(move |ecx| {
                ecx.unblock_thread(active_thread);
                Ok(())
            }),
        );

        Ok(())
    }
}
<commit_msg>add some missing assert_target_os<commit_after>use std::time::{Duration, Instant, SystemTime};

use crate::*;
use thread::Time;

/// Returns the time elapsed between the provided time and the unix epoch as a `Duration`.
pub fn system_time_to_duration<'tcx>(time: &SystemTime) -> InterpResult<'tcx, Duration> {
    time.duration_since(SystemTime::UNIX_EPOCH)
        .map_err(|_| err_unsup_format!("times before the Unix epoch are not supported").into())
}

impl<'mir, 'tcx: 'mir> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn clock_gettime(
        &mut self,
        clk_id_op: &OpTy<'tcx, Provenance>,
        tp_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx, Scalar<Provenance>> {
        // This clock support is deliberately minimal because a lot of clock types have fiddly
        // properties (is it possible for Miri to be suspended independently of the host?). If you
        // have a use for another clock type, please open an issue.

        let this = self.eval_context_mut();

        this.assert_target_os("linux", "clock_gettime");
        this.check_no_isolation("`clock_gettime`")?;

        let clk_id = this.read_scalar(clk_id_op)?.to_i32()?;

        // Linux has two main kinds of clocks. REALTIME clocks return the actual time since the
        // Unix epoch, including effects which may cause time to move backwards such as NTP.
        // Linux further distinguishes regular and "coarse" clocks, but the "coarse" version
        // is just specified to be "faster and less precise", so we implement both the same way.
        let absolute_clocks =
            [this.eval_libc_i32("CLOCK_REALTIME")?, this.eval_libc_i32("CLOCK_REALTIME_COARSE")?];
        // The second kind is MONOTONIC clocks for which 0 is an arbitrary time point, but they are
        // never allowed to go backwards. We don't need to do any additonal monotonicity
        // enforcement because std::time::Instant already guarantees that it is monotonic.
        let relative_clocks =
            [this.eval_libc_i32("CLOCK_MONOTONIC")?, this.eval_libc_i32("CLOCK_MONOTONIC_COARSE")?];

        let duration = if absolute_clocks.contains(&clk_id) {
            system_time_to_duration(&SystemTime::now())?
        } else if relative_clocks.contains(&clk_id) {
            Instant::now().duration_since(this.machine.time_anchor)
        } else {
            let einval = this.eval_libc("EINVAL")?;
            this.set_last_error(einval)?;
            return Ok(Scalar::from_i32(-1));
        };

        let tv_sec = duration.as_secs();
        let tv_nsec = duration.subsec_nanos();

        this.write_int_fields(&[tv_sec.into(), tv_nsec.into()], &this.deref_operand(tp_op)?)?;

        Ok(Scalar::from_i32(0))
    }

    fn gettimeofday(
        &mut self,
        tv_op: &OpTy<'tcx, Provenance>,
        tz_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.assert_target_os_is_unix("gettimeofday");
        this.check_no_isolation("`gettimeofday`")?;

        // Using tz is obsolete and should always be null
        let tz = this.read_pointer(tz_op)?;
        if !this.ptr_is_null(tz)? {
            let einval = this.eval_libc("EINVAL")?;
            this.set_last_error(einval)?;
            return Ok(-1);
        }

        let duration = system_time_to_duration(&SystemTime::now())?;
        let tv_sec = duration.as_secs();
        let tv_usec = duration.subsec_micros();

        this.write_int_fields(&[tv_sec.into(), tv_usec.into()], &this.deref_operand(tv_op)?)?;

        Ok(0)
    }

    #[allow(non_snake_case)]
    fn GetSystemTimeAsFileTime(
        &mut self,
        LPFILETIME_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx> {
        let this = self.eval_context_mut();

        this.assert_target_os("windows", "GetSystemTimeAsFileTime");
        this.check_no_isolation("`GetSystemTimeAsFileTime`")?;

        let NANOS_PER_SEC = this.eval_windows_u64("time", "NANOS_PER_SEC")?;
        let INTERVALS_PER_SEC = this.eval_windows_u64("time", "INTERVALS_PER_SEC")?;
        let INTERVALS_TO_UNIX_EPOCH = this.eval_windows_u64("time", "INTERVALS_TO_UNIX_EPOCH")?;
        let NANOS_PER_INTERVAL = NANOS_PER_SEC / INTERVALS_PER_SEC;
        let SECONDS_TO_UNIX_EPOCH = INTERVALS_TO_UNIX_EPOCH / INTERVALS_PER_SEC;

        let duration = system_time_to_duration(&SystemTime::now())?
            + Duration::from_secs(SECONDS_TO_UNIX_EPOCH);
        let duration_ticks = u64::try_from(duration.as_nanos() / u128::from(NANOS_PER_INTERVAL))
            .map_err(|_| err_unsup_format!("programs running more than 2^64 Windows ticks after the Windows epoch are not supported"))?;

        let dwLowDateTime = u32::try_from(duration_ticks & 0x00000000FFFFFFFF).unwrap();
        let dwHighDateTime = u32::try_from((duration_ticks & 0xFFFFFFFF00000000) >> 32).unwrap();
        this.write_int_fields(
            &[dwLowDateTime.into(), dwHighDateTime.into()],
            &this.deref_operand(LPFILETIME_op)?,
        )?;

        Ok(())
    }

    #[allow(non_snake_case)]
    fn QueryPerformanceCounter(
        &mut self,
        lpPerformanceCount_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.assert_target_os("windows", "QueryPerformanceCounter");
        this.check_no_isolation("`QueryPerformanceCounter`")?;

        // QueryPerformanceCounter uses a hardware counter as its basis.
        // Miri will emulate a counter with a resolution of 1 nanosecond.
        let duration = Instant::now().duration_since(this.machine.time_anchor);
        let qpc = i64::try_from(duration.as_nanos()).map_err(|_| {
            err_unsup_format!("programs running longer than 2^63 nanoseconds are not supported")
        })?;
        this.write_scalar(
            Scalar::from_i64(qpc),
            &this.deref_operand(lpPerformanceCount_op)?.into(),
        )?;
        Ok(-1) // return non-zero on success
    }

    #[allow(non_snake_case)]
    fn QueryPerformanceFrequency(
        &mut self,
        lpFrequency_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.assert_target_os("windows", "QueryPerformanceFrequency");
        this.check_no_isolation("`QueryPerformanceFrequency`")?;

        // Retrieves the frequency of the hardware performance counter.
        // The frequency of the performance counter is fixed at system boot and
        // is consistent across all processors.
        // Miri emulates a "hardware" performance counter with a resolution of 1ns,
        // and thus 10^9 counts per second.
        this.write_scalar(
            Scalar::from_i64(1_000_000_000),
            &this.deref_operand(lpFrequency_op)?.into(),
        )?;
        Ok(-1) // Return non-zero on success
    }

    fn mach_absolute_time(&self) -> InterpResult<'tcx, Scalar<Provenance>> {
        let this = self.eval_context_ref();

        this.assert_target_os("macos", "mach_absolute_time");
        this.check_no_isolation("`mach_absolute_time`")?;

        // This returns a u64, with time units determined dynamically by `mach_timebase_info`.
        // We return plain nanoseconds.
        let duration = Instant::now().duration_since(this.machine.time_anchor);
        let res = u64::try_from(duration.as_nanos()).map_err(|_| {
            err_unsup_format!("programs running longer than 2^64 nanoseconds are not supported")
        })?;
        Ok(Scalar::from_u64(res))
    }

    fn mach_timebase_info(
        &mut self,
        info_op: &OpTy<'tcx, Provenance>,
    ) -> InterpResult<'tcx, Scalar<Provenance>> {
        let this = self.eval_context_mut();

        this.assert_target_os("macos", "mach_timebase_info");
        this.check_no_isolation("`mach_timebase_info`")?;

        let info = this.deref_operand(info_op)?;

        // Since our emulated ticks in `mach_absolute_time` *are* nanoseconds,
        // no scaling needs to happen.
        let (numer, denom) = (1, 1);
        this.write_int_fields(&[numer.into(), denom.into()], &info)?;

        Ok(Scalar::from_i32(0)) // KERN_SUCCESS
    }

    fn nanosleep(
        &mut self,
        req_op: &OpTy<'tcx, Provenance>,
        _rem: &OpTy<'tcx, Provenance>, // Signal handlers are not supported, so rem will never be written to.
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.assert_target_os_is_unix("nanosleep");
        this.check_no_isolation("`nanosleep`")?;

        let duration = match this.read_timespec(&this.deref_operand(req_op)?)? {
            Some(duration) => duration,
            None => {
                let einval = this.eval_libc("EINVAL")?;
                this.set_last_error(einval)?;
                return Ok(-1);
            }
        };
        // If adding the duration overflows, let's just sleep for an hour. Waking up early is always acceptable.
        let timeout_time = Instant::now()
            .checked_add(duration)
            .unwrap_or_else(|| Instant::now().checked_add(Duration::from_secs(3600)).unwrap());
        let timeout_time = Time::Monotonic(timeout_time);

        let active_thread = this.get_active_thread();
        this.block_thread(active_thread);

        this.register_timeout_callback(
            active_thread,
            timeout_time,
            Box::new(move |ecx| {
                ecx.unblock_thread(active_thread);
                Ok(())
            }),
        );

        Ok(0)
    }

    #[allow(non_snake_case)]
    fn Sleep(&mut self, timeout: &OpTy<'tcx, Provenance>) -> InterpResult<'tcx> {
        let this = self.eval_context_mut();

        this.assert_target_os("windows", "Sleep");
        this.check_no_isolation("`Sleep`")?;

        let timeout_ms = this.read_scalar(timeout)?.to_u32()?;

        let duration = Duration::from_millis(timeout_ms.into());
        let timeout_time = Time::Monotonic(Instant::now().checked_add(duration).unwrap());

        let active_thread = this.get_active_thread();
        this.block_thread(active_thread);

        this.register_timeout_callback(
            active_thread,
            timeout_time,
            Box::new(move |ecx| {
                ecx.unblock_thread(active_thread);
                Ok(())
            }),
        );

        Ok(())
    }
}
<|endoftext|>


<commit_before>//! This module implements a synchronous transport over a raw TcpListener.

use std::os::unix::net::UnixStream;
use std::{error, fmt, io, path, time};

use serde;
use serde_json;

use crate::client::Transport;
use crate::{Request, Response};

/// Error that can occur while using the UDS transport.
#[derive(Debug)]
pub enum Error {
    /// An error occurred on the socket layer
    SocketError(io::Error),
    /// We didn't receive a complete response till the deadline ran out
    Timeout,
    /// JSON parsing error.
    Json(serde_json::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::SocketError(ref e) => write!(f, "Couldn't connect to host: {}", e),
            Error::Timeout => f.write_str("Didn't receive response data in time, timed out."),
            Error::Json(ref e) => write!(f, "JSON error: {}", e),
        }
    }
}

impl error::Error for Error {
    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
        use self::Error::*;

        match *self {
            SocketError(ref e) => Some(e),
            Timeout => None,
            Json(ref e) => Some(e),
        }
    }
}

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self {
        Error::SocketError(e)
    }
}

impl From<serde_json::Error> for Error {
    fn from(e: serde_json::Error) -> Self {
        Error::Json(e)
    }
}

impl From<Error> for crate::error::Error {
    fn from(e: Error) -> crate::error::Error {
        match e {
            Error::Json(e) => crate::error::Error::Json(e),
            e => crate::error::Error::Transport(Box::new(e)),
        }
    }
}

/// Simple synchronous UDS transport.
#[derive(Debug, Clone)]
pub struct UdsTransport {
    /// The path to the Unix Domain Socket
    pub sockpath: path::PathBuf,
    /// The read and write timeout to use
    pub timeout: Option<time::Duration>,
}

impl UdsTransport {
    /// Create a new UdsTransport without timeouts to use
    pub fn new<P: AsRef<path::Path>>(sockpath: P) -> UdsTransport {
        UdsTransport {
            sockpath: sockpath.as_ref().to_path_buf(),
            timeout: None,
        }
    }

    fn request<R>(&self, req: impl serde::Serialize) -> Result<R, Error>
    where
        R: for<'a> serde::de::Deserialize<'a>,
    {
        let mut sock = UnixStream::connect(&self.sockpath)?;
        sock.set_read_timeout(self.timeout)?;
        sock.set_write_timeout(self.timeout)?;

        serde_json::to_writer(&mut sock, &req)?;

        // NOTE: we don't check the id there, so it *must* be synchronous
        let resp: R = serde_json::Deserializer::from_reader(&mut sock)
            .into_iter()
            .next()
            .ok_or(Error::Timeout)??;
        Ok(resp)
    }
}

impl Transport for UdsTransport {
    fn send_request(&self, req: Request) -> Result<Response, crate::error::Error> {
        Ok(self.request(req)?)
    }

    fn send_batch(&self, reqs: &[Request]) -> Result<Vec<Response>, crate::error::Error> {
        Ok(self.request(reqs)?)
    }

    fn fmt_target(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.sockpath.to_string_lossy())
    }
}

#[cfg(test)]
mod tests {
    use std::{
        fs,
        io::{Read, Write},
        os::unix::net::UnixListener,
        process,
        thread,
    };

    use super::*;
    use crate::Client;

    // Test a dummy request / response over an UDS
    #[test]
    fn sanity_check_uds_transport() {
        let socket_path: path::PathBuf = format!("uds_scratch_{}.socket", process::id()).into();
        // Any leftover?
        fs::remove_file(&socket_path).unwrap_or_else(|_| ());

        let server = UnixListener::bind(&socket_path).unwrap();
        let dummy_req = Request {
            method: "getinfo",
            params: &[],
            id: serde_json::Value::Number(111.into()),
            jsonrpc: Some("2.0".into()),
        };
        let dummy_req_ser = serde_json::to_vec(&dummy_req).unwrap();
        let dummy_resp = Response {
            result: None,
            error: None,
            id: serde_json::Value::Number(111.into()),
            jsonrpc: Some("2.0".into()),
        };
        let dummy_resp_ser = serde_json::to_vec(&dummy_resp).unwrap();

        let cli_socket_path = socket_path.clone();
        let client_thread = thread::spawn(move || {
            let transport = UdsTransport {
                sockpath: cli_socket_path,
                timeout: Some(time::Duration::from_secs(5)),
            };
            let client = Client::with_transport(transport);

            client.send_request(dummy_req.clone()).unwrap()
        });

        let (mut stream, _) = server.accept().unwrap();
        stream.set_read_timeout(Some(time::Duration::from_secs(5))).unwrap();
        let mut recv_req = vec![0; dummy_req_ser.len()];
        let mut read = 0;
        while read < dummy_req_ser.len() {
            read += stream.read(&mut recv_req[read..]).unwrap();
        }
        assert_eq!(recv_req, dummy_req_ser);

        stream.write_all(&dummy_resp_ser).unwrap();
        stream.flush().unwrap();
        let recv_resp = client_thread.join().unwrap();
        assert_eq!(serde_json::to_vec(&recv_resp).unwrap(), dummy_resp_ser);

        // Clean up
        drop(server);
        fs::remove_file(&socket_path).unwrap();
    }
}
<commit_msg>Remove unnecessary closure<commit_after>//! This module implements a synchronous transport over a raw TcpListener.

use std::os::unix::net::UnixStream;
use std::{error, fmt, io, path, time};

use serde;
use serde_json;

use crate::client::Transport;
use crate::{Request, Response};

/// Error that can occur while using the UDS transport.
#[derive(Debug)]
pub enum Error {
    /// An error occurred on the socket layer
    SocketError(io::Error),
    /// We didn't receive a complete response till the deadline ran out
    Timeout,
    /// JSON parsing error.
    Json(serde_json::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::SocketError(ref e) => write!(f, "Couldn't connect to host: {}", e),
            Error::Timeout => f.write_str("Didn't receive response data in time, timed out."),
            Error::Json(ref e) => write!(f, "JSON error: {}", e),
        }
    }
}

impl error::Error for Error {
    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
        use self::Error::*;

        match *self {
            SocketError(ref e) => Some(e),
            Timeout => None,
            Json(ref e) => Some(e),
        }
    }
}

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self {
        Error::SocketError(e)
    }
}

impl From<serde_json::Error> for Error {
    fn from(e: serde_json::Error) -> Self {
        Error::Json(e)
    }
}

impl From<Error> for crate::error::Error {
    fn from(e: Error) -> crate::error::Error {
        match e {
            Error::Json(e) => crate::error::Error::Json(e),
            e => crate::error::Error::Transport(Box::new(e)),
        }
    }
}

/// Simple synchronous UDS transport.
#[derive(Debug, Clone)]
pub struct UdsTransport {
    /// The path to the Unix Domain Socket
    pub sockpath: path::PathBuf,
    /// The read and write timeout to use
    pub timeout: Option<time::Duration>,
}

impl UdsTransport {
    /// Create a new UdsTransport without timeouts to use
    pub fn new<P: AsRef<path::Path>>(sockpath: P) -> UdsTransport {
        UdsTransport {
            sockpath: sockpath.as_ref().to_path_buf(),
            timeout: None,
        }
    }

    fn request<R>(&self, req: impl serde::Serialize) -> Result<R, Error>
    where
        R: for<'a> serde::de::Deserialize<'a>,
    {
        let mut sock = UnixStream::connect(&self.sockpath)?;
        sock.set_read_timeout(self.timeout)?;
        sock.set_write_timeout(self.timeout)?;

        serde_json::to_writer(&mut sock, &req)?;

        // NOTE: we don't check the id there, so it *must* be synchronous
        let resp: R = serde_json::Deserializer::from_reader(&mut sock)
            .into_iter()
            .next()
            .ok_or(Error::Timeout)??;
        Ok(resp)
    }
}

impl Transport for UdsTransport {
    fn send_request(&self, req: Request) -> Result<Response, crate::error::Error> {
        Ok(self.request(req)?)
    }

    fn send_batch(&self, reqs: &[Request]) -> Result<Vec<Response>, crate::error::Error> {
        Ok(self.request(reqs)?)
    }

    fn fmt_target(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.sockpath.to_string_lossy())
    }
}

#[cfg(test)]
mod tests {
    use std::{
        fs,
        io::{Read, Write},
        os::unix::net::UnixListener,
        process,
        thread,
    };

    use super::*;
    use crate::Client;

    // Test a dummy request / response over an UDS
    #[test]
    fn sanity_check_uds_transport() {
        let socket_path: path::PathBuf = format!("uds_scratch_{}.socket", process::id()).into();
        // Any leftover?
        fs::remove_file(&socket_path).unwrap_or(());

        let server = UnixListener::bind(&socket_path).unwrap();
        let dummy_req = Request {
            method: "getinfo",
            params: &[],
            id: serde_json::Value::Number(111.into()),
            jsonrpc: Some("2.0".into()),
        };
        let dummy_req_ser = serde_json::to_vec(&dummy_req).unwrap();
        let dummy_resp = Response {
            result: None,
            error: None,
            id: serde_json::Value::Number(111.into()),
            jsonrpc: Some("2.0".into()),
        };
        let dummy_resp_ser = serde_json::to_vec(&dummy_resp).unwrap();

        let cli_socket_path = socket_path.clone();
        let client_thread = thread::spawn(move || {
            let transport = UdsTransport {
                sockpath: cli_socket_path,
                timeout: Some(time::Duration::from_secs(5)),
            };
            let client = Client::with_transport(transport);

            client.send_request(dummy_req.clone()).unwrap()
        });

        let (mut stream, _) = server.accept().unwrap();
        stream.set_read_timeout(Some(time::Duration::from_secs(5))).unwrap();
        let mut recv_req = vec![0; dummy_req_ser.len()];
        let mut read = 0;
        while read < dummy_req_ser.len() {
            read += stream.read(&mut recv_req[read..]).unwrap();
        }
        assert_eq!(recv_req, dummy_req_ser);

        stream.write_all(&dummy_resp_ser).unwrap();
        stream.flush().unwrap();
        let recv_resp = client_thread.join().unwrap();
        assert_eq!(serde_json::to_vec(&recv_resp).unwrap(), dummy_resp_ser);

        // Clean up
        drop(server);
        fs::remove_file(&socket_path).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::io::{Error,ErrorKind,Read,Result,Write};
use std::{result,str};
use std::iter::repeat;
use std::collections::{HashMap,VecDeque};
use amq_protocol_types::AMQPValue;
use nom::{HexDisplay,IResult,Offset};
use sasl::{ChannelBinding, Credentials, Secret, Mechanism};
use sasl::mechanisms::Plain;
use cookie_factory::GenError;

use format::frame::*;
use format::field::*;
use format::content::*;
use channel::Channel;
use api::{Answer,ChannelState};
use buffer::Buffer;
use generated::*;
use error;

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ConnectionState {
  Initial,
  Connecting(ConnectingState),
  Connected,
  Closing(ClosingState),
  Closed,
  Error,
}

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ConnectingState {
  Initial,
  SentProtocolHeader,
  ReceivedStart,
  SentStartOk,
  ReceivedSecure,
  SentSecure,
  ReceivedSecondSecure,
  ReceivedTune,
  SentTuneOk,
  SentOpen,
  Error,
}

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ClosingState {
  Initial,
  SentClose,
  ReceivedClose,
  SentCloseOk,
  ReceivedCloseOk,
  Error,
}

#[derive(Clone,Debug,PartialEq)]
pub struct Configuration {
  channel_max: u16,
  frame_max:   u32,
  heartbeat:   u16,
}

#[derive(Clone,Debug,PartialEq)]
pub struct Connection<'a> {
  pub state:            ConnectionState,
  pub channels:         HashMap<u16, Channel<'a>>,
  pub configuration:    Configuration,
  pub channel_index:    u16,
  pub prefetch_size:    u32,
  pub prefetch_count:   u16,
  pub frame_queue:      VecDeque<Frame>,
}

impl<'a> Connection<'a> {
  pub fn new() -> Connection<'a> {
    let mut h = HashMap::new();
    h.insert(0, Channel::global());

    let configuration = Configuration {
      channel_max: 0,
      frame_max:   8192,
      heartbeat:   60,
    };

    Connection {
      state:          ConnectionState::Initial,
      channels:       h,
      configuration:  configuration,
      channel_index:  1,
      prefetch_size:  0,
      prefetch_count: 0,
      frame_queue:    VecDeque::new(),
    }
  }

  pub fn create_channel(&mut self) -> u16 {
    let c  = Channel::new(self.channel_index);
    self.channels.insert(self.channel_index, c);
    self.channel_index += 1;

    self.channel_index - 1
  }

  pub fn set_channel_state(&mut self, channel_id: u16, new_state: ChannelState) {
    self.channels.get_mut(&channel_id).map(|c| c.state = new_state);
  }

  pub fn check_state(&self, channel_id: u16, state: ChannelState) -> Option<bool> {
    self.channels
          .get(&channel_id)
          .map(|c| c.state == state)
  }

  pub fn get_state(&self, channel_id: u16) -> Option<ChannelState> {
    self.channels
          .get(&channel_id)
          .map(|c| c.state.clone())
  }

  pub fn push_back_answer(&mut self, channel_id: u16, answer: Answer) {
    self.channels
      .get_mut(&channel_id)
      .map(|c| c.awaiting.push_back(answer));
  }

  pub fn check_next_answer(&self, channel_id: u16, answer: Answer) -> bool {
    self.channels
          .get(&channel_id)
          .map(|c| c.awaiting.front() == Some(&answer)).unwrap_or(false)
  }

  pub fn get_next_answer(&mut self, channel_id: u16) -> Option<Answer> {
    self.channels
          .get_mut(&channel_id)
          .and_then(|c| c.awaiting.pop_front())
  }

  pub fn is_connected(&self, channel_id: u16) -> bool {
    self.channels
          .get(&channel_id)
          .map(|c| c.is_connected()).unwrap_or(false)
  }

  pub fn connect(&mut self) -> Result<ConnectionState> {
    if self.state != ConnectionState::Initial {
      self.state = ConnectionState::Error;
      return Err(Error::new(ErrorKind::Other, "invalid state"))
    }

    self.frame_queue.push_back(Frame::ProtocolHeader);
    self.state = ConnectionState::Connecting(ConnectingState::SentProtocolHeader);
    Ok(self.state)
  }

  pub fn serialize(&mut self, send_buffer: &mut [u8]) -> Result<(usize, ConnectionState)> {
    let next_msg = self.frame_queue.pop_front();
    if next_msg == None {
      return Err(Error::new(ErrorKind::WouldBlock, "no new message"));
    }

    let next_msg = next_msg.unwrap();
    println!("will write to buffer: {:?}", next_msg);

    let gen_res = match &next_msg {
      &Frame::ProtocolHeader => {
        gen_protocol_header((send_buffer, 0)).map(|tup| tup.1)
      },
      &Frame::Heartbeat(_) => {
        gen_heartbeat_frame((send_buffer, 0)).map(|tup| tup.1)
      },
      &Frame::Method(channel, ref method) => {
        gen_method_frame((send_buffer, 0), channel, method).map(|tup| tup.1)
      },
      &Frame::Header(channel_id, class_id, ref header) => {
        gen_content_header_frame((send_buffer, 0), channel_id, class_id, header.body_size).map(|tup| tup.1)
      },
      &Frame::Body(channel_id, ref data) => {
        gen_content_body_frame((send_buffer, 0), channel_id, data).map(|tup| tup.1)
      }
    };

    match gen_res {
      Ok(sz) => {
        Ok((sz, self.state))
      },
      Err(e) => {
        println!("error generating frame: {:?}", e);
        self.state = ConnectionState::Error;
        match e {
          GenError::BufferTooSmall(_) => {
            self.frame_queue.push_front(next_msg);
            return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
          },
          GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
            return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
          }
        }
      }
    }
  }

  pub fn parse(&mut self, data: &[u8]) -> Result<(usize,ConnectionState)> {
    let parsed_frame = frame(data);
    match parsed_frame {
      IResult::Done(i,_)     => {},
      IResult::Incomplete(_) => {
        return Ok((0,self.state));
      },
      IResult::Error(e) => {
        //FIXME: should probably disconnect on error here
        let err = format!("parse error: {:?}", e);
        self.state = ConnectionState::Error;
        return Err(Error::new(ErrorKind::Other, err))
      }
    }

    let (i, f) = parsed_frame.unwrap();

    //FIXME: what happens if we fail to parse a packet in a channel?
    // do we continue?
    let consumed = data.offset(i);

    self.handle_frame(f);

    return Ok((consumed, self.state));
  }

  pub fn handle_frame(&mut self, f: Frame) {
    match f {
      Frame::ProtocolHeader => {
        println!("error: the client should not receive a protocol header");
        self.state = ConnectionState::Error;
      },
      Frame::Method(channel_id, method) => {
        if channel_id == 0 {
          self.handle_global_method(method);
        } else {
          self.receive_method(channel_id, method);
        }
      },
      Frame::Heartbeat(channel_id) => {
        self.frame_queue.push_back(Frame::Heartbeat(0));
      },
      Frame::Header(channel_id, class_id, header) => {
        self.handle_content_header_frame(channel_id, header.body_size);
      },
      Frame::Body(channel_id, payload) => {
        self.handle_body_frame(channel_id, payload);
      }
    };
  }

  pub fn handle_global_method(&mut self, c: Class) {
    match self.state {
      ConnectionState::Initial | ConnectionState::Closed | ConnectionState::Error => {
        self.state = ConnectionState::Error
      },
      ConnectionState::Connecting(connecting_state) => {
        match connecting_state {
          ConnectingState::Initial | ConnectingState::Error => {
            self.state = ConnectionState::Error
          },
          ConnectingState::SentProtocolHeader => {
            if let Class::Connection(connection::Methods::Start(s)) = c {
              println!("Server sent Connection::Start: {:?}", s);
              self.state = ConnectionState::Connecting(ConnectingState::ReceivedStart);

              let mut h = HashMap::new();
              h.insert("product".to_string(), AMQPValue::LongString("lapin".to_string()));

              let creds = Credentials {
                username: Some("guest".to_owned()),
                secret: Secret::Password("guest".to_owned()),
                channel_binding: ChannelBinding::None,
              };

              let mut mechanism = Plain::from_credentials(creds).unwrap();

              let initial_data = mechanism.initial().unwrap();
              let s = str::from_utf8(&initial_data).unwrap();

              //FIXME: fill with user configured data
              //we need to handle the server properties, and have some client properties
              let start_ok = Class::Connection(connection::Methods::StartOk(
                connection::StartOk {
                  client_properties: h,
                  mechanism: "PLAIN".to_string(),
                  locale:    "en_US".to_string(), // FIXME: comes from the server
                  response:  s.to_string(),     //FIXME: implement SASL
                }
              ));

              println!("client sending Connection::StartOk: {:?}", start_ok);
              self.frame_queue.push_back(Frame::Method(0, start_ok));
              self.state = ConnectionState::Connecting(ConnectingState::SentStartOk);
            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          /*ConnectingState::ReceivedStart => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },*/
          ConnectingState::SentStartOk => {
            if let Class::Connection(connection::Methods::Tune(t)) = c {
              println!("Server sent Connection::Tune: {:?}", t);
              self.state = ConnectionState::Connecting(ConnectingState::ReceivedTune);

              self.configuration.channel_max = t.channel_max;
              self.configuration.heartbeat   = t.heartbeat;

              if t.frame_max > self.configuration.frame_max {
                //FIXME: what do we do without the buffers available?
                //self.send_buffer.grow(t.frame_max as usize);
                //self.receive_buffer.grow(t.frame_max as usize);
              }

              self.configuration.frame_max = t.frame_max;

              let tune_ok = Class::Connection(connection::Methods::TuneOk(
                connection::TuneOk {
                  channel_max : self.configuration.channel_max,
                  frame_max   : self.configuration.frame_max,
                  heartbeat   : self.configuration.heartbeat,
                }
              ));

              println!("client sending Connection::TuneOk: {:?}", tune_ok);

              self.frame_queue.push_back(Frame::Method(0, tune_ok));
              self.state = ConnectionState::Connecting(ConnectingState::SentTuneOk);

              let open = Class::Connection(connection::Methods::Open(
                  connection::Open {
                    virtual_host: "/".to_string(),
                    capabilities: "".to_string(),
                    insist:       false,
                  }
                  ));

              println!("client sending Connection::Open: {:?}", open);
              self.frame_queue.push_back(Frame::Method(0,open));
              self.state = ConnectionState::Connecting(ConnectingState::SentOpen);

            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          ConnectingState::ReceivedSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::SentSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::ReceivedSecondSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::ReceivedTune => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::SentOpen => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
            if let Class::Connection(connection::Methods::OpenOk(o)) = c {
              println!("Server sent Connection::OpenOk: {:?}, client now connected", o);
              self.state = ConnectionState::Connected;
            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          ConnectingState::Error => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          s => {
            println!("invalid state {:?}", s);
            self.state = ConnectionState::Error;
          }
        }
      },
      ConnectionState::Connected => {},
      ConnectionState::Closing(ClosingState) => {},
    };
  }

  pub fn handle_content_header_frame(&mut self, channel_id: u16, size: u64) {
    let state = self.channels.get_mut(&channel_id).map(|channel| {
      channel.state.clone()
    }).unwrap();
    if let ChannelState::WillReceiveContent(queue_name, consumer_tag) = state {
      self.set_channel_state(channel_id, ChannelState::ReceivingContent(queue_name, consumer_tag.clone(), size as usize));
    } else {
      self.set_channel_state(channel_id, ChannelState::Error);
    }
  }

  pub fn handle_body_frame(&mut self, channel_id: u16, payload: Vec<u8>) {
    let state = self.channels.get_mut(&channel_id).map(|channel| {
      channel.state.clone()
    }).unwrap();

    let payload_size = payload.len();

    if let ChannelState::ReceivingContent(queue_name, consumer_tag, remaining_size) = state {
      if remaining_size >= payload_size {

        if let Some(ref mut c) = self.channels.get_mut(&channel_id) {
          if let Some(ref mut q) = c.queues.get_mut(&queue_name) {
            if let Some(ref mut cs) = q.consumers.get_mut(&consumer_tag) {
              (*cs.callback).receive_content(payload);
              if remaining_size == payload_size {
                (*cs.callback).done();
              }
            }
          }
        }

        if remaining_size == payload_size {
          self.set_channel_state(channel_id, ChannelState::Connected);
        } else {
          self.set_channel_state(channel_id, ChannelState::ReceivingContent(queue_name, consumer_tag, remaining_size - payload_size));
        }
      } else {
        println!("body frame too large");
        self.set_channel_state(channel_id, ChannelState::Error);
      }
    } else {
      self.set_channel_state(channel_id, ChannelState::Error);
    }
  }

  //FIXME: no chunking for now
  pub fn send_content_frames(&mut self, channel_id: u16, class_id: u16, slice: &[u8]) {
    let header = ContentHeader {
      class_id:       class_id,
      weight:         0,
      body_size:      slice.len() as u64,
      property_flags: 0x2000,
      property_list:  HashMap::new(),
    };
    self.frame_queue.push_back(Frame::Header(channel_id, class_id, header));
    self.frame_queue.push_back(Frame::Body(channel_id,Vec::from(slice)));
  }

  pub fn send_method_frame(&mut self, channel: u16, method: Class) -> result::Result<(), error::Error> {
    self.frame_queue.push_back(Frame::Method(channel,method));
    Ok(())
  }
}
<commit_msg>next frame method<commit_after>use std::io::{Error,ErrorKind,Read,Result,Write};
use std::{result,str};
use std::iter::repeat;
use std::collections::{HashMap,VecDeque};
use amq_protocol_types::AMQPValue;
use nom::{HexDisplay,IResult,Offset};
use sasl::{ChannelBinding, Credentials, Secret, Mechanism};
use sasl::mechanisms::Plain;
use cookie_factory::GenError;

use format::frame::*;
use format::field::*;
use format::content::*;
use channel::Channel;
use api::{Answer,ChannelState};
use buffer::Buffer;
use generated::*;
use error;

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ConnectionState {
  Initial,
  Connecting(ConnectingState),
  Connected,
  Closing(ClosingState),
  Closed,
  Error,
}

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ConnectingState {
  Initial,
  SentProtocolHeader,
  ReceivedStart,
  SentStartOk,
  ReceivedSecure,
  SentSecure,
  ReceivedSecondSecure,
  ReceivedTune,
  SentTuneOk,
  SentOpen,
  Error,
}

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ClosingState {
  Initial,
  SentClose,
  ReceivedClose,
  SentCloseOk,
  ReceivedCloseOk,
  Error,
}

#[derive(Clone,Debug,PartialEq)]
pub struct Configuration {
  channel_max: u16,
  frame_max:   u32,
  heartbeat:   u16,
}

#[derive(Clone,Debug,PartialEq)]
pub struct Connection<'a> {
  pub state:            ConnectionState,
  pub channels:         HashMap<u16, Channel<'a>>,
  pub configuration:    Configuration,
  pub channel_index:    u16,
  pub prefetch_size:    u32,
  pub prefetch_count:   u16,
  pub frame_queue:      VecDeque<Frame>,
}

impl<'a> Connection<'a> {
  pub fn new() -> Connection<'a> {
    let mut h = HashMap::new();
    h.insert(0, Channel::global());

    let configuration = Configuration {
      channel_max: 0,
      frame_max:   8192,
      heartbeat:   60,
    };

    Connection {
      state:          ConnectionState::Initial,
      channels:       h,
      configuration:  configuration,
      channel_index:  1,
      prefetch_size:  0,
      prefetch_count: 0,
      frame_queue:    VecDeque::new(),
    }
  }

  pub fn create_channel(&mut self) -> u16 {
    let c  = Channel::new(self.channel_index);
    self.channels.insert(self.channel_index, c);
    self.channel_index += 1;

    self.channel_index - 1
  }

  pub fn set_channel_state(&mut self, channel_id: u16, new_state: ChannelState) {
    self.channels.get_mut(&channel_id).map(|c| c.state = new_state);
  }

  pub fn check_state(&self, channel_id: u16, state: ChannelState) -> Option<bool> {
    self.channels
          .get(&channel_id)
          .map(|c| c.state == state)
  }

  pub fn get_state(&self, channel_id: u16) -> Option<ChannelState> {
    self.channels
          .get(&channel_id)
          .map(|c| c.state.clone())
  }

  pub fn push_back_answer(&mut self, channel_id: u16, answer: Answer) {
    self.channels
      .get_mut(&channel_id)
      .map(|c| c.awaiting.push_back(answer));
  }

  pub fn check_next_answer(&self, channel_id: u16, answer: Answer) -> bool {
    self.channels
          .get(&channel_id)
          .map(|c| c.awaiting.front() == Some(&answer)).unwrap_or(false)
  }

  pub fn get_next_answer(&mut self, channel_id: u16) -> Option<Answer> {
    self.channels
          .get_mut(&channel_id)
          .and_then(|c| c.awaiting.pop_front())
  }

  pub fn is_connected(&self, channel_id: u16) -> bool {
    self.channels
          .get(&channel_id)
          .map(|c| c.is_connected()).unwrap_or(false)
  }

  pub fn connect(&mut self) -> Result<ConnectionState> {
    if self.state != ConnectionState::Initial {
      self.state = ConnectionState::Error;
      return Err(Error::new(ErrorKind::Other, "invalid state"))
    }

    self.frame_queue.push_back(Frame::ProtocolHeader);
    self.state = ConnectionState::Connecting(ConnectingState::SentProtocolHeader);
    Ok(self.state)
  }

  pub fn next_frame(&mut self) -> Option<Frame> {
    self.frame_queue.pop_front()
  }

  pub fn serialize(&mut self, send_buffer: &mut [u8]) -> Result<(usize, ConnectionState)> {
    let next_msg = self.frame_queue.pop_front();
    if next_msg == None {
      return Err(Error::new(ErrorKind::WouldBlock, "no new message"));
    }

    let next_msg = next_msg.unwrap();
    println!("will write to buffer: {:?}", next_msg);

    let gen_res = match &next_msg {
      &Frame::ProtocolHeader => {
        gen_protocol_header((send_buffer, 0)).map(|tup| tup.1)
      },
      &Frame::Heartbeat(_) => {
        gen_heartbeat_frame((send_buffer, 0)).map(|tup| tup.1)
      },
      &Frame::Method(channel, ref method) => {
        gen_method_frame((send_buffer, 0), channel, method).map(|tup| tup.1)
      },
      &Frame::Header(channel_id, class_id, ref header) => {
        gen_content_header_frame((send_buffer, 0), channel_id, class_id, header.body_size).map(|tup| tup.1)
      },
      &Frame::Body(channel_id, ref data) => {
        gen_content_body_frame((send_buffer, 0), channel_id, data).map(|tup| tup.1)
      }
    };

    match gen_res {
      Ok(sz) => {
        Ok((sz, self.state))
      },
      Err(e) => {
        println!("error generating frame: {:?}", e);
        self.state = ConnectionState::Error;
        match e {
          GenError::BufferTooSmall(_) => {
            self.frame_queue.push_front(next_msg);
            return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
          },
          GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
            return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
          }
        }
      }
    }
  }

  pub fn parse(&mut self, data: &[u8]) -> Result<(usize,ConnectionState)> {
    let parsed_frame = frame(data);
    match parsed_frame {
      IResult::Done(i,_)     => {},
      IResult::Incomplete(_) => {
        return Ok((0,self.state));
      },
      IResult::Error(e) => {
        //FIXME: should probably disconnect on error here
        let err = format!("parse error: {:?}", e);
        self.state = ConnectionState::Error;
        return Err(Error::new(ErrorKind::Other, err))
      }
    }

    let (i, f) = parsed_frame.unwrap();

    //FIXME: what happens if we fail to parse a packet in a channel?
    // do we continue?
    let consumed = data.offset(i);

    self.handle_frame(f);

    return Ok((consumed, self.state));
  }

  pub fn handle_frame(&mut self, f: Frame) {
    match f {
      Frame::ProtocolHeader => {
        println!("error: the client should not receive a protocol header");
        self.state = ConnectionState::Error;
      },
      Frame::Method(channel_id, method) => {
        if channel_id == 0 {
          self.handle_global_method(method);
        } else {
          self.receive_method(channel_id, method);
        }
      },
      Frame::Heartbeat(channel_id) => {
        self.frame_queue.push_back(Frame::Heartbeat(0));
      },
      Frame::Header(channel_id, class_id, header) => {
        self.handle_content_header_frame(channel_id, header.body_size);
      },
      Frame::Body(channel_id, payload) => {
        self.handle_body_frame(channel_id, payload);
      }
    };
  }

  pub fn handle_global_method(&mut self, c: Class) {
    match self.state {
      ConnectionState::Initial | ConnectionState::Closed | ConnectionState::Error => {
        self.state = ConnectionState::Error
      },
      ConnectionState::Connecting(connecting_state) => {
        match connecting_state {
          ConnectingState::Initial | ConnectingState::Error => {
            self.state = ConnectionState::Error
          },
          ConnectingState::SentProtocolHeader => {
            if let Class::Connection(connection::Methods::Start(s)) = c {
              println!("Server sent Connection::Start: {:?}", s);
              self.state = ConnectionState::Connecting(ConnectingState::ReceivedStart);

              let mut h = HashMap::new();
              h.insert("product".to_string(), AMQPValue::LongString("lapin".to_string()));

              let creds = Credentials {
                username: Some("guest".to_owned()),
                secret: Secret::Password("guest".to_owned()),
                channel_binding: ChannelBinding::None,
              };

              let mut mechanism = Plain::from_credentials(creds).unwrap();

              let initial_data = mechanism.initial().unwrap();
              let s = str::from_utf8(&initial_data).unwrap();

              //FIXME: fill with user configured data
              //we need to handle the server properties, and have some client properties
              let start_ok = Class::Connection(connection::Methods::StartOk(
                connection::StartOk {
                  client_properties: h,
                  mechanism: "PLAIN".to_string(),
                  locale:    "en_US".to_string(), // FIXME: comes from the server
                  response:  s.to_string(),     //FIXME: implement SASL
                }
              ));

              println!("client sending Connection::StartOk: {:?}", start_ok);
              self.frame_queue.push_back(Frame::Method(0, start_ok));
              self.state = ConnectionState::Connecting(ConnectingState::SentStartOk);
            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          /*ConnectingState::ReceivedStart => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },*/
          ConnectingState::SentStartOk => {
            if let Class::Connection(connection::Methods::Tune(t)) = c {
              println!("Server sent Connection::Tune: {:?}", t);
              self.state = ConnectionState::Connecting(ConnectingState::ReceivedTune);

              self.configuration.channel_max = t.channel_max;
              self.configuration.heartbeat   = t.heartbeat;

              if t.frame_max > self.configuration.frame_max {
                //FIXME: what do we do without the buffers available?
                //self.send_buffer.grow(t.frame_max as usize);
                //self.receive_buffer.grow(t.frame_max as usize);
              }

              self.configuration.frame_max = t.frame_max;

              let tune_ok = Class::Connection(connection::Methods::TuneOk(
                connection::TuneOk {
                  channel_max : self.configuration.channel_max,
                  frame_max   : self.configuration.frame_max,
                  heartbeat   : self.configuration.heartbeat,
                }
              ));

              println!("client sending Connection::TuneOk: {:?}", tune_ok);

              self.frame_queue.push_back(Frame::Method(0, tune_ok));
              self.state = ConnectionState::Connecting(ConnectingState::SentTuneOk);

              let open = Class::Connection(connection::Methods::Open(
                  connection::Open {
                    virtual_host: "/".to_string(),
                    capabilities: "".to_string(),
                    insist:       false,
                  }
                  ));

              println!("client sending Connection::Open: {:?}", open);
              self.frame_queue.push_back(Frame::Method(0,open));
              self.state = ConnectionState::Connecting(ConnectingState::SentOpen);

            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          ConnectingState::ReceivedSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::SentSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::ReceivedSecondSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::ReceivedTune => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::SentOpen => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
            if let Class::Connection(connection::Methods::OpenOk(o)) = c {
              println!("Server sent Connection::OpenOk: {:?}, client now connected", o);
              self.state = ConnectionState::Connected;
            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          ConnectingState::Error => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          s => {
            println!("invalid state {:?}", s);
            self.state = ConnectionState::Error;
          }
        }
      },
      ConnectionState::Connected => {},
      ConnectionState::Closing(ClosingState) => {},
    };
  }

  pub fn handle_content_header_frame(&mut self, channel_id: u16, size: u64) {
    let state = self.channels.get_mut(&channel_id).map(|channel| {
      channel.state.clone()
    }).unwrap();
    if let ChannelState::WillReceiveContent(queue_name, consumer_tag) = state {
      self.set_channel_state(channel_id, ChannelState::ReceivingContent(queue_name, consumer_tag.clone(), size as usize));
    } else {
      self.set_channel_state(channel_id, ChannelState::Error);
    }
  }

  pub fn handle_body_frame(&mut self, channel_id: u16, payload: Vec<u8>) {
    let state = self.channels.get_mut(&channel_id).map(|channel| {
      channel.state.clone()
    }).unwrap();

    let payload_size = payload.len();

    if let ChannelState::ReceivingContent(queue_name, consumer_tag, remaining_size) = state {
      if remaining_size >= payload_size {

        if let Some(ref mut c) = self.channels.get_mut(&channel_id) {
          if let Some(ref mut q) = c.queues.get_mut(&queue_name) {
            if let Some(ref mut cs) = q.consumers.get_mut(&consumer_tag) {
              (*cs.callback).receive_content(payload);
              if remaining_size == payload_size {
                (*cs.callback).done();
              }
            }
          }
        }

        if remaining_size == payload_size {
          self.set_channel_state(channel_id, ChannelState::Connected);
        } else {
          self.set_channel_state(channel_id, ChannelState::ReceivingContent(queue_name, consumer_tag, remaining_size - payload_size));
        }
      } else {
        println!("body frame too large");
        self.set_channel_state(channel_id, ChannelState::Error);
      }
    } else {
      self.set_channel_state(channel_id, ChannelState::Error);
    }
  }

  //FIXME: no chunking for now
  pub fn send_content_frames(&mut self, channel_id: u16, class_id: u16, slice: &[u8]) {
    let header = ContentHeader {
      class_id:       class_id,
      weight:         0,
      body_size:      slice.len() as u64,
      property_flags: 0x2000,
      property_list:  HashMap::new(),
    };
    self.frame_queue.push_back(Frame::Header(channel_id, class_id, header));
    self.frame_queue.push_back(Frame::Body(channel_id,Vec::from(slice)));
  }

  pub fn send_method_frame(&mut self, channel: u16, method: Class) -> result::Result<(), error::Error> {
    self.frame_queue.push_back(Frame::Method(channel,method));
    Ok(())
  }
}
<|endoftext|>


<commit_before>extern mod std;

#[deriving(Eq)]
#[deriving(Clone)]
pub enum Value {
	List(~[Value]),
	Atom(~str),
	Str(~str),
	Num(float)
}

#[inline]
fn isws(c: char) -> bool {
	c == ' ' || c == '\t' || c == '\n'
}

/// Determines if the input begins a number
#[inline]
fn beginning_is_num(input: &str, i: uint) -> bool {
	let (c, c2) = (input[i] as char, input[i+1] as char);
	if c.is_digit() || c == '.' { true }
	else if c == '-' && (c2.is_digit() || c2 == '.') { true }
	else { false }
}

/// Reads a character after skipping whitespace
fn read_nows(input: &str, start: uint) -> Option<char> {
	input.iter().skip(start).skip_while(|&c| isws(c)).next()
}

/// Reads a numeric literal (integers or floats)
fn read_number(input: &str, start: uint) -> Option<(~Value, uint)> {
	let end = input.iter().skip(start).position(|c| isws(c) || c == ')');
	let pos = match end {
		Some(end) => start+end,
		None => input.len()
	};
	let acc = input.slice(start, pos);
	match std::float::from_str(acc) {
		Some(f) => Some((~Num(f), pos)),
		None => None
	}
}

/// Reads a string literal
fn read_string(input: &str, start: uint) -> Option<(~Value, uint)> {
	let mut acc = ~"";
	let mut was_escape = false;
	for i in range(start+1, input.len()) {
		match(input[i] as char) {
			'"' if !was_escape =>
				return Some((~Str(acc), i+1)), // got closing quote
			'\\' if !was_escape => {
				// escape codes
				was_escape = true;
				match input[i+1] as char {
					'\\' => acc.push_char('\\'),
					'"' => acc.push_char('"'),
					't' => acc.push_char('\t'),
					'n' => acc.push_char('\n'),
					c => fail!("unknown escape code: \\%c", c)
				}
			}
			c if !was_escape => acc.push_char(c),
			_ => { was_escape = false }
		}
	}

	// we didn't hit the end " so there must be a parse error
	None
}

/// Reads a list
fn read_list(input: &str, start: uint) -> Option<(~Value, uint)> {
	let mut list = ~[];
	let mut i = start;
	loop {
		match read_value(input, i) {
			Some((v,new_i)) => {
				list.push(*v);
				i = new_i;
			}
			None if read_nows(input, i) == Some(')') => return Some((~List(list), i)),
			None => return None
		}
	}
}

/// Reads an atom (symbol/identifier)
fn read_atom(input: &str, start: uint) -> Option<(~Value, uint)> {
	let atom = std::str::from_chars(input.iter().skip(start).take_while(|&c| !isws(c) && c != ')').to_owned_vec());
	let len = atom.len();
	Some((~Atom(atom), start + len))
}

/// Reads any S-expression value
fn read_value(input: &str, start: uint) -> Option<(~Value, uint)> {
	if(start >= input.len()) {
		return None
	}

	match input[start] as char {
		c if isws(c) => read_value(input, start+1), // ignore whitespace
		'(' => read_list(input, start+1),
		')' => None,
		'"' => read_string(input, start),
		_ if beginning_is_num(input, start) => read_number(input, start),
		_ => read_atom(input, start)
	}
}

/// Public function for parsing a value from a string
pub fn from_str(input: &str) -> Option<Value> {
	match read_value(input, 0) {
		Some((v,_)) => Some(*v),
		None => None
	}
}

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

	fn parse_some(input: &str, expr: Value) {
		match from_str(input) {
			Some(ref p) if p == &expr => {}
			Some(ref p) => {
				if(p != &expr) {
					fail!("Expected %?, not %?, for %?", expr, *p, input)
				}
			}
			None => fail!("Got None for %?", input)
		}
	}

	#[test]
	fn test_sexpr() {
		assert_eq!(from_str(""), None);
		parse_some("123", Num(123.0));
		parse_some("-123", Num(-123.0));
		parse_some("()", List(~[]));
		parse_some("3.14159265358", Num(3.14159265358));
		parse_some("(hi)", List(~[Atom(~"hi")]));
		parse_some("(-hi)", List(~[Atom(~"-hi")]));
		parse_some("(-)", List(~[Atom(~"-")]));
		parse_some("(1)", List(~[Num(1.0)]));
		parse_some("(hi there)", List(~[Atom(~"hi"), Atom(~"there")]));
		parse_some("(hi (there))", List(~[Atom(~"hi"), List(~[Atom(~"there")])]));
		parse_some("(hi 123456)", List(~[Atom(~"hi"), Num(123456.0)]));
		parse_some("(())", List(~[List(~[])]));
		parse_some("(hi (there (fellow (human-bot!))))",
			List(~[Atom(~"hi"),
				List(~[Atom(~"there"),
					List(~[Atom(~"fellow"),
						List(~[Atom(~"human-bot!")])])])])
		);
	}

	#[test]
	fn test_string() {
		parse_some("\"\"", Str(~""));
		parse_some("\"string\"", Str(~"string"));
		parse_some("\"hi\tthere\n\"", Str(~"hi\tthere\n"));
		parse_some("\"a\\\\b\"", Str(~"a\\b"));
		parse_some("(\"hello\" \"world\")", List(~[Str(~"hello"), Str(~"world")]));
	}

	#[test]
	fn test_whitespace() {
		// spaces
		parse_some("  (  1  )  ", List(~[Num(1.0)]));
		parse_some("  (  (  )  )  )", List(~[List(~[])]));

		// newlines
		parse_some("\n\n( \n\n )\n\n", List(~[]));
		parse_some("\n\n(\n\n123\n\n)\n\n", List(~[Num(123.0)]));
	}
}<commit_msg>sexpr: add tests for lists in the middle of lists; fix the parser to parse them<commit_after>extern mod std;

#[deriving(Eq)]
#[deriving(Clone)]
pub enum Value {
	List(~[Value]),
	Atom(~str),
	Str(~str),
	Num(float)
}

#[inline]
fn isws(c: char) -> bool {
	c == ' ' || c == '\t' || c == '\n'
}

/// Determines if the input begins a number
#[inline]
fn beginning_is_num(input: &str, i: uint) -> bool {
	let (c, c2) = (input[i] as char, input[i+1] as char);
	if c.is_digit() || c == '.' { true }
	else if c == '-' && (c2.is_digit() || c2 == '.') { true }
	else { false }
}

/// Reads a character after skipping whitespace
fn read_nows(input: &str, start: uint) -> Option<(uint, char)> {
	input.iter().enumerate().skip(start).skip_while(|&(_,c)| isws(c)).next()
}

/// Reads a numeric literal (integers or floats)
fn read_number(input: &str, start: uint) -> Option<(~Value, uint)> {
	let end = input.iter().skip(start).position(|c| isws(c) || c == ')');
	let pos = match end {
		Some(end) => start+end,
		None => input.len()
	};
	let acc = input.slice(start, pos);
	match std::float::from_str(acc) {
		Some(f) => Some((~Num(f), pos)),
		None => None
	}
}

/// Reads a string literal
fn read_string(input: &str, start: uint) -> Option<(~Value, uint)> {
	let mut acc = ~"";
	let mut was_escape = false;
	for i in range(start+1, input.len()) {
		match(input[i] as char) {
			'"' if !was_escape =>
				return Some((~Str(acc), i+1)), // got closing quote
			'\\' if !was_escape => {
				// escape codes
				was_escape = true;
				match input[i+1] as char {
					'\\' => acc.push_char('\\'),
					'"' => acc.push_char('"'),
					't' => acc.push_char('\t'),
					'n' => acc.push_char('\n'),
					c => fail!("unknown escape code: \\%c", c)
				}
			}
			c if !was_escape => acc.push_char(c),
			_ => { was_escape = false }
		}
	}

	// we didn't hit the end " so there must be a parse error
	None
}

/// Reads a list
fn read_list(input: &str, start: uint) -> Option<(~Value, uint)> {
	let mut list = ~[];
	let mut i = start;
	loop {
		match read_value(input, i) {
			Some((v,new_i)) => {
				list.push(*v);
				i = new_i;
			}
			None => {
				// probably a terminating ')' - if so, return the list and seek past it
				match read_nows(input, i) {
					Some((new_i, ')')) => return Some((~List(list), new_i+1)),
					_ => return None
				}
			}
		}
	}
}

/// Reads an atom (symbol/identifier)
fn read_atom(input: &str, start: uint) -> Option<(~Value, uint)> {
	let atom = std::str::from_chars(input.iter().skip(start).take_while(|&c| !isws(c) && c != ')').to_owned_vec());
	let len = atom.len();
	Some((~Atom(atom), start + len))
}

/// Reads any S-expression value
fn read_value(input: &str, start: uint) -> Option<(~Value, uint)> {
	if(start >= input.len()) {
		return None
	}

	match input[start] as char {
		c if isws(c) => read_value(input, start+1), // ignore whitespace
		'(' => read_list(input, start+1),
		')' => None,
		'"' => read_string(input, start),
		_ if beginning_is_num(input, start) => read_number(input, start),
		_ => read_atom(input, start)
	}
}

/// Public function for parsing a value from a string
pub fn from_str(input: &str) -> Option<Value> {
	match read_value(input, 0) {
		Some((v,_)) => Some(*v),
		None => None
	}
}

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

	fn parse_some(input: &str, expr: Value) {
		match from_str(input) {
			Some(ref p) if p == &expr => {}
			Some(ref p) => {
				if(p != &expr) {
					fail!("Expected %?, not %?, for %?", expr, *p, input)
				}
			}
			None => fail!("Got None for %?", input)
		}
	}

	#[test]
	fn test_sexpr() {
		assert_eq!(from_str(""), None);
		parse_some("123", Num(123.0));
		parse_some("-123", Num(-123.0));
		parse_some("()", List(~[]));
		parse_some("3.14159265358", Num(3.14159265358));
		parse_some("(hi)", List(~[Atom(~"hi")]));
		parse_some("(-hi)", List(~[Atom(~"-hi")]));
		parse_some("(-)", List(~[Atom(~"-")]));
		parse_some("(1)", List(~[Num(1.0)]));
		parse_some("(hi there)", List(~[Atom(~"hi"), Atom(~"there")]));
		parse_some("(hi (there))", List(~[Atom(~"hi"), List(~[Atom(~"there")])]));
		parse_some("(1 (2) 3)", List(~[Num(1f), List(~[Num(2f)]), Num(3f)]));
		parse_some("(hi (there) you!)", List(~[Atom(~"hi"), List(~[Atom(~"there")]), Atom(~"you!")]));
		parse_some("(hi 123456)", List(~[Atom(~"hi"), Num(123456.0)]));
		parse_some("(())", List(~[List(~[])]));
		parse_some("(hi (there (fellow (human-bot!))))",
			List(~[Atom(~"hi"),
				List(~[Atom(~"there"),
					List(~[Atom(~"fellow"),
						List(~[Atom(~"human-bot!")])])])])
		);
	}

	#[test]
	fn test_string() {
		parse_some("\"\"", Str(~""));
		parse_some("\"string\"", Str(~"string"));
		parse_some("\"hi\tthere\n\"", Str(~"hi\tthere\n"));
		parse_some("\"a\\\\b\"", Str(~"a\\b"));
		parse_some("(\"hello\" \"world\")", List(~[Str(~"hello"), Str(~"world")]));
	}

	#[test]
	fn test_whitespace() {
		// spaces
		parse_some("  (  1  )  ", List(~[Num(1.0)]));
		parse_some("  (  (  )  )  )", List(~[List(~[])]));

		// newlines
		parse_some("\n\n( \n\n )\n\n", List(~[]));
		parse_some("\n\n(\n\n123\n\n)\n\n", List(~[Num(123.0)]));
	}
}<|endoftext|>


<commit_before>use std::io::stdio;
use std::io::Process;
use std::io::process;
use std::path::posix::Path;

struct Shell {
    cmd_prompt: ~str,
    cwd: Path
}

impl Shell {
    fn new(prompt_str: &str, cwd: Path) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
            cwd: cwd
        }
    }

    fn start(&mut self) {
        let mut stdin = stdio::stdin();
        loop {
            stdio::print(self.cmd_prompt);
            stdio::flush();

            let line = stdin.read_line().unwrap();
            let cmd_line = line.trim().to_owned();
            let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");

            match program {
                ""      =>  { continue; }
                "\u274c"  =>  { return; }
                _       =>  { self.cmd(cmd_line.clone()); }
                // _ => { println!("You just ran {}", cmd_line) }
            }
        }
    }

    fn cmd(&mut self, cmd_line: ~str) {
        let mut arg_vec = ~[];
        let mut cmd = ~"";
        for (index, s) in cmd_line.split_str(" ").enumerate() {
            if s == "" { continue; }
            if index == 0 { cmd = s.to_owned(); } else { arg_vec.push(s.to_owned()); }
        }
        self.run_cmd(cmd, arg_vec);
    }

    fn run_cmd(&mut self, cmd: ~str, argv: ~[~str]) {
        match cmd.as_slice() {
            &"\U0001F697" => self.cd(argv),
            &"\U0001F4CD" => self.pwd(),
            _ => self.execute_program(cmd, argv)
        }
    }

    fn execute_program(&mut self, cmd: ~str, argv: ~[~str]) {
        if self.cmd_exists(cmd) {
            println!("\U0001F3C3 {}", cmd);
            let config = process::ProcessConfig {
                program: cmd,
                args: argv,
                cwd: Some(&self.cwd),
                stdin: process::InheritFd(0),
                stdout: process::InheritFd(1),
                stderr: process::InheritFd(2),
                .. process::ProcessConfig::new()
            };

            let p = Process::configure(config);
            let status = p.unwrap().wait();
            match status {
                process::ExitSignal(st) => println!("Exited with signal {}", st),
                process::ExitStatus(st) => println!("Exited with status {}", st)
            }
        }
        else {
            println!("Command {} not found.", cmd)
        }
    }

    fn cd(&mut self, path: ~[~str]) {
        if path.len() == 0 {
            println!("Please specify a path to \U0001F697  to.");
            return;
        }
        let p = Path::new(path[0]);
        if !p.is_dir() {
            println!("\U0001F697  could not find that, \U0001F62D  ")
        }
        self.cwd = p;
    }

    fn pwd(&mut self) {
        println!("\u27A1 {} \u2B05  \U0001F4CD", self.cwd.as_str().unwrap());
    }

    fn cmd_exists(&mut self, cmd: &str) -> bool {
        let p = Process::new("which", [cmd.to_owned()]);
        let status = p.unwrap().wait();
        match status {
            process::ExitSignal(st) => st == 0,
            process::ExitStatus(st) => st == 0
        }
    }
}

fn main() {
    Shell::new("\U0001f41a  ", std::os::getcwd()).start();
}
<commit_msg>Change around some emojis, trim inputs, tomatoes to separate parameters<commit_after>use std::io::stdio;
use std::io::Process;
use std::io::process;
use std::path::posix::Path;

struct Shell {
    cmd_prompt: ~str,
    cwd: Path
}

impl Shell {
    fn new(prompt_str: &str, cwd: Path) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
            cwd: cwd
        }
    }

    fn start(&mut self) {
        let mut stdin = stdio::stdin();
        loop {
            stdio::print(self.cmd_prompt);
            stdio::flush();

            let line = stdin.read_line().unwrap();
            let cmd_line = line.trim().to_owned();
            let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");

            match program {
                ""      =>  { continue; }
                "\u274c"  =>  { return; }
                _       =>  { self.cmd(cmd_line.clone()); }
                // _ => { println!("You just ran {}", cmd_line) }
            }
        }
    }

    fn cmd(&mut self, cmd_line: ~str) {
        let mut arg_vec = ~[];
        let mut cmd = ~"";
        for (index, raw_s) in cmd_line.split_str("\U0001F345").enumerate() {
            let s = raw_s.trim();
            if s == "" { continue; }
            if index == 0 { cmd = s.to_owned(); } else { arg_vec.push(s.to_owned()); }
        }
        self.run_cmd(cmd, arg_vec);
    }

    fn run_cmd(&mut self, cmd: ~str, argv: ~[~str]) {
        match cmd.as_slice() {
            &"\U0001F697" => self.cd(argv),
            &"\U0001F4CD" => self.pwd(),
            _ => self.execute_program(cmd, argv)
        }
    }

    fn execute_program(&mut self, cmd: ~str, argv: ~[~str]) {
        if self.cmd_exists(cmd) {
            println!("\U0001F3C3  {}", cmd);
            let config = process::ProcessConfig {
                program: cmd,
                args: argv,
                cwd: Some(&self.cwd),
                stdin: process::InheritFd(0),
                stdout: process::InheritFd(1),
                stderr: process::InheritFd(2),
                .. process::ProcessConfig::new()
            };

            let p = Process::configure(config);
            let status = p.unwrap().wait();
            match status {
                process::ExitSignal(st) => println!("\U0001F6B7  {} \u27A1  {}", cmd, st),
                process::ExitStatus(st) => println!("\U0001F6B7  {} \u27A1  {}", cmd, st)
            }
        }
        else {
            println!("Command {} not found.", cmd)
        }
    }

    fn cd(&mut self, path: ~[~str]) {
        if path.len() == 0 {
            println!("Please specify a path to \U0001F697  to.");
            return;
        }
        let p = Path::new(path[0]);
        if !p.is_dir() {
            println!("\U0001F697  could not find that, \U0001F62D  ")
        }
        self.cwd = p;
    }

    fn pwd(&mut self) {
        println!("\u27A1 {} \u2B05  \U0001F4CD", self.cwd.as_str().unwrap());
    }

    fn cmd_exists(&mut self, cmd: &str) -> bool {
        let p = Process::new("which", [cmd.to_owned()]);
        let status = p.unwrap().wait();
        match status {
            process::ExitSignal(st) => st == 0,
            process::ExitStatus(st) => st == 0
        }
    }
}

fn main() {
    Shell::new("\U0001f41a  ", std::os::getcwd()).start();
}
<|endoftext|>