<commit_before>// Copied from nom master:
// https://github.com/Geal/nom/blob/a38188f333c29d00c32a3082bec5491d2eefa33f/src/sequence.rs#L591-L687
// Will be released in nom 2.0.

#![cfg(feature = "parsing")]

#[macro_export]
macro_rules! do_parse (
  ($i:expr, $($rest:tt)*) => (
    {
      do_parse_impl!($i, 0usize, $($rest)*)
    }
  );
);

/// Internal parser, do not use directly
#[doc(hidden)]
#[macro_export]
macro_rules! do_parse_impl (

  ($i:expr, $consumed:expr, ( $($rest:expr),* )) => (
    ::nom::IResult::Done($i, ( $($rest),* ))
  );

  ($i:expr, $consumed:expr, $e:ident >> $($rest:tt)*) => (
    do_parse_impl!($i, $consumed, call!($e) >> $($rest)*);
  );
  ($i:expr, $consumed:expr, $submac:ident!( $($args:tt)* ) >> $($rest:tt)*) => (
    {
      match $submac!($i, $($args)*) {
        ::nom::IResult::Error(e)      => ::nom::IResult::Error(e),
        ::nom::IResult::Incomplete(::nom::Needed::Unknown) =>
          ::nom::IResult::Incomplete(::nom::Needed::Unknown),
        ::nom::IResult::Incomplete(::nom::Needed::Size(i)) =>
          ::nom::IResult::Incomplete(::nom::Needed::Size($consumed + i)),
        ::nom::IResult::Done(i,_)     => {
          do_parse_impl!(i,
            $consumed + (::nom::InputLength::input_len(&($i)) -
                         ::nom::InputLength::input_len(&i)), $($rest)*)
        },
      }
    }
  );

  ($i:expr, $consumed:expr, $field:ident : $e:ident >> $($rest:tt)*) => (
    do_parse_impl!($i, $consumed, $field: call!($e) >> $($rest)*);
  );

  ($i:expr, $consumed:expr, $field:ident : $submac:ident!( $($args:tt)* ) >> $($rest:tt)*) => (
    {
      match  $submac!($i, $($args)*) {
        ::nom::IResult::Error(e)      => ::nom::IResult::Error(e),
        ::nom::IResult::Incomplete(::nom::Needed::Unknown) =>
          ::nom::IResult::Incomplete(::nom::Needed::Unknown),
        ::nom::IResult::Incomplete(::nom::Needed::Size(i)) =>
          ::nom::IResult::Incomplete(::nom::Needed::Size($consumed + i)),
        ::nom::IResult::Done(i,o)     => {
          let $field = o;
          do_parse_impl!(i,
            $consumed + (::nom::InputLength::input_len(&($i)) -
                         ::nom::InputLength::input_len(&i)), $($rest)*)
        },
      }
    }
  );

  // ending the chain
  ($i:expr, $consumed:expr, $e:ident >> ( $($rest:tt)* )) => (
    do_parse_impl!($i, $consumed, call!($e) >> ( $($rest)* ));
  );

  ($i:expr, $consumed:expr, $submac:ident!( $($args:tt)* ) >> ( $($rest:tt)* )) => (
    match $submac!($i, $($args)*) {
      ::nom::IResult::Error(e)      => ::nom::IResult::Error(e),
      ::nom::IResult::Incomplete(::nom::Needed::Unknown) =>
        ::nom::IResult::Incomplete(::nom::Needed::Unknown),
      ::nom::IResult::Incomplete(::nom::Needed::Size(i)) =>
        ::nom::IResult::Incomplete(::nom::Needed::Size($consumed + i)),
      ::nom::IResult::Done(i,_)     => {
        ::nom::IResult::Done(i, ( $($rest)* ))
      },
    }
  );

  ($i:expr, $consumed:expr, $field:ident : $e:ident >> ( $($rest:tt)* )) => (
    do_parse_impl!($i, $consumed, $field: call!($e) >> ( $($rest)* ) );
  );

  ($i:expr, $consumed:expr, $field:ident : $submac:ident!( $($args:tt)* ) >> ( $($rest:tt)* )) => (
    match $submac!($i, $($args)*) {
      ::nom::IResult::Error(e)      => ::nom::IResult::Error(e),
      ::nom::IResult::Incomplete(::nom::Needed::Unknown) =>
        ::nom::IResult::Incomplete(::nom::Needed::Unknown),
      ::nom::IResult::Incomplete(::nom::Needed::Size(i)) =>
        ::nom::IResult::Incomplete(::nom::Needed::Size($consumed + i)),
      ::nom::IResult::Done(i,o)     => {
        let $field = o;
        ::nom::IResult::Done(i, ( $($rest)* ))
      },
    }
  );

);
<commit_msg>Hide do_parse macro from docs<commit_after>// Copied from nom master:
// https://github.com/Geal/nom/blob/a38188f333c29d00c32a3082bec5491d2eefa33f/src/sequence.rs#L591-L687
// Will be released in nom 2.0.

#![cfg(feature = "parsing")]

#[macro_export]
#[doc(hidden)]
macro_rules! do_parse (
  ($i:expr, $($rest:tt)*) => (
    {
      do_parse_impl!($i, 0usize, $($rest)*)
    }
  );
);

/// Internal parser, do not use directly
#[doc(hidden)]
#[macro_export]
macro_rules! do_parse_impl (

  ($i:expr, $consumed:expr, ( $($rest:expr),* )) => (
    ::nom::IResult::Done($i, ( $($rest),* ))
  );

  ($i:expr, $consumed:expr, $e:ident >> $($rest:tt)*) => (
    do_parse_impl!($i, $consumed, call!($e) >> $($rest)*);
  );
  ($i:expr, $consumed:expr, $submac:ident!( $($args:tt)* ) >> $($rest:tt)*) => (
    {
      match $submac!($i, $($args)*) {
        ::nom::IResult::Error(e)      => ::nom::IResult::Error(e),
        ::nom::IResult::Incomplete(::nom::Needed::Unknown) =>
          ::nom::IResult::Incomplete(::nom::Needed::Unknown),
        ::nom::IResult::Incomplete(::nom::Needed::Size(i)) =>
          ::nom::IResult::Incomplete(::nom::Needed::Size($consumed + i)),
        ::nom::IResult::Done(i,_)     => {
          do_parse_impl!(i,
            $consumed + (::nom::InputLength::input_len(&($i)) -
                         ::nom::InputLength::input_len(&i)), $($rest)*)
        },
      }
    }
  );

  ($i:expr, $consumed:expr, $field:ident : $e:ident >> $($rest:tt)*) => (
    do_parse_impl!($i, $consumed, $field: call!($e) >> $($rest)*);
  );

  ($i:expr, $consumed:expr, $field:ident : $submac:ident!( $($args:tt)* ) >> $($rest:tt)*) => (
    {
      match  $submac!($i, $($args)*) {
        ::nom::IResult::Error(e)      => ::nom::IResult::Error(e),
        ::nom::IResult::Incomplete(::nom::Needed::Unknown) =>
          ::nom::IResult::Incomplete(::nom::Needed::Unknown),
        ::nom::IResult::Incomplete(::nom::Needed::Size(i)) =>
          ::nom::IResult::Incomplete(::nom::Needed::Size($consumed + i)),
        ::nom::IResult::Done(i,o)     => {
          let $field = o;
          do_parse_impl!(i,
            $consumed + (::nom::InputLength::input_len(&($i)) -
                         ::nom::InputLength::input_len(&i)), $($rest)*)
        },
      }
    }
  );

  // ending the chain
  ($i:expr, $consumed:expr, $e:ident >> ( $($rest:tt)* )) => (
    do_parse_impl!($i, $consumed, call!($e) >> ( $($rest)* ));
  );

  ($i:expr, $consumed:expr, $submac:ident!( $($args:tt)* ) >> ( $($rest:tt)* )) => (
    match $submac!($i, $($args)*) {
      ::nom::IResult::Error(e)      => ::nom::IResult::Error(e),
      ::nom::IResult::Incomplete(::nom::Needed::Unknown) =>
        ::nom::IResult::Incomplete(::nom::Needed::Unknown),
      ::nom::IResult::Incomplete(::nom::Needed::Size(i)) =>
        ::nom::IResult::Incomplete(::nom::Needed::Size($consumed + i)),
      ::nom::IResult::Done(i,_)     => {
        ::nom::IResult::Done(i, ( $($rest)* ))
      },
    }
  );

  ($i:expr, $consumed:expr, $field:ident : $e:ident >> ( $($rest:tt)* )) => (
    do_parse_impl!($i, $consumed, $field: call!($e) >> ( $($rest)* ) );
  );

  ($i:expr, $consumed:expr, $field:ident : $submac:ident!( $($args:tt)* ) >> ( $($rest:tt)* )) => (
    match $submac!($i, $($args)*) {
      ::nom::IResult::Error(e)      => ::nom::IResult::Error(e),
      ::nom::IResult::Incomplete(::nom::Needed::Unknown) =>
        ::nom::IResult::Incomplete(::nom::Needed::Unknown),
      ::nom::IResult::Incomplete(::nom::Needed::Size(i)) =>
        ::nom::IResult::Incomplete(::nom::Needed::Size($consumed + i)),
      ::nom::IResult::Done(i,o)     => {
        let $field = o;
        ::nom::IResult::Done(i, ( $($rest)* ))
      },
    }
  );

);
<|endoftext|>


<commit_before>use crate::base::ExtCtxt;
use crate::expand::{AstFragment, AstFragmentKind};

use rustc_ast as ast;
use rustc_ast::mut_visit::*;
use rustc_ast::ptr::P;
use rustc_span::source_map::DUMMY_SP;
use rustc_span::symbol::Ident;

use smallvec::{smallvec, SmallVec};

use rustc_data_structures::fx::FxHashMap;

pub fn placeholder(
    kind: AstFragmentKind,
    id: ast::NodeId,
    vis: Option<ast::Visibility>,
) -> AstFragment {
    fn mac_placeholder() -> ast::MacCall {
        ast::MacCall {
            path: ast::Path { span: DUMMY_SP, segments: Vec::new(), tokens: None },
            args: P(ast::MacArgs::Empty),
            prior_type_ascription: None,
        }
    }

    let ident = Ident::invalid();
    let attrs = Vec::new();
    let vis = vis.unwrap_or(ast::Visibility {
        span: DUMMY_SP,
        kind: ast::VisibilityKind::Inherited,
        tokens: None,
    });
    let span = DUMMY_SP;
    let expr_placeholder = || {
        P(ast::Expr {
            id,
            span,
            attrs: ast::AttrVec::new(),
            kind: ast::ExprKind::MacCall(mac_placeholder()),
            tokens: None,
        })
    };
    let ty =
        || P(ast::Ty { id, kind: ast::TyKind::MacCall(mac_placeholder()), span, tokens: None });
    let pat =
        || P(ast::Pat { id, kind: ast::PatKind::MacCall(mac_placeholder()), span, tokens: None });

    match kind {
        AstFragmentKind::Expr => AstFragment::Expr(expr_placeholder()),
        AstFragmentKind::OptExpr => AstFragment::OptExpr(Some(expr_placeholder())),
        AstFragmentKind::Items => AstFragment::Items(smallvec![P(ast::Item {
            id,
            span,
            ident,
            vis,
            attrs,
            kind: ast::ItemKind::MacCall(mac_placeholder()),
            tokens: None,
        })]),
        AstFragmentKind::TraitItems => AstFragment::TraitItems(smallvec![P(ast::AssocItem {
            id,
            span,
            ident,
            vis,
            attrs,
            kind: ast::AssocItemKind::MacCall(mac_placeholder()),
            tokens: None,
        })]),
        AstFragmentKind::ImplItems => AstFragment::ImplItems(smallvec![P(ast::AssocItem {
            id,
            span,
            ident,
            vis,
            attrs,
            kind: ast::AssocItemKind::MacCall(mac_placeholder()),
            tokens: None,
        })]),
        AstFragmentKind::ForeignItems => {
            AstFragment::ForeignItems(smallvec![P(ast::ForeignItem {
                id,
                span,
                ident,
                vis,
                attrs,
                kind: ast::ForeignItemKind::MacCall(mac_placeholder()),
                tokens: None,
            })])
        }
        AstFragmentKind::Pat => AstFragment::Pat(P(ast::Pat {
            id,
            span,
            kind: ast::PatKind::MacCall(mac_placeholder()),
            tokens: None,
        })),
        AstFragmentKind::Ty => AstFragment::Ty(P(ast::Ty {
            id,
            span,
            kind: ast::TyKind::MacCall(mac_placeholder()),
            tokens: None,
        })),
        AstFragmentKind::Stmts => AstFragment::Stmts(smallvec![{
            let mac = P(ast::MacCallStmt {
                mac: mac_placeholder(),
                style: ast::MacStmtStyle::Braces,
                attrs: ast::AttrVec::new(),
            });
            ast::Stmt { id, span, kind: ast::StmtKind::MacCall(mac), tokens: None }
        }]),
        AstFragmentKind::Arms => AstFragment::Arms(smallvec![ast::Arm {
            attrs: Default::default(),
            body: expr_placeholder(),
            guard: None,
            id,
            pat: pat(),
            span,
            is_placeholder: true,
        }]),
        AstFragmentKind::Fields => AstFragment::Fields(smallvec![ast::Field {
            attrs: Default::default(),
            expr: expr_placeholder(),
            id,
            ident,
            is_shorthand: false,
            span,
            is_placeholder: true,
        }]),
        AstFragmentKind::FieldPats => AstFragment::FieldPats(smallvec![ast::FieldPat {
            attrs: Default::default(),
            id,
            ident,
            is_shorthand: false,
            pat: pat(),
            span,
            is_placeholder: true,
        }]),
        AstFragmentKind::GenericParams => AstFragment::GenericParams(smallvec![{
            ast::GenericParam {
                attrs: Default::default(),
                bounds: Default::default(),
                id,
                ident,
                is_placeholder: true,
                kind: ast::GenericParamKind::Lifetime,
            }
        }]),
        AstFragmentKind::Params => AstFragment::Params(smallvec![ast::Param {
            attrs: Default::default(),
            id,
            pat: pat(),
            span,
            ty: ty(),
            is_placeholder: true,
        }]),
        AstFragmentKind::StructFields => AstFragment::StructFields(smallvec![ast::StructField {
            attrs: Default::default(),
            id,
            ident: None,
            span,
            ty: ty(),
            vis,
            is_placeholder: true,
        }]),
        AstFragmentKind::Variants => AstFragment::Variants(smallvec![ast::Variant {
            attrs: Default::default(),
            data: ast::VariantData::Struct(Default::default(), false),
            disr_expr: None,
            id,
            ident,
            span,
            vis,
            is_placeholder: true,
        }]),
    }
}

pub struct PlaceholderExpander<'a, 'b> {
    expanded_fragments: FxHashMap<ast::NodeId, AstFragment>,
    cx: &'a mut ExtCtxt<'b>,
    monotonic: bool,
}

impl<'a, 'b> PlaceholderExpander<'a, 'b> {
    pub fn new(cx: &'a mut ExtCtxt<'b>, monotonic: bool) -> Self {
        PlaceholderExpander { cx, expanded_fragments: FxHashMap::default(), monotonic }
    }

    pub fn add(&mut self, id: ast::NodeId, mut fragment: AstFragment) {
        fragment.mut_visit_with(self);
        self.expanded_fragments.insert(id, fragment);
    }

    fn remove(&mut self, id: ast::NodeId) -> AstFragment {
        self.expanded_fragments.remove(&id).unwrap()
    }
}

impl<'a, 'b> MutVisitor for PlaceholderExpander<'a, 'b> {
    fn flat_map_arm(&mut self, arm: ast::Arm) -> SmallVec<[ast::Arm; 1]> {
        if arm.is_placeholder {
            self.remove(arm.id).make_arms()
        } else {
            noop_flat_map_arm(arm, self)
        }
    }

    fn flat_map_field(&mut self, field: ast::Field) -> SmallVec<[ast::Field; 1]> {
        if field.is_placeholder {
            self.remove(field.id).make_fields()
        } else {
            noop_flat_map_field(field, self)
        }
    }

    fn flat_map_field_pattern(&mut self, fp: ast::FieldPat) -> SmallVec<[ast::FieldPat; 1]> {
        if fp.is_placeholder {
            self.remove(fp.id).make_field_patterns()
        } else {
            noop_flat_map_field_pattern(fp, self)
        }
    }

    fn flat_map_generic_param(
        &mut self,
        param: ast::GenericParam,
    ) -> SmallVec<[ast::GenericParam; 1]> {
        if param.is_placeholder {
            self.remove(param.id).make_generic_params()
        } else {
            noop_flat_map_generic_param(param, self)
        }
    }

    fn flat_map_param(&mut self, p: ast::Param) -> SmallVec<[ast::Param; 1]> {
        if p.is_placeholder {
            self.remove(p.id).make_params()
        } else {
            noop_flat_map_param(p, self)
        }
    }

    fn flat_map_struct_field(&mut self, sf: ast::StructField) -> SmallVec<[ast::StructField; 1]> {
        if sf.is_placeholder {
            self.remove(sf.id).make_struct_fields()
        } else {
            noop_flat_map_struct_field(sf, self)
        }
    }

    fn flat_map_variant(&mut self, variant: ast::Variant) -> SmallVec<[ast::Variant; 1]> {
        if variant.is_placeholder {
            self.remove(variant.id).make_variants()
        } else {
            noop_flat_map_variant(variant, self)
        }
    }

    fn flat_map_item(&mut self, item: P<ast::Item>) -> SmallVec<[P<ast::Item>; 1]> {
        match item.kind {
            ast::ItemKind::MacCall(_) => return self.remove(item.id).make_items(),
            ast::ItemKind::MacroDef(_) => return smallvec![item],
            _ => {}
        }

        noop_flat_map_item(item, self)
    }

    fn flat_map_trait_item(&mut self, item: P<ast::AssocItem>) -> SmallVec<[P<ast::AssocItem>; 1]> {
        match item.kind {
            ast::AssocItemKind::MacCall(_) => self.remove(item.id).make_trait_items(),
            _ => noop_flat_map_assoc_item(item, self),
        }
    }

    fn flat_map_impl_item(&mut self, item: P<ast::AssocItem>) -> SmallVec<[P<ast::AssocItem>; 1]> {
        match item.kind {
            ast::AssocItemKind::MacCall(_) => self.remove(item.id).make_impl_items(),
            _ => noop_flat_map_assoc_item(item, self),
        }
    }

    fn flat_map_foreign_item(
        &mut self,
        item: P<ast::ForeignItem>,
    ) -> SmallVec<[P<ast::ForeignItem>; 1]> {
        match item.kind {
            ast::ForeignItemKind::MacCall(_) => self.remove(item.id).make_foreign_items(),
            _ => noop_flat_map_foreign_item(item, self),
        }
    }

    fn visit_expr(&mut self, expr: &mut P<ast::Expr>) {
        match expr.kind {
            ast::ExprKind::MacCall(_) => *expr = self.remove(expr.id).make_expr(),
            _ => noop_visit_expr(expr, self),
        }
    }

    fn filter_map_expr(&mut self, expr: P<ast::Expr>) -> Option<P<ast::Expr>> {
        match expr.kind {
            ast::ExprKind::MacCall(_) => self.remove(expr.id).make_opt_expr(),
            _ => noop_filter_map_expr(expr, self),
        }
    }

    fn flat_map_stmt(&mut self, stmt: ast::Stmt) -> SmallVec<[ast::Stmt; 1]> {
        let (style, mut stmts) = match stmt.kind {
            ast::StmtKind::MacCall(mac) => (mac.style, self.remove(stmt.id).make_stmts()),
            _ => return noop_flat_map_stmt(stmt, self),
        };

        if style == ast::MacStmtStyle::Semicolon {
            // Implement the proposal described in
            // https://github.com/rust-lang/rust/issues/61733#issuecomment-509626449
            //
            // The macro invocation expands to the list of statements.
            // If the list of statements is empty, then 'parse'
            // the trailing semicolon on the original invocation
            // as an empty statement. That is:
            //
            // `empty();` is parsed as a single `StmtKind::Empty`
            //
            // If the list of statements is non-empty, see if the
            // final statement alreayd has a trailing semicolon.
            //
            // If it doesn't have a semicolon, then 'parse' the trailing semicolon
            // from the invocation as part of the final statement,
            // using `stmt.add_trailing_semicolon()`
            //
            // If it does have a semicolon, then 'parse' the trailing semicolon
            // from the invocation as a new StmtKind::Empty

            // FIXME: We will need to preserve the original
            // semicolon token and span as part of #15701
            let empty_stmt = ast::Stmt {
                id: ast::DUMMY_NODE_ID,
                kind: ast::StmtKind::Empty,
                span: DUMMY_SP,
                tokens: None,
            };

            if let Some(stmt) = stmts.pop() {
                if stmt.has_trailing_semicolon() {
                    stmts.push(stmt);
                    stmts.push(empty_stmt);
                } else {
                    stmts.push(stmt.add_trailing_semicolon());
                }
            } else {
                stmts.push(empty_stmt);
            }
        }

        stmts
    }

    fn visit_pat(&mut self, pat: &mut P<ast::Pat>) {
        match pat.kind {
            ast::PatKind::MacCall(_) => *pat = self.remove(pat.id).make_pat(),
            _ => noop_visit_pat(pat, self),
        }
    }

    fn visit_ty(&mut self, ty: &mut P<ast::Ty>) {
        match ty.kind {
            ast::TyKind::MacCall(_) => *ty = self.remove(ty.id).make_ty(),
            _ => noop_visit_ty(ty, self),
        }
    }

    fn visit_block(&mut self, block: &mut P<ast::Block>) {
        noop_visit_block(block, self);

        for stmt in block.stmts.iter_mut() {
            if self.monotonic {
                assert_eq!(stmt.id, ast::DUMMY_NODE_ID);
                stmt.id = self.cx.resolver.next_node_id();
            }
        }
    }

    fn visit_mod(&mut self, module: &mut ast::Mod) {
        noop_visit_mod(module, self);
        // remove macro definitions
        module.items.retain(
            |item| !matches!(item.kind, ast::ItemKind::MacCall(_) if !self.cx.ecfg.keep_macs),
        );
    }

    fn visit_mac(&mut self, _mac: &mut ast::MacCall) {
        // Do nothing.
    }
}
<commit_msg>Rollup merge of #78861 - o752d:patch-1, r=jonas-schievink<commit_after>use crate::base::ExtCtxt;
use crate::expand::{AstFragment, AstFragmentKind};

use rustc_ast as ast;
use rustc_ast::mut_visit::*;
use rustc_ast::ptr::P;
use rustc_span::source_map::DUMMY_SP;
use rustc_span::symbol::Ident;

use smallvec::{smallvec, SmallVec};

use rustc_data_structures::fx::FxHashMap;

pub fn placeholder(
    kind: AstFragmentKind,
    id: ast::NodeId,
    vis: Option<ast::Visibility>,
) -> AstFragment {
    fn mac_placeholder() -> ast::MacCall {
        ast::MacCall {
            path: ast::Path { span: DUMMY_SP, segments: Vec::new(), tokens: None },
            args: P(ast::MacArgs::Empty),
            prior_type_ascription: None,
        }
    }

    let ident = Ident::invalid();
    let attrs = Vec::new();
    let vis = vis.unwrap_or(ast::Visibility {
        span: DUMMY_SP,
        kind: ast::VisibilityKind::Inherited,
        tokens: None,
    });
    let span = DUMMY_SP;
    let expr_placeholder = || {
        P(ast::Expr {
            id,
            span,
            attrs: ast::AttrVec::new(),
            kind: ast::ExprKind::MacCall(mac_placeholder()),
            tokens: None,
        })
    };
    let ty =
        || P(ast::Ty { id, kind: ast::TyKind::MacCall(mac_placeholder()), span, tokens: None });
    let pat =
        || P(ast::Pat { id, kind: ast::PatKind::MacCall(mac_placeholder()), span, tokens: None });

    match kind {
        AstFragmentKind::Expr => AstFragment::Expr(expr_placeholder()),
        AstFragmentKind::OptExpr => AstFragment::OptExpr(Some(expr_placeholder())),
        AstFragmentKind::Items => AstFragment::Items(smallvec![P(ast::Item {
            id,
            span,
            ident,
            vis,
            attrs,
            kind: ast::ItemKind::MacCall(mac_placeholder()),
            tokens: None,
        })]),
        AstFragmentKind::TraitItems => AstFragment::TraitItems(smallvec![P(ast::AssocItem {
            id,
            span,
            ident,
            vis,
            attrs,
            kind: ast::AssocItemKind::MacCall(mac_placeholder()),
            tokens: None,
        })]),
        AstFragmentKind::ImplItems => AstFragment::ImplItems(smallvec![P(ast::AssocItem {
            id,
            span,
            ident,
            vis,
            attrs,
            kind: ast::AssocItemKind::MacCall(mac_placeholder()),
            tokens: None,
        })]),
        AstFragmentKind::ForeignItems => {
            AstFragment::ForeignItems(smallvec![P(ast::ForeignItem {
                id,
                span,
                ident,
                vis,
                attrs,
                kind: ast::ForeignItemKind::MacCall(mac_placeholder()),
                tokens: None,
            })])
        }
        AstFragmentKind::Pat => AstFragment::Pat(P(ast::Pat {
            id,
            span,
            kind: ast::PatKind::MacCall(mac_placeholder()),
            tokens: None,
        })),
        AstFragmentKind::Ty => AstFragment::Ty(P(ast::Ty {
            id,
            span,
            kind: ast::TyKind::MacCall(mac_placeholder()),
            tokens: None,
        })),
        AstFragmentKind::Stmts => AstFragment::Stmts(smallvec![{
            let mac = P(ast::MacCallStmt {
                mac: mac_placeholder(),
                style: ast::MacStmtStyle::Braces,
                attrs: ast::AttrVec::new(),
            });
            ast::Stmt { id, span, kind: ast::StmtKind::MacCall(mac), tokens: None }
        }]),
        AstFragmentKind::Arms => AstFragment::Arms(smallvec![ast::Arm {
            attrs: Default::default(),
            body: expr_placeholder(),
            guard: None,
            id,
            pat: pat(),
            span,
            is_placeholder: true,
        }]),
        AstFragmentKind::Fields => AstFragment::Fields(smallvec![ast::Field {
            attrs: Default::default(),
            expr: expr_placeholder(),
            id,
            ident,
            is_shorthand: false,
            span,
            is_placeholder: true,
        }]),
        AstFragmentKind::FieldPats => AstFragment::FieldPats(smallvec![ast::FieldPat {
            attrs: Default::default(),
            id,
            ident,
            is_shorthand: false,
            pat: pat(),
            span,
            is_placeholder: true,
        }]),
        AstFragmentKind::GenericParams => AstFragment::GenericParams(smallvec![{
            ast::GenericParam {
                attrs: Default::default(),
                bounds: Default::default(),
                id,
                ident,
                is_placeholder: true,
                kind: ast::GenericParamKind::Lifetime,
            }
        }]),
        AstFragmentKind::Params => AstFragment::Params(smallvec![ast::Param {
            attrs: Default::default(),
            id,
            pat: pat(),
            span,
            ty: ty(),
            is_placeholder: true,
        }]),
        AstFragmentKind::StructFields => AstFragment::StructFields(smallvec![ast::StructField {
            attrs: Default::default(),
            id,
            ident: None,
            span,
            ty: ty(),
            vis,
            is_placeholder: true,
        }]),
        AstFragmentKind::Variants => AstFragment::Variants(smallvec![ast::Variant {
            attrs: Default::default(),
            data: ast::VariantData::Struct(Default::default(), false),
            disr_expr: None,
            id,
            ident,
            span,
            vis,
            is_placeholder: true,
        }]),
    }
}

pub struct PlaceholderExpander<'a, 'b> {
    expanded_fragments: FxHashMap<ast::NodeId, AstFragment>,
    cx: &'a mut ExtCtxt<'b>,
    monotonic: bool,
}

impl<'a, 'b> PlaceholderExpander<'a, 'b> {
    pub fn new(cx: &'a mut ExtCtxt<'b>, monotonic: bool) -> Self {
        PlaceholderExpander { cx, expanded_fragments: FxHashMap::default(), monotonic }
    }

    pub fn add(&mut self, id: ast::NodeId, mut fragment: AstFragment) {
        fragment.mut_visit_with(self);
        self.expanded_fragments.insert(id, fragment);
    }

    fn remove(&mut self, id: ast::NodeId) -> AstFragment {
        self.expanded_fragments.remove(&id).unwrap()
    }
}

impl<'a, 'b> MutVisitor for PlaceholderExpander<'a, 'b> {
    fn flat_map_arm(&mut self, arm: ast::Arm) -> SmallVec<[ast::Arm; 1]> {
        if arm.is_placeholder {
            self.remove(arm.id).make_arms()
        } else {
            noop_flat_map_arm(arm, self)
        }
    }

    fn flat_map_field(&mut self, field: ast::Field) -> SmallVec<[ast::Field; 1]> {
        if field.is_placeholder {
            self.remove(field.id).make_fields()
        } else {
            noop_flat_map_field(field, self)
        }
    }

    fn flat_map_field_pattern(&mut self, fp: ast::FieldPat) -> SmallVec<[ast::FieldPat; 1]> {
        if fp.is_placeholder {
            self.remove(fp.id).make_field_patterns()
        } else {
            noop_flat_map_field_pattern(fp, self)
        }
    }

    fn flat_map_generic_param(
        &mut self,
        param: ast::GenericParam,
    ) -> SmallVec<[ast::GenericParam; 1]> {
        if param.is_placeholder {
            self.remove(param.id).make_generic_params()
        } else {
            noop_flat_map_generic_param(param, self)
        }
    }

    fn flat_map_param(&mut self, p: ast::Param) -> SmallVec<[ast::Param; 1]> {
        if p.is_placeholder {
            self.remove(p.id).make_params()
        } else {
            noop_flat_map_param(p, self)
        }
    }

    fn flat_map_struct_field(&mut self, sf: ast::StructField) -> SmallVec<[ast::StructField; 1]> {
        if sf.is_placeholder {
            self.remove(sf.id).make_struct_fields()
        } else {
            noop_flat_map_struct_field(sf, self)
        }
    }

    fn flat_map_variant(&mut self, variant: ast::Variant) -> SmallVec<[ast::Variant; 1]> {
        if variant.is_placeholder {
            self.remove(variant.id).make_variants()
        } else {
            noop_flat_map_variant(variant, self)
        }
    }

    fn flat_map_item(&mut self, item: P<ast::Item>) -> SmallVec<[P<ast::Item>; 1]> {
        match item.kind {
            ast::ItemKind::MacCall(_) => return self.remove(item.id).make_items(),
            ast::ItemKind::MacroDef(_) => return smallvec![item],
            _ => {}
        }

        noop_flat_map_item(item, self)
    }

    fn flat_map_trait_item(&mut self, item: P<ast::AssocItem>) -> SmallVec<[P<ast::AssocItem>; 1]> {
        match item.kind {
            ast::AssocItemKind::MacCall(_) => self.remove(item.id).make_trait_items(),
            _ => noop_flat_map_assoc_item(item, self),
        }
    }

    fn flat_map_impl_item(&mut self, item: P<ast::AssocItem>) -> SmallVec<[P<ast::AssocItem>; 1]> {
        match item.kind {
            ast::AssocItemKind::MacCall(_) => self.remove(item.id).make_impl_items(),
            _ => noop_flat_map_assoc_item(item, self),
        }
    }

    fn flat_map_foreign_item(
        &mut self,
        item: P<ast::ForeignItem>,
    ) -> SmallVec<[P<ast::ForeignItem>; 1]> {
        match item.kind {
            ast::ForeignItemKind::MacCall(_) => self.remove(item.id).make_foreign_items(),
            _ => noop_flat_map_foreign_item(item, self),
        }
    }

    fn visit_expr(&mut self, expr: &mut P<ast::Expr>) {
        match expr.kind {
            ast::ExprKind::MacCall(_) => *expr = self.remove(expr.id).make_expr(),
            _ => noop_visit_expr(expr, self),
        }
    }

    fn filter_map_expr(&mut self, expr: P<ast::Expr>) -> Option<P<ast::Expr>> {
        match expr.kind {
            ast::ExprKind::MacCall(_) => self.remove(expr.id).make_opt_expr(),
            _ => noop_filter_map_expr(expr, self),
        }
    }

    fn flat_map_stmt(&mut self, stmt: ast::Stmt) -> SmallVec<[ast::Stmt; 1]> {
        let (style, mut stmts) = match stmt.kind {
            ast::StmtKind::MacCall(mac) => (mac.style, self.remove(stmt.id).make_stmts()),
            _ => return noop_flat_map_stmt(stmt, self),
        };

        if style == ast::MacStmtStyle::Semicolon {
            // Implement the proposal described in
            // https://github.com/rust-lang/rust/issues/61733#issuecomment-509626449
            //
            // The macro invocation expands to the list of statements. If the
            // list of statements is empty, then 'parse' the trailing semicolon
            // on the original invocation as an empty statement. That is:
            //
            // `empty();` is parsed as a single `StmtKind::Empty`
            //
            // If the list of statements is non-empty, see if the final
            // statement already has a trailing semicolon.
            //
            // If it doesn't have a semicolon, then 'parse' the trailing
            // semicolon from the invocation as part of the final statement,
            // using `stmt.add_trailing_semicolon()`
            //
            // If it does have a semicolon, then 'parse' the trailing semicolon
            // from the invocation as a new StmtKind::Empty

            // FIXME: We will need to preserve the original semicolon token and
            // span as part of #15701
            let empty_stmt = ast::Stmt {
                id: ast::DUMMY_NODE_ID,
                kind: ast::StmtKind::Empty,
                span: DUMMY_SP,
                tokens: None,
            };

            if let Some(stmt) = stmts.pop() {
                if stmt.has_trailing_semicolon() {
                    stmts.push(stmt);
                    stmts.push(empty_stmt);
                } else {
                    stmts.push(stmt.add_trailing_semicolon());
                }
            } else {
                stmts.push(empty_stmt);
            }
        }

        stmts
    }

    fn visit_pat(&mut self, pat: &mut P<ast::Pat>) {
        match pat.kind {
            ast::PatKind::MacCall(_) => *pat = self.remove(pat.id).make_pat(),
            _ => noop_visit_pat(pat, self),
        }
    }

    fn visit_ty(&mut self, ty: &mut P<ast::Ty>) {
        match ty.kind {
            ast::TyKind::MacCall(_) => *ty = self.remove(ty.id).make_ty(),
            _ => noop_visit_ty(ty, self),
        }
    }

    fn visit_block(&mut self, block: &mut P<ast::Block>) {
        noop_visit_block(block, self);

        for stmt in block.stmts.iter_mut() {
            if self.monotonic {
                assert_eq!(stmt.id, ast::DUMMY_NODE_ID);
                stmt.id = self.cx.resolver.next_node_id();
            }
        }
    }

    fn visit_mod(&mut self, module: &mut ast::Mod) {
        noop_visit_mod(module, self);
        // remove macro definitions
        module.items.retain(
            |item| !matches!(item.kind, ast::ItemKind::MacCall(_) if !self.cx.ecfg.keep_macs),
        );
    }

    fn visit_mac(&mut self, _mac: &mut ast::MacCall) {
        // Do nothing.
    }
}
<|endoftext|>


<commit_before>use std::{convert, io};
use std::fmt;
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

const BLOCK_SIZE: usize = 16384;
type Block = Vec<u8>;

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    downloaded: Vec<Option<Block>>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.pieces.len();

        println!("num_pieces = {}", num_pieces);
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces],
            downloaded: vec![None; num_pieces],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.handshake(&metainfo.info_hash));
        loop {
            let message = try!(conn.receive_message());
            println!("{:?}", message);
            try!(conn.process(message));
        }
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(19);
        message.extend("BitTorrent protocol".as_bytes().iter().cloned());
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = convert_big_endian_to_integer(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            let message_id = &message[0];
            let message_body = &message[1..];
            Ok(Message::new(message_id, message_body))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_piece());
                }
            }
            Message::Piece(index, begin, data) => {
                self.downloaded[index] = Some(data);
                  try!(self.request_next_piece());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            let mut message = vec![];
            message.push(0);
            message.push(0);
            message.push(0);
            message.push(1);
            message.push(2);
            try!(self.stream.write_all(&message));
            self.am_i_interested = true;
        }
        Ok(())
    }

    fn request_next_piece(&mut self) -> Result<(), Error> {
        match self.next_piece_to_request() {
            Some(piece_index) => self.send_request(piece_index),
            None => Ok(())
        }
    }

    fn next_piece_to_request(&self) -> Option<usize> {
        for i in 0..self.downloaded.len() {
            if self.downloaded[i].is_none() {
                return Some(i)
            }
        }
        println!("Done!");
        None
    }

    fn send_request(&mut self, piece: usize) -> Result<(), Error> {
        println!("Sending request for {}", piece);
        let num_pieces = self.downloaded.len();
        let request_size = if piece == num_pieces - 1 {
            self.metainfo.info.length as usize - (self.metainfo.info.piece_length as usize * (num_pieces - 1))
        } else {
            BLOCK_SIZE
        };
        println!("Request size: {}", request_size);

        let mut message = vec![];
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(13);
        message.push(6);
        message.extend(convert_usize_to_bytes(piece).iter().cloned());
        message.extend(convert_usize_to_bytes(0).iter().cloned());
        message.extend(convert_usize_to_bytes(request_size).iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request, // TODO add params
    Piece(usize, usize, Block),
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(convert_big_endian_to_integer(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => Message::Request,
            7 => {
                let index = convert_big_endian_to_integer(&body[0..4]);
                let begin = convert_big_endian_to_integer(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, begin, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request => write!(f, "Request"),
             Message::Piece(ref index, ref begin, ref data) => write!(f, "Piece({}, {}, size={})", index, begin, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

const BYTE_0: usize = 256 * 256 * 256;
const BYTE_1: usize = 256 * 256;
const BYTE_2: usize = 256;
const BYTE_3: usize = 1;

fn convert_big_endian_to_integer(bytes: &[u8]) -> usize {
    bytes[0] as usize * BYTE_0 +
    bytes[1] as usize * BYTE_1 +
    bytes[2] as usize * BYTE_2 +
    bytes[3] as usize * BYTE_3
}

fn convert_usize_to_bytes(integer: usize) -> Vec<u8> {
    let mut rest = integer;
    let first = rest / BYTE_0;
    rest -= first * BYTE_0;
    let second = rest / BYTE_1;
    rest -= second * BYTE_1;
    let third = rest / BYTE_2;
    rest -= third * BYTE_2;
    let fourth = rest;
    vec![first as u8, second as u8, third as u8, fourth as u8]
}

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

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<commit_msg>Cleanup.<commit_after>use std::{convert, io};
use std::fmt;
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

const PROTOCOL: &'static str = "BitTorrent protocol";
const BLOCK_SIZE: usize = 16384;
type Block = Vec<u8>;

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    downloaded: Vec<Option<Block>>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.pieces.len();

        println!("num_pieces = {}", num_pieces);
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces],
            downloaded: vec![None; num_pieces],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.handshake(&metainfo.info_hash));
        loop {
            let message = try!(conn.receive_message());
            println!("{:?}", message);
            try!(conn.process(message));
        }
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(PROTOCOL.len() as u8);
        message.extend(PROTOCOL.as_bytes().iter().cloned());
        message.extend((&[0; 8]).iter().cloned());
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = convert_big_endian_to_integer(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            let message_id = &message[0];
            let message_body = &message[1..];
            Ok(Message::new(message_id, message_body))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_piece());
                }
            }
            Message::Piece(index, begin, data) => {
                self.downloaded[index] = Some(data);
                  try!(self.request_next_piece());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            let mut message = vec![];
            message.push(0);
            message.push(0);
            message.push(0);
            message.push(1);
            message.push(2);
            try!(self.stream.write_all(&message));
            self.am_i_interested = true;
        }
        Ok(())
    }

    fn request_next_piece(&mut self) -> Result<(), Error> {
        match self.next_piece_to_request() {
            Some(piece_index) => self.send_request(piece_index),
            None => Ok(())
        }
    }

    fn next_piece_to_request(&self) -> Option<usize> {
        for i in 0..self.downloaded.len() {
            if self.downloaded[i].is_none() {
                return Some(i)
            }
        }
        println!("Done!");
        None
    }

    fn send_request(&mut self, piece: usize) -> Result<(), Error> {
        println!("Sending request for {}", piece);
        let num_pieces = self.downloaded.len();
        let request_size = if piece == num_pieces - 1 {
            self.metainfo.info.length as usize - (self.metainfo.info.piece_length as usize * (num_pieces - 1))
        } else {
            BLOCK_SIZE
        };
        println!("Request size: {}", request_size);

        let mut message = vec![];
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(13);
        message.push(6);
        message.extend(convert_usize_to_bytes(piece).iter().cloned());
        message.extend(convert_usize_to_bytes(0).iter().cloned());
        message.extend(convert_usize_to_bytes(request_size).iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request, // TODO add params
    Piece(usize, usize, Block),
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(convert_big_endian_to_integer(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => Message::Request,
            7 => {
                let index = convert_big_endian_to_integer(&body[0..4]);
                let begin = convert_big_endian_to_integer(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, begin, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request => write!(f, "Request"),
             Message::Piece(ref index, ref begin, ref data) => write!(f, "Piece({}, {}, size={})", index, begin, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

const BYTE_0: usize = 256 * 256 * 256;
const BYTE_1: usize = 256 * 256;
const BYTE_2: usize = 256;
const BYTE_3: usize = 1;

fn convert_big_endian_to_integer(bytes: &[u8]) -> usize {
    bytes[0] as usize * BYTE_0 +
    bytes[1] as usize * BYTE_1 +
    bytes[2] as usize * BYTE_2 +
    bytes[3] as usize * BYTE_3
}

fn convert_usize_to_bytes(integer: usize) -> Vec<u8> {
    let mut rest = integer;
    let first = rest / BYTE_0;
    rest -= first * BYTE_0;
    let second = rest / BYTE_1;
    rest -= second * BYTE_1;
    let third = rest / BYTE_2;
    rest -= third * BYTE_2;
    let fourth = rest;
    vec![first as u8, second as u8, third as u8, fourth as u8]
}

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

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<|endoftext|>


<commit_before>use reqwest::{ self, header };
use std::fs::File;
use std::io::{ BufWriter, Read, Write };
use std::path::Path;

use config::Config;
use error::PackageError;

pub fn download_file<F>(url: &str, to_path: &Path, callback: F) -> Result<u64, PackageError>
    where F: Fn(u64, u64)
{
    // GET the URL, returning an error if necessary.
    let response = reqwest::get(url)?;

    // Try to retrieve the content length.
    let content_length =
        response
        .headers()
        .get::<header::ContentLength>()
        // If the Response contains no content-length, default to 0.
        .map_or(0u64, |x| {
            let header::ContentLength(len) = *x;
            len
        });

    // For every `buffer_size` bytes, call the callback once.
    let Config { buffer_size: buffer_size, .. } = Config::default();

    // `buffer_size` decides how frequent the function calls the callback by
    // checking whether (downloaded bytes) % `buffer_size` is 0.
    let buffer_size = buffer_size.unwrap_or(65536);

    // Obviously, the target file to be written to by the BufWriter.
    let mut target = BufWriter::new(File::create(to_path)?);

    for (index, byte) in response.bytes().enumerate() {
        // Write to the target, returns an error if necessary.
        target.write(&[byte?])?;

        // Checks whether it's time to call the callback, and call it if it is.
        // Plus one to `index` as `index` itself is zero-indexed.
        if index as u64 % buffer_size == 0 {
            callback((index + 1) as u64, content_length);
        };
    }

    // In case the callback never got a chance to be called due to the package
    // itself being too small.
    callback(content_length, content_length);

    Ok(content_length)
}<commit_msg>Allow downloading into memory<commit_after>use reqwest::{ self, header };
use std::fs::File;
use std::io::{ BufWriter, Read, Write };
use std::path::Path;

use config::Config;
use error::PackageError;

fn get_content_length(url: &str) -> Result<u64, PackageError> {
    // GET the URL, returning an error if necessary.
    let response = reqwest::get(url)?;

    // Try to retrieve the content length, and return it.
    Ok(response
    .headers()
    .get::<header::ContentLength>()
    // If the Response contains no content-length, default to 0.
    .map_or(0u64, |x| {
        let header::ContentLength(len) = *x;
        len
    }))
}

pub fn download_file<F>(url: &str, to_path: &Path, callback: F) -> Result<u64, PackageError>
    where F: Fn(u64, u64)
{
    // GET the URL, returning an error if necessary.
    let response = reqwest::get(url)?;

    // Try to retrieve the content length.
    let content_length = get_content_length(url)?;

    // For every `buffer_size` bytes, call the callback once.
    // `buffer_size` decides how frequent the function calls the callback by
    // checking whether (downloaded bytes) % `buffer_size` is 0.
    let Config { buffer_size: buffer_size, .. } = Config::default();
    let buffer_size = buffer_size.unwrap_or(65536);

    // Obviously, the target file to be written to by the BufWriter.
    let mut target = BufWriter::new(File::create(to_path)?);

    for (index, byte) in response.bytes().enumerate() {
        // Write to the target, returns an error if necessary.
        target.write(&[byte?])?;

        // Checks whether it's time to call the callback, and call it if it is.
        // Plus one to `index` as `index` itself is zero-indexed.
        if index as u64 % buffer_size == 0 {
            callback((index + 1) as u64, content_length);
        };
    }

    // In case the callback never got a chance to be called due to the package
    // itself being too small.
    callback(content_length, content_length);

    Ok(content_length)
}

pub fn download_buf(url: &str) -> Result<Vec<u8>, PackageError> {
    // GET the URL, returning an error if necessary.
    let mut response = reqwest::get(url)?;

    // Try to read the Response into memory
    let mut content = Vec::new();
    response.read_to_end(&mut content)?;

    Ok(content)
}<|endoftext|>


<commit_before>use std::{convert, io};
use std::fmt;
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

const PROTOCOL: &'static str = "BitTorrent protocol";
const BLOCK_SIZE: usize = 16384;
type Block = Vec<u8>;

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    downloaded: Vec<Option<Block>>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.pieces.len();
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces],
            downloaded: vec![None; num_pieces],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.handshake(&metainfo.info_hash));
        loop {
            let message = try!(conn.receive_message());
            println!("Recieved: {:?}", message);
            try!(conn.process(message));
        }
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(PROTOCOL.len() as u8);
        message.extend(PROTOCOL.as_bytes().iter().cloned());
        message.extend((&[0; 8]).iter().cloned());
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn send_message(&mut self, message: Message) -> Result<(), Error> {
        println!("Sending: {:?}", message);
        try!(self.stream.write_all(&message.encode()));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = convert_big_endian_to_integer(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            Ok(Message::new(&message[0], &message[1..]))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_piece());
                }
            }
            Message::Piece(index, begin, data) => {
                self.downloaded[index] = Some(data);
                try!(self.request_next_piece());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            self.am_i_interested = true;
            try!(self.send_message(Message::Interested));
        }
        Ok(())
    }

    fn request_next_piece(&mut self) -> Result<(), Error> {
        match self.next_piece_to_request() {
            Some(piece_index) => self.send_request(piece_index),
            None => Ok(())
        }
    }

    fn next_piece_to_request(&self) -> Option<usize> {
        for i in 0..self.downloaded.len() {
            if self.downloaded[i].is_none() {
                return Some(i)
            }
        }
        println!("Done downloading file!");
        None
    }

    fn send_request(&mut self, piece: usize) -> Result<(), Error> {
        let num_pieces = self.downloaded.len();
        let request_size = if piece == num_pieces - 1 {
            self.metainfo.info.length as usize - (self.metainfo.info.piece_length as usize * (num_pieces - 1))
        } else {
            BLOCK_SIZE
        };
        self.send_message(Message::Request(piece, 0, request_size))
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request(usize, usize, usize),
    Piece(usize, usize, Block),
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(convert_big_endian_to_integer(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => {
                let index = convert_big_endian_to_integer(&body[0..4]);
                let begin = convert_big_endian_to_integer(&body[4..8]);
                let offset = convert_big_endian_to_integer(&body[8..12]);
                Message::Request(index, begin, offset)
            },
            7 => {
                let index = convert_big_endian_to_integer(&body[0..4]);
                let begin = convert_big_endian_to_integer(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, begin, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }

    fn encode(self) -> Vec<u8> {
        let mut payload = vec![];
        match self {
            Message::KeepAlive => {},
            Message::Choke => payload.push(0),
            Message::Unchoke => payload.push(1),
            Message::Interested => payload.push(2),
            Message::NotInterested => payload.push(3),
            Message::Have(index) => {
                payload.push(4);
                payload.extend(convert_usize_to_bytes(index).iter().cloned());
            },
            Message::Bitfield(bytes) => {
                payload.push(5);
                payload.extend(bytes);
            },
            Message::Request(index, begin, amount) => {
                payload.push(6);
                payload.extend(convert_usize_to_bytes(index).iter().cloned());
                payload.extend(convert_usize_to_bytes(begin).iter().cloned());
                payload.extend(convert_usize_to_bytes(amount).iter().cloned());
            },
            Message::Piece(index, begin, data) => {
                payload.push(6);
                payload.extend(convert_usize_to_bytes(index).iter().cloned());
                payload.extend(convert_usize_to_bytes(begin).iter().cloned());
                payload.extend(data);
            },
            Message::Cancel => payload.push(8),
            Message::Port => payload.push(9),
        };

        // prepend size
        let mut size = convert_usize_to_bytes(payload.len());
        size.extend(payload);
        size
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request(ref index, ref begin, ref offset) => write!(f, "Request({}, {}, {})", index, begin, offset),
             Message::Piece(ref index, ref begin, ref data) => write!(f, "Piece({}, {}, size={})", index, begin, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

const BYTE_0: usize = 256 * 256 * 256;
const BYTE_1: usize = 256 * 256;
const BYTE_2: usize = 256;
const BYTE_3: usize = 1;

fn convert_big_endian_to_integer(bytes: &[u8]) -> usize {
    bytes[0] as usize * BYTE_0 +
    bytes[1] as usize * BYTE_1 +
    bytes[2] as usize * BYTE_2 +
    bytes[3] as usize * BYTE_3
}

fn convert_usize_to_bytes(integer: usize) -> Vec<u8> {
    let mut rest = integer;
    let first = rest / BYTE_0;
    rest -= first * BYTE_0;
    let second = rest / BYTE_1;
    rest -= second * BYTE_1;
    let third = rest / BYTE_2;
    rest -= third * BYTE_2;
    let fourth = rest;
    vec![first as u8, second as u8, third as u8, fourth as u8]
}

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

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<commit_msg>Cleanup.<commit_after>use std::{convert, io};
use std::fmt;
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

const PROTOCOL: &'static str = "BitTorrent protocol";
const BLOCK_SIZE: usize = 16384;
type Block = Vec<u8>;

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    downloaded: Vec<Option<Block>>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.pieces.len();
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces],
            downloaded: vec![None; num_pieces],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.run());
        Ok(conn)
    }

    fn run(&mut self) -> Result<(), Error> {
        try!(self.send_handshake());
        try!(self.receive_handshake());
        loop {
            let message = try!(self.receive_message());
            println!("Recieved: {:?}", message);
            try!(self.process(message));
        }
        Ok(())
    }

    fn send_handshake(&mut self) -> Result<(), Error> {
        let mut message = vec![];
        message.push(PROTOCOL.len() as u8);
        message.extend(PROTOCOL.bytes());
        message.extend(vec![0; 8].into_iter());
        message.extend(self.metainfo.info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".bytes());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn send_message(&mut self, message: Message) -> Result<(), Error> {
        println!("Sending: {:?}", message);
        try!(self.stream.write_all(&message.serialize()));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = bytes_to_usize(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            Ok(Message::new(&message[0], &message[1..]))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_piece());
                }
            }
            Message::Piece(index, begin, data) => {
                self.downloaded[index] = Some(data);
                try!(self.request_next_piece());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            self.am_i_interested = true;
            try!(self.send_message(Message::Interested));
        }
        Ok(())
    }

    fn request_next_piece(&mut self) -> Result<(), Error> {
        match self.next_piece_to_request() {
            Some(piece_index) => self.send_request(piece_index),
            None => Ok(())
        }
    }

    fn next_piece_to_request(&self) -> Option<usize> {
        for i in 0..self.downloaded.len() {
            if self.downloaded[i].is_none() {
                return Some(i)
            }
        }
        println!("Done downloading file!");
        None
    }

    fn send_request(&mut self, piece: usize) -> Result<(), Error> {
        let num_pieces = self.downloaded.len();
        let request_size = if piece == num_pieces - 1 {
            self.metainfo.info.length as usize - (self.metainfo.info.piece_length as usize * (num_pieces - 1))
        } else {
            BLOCK_SIZE
        };
        self.send_message(Message::Request(piece, 0, request_size))
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request(usize, usize, usize),
    Piece(usize, usize, Block),
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(bytes_to_usize(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => {
                let index = bytes_to_usize(&body[0..4]);
                let begin = bytes_to_usize(&body[4..8]);
                let offset = bytes_to_usize(&body[8..12]);
                Message::Request(index, begin, offset)
            },
            7 => {
                let index = bytes_to_usize(&body[0..4]);
                let begin = bytes_to_usize(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, begin, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }

    fn serialize(self) -> Vec<u8> {
        let mut payload = vec![];
        match self {
            Message::KeepAlive => {},
            Message::Choke => payload.push(0),
            Message::Unchoke => payload.push(1),
            Message::Interested => payload.push(2),
            Message::NotInterested => payload.push(3),
            Message::Have(index) => {
                payload.push(4);
                payload.extend(usize_to_bytes(index).into_iter());
            },
            Message::Bitfield(bytes) => {
                payload.push(5);
                payload.extend(bytes);
            },
            Message::Request(index, begin, amount) => {
                payload.push(6);
                payload.extend(usize_to_bytes(index).into_iter());
                payload.extend(usize_to_bytes(begin).into_iter());
                payload.extend(usize_to_bytes(amount).into_iter());
            },
            Message::Piece(index, begin, data) => {
                payload.push(6);
                payload.extend(usize_to_bytes(index).into_iter());
                payload.extend(usize_to_bytes(begin).into_iter());
                payload.extend(data);
            },
            Message::Cancel => payload.push(8),
            Message::Port => payload.push(9),
        };

        // prepend size
        let mut size = usize_to_bytes(payload.len());
        size.extend(payload);
        size
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request(ref index, ref begin, ref offset) => write!(f, "Request({}, {}, {})", index, begin, offset),
             Message::Piece(ref index, ref begin, ref data) => write!(f, "Piece({}, {}, size={})", index, begin, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

const BYTE_0: usize = 256 * 256 * 256;
const BYTE_1: usize = 256 * 256;
const BYTE_2: usize = 256;
const BYTE_3: usize = 1;

fn bytes_to_usize(bytes: &[u8]) -> usize {
    bytes[0] as usize * BYTE_0 +
    bytes[1] as usize * BYTE_1 +
    bytes[2] as usize * BYTE_2 +
    bytes[3] as usize * BYTE_3
}

fn usize_to_bytes(integer: usize) -> Vec<u8> {
    let mut rest = integer;
    let first = rest / BYTE_0;
    rest -= first * BYTE_0;
    let second = rest / BYTE_1;
    rest -= second * BYTE_1;
    let third = rest / BYTE_2;
    rest -= third * BYTE_2;
    let fourth = rest;
    vec![first as u8, second as u8, third as u8, fourth as u8]
}

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

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<|endoftext|>


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

use super::SharedBuf;
use std::io::{Read, Write};
use tink::keyset::{Reader, Writer};

#[test]
fn test_json_io_unencrypted() {
    tink_mac::init();
    let buf = SharedBuf::new();
    let mut w = tink::keyset::JsonWriter::new(buf.clone());
    let mut r = tink::keyset::JsonReader::new(buf);

    let manager = tink_testutil::new_hmac_keyset_manager();
    let h = manager.handle().expect("cannot get keyset handle");

    let ks1 = tink::keyset::insecure::keyset_material(&h);
    w.write(&ks1).expect("cannot write keyset");

    let ks2 = r.read().expect("cannot read keyset");
    assert_eq!(
        ks1, ks2,
        "written keyset ({:?}) doesn't match read keyset ({:?})",
        ks1, ks2
    );
}

// TODO: get this in sync with key encoding - the Go *test* code
// does a one-line compact text format of the key protobuf (which is
// then base64-encoded by the standard Go JSON encoder).  The code
// below uses a Rust `Debug` format (not text proto), but this may
// be OK for a test -- would a 'real' Keyset have serialized binary
// proto data?
#[test]
fn test_json_reader() {
    tink_mac::init();
    let gcm_key = tink_testutil::proto_encode(&tink_testutil::new_aes_gcm_key(0, 16));
    let eax_key = tink_testutil::proto_encode(&tink_testutil::new_hmac_key(
        tink::proto::HashType::Sha512,
        32,
    ));
    let json_keyset = format!(
        r#"{{
         "primaryKeyId":42,
         "key":[
            {{
               "keyData":{{
                  "typeUrl":"type.googleapis.com/google.crypto.tink.AesGcmKey",
                  "keyMaterialType":"SYMMETRIC",
                  "value": "{}"
               }},
               "outputPrefixType":"TINK",
               "keyId":42,
               "status":"ENABLED"
            }},
            {{
               "keyData":{{
                  "typeUrl":"type.googleapis.com/google.crypto.tink.AesEaxKey",
                  "keyMaterialType":"SYMMETRIC",
                  "value": "{}"
               }},
               "outputPrefixType":"RAW",
               "keyId":711,
               "status":"ENABLED"
            }}
         ]
      }}"#,
        base64::encode(&gcm_key),
        base64::encode(&eax_key)
    );
    let mut buf = SharedBuf::new();
    buf.write_all(json_keyset.as_bytes()).unwrap();
    let mut r = tink::keyset::JsonReader::new(buf);

    let got = r.read().expect("cannot read keyset");

    let want = tink::proto::Keyset {
        primary_key_id: 42,
        key: vec![
            tink::proto::keyset::Key {
                key_data: Some(tink::proto::KeyData {
                    type_url: "type.googleapis.com/google.crypto.tink.AesGcmKey".to_string(),
                    key_material_type: tink::proto::key_data::KeyMaterialType::Symmetric as i32,
                    value: gcm_key,
                }),
                output_prefix_type: tink::proto::OutputPrefixType::Tink as i32,
                key_id: 42,
                status: tink::proto::KeyStatusType::Enabled as i32,
            },
            tink::proto::keyset::Key {
                key_data: Some(tink::proto::KeyData {
                    type_url: "type.googleapis.com/google.crypto.tink.AesEaxKey".to_string(),
                    key_material_type: tink::proto::key_data::KeyMaterialType::Symmetric as i32,
                    value: eax_key,
                }),
                output_prefix_type: tink::proto::OutputPrefixType::Raw as i32,
                key_id: 711,
                status: tink::proto::KeyStatusType::Enabled as i32,
            },
        ],
    };
    assert_eq!(got, want, "written keyset doesn't match expected keyset");
}

#[test]
fn test_json_reader_large_ids() {
    tink_mac::init();
    let gcm_key = tink_testutil::proto_encode(&tink_testutil::new_aes_gcm_key(0, 16));
    let json_keyset = format!(
        r#"{{
         "primaryKeyId":4294967275,
         "key":[
            {{
               "keyData":{{
                  "typeUrl":"type.googleapis.com/google.crypto.tink.AesGcmKey",
                  "keyMaterialType":"SYMMETRIC",
                  "value": "{}"
               }},
               "outputPrefixType":"TINK",
               "keyId":4294967275,
               "status":"ENABLED"
            }}
         ]
      }}"#,
        base64::encode(&gcm_key),
    );
    let mut buf = SharedBuf::new();
    buf.write_all(json_keyset.as_bytes()).unwrap();
    let mut r = tink::keyset::JsonReader::new(buf);

    let got = r.read().expect("cannot read keyset");

    let want = tink::proto::Keyset {
        primary_key_id: 4294967275,
        key: vec![tink::proto::keyset::Key {
            key_data: Some(tink::proto::KeyData {
                type_url: "type.googleapis.com/google.crypto.tink.AesGcmKey".to_string(),
                key_material_type: tink::proto::key_data::KeyMaterialType::Symmetric as i32,
                value: gcm_key,
            }),
            output_prefix_type: tink::proto::OutputPrefixType::Tink as i32,
            key_id: 4294967275,
            status: tink::proto::KeyStatusType::Enabled as i32,
        }],
    };

    assert_eq!(got, want, "written keyset doesn't match expected keyset");
}

#[test]
fn test_json_reader_negative_ids() {
    tink_mac::init();
    let gcm_key = tink_testutil::proto_encode(&tink_testutil::new_aes_gcm_key(0, 16));
    let json_keyset = format!(
        r#"{{
         "primaryKeyId": -10,
         "key":[
            {{
               "keyData":{{
                  "typeUrl":"type.googleapis.com/google.crypto.tink.AesGcmKey",
                  "keyMaterialType":"SYMMETRIC",
                  "value": "{}"
               }},
               "outputPrefixType":"TINK",
               "keyId": -10,
               "status":"ENABLED"
            }}
         ]
      }}"#,
        base64::encode(&gcm_key),
    );
    let mut buf = SharedBuf::new();
    buf.write_all(json_keyset.as_bytes()).unwrap();
    let mut r = tink::keyset::JsonReader::new(buf);

    assert!(r.read().is_err(), "Expected failure due to negative key id");
}

// Tests that large IDs (>2^31) are written correctly.
#[test]
fn test_json_writer_large_id() {
    tink_mac::init();
    let eax_key = tink_testutil::proto_encode(&tink_testutil::new_hmac_key(
        tink::proto::HashType::Sha512,
        32,
    ));

    let ks = tink::proto::Keyset {
        primary_key_id: 4294967275,
        key: vec![tink::proto::keyset::Key {
            key_data: Some(tink::proto::KeyData {
                type_url: "type.googleapis.com/google.crypto.tink.AesEaxKey".to_string(),
                key_material_type: tink::proto::key_data::KeyMaterialType::Symmetric as i32,
                value: eax_key,
            }),
            output_prefix_type: tink::proto::OutputPrefixType::Raw as i32,
            key_id: 4294967275,
            status: tink::proto::KeyStatusType::Enabled as i32,
        }],
    };

    let mut buf = SharedBuf::new();
    let mut w = tink::keyset::JsonWriter::new(buf.clone());
    w.write(&ks).expect("cannot write keyset");

    let mut contents = String::new();
    buf.read_to_string(&mut contents).unwrap();
    assert!(
        contents.contains("\"keyId\":4294967275"),
        "written keyset does not contain a key with keyId 4294967275"
    );
    assert!(
        contents.contains("\"primaryKeyId\":4294967275"),
        "written keyset does not contain have primaryKeyId 4294967275"
    );
}

#[test]
fn test_json_io_encrypted() {
    tink_mac::init();
    let buf = SharedBuf::new();
    let mut w = tink::keyset::JsonWriter::new(buf.clone());
    let mut r = tink::keyset::JsonReader::new(buf);

    let kse1 = tink::proto::EncryptedKeyset {
        encrypted_keyset: vec![b'A'; 32],
        keyset_info: None,
    };
    w.write_encrypted(&kse1)
        .expect("cannot write encrypted keyset");

    let kse2 = r.read_encrypted().expect("cannot read encrypted keyset");
    assert_eq!(
        kse1, kse2,
        "written encrypted keyset ({:?}) doesn't match read encrypted keyset ({:?})",
        kse1, kse2
    );
}
<commit_msg>tink::keyset: drop out of date comment<commit_after>// Copyright 2020 The Tink-Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
////////////////////////////////////////////////////////////////////////////////

use super::SharedBuf;
use std::io::{Read, Write};
use tink::keyset::{Reader, Writer};

#[test]
fn test_json_io_unencrypted() {
    tink_mac::init();
    let buf = SharedBuf::new();
    let mut w = tink::keyset::JsonWriter::new(buf.clone());
    let mut r = tink::keyset::JsonReader::new(buf);

    let manager = tink_testutil::new_hmac_keyset_manager();
    let h = manager.handle().expect("cannot get keyset handle");

    let ks1 = tink::keyset::insecure::keyset_material(&h);
    w.write(&ks1).expect("cannot write keyset");

    let ks2 = r.read().expect("cannot read keyset");
    assert_eq!(
        ks1, ks2,
        "written keyset ({:?}) doesn't match read keyset ({:?})",
        ks1, ks2
    );
}

#[test]
fn test_json_reader() {
    tink_mac::init();
    let gcm_key = tink_testutil::proto_encode(&tink_testutil::new_aes_gcm_key(0, 16));
    let eax_key = tink_testutil::proto_encode(&tink_testutil::new_hmac_key(
        tink::proto::HashType::Sha512,
        32,
    ));
    let json_keyset = format!(
        r#"{{
         "primaryKeyId":42,
         "key":[
            {{
               "keyData":{{
                  "typeUrl":"type.googleapis.com/google.crypto.tink.AesGcmKey",
                  "keyMaterialType":"SYMMETRIC",
                  "value": "{}"
               }},
               "outputPrefixType":"TINK",
               "keyId":42,
               "status":"ENABLED"
            }},
            {{
               "keyData":{{
                  "typeUrl":"type.googleapis.com/google.crypto.tink.AesEaxKey",
                  "keyMaterialType":"SYMMETRIC",
                  "value": "{}"
               }},
               "outputPrefixType":"RAW",
               "keyId":711,
               "status":"ENABLED"
            }}
         ]
      }}"#,
        base64::encode(&gcm_key),
        base64::encode(&eax_key)
    );
    let mut buf = SharedBuf::new();
    buf.write_all(json_keyset.as_bytes()).unwrap();
    let mut r = tink::keyset::JsonReader::new(buf);

    let got = r.read().expect("cannot read keyset");

    let want = tink::proto::Keyset {
        primary_key_id: 42,
        key: vec![
            tink::proto::keyset::Key {
                key_data: Some(tink::proto::KeyData {
                    type_url: "type.googleapis.com/google.crypto.tink.AesGcmKey".to_string(),
                    key_material_type: tink::proto::key_data::KeyMaterialType::Symmetric as i32,
                    value: gcm_key,
                }),
                output_prefix_type: tink::proto::OutputPrefixType::Tink as i32,
                key_id: 42,
                status: tink::proto::KeyStatusType::Enabled as i32,
            },
            tink::proto::keyset::Key {
                key_data: Some(tink::proto::KeyData {
                    type_url: "type.googleapis.com/google.crypto.tink.AesEaxKey".to_string(),
                    key_material_type: tink::proto::key_data::KeyMaterialType::Symmetric as i32,
                    value: eax_key,
                }),
                output_prefix_type: tink::proto::OutputPrefixType::Raw as i32,
                key_id: 711,
                status: tink::proto::KeyStatusType::Enabled as i32,
            },
        ],
    };
    assert_eq!(got, want, "written keyset doesn't match expected keyset");
}

#[test]
fn test_json_reader_large_ids() {
    tink_mac::init();
    let gcm_key = tink_testutil::proto_encode(&tink_testutil::new_aes_gcm_key(0, 16));
    let json_keyset = format!(
        r#"{{
         "primaryKeyId":4294967275,
         "key":[
            {{
               "keyData":{{
                  "typeUrl":"type.googleapis.com/google.crypto.tink.AesGcmKey",
                  "keyMaterialType":"SYMMETRIC",
                  "value": "{}"
               }},
               "outputPrefixType":"TINK",
               "keyId":4294967275,
               "status":"ENABLED"
            }}
         ]
      }}"#,
        base64::encode(&gcm_key),
    );
    let mut buf = SharedBuf::new();
    buf.write_all(json_keyset.as_bytes()).unwrap();
    let mut r = tink::keyset::JsonReader::new(buf);

    let got = r.read().expect("cannot read keyset");

    let want = tink::proto::Keyset {
        primary_key_id: 4294967275,
        key: vec![tink::proto::keyset::Key {
            key_data: Some(tink::proto::KeyData {
                type_url: "type.googleapis.com/google.crypto.tink.AesGcmKey".to_string(),
                key_material_type: tink::proto::key_data::KeyMaterialType::Symmetric as i32,
                value: gcm_key,
            }),
            output_prefix_type: tink::proto::OutputPrefixType::Tink as i32,
            key_id: 4294967275,
            status: tink::proto::KeyStatusType::Enabled as i32,
        }],
    };

    assert_eq!(got, want, "written keyset doesn't match expected keyset");
}

#[test]
fn test_json_reader_negative_ids() {
    tink_mac::init();
    let gcm_key = tink_testutil::proto_encode(&tink_testutil::new_aes_gcm_key(0, 16));
    let json_keyset = format!(
        r#"{{
         "primaryKeyId": -10,
         "key":[
            {{
               "keyData":{{
                  "typeUrl":"type.googleapis.com/google.crypto.tink.AesGcmKey",
                  "keyMaterialType":"SYMMETRIC",
                  "value": "{}"
               }},
               "outputPrefixType":"TINK",
               "keyId": -10,
               "status":"ENABLED"
            }}
         ]
      }}"#,
        base64::encode(&gcm_key),
    );
    let mut buf = SharedBuf::new();
    buf.write_all(json_keyset.as_bytes()).unwrap();
    let mut r = tink::keyset::JsonReader::new(buf);

    assert!(r.read().is_err(), "Expected failure due to negative key id");
}

// Tests that large IDs (>2^31) are written correctly.
#[test]
fn test_json_writer_large_id() {
    tink_mac::init();
    let eax_key = tink_testutil::proto_encode(&tink_testutil::new_hmac_key(
        tink::proto::HashType::Sha512,
        32,
    ));

    let ks = tink::proto::Keyset {
        primary_key_id: 4294967275,
        key: vec![tink::proto::keyset::Key {
            key_data: Some(tink::proto::KeyData {
                type_url: "type.googleapis.com/google.crypto.tink.AesEaxKey".to_string(),
                key_material_type: tink::proto::key_data::KeyMaterialType::Symmetric as i32,
                value: eax_key,
            }),
            output_prefix_type: tink::proto::OutputPrefixType::Raw as i32,
            key_id: 4294967275,
            status: tink::proto::KeyStatusType::Enabled as i32,
        }],
    };

    let mut buf = SharedBuf::new();
    let mut w = tink::keyset::JsonWriter::new(buf.clone());
    w.write(&ks).expect("cannot write keyset");

    let mut contents = String::new();
    buf.read_to_string(&mut contents).unwrap();
    assert!(
        contents.contains("\"keyId\":4294967275"),
        "written keyset does not contain a key with keyId 4294967275"
    );
    assert!(
        contents.contains("\"primaryKeyId\":4294967275"),
        "written keyset does not contain have primaryKeyId 4294967275"
    );
}

#[test]
fn test_json_io_encrypted() {
    tink_mac::init();
    let buf = SharedBuf::new();
    let mut w = tink::keyset::JsonWriter::new(buf.clone());
    let mut r = tink::keyset::JsonReader::new(buf);

    let kse1 = tink::proto::EncryptedKeyset {
        encrypted_keyset: vec![b'A'; 32],
        keyset_info: None,
    };
    w.write_encrypted(&kse1)
        .expect("cannot write encrypted keyset");

    let kse2 = r.read_encrypted().expect("cannot read encrypted keyset");
    assert_eq!(
        kse1, kse2,
        "written encrypted keyset ({:?}) doesn't match read encrypted keyset ({:?})",
        kse1, kse2
    );
}
<|endoftext|>