<commit_before>pub use self::tokens::{Doctype, Attributes, TagKind, StartTag, EndTag, Tag, Token};
pub use self::tokens::{DoctypeToken, TagToken, CommentToken, CharacterToken};

use self::states::{RawLessThanSign, RawEndTagOpen, RawEndTagName};
use self::states::{Rcdata, ScriptData, ScriptDataEscaped};

mod tokens;
mod states;

pub trait TokenSink {
    fn process_token(&mut self, token: Token);
}

fn ascii_letter(c: char) -> Option<char> {
    c.to_ascii_opt()
        .filtered(|a| a.is_alpha())
        .map(|a| a.to_lower().to_char())
}

pub struct Tokenizer<'sink, Sink> {
    priv sink: &'sink mut Sink,
    priv state: states::State,

    // FIXME: The state machine guarantees the tag exists when
    // we need it, so we could eliminate the Option overhead.
    // Leaving it as Option for now, to find bugs.
    priv current_tag: Option<Tag>,

    /// Last start tag name, for use in checking "appropriate end tag".
    priv last_start_tag_name: Option<~str>,

    /// The "temporary buffer" mentioned in the spec.
    priv temp_buf: ~str
}

#[deriving(Eq)]
enum ConsumeCharResult {
    Reconsume,
    Finished,
}

impl<'sink, Sink: TokenSink> Tokenizer<'sink, Sink> {
    pub fn new(sink: &'sink mut Sink) -> Tokenizer<'sink, Sink> {
        Tokenizer {
            sink: sink,
            state: states::Data,
            current_tag: None,
            last_start_tag_name: None,
            temp_buf: ~"",
        }
    }

    pub fn feed(&mut self, input: &str) {
        debug!("feeding {:s}", input);
        let mut it = input.chars(); //.peekable();
        loop {
            match self.state {
                // These states do something other than consume a single character.
                states::CharacterReferenceInData | states::CharacterReferenceInRcdata
                | states::CharacterReferenceInAttributeValue | states::BogusComment
                | states::MarkupDeclarationOpen | states::CdataSection => {
                    fail!("FIXME: state {:?} not implemented", self.state);
                }

                _ => match it.next() {
                    None => return,
                    Some(c) => {
                        while self.process_char(c) == Reconsume {
                            // FIXME: this is not correct when state changes to one
                            // of the above!
                        }
                    }
                }
            }
        }
    }

    fn parse_error(&self, c: char) {
        error!("Parse error: saw {:?} in state {:?}", c, self.state);
    }

    fn emit_char(&self, c: char) {
        self.sink.process_token(CharacterToken(c));
    }

    fn emit_current_tag(&mut self) {
        let tag = self.current_tag.take().unwrap();
        match tag.kind {
            StartTag => self.last_start_tag_name = Some(tag.name.clone()),
            _ => ()
        }
        self.sink.process_token(TagToken(tag));
    }

    fn emit_temp_buf(&mut self) {
        // FIXME: Add a multi-character token and move temp_buf into it, something like
        //     self.sink.process_token(StringToken(
        //         replace(&mut self.temp_buf, ~"")))
        //
        // Need to make sure that clearing on emit is spec-compatible.
        //
        // Until then we reuse the same buffer allocation forever.
        for c in self.temp_buf.chars() {
            self.emit_char(c);
        }
    }

    fn clear_temp_buf(&mut self) {
        // Do this without a new allocation.
        self.temp_buf.truncate(0);
    }

    fn create_tag(&mut self, kind: TagKind, c: char) {
        assert!(self.current_tag.is_none());
        let mut t = Tag::new(kind);
        t.name.push_char(c);
        self.current_tag = Some(t);
    }

    fn append_to_tag_name(&mut self, c: char) {
        self.current_tag.get_mut_ref().name.push_char(c);
    }

    fn have_appropriate_end_tag(&self) -> bool {
        match (self.last_start_tag_name.as_ref(), self.current_tag.as_ref()) {
            (Some(last), Some(tag)) =>
                (tag.kind == EndTag) && (tag.name.as_slice() == last.as_slice()),
            _ => false
        }
    }
}

// Shorthand for common state machine behaviors.
macro_rules! shorthand (
    ( to          $state:ident             ) => ( self.state = states::$state;        );
    ( to_raw      $state:ident $kind:ident ) => ( self.state = states::$state($kind); );
    ( emit        $c:expr                  ) => ( self.emit_char($c);                 );
    ( create_tag  $kind:expr   $c:expr     ) => ( self.create_tag($kind, $c);         );
    ( append_tag  $c:expr                  ) => ( self.append_to_tag_name($c);        );
    ( emit_tag                             ) => ( self.emit_current_tag();            );
    ( append_temp $c:expr                  ) => ( self.temp_buf.push_char($c);        );
    ( emit_temp                            ) => ( self.emit_temp_buf();               );
    ( clear_temp                           ) => ( self.clear_temp_buf();              );

    // NB: Deliberate capture of 'c'
    ( error                                ) => ( self.parse_error(c);                );
)

// A little DSL for sequencing shorthand actions.
macro_rules! go (
    // A pattern like $($cmd:tt)* ; $($rest:tt)* causes parse ambiguity.
    // We have to tell the parser how much lookahead we need.
    ( $a:tt             ; $($rest:tt)* ) => ({ shorthand!($a);       go!($($rest)*); });
    ( $a:tt $b:tt       ; $($rest:tt)* ) => ({ shorthand!($a $b);    go!($($rest)*); });
    ( $a:tt $b:tt $c:tt ; $($rest:tt)* ) => ({ shorthand!($a $b $c); go!($($rest)*); });

    // These can only come at the end.
    // FIXME: Come up with a better name for 'finish'.
    ( reconsume ) => ( return Reconsume; );
    ( finish    ) => ( return Finished;  );

    // If nothing else matched, it's a single command
    ( $($cmd:tt)+ ) => ( shorthand!($($cmd)+); );

    // or nothing.
    () => (());
)

impl<'sink, Sink: TokenSink> Tokenizer<'sink, Sink> {
    // FIXME: explicitly represent the EOF character?
    // For now the plan is to handle EOF in a separate match.
    fn process_char(&mut self, c: char) -> ConsumeCharResult {
        debug!("Processing {:?} in state {:?}", c, self.state);
        match self.state {
            states::Data => match c {
                '&'  => go!(to CharacterReferenceInData),
                '<'  => go!(to TagOpen),
                '\0' => go!(error; emit '\0'),
                _    => go!(emit c),
            },

            // RCDATA, RAWTEXT, script, or script escaped
            states::RawData(kind) => match c {
                '&' if kind == Rcdata            => go!(to CharacterReferenceInRcdata),
                '-' if kind == ScriptDataEscaped => go!(to ScriptDataEscapedDash; emit '-'),
                '<'  => go!(to_raw RawLessThanSign kind),
                '\0' => go!(error; emit '\ufffd'),
                _    => go!(emit c),
            },

            states::Plaintext => match c {
                '\0' => go!(error; emit '\ufffd'),
                _    => go!(emit c),
            },

            states::TagOpen => match c {
                '!' => go!(to MarkupDeclarationOpen),
                '/' => go!(to EndTagOpen),
                '?' => go!(error; to BogusComment),
                _ => match ascii_letter(c) {
                    Some(cl) => go!(create_tag StartTag cl; to TagName),
                    None     => go!(error; emit '<'; to Data; reconsume),
                }
            },

            states::EndTagOpen => match c {
                '>' => go!(error; to Data),
                _ => match ascii_letter(c) {
                    Some(cl) => go!(create_tag EndTag cl; to TagName),
                    None     => go!(error; to BogusComment),
                }
            },

            states::TagName => match c {
                '\t' | '\n' | '\x0C' | ' '
                     => go!(to BeforeAttributeName),
                '/'  => go!(to SelfClosingStartTag),
                '>'  => go!(emit_tag; to Data),
                '\0' => go!(error; append_tag '\ufffd'),
                _    => go!(append_tag (ascii_letter(c).unwrap_or(c))),
            },

            // kind = ScriptDataEscaped
            states::RawLessThanSign(ScriptDataEscaped) => match c {
                '/' => go!(clear_temp; to_raw RawEndTagOpen ScriptDataEscaped),
                _ => match ascii_letter(c) {
                    Some(cl) => go!(clear_temp; append_temp cl;
                                    to ScriptDataDoubleEscapeStart; emit '<'; emit c),
                    None => go!(to_raw RawData ScriptDataEscaped; emit '<'; reconsume),
                }
            },

            // otherwise
            states::RawLessThanSign(kind) => match c {
                '/' => go!(clear_temp; to_raw RawEndTagOpen kind),
                '!' if kind == ScriptData => go!(to ScriptDataEscapeStart; emit '<'; emit '!'),
                _   => go!(to_raw RawData Rcdata; emit '<'; reconsume),
            },

            states::RawEndTagOpen(kind) => match ascii_letter(c) {
                Some(cl) => go!(create_tag EndTag cl; append_temp c; to_raw RawEndTagName kind),
                None     => go!(to_raw RawData kind; emit '<'; emit '/'; reconsume),
            },

            states::RawEndTagName(kind) => {
                if self.have_appropriate_end_tag() {
                    match c {
                        '\t' | '\n' | '\x0C' | ' '
                            => go!(to BeforeAttributeName; finish),
                        '/' => go!(to SelfClosingStartTag; finish),
                        '>' => go!(emit_tag; to Data; finish),
                        // All of the above end with a return from this function.

                        _ => (),
                    }
                }

                match ascii_letter(c) {
                    Some(cl) => go!(append_tag cl; append_temp c),
                    None     => go!(emit '<'; emit '/'; emit_temp; to_raw RawData kind; reconsume),
                }
            },

            states::ScriptDataEscapeStart => match c {
                '-' => go!(to ScriptDataEscapeStartDash; emit '-'),
                _   => go!(to_raw RawData ScriptData; reconsume),
            },

            states::ScriptDataEscapeStartDash => match c {
                '-' => go!(to ScriptDataEscapedDashDash; emit '-'),
                _   => go!(to_raw RawData ScriptData; reconsume),
            },

            states::ScriptDataEscapedDash => match c {
                '-'  => go!(to ScriptDataEscapedDashDash; emit '-'),
                '<'  => go!(to_raw RawLessThanSign ScriptDataEscaped),
                '\0' => go!(error; to_raw RawData ScriptDataEscaped; emit '\ufffd'),
                _    => go!(to_raw RawData ScriptDataEscaped; emit c),
            },

            states::ScriptDataEscapedDashDash => match c {
                '-'  => go!(emit '-'),
                '<'  => go!(to_raw RawLessThanSign ScriptDataEscaped),
                '>'  => go!(to_raw RawData ScriptData; emit '>'),
                '\0' => go!(error; to_raw RawData ScriptDataEscaped; emit '\ufffd'),
                _    => go!(to_raw RawData ScriptDataEscaped; emit c),
            },

            states::ScriptDataDoubleEscapeStart |
            states::ScriptDataDoubleEscaped |
            states::ScriptDataDoubleEscapedDash |
            states::ScriptDataDoubleEscapedDashDash |
            states::ScriptDataDoubleEscapedLessThanSign |
            states::ScriptDataDoubleEscapeEnd |
            states::BeforeAttributeName |
            states::AttributeName |
            states::AfterAttributeName |
            states::BeforeAttributeValue |
            states::AttributeValueDoubleQuoted |
            states::AttributeValueSingleQuoted |
            states::AttributeValueUnquoted |
            states::AfterAttributeValueQuoted |
            states::SelfClosingStartTag |
            states::CommentStart |
            states::CommentStartDash |
            states::Comment |
            states::CommentEndDash |
            states::CommentEnd |
            states::CommentEndBang |
            states::Doctype |
            states::BeforeDoctypeName |
            states::DoctypeName |
            states::AfterDoctypeName |
            states::AfterDoctypePublicKeyword |
            states::BeforeDoctypePublicIdentifier |
            states::DoctypePublicIdentifierDoubleQuoted |
            states::DoctypePublicIdentifierSingleQuoted |
            states::AfterDoctypePublicIdentifier |
            states::BetweenDoctypePublicAndSystemIdentifiers |
            states::AfterDoctypeSystemKeyword |
            states::BeforeDoctypeSystemIdentifier |
            states::DoctypeSystemIdentifierDoubleQuoted |
            states::DoctypeSystemIdentifierSingleQuoted |
            states::AfterDoctypeSystemIdentifier |
            states::BogusDoctype |
            states::CharacterReferenceInData |
            states::CharacterReferenceInRcdata |
            states::CharacterReferenceInAttributeValue |
            states::BogusComment |
            states::MarkupDeclarationOpen |
            states::CdataSection
                => fail!("FIXME: state {:?} not implemented", self.state),
        }

        Finished

    }
}
<commit_msg>Generalize 'to_raw' and merge it with 'to'<commit_after>pub use self::tokens::{Doctype, Attributes, TagKind, StartTag, EndTag, Tag, Token};
pub use self::tokens::{DoctypeToken, TagToken, CommentToken, CharacterToken};

use self::states::{RawLessThanSign, RawEndTagOpen, RawEndTagName};
use self::states::{Rcdata, ScriptData, ScriptDataEscaped};

mod tokens;
mod states;

pub trait TokenSink {
    fn process_token(&mut self, token: Token);
}

fn ascii_letter(c: char) -> Option<char> {
    c.to_ascii_opt()
        .filtered(|a| a.is_alpha())
        .map(|a| a.to_lower().to_char())
}

pub struct Tokenizer<'sink, Sink> {
    priv sink: &'sink mut Sink,
    priv state: states::State,

    // FIXME: The state machine guarantees the tag exists when
    // we need it, so we could eliminate the Option overhead.
    // Leaving it as Option for now, to find bugs.
    priv current_tag: Option<Tag>,

    /// Last start tag name, for use in checking "appropriate end tag".
    priv last_start_tag_name: Option<~str>,

    /// The "temporary buffer" mentioned in the spec.
    priv temp_buf: ~str
}

#[deriving(Eq)]
enum ConsumeCharResult {
    Reconsume,
    Finished,
}

impl<'sink, Sink: TokenSink> Tokenizer<'sink, Sink> {
    pub fn new(sink: &'sink mut Sink) -> Tokenizer<'sink, Sink> {
        Tokenizer {
            sink: sink,
            state: states::Data,
            current_tag: None,
            last_start_tag_name: None,
            temp_buf: ~"",
        }
    }

    pub fn feed(&mut self, input: &str) {
        debug!("feeding {:s}", input);
        let mut it = input.chars(); //.peekable();
        loop {
            match self.state {
                // These states do something other than consume a single character.
                states::CharacterReferenceInData | states::CharacterReferenceInRcdata
                | states::CharacterReferenceInAttributeValue | states::BogusComment
                | states::MarkupDeclarationOpen | states::CdataSection => {
                    fail!("FIXME: state {:?} not implemented", self.state);
                }

                _ => match it.next() {
                    None => return,
                    Some(c) => {
                        while self.process_char(c) == Reconsume {
                            // FIXME: this is not correct when state changes to one
                            // of the above!
                        }
                    }
                }
            }
        }
    }

    fn parse_error(&self, c: char) {
        error!("Parse error: saw {:?} in state {:?}", c, self.state);
    }

    fn emit_char(&self, c: char) {
        self.sink.process_token(CharacterToken(c));
    }

    fn emit_current_tag(&mut self) {
        let tag = self.current_tag.take().unwrap();
        match tag.kind {
            StartTag => self.last_start_tag_name = Some(tag.name.clone()),
            _ => ()
        }
        self.sink.process_token(TagToken(tag));
    }

    fn emit_temp_buf(&mut self) {
        // FIXME: Add a multi-character token and move temp_buf into it, something like
        //     self.sink.process_token(StringToken(
        //         replace(&mut self.temp_buf, ~"")))
        //
        // Need to make sure that clearing on emit is spec-compatible.
        //
        // Until then we reuse the same buffer allocation forever.
        for c in self.temp_buf.chars() {
            self.emit_char(c);
        }
    }

    fn clear_temp_buf(&mut self) {
        // Do this without a new allocation.
        self.temp_buf.truncate(0);
    }

    fn create_tag(&mut self, kind: TagKind, c: char) {
        assert!(self.current_tag.is_none());
        let mut t = Tag::new(kind);
        t.name.push_char(c);
        self.current_tag = Some(t);
    }

    fn append_to_tag_name(&mut self, c: char) {
        self.current_tag.get_mut_ref().name.push_char(c);
    }

    fn have_appropriate_end_tag(&self) -> bool {
        match (self.last_start_tag_name.as_ref(), self.current_tag.as_ref()) {
            (Some(last), Some(tag)) =>
                (tag.kind == EndTag) && (tag.name.as_slice() == last.as_slice()),
            _ => false
        }
    }
}

// Shorthand for common state machine behaviors.
macro_rules! shorthand (
    ( to          $state:ident             ) => ( self.state = states::$state;        );
    ( to          $state:ident $kind:expr  ) => ( self.state = states::$state($kind); );
    ( emit        $c:expr                  ) => ( self.emit_char($c);                 );
    ( create_tag  $kind:expr   $c:expr     ) => ( self.create_tag($kind, $c);         );
    ( append_tag  $c:expr                  ) => ( self.append_to_tag_name($c);        );
    ( emit_tag                             ) => ( self.emit_current_tag();            );
    ( append_temp $c:expr                  ) => ( self.temp_buf.push_char($c);        );
    ( emit_temp                            ) => ( self.emit_temp_buf();               );
    ( clear_temp                           ) => ( self.clear_temp_buf();              );

    // NB: Deliberate capture of 'c'
    ( error                                ) => ( self.parse_error(c);                );
)

// A little DSL for sequencing shorthand actions.
macro_rules! go (
    // A pattern like $($cmd:tt)* ; $($rest:tt)* causes parse ambiguity.
    // We have to tell the parser how much lookahead we need.
    ( $a:tt             ; $($rest:tt)* ) => ({ shorthand!($a);       go!($($rest)*); });
    ( $a:tt $b:tt       ; $($rest:tt)* ) => ({ shorthand!($a $b);    go!($($rest)*); });
    ( $a:tt $b:tt $c:tt ; $($rest:tt)* ) => ({ shorthand!($a $b $c); go!($($rest)*); });

    // These can only come at the end.
    // FIXME: Come up with a better name for 'finish'.
    ( reconsume ) => ( return Reconsume; );
    ( finish    ) => ( return Finished;  );

    // If nothing else matched, it's a single command
    ( $($cmd:tt)+ ) => ( shorthand!($($cmd)+); );

    // or nothing.
    () => (());
)

impl<'sink, Sink: TokenSink> Tokenizer<'sink, Sink> {
    // FIXME: explicitly represent the EOF character?
    // For now the plan is to handle EOF in a separate match.
    fn process_char(&mut self, c: char) -> ConsumeCharResult {
        debug!("Processing {:?} in state {:?}", c, self.state);
        match self.state {
            states::Data => match c {
                '&'  => go!(to CharacterReferenceInData),
                '<'  => go!(to TagOpen),
                '\0' => go!(error; emit '\0'),
                _    => go!(emit c),
            },

            // RCDATA, RAWTEXT, script, or script escaped
            states::RawData(kind) => match c {
                '&' if kind == Rcdata            => go!(to CharacterReferenceInRcdata),
                '-' if kind == ScriptDataEscaped => go!(to ScriptDataEscapedDash; emit '-'),
                '<'  => go!(to RawLessThanSign kind),
                '\0' => go!(error; emit '\ufffd'),
                _    => go!(emit c),
            },

            states::Plaintext => match c {
                '\0' => go!(error; emit '\ufffd'),
                _    => go!(emit c),
            },

            states::TagOpen => match c {
                '!' => go!(to MarkupDeclarationOpen),
                '/' => go!(to EndTagOpen),
                '?' => go!(error; to BogusComment),
                _ => match ascii_letter(c) {
                    Some(cl) => go!(create_tag StartTag cl; to TagName),
                    None     => go!(error; emit '<'; to Data; reconsume),
                }
            },

            states::EndTagOpen => match c {
                '>' => go!(error; to Data),
                _ => match ascii_letter(c) {
                    Some(cl) => go!(create_tag EndTag cl; to TagName),
                    None     => go!(error; to BogusComment),
                }
            },

            states::TagName => match c {
                '\t' | '\n' | '\x0C' | ' '
                     => go!(to BeforeAttributeName),
                '/'  => go!(to SelfClosingStartTag),
                '>'  => go!(emit_tag; to Data),
                '\0' => go!(error; append_tag '\ufffd'),
                _    => go!(append_tag (ascii_letter(c).unwrap_or(c))),
            },

            // kind = ScriptDataEscaped
            states::RawLessThanSign(ScriptDataEscaped) => match c {
                '/' => go!(clear_temp; to RawEndTagOpen ScriptDataEscaped),
                _ => match ascii_letter(c) {
                    Some(cl) => go!(clear_temp; append_temp cl;
                                    to ScriptDataDoubleEscapeStart; emit '<'; emit c),
                    None => go!(to RawData ScriptDataEscaped; emit '<'; reconsume),
                }
            },

            // otherwise
            states::RawLessThanSign(kind) => match c {
                '/' => go!(clear_temp; to RawEndTagOpen kind),
                '!' if kind == ScriptData => go!(to ScriptDataEscapeStart; emit '<'; emit '!'),
                _   => go!(to RawData Rcdata; emit '<'; reconsume),
            },

            states::RawEndTagOpen(kind) => match ascii_letter(c) {
                Some(cl) => go!(create_tag EndTag cl; append_temp c; to RawEndTagName kind),
                None     => go!(to RawData kind; emit '<'; emit '/'; reconsume),
            },

            states::RawEndTagName(kind) => {
                if self.have_appropriate_end_tag() {
                    match c {
                        '\t' | '\n' | '\x0C' | ' '
                            => go!(to BeforeAttributeName; finish),
                        '/' => go!(to SelfClosingStartTag; finish),
                        '>' => go!(emit_tag; to Data; finish),
                        // All of the above end with a return from this function.

                        _ => (),
                    }
                }

                match ascii_letter(c) {
                    Some(cl) => go!(append_tag cl; append_temp c),
                    None     => go!(emit '<'; emit '/'; emit_temp; to RawData kind; reconsume),
                }
            },

            states::ScriptDataEscapeStart => match c {
                '-' => go!(to ScriptDataEscapeStartDash; emit '-'),
                _   => go!(to RawData ScriptData; reconsume),
            },

            states::ScriptDataEscapeStartDash => match c {
                '-' => go!(to ScriptDataEscapedDashDash; emit '-'),
                _   => go!(to RawData ScriptData; reconsume),
            },

            states::ScriptDataEscapedDash => match c {
                '-'  => go!(to ScriptDataEscapedDashDash; emit '-'),
                '<'  => go!(to RawLessThanSign ScriptDataEscaped),
                '\0' => go!(error; to RawData ScriptDataEscaped; emit '\ufffd'),
                _    => go!(to RawData ScriptDataEscaped; emit c),
            },

            states::ScriptDataEscapedDashDash => match c {
                '-'  => go!(emit '-'),
                '<'  => go!(to RawLessThanSign ScriptDataEscaped),
                '>'  => go!(to RawData ScriptData; emit '>'),
                '\0' => go!(error; to RawData ScriptDataEscaped; emit '\ufffd'),
                _    => go!(to RawData ScriptDataEscaped; emit c),
            },

            states::ScriptDataDoubleEscapeStart |
            states::ScriptDataDoubleEscaped |
            states::ScriptDataDoubleEscapedDash |
            states::ScriptDataDoubleEscapedDashDash |
            states::ScriptDataDoubleEscapedLessThanSign |
            states::ScriptDataDoubleEscapeEnd |
            states::BeforeAttributeName |
            states::AttributeName |
            states::AfterAttributeName |
            states::BeforeAttributeValue |
            states::AttributeValueDoubleQuoted |
            states::AttributeValueSingleQuoted |
            states::AttributeValueUnquoted |
            states::AfterAttributeValueQuoted |
            states::SelfClosingStartTag |
            states::CommentStart |
            states::CommentStartDash |
            states::Comment |
            states::CommentEndDash |
            states::CommentEnd |
            states::CommentEndBang |
            states::Doctype |
            states::BeforeDoctypeName |
            states::DoctypeName |
            states::AfterDoctypeName |
            states::AfterDoctypePublicKeyword |
            states::BeforeDoctypePublicIdentifier |
            states::DoctypePublicIdentifierDoubleQuoted |
            states::DoctypePublicIdentifierSingleQuoted |
            states::AfterDoctypePublicIdentifier |
            states::BetweenDoctypePublicAndSystemIdentifiers |
            states::AfterDoctypeSystemKeyword |
            states::BeforeDoctypeSystemIdentifier |
            states::DoctypeSystemIdentifierDoubleQuoted |
            states::DoctypeSystemIdentifierSingleQuoted |
            states::AfterDoctypeSystemIdentifier |
            states::BogusDoctype |
            states::CharacterReferenceInData |
            states::CharacterReferenceInRcdata |
            states::CharacterReferenceInAttributeValue |
            states::BogusComment |
            states::MarkupDeclarationOpen |
            states::CdataSection
                => fail!("FIXME: state {:?} not implemented", self.state),
        }

        Finished

    }
}
<|endoftext|>


<commit_before>use regex::Regex;
use http::method::{Method, Options};
use http::status::{InternalServerError};
use iron::{Ingot, Request, Response, Alloy};
use iron::ingot::{Status, Continue, Unwind};

pub type Handler<Rq, Rs> = fn(&mut Rq, &mut Rs, &mut Alloy) -> ();

#[deriving(Clone)]
pub struct Router<Rq, Rs> {
    options: Vec<Method>,
    routes: Vec<Route<Rq, Rs>>
}

struct Route<Rq, Rs> {
    method: Method,
    glob: String,
    matches: Regex,
    handler: Handler<Rq, Rs>,
    params: Vec<String>
}

impl<Rq, Rs> Clone for Route<Rq, Rs> {
    fn clone(&self) -> Route<Rq, Rs> {
        Route {
            method: self.method.clone(),
            glob: self.glob.clone(),
            matches: self.matches.clone(),
            handler: self.handler,
            params: self.params.clone()
        }
    }
}

impl<Rq, Rs> Router<Rq, Rs> {
    pub fn new() -> Router<Rq, Rs> { Router { options: Vec::new(), routes: Vec::new() } }
    pub fn route(&mut self, method: Method, glob: String,
                 params: Vec<String>, handler: Handler<Rq, Rs>) {
        self.add_route(Route {
            method: method,
            matches: deglob(glob.clone()),
            params: params,
            handler: handler,
            glob: glob
        });
    }

    fn add_route(&mut self, route: Route<Rq, Rs>) {
        if !self.options.contains(&route.method) {
            self.options.push(route.method.clone())
        }
        self.routes.push(route);
    }
}

impl<Rq: Request, Rs: Response> Ingot<Rq, Rs> for Router {
    fn enter(&mut self, req: &mut Rq, res: &mut Rs, alloy: &mut Alloy) -> Status {
        if *req.method() == Options {
            res.write(self.options.iter().map(|p| p.show()).collect().join(" "));
            return Unwind;
        }
        for route in self.routes.iter() {
            if route.matches.is_match(req.reques_uri()) {
                route.handler(req, res, alloy);
                return Unwind;
            }
        }
        Continue
    }
}

static PARAMS: Regex = regex!(r":([a-zA-Z0-9_-]*)");

fn deglob(glob: String) -> Regex {
    // Replace glob patterns with corresponding regexs.
    let deglobbed = glob
        // Have to do this because the ** regex contains *
        .replace("**", "___DOUBLE_WILDCARD___")
        // Now only __DOUBLE_WILDCARD___ remains.
        .replace("*", "[a-zA-Z0-9_-]*")
        // Replace ** with its associated regex.
        .replace("___DOUBLE_WILDCARD___", "[/a-zA-Z0-9_-]*");
    // Replace :param patterns with corresponding regexs.
    let debound = PARAMS
        .replace_all(deglobbed.as_slice(), "([a-zA-Z0-9_-]*)");
    Regex::new("^".to_string().append(debound.as_slice()).append("$").as_slice()).unwrap()
}

<commit_msg>(feat) Implemented Ingot for Router.<commit_after>use regex::Regex;
use http::method::{Method, Options};
use http::status::{InternalServerError};
use iron::{Ingot, Request, Response, Alloy};
use iron::ingot::{Status, Continue, Unwind};

pub type Handler<Rq, Rs> = fn(&mut Rq, &mut Rs, &mut Alloy) -> ();

#[deriving(Clone)]
pub struct Router<Rq, Rs> {
    options: Vec<Method>,
    routes: Vec<Route<Rq, Rs>>
}

struct Route<Rq, Rs> {
    method: Method,
    glob: String,
    matches: Regex,
    handler: Handler<Rq, Rs>,
    params: Vec<String>
}

impl<Rq, Rs> Clone for Route<Rq, Rs> {
    fn clone(&self) -> Route<Rq, Rs> {
        Route {
            method: self.method.clone(),
            glob: self.glob.clone(),
            matches: self.matches.clone(),
            handler: self.handler,
            params: self.params.clone()
        }
    }
}

impl<Rq, Rs> Router<Rq, Rs> {
    pub fn new() -> Router<Rq, Rs> { Router { options: Vec::new(), routes: Vec::new() } }
    pub fn route(&mut self, method: Method, glob: String,
                 params: Vec<String>, handler: Handler<Rq, Rs>) {
        self.add_route(Route {
            method: method,
            matches: deglob(glob.clone()),
            params: params,
            handler: handler,
            glob: glob
        });
    }

    fn add_route(&mut self, route: Route<Rq, Rs>) {
        if !self.options.contains(&route.method) {
            self.options.push(route.method.clone())
        }
        self.routes.push(route);
    }
}

impl<Rq: Request + Clone, Rs: Response + Clone> Ingot<Rq, Rs> for Router<Rq, Rs> {
    fn enter(&mut self, req: &mut Rq, res: &mut Rs, alloy: &mut Alloy) -> Status {
        if *req.method() == Options {
            match res.write(
                self.options.iter()
                    .map(|p| format!("{}", p))
                    .collect::<Vec<String>>()
                    .connect(" ").as_bytes()) {
                Ok(_) => {},
                Err(err) => {
                    error!("Failed to write response: {}", err);
                    *res.status_mut() = InternalServerError;
                }
            }
            return Unwind;
        }
        for route in self.routes.iter() {
            if route.matches.is_match(format!("{}", req.uri()).as_slice()) {
                (route.handler)(req, res, alloy);
                return Unwind;
            }
        }
        Continue
    }
}

static PARAMS: Regex = regex!(r":([a-zA-Z0-9_-]*)");

fn deglob(glob: String) -> Regex {
    // Replace glob patterns with corresponding regexs.
    let deglobbed = glob
        // Have to do this because the ** regex contains *
        .replace("**", "___DOUBLE_WILDCARD___")
        // Now only __DOUBLE_WILDCARD___ remains.
        .replace("*", "[a-zA-Z0-9_-]*")
        // Replace ** with its associated regex.
        .replace("___DOUBLE_WILDCARD___", "[/a-zA-Z0-9_-]*");
    // Replace :param patterns with corresponding regexs.
    let debound = PARAMS
        .replace_all(deglobbed.as_slice(), "([a-zA-Z0-9_-]*)");
    Regex::new("^".to_string().append(debound.as_slice()).append("$").as_slice()).unwrap()
}

<|endoftext|>


<commit_before>#[macro_use]
extern crate adapton;

mod engine_is_from_scratch_consistent {
    //! This module tests that change propagation is from-scratch consistent

    /// Compare to this sound behavior to the unsound behavior from
    /// JaneStreet 'Incremental' libary.  In particular, Yit's gist
    /// shows the problematic change propagation behavior for a
    /// similar example program (dividing by zero when it should not):
    /// https://gist.github.com/khooyp/98abc0e64dc296deaa48
    #[test]
    fn avoid_divide_by_zero () {

        use std::rc::Rc;
        use adapton::macros::*;
        use adapton::engine::*;
        manage::init_dcg();

        // Construct two mutable inputs, `nom` and `den`, a
        // computation `div` that divides the nominator in `nom` by
        // the denominator in `den`, and a root thunk `t` that first
        // checks whether the denominator is zero (returning zero if
        // so) and if non-zero, returns the value of the division.

        let nom  = cell!(0);
        let den  = cell!(1);

        // In Rust, cloning is explicit:
        let den2 = den.clone(); // here, we clone the _global reference_ to the cell.
        let den3 = den.clone(); // here, we clone the _global reference_ to the cell, again.

        let div  = thunk![ get!(nom) / get!(den) ];
        let root = thunk![ if get!(den2) == 0 { 0 } else { get!(div) } ];

        assert_eq!(get!(root), 0);
        set(&den3, 0);
        assert_eq!(get!(root), 0);
    }

    #[test]
    fn avoid_expensive_subcomp () {

        use std::rc::Rc;
        use adapton::macros::*;
        use adapton::engine::*;
        use std::hash::{Hash, SipHasher, Hasher};

        manage::init_dcg();

        let flag = cell!(true); // flag: whether or not to do "iterative hash"
        let iter = cell!(3);    // iter: iterations for hashing
        let inp  = cell!(42);   // inp1: input value to hash

        let inp1  = inp.clone();  // we use inp non-linearly below
        let inp2  = inp.clone();  // we use inp non-linearly below
        let flag1 = flag.clone(); // we use flag non-linearly below
        let iter1 = iter.clone(); // we use iter non-linearly below

        let iter_hash = thunk![{
            let inp_val = get!(inp); // value to hash iteratively
            let mut hash_state = SipHasher::new();
            for _ in 0..get!(iter) { // hash iteratively, `iter` number of times
                inp_val.hash(&mut hash_state);
            };
            assert!(get!(iter) < 100); // For Testing: Too many iters is an error!
            hash_state.finish() // returns final hash value
        }];
        let root = thunk![ // switches between hashing iteratively, or not
            if get!(flag) { get!(iter_hash) }
            else { get!(inp1) }
        ];

        assert!(get!(root) != get!(inp2));
        set(&flag1, false);
        set(&iter1, 999999999);
        assert!(get!(root) == get!(inp2));
    }
}

mod engine_is_demand_driven {

    #[test]
    fn test1() {
        // TODO
    }

}

mod engine_api {
    //! This module tests gives unit tests for the engine's core API

    #[test]
    fn force_cell () {
        use adapton::engine::*;
        manage::init_dcg();
        let a : u32      = 1234;
        let b : Art<u32> = cell(name_of_usize(0), a);
        let c : u32      = force(&b);
        assert_eq!(a, c);
    }

    #[test]
    fn force_map_cell () {
        use adapton::engine::*;
        manage::init_dcg();
        let a : u32      = 1234;
        let b : Art<u32> = cell(name_of_usize(0), a);
        let c : u64      = force_map(&b, |_,x| x as u64);
        assert_eq!(a as u64, c);
    }

    #[test]
    fn force_map_cell_project () {
        use adapton::engine::*;
        manage::init_dcg();
        let pair = (1234 as usize, 5678 as usize);
        let c    = cell(name_of_usize(0), pair);
        let fst  = force_map(&c, |_,x| x.0);
        let snd  = force_map(&c, |_,x| x.1);
        assert_eq!(pair.0, fst);
        assert_eq!(pair.1, snd);
    }

    #[test]
    fn force_map_prunes_dirty_traversal () {
        // Test whether using force_map correctly prunes dirtying;
        // this test traces the engine, counts the number of dirtying
        // steps, and ensures that this count is zero, as expected.
        use std::rc::Rc;
        use adapton::macros::*;
        use adapton::engine::*;
        manage::init_dcg();
        reflect::dcg_reflect_begin();
        let pair = (1234 as usize, 5678 as usize);
        let c    = cell(name_of_usize(0), pair);
        let t = thunk![ name_of_usize(1) =>> {
            let fst = force_map(&c, |_,x| x.0);
            fst + 100
        }];
        assert_eq!(force(&t), 1334);
        let pair = (1234 as usize, 8765 as usize);
        let _    = cell(name_of_usize(0), pair);
        assert_eq!(force(&t), 1334);
        let traces = reflect::dcg_reflect_end();
        let counts = reflect::trace::trace_count(&traces, Some(1));
        assert_eq!(counts.dirty.0, 0);
        assert_eq!(counts.dirty.1, 0);
    }

    #[test]
    fn force_map_thunk () {
        use std::rc::Rc;
        use adapton::macros::*;
        use adapton::engine::*;
        manage::init_dcg();
        let a : u32      = 1234;
        let b : Art<u32> = thunk![ name_of_usize(0) =>> a];
        let c : u64      = force_map(&b, |_,x| x as u64);
        assert_eq!(a as u64, c);
    }
}
<commit_msg>engine tests: minor clean up and simplification<commit_after>#[macro_use]
extern crate adapton;

mod engine_is_from_scratch_consistent {
    //! This module tests that change propagation is from-scratch consistent

    /// Compare to this sound behavior to the unsound behavior from
    /// JaneStreet 'Incremental' libary.  In particular, Yit's gist
    /// shows the problematic change propagation behavior for a
    /// similar example program (dividing by zero when it should not):
    /// https://gist.github.com/khooyp/98abc0e64dc296deaa48
    #[test]
    fn avoid_divide_by_zero () {

        use std::rc::Rc;
        use adapton::macros::*;
        use adapton::engine::*;
        manage::init_dcg();

        // Construct two mutable inputs, `nom` and `den`, a
        // computation `div` that divides the nominator in `nom` by
        // the denominator in `den`, and a root thunk `t` that first
        // checks whether the denominator is zero (returning zero if
        // so) and if non-zero, returns the value of the division.

        let nom  = cell!(0);
        let den  = cell!(1);

        // In Rust, cloning is explicit:
        let den2 = den.clone(); // here, we clone the _global reference_ to the cell.
        let den3 = den.clone(); // here, we clone the _global reference_ to the cell, again.

        let div  = thunk![ get!(nom) / get!(den) ];
        let root = thunk![ if get!(den2) == 0 { 0 } else { get!(div) } ];

        assert_eq!(get!(root), 0);
        set(&den3, 0);
        assert_eq!(get!(root), 0);
    }

    #[test]
    fn avoid_expensive_subcomp () {

        use std::rc::Rc;
        use adapton::macros::*;
        use adapton::engine::*;
        use std::hash::{Hash, SipHasher, Hasher};

        manage::init_dcg();

        let flag = cell!(true); // flag: whether or not to do "iterative hash"
        let iter = cell!(3);    // iter: iterations for hashing
        let inp  = cell!(42);   // inp1: input value to hash

        let inp1  = inp.clone();  // we use inp non-linearly below
        let inp2  = inp.clone();  // we use inp non-linearly below
        let flag1 = flag.clone(); // we use flag non-linearly below
        let iter1 = iter.clone(); // we use iter non-linearly below

        let iter_hash = thunk![{
            let inp_val = get!(inp); // value to hash iteratively
            let mut hash_state = SipHasher::new();
            for _ in 0..get!(iter) { // hash iteratively, `iter` number of times
                inp_val.hash(&mut hash_state);
            };
            assert!(get!(iter) < 100); // For Testing: Too many iters is an error!
            hash_state.finish() // returns final hash value
        }];
        let root = thunk![ // switches between hashing iteratively, or not
            if get!(flag) { get!(iter_hash) }
            else { get!(inp1) }
        ];

        assert!(get!(root) != get!(inp2));
        set(&flag1, false);
        set(&iter1, 999999999);
        assert!(get!(root) == get!(inp2));
    }
}

mod engine_api {
    //! This module tests gives unit tests for the engine's core API

    #[test]
    fn force_cell () {
        use std::rc::Rc;
        use adapton::macros::*;
        use adapton::engine::*;
        manage::init_dcg();
        let a : u32      = 1234;
        let b : Art<u32> = cell!(a);
        let c : u32      = force(&b);
        assert_eq!(a, c);
    }

    #[test]
    fn force_map_cell () {
        use std::rc::Rc;
        use adapton::macros::*;
        use adapton::engine::*;
        manage::init_dcg();
        let a : u32      = 1234;
        let b : Art<u32> = cell!(a);
        let c : u64      = force_map(&b, |_,x| x as u64);
        assert_eq!(a as u64, c);
    }

    #[test]
    fn force_map_cell_project () {
        use std::rc::Rc;
        use adapton::macros::*;
        use adapton::engine::*;
        manage::init_dcg();
        let pair = (1234 as usize, 5678 as usize);
        let c    = cell!(pair);
        let fst  = force_map(&c, |_,x| x.0);
        let snd  = force_map(&c, |_,x| x.1);
        assert_eq!(pair.0, fst);
        assert_eq!(pair.1, snd);
    }

    #[test]
    fn force_map_prunes_dirty_traversal () {
        // Test whether using force_map correctly prunes dirtying;
        // this test traces the engine, counts the number of dirtying
        // steps, and ensures that this count is zero, as expected.
        use std::rc::Rc;
        use adapton::macros::*;
        use adapton::engine::*;
        manage::init_dcg();
        reflect::dcg_reflect_begin();
        let pair = (1234 as usize, 5678 as usize);
        let c    = cell!(pair);
        let t = thunk![{
            let fst = force_map(&c, |_,x| x.0);
            fst + 100
        }];
        assert_eq!(force(&t), 1334);
        let pair = (1234 as usize, 8765 as usize);
        let _    = cell!(pair);
        assert_eq!(force(&t), 1334);
        let traces = reflect::dcg_reflect_end();
        let counts = reflect::trace::trace_count(&traces, Some(1));
        assert_eq!(counts.dirty.0, 0);
        assert_eq!(counts.dirty.1, 0);
    }

    #[test]
    fn force_map_thunk () {
        use std::rc::Rc;
        use adapton::macros::*;
        use adapton::engine::*;
        manage::init_dcg();
        let a : u32      = 1234;
        let b : Art<u32> = thunk![a];
        let c : u64      = force_map(&b, |_,x| x as u64);
        assert_eq!(a as u64, c);
    }
}
<|endoftext|>


<commit_before>extern crate redis;

extern crate futures;
extern crate partial_io;
#[macro_use]
extern crate quickcheck;
extern crate tokio;

mod support;

use std::io::{self, BufReader};

use partial_io::{GenWouldBlock, PartialAsyncRead, PartialWithErrors};

use tokio::executor::current_thread::block_on_all;

use futures::Future;

use redis::Value;

use support::encode_value;

#[derive(Clone, Debug)]
struct ArbitraryValue(Value);
impl ::quickcheck::Arbitrary for ArbitraryValue {
    fn arbitrary<G: ::quickcheck::Gen>(g: &mut G) -> Self {
        let size = g.size();
        ArbitraryValue(arbitrary_value(g, size))
    }
    fn shrink(&self) -> Box<Iterator<Item = Self>> {
        match self.0 {
            Value::Nil | Value::Okay => Box::new(None.into_iter()),
            Value::Int(i) => Box::new(i.shrink().map(Value::Int).map(ArbitraryValue)),
            Value::Data(ref xs) => Box::new(xs.shrink().map(Value::Data).map(ArbitraryValue)),
            Value::Bulk(ref xs) => {
                let ys = xs
                    .iter()
                    .map(|x| ArbitraryValue(x.clone()))
                    .collect::<Vec<_>>();
                Box::new(
                    ys.shrink()
                        .map(|xs| xs.into_iter().map(|x| x.0).collect())
                        .map(Value::Bulk)
                        .map(ArbitraryValue),
                )
            }
            Value::Status(ref status) => {
                Box::new(status.shrink().map(Value::Status).map(ArbitraryValue))
            }
        }
    }
}

fn arbitrary_value<G: ::quickcheck::Gen>(g: &mut G, recursive_size: usize) -> Value {
    use quickcheck::Arbitrary;
    if recursive_size == 0 {
        Value::Nil
    } else {
        match g.gen_range(0, 6) {
            0 => Value::Nil,
            1 => Value::Int(Arbitrary::arbitrary(g)),
            2 => Value::Data(Arbitrary::arbitrary(g)),
            3 => {
                let size = {
                    let s = g.size();
                    g.gen_range(0, s)
                };
                Value::Bulk(
                    (0..size)
                        .map(|_| arbitrary_value(g, recursive_size / size))
                        .collect(),
                )
            }
            4 => {
                let size = {
                    let s = g.size();
                    g.gen_range(0, s)
                };
                let status = g.gen_ascii_chars().take(size).collect();
                if status == "OK" {
                    Value::Okay
                } else {
                    Value::Status(status)
                }
            }
            5 => Value::Okay,
            _ => unreachable!(),
        }
    }
}

quickcheck! {
    fn partial_io_parse(input: ArbitraryValue, seq: PartialWithErrors<GenWouldBlock>) -> () {
        let mut encoded_input = Vec::new();
        encode_value(&input.0, &mut encoded_input).unwrap();

        let mut reader = &encoded_input[..];
        let partial_reader = PartialAsyncRead::new(&mut reader, seq);

        let result = block_on_all(redis::parse_async(BufReader::new(partial_reader))
            .map(|t| t.1));
        assert!(result.as_ref().is_ok(), "{}", result.unwrap_err());
        assert_eq!(
            result.unwrap(),
            input.0,
        );
    }
}
<commit_msg>Remove unused import<commit_after>extern crate redis;

extern crate futures;
extern crate partial_io;
#[macro_use]
extern crate quickcheck;
extern crate tokio;

mod support;

use std::io::BufReader;

use partial_io::{GenWouldBlock, PartialAsyncRead, PartialWithErrors};

use tokio::executor::current_thread::block_on_all;

use futures::Future;

use redis::Value;

use support::encode_value;

#[derive(Clone, Debug)]
struct ArbitraryValue(Value);
impl ::quickcheck::Arbitrary for ArbitraryValue {
    fn arbitrary<G: ::quickcheck::Gen>(g: &mut G) -> Self {
        let size = g.size();
        ArbitraryValue(arbitrary_value(g, size))
    }
    fn shrink(&self) -> Box<Iterator<Item = Self>> {
        match self.0 {
            Value::Nil | Value::Okay => Box::new(None.into_iter()),
            Value::Int(i) => Box::new(i.shrink().map(Value::Int).map(ArbitraryValue)),
            Value::Data(ref xs) => Box::new(xs.shrink().map(Value::Data).map(ArbitraryValue)),
            Value::Bulk(ref xs) => {
                let ys = xs
                    .iter()
                    .map(|x| ArbitraryValue(x.clone()))
                    .collect::<Vec<_>>();
                Box::new(
                    ys.shrink()
                        .map(|xs| xs.into_iter().map(|x| x.0).collect())
                        .map(Value::Bulk)
                        .map(ArbitraryValue),
                )
            }
            Value::Status(ref status) => {
                Box::new(status.shrink().map(Value::Status).map(ArbitraryValue))
            }
        }
    }
}

fn arbitrary_value<G: ::quickcheck::Gen>(g: &mut G, recursive_size: usize) -> Value {
    use quickcheck::Arbitrary;
    if recursive_size == 0 {
        Value::Nil
    } else {
        match g.gen_range(0, 6) {
            0 => Value::Nil,
            1 => Value::Int(Arbitrary::arbitrary(g)),
            2 => Value::Data(Arbitrary::arbitrary(g)),
            3 => {
                let size = {
                    let s = g.size();
                    g.gen_range(0, s)
                };
                Value::Bulk(
                    (0..size)
                        .map(|_| arbitrary_value(g, recursive_size / size))
                        .collect(),
                )
            }
            4 => {
                let size = {
                    let s = g.size();
                    g.gen_range(0, s)
                };
                let status = g.gen_ascii_chars().take(size).collect();
                if status == "OK" {
                    Value::Okay
                } else {
                    Value::Status(status)
                }
            }
            5 => Value::Okay,
            _ => unreachable!(),
        }
    }
}

quickcheck! {
    fn partial_io_parse(input: ArbitraryValue, seq: PartialWithErrors<GenWouldBlock>) -> () {
        let mut encoded_input = Vec::new();
        encode_value(&input.0, &mut encoded_input).unwrap();

        let mut reader = &encoded_input[..];
        let partial_reader = PartialAsyncRead::new(&mut reader, seq);

        let result = block_on_all(redis::parse_async(BufReader::new(partial_reader))
            .map(|t| t.1));
        assert!(result.as_ref().is_ok(), "{}", result.unwrap_err());
        assert_eq!(
            result.unwrap(),
            input.0,
        );
    }
}
<|endoftext|>


<commit_before>use std::{
    collections::VecDeque,
    fmt, io,
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc, Mutex,
    },
    time::{Duration, Instant},
};

use crate::{
    config::ServerConfig,
    context::SharedContext,
    relay::tcprelay::EncryptedWrite,
    relay::{loadbalancing::server::LoadBalancer, socks5::Address, tcprelay::client::ServerClient},
};

use futures::{Future, Stream};
use log::{debug, error, info};
use tokio::{
    self,
    io::read_exact,
    timer::{Interval, Timeout},
};

struct Server {
    config: Arc<ServerConfig>,
    score: AtomicU64,
}

impl Server {
    fn new(config: ServerConfig) -> Server {
        Server {
            config: Arc::new(config),
            score: AtomicU64::new(0),
        }
    }

    fn score(&self) -> u64 {
        self.score.load(Ordering::Acquire)
    }
}

const MAX_LATENCY_QUEUE_SIZE: usize = 37;

struct ServerLatencyInner {
    latency_queue: VecDeque<u64>,
}

impl ServerLatencyInner {
    fn new() -> ServerLatencyInner {
        ServerLatencyInner {
            latency_queue: VecDeque::with_capacity(MAX_LATENCY_QUEUE_SIZE),
        }
    }

    fn push(&mut self, lat: u64) -> u64 {
        self.latency_queue.push_back(lat);
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.score()
    }

    fn score(&self) -> u64 {
        if self.latency_queue.is_empty() {
            return u64::max_value();
        }

        let sum: u64 = self.latency_queue.iter().sum();
        let avg = sum as f64 / self.latency_queue.len() as f64;
        let s_sum: u64 = self.latency_queue.iter().map(|n| *n * *n).sum();
        let dev = (s_sum as f64 / self.latency_queue.len() as f64 - avg * avg).sqrt();

        let mut v = self
            .latency_queue
            .iter()
            .cloned()
            .filter(|n| *n as f64 >= (avg - 3.0 * dev) && *n as f64 <= (avg + 3.0 * dev))
            .collect::<Vec<u64>>();
        v.sort();

        if v.is_empty() {
            return u64::max_value();
        }

        let mid = v.len() / 2;
        let median = if (v.len() & 1) == 0 {
            (v[mid - 1] + v[mid]) / 2
        } else {
            v[mid]
        };

        debug!(
            "latency queue {:?}, avg {}, dev {}, sorted {:?}, median {}",
            self.latency_queue, avg, dev, v, median
        );

        (((median as f64 * 2.0) + (dev * 1.0)) / 2.0) as u64
    }
}

#[derive(Clone)]
struct ServerLatency {
    inner: Arc<Mutex<ServerLatencyInner>>,
}

impl ServerLatency {
    fn new() -> ServerLatency {
        ServerLatency {
            inner: Arc::new(Mutex::new(ServerLatencyInner::new())),
        }
    }

    fn push(&self, lat: u64) -> u64 {
        let mut inner = self.inner.lock().unwrap();
        inner.push(lat)
    }
}

impl fmt::Debug for ServerLatency {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.inner.lock().unwrap();
        write!(f, "ServerLatency {{ latency: {:?} }}", inner.latency_queue)
    }
}

const DEFAULT_CHECK_INTERVAL_SEC: u64 = 6;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 5;

struct Inner {
    servers: Vec<Arc<Server>>,
    best_idx: AtomicUsize,
}

impl Inner {
    fn new(context: SharedContext) -> Inner {
        let config = context.config();

        if config.server.is_empty() {
            panic!("no servers");
        }

        let check_required = config.server.len() > 1;

        let mut servers = Vec::new();
        for svr in &config.server {
            let sc = Arc::new(Server::new(svr.clone()));
            servers.push(sc.clone());

            if !check_required {
                continue;
            }

            let context = context.clone();
            let latency = ServerLatency::new();

            tokio::spawn(
                // Check every DEFAULT_CHECK_INTERVAL_SEC seconds
                Interval::new(
                    Instant::now() + Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
                    Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
                )
                .for_each(move |_| {
                    let sc = sc.clone();
                    let lat = latency.clone();

                    Inner::check_delay(sc.clone(), context.clone()).then(move |res| {
                        let score = match res {
                            Ok(d) => lat.push(d),
                            Err(..) => lat.push(DEFAULT_CHECK_TIMEOUT_SEC * 2 * 1000), // Penalty
                        };
                        debug!("updated remote server {} (score: {})", sc.config.addr(), score);
                        sc.score.store(score, Ordering::Release);

                        Ok(())
                    })
                })
                .map_err(|_| ()),
            );
        }

        Inner {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    fn check_delay(sc: Arc<Server>, context: SharedContext) -> impl Future<Item = u64, Error = io::Error> {
        static GET_BODY: &[u8] = b"GET / HTTP/1.1\r\nHost: www.example.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("www.example.com".to_owned(), 80);

        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let fut = ServerClient::connect(context, addr, sc.config.clone()).and_then(|ServerClient { r, w }| {
            w.write_all(GET_BODY)
                .and_then(|(mut w, _)| w.flush())
                .and_then(move |_| read_exact(r, [0u8, 1]))
        });

        Timeout::new(fut, Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC)).then(move |res| {
            let elapsed = Instant::now() - start;
            let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
            match res {
                Ok(..) => {
                    // Got the result ... record its time
                    debug!("checked remote server {} latency with {} ms", sc.config.addr(), elapsed);
                    Ok(elapsed)
                }
                Err(err) => {
                    match err.into_inner() {
                        Some(err) => {
                            error!("failed to check server {}, error: {}", sc.config.addr(), err);

                            // NOTE: connection / handshake error, server is down
                            Err(err)
                        }
                        None => {
                            // Timeout
                            debug!(
                                "checked remote server {} latency timeout, elapsed {} ms",
                                sc.config.addr(),
                                elapsed
                            );

                            // NOTE: timeout is still available, but server is too slow
                            Ok(elapsed)
                        }
                    }
                }
            }
        })
    }

    fn best_idx(&self) -> usize {
        self.best_idx.load(Ordering::Acquire)
    }

    fn set_best_idx(&self, idx: usize) {
        self.best_idx.store(idx, Ordering::Release)
    }

    fn best_server(&self) -> &Arc<Server> {
        &self.servers[self.best_idx()]
    }
}

#[derive(Clone)]
pub struct PingBalancer {
    inner: Arc<Inner>,
}

impl PingBalancer {
    pub fn new(context: SharedContext) -> PingBalancer {
        let inner = Arc::new(Inner::new(context));

        let inner_cloned = inner.clone();
        tokio::spawn(
            Interval::new(
                Instant::now() + Duration::from_secs(2),
                Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
            )
            .for_each(move |_| {
                let inner = inner_cloned.clone();

                if inner.servers.is_empty() {
                    panic!("No server");
                }

                // Choose the best one
                let mut svr_idx = 0;

                for (idx, svr) in inner.servers.iter().enumerate() {
                    let choosen_svr = &inner.servers[svr_idx];
                    if svr.score() < choosen_svr.score() {
                        svr_idx = idx;
                    }
                }

                let choosen_svr = &inner.servers[svr_idx];
                debug!(
                    "chosen the best server {} (score: {})",
                    choosen_svr.config.addr(),
                    choosen_svr.score()
                );

                if inner.best_idx() != svr_idx {
                    info!(
                        "switched server from {} (score: {}) to {} (score: {})",
                        inner.best_server().config.addr(),
                        inner.best_server().score(),
                        choosen_svr.config.addr(),
                        choosen_svr.score(),
                    );
                }

                inner.set_best_idx(svr_idx);

                Ok(())
            })
            .map_err(|_| ()),
        );

        PingBalancer { inner }
    }
}

impl LoadBalancer for PingBalancer {
    fn pick_server(&mut self) -> Arc<ServerConfig> {
        if self.inner.servers.is_empty() {
            panic!("no server in configuration");
        }
        self.inner.best_server().config.clone()
    }

    fn total(&self) -> usize {
        self.inner.servers.len()
    }
}
<commit_msg>Removed load balancing filter by deviation logic<commit_after>use std::{
    collections::VecDeque,
    fmt, io,
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc, Mutex,
    },
    time::{Duration, Instant},
};

use crate::{
    config::ServerConfig,
    context::SharedContext,
    relay::tcprelay::EncryptedWrite,
    relay::{loadbalancing::server::LoadBalancer, socks5::Address, tcprelay::client::ServerClient},
};

use futures::{Future, Stream};
use log::{debug, error, info};
use tokio::{
    self,
    io::read_exact,
    timer::{Interval, Timeout},
};

struct Server {
    config: Arc<ServerConfig>,
    score: AtomicU64,
}

impl Server {
    fn new(config: ServerConfig) -> Server {
        Server {
            config: Arc::new(config),
            score: AtomicU64::new(0),
        }
    }

    fn score(&self) -> u64 {
        self.score.load(Ordering::Acquire)
    }
}

const MAX_LATENCY_QUEUE_SIZE: usize = 37;

struct ServerLatencyInner {
    latency_queue: VecDeque<u64>,
}

impl ServerLatencyInner {
    fn new() -> ServerLatencyInner {
        ServerLatencyInner {
            latency_queue: VecDeque::with_capacity(MAX_LATENCY_QUEUE_SIZE),
        }
    }

    fn push(&mut self, lat: u64) -> u64 {
        self.latency_queue.push_back(lat);
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.score()
    }

    fn score(&self) -> u64 {
        if self.latency_queue.is_empty() {
            return u64::max_value();
        }

        let mut v = self.latency_queue.iter().cloned().collect::<Vec<u64>>();
        v.sort();

        let mid = v.len() / 2;
        if (v.len() & 1) == 0 {
            (v[mid - 1] + v[mid]) / 2
        } else {
            v[mid]
        }
    }
}

#[derive(Clone)]
struct ServerLatency {
    inner: Arc<Mutex<ServerLatencyInner>>,
}

impl ServerLatency {
    fn new() -> ServerLatency {
        ServerLatency {
            inner: Arc::new(Mutex::new(ServerLatencyInner::new())),
        }
    }

    fn push(&self, lat: u64) -> u64 {
        let mut inner = self.inner.lock().unwrap();
        inner.push(lat)
    }
}

impl fmt::Debug for ServerLatency {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.inner.lock().unwrap();
        write!(f, "ServerLatency {{ latency: {:?} }}", inner.latency_queue)
    }
}

const DEFAULT_CHECK_INTERVAL_SEC: u64 = 6;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 5;

struct Inner {
    servers: Vec<Arc<Server>>,
    best_idx: AtomicUsize,
}

impl Inner {
    fn new(context: SharedContext) -> Inner {
        let config = context.config();

        if config.server.is_empty() {
            panic!("no servers");
        }

        let check_required = config.server.len() > 1;

        let mut servers = Vec::new();
        for svr in &config.server {
            let sc = Arc::new(Server::new(svr.clone()));
            servers.push(sc.clone());

            if !check_required {
                continue;
            }

            let context = context.clone();
            let latency = ServerLatency::new();

            tokio::spawn(
                // Check every DEFAULT_CHECK_INTERVAL_SEC seconds
                Interval::new(
                    Instant::now() + Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
                    Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
                )
                .for_each(move |_| {
                    let sc = sc.clone();
                    let lat = latency.clone();

                    Inner::check_delay(sc.clone(), context.clone()).then(move |res| {
                        let score = match res {
                            Ok(d) => lat.push(d),
                            Err(..) => lat.push(DEFAULT_CHECK_TIMEOUT_SEC * 2 * 1000), // Penalty
                        };
                        debug!("updated remote server {} (score: {})", sc.config.addr(), score);
                        sc.score.store(score, Ordering::Release);

                        Ok(())
                    })
                })
                .map_err(|_| ()),
            );
        }

        Inner {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    fn check_delay(sc: Arc<Server>, context: SharedContext) -> impl Future<Item = u64, Error = io::Error> {
        static GET_BODY: &[u8] = b"GET / HTTP/1.1\r\nHost: www.example.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("www.example.com".to_owned(), 80);

        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let fut = ServerClient::connect(context, addr, sc.config.clone()).and_then(|ServerClient { r, w }| {
            w.write_all(GET_BODY)
                .and_then(|(mut w, _)| w.flush())
                .and_then(move |_| read_exact(r, [0u8, 1]))
        });

        Timeout::new(fut, Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC)).then(move |res| {
            let elapsed = Instant::now() - start;
            let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
            match res {
                Ok(..) => {
                    // Got the result ... record its time
                    debug!("checked remote server {} latency with {} ms", sc.config.addr(), elapsed);
                    Ok(elapsed)
                }
                Err(err) => {
                    match err.into_inner() {
                        Some(err) => {
                            error!("failed to check server {}, error: {}", sc.config.addr(), err);

                            // NOTE: connection / handshake error, server is down
                            Err(err)
                        }
                        None => {
                            // Timeout
                            debug!(
                                "checked remote server {} latency timeout, elapsed {} ms",
                                sc.config.addr(),
                                elapsed
                            );

                            // NOTE: timeout is still available, but server is too slow
                            Ok(elapsed)
                        }
                    }
                }
            }
        })
    }

    fn best_idx(&self) -> usize {
        self.best_idx.load(Ordering::Acquire)
    }

    fn set_best_idx(&self, idx: usize) {
        self.best_idx.store(idx, Ordering::Release)
    }

    fn best_server(&self) -> &Arc<Server> {
        &self.servers[self.best_idx()]
    }
}

#[derive(Clone)]
pub struct PingBalancer {
    inner: Arc<Inner>,
}

impl PingBalancer {
    pub fn new(context: SharedContext) -> PingBalancer {
        let inner = Arc::new(Inner::new(context));

        let inner_cloned = inner.clone();
        tokio::spawn(
            Interval::new(
                Instant::now() + Duration::from_secs(2),
                Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
            )
            .for_each(move |_| {
                let inner = inner_cloned.clone();

                if inner.servers.is_empty() {
                    panic!("No server");
                }

                // Choose the best one
                let mut svr_idx = 0;

                for (idx, svr) in inner.servers.iter().enumerate() {
                    let choosen_svr = &inner.servers[svr_idx];
                    if svr.score() < choosen_svr.score() {
                        svr_idx = idx;
                    }
                }

                let choosen_svr = &inner.servers[svr_idx];
                debug!(
                    "chosen the best server {} (score: {})",
                    choosen_svr.config.addr(),
                    choosen_svr.score()
                );

                if inner.best_idx() != svr_idx {
                    info!(
                        "switched server from {} (score: {}) to {} (score: {})",
                        inner.best_server().config.addr(),
                        inner.best_server().score(),
                        choosen_svr.config.addr(),
                        choosen_svr.score(),
                    );
                }

                inner.set_best_idx(svr_idx);

                Ok(())
            })
            .map_err(|_| ()),
        );

        PingBalancer { inner }
    }
}

impl LoadBalancer for PingBalancer {
    fn pick_server(&mut self) -> Arc<ServerConfig> {
        if self.inner.servers.is_empty() {
            panic!("no server in configuration");
        }
        self.inner.best_server().config.clone()
    }

    fn total(&self) -> usize {
        self.inner.servers.len()
    }
}
<|endoftext|>


<commit_before>extern crate bazel_protos;
extern crate boxfuture;
extern crate clap;
extern crate fs;
extern crate futures;
extern crate protobuf;

use boxfuture::{Boxable, BoxFuture};
use clap::{App, Arg, SubCommand};
use fs::hash::Fingerprint;
use fs::store::Store;
use futures::future::{Future, join_all};
use std::error::Error;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::process::exit;
use std::sync::Arc;

fn main() {
  match execute(
    App::new("fs_util")
      .subcommand(
        SubCommand::with_name("file")
          .subcommand(
            SubCommand::with_name("cat")
              .about("Output the contents of a file by fingerprint.")
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              )),
          )
          .subcommand(
            SubCommand::with_name("save")
              .about(
                "Ingest a file by path, which allows it to be used in Directories/Snapshots. \
Outputs a fingerprint of its contents and its size in bytes, separated by a space.",
              )
              .arg(Arg::with_name("path").required(true).takes_value(true)),
          ),
      )
      .subcommand(
        SubCommand::with_name("directory").subcommand(
          SubCommand::with_name("save")
            .about(
              "Ingest a directory recursively. Saves all files found therein and saves Directory \
protos for each directory found. Outputs a fingerprint of the canonical top-level Directory proto \
and the size of the serialized proto in bytes, separated by a space.",
            )
            .arg(Arg::with_name("source").required(true).takes_value(true)),
        ),
      )
      .arg(
        Arg::with_name("local_store_path")
          .takes_value(true)
          .long("local_store_path")
          .required(true),
      )
      .get_matches(),
  ) {
    Ok(_) => {}
    Err(err) => {
      eprintln!("{}", err);
      exit(1)
    }
  };
}

fn execute(top_match: clap::ArgMatches) -> Result<(), String> {
  let store_dir = top_match.value_of("local_store_path").unwrap();
  let store = Arc::new(Store::new(store_dir).map_err(|e| {
    format!(
      "Failed to open/create store for directory {}: {}",
      store_dir,
      e
    )
  })?);

  match top_match.subcommand() {
    ("file", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("cat", Some(args)) => {
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          match store.load_bytes(&fingerprint)? {
            Some(bytes) => {
              io::stdout().write(&bytes).unwrap();
              Ok(())
            }
            None => Err(format!("File with fingerprint {} not found", fingerprint)),
          }
        }
        ("save", Some(args)) => {
          let path = PathBuf::from(args.value_of("path").unwrap());
          let posix_fs = make_posix_fs(path.parent().unwrap());
          let file = posix_fs
            .stat(PathBuf::from(path.file_name().unwrap()))
            .unwrap();
          match file {
            fs::Stat::File(f) => {
              let (fingerprint, size_bytes) = save_file(store, &posix_fs, f).wait().unwrap();
              Ok(println!("{} {}", fingerprint, size_bytes))
            }
            o => Err(format!(
              "Tried to save file {:?} but it was not a file, was a {:?}",
              path,
              o
            )),
          }

        }
        (_, _) => unimplemented!(),
      }
    }
    ("directory", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("save", Some(args)) => {
          let posix_fs = Arc::new(make_posix_fs(args.value_of("source").unwrap()));
          let (fingerprint, size_bytes) =
            save_directory(store, posix_fs, Arc::new(fs::Dir(PathBuf::from("."))))
              .wait()?;
          Ok(println!("{} {}", fingerprint, size_bytes))
        }
        (_, _) => unimplemented!(),
      }
    }
    (_, _) => unimplemented!(),
  }
}

fn make_posix_fs<P: AsRef<Path>>(root: P) -> fs::PosixFS {
  fs::PosixFS::new(&root, vec![]).unwrap()
}

fn save_file(
  store: Arc<Store>,
  posix_fs: &fs::PosixFS,
  file: fs::File,
) -> BoxFuture<(Fingerprint, usize), String> {
  posix_fs
    .read_file(&file)
    .map_err(move |err| {
      format!("Error reading file {:?}: {}", file, err.description())
    })
    .and_then(move |content| {
      Ok((
        store.store_file_bytes(&content.content)?,
        content.content.len(),
      ))
    })
    .to_boxed()
}

fn save_directory(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  base_dir: Arc<fs::Dir>,
) -> BoxFuture<(Fingerprint, usize), String> {
  let store_copy = store.clone();
  Box::new(
    posix_fs
      .scandir(&base_dir)
      .map_err(move |err| {
        format!("Failed to scandir {:?}: {}", &base_dir, err.description())
      })
      .and_then(move |stats| {
        let mut dirs = Vec::new();
        let mut files = Vec::new();
        for stat in stats.into_iter() {
          match stat {
            fs::Stat::Dir(dir) => dirs.push(dir),
            fs::Stat::File(file) => files.push(file),
            fs::Stat::Link(fs::Link(path)) => {
              let err: BoxFuture<_, _> =
                futures::future::err(format!("Don't yet know how to handle symlinks: {:?}", path))
                  .to_boxed();
              return err;
            }
          }
        }

        let store_inner_copy = store_copy.clone();
        let posix_fs_copy = posix_fs.clone();
        let dir_futures: Vec<_> = dirs
          .into_iter()
          .map(move |dir| {
            let dir = Arc::new(dir);
            save_directory(store_inner_copy.clone(), posix_fs_copy.clone(), dir.clone())
              .and_then(move |(fingerprint, proto_size_bytes)| {
                let mut dir_node = bazel_protos::remote_execution::DirectoryNode::new();
                dir_node.set_name(file_name_as_utf8(&dir.0)?);
                dir_node.set_digest(fingerprint_to_digest(&fingerprint, proto_size_bytes as i64));
                Ok(dir_node)
              })
          })
          .collect();

        let file_futures: Vec<_> = files
          .into_iter()
          .map(move |file| {
            file_to_file_node(store_copy.clone(), posix_fs.clone(), file)
          })
          .collect();

        join_all(dir_futures)
          .join(join_all(file_futures))
          .to_boxed()
      })
      .and_then(move |(dirs, files)| {
        let mut directory = bazel_protos::remote_execution::Directory::new();
        directory.set_directories(protobuf::RepeatedField::from_vec(dirs));
        directory.set_files(protobuf::RepeatedField::from_vec(files));
        store.record_directory(&directory)
      }),
  )
}

fn file_to_file_node(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  file: fs::File,
) -> BoxFuture<bazel_protos::remote_execution::FileNode, String> {
  let file_copy = file.clone();
  save_file(store, &posix_fs, file)
    .and_then(move |(fingerprint, size_bytes)| {
      let mut file_node = bazel_protos::remote_execution::FileNode::new();
      file_node.set_name(file_name_as_utf8(&file_copy.path)?);
      file_node.set_digest(fingerprint_to_digest(&fingerprint, size_bytes as i64));
      file_node.set_is_executable(file_copy.is_executable);
      Ok(file_node)
    })
    .to_boxed()
}

pub fn fingerprint_to_digest(
  fingerprint: &Fingerprint,
  size_bytes: i64,
) -> bazel_protos::remote_execution::Digest {
  let mut digest = bazel_protos::remote_execution::Digest::new();
  digest.set_hash(fingerprint.to_hex());
  digest.set_size_bytes(size_bytes);
  digest
}

fn file_name_as_utf8(path: &Path) -> Result<String, String> {
  match path.file_name() {
    Some(name) => {
      match name.to_str() {
        Some(name_utf8) => Ok(name_utf8.to_string()),
        None => Err(format!(
          "{:?}'s file_name is not representable in UTF8",
          path
        )),
      }
    }
    None => Err(format!("{:?} did not have a file_name", path)),
  }
}
<commit_msg>fs_util exits 2 for ENOENT (#5064)<commit_after>extern crate bazel_protos;
extern crate boxfuture;
extern crate clap;
extern crate fs;
extern crate futures;
extern crate protobuf;

use boxfuture::{Boxable, BoxFuture};
use clap::{App, Arg, SubCommand};
use fs::hash::Fingerprint;
use fs::store::Store;
use futures::future::{Future, join_all};
use std::error::Error;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::process::exit;
use std::sync::Arc;

#[derive(Debug)]
enum ExitCode {
  UnknownError = 1,
  NotFound = 2,
}

#[derive(Debug)]
struct ExitError(pub String, pub ExitCode);

impl From<String> for ExitError {
  fn from(s: String) -> Self {
    ExitError(s, ExitCode::UnknownError)
  }
}

fn main() {
  match execute(
    App::new("fs_util")
      .subcommand(
        SubCommand::with_name("file")
          .subcommand(
            SubCommand::with_name("cat")
              .about("Output the contents of a file by fingerprint.")
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              )),
          )
          .subcommand(
            SubCommand::with_name("save")
              .about(
                "Ingest a file by path, which allows it to be used in Directories/Snapshots. \
Outputs a fingerprint of its contents and its size in bytes, separated by a space.",
              )
              .arg(Arg::with_name("path").required(true).takes_value(true)),
          ),
      )
      .subcommand(
        SubCommand::with_name("directory").subcommand(
          SubCommand::with_name("save")
            .about(
              "Ingest a directory recursively. Saves all files found therein and saves Directory \
protos for each directory found. Outputs a fingerprint of the canonical top-level Directory proto \
and the size of the serialized proto in bytes, separated by a space.",
            )
            .arg(Arg::with_name("source").required(true).takes_value(true)),
        ),
      )
      .arg(
        Arg::with_name("local_store_path")
          .takes_value(true)
          .long("local_store_path")
          .required(true),
      )
      .get_matches(),
  ) {
    Ok(_) => {}
    Err(err) => {
      eprintln!("{}", err.0);
      exit(err.1 as i32)
    }
  };
}

fn execute(top_match: clap::ArgMatches) -> Result<(), ExitError> {
  let store_dir = top_match.value_of("local_store_path").unwrap();
  let store = Arc::new(Store::new(store_dir).map_err(|e| {
    format!(
      "Failed to open/create store for directory {}: {}",
      store_dir,
      e
    )
  })?);

  match top_match.subcommand() {
    ("file", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("cat", Some(args)) => {
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          match store.load_bytes(&fingerprint)? {
            Some(bytes) => {
              io::stdout().write(&bytes).unwrap();
              Ok(())
            }
            None => Err(ExitError(
              format!("File with fingerprint {} not found", fingerprint),
              ExitCode::NotFound,
            )),
          }
        }
        ("save", Some(args)) => {
          let path = PathBuf::from(args.value_of("path").unwrap());
          let posix_fs = make_posix_fs(path.parent().unwrap());
          let file = posix_fs
            .stat(PathBuf::from(path.file_name().unwrap()))
            .unwrap();
          match file {
            fs::Stat::File(f) => {
              let (fingerprint, size_bytes) = save_file(store, &posix_fs, f).wait().unwrap();
              Ok(println!("{} {}", fingerprint, size_bytes))
            }
            o => Err(ExitError(
              format!(
                "Tried to save file {:?} but it was not a file, was a {:?}",
                path,
                o
              ),
              ExitCode::UnknownError,
            )),
          }

        }
        (_, _) => unimplemented!(),
      }
    }
    ("directory", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("save", Some(args)) => {
          let posix_fs = Arc::new(make_posix_fs(args.value_of("source").unwrap()));
          let (fingerprint, size_bytes) =
            save_directory(store, posix_fs, Arc::new(fs::Dir(PathBuf::from("."))))
              .wait()?;
          Ok(println!("{} {}", fingerprint, size_bytes))
        }
        (_, _) => unimplemented!(),
      }
    }
    (_, _) => unimplemented!(),
  }
}

fn make_posix_fs<P: AsRef<Path>>(root: P) -> fs::PosixFS {
  fs::PosixFS::new(&root, vec![]).unwrap()
}

fn save_file(
  store: Arc<Store>,
  posix_fs: &fs::PosixFS,
  file: fs::File,
) -> BoxFuture<(Fingerprint, usize), String> {
  posix_fs
    .read_file(&file)
    .map_err(move |err| {
      format!("Error reading file {:?}: {}", file, err.description())
    })
    .and_then(move |content| {
      Ok((
        store.store_file_bytes(&content.content)?,
        content.content.len(),
      ))
    })
    .to_boxed()
}

fn save_directory(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  base_dir: Arc<fs::Dir>,
) -> BoxFuture<(Fingerprint, usize), String> {
  let store_copy = store.clone();
  Box::new(
    posix_fs
      .scandir(&base_dir)
      .map_err(move |err| {
        format!("Failed to scandir {:?}: {}", &base_dir, err.description())
      })
      .and_then(move |stats| {
        let mut dirs = Vec::new();
        let mut files = Vec::new();
        for stat in stats.into_iter() {
          match stat {
            fs::Stat::Dir(dir) => dirs.push(dir),
            fs::Stat::File(file) => files.push(file),
            fs::Stat::Link(fs::Link(path)) => {
              let err: BoxFuture<_, _> =
                futures::future::err(format!("Don't yet know how to handle symlinks: {:?}", path))
                  .to_boxed();
              return err;
            }
          }
        }

        let store_inner_copy = store_copy.clone();
        let posix_fs_copy = posix_fs.clone();
        let dir_futures: Vec<_> = dirs
          .into_iter()
          .map(move |dir| {
            let dir = Arc::new(dir);
            save_directory(store_inner_copy.clone(), posix_fs_copy.clone(), dir.clone())
              .and_then(move |(fingerprint, proto_size_bytes)| {
                let mut dir_node = bazel_protos::remote_execution::DirectoryNode::new();
                dir_node.set_name(file_name_as_utf8(&dir.0)?);
                dir_node.set_digest(fingerprint_to_digest(&fingerprint, proto_size_bytes as i64));
                Ok(dir_node)
              })
          })
          .collect();

        let file_futures: Vec<_> = files
          .into_iter()
          .map(move |file| {
            file_to_file_node(store_copy.clone(), posix_fs.clone(), file)
          })
          .collect();

        join_all(dir_futures)
          .join(join_all(file_futures))
          .to_boxed()
      })
      .and_then(move |(dirs, files)| {
        let mut directory = bazel_protos::remote_execution::Directory::new();
        directory.set_directories(protobuf::RepeatedField::from_vec(dirs));
        directory.set_files(protobuf::RepeatedField::from_vec(files));
        store.record_directory(&directory)
      }),
  )
}

fn file_to_file_node(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  file: fs::File,
) -> BoxFuture<bazel_protos::remote_execution::FileNode, String> {
  let file_copy = file.clone();
  save_file(store, &posix_fs, file)
    .and_then(move |(fingerprint, size_bytes)| {
      let mut file_node = bazel_protos::remote_execution::FileNode::new();
      file_node.set_name(file_name_as_utf8(&file_copy.path)?);
      file_node.set_digest(fingerprint_to_digest(&fingerprint, size_bytes as i64));
      file_node.set_is_executable(file_copy.is_executable);
      Ok(file_node)
    })
    .to_boxed()
}

pub fn fingerprint_to_digest(
  fingerprint: &Fingerprint,
  size_bytes: i64,
) -> bazel_protos::remote_execution::Digest {
  let mut digest = bazel_protos::remote_execution::Digest::new();
  digest.set_hash(fingerprint.to_hex());
  digest.set_size_bytes(size_bytes);
  digest
}

fn file_name_as_utf8(path: &Path) -> Result<String, String> {
  match path.file_name() {
    Some(name) => {
      match name.to_str() {
        Some(name_utf8) => Ok(name_utf8.to_string()),
        None => Err(format!(
          "{:?}'s file_name is not representable in UTF8",
          path
        )),
      }
    }
    None => Err(format!("{:?} did not have a file_name", path)),
  }
}
<|endoftext|>


<commit_before><commit_msg>Implement the entire range of Iterator<commit_after><|endoftext|>


<commit_before>// Rust JSON-RPC Library
// Written in 2014 by
//   Andrew Poelstra <apoelstra@wpsoftware.net>
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # JSON RPC-Server Support
//!

use http::server;
use http::headers::content_type::MediaType;

use std::io::net::ip::{SocketAddr, Ipv4Addr};

use serialize::json;
use time;

use error::{Error, result_to_response};
use Request;
use decode::{json_decode_field, json_decode_field_list, json_decode_field_string};

/// A Server which reacts to JSONRPC requests by passing the request,
/// along with a response channel, across a channel.
#[deriving(Clone)]
pub struct JsonRpcServer {
  req_tx: Sender<(Request, Sender<Result<json::Json, Error>>)>,
}

impl JsonRpcServer {
  /// Constructor: returns a new `JsonRpcServer` along with a `Receiver`
  /// which should be listened on for new requests from peers
  pub fn new() -> (JsonRpcServer, Receiver<(Request, Sender<Result<json::Json, Error>>)>) {
    let (tx, rx) = channel();
    (JsonRpcServer { req_tx: tx }, rx)
  }
}

impl server::Server for JsonRpcServer {
  fn get_config (&self) -> server::Config {
    server::Config { bind_address: SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 8001 } }
  }

  fn handle_request (&self, r: server::Request, w: &mut server::ResponseWriter) {
    w.headers.date = Some (time::now_utc());
    w.headers.content_type = Some (MediaType {
      type_: "application".to_string(),
      subtype: "json".to_string(),
      parameters: vec![("charset".to_string(), "UTF-8".to_string())]
    });
    w.headers.server = Some("rust-bitcoin-rpc".to_string());

    // Decode the request
    let js = json::from_str(r.body.as_slice());
    if js.is_err() { return; }
    let js = js.unwrap();
    let id = json_decode_field(&js, "id");
    let method = json_decode_field_string(&js, "method");
    let params = json_decode_field_list(&js, "params");

    // If it was a valid request, pass it to the main thread to get a response
    let response = if id.is_ok() && method.is_ok() && params.is_ok() {
      let (resp_tx, resp_rx) = channel();
      let request = Request {
        id: id.clone().unwrap(),
        method: method.unwrap(),
        params: params.unwrap()
      };
      // Request response from main thread
      self.req_tx.send((request, resp_tx));
      result_to_response(resp_rx.recv(), id.unwrap())
    // Otherwise use the error as a response
    } else if id.is_err() {
      result_to_response(id, json::Null)
    } else if method.is_err() {
      result_to_response(method.map(|meth| json::String(meth)), id.unwrap())
    } else if params.is_err() {
      result_to_response(params.map(|parms| json::List(parms)), id.unwrap())
    // Romy says, `else` as code for `else if [last thing i can think of]` is a
    // subtle bug waiting to happen. So actually use `else if` and add an unreachable
    // clause to satisfy the code-path checker
    } else {
      unreachable!()
    };

    // Pass the response to the peer
    let reply_str = json::encode(&response);
    let reply_bytes = reply_str.as_bytes();

    w.headers.content_length = Some(reply_bytes.len());
    if w.write(reply_bytes).is_err() {
      return;
    }
  }
}

<commit_msg>Make address/port configurable<commit_after>// Rust JSON-RPC Library
// Written in 2014 by
//   Andrew Poelstra <apoelstra@wpsoftware.net>
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # JSON RPC-Server Support
//!

use http::server;
use http::headers::content_type::MediaType;

use std::from_str::FromStr;
use std::io::net::ip::{SocketAddr, IpAddr};

use serialize::json;
use time;

use error::{Error, result_to_response};
use Request;
use decode::{json_decode_field, json_decode_field_list, json_decode_field_string};

/// A Server which reacts to JSONRPC requests by passing the request,
/// along with a response channel, across a channel.
#[deriving(Clone)]
pub struct JsonRpcServer {
  address: IpAddr,
  port: u16,
  req_tx: Sender<(Request, Sender<Result<json::Json, Error>>)>,
}

impl JsonRpcServer {
  /// Constructor: returns a new `JsonRpcServer` along with a `Receiver`
  /// which should be listened on for new requests from peers
  pub fn new(addr: &str, port: u16) -> Result<(JsonRpcServer, Receiver<(Request, Sender<Result<json::Json, Error>>)>), String> {
    let (tx, rx) = channel();
    match FromStr::from_str(addr) {
      Some(addr) => Ok((JsonRpcServer { address: addr, port: port, req_tx: tx }, rx)),
      None => Err(format!("Unable to get IP address from `{}`", addr))
    }
  }
}

impl server::Server for JsonRpcServer {
  fn get_config (&self) -> server::Config {
    server::Config { bind_address: SocketAddr { ip: self.address, port: self.port } }
  }

  fn handle_request (&self, r: server::Request, w: &mut server::ResponseWriter) {
    w.headers.date = Some (time::now_utc());
    w.headers.content_type = Some (MediaType {
      type_: "application".to_string(),
      subtype: "json".to_string(),
      parameters: vec![("charset".to_string(), "UTF-8".to_string())]
    });
    w.headers.server = Some("rust-bitcoin-rpc".to_string());

    // Decode the request
    let js = json::from_str(r.body.as_slice());
    if js.is_err() { return; }
    let js = js.unwrap();
    let id = json_decode_field(&js, "id");
    let method = json_decode_field_string(&js, "method");
    let params = json_decode_field_list(&js, "params");

    // If it was a valid request, pass it to the main thread to get a response
    let response = if id.is_ok() && method.is_ok() && params.is_ok() {
      let (resp_tx, resp_rx) = channel();
      let request = Request {
        id: id.clone().unwrap(),
        method: method.unwrap(),
        params: params.unwrap()
      };
      // Request response from main thread
      self.req_tx.send((request, resp_tx));
      result_to_response(resp_rx.recv(), id.unwrap())
    // Otherwise use the error as a response
    } else if id.is_err() {
      result_to_response(id, json::Null)
    } else if method.is_err() {
      result_to_response(method.map(|meth| json::String(meth)), id.unwrap())
    } else if params.is_err() {
      result_to_response(params.map(|parms| json::List(parms)), id.unwrap())
    // Romy says, `else` as code for `else if [last thing i can think of]` is a
    // subtle bug waiting to happen. So actually use `else if` and add an unreachable
    // clause to satisfy the code-path checker
    } else {
      unreachable!()
    };

    // Pass the response to the peer
    let reply_str = json::encode(&response);
    let reply_bytes = reply_str.as_bytes();

    w.headers.content_length = Some(reply_bytes.len());
    if w.write(reply_bytes).is_err() {
      return;
    }
  }
}

<|endoftext|>