<commit_before>use borrow_check::borrow_set::BorrowSet;
use borrow_check::location::LocationTable;
use borrow_check::{JustWrite, WriteAndRead};
use borrow_check::{AccessDepth, Deep, Shallow};
use borrow_check::{ReadOrWrite, Activation, Read, Reservation, Write};
use borrow_check::{Context, ContextKind};
use borrow_check::{LocalMutationIsAllowed, MutateMode};
use borrow_check::ArtificialField;
use borrow_check::{ReadKind, WriteKind};
use borrow_check::nll::facts::AllFacts;
use borrow_check::path_utils::*;
use dataflow::move_paths::indexes::BorrowIndex;
use rustc::ty::TyCtxt;
use rustc::mir::visit::Visitor;
use rustc::mir::{BasicBlock, Location, Mir, Place, Rvalue};
use rustc::mir::{Statement, StatementKind};
use rustc::mir::{Terminator, TerminatorKind};
use rustc::mir::{Operand, BorrowKind};
use rustc_data_structures::graph::dominators::Dominators;

pub(super) fn generate_invalidates<'cx, 'gcx, 'tcx>(
    tcx: TyCtxt<'cx, 'gcx, 'tcx>,
    all_facts: &mut Option<AllFacts>,
    location_table: &LocationTable,
    mir: &Mir<'tcx>,
    borrow_set: &BorrowSet<'tcx>,
) {
    if all_facts.is_none() {
        // Nothing to do if we don't have any facts
        return;
    }

    if let Some(all_facts) = all_facts {
        let dominators = mir.dominators();
        let mut ig = InvalidationGenerator {
            all_facts,
            borrow_set,
            tcx,
            location_table,
            mir,
            dominators,
        };
        ig.visit_mir(mir);
    }
}

struct InvalidationGenerator<'cx, 'tcx: 'cx, 'gcx: 'tcx> {
    tcx: TyCtxt<'cx, 'gcx, 'tcx>,
    all_facts: &'cx mut AllFacts,
    location_table: &'cx LocationTable,
    mir: &'cx Mir<'tcx>,
    dominators: Dominators<BasicBlock>,
    borrow_set: &'cx BorrowSet<'tcx>,
}

/// Visits the whole MIR and generates invalidates() facts
/// Most of the code implementing this was stolen from borrow_check/mod.rs
impl<'cx, 'tcx, 'gcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx, 'gcx> {
    fn visit_statement(&mut self,
                       block: BasicBlock,
                       statement: &Statement<'tcx>,
                       location: Location) {
        match statement.kind {
            StatementKind::Assign(ref lhs, ref rhs) => {
                self.consume_rvalue(
                    ContextKind::AssignRhs.new(location),
                    rhs,
                );

                self.mutate_place(
                    ContextKind::AssignLhs.new(location),
                    lhs,
                    Shallow(None),
                    JustWrite
                );
            }
            StatementKind::FakeRead(_, ref place) => {
                self.access_place(
                    ContextKind::FakeRead.new(location),
                    place,
                    (Deep, Read(ReadKind::Borrow(BorrowKind::Shared))),
                    LocalMutationIsAllowed::No,
                );
            }
            StatementKind::SetDiscriminant {
                ref place,
                variant_index: _,
            } => {
                self.mutate_place(
                    ContextKind::SetDiscrim.new(location),
                    place,
                    Shallow(None),
                    JustWrite,
                );
            }
            StatementKind::InlineAsm {
                ref asm,
                ref outputs,
                ref inputs,
            } => {
                let context = ContextKind::InlineAsm.new(location);
                for (o, output) in asm.outputs.iter().zip(outputs.iter()) {
                    if o.is_indirect {
                        // FIXME(eddyb) indirect inline asm outputs should
                        // be encoeded through MIR place derefs instead.
                        self.access_place(
                            context,
                            output,
                            (Deep, Read(ReadKind::Copy)),
                            LocalMutationIsAllowed::No,
                        );
                    } else {
                        self.mutate_place(
                            context,
                            output,
                            if o.is_rw { Deep } else { Shallow(None) },
                            if o.is_rw { WriteAndRead } else { JustWrite },
                        );
                    }
                }
                for (_, input) in inputs.iter() {
                    self.consume_operand(context, input);
                }
            }
            StatementKind::Nop |
            StatementKind::AscribeUserType(..) |
            StatementKind::Retag { .. } |
            StatementKind::StorageLive(..) => {
                // `Nop`, `AscribeUserType`, `Retag`, and `StorageLive` are irrelevant
                // to borrow check.
            }
            StatementKind::StorageDead(local) => {
                self.access_place(
                    ContextKind::StorageDead.new(location),
                    &Place::Local(local),
                    (Shallow(None), Write(WriteKind::StorageDeadOrDrop)),
                    LocalMutationIsAllowed::Yes,
                );
            }
        }

        self.super_statement(block, statement, location);
    }

    fn visit_terminator(
        &mut self,
        block: BasicBlock,
        terminator: &Terminator<'tcx>,
        location: Location
    ) {
        match terminator.kind {
            TerminatorKind::SwitchInt {
                ref discr,
                switch_ty: _,
                values: _,
                targets: _,
            } => {
                self.consume_operand(ContextKind::SwitchInt.new(location), discr);
            }
            TerminatorKind::Drop {
                location: ref drop_place,
                target: _,
                unwind: _,
            } => {
                self.access_place(
                    ContextKind::Drop.new(location),
                    drop_place,
                    (AccessDepth::Drop, Write(WriteKind::StorageDeadOrDrop)),
                    LocalMutationIsAllowed::Yes,
                );
            }
            TerminatorKind::DropAndReplace {
                location: ref drop_place,
                value: ref new_value,
                target: _,
                unwind: _,
            } => {
                self.mutate_place(
                    ContextKind::DropAndReplace.new(location),
                    drop_place,
                    Deep,
                    JustWrite,
                );
                self.consume_operand(
                    ContextKind::DropAndReplace.new(location),
                    new_value,
                );
            }
            TerminatorKind::Call {
                ref func,
                ref args,
                ref destination,
                cleanup: _,
                from_hir_call: _,
            } => {
                self.consume_operand(ContextKind::CallOperator.new(location), func);
                for arg in args {
                    self.consume_operand(ContextKind::CallOperand.new(location), arg);
                }
                if let Some((ref dest, _ /*bb*/)) = *destination {
                    self.mutate_place(
                        ContextKind::CallDest.new(location),
                        dest,
                        Deep,
                        JustWrite,
                    );
                }
            }
            TerminatorKind::Assert {
                ref cond,
                expected: _,
                ref msg,
                target: _,
                cleanup: _,
            } => {
                self.consume_operand(ContextKind::Assert.new(location), cond);
                use rustc::mir::interpret::EvalErrorKind::BoundsCheck;
                if let BoundsCheck { ref len, ref index } = *msg {
                    self.consume_operand(ContextKind::Assert.new(location), len);
                    self.consume_operand(ContextKind::Assert.new(location), index);
                }
            }
            TerminatorKind::Yield {
                ref value,
                resume,
                drop: _,
            } => {
                self.consume_operand(ContextKind::Yield.new(location), value);

                // Invalidate all borrows of local places
                let borrow_set = self.borrow_set.clone();
                let resume = self.location_table.start_index(resume.start_location());
                for i in borrow_set.borrows.indices() {
                    if borrow_of_local_data(&borrow_set.borrows[i].borrowed_place) {
                        self.all_facts.invalidates.push((resume, i));
                    }
                }
            }
            TerminatorKind::Resume | TerminatorKind::Return | TerminatorKind::GeneratorDrop => {
                // Invalidate all borrows of local places
                let borrow_set = self.borrow_set.clone();
                let start = self.location_table.start_index(location);
                for i in borrow_set.borrows.indices() {
                    if borrow_of_local_data(&borrow_set.borrows[i].borrowed_place) {
                        self.all_facts.invalidates.push((start, i));
                    }
                }
            }
            TerminatorKind::Goto { target: _ }
            | TerminatorKind::Abort
            | TerminatorKind::Unreachable
            | TerminatorKind::FalseEdges {
                real_target: _,
                imaginary_targets: _,
            }
            | TerminatorKind::FalseUnwind {
                real_target: _,
                unwind: _,
            } => {
                // no data used, thus irrelevant to borrowck
            }
        }

        self.super_terminator(block, terminator, location);
    }
}

impl<'cg, 'cx, 'tcx, 'gcx> InvalidationGenerator<'cx, 'tcx, 'gcx> {
    /// Simulates mutation of a place
    fn mutate_place(
        &mut self,
        context: Context,
        place: &Place<'tcx>,
        kind: AccessDepth,
        _mode: MutateMode,
    ) {
        self.access_place(
            context,
            place,
            (kind, Write(WriteKind::Mutate)),
            LocalMutationIsAllowed::ExceptUpvars,
        );
    }

    /// Simulates consumption of an operand
    fn consume_operand(
        &mut self,
        context: Context,
        operand: &Operand<'tcx>,
    ) {
        match *operand {
            Operand::Copy(ref place) => {
                self.access_place(
                    context,
                    place,
                    (Deep, Read(ReadKind::Copy)),
                    LocalMutationIsAllowed::No,
                );
            }
            Operand::Move(ref place) => {
                self.access_place(
                    context,
                    place,
                    (Deep, Write(WriteKind::Move)),
                    LocalMutationIsAllowed::Yes,
                );
            }
            Operand::Constant(_) => {}
        }
    }

    // Simulates consumption of an rvalue
    fn consume_rvalue(
        &mut self,
        context: Context,
        rvalue: &Rvalue<'tcx>,
    ) {
        match *rvalue {
            Rvalue::Ref(_ /*rgn*/, bk, ref place) => {
                let access_kind = match bk {
                    BorrowKind::Shallow => {
                        (Shallow(Some(ArtificialField::ShallowBorrow)), Read(ReadKind::Borrow(bk)))
                    },
                    BorrowKind::Shared => (Deep, Read(ReadKind::Borrow(bk))),
                    BorrowKind::Unique | BorrowKind::Mut { .. } => {
                        let wk = WriteKind::MutableBorrow(bk);
                        if allow_two_phase_borrow(&self.tcx, bk) {
                            (Deep, Reservation(wk))
                        } else {
                            (Deep, Write(wk))
                        }
                    }
                };

                self.access_place(
                    context,
                    place,
                    access_kind,
                    LocalMutationIsAllowed::No,
                );
            }

            Rvalue::Use(ref operand)
            | Rvalue::Repeat(ref operand, _)
            | Rvalue::UnaryOp(_ /*un_op*/, ref operand)
            | Rvalue::Cast(_ /*cast_kind*/, ref operand, _ /*ty*/) => {
                self.consume_operand(context, operand)
            }

            Rvalue::Len(ref place) | Rvalue::Discriminant(ref place) => {
                let af = match *rvalue {
                    Rvalue::Len(..) => Some(ArtificialField::ArrayLength),
                    Rvalue::Discriminant(..) => None,
                    _ => unreachable!(),
                };
                self.access_place(
                    context,
                    place,
                    (Shallow(af), Read(ReadKind::Copy)),
                    LocalMutationIsAllowed::No,
                );
            }

            Rvalue::BinaryOp(_bin_op, ref operand1, ref operand2)
            | Rvalue::CheckedBinaryOp(_bin_op, ref operand1, ref operand2) => {
                self.consume_operand(context, operand1);
                self.consume_operand(context, operand2);
            }

            Rvalue::NullaryOp(_op, _ty) => {
            }

            Rvalue::Aggregate(_, ref operands) => {
                for operand in operands {
                    self.consume_operand(context, operand);
                }
            }
        }
    }

    /// Simulates an access to a place
    fn access_place(
        &mut self,
        context: Context,
        place: &Place<'tcx>,
        kind: (AccessDepth, ReadOrWrite),
        _is_local_mutation_allowed: LocalMutationIsAllowed,
    ) {
        let (sd, rw) = kind;
        // note: not doing check_access_permissions checks because they don't generate invalidates
        self.check_access_for_conflict(context, place, sd, rw);
    }

    fn check_access_for_conflict(
        &mut self,
        context: Context,
        place: &Place<'tcx>,
        sd: AccessDepth,
        rw: ReadOrWrite,
    ) {
        debug!(
            "invalidation::check_access_for_conflict(context={:?}, place={:?}, sd={:?}, \
             rw={:?})",
            context,
            place,
            sd,
            rw,
        );
        let tcx = self.tcx;
        let mir = self.mir;
        let borrow_set = self.borrow_set.clone();
        let indices = self.borrow_set.borrows.indices();
        each_borrow_involving_path(
            self,
            tcx,
            mir,
            context,
            (sd, place),
            &borrow_set.clone(),
            indices,
            |this, borrow_index, borrow| {
                match (rw, borrow.kind) {
                    // Obviously an activation is compatible with its own
                    // reservation (or even prior activating uses of same
                    // borrow); so don't check if they interfere.
                    //
                    // NOTE: *reservations* do conflict with themselves;
                    // thus aren't injecting unsoundenss w/ this check.)
                    (Activation(_, activating), _) if activating == borrow_index => {
                        // Activating a borrow doesn't generate any invalidations, since we
                        // have already taken the reservation
                    }

                    (Read(_), BorrowKind::Shallow) | (Reservation(..), BorrowKind::Shallow)
                    | (Read(_), BorrowKind::Shared) | (Reservation(..), BorrowKind::Shared) => {
                        // Reads/reservations don't invalidate shared or shallow borrows
                    }

                    (Read(_), BorrowKind::Unique) | (Read(_), BorrowKind::Mut { .. }) => {
                        // Reading from mere reservations of mutable-borrows is OK.
                        if !is_active(&this.dominators, borrow, context.loc) {
                            // If the borrow isn't active yet, reads don't invalidate it
                            assert!(allow_two_phase_borrow(&this.tcx, borrow.kind));
                            return Control::Continue;
                        }

                        // Unique and mutable borrows are invalidated by reads from any
                        // involved path
                        this.generate_invalidates(borrow_index, context.loc);
                    }

                    (Reservation(_), BorrowKind::Unique)
                        | (Reservation(_), BorrowKind::Mut { .. })
                        | (Activation(_, _), _)
                        | (Write(_), _) => {
                            // unique or mutable borrows are invalidated by writes.
                            // Reservations count as writes since we need to check
                            // that activating the borrow will be OK
                            // FIXME(bob_twinkles) is this actually the right thing to do?
                            this.generate_invalidates(borrow_index, context.loc);
                        }
                }
                Control::Continue
            },
        );
    }

    /// Generate a new invalidates(L, B) fact
    fn generate_invalidates(&mut self, b: BorrowIndex, l: Location) {
        let lidx = self.location_table.start_index(l);
        self.all_facts.invalidates.push((lidx, b));
    }
}

<commit_msg>generate invalidations from 2-phase-borrow activations<commit_after>use borrow_check::borrow_set::BorrowSet;
use borrow_check::location::LocationTable;
use borrow_check::{JustWrite, WriteAndRead};
use borrow_check::{AccessDepth, Deep, Shallow};
use borrow_check::{ReadOrWrite, Activation, Read, Reservation, Write};
use borrow_check::{Context, ContextKind};
use borrow_check::{LocalMutationIsAllowed, MutateMode};
use borrow_check::ArtificialField;
use borrow_check::{ReadKind, WriteKind};
use borrow_check::nll::facts::AllFacts;
use borrow_check::path_utils::*;
use dataflow::move_paths::indexes::BorrowIndex;
use rustc::ty::TyCtxt;
use rustc::mir::visit::Visitor;
use rustc::mir::{BasicBlock, Location, Mir, Place, Rvalue};
use rustc::mir::{Statement, StatementKind};
use rustc::mir::{Terminator, TerminatorKind};
use rustc::mir::{Operand, BorrowKind};
use rustc_data_structures::graph::dominators::Dominators;

pub(super) fn generate_invalidates<'cx, 'gcx, 'tcx>(
    tcx: TyCtxt<'cx, 'gcx, 'tcx>,
    all_facts: &mut Option<AllFacts>,
    location_table: &LocationTable,
    mir: &Mir<'tcx>,
    borrow_set: &BorrowSet<'tcx>,
) {
    if all_facts.is_none() {
        // Nothing to do if we don't have any facts
        return;
    }

    if let Some(all_facts) = all_facts {
        let dominators = mir.dominators();
        let mut ig = InvalidationGenerator {
            all_facts,
            borrow_set,
            tcx,
            location_table,
            mir,
            dominators,
        };
        ig.visit_mir(mir);
    }
}

struct InvalidationGenerator<'cx, 'tcx: 'cx, 'gcx: 'tcx> {
    tcx: TyCtxt<'cx, 'gcx, 'tcx>,
    all_facts: &'cx mut AllFacts,
    location_table: &'cx LocationTable,
    mir: &'cx Mir<'tcx>,
    dominators: Dominators<BasicBlock>,
    borrow_set: &'cx BorrowSet<'tcx>,
}

/// Visits the whole MIR and generates invalidates() facts
/// Most of the code implementing this was stolen from borrow_check/mod.rs
impl<'cx, 'tcx, 'gcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx, 'gcx> {
    fn visit_statement(
        &mut self,
        block: BasicBlock,
        statement: &Statement<'tcx>,
        location: Location,
    ) {
        self.check_activations(location);

        match statement.kind {
            StatementKind::Assign(ref lhs, ref rhs) => {
                self.consume_rvalue(
                    ContextKind::AssignRhs.new(location),
                    rhs,
                );

                self.mutate_place(
                    ContextKind::AssignLhs.new(location),
                    lhs,
                    Shallow(None),
                    JustWrite
                );
            }
            StatementKind::FakeRead(_, ref place) => {
                self.access_place(
                    ContextKind::FakeRead.new(location),
                    place,
                    (Deep, Read(ReadKind::Borrow(BorrowKind::Shared))),
                    LocalMutationIsAllowed::No,
                );
            }
            StatementKind::SetDiscriminant {
                ref place,
                variant_index: _,
            } => {
                self.mutate_place(
                    ContextKind::SetDiscrim.new(location),
                    place,
                    Shallow(None),
                    JustWrite,
                );
            }
            StatementKind::InlineAsm {
                ref asm,
                ref outputs,
                ref inputs,
            } => {
                let context = ContextKind::InlineAsm.new(location);
                for (o, output) in asm.outputs.iter().zip(outputs.iter()) {
                    if o.is_indirect {
                        // FIXME(eddyb) indirect inline asm outputs should
                        // be encoeded through MIR place derefs instead.
                        self.access_place(
                            context,
                            output,
                            (Deep, Read(ReadKind::Copy)),
                            LocalMutationIsAllowed::No,
                        );
                    } else {
                        self.mutate_place(
                            context,
                            output,
                            if o.is_rw { Deep } else { Shallow(None) },
                            if o.is_rw { WriteAndRead } else { JustWrite },
                        );
                    }
                }
                for (_, input) in inputs.iter() {
                    self.consume_operand(context, input);
                }
            }
            StatementKind::Nop |
            StatementKind::AscribeUserType(..) |
            StatementKind::Retag { .. } |
            StatementKind::StorageLive(..) => {
                // `Nop`, `AscribeUserType`, `Retag`, and `StorageLive` are irrelevant
                // to borrow check.
            }
            StatementKind::StorageDead(local) => {
                self.access_place(
                    ContextKind::StorageDead.new(location),
                    &Place::Local(local),
                    (Shallow(None), Write(WriteKind::StorageDeadOrDrop)),
                    LocalMutationIsAllowed::Yes,
                );
            }
        }

        self.super_statement(block, statement, location);
    }

    fn visit_terminator(
        &mut self,
        block: BasicBlock,
        terminator: &Terminator<'tcx>,
        location: Location
    ) {
        self.check_activations(location);

        match terminator.kind {
            TerminatorKind::SwitchInt {
                ref discr,
                switch_ty: _,
                values: _,
                targets: _,
            } => {
                self.consume_operand(ContextKind::SwitchInt.new(location), discr);
            }
            TerminatorKind::Drop {
                location: ref drop_place,
                target: _,
                unwind: _,
            } => {
                self.access_place(
                    ContextKind::Drop.new(location),
                    drop_place,
                    (AccessDepth::Drop, Write(WriteKind::StorageDeadOrDrop)),
                    LocalMutationIsAllowed::Yes,
                );
            }
            TerminatorKind::DropAndReplace {
                location: ref drop_place,
                value: ref new_value,
                target: _,
                unwind: _,
            } => {
                self.mutate_place(
                    ContextKind::DropAndReplace.new(location),
                    drop_place,
                    Deep,
                    JustWrite,
                );
                self.consume_operand(
                    ContextKind::DropAndReplace.new(location),
                    new_value,
                );
            }
            TerminatorKind::Call {
                ref func,
                ref args,
                ref destination,
                cleanup: _,
                from_hir_call: _,
            } => {
                self.consume_operand(ContextKind::CallOperator.new(location), func);
                for arg in args {
                    self.consume_operand(ContextKind::CallOperand.new(location), arg);
                }
                if let Some((ref dest, _ /*bb*/)) = *destination {
                    self.mutate_place(
                        ContextKind::CallDest.new(location),
                        dest,
                        Deep,
                        JustWrite,
                    );
                }
            }
            TerminatorKind::Assert {
                ref cond,
                expected: _,
                ref msg,
                target: _,
                cleanup: _,
            } => {
                self.consume_operand(ContextKind::Assert.new(location), cond);
                use rustc::mir::interpret::EvalErrorKind::BoundsCheck;
                if let BoundsCheck { ref len, ref index } = *msg {
                    self.consume_operand(ContextKind::Assert.new(location), len);
                    self.consume_operand(ContextKind::Assert.new(location), index);
                }
            }
            TerminatorKind::Yield {
                ref value,
                resume,
                drop: _,
            } => {
                self.consume_operand(ContextKind::Yield.new(location), value);

                // Invalidate all borrows of local places
                let borrow_set = self.borrow_set.clone();
                let resume = self.location_table.start_index(resume.start_location());
                for i in borrow_set.borrows.indices() {
                    if borrow_of_local_data(&borrow_set.borrows[i].borrowed_place) {
                        self.all_facts.invalidates.push((resume, i));
                    }
                }
            }
            TerminatorKind::Resume | TerminatorKind::Return | TerminatorKind::GeneratorDrop => {
                // Invalidate all borrows of local places
                let borrow_set = self.borrow_set.clone();
                let start = self.location_table.start_index(location);
                for i in borrow_set.borrows.indices() {
                    if borrow_of_local_data(&borrow_set.borrows[i].borrowed_place) {
                        self.all_facts.invalidates.push((start, i));
                    }
                }
            }
            TerminatorKind::Goto { target: _ }
            | TerminatorKind::Abort
            | TerminatorKind::Unreachable
            | TerminatorKind::FalseEdges {
                real_target: _,
                imaginary_targets: _,
            }
            | TerminatorKind::FalseUnwind {
                real_target: _,
                unwind: _,
            } => {
                // no data used, thus irrelevant to borrowck
            }
        }

        self.super_terminator(block, terminator, location);
    }
}

impl<'cg, 'cx, 'tcx, 'gcx> InvalidationGenerator<'cx, 'tcx, 'gcx> {
    /// Simulates mutation of a place
    fn mutate_place(
        &mut self,
        context: Context,
        place: &Place<'tcx>,
        kind: AccessDepth,
        _mode: MutateMode,
    ) {
        self.access_place(
            context,
            place,
            (kind, Write(WriteKind::Mutate)),
            LocalMutationIsAllowed::ExceptUpvars,
        );
    }

    /// Simulates consumption of an operand
    fn consume_operand(
        &mut self,
        context: Context,
        operand: &Operand<'tcx>,
    ) {
        match *operand {
            Operand::Copy(ref place) => {
                self.access_place(
                    context,
                    place,
                    (Deep, Read(ReadKind::Copy)),
                    LocalMutationIsAllowed::No,
                );
            }
            Operand::Move(ref place) => {
                self.access_place(
                    context,
                    place,
                    (Deep, Write(WriteKind::Move)),
                    LocalMutationIsAllowed::Yes,
                );
            }
            Operand::Constant(_) => {}
        }
    }

    // Simulates consumption of an rvalue
    fn consume_rvalue(
        &mut self,
        context: Context,
        rvalue: &Rvalue<'tcx>,
    ) {
        match *rvalue {
            Rvalue::Ref(_ /*rgn*/, bk, ref place) => {
                let access_kind = match bk {
                    BorrowKind::Shallow => {
                        (Shallow(Some(ArtificialField::ShallowBorrow)), Read(ReadKind::Borrow(bk)))
                    },
                    BorrowKind::Shared => (Deep, Read(ReadKind::Borrow(bk))),
                    BorrowKind::Unique | BorrowKind::Mut { .. } => {
                        let wk = WriteKind::MutableBorrow(bk);
                        if allow_two_phase_borrow(&self.tcx, bk) {
                            (Deep, Reservation(wk))
                        } else {
                            (Deep, Write(wk))
                        }
                    }
                };

                self.access_place(
                    context,
                    place,
                    access_kind,
                    LocalMutationIsAllowed::No,
                );
            }

            Rvalue::Use(ref operand)
            | Rvalue::Repeat(ref operand, _)
            | Rvalue::UnaryOp(_ /*un_op*/, ref operand)
            | Rvalue::Cast(_ /*cast_kind*/, ref operand, _ /*ty*/) => {
                self.consume_operand(context, operand)
            }

            Rvalue::Len(ref place) | Rvalue::Discriminant(ref place) => {
                let af = match *rvalue {
                    Rvalue::Len(..) => Some(ArtificialField::ArrayLength),
                    Rvalue::Discriminant(..) => None,
                    _ => unreachable!(),
                };
                self.access_place(
                    context,
                    place,
                    (Shallow(af), Read(ReadKind::Copy)),
                    LocalMutationIsAllowed::No,
                );
            }

            Rvalue::BinaryOp(_bin_op, ref operand1, ref operand2)
            | Rvalue::CheckedBinaryOp(_bin_op, ref operand1, ref operand2) => {
                self.consume_operand(context, operand1);
                self.consume_operand(context, operand2);
            }

            Rvalue::NullaryOp(_op, _ty) => {
            }

            Rvalue::Aggregate(_, ref operands) => {
                for operand in operands {
                    self.consume_operand(context, operand);
                }
            }
        }
    }

    /// Simulates an access to a place
    fn access_place(
        &mut self,
        context: Context,
        place: &Place<'tcx>,
        kind: (AccessDepth, ReadOrWrite),
        _is_local_mutation_allowed: LocalMutationIsAllowed,
    ) {
        let (sd, rw) = kind;
        // note: not doing check_access_permissions checks because they don't generate invalidates
        self.check_access_for_conflict(context, place, sd, rw);
    }

    fn check_access_for_conflict(
        &mut self,
        context: Context,
        place: &Place<'tcx>,
        sd: AccessDepth,
        rw: ReadOrWrite,
    ) {
        debug!(
            "invalidation::check_access_for_conflict(context={:?}, place={:?}, sd={:?}, \
             rw={:?})",
            context,
            place,
            sd,
            rw,
        );
        let tcx = self.tcx;
        let mir = self.mir;
        let borrow_set = self.borrow_set.clone();
        let indices = self.borrow_set.borrows.indices();
        each_borrow_involving_path(
            self,
            tcx,
            mir,
            context,
            (sd, place),
            &borrow_set.clone(),
            indices,
            |this, borrow_index, borrow| {
                match (rw, borrow.kind) {
                    // Obviously an activation is compatible with its own
                    // reservation (or even prior activating uses of same
                    // borrow); so don't check if they interfere.
                    //
                    // NOTE: *reservations* do conflict with themselves;
                    // thus aren't injecting unsoundenss w/ this check.)
                    (Activation(_, activating), _) if activating == borrow_index => {
                        // Activating a borrow doesn't generate any invalidations, since we
                        // have already taken the reservation
                    }

                    (Read(_), BorrowKind::Shallow) | (Reservation(..), BorrowKind::Shallow)
                    | (Read(_), BorrowKind::Shared) | (Reservation(..), BorrowKind::Shared) => {
                        // Reads/reservations don't invalidate shared or shallow borrows
                    }

                    (Read(_), BorrowKind::Unique) | (Read(_), BorrowKind::Mut { .. }) => {
                        // Reading from mere reservations of mutable-borrows is OK.
                        if !is_active(&this.dominators, borrow, context.loc) {
                            // If the borrow isn't active yet, reads don't invalidate it
                            assert!(allow_two_phase_borrow(&this.tcx, borrow.kind));
                            return Control::Continue;
                        }

                        // Unique and mutable borrows are invalidated by reads from any
                        // involved path
                        this.generate_invalidates(borrow_index, context.loc);
                    }

                    (Reservation(_), BorrowKind::Unique)
                        | (Reservation(_), BorrowKind::Mut { .. })
                        | (Activation(_, _), _)
                        | (Write(_), _) => {
                            // unique or mutable borrows are invalidated by writes.
                            // Reservations count as writes since we need to check
                            // that activating the borrow will be OK
                            // FIXME(bob_twinkles) is this actually the right thing to do?
                            this.generate_invalidates(borrow_index, context.loc);
                        }
                }
                Control::Continue
            },
        );
    }

    /// Generate a new invalidates(L, B) fact
    fn generate_invalidates(&mut self, b: BorrowIndex, l: Location) {
        let lidx = self.location_table.start_index(l);
        self.all_facts.invalidates.push((lidx, b));
    }

    fn check_activations(
        &mut self,
        location: Location,
    ) {
        if !self.tcx.two_phase_borrows() {
            return;
        }

        // Two-phase borrow support: For each activation that is newly
        // generated at this statement, check if it interferes with
        // another borrow.
        for &borrow_index in self.borrow_set.activations_at_location(location) {
            let borrow = &self.borrow_set[borrow_index];

            // only mutable borrows should be 2-phase
            assert!(match borrow.kind {
                BorrowKind::Shared | BorrowKind::Shallow => false,
                BorrowKind::Unique | BorrowKind::Mut { .. } => true,
            });

            self.access_place(
                ContextKind::Activation.new(location),
                &borrow.borrowed_place,
                (
                    Deep,
                    Activation(WriteKind::MutableBorrow(borrow.kind), borrow_index),
                ),
                LocalMutationIsAllowed::No,
            );

            // We do not need to call `check_if_path_or_subpath_is_moved`
            // again, as we already called it when we made the
            // initial reservation.
        }
    }
}

<|endoftext|>


<commit_before>use byteorder::{BigEndian, ReaderBytesExt, WriterBytesExt};

use rustc_serialize::Encodable;
use rustc_serialize::Encoder;
use rustc_serialize::json;

use openssl;
use openssl::ssl;
use openssl::ssl::SslStream;
use openssl::ssl::error::SslError;

use std::cell::RefCell;
use std::num::Int;
use std::ops::{Range, Index};
use std::old_io::net;
use std::old_io::net::ip::SocketAddr;
use std::old_io::{TcpStream, MemReader};
use std::old_path::posix::Path;
use std::collections::HashMap;

use rand;
use time;

#[derive(Debug)]
pub struct Payload<'a> {
	pub aps: PayloadAPS<'a>,
	pub info: Option<HashMap<&'a str, &'a str>>
}

#[derive(Debug)]
pub struct PayloadAPS<'a> {
	pub alert: PayloadAPSAlert<'a>,
	pub badge: Option<i32>,
	pub sound: Option<&'a str>,
	pub content_available: Option<i32>
}

#[derive(Debug)]
pub enum PayloadAPSAlert<'a> {
	Plain(&'a str),
	Localized(&'a str, Vec<&'a str>)
}

impl<'a> Encodable for Payload<'a> {
	fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
	    match *self {
			Payload{ref aps, ref info} => {
				if let Some(ref map) = *info {
	    	        encoder.emit_struct("Payload", 1 + map.len(), |encoder| {
						try!(encoder.emit_struct_field( "aps", 0usize, |encoder| aps.encode(encoder)));
						let mut index = 1usize;
						for (key, val) in map.iter() {
							try!(encoder.emit_struct_field(key, index, |encoder| val.encode(encoder)));
							index = index + 1;
						}
						Ok(())
	    	        })
				}
				else {
	    	        encoder.emit_struct("Payload", 1, |encoder| {
						try!(encoder.emit_struct_field( "aps", 0usize, |encoder| aps.encode(encoder)));
						Ok(())
	    	        })
				}
			}
		}
	}
}

impl<'a> Encodable for PayloadAPS<'a> {
	fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
	    match *self {
			PayloadAPS{ref alert, ref badge, ref sound, ref content_available} => {
				let mut count = 1;
				if badge.is_some() { count = count + 1; }
				if sound.is_some() { count = count + 1; }
				if content_available.is_some() { count = count + 1; }

				let mut index = 0usize;
    	        encoder.emit_struct("PayloadAPS", count, |encoder| {
					try!(encoder.emit_struct_field( "alert", index, |encoder| alert.encode(encoder)));
					index = index + 1;
					if badge.is_some() {
						try!(encoder.emit_struct_field( "badge", index, |encoder| badge.unwrap().encode(encoder)));
						index = index + 1;
					}
					if sound.is_some() {
						try!(encoder.emit_struct_field( "sound", index, |encoder| sound.unwrap().encode(encoder)));
						index = index + 1;
					}
					if content_available.is_some() {
						try!(encoder.emit_struct_field( "content-available", index, |encoder| content_available.unwrap().encode(encoder)));
						index = index + 1;
					}
					Ok(())
    	        })
			}
		}
	}
}

impl<'a> Encodable for PayloadAPSAlert<'a> {
	fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
	    match *self {
			PayloadAPSAlert::Plain(ref str) => {
    	        encoder.emit_str(str)
			},
			PayloadAPSAlert::Localized(ref key, ref args) => {
    	        encoder.emit_struct("PayloadAPSAlert", 2, |encoder| {
					try!(encoder.emit_struct_field( "loc-key", 0usize, |encoder| key.encode(encoder)));
					try!(encoder.emit_struct_field( "loc-args", 1usize, |encoder| args.encode(encoder)));
					Ok(())
				})
			}
		}
	}
}

#[allow(dead_code)]
fn hex_to_int(hex: &str) -> u32 {
	let mut total = 0u32;
	let mut n = hex.to_string().len() - 1;

	for c in hex.chars() {
		match c {
			'0'...'9' => {
				total += 16.pow(n) * ((c as u32) - ('0' as u32));
			},
			'a'...'f' => {
				total += 16.pow(n) * ((c as u32) - ('a' as u32) + 10);
			},
			_ => {
			}
		}
		n = n - 1;
	}

	return total;
}

#[allow(dead_code)]
fn convert_to_binary(device_token: &str) -> Vec<u8> {
	let mut device_token_bytes: Vec<u8> = vec![];
	for i in 0..8 {
		let string = device_token.to_string();
		let range = Range{start:i*8, end:i*8+8};
		let sub_str = string.index(&range);

		let sub_str_num = hex_to_int(sub_str);
		let mut sub_str_bytes = vec![];
		let _ = sub_str_bytes.write_u32::<BigEndian>(sub_str_num);

		device_token_bytes.push_all(sub_str_bytes.as_slice());
	}

	return device_token_bytes;
}

#[allow(dead_code)]
pub fn convert_to_token(binary: &[u8]) -> String {
	let mut token = "".to_string();
	for i in 0..8 {
		let range = Range{start:i*4, end:i*4+4};
		let sub_slice = binary.index(&range);

		let mut rdr = MemReader::new(sub_slice.to_vec());
		let num = rdr.read_u32::<BigEndian>().unwrap();

		token = format!("{}{:x}", token, num);
	}
	return token;
}

#[allow(dead_code)]
pub fn convert_to_timestamp(binary: &[u8]) -> u32 {
	let mut rdr = MemReader::new(binary.to_vec());
	let num = rdr.read_u32::<BigEndian>().unwrap();

	return num;
}

pub struct APNS {
	pub sandbox: bool,
	pub certificate: Path,
	pub private_key: Path,
	pub ca_certificate: Path,
	pub ssl_stream: RefCell<SslStream<TcpStream>>
}

impl APNS {
	pub fn new(sandbox: bool, cert_file: Path, private_key_file: Path, ca_file: Path) -> APNS {
		let apns_url_production = "gateway.push.apple.com";
		let apns_url_development = "gateway.sandbox.push.apple.com";
		let apns_port = 2195;

		let apns_url = if sandbox {
			apns_url_development
		}
		else {
			apns_url_production
		};

		let ssl_stream = RefCell::new(get_ssl_stream(apns_url, apns_port, &cert_file, &private_key_file, &ca_file).unwrap());

		APNS{sandbox: sandbox, certificate: cert_file, private_key: private_key_file, ca_certificate: ca_file, ssl_stream: ssl_stream}
	}

	#[allow(dead_code)]
    pub fn get_feedback(&self) -> Vec<(u32, String)> {
		let apns_feedback_production = "feedback.push.apple.com";
		let apns_feedback_development = "feedback.sandbox.push.apple.com";
		let feedback_port = 2196;

		let apns_feedback_url = if self.sandbox {
			apns_feedback_development
		}
		else {
			apns_feedback_production
		};

		let mut stream = get_ssl_stream(apns_feedback_url, feedback_port, &self.certificate, &self.private_key, &self.ca_certificate).unwrap();

		let mut tokens: Vec<(u32, String)> = Vec::new();
        loop {
			let mut read_buffer = [0u8; 38];
			if let Err(..) = stream.read(&mut read_buffer) {
/*                println!("Feedback read error {:?}", error);*/
				break;
            }

			let time_range = Range{start:0, end:4};
			let time_slice = read_buffer.index(&time_range);
			let time = convert_to_timestamp(time_slice);

			let token_range = Range{start:6, end:38};
			let token_slice = read_buffer.index(&token_range);
            let token = convert_to_token(token_slice);

			tokens.push((time, token));
        }

        return tokens;
    }

	#[allow(dead_code)]
	pub fn send_payload(&self, payload: Payload, device_token: &str) {
		let payload_str = match json::encode(&payload) {
			Ok(json_str) => {
				json_str.to_string()
			}
			_ => { return; }
		};

		let payload_bytes = payload_str.into_bytes();
		let device_token_bytes: Vec<u8> = convert_to_binary(device_token);

		let mut notification_buffer: Vec<u8> = vec![];
		let mut message_buffer: Vec<u8> = vec![];

		// Device token
		let mut device_token_length = vec![];
		let _ = device_token_length.write_u16::<BigEndian>(device_token_bytes.len() as u16);

		message_buffer.push(1u8);
		message_buffer.push_all(device_token_length.as_slice());
		message_buffer.push_all(device_token_bytes.as_slice());

		// Payload
		let mut payload_length = vec![];
		let _ = payload_length.write_u16::<BigEndian>(payload_bytes.len() as u16);

		message_buffer.push(2u8);
		message_buffer.push_all(payload_length.as_slice());
		message_buffer.push_all(payload_bytes.as_slice());

		// Notification identifier
		let payload_id = rand::random::<u32>();
		let mut payload_id_be = vec![];
		let _ = payload_id_be.write_u32::<BigEndian>(payload_id);

		let mut payload_id_length = vec![];
		let _ = payload_id_length.write_u16::<BigEndian>(payload_id_be.len() as u16);

		message_buffer.push(3u8);
		message_buffer.push_all(payload_id_length.as_slice());
		message_buffer.push_all(payload_id_be.as_slice());

		//	Expiration date
		let time = time::now().to_timespec().sec + 86400;	// expired after one day
		let mut exp_date_be = vec![];
		let _ = exp_date_be.write_u32::<BigEndian>(time as u32);

		let mut exp_date_length = vec![];
		let _ = exp_date_length.write_u16::<BigEndian>(exp_date_be.len() as u16);

		message_buffer.push(4u8);
		message_buffer.push_all(exp_date_length.as_slice());
		message_buffer.push_all(exp_date_be.as_slice());

		// Priority
		let mut priority_length = vec![];
		let _ = priority_length.write_u16::<BigEndian>(1u16);

		message_buffer.push(5u8);
		message_buffer.push_all(priority_length.as_slice());
		message_buffer.push(10u8);

		let mut message_buffer_length = vec![];
		let _ = message_buffer_length.write_u32::<BigEndian>(message_buffer.len() as u32);

		let command = 2u8;
		notification_buffer.push(command);
		notification_buffer.push_all(message_buffer_length.as_slice());
		notification_buffer.push_all(message_buffer.as_slice());

	    let mut retry_count = 3;
		let mut borrow_ssl_stream = self.ssl_stream.borrow_mut();

		while let Err(error) = borrow_ssl_stream.write_all(&notification_buffer) {
			println!("ssl_stream write error {:?}", error);

            retry_count = retry_count - 1;
            if retry_count <= 0 {
                break;
            }

            // try to recreate ssl stream
			let apns_url_production = "gateway.push.apple.com";
			let apns_url_development = "gateway.sandbox.push.apple.com";
			let apns_port = 2195;

			let apns_url = if self.sandbox {
				apns_url_development
			}
			else {
				apns_url_production
			};

			let ssl_result = get_ssl_stream(apns_url, apns_port, &self.certificate, &self.private_key, &self.ca_certificate);
			*borrow_ssl_stream = match ssl_result {
				Ok(ssl_stream) => {
					ssl_stream
				},
				Err(error) => {
                    println!("failed to get_ssl_stream error {:?}", error);
                    continue;
				}
			};
		}

		// Response error code
/*
		let mut read_buffer = [0u8; 6];
		println!("SslStream read {:?}", self.ssl_stream.read(&mut read_buffer));

		for c in read_buffer.iter() {
			print!("{}", c);
		}
		println!("");
*/
	}
}

fn get_ssl_stream(url: &str, port: u16, cert_file: &Path, private_key_file: &Path, ca_file: &Path) -> Result<SslStream<TcpStream>, SslError> {
	let mut context = match ssl::SslContext::new(ssl::SslMethod::Sslv23) {
		Ok(ctx) => {
			ctx
		},
		Err(error) => {
			panic!("SslContext new error {:?}", error);
		}
	};
	if let Some(error) = context.set_CA_file(&ca_file) {
		println!("set_CA_file error {:?}", error);
	}
	if let Some(error) = context.set_certificate_file(&cert_file, openssl::x509::X509FileType::PEM) {
		println!("set_certificate_file error {:?}", error);
	}
	if let Some(error) = context.set_private_key_file(&private_key_file, openssl::x509::X509FileType::PEM) {
		println!("set_private_key_file error {:?}", error);
	}

	let apns_ip = match net::addrinfo::get_host_addresses(url) {
		Ok(results) => {
			results[0]
		},
		Err(error) => {
			panic!("unable to resolve host address: {:?}", error);
		}
	};

	let sock_addr = SocketAddr{ip: apns_ip, port: port};
    let tcp_conn = TcpStream::connect(sock_addr).unwrap();
	let ssl = match ssl::Ssl::new(&context) {
		Ok(conn) => {
			conn
		},
		Err(error) => {
			panic!("Ssl new error {:?}", error);
		}
	};

	return SslStream::new_from(ssl, tcp_conn);
}

<commit_msg>lazy creating ssl_stream for apns<commit_after>use byteorder::{BigEndian, ReaderBytesExt, WriterBytesExt};

use rustc_serialize::Encodable;
use rustc_serialize::Encoder;
use rustc_serialize::json;

use openssl;
use openssl::ssl;
use openssl::ssl::SslStream;
use openssl::ssl::error::SslError;

use std::cell::RefCell;
use std::num::Int;
use std::ops::{Range, Index};
use std::old_io::net;
use std::old_io::net::ip::SocketAddr;
use std::old_io::{TcpStream, MemReader};
use std::old_path::posix::Path;
use std::collections::HashMap;

use rand;
use time;

#[derive(Debug)]
pub struct Payload<'a> {
	pub aps: PayloadAPS<'a>,
	pub info: Option<HashMap<&'a str, &'a str>>
}

#[derive(Debug)]
pub struct PayloadAPS<'a> {
	pub alert: PayloadAPSAlert<'a>,
	pub badge: Option<i32>,
	pub sound: Option<&'a str>,
	pub content_available: Option<i32>
}

#[derive(Debug)]
pub enum PayloadAPSAlert<'a> {
	Plain(&'a str),
	Localized(&'a str, Vec<&'a str>)
}

impl<'a> Encodable for Payload<'a> {
	fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
	    match *self {
			Payload{ref aps, ref info} => {
				if let Some(ref map) = *info {
	    	        encoder.emit_struct("Payload", 1 + map.len(), |encoder| {
						try!(encoder.emit_struct_field( "aps", 0usize, |encoder| aps.encode(encoder)));
						let mut index = 1usize;
						for (key, val) in map.iter() {
							try!(encoder.emit_struct_field(key, index, |encoder| val.encode(encoder)));
							index = index + 1;
						}
						Ok(())
	    	        })
				}
				else {
	    	        encoder.emit_struct("Payload", 1, |encoder| {
						try!(encoder.emit_struct_field( "aps", 0usize, |encoder| aps.encode(encoder)));
						Ok(())
	    	        })
				}
			}
		}
	}
}

impl<'a> Encodable for PayloadAPS<'a> {
	fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
	    match *self {
			PayloadAPS{ref alert, ref badge, ref sound, ref content_available} => {
				let mut count = 1;
				if badge.is_some() { count = count + 1; }
				if sound.is_some() { count = count + 1; }
				if content_available.is_some() { count = count + 1; }

				let mut index = 0usize;
    	        encoder.emit_struct("PayloadAPS", count, |encoder| {
					try!(encoder.emit_struct_field( "alert", index, |encoder| alert.encode(encoder)));
					index = index + 1;
					if badge.is_some() {
						try!(encoder.emit_struct_field( "badge", index, |encoder| badge.unwrap().encode(encoder)));
						index = index + 1;
					}
					if sound.is_some() {
						try!(encoder.emit_struct_field( "sound", index, |encoder| sound.unwrap().encode(encoder)));
						index = index + 1;
					}
					if content_available.is_some() {
						try!(encoder.emit_struct_field( "content-available", index, |encoder| content_available.unwrap().encode(encoder)));
						index = index + 1;
					}
					Ok(())
    	        })
			}
		}
	}
}

impl<'a> Encodable for PayloadAPSAlert<'a> {
	fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
	    match *self {
			PayloadAPSAlert::Plain(ref str) => {
    	        encoder.emit_str(str)
			},
			PayloadAPSAlert::Localized(ref key, ref args) => {
    	        encoder.emit_struct("PayloadAPSAlert", 2, |encoder| {
					try!(encoder.emit_struct_field( "loc-key", 0usize, |encoder| key.encode(encoder)));
					try!(encoder.emit_struct_field( "loc-args", 1usize, |encoder| args.encode(encoder)));
					Ok(())
				})
			}
		}
	}
}

#[allow(dead_code)]
fn hex_to_int(hex: &str) -> u32 {
	let mut total = 0u32;
	let mut n = hex.to_string().len() - 1;

	for c in hex.chars() {
		match c {
			'0'...'9' => {
				total += 16.pow(n) * ((c as u32) - ('0' as u32));
			},
			'a'...'f' => {
				total += 16.pow(n) * ((c as u32) - ('a' as u32) + 10);
			},
			_ => {
			}
		}
		n = n - 1;
	}

	return total;
}

#[allow(dead_code)]
fn convert_to_binary(device_token: &str) -> Vec<u8> {
	let mut device_token_bytes: Vec<u8> = vec![];
	for i in 0..8 {
		let string = device_token.to_string();
		let range = Range{start:i*8, end:i*8+8};
		let sub_str = string.index(&range);

		let sub_str_num = hex_to_int(sub_str);
		let mut sub_str_bytes = vec![];
		let _ = sub_str_bytes.write_u32::<BigEndian>(sub_str_num);

		device_token_bytes.push_all(sub_str_bytes.as_slice());
	}

	return device_token_bytes;
}

#[allow(dead_code)]
pub fn convert_to_token(binary: &[u8]) -> String {
	let mut token = "".to_string();
	for i in 0..8 {
		let range = Range{start:i*4, end:i*4+4};
		let sub_slice = binary.index(&range);

		let mut rdr = MemReader::new(sub_slice.to_vec());
		let num = rdr.read_u32::<BigEndian>().unwrap();

		token = format!("{}{:x}", token, num);
	}
	return token;
}

#[allow(dead_code)]
pub fn convert_to_timestamp(binary: &[u8]) -> u32 {
	let mut rdr = MemReader::new(binary.to_vec());
	let num = rdr.read_u32::<BigEndian>().unwrap();

	return num;
}

pub struct APNS {
	pub sandbox: bool,
	pub certificate: Path,
	pub private_key: Path,
	pub ca_certificate: Path,
	pub ssl_stream: RefCell<Option<SslStream<TcpStream>>>
}

impl APNS {
	pub fn new(sandbox: bool, cert_file: Path, private_key_file: Path, ca_file: Path) -> APNS {
		APNS{sandbox: sandbox, certificate: cert_file, private_key: private_key_file, ca_certificate: ca_file, ssl_stream: RefCell::new(None)}
	}

	#[allow(dead_code)]
    pub fn get_feedback(&self) -> Result<Vec<(u32, String)>, SslError> {
		let apns_feedback_production = "feedback.push.apple.com";
		let apns_feedback_development = "feedback.sandbox.push.apple.com";
		let feedback_port = 2196;

		let apns_feedback_url = if self.sandbox { apns_feedback_development } else { apns_feedback_production };
		let mut stream = try!(get_ssl_stream(apns_feedback_url, feedback_port, &self.certificate, &self.private_key, &self.ca_certificate));

		let mut tokens: Vec<(u32, String)> = Vec::new();
        loop {
			let mut read_buffer = [0u8; 38];
			if let Err(..) = stream.read(&mut read_buffer) {
				break;
/*				return Result::Err(SslError::StreamError(error));*/
            }

			let time_range = Range{start:0, end:4};
			let time_slice = read_buffer.index(&time_range);
			let time = convert_to_timestamp(time_slice);

			let token_range = Range{start:6, end:38};
			let token_slice = read_buffer.index(&token_range);
            let token = convert_to_token(token_slice);

			tokens.push((time, token));
        }

        return Result::Ok(tokens);
    }

	#[allow(dead_code)]
	pub fn send_payload(&self, payload: Payload, device_token: &str) {
		let payload_str = match json::encode(&payload) {
			Ok(json_str) => {
				json_str.to_string()
			}
			_ => { return; }
		};

		let payload_bytes = payload_str.into_bytes();
		let device_token_bytes: Vec<u8> = convert_to_binary(device_token);

		let mut notification_buffer: Vec<u8> = vec![];
		let mut message_buffer: Vec<u8> = vec![];

		// Device token
		let mut device_token_length = vec![];
		let _ = device_token_length.write_u16::<BigEndian>(device_token_bytes.len() as u16);

		message_buffer.push(1u8);
		message_buffer.push_all(device_token_length.as_slice());
		message_buffer.push_all(device_token_bytes.as_slice());

		// Payload
		let mut payload_length = vec![];
		let _ = payload_length.write_u16::<BigEndian>(payload_bytes.len() as u16);

		message_buffer.push(2u8);
		message_buffer.push_all(payload_length.as_slice());
		message_buffer.push_all(payload_bytes.as_slice());

		// Notification identifier
		let payload_id = rand::random::<u32>();
		let mut payload_id_be = vec![];
		let _ = payload_id_be.write_u32::<BigEndian>(payload_id);

		let mut payload_id_length = vec![];
		let _ = payload_id_length.write_u16::<BigEndian>(payload_id_be.len() as u16);

		message_buffer.push(3u8);
		message_buffer.push_all(payload_id_length.as_slice());
		message_buffer.push_all(payload_id_be.as_slice());

		//	Expiration date
		let time = time::now().to_timespec().sec + 86400;	// expired after one day
		let mut exp_date_be = vec![];
		let _ = exp_date_be.write_u32::<BigEndian>(time as u32);

		let mut exp_date_length = vec![];
		let _ = exp_date_length.write_u16::<BigEndian>(exp_date_be.len() as u16);

		message_buffer.push(4u8);
		message_buffer.push_all(exp_date_length.as_slice());
		message_buffer.push_all(exp_date_be.as_slice());

		// Priority
		let mut priority_length = vec![];
		let _ = priority_length.write_u16::<BigEndian>(1u16);

		message_buffer.push(5u8);
		message_buffer.push_all(priority_length.as_slice());
		message_buffer.push(10u8);

		let mut message_buffer_length = vec![];
		let _ = message_buffer_length.write_u32::<BigEndian>(message_buffer.len() as u32);

		let command = 2u8;
		notification_buffer.push(command);
		notification_buffer.push_all(message_buffer_length.as_slice());
		notification_buffer.push_all(message_buffer.as_slice());

		if self.ssl_stream.borrow().is_none() {
			self.init_push_ssl_stream();
		}
		else {
			if let Some(ssl_stream) = self.ssl_stream.borrow_mut().as_mut() {
				while let Err(error) = ssl_stream.write_all(&notification_buffer) {
					println!("ssl_stream write error {:?}", error);
		            self.init_push_ssl_stream();
				}
			}
		}

		// Response error code
/*
		let mut read_buffer = [0u8; 6];
		println!("SslStream read {:?}", self.ssl_stream.read(&mut read_buffer));

		for c in read_buffer.iter() {
			print!("{}", c);
		}
		println!("");
*/
	}

	fn init_push_ssl_stream(&self) {
		let mut borrow_ssl_stream = self.ssl_stream.borrow_mut();

		let apns_url_production = "gateway.push.apple.com";
		let apns_url_development = "gateway.sandbox.push.apple.com";
		let apns_port = 2195;

		let apns_url = if self.sandbox { apns_url_development } else { apns_url_production };

		*borrow_ssl_stream = Some(get_ssl_stream(apns_url, apns_port, &self.certificate, &self.private_key, &self.ca_certificate).unwrap());
	}
}

fn get_ssl_stream(url: &str, port: u16, cert_file: &Path, private_key_file: &Path, ca_file: &Path) -> Result<SslStream<TcpStream>, SslError> {
	let mut context = try!(ssl::SslContext::new(ssl::SslMethod::Sslv23));

	if let Some(error) = context.set_CA_file(&ca_file) {
		println!("set_CA_file error {:?}", error);
	}
	if let Some(error) = context.set_certificate_file(&cert_file, openssl::x509::X509FileType::PEM) {
		println!("set_certificate_file error {:?}", error);
	}
	if let Some(error) = context.set_private_key_file(&private_key_file, openssl::x509::X509FileType::PEM) {
		println!("set_private_key_file error {:?}", error);
	}

	let apns_ip = match net::addrinfo::get_host_addresses(url) {
		Ok(results) => { results[0] },
		Err(error) => {
			return Result::Err(SslError::StreamError(error));
		}
	};

	let sock_addr = SocketAddr{ip: apns_ip, port: port};
    let tcp_conn = match TcpStream::connect(sock_addr) {
		Ok(conn) => { conn },
		Err(error) => {
			return Result::Err(SslError::StreamError(error));
		}
	};

	let ssl = try!(ssl::Ssl::new(&context));

	return SslStream::new_from(ssl, tcp_conn);
}

<|endoftext|>


<commit_before>#![feature(rustc_private)]

extern crate getopts;
extern crate miri;
extern crate rustc;
extern crate rustc_driver;
extern crate rustc_plugin;
extern crate env_logger;
extern crate log_settings;
extern crate syntax;
#[macro_use] extern crate log;

use miri::{eval_main, run_mir_passes};
use rustc::session::Session;
use rustc::mir::mir_map::MirMap;
use rustc_driver::{driver, CompilerCalls, Compilation};
use syntax::ast::MetaItemKind;

struct MiriCompilerCalls;

impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
    fn build_controller(
        &mut self,
        _: &Session,
        _: &getopts::Matches
    ) -> driver::CompileController<'a> {
        let mut control = driver::CompileController::basic();
        control.after_hir_lowering.callback = Box::new(|state| {
            state.session.plugin_attributes.borrow_mut().push(("miri".to_owned(), syntax::feature_gate::AttributeType::Whitelisted));
        });
        control.after_analysis.stop = Compilation::Stop;
        control.after_analysis.callback = Box::new(|state| {
            state.session.abort_if_errors();

            let tcx = state.tcx.unwrap();
            let mir_map = state.mir_map.unwrap();
            let (node_id, _) = state.session.entry_fn.borrow()
                .expect("no main or start function found");

            let krate = state.hir_crate.as_ref().unwrap();
            let mut memory_size = 100*1024*1024; // 100MB
            let mut step_limit = 1000_000;
            let mut stack_limit = 100;
            fn extract_str(lit: &syntax::ast::Lit) -> syntax::parse::token::InternedString {
                match lit.node {
                    syntax::ast::LitKind::Str(ref s, _) => s.clone(),
                    _ => panic!("attribute values need to be strings"),
                }
            }
            for attr in krate.attrs.iter() {
                match attr.node.value.node {
                    MetaItemKind::List(ref name, _) if name != "miri" => {}
                    MetaItemKind::List(_, ref items) => for item in items {
                        match item.node {
                            MetaItemKind::NameValue(ref name, ref value) => {
                                match &**name {
                                    "memory_size" => memory_size = extract_str(value).parse::<u64>().expect("not a number"),
                                    "step_limit" => step_limit = extract_str(value).parse::<u64>().expect("not a number"),
                                    "stack_limit" => stack_limit = extract_str(value).parse::<u64>().expect("not a number"),
                                    _ => state.session.span_err(item.span, "unknown miri attribute"),
                                }
                            }
                            _ => state.session.span_err(item.span, "miri attributes need to be of key = value kind"),
                        }
                    },
                    _ => {},
                }
            }

            let mut mir_map = MirMap { map: mir_map.map.clone() };
            run_mir_passes(tcx, &mut mir_map);
            eval_main(tcx, &mir_map, node_id, memory_size, step_limit, stack_limit);

            state.session.abort_if_errors();
        });

        control
    }
}

fn init_logger() {
    const MAX_INDENT: usize = 40;

    let format = |record: &log::LogRecord| {
        if record.level() == log::LogLevel::Trace {
            // prepend spaces to indent the final string
            let indentation = log_settings::settings().indentation;
            format!("{lvl}:{module}{depth:2}{indent:<indentation$} {text}",
                lvl = record.level(),
                module = record.location().module_path(),
                depth = indentation / MAX_INDENT,
                indentation = indentation % MAX_INDENT,
                indent = "",
                text = record.args())
        } else {
            format!("{lvl}:{module}: {text}",
                lvl = record.level(),
                module = record.location().module_path(),
                text = record.args())
        }
    };

    let mut builder = env_logger::LogBuilder::new();
    builder.format(format).filter(None, log::LogLevelFilter::Info);

    if std::env::var("MIRI_LOG").is_ok() {
        builder.parse(&std::env::var("MIRI_LOG").unwrap());
    }

    builder.init().unwrap();
}

fn find_sysroot() -> String {
    // Taken from https://github.com/Manishearth/rust-clippy/pull/911.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => option_env!("RUST_SYSROOT")
            .expect("need to specify RUST_SYSROOT env var or use rustup or multirust")
            .to_owned(),
    }
}

fn main() {
    init_logger();
    let mut args: Vec<String> = std::env::args().collect();

    let sysroot_flag = String::from("--sysroot");
    if !args.contains(&sysroot_flag) {
        args.push(sysroot_flag);
        args.push(find_sysroot());
    }

    rustc_driver::run_compiler(&args, &mut MiriCompilerCalls);
}
<commit_msg>remove unused extern crate<commit_after>#![feature(rustc_private)]

extern crate getopts;
extern crate miri;
extern crate rustc;
extern crate rustc_driver;
extern crate env_logger;
extern crate log_settings;
extern crate syntax;
#[macro_use] extern crate log;

use miri::{eval_main, run_mir_passes};
use rustc::session::Session;
use rustc::mir::mir_map::MirMap;
use rustc_driver::{driver, CompilerCalls, Compilation};
use syntax::ast::MetaItemKind;

struct MiriCompilerCalls;

impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
    fn build_controller(
        &mut self,
        _: &Session,
        _: &getopts::Matches
    ) -> driver::CompileController<'a> {
        let mut control = driver::CompileController::basic();
        control.after_hir_lowering.callback = Box::new(|state| {
            state.session.plugin_attributes.borrow_mut().push(("miri".to_owned(), syntax::feature_gate::AttributeType::Whitelisted));
        });
        control.after_analysis.stop = Compilation::Stop;
        control.after_analysis.callback = Box::new(|state| {
            state.session.abort_if_errors();

            let tcx = state.tcx.unwrap();
            let mir_map = state.mir_map.unwrap();
            let (node_id, _) = state.session.entry_fn.borrow()
                .expect("no main or start function found");

            let krate = state.hir_crate.as_ref().unwrap();
            let mut memory_size = 100*1024*1024; // 100MB
            let mut step_limit = 1000_000;
            let mut stack_limit = 100;
            fn extract_str(lit: &syntax::ast::Lit) -> syntax::parse::token::InternedString {
                match lit.node {
                    syntax::ast::LitKind::Str(ref s, _) => s.clone(),
                    _ => panic!("attribute values need to be strings"),
                }
            }
            for attr in krate.attrs.iter() {
                match attr.node.value.node {
                    MetaItemKind::List(ref name, _) if name != "miri" => {}
                    MetaItemKind::List(_, ref items) => for item in items {
                        match item.node {
                            MetaItemKind::NameValue(ref name, ref value) => {
                                match &**name {
                                    "memory_size" => memory_size = extract_str(value).parse::<u64>().expect("not a number"),
                                    "step_limit" => step_limit = extract_str(value).parse::<u64>().expect("not a number"),
                                    "stack_limit" => stack_limit = extract_str(value).parse::<u64>().expect("not a number"),
                                    _ => state.session.span_err(item.span, "unknown miri attribute"),
                                }
                            }
                            _ => state.session.span_err(item.span, "miri attributes need to be of key = value kind"),
                        }
                    },
                    _ => {},
                }
            }

            let mut mir_map = MirMap { map: mir_map.map.clone() };
            run_mir_passes(tcx, &mut mir_map);
            eval_main(tcx, &mir_map, node_id, memory_size, step_limit, stack_limit);

            state.session.abort_if_errors();
        });

        control
    }
}

fn init_logger() {
    const MAX_INDENT: usize = 40;

    let format = |record: &log::LogRecord| {
        if record.level() == log::LogLevel::Trace {
            // prepend spaces to indent the final string
            let indentation = log_settings::settings().indentation;
            format!("{lvl}:{module}{depth:2}{indent:<indentation$} {text}",
                lvl = record.level(),
                module = record.location().module_path(),
                depth = indentation / MAX_INDENT,
                indentation = indentation % MAX_INDENT,
                indent = "",
                text = record.args())
        } else {
            format!("{lvl}:{module}: {text}",
                lvl = record.level(),
                module = record.location().module_path(),
                text = record.args())
        }
    };

    let mut builder = env_logger::LogBuilder::new();
    builder.format(format).filter(None, log::LogLevelFilter::Info);

    if std::env::var("MIRI_LOG").is_ok() {
        builder.parse(&std::env::var("MIRI_LOG").unwrap());
    }

    builder.init().unwrap();
}

fn find_sysroot() -> String {
    // Taken from https://github.com/Manishearth/rust-clippy/pull/911.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => option_env!("RUST_SYSROOT")
            .expect("need to specify RUST_SYSROOT env var or use rustup or multirust")
            .to_owned(),
    }
}

fn main() {
    init_logger();
    let mut args: Vec<String> = std::env::args().collect();

    let sysroot_flag = String::from("--sysroot");
    if !args.contains(&sysroot_flag) {
        args.push(sysroot_flag);
        args.push(find_sysroot());
    }

    rustc_driver::run_compiler(&args, &mut MiriCompilerCalls);
}
<|endoftext|>


<commit_before>use std::slice;
use std::mem;
use std::ptr;
use std::os::raw;
use std::os::raw::c_char;

/// An enum identifying supported signature algorithms. Currently only ECDSA with SHA256 (ES256) and
/// RSASSA-PSS with SHA-256 (PS256) are supported. Note that with PS256, the salt length is defined
/// to be 32 bytes.
#[derive(Debug)]
pub enum SignatureAlgorithm {
    ES256,
    PS256,
}

type SECItemType = raw::c_uint; // TODO: actually an enum - is this the right size?
const SI_BUFFER: SECItemType = 0; // called siBuffer in NSS

#[repr(C)]
struct SECItem {
    typ: SECItemType,
    data: *const u8, // ugh it's not really const...
    len: raw::c_uint,
}

impl SECItem {
    fn maybe_new(data: &[u8]) -> Result<SECItem, NSSError> {
        if data.len() > u32::max_value() as usize {
            return Err(NSSError::InputTooLarge);
        }
        Ok(SECItem {
            typ: SI_BUFFER,
            data: data.as_ptr(),
            len: data.len() as u32,
        })
    }
}

#[repr(C)]
struct CkRsaPkcsPssParams {
    // Called CK_RSA_PKCS_PSS_PARAMS in NSS
    hash_alg: CkMechanismType, // Called hashAlg in NSS
    mgf: CkRsaPkcsMgfType,
    s_len: raw::c_ulong, // Called sLen in NSS
}

impl CkRsaPkcsPssParams {
    fn new() -> CkRsaPkcsPssParams {
        CkRsaPkcsPssParams {
            hash_alg: CKM_SHA256,
            mgf: CKG_MGF1_SHA256,
            s_len: 32,
        }
    }

    fn get_params_item(&self, alg: &SignatureAlgorithm) -> Result<*const SECItem, NSSError> {
        // This isn't entirely NSS' fault, but it mostly is.
        let rsa_pss_params_ptr: *const CkRsaPkcsPssParams = self;
        let rsa_pss_params_ptr: *const u8 = rsa_pss_params_ptr as *const u8;
        let rsa_pss_params_bytes = unsafe {
            slice::from_raw_parts(rsa_pss_params_ptr, mem::size_of::<CkRsaPkcsPssParams>())
        };
        let rsa_pss_params_secitem = SECItem::maybe_new(rsa_pss_params_bytes)?;
        let params_item: *const SECItem = match *alg {
            SignatureAlgorithm::ES256 => ptr::null(),
            SignatureAlgorithm::PS256 => &rsa_pss_params_secitem,
        };
        Ok(params_item)
    }
}

// TODO: link to NSS source where these are defined
type SECOidTag = raw::c_uint; // TODO: actually an enum - is this the right size?
const SEC_OID_SHA256: SECOidTag = 191;

type CkMechanismType = raw::c_ulong; // called CK_MECHANISM_TYPE in NSS
const CKM_ECDSA: CkMechanismType = 0x0000_1041;
const CKM_RSA_PKCS_PSS: CkMechanismType = 0x0000_000D;
const CKM_SHA256: CkMechanismType = 0x0000_0250;

type CkRsaPkcsMgfType = raw::c_ulong; // called CK_RSA_PKCS_MGF_TYPE in NSS
const CKG_MGF1_SHA256: CkRsaPkcsMgfType = 0x0000_0002;

type SECStatus = raw::c_int; // TODO: enum - right size?
const SEC_SUCCESS: SECStatus = 0; // Called SECSuccess in NSS
const SEC_FAILURE: SECStatus = -1; // Called SECFailure in NSS

enum SECKEYPublicKey {}
enum SECKEYPrivateKey {}
enum PK11SlotInfo {}
enum CERTCertificate {}
enum CERTCertDBHandle {}

const SHA256_LENGTH: usize = 32;

// TODO: ugh this will probably have a platform-specific name...
#[link(name = "nss3")]
extern "C" {
    fn PK11_HashBuf(
        hashAlg: SECOidTag,
        out: *mut u8,
        data_in: *const u8, // called "in" in NSS
        len: raw::c_int,
    ) -> SECStatus;
    fn PK11_VerifyWithMechanism(
        key: *const SECKEYPublicKey,
        mechanism: CkMechanismType,
        param: *const SECItem,
        sig: *const SECItem,
        hash: *const SECItem,
        wincx: *const raw::c_void,
    ) -> SECStatus;

    fn SECKEY_DestroyPublicKey(pubk: *const SECKEYPublicKey);

    fn CERT_GetDefaultCertDB() -> *const CERTCertDBHandle;
    fn CERT_DestroyCertificate(cert: *mut CERTCertificate);
    fn CERT_NewTempCertificate(
        handle: *const CERTCertDBHandle,
        derCert: *const SECItem,
        nickname: *const c_char,
        isperm: bool,
        copyDER: bool,
    ) -> *mut CERTCertificate;
    fn CERT_ExtractPublicKey(cert: *const CERTCertificate) -> *const SECKEYPublicKey;

    fn PK11_ImportDERPrivateKeyInfoAndReturnKey(
        slot: *mut PK11SlotInfo,
        derPKI: *const SECItem,
        nickname: *const SECItem,
        publicValue: *const SECItem,
        isPerm: bool,
        isPrivate: bool,
        keyUsage: u32,
        privk: *mut *mut SECKEYPrivateKey,
        wincx: *const u8,
    ) -> SECStatus;
    fn PK11_GetInternalSlot() -> *mut PK11SlotInfo;
    fn PK11_FreeSlot(slot: *mut PK11SlotInfo);
    fn PK11_SignatureLen(key: *const SECKEYPrivateKey) -> usize;
    fn PK11_SignWithMechanism(
        key: *const SECKEYPrivateKey,
        mech: CkMechanismType,
        param: *const SECItem,
        sig: *mut SECItem,
        hash: *const SECItem,
    ) -> SECStatus;
}

/// An error type describing errors that may be encountered during verification.
#[derive(Debug, PartialEq)]
pub enum NSSError {
    ImportCertError,
    DecodingPKCS8Failed,
    InputTooLarge,
    LibraryFailure,
    SignatureVerificationFailed,
    SigningFailed,
    ExtractPublicKeyFailed,
}

// XXX: make this work with other hashe algos.
fn hash(payload: &[u8]) -> Result<Vec<u8>, NSSError> {
    if payload.len() > raw::c_int::max_value() as usize {
        return Err(NSSError::InputTooLarge);
    }
    let mut hash_buf = vec![0; SHA256_LENGTH];
    let len: raw::c_int = payload.len() as raw::c_int;
    let hash_result =
        unsafe { PK11_HashBuf(SEC_OID_SHA256, hash_buf.as_mut_ptr(), payload.as_ptr(), len) };
    if hash_result != SEC_SUCCESS {
        return Err(NSSError::LibraryFailure);
    }
    Ok(hash_buf)
}

/// Main entrypoint for verification. Given a signature algorithm, the bytes of a subject public key
/// info, a payload, and a signature over the payload, returns a result based on the outcome of
/// decoding the subject public key info and running the signature verification algorithm on the
/// signed data.
pub fn verify_signature(
    signature_algorithm: &SignatureAlgorithm,
    cert: &[u8],
    payload: &[u8],
    signature: &[u8],
) -> Result<(), NSSError> {
    let slot = unsafe { PK11_GetInternalSlot() };
    if slot.is_null() {
        return Err(NSSError::LibraryFailure);
    }
    defer!(unsafe {
        PK11_FreeSlot(slot);
    });

    let hash_buf = hash(payload).unwrap();
    let hash_item = SECItem::maybe_new(hash_buf.as_slice())?;

    // Import DER cert into NSS.
    let der_cert = SECItem::maybe_new(cert)?;
    let db_handle = unsafe { CERT_GetDefaultCertDB() };
    if db_handle.is_null() {
        // TODO #28
        return Err(NSSError::LibraryFailure);
    }
    let nss_cert =
        unsafe { CERT_NewTempCertificate(db_handle, &der_cert, ptr::null(), false, true) };
    if nss_cert.is_null() {
        return Err(NSSError::ImportCertError);
    }
    defer!(unsafe {
        CERT_DestroyCertificate(nss_cert);
    });

    let key = unsafe { CERT_ExtractPublicKey(nss_cert) };
    if key.is_null() {
        return Err(NSSError::ExtractPublicKeyFailed);
    }
    defer!(unsafe {
        SECKEY_DestroyPublicKey(key);
    });
    let signature_item = SECItem::maybe_new(signature)?;
    let mechanism = match *signature_algorithm {
        SignatureAlgorithm::ES256 => CKM_ECDSA,
        SignatureAlgorithm::PS256 => CKM_RSA_PKCS_PSS,
    };
    let rsa_pss_params = CkRsaPkcsPssParams::new();
    let params_item = rsa_pss_params.get_params_item(signature_algorithm)?;
    let null_cx_ptr: *const raw::c_void = ptr::null();
    let result = unsafe {
        PK11_VerifyWithMechanism(
            key,
            mechanism,
            params_item,
            &signature_item,
            &hash_item,
            null_cx_ptr,
        )
    };
    match result {
        SEC_SUCCESS => Ok(()),
        SEC_FAILURE => Err(NSSError::SignatureVerificationFailed),
        _ => Err(NSSError::LibraryFailure),
    }
}

pub fn sign(
    signature_algorithm: &SignatureAlgorithm,
    pk8: &[u8],
    payload: &[u8],
) -> Result<Vec<u8>, NSSError> {
    let slot = unsafe { PK11_GetInternalSlot() };
    if slot.is_null() {
        return Err(NSSError::LibraryFailure);
    }
    defer!(unsafe {
        PK11_FreeSlot(slot);
    });
    let pkcs8item = SECItem::maybe_new(pk8)?;
    let mut key: *mut SECKEYPrivateKey = ptr::null_mut();
    let ku_all = 0xFF;
    let rv = unsafe {
        PK11_ImportDERPrivateKeyInfoAndReturnKey(
            slot,
            &pkcs8item,
            ptr::null(),
            ptr::null(),
            false,
            false,
            ku_all,
            &mut key,
            ptr::null(),
        )
    };
    if rv != SEC_SUCCESS || key.is_null() {
        println!("Decoding the PKCS8 failed.");
        return Err(NSSError::DecodingPKCS8Failed);
    }
    let hash_buf = hash(payload).unwrap();
    let hash_item = SECItem::maybe_new(hash_buf.as_slice())?;
    let signature_len = unsafe { PK11_SignatureLen(key) };
    // We indirectly pass this to PK11_SignWithMechanism and it gets modified,
    // so this has to be mut.
#[allow(unused_mut)]
    let mut signature: Vec<u8> = vec![0; signature_len];
    let mut signature_item = SECItem::maybe_new(signature.as_slice())?;
    let mechanism = match *signature_algorithm {
        SignatureAlgorithm::ES256 => CKM_ECDSA,
        SignatureAlgorithm::PS256 => CKM_RSA_PKCS_PSS,
    };
    let rsa_pss_params = CkRsaPkcsPssParams::new();
    let params_item = rsa_pss_params.get_params_item(signature_algorithm)?;
    let rv = unsafe {
        PK11_SignWithMechanism(key, mechanism, params_item, &mut signature_item, &hash_item)
    };

    if rv != SEC_SUCCESS || signature_item.len != signature_len as u32 {
        println!("Signing failed.");
        return Err(NSSError::SigningFailed);
    }
    Ok(signature)
}
<commit_msg>sort out the RSA-PSS params passing to NSS<commit_after>use std::marker::PhantomData;
use std::mem;
use std::os::raw::c_char;
use std::os::raw;
use std::ptr;

/// An enum identifying supported signature algorithms. Currently only ECDSA with SHA256 (ES256) and
/// RSASSA-PSS with SHA-256 (PS256) are supported. Note that with PS256, the salt length is defined
/// to be 32 bytes.
#[derive(Debug)]
pub enum SignatureAlgorithm {
    ES256,
    PS256,
}

type SECItemType = raw::c_uint; // TODO: actually an enum - is this the right size?
const SI_BUFFER: SECItemType = 0; // called siBuffer in NSS

#[repr(C)]
struct SECItem {
    typ: SECItemType,
    data: *const u8,
    len: raw::c_uint,
}

impl SECItem {
    fn maybe_new(data: &[u8]) -> Result<SECItem, NSSError> {
        if data.len() > u32::max_value() as usize {
            return Err(NSSError::InputTooLarge);
        }
        Ok(SECItem {
            typ: SI_BUFFER,
            data: data.as_ptr(),
            len: data.len() as u32,
        })
    }

    fn maybe_from_parts(data: *const u8, len: usize) -> Result<SECItem, NSSError> {
        if len > u32::max_value() as usize {
            return Err(NSSError::InputTooLarge);
        }
        Ok(SECItem {
            typ: SI_BUFFER,
            data: data,
            len: len as u32,
        })
    }
}

/// Many NSS APIs take constant data input as SECItems. Some, however, output data as SECItems.
/// To represent this, we define another type of mutable SECItem.
#[repr(C)]
struct SECItemMut<'a> {
    typ: SECItemType,
    data: *mut u8,
    len: raw::c_uint,
    _marker: PhantomData<&'a mut Vec<u8>>,
}

impl<'a> SECItemMut<'a> {
    /// Given a mutable reference to a Vec<u8> that has a particular allocated size, create a
    /// SECItemMut that points to the vec and has the same capacity.
    /// The input vec is not expected to have any actual contents, and in any case is cleared.
    fn maybe_from_vec(vec: &'a mut Vec<u8>) -> Result<SECItemMut<'a>, NSSError> {
        if vec.len() > u32::max_value() as usize {
            return Err(NSSError::InputTooLarge);
        }
        vec.clear();
        Ok(SECItemMut {
            typ: SI_BUFFER,
            data: vec.as_mut_ptr(),
            len: vec.capacity() as u32,
            _marker: PhantomData,
        })
    }
}

#[repr(C)]
struct CkRsaPkcsPssParams {
    // Called CK_RSA_PKCS_PSS_PARAMS in NSS
    hash_alg: CkMechanismType, // Called hashAlg in NSS
    mgf: CkRsaPkcsMgfType,
    s_len: raw::c_ulong, // Called sLen in NSS
}

impl CkRsaPkcsPssParams {
    fn new() -> CkRsaPkcsPssParams {
        CkRsaPkcsPssParams {
            hash_alg: CKM_SHA256,
            mgf: CKG_MGF1_SHA256,
            s_len: 32,
        }
    }

    fn get_params_item(&self) -> Result<SECItem, NSSError> {
        // This isn't entirely NSS' fault, but it mostly is.
        let params_ptr: *const CkRsaPkcsPssParams = self;
        let params_ptr: *const u8 = params_ptr as *const u8;
        let params_secitem =
            SECItem::maybe_from_parts(params_ptr, mem::size_of::<CkRsaPkcsPssParams>())?;
        Ok(params_secitem)
    }
}

// TODO: link to NSS source where these are defined
type SECOidTag = raw::c_uint; // TODO: actually an enum - is this the right size?
const SEC_OID_SHA256: SECOidTag = 191;

type CkMechanismType = raw::c_ulong; // called CK_MECHANISM_TYPE in NSS
const CKM_ECDSA: CkMechanismType = 0x0000_1041;
const CKM_RSA_PKCS_PSS: CkMechanismType = 0x0000_000D;
const CKM_SHA256: CkMechanismType = 0x0000_0250;

type CkRsaPkcsMgfType = raw::c_ulong; // called CK_RSA_PKCS_MGF_TYPE in NSS
const CKG_MGF1_SHA256: CkRsaPkcsMgfType = 0x0000_0002;

type SECStatus = raw::c_int; // TODO: enum - right size?
const SEC_SUCCESS: SECStatus = 0; // Called SECSuccess in NSS
const SEC_FAILURE: SECStatus = -1; // Called SECFailure in NSS

enum SECKEYPublicKey {}
enum SECKEYPrivateKey {}
enum PK11SlotInfo {}
enum CERTCertificate {}
enum CERTCertDBHandle {}

const SHA256_LENGTH: usize = 32;

// TODO: ugh this will probably have a platform-specific name...
#[link(name = "nss3")]
extern "C" {
    fn PK11_HashBuf(
        hashAlg: SECOidTag,
        out: *mut u8,
        data_in: *const u8, // called "in" in NSS
        len: raw::c_int,
    ) -> SECStatus;
    fn PK11_VerifyWithMechanism(
        key: *const SECKEYPublicKey,
        mechanism: CkMechanismType,
        param: *const SECItem,
        sig: *const SECItem,
        hash: *const SECItem,
        wincx: *const raw::c_void,
    ) -> SECStatus;

    fn SECKEY_DestroyPublicKey(pubk: *const SECKEYPublicKey);

    fn CERT_GetDefaultCertDB() -> *const CERTCertDBHandle;
    fn CERT_DestroyCertificate(cert: *mut CERTCertificate);
    fn CERT_NewTempCertificate(
        handle: *const CERTCertDBHandle,
        derCert: *const SECItem,
        nickname: *const c_char,
        isperm: bool,
        copyDER: bool,
    ) -> *mut CERTCertificate;
    fn CERT_ExtractPublicKey(cert: *const CERTCertificate) -> *const SECKEYPublicKey;

    fn PK11_ImportDERPrivateKeyInfoAndReturnKey(
        slot: *mut PK11SlotInfo,
        derPKI: *const SECItem,
        nickname: *const SECItem,
        publicValue: *const SECItem,
        isPerm: bool,
        isPrivate: bool,
        keyUsage: u32,
        privk: *mut *mut SECKEYPrivateKey,
        wincx: *const u8,
    ) -> SECStatus;
    fn PK11_GetInternalSlot() -> *mut PK11SlotInfo;
    fn PK11_FreeSlot(slot: *mut PK11SlotInfo);
    fn PK11_SignatureLen(key: *const SECKEYPrivateKey) -> usize;
    fn PK11_SignWithMechanism(
        key: *const SECKEYPrivateKey,
        mech: CkMechanismType,
        param: *const SECItem,
        sig: *mut SECItemMut,
        hash: *const SECItem,
    ) -> SECStatus;
}

/// An error type describing errors that may be encountered during verification.
#[derive(Debug, PartialEq)]
pub enum NSSError {
    ImportCertError,
    DecodingPKCS8Failed,
    InputTooLarge,
    LibraryFailure,
    SignatureVerificationFailed,
    SigningFailed,
    ExtractPublicKeyFailed,
}

// XXX: make this work with other hashe algos.
fn hash(payload: &[u8]) -> Result<Vec<u8>, NSSError> {
    if payload.len() > raw::c_int::max_value() as usize {
        return Err(NSSError::InputTooLarge);
    }
    let mut hash_buf = vec![0; SHA256_LENGTH];
    let len: raw::c_int = payload.len() as raw::c_int;
    let hash_result =
        unsafe { PK11_HashBuf(SEC_OID_SHA256, hash_buf.as_mut_ptr(), payload.as_ptr(), len) };
    if hash_result != SEC_SUCCESS {
        return Err(NSSError::LibraryFailure);
    }
    Ok(hash_buf)
}

/// Main entrypoint for verification. Given a signature algorithm, the bytes of a subject public key
/// info, a payload, and a signature over the payload, returns a result based on the outcome of
/// decoding the subject public key info and running the signature verification algorithm on the
/// signed data.
pub fn verify_signature(
    signature_algorithm: &SignatureAlgorithm,
    cert: &[u8],
    payload: &[u8],
    signature: &[u8],
) -> Result<(), NSSError> {
    let slot = unsafe { PK11_GetInternalSlot() };
    if slot.is_null() {
        return Err(NSSError::LibraryFailure);
    }
    defer!(unsafe {
        PK11_FreeSlot(slot);
    });

    let hash_buf = hash(payload).unwrap();
    let hash_item = SECItem::maybe_new(hash_buf.as_slice())?;

    // Import DER cert into NSS.
    let der_cert = SECItem::maybe_new(cert)?;
    let db_handle = unsafe { CERT_GetDefaultCertDB() };
    if db_handle.is_null() {
        // TODO #28
        return Err(NSSError::LibraryFailure);
    }
    let nss_cert =
        unsafe { CERT_NewTempCertificate(db_handle, &der_cert, ptr::null(), false, true) };
    if nss_cert.is_null() {
        return Err(NSSError::ImportCertError);
    }
    defer!(unsafe {
        CERT_DestroyCertificate(nss_cert);
    });

    let key = unsafe { CERT_ExtractPublicKey(nss_cert) };
    if key.is_null() {
        return Err(NSSError::ExtractPublicKeyFailed);
    }
    defer!(unsafe {
        SECKEY_DestroyPublicKey(key);
    });
    let signature_item = SECItem::maybe_new(signature)?;
    let mechanism = match *signature_algorithm {
        SignatureAlgorithm::ES256 => CKM_ECDSA,
        SignatureAlgorithm::PS256 => CKM_RSA_PKCS_PSS,
    };
    let rsa_pss_params = CkRsaPkcsPssParams::new();
    let rsa_pss_params_item = rsa_pss_params.get_params_item()?;
    let params_item = match *signature_algorithm {
        SignatureAlgorithm::ES256 => ptr::null(),
        SignatureAlgorithm::PS256 => &rsa_pss_params_item,
    };
    let null_cx_ptr: *const raw::c_void = ptr::null();
    let result = unsafe {
        PK11_VerifyWithMechanism(
            key,
            mechanism,
            params_item,
            &signature_item,
            &hash_item,
            null_cx_ptr,
        )
    };
    match result {
        SEC_SUCCESS => Ok(()),
        SEC_FAILURE => Err(NSSError::SignatureVerificationFailed),
        _ => Err(NSSError::LibraryFailure),
    }
}

pub fn sign(
    signature_algorithm: &SignatureAlgorithm,
    pk8: &[u8],
    payload: &[u8],
) -> Result<Vec<u8>, NSSError> {
    let slot = unsafe { PK11_GetInternalSlot() };
    if slot.is_null() {
        return Err(NSSError::LibraryFailure);
    }
    defer!(unsafe {
        PK11_FreeSlot(slot);
    });
    let pkcs8item = SECItem::maybe_new(pk8)?;
    let mut key: *mut SECKEYPrivateKey = ptr::null_mut();
    let ku_all = 0xFF;
    let rv = unsafe {
        PK11_ImportDERPrivateKeyInfoAndReturnKey(
            slot,
            &pkcs8item,
            ptr::null(),
            ptr::null(),
            false,
            false,
            ku_all,
            &mut key,
            ptr::null(),
        )
    };
    if rv != SEC_SUCCESS || key.is_null() {
        return Err(NSSError::DecodingPKCS8Failed);
    }
    let mechanism = match *signature_algorithm {
        SignatureAlgorithm::ES256 => CKM_ECDSA,
        SignatureAlgorithm::PS256 => CKM_RSA_PKCS_PSS,
    };
    let rsa_pss_params = CkRsaPkcsPssParams::new();
    let rsa_pss_params_item = rsa_pss_params.get_params_item()?;
    let params_item = match *signature_algorithm {
        SignatureAlgorithm::ES256 => ptr::null(),
        SignatureAlgorithm::PS256 => &rsa_pss_params_item,
    };
    let signature_len = unsafe { PK11_SignatureLen(key) };
    // Allocate enough space for the signature.
    let mut signature: Vec<u8> = Vec::with_capacity(signature_len);
    let hash_buf = hash(payload).unwrap();
    let hash_item = SECItem::maybe_new(hash_buf.as_slice())?;
    {
        // Get a mutable SECItem on the preallocated signature buffer. PK11_SignWithMechanism will
        // fill the SECItem's buf with the bytes of the signature.
        let mut signature_item = SECItemMut::maybe_from_vec(&mut signature)?;
        let rv = unsafe {
            PK11_SignWithMechanism(key, mechanism, params_item, &mut signature_item, &hash_item)
        };
        if rv != SEC_SUCCESS || signature_item.len as usize != signature_len {
            return Err(NSSError::SigningFailed);
        }
    }
    unsafe {
        // Now that the bytes of the signature have been filled out, set its length.
        signature.set_len(signature_len);
    }
    Ok(signature)
}
<|endoftext|>