<commit_before>/// Decodes an integer encoded with a given prefix size (in bits).
/// Assumes that the buffer `buf` contains the integer to be decoded,
/// with the first byte representing the octet that contains the
/// prefix.
///
/// Returns a tuple representing the decoded integer and the number
/// of bytes from the buffer that were used.
fn decode_integer(buf: &[u8], prefix_size: u8) -> (usize, usize) {
    assert!(prefix_size >= 1 && prefix_size <= 8);
    assert!(buf.len() >= 1);
    let mask = (1u8 << prefix_size) - 1;
    let mut value: usize = (buf[0] & mask) as usize;
    if value < (mask as usize) {
        return (value, 1);
    }

    // Already one byte used (the prefix)
    let mut total = 1;
    let mut m = 0;
    for &b in buf[1..].iter() {
        total += 1;
        value += ((b & 127) as usize) * (1us << m);
        m += 7;
        if b & 128 != 128 {
            break;
        }
    }

    // TODO Error situation -- what happens if we reach here w/o having
    //      `b` byte-sized...
    (value, total)
}

/// Encode an integer to the representation defined by HPACK.
///
/// Returns a newly allocated `Vec` containing the encoded bytes.
/// Only `prefix_size` lowest-order bits of the first byte in the
/// array are guaranteed to be used.
fn encode_integer(mut value: usize, prefix_size: u8) -> Vec<u8> {
    let mask: usize = ((1u8 << prefix_size) - 1) as usize;
    if value < mask {
        // Right now, the caller would need to be the one to combine
        // the other part of the prefix byte (but would know that it's
        // safe to do so using the bit-wise or).
        return vec![value as u8];
    }

    let mut res: Vec<u8> = Vec::new();
    res.push(mask as u8);
    value -= mask;
    while value >= 128 {
        res.push(((value % 128) + 128) as u8);
        value = value / 128;
    }
    res.push(value as u8);

    res
}

/// The table represents the static header table defined by the HPACK spec.
/// (HPACK, Appendix A)
static STATIC_TABLE: &'static [(&'static [u8], &'static [u8])] = &[
  (b":authority", b""),
  (b":method", b"GET"),
  (b":method", b"POST"),
  (b":path", b"/"),
  (b":path", b"/index.html"),
  (b":scheme", b"http"),
  (b":scheme", b"https"),
  (b":status", b"200"),
  (b":status", b"204"),
  (b":status", b"206"),
  (b":status", b"304"),
  (b":status", b"400"),
  (b":status", b"404"),
  (b":status", b"500"),
  (b"accept-", b""),
  (b"accept-encoding", b"gzip, deflate"),
  (b"accept-language", b""),
  (b"accept-ranges", b""),
  (b"accept", b""),
  (b"access-control-allow-origin", b""),
  (b"age", b""),
  (b"allow", b""),
  (b"authorization", b""),
  (b"cache-control", b""),
  (b"content-disposition", b""),
  (b"content-encoding", b""),
  (b"content-language", b""),
  (b"content-length", b""),
  (b"content-location", b""),
  (b"content-range", b""),
  (b"content-type", b""),
  (b"cookie", b""),
  (b"date", b""),
  (b"etag", b""),
  (b"expect", b""),
  (b"expires", b""),
  (b"from", b""),
  (b"host", b""),
  (b"if-match", b""),
  (b"if-modified-since", b""),
  (b"if-none-match", b""),
  (b"if-range", b""),
  (b"if-unmodified-since", b""),
  (b"last-modified", b""),
  (b"link", b""),
  (b"location", b""),
  (b"max-forwards", b""),
  (b"proxy-authenticate", b""),
  (b"proxy-authorization", b""),
  (b"range", b""),
  (b"referer", b""),
  (b"refresh", b""),
  (b"retry-after", b""),
  (b"server", b""),
  (b"set-cookie", b""),
  (b"strict-transport-security", b""),
  (b"transfer-encoding", b""),
  (b"user-agent", b""),
  (b"vary", b""),
  (b"via", b""),
  (b"www-authenticate", b""),
];

mod tests {
    use super::{decode_integer};
    use super::{encode_integer};

    #[test]
    fn test_decode_integer() {
        assert_eq!((10us, 1), decode_integer(&[10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 32, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 64, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10, 342, 22], 5));

        assert_eq!((127us, 2), decode_integer(&[255, 0], 7));
        assert_eq!((127us, 2), decode_integer(&[127, 0], 7));
    }

    #[test]
    fn test_encode_integer() {
        assert_eq!(encode_integer(10, 5), [10]);
        assert_eq!(encode_integer(1337, 5), [31, 154, 10]);
        assert_eq!(encode_integer(127, 7), [127, 0]);
    }
}
<commit_msg>Add a struct representing the dynamic table<commit_after>use std::fmt;
use std::collections::RingBuf;

/// Decodes an integer encoded with a given prefix size (in bits).
/// Assumes that the buffer `buf` contains the integer to be decoded,
/// with the first byte representing the octet that contains the
/// prefix.
///
/// Returns a tuple representing the decoded integer and the number
/// of bytes from the buffer that were used.
fn decode_integer(buf: &[u8], prefix_size: u8) -> (usize, usize) {
    assert!(prefix_size >= 1 && prefix_size <= 8);
    assert!(buf.len() >= 1);
    let mask = (1u8 << prefix_size) - 1;
    let mut value: usize = (buf[0] & mask) as usize;
    if value < (mask as usize) {
        return (value, 1);
    }

    // Already one byte used (the prefix)
    let mut total = 1;
    let mut m = 0;
    for &b in buf[1..].iter() {
        total += 1;
        value += ((b & 127) as usize) * (1us << m);
        m += 7;
        if b & 128 != 128 {
            break;
        }
    }

    // TODO Error situation -- what happens if we reach here w/o having
    //      `b` byte-sized...
    (value, total)
}

/// Encode an integer to the representation defined by HPACK.
///
/// Returns a newly allocated `Vec` containing the encoded bytes.
/// Only `prefix_size` lowest-order bits of the first byte in the
/// array are guaranteed to be used.
fn encode_integer(mut value: usize, prefix_size: u8) -> Vec<u8> {
    let mask: usize = ((1u8 << prefix_size) - 1) as usize;
    if value < mask {
        // Right now, the caller would need to be the one to combine
        // the other part of the prefix byte (but would know that it's
        // safe to do so using the bit-wise or).
        return vec![value as u8];
    }

    let mut res: Vec<u8> = Vec::new();
    res.push(mask as u8);
    value -= mask;
    while value >= 128 {
        res.push(((value % 128) + 128) as u8);
        value = value / 128;
    }
    res.push(value as u8);

    res
}

/// The table represents the static header table defined by the HPACK spec.
/// (HPACK, Appendix A)
static STATIC_TABLE: &'static [(&'static [u8], &'static [u8])] = &[
  (b":authority", b""),
  (b":method", b"GET"),
  (b":method", b"POST"),
  (b":path", b"/"),
  (b":path", b"/index.html"),
  (b":scheme", b"http"),
  (b":scheme", b"https"),
  (b":status", b"200"),
  (b":status", b"204"),
  (b":status", b"206"),
  (b":status", b"304"),
  (b":status", b"400"),
  (b":status", b"404"),
  (b":status", b"500"),
  (b"accept-", b""),
  (b"accept-encoding", b"gzip, deflate"),
  (b"accept-language", b""),
  (b"accept-ranges", b""),
  (b"accept", b""),
  (b"access-control-allow-origin", b""),
  (b"age", b""),
  (b"allow", b""),
  (b"authorization", b""),
  (b"cache-control", b""),
  (b"content-disposition", b""),
  (b"content-encoding", b""),
  (b"content-language", b""),
  (b"content-length", b""),
  (b"content-location", b""),
  (b"content-range", b""),
  (b"content-type", b""),
  (b"cookie", b""),
  (b"date", b""),
  (b"etag", b""),
  (b"expect", b""),
  (b"expires", b""),
  (b"from", b""),
  (b"host", b""),
  (b"if-match", b""),
  (b"if-modified-since", b""),
  (b"if-none-match", b""),
  (b"if-range", b""),
  (b"if-unmodified-since", b""),
  (b"last-modified", b""),
  (b"link", b""),
  (b"location", b""),
  (b"max-forwards", b""),
  (b"proxy-authenticate", b""),
  (b"proxy-authorization", b""),
  (b"range", b""),
  (b"referer", b""),
  (b"refresh", b""),
  (b"retry-after", b""),
  (b"server", b""),
  (b"set-cookie", b""),
  (b"strict-transport-security", b""),
  (b"transfer-encoding", b""),
  (b"user-agent", b""),
  (b"vary", b""),
  (b"via", b""),
  (b"www-authenticate", b""),
];

/// A struct representing the dynamic table that needs to be maintained by the
/// coder.
struct DynamicTable {
    table: RingBuf<(Vec<u8>, Vec<u8>)>,
    size: usize,
    max_size: usize,
}

impl DynamicTable {
    /// Creates a new empty dynamic table with default size.
    fn new() -> DynamicTable {
        DynamicTable {
            table: RingBuf::new(),
            size: 0,
            // The default maximum size corresponds to the default HTTP/2
            // setting
            max_size: 4096,
        }
    }

    fn get_size(&self) -> usize {
        self.size
    }

    /// Converts the current state of the table to a `Vec`
    fn get_table_as_list(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
        let mut ret: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
        for elem in self.table.iter() {
            ret.push(elem.clone());
        }

        ret
    }

}

impl fmt::Show for DynamicTable {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "{:?}", self.table)
    }
}

mod tests {
    use super::{decode_integer};
    use super::{encode_integer};

    #[test]
    fn test_decode_integer() {
        assert_eq!((10us, 1), decode_integer(&[10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 32, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 64, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10, 342, 22], 5));

        assert_eq!((127us, 2), decode_integer(&[255, 0], 7));
        assert_eq!((127us, 2), decode_integer(&[127, 0], 7));
    }

    #[test]
    fn test_encode_integer() {
        assert_eq!(encode_integer(10, 5), [10]);
        assert_eq!(encode_integer(1337, 5), [31, 154, 10]);
        assert_eq!(encode_integer(127, 7), [127, 0]);
    }
}
<|endoftext|>


<commit_before>//! This module contains the machinery necessary to detect infinite loops
//! during const-evaluation by taking snapshots of the state of the interpreter
//! at regular intervals.

// This lives in `interpret` because it needs access to all sots of private state.  However,
// it is not used by the general miri engine, just by CTFE.

use std::hash::{Hash, Hasher};

use rustc::ich::StableHashingContextProvider;
use rustc::mir;
use rustc::mir::interpret::{
    AllocId, Pointer, Scalar,
    Relocations, Allocation, UndefMask,
    EvalResult, EvalErrorKind,
};

use rustc::ty::{self, TyCtxt};
use rustc::ty::layout::Align;
use rustc_data_structures::fx::FxHashSet;
use rustc_data_structures::indexed_vec::IndexVec;
use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
use syntax::ast::Mutability;
use syntax::source_map::Span;

use super::eval_context::{LocalState, StackPopCleanup};
use super::{Frame, Memory, Operand, MemPlace, Place, Immediate, ScalarMaybeUndef, LocalValue};
use crate::const_eval::CompileTimeInterpreter;

#[derive(Default)]
pub(crate) struct InfiniteLoopDetector<'a, 'mir, 'tcx: 'a + 'mir> {
    /// The set of all `EvalSnapshot` *hashes* observed by this detector.
    ///
    /// When a collision occurs in this table, we store the full snapshot in
    /// `snapshots`.
    hashes: FxHashSet<u64>,

    /// The set of all `EvalSnapshot`s observed by this detector.
    ///
    /// An `EvalSnapshot` will only be fully cloned once it has caused a
    /// collision in `hashes`. As a result, the detector must observe at least
    /// *two* full cycles of an infinite loop before it triggers.
    snapshots: FxHashSet<EvalSnapshot<'a, 'mir, 'tcx>>,
}

impl<'a, 'mir, 'tcx> InfiniteLoopDetector<'a, 'mir, 'tcx>
{
    pub fn observe_and_analyze<'b>(
        &mut self,
        tcx: &TyCtxt<'b, 'tcx, 'tcx>,
        span: Span,
        memory: &Memory<'a, 'mir, 'tcx, CompileTimeInterpreter<'a, 'mir, 'tcx>>,
        stack: &[Frame<'mir, 'tcx>],
    ) -> EvalResult<'tcx, ()> {
        // Compute stack's hash before copying anything
        let mut hcx = tcx.get_stable_hashing_context();
        let mut hasher = StableHasher::<u64>::new();
        stack.hash_stable(&mut hcx, &mut hasher);
        let hash = hasher.finish();

        // Check if we know that hash already
        if self.hashes.is_empty() {
            // FIXME(#49980): make this warning a lint
            tcx.sess.span_warn(span,
                "Constant evaluating a complex constant, this might take some time");
        }
        if self.hashes.insert(hash) {
            // No collision
            return Ok(())
        }

        // We need to make a full copy. NOW things that to get really expensive.
        info!("snapshotting the state of the interpreter");

        if self.snapshots.insert(EvalSnapshot::new(memory, stack)) {
            // Spurious collision or first cycle
            return Ok(())
        }

        // Second cycle
        Err(EvalErrorKind::InfiniteLoop.into())
    }
}

trait SnapshotContext<'a> {
    fn resolve(&'a self, id: &AllocId) -> Option<&'a Allocation>;
}

/// Taking a snapshot of the evaluation context produces a view of
/// the state of the interpreter that is invariant to `AllocId`s.
trait Snapshot<'a, Ctx: SnapshotContext<'a>> {
    type Item;
    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item;
}

macro_rules! __impl_snapshot_field {
    ($field:ident, $ctx:expr) => ($field.snapshot($ctx));
    ($field:ident, $ctx:expr, $delegate:expr) => ($delegate);
}

// This assumes the type has two type parameters, first for the tag (set to `()`),
// then for the id
macro_rules! impl_snapshot_for {
    // FIXME(mark-i-m): Some of these should be `?` rather than `*`.
    (enum $enum_name:ident {
        $( $variant:ident $( ( $($field:ident $(-> $delegate:expr)*),* ) )* ),* $(,)*
    }) => {

        impl<'a, Ctx> self::Snapshot<'a, Ctx> for $enum_name
            where Ctx: self::SnapshotContext<'a>,
        {
            type Item = $enum_name<(), AllocIdSnapshot<'a>>;

            #[inline]
            fn snapshot(&self, __ctx: &'a Ctx) -> Self::Item {
                match *self {
                    $(
                        $enum_name::$variant $( ( $(ref $field),* ) )* =>
                            $enum_name::$variant $(
                                ( $( __impl_snapshot_field!($field, __ctx $(, $delegate)*) ),* ),
                            )*
                    )*
                }
            }
        }
    };

    // FIXME(mark-i-m): same here.
    (struct $struct_name:ident { $($field:ident $(-> $delegate:expr)*),*  $(,)* }) => {
        impl<'a, Ctx> self::Snapshot<'a, Ctx> for $struct_name
            where Ctx: self::SnapshotContext<'a>,
        {
            type Item = $struct_name<(), AllocIdSnapshot<'a>>;

            #[inline]
            fn snapshot(&self, __ctx: &'a Ctx) -> Self::Item {
                let $struct_name {
                    $(ref $field),*
                } = *self;

                $struct_name {
                    $( $field: __impl_snapshot_field!($field, __ctx $(, $delegate)*) ),*
                }
            }
        }
    };
}

impl<'a, Ctx, T> Snapshot<'a, Ctx> for Option<T>
    where Ctx: SnapshotContext<'a>,
          T: Snapshot<'a, Ctx>
{
    type Item = Option<<T as Snapshot<'a, Ctx>>::Item>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        match self {
            Some(x) => Some(x.snapshot(ctx)),
            None => None,
        }
    }
}

#[derive(Eq, PartialEq)]
struct AllocIdSnapshot<'a>(Option<AllocationSnapshot<'a>>);

impl<'a, Ctx> Snapshot<'a, Ctx> for AllocId
    where Ctx: SnapshotContext<'a>,
{
    type Item = AllocIdSnapshot<'a>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        AllocIdSnapshot(ctx.resolve(self).map(|alloc| alloc.snapshot(ctx)))
    }
}

impl_snapshot_for!(struct Pointer {
    alloc_id,
    offset -> *offset, // just copy offset verbatim
    tag -> *tag, // just copy tag
});

impl<'a, Ctx> Snapshot<'a, Ctx> for Scalar
    where Ctx: SnapshotContext<'a>,
{
    type Item = Scalar<(), AllocIdSnapshot<'a>>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        match self {
            Scalar::Ptr(p) => Scalar::Ptr(p.snapshot(ctx)),
            Scalar::Bits{ size, bits } => Scalar::Bits {
                size: *size,
                bits: *bits,
            },
        }
    }
}

impl_snapshot_for!(enum ScalarMaybeUndef {
    Scalar(s),
    Undef,
});

impl_stable_hash_for!(struct crate::interpret::MemPlace {
    ptr,
    align,
    meta,
});
impl_snapshot_for!(struct MemPlace {
    ptr,
    meta,
    align -> *align, // just copy alignment verbatim
});

impl_stable_hash_for!(enum crate::interpret::Place {
    Ptr(mem_place),
    Local { frame, local },
});
impl<'a, Ctx> Snapshot<'a, Ctx> for Place
    where Ctx: SnapshotContext<'a>,
{
    type Item = Place<(), AllocIdSnapshot<'a>>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        match self {
            Place::Ptr(p) => Place::Ptr(p.snapshot(ctx)),

            Place::Local{ frame, local } => Place::Local{
                frame: *frame,
                local: *local,
            },
        }
    }
}

impl_stable_hash_for!(enum crate::interpret::Immediate {
    Scalar(x),
    ScalarPair(x, y),
});
impl_snapshot_for!(enum Immediate {
    Scalar(s),
    ScalarPair(s, t),
});

impl_stable_hash_for!(enum crate::interpret::Operand {
    Immediate(x),
    Indirect(x),
});
impl_snapshot_for!(enum Operand {
    Immediate(v),
    Indirect(m),
});

impl_stable_hash_for!(enum crate::interpret::LocalValue {
    Dead,
    Live(x),
});
impl_snapshot_for!(enum LocalValue {
    Live(v),
    Dead,
});

impl<'a, Ctx> Snapshot<'a, Ctx> for Relocations
    where Ctx: SnapshotContext<'a>,
{
    type Item = Relocations<(), AllocIdSnapshot<'a>>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        Relocations::from_presorted(self.iter()
            .map(|(size, ((), id))| (*size, ((), id.snapshot(ctx))))
            .collect())
    }
}

#[derive(Eq, PartialEq)]
struct AllocationSnapshot<'a> {
    bytes: &'a [u8],
    relocations: Relocations<(), AllocIdSnapshot<'a>>,
    undef_mask: &'a UndefMask,
    align: &'a Align,
    mutability: &'a Mutability,
}

impl<'a, Ctx> Snapshot<'a, Ctx> for &'a Allocation
    where Ctx: SnapshotContext<'a>,
{
    type Item = AllocationSnapshot<'a>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        let Allocation { bytes, relocations, undef_mask, align, mutability, extra: () } = self;

        AllocationSnapshot {
            bytes,
            undef_mask,
            align,
            mutability,
            relocations: relocations.snapshot(ctx),
        }
    }
}

impl_stable_hash_for!(enum crate::interpret::eval_context::StackPopCleanup {
    Goto(block),
    None { cleanup },
});

#[derive(Eq, PartialEq)]
struct FrameSnapshot<'a, 'tcx: 'a> {
    instance: &'a ty::Instance<'tcx>,
    span: &'a Span,
    return_to_block: &'a StackPopCleanup,
    return_place: Option<Place<(), AllocIdSnapshot<'a>>>,
    locals: IndexVec<mir::Local, LocalValue<(), AllocIdSnapshot<'a>>>,
    block: &'a mir::BasicBlock,
    stmt: usize,
}

impl_stable_hash_for!(impl<'mir, 'tcx: 'mir> for struct Frame<'mir, 'tcx> {
    mir,
    instance,
    span,
    return_to_block,
    return_place -> (return_place.as_ref().map(|r| &**r)),
    locals,
    block,
    stmt,
    extra,
});

impl<'a, 'mir, 'tcx, Ctx> Snapshot<'a, Ctx> for &'a Frame<'mir, 'tcx>
    where Ctx: SnapshotContext<'a>,
{
    type Item = FrameSnapshot<'a, 'tcx>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        let Frame {
            mir: _,
            instance,
            span,
            return_to_block,
            return_place,
            locals,
            block,
            stmt,
            extra: _,
        } = self;

        FrameSnapshot {
            instance,
            span,
            return_to_block,
            block,
            stmt: *stmt,
            return_place: return_place.map(|r| r.snapshot(ctx)),
            locals: locals.iter().map(|local| local.snapshot(ctx)).collect(),
        }
    }
}

impl<'a, 'tcx, Ctx> Snapshot<'a, Ctx> for &'a LocalState<'tcx>
    where Ctx: SnapshotContext<'a>,
{
    type Item = LocalValue<(), AllocIdSnapshot<'a>>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        let LocalState { state, layout: _ } = self;
        state.snapshot(ctx)
    }
}

impl_stable_hash_for!(struct LocalState<'tcx> {
    state,
    layout -> _,
});

impl<'a, 'b, 'mir, 'tcx: 'a+'mir> SnapshotContext<'b>
    for Memory<'a, 'mir, 'tcx, CompileTimeInterpreter<'a, 'mir, 'tcx>>
{
    fn resolve(&'b self, id: &AllocId) -> Option<&'b Allocation> {
        self.get(*id).ok()
    }
}

/// The virtual machine state during const-evaluation at a given point in time.
/// We assume the `CompileTimeInterpreter` has no interesting extra state that
/// is worth considering here.
struct EvalSnapshot<'a, 'mir, 'tcx: 'a + 'mir> {
    memory: Memory<'a, 'mir, 'tcx, CompileTimeInterpreter<'a, 'mir, 'tcx>>,
    stack: Vec<Frame<'mir, 'tcx>>,
}

impl<'a, 'mir, 'tcx: 'a + 'mir> EvalSnapshot<'a, 'mir, 'tcx>
{
    fn new(
        memory: &Memory<'a, 'mir, 'tcx, CompileTimeInterpreter<'a, 'mir, 'tcx>>,
        stack: &[Frame<'mir, 'tcx>]
    ) -> Self {
        EvalSnapshot {
            memory: memory.clone(),
            stack: stack.into(),
        }
    }

    // Used to compare two snapshots
    fn snapshot(&'b self)
        -> Vec<FrameSnapshot<'b, 'tcx>>
    {
        // Start with the stack, iterate and recursively snapshot
        self.stack.iter().map(|frame| frame.snapshot(&self.memory)).collect()
    }

}

impl<'a, 'mir, 'tcx> Hash for EvalSnapshot<'a, 'mir, 'tcx>
{
    fn hash<H: Hasher>(&self, state: &mut H) {
        // Implement in terms of hash stable, so that k1 == k2 -> hash(k1) == hash(k2)
        let mut hcx = self.memory.tcx.get_stable_hashing_context();
        let mut hasher = StableHasher::<u64>::new();
        self.hash_stable(&mut hcx, &mut hasher);
        hasher.finish().hash(state)
    }
}

impl_stable_hash_for!(impl<'tcx, 'b, 'mir> for struct EvalSnapshot<'b, 'mir, 'tcx> {
    // Not hashing memory: Avoid hashing memory all the time during execution
    memory -> _,
    stack,
});

impl<'a, 'mir, 'tcx> Eq for EvalSnapshot<'a, 'mir, 'tcx>
{}

impl<'a, 'mir, 'tcx> PartialEq for EvalSnapshot<'a, 'mir, 'tcx>
{
    fn eq(&self, other: &Self) -> bool {
        // FIXME: This looks to be a *ridicolously expensive* comparison operation.
        // Doesn't this make tons of copies?  Either `snapshot` is very badly named,
        // or it does!
        self.snapshot() == other.snapshot()
    }
}
<commit_msg>Rollup merge of #58359 - taiki-e:impl_snapshot_for, r=oli-obk<commit_after>//! This module contains the machinery necessary to detect infinite loops
//! during const-evaluation by taking snapshots of the state of the interpreter
//! at regular intervals.

// This lives in `interpret` because it needs access to all sots of private state.  However,
// it is not used by the general miri engine, just by CTFE.

use std::hash::{Hash, Hasher};

use rustc::ich::StableHashingContextProvider;
use rustc::mir;
use rustc::mir::interpret::{
    AllocId, Pointer, Scalar,
    Relocations, Allocation, UndefMask,
    EvalResult, EvalErrorKind,
};

use rustc::ty::{self, TyCtxt};
use rustc::ty::layout::Align;
use rustc_data_structures::fx::FxHashSet;
use rustc_data_structures::indexed_vec::IndexVec;
use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
use syntax::ast::Mutability;
use syntax::source_map::Span;

use super::eval_context::{LocalState, StackPopCleanup};
use super::{Frame, Memory, Operand, MemPlace, Place, Immediate, ScalarMaybeUndef, LocalValue};
use crate::const_eval::CompileTimeInterpreter;

#[derive(Default)]
pub(crate) struct InfiniteLoopDetector<'a, 'mir, 'tcx: 'a + 'mir> {
    /// The set of all `EvalSnapshot` *hashes* observed by this detector.
    ///
    /// When a collision occurs in this table, we store the full snapshot in
    /// `snapshots`.
    hashes: FxHashSet<u64>,

    /// The set of all `EvalSnapshot`s observed by this detector.
    ///
    /// An `EvalSnapshot` will only be fully cloned once it has caused a
    /// collision in `hashes`. As a result, the detector must observe at least
    /// *two* full cycles of an infinite loop before it triggers.
    snapshots: FxHashSet<EvalSnapshot<'a, 'mir, 'tcx>>,
}

impl<'a, 'mir, 'tcx> InfiniteLoopDetector<'a, 'mir, 'tcx>
{
    pub fn observe_and_analyze<'b>(
        &mut self,
        tcx: &TyCtxt<'b, 'tcx, 'tcx>,
        span: Span,
        memory: &Memory<'a, 'mir, 'tcx, CompileTimeInterpreter<'a, 'mir, 'tcx>>,
        stack: &[Frame<'mir, 'tcx>],
    ) -> EvalResult<'tcx, ()> {
        // Compute stack's hash before copying anything
        let mut hcx = tcx.get_stable_hashing_context();
        let mut hasher = StableHasher::<u64>::new();
        stack.hash_stable(&mut hcx, &mut hasher);
        let hash = hasher.finish();

        // Check if we know that hash already
        if self.hashes.is_empty() {
            // FIXME(#49980): make this warning a lint
            tcx.sess.span_warn(span,
                "Constant evaluating a complex constant, this might take some time");
        }
        if self.hashes.insert(hash) {
            // No collision
            return Ok(())
        }

        // We need to make a full copy. NOW things that to get really expensive.
        info!("snapshotting the state of the interpreter");

        if self.snapshots.insert(EvalSnapshot::new(memory, stack)) {
            // Spurious collision or first cycle
            return Ok(())
        }

        // Second cycle
        Err(EvalErrorKind::InfiniteLoop.into())
    }
}

trait SnapshotContext<'a> {
    fn resolve(&'a self, id: &AllocId) -> Option<&'a Allocation>;
}

/// Taking a snapshot of the evaluation context produces a view of
/// the state of the interpreter that is invariant to `AllocId`s.
trait Snapshot<'a, Ctx: SnapshotContext<'a>> {
    type Item;
    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item;
}

macro_rules! __impl_snapshot_field {
    ($field:ident, $ctx:expr) => ($field.snapshot($ctx));
    ($field:ident, $ctx:expr, $delegate:expr) => ($delegate);
}

// This assumes the type has two type parameters, first for the tag (set to `()`),
// then for the id
macro_rules! impl_snapshot_for {
    (enum $enum_name:ident {
        $( $variant:ident $( ( $($field:ident $(-> $delegate:expr)?),* ) )? ),* $(,)?
    }) => {

        impl<'a, Ctx> self::Snapshot<'a, Ctx> for $enum_name
            where Ctx: self::SnapshotContext<'a>,
        {
            type Item = $enum_name<(), AllocIdSnapshot<'a>>;

            #[inline]
            fn snapshot(&self, __ctx: &'a Ctx) -> Self::Item {
                match *self {
                    $(
                        $enum_name::$variant $( ( $(ref $field),* ) )? =>
                            $enum_name::$variant $(
                                ( $( __impl_snapshot_field!($field, __ctx $(, $delegate)?) ),* ),
                            )?
                    )*
                }
            }
        }
    };

    (struct $struct_name:ident { $($field:ident $(-> $delegate:expr)?),*  $(,)? }) => {
        impl<'a, Ctx> self::Snapshot<'a, Ctx> for $struct_name
            where Ctx: self::SnapshotContext<'a>,
        {
            type Item = $struct_name<(), AllocIdSnapshot<'a>>;

            #[inline]
            fn snapshot(&self, __ctx: &'a Ctx) -> Self::Item {
                let $struct_name {
                    $(ref $field),*
                } = *self;

                $struct_name {
                    $( $field: __impl_snapshot_field!($field, __ctx $(, $delegate)?) ),*
                }
            }
        }
    };
}

impl<'a, Ctx, T> Snapshot<'a, Ctx> for Option<T>
    where Ctx: SnapshotContext<'a>,
          T: Snapshot<'a, Ctx>
{
    type Item = Option<<T as Snapshot<'a, Ctx>>::Item>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        match self {
            Some(x) => Some(x.snapshot(ctx)),
            None => None,
        }
    }
}

#[derive(Eq, PartialEq)]
struct AllocIdSnapshot<'a>(Option<AllocationSnapshot<'a>>);

impl<'a, Ctx> Snapshot<'a, Ctx> for AllocId
    where Ctx: SnapshotContext<'a>,
{
    type Item = AllocIdSnapshot<'a>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        AllocIdSnapshot(ctx.resolve(self).map(|alloc| alloc.snapshot(ctx)))
    }
}

impl_snapshot_for!(struct Pointer {
    alloc_id,
    offset -> *offset, // just copy offset verbatim
    tag -> *tag, // just copy tag
});

impl<'a, Ctx> Snapshot<'a, Ctx> for Scalar
    where Ctx: SnapshotContext<'a>,
{
    type Item = Scalar<(), AllocIdSnapshot<'a>>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        match self {
            Scalar::Ptr(p) => Scalar::Ptr(p.snapshot(ctx)),
            Scalar::Bits{ size, bits } => Scalar::Bits {
                size: *size,
                bits: *bits,
            },
        }
    }
}

impl_snapshot_for!(enum ScalarMaybeUndef {
    Scalar(s),
    Undef,
});

impl_stable_hash_for!(struct crate::interpret::MemPlace {
    ptr,
    align,
    meta,
});
impl_snapshot_for!(struct MemPlace {
    ptr,
    meta,
    align -> *align, // just copy alignment verbatim
});

impl_stable_hash_for!(enum crate::interpret::Place {
    Ptr(mem_place),
    Local { frame, local },
});
impl<'a, Ctx> Snapshot<'a, Ctx> for Place
    where Ctx: SnapshotContext<'a>,
{
    type Item = Place<(), AllocIdSnapshot<'a>>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        match self {
            Place::Ptr(p) => Place::Ptr(p.snapshot(ctx)),

            Place::Local{ frame, local } => Place::Local{
                frame: *frame,
                local: *local,
            },
        }
    }
}

impl_stable_hash_for!(enum crate::interpret::Immediate {
    Scalar(x),
    ScalarPair(x, y),
});
impl_snapshot_for!(enum Immediate {
    Scalar(s),
    ScalarPair(s, t),
});

impl_stable_hash_for!(enum crate::interpret::Operand {
    Immediate(x),
    Indirect(x),
});
impl_snapshot_for!(enum Operand {
    Immediate(v),
    Indirect(m),
});

impl_stable_hash_for!(enum crate::interpret::LocalValue {
    Dead,
    Live(x),
});
impl_snapshot_for!(enum LocalValue {
    Live(v),
    Dead,
});

impl<'a, Ctx> Snapshot<'a, Ctx> for Relocations
    where Ctx: SnapshotContext<'a>,
{
    type Item = Relocations<(), AllocIdSnapshot<'a>>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        Relocations::from_presorted(self.iter()
            .map(|(size, ((), id))| (*size, ((), id.snapshot(ctx))))
            .collect())
    }
}

#[derive(Eq, PartialEq)]
struct AllocationSnapshot<'a> {
    bytes: &'a [u8],
    relocations: Relocations<(), AllocIdSnapshot<'a>>,
    undef_mask: &'a UndefMask,
    align: &'a Align,
    mutability: &'a Mutability,
}

impl<'a, Ctx> Snapshot<'a, Ctx> for &'a Allocation
    where Ctx: SnapshotContext<'a>,
{
    type Item = AllocationSnapshot<'a>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        let Allocation { bytes, relocations, undef_mask, align, mutability, extra: () } = self;

        AllocationSnapshot {
            bytes,
            undef_mask,
            align,
            mutability,
            relocations: relocations.snapshot(ctx),
        }
    }
}

impl_stable_hash_for!(enum crate::interpret::eval_context::StackPopCleanup {
    Goto(block),
    None { cleanup },
});

#[derive(Eq, PartialEq)]
struct FrameSnapshot<'a, 'tcx: 'a> {
    instance: &'a ty::Instance<'tcx>,
    span: &'a Span,
    return_to_block: &'a StackPopCleanup,
    return_place: Option<Place<(), AllocIdSnapshot<'a>>>,
    locals: IndexVec<mir::Local, LocalValue<(), AllocIdSnapshot<'a>>>,
    block: &'a mir::BasicBlock,
    stmt: usize,
}

impl_stable_hash_for!(impl<'mir, 'tcx: 'mir> for struct Frame<'mir, 'tcx> {
    mir,
    instance,
    span,
    return_to_block,
    return_place -> (return_place.as_ref().map(|r| &**r)),
    locals,
    block,
    stmt,
    extra,
});

impl<'a, 'mir, 'tcx, Ctx> Snapshot<'a, Ctx> for &'a Frame<'mir, 'tcx>
    where Ctx: SnapshotContext<'a>,
{
    type Item = FrameSnapshot<'a, 'tcx>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        let Frame {
            mir: _,
            instance,
            span,
            return_to_block,
            return_place,
            locals,
            block,
            stmt,
            extra: _,
        } = self;

        FrameSnapshot {
            instance,
            span,
            return_to_block,
            block,
            stmt: *stmt,
            return_place: return_place.map(|r| r.snapshot(ctx)),
            locals: locals.iter().map(|local| local.snapshot(ctx)).collect(),
        }
    }
}

impl<'a, 'tcx, Ctx> Snapshot<'a, Ctx> for &'a LocalState<'tcx>
    where Ctx: SnapshotContext<'a>,
{
    type Item = LocalValue<(), AllocIdSnapshot<'a>>;

    fn snapshot(&self, ctx: &'a Ctx) -> Self::Item {
        let LocalState { state, layout: _ } = self;
        state.snapshot(ctx)
    }
}

impl_stable_hash_for!(struct LocalState<'tcx> {
    state,
    layout -> _,
});

impl<'a, 'b, 'mir, 'tcx: 'a+'mir> SnapshotContext<'b>
    for Memory<'a, 'mir, 'tcx, CompileTimeInterpreter<'a, 'mir, 'tcx>>
{
    fn resolve(&'b self, id: &AllocId) -> Option<&'b Allocation> {
        self.get(*id).ok()
    }
}

/// The virtual machine state during const-evaluation at a given point in time.
/// We assume the `CompileTimeInterpreter` has no interesting extra state that
/// is worth considering here.
struct EvalSnapshot<'a, 'mir, 'tcx: 'a + 'mir> {
    memory: Memory<'a, 'mir, 'tcx, CompileTimeInterpreter<'a, 'mir, 'tcx>>,
    stack: Vec<Frame<'mir, 'tcx>>,
}

impl<'a, 'mir, 'tcx: 'a + 'mir> EvalSnapshot<'a, 'mir, 'tcx>
{
    fn new(
        memory: &Memory<'a, 'mir, 'tcx, CompileTimeInterpreter<'a, 'mir, 'tcx>>,
        stack: &[Frame<'mir, 'tcx>]
    ) -> Self {
        EvalSnapshot {
            memory: memory.clone(),
            stack: stack.into(),
        }
    }

    // Used to compare two snapshots
    fn snapshot(&'b self)
        -> Vec<FrameSnapshot<'b, 'tcx>>
    {
        // Start with the stack, iterate and recursively snapshot
        self.stack.iter().map(|frame| frame.snapshot(&self.memory)).collect()
    }

}

impl<'a, 'mir, 'tcx> Hash for EvalSnapshot<'a, 'mir, 'tcx>
{
    fn hash<H: Hasher>(&self, state: &mut H) {
        // Implement in terms of hash stable, so that k1 == k2 -> hash(k1) == hash(k2)
        let mut hcx = self.memory.tcx.get_stable_hashing_context();
        let mut hasher = StableHasher::<u64>::new();
        self.hash_stable(&mut hcx, &mut hasher);
        hasher.finish().hash(state)
    }
}

impl_stable_hash_for!(impl<'tcx, 'b, 'mir> for struct EvalSnapshot<'b, 'mir, 'tcx> {
    // Not hashing memory: Avoid hashing memory all the time during execution
    memory -> _,
    stack,
});

impl<'a, 'mir, 'tcx> Eq for EvalSnapshot<'a, 'mir, 'tcx>
{}

impl<'a, 'mir, 'tcx> PartialEq for EvalSnapshot<'a, 'mir, 'tcx>
{
    fn eq(&self, other: &Self) -> bool {
        // FIXME: This looks to be a *ridicolously expensive* comparison operation.
        // Doesn't this make tons of copies?  Either `snapshot` is very badly named,
        // or it does!
        self.snapshot() == other.snapshot()
    }
}
<|endoftext|>


<commit_before>use super::plumbing::*;
use super::*;

use std::fmt::{self, Debug};

/// The `split` function takes arbitrary data and a closure that knows how to
/// split it, and turns this into a `ParallelIterator`.
///
/// # Examples
///
/// As a simple example, Rayon can recursively split ranges of indices
///
/// ```
/// use rayon::iter;
/// use rayon::prelude::*;
/// use std::ops::Range;
///
///
/// // We define a range of indices as follows
/// type Range1D = Range<usize>;
///
/// // Splitting it in two can be done like this
/// fn split_range1(r: Range1D) -> (Range1D, Option<Range1D>) {
///     // We are mathematically unable to split the range if there is only
///     // one point inside of it, but we could stop splitting before that.
///     if r.end - r.start <= 1 { return (r, None); }
///
///     // Here, our range is considered large enough to be splittable
///     let midpoint = r.start + (r.end - r.start) / 2;
///     let out1 = Range1D {
///         start: r.start,
///         end: midpoint,
///     };
///     let out2 = Range1D {
///         start: midpoint,
///         end: r.end,
///     };
///     (out1, Some(out2))
/// }
///
/// // By using iter::split, Rayon will split the range until it has enough work
/// // to feed the CPU cores, then give us the resulting sub-ranges
/// iter::split(0..4096, split_range1).for_each(|sub_range| {
///     // As our initial range had a power-of-two size, the final sub-ranges
///     // should have power-of-two sizes too
///     assert!((sub_range.end - sub_range.start).is_power_of_two());
/// });
/// ```
///
/// This recursive splitting can be extended to two or three dimensions,
/// to reproduce a classic "block-wise" parallelization scheme of graphics and
/// numerical simulations:
///
/// ```
/// # type Range1D = Range<usize>;
/// # fn split_range1(r: Range1D) -> (Range1D, Option<Range1D>) {
/// #     if r.end - r.start <= 1 { return (r, None); }
/// #     let midpoint = r.start + (r.end - r.start) / 2;
/// #     let out1 = Range1D {
/// #         start: r.start,
/// #         end: midpoint,
/// #     };
/// #     let out2 = Range1D {
/// #         start: midpoint,
/// #         end: r.end,
/// #     };
/// #     (out1, Some(out2))
/// # }
/// #
/// // A two-dimensional range of indices can be built out of two 1D ones
/// #[derive(Debug)]
/// struct Range2D {
///     // Range of horizontal indices
///     pub rx: Range1D,
///
///     // Range of vertical indices
///     pub ry: Range1D,
/// }
///
/// // We want to recursively split them by the largest dimension until we have
/// // enough sub-ranges to feed our mighty multi-core CPU. This function
/// // carries out one such split.
/// fn split_range2(r2: Range2D) -> (Range2D, Option<Range2D>) {
///     // Decide on which axis (horizontal/vertical) the range should be split
///     let width = r2.rx.end - r2.rx.start;
///     let height = r2.ry.end - r2.ry.start;
///     if width >= height {
///         // This is a wide rectangle, split it on the horizontal axis
///         let (split_rx, ry) = (split_range1(r2.rx), r2.ry);
///         let out1 = Range2D {
///             rx: split_rx.0,
///             ry: ry.clone(),
///         };
///         let out2 = split_rx.1.map(|rx| Range2D { rx, ry });
///         (out1, out2)
///     } else {
///         // This is a tall rectangle, split it on the vertical axis
///         let (rx, split_ry) = (r2.rx, split_range1(r2.ry));
///         let out1 = Range2D {
///             rx: rx.clone(),
///             ry: split_ry.0,
///         };
///         let out2 = split_ry.1.map(|ry| Range2D { rx, ry, });
///         (out1, out2)
///     }
/// }
///
/// // Again, rayon can handle the recursive splitting for us
/// let range = Range2D { rx: 0..800, ry: 0..600 };
/// iter::split(range, split_range2).for_each(|sub_range| {
///     // If the sub-ranges were indeed split by the largest dimension, then
///     // if no dimension was twice larger than the other initially, this
///     // property will remain true in the final sub-ranges.
///     let width = sub_range.rx.end - sub_range.rx.start;
///     let height = sub_range.ry.end - sub_range.ry.start;
///     assert!((width < 2 * height) && (height < 2 * width));
/// });
/// ```
///
pub fn split<D, S>(data: D, splitter: S) -> Split<D, S>
    where D: Send,
          S: Fn(D) -> (D, Option<D>) + Sync
{
    Split {
        data: data,
        splitter: splitter,
    }
}

/// `Split` is a parallel iterator using arbitrary data and a splitting function.
/// This struct is created by the [`split()`] function.
///
/// [`split()`]: fn.split.html
#[derive(Clone)]
pub struct Split<D, S> {
    data: D,
    splitter: S,
}

impl<D: Debug, S> Debug for Split<D, S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Split")
            .field("data", &self.data)
            .finish()
    }
}

impl<D, S> ParallelIterator for Split<D, S>
    where D: Send,
          S: Fn(D) -> (D, Option<D>) + Sync + Send
{
    type Item = D;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        let producer = SplitProducer {
            data: self.data,
            splitter: &self.splitter,
        };
        bridge_unindexed(producer, consumer)
    }
}

struct SplitProducer<'a, D, S: 'a> {
    data: D,
    splitter: &'a S,
}

impl<'a, D, S> UnindexedProducer for SplitProducer<'a, D, S>
    where D: Send,
          S: Fn(D) -> (D, Option<D>) + Sync
{
    type Item = D;

    fn split(mut self) -> (Self, Option<Self>) {
        let splitter = self.splitter;
        let (left, right) = splitter(self.data);
        self.data = left;
        (self,
         right.map(|data| {
                       SplitProducer {
                           data: data,
                           splitter: splitter,
                       }
                   }))
    }

    fn fold_with<F>(self, folder: F) -> F
        where F: Folder<Self::Item>
    {
        folder.consume(self.data)
    }
}
<commit_msg>Remove Debug impl<commit_after>use super::plumbing::*;
use super::*;

use std::fmt::{self, Debug};

/// The `split` function takes arbitrary data and a closure that knows how to
/// split it, and turns this into a `ParallelIterator`.
///
/// # Examples
///
/// As a simple example, Rayon can recursively split ranges of indices
///
/// ```
/// use rayon::iter;
/// use rayon::prelude::*;
/// use std::ops::Range;
///
///
/// // We define a range of indices as follows
/// type Range1D = Range<usize>;
///
/// // Splitting it in two can be done like this
/// fn split_range1(r: Range1D) -> (Range1D, Option<Range1D>) {
///     // We are mathematically unable to split the range if there is only
///     // one point inside of it, but we could stop splitting before that.
///     if r.end - r.start <= 1 { return (r, None); }
///
///     // Here, our range is considered large enough to be splittable
///     let midpoint = r.start + (r.end - r.start) / 2;
///     let out1 = Range1D {
///         start: r.start,
///         end: midpoint,
///     };
///     let out2 = Range1D {
///         start: midpoint,
///         end: r.end,
///     };
///     (out1, Some(out2))
/// }
///
/// // By using iter::split, Rayon will split the range until it has enough work
/// // to feed the CPU cores, then give us the resulting sub-ranges
/// iter::split(0..4096, split_range1).for_each(|sub_range| {
///     // As our initial range had a power-of-two size, the final sub-ranges
///     // should have power-of-two sizes too
///     assert!((sub_range.end - sub_range.start).is_power_of_two());
/// });
/// ```
///
/// This recursive splitting can be extended to two or three dimensions,
/// to reproduce a classic "block-wise" parallelization scheme of graphics and
/// numerical simulations:
///
/// ```
/// # type Range1D = Range<usize>;
/// # fn split_range1(r: Range1D) -> (Range1D, Option<Range1D>) {
/// #     if r.end - r.start <= 1 { return (r, None); }
/// #     let midpoint = r.start + (r.end - r.start) / 2;
/// #     let out1 = Range1D {
/// #         start: r.start,
/// #         end: midpoint,
/// #     };
/// #     let out2 = Range1D {
/// #         start: midpoint,
/// #         end: r.end,
/// #     };
/// #     (out1, Some(out2))
/// # }
/// #
/// // A two-dimensional range of indices can be built out of two 1D ones
/// struct Range2D {
///     // Range of horizontal indices
///     pub rx: Range1D,
///
///     // Range of vertical indices
///     pub ry: Range1D,
/// }
///
/// // We want to recursively split them by the largest dimension until we have
/// // enough sub-ranges to feed our mighty multi-core CPU. This function
/// // carries out one such split.
/// fn split_range2(r2: Range2D) -> (Range2D, Option<Range2D>) {
///     // Decide on which axis (horizontal/vertical) the range should be split
///     let width = r2.rx.end - r2.rx.start;
///     let height = r2.ry.end - r2.ry.start;
///     if width >= height {
///         // This is a wide rectangle, split it on the horizontal axis
///         let (split_rx, ry) = (split_range1(r2.rx), r2.ry);
///         let out1 = Range2D {
///             rx: split_rx.0,
///             ry: ry.clone(),
///         };
///         let out2 = split_rx.1.map(|rx| Range2D { rx, ry });
///         (out1, out2)
///     } else {
///         // This is a tall rectangle, split it on the vertical axis
///         let (rx, split_ry) = (r2.rx, split_range1(r2.ry));
///         let out1 = Range2D {
///             rx: rx.clone(),
///             ry: split_ry.0,
///         };
///         let out2 = split_ry.1.map(|ry| Range2D { rx, ry, });
///         (out1, out2)
///     }
/// }
///
/// // Again, rayon can handle the recursive splitting for us
/// let range = Range2D { rx: 0..800, ry: 0..600 };
/// iter::split(range, split_range2).for_each(|sub_range| {
///     // If the sub-ranges were indeed split by the largest dimension, then
///     // if no dimension was twice larger than the other initially, this
///     // property will remain true in the final sub-ranges.
///     let width = sub_range.rx.end - sub_range.rx.start;
///     let height = sub_range.ry.end - sub_range.ry.start;
///     assert!((width < 2 * height) && (height < 2 * width));
/// });
/// ```
///
pub fn split<D, S>(data: D, splitter: S) -> Split<D, S>
    where D: Send,
          S: Fn(D) -> (D, Option<D>) + Sync
{
    Split {
        data: data,
        splitter: splitter,
    }
}

/// `Split` is a parallel iterator using arbitrary data and a splitting function.
/// This struct is created by the [`split()`] function.
///
/// [`split()`]: fn.split.html
#[derive(Clone)]
pub struct Split<D, S> {
    data: D,
    splitter: S,
}

impl<D: Debug, S> Debug for Split<D, S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Split")
            .field("data", &self.data)
            .finish()
    }
}

impl<D, S> ParallelIterator for Split<D, S>
    where D: Send,
          S: Fn(D) -> (D, Option<D>) + Sync + Send
{
    type Item = D;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        let producer = SplitProducer {
            data: self.data,
            splitter: &self.splitter,
        };
        bridge_unindexed(producer, consumer)
    }
}

struct SplitProducer<'a, D, S: 'a> {
    data: D,
    splitter: &'a S,
}

impl<'a, D, S> UnindexedProducer for SplitProducer<'a, D, S>
    where D: Send,
          S: Fn(D) -> (D, Option<D>) + Sync
{
    type Item = D;

    fn split(mut self) -> (Self, Option<Self>) {
        let splitter = self.splitter;
        let (left, right) = splitter(self.data);
        self.data = left;
        (self,
         right.map(|data| {
                       SplitProducer {
                           data: data,
                           splitter: splitter,
                       }
                   }))
    }

    fn fold_with<F>(self, folder: F) -> F
        where F: Folder<Self::Item>
    {
        folder.consume(self.data)
    }
}
<|endoftext|>


<commit_before>use std::io::stdin;
use std::sync::mpsc;
use std::thread;

use termion;
use termion::input::TermRead;

pub struct Input {
    tx: mpsc::Sender<termion::event::Event>,
    rx: mpsc::Receiver<termion::event::Event>,
}

impl Input {
    pub fn new() -> Input {
        let (tx, rx) = mpsc::channel();
        Input {
            tx: tx,
            rx: rx,
        }
    }

    pub fn run(&mut self) {
        let tx = self.tx.clone();
        thread::spawn(move || {
            for event_res in stdin().events() {
                match event_res {
                    Ok(event) => {
                        tx.send(event).unwrap();
                    },
                    Err(err) => {
                        error!("{:?}", err);
                    }
                }
            }
        });
    }

    pub fn try_recv(&mut self) -> Result<termion::event::Event, mpsc::TryRecvError> {
        self.rx.try_recv()
    }
}

impl Default for Input {
	fn default() -> Self {
		Self::new()
	}
}
<commit_msg>Removes implementation of Default trait for Input<commit_after>use std::io::stdin;
use std::sync::mpsc;
use std::thread;

use termion;
use termion::input::TermRead;

pub struct Input {
    tx: mpsc::Sender<termion::event::Event>,
    rx: mpsc::Receiver<termion::event::Event>,
}

impl Input {
    pub fn new() -> Input {
        let (tx, rx) = mpsc::channel();
        Input {
            tx: tx,
            rx: rx,
        }
    }

    pub fn run(&mut self) {
        let tx = self.tx.clone();
        thread::spawn(move || {
            for event_res in stdin().events() {
                match event_res {
                    Ok(event) => {
                        tx.send(event).unwrap();
                    },
                    Err(err) => {
                        error!("{:?}", err);
                    }
                }
            }
        });
    }

    pub fn try_recv(&mut self) -> Result<termion::event::Event, mpsc::TryRecvError> {
        self.rx.try_recv()
    }
}
<|endoftext|>


<commit_before>use std::io;

use tokio_core::io::Io;

use flushed::{Flushed, flushed};
use {Buf};

/// A wrapper for full-duplex stream
pub struct IoBuf<S: Io> {
    pub in_buf: Buf,
    pub out_buf: Buf,
    socket: S,
    done: bool,
}

/// Main trait of a stream (meaning socket) with input and output buffers
///
/// This is ought to be similar to `tokio_core::Io` but with buffers
impl<S: Io> IoBuf<S> {
    /// Create a new IoBuf object with empty buffers
    pub fn new(sock: S) -> IoBuf<S> {
        IoBuf {
            in_buf: Buf::new(),
            out_buf: Buf::new(),
            socket: sock,
            done: false,
        }
    }
    /// Read a chunk of data into a buffer
    ///
    /// The data just read can then be found in `self.in_buf`.
    ///
    /// This method does just one read. Because you are ought to try parse
    /// request after every read rather than reading a lot of the data in
    /// memory.
    ///
    /// This method returns `0` when no bytes are read, both when WouldBlock
    /// occurred and when connection has been closed. You may then use
    /// `self.done()` to distinguish from these conditions.
    pub fn read(&mut self) -> Result<usize, io::Error> {
        match self.in_buf.read_from(&mut self.socket) {
            Ok(0) => {
                self.done = true;
                Ok(0)
            }
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => Ok(0),
            Err(ref e)
                if e.kind() == io::ErrorKind::BrokenPipe ||
                   e.kind() == io::ErrorKind::ConnectionReset
            => {
                self.done = true;
                Ok(0)
            }
            result => result,
        }
    }

    /// Write data in the output buffer to actual stream
    ///
    /// You should put the data to be sent into `self.out_buf` before flush
    pub fn flush(&mut self) -> Result<(), io::Error> {
        loop {
            if self.out_buf.len() == 0 {
                break;
            }
            match self.out_buf.write_to(&mut self.socket) {
                Ok(0) => break,
                Ok(_) => continue,
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    break;
                }
                Err(ref e)
                    if e.kind() == io::ErrorKind::BrokenPipe ||
                       e.kind() == io::ErrorKind::ConnectionReset
                => {
                    self.done = true;
                    break;
                }
                Err(e) => {
                    return Err(e);
                },
            }
        }
        // This probably aways does nothing, but we have to support the full
        // Io protocol
        match self.socket.flush() {
            Ok(()) => Ok(()),
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => Ok(()),
            Err(ref e) if e.kind() == io::ErrorKind::BrokenPipe => {
                self.done = true;
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    /// Returns true when connection is closed by peer
    pub fn done(&self) -> bool {
        return self.done;
    }

    /// Returns a future which resolves to this stream when output buffer is
    /// flushed
    pub fn flushed(self) -> Flushed<S> {
        flushed(self)
    }
}
<commit_msg>Implement `AsRawFd` for `IoBuf`<commit_after>use std::io;
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, RawFd};

use tokio_core::io::Io;

use flushed::{Flushed, flushed};
use {Buf};

/// A wrapper for full-duplex stream
pub struct IoBuf<S: Io> {
    pub in_buf: Buf,
    pub out_buf: Buf,
    socket: S,
    done: bool,
}

/// Main trait of a stream (meaning socket) with input and output buffers
///
/// This is ought to be similar to `tokio_core::Io` but with buffers
impl<S: Io> IoBuf<S> {
    /// Create a new IoBuf object with empty buffers
    pub fn new(sock: S) -> IoBuf<S> {
        IoBuf {
            in_buf: Buf::new(),
            out_buf: Buf::new(),
            socket: sock,
            done: false,
        }
    }
    /// Read a chunk of data into a buffer
    ///
    /// The data just read can then be found in `self.in_buf`.
    ///
    /// This method does just one read. Because you are ought to try parse
    /// request after every read rather than reading a lot of the data in
    /// memory.
    ///
    /// This method returns `0` when no bytes are read, both when WouldBlock
    /// occurred and when connection has been closed. You may then use
    /// `self.done()` to distinguish from these conditions.
    pub fn read(&mut self) -> Result<usize, io::Error> {
        match self.in_buf.read_from(&mut self.socket) {
            Ok(0) => {
                self.done = true;
                Ok(0)
            }
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => Ok(0),
            Err(ref e)
                if e.kind() == io::ErrorKind::BrokenPipe ||
                   e.kind() == io::ErrorKind::ConnectionReset
            => {
                self.done = true;
                Ok(0)
            }
            result => result,
        }
    }

    /// Write data in the output buffer to actual stream
    ///
    /// You should put the data to be sent into `self.out_buf` before flush
    pub fn flush(&mut self) -> Result<(), io::Error> {
        loop {
            if self.out_buf.len() == 0 {
                break;
            }
            match self.out_buf.write_to(&mut self.socket) {
                Ok(0) => break,
                Ok(_) => continue,
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    break;
                }
                Err(ref e)
                    if e.kind() == io::ErrorKind::BrokenPipe ||
                       e.kind() == io::ErrorKind::ConnectionReset
                => {
                    self.done = true;
                    break;
                }
                Err(e) => {
                    return Err(e);
                },
            }
        }
        // This probably aways does nothing, but we have to support the full
        // Io protocol
        match self.socket.flush() {
            Ok(()) => Ok(()),
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => Ok(()),
            Err(ref e) if e.kind() == io::ErrorKind::BrokenPipe => {
                self.done = true;
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    /// Returns true when connection is closed by peer
    pub fn done(&self) -> bool {
        return self.done;
    }

    /// Returns a future which resolves to this stream when output buffer is
    /// flushed
    pub fn flushed(self) -> Flushed<S> {
        flushed(self)
    }
}

#[cfg(unix)]
impl<S: AsRawFd + Io> AsRawFd for IoBuf<S> {
    fn as_raw_fd(&self) -> RawFd {
        self.socket.as_raw_fd()
    }
}
<|endoftext|>


<commit_before>use crate::rparser::*;
use crate::{gen_get_variants, Variant};
use ipsec_parser::*;
use itertools::Itertools;

impl<'a> From<IkeTransformDHType> for Variant<'a> {
    fn from(input: IkeTransformDHType) -> Self {
        Variant::U16(input.0)
    }
}

pub struct IPsecBuilder {}
impl RBuilder for IPsecBuilder {
    fn build(&self) -> Box<dyn RParser> {
        Box::new(IPsecParser::new(b"IKEv2"))
    }
    fn get_l4_probe(&self) -> Option<ProbeL4> {
        Some(ipsec_probe)
    }
}

pub struct IPsecParser<'a> {
    _name: Option<&'a [u8]>,

    /// The transforms proposed by the initiator
    pub client_proposals: Vec<Vec<IkeV2Transform>>,

    /// The transforms selected by the responder
    pub server_proposals: Vec<Vec<IkeV2Transform>>,

    /// The Diffie-Hellman group from the server KE message, if present.
    pub dh_group: IkeTransformDHType,

    /// JA3-like hash
    pub fingerprint: Option<String>,
}

impl<'a> RParser for IPsecParser<'a> {
    fn parse_l4(&mut self, data: &[u8], direction: Direction) -> ParseResult {
        match parse_ikev2_header(data) {
            Ok((rem, ref hdr)) => {
                debug!("parse_ikev2_header: {:?}", hdr);
                if rem.is_empty() && hdr.length == 28 {
                    return ParseResult::Ok;
                }
                // Rule 0: check version
                if hdr.maj_ver != 2 || hdr.min_ver != 0 {
                    warn!("Unknown header version: {}.{}", hdr.maj_ver, hdr.min_ver);
                }
                match parse_ikev2_payload_list(rem, hdr.next_payload) {
                    Ok((_, Ok(p))) => {
                        debug!("parse_ikev2_payload_with_type: {:?}", p);
                        let fingerprint = build_ipsec_fingerprint(hdr, &p);
                        let digest = md5::compute(&fingerprint);
                        debug!("Fingerprint: {} --> {:x}", fingerprint, digest);
                        self.fingerprint = Some(format!("{:?}", digest));
                        for payload in p {
                            match payload.content {
                                IkeV2PayloadContent::SA(ref prop) => {
                                    let is_initiator = hdr.flags & IKEV2_FLAG_INITIATOR != 0;
                                    self.add_proposals(prop, is_initiator);
                                }
                                IkeV2PayloadContent::KE(ref kex) => {
                                    debug!("KEX {:?}", kex.dh_group);
                                    if direction == Direction::ToClient {
                                        self.dh_group = kex.dh_group;
                                    }
                                }
                                IkeV2PayloadContent::Nonce(ref n) => {
                                    debug!("Nonce: {:?}", n);
                                }
                                IkeV2PayloadContent::Notify(ref n) => {
                                    debug!("Notify: {:?}", n);
                                }
                                _ => {
                                    debug!("Unknown payload content {:?}", payload.content);
                                }
                            }
                        }
                    }
                    e => warn!("parse_ikev2_payload_with_type: {:?}", e),
                };
            }
            e => warn!("parse_ikev2_header: {:?}", e),
        };
        ParseResult::Ok
    }

    gen_get_variants! {IPsecParser, "ikev2.",
        dh_group    => into,
        enc         => |s| { s.get_server_proposal_enc().map(|x| x.into()) },
        prf         => |s| { s.get_server_proposal_prf().map(|x| x.into()) },
        auth        => |s| { s.get_server_proposal_auth().map(|x| x.into()) },
        dh          => |s| { s.get_server_proposal_dh().map(|x| x.into()) },
        esn         => |s| { s.get_server_proposal_esn().map(|x| x.into()) },
        fingerprint => map_as_ref,
    }
}

macro_rules! get_server_proposal {
    ( $t:tt, $n:ident ) => {
        pub fn $n(&self) -> Option<u16> {
            if let Some(xform) = self.server_proposals.first() {
                xform.iter().find_map(|x| {
                    if let IkeV2Transform::$t(e) = x {
                        Some(e.0)
                    } else {
                        None
                    }
                })
            } else {
                None
            }
        }
    };
}

impl<'a> IPsecParser<'a> {
    pub fn new(name: &'a [u8]) -> IPsecParser<'a> {
        IPsecParser {
            _name: Some(name),
            client_proposals: Vec::new(),
            server_proposals: Vec::new(),
            dh_group: IkeTransformDHType::None,
            fingerprint: None,
        }
    }

    get_server_proposal! {Encryption, get_server_proposal_enc}
    get_server_proposal! {PRF, get_server_proposal_prf}
    get_server_proposal! {Auth, get_server_proposal_auth}
    get_server_proposal! {DH, get_server_proposal_dh}
    get_server_proposal! {ESN, get_server_proposal_esn}

    #[allow(clippy::cognitive_complexity)]
    fn add_proposals(&mut self, prop: &[IkeV2Proposal], is_initiator: bool) {
        debug!("num_proposals: {}", prop.len());
        for p in prop {
            debug!("proposal: {:?}", p);
            debug!("num_transforms: {}", p.num_transforms);
            for xform in &p.transforms {
                debug!("transform: {:?}", xform);
                debug!("\ttype: {:?}", xform.transform_type);
                match xform.transform_type {
                    IkeTransformType::EncryptionAlgorithm => {
                        debug!(
                            "\tEncryptionAlgorithm: {:?}",
                            IkeTransformEncType(xform.transform_id)
                        );
                    }
                    IkeTransformType::PseudoRandomFunction => {
                        debug!(
                            "\tPseudoRandomFunction: {:?}",
                            IkeTransformPRFType(xform.transform_id)
                        );
                    }
                    IkeTransformType::IntegrityAlgorithm => {
                        debug!(
                            "\tIntegrityAlgorithm: {:?}",
                            IkeTransformAuthType(xform.transform_id)
                        );
                    }
                    IkeTransformType::DiffieHellmanGroup => {
                        debug!(
                            "\tDiffieHellmanGroup: {:?}",
                            IkeTransformDHType(xform.transform_id)
                        );
                    }
                    IkeTransformType::ExtendedSequenceNumbers => {
                        debug!(
                            "\tExtendedSequenceNumbers: {:?}",
                            IkeTransformESNType(xform.transform_id)
                        );
                    }
                    _ => warn!("\tUnknown transform type {:?}", xform.transform_type),
                }
                if xform.transform_id == 0 {
                    warn!("\tTransform ID == 0 (choice left to responder)");
                };
            }
            let proposals: Vec<IkeV2Transform> = p.transforms.iter().map(|x| x.into()).collect();
            debug!("Proposals\n{:?}", proposals);
            // Rule 1: warn on weak or unknown transforms
            for prop in &proposals {
                match *prop {
                    IkeV2Transform::Encryption(ref enc) => match *enc {
                        IkeTransformEncType::ENCR_DES_IV64
                        | IkeTransformEncType::ENCR_DES
                        | IkeTransformEncType::ENCR_3DES
                        | IkeTransformEncType::ENCR_RC5
                        | IkeTransformEncType::ENCR_IDEA
                        | IkeTransformEncType::ENCR_CAST
                        | IkeTransformEncType::ENCR_BLOWFISH
                        | IkeTransformEncType::ENCR_3IDEA
                        | IkeTransformEncType::ENCR_DES_IV32
                        | IkeTransformEncType::ENCR_NULL => {
                            warn!("Weak Encryption: {:?}", enc);
                        }
                        _ => (),
                    },
                    IkeV2Transform::Auth(ref auth) => {
                        match *auth {
                            IkeTransformAuthType::NONE => {
                                // Note: this could be expected with an AEAD encryption alg.
                                // See rule 4
                            }
                            IkeTransformAuthType::AUTH_HMAC_MD5_96
                            | IkeTransformAuthType::AUTH_HMAC_SHA1_96
                            | IkeTransformAuthType::AUTH_DES_MAC
                            | IkeTransformAuthType::AUTH_KPDK_MD5
                            | IkeTransformAuthType::AUTH_AES_XCBC_96
                            | IkeTransformAuthType::AUTH_HMAC_MD5_128
                            | IkeTransformAuthType::AUTH_HMAC_SHA1_160 => {
                                warn!("Weak auth: {:?}", auth);
                            }
                            _ => (),
                        }
                    }
                    IkeV2Transform::DH(ref dh) => match *dh {
                        IkeTransformDHType::None => {
                            warn!("'None' DH transform proposed");
                        }
                        IkeTransformDHType::Modp768
                        | IkeTransformDHType::Modp1024
                        | IkeTransformDHType::Modp1024s160
                        | IkeTransformDHType::Modp1536 => {
                            warn!("Weak DH: {:?}", dh);
                        }
                        _ => (),
                    },
                    IkeV2Transform::Unknown(tx_type, tx_id) => {
                        warn!("Unknown proposal: type={}, id={}", tx_type.0, tx_id);
                    }
                    _ => (),
                }
            }
            // Rule 2: check if no DH was proposed
            fn has_dh(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| matches!(x, IkeV2Transform::DH(_)))
            }
            if !has_dh(&proposals) {
                warn!("No DH transform found");
            }
            // Rule 3: check if proposing AH ([RFC7296] section 3.3.1)
            if p.protocol_id == ProtocolID::AH {
                warn!("Proposal uses protocol AH - no confidentiality");
            }
            // Rule 4: lack of integrity is accepted only if using an AEAD proposal
            // Look if no auth was proposed, including if proposal is Auth::None
            fn has_auth(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| match *x {
                    IkeV2Transform::Auth(IkeTransformAuthType::NONE) => false,
                    IkeV2Transform::Auth(_) => true,
                    _ => false,
                })
            }
            fn has_gcm(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| {
                    if let IkeV2Transform::Encryption(ref enc) = x {
                        enc.is_aead()
                    } else {
                        false
                    }
                })
            }
            if !has_auth(&proposals) && !has_gcm(&proposals) {
                warn!("No integrity transform found");
            }
            // Rule 5: Check if an integrity and no integrity are part of the same proposal ?
            // XXX
            // Finally
            if is_initiator {
                self.client_proposals.push(proposals);
            } else {
                self.server_proposals.push(proposals);
            }
        }
        // The server must accept one proposal, or reject them all (RFC 5996 2.7)
        if self.server_proposals.len() > 1 {
            warn!("more than one server proposals")
        }
        debug!("client_proposals: {:?}", self.client_proposals);
        debug!("server_proposals: {:?}", self.server_proposals);
    }
}

/// IPsec version,exchnage type,payload types,notify types,transform proposals
pub fn build_ipsec_fingerprint(hdr: &IkeV2Header, payload_list: &[IkeV2Payload]) -> String {
    // version is fix
    let mut fingerprint = String::from("2.0,");
    // exchange type
    fingerprint += &format!("{},", hdr.exch_type.0);
    // payload types
    let payload_types_str = payload_list
        .iter()
        .fold(vec![hdr.next_payload.0], |mut acc, payload| {
            acc.push(payload.hdr.next_payload_type.0);
            acc
        })
        .iter()
        .join("-");
    fingerprint.push_str(&payload_types_str);
    // notify types
    fingerprint.push(',');
    let notify_types_str = payload_list
        .iter()
        .filter_map(|payload| {
            if let IkeV2PayloadContent::Notify(notify) = &payload.content {
                Some(notify.notify_type.0)
            } else {
                None
            }
        })
        .join("-");
    fingerprint.push_str(&notify_types_str);
    // proposals and transforms
    fingerprint.push(',');
    let mut transforms = Vec::new();
    for payload in payload_list {
        if let IkeV2PayloadContent::SA(sa_list) = &payload.content {
            for sa in sa_list {
                for xform in &sa.transforms {
                    transforms.push(format!("{}:{}", xform.transform_type.0, xform.transform_id));
                }
            }
        }
    }
    let transforms_str = transforms.iter().join("-");
    fingerprint.push_str(&transforms_str);
    // finish
    fingerprint
}

pub fn ipsec_probe(i: &[u8], _l4info: &L4Info) -> ProbeResult {
    if i.len() <= 20 {
        return ProbeResult::Unsure;
    }
    match parse_ikev2_header(i) {
        Ok((_, ref hdr)) => {
            if hdr.maj_ver != 2 || hdr.min_ver != 0 {
                debug!(
                    "ipsec_probe: could be ipsec, but with unsupported/invalid version {}.{}",
                    hdr.maj_ver, hdr.min_ver
                );
                return ProbeResult::NotForUs;
            }
            if hdr.exch_type.0 < 34 || hdr.exch_type.0 > 37 {
                debug!(
                    "ipsec_probe: could be ipsec, but with unsupported/invalid exchange type {}",
                    hdr.exch_type.0
                );
                return ProbeResult::NotForUs;
            }
            ProbeResult::Certain
        }
        _ => ProbeResult::NotForUs,
    }
}
<commit_msg>IKEv2: reduce verbosity of probing function<commit_after>use crate::rparser::*;
use crate::{gen_get_variants, Variant};
use ipsec_parser::*;
use itertools::Itertools;

impl<'a> From<IkeTransformDHType> for Variant<'a> {
    fn from(input: IkeTransformDHType) -> Self {
        Variant::U16(input.0)
    }
}

pub struct IPsecBuilder {}
impl RBuilder for IPsecBuilder {
    fn build(&self) -> Box<dyn RParser> {
        Box::new(IPsecParser::new(b"IKEv2"))
    }
    fn get_l4_probe(&self) -> Option<ProbeL4> {
        Some(ipsec_probe)
    }
}

pub struct IPsecParser<'a> {
    _name: Option<&'a [u8]>,

    /// The transforms proposed by the initiator
    pub client_proposals: Vec<Vec<IkeV2Transform>>,

    /// The transforms selected by the responder
    pub server_proposals: Vec<Vec<IkeV2Transform>>,

    /// The Diffie-Hellman group from the server KE message, if present.
    pub dh_group: IkeTransformDHType,

    /// JA3-like hash
    pub fingerprint: Option<String>,
}

impl<'a> RParser for IPsecParser<'a> {
    fn parse_l4(&mut self, data: &[u8], direction: Direction) -> ParseResult {
        match parse_ikev2_header(data) {
            Ok((rem, ref hdr)) => {
                debug!("parse_ikev2_header: {:?}", hdr);
                if rem.is_empty() && hdr.length == 28 {
                    return ParseResult::Ok;
                }
                // Rule 0: check version
                if hdr.maj_ver != 2 || hdr.min_ver != 0 {
                    warn!("Unknown header version: {}.{}", hdr.maj_ver, hdr.min_ver);
                }
                match parse_ikev2_payload_list(rem, hdr.next_payload) {
                    Ok((_, Ok(p))) => {
                        debug!("parse_ikev2_payload_with_type: {:?}", p);
                        let fingerprint = build_ipsec_fingerprint(hdr, &p);
                        let digest = md5::compute(&fingerprint);
                        debug!("Fingerprint: {} --> {:x}", fingerprint, digest);
                        self.fingerprint = Some(format!("{:?}", digest));
                        for payload in p {
                            match payload.content {
                                IkeV2PayloadContent::SA(ref prop) => {
                                    let is_initiator = hdr.flags & IKEV2_FLAG_INITIATOR != 0;
                                    self.add_proposals(prop, is_initiator);
                                }
                                IkeV2PayloadContent::KE(ref kex) => {
                                    debug!("KEX {:?}", kex.dh_group);
                                    if direction == Direction::ToClient {
                                        self.dh_group = kex.dh_group;
                                    }
                                }
                                IkeV2PayloadContent::Nonce(ref n) => {
                                    debug!("Nonce: {:?}", n);
                                }
                                IkeV2PayloadContent::Notify(ref n) => {
                                    debug!("Notify: {:?}", n);
                                }
                                _ => {
                                    debug!("Unknown payload content {:?}", payload.content);
                                }
                            }
                        }
                    }
                    e => warn!("parse_ikev2_payload_with_type: {:?}", e),
                };
            }
            e => warn!("parse_ikev2_header: {:?}", e),
        };
        ParseResult::Ok
    }

    gen_get_variants! {IPsecParser, "ikev2.",
        dh_group    => into,
        enc         => |s| { s.get_server_proposal_enc().map(|x| x.into()) },
        prf         => |s| { s.get_server_proposal_prf().map(|x| x.into()) },
        auth        => |s| { s.get_server_proposal_auth().map(|x| x.into()) },
        dh          => |s| { s.get_server_proposal_dh().map(|x| x.into()) },
        esn         => |s| { s.get_server_proposal_esn().map(|x| x.into()) },
        fingerprint => map_as_ref,
    }
}

macro_rules! get_server_proposal {
    ( $t:tt, $n:ident ) => {
        pub fn $n(&self) -> Option<u16> {
            if let Some(xform) = self.server_proposals.first() {
                xform.iter().find_map(|x| {
                    if let IkeV2Transform::$t(e) = x {
                        Some(e.0)
                    } else {
                        None
                    }
                })
            } else {
                None
            }
        }
    };
}

impl<'a> IPsecParser<'a> {
    pub fn new(name: &'a [u8]) -> IPsecParser<'a> {
        IPsecParser {
            _name: Some(name),
            client_proposals: Vec::new(),
            server_proposals: Vec::new(),
            dh_group: IkeTransformDHType::None,
            fingerprint: None,
        }
    }

    get_server_proposal! {Encryption, get_server_proposal_enc}
    get_server_proposal! {PRF, get_server_proposal_prf}
    get_server_proposal! {Auth, get_server_proposal_auth}
    get_server_proposal! {DH, get_server_proposal_dh}
    get_server_proposal! {ESN, get_server_proposal_esn}

    #[allow(clippy::cognitive_complexity)]
    fn add_proposals(&mut self, prop: &[IkeV2Proposal], is_initiator: bool) {
        debug!("num_proposals: {}", prop.len());
        for p in prop {
            debug!("proposal: {:?}", p);
            debug!("num_transforms: {}", p.num_transforms);
            for xform in &p.transforms {
                debug!("transform: {:?}", xform);
                debug!("\ttype: {:?}", xform.transform_type);
                match xform.transform_type {
                    IkeTransformType::EncryptionAlgorithm => {
                        debug!(
                            "\tEncryptionAlgorithm: {:?}",
                            IkeTransformEncType(xform.transform_id)
                        );
                    }
                    IkeTransformType::PseudoRandomFunction => {
                        debug!(
                            "\tPseudoRandomFunction: {:?}",
                            IkeTransformPRFType(xform.transform_id)
                        );
                    }
                    IkeTransformType::IntegrityAlgorithm => {
                        debug!(
                            "\tIntegrityAlgorithm: {:?}",
                            IkeTransformAuthType(xform.transform_id)
                        );
                    }
                    IkeTransformType::DiffieHellmanGroup => {
                        debug!(
                            "\tDiffieHellmanGroup: {:?}",
                            IkeTransformDHType(xform.transform_id)
                        );
                    }
                    IkeTransformType::ExtendedSequenceNumbers => {
                        debug!(
                            "\tExtendedSequenceNumbers: {:?}",
                            IkeTransformESNType(xform.transform_id)
                        );
                    }
                    _ => warn!("\tUnknown transform type {:?}", xform.transform_type),
                }
                if xform.transform_id == 0 {
                    warn!("\tTransform ID == 0 (choice left to responder)");
                };
            }
            let proposals: Vec<IkeV2Transform> = p.transforms.iter().map(|x| x.into()).collect();
            debug!("Proposals\n{:?}", proposals);
            // Rule 1: warn on weak or unknown transforms
            for prop in &proposals {
                match *prop {
                    IkeV2Transform::Encryption(ref enc) => match *enc {
                        IkeTransformEncType::ENCR_DES_IV64
                        | IkeTransformEncType::ENCR_DES
                        | IkeTransformEncType::ENCR_3DES
                        | IkeTransformEncType::ENCR_RC5
                        | IkeTransformEncType::ENCR_IDEA
                        | IkeTransformEncType::ENCR_CAST
                        | IkeTransformEncType::ENCR_BLOWFISH
                        | IkeTransformEncType::ENCR_3IDEA
                        | IkeTransformEncType::ENCR_DES_IV32
                        | IkeTransformEncType::ENCR_NULL => {
                            warn!("Weak Encryption: {:?}", enc);
                        }
                        _ => (),
                    },
                    IkeV2Transform::Auth(ref auth) => {
                        match *auth {
                            IkeTransformAuthType::NONE => {
                                // Note: this could be expected with an AEAD encryption alg.
                                // See rule 4
                            }
                            IkeTransformAuthType::AUTH_HMAC_MD5_96
                            | IkeTransformAuthType::AUTH_HMAC_SHA1_96
                            | IkeTransformAuthType::AUTH_DES_MAC
                            | IkeTransformAuthType::AUTH_KPDK_MD5
                            | IkeTransformAuthType::AUTH_AES_XCBC_96
                            | IkeTransformAuthType::AUTH_HMAC_MD5_128
                            | IkeTransformAuthType::AUTH_HMAC_SHA1_160 => {
                                warn!("Weak auth: {:?}", auth);
                            }
                            _ => (),
                        }
                    }
                    IkeV2Transform::DH(ref dh) => match *dh {
                        IkeTransformDHType::None => {
                            warn!("'None' DH transform proposed");
                        }
                        IkeTransformDHType::Modp768
                        | IkeTransformDHType::Modp1024
                        | IkeTransformDHType::Modp1024s160
                        | IkeTransformDHType::Modp1536 => {
                            warn!("Weak DH: {:?}", dh);
                        }
                        _ => (),
                    },
                    IkeV2Transform::Unknown(tx_type, tx_id) => {
                        warn!("Unknown proposal: type={}, id={}", tx_type.0, tx_id);
                    }
                    _ => (),
                }
            }
            // Rule 2: check if no DH was proposed
            fn has_dh(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| matches!(x, IkeV2Transform::DH(_)))
            }
            if !has_dh(&proposals) {
                warn!("No DH transform found");
            }
            // Rule 3: check if proposing AH ([RFC7296] section 3.3.1)
            if p.protocol_id == ProtocolID::AH {
                warn!("Proposal uses protocol AH - no confidentiality");
            }
            // Rule 4: lack of integrity is accepted only if using an AEAD proposal
            // Look if no auth was proposed, including if proposal is Auth::None
            fn has_auth(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| match *x {
                    IkeV2Transform::Auth(IkeTransformAuthType::NONE) => false,
                    IkeV2Transform::Auth(_) => true,
                    _ => false,
                })
            }
            fn has_gcm(proposals: &[IkeV2Transform]) -> bool {
                proposals.iter().any(|x| {
                    if let IkeV2Transform::Encryption(ref enc) = x {
                        enc.is_aead()
                    } else {
                        false
                    }
                })
            }
            if !has_auth(&proposals) && !has_gcm(&proposals) {
                warn!("No integrity transform found");
            }
            // Rule 5: Check if an integrity and no integrity are part of the same proposal ?
            // XXX
            // Finally
            if is_initiator {
                self.client_proposals.push(proposals);
            } else {
                self.server_proposals.push(proposals);
            }
        }
        // The server must accept one proposal, or reject them all (RFC 5996 2.7)
        if self.server_proposals.len() > 1 {
            warn!("more than one server proposals")
        }
        debug!("client_proposals: {:?}", self.client_proposals);
        debug!("server_proposals: {:?}", self.server_proposals);
    }
}

/// IPsec version,exchnage type,payload types,notify types,transform proposals
pub fn build_ipsec_fingerprint(hdr: &IkeV2Header, payload_list: &[IkeV2Payload]) -> String {
    // version is fix
    let mut fingerprint = String::from("2.0,");
    // exchange type
    fingerprint += &format!("{},", hdr.exch_type.0);
    // payload types
    let payload_types_str = payload_list
        .iter()
        .fold(vec![hdr.next_payload.0], |mut acc, payload| {
            acc.push(payload.hdr.next_payload_type.0);
            acc
        })
        .iter()
        .join("-");
    fingerprint.push_str(&payload_types_str);
    // notify types
    fingerprint.push(',');
    let notify_types_str = payload_list
        .iter()
        .filter_map(|payload| {
            if let IkeV2PayloadContent::Notify(notify) = &payload.content {
                Some(notify.notify_type.0)
            } else {
                None
            }
        })
        .join("-");
    fingerprint.push_str(&notify_types_str);
    // proposals and transforms
    fingerprint.push(',');
    let mut transforms = Vec::new();
    for payload in payload_list {
        if let IkeV2PayloadContent::SA(sa_list) = &payload.content {
            for sa in sa_list {
                for xform in &sa.transforms {
                    transforms.push(format!("{}:{}", xform.transform_type.0, xform.transform_id));
                }
            }
        }
    }
    let transforms_str = transforms.iter().join("-");
    fingerprint.push_str(&transforms_str);
    // finish
    fingerprint
}

pub fn ipsec_probe(i: &[u8], _l4info: &L4Info) -> ProbeResult {
    if i.len() <= 20 {
        return ProbeResult::Unsure;
    }
    match parse_ikev2_header(i) {
        Ok((_, ref hdr)) => {
            if hdr.maj_ver != 2 || hdr.min_ver != 0 {
                // trace!(
                //     "ipsec_probe: could be ipsec, but with unsupported/invalid version {}.{}",
                //     hdr.maj_ver, hdr.min_ver
                // );
                return ProbeResult::NotForUs;
            }
            if hdr.exch_type.0 < 34 || hdr.exch_type.0 > 37 {
                // trace!(
                //     "ipsec_probe: could be ipsec, but with unsupported/invalid exchange type {}",
                //     hdr.exch_type.0
                // );
                return ProbeResult::NotForUs;
            }
            ProbeResult::Certain
        }
        _ => ProbeResult::NotForUs,
    }
}
<|endoftext|>